log_softmax produces -inf on Apple Neural Engine in fp16 when one class dominates
- Dominant language
- Python
- Stars
- 5.4k
- Forks
- 850
- Avg merge
- 4d 5h
- Merged PRs (30d)
- 10
Description
## Problem
The PyTorch `log_softmax` converter produces `-inf` values on Apple Neural Engine (ANE) in fp16 when the input has a dominant class with large logit values. This silently corrupts the output of every classification model using `log_softmax`, `F.log_softmax`, or `F.cross_entropy` on ANE.
## Root Cause
The current converter at `converters/mil/frontend/torch/ops.py` line 5904 computes:
`python
res = mb.softmax(x=x, axis=axis)
res = mb.log(x=res)
`
This is a naive `log(softmax(x))` decomposition. While `softmax` itself uses max-shift stabilization internally, the resulting probabilities for non-dominant classes underflow to **0** in fp16 (any probability below ~6e-5). Then `log(0) -> -inf`.
In fp32, these tiny probabilities are representable (e.g., 1e-38), so CPU and GPU compute units are unaffected. The bug is specific to fp16 on ANE.
## Reproduction
`python
import torch
import coremltools as ct
import numpy as np
class LogSoftmaxModel(torch.nn.Module):
def forward(self, x):
return torch.nn.functional.log_softmax(x, dim=-1)
model = LogSoftmaxModel().eval()
# One dominant class at index 3 with large logit
x = torch.tensor([[0.0, 0.0, 0.0, 50.0, 0.0, 0.0, 0.0, 0.0]])
traced = torch.jit.trace(model, x)
mlmodel = ct.convert(traced,
inputs=[ct.TensorType(shape=x.shape)],
compute_precision=ct.precision.FLOAT16)
pytorch_out = model(x).detach().numpy()
coreml_out = list(mlmodel.predict({'x_1': x.numpy()}).values())[0]
print('PyTorch:', pytorch_out)
# Expected: [[-50., -50., -50., 0., -50., -50., -50., -50.]]
print('CoreML:', coreml_out)
# Broken: [[-inf, -inf, -inf, 0., -inf, -inf, -inf, -inf]]
`
## Expected vs Actual
| Input logit | PyTorch (correct) | CoreML fp16 (broken) |
|-------------|------------------|---------------------|
| 0.0 (non-dominant) | -50.0 | **-inf** (WRONG) |
| 50.0 (dominant) | 0.0 | 0.0 (correct) |
## Fix
Use the standard numerically stable log-softmax decomposition:
`
log_softmax(x) = x - max(x) - log(sum(exp(x - max(x))))
`
By subtracting `max(x)` first, all `exp()` arguments are <= 0, so values are in (0, 1]. The log of the sum is computed directly, avoiding the underflow-prone intermediate softmax probabilities.
This is the formula used by:
- PyTorch's own fused `log_softmax` CUDA kernel
- coremltools' TensorFlow frontend for `_softmax_cross_entropy_with_logits`
- JAX's `jax.nn.log_softmax`
## Impact
Every classification model using `nn.LogSoftmax`, `F.log_softmax`, or `F.cross_entropy` on ANE with fp16 precision. This includes BERT, ResNet, ViT, and most models that compute cross-entropy loss or log-probability outputs.
## Environment
- coremltools version: 9.0 (main branch, commit as of 2026-05-29)
- Affected compute unit: Neural Engine (fp16)
- Unaffected: CPU, GPU
## Related Issues
- #2687 -- Softplus fp16 overflow (same class of bug, fixed in PR #2725)
- #2690 -- LogSumExp fp16 overflow (same class of bug, fixed in PR #2726)
- #2359 -- Mish fp16 errors (related, softplus-dependent)
- #2625 -- MobileNetV3 fp16 errors (related)
Contributor guide
Research direction
Start in converters/mil/frontend/torch/ops.py at the log_softmax converter around line 5904, then run the provided fp16 Apple Neural Engine reproduction. Replace the underflow-prone decomposition with the stated stable formulation and verify that non-dominant classes produce finite values such as -50 rather than -inf while the dominant class remains 0.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch
- Domain
- machine-learning, tooling
- Issue type
- Bug
- Difficulty
- 2/5
- Estimated time
- 1-3 hours
- Activity status
- Quiet
- Clarity
- Clearly specified
- Newbie friendliness
- 75/100