Calculate individual gradients for each example in mini-batch?

username: garland
cooked: <p>Say I have an the output of <code>y=model(x)</code> where <code>y.size()</code> is [batch_dimensions, 1]. I would then like to calculate the individual gradients of each example in the output with respect to the <code>model.parameters()</code>. Naively, I could do:</p>
<pre><code class="lang-auto">grads = [torch.autograd.grad(y[batch], model.parameters(), retain_graph=True) batch in range(batch_dimensions)]
</code></pre>
<p>but this is serial and slow.</p>
<p>I know I can, however, do:</p>
<pre><code class="lang-auto">grads = torch.autograd.grad(y, model.parameters(), grad_outputs=torch.ones(y.size()))
</code></pre>
<p>which is an awful lot faster. However, this seems to sum the gradients <em>across</em> the batches, which is not quite what I want. How could I possibly achieve what I require efficiently?</p>
<p>Cheers.</p>

username: Yaroslav_Bulatov
cooked: <p>Try this:</p>
<blockquote>
<p><a href="https://github.com/cybertronai/autograd-hacks#per-example-gradients" class="inline-onebox" rel="noopener nofollow ugc">GitHub - cybertronai/autograd-hacks</a></p>
</blockquote>

username: roylight
cooked: <p><a class="mention" href="/u/yaroslav_bulatov">@Yaroslav_Bulatov</a>  Hi Can this module support LSTM? Thank you!</p>

username: richard
cooked: <p><a class="mention" href="/u/roylight">@roylight</a>, could I ask about your use case for individual per-example gradients?</p>
<p>To answer your question, cybertronai/autograd-hacks don’t support PyTorch’s nn.LSTM out of the box, but if you write a custom LSTM (using nn.Linear, some activation functions, and a for-loop), cybertronai/autograd-hacks should be able to help (there is an implementation for per-sample-gradients of nn.Linear in there).</p>

username: Yaroslav_Bulatov
cooked: <p>Indeed, it doesn’t support anything other than Conv/Linear, so would need some work. In ideal world, an autovectorization operation would be able to do this automatically, but right now it seems manual algebra is required.</p>
<p><a class="mention" href="/u/richard">@richard</a> one common demand for per-example gradient norms is from privacy people – a datapoint with a large gradient could influence the final solution strongly, so you could “leak” information about that datapoint through the final model.</p>

username: richard
cooked: <p><a class="mention" href="/u/yaroslav_bulatov">@Yaroslav_Bulatov</a> we’re working on prototyping an auto-vectorization operation like JAX’s vmap (<a href="https://github.com/pytorch/pytorch/issues/42368">https://github.com/pytorch/pytorch/issues/42368</a>) in the hopes that it may be able to compute per-sample-gradients (among other use cases). There’s a slight problem where it doesn’t quite work for the per-sample-gradient use case, because of how PyTorch internals are organized, but we’re figuring that out.</p>
<p>I’m curious if you’ve come across other demands for per-example gradients. I’ve spoken with some folks in the differential privacy space but haven’t been able to find other use cases.</p>

username: Yaroslav_Bulatov
cooked: <p>The other place is optimization research. There are thousands of papers whose implementation would benefit from efficient access to per-example gradients. Few examples –</p>
<ul>
<li>Optimize faster: cluster gradients, then undersample large clusters, <a href="https://arxiv.org/abs/1405.3080" rel="nofollow noopener">paper</a>.</li>
<li>Find potentially mislabeled examples by checking if gradients change rapidly, <a href="https://arxiv.org/abs/2008.11600" rel="nofollow noopener">paper</a>.</li>
<li>Predict generalization by measuring sharpness of feasible solution. This needs variance of per-example Hessians, which are gradients of the gradient function. <a href="https://papers.nips.cc/paper/8049-how-sgd-selects-the-global-minima-in-over-parameterized-learning-a-dynamical-stability-perspective.pdf" rel="nofollow noopener">paper</a>.</li>
</ul>
<p>Popular optimizers like AdaGrad and natural gradient were formulated in terms of per-example gradients. Because frameworks didn’t provide efficient way to do this, all practical implementations substituted batch gradients as an approximation.<br>
Their derivatives (Adam, <a href="https://arxiv.org/abs/1503.05671" rel="nofollow noopener">KFAC</a>, <a href="https://arxiv.org/abs/1802.09568" rel="nofollow noopener">Shampoo</a>, <a href="http://proceedings.mlr.press/v97/agarwal19b/agarwal19b.pdf" rel="nofollow noopener">full AdaGrad</a>) ended up inheriting this approximation. What if this approximation is bad? Without per-example gradient support, it’s too much work to check.</p>
<p>Basically this feature would make PyTorch a more friendly research platform.</p>
<p>PS: efficient per-example gradients (as well as per-example Hessians, Hessian vector products, per-example gradient norms, etc) are all just special cases of einsum optimization. Computation for these quantities comes down to a series of multiplications and additions, and einsum optimizer is what tells you how to properly rearrange them.</p>
<p>An example of using this perspective to <a href="https://colab.research.google.com/drive/1VlirNQxEr3I36s7ra7wzYIapNpcB-l97" rel="nofollow noopener">improve on</a> PyTorch’s Hessian vector product and to compute batch of Hessian <a href="https://github.com/pytorch/pytorch/issues/28341#issuecomment-545092466" rel="nofollow noopener">norms efficiently</a></p>

username: richard
cooked: <p>Thank you for the detailed reply! I’ll take a look through the papers. This is really helpful - we’re also exploring other mechanisms to compute per-example-gradients (that aren’t through an auto-vectorization operator) and investigating use cases will help out with the decision making.</p>

username: christianperone
cooked: <p>Another project that is worth mentioning: <a href="https://backpack.pt" rel="noopener nofollow ugc">https://backpack.pt</a> (from "<a href="https://arxiv.org/abs/1912.10985%22" rel="noopener nofollow ugc">BackPACK: Packing more into backprop</a>).</p>