Higher-order gradients w.r.t. different functions

lobachevzky: <p>I am writing a program with essentially the following update rules:</p>
<p><img src="https://discuss.pytorch.org/uploads/default/original/2X/0/09cd6f185642b2df96c41d179070c59add526bf1.png" alt="tex" data-base62-sha1="1oIr7Eog9780Jos9yh0JQBHpSyl" width="337" height="171"></p>
<p>I have the following simplified code:</p>
<pre><code class="lang-auto">import torch
from torch import optim

def f(x, y):
    return (x + y)**2

def g(x, y):
    return x * y

x = torch.tensor([3.], requires_grad=True)
y = torch.tensor([4.], requires_grad=True)
x_optim = optim.SGD([x], lr=1.)
y_optim = optim.SGD([y], lr=1.)
ddx, = torch.autograd.grad(f(x, y).mean(), x, create_graph=True)
# 2(x + y) = 14

ddx.mean().backward()
# x.grad = d^2/dx^2 f(x, y) = 2
# y.grad = d/dy d/dx f(x, y) = 2

ddx, = torch.autograd.grad(g(x, y).mean(), x)
x.grad = ddx
# x.grad = d/dx g(x, y) = y = 4

y_optim.step()
# y = 2
# x = 3

x_optim.step()
# y = 2
# x = -1
</code></pre>
<p>My question is: is this the best (and most performant) way to do this? <code>x.grad = ddx</code> is not so bad, but when it’s all the parameters of several neural networks, this involves a lot of careful matching of gradients to variables and a lot of room for error.<br>
Thanks!</p>

tom: <p>I won’t say anything about performance, but don’t do this! As you rightly point out, <code>x.grad = ...</code> is playing with the fire.<br>
A better way to use PyTorch’s capabilities properly is to detach the things you don’t want to backward into. In order to do the second derivative, you need to have a copy of <code>x</code> that is detached, but requires gradient.<br>
So this is what your sample would look like:</p>
<pre><code class="lang-python">import torch
from torch import optim

def f(x, y):
    return (x + y)**2

def g(x, y):
    return x * y

x = torch.tensor([3.], requires_grad=True)
x_detached = x.detach().requires_grad_()
y = torch.tensor([4.], requires_grad=True)
x_optim = optim.SGD([x], lr=1.)
y_optim = optim.SGD([y], lr=1.)
ddx, = torch.autograd.grad(f(x_detached, y).mean(), x_detached, create_graph=True)
# 2(x + y) = 14

ddx.mean().backward()
# x_detached.grad = d^2/dx^2 f(x, y) = 2
# x.grad unaffected
# y.grad = d/dy d/dx f(x, y) = 2

g(x, y.detach()).mean().backward()
# x.grad = d/dx g(x, y) = y = 4

y_optim.step()
# y = 2
# x = 3

x_optim.step()
# y = 2
# x = -1
</code></pre>
<p>Best regards</p>
<p>Thomas</p>

lobachevzky: <p>Thanks for the response. So in the actual code, <code>f()</code> is a neural network (actually several) and <code>x</code> is the parameters of these networks.</p>
<p>I thought that the way to apply <code>.detached</code> to the parameters would be something like this:</p>
<pre><code class="lang-auto">import torch

x = torch.ones(1, requires_grad=True)
dense = torch.nn.Linear(1, 1)
for var in dense.parameters():
    var.requires_grad = False
y = dense(x)
print(torch.autograd.grad(y, dense.parameters()))
</code></pre>
<p>This raises</p>
<pre><code class="lang-auto">RuntimeError: One of the differentiated Tensors does not require grad
</code></pre>
<p>because <code>torch.no_grad()</code> seems to prevent gradients from flowing back to <code>x</code>.</p>
<p>So I assume that this is not the correct way to detach the parameters of the networks, right?</p>

tom: <p>I must admit that I don’t even know what that bit of code might try to achieve (I prefer <code>.requires_grad_(False)</code>, but that obviously is in conflict with takimg the grad?). Maybe just leaving the nets as they are is a good strategy or you can describe what you are trying to compute or how it compares to your initial posting.</p>
<p>Best regards</p>
<p>Thomas</p>

lobachevzky: <p>Sure. Referencing the equations in my initial post, f is essentially my objective function. Both x and y are parameters of neural networks. I am trying to learn x values that maximize f(x, y). I am trying to learn y values that yield large gradients for f(x,y). Thus the update rule for y is in the direction of greater df/dx.</p>
<p>The code in my previous post was an attempt at applying <code>.detach</code> to the parameters of a neural network, which it seemed like you were suggesting in your post.</p>
<p>Hope that clarifies things a little. Also, thank you for your help.</p>

tom: <p>Hi,</p>
<p>thanks for elaborating on your use case.<br>
When you say <em>large gradients</em>, do you have a specific norm in mind? I think you need to specify that to go multi-dimensional (if x isn’t scalar, I think that the update rule in your original runs into trouble with dimensionality.)<br>
I’m not entirely sure what is the best way to “isolate” the two backwards passes in full generality. Does the function f factor into “application of NNs” and “compute loss”? Then detaching the outputs of the NNs might work well.</p>
<p>Best regards</p>
<p>Thomas</p>

lobachevzky: <blockquote>
<p>When you say  <em>large gradients</em> , do you have a specific norm in mind?</p>
</blockquote>
<p>Currently, I am using the global l2-norm:</p>
<pre><code class="lang-auto">def global_norm(grads):
    norm = 0
    for grad in grads:
        norm += grad.norm(2)**2
    return norm**.5
</code></pre>
<blockquote>
<p>Does the function f factor into “application of NNs” and “compute loss”?</p>
</blockquote>
<p>When you say “factor” here, I’m assuming you don’t mean multiplicative factorization. Perhaps it will clarify to mention that I have settled on an intermediate solution that is not very performant. Essentially:</p>
<pre><code class="lang-auto">g_optimizer.zero_grad()
f_optimizer = optim.Adam(f.parameters())
g_optimizer = optim.Adam(g.parameters())
f_grad = torch.autograd.grad(compute_f_loss(), f.parameters())
global_norm(f_grad).backward()  # this assigns grads to both f.parameters() and g.parameters()
g_optimizer.step()  # this only steps g.parameters()

f_optimizer.zero_grad()  # this clears the gradients assigned to f.parameters()
compute_f_loss().backward()
f_optimizer.step()  # this only steps f.parameters()
</code></pre>
<p>There are two performance problems:</p>
<ol>
<li>It requires two backward passes, which is very slow. It would be great if I could reuse <code>f_grad</code> instead of recomputing it with <code>compute_f_loss().backward()</code>.</li>
<li>It requires two calls to <code>compute_f_loss()</code>. This might be remedied with <code>retain_graph=True</code>.</li>
</ol>

tom: <p>By factor I mean that you spell out the call to f and g in the above code snippet and separate it from the loss.<br>
In that case you can adapt the thing I wrote about your specific case by detaching the outputs of f and g before feeding them into the loss similar to what I did with your original example and get the right thing.</p>
<p>Best regards</p>
<p>Thomas</p>

lobachevzky: <p>I realize that one thing that was not clear is that the <code>compute_f_loss()</code> actually computes a loss that is a function of the output of <code>f</code> <em>and</em> of <code>g</code>. Essentially, the output of <code>g</code> is fed into <code>f</code>.</p>
<p>The problem with detaching the outputs of <code>f</code> and <code>g</code> before feeding them into the loss is that I need the gradient of <code>calculate_f_loss()</code> with respect to <code>f.parameters()</code> and then I need the 2nd order gradient with respect to <code>g.parameters()</code>. If I detached the outputs of <code>f</code> and <code>g</code>, then <code>calculate_f_loss().backward()</code> would not populate the <code>.grad</code> fields of <code>f.parameters()</code> or <code>g.parameters()</code>.</p>

tom: <p>Yes, that is why I made a detached copy to use in the second order computation. Did you try that and it gave unsatisfactory results?</p>
<p>Best regards</p>
<p>Thomas</p>

lobachevzky: <p>One thing that we didn’t quite clear up from earlier was how exactly to <code>.detach()</code> the parameters of a neural network. Would I have to create a copy of the neural network somehow? I think as soon as we clear that up, I could implement something like what you first proposed. Or were you suggesting detaching the output of the network instead of it’s parameters?</p>
<p>I should add that, as a newcomer to Torch, I’m not exactly sure what <code>.detach()</code> does. My interpretation of the <a href="https://pytorch.org/docs/stable/autograd.html#torch.Tensor.detach" rel="nofollow noopener">docs</a> is that it acts something like <code>tf.stop_gradient</code> in Tensorflow, preventing gradients from flowing through the variable. But if that is the case, then what does calling <code>x.detach().requires_grad_()</code> do? Does this essentially create a copy of the <code>x</code> node in the graph with all the same in- and out-edges? And what is the performance benefit to doing this?</p>

tom: <p>Yes, exactly. If you detach the output, nothing will be flowing into the parameters anyways (except when you use the pre-detached output)…</p>
<p>Best regards</p>
<p>Thomas</p>

lobachevzky: <p>Hi Thomas,<br>
So I had a chance to implement this and test it against the original version. The two versions seem to perform identically, both in terms of outputs and runtime. Both are very slow with the forward pass and backward pass accounting for the majority of the runtime.</p>
<p>So my questions are:</p>
<ol>
<li>Is it possible to avoid the <code>compute_f_loss().backward()</code> call by somehow reusing <code>f_grad</code>? (This would avoid a second backward pass.)</li>
<li>Is it possible to reuse the first call to <code>compute_f_loss()</code> so that the forward pass only needs to be done once?</li>
</ol>
<p>In principle this should be possible.</p>

tom: <p>It’s hard to debug this without looking at the code. You can definitely avoid .backward through the nets several times by just summing the two loss components. The bits that are necessarily evaluated twice in the forward won’t magically go away for the backward, though.</p>
<p>Best regards</p>
<p>Thomas</p>

lobachevzky: <aside class="quote no-group" data-username="tom" data-post="14" data-topic="33377">
<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/tom/48/3162_2.png" class="avatar"> tom:</div>
<blockquote>
<p>You can definitely avoid .backward through the nets several times by just summing the two loss components.</p>
</blockquote>
</aside>
<p>The issue is not calling <code>.backward</code> twice. The issue is redundant backward passes. Specifically: <code>f_grad</code> requires a backward pass and <code>compute_f_loss().backward()</code> requires a backward pass.</p>
<p>I can post my code but it is long and complicated. The <a href="https://discuss.pytorch.org/t/higher-order-gradients-w-r-t-different-functions/33377/7">small example</a> I gave in an earlier post had all the same issues.</p>
<p>I don’t know if this is helpful, but if I were working in tensorflow, I would compute <code>f_grad</code> with <code>tf.gradients</code> and then I could feed those same gradients to the optimizer using <code>opimizer.apply_gradients(variables, gradients)</code>. If torch has some kind of similar mechanism, it would solve my issue.</p>

pinocchio: <p><a class="mention" href="/u/lobachevzky">@lobachevzky</a></p>
<p>Why did you not use the pytorch higher library?</p>
<p><a href="https://higher.readthedocs.io/en/latest/toplevel.html?highlight=innerloop_ctx#higher.innerloop_ctx" class="onebox" target="_blank" rel="nofollow noopener">https://higher.readthedocs.io/en/latest/toplevel.html?highlight=innerloop_ctx#higher.innerloop_ctx</a></p>

Enamex: <p>Would <code>retain_graph</code> in-between optim steps here save you the recompute of <code>f</code>? The semantics will be different, though (the optim steps for f, g would fall in lockstep).</p>