Autograd For Intermediate Tensors

nreHieW: <p>Hi, a question about how the autograd engine works for intermediate values.</p>
<p>Suppose I have an equation <code>(a*b + c*d) = e</code>, pretty common in matrix multiplication. My understanding is that under the hood the calculation looks something like</p>
<pre><code class="lang-auto">_t1 = a * b
_t2 = c * d
e = _t1 + _t2
</code></pre>
<p>Then when <code>.backward()</code> is called, the gradient has to flow through <code>e</code> to <code>_t1</code> and <code>_t2</code> before reaching <code>a,b,c,d</code>.</p>
<p>For large matrices, it seems that storing the gradient information of all these temporary values will be expensive memory wise but they also (?) seem necessary to be able to backpropagate to the a,b,c and d tensors.</p>
<p>I am guessing this has something to do with <code>.retain_grad()</code> but am not too sure. Any help is appreciated to help clarify my understanding.</p>
<p>Thanks!</p>

KFrank: <p>Hi Wei!</p>
<aside class="quote no-group" data-username="nreHieW" data-post="1" data-topic="205501" 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/nrehiew/48/67709_2.png" class="avatar"> nreHieW:</div>
<blockquote>
<p>Suppose I have an equation <code>(a*b + c*d) = e</code></p>
</blockquote>
</aside>
<p>Note that <code>*</code> denotes element-wise multiplication, where for matrix<br>
multiplication, one would use <code>@</code>.  (However, my answer applies equally<br>
well to both cases.)</p>
<aside class="quote no-group">
<blockquote>
<p>My understanding is that under the hood the calculation looks something like</p>
<pre><code class="lang-auto">_t1 = a * b
_t2 = c * d
e = _t1 + _t2
</code></pre>
</blockquote>
</aside>
<p>Yes, if you write the python statement <code>e = a * b + c * d</code>, python will do<br>
essentially what you’ve written down, except that <code>a * b</code> and <code>c * d</code> won’t<br>
have explicit references bound to them <em>unless</em> autograd decides (which it<br>
doesn’t) that it needs to save those tensors for the backward pass.</p>
<aside class="quote no-group">
<blockquote>
<p>Then when <code>.backward()</code> is called, the gradient has to flow through <code>e</code> to <code>_t1</code> and <code>_t2</code> before reaching <code>a,b,c,d</code>.</p>
</blockquote>
</aside>
<p>Yes (depending upon which of the tensors <code>a</code>, <code>b</code>, <code>c</code>, and <code>d</code> carry<br>
<code>requires_grad = True</code>).</p>
<aside class="quote no-group">
<blockquote>
<p>For large matrices, it seems that storing the gradient information of all these temporary values will be expensive memory wise</p>
</blockquote>
</aside>
<p>Roughly speaking, yes.  The various temporary tensors that autograd<br>
stores during the forward pass (because it needs them to perform the<br>
backward pass) are almost always the dominant cost in memory when<br>
training a model (forward pass, backward pass, optimization step).</p>
<p>In general, evaluation (forward pass, but with no need to save anything<br>
for a backward pass) requires significantly less memory than training.</p>
<aside class="quote no-group">
<blockquote>
<p>but they also (?) seem necessary to be able to backpropagate to the a,b,c and d tensors.</p>
</blockquote>
</aside>
<p>Let’s look at what autograd will save in your example.  (In general<br>
autograd doesn’t save any and all intermediate tensors – it’s pretty<br>
smart about only saving what it actually needs for the backward pass.)</p>
<p>Let me assume that all of <code>a</code>, <code>b</code>, <code>c</code>, and <code>d</code> carry <code>requires_grad = True</code>.</p>
<p>To compute the derivative of <code>a * b</code> with respect to <code>a</code>, autograd needs<br>
to save <code>b</code>.  Similarly, to compute the derivative of <code>a * b</code> with respect to<br>
<code>b</code>, autograd needs to save <code>a</code>.  Likewise, autograd will save both <code>c</code> and<br>
<code>d</code>.  Note that for your particular example, autograd does not need to<br>
save <code>_t1</code> nor <code>_t2</code> (nor <code>e</code>).  (It needs to know that <code>_t1</code> was obtained<br>
from <code>a</code> and <code>b</code> using <code>*</code>, but that is a much smaller amount of information<br>
that is independent of the size of <code>_t1</code>.)</p>
<p>(Also note that the autograd system may sometimes save auxiliary<br>
intermediate tensors that never appear explicitly in the forward pass<br>
if they support the most efficient way to perform the backward pass.)</p>
<p>So yes, autograd stores various intermediate tensors for use in its<br>
backward pass and these tensors are a big part of the memory cost.<br>
But the details of which particular tensors it needs to store depend on<br>
the specific details of your forward-pass computation.</p>
<p>Take a look at how one implements a <a href="https://pytorch.org/docs/2.3/notes/extending.html#extending-torch-autograd" rel="noopener nofollow ugc">custom autograd function</a> for<br>
further detail on how this all works.</p>
<p>Best.</p>
<p>K. Frank</p>

nreHieW: <p>Hi Frank thanks for the detailed answer!</p>
<p>Firstly, I just wanted to clarify my understanding of how autograd works for tensors/matrices . My understanding is that each scalar element in a tensor can be thought of as independent (that is the nodes it creates will naturally be different from the nodes the other values in the tensor creates). Hence <code>(a*b + c*d) = e</code> was my example where <code>a,b</code> are in the same row of the first matrix, <code>c,d</code> are in the same column of the second matrix and <code>e</code> is the result. Is my understanding accurate?</p>
<blockquote>
<p>To compute the derivative of <code>a * b</code> with respect to <code>a</code> , autograd needs to save <code>b</code> . Similarly, to compute the derivative of <code>a * b</code> with respect to<code>b</code> , autograd needs to save <code>a</code> . Likewise, autograd will save both <code>c</code> and <code>d</code> . Note that for your particular example, autograd does not need to save <code>_t1</code> nor <code>_t2</code> (nor <code>e</code> ). (It needs to know that <code>_t1</code> was obtainedfrom <code>a</code> and <code>b</code> using <code>*</code> , but that is a much smaller amount of information that is independent of the size of <code>_t1</code> .)</p>
</blockquote>
<p>I don’t really understand this. Using the same expression</p>
<pre><code class="lang-auto">_t1 = a * b
_t2 = c * d
e = _t1 + _t2
</code></pre>
<p>The derivative of e with respect to a which is needed for gradient descent for eg, is a function of both the local derivative (b) and the derivative of e with respect to _t1. In that case, doesn’t autograd need to store the intermediate value _t1 as well?</p>
<p>Thank you!</p>

srishti-git1110: <p>Hey,</p>
<blockquote>
<p>In that case, doesn’t autograd need to store the intermediate value _t1 as well?</p>
</blockquote>
<p>Let’s emphasize on the particular word “need”: In the example that you’ve shared, there’s no need to store either of <code>_t1</code> or <code>_t2</code>.</p>
<p>And why is that? Because the derivative of <code>e</code> wrt any (leaf) tensor <code>a</code>, <code>b</code>, <code>c</code>, or <code>d</code> does not require <code>_t1</code> or <code>_t2</code> as you can easily see:</p>
<pre><code class="lang-auto">de/da = b
</code></pre>
<p>and so on.</p>
<p>(Excuse my use of bad differentiation notation.)</p>
<p>To reinforce your understanding, let’s look at another example where</p>
<pre><code class="lang-auto">e = 5 _t1_t2 + 6 _t2
</code></pre>
<p>Now, here there’s a need to store <code>_t1</code> or <code>_t2</code> because:</p>
<pre><code class="lang-auto">de/da = de/d_t1 x d_t1/da = 5_t2 . b
</code></pre>
<p>which requires <code>_t2</code> to be stored in the memory.</p>
<p>If it was clear till here, then think about this: What if we do not store <code>_t2</code> in the memory?<br>
Mathematically speaking, it should be just fine because <code>_t2 = cd</code> and as long as <code>c</code> and <code>d</code> are stored in the memory, you can always sucessfully “re-compute” <code>_t2</code> and hence use it in the formula above.</p>
<p>What this means is that it’s a tradeoff between speed and memory: having <code>_t2</code> stored in the memory would allow for a faster backprop as compared to re-calculating it during the backward pass.<br>
Please also see <a href="https://pytorch.org/docs/stable/checkpoint.html" rel="noopener nofollow ugc">activation checkpointing</a>.</p>
<p>Hope it helps,<br>
S</p>