To anyone in the future looking for a solution to getting the diagonal of the Hessian: consider using the second directional derivative. Suppose f(x) is the objective function, then the second directional derivative can be calculated like this:
(d/dh(d/dh f(x+a*h))) | h=0
Here the “a” vector indicates the direction of interest and “h” is a scalar value.
To retrieve an element of the diagonal just plug in the corresponding one-hot vector for “a”.
I was able to easily find the second directional derivative using http://www.matrixcalculus.org/ 40
Just plug in f(x+a*h) and differentiate with respect to h. This gives the first directional derivative. Take the answer and plug it in again to obtain the second directional derivative. If for some reason you cannot find the second directional derivative, but can find the first, then PyTorch should be able to obtain it for you anyway via autograd.