Gradients for torch.nn.functional.softplus() are numerically unstable

KFrank: <p>Hello Forum!</p>
<p>It appears that <code>backward()</code> for <code>softplus()</code> has a numerical issue<br>
(that does not appear to be caused directly by a related <code>forward()</code><br>
issue) for large negative arguments.</p>
<p>A straightforward (but not completely naive) implementation of<br>
<code>softplus()</code> – that has no explicit <code>backward()</code> –  seems to work<br>
fine with autograd.</p>
<p>Here is an illustrative test script:</p>
<pre data-code-wrap="python"><code class="lang-python">import torch
print (torch.__version__)

def test_softplus (x):   # use log-sum-exp trick and loq1p
    return torch.where (x &gt; 0, x + torch.exp (-x).log1p(), torch.exp (x).log1p())

def log1p_softplus (x):  # log1p version works well enough
    return torch.exp (x).log1p()

xpt = torch.arange (-500, 101, dtype = torch.float) / 10

relerr = ((torch.nn.functional.softplus (xpt) - test_softplus (xpt)) / test_softplus (xpt)).abs().max()

# forward looks okay
print ('forward looks okay:                   relerr =', relerr)

# look at gradients

sig = torch.sigmoid (xpt)   # sigmoid is the derivative of softplus

xtt = torch.arange (-500, 101, dtype = torch.float) / 10

xpt.requires_grad = True
torch.nn.functional.softplus (xpt).sum().backward()
xtt.requires_grad = True
test_softplus (xtt).sum().backward()

relerrpt = ((xpt.grad - sig) / sig).abs().max()
relerrtt = ((xtt.grad - sig) / sig).abs().max()

# built-in softplus-backward looks numerically unstable
print ('built-in backward looks unstable:     relerrpt =', relerrpt)
print ('backward for hand-rolled looks okay:  relerrtt =', relerrtt)

print ('sig[:5] =', sig[:5])
print ('xpt.grad[:5] =', xpt.grad[:5])
print ('xtt.grad[:5] =', xtt.grad[:5])
</code></pre>
<p>And here is its output:</p>
<pre><code class="lang-plaintext">1.7.1
forward looks okay:                   relerr = tensor(1.0752e-07)
built-in backward looks unstable:     relerrpt = tensor(1.)
backward for hand-rolled looks okay:  relerrtt = tensor(1.8833e-07)
sig[:5] = tensor([1.9287e-22, 2.1316e-22, 2.3558e-22, 2.6035e-22, 2.8774e-22])
xpt.grad[:5] = tensor([0., 0., 0., 0., 0.])
xtt.grad[:5] = tensor([1.9287e-22, 2.1316e-22, 2.3558e-22, 2.6035e-22, 2.8774e-22])
</code></pre>
<p>A possibly-related issue is discussed in this thread:</p>
<aside class="quote" data-post="1" data-topic="119909">
  <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/stevethesteve/48/12812_2.png" class="avatar">
    <a href="https://discuss.pytorch.org/t/computationally-stable-log-softplus/119909">Computationally stable log-softplus</a> 
  </div>
  <blockquote>
    At some point during computation, my model needs to compute the logarithm of softplus of a parameter. 
Currently, I implement this via: 
torch.nn.functional.softplus(theta).log()

Due to the log() call I fear that there might be issues with computational stability. 
Is there a computationally more stable way of computing log-softplus? 
I do not want to change my parameterization. 
Thank you!
  </blockquote>
</aside>

<p>Best.</p>
<p>K. Frank</p>

tom: <p>Great analysis <a class="mention" href="/u/kfrank">@KFrank</a>!</p>
<p><a class="mention" href="/u/ptrblck">@ptrblck</a>  and I were chatting about this and we found that it appears to be solved in the nightlies (<a class="mention" href="/u/ptrblck">@ptrblck</a> found the probable cause: <a href="https://github.com/pytorch/pytorch/pull/56484" class="inline-onebox" rel="noopener nofollow ugc">Fix: Compare input against beta * threshold in softplus backwards by jbschlosser · Pull Request #56484 · pytorch/pytorch · GitHub</a> fixing <a href="https://github.com/pytorch/pytorch/issues/55587" class="inline-onebox" rel="noopener nofollow ugc">`Softplus` forward and backward discrepancy · Issue #55587 · pytorch/pytorch · GitHub</a>).</p>
<p>Best regards</p>
<p>Thomas</p>

KFrank: <p>Hi Thomas!</p>
<aside class="quote no-group" data-username="tom" data-post="2" data-topic="120169" 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/tom/48/3162_2.png" class="avatar"> tom:</div>
<blockquote>
<p>it appears to be solved in the nightlies</p>
</blockquote>
</aside>
<p>I can confirm that this is fixed in today’s nightly, 1.9.0.dev20210504<br>
(but not yet in my April nightly, 1.9.0.dev20210416):</p>
<pre><code class="lang-plaintext">1.9.0.dev20210504
forward looks okay:                   relerr = tensor(1.0752e-07)
built-in backward looks unstable:     relerrpt = tensor(1.5066e-07)
backward for hand-rolled looks okay:  relerrtt = tensor(1.8833e-07)
sig[:5] = tensor([1.9287e-22, 2.1316e-22, 2.3558e-22, 2.6035e-22, 2.8774e-22])
xpt.grad[:5] = tensor([1.9287e-22, 2.1316e-22, 2.3558e-22, 2.6035e-22, 2.8774e-22])
xtt.grad[:5] = tensor([1.9287e-22, 2.1316e-22, 2.3558e-22, 2.6035e-22, 2.8774e-22])
</code></pre>
<p>(The “looks unstable” is just hard-wired text; the result is correct.)</p>
<p>Thanks.</p>
<p>K. Frank</p>