Applying an input-dependent transformation to the gradients during the backward pass

abenechehab: <p>Hi everyone,</p>
<p>I have a specific usage of autograd and I want to know your opinion about the efficient way to make it work.</p>
<p>For context, I want to apply a transformation to the gradients before executing the optimization step. This transformation multiplies the gradient of the loss with the gradient of the prediction function.</p>
<p>For a parameter <em>p</em>, a prediction function (neural network) <em>f</em>, and a loss <em>L</em>, the transformation is the following:</p>
<pre><code class="lang-auto">grad_p = d(f_p)/d(p) * d(L_p)/d(p)
</code></pre>
<p>The problem is that both gradients are input-dependent, which means that the transformation must be applied before the gradient accumulation step during the backward pass.</p>
<p>The naive solution I can think of is using a batch_size of 1, in this case I can proceed as follows:</p>
<pre><code class="lang-auto"># batch input X, batch groundtruth y
prediction = model(X)
prediction.backward(retain_graph=True)
gradient = []
For param in model.parameters():
    gradients.append(param.grad)
model.zero_grad()
loss = loss_fn(prediction, y)  # MSE
loss.backward()
For i, param in enumerate(model.parameters()):
    param.grad *= gradients[i]
</code></pre>
<p>Now this is very compute-inefficient and I would like to ask you if there is a way to do it by batch applying the transformation on the fly during the backward pass?</p>
<p>I have thought about registering a backward hook but it’s not clear to me how to pass the first gradient to the hook.</p>
<p>Thank you.</p>

yiftach: <p>That is an interesting question, and thanks for the clear code example.</p>
<p>Using a backward hook as you suggest is how I would approach it. I’ve not tested it, but have you tried something like adding:</p>
<pre><code class="lang-auto">loss.register_hook(lambda param_grad: param_grad * gradients.pop(0))
</code></pre>
<p>Right before <code>loss.backward()</code>? What I do here is consume the first list of gradients in the same order (this relies on both backward passes to occur in the same order).</p>

abenechehab: <p>Hi yiftach,</p>
<p>Thank you for your answer.</p>
<p>In the case of a batch size &gt; 1, and if we first compute the loss of the prediction <code>prediction.backward(retain_graph=True)</code>, there will be gradient accumulation across the batch.</p>
<p>Then if we apply the hook to the loss as you suggested we will again do the multiplication (call the hook) after the accumulation across the batch.</p>
<p>So for a set of batches <code>B</code>, parameter <code>\theta</code>, model <code>f_\theta</code>, and loss function <code>L</code>, what we would have computed at the end is the first equation, while the expected behavior is the second equation:</p>
<p><img src="https://latex.codecogs.com/png.image?\dpi{110}$$new\_grad_\theta=\frac{1}{|B|}\sum_{b\in&amp;space;B}\big(\frac{1}{|b|}\sum_{(X,y)\in&amp;space;b}\nabla_\theta&amp;space;L(f_\theta(X),y)\big)\big(\frac{1}{|b|}\sum_{(X,y)\in&amp;space;b}\nabla_\theta&amp;space;f_\theta(X)\big)$$\newline\begin{equation}new\_grad_\theta=\frac{1}{|B|}\sum_{b\in&amp;space;B}\frac{1}{|b|}\sum_{(X,y)\in&amp;space;b}\big(\nabla_\theta&amp;space;L(f_\theta(X),y)\big)\big(\nabla_\theta&amp;space;f_\theta(X)\big)\end{equation}" title="$$new\_grad_\theta=\frac{1}{|B|}\sum_{b\in B}\big(\frac{1}{|b|}\sum_{(X,y)\in b}\nabla_\theta L(f_\theta(X),y)\big)\big(\frac{1}{|b|}\sum_{(X,y)\in b}\nabla_\theta f_\theta(X)\big)$$\newline\begin{equation}new\_grad_\theta=\frac{1}{|B|}\sum_{b\in B}\frac{1}{|b|}\sum_{(X,y)\in b}\big(\nabla_\theta L(f_\theta(X),y)\big)\big(\nabla_\theta f_\theta(X)\big)\end{equation}" width="585" height="52"></p>
<p>Am I right? or this not the usage of the hook that you suggested?</p>