[ADD] Support Kronecker Hessian approximations for `Conv{1,2,3}d`

f-dangel / backpack

BackPACK - a backpropagation package built on top of PyTorch which efficiently computes quantities other than the gradient.

https://backpack.pt/

MIT License

549 stars 55 forks source link

[ADD] Support Kronecker Hessian approximations for `Conv{1,2,3}d` #291

Closed f-dangel closed 1 year ago

f-dangel commented 1 year ago

This PR generalizes the Kronecker-factored approximations of Hessian diagonal blocks (KFRA, KFLR, KFAC) for Conv2d to Conv1d and Conv3d. Adds a test for the KFRA approximation under specific limits.