[austin362667]

• ☑️ I understand it is strictly prohibited to use AI to write PR description
• AI usage disclosure: Reviewed by GPT 6.1 Sol High

Why?

When the target logit dominates, true gradient rounds to zero and the sum of the gradients is not.

import mlx.core as mx
import mlx.nn as nn

z = mx.array([9.0, -9.0, -9.0, -9.0])
t = mx.array([1.0, 0.0, 0.0, 0.0]) # target as probs

loss, grads = mx.value_and_grad(
    lambda z: nn.losses.cross_entropy(z, t, weights=mx.array(1.0))
)(z)

print(loss) # array(0, dtype=float32)
print(grads) # array([0, 1.523e-08, 1.523e-08, 1.523e-08], dtype=float32)

It happens in softmax, log-softmax, and cross-entropy.
This numerical bug is called Softmax Collapse and paper proposed a neat fix.

For cross-entrop loss L, avoid 1 + tiny fractions that rounds to zero by separating the 1 (1 = e^0 = shifted true label logits) from others during logsumexp.

What?

MLX non-fast mode has two paths for computing cross entropy loss.

1. Targets are probabilities
   Apply log-softmax fix as introduced in the paper. Ref JAX code:

   import jax, jax.numpy as jnp, optax
   
   def log_softmax(z):
     i_max = jnp.argmax(z, axis=-1)
     z_max = jnp.take_along_axis(z, i_max[..., None], axis=-1)  # Max logit.
     is_max = jax.nn.one_hot(i_max, z.shape[-1], dtype=bool)    # Max mask.
     shifted = jnp.where(is_max, 0.0, z - z_max)                # Exact 0 at argmax.
     exp_others = jnp.where(is_max, 0.0, jnp.exp(shifted))      # Non-max exp.
     s = jnp.sum(exp_others, axis=-1, keepdims=True)            # Sum non-max exp.
     return shifted - jnp.log1p(s)                              # log_softmax(z).

2. Targets are class indices
   Apply hard-target cross-entropy fix as introduced in the paper. Ref JAX code:

   import jax, jax.numpy as jnp, optax
   
   def cross_entropy(z, t):
     z_t = jnp.take_along_axis(z, t[..., None], axis=-1)        # Target logit.
     is_t = jax.nn.one_hot(t, z.shape[-1], dtype=bool)          # Target mask.
     shifted = jnp.where(is_t, 0.0, z - z_t)                    # Exact 0 at target.
     b = jnp.max(shifted, axis=-1, keepdims=True)               # Max shift >= 0.
     s = jnp.sum(jnp.where(is_t, 0.0, jnp.exp(shifted - b)), axis=-1)
     return b[..., 0] + jnp.log1p(s + jnp.expm1(-b[..., 0]))    # -log_softmax(z)_t.

Verification?

The public interface is unchanged. The loss and gradients are correct.

print(loss) # array(4.56899e-08, dtype=float32)
print(grads) # array([-4.56899e-08, 1.523e-08, 1.523e-08, 1.523e-08], dtype=float32)

Todo?

We need numerically correct implementations in fast mode (i.e., all the related Metal and CUDA kernels)