Hessian calculation on specific components

Leonard_Deroose: <p>Hello,</p>
<p>I need to calculate the gradients of each layer of my model. Then, for each gradient, I need to select k components of it and finally I calculate the Hessian over these k components.</p>
<p>example:<br>
If my gradient has d components, I select k components from d (k&lt;d) and calculate the hessian (sub-hessian) on these k selected components. This means that the hessian (sub-hessian) does not have a form (d,d) but (k, k).</p>
<p>Here’s a small code of how I currently proceed:</p>
<pre data-code-wrap="python"><code class="lang-python">import torch
import numpy as np
import torch.nn as nn

def compute_sub_hessian(grad, weights, idxs):
    n_elems = np.prod(grad.shape)
    flat_grad = grad.reshape(-1)
    sub_hessian = torch.zeros(size=(len(idxs),n_elems))
    for i, idx in enumerate(idxs): 
        row_hessian = torch.autograd.grad(flat_grad[idx], weights, retain_graph=True)[0]
        row_hessian = row_hessian.reshape(-1)
        sub_hessian[i] = row_hessian
    sub_hessian = sub_hessian[:, idxs]
    return sub_hessian

class MnistModel(nn.Module):
    def __init__(self):
        super(MnistModel, self).__init__()
        self.fc = nn.Linear(28*28, 10, bias=True)
        return None

    def forward(self, x):
        x = x.reshape(-1, 28*28)
        x = self.fc(x)
        return x

model = MnistModel()
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
criterion = nn.CrossEntropyLoss()

X = torch.rand(100, 1, 28, 28)  # 100 images of size 28*28
y = torch.randint(0, 10, size=(100,))

optimizer.zero_grad()
outputs = model(X)
loss = criterion(outputs, y)
loss.backward(create_graph=True)

for name, param in model.named_parameters():
    if param.grad is not None:
        idxs = [0, 4]  # not fixed in reallity but doesn't matter
        sub_hessian = compute_sub_hessian(param.grad, param, idxs)
        print(sub_hessian) 
</code></pre>
<p>It works, but it is very inefficient. Indeed, the row:</p>
<pre data-code-wrap="python"><code class="lang-python">row_hessian = torch.autograd.grad(flat_grad[idx], weights, retain_graph=True)[0]
</code></pre>
<p>calculates “d” derivatives, one for each parameter belonging to the weight tensor, instead of “k”.</p>
<p>Does anyone see a better way of doing this?<br>
Thanks for your help.</p>

KFrank: <p>Hi Leonard!</p>
<aside class="quote no-group quote-modified" data-username="Leonard_Deroose" data-post="1" data-topic="198808" data-full="true">
<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/leonard_deroose/48/63030_2.png" class="avatar"> Leonard_Deroose:</div>
<blockquote>
<p>If my gradient has d components, I select k components from d (k&lt;d) and calculate the hessian (sub-hessian) on these k selected components. This means that the hessian (sub-hessian) does not have a form (d,d) but (k, k).<br>
…<br>
It works, but it is very inefficient. Indeed, the row:</p>
<pre data-code-wrap="python"><code class="lang-python">row_hessian = torch.autograd.grad(flat_grad[idx], weights, retain_graph=True)[0]
</code></pre>
<p>calculates “d” derivatives, one for each parameter belonging to the weight tensor, instead of “k”.</p>
<p>Does anyone see a better way of doing this?</p>
</blockquote>
</aside>
<p>There’s really no way to avoid computing all “d” derivatives.  Pytorch is<br>
designed to work with entire tensors, so it naturally computes the gradient<br>
with respect to the full <code>weights</code> tensor.</p>
<p>If you were to subvert this system somehow so that you only actually<br>
computed “k” derivatives and only actually performed the smaller number<br>
of floating-point operations required to compute the subset of derivatives,<br>
your code would almost certainly run <em>more slowly,</em> even while performing<br>
fewer floating-point operations.  This is because streaming the full tensor<br>
computations through the floating-point pipeline and ignoring the unused<br>
computations turns out to be much faster than interrupting those streams<br>
in order to perform just the required computations.</p>
<p>However, you can probably speed things up by pushing your python <code>for</code><br>
loop in <code>compute_sub_hessian()</code> down into pytorch’s infrastructure by<br>
using “batched” gradients with something along the lines of:</p>
<pre data-code-wrap="python"><code class="lang-python">sub_hessian = torch.autograd.grad (flat_grad, weights, grad_outputs = torch.nn.functional.one_hot (torch.tensor (idxs)).float(), is_grads_batched = True)
sub_hessian = sub_hessian[:, idxs]
</code></pre>
<p>Note that both this version and your for-loop version compute <code>k * d</code><br>
individual derivatives and discard the ones that aren’t desired.  This<br>
version just (hopefully) speeds things up by using the <code>is_grads_batched</code><br>
feature (which uses <code>vmap</code> under the hood).</p>
<p>Best.</p>
<p>K. Frank</p>

Leonard_Deroose: <p>Hi KFrank,</p>
<p>Thanks for the answer and also for the tip to speed things up. This will be particularly useful when k is large.</p>
<p>Best,<br>
Léonard</p>