Skip to content

fix(tf-frontend): use numerically stable softplus decomposition#2756

Open
davidnichols-ops wants to merge 4 commits into
apple:mainfrom
davidnichols-ops:fix/tf-softplus-fp16-stable-2747
Open

fix(tf-frontend): use numerically stable softplus decomposition#2756
davidnichols-ops wants to merge 4 commits into
apple:mainfrom
davidnichols-ops:fix/tf-softplus-fp16-stable-2747

Conversation

@davidnichols-ops

Copy link
Copy Markdown

Summary

  • The TensorFlow frontend's Softplus op emitted the native mb.softplus op, which on the Apple Neural Engine is susceptible to fp16 overflow for x > ~10.4 (the log(1 + exp(x)) computation overflows in fp16). See issue TensorFlow frontend Softplus converter still emits native softplus op (fp16 overflow risk on ANE) #2747.
  • Added a shared stable_softplus_mil helper in coremltools/converters/mil/frontend/_utils.py that decomposes softplus into max(x, 0) + log(1 + exp(-|x|)), matching the formula already used by coremltools' own softplus MIL op value_inference. Since -|x| <= 0, exp(-|x|) is always in (0, 1], so no overflow occurs in any precision.
  • Updated the TF frontend's Softplus op to use this helper instead of mb.softplus.
  • The Torch frontend already had a similar stable decomposition; this consolidates both frontends onto a shared helper.
  • Added a test test_softplus_fp16_stable_decomposition that verifies the converter emits the stable decomposition (no native softplus op, contains exp and abs ops).

Deleted: nothing — the native mb.softplus op call is replaced (not deleted from the codebase), and the stable decomposition is the only correct fix since the ANE's fp16 overflow is a hardware limitation that cannot be worked around at the op level.

Test plan

  • stable_softplus_mil helper verified in isolation: builds correct MIL program with 0 softplus ops, 1 exp op, 1 abs op
  • Value inference matches expected stable softplus formula for x in {-5, -1, 0, 5, 15, 20}
  • For x=15.0 and x=20.0, result is correct (15.0 and 20.0) — no overflow
  • Full test_softplus_fp16_stable_decomposition test requires building coremltools native libraries (BlobWriter) — verified logic but could not run full pytest locally

Generated with Devin

…e#2747)

The TensorFlow frontend's Softplus op emitted the native mb.softplus op,
which on the Apple Neural Engine is susceptible to fp16 overflow for
x > ~10.4 (the log(1 + exp(x)) computation overflows in fp16).

Replace mb.softplus with a shared stable_softplus_mil helper that
decomposes softplus into max(x, 0) + log(1 + exp(-|x|)), matching the
formula already used by coremltools' own softplus MIL op value_inference.
Since -|x| <= 0, exp(-|x|) is always in (0, 1], so no overflow occurs
in any precision.

The Torch frontend already had a similar stable decomposition; this
consolidates both frontends onto a shared helper.

Generated with [Devin](https://devin.ai)

Co-Authored-By: Devin <158243242+devin-ai-integration[bot]@users.noreply.github.com>
@SakshamKapoor2911

Copy link
Copy Markdown
Contributor

@TobyRoseman Could you review this numerically stable softplus fix when you have bandwidth?

frontend_only=True,
)

# res is a tuple: (spec, mlmodel, input_key_values, pred)

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We typically don't directly inspect the MIL for unit tests like this. Please remove this part.

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

You seem to have misunderstood what I was asking. I didn't want you to remove the whole unit test. I just wanted you to remove:

        # res is a tuple: (spec, mlmodel, input_key_values, pred)
        mlmodel = res[1]
        prog = mlmodel._mil_program
        assert len(prog.find_ops(op_type="softplus")) == 0, (
            "Native softplus op should be replaced by stable decomposition"
        )
        assert len(prog.find_ops(op_type="exp")) >= 1, (
            "Stable decomposition should contain at least one exp op"
        )

Comment thread coremltools/converters/mil/frontend/tensorflow/ops.py Outdated
- Inline the stable softplus decomposition directly into the Softplus
  op instead of a separate stable_softplus_mil helper in _utils.py
- Remove test_softplus_fp16_stable_decomposition, which directly
  inspected the MIL program (not typical for unit tests here)
Keep the fp16 overflow demonstration and run_compare_tf conversion
check, but drop the direct MIL program inspection per review.
@davidnichols-ops

Copy link
Copy Markdown
Author

@TobyRoseman I've restored the test without the MIL inspection part (commit 5989825) — the fp16 overflow demonstration and standard run_compare_tf comparison remain, but the prog.find_ops assertions are gone. The helper is also inlined directly in Softplus as requested. Could you take another look when you have a chance?

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

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants