Per-sample gradient, should we design each layer differently?

[AlphaBetaGamma96]

Hi @Yaroslav_Bulatov,

I’ve just made a relatively short example that illustrates my point where this method fails if your loss function depends on the derivative of other terms. I’ve made an example script here on Github and I think this explains my current issue with hooks for per-sample gradients at the moment.

Could this issue be a potential bug of using hooks for such derivatives or is it a limitation when using the batch-supported Laplacian trick I reference in my previous response?

Thank you for all your help! It’s greatly appreciated!

[AlphaBetaGamma96]

@AlphaBetaGamma96

As I store the grad_output within a dictionary with the key ['e'] every time I call torch.autograd.grad (which happens N+1 times for an NxN Hessian matrix) it’ll overwrite the grad_output for each layer with the final component of the laplacian as opposed to the entire laplacian itself. Is this method (as shown here) compatible with getting a per-sample gradient via hooks if the per-sample gradients depend on other derivative information as well, which with the code in the link is the laplacian of the output of the network w.r.t the input! (If this makes sense?)

For completeness, it seems that you can’t use this method if your loss function contains multiple derivatives (unless I’m mistaken?)

A useful example is here → https://github.com/AlphaBetaGamma96/per-sample-gradient-limitation/blob/main/example.py . This might just be a limitation with how I store the grad_output variables if anyone does figure out a way around this, do let me know!