Torch.backward not allowing second gradients

Shubh_Goyal: <p>I want to implement some changes in the adam optimizer which requires second derivatives along with the first derivatives (gradients) of loss.</p>
<p>But due to loss.backwards implementation, I am unable to calculate second derivatives. Is there any way I achieve so?</p>
<p>Thank You.</p>

AlphaBetaGamma96: <p>Hi <a class="mention" href="/u/shubh_goyal">@Shubh_Goyal</a>,</p>
<p>You can have a look at <code>torch.autograd.functional.hessian</code> <a href="https://pytorch.org/docs/stable/generated/torch.autograd.functional.hessian.html" class="inline-onebox" rel="noopener nofollow ugc">torch.autograd.functional.hessian — PyTorch 2.3 documentation</a> to compute 2nd order derivatives (although it’s pretty computationally expensive).</p>
<p>Or you could have a look at the <code>torch.func</code> namespace and compute higher-order gradients there: <a href="https://pytorch.org/docs/stable/func.api.html" class="inline-onebox" rel="noopener nofollow ugc">torch.func API Reference — PyTorch 2.3 documentation</a></p>
<p>You can then pass these gradients to the <code>step()</code> method of your custom optimizer.</p>

Shubh_Goyal: <p>Hey <a class="mention" href="/u/alphabetagamma96">@AlphaBetaGamma96</a></p>
<p>Is it not necessary to do loss.backward().</p>
<p>Like what I can understand from the second reference you gave is that I can compute the first gradients as well as second gradients through grad function of torch.func and pass it while taking step in optimizer??</p>

AlphaBetaGamma96: <p>Yes, you can define your own custom <code>torch.optim.Optimizer</code> object, which takes the second-order derivatives as an input in the <code>optim.step()</code> method. Then you can define the optimizer however you want with the second-derivative terms.</p>

Shubh_Goyal: <p>Hey <a class="mention" href="/u/alphabetagamma96">@AlphaBetaGamma96</a></p>
<p>I found this <a href="https://towardsdatascience.com/introduction-to-functional-pytorch-b5bf739e1e6e" rel="noopener nofollow ugc">article</a>. Do you suggest doing this. This is completely functional programming. But the torchoptim library used is not official from pytorch I suppose.</p>
<p>Can you guide me more on how to use the func.grad with out torch models.</p>
<p>It would be of great help.</p>

AlphaBetaGamma96: <p>For the <code>torch.func</code> library, I’d recommend reading the docs <a href="https://pytorch.org/docs/stable/func.html" rel="noopener nofollow ugc">here</a></p>
<p>And for the custom <code>torch.optim.Optimizer</code> object, I’d recommend reading the docs <a href="https://pytorch.org/docs/stable/optim.html" rel="noopener nofollow ugc">here</a></p>

Shubh_Goyal: <p>Hey, I have read the docs and wrote some code, but the loss is not at all decreasing, I didn’t change the Adam formula yet. Can you look at my code?</p>
<p>I checked the calculated gradients and they are all zero.</p>

AlphaBetaGamma96: <p>If you share a minimal reproducible example, I can have a look.</p>
<p>Also, it should be noted that <code>torch.func</code> returns gradients from a function, whereas <code>torch.autograd</code> returns them to the <code>param.grad</code> attribute. So, you might just need to attach them the <code>.grad</code> attribute of your parameters.</p>

Shubh_Goyal: <pre><code class="lang-auto"># Function to compute the loss
def compute_loss(data, target, params):
    output = model_func(data, params)
    loss = loss_fn(output, target)
    return loss

for i in range(num_epochs):
    # Compute the gradients using torch.func.grad
    loss = compute_loss(x_train, y_train, params)
    grads = grad(compute_loss, argnums=2)(x_train, y_train, params)

    first_order_grads = grad(compute_loss, argnums=2)(x_train, y_train, params)

    # Compute second-order gradients
    def compute_gradients(params):
        loss = compute_loss(x_train, y_train, params)
        grads = grad(compute_loss, argnums=2)(x_train, y_train, params)
        return grads

    # Use the first-order gradients to compute second-order derivatives
    second_order_grads = grad(lambda p: torch.sum(torch.cat([g.view(-1) for g in compute_gradients(p)])), argnums=0)(params)
</code></pre>
<p>Do you think this is correct?</p>

AlphaBetaGamma96: <p>I’d advise you use <code>torch.func.hessian</code> as it’ll be significantly more efficient than an <code>torch.autograd</code> approach, I have an example of the forums here: <a href="https://discuss.pytorch.org/t/efficient-computation-of-hessian-with-respect-to-network-weights-using-autograd-grad-and-symmetry-of-hessian-matrix/156222/8" class="inline-onebox">Efficient computation of Hessian with respect to network weights using autograd.grad and symmetry of Hessian matrix - #8 by AlphaBetaGamma96</a></p>

Shubh_Goyal: <p>Hey. I did go through it. Is there any efficient way to do it. I can’t use it for fully connected layers at all.</p>

AlphaBetaGamma96: <p><code>torch.func</code> works directly on torch primitives, so <code>nn.Linear</code> layers are supported. Just make sure to use <code>torch.func.functional_call</code> when interacting with <code>nn.Module</code> objects to create a ‘functional’ version.</p>

Shubh_Goyal: <p>Actually what I meant is that the number of parameters is a problem for me. I have computational constraint so am unable to calculate hessian.</p>

AlphaBetaGamma96: <p>The Hessian scales quadratically in the number of parameters, so it can become intractable quite quickly. How many parameters does your model have?</p>
<p>Also, if you’re going to multiply the Hessian with a vector (which might not necessarily be the case), it’d be best to directly compute the Hessian-vector product, which avoids computing the Hessian directly and scales much better than naively computing the Hessian then multiplying it by a vector (tutorial here: <a href="https://pytorch.org/tutorials/intermediate/jacobians_hessians.html#computing-hessian-vector-products" class="inline-onebox" rel="noopener nofollow ugc">Jacobians, Hessians, hvp, vhp, and more: composing function transforms — PyTorch Tutorials 2.4.0+cu124 documentation</a>).</p>