Shift by the max in cross_entropy

[ayaangazali]

Proposed changes

cross_entropy only depends on the gaps between logits, but a large shared offset changes the answer and eventually zeroes it:

>>> logits, targets = mx.array([[2.0, -1.0]]), mx.array([0])
>>> [nn.losses.cross_entropy(logits + off, targets).item()
...  for off in (0, 1e4, 1e6, 1e8)]
[0.04858734, 0.048828125, 0.0625, 0.0]

Equal logits are the clearest case, since the answer has to be log(2) at any magnitude:

>>> [nn.losses.cross_entropy(mx.array([[v, v]]), mx.array([0])).item()
...  for v in (1e0, 1e4, 1e6, 1e8)]
[0.6931472, 0.693359375, 0.6875, 0.0]

pytorch returns 0.6931472 for all four. A loss of 0.0 is the bad one: it reports a perfect prediction and hands back a zero gradient.

The cause is the composition, not any single op:

logsumexp_logits = mx.logsumexp(logits, axis=axis)
...
loss = logsumexp_logits - score

mx.logsumexp is right. logsumexp([1e8, 1e8]) is 1e8 + log(2), and the ulp of 1e8 in float32 is 8, so 1e8 is the correctly rounded answer. The problem is that the gap between the logits gets added to a large magnitude and rounds away before it is subtracted back off.

Shifting the logits by their max first keeps the gap at a magnitude where it survives. The loss is invariant to that shift, so nothing else has to change, and stop_gradient keeps the graph the same function it was:

logits = logits - mx.stop_gradient(mx.max(logits, axis=axis, keepdims=True))

The shift is applied before score is taken so every path benefits, including label smoothing and probability targets, which were also returning 0.0 at 1e8. It cancels for label smoothing too, since the smoothing term subtracts the mean logit and the shift moves the mean by the same amount.

Checked against pytorch after the change with zero mismatches: every magnitude from 1e0 to 1e36, the offset cases above, random inputs for all three reductions, label_smoothing of 0, 0.1 and 0.5, probability targets, and gradients (max difference 9e-10).

Checklist

Put an x in the boxes that apply.

• I have read the CONTRIBUTING document
• I have run pre-commit run --all-files to format my code / installed pre-commit prior to committing changes
• I have added tests that prove my fix is effective or that my feature works
• I have updated the necessary documentation (if needed)

CPU only build (MLX_BUILD_METAL=OFF) at 596dc79. test_losses.py, test_nn.py and test_optimizers.py pass, and the new assertions fail on main from 1e4 up. The existing -inf logits case is unaffected, since the row max is already 0 there. Python only, no rebuild involved.

---

freshman contributor, i use Claude Code while digging. found this by feeding a large constant offset through the losses and watching which ones stopped being shift invariant. the thing that took me longest was convincing myself logsumexp was not at fault: at 1e8 the ulp is 8, so it is returning the best float32 value there is and the loss has to be assembled differently instead.