Question on number of Backwards

[edo-urettini]

Thank you very much for your work!

I have a question about the required number of backward passes needed to compute the FIM. In the original KFAC paper (https://arxiv.org/pdf/1503.05671) and in other implementations (https://openreview.net/pdf?id=BJlrF24twB), it seems like the FIM can be computed in a single backward pass. In the first paper, this is due to a single MC sample, while in the second, it seems possible to compute the exact FIM in a single pass.

In nngeometry, when sampling a single MC example, we still explicitly compute the Jacobian with "for i in range(n_output)". It seems to me that this approach is in contrast with the two papers, and can be avoided by doing a standard backpropagation from the loss computed on the new sampled target.
What am I getting wrong?
In particular, I am interested in using nngeometry for a regression problem with MANY output units, but the computation time is very high.

Thanks again.

[tfjgeorge]

Hi Edoardo, When performing a MC estimate of the FIM, n_output is actually the number of MC samples, which is typically smaller than the dimension of the output space. In the case of regression, I did not implement the MC estimator yet. Actually I am not 100% sure what would be the correct way : my first guess would be to add a scalar \sigma^2 variance as an argument, then treat the output of the model as the mean of a normal distribution with diagonal covariance \sigma^2 * I . Or in more generality allow to pass a full covariance matrix. If you know of a paper that describes another way of doing this I can have a look and implement it as well. What do you think? I don't have experience with BackPACK, but as far as I can tell they actually re-implement the whole backward pass, which makes it possible to compute jacobians of >1d vectors with respect to intermediate activations. This is why this is done in a single backward pass, although a more compute intensive one. Some while ago I was planning on using BackPACK as a backend in nng, but now that the torch.func api is released I no longer think this is relevant for nng (but I might still be useful for other cases that I am not aware of). In contrast, in nng, the backward pass (gradients w.r.t. activations) is performed by torch autodiff as usually done when training neural nets, then the custom (tricky) part is done when computing indvidiual gradients w.r.t parameters, as well as KFAC/EKFAC coefficients. Thomas

…

On Sat, May 24, 2025 at 4:08 PM Edoardo Urettini ***@***.***> wrote: *edo-urettini* created an issue (tfjgeorge/nngeometry#134) <#134> Thank you very much for your work! I have a question about the required number of backward passes needed to compute the FIM. In the original KFAC paper ( https://arxiv.org/pdf/1503.05671) and in other implementations ( https://openreview.net/pdf?id=BJlrF24twB), it seems like the FIM can be computed in a single backward pass. In the first paper, this is due to a single MC sample, while in the second, it seems possible to compute the exact FIM in a single pass. In nngeometry, when sampling a single MC example, we still explicitly compute the Jacobian with "for i in range(n_output)". It seems to me that this approach is in contrast with the two papers, and can be avoided by doing a standard backpropagation from the loss computed on the new sampled target. What am I getting wrong? In particular, I am interested in using nngeometry for a regression problem with MANY output units, but the computation time is very high. Thanks again. — Reply to this email directly, view it on GitHub <#134>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/AALTMWMUWXG6242YWNTKXPT3AB4OVAVCNFSM6AAAAAB52TCH6OVHI2DSMVQWIX3LMV43ASLTON2WKOZTGA4DQNBZGMYTKMI> . You are receiving this because you are subscribed to this thread.Message ID: ***@***.***>

[edo-urettini]

Hi Thomas, thanks for the answer.

The implementation of the MC estimator for regression would be very nice. To my understanding, the reason the "for i in range(n_output)" has fewer iterations with the MC estimator is that only the sampled labels are used as outputs. In regression, I think this would not be possible because when sampling from the distribution, you still have all the output units that are "active" in the backpropagation. I think that the correct approach would be to take some samples from the predictive distribution, computing separate losses for all of them and backpropagating from these losses. I agree that minimizing the MSE corresponds to maximizing the log-likelihood of a normal distribution with covariance $\sigma^2$ (Deep Learning by Bishop page 195, "error functions").

Thank you so much for the clarification on the differences between Backpack and nng. I agree that with torch.func it seems to have lost some relevance.

Best,
Edoardo

[tfjgeorge]

This should be ok with #138 You can specify a scalar parameter "covariance", I have no idea what would be a good heuristic for choosing its value though.