Gradient of intermediate variable as loss input

rubencart: <p>I have a tensor <code>logZ</code> that’s the result of a computation. I want to use this tensor to compute 2 different losses used to update the parameters of <code>self.forward(input)</code>:</p>
<ul>
<li>One loss is computed as a function of the gradient of <code>logZ</code> w.r.t. another tensor, say <code>A</code>.</li>
<li>Another loss is computed directly from <code>logZ</code>.</li>
</ul>
<p>Would this be the right way to do this (especially regarding creating and retaining the differentiation graph correctly)?</p>
<pre><code class="lang-auto">A = self.forward(input)
logZ = some_differentiable_fn(A)

# compute input for first loss
logZ.sum().backward(retain_graph=True, create_graph=True)
marginals = A.grad
loss1 = some_other_differentiable_fn(marginals).mean()

# reset param gradients that got filled because of first .backward()
for p in self.parameters():
    p.grad = None

# second loss and backward on sum of losses, update params
loss2 = (-logZ).mean()
(loss1 + loss2).backward()
optimizer.step()
</code></pre>
<p>I tried computing <code>marginals</code> as <code>marginals = autograd.grad(logZ, A, retain_graph=True, create_graph=True)</code>, but I wasn’t sure whether the newly created graph out of this operation would reach the network’s parameters (since A is an output of the network and not an input).</p>

KFrank: <p>Hi Ruben!</p>
<aside class="quote no-group" data-username="rubencart" data-post="1" data-topic="206887" 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/rubencart/48/25961_2.png" class="avatar"> rubencart:</div>
<blockquote>
<p>Would this be the right way to do this (especially regarding creating and retaining the differentiation graph correctly)?</p>
</blockquote>
</aside>
<p>Yes, this is right (but I haven’t looked at your code in detail).</p>
<aside class="quote no-group">
<blockquote>
<pre><code class="lang-auto"># compute input for first loss
logZ.sum().backward(retain_graph=True, create_graph=True)
marginals = A.grad
loss1 = some_other_differentiable_fn(marginals).mean()

# reset param gradients that got filled because of first .backward()
for p in self.parameters():
    p.grad = None
</code></pre>
</blockquote>
</aside>
<p>It’s probably a little more convenient to:</p>
<pre data-code-wrap="python"><code class="lang-python">marginals = torch.autograd.grad (logZ.sum(), A, create_graph = True)[0]
</code></pre>
<p>because <code>autograd.grad()</code> doesn’t set the parameter gradients, so you<br>
don’t need to then reset them.</p>
<aside class="quote no-group">
<blockquote>
<p>I tried computing <code>marginals</code> as <code>marginals = autograd.grad(logZ, A, retain_graph=True, create_graph=True)</code></p>
</blockquote>
</aside>
<p>Yes, as noted above, this will work (provided you <code>.sum()</code> <code>logZ</code> to a<br>
scalar and use <code>[0]</code> to extract the first (and only) term of the tuple<br>
returned by <code>autograd.grad()</code>).</p>
<aside class="quote no-group">
<blockquote>
<p>but I wasn’t sure whether the newly created graph out of this operation would reach the network’s parameters (since A is an output of the network and not an input).</p>
</blockquote>
</aside>
<p>Yes, the fact that <code>A</code> is an output of the network is fine.  <code>A</code> carries<br>
<code>requires_grad = True</code>, so the computation graph continues forward<br>
from <code>A</code> to <code>logZ</code>.  It doesn’t matter that <code>A</code> is an “intermediate variable”<br>
and not a “leaf” of the computation graph.</p>
<p>As an aside, when <code>create_graph = True</code>, <code>retain_graph</code> will default<br>
to <code>True</code>, so you don’t have to specify it explicitly (but it doesn’t hurt if<br>
you do).</p>
<p>Best.</p>
<p>K. Frank</p>

rubencart: <p>Thank you for your answer, K. Frank :-). Indeed it seems a little more practical to use <code>autograd.grad</code> instead of <code>.backward</code>.</p>
<p>I was a little confused, since I guess <code>marginals = autograd.grad(logZ, A, create_graph=True)</code> backpropagates the gradient of <code>logZ</code> up until <code>A</code> and stops there, while I want <code>self</code>’s parameters to be updated. However, since we use <code>create_graph=True</code>, this differentiation gets attached to the computation graph, which does include the network’s parameters, so we can later do <code>.backward()</code> on the computed gradient of <code>logZ</code> and this time the new gradient will reach the parameters. Pretty neat.</p>
<p>Am I right to conclude then that <code>marginals = autograd.grad(logZ, A, create_graph=True)</code> should also be more efficient in time and space than <code>logZ.sum().backward()</code>, since we only actually differentiate up until the variable that we care about, and not the entire graph?</p>

KFrank: <p>Hi Ruben!</p>
<aside class="quote no-group" data-username="rubencart" data-post="3" data-topic="206887" 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/rubencart/48/25961_2.png" class="avatar"> rubencart:</div>
<blockquote>
<p>Am I right to conclude then that <code>marginals = autograd.grad(logZ, A, create_graph=True)</code> should also be more efficient in time and space than <code>logZ.sum().backward()</code>, since we only actually differentiate up until the variable that we care about, and not the entire graph?</p>
</blockquote>
</aside>
<p>This is correct.  Furthermore, <code>.backward()</code> (unlike <code>autograd.grad()</code>)<br>
will populate the <code>.grad</code> properties of the leaf tensors (typically your<br>
model parameters) of the computation graph.  This also takes time and<br>
space.</p>
<p>Best.</p>
<p>K. Frank</p>