Efficient computation of Hessian with respect to network weights using autograd.grad and symmetry of Hessian matrix

username: cijerezg
cooked: <p>As the title says, I am trying to compute the Hessian with respect to the weights (not the inputs) of the network. Right now, my code is as follows:</p>
<pre><code class="lang-auto">def hessian(loss, params):
      grads = autograd.grad(loss, params, retain_graph=True, create_graph=True)
      flattened_grads = torch.cat(([grad.flatten() for grad in grads]))

      hessian = torch.zeros(flattened_grads.shape[0], flattened_grads.shape[0])

      for idx, grad in enumerate(grads):
           second_der = autograd.grad(grad, params, retain_graph=True, allow_unused=True)
           second_der = torch.cat(([grad.flatten() for grad in second_der]))
           hessian[idx, :] = second_der
</code></pre>
<p>That works fine, but it is inefficient because the hessian matrix is symmetric, so the whole matrix does not need to be calculated.</p>
<p>To make it more efficient I thought of the following:</p>
<pre><code class="lang-auto">flat_params = torch.cat(([par.flatten() for par in params]))
for idx, grad in enumerate(grads):
     second_der = autograd.grad(grad, flat_params[idx::], retain_graph=True, allow_unused=True)
     second_der = torch.cat(([grad.flatten() for grad in second_der]))
     hessian[idx, :] = second_der
</code></pre>
<p>Namely, I calculate the upper triangular matrix and the lower half can be copied later, but unfortunately <code>second_der = autograd.grad(grad, flat_params[idx::], retain_graph=True, allow_unused=True)</code> throws <code>None</code>. How to fix this? is there I a way can do what I want?</p>

username: AlphaBetaGamma96
cooked: <p>If you’re trying to compute the Hessian of a matrix effeciently, I’d recommend using the <code>functorch</code> package. You can find their repo <a href="https://github.com/pytorch/functorch" rel="noopener nofollow ugc">here</a>, and their documentation on higher-order gradients <a href="https://github.com/pytorch/functorch#what-are-the-transforms" rel="noopener nofollow ugc">here</a>.</p>

username: cijerezg
cooked: <p>is <code>functorch</code> much more optimized than <code>autograd.grad</code>? I thought <code>functorch</code> was more to enable per batch analysis, i.e., getting the gradient per sample instead of the sum of all gradients, but it was as fast as <code>autograd.grad</code>.</p>

username: AlphaBetaGamma96
cooked: <p>It’s not just for per-sample gradients (although functorch’s vmap will handle that very easily). I’d recommend giving the documentation a brief read, one of its main features is efficient Hessian calculation. The documentation (with example code) can be found here <a href="https://pytorch.org/functorch/stable/notebooks/jacobians_hessians.html" class="inline-onebox" rel="noopener nofollow ugc">Jacobians, Hessians, hvp, vhp, and more: composing functorch transforms — functorch 0.2.0 documentation</a></p>

username: maheshakya
cooked: <p>We can technically use this if it is a simple function that takes some <code>dtype</code> objects as parameters so that we can compute Hessian with respect to these parameters. However, I’m not sure how to use this when the parameters are encapsulated in a model.<br>
For example, if the function we consider is a loss function where we first compute the outputs of the model by doing a forward pass, and then the loss using the outputs, we may pass model parameters to the function, but is there a way to perform forward pass by model parameters themselves?</p>

username: AlphaBetaGamma96
cooked: <p>Hi <a class="mention" href="/u/maheshakya">@maheshakya</a>,</p>
<p>Can you give a clear example of this? I <em>think</em> what you’re referring to is <code>torch.func.functional_call</code>. If your loss function is written in the standard PyTorch syntax is defined as,</p>
<pre><code class="lang-auto">outputs = net(inputs)
loss = loss_fn(outputs, targets)
</code></pre>
<p>In the <code>torch.func</code> namespace, you would define it along the lines of,</p>
<pre><code class="lang-auto">params = dict(net.named_parameters())
outputs = torch.func.functional_call(net, params, inputs)
loss = loss_fn(outputs, targets)
</code></pre>
<p>This example is for a single sample, for multiple samples you’ll need to wrap the <code>outputs</code> calculations in a <code>torch.func.vmap</code> call to vectorize over the batch-dimension.</p>

username: maheshakya
cooked: <p>Thanks for your reply <a class="mention" href="/u/alphabetagamma96">@AlphaBetaGamma96</a>. I have the following example in standard PyTorch syntax:<br>
Suppose <code>net</code> is some pretrained model and <code>loss_fn</code> is some loss function.</p>
<pre><code class="lang-auto">outputs = net(inputs)
loss = loss_fn(outputs, targets)
</code></pre>
<p>Now I want to compute the Hessian of <code>loss</code> w.r.t. parameters of <code>net</code>. Presumably, I would do something like below using <code>functorch</code>.</p>
<pre><code class="lang-auto">from functorch import hessian

def compute_loss(net, inputs, targets):
    outputs = net(inputs)
    return loss_fn(outputs, targets)

h = hessian(compute_loss, argnums=0)(net, inputs, targets)
</code></pre>
<p>Can I use <code>hessian</code> from <code>functorch</code> like this?</p>

username: AlphaBetaGamma96
cooked: <p>The documentation of the <code>hessian</code> function for <code>functorch</code> can be found here: <a href="https://pytorch.org/functorch/stable/notebooks/jacobians_hessians.html" class="inline-onebox" rel="noopener nofollow ugc">Jacobians, Hessians, hvp, vhp, and more: composing functorch transforms — functorch nightly documentation</a></p>
<p>I would advise using <code>torch.func</code> instead (if you’re using a more recent version of PyTorch), but something similar is possible in <code>functorch</code>. In <code>torch.func</code> an example of computing the Hessian with respect to parameters would be something like,</p>
<pre><code class="lang-auto">import torch
from torch import nn
from torch.func import functional_call, vmap, hessian

class Model(nn.Module):
  def __init__(self):
    super(Model, self).__init__()
    self.fc1=nn.Linear(1,16)
    self.fc2=nn.Linear(16,1)
    self.af=nn.Tanh()
  def forward(self, x):
    x=self.fc1(x)
    x=self.af(x)
    x=self.fc2(x)
    return x.squeeze(-1)

net = Model()

batch_size=100

targets = torch.randn(batch_size)
inputs = torch.randn(batch_size, 1)
params = dict(net.named_parameters())

def fcall(params, inputs):
  outputs = functional_call(net, params, inputs)
  return outputs

def loss_fn(outputs, targets):
  return torch.mean((outputs - targets)**2, dim=0)

def compute_loss(params, inputs, targets):
  outputs = vmap(fcall, in_dims=(None,0))(params, inputs) #vectorize over batch
  return loss_fn(outputs, targets)
  
def compute_hessian_loss(params, inputs, targets):
  return hessian(compute_loss, argnums=(0))(params, inputs, targets)

loss = compute_loss(params, inputs, targets)
print(loss)

hess = compute_hessian_loss(params, inputs, targets)
key=list(params.keys())[0] #take weight in first layer as example key
print(hess[key][key].shape) #Hessian of loss w.r.t first weight (shape [16, 1, 16, 1])

</code></pre>

username: maheshakya
cooked: <p>This approach works. Thanks!</p>