apple / apple/coremltools

log_softmax produces -inf on Apple Neural Engine in fp16 when one class dominates

Open Beginner friendly
#2,728 0 comments 0 reactions 0 assignees View on GitHub
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

Open the contributing 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.