dwavesystems / dwavesystems/dwave-optimization
`SoftMax` produces `nan` when any input element exceeds ~710
- 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
Assessment
This issue has not been assessed yet.