Repository: pytorch
Issue: Implementation of a numerically stable log(1 - softmax) function in PyTorch
Labels: module: numerical-stability,triaged,module: python frontend

### 🚀 The feature, motivation and pitch

Hello,

I am working on editing knowledge in transformers which, in the case of knowledge deletion, requires the minimization of the likelihood of targeted sequences of tokens, and therefore, I need in my work, a numerically stable log(1-softmax) function. Details are below.

The maximization of the likelihood of correct tokens is extensively used when training LLMs, and it requires an efficient and numerically stable implementation of the log(softmax(x)) function (to compute the CrossEntropy Loss) which was available in Pytorch for as long as I can remember. On the other hand, the minimization requires the implementation of a stable log(1-softmax(x)) function **which is not available at the moment**.

The naive implementation of log(1-softmax) using torch.log and torch.softmax is highly unstable when the softmax(x) is close to 1 for one element of its output. I used it myself and I found that my loss sometimes equaled "-inf." which consequently, produced "nan" values when back-propagating.

I consider the log(1-softmax) to be an important function in my use-case but also more generally in classification tasks. One could assume that someone would like to maximize some classes while minimizing others in some use-case.

After asking for a numerically stable log(1-softmax) function in a [Pytorch forum](https://discuss.pytorch.org/t/how-to-calculate-log-1-softmax-x-numerically-stably/169007/12?u=hichem), a user named KFrank proposed an elegant solution:

```
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
```

I generalized it to multi-dimensional tensors:

```
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
```

Finally I tested its numerical stability ; the methodology and results are explained in the previous [Pytorch forum](https://discuss.pytorch.org/t/how-to-calculate-log-1-softmax-x-numerically-stably/169007/12?u=hichem).

## Request

Is it possible to add this function to the list of native functions of torch, such that it can be used through torch.log1m_softmax(x, dim) ?

I'll be happy to contribute to the code. I imagine this requires rewriting this function in C++, and placing somewhere closed the log_softmax function located in "pytorch/aten/src/ATen/native/SoftMax.cpp"

### Alternatives

The only alternative is to compute log(1-softmax) using a combination of torch.log and torch.softmax which is highly unstable.

### Additional context

_No response_

cc @albanD