Вычисление гессиана функции потерь, включая NN по параметрам в Julia, с использованием Zygote

Question

Как бы вы вычислили гессиан функции потерь, которая состоит из нейронной сети по параметрам нейронной сети?

Например, рассмотрим функцию потерь ниже

      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)

Я могу получить градиент по параметрам двумя способами…

      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 не принимает на вход)

Недавно в ветку master (issue # 910) добавлена функция , которая понимает <tcode id="2584626"></tcode>как вход .

Я пытался совместить gradient и jacobianчтобы получить гессиан (потому что гессиан - это якобиан градиента функции), но безрезультатно. Я думаю проблема в том, что model это Chain объект, который включает общие функции, такие как reshape и σ. в которых отсутствуют параметры, но я не могу обойтись без этого.

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

Answer 1

Не уверен, что это поможет в вашем конкретном случае использования, но вы можете работать с аппроксимацией гессиана, например, эмпирическим Фишером (EF). Я работал с этим подходом для реализации аппроксимации Лапласа для моделей Flux (см. здесь ), вдохновленный этой реализацией PyTorch . Ниже я применил подход к вашему примеру.

      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 для полного гессиана по-прежнему имеет квадратичное масштабирование по количеству параметров, но, как показано в этой статье NeurIPS 2021 , вы можете дополнительно аппроксимировать гессиан, используя (блог-)диагональную факторизацию. В документе также показано, что в контексте байесовского глубокого обучения вероятностная обработка только последнего слоя обычно дает хорошие результаты, но опять же не уверен, что это уместно в вашем случае.