Differing gradients with same algorithm

username: legoh
cooked: <p>Hi all, I need a little bit of help figuring out where my implementation went wrong. I have an algorithm that runs on a 3D tensor, given by <code>[batch, row, col]</code>. I used to use a <code>for</code> loop to iterate through each <code>n</code>, where <code>n</code> is the number of rows/cols (n x n square matrix here). I figured out how to use the <code>torch.scatter_add_</code> function to remove the for-loop.</p>
<p>Here’s the input to the function:</p>
<pre><code class="lang-auto">input = torch.Tensor([[[7,  8, 7, 20, 7],
                                   [15,  2,  2,  4,  8],
                                   [ 6, 17,  6,  7,  8],
                                   [ 1,  6, 12,  6,  2],
                                   [ 8,  1,  7,  2, 10]]]
</code></pre>
<p>Here’s the for-loop implementation:</p>
<pre><code class="lang-auto">def iter_simp_min_sum_batch(m_alpha_beta, m_beta_alpha, weights, n_iter):

    for _ in range(n_iter):

        # Message passing
        for i in range(n):
            alpha_beta_max = torch.max(torch.cat((m_alpha_beta[:, :i, :], m_alpha_beta[:, (i + 1):, :]), dim=1), dim=1)[0]
            beta_alpha_max = torch.max(torch.cat((m_beta_alpha[:, :i, :], m_beta_alpha[:, (i + 1):, :]), dim=1), dim=1)[0]
            if i == 0:
                m_beta_alpha_k = (weights[:, i, :] - alpha_beta_max).unsqueeze(2)
                m_alpha_beta_k = (weights[:, :, i] - beta_alpha_max).unsqueeze(2)
            else:
                m_beta_alpha_k = torch.cat((m_beta_alpha_k, (weights[:, i, :] - alpha_beta_max).unsqueeze(2)), dim=2)
                m_alpha_beta_k = torch.cat((m_alpha_beta_k, (weights[:, :, i] - beta_alpha_max).unsqueeze(2)), dim=2)

        m_alpha_beta = m_alpha_beta_k
        m_beta_alpha = m_beta_alpha_k

    return m_alpha_beta, m_beta_alpha

</code></pre>
<p>Here’s my scatter implementation of the above:</p>
<pre><code class="lang-auto">def iter_simp_min_sum_batch_scatter(m_alpha_beta, m_beta_alpha, weights, n_iter):

    for _ in range(n_iter):
        # Message passing
    

        beta_alpha_maxes, beta_alpha_indices = torch.topk(m_beta_alpha, 2, dim=1)
        print(beta_alpha_maxes[:, 0, :].requires_grad)
        m_alpha_beta_k = weights.permute(0, 2, 1) - beta_alpha_maxes[:, 0, :].unsqueeze(1)
        m_alpha_beta_k = m_alpha_beta_k.scatter_add_(dim=1, index=beta_alpha_indices[:, 0, :].unsqueeze(1), src=beta_alpha_maxes[:, 0, :].unsqueeze(1))
        m_alpha_beta_k = m_alpha_beta_k.scatter_add_(dim=1, index=beta_alpha_indices[:, 0, :].unsqueeze(1), src=beta_alpha_maxes[:, 1, :].unsqueeze(1) * -1).permute(0, 2, 1)
        
        alpha_beta_maxes, alpha_beta_indices = torch.topk(m_alpha_beta, 2, dim=1)
        m_beta_alpha_k = weights - alpha_beta_maxes[:, 0, :].unsqueeze(1)
        m_beta_alpha_k = m_beta_alpha_k.scatter_add_(dim=1, index=alpha_beta_indices[:, 0, :].unsqueeze(1), src=alpha_beta_maxes[:, 0, :].unsqueeze(1))
        m_beta_alpha_k = m_beta_alpha_k.scatter_add_(dim=1, index=alpha_beta_indices[:, 0, :].unsqueeze(1), src=alpha_beta_maxes[:, 1, :].unsqueeze(1) * -1).permute(0, 2, 1)
        
        m_alpha_beta = m_alpha_beta_k
        m_beta_alpha = m_beta_alpha_k

    return m_alpha_beta, m_beta_alpha
</code></pre>
<p>Finally, here’s how I test the above implementations:</p>
<pre><code class="lang-auto"># Scatter implementation
a = torch.nn.Parameter(input)

scatter_b_a, scatter_b_b = iter_simp_min_sum_batch_scatter(a, a.permute(0, 2, 1), a.clone().detach(), 3)

scatter_b_b.backward(gradient=torch.ones((scatter_b_b.size(0), scatter_b_b.size(1), scatter_b_b.size(2))))

# For-loop implementation
b = torch.nn.Parameter(input)

iter_b_a, iter_b_b = iter_simp_min_sum_batch(b, b.permute(0, 2, 1), b.clone().detach(), 3)

iter_b_b.backward(gradient=torch.ones((iter_b_b.size(0), iter_b_b.size(1), iter_b_b.size(2))))

</code></pre>
<p>Here’s the results:</p>
<pre><code class="lang-auto">a.grad

tensor([[[  0.,   0.,  -2.,   0.,   0.],
         [  0.,   0.,   0.,   0.,   0.],
         [  0.,   0.,   0.,   0.,  -2.],
         [  0.,   0.,  -4.,   0.,   0.],
         [ -1.,   0.,   0.,   0., -16.]]])


b.grad

tensor([[[  0.,   0.,   0.,   0.,   0.],
         [  0.,   0.,   0.,   0.,   0.],
         [  0.,   0.,   0.,   0.,  -2.],
         [  0.,   0.,  -4.,   0.,   0.],
         [ -1.,   0.,  -2.,   0., -16.]]])

scatter_b_b - iter_b_b

tensor([[[0., 0., 0., 0., 0.],
         [0., 0., 0., 0., 0.],
         [0., 0., 0., 0., 0.],
         [0., 0., 0., 0., 0.],
         [0., 0., 0., 0., 0.]]], grad_fn=&lt;SubBackward0&gt;)

</code></pre>
<p>As shown in <code>scatter_b_b - iter_b_b</code>, the solutions are the same. The difference is in the gradients which I don’t really understand where the difference is.</p>
<p>Thanks a lot for any help!</p>

username: albanD
cooked: <p>Hi,</p>
<p>What is the content of <code>m_alpha_beta</code> and <code>m_beta_alpha</code>?<br>
One case that could lead to this is if they contain entries that are equal within them</p>
<p>Indeed, the topk and the indices selected can be any of the elements when there is a tie.</p>
<p>Similarly, in your for-loop version, the max will select one element “at random” if there are multiple elements equal to the max.</p>
<p>As you saw, this does not change the output result at all as these values are the same.<br>
But since these values come from different entries, the gradients (that flow back to these selected entries only) can be different!</p>
<p>Hope that helps!</p>

username: legoh
cooked: <p>Hey! <code>m_alpha_beta</code> contains <code>input</code>, which is basically just a <code>b x n x n</code> parameter tensor, where <code>b</code> is the batch size. <code>m_beta_alpha</code> contains <code>input.permute(0, 2, 1)</code>, which is basically the transpose of <code>m_alpha_beta</code>.</p>
<p>I tested it with the input above, which is</p>
<pre><code class="lang-auto">input = torch.Tensor([[[7,  8, 7, 20, 7],
                       [15,  2,  2,  4,  8],
                       [ 6, 17,  6,  7,  8],
                       [ 1,  6, 12,  6,  2],
                       [ 8,  1,  7,  2, 10]]]
</code></pre>
<p>I have also played around with the number of iterations and it always seems like the gradient is slightly off in terms of the positions.</p>

username: albanD
cooked: <p>As you can see in your example, your gradients don’t match in the third column which is exactly the oclumn that has duplicatevalues. With one version putting grad on the top <code>7</code> and the other one on the bottom <code>7</code> so looks like this is what happens (and is expected behavior!)</p>

username: legoh
cooked: <p>Oh, right! I just tested with a new tensor that has unique values in all rows and columns, and the gradients are the same! Thanks a lot. It appears as the <code>max</code> and <code>topk</code> functions were returning different indices because of the same values, just as you suspected. Here are some results for reference:</p>
<pre><code class="lang-auto">tens = torch.Tensor([[[ 2.,  8., 10., 20.,  7.],
                      [15.,  3.,  2.,  4.,  9.],
                      [10., 17.,  6.,  7.,  8.],
                      [ 1.,  7., 12.,  6.,  2.],
                      [ 8.,  1.,  7.,  2., 10.]]])

a.grad
tensor([[[ 0., -1., -1., -4.,  0.],
         [-4.,  0.,  0.,  0., -1.],
         [-1., -4.,  0., -1.,  0.],
         [ 0.,  0., -4.,  0.,  0.],
         [ 0.,  0.,  0.,  0., -4.]]])

b.grad
tensor([[[ 0., -1., -1., -4.,  0.],
         [-4.,  0.,  0.,  0., -1.],
         [-1., -4.,  0., -1.,  0.],
         [ 0.,  0., -4.,  0.,  0.],
         [ 0.,  0.,  0.,  0., -4.]]])

</code></pre>