dwavesystems / dwavesystems/dwave-optimization

`SoftMax` produces `nan` when any input element exceeds ~710

Open
#555 1 comment 0 reactions 0 assignees View on GitHub
bug
Dominant language
C++
Stars
31
Forks
36
Avg merge
16h 55m
Merged PRs (30d)
8

Description

### Summary

`SoftMax` computes `exp(x_i)` directly for each element.
For any `x_i ≥ 710` (approximately), `exp(x_i)` overflows to `inf` in IEEE 754 double precision.
The result is `inf / inf = nan` for the overflowing elements and `0 / inf = 0` for the rest —
a silent, meaningless output with no error or warning.

The standard fix is the **max-shift trick**: subtract `max(x)` before exponentiating.
This is algebraically identical to the original formula but numerically stable because
all shifted values are `≤ 0`, which cannot overflow.

The documentation links to `scipy.special.softmax` as an equivalent — SciPy's implementation
is numerically stable and does not exhibit this issue.

---

### Expected behaviour

```python
x = model.integer(shape=(3,), lower_bound=0, upper_bound=1000)
sm = SoftMax(x)
# For x = [800, 1, 1]: result should be approximately [1.0, 0.0, 0.0]
```

`softmax([800, 1, 1])` should return a valid probability distribution.
The largest element dominates, so the result should be close to `[1, 0, 0]`.

---

### Actual behaviour

`nan` is returned silently. No error, no warning.

```python
from dwave.optimization.model import Model
from dwave.optimization.symbols import SoftMax

model = Model()
x = model.integer(shape=(3,), lower_bound=0, upper_bound=1000)
sm = SoftMax(x)

with model.lock():
model.states.resize(1)
x.set_state(0, [710, 1, 1])
print(sm.state(0)) # [nan, 0., 0.] <- silent wrong result
x.set_state(0, [800, 1, 1])
print(sm.state(0)) # [nan, 0., 0.] <- silent wrong result
```

NumPy stable reference:
```python
import numpy as np
x = np.array([800., 1., 1.])
result = np.exp(x - x.max()) / np.exp(x - x.max()).sum()
# array([1., 0., 0.]) <- correct
```

---

### Affected scope

Any `SoftMaxNode` whose predecessor can produce values above ~709.78.
`integer` variables with `upper_bound ≥ 710` are at risk.
The overflow threshold is `exp(709.78...) ≈ DBL_MAX`.

---

### Proposed fix (recommended by AI ([ChatGPT Enterprise](https://chatgpt.com/business/)))

In `src/nodes/softmax.cpp`, subtract the maximum element before exponentiating.
This is the same approach used by `scipy.special.softmax` — whose own docs cite [1] as the reference
for using shifting to avoid overflow.

```cpp
// Before (naive, overflows for large inputs):
for (double& val : values) {
val = std::exp(val);
denominator += val;
}

// After (numerically stable, identical to scipy.special.softmax):
double max_val = *std::max_element(values.begin(), values.end());
for (double& val : values) {
val = std::exp(val - max_val);
denominator += val;
}
```

The same one-line change is needed in the `propagate` path.

> [1] P. Blanchard, D.J. Higham, N.J. Higham,
> *"Accurately computing the log-sum-exp and softmax functions"*,
> IMA Journal of Numerical Analysis, Vol. 41(4), 2021.
> https://doi.org/10.1093/imanum/draa038

---

### Workaround

There is no safe workaround within the model graph itself.
Avoid `SoftMax` when any predecessor variable has `upper_bound ≥ 710`,
or pre-scale the inputs to a safe range before passing them to `SoftMax`.

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.