fix(torch-frontend): numerically stable log_softmax and logcumsumexp - #2763
Open
davidnichols-ops wants to merge 3 commits into
Open
fix(torch-frontend): numerically stable log_softmax and logcumsumexp#2763davidnichols-ops wants to merge 3 commits into
davidnichols-ops wants to merge 3 commits into
Conversation
The TensorFlow frontend's Softplus op emitted the native mb.softplus MIL op, which on the Apple Neural Engine overflows in fp16 for x > ~10.4 (exp(x) exceeds 65504, the fp16 max). This caused output collapse to inf/0 instead of the correct ~x value. Replace mb.softplus with the numerically stable decomposition max(x, 0) + log(1 + exp(-|x|)), inlined directly in the converter. Since -|x| <= 0, exp(-|x|) is always in (0, 1], so no overflow occurs in any precision. Add test_softplus_large_values which converts a TF softplus model with large inputs (x=15, x=20) and verifies correct output by running the model (not frontend-only). The test hardcodes ComputeUnit.ALL to make the ANE eligible for execution. Without the fix, the native mb.softplus overflows in fp16 on the ANE for these values, failing the output comparison. With the fix, the stable decomposition produces correct results. Closes apple#2747
…er request The test now only verifies runtime output correctness for large values. The MIL op inspection was removed per reviewer feedback. Note that on CPU, the native softplus produces correct output regardless of the fix — the overflow only manifests on ANE with fp16 compute.
log_softmax: replace naive softmax -> log decomposition with x - reduce_log_sum_exp(x, axis). The naive version overflows in fp16 on the Apple Neural Engine when one class dominates, producing -inf for non-dominant classes. reduce_log_sum_exp applies the max-shift trick internally, keeping exp arguments <= 0. logcumsumexp: shift by global max along the cumulative dimension before exp, then shift back. The naive exp -> cumsum -> log overflows in fp16 for x > ~10.4. This uses the global max rather than a cumulative max (not available as a MIL op), so extreme underflow may still produce -inf for early elements — but this is strictly better than the naive version which overflows to +inf/NaN for all elements. Adds large-value regression tests for both ops. Fixes apple#2728, apple#2729
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
softmax -> logdecomposition inlog_softmaxconverter withx - reduce_log_sum_exp(x, axis), which uses the max-shift trick internally to prevent fp16 overflow on the Apple Neural Engine.exp -> cumsum -> logdecomposition inlogcumsumexpconverter with a global-max-shifted version that prevents overflow for large inputs.Why
The naive decompositions overflow in fp16 on the ANE:
log_softmax:softmax -> logoverflows when one class dominates, producing-inffor non-dominant classes (exp(x) > 65504whenx > ~10.4).logcumsumexp:exp -> cumsum -> logoverflows for the same reason.reduce_log_sum_expalready exists as a MIL op and applies the numerically stable log-sum-exp trick:max(x) + log(sum(exp(x - max(x)))), keeping exp arguments ≤ 0.For
logcumsumexp, there is no cumulative-max MIL op, so we shift by the global max along the dimension. This prevents overflow (the catastrophic failure) though extreme underflow may still produce-inffor early elements far below the max — strictly better than the naive version which overflows to+inf/NaNfor all elements.What Changed
coremltools/converters/mil/frontend/torch/ops.py:log_softmax:mb.softmax -> mb.logreplaced withmb.reduce_log_sum_exp -> mb.sublogcumsumexp:mb.exp -> mb.cumsum -> mb.logreplaced withmb.reduce_max -> mb.sub -> mb.exp -> mb.cumsum -> mb.log -> mb.addcoremltools/converters/mil/frontend/torch/test/test_torch_ops.py:test_log_softmax: basic log_softmax test (did not exist before)test_log_softmax_large_values: regression test withx * 15.0inputstest_logcumsumexp_large_values: regression test withx * 15.0inputsScope
Only the
log_softmaxandlogcumsumexptorch frontend converters and their tests. No MIL op definitions, no other converters, no API changes.Deleted: nothing (bugfix adding stable decompositions — net-negative justification: the fix replaces 2-line naive decompositions with 3-4-line stable ones; deletion/relocation impossible because the numerical instability is in the decomposition itself, not in misplaced logic).