Robustness of parameters to different perturbations

[adamhaber]

I wonder if we can use the Hessian of the models to ask questions about "robustness" of different (classes of) solutions to different kinds of parameter perturbations. For example, we can ask what can we learn from the Hessians of different models about "stiff" and "sloppy" directions in parameters space, about the curvature of the loss functions at different solutions, etc.

Computing these Hessians should be very easy with tools like PyTorch and JAX, and for models with relatively small number of parameters we can even try to visualize them (which can be useful sometimes).

Would be happy to hear people's thoughts on this.

[thesamovar]

I have no idea if this would work with surrogate gradients but would be a very interesting question to ask.

@fzenke do you know of anyone who's tried that?

[fzenke]

Not aware of anyone having tried this with surrogate gradients. Evaluating the full Hessian with JAX should indeed be relatively straight forward. With PyTorch it was only easy to do Hessian vector products (last time I checked). One will need to take a little bit of care interpreting the results since these will strictly be "surrogate Hessians," but it would be something interesting to explore.

[adamhaber]

Maybe we can use something like PyHessian? It's built on top of PyTorch...

Also - not sure if this has been done before, but we can also consider something like finite-differences approach for computing the Hessian, as in here.

[adamhaber]

@thesamovar is this something you think we should try? Would be happy to give this a go if you could point me where to start :-)

[adamhaber]

I think it would be interesting to compute the Hessian with respect to parameters that have immediate biological correspondence; so, if we think of the weights as a possible "wiring diagram" which implements this computation, that's relevant - and the neuronal parameters are obviously relevant, as well. So - both? :-)

What I had in mind is something along these lines:

1. Train a network on the sound localization task - obtain "trained" parameters $\theta_0\in R^D$ which map each IPD to a probability vector: $f_{\theta_0}(ipd) = (p_1, ..., p_{N_c})$
2. Construct a "sensitivity function" which measures how, on average, changing $\theta_0$ (in the high dimensional parameter space) changes the outputs of the network. Maybe use KL-divergence here?
   $g(\Delta\theta) = \sum_{i}{KL[f_{\theta_0}(ipd_i),f_{\theta_0+\Delta\theta}(ipd_i)}]$
   we have $g : R^D \to R^+$, $g(0) = 0$ and $g\ge 0$ so the Hessian at 0 is PSD. Now we can about the geometry at this minimum - which parameters are more/less sensitive (but we need to be careful with parameters on different scales), how the average curvature (trace of Hessian) varies by architecture, etc...

What do you think?

[tfiers]

That makes sense, thanks! Curious to see what this will give

[thesamovar]

I think it's definitely worth trying! Are you coming along to the upcoming virtual hackathon (some time in the next couple of weeks I think)? If so, you could lead a group to work on it there?

[adamhaber]

I'll try! Can you share a link with the details?

In any case, I'll try to prepare a notebook with a proof-of-concept. How should I share it? As a PR?

[thesamovar]

join the discord, i think all the details are posted there! https://discord.gg/vMAVctwY