Computationally stable log-softplus

username: stevethesteve
cooked: <p>At some point during computation, my model needs to compute the logarithm of softplus of a parameter.<br>
Currently, I implement this via:</p>
<pre><code class="lang-auto">torch.nn.functional.softplus(theta).log()
</code></pre>
<p>Due to the log() call I fear that there might be issues with computational stability.<br>
Is there a computationally more stable way of computing log-softplus?<br>
I do not want to change my parameterization.<br>
Thank you!</p>

username: googlebot
cooked: <p>I don’t think this can be simplified. As for “stability”, you’re limited by float32 representation bounds, it is not like you’ll get biased errors or something (except inherent log() imprecision).</p>

username: stevethesteve
cooked: <p>quick update:<br>
for now, I am using the approximation as in this pseudo-code:</p>
<pre><code class="lang-auto">def log_softplus(x):
  if x &gt; -5:
    return torch.nn.functional.softplus(x).log()
  else:
    return x
</code></pre>
<p>which is numerically stable for small x. If done naively (softplus then log), softplus would return a very small number with low precision.</p>

username: googlebot
cooked: <p>Ah, sorry, you’re right, I somehow assumed that lim(log(x))=0, and problem in that area. You may just want to vectorize your approach with torch.where.</p>

username: KFrank
cooked: <p>Hi Steve!</p>
<aside class="quote no-group" data-username="stevethesteve" data-post="3" data-topic="119909" 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/stevethesteve/48/12812_2.png" class="avatar"> stevethesteve:</div>
<blockquote>
<pre><code class="lang-auto">def log_softplus(x):
  if x &gt; -5:
    return torch.nn.functional.softplus(x).log()
  else:
    return x
</code></pre>
<p>which is numerically stable for small x. If done naively (softplus then log), softplus would return a very small number with low precision.</p>
</blockquote>
</aside>
<p>I think your concern about loss of precision is a partial red herring.</p>
<p>Because we’re working with floating-point numbers, we don’t lose<br>
precision when <code>softplus (x)</code> becomes small (say, on the order of<br>
10**-7).  We only start to lose precision when <code>softplus (x)</code> becomes<br>
<em>very</em> small and starts to denormalize and then underflows to zero<br>
(at which point the subsequent <code>log()</code> will return <code>-inf</code>). This starts to<br>
happen around 10**-38 (which corresponds to an <code>x</code> of about -90.0).</p>
<p>Whether this matters depends on your specific use case.  But, unless<br>
your argument to log-softplus can reasonably become as negative as,<br>
say, -90, your “naive” implementation will be fine.</p>
<p>On the other hand, if your argument to log-softplus can reasonably<br>
become that negative, then your suggested version does make<br>
sense, except that you should use a 'breakpoint" (<code>x &gt; -5</code>) that is<br>
significantly more negative than <code>-5</code> (perhaps something like -40).<br>
As it stands, your breakpoint of <code>-5</code> does more harm than good,<br>
because in this range <code>x</code> is an imperfect approximation to log-softplus.</p>
<p>The following illustrative script uses a double-precision log-softplus<br>
computation as a surrogate for the exact result:</p>
<pre data-code-wrap="python"><code class="lang-python">import torch
print (torch.__version__)

def log_softplusA (x):
    return torch.nn.functional.softplus (x).log()

def log_softplusB (x, brk = -5.0):
    return torch.where (x &gt; brk, torch.nn.functional.softplus (x).log(), x)

x = torch.arange (-100, 101, 1).double() / 10.0

lsp = log_softplusA (x)   # double-precision "ground truth"

diffA = log_softplusA (x.float()) - lsp
diffB = log_softplusB (x.float()) - lsp

print ('diffA.abs().max() =', diffA.abs().max())
print ('diffB.abs().max() =', diffB.abs().max())
print ('diffB.abs().argmax() =', diffB.abs().argmax())
print ('x[diffB.abs().argmax()] =', x[diffB.abs().argmax()])

print ('log_softplusA (torch.tensor ([-90.0])) =', log_softplusA (torch.tensor ([-90.0])))
print ('log_softplusA (torch.tensor ([-100.0])) =', log_softplusA (torch.tensor ([-100.0])))
print ('log_softplusA (torch.tensor ([-110.0])) =', log_softplusA (torch.tensor ([-110.0])))
print ('log_softplusA (torch.tensor ([-110.0], dtype = torch.double)) =', log_softplusA (torch.tensor ([-110.0], dtype = torch.double)))
</code></pre>
<p>And here is its result:</p>
<pre><code class="lang-plaintext">1.7.1
diffA.abs().max() = tensor(8.1795e-07, dtype=torch.float64)
diffB.abs().max() = tensor(0.0034, dtype=torch.float64)
diffB.abs().argmax() = tensor(50)
x[diffB.abs().argmax()] = tensor(-5., dtype=torch.float64)
log_softplusA (torch.tensor ([-90.0])) = tensor([-90.])
log_softplusA (torch.tensor ([-100.0])) = tensor([-99.9831])
log_softplusA (torch.tensor ([-110.0])) = tensor([-inf])
log_softplusA (torch.tensor ([-110.0], dtype = torch.double)) = tensor([-110.], dtype=torch.float64)
</code></pre>
<p>Best.</p>
<p>K. Frank</p>

username: stevethesteve
cooked: <p>Well, thank you for your insightful comment!</p>
<p>I think you are right, for precise computation of log-softmax, a much lower breakpoint is sufficient.</p>
<p>What I did not mention (my bad) is that I want to apply gradient backprop through this computation.<br>
This changes things a lot: If setting the breakpoint too low, the error in the gradient can be 7 orders of magnitude bigger than the error in the log-softplus output.<br>
The following code is an extension of yours.</p>
<pre><code class="lang-auto">import torch
print (torch.__version__)

def log_softplusA (x):
    return torch.nn.functional.softplus (x).log()

def log_softplusB (x, brk = -15.0):
    return torch.where (x &gt; brk, torch.nn.functional.softplus (x).log(), x)

x = torch.arange (-200, 100, 1).double() / 10.0
x.requires_grad = True
lsp = log_softplusA (x)   # double-precision "ground truth"
lsp.sum().backward()
lspg = x.grad

xA = torch.arange (-200, 100, 1) / 10.0
xA.requires_grad = True
lspA = log_softplusA (xA)
diffA = lspA - lsp
lspA.sum().backward()
lspgA = xA.grad
diffgA = lspgA - lspg

xB = torch.arange (-200, 100, 1) / 10.0
xB.requires_grad = True
lspB = log_softplusB (xB)
diffB = lspB - lsp
lspB.sum().backward()

lspgB = xB.grad
diffgB = lspgB - lspg


print ('diffA.abs().max() =', diffA.abs().max())
print ('diffB.abs().max() =', diffB.abs().max())
print ('diffB.abs().argmax() =', diffB.abs().argmax())
print ('x[diffB.abs().argmax()] =', x[diffB.abs().argmax()])

print ('diffgA.abs().max() =', diffgA.abs().max())
print ('diffgB.abs().max() =', diffgB.abs().max())
print ('diffgB.abs().argmax() =', diffgB.abs().argmax())
print ('x[diffgB.abs().argmax()] =', x[diffgB.abs().argmax()])
print ('lspgB[diffgB.abs().argmax()], lspg[diffgB.abs().argmax()] =', lspgB[diffgB.abs().argmax()], lspg[diffgB.abs().argmax()])
</code></pre>
<p>yields</p>
<pre><code class="lang-auto">1.8.1
diffA.abs().max() = tensor(8.4096e-07, dtype=torch.float64, grad_fn=&lt;MaxBackward1&gt;)
diffB.abs().max() = tensor(8.4096e-07, dtype=torch.float64, grad_fn=&lt;MaxBackward1&gt;)
diffB.abs().argmax() = tensor(61)
x[diffB.abs().argmax()] = tensor(-13.9000, dtype=torch.float64, grad_fn=&lt;SelectBackward&gt;)
diffgA.abs().max() = tensor(1.0000, dtype=torch.float64)
diffgB.abs().max() = tensor(0.1339, dtype=torch.float64)
diffgB.abs().argmax() = tensor(53)
x[diffgB.abs().argmax()] = tensor(-14.7000, dtype=torch.float64, grad_fn=&lt;SelectBackward&gt;)
lspgB[diffgB.abs().argmax()], lspg[diffgB.abs().argmax()] = tensor(0.8661) tensor(1.0000, dtype=torch.float64)
</code></pre>
<p>Upon variating the breakpoint and re-running the code, I empirically found the value -8.0 to perform best.</p>

username: KFrank
cooked: <p>Hi Steve!</p>
<aside class="quote no-group" data-username="stevethesteve" data-post="6" data-topic="119909" 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/stevethesteve/48/12812_2.png" class="avatar"> stevethesteve:</div>
<blockquote>
<p>This changes things a lot: If setting the breakpoint too low, the error in the gradient can be 7 orders of magnitude bigger than the error in the log-softplus output.</p>
</blockquote>
</aside>
<p>This may be related to some weirdness in pytorch’s built-in<br>
<code>torch.nn.functional.softplus()</code>.  Here is the relevant thread:</p>
<aside class="quote quote-modified" data-post="1" data-topic="120169">
  <div class="title">
    <div class="quote-controls"></div>
    <img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/k/ecb155/48.png" class="avatar">
    <a href="https://discuss.pytorch.org/t/gradients-for-torch-nn-functional-softplus-are-numerically-unstable/120169">Gradients for torch.nn.functional.softplus() are numerically unstable</a> 
  </div>
  <blockquote>
    Hello Forum! 
It appears that backward() for softplus() has a numerical issue 
(that does not appear to be caused directly by a related forward() 
issue) for large negative arguments. 
A straightforward (but not completely naive) implementation of 
softplus() – that has no explicit backward() –  seems to work 
fine with autograd. 
Here is an illustrative test script: 
import torch
print (torch.__version__)

def test_softplus (x):   # use log-sum-exp trick and loq1p
    return torch.where (x &gt; 0,…
  </blockquote>
</aside>

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

username: KFrank
cooked: <p>Hi Steve!</p>
<p>I haven’t worked through this in detail, but a first look suggests that<br>
the <code>softplus()</code> fix in the recent nightlies addresses your issue.</p>
<aside class="quote no-group" data-username="stevethesteve" data-post="6" data-topic="119909" 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/stevethesteve/48/12812_2.png" class="avatar"> stevethesteve:</div>
<blockquote>
<p>yields</p>
<pre><code class="lang-auto">1.8.1
diffA.abs().max() = tensor(8.4096e-07, dtype=torch.float64, grad_fn=&lt;MaxBackward1&gt;)
diffB.abs().max() = tensor(8.4096e-07, dtype=torch.float64, grad_fn=&lt;MaxBackward1&gt;)
diffB.abs().argmax() = tensor(61)
x[diffB.abs().argmax()] = tensor(-13.9000, dtype=torch.float64, grad_fn=&lt;SelectBackward&gt;)
diffgA.abs().max() = tensor(1.0000, dtype=torch.float64)
diffgB.abs().max() = tensor(0.1339, dtype=torch.float64)
diffgB.abs().argmax() = tensor(53)
x[diffgB.abs().argmax()] = tensor(-14.7000, dtype=torch.float64, grad_fn=&lt;SelectBackward&gt;)
lspgB[diffgB.abs().argmax()], lspg[diffgB.abs().argmax()] = tensor(0.8661) tensor(1.0000, dtype=torch.float64)
</code></pre>
</blockquote>
</aside>
<p>Running the script you posted on today’s nightly, 1.9.0.dev20210504,<br>
seems to show that the issue is gone:</p>
<pre><code class="lang-plaintext">1.9.0.dev20210504
diffA.abs().max() = tensor(8.4096e-07, dtype=torch.float64, grad_fn=&lt;MaxBackward1&gt;)
diffB.abs().max() = tensor(8.4096e-07, dtype=torch.float64, grad_fn=&lt;MaxBackward1&gt;)
diffB.abs().argmax() = tensor(61)
x[diffB.abs().argmax()] = tensor(-13.9000, dtype=torch.float64, grad_fn=&lt;SelectBackward&gt;)
diffgA.abs().max() = tensor(1.6522e-07, dtype=torch.float64)
diffgB.abs().max() = tensor(1.6522e-07, dtype=torch.float64)
diffgB.abs().argmax() = tensor(162)
x[diffgB.abs().argmax()] = tensor(-3.8000, dtype=torch.float64, grad_fn=&lt;SelectBackward&gt;)
lspgB[diffgB.abs().argmax()], lspg[diffgB.abs().argmax()] = tensor(0.9890) tensor(0.9890, dtype=torch.float64)
</code></pre>
<p>Best.</p>
<p>K. Frank</p>