Get loss, gradient and hessian in one go

username: qq-me
cooked: <p>Currently I do this:</p>
<pre><code class="lang-auto">from collections.abc import Sequence
import torch

def jacobian(input: Sequence[torch.Tensor], wrt: Sequence[torch.Tensor], create_graph=False):
    flat_input = torch.cat([i.reshape(-1) for i in input])
    return torch.autograd.grad(
        flat_input,
        wrt,
        torch.eye(len(flat_input), device=input[0].device, dtype=input[0].dtype),
        retain_graph=True,
        create_graph=create_graph,
        allow_unused=True,
        is_grads_batched=True,
    )

input = torch.randn(3, requires_grad=True)
loss = (input**2).mean()
grad = torch.autograd.grad(loss, input, create_graph=True)
hessian = jacobian(grad, [input])
</code></pre>
<p>However as I understand the most efficient way of computing hessian is <a href="https://pytorch.org/docs/stable/generated/torch.func.hessian.html" class="inline-onebox" rel="noopener nofollow ugc">torch.func.hessian — PyTorch 2.7 documentation</a></p>
<p>But my trouble is that when I use it, I do this:</p>
<pre><code class="lang-auto">def fn(x): return x.pow(2).mean()
loss = fn(input)
grad = torch.autograd.grad(loss, input)
hessian = torch.func.hessian(fn)(input)
</code></pre>
<p>Doesn’t hessian then need to re-evaluate the function value and the gradients again even though I have already computed them? My function might be quite expensive. What would be the best solution to calculating hessian?</p>

username: KFrank
cooked: <p>Hi Ivan!</p>
<aside class="quote no-group" data-username="qq-me" data-post="1" data-topic="219601" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/qq-me/48/70055_2.png" class="avatar"> qq-me:</div>
<blockquote>
<p>But my trouble is that when I use it, I do this:</p>
<pre><code class="lang-auto">def fn(x): return x.pow(2).mean()
loss = fn(input)
grad = torch.autograd.grad(loss, input)
hessian = torch.func.hessian(fn)(input)
</code></pre>
<p>Doesn’t hessian then need to re-evaluate the function value and the gradients again even though I have already computed them?</p>
</blockquote>
</aside>
<p>I believe that you are correct about this.  I am not aware of any pre-packaged pytorch<br>
<code>hessian</code> functionality that also gives you access to the <code>loss</code> and <code>grad</code> it would have<br>
computed under the hood.</p>
<p>So I do believe that <code>func.hessian()</code> does require duplicative computation in your use<br>
case.</p>
<p>On the other hand, if your <code>grad</code> consists of <code>n</code> components, computing the hessian<br>
(e.g., with <code>func.hessian()</code>) requires <code>n</code> autograd passes (in addition to the first autograd<br>
pass used to compute <code>grad</code>).  So the duplicative computation may well be relatively<br>
insignificant.</p>
<p>It is also conceivable that <code>func.hessian()</code> contains some minor internal efficiencies<br>
that are sufficient to overcome the cost of the unnecessary computation.</p>
<aside class="quote no-group">
<blockquote>
<p>My function might be quite expensive.</p>
</blockquote>
</aside>
<p>Note that if your function is expensive, it is also likely the first-derivative backward pass<br>
for <code>grad</code> and the second-derivative backward passes for <code>hessian</code> will be expensive.<br>
So, again, the cost of the <code>n</code> second-derivative backward passes may dominate the<br>
overall cost, with cost of the redundant computation of <code>loss</code> and <code>grad</code> being relatively<br>
minor.</p>
<p>If the cost of <code>hessian</code> is important to your use case, it would probably make sense to<br>
time both approaches.</p>
<p>Best.</p>
<p>K. Frank.</p>