Skip to content

Make aten::logaddexp, logaddexp2 and log_sigmoid numerically stable - #3068

Open
Raashish Aggarwal (raashish1601) wants to merge 1 commit into
microsoft:mainfrom
raashish1601:fix/stable-logaddexp-logsigmoid
Open

Raashish Aggarwal (raashish1601) wants to merge 1 commit into
microsoft:mainfrom
raashish1601:fix/stable-logaddexp-logsigmoid

Conversation

@raashish1601

Copy link
Copy Markdown

aten_logaddexp, aten_logaddexp2 and aten_log_sigmoid used the direct formulas, which overflow or underflow well inside the float32 range:

  • logaddexp(a, b) = log(exp(a) + exp(b)): for a, b around 100 or more exp overflows and the result is inf (PyTorch returns e.g. 1000.69 for logaddexp(1000, 1000)).
  • logaddexp2 has the same problem with 2**a + 2**b.
  • log_sigmoid(x) = log(sigmoid(x)): for x below about -100, sigmoid underflows to 0 and the result is -inf instead of about x.

They now use the usual stable forms (the ones PyTorch uses):

  • logaddexp: max(a, b) + log(1 + exp(-|a - b|))
  • logaddexp2: max(a, b) + log2(1 + 2**-|a - b|)
  • log_sigmoid: min(x, 0) - log(1 + exp(-|x|))

For a == b == +/-inf, a - b is NaN, so logaddexp/logaddexp2 return a + log(2) / a + 1 when the inputs are equal, which gives the right infinity (and the same value as the formula for finite inputs).

Testing:

  • Added ops.aten.logaddexp.large, ops.aten.logaddexp2.large and ops.aten.log_sigmoid.large OpInfos with values from -1000 to 1000 and +/-inf (all pairs for the binary ops). All three fail on main and pass with the change.
  • pytest tests/function_libs/torch_lib/ops_test.py -k "logaddexp or log_sigmoid or logsigmoid" passes locally, and an exported model with these ops matches PyTorch on large inputs through onnxruntime. ruff check / format are clean on the changed files.

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

Development

Successfully merging this pull request may close these issues.

1 participant