How to calculate log(1 - softmax(X)) numerically stably

Zwei: <p>Hi all,</p>
<p>I have a multiclass classification problem and there are some inter-class relationship. Then a modified version of Cross-Entropy Loss Function is used.</p>
<p><img src="https://discuss.pytorch.org/uploads/default/original/3X/c/4/c4a5056772a467add5598ef72f02d102c3bdcb90.png" alt="image" data-base62-sha1="s3B8HkCuZ9254u0bnfojukKHRfi" width="683" height="93"><br>
, where <strong>P</strong> is the probability and <strong>M</strong> is the label.</p>
<p>For short, in addtion to log_softmax(), I need to implement log(1 - softmax(X)), let’s call it <strong>log1m_softmax()</strong>. However, log1m_softmax() is <strong>numerically unstable</strong> even with LogSumExp trick.</p>
<p>I believe that I understand the implemention of numerically stable log_softmax() well, which has been well explained <a href="https://stackoverflow.com/questions/61567597/how-is-log-softmax-implemented-to-compute-its-value-and-gradient-with-better" rel="noopener nofollow ugc">here</a>.</p>
<p>A related disscusion can be found <a href="https://stats.stackexchange.com/questions/469706/log1-softmaxx" rel="noopener nofollow ugc">here</a>, but it leads to <strong>another numerical problem of log(1 - exp(x)) when x underflowes to zero</strong>.</p>
<p>Please help me to find a stable log1m_softmax(), thank you!</p>

ConvolutionalAtom: <p>Would this help?</p><aside class="onebox githubissue" data-onebox-src="https://github.com/pytorch/pytorch/issues/39242#issuecomment-1349202573">
  <header class="source">

      <a href="https://github.com/pytorch/pytorch/issues/39242#issuecomment-1349202573" target="_blank" rel="noopener nofollow ugc">github.com/pytorch/pytorch</a>
  </header>

  <article class="onebox-body">
    <div class="github-row">
  <div class="github-icon-container" title="Issue">
	  <svg width="60" height="60" class="github-icon" viewbox="0 0 14 16" aria-hidden="true"><path fill-rule="evenodd" d="M7 2.3c3.14 0 5.7 2.56 5.7 5.7s-2.56 5.7-5.7 5.7A5.71 5.71 0 0 1 1.3 8c0-3.14 2.56-5.7 5.7-5.7zM7 1C3.14 1 0 4.14 0 8s3.14 7 7 7 7-3.14 7-7-3.14-7-7-7zm1 3H6v5h2V4zm0 6H6v2h2v-2z"></path></svg>
  </div>

  <div class="github-info-container">
    <h4>
      <a href="https://github.com/pytorch/pytorch/issues/39242#issuecomment-1349202573" target="_blank" rel="noopener nofollow ugc">Add numerically stable log1mexp = log(1 - exp(-|x|)) function</a>
    </h4>

    <div class="github-info">
      <div class="date">
        opened <span class="discourse-local-date" data-format="ll" data-date="2020-05-29" data-time="13:30:02" data-timezone="UTC">01:30PM - 29 May 20 UTC</span>
      </div>


      <div class="user">
        <a href="https://github.com/simonepri" target="_blank" rel="noopener nofollow ugc">
          <img alt="simonepri" src="https://avatars.githubusercontent.com/u/3505087?v=4" class="onebox-avatar-inline" width="20" height="20">
          simonepri
        </a>
      </div>
    </div>

    <div class="labels">
        <span style="display:inline-block;margin-top:2px;background-color: #B8B8B8;padding: 2px;border-radius: 4px;color: #fff;margin-left: 3px;">
          module: numerical-stability
        </span>
        <span style="display:inline-block;margin-top:2px;background-color: #B8B8B8;padding: 2px;border-radius: 4px;color: #fff;margin-left: 3px;">
          triaged
        </span>
        <span style="display:inline-block;margin-top:2px;background-color: #B8B8B8;padding: 2px;border-radius: 4px;color: #fff;margin-left: 3px;">
          enhancement
        </span>
        <span style="display:inline-block;margin-top:2px;background-color: #B8B8B8;padding: 2px;border-radius: 4px;color: #fff;margin-left: 3px;">
          function request
        </span>
    </div>
  </div>
</div>

  <div class="github-row">
    <p class="github-body-container">## 🚀 Feature
Add the implementation for a numerically stable `log(1 - exp(-|x|)<span class="show-more-container"><a href="" rel="noopener" class="show-more">…</a></span><span class="excerpt hidden">)` function.

See https://cran.r-project.org/web/packages/Rmpfr/vignettes/log1mexp-note.pdf

Ref https://www.rdocumentation.org/packages/VGAM/versions/1.1-3/topics/log1mexp

## Motivation

The function can be implemented using other pytorch functions but it might be tricky to get it right.

Some implementations found on GitHub:

[@wouterkool/estimating-gradients-without-replacement/blob/9d8bf8b/bernoulli/gumbel.py#L7-L11](https://github.com/wouterkool/estimating-gradients-without-replacement/blob/9d8bf8b/bernoulli/gumbel.py#L7-L11)

[@visinf/n3net/blob/5d5883a/src_denoising/models/non_local.py#L94-L108](https://github.com/visinf/n3net/blob/5d5883a/src_denoising/models/non_local.py#L94-L108)

cc: @wouterkool @visinf</span></p>
  </div>

  </article>

  <div class="onebox-metadata">
    
    
  </div>

  <div style="clear: both"></div>
</aside>


Zwei: <aside class="quote no-group" data-username="ConvolutionalAtom" data-post="2" data-topic="169007">
<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/convolutionalatom/48/45668_2.png" class="avatar"> ConvolutionalAtom:</div>
<blockquote>
<p>Would this help?</p>
</blockquote>
</aside>
<p>Not exactly, because when x=0, the result of this log1mexp() is still -Inf.</p>
<p>This is real case when one value <strong>x<sub>i</sub></strong> in <strong>X</strong> dominates LogSumExp(<strong>X</strong>), where <em><strong>x<sub>i</sub></strong> - LogSumExp(<strong>X</strong>) = 0</em> with limited precision of float.</p>

ConvolutionalAtom: <p>Can you also not get away with using the original trick for log softmax, trick of subtracting max value?</p>
<p>In the original trick of logsoftmax, the denomitor was numerically stable while summing over all exponential values (with max value subtracted in the power)</p>
<p>The denominator for your case is the same, the numerator will be the same except for the one term which is absent. So, it seems that numerator is also stable ??</p>

Zwei: <p>Let’s consider the case where <strong>X</strong> = [-1000, -1000, 1000] and <strong>c</strong> = <strong>X</strong>.max() = 1000.</p>
<p>After subtracting the maxium, <strong>X - c</strong> = [-2000, -2000, 0] and exp(<strong>X - c</strong>) = [0.0, 0.0, 1].</p>
<p>Therefore, <strong>LSE</strong> = LogSumExp(<strong>X - c</strong>)  = 0, and <strong>LSE</strong> - (<strong>X - c</strong>) = [2000, 2000, 0].</p>
<p>This lead to log(<strong>LSE</strong> - (<strong>X - c</strong>)) = [7.6, 7.6, -Inf].</p>
<aside class="quote no-group" data-username="ConvolutionalAtom" data-post="4" data-topic="169007">
<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/convolutionalatom/48/45668_2.png" class="avatar"> ConvolutionalAtom:</div>
<blockquote>
<p>The denominator for your case is the same, the numerator will be the same except for the one term which is absent.</p>
</blockquote>
</aside>
<p>I think the denominator you mentioned is <strong>LSE</strong>, and after we subtract the dominant term from it, it gets to zero.</p>

KFrank: <p>Hi Zhengwei!</p>
<aside class="quote no-group" data-username="Zwei" data-post="1" data-topic="169007" 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/zwei/48/48403_2.png" class="avatar"> Zwei:</div>
<blockquote>
<p>I need to implement log(1 - softmax(X)), let’s call it <strong>log1m_softmax()</strong>. However, log1m_softmax() is <strong>numerically unstable</strong> even with LogSumExp trick.</p>
</blockquote>
</aside>
<p>As noted by Mert:</p>
<aside class="quote no-group" data-username="ConvolutionalAtom" data-post="4" data-topic="169007">
<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/convolutionalatom/48/45668_2.png" class="avatar"> ConvolutionalAtom:</div>
<blockquote>
<p>The denominator for your case is the same, the numerator will be the same except for the one term which is absent.</p>
</blockquote>
</aside>
<p>That is, <code>1 - softmax (X)_i</code> is</p>
<pre><code class="lang-auto">(sum_j (exp (X_j)) - exp (X_i)) / sum_j (exp (X_j) = (sum_(j != i) (exp (X_j)) / sum_j (exp (X_j)
</code></pre>
<p>This is also discussed in the stats.stackexchange thread you linked to.<br>
You can implement your numerically-stable <code>log1m_softmax()</code> “by hand”<br>
along the lines discussed in that thread.</p>
<p>You can also use pytorch’s <code>logsumexp()</code> to compute <code>log1m_softmax()</code><br>
without, in effect, reimplementing the log-sum-exp trick.</p>
<p>With a little manipulation, you can zero out the <code>i == j</code> term in probability<br>
space (i.e., in “exp” space) by replacing the term with <code>-inf</code> (or a very<br>
large negative number) in log space (i.e., the space of your original <code>X</code>)<br>
and then apply pytorch’s <code>logsumexp()</code> to both the numerator and<br>
denominator of the above expression for <code>1 - softmax (X)</code>.</p>
<p>Consider:</p>
<pre><code class="lang-plaintext">&gt;&gt;&gt; import torch
&gt;&gt;&gt; print (torch.__version__)
1.13.0
&gt;&gt;&gt;
&gt;&gt;&gt; _ = torch.manual_seed (2022)
&gt;&gt;&gt;
&gt;&gt;&gt; X = torch.tensor ([-1000, -1000, 1000.0])
&gt;&gt;&gt;
&gt;&gt;&gt; X.softmax (0)               # softmax underflows to zero
tensor([0., 0., 1.])
&gt;&gt;&gt; X.softmax (0).log()         # log gives -inf
tensor([-inf, -inf, 0.])
&gt;&gt;&gt; X.log_softmax (0)           # numerically-stable version uses log-sum-exp trick internally
tensor([-2000., -2000.,     0.])
&gt;&gt;&gt;
&gt;&gt;&gt; 1 - X.softmax (0)           # (1 - softmax) underflows to zero
tensor([1., 1., 0.])
&gt;&gt;&gt; (1 - X.softmax (0)).log()   # log gives -inf
tensor([0., 0., -inf])
&gt;&gt;&gt; # use pytorch's logsumexp() to implement numerically-stable version
&gt;&gt;&gt; n = X.size (0)
&gt;&gt;&gt; (X.unsqueeze (0).expand (n, n) + torch.diag (-torch.ones_like (X) / 0)).logsumexp (1) - X.logsumexp (0)
tensor([    0.0000,     0.0000, -1999.3069])
&gt;&gt;&gt;
&gt;&gt;&gt; Y = torch.randn (5)         # check against unstable version with random data (that doesn't underflow)
&gt;&gt;&gt; Y
tensor([ 0.1915,  0.3306,  0.2306,  0.8936, -0.2044])
&gt;&gt;&gt; (1 - Y.softmax (0)).log()
tensor([-0.1864, -0.2175, -0.1946, -0.4203, -0.1216])
&gt;&gt;&gt; n = Y.size (0)
&gt;&gt;&gt; (Y.unsqueeze (0).expand (n, n) + torch.diag (-torch.ones_like (Y) / 0)).logsumexp (1) - Y.logsumexp (0)
tensor([-0.1864, -0.2175, -0.1946, -0.4203, -0.1216])
</code></pre>
<p>Best.</p>
<p>K. Frank</p>

Zwei: <p>Hi Frank!</p>
<p>Thanks a lot for your answer code and explaination.</p>
<aside class="quote no-group" data-username="KFrank" data-post="6" data-topic="169007">
<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"> KFrank:</div>
<blockquote>
<p>With a little manipulation, you can zero out the <code>i == j</code> term in probability<br>
space (i.e., in “exp” space) by replacing the term with <code>-inf</code> (or a very<br>
large negative number) in log space (i.e., the space of your original <code>X</code>)<br>
and then apply pytorch’s <code>logsumexp()</code> to both the numerator and<br>
denominator of the above expression for <code>1 - softmax (X)</code>.</p>
</blockquote>
</aside>
<p>Adding <code>-inf</code> to the term <code>j== i</code> or making <strong>x<sub>i</sub> = - inf</strong> can remove the term perfectly in probability space (i.e., <code>exp(-inf) = 0</code>), and then combine this with <code>logsumexp()</code>, a numerically stable log1m_softmax() can be achived. This is the solution to this question!</p>
<p>However, I cannot find the solution in the stackexchange thread in all answers, which don’t introduce this -inf trick and lead to other unstable problem. Could you please point out the one where you found the answer?</p>
<aside class="quote no-group" data-username="KFrank" data-post="6" data-topic="169007">
<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"> KFrank:</div>
<blockquote>
<p>This is also discussed in the stats.stackexchange thread you linked to.<br>
You can implement your numerically-stable <code>log1m_softmax()</code> “by hand”<br>
along the lines discussed in that thread.</p>
</blockquote>
</aside>

KFrank: <p>Hi Zhengwei!</p>
<aside class="quote no-group" data-username="Zwei" data-post="7" data-topic="169007" 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/zwei/48/48403_2.png" class="avatar"> Zwei:</div>
<blockquote>
<p>However, I cannot find the solution in the stackexchange thread in all answers, which don’t introduce this -inf trick and lead to other unstable problem.</p>
</blockquote>
</aside>
<p>I didn’t mean to suggest that the stackexchange thread proposes using<br>
<code>-inf</code> to “zero out” the extra term in log-probability space.  I chose to go<br>
this route in order to use pytorch’s <code>logsumexp()</code> and avoid implementing<br>
a version of the log-sum-exp trick that works for the terms that show up<br>
in the expression for <code>log (1 - softmax (X))</code>.</p>
<p>Best.</p>
<p>K. Frank</p>

Zwei: <aside class="quote no-group quote-modified" data-username="Zwei" data-post="5" data-topic="169007">
<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/zwei/48/48403_2.png" class="avatar"> Zwei:</div>
<blockquote>
<p>Therefore, <strong>LSE</strong> = LogSumExp(<strong>X - c</strong>) = 0, and <strong>LSE</strong> - (<strong>X - c</strong>) = [2000, 2000, 0].</p>
<p>This lead to log(<strong>LSE</strong> - (<strong>X - c</strong>)) = [7.6, 7.6, -Inf].</p>
<aside class="quote no-group" data-username="ConvolutionalAtom" data-post="4" data-topic="169007">
<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/convolutionalatom/48/45668_2.png" class="avatar"> ConvolutionalAtom:</div>
<blockquote>
<p>The denominator for your case is the same, the numerator will be the same except for the one term which is absent.</p>
</blockquote>
</aside>
<p>I think the denominator you mentioned is <strong>LSE</strong>, and after we subtract the dominant term from it, it gets to zero.</p>
</blockquote>
</aside>
<p>Sorry for a mistake here, it should be:</p>
<p>Therefore, sum of exp(<strong>X</strong> - <strong>c</strong>) is <strong>SE</strong> = 1, and <strong>SE</strong> - exp(<strong>X</strong> - <strong>c</strong>) = [1, 1, 0].</p>
<p>This leads to  log(<strong>SE</strong> - exp(<strong>X</strong> - <strong>c</strong>)) = [0, 0, -inf], so although the denominator is stable with LSE, the numerator is unstable on dominant term.</p>

Hichem: <p>I tried for days to develop an efficient way of computing a stable log(1-softmax) but with no success.</p>
<p>I searched in many forums but it seems your solution is the only one that is really stable : your method outputs log(1-softmax(x)) = [0., 0., -1001.3068] when x = [-1000., -1000., 0], which is very accurate (validated using the mpmath package), while all other solutions still output log(1-softmax(x)) = [0., 0., -inf]. In addition, the backward function does not produce any nan values which is nice.</p>
<p>The problem with your approach is that memory and compute requirements scale quadratically with the size of x.</p>
<p>Do you see any possible improvements to make its complexity linear?</p>
<p>Thank you in advance,<br>
A.K. Hichem</p>

KFrank: <p>Hi Ammar!</p>
<aside class="quote no-group" data-username="Hichem" data-post="10" data-topic="169007" 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/hichem/48/64563_2.png" class="avatar"> Hichem:</div>
<blockquote>
<p>your method outputs log(1-softmax(x)) = [0., 0., -1001.3068] when x = [-1000., -1000., 0]</p>
</blockquote>
</aside>
<p>I do not believe that the method I posted above outputs this specific result.<br>
(See the example script, below.)  I think you must have a mistake in this<br>
test case you posted.</p>
<aside class="quote no-group">
<blockquote>
<p>The problem with your approach is that memory and compute requirements scale quadratically with the size of x.</p>
<p>Do you see any possible improvements to make its complexity linear?</p>
</blockquote>
</aside>
<p>The singularity / numerical instability in <code>log (1-softmax (X))</code> occurs<br>
only when largest element of <code>X</code> is sufficiently larger than the others so<br>
that <code>softmax (X)</code> contains an element equal (or very nearly equal) to<br>
one.  Then the corresponding element of <code>1-softmax (X)</code> approaches<br>
zero and the <code>log()</code> diverges.</p>
<p>In the approach I posted above, it was convenient to “zero out” all of the<br>
terms in <code>X</code> – whether or not they were potentially problematic – at the<br>
quadratic cost of <code>expand()</code>ing <code>X</code> into a square matrix with <code>-inf</code> along its<br>
diagonal.</p>
<p>It is indeed possible to deal only with largest term in <code>X</code> – the only one<br>
which could be problematic – in linear time and space, but at the cost<br>
of additional code complexity and, roughly speaking, implementing the<br>
log-sum-exp trick “by hand.”</p>
<p>This is illustrated in the following script that implements a linear algorithm<br>
and compares its results with the quadratic algorithm I posted above:</p>
<pre data-code-wrap="python"><code class="lang-python">import torch
print (torch.__version__)

_ = torch.manual_seed (2024)

# algorithm linear in length of X

def log1msm (X):                                        # for one-dimensional tensor X
    xm, im = X.max (0)                                  # largest value in X is the potential problem
    X_bar = X - xm                                      # uniform shift doesn't affect softmax (except numerically)
    lse = X_bar.logsumexp (0)                           # denominator for final result
    sumexp = X_bar.exp().sum() - X_bar.exp()            # sumexp[im] potentially zero
    sumexp[im] = 1.0                                    # protect against log (0)
    log1msm = sumexp.log()                              # good for all but i = im
    X_bar = X_bar.clone()                               # to support backward pass
    X_bar[im] = -float ('inf')                          # "zero out" xm in log space
    log1msm[im] = X_bar.logsumexp (0)                   # replace bad xm term
    log1msm -= lse                                      # final result
    return  log1msm

def log1msmB (X):                                       # previous algorithm quadratic in length of X
    n = X.size (0)
    return  (X.unsqueeze (0).expand (n, n) + torch.diag (-torch.ones_like (X) / 0)).logsumexp (1) - X.logsumexp (0)

y = torch.randn (8)
y[1] = -1000.0
y.requires_grad = True
z = torch.randn (8)
z[1] = 1000.0
z.requires_grad = True

xl = [   # list of test tensors for log1msm()
    torch.tensor ([-1000.0, -1000.0,     0.0], requires_grad = True),
    torch.tensor ([-1000.0, -1000.0,  1000.0], requires_grad = True),
    torch.tensor ([-1000.0,  1000.0,  1000.0], requires_grad = True),
    torch.randn (8, requires_grad = True),
    y,
    z
]

for  X in xl:   # compare log1msm() with previous quadratic version
    print ('X = ...')
    print (X)
    lg1m = log1msm (X)
    lg1m.sum().backward()
    X_grad = X.grad
    X.grad = None
    lg1mB = log1msmB (X)
    lg1mB.sum().backward()
    X_grad_B = X.grad
    print ('log1msm (X) = ...')
    print (lg1m)
    print ('log1msmB (X) = ...')
    print (lg1mB)
    print ('torch.allclose (lg1m, lg1mB, atol = 1.e-4):', torch.allclose (lg1m, lg1mB, atol = 1.e-4))
    print ('X_grad = ...')
    print (X_grad)
    print ('X_grad_B = ...')
    print (X_grad_B)
    print ('torch.allclose (X_grad, X_grad_B, atol = 1.e-4):', torch.allclose (X_grad, X_grad_B, atol = 1.e-4))
</code></pre>
<p>Here is the script’s output:</p>
<pre><code class="lang-plaintext">2.1.2
X = ...
tensor([-1000., -1000.,     0.], requires_grad=True)
log1msm (X) = ...
tensor([   0.0000,    0.0000, -999.3068], grad_fn=&lt;SubBackward0&gt;)
log1msmB (X) = ...
tensor([   0.0000,    0.0000, -999.3068], grad_fn=&lt;SubBackward0&gt;)
torch.allclose (lg1m, lg1mB, atol = 1.e-4): True
X_grad = ...
tensor([ 0.5000,  0.5000, -1.0000])
X_grad_B = ...
tensor([ 0.5000,  0.5000, -1.0000])
torch.allclose (X_grad, X_grad_B, atol = 1.e-4): True
X = ...
tensor([-1000., -1000.,  1000.], requires_grad=True)
log1msm (X) = ...
tensor([    0.0000,     0.0000, -1999.3069], grad_fn=&lt;SubBackward0&gt;)
log1msmB (X) = ...
tensor([    0.0000,     0.0000, -1999.3069], grad_fn=&lt;SubBackward0&gt;)
torch.allclose (lg1m, lg1mB, atol = 1.e-4): True
X_grad = ...
tensor([ 0.5000,  0.5000, -1.0000])
X_grad_B = ...
tensor([ 0.5000,  0.5000, -1.0000])
torch.allclose (X_grad, X_grad_B, atol = 1.e-4): True
X = ...
tensor([-1000.,  1000.,  1000.], requires_grad=True)
log1msm (X) = ...
tensor([ 0.0000, -0.6931, -0.6931], grad_fn=&lt;SubBackward0&gt;)
log1msmB (X) = ...
tensor([ 0.0000, -0.6932, -0.6932], grad_fn=&lt;SubBackward0&gt;)
torch.allclose (lg1m, lg1mB, atol = 1.e-4): True
X_grad = ...
tensor([0., 0., 0.])
X_grad_B = ...
tensor([0.0000e+00, 2.9087e-05, 2.9087e-05])
torch.allclose (X_grad, X_grad_B, atol = 1.e-4): True
X = ...
tensor([ 0.7893, -0.9312, -1.9368, -0.2072, -1.0539,  1.0529,  1.3023, -0.4860],
       requires_grad=True)
log1msm (X) = ...
tensor([-0.2220, -0.0363, -0.0131, -0.0763, -0.0320, -0.2999, -0.4042, -0.0572],
       grad_fn=&lt;SubBackward0&gt;)
log1msmB (X) = ...
tensor([-0.2220, -0.0363, -0.0131, -0.0763, -0.0320, -0.2999, -0.4042, -0.0572],
       grad_fn=&lt;SubBackward0&gt;)
torch.allclose (lg1m, lg1mB, atol = 1.e-4): True
X_grad = ...
tensor([ 0.0137,  0.0100,  0.0040,  0.0175,  0.0090, -0.0084, -0.0601,  0.0144])
X_grad_B = ...
tensor([ 0.0137,  0.0100,  0.0040,  0.0175,  0.0090, -0.0084, -0.0601,  0.0144])
torch.allclose (X_grad, X_grad_B, atol = 1.e-4): True
X = ...
tensor([-4.0429e-02, -1.0000e+03, -8.1404e-01,  1.3722e+00,  5.0596e-01,
        -4.8234e-01, -7.8525e-01,  6.6808e-01], requires_grad=True)
log1msm (X) = ...
tensor([-0.1007,  0.0000, -0.0452, -0.4996, -0.1808, -0.0635, -0.0465, -0.2163],
       grad_fn=&lt;SubBackward0&gt;)
log1msmB (X) = ...
tensor([-0.1007,  0.0000, -0.0452, -0.4996, -0.1808, -0.0635, -0.0465, -0.2163],
       grad_fn=&lt;SubBackward0&gt;)
torch.allclose (lg1m, lg1mB, atol = 1.e-4): True
X_grad = ...
tensor([ 0.0237,  0.0000,  0.0135, -0.1160,  0.0256,  0.0177,  0.0139,  0.0217])
X_grad_B = ...
tensor([ 0.0237,  0.0000,  0.0135, -0.1160,  0.0256,  0.0177,  0.0139,  0.0217])
torch.allclose (X_grad, X_grad_B, atol = 1.e-4): True
X = ...
tensor([-4.4394e-01,  1.0000e+03,  5.9861e-01,  6.4581e-01,  6.3061e-01,
        -1.4668e+00, -6.7976e-01, -6.9703e-01], requires_grad=True)
log1msm (X) = ...
tensor([   0.0000, -997.9874,    0.0000,    0.0000,    0.0000,    0.0000,
           0.0000,    0.0000], grad_fn=&lt;SubBackward0&gt;)
log1msmB (X) = ...
tensor([   0.0000, -997.9874,    0.0000,    0.0000,    0.0000,    0.0000,
           0.0000,    0.0000], grad_fn=&lt;SubBackward0&gt;)
torch.allclose (lg1m, lg1mB, atol = 1.e-4): True
X_grad = ...
tensor([ 0.0857, -1.0000,  0.2432,  0.2549,  0.2511,  0.0308,  0.0677,  0.0666])
X_grad_B = ...
tensor([ 0.0857, -1.0000,  0.2432,  0.2549,  0.2511,  0.0308,  0.0677,  0.0666])
torch.allclose (X_grad, X_grad_B, atol = 1.e-4): True
</code></pre>
<p>Best.</p>
<p>K. Frank</p>

Hichem: <p>Hi Frank!</p>
<p>Thank you for your answer! it was very helpful.</p>
<aside class="quote no-group quote-modified" data-username="KFrank" data-post="11" data-topic="169007">
<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"> KFrank:</div>
<blockquote>
<aside class="quote no-group" data-username="Hichem" data-post="10" data-topic="169007">
<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/hichem/48/64563_2.png" class="avatar"> Hichem:</div>
<blockquote>
<p>your method outputs log(1-softmax(x)) = [0., 0., -1001.3068] when x = [-1000., -1000., 0]</p>
</blockquote>
</aside>
<p>I do not believe that the method I posted above outputs this specific result.<br>
(See the example script, below.) I think you must have a mistake in this<br>
test case you posted.</p>
</blockquote>
</aside>
<p>Yes! I made a little mistake. This result appears when x = [-1000., -1000., <strong>2</strong>].</p>
<h1><a name="stability-and-precision-study-1" class="anchor" href="#stability-and-precision-study-1"></a>Stability and precision study</h1>
<p>To confirm the stability and precision of your function (let’s call it KF), I compared its output and its gradient with the “real” values of the log(1-softmax).</p>
<p>The “real” values of log(1-softmax) are computed using the mpmath package using a ridiculously high precision (2000 bits).</p>
<p>The input I used is a randomly generated vector X of 100 elements on a normal distribution N(0,1) with a little modification : I set an element of X to be equal to C and C varies from 0 to 100. As C grows, the sum of the exponents of X will be almost equal to exp(C) which creates the numerical instability.</p>
<p>The next image is a plot of the maximum absolute error between KF’s output and the “real” output of log(1-softmax), noted max(abs(KF(X) - real_values)), with respect to C. I also plotted the naive implementation torch.log(1-torch.softmax(X)) for reference. As we can see, KF stays stable and accurate while the “naive” approach diverges to infinity after C ≈ 25. The maximum absolute error of KF, reported in this test, is 6.56e-6.</p>
<p><img src="https://discuss.pytorch.org/uploads/default/original/3X/9/7/97deb91c89723e51fc43700b8af51fc5036b6c93.png" alt="image" data-base62-sha1="lFvhmynA1zpqf79P8zjDEAcgd8L" width="567" height="455"></p>
<p>To test the gradient, the next image is a plot of max(abs(J(KF(X)) - real_jacobian)) with respect to C, where J(KF(X)) is the Jacobian matrix of KF with respect to X. We can see that the absolute error stays small even for large values of C. The maximum absolute error reported in this test is 6.79e-6.</p>
<p><img src="https://discuss.pytorch.org/uploads/default/original/3X/a/4/a4024e9094c77d99c97afbf3ef23a6cf393626ee.png" alt="image" data-base62-sha1="noTflY2phEAujPPrR8aqd1OCheS" width="554" height="455"></p>
<h1><a name="generalization-to-multi-dimensional-tensors-2" class="anchor" href="#generalization-to-multi-dimensional-tensors-2"></a>Generalization to multi-dimensional tensors</h1>
<p>I also generalized your function to support the “dim” argument present in torch.log_softmax to be able to use it with multi-dimensional tensors:</p>
<pre><code class="lang-auto">def log1m_softmax(X : torch.Tensor, dim : int):
    xm, im = X.max(dim, keepdim=True)                                # largest value in X is the potential problem
    X_bar = X - xm                                                   # uniform shift doesn't affect softmax (except numerically)
    lse = X_bar.logsumexp(dim, keepdim=True)                         # denominator for final result
    sumexp = X_bar.exp().sum(dim, keepdim=True) - X_bar.exp()        # sumexp[im] potentially zero
    sumexp.scatter_(dim, im, 1.0)                                    # protect against log (0)
    log1msm = sumexp.log()                                           # good for all but i = im
    X_bar = X_bar.clone()                                            # to support backward pass
    X_bar.scatter_(dim, im, -float ('inf'))                          # "zero out" xm in log space
    log1msm.scatter_(dim, im, X_bar.logsumexp (dim).view(im.shape))  # replace bad xm term
    log1msm -= lse                                                   # final result
    return log1msm
</code></pre>
<p>Thanks again,<br>
Hichem A.K.</p>

Qi-Zhan: <p>Hello, could you share the mpmath script that implement log(1-softmax), I’m interested in it. Thank you!</p>

Hichem: <p>Hello Qi-Zhan</p>
<p>Yes sure ! Here is a <a href="https://github.com/HichemAK/stable-log1msoftmax" rel="noopener nofollow ugc">repository</a> that contains the mpmath script inside the notebook named <em>test.ipynb</em></p>
<p>Best,<br>
Hichem</p>

Qi-Zhan: <p>Thank you, already stared <img src="https://discuss.pytorch.org/images/emoji/apple/wink.png?v=12" title=":wink:" class="emoji" alt=":wink:" loading="lazy" width="20" height="20"></p>