Computing Hessian for loss function

username: blade
cooked: <p>I’m looking at an implementation for calculating the Hessian matrix of the loss function.</p>
<pre><code class="lang-auto">loss = self.loss_function()
loss.backward(retain_graph=True)
grad_params = torch.autograd.grad(loss, p, create_graph=True)  # p is the weight matrix for a particular layer 
hess_params = torch.zeros_like(grad_params[0])

for i in range(grad_params[0].size(0)):
    for j in range(grad_params[0].size(1)):
        hess_params[i, j] = torch.autograd.grad(grad_params[0][i][j], p, retain_graph=True)[0][i, j]
</code></pre>
<p>I had 3 questions:</p>
<ol>
<li>Why do we compute hessian in a loop?Can’t we use something in the lines of</li>
</ol>
<pre><code class="lang-auto">hess_params = torch.autograd.grad(grad_params, p, retain_graph=True)
</code></pre>
<ol start="2">
<li>
<p>Current setting takes hours to run when it comes to larger weight matrices. What can I do to enhance the code?</p>
</li>
<li>
<p>I have seen a <code>hessian</code> function has been implemented in <code>autograd</code> package. How can we use that in this case?</p>
</li>
</ol>
<p>Pointing to reading resources and similar questions would also be highly appreciated.</p>

username: Yaroslav_Bulatov
cooked: <p>There’s a solution from Adam which does regular for-loop instead of nested for-loop – <a href="https://gist.github.com/apaszke/226abdf867c4e9d6698bd198f3b45fb7" rel="nofollow noopener">https://gist.github.com/apaszke/226abdf867c4e9d6698bd198f3b45fb7</a></p>
<p>A proposal to provide hessians native in PyTorch by <a class="mention" href="/u/alband">@albanD</a>  <a href="https://github.com/pytorch/pytorch/issues/30632" rel="nofollow noopener">https://github.com/pytorch/pytorch/issues/30632</a></p>
<p>You can replace k-backward calls with a single backward call by putting k things in the batch dimension, however, this also increases your memory usage by a factor of k, see this thread – <a href="https://discuss.pytorch.org/t/efficient-computation-with-multiple-grad-outputs-in-autograd-grad/66594" class="inline-onebox">Efficient computation with multiple grad_output's in autograd.grad</a></p>
<p>There’s a fundamental problem of Hessian being large. Take resnet-50 which has 25M parameters, Hessian then has 625 trillion entries. This means for large networks you have to deal with factorized approximations or consider a subset of the entries like the diagonal, which can be obtained at a similar cost to the gradient.</p>
<p>IE, for ReLU networks, you can get diagonal exactly using Gauss-Newton trick implemented <a href="https://github.com/cybertronai/autograd-lib#example-2-hessian-quantities" rel="nofollow noopener">here</a> and for more general networks you can use Hutchison estimator through Hessian-vector products like in  <a href="https://colab.research.google.com/drive/1JBWTY-2tUSk736yGxRdlYpH9Gh6GGBZa#scrollTo=_0viM3fQlnHZ" rel="nofollow noopener">this colab</a></p>

username: blade
cooked: <p>Thanks for the detailed response. I’ll look into these. My network has about 460k parameters, so not as large as resnet-50, but definitely not small.</p>
<p>Interestingly enough, the comment at the beginning of this <code>getHessian()</code> function reads</p>
<blockquote>
<p>This function computes the diagonal entries of the Hessian matrix of the<br>
decoding NN parameters</p>
</blockquote>
<p>But to me, this looks a full Hessian matrix. Doesn’t it? I’ll look into that possibility as well.</p>

username: tjlane
cooked: <p><a class="mention" href="/u/yaroslav_bulatov">@Yaroslav_Bulatov</a> and others helped me formulate more efficient Jacobian calculators:</p><aside class="quote quote-modified" data-post="5" data-topic="66594">
  <div class="title">
    <div class="quote-controls"></div>
    <img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/tjlane/48/19625_2.png" class="avatar">
    <a href="https://discuss.pytorch.org/t/efficient-computation-with-multiple-grad-outputs-in-autograd-grad/66594/5">Efficient computation with multiple grad_output's in autograd.grad</a> <a class="badge-category__wrapper " href="/c/autograd/7"><span data-category-id="7" data-drop-close="true" class="badge-category " title="A category of posts relating to the autograd engine itself."><span class="badge-category__name">autograd</span></span></a>
  </div>
  <blockquote>
    Right. I suppose I was operating under the assumption pytorch’s implementation was flexible enough to accept an identity matrix as “v”. That may be wrong. 
So if you compare these two implementations, the first gives significantly faster run times (in my hands) than the second. Factor of 10-100x depending on the case. Perhaps this is just to some skipped CUDA overhead? I am trying to understand this behavior and see if I can exploit the same kind of speedup using the “forward mode trick”. 
def _…
  </blockquote>
</aside>

<p>Essentially you can eliminate for loops (and significantly speed computation time) at the cost of memory. Probably too much memory cost for your large Hessian, but thought I’d add this in case useful.</p>