Compute expected transformation of gradient

username: filthyswit
cooked: <p>Consider an additive loss (e.g., MSE) of the form <code>F(y_hat, y) = sum f(y_hat[i], y[i])</code>. We can think of this as an expectation over the empirical distribution induced by the training set. That is, the above can also be written as the expectation <code>E[ f(Y_hat, Y) ]</code>.</p>
<p>When we execute the following code</p>
<pre><code class="lang-auto">optim.zero_grad()
loss.backward()
optim.step()
</code></pre>
<p>the result is to step in the negative direction of <code>grad F(y_hat, y) = sum grad f(y_hat[i], y[i])</code>. Again, in expectation notation, this is <code>E[ grad f(Y_hat, Y) ]</code>.</p>
<p>I would like to compute the expectation of the transformation <code>T</code> of the gradient. That is, <code>E[ T(grad f(Y_hat, Y)) ]</code>. How can I do this in PyTorch?</p>

username: AlphaBetaGamma96
cooked: <p>What do you mean by transformation <code>T</code>? If it’s an arbitrary function, surely you could just apply that to the individual losses pre-expectation, then backpropagate and you’d get what you want?</p>
<p>Also, your resultant expectation is only valid if the training data is independent of your model! (In most cases this is true, but something to be aware of in case it isn’t in the future)</p>

username: filthyswit
cooked: <p>Yes, I would like to apply <code>T</code> before the expectation is taken, but in PyTorch and other deep learning frameworks, gradients are aggregated over samples. So you can’t do what you suggest (at least not without some elbow grease). That’s effectively what I’m asking here: how/where do I apply the elbow grease to achieve this?</p>
<p>To convince you that this is non-standard, what I am describing has been implemented for TensorFlow here: <a href="https://github.com/ProbabilisticNumerics/cabs" class="inline-onebox" rel="noopener nofollow ugc">GitHub - ProbabilisticNumerics/cabs: Tensorflow implementation of SGD with Coupled Adaptive Batch Size (CABS)</a></p>

username: AlphaBetaGamma96
cooked: <p>So, if <code>loss</code> is defined as the <code>sum</code> of some individual loss values then would you be able to apply <code>T</code> there before you define <code>loss</code>? So,</p>
<pre><code class="lang-auto">x = torch.randn(10, 2) #batch size = 10
yhat = model(x) #R^2 -&gt; R^1 function represented by ANN
y = torch.randn(10, 1) #target values 

individual_losses = (y - yhat).pow(2)
individual_losses_transformed = T(individual_losses)
loss = individual_losses_transformed.mean()

#Then call optimizer...
optim.zero_grad()
loss.backward()
optim.step()
</code></pre>
<p>So along the lines of define your own custom <code>MSELoss</code> function and let autograd handle the rest?</p>

username: filthyswit
cooked: <p>Your code is computing the <strong>expected gradient of the transformation</strong>, not the <strong>expected transformation of the gradient</strong>.</p>
<p>In other words, your code computes</p>
<pre><code class="lang-auto">E[ grad T ( f ( Y_hat, Y ) ) ]
</code></pre>
<p>I want</p>
<pre><code class="lang-auto">E[ T ( grad f( Y_hat, Y ) ) ] 
</code></pre>
<p>(The order of operations differs in the above two)</p>

username: AlphaBetaGamma96
cooked: <p>Ah yes, those operations won’t commute. In that case, an efficient way of doing would be per-sample gradients.</p>
<p>There are 2 ways to compute per-sample gradients</p>
<p>1.Via hooks (discussion <a href="https://discuss.pytorch.org/t/per-sample-gradient-should-we-design-each-layer-differently/57293">here</a>)<br>
2. FuncTorch (Repo <a href="https://github.com/pytorch/functorch" rel="noopener nofollow ugc">here</a>)</p>
<p>As you’re wanting to apply a transformation, FuncTorch will probably be best way as that allows for higher-order gradients of per-sample gradients whereas the hooks method is purely to get the gradients and you won’t be able to differentiate them if your transformation requires it, you’ll also need to define some manual derivatives too which can get messy).</p>

username: filthyswit
cooked: <p>Thank you. I will check that out and get back to you.</p>