Getting elementwise gradients is slow

Liran_Ringel: <p>I’m trying to get the gradient of each sample like that:</p>
<blockquote>
<p>grads = torch.autograd.grad(losses, model.parameters(), torch.eye(len(losses), device=losses.device), retain_graph=True, is_grads_batched=True, allow_unused=True)</p>
</blockquote>
<p>where <code>losses</code> is a vector of losses, one for each sample.<br>
It works. However, it is significantly slower than computing the gradient of the mean loss like that:</p>
<blockquote>
<p>torch.autograd.grad(losses.mean(), model.parameters(), retain_graph=True, allow_unused=True)</p>
</blockquote>
<p>My guess was that these should execute with similar runtime because both do the same calculations. The second code snippet also has to do the average.</p>
<p>It is suppose to be slower?<br>
Can I achieve my goal in a faster way?</p>
<p>Thanks!</p>

KFrank: <p>Hi Liran!</p>
<aside class="quote no-group quote-modified" data-username="Liran_Ringel" data-post="1" data-topic="196935" 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/liran_ringel/48/67607_2.png" class="avatar"> Liran_Ringel:</div>
<blockquote>
<blockquote>
<p>grads = torch.autograd.grad(losses, model.parameters(), torch.eye(len(losses), device=losses.device), retain_graph=True, is_grads_batched=True, allow_unused=True)</p>
</blockquote>
</blockquote>
</aside>
<p>This line of code computes a <em>batch</em> of gradients, one for each element of<br>
<code>losses</code>.  (It should be faster than a python loop because <code>autograd.grad</code><br>
presumably uses <code>vmap()</code> vectorization under the hood, but it still performs<br>
multiple gradient computations.)</p>
<aside class="quote no-group">
<blockquote>
<blockquote>
<p>torch.autograd.grad(losses.mean(), model.parameters(), retain_graph=True, allow_unused=True)</p>
</blockquote>
</blockquote>
</aside>
<p>This line of code, however, computes just a <em>single</em> gradient of the scalar<br>
quantity <code>losses.mean()</code>, so it really is doing less work.</p>
<aside class="quote no-group">
<blockquote>
<p>My guess was that these should execute with similar runtime because both do the same calculations.</p>
</blockquote>
</aside>
<p>To reiterate, the two versions are not doing the same calculations.  The<br>
first computes a batch of gradients, while the second computes just a<br>
single gradient.</p>
<aside class="quote no-group">
<blockquote>
<p>The second code snippet also has to do the average.</p>
</blockquote>
</aside>
<p>This is true.  But by averaging the <code>losses</code> together, you reduce them to a<br>
single scalar for which you compute a single gradient, rather than a batch.</p>
<aside class="quote no-group">
<blockquote>
<p>Can I achieve my goal in a faster way?</p>
</blockquote>
</aside>
<p>Possibly, but no guarantees.  Depending on your use case it might be<br>
possible to rearrange your computation to use pytorch’s so-called<br>
<a href="https://pytorch.org/tutorials/intermediate/forward_ad_usage.html" rel="noopener nofollow ugc">forward-mode automatic differentiation.</a>  This could be faster, but could<br>
also be slower.</p>
<p>Best.</p>
<p>K. Frank</p>