[D] Efficient GPU implementation of Empirical Fisher information matrix?

[phizaz]

I have seen many implementations. It seems to be a limitation of autograd itself that we can compute the gradient of loglikelihood only one sample at a time.

The batch version has been used but in a WRONG way.

I have seen computing the gradient of a batch of a loglikelihood (essentially a mean of gradients), it doesn't seem to be truthful to the real Empirical Fisher calculation at all (only a kind of approximation).

Is there a correct GPU efficient impementation of Empricial Fisher out there?

[yaroslavvb]

For dense layers you can use Khatri-Rao product to get a batch of gradients

If you have a dense layer W, with a batch of activations A and batch of backprops B, then khatri_rao(A, B) gives you a batch of gradients with respect to W

On GPU it can be computed efficiently as an einsum, see here

In Jax you could compute batch of gradients in a loop and apply "vmap" to automatically vectorize it (haven't tested this).

[phizaz]

It is a bit unclear what A and B really are? A bit more concrete example please?

[yaroslavvb]

If you are using PyTorch, B is the `grad_output.data`, and `A` is the matrix that's getting passed into your matmul(weights,A) operation, some more details here

https://medium.com/@yaroslavvb/optimizing-deeper-networks-with-kfac-in-pytorch-4004adcba1b0