Repository: pytorch
Issue: cure unnecessary NaN values that arise from ±∞ arguments
Labels: module: numerical-stability,module: bootcamp,module: autograd,triaged,module: NaNs and Infs

## 🐛 Bug

### First problem
`logaddexp` (⊕) and its generalization `logsumexp` are supposed to be associative operations.
Indeed ((-∞ ⊕ -∞) ⊕ 2) returns the same forward value as (-∞ ⊕ (-∞ ⊕ 2)) or `logsumexp`(-∞, -∞, 2), namely 2.
But unfortunately it returns `nan` values on the backward pass, where it should return `0` values as the others do.

**At the end of this ticket, I give a detailed discussion including real consequences suffered by my classroom students, as well as a solution including a monkey patch as proof of principle.**

### Second problem
In addition, these functions are supposed to agree, but they have different `nan` behavior for x < ∞ and y = ∞:
* `∂logaddexp(x,y)/∂(x,y) = (0,1)`
* `∂logsumexp([x,y])/∂(x,y) = (0, nan)`
* `softmax([x,y]) = [nan, nan]`

### Third problem
And these calls unnecessarily return `nan` when x = -∞ or x = ∞:
* `∂logsumexp([x])/∂x = nan`
* `softmax([x]) = [nan]`

At the end of the ticket, I'll explain why none of these bullet point cases actually need to return `nan`.  

### Potential other problems
Should other PyTorch library functions be checked for similar bugs?

## To Reproduce
```
# SETUP
from torch import tensor, stack, softmax, logsumexp
from torch import logaddexp as lae
def lse(*inputs): return logsumexp(stack(inputs), dim=0)  # logsumexp applied to scalars
inf = tensor(float('inf'))
two = tensor(2.)
def leaf(value): return tensor(value, requires_grad=True)

# FIRST PROBLEM
x = leaf(float('-inf'))
lae(lae(x, -inf), two).backward()
print(x.grad.item())

x = leaf(float('-inf'))
lae(lse(x, -inf), two).backward()
print(x.grad.item())

x = leaf(float('-inf'))
lae(x, lae(-inf, two)).backward()
print(x.grad.item())

# SECOND PROBLEM
x = leaf(2.)
y = leaf(float('inf'))
lae(x, y).backward()
print((x.grad.item(), y.grad.item()))

x = leaf(2.)
y = leaf(float('inf'))
lse(x, y).backward()
print((x.grad.item(), y.grad.item()))

softmax(stack((x,y)), dim=0).tolist()

# THIRD PROBLEM
for x in leaf(float('-inf')), leaf(float('inf')):
    lse(x).backward()
    print(x.grad.item())
    softmax(stack((x,)), dim=0).backward()
    print(x.grad.item())
```

## Expected behavior

None of these examples should print any `nan` values.  
* Under `FIRST PROBLEM`, all `nan` values should be 0.0.  
* Under `SECOND PROBLEM`, all printed pairs should be `(0.0, 1.0)` (or `[0.0, 1.0]`).
* Under `THIRD PROBLEM`, all print statements should print `1.0`.

## Environment

[PyTorch version: 1.6.0
Is debug build: False
CUDA used to build PyTorch: 10.2
ROCM used to build PyTorch: N/A

OS: Linux Mint 20 (x86_64)
GCC version: (Ubuntu 9.3.0-17ubuntu1~20.04) 9.3.0
Clang version: Could not collect
CMake version: Could not collect

Python version: 3.8 (64-bit runtime)
Is CUDA available: False
CUDA runtime version: No CUDA
GPU models and configuration: No CUDA
Nvidia driver version: No CUDA
cuDNN version: No CUDA
HIP runtime version: N/A
MIOpen runtime version: N/A

Versions of relevant libraries:
[pip3] numpy==1.19.2
[pip3] torch==1.6.0
[conda] blas                      1.0                         mkl  
[conda] cudatoolkit               10.2.89              hfd86e86_1  
[conda] mkl                       2020.2                      256  
[conda] mkl-service               2.3.0            py38he904b0f_0  
[conda] mkl_fft                   1.2.0            py38h23d657b_0  
[conda] mkl_random                1.1.1            py38h0573a6f_0  
[conda] numpy                     1.19.2           py38h54aff64_0  
[conda] numpy-base                1.19.2           py38hfa32c7d_0  
[conda] pytorch                   1.6.0           py3.8_cuda10.2.89_cudnn7.6.5_0    pytorch

## Additional context

* `logsumexp` on a collection of scalars currently returns `-inf` if all of its arguments are `-inf`, and returns `inf` if any of its arguments are `inf`.  _This is the correct behavior._
* In these cases, the partials with respect to the `inf` and `-inf` scalars are currently returned as `nan`.  _This is *usually* the correct behavior (with exceptions discussed below)._   
* In these cases, the partials with respect to any finite inputs are returned as 0.  _This is the correct behavior_.

### Real-world impact of first problem

Recently I gave a class assignment about HMMs.  The HMM in question included some zero probabilities for impossible transitions or emissions.  Unfortunately, students who represented probabilities in the log domain were unable to train their model because they got `nan` values in the gradient, due to adding log-zeros in the forward pass, as shown at the start of this post.

### Solution to first problem 

These `nan` values should have been 0.  Let  x = y = -∞.  Even though the gradient of (x ⊕ y) = -∞ really is `[nan, nan]`, the gradient of ((x ⊕ y) ⊕ 2) = 2 should be `[0, 0, 1]`.  After all, increasing x or y by a small finite epsilon has 0 effect on the result, right?

What's going on is that the `nan` partials for `logsumexp` represent indeterminate _finite_ quantities.  Thus, `0 * nan` should yield `0`.  In particular, when `grad_output = ∂loss/∂logsumexp(x, ...) = 0`, then we want `grad_input = ∂loss/∂x = 0 * ∂logsumexp(x, ...)/∂x` to be `0` as well, even when `∂logsumexp(x, ...)/∂x = nan`.

I've placed a monkey patch [here](http://cs.jhu.edu/~jason/tmp/logsumexp_safe.py) that shows how this can be arranged by looking at the zeroes in `grad_output`.  If you do `import logsumexp_safe`, you'll get a `safe_inf` keyword argument to `logaddexp` and `logsumexp` that turns on this `0 * nan = 0` behavior by checking for zeroes in the `grad_output` tensor.  The keyword argument is so that you can easily compare the legacy behavior to the new behavior.

Of course, it would be faster to build this into the C++ built-in functions.  Is that reasonable to request?  The new behavior could be turned on all the time -- I don't think there's much backward compatibility concern, since it's unlikely that anyone ever really _wanted_ a `nan` in their gradient!  However, if it creates a slowdown, the keyword argument could be retained so that users would turn it on only when necessary.

### Discussion of solution to first problem

In general, you can think of `grad_output = 0` as license to abort the backward computation through a node, since it means that no change to the inputs (at least, no finite change) could affect the outputs.  If PyTorch ran on scalars, it might test for this case as an optimization.  In practice, however,  `grad_output` is often a tensor that mixes 0 and non-0 values.  In that case, only part of the backward pass can be skipped, and there would be considerable overhead to trying to skip the irrelevant parts of each tensor.

When I say that these `nan`s are _finite_ indeterminate quantities, I mean that they are like 0/0, which denotes any finite number.  To be precise, 0/0 denotes "a number whose product with 0 is 0," and thus we want its product with 0 to indeed be 0.  (This is different from the `nan` yielded by ∞ - ∞ or 0 * ∞.  Unfortunately, the IEEE floating-point standard only defines one kind of `nan`, so it conservatively defines `0 * nan = nan`, which is what causes the bug.)

But why finite?  A simple intuition is that in the computation (x ⊕ y) = -∞, increasing x or y by a small finite epsilon changes the output -∞ to -∞, and that could be a change of any finite size.  Hence the partial derivative could be any finite number.

A better perspective is that the value and gradient of x ⊕ y at x=y=-∞ should be determined by taking limits as finite x → -∞ and finite y → -∞.  The reason we get `nan`s in the gradient is that this double limit is not well defined.  We get different answers depending on whether x or y is the outer limit, and other answers if, for example, we say that x → -∞ with y = x+3.  Even so, the gradient is always in the interval (0,1) for any finite values of x and y, and therefore all of these limits will give values in the interval [0,1].  Thus, the `nan` values represent indeterminate finite quantities.

### Explanation of second problem

I wrote above;
> * `logsumexp` on a collection of scalars currently returns `-inf` if all of its arguments are `-inf`, and returns `inf` if any of its arguments are `inf`.  _This is the correct behavior._
> * In these cases, the partials with respect to the `inf` and `-inf` scalars are currently returned as `nan`.  _This is *usually* the correct behavior (with exceptions discussed below)._   

If `logsumexp` has multiple `inf` arguments, then the gradient should be `nan` with respect to these arguments (by the multiple limit explanation above) and should be `0` with respect to the other arguments.  This is currently implemented correctly.  

However, if it has only a single `inf` argument, then the gradient with respect to that argument should really be 1 (there's a single limit operation and the result is 1).  Currently `logaddexp` gets this right and returns a one-hot vector, but `logsumexp` gets it wrong.  `softmax` gets it even wronger.  It would be nice to fix these.

### Explanation of third problem

Similarly, if `logsumexp` has multiple `-inf` arguments and no other arguments, then the gradient should be `nan` with respect to these arguments (by the multiple limit explanation above).  This is currently implemented correctly.

However, in the corner case where `logsumexp` has only a single `-inf` argument and no other arguments, then the gradient should be `1` (again, there's only a single limit).  Again, `logsumexp` and `softmax` get these cases wrong.  It would be nice to fix them.

As a practical matter, the first problem is the most important, and the third problem is the least important.

cc @ezyang @albanD @zou3519 @gqchen @pearu @nikitaved @soulitzer