Second Gradient Computation with autograd yield zeros

username: katonapng
cooked: <p>Hi everyone,</p>
<p>When I compute the gradient using <code>torch.autograd.grad</code> and then try to compute the second derivative via another call to <code>autograd.grad</code>, I unexpectedly get a tensor of <strong>all zeros</strong>. However, if I compute the second derivative using either:</p>
<ol>
<li>The <strong>Jacobian</strong> of the gradient and take its trace (i.e., sum of second partials), or</li>
<li><strong>Loop through each input component and compute its second derivative individually</strong>,</li>
</ol>
<p>I get the correct analytical result. The problem is:</p>
<ul>
<li><strong>Using Jacobians</strong>: Although I get the correct result, the model’s loss behaves strangely during training and doesn’t converge.</li>
<li><strong>Per-input gradient computation</strong>: This works, but it is <strong>extremely slow</strong>—not feasible for my application.</li>
</ul>
<p>Has anyone encountered this issue with <code>autograd.grad</code> returning zero second derivatives?</p>
<p>Here is a small test case:</p>
<pre><code class="lang-auto">def minimal_test(x):
    xi = x[:, 0, :]
    xj = x[:, 1, :]
    diff = xi - xj
    r = torch.sqrt(torch.sum(diff ** 2, dim=-1) + 1e-10)
    sigma = torch.exp(torch.tensor(0.15))
    k = torch.sigmoid(torch.tensor(0.5))
    phi = torch.pow(sigma / r, 2 / k)
    
    energy = phi.sum(dim=-1)
    return - energy

# Forward pass
x_t = torch.tensor([[[ 0.4806, -0.5267], [ 0.5513,  0.6484]]], requires_grad=True)  # Simplified input, shape (B, N, D)

def psi(x_t):
    output = minimal_test(x_t)
    gradients = torch.autograd.grad(
        outputs=output, inputs=x_t, grad_outputs=torch.ones_like(output),
        retain_graph=True, create_graph=True,
    )[0]
    return gradients

gradients = psi(x_t)
grad_outputs = torch.ones_like(gradients)
divergence = torch.autograd.grad(outputs=gradients, inputs=x_t, grad_outputs=grad_outputs, create_graph=True)[0]
print("Gradients:", gradients)
print("Second Gradients (autograd):", divergence)

jac = jacobian(psi, x_t)

divergence = torch.zeros(x_t.shape[:-1], device=x_t.device)  # shape: (B, N)
for b in range(x_t.shape[0]):
    for n in range(x_t.shape[1]):
        # Jacobian at point (b, n) is D x D
        J = jac[b, n, :, b, n, :]  # shape: (D, D)
        divergence[b, n] = torch.trace(J)

print("Second Gradients (jacobian):", divergence)

divergence = torch.zeros(x_t.shape[:-1], device=x_t.device)  # shape: (B, N)
for b in range(x_t.shape[0]):
    for n in range(x_t.shape[1]):
        for d in range(x_t.shape[2]):
            divergence[b, n] += torch.autograd.grad(
                gradients[b, n, d], x_t, retain_graph=True, create_graph=True,
            )[0][b, n, d]


print("Second Gradients (autograd per input):", divergence)
</code></pre>
<p>This is the output:</p>
<pre><code class="lang-auto">Gradients: tensor([[[-0.1571, -2.6116],
         [ 0.1571,  2.6116]]], grad_fn=&lt;AddBackward0&gt;)
Second Gradients (autograd): tensor([[[0., 0.],
         [0., 0.]]], grad_fn=&lt;AddBackward0&gt;)
Second Gradients (jacobian): tensor([[-7.1409, -7.1409]])
Second Gradients (autograd per input): tensor([[-7.1409, -7.1409]], grad_fn=&lt;CopySlices&gt;)
</code></pre>
<p>P.S.<br>
Setting any entry in <code>grad_outputs</code> to 0 will return a non-zero result, the problem is only when trying to get the full second derivative.</p>

username: valerian.rey
cooked: <p>The second (and probably third, since it gives the same results) time you’re computing <code>divergence</code>, you start by computing the hessian: <code>jac = jacobian(psi, x_t)</code>. This is the jacobian of the gradients with respect to the input, i.e. the hessian. Note that the result will be the same with <code>hessian = torch.autograd.functional.hessian(minimal_test, x_t)</code></p>
<p>This hessian is of shape [B, N, D, B, N, D], which is why it’s so expensive to compute, no matter how you do it.</p>
<p>You’re then computing something which is the trace of parts of the hessian, and using this for divergence.</p>
<p>The first time you’re computing the divergence, however, you’re computing directly something that is equal to <code>hessian.sum(dim=[0, 1, 2])</code>. This is not the divergence, and I think it’s normal for this to be equal to zero in your case.</p>

username: katonapng
cooked: <p>Ah, I see — the code worked for the Poisson process, but everything broke with Gibbs. I guess in the Poisson case, <code>hessian.sum(dim=[0, 1, 2])</code> accidentally gives the correct result, even though the logic behind it is incorrect. Thanks for the explanation! <img src="https://discuss.pytorch.org/images/emoji/apple/wink.png?v=14" title=":wink:" class="emoji" alt=":wink:" loading="lazy" width="20" height="20"></p>