Torch.logsumexp returning nan gradients when inputs are -inf

username: shivammehta007
cooked: <p>Hello,</p>
<p>When I try to backpropagate on a tensor full of -inf and I have a torch.logsumexp , the gradients of that becomes nan. Like the code below:</p>
<pre><code class="lang-auto">&gt;&gt;&gt; a = torch.nn.Parameter(torch.tensor([-float("inf"), -float("inf"), -float("inf")]))
&gt;&gt;&gt; b = 2 + a
&gt;&gt;&gt; torch.logsumexp(a, dim=0)
tensor(-inf, grad_fn=&lt;LogsumexpBackward&gt;)
&gt;&gt;&gt; torch.logsumexp(a, dim=0).backward()
&gt;&gt;&gt; a.grad
tensor([nan, nan, nan])
</code></pre>
<p>But even if one of the elements in the tensor is non -inf the gradients are propagated properly:</p>
<pre><code class="lang-auto">&gt;&gt;&gt; a = torch.nn.Parameter(torch.tensor([1.0, -float("inf"), -float("inf")]))
&gt;&gt;&gt; b = 2 + a
&gt;&gt;&gt; torch.logsumexp(a, dim=0)
tensor(1., grad_fn=&lt;LogsumexpBackward&gt;)
&gt;&gt;&gt; torch.logsumexp(a, dim=0).backward()
&gt;&gt;&gt; a.grad
tensor([1., 0., 0.])

</code></pre>
<p>How do I ensure that even in case of all -inf my gradients should be zero?</p>

username: albanD
cooked: <p>The problem is that at the point where the final result is <code>-inf</code>, the gradient is infinite. So during backprop, the gradient becomes nan.<br>
In your second example, the gradient at point 1. is finite and everything works fine.</p>
<p>So I would say this is expected behavior no?</p>

username: shivammehta007
cooked: <p>But in a larger graphical models like CTC or HMMs, this is introducing nans in my backpropagation, What could be the best possible way to avoid this behavior; as I want my model to be training on the non-infinity parameters in the different locations what I mean is that consider a matrix;</p>
<pre><code class="lang-auto">x = torch.tensor([[ 1.0, -float("inf")], [-float("inf"), -float("inf")]])
&gt;&gt;&gt; x.shape
torch.Size([2, 2])
&gt;&gt;&gt; x
tensor([[1., -inf],
        [-inf, -inf]])
&gt;&gt;&gt; x = nn.Parameter(x)
&gt;&gt;&gt; y = torch.logsumexp(x, dim=1)
&gt;&gt;&gt; y
tensor([1., -inf], grad_fn=&lt;LogsumexpBackward&gt;)
&gt;&gt;&gt; z = torch.sum(y, dim=0)
&gt;&gt;&gt; z.backward()
&gt;&gt;&gt; x.grad
tensor([[1., 0.],
        [nan, nan]])
</code></pre>
<p>Take this example for say; how can I make changes that instead of nan I get zero gradients at x.grad?</p>

username: albanD
cooked: <p>I would argue that you need to make sure not to get a whole row of infinites. Because at this point, your logsumexp function is not properly differentiable and so it’s gradients will be nan.</p>

username: shivammehta007
cooked: <p>But incase I don’t want to do gradient update of those places where there are nan’s. What could be the possible work around this?</p>

username: albanD
cooked: <p>It depends on your optimizer.<br>
If you don’t have momentum/accumulated terms, then you can simply set these gradients to 0 and your optimizer won’t change the values.<br>
If you have a fancy optimizer that will update the weights even for a 0 gradient, the simplest solution might be to save the original value of the weights before performing the step and then restoring them after the optimizer step.</p>

username: shivammehta007
cooked: <p>I think my use case is somewhat similar to that of this issue on Github: <a href="https://github.com/pytorch/pytorch/issues/31829" rel="noopener nofollow ugc">https://github.com/pytorch/pytorch/issues/31829</a></p>
<p>Do we have any workaround for this?</p>

username: albanD
cooked: <p>It think what I mentioned in my previous message is the right thing.</p>