A Fix for Softmax and Cross-Entropy

Standard implementations of softmax, log-softmax, and cross-entropy have two numerical bugs. The first bug, Softmax Collapse, rounds the target logit's gradient to zero if that one logit is much larger than the others, resulting in gradients that no longer sum to zero. The second bug, Softmax Cancellation, occurs when student logits are close to teacher logits, and produces non-zero gradients even when the student and teacher logits are equal. We present simple implementations of log-softmax, softmax, and cross-entropy that fix both bugs. Our fix eliminates the previously observed “slingshot” phenomenon, speeds up grokking with weight decay, prevents accuracy collapse without weight decay, stops student networks from drifting away from the teacher during distillation, and shows that after memorization, logit growth is bounded by an equilibrium between weight decay and Adam's (\epsilon) hyperparameter.

Naively computing softmax or cross-entropy on logits (\mathbf{z} \in \mathbb{R}^{K}) as (e^{z_k} / \sum_j e^{z_j}) overflows if any logit exceeds the floating-point exponent limit ((\approx 88) in fp32), which is why standard implementations in all deep learning frameworks (JAX, PyTorch, TensorFlow, etc.) shift the input logits by subtracting their maximum value (z_{\max} = \max_j z_j) (Higham, 2002; Blanchard et al., 2021):

For a target distribution (\mathbf{p} \in \Delta^{K-1}), cross-entropy (CE) is (\mathcal{L}(\mathbf{z}, \mathbf{p}) = -\sum_k p_k,\operatorname{logsoftmax}(\mathbf{z})k), with exact gradient (\mathbf{g} = \nabla{\mathbf{z}} \mathcal{L}(\mathbf{z}, \mathbf{p}) = \mathbf{q} - \mathbf{p}). Though subtracting (z_{\max}) prevents overflow, two numerical bugs remain.

The first bug, Softmax Collapse (Figure 1a, Prieto et al., 2025), occurs when a model correctly predicts the target class (t \in {1, \dots, K}) with high confidence, so that (z_{t} = z_{\max}) is much larger than all non-target logits (z_j) ((j \neq t)). We measure this confidence by the logit margin, which is how much larger the target logit is than the log-sum-exp of all non-target logits:

In Equation 1, subtracting (z_{\max} = z_{t}) sets the shifted target logit to (\tilde{z}{t} = 0) (so (e^{\tilde{z}{t}} = 1) in the denominator), while the (K

• 1) non-target exponentials sum to a small positive number (e^{-m} = \sum_{j \neq t} e^{z_j - z_{t}} \ll 1). Floating-point numbers have much finer precision near (0) than near (1): in fp32, values near (0) can be as small as ({\sim}10^{-38}), whereas values added to (1.0) are lost to rounding below ({\sim}10^{-7}). Because standard implementations form the denominator (1 + e^{-m}) by adding (e^{-m}) to (1.0), the (e^{-m}) term is lost to rounding once the margin (m) exceeds a threshold (c_{d}) that we call the precision cliff ((c_{16} \approx 5.5) in bf16, (c_{32} \approx 16.6) in fp32, and (c_{64} \approx 36.7) in fp64). When (m > c_{d}), the denominator (1 + e^{-m}) evaluates to exactly (1.0), which causes (q_{t} = 1/(1 + e^{-m})) to evaluate to (1.0) and zeroes out both the loss and the target gradient (\tilde{g}{t} = q{t} - 1 = 0.0). However, the non-target gradients (\tilde{g}j = q_j = e^{z_j - z{t}} / (1 + e^{-m})) ((j \neq t)) are evaluated near (0) without adding (1.0) to their numerators, so they remain positive all the way down to exponent underflow ((m \approx 88)). Because the target gradient rounds to (0.0) while the non-target gradients do not, the gradients no longer sum to zero ((\sum_k \tilde{g}_k > 0)). When training with adaptive optimizers such as Adam and AdamW (Kingma & Ba, 2015; Loshchilov & Hutter, 2019), which normalize each gradient by a moving average of its magnitude, this asymmetry causes the recurring “slingshot” loss spikes observed by Thilak et al. (2022) and the grokking failures studied by Prieto et al. (2025). While prior works used surrogate losses (Prieto et al., 2025), gradient projection (Liu et al., 2026), or fp64 casting (Nanda et al., 2023; Varma et al., 2023) to ameliorate these issues, we instead fix log-softmax, softmax, and cross-entropy directly (Section 2).

The second bug, Softmax Cancellation (Figure 1b), occurs in knowledge distillation (Hinton et al., 2015) and KL regularization, where the target (\mathbf{p} = \operatorname{softmax}(\mathbf{z}{\mathrm{T}})) comes from teacher logits (\mathbf{z}{\mathrm{T}}). Standard implementations compute (\mathbf{q} = \operatorname{softmax}(\mathbf{z}{\mathrm{S}})) and (\mathbf{p} = \operatorname{softmax}(\mathbf{z}{\mathrm{T}})) separately and subtract them to get (\tilde{\mathbf{g}} = \mathbf{q} - \mathbf{p}). This means that if (\mathbf{z}{\mathrm{S}}) is close to (\mathbf{z}{\mathrm{T}}), subtracting these nearly equal probabilities loses precision, and slight differences between the forward-pass softmax and backward-pass log-softmax formulas leave a non-zero gradient even when (\mathbf{z}{\mathrm{S}} = \mathbf{z}{\mathrm{T}}). We fix this by computing (\mathbf{q} - \mathbf{p}) directly from (\mathbf{z}{\mathrm{S}} - \mathbf{z}{\mathrm{T}}) in log space (Section 3).

When the target logit (z_{t} = z_{\max}) is the largest logit, shifting by (z_{\max}) in Equation 1 sets (\tilde{z}{t} = 0) and (\tilde{z}j = z_j - z{t} < 0) for (j \neq t), so the denominator is (\sum_i e^{\tilde{z}i} = 1 + \sum{j \neq t} e^{z_j - z{t}} = 1 + e^{-m}). Mathematically, the hard-target loss is (\mathcal{L}(\mathbf{z}, t) = -\operatorname{logsoftmax}(\mathbf{z})_{t} = \log(1 + e^{-m})), with gradients

which sum to zero ((\sum_k g_k = 0)). In floating point, however, the two expressions for (g_{t}) in Equation 3 behave very differently once (m > c_{d}), where (e^{-m}) is too small to change (1.0) when added to it (Figure 1a). Standard implementations add (1) and (e^{-m}) to form the denominator (1 + e^{-m}), which rounds to (1) when (m > c_{d}). In the forward pass, (\log(1 + e^{-m})) then evaluates to (\log(1) = 0), zeroing out both (\operatorname{logsoftmax}(\mathbf{z}){t}) and (\mathcal{L}(\mathbf{z}, t)). In the backward pass, standard cross-entropy implementations compute (q{t} = 1/(1 + e^{-m}) = 1) and subtract (1) to get (\tilde{g}{t} = q{t} - 1 = 0). In contrast, our fixed expression (-e^{-m}/(1 + e^{-m})) keeps (e^{-m}) in the numerator: even when the denominator (1 + e^{-m}) rounds to (1.0), dividing (-e^{-m}) by (1.0) preserves (g_{t} \approx -e^{-m}) down to exponent underflow ((m \approx 88)).

While (\tilde{g}{t}) rounds to zero when (m > c{d}), the non-target gradients do not. For each non-target class (j \neq t), standard cross-entropy computes (\tilde{g}j = q_j = e^{z_j - z{t}} / (1 + e^{-m})). Because the numerator (e^{z_j - z_{t}}) is evaluated directly from the negative difference (z_j - z_{t}) without adding (1), it remains accurate down to exponent underflow while the denominator rounds to (1). When (m > c_{d}), standard cross-entropy therefore zeroes out only the target entry (t) while all (K - 1) non-target gradients remain positive:

The same bug affects standalone (\operatorname{softmax}(\mathbf{z})) (Qiu & Yao, 2026): once (m > c_{d}), its backward pass rounds (q_{i_{\max}}) ((i_{\max} = \arg\max_i z_i)) to (1) and zeroes (\tilde{g}{i{\max}}) while the other gradients stay non-zero.

This bug can be fixed by never adding (1) to (e^{-m}) directly. In (\operatorname{logsoftmax}(\mathbf{z})) and (\operatorname{softmax}(\mathbf{z})), only the largest logit (z_{i_{\max}}) is shifted to (0) (where (e^0 = 1) in the denominator). Instead of including (e^0 = 1) in the sum of exponentials, we set (\tilde{z}{i{\max}} = 0), compute the margin (m) by excluding (i_{\max}) from the log-sum-exp (Blanchard et al., 2021), and combine them with a numerically stable (\operatorname{softplus}(x) = \log(1 + e^x)) (with (\operatorname{softmax}(\mathbf{z})_k = \exp(\operatorname{logsoftmax}(\mathbf{z})_k))):

For hard-target cross-entropy with target (t), we can use (t) directly in place of (i_{\max}) without an (\arg\max):

Because the (t)-th slot is replaced with the constant (0) rather than (z_{t} - z_{t}), (z_{t}) only appears in (\mathcal{L}) through the (K - 1) differences (z_j - z_{t}) ((j \neq t)), so autodiff correctly computes the target gradient as the negated sum of the non-target gradients:

Here is a JAX implementation of log-softmax, softmax, and hard-target cross-entropy, which remains accurate down to exponent underflow and satisfies (\sum_k g_k = 0):

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. others = jnp.where(is_max, -jnp.inf, shifted) # Exclude argmax from sum. m = -jax.nn.logsumexp(others, axis=-1, keepdims=True) # Logit margin. return shifted - jax.nn.softplus(-m) # log_softmax(z).

softmax = lambda z: jnp.exp(log_softmax(z))

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. others = jnp.where(is_t, -jnp.inf, z - z_t) # Exclude target from sum. m = -jax.nn.logsumexp(others, axis=-1) # Logit margin. return jax.nn.softplus(-m) # -log_softmax(z)_t.

As an example, consider (\mathbf{z} = [9.0, -9.0, -9.0]) and (t = 0). The standard jax.nn and optax implementations return zero for the target log-probability, loss, and target gradient, while our functions return accurate values whose gradients sum to zero:

z, t = jnp.array([9.0, -9.0, -9.0], dtype=jnp.float32), jnp.array(0) std_ce = optax.softmax_cross_entropy_with_integer_labels jax.nn.log_softmax(z) # -> [ 0.0, -18.0, -18.0 ] jax.grad(lambda z: jax.nn.softmax(z)[0])(z) # -> [ 0.0, -1.523e-8, -1.523e-8] jax.value_and_grad(std_ce)(z, t) # -> 0.0, [ 0.0, 1.523e-8, 1.523e-8] log_softmax(z) # -> [-3.046e-8, -18.0, -18.0 ] jax.grad(lambda z: softmax(z)[0])(z) # -> [ 3.046e-8, -1.523e-8, -1.523e-8] jax.value_and_grad(cross_entropy)(z, t) # -> 3.046e-8, [-3.046e-8, 1.523e-8, 1.523e-8]

3. Softmax Cancellation#

Given student logits (\mathbf{z}{\mathrm{S}} \in \mathbb{R}^{K}) and teacher logits (\mathbf{z}{\mathrm{T}} \in \mathbb{R}^{K}), standard soft-target cross-entropy computes (\mathbf{p} = \operatorname{softmax}(\mathbf{z}_{\mathrm{T}})) and evaluates

As the student logits approach the teacher logits ((\mathbf{z}{\mathrm{S}} \to \mathbf{z}{\mathrm{T}})), the true gradient (\mathbf{g} = \mathbf{q} - \mathbf{p}) shrinks toward (\mathbf{0}) and equals (\mathbf{0}) when (\mathbf{z}{\mathrm{S}} = \mathbf{z}{\mathrm{T}}). In floating point, however, computing (\mathbf{q}) and (\mathbf{p}) separately and subtracting them (Equation 8) loses precision when (\mathbf{z}{\mathrm{S}}) is close to (\mathbf{z}{\mathrm{T}}), because the small logit difference (\boldsymbol{\delta} = \mathbf{z}{\mathrm{S}} - \mathbf{z}{\mathrm{T}}) is lost when (\mathbf{q}) and (\mathbf{p}) are rounded to probabilities before being subtracted (Figure 1b). Furthermore, because (\mathbf{p} = \operatorname{softmax}(\mathbf{z}{\mathrm{T}})) in the forward pass and (\mathbf{q} = \exp(\operatorname{logsoftmax}(\mathbf{z}{\mathrm{S}}))) in the backward pass are evaluated using slightly different sequences of operations, (\mathbf{q}) and (\mathbf{p}) do not match to machine precision even when (\mathbf{z}{\mathrm{S}} = \mathbf{z}{\mathrm{T}}), leaving a non-zero gradient (\tilde{\mathbf{g}} = \mathbf{q} - \mathbf{p} \neq \mathbf{0}).

To avoid subtracting nearly equal probabilities, we pass (\mathbf{z}_{\mathrm{T}}) directly to our