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

Open
#555 1 comment 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

Assessment

Difficulty
3/5
Estimated time
1-2 days
Newbie friendliness
66/100
Issue type
Bug
Clarity
Clearly specified
Activity status
Quiet
Tech stack
cpp, numpy, python
Domain
backend

Research direction

Start in src/nodes/softmax.cpp and inspect both the evaluation and propagate paths for SoftMaxNode. Reproduce the issue with [710, 1, 1] or [800, 1, 1], then verify that both paths return a valid probability distribution rather than nan for large inputs. Done means the overflowing inputs produce the expected dominant-element result without changing ordinary behavior.

Written by the indexing model from the issue text.

Description

bug
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
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.

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:

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))

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.

// 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.

Dominant language
C++
Stars
31
Forks
36
Avg merge
1d 9h
Merged PRs (30d)
4

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

More from dwavesystems/dwave-optimization

All issues in dwavesystems/dwave-optimization

Similar issues

More C++ issues

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.