Computing output gradients for NTK

username: Simone256
cooked: <p>Hello everyone, sorry in advance for the beginner question.</p>
<p>Let’s say I have a small model <code>net</code>, with d parameters, with dimension c as output (like a classifier for MNIST). Let’s also say that I feed this model with a minibatch of images (batchsize B). At this point I would have something like</p>
<pre><code class="lang-auto">outputs = net(images)
</code></pre>
<p>where the shape of output will be (B, c). I would like to take the gradients with respect to the parameters of the network, something like this:</p>
<pre><code class="lang-auto">outputs.backward()
feature_maps = [param.grad for param in iter(net.parameters())]
</code></pre>
<p>And having as output this <code>feature_maps</code> with shape (B, c, d). This is not doable for many reasons I guess, in particular pytorch complains that when I call <code>.backward()</code> I need to have a scalar tensor in front of it.</p>
<p>This problem is solvable, I just have to feed one image at a time in the model, and compute the gradient for each output separately, and then putting everything together again. Is there a smarter/faster way to achieve this?</p>
<p>I’m basically computing the feature maps for computing the neural tangent kernel later (even if I’m doing it for each output separately for this classifier), so I believe there should be a better usual way for doing this.</p>

username: ptrblck
cooked: <p>You could either reduce the <code>ouputs</code> to a scalar value and call <code>.backward()</code> on it or you could alternatively pass the gradients to the <code>backward</code> operation.<br>
<a class="mention" href="/u/alband">@albanD</a> explains the reasoning behind this in <a href="https://discuss.pytorch.org/t/why-do-we-need-to-pass-the-gradient-parameter-to-the-backward-function-in-pytorch/116079/4">this post</a>.</p>

username: Simone256
cooked: <p>Thank you for your answer!</p>
<p>Then, it appears that <code>.backward()</code> always “returns” a vector with the same dimension of the parameters (to be more specific, it saves in the <code>.grad</code> attribute of each parameter only one number. Either the derivative of a scalar loss with respect to this parameter, either the scalar product between a fixed vector <code>vec</code> (given as input in <code>.backward()</code>), and the vector made of all the derivatives with respect to this parameter of a vectorial loss).</p>
<p>In particular, there is no way to parallelize the computation of the Jacobian? I guess there’s not…</p>
<p>Furthermore, let’s say my <code>vec</code> I give as input in <code>.backward()</code> is just a vector of ones, since I’m interested in the sum of the gradients. Is it better to approach the problem this way, or to sum up the outputs of my vectorial loss before? (This should be equivalent for linearity, but maybe one method is computationally / numerically better…</p>

username: albanD
cooked: <aside class="quote no-group" data-username="Simone256" data-post="3" data-topic="117805">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/s/5f8ce5/48.png" class="avatar"> Simone256:</div>
<blockquote>
<p>In particular, there is no way to parallelize the computation of the Jacobian? I guess there’s not…</p>
</blockquote>
</aside>
<p>There is an experimental way to do it <img src="https://discuss.pytorch.org/images/emoji/apple/wink.png?v=12" title=":wink:" class="emoji" alt=":wink:" loading="lazy" width="20" height="20"> see the vectorize flag on jacobian <a href="https://pytorch.org/docs/stable/autograd.html#torch.autograd.functional.jacobian">here</a> or using the <code>vmap</code> prototype directly in your code <a href="https://pytorch.org/docs/master/generated/torch.vmap.html?highlight=vmap#torch.vmap">doc here</a>.</p>
<blockquote>
<p>Furthermore, let’s say my <code>vec</code> I give as input in <code>.backward()</code> is just a vector of ones, since I’m interested in the sum of the gradients. Is it better to approach the problem this way, or to sum up the outputs of my vectorial loss before? (This should be equivalent for linearity, but maybe one method is computationally / numerically better…</p>
</blockquote>
<p>The two are the same really.<br>
Giving the vector of ones will be marginally faster because you create one less Node in the graph, but tha won’t make any measurable difference really.</p>