Unexpected gradient of sum - doesn't match sum of gradients

ruben: <p>Alright so this looks like such a simple issue that I’m afraid I’ve misunderstood some of the core workings of PyTorch/autograd.</p>
<p>My problem is straightforward: I have a gradient of a sum that turns out 0 while it shouldn’t be, and it doesn’t match the sum of the gradients separately.</p>
<p>I’ve concocted a simplified example to illustrate. We’re in the probability domain, or rather, the logprob domain, so this is the straightforward approximation of probabilities using a softmax on top of a linear layer. Well, not quite a softmax, because logprobs. There are two output classes.</p>
<pre><code class="lang-auto">import torch
import torch.nn as nn
import torch.nn.functional as F

feat_dim, out_dim = 1, 2
encoder = nn.Linear(feat_dim, out_dim, bias=False)
nn.init.zeros_(encoder.weight.data)
# Fabricate some one-hot encoded data of batch size 1, good enough
batch_size = 1
input = torch.eye(feat_dim)[:batch_size]
</code></pre>
<p>As I said, I have the sum of gradients, or rather a difference really. Below I’ll show the gradients of the linear layer’s weights after backpropagating first the difference, then the summands themselves separately.</p>
<pre><code class="lang-auto">logits = encoder.forward(input)
norm = torch.logsumexp(logits, dim=1, keepdim=True)
(logits-norm).sum().backward()
print(encoder.weight.grad)
nn.init.zeros_(encoder.weight.grad.data)

logits = encoder.forward(input)
logits.sum().backward()
print(encoder.weight.grad)
nn.init.zeros_(encoder.weight.grad.data)

logits = encoder.forward(input)
norm = torch.logsumexp(logits, dim=1, keepdim=True)
norm.sum().backward()
print(encoder.weight.grad)
</code></pre>
<p>This gives (annotated):</p>
<pre><code class="lang-auto"># grad of logits - norm
tensor([[0.],
        [0.]])
# grad of logits
tensor([[1.],
        [1.]])
# grad of norm
tensor([[0.5000],
        [0.5000]])
</code></pre>
<p>The gradients for <code>logits</code> and <code>norm</code> are exactly what I expected. I expected the gradient for <code>logprobs</code> to be the difference of the gradients of <code>logits</code> and <code>norm</code>, but that’s not the case. Instead it’s zeros, for no conceivable reason.</p>
<p>I must be missing something incredibly simple. Maybe the gradient of <code>norm</code> gets backpropagated multiple times, but I don’t see it. Could anyone shed some light on the issue?</p>

Yaroslav_Bulatov: <p>You can you look at your situation as finding derivative of f(g(h(w))), and the values you get are h’ = 1, g’=1/2 and f’=0 . Since this is composition, f’!=h’+g’</p>
<p>To see why specifically you get zero you can decorate your graph to print backward values</p>
<pre><code class="lang-auto">import torch
import torch.nn as nn
import torch.nn.functional as F

feat_dim, out_dim = 1, 2
encoder = nn.Linear(feat_dim, out_dim, bias=False)
nn.init.zeros_(encoder.weight.data)
batch_size = 1
input = torch.eye(feat_dim)[:batch_size]

def decorate(tensor, name):
    def print_backward(grad):
        print(f"({name}): ", grad)
    tensor.register_hook(print_backward)
    
logits = encoder.forward(input)
norm = torch.logsumexp(logits, dim=1, keepdim=True)
decorate(logits, 'logits')
decorate(logits, 'norm')
norm.expand_as(logits).backward(torch.ones_like(logits))
</code></pre>
<p>This prints following</p>
<pre><code class="lang-auto">(logits):  tensor([[1., 1.]])
(norm):  tensor([[1., 1.]])
</code></pre>
<p>Since both logits and norm vectors are 1’s, their difference will cancel out. The following is an illustration of how 1’s come about</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/2X/9/9a188cb5286a9c956bfd35255706c4424ee687c9.png" data-download-href="https://discuss.pytorch.org/uploads/default/9a188cb5286a9c956bfd35255706c4424ee687c9" title="41"><img src="https://discuss.pytorch.org/uploads/default/optimized/2X/9/9a188cb5286a9c956bfd35255706c4424ee687c9_2_690x405.png" alt="41" data-base62-sha1="lZc7VinReU3b36hGmzSLPLBQBh7" width="690" height="405" srcset="https://discuss.pytorch.org/uploads/default/optimized/2X/9/9a188cb5286a9c956bfd35255706c4424ee687c9_2_690x405.png, https://discuss.pytorch.org/uploads/default/optimized/2X/9/9a188cb5286a9c956bfd35255706c4424ee687c9_2_1035x607.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/2X/9/9a188cb5286a9c956bfd35255706c4424ee687c9_2_1380x810.png 2x" data-dominant-color="FCFCFC"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">41</span><span class="informations">1416×832 18.9 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>

ruben: <p>That is a beautifully crafted reply <a class="mention" href="/u/yaroslav_bulatov">@Yaroslav_Bulatov</a>. Many thanks for the informative graphic! And the illustration of how to use <code>register_hook</code>, that should come in useful in hte future.</p>
<p>It shows what really happens: PyTorch has summed up the partial derivatives over the different outputs. Which is why you’ve got 0.5 in the graphic, but only a 1D vector of derivatives in the backward hook.</p>
<p>I would still argue my original example of <code>logits - norm</code> is a sum however.  And if I were to look at the partial derivatives individually, that sum would match. The problem was that I was looking at the sum of the derivatives (summed over outputs), which is what PyTorch gives us.</p>
<p>Anyway, cheers!</p>