EMD gradient regarding softmax input from autograd different from theoretical value

XZLeo: <p>I’m playing with simplified Wasserstein distance (also known as earth mover distance) as the loss function for N classification task.  Since the gnd is a one-hot distribution, the loss is the weighted sum of the absolute value of each class id minus the gnd class id.<br>
<img src="https://discuss.pytorch.org/uploads/default/original/3X/6/2/624ea9c0674afe64f13eae8ed4d9232acaca5b1f.png" alt="image" data-base62-sha1="e1FkgZCH0Uib9DhBKEZyN6zlPlR" width="95" height="62"></p>
<p>p_i is the softmax output.<br>
It is defined as follows:</p>
<pre><code class="lang-auto">class WassersteinClass(nn.Module):
    def __init__(self) -&gt; None:
        super().__init__()

    def forward(self, likelihood, gnd_idx):
        batch_size = 1
        l = likelihood.shape[1]  # number of bins  100

        gnd_idx = gnd_idx.reshape((batch_size, 1)) 
        idxs = torch.arange(0, l, dtype=torch.float32).to(device=config.device, non_blocking=True)
        batch_idxs = idxs.repeat(batch_size, 1)
        D = torch.abs(batch_idxs-gnd_idx) # broadcast
        
        loss = torch.sum(likelihood*D)
        return loss
</code></pre>
<p>Now I noticed that the gradient vanishes so I want to check its gradient regarding softmax input by creating a toy example. In this example, gnd (4) is one-hot and likelihood is nearly one-hot (peak at 2).</p>
<pre><code class="lang-auto">criterion = WassersteinClass()
# check gradient of softmax input
gnd_idx = torch.full([1], 4)
likelihood = torch.ones(1, 5)*1e-3
likelihood[:, 2] = 100
likelihood.requires_grad_()

softmax = nn.Softmax(1)
prob = softmax(likelihood)
prob.requires_grad_()
print('sfotmax output:', prob)
prob.retain_grad()
&gt;&gt;&gt;sfotmax output: tensor([[3.7835e-44, 1.0089e-43, 1.0000e+00, 3.7835e-44, 3.7835e-44]],
       grad_fn=&lt;SoftmaxBackward0&gt;)

wass = criterion(prob, gnd_idx)
print('loss:', was)
&gt;&gt;&gt; loss: tensor(2., grad_fn=&lt;SumBackward0&gt;)
</code></pre>
<p>Its gradient regarding softmax output is quite neat.</p>
<pre><code class="lang-auto">b = prob.grad
print('loss gradient after softmax:', b)
&gt;&gt;&gt;loss gradient after softmax: tensor([[4., 3., 2., 1., 0.]])
</code></pre>
<p>The gradient of softmax input from autograd:</p>
<pre><code class="lang-auto">a = torch.autograd.grad(wass, [likelihood])
print('loss gradient before softmax:', a) # the gradient is the distance to gnd
&gt;&gt;&gt; loss gradient before softmax: (tensor([[ 7.5670e-44,  1.0089e-43,  0.0000e+00, -3.7835e-44, -7.5670e-44]]),)
</code></pre>
<p>According to the chain rule of backpropagation<br>
<img src="https://discuss.pytorch.org/uploads/default/original/3X/5/a/5a44ce7b1d7f83637d73d556860917ccec00a9fd.png" alt="image" data-base62-sha1="cSyoa7btSfik5qLrE1aNFeaOEhL" width="360" height="84"></p>
<p><img src="https://discuss.pytorch.org/uploads/default/original/3X/f/6/f6211c0b8e697715b07c2b97a922a4b7d6ffdde4.png" alt="image" data-base62-sha1="z7mosUi31p5OotzrEwVww0Nt2ny" width="254" height="97"></p>
<pre><code class="lang-auto"># gradient of loss regarding softmax input a[i]
for i in range(5):
    grad_softmax = -prob[:, i]*prob
    grad_softmax[:, i] = prob[:, i] * (1-prob[:, i])
    print(f'grad L/a{i}:{torch.sum(b*grad_softmax)}')
&gt;&gt;&gt;grad L/a0:7.567011707354012e-44
&gt;&gt;&gt;grad L/a1:1.0089348943138683e-43
&gt;&gt;&gt;grad L/a2:-4.918557609780108e-43 !!!!!!!!!!Different
&gt;&gt;&gt;grad L/a3:-3.783505853677006e-44
&gt;&gt;&gt;grad L/a4:-7.567011707354012e-44
</code></pre>
<p>All of the gradients are correct except L/a2, which is Wasserstein distance regarding the second input of softmax. Why does this happen?</p>