4

您将如何计算由神经网络 wrt NN 参数组成的损失函数的粗麻布？

例如，考虑下面的损失函数

using Flux: Chain, Dense, σ, crossentropy, params
using Zygote
model = Chain(
    x -> reshape(x, :, size(x, 4)),
    Dense(2, 5),
    Dense(5, 1),
    x -> σ.(x)
)
n_data = 5
input = randn(2, 1, 1, n_data)
target = randn(1, n_data)
loss = model -> crossentropy(model(input), target)

我可以通过两种方式获得渐变 wrt 参数……</p> 

Zygote.gradient(model -> loss(model), model)

或者

grad = Zygote.gradient(() -> loss(model), params(model))
grad[params(model)[1]]

但是，我找不到一种方法来获取粗麻布的参数。（我想做类似的事情Zygote.hessian(model -> loss(model), model)，但Zygote.hessian不::Params作为输入）

最近，一个jacobian函数被添加到 master 分支（issue #910），它被理解::Params为一个 input。

我一直在尝试组合gradient并jacobian获得粗麻布（因为粗麻布是函数梯度的雅可比），但无济于事。我认为问题在于这model是一个Chain包含通用函数的对象，reshape并且σ.缺少参数，但我无法超越这一点。

grad = model -> Zygote.gradient(model -> loss(model), model)
jacob = model -> Zygote.jacobian(grad, model)
jacob(model) ## does not work

编辑：作为参考，我之前在 pytorch 中创建了这个

4

1 回答 1

0

不确定这是否会对您的特定用例有所帮助，但您可以使用 Hessian 的近似值，例如经验费舍尔 (EF)。我已经使用这种方法来实现受此PyTorch 实现启发的Flux 模型的拉普拉斯近似（请参阅此处） 。下面我将该方法应用于您的示例。

using Flux: Chain, Dense, σ, crossentropy, params, DataLoader
using Zygote
using Random

Random.seed!(2022)
model = Chain(
    x -> reshape(x, :, size(x, 4)),
    Dense(2, 5),
    Dense(5, 1),
    x -> σ.(x)
)
n_data = 5
input = randn(2, 1, 1, n_data)
target = randn(1, n_data)
loss(x, y) = crossentropy(model(x), y)

n_params = length(reduce(vcat, [vec(θ) for θ ∈ params(model)]))
 = zeros(n_params,n_params)
data = DataLoader((input, target))

for d in data
  x, y = d
   = gradient(() -> loss(x,y),params(model))  
   = reduce(vcat,[vec([θ]) for θ ∈ params(model)])
   +=  * ' # empirical fisher
end

如果有办法直接（并且更有效地）使用 Zygote autodiff，我也有兴趣看到这一点。将 EF 用于完整 Hessian 的参数数量仍然是二次方的，但如本 NeurIPS 2021论文中所示，您可以使用（博客）对角分解进一步近似 Hessian。该论文还表明，在贝叶斯深度学习的背景下，仅在概率上仅处理最后一层通常会产生良好的结果，但再次不确定是否与您的情况相关。

于 2022-02-21T13:03:22.977 回答