A Fix for Softmax and Cross‑Entropy

1. Introduction#

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.

(a) Softmax Collapse: $|g_{t}|$ vs. margin $m$

(b) Softmax Cancellation: $\|\nabla_{\mathbf{z}_{\mathrm{S}}}\mathcal{L}\|_{2}$ as $\mathbf{z}_{\mathrm{S}}\to\mathbf{z}_{\mathrm{T}}$

The first bug, identified by Prieto et al. (2025) as Softmax Collapse (Figure 1a), 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 will 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 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 (Thilak et al., 2022; Liu et al., 2026) and stalls grokking (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, which we will call 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}$. When $\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).

2. Softmax Collapse#

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 form the denominator $1+e^{-m}$, which rounds to $1$ when $m>c_{d}$ (Prieto et al., 2025). In the forward pass, $\log(1+e^{-m})$ evaluates to $\log(1)=0$, so both $\operatorname{logsoftmax}(\mathbf{z})_{t}$ and $\mathcal{L}(\mathbf{z},t)$ round to zero. In the backward pass, standard 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 $z_{j}-z_{t}<0$ without adding $1$, it remains accurate down to exponent underflow while the denominator rounds to $1$. When $m>c_{d}$, standard cross-entropy zeroes out only the target entry $t$ while all $K-1$ non-target gradients remain positive (Liu et al., 2026):

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 evaluates $\log(1+e^{-m})$ via $\operatorname{log1p}$, 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.
  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).

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.
  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.

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 loss function and compute $\mathbf{q}-\mathbf{p}$ in log space using $\operatorname{expm1}(x)=e^{x}-1$ and $\operatorname{log1p}(x)=\log(1+x)$. Given $\boldsymbol{\delta}=\mathbf{z}_{\mathrm{S}}-\mathbf{z}_{\mathrm{T}}$, we center $\boldsymbol{\delta}$ under $\mathbf{p}$ and write the log-ratio $s_{k}=\log(q_{k}/p_{k})$ and difference $u_{k}=q_{k}-p_{k}$ as:

Because $\sum_{i}p_{i}d_{i}=0$ and $\operatorname{expm1}(x)\geq x$, the input to $\operatorname{log1p}$ is always non-negative. When $|s_{k}|\geq 1$, we fall back to $u_{k}=q_{k}-p_{k}$ to avoid $0.0\times\infty$ if $p_{k}$ rounds to $0.0$ while $\operatorname{expm1}$ overflows, and we set $g_{k}=u_{k}-q_{k}\sum_{i}u_{i}$ so that $\sum_{k}g_{k}=0$. Because Equation 9 avoids adding or subtracting $1$, $\mathbf{g}$ remains accurate as $\mathbf{z}_{\mathrm{S}}\to\mathbf{z}_{\mathrm{T}}$ and is $\mathbf{0}$ when $\mathbf{z}_{\mathrm{S}}=\mathbf{z}_{\mathrm{T}}$. Here is a JAX implementation of soft-target cross-entropy (using log_softmax and softmax from Section 2), which uses jax.lax.stop_gradient (sg) to return the forward loss while substituting $\mathbf{g}$ in the backward pass:

def soft_target_cross_entropy(z_S, z_T):
  log_q = log_softmax(z_S)
  q, p = jnp.exp(log_q), softmax(z_T)
  delta = z_S - z_T                                                    # Logit difference.
  d = delta - jnp.sum(p * delta, axis=-1, keepdims=True)               # Center under p.
  s = d - jnp.log1p(jnp.sum(p * jnp.expm1(d), axis=-1, keepdims=True)) # log(q / p).
  u = jnp.where(jnp.abs(s) < 1, p * jnp.expm1(s), q - p)               # Difference q - p.
  g = u - q * jnp.sum(u, axis=-1, keepdims=True)                       # Ensure sum(g) = 0.
  sg = jax.lax.stop_gradient
  loss = -jnp.sum(p * sg(log_q), axis=-1)                              # Forward loss.
  return loss + jnp.sum(sg(g) * (z_S - sg(z_S)), axis=-1)              # Custom gradient g.

As an example, let $\mathbf{z}_{\mathrm{S}}$ be a small fp32 perturbation of $\mathbf{z}_{\mathrm{T}}$. Standard soft-target cross-entropy shrinks the second gradient entry by nearly $3\times$, flips the sign of the third, and fails to sum to zero, while our implementation matches the exact fp64 gradient:

z_T = jnp.array([2.0, 1.0, 0.0], dtype=jnp.float32)
z_S = z_T + jnp.array([-1e-7, 1e-7, 0.0], dtype=jnp.float32)
std_soft_ce = lambda z_S, z_T: -jnp.sum(jax.nn.softmax(z_T) * jax.nn.log_softmax(z_S))
jax.grad(std_soft_ce)(z_S, z_T)               # -> [-5.960e-8, 1.452e-8, -1.123e-9]
jax.grad(soft_target_cross_entropy)(z_S, z_T) # -> [-4.596e-8, 4.144e-8,  4.513e-9]
# Exact (float64) gradient:                        [-4.596e-8, 4.144e-8,  4.513e-9]

4. The Slingshot Mechanism#

In grokking (Power et al., 2022), a network quickly memorizes the training set ($q_{t}\to 1$) and generalizes much later. While studying grokking on modular addition with Adam (Kingma & Ba, 2015; Loshchilov & Hutter, 2019), Thilak et al. (2022) observed the slingshot mechanism, where training after memorization alternates between flat loss plateaus near zero and sudden loss spikes that coincide with jumps in test accuracy.

These slingshot spikes are a direct consequence of Softmax Collapse (Section 2, Prieto et al. (2025); Liu et al. (2026)). Recall that AdamW updates parameters $\boldsymbol{\theta}$ by subtracting $\eta\big{(}\hat{\mathbf{m}}/(\sqrt{\hat{\mathbf{v}}}+\epsilon)+\lambda\boldsymbol{\theta}\big{)}$, where $\hat{\mathbf{m}}$ and $\hat{\mathbf{v}}$ are bias-corrected moving averages of the parameter gradient $\mathbf{g}=\nabla_{\boldsymbol{\theta}}\mathcal{L}(\boldsymbol{\theta})$ and its square (with decay rates $\beta_{1},\beta_{2}$). After memorization, the logit margins $m$ grow as the network pushes its target probabilities toward $1$. Once an example’s margin exceeds the precision cliff $c_{d}$, standard cross-entropy rounds its target gradient $\tilde{g}_{t}$ to zero while its non-target gradients $\tilde{g}_{j}>0$ ($j\neq t$) remain positive. When all examples in a batch sit above $c_{d}$, the parameter gradients nearly vanish and Adam’s second moment $\hat{v}_{i}$ decays exponentially toward zero while weight decay slowly shrinks the parameters. Eventually, weight decay or the residual non-target updates pull a few examples back below $c_{d}$, where their target gradients $\tilde{g}_{t}\approx-e^{-m}$ abruptly turn back on. Because $\sqrt{\hat{v}_{i}}$ decayed while those examples were above $c_{d}$, Adam divides the restored gradients by a very small denominator $\sqrt{\hat{v}_{i}}+\epsilon$, producing a large parameter step that spikes the training loss, resets $\hat{v}_{i}$, and repeats the cycle. With our fixed cross-entropy, the target-logit gradient $g_{t}=-e^{-m}/(1+e^{-m})$ does not round to zero at $c_{d}$, so $\hat{m}_{i}$ and $\sqrt{\hat{v}_{i}}$ decay smoothly together and slingshot spikes disappear (Figure 2).

(a) Mean logit margin $m$

(b) Training loss

(c) Test accuracy

Whether slingshot spikes appear to help or hurt generalization depends on how the network’s weight norm is controlled (Figure 2). When a network uses LayerNorm before the output layer and no weight decay (Thilak et al., 2022; Prieto et al., 2025), growing the pre-LayerNorm weights does not change the logits, but it does shrink the effective step size $\eta/\|\boldsymbol{\theta}\|_{2}$; each slingshot spike inflates $\|\boldsymbol{\theta}\|_{2}$ and acts as an accidental learning-rate decay. When weight norms are already controlled by weight decay (Varma et al., 2023), slingshot spikes at $c_{32}$ only slow down and hurt grokking ($94\pm 5\%$ vs. $100\pm 0\%$ test accuracy on $a+b\bmod 31$; Figure 2, top). And when neither LayerNorm nor weight decay is used (Doshi et al., 2024), a single slingshot spike at $c_{32}$ is destructive: on $a^{2}+b^{2}\bmod 97$ ($\lambda=0$, $\epsilon=10^{-16}$), standard cross-entropy tracks fixed cross-entropy up to $61\%$ test accuracy until $m$ hits $c_{32}$ around step $1250$, where a spike drops test accuracy to $23\%$ (and from $60\%$ to $21\%$ on $a^{3}+b\bmod 97$), while fixed cross-entropy continues smoothly to $85\%$ (Figure 2, bottom).

5. Margin Equilibria#

In Figure 2a, our fixed cross-entropy eliminates the precision cliff at $c_{32}\approx 16.6$, yet with weight decay ($\lambda=1$, top) its margin still saturates just above $c_{32}$ near $18$, whereas without weight decay ($\lambda=0$, bottom) its margin continues to climb. This plateau is the equilibrium where weight decay balances AdamW’s normalized gradient step. At a stationary point where the parameters stop changing, the moving averages converge to $\hat{m}_{i}=g_{i}$ and $\hat{v}_{i}=g_{i}^{2}$ (so $\sqrt{\hat{v}_{i}}=|g_{i}|$), and setting the update to zero gives:

Here $\eta$, $\beta_{1}$, and $\beta_{2}$ cancel out, so RMSProp (Tieleman & Hinton, 2012) and LaProp (Ziyin et al., 2020) share the same stationary point. After memorization, the gradients shrink exponentially with the margin ($|g_{i}|\propto e^{-m}$). If $\epsilon=0$, the normalized gradient $g_{i}/|g_{i}|=\mathrm{sign}(g_{i})$ has a magnitude of $1$ no matter how small $|g_{i}|$ becomes, which causes the margin to grow until the gradients underflow. When $\epsilon>0$, however, the normalized gradient has magnitude $\approx 1$ while $|g_{i}|\gg\epsilon$, and shrinks as $g_{i}/\epsilon\propto e^{-m}$ once $|g_{i}|\ll\epsilon$ (Figure 3). Balancing this shrinking step against weight decay $\lambda\theta_{i}$ requires $|g_{i}|\approx\epsilon\lambda|\theta_{i}|$ (Wortsman et al., 2024; Molybog et al., 2023; Everett et al., 2024; Kosson et al., 2024), which acts like an $L_{2}$ regularizer with weight $\epsilon\lambda$ (Rosset et al., 2004; Soudry et al., 2018). Since $|g_{i}|\propto e^{-m}$, the margin stops growing at:

Equation 11 explains why the default $\epsilon=10^{-8}$ causes slingshot spikes in fp32: its equilibrium $\log(\nicefrac{{1}}{{\epsilon}})\approx 18.4$ sits just above $c_{32}\approx 16.6$ (Figure 3), so standard cross-entropy hits the cliff before equilibrium (whereas for $\epsilon\geq 10^{-7}$, $\log(\nicefrac{{1}}{{\epsilon}})\leq 16.1<c_{32}$ and standard cross-entropy settles without spiking). To test Equation 11 across many orders of magnitude down to $\epsilon=0$ in fp32 (Figure 4), we train with our fixed cross-entropy and LaProp on $a+b\bmod 31$, updating $\sqrt{\hat{\mathbf{v}}}$ directly without squaring $g_{i}$ (which underflows in fp32 once $|g_{i}|<2^{-63}$, or $m\approx 38$). With weight decay ($\lambda=1$, Figure 4a), the margin for each $\epsilon>0$ climbs along the $\epsilon=0$ trajectory and levels off at $m^{\star}(\epsilon)=\log(\nicefrac{{1}}{{\epsilon}})$, while $\epsilon=0$ grows until the fp32 gradients underflow at $m\approx 81.2$. Without weight decay ($\lambda=0$, Figure 4b), nothing balances the $g_{i}/\epsilon$ step, so margins keep rising past $\log(\nicefrac{{1}}{{\epsilon}})$.

(a) Weight decay ($\lambda=1$)

(b) No weight decay ($\lambda=0$)

6. Grokking Across Precisions and Benchmarks#

To see how the precision cliff behaves across floating-point formats, we follow Varma et al. (2023) and Lyle et al. (2025) and train a 1-layer, 4-head, width-$64$ ReLU Transformer without LayerNorm on $a+b\bmod 31$ ($50\%$ training split) for $10{,}000$ epochs of full-batch AdamW ($\eta=10^{-3}$, $\beta_{1}=0.9$, $\beta_{2}=0.98$, $\lambda=1$), varying the loss (standard vs. fixed cross-entropy), logit precision at the loss (bf16, fp32, fp64), and $\epsilon\in\{10^{-8},10^{-12},10^{-16}\}$ (Figure 5, Table 1). At $\epsilon=10^{-8}$ (Figure 5, top), standard cross-entropy stalls at $c_{16}\approx 5.5$ in bf16 and repeatedly spikes at $c_{32}\approx 16.6$ in fp32, whereas fixed cross-entropy crosses both cliffs smoothly to $\log(\nicefrac{{1}}{{\epsilon}})\approx 18.4$ and reaches $100\%$ test accuracy in all three precisions. Reducing $\epsilon$ to $10^{-16}$ (Figure 5, bottom) raises the equilibrium to $\log(\nicefrac{{1}}{{\epsilon}})\approx 36.8>c_{64}$, breaking even fp64 standard cross-entropy at $c_{64}\approx 36.7$ while fixed cross-entropy converges smoothly across all precisions.

Table 1 summarizes these results and compares against two recent methods that modify the gradients or loss to mitigate Softmax Collapse. Liu et al. (2026) subtract the mean logit gradient from each entry so the gradients sum to zero (“Zero-sum Proj.”), which eliminates slingshot spikes. However, once $m>c_{32}$ and $\tilde{g}_{t}=0$, subtracting the mean only gives the target logit $-\nicefrac{{1}}{{K}}$ of the total non-target gradient, stalling $m$ near $18$ and dropping test accuracy below $50\%$ for $\epsilon\leq 10^{-12}$. Prieto et al. (2025) propose “StableMax,” which replaces $\exp(x)$ in softmax with a piecewise rational function ($x+1$ for $x\geq 0$ and $1/(1-x)$ for $x<0$). Because this function grows linearly instead of exponentially, logit margins blow up to ${\sim}10^{4}$ and test accuracy lags our fixed cross-entropy across all $\epsilon$.

We observe the same behavior across four grokking benchmarks at $\epsilon=10^{-12}$ (Figure 6): 1-layer Transformers on $a+b\bmod 113$ (Nanda et al., 2023) and $a+b\bmod 31$ (Varma et al., 2023) with $\lambda=1$, and 2-layer MLPs on $a^{2}+b^{2}\bmod 97$ and $a^{3}+b\bmod 97$ (Doshi et al., 2024) with $\lambda=0.2$. On all four tasks, fp32 standard cross-entropy spikes once it hits $c_{32}$, whereas fp32 fixed cross-entropy matches fp64 standard cross-entropy and reaches $100\%$ test accuracy.

(a) Mean logit margin $m$

(b) Training loss

(c) Test accuracy

(a) $a{+}b\bmod 113$
Nanda et al. (2023)

(b) $a{+}b\bmod 31$
Varma et al. (2023)

(c) $a^{2}{+}b^{2}\bmod 97$
Doshi et al. (2024)

(d) $a^{3}{+}b\bmod 97$
Doshi et al. (2024)

7. Soft-Target Distillation#

To see how Softmax Cancellation (Section 3) affects optimization, we initialize a 2-layer ReLU MLP student with the fp32 weights of a teacher trained to $100\%$ accuracy on $a+b\bmod 53$ ($\boldsymbol{\theta}_{\mathrm{S}}^{(0)}=\boldsymbol{\theta}_{\mathrm{T}}$), and run $2{,}000$ full-batch steps of Adam ($\eta=2\times 10^{-3}$, $\beta_{1}=0.9$, $\beta_{2}=0.98$, $\lambda=0$, $\epsilon=10^{-8}$) on the $50\%$ training split at temperature $T=4$ (Hinton et al., 2015) (which keeps $m/T<c_{16}$ so $\mathbf{p}$ does not round to one-hot), computing the forward pass and loss in bf16, fp32, or fp64. Because $\mathbf{z}_{\mathrm{S}}=\mathbf{z}_{\mathrm{T}}$ initially, the true gradient is zero and the student should not move. However, standard soft-target cross-entropy has a non-zero residual gradient at $\mathbf{z}_{\mathrm{S}}=\mathbf{z}_{\mathrm{T}}$ due to the mismatch between its forward softmax and backward log-softmax (Figure 1b). Because Adam divides by $\sqrt{\hat{\mathbf{v}}}+\epsilon$, it amplifies even tiny residual gradients (Figure 7): in bf16 (${\sim}10^{-3}$) it drives the student away from the teacher and degrades test accuracy, in fp32 (${\sim}10^{-8}$) it causes an initial $\mathrm{KL}$ spike and a ${\sim}10^{-6}$ floor, and in fp64 (${\sim}10^{-17}$) the update is too small to change the fp32 weights. Our fixed soft-target cross-entropy returns $\mathbf{g}=\mathbf{0}$ at $\mathbf{z}_{\mathrm{S}}=\mathbf{z}_{\mathrm{T}}$, so the student stays at $\mathrm{KL}=0$ and $100\%$ accuracy across all precisions.

(a) Held-out $\mathrm{KL}(\mathbf{p}_{\mathrm{T}}\parallel\mathbf{q}_{\mathrm{S}})$

(b) Test accuracy

8. Conclusion#

We have presented simple fixes for two numerical bugs in standard implementations of softmax, log-softmax, and cross-entropy: Softmax Collapse, which zeroes out the dominant logit’s gradient at high confidence while leaving non-dominant gradients non-zero, and Softmax Cancellation, which loses precision when student logits approach teacher logits in soft-target distillation. Rewriting these functions in log space using $\operatorname{softplus}$, $\operatorname{expm1}$, and $\operatorname{log1p}$ restores accurate, zero-sum gradients across bf16, fp32, and fp64 without internally casting to higher precision. Fixing these bugs eliminates slingshot loss spikes during grokking, prevents student networks from drifting away from the teacher during distillation, and reveals that logit margins under AdamW and weight decay naturally converge to $\log(\nicefrac{{1}}{{\epsilon}})$. Though the extreme margins and near-zero student–teacher gaps probed by our experiments are uncommon in standard large-scale training, our fixed functions are easy drop-in replacements that can be implemented without meaningful slowdown, and may prove useful whenever models make high-confidence predictions or distill from close teachers.

Acknowledgments: We thank Jean-Baptiste Alayrac, Forrester Cole, Roy Frostig, Mark Hamilton, Muhammad Maaz, Miki Rubinstein, Keith Rush, and Brandon Zhao for their helpful advice.

Errata: In an earlier version of this paper, the code in Section 2 produced $\mathrm{NaN}$ gradients when all but one logit was effectively $-\infty$, as happens with masking in causal attention or with extreme logits. This bug was not present in our internal codebase, which produced all results in this paper. It appears to have been introduced by an over-eager agent that simplified our code for this paper, likely in an attempt to match Equations 5 and 6 (with $s=\sum_{j\neq i_{\max}}e^{\tilde{z}_{j}}$, $\operatorname{softplus}(-m)$ can be written as $\operatorname{log1p}(s)$ or equivalently as $\operatorname{softplus}(\log s)$, but the latter makes autodiff differentiate through $\log 0$ when $s=0$). We thank Maaz (2026) for identifying this issue.

References#

1. [Blanchard et al., 2021] Pierre Blanchard, Desmond J. Higham, and Nicholas J. Higham. Accurately computing the log-sum-exp and softmax functions. IMA Journal of Numerical Analysis, 2021.
2. [Doshi et al., 2024] Darshil Doshi, Tianyu He, Aritra Das, and Andrey Gromov. Grokking modular polynomials. arXiv:2406.03495, 2024.
3. [Everett et al., 2024] Katie Everett, Lechao Xiao, Mitchell Wortsman, Alexander A. Alemi, Roman Novak, Peter J. Liu, Izzeddin Gur, Jascha Sohl-Dickstein, Leslie Pack Kaelbling, Jaehoon Lee, and Jeffrey Pennington. Scaling exponents across parameterizations and optimizers. ICML, 2024.
4. [Higham, 2002] Nicholas J. Higham. Accuracy and stability of numerical algorithms. SIAM, 2nd edition, 2002.
5. [Hinton et al., 2015] Geoffrey Hinton, Oriol Vinyals, and Jeff Dean. Distilling the knowledge in a neural network. arXiv:1503.02531, 2015.
6. [Kingma & Ba, 2015] Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization. ICLR, 2015.
7. [Kosson et al., 2024] Atli Kosson, Bettina Messmer, and Martin Jaggi. Rotational equilibrium: How weight decay balances learning across neural networks. ICML, 2024.
8. [Liu et al., 2026] Hanqing Liu, Jianjun Cao, Yuanze Li, and Zijian Zhou. Grokking or glitching? How low-precision drives Slingshot loss spikes. arXiv:2605.06152, 2026.
9. [Loshchilov & Hutter, 2019] Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. ICLR, 2019.
10. [Lyle et al., 2025] Clare Lyle, Ghada Sokar, András György, and Razvan Pascanu. What can grokking teach us about learning under non-stationarity? Conference on Lifelong Learning Agents, 2025.
11. [Maaz, 2026] Muhammad Maaz. Property-based testing for robust softmax and cross-entropy. arXiv:2610.36427, 2026.
12. [Molybog et al., 2023] Igor Molybog, Peter Albert, Moya Chen, Zachary DeVito, David Esiobu, Naman Goyal, Punit Singh Koura, Sharan Narang, Andrew Poulton, Ruan Silva, et al. A theory on Adam instability in large-scale machine learning. arXiv:2304.09871, 2023.
13. [Nanda et al., 2023] Neel Nanda, Lawrence Chan, Tom Lieberum, Jess Smith, and Jacob Steinhardt. Progress measures for grokking via mechanistic interpretability. ICLR, 2023.
14. [Power et al., 2022] Alethea Power, Yuri Burda, Harri Edwards, Igor Babuschkin, and Vedant Misra. Grokking: Generalization beyond overfitting on small algorithmic datasets. arXiv:2201.02177, 2022.
15. [Prieto et al., 2025] Lucas Prieto, Melih Barsbey, Pedro A. M. Mediano, and Tolga Birdal. Grokking at the edge of numerical stability. ICLR, 2025.
16. [Qiu & Yao, 2026] Haiquan Qiu and Quanming Yao. Why low-precision Transformer training fails: An analysis on Flash Attention. ICLR, 2026.
17. [Rosset et al., 2004] Saharon Rosset, Ji Zhu, and Trevor Hastie. Margin maximizing loss functions. NeurIPS, 2004.
18. [Soudry et al., 2018] Daniel Soudry, Elad Hoffer, Mor Shpigel Nacson, Suriya Gunasekar, and Nathan Srebro. The implicit bias of gradient descent on separable data. JMLR, 2018.
19. [Thilak et al., 2022] Vimal Thilak, Etai Littwin, Shuangfei Zhai, Omid Saremi, Roni Paiss, and Joshua Susskind. The Slingshot mechanism: An empirical study of adaptive optimizers and the grokking phenomenon. arXiv:2206.04817, 2022.
20. [Tieleman & Hinton, 2012] Tijmen Tieleman and Geoffrey Hinton. Lecture 6.5—RMSProp: Divide the gradient by a running average of its recent magnitude. COURSERA: Neural Networks for Machine Learning, 2012.
21. [Varma et al., 2023] Vikrant Varma, Rohin Shah, Zachary Kenton, János Kramár, and Ramana Kumar. Explaining grokking through circuit efficiency. arXiv:2309.02390, 2023.
22. [Wortsman et al., 2024] Mitchell Wortsman, Peter J. Liu, Lechao Xiao, Katie Everett, Alexander Alemi, Ben Adlam, John D. Co-Reyes, Izzeddin Gur, Abhishek Kumar, Roman Novak, et al. Small-scale proxies for large-scale Transformer training instabilities. ICLR, 2024.
23. [Ziyin et al., 2020] Liu Ziyin, Zhikang T. Wang, and Masahito Ueda. LaProp: Separating momentum and adaptivity in Adam. arXiv:2002.04839, 2020.