![\"paper_fig2_legend.png\"](\"https://jonbarron.info/softmaxfix/img/paper_fig2_legend.png\")
Jonathan T. Barron
\nGoogle DeepMind
\nDan Gnanapragasam
\nGoogle DeepMind
\nMichael Tschannen
\nGoogle DeepMind
\nAbstract
\nStandard implementations of softmax, log-softmax, and\ncross-entropy have two numerical bugs.\nThe first bug, Softmax Collapse, rounds the target logit's gradient to\nzero if that one logit is much larger than the others, resulting in\ngradients that no longer sum to zero.\nThe second bug, Softmax Cancellation,\noccurs when student logits are close to teacher logits,\nand produces non-zero gradients even when the student and teacher logits are\nequal.\nWe present simple implementations of log-softmax, softmax, and\ncross-entropy that fix both bugs.\nOur fix eliminates the previously observed “slingshot” phenomenon,\nspeeds up grokking with weight decay, prevents accuracy collapse without weight\ndecay, stops student networks from drifting away from the teacher during\ndistillation, and shows that after memorization, logit growth is bounded by an\nequilibrium between weight decay and Adam's (\\epsilon) hyperparameter.
\n![\"paper_fig1a_hard_cliff.png\"](\"https://jonbarron.info/softmaxfix/img/paper_fig1a_hard_cliff.png\")
(a) Softmax Collapse: (|g_{t}|) vs. margin (m)
\n![\"paper_fig4a_floor.png\"](\"https://jonbarron.info/softmaxfix/img/paper_fig4a_floor.png\")
(b) Softmax Cancellation: (|\\nabla_{\\mathbf{z}_{\\mathrm{S}}}\\mathcal{L}|2) as (\\mathbf{z}{\\mathrm{S}}\n\\to \\mathbf{z}_{\\mathrm{T}})
\nFigure 1: The two bugs in standard softmax and cross-entropy.\n(a) Softmax Collapse: As a model becomes more confident in the\ncorrect class (thereby increasing the logit margin (m)), the true\ngradient on the target logit shrinks exponentially, but standard\nimplementations incorrectly round that gradient to zero once (m)\ncrosses a precision cliff ((\\approx 5.5) in bf16, (16.6) in\nfp32, and (36.7) in fp64).\n(b) Softmax Cancellation: In soft-target distillation, as\nstudent logits (\\mathbf{z}{\\mathrm{S}}) approach teacher logits (\\mathbf{z}{\\mathrm{T}}), the true\ngradient shrinks toward zero, but standard implementations lose precision\nand plateau at a non-zero noise floor.\nOur fix eliminates both bugs and yields accurate gradients across all\nprecisions.
\nNaively computing softmax or cross-entropy on logits (\\mathbf{z} \\in\n\\mathbb{R}^{K}) as (e^{z_k} / \\sum_j e^{z_j}) overflows if\nany logit exceeds the floating-point exponent limit ((\\approx 88) in\nfp32), which is why standard implementations in all deep learning\nframeworks (JAX, PyTorch, TensorFlow, etc.) shift the input logits by\nsubtracting their maximum value (z_{\\max} = \\max_j z_j)\n(Higham, 2002; Blanchard et al., 2021):
\n[\\tilde{z}_k = z_k - z_{\\max}\\,, \\quad\n\\operatorname{logsoftmax}(\\mathbf{z})_k\n= \\tilde{z}_k - \\log \\sum_j e^{\\tilde{z}_j}\\,,\n\\quad\nq_k = \\operatorname{softmax}(\\mathbf{z})_k\n= \\frac{e^{\\tilde{z}_k}}\n{\\sum_j e^{\\tilde{z}_j}}\\,.]
\n(1)
\nFor a target distribution (\\mathbf{p} \\in \\Delta^{K-1}),\ncross-entropy (CE) is (\\mathcal{L}(\\mathbf{z}, \\mathbf{p}) =\n-\\sum_k p_k\\,\\operatorname{logsoftmax}(\\mathbf{z})_k), with exact gradient\n(\\mathbf{g} = \\nabla_{\\mathbf{z}} \\mathcal{L}(\\mathbf{z}, \\mathbf{p}) = \\mathbf{q} -\n\\mathbf{p}).\nThough subtracting (z_{\\max}) prevents overflow, two numerical bugs remain.
\nThe first bug, Softmax Collapse (Figure 1a,\nPrieto et al., 2025), occurs when a\nmodel correctly predicts the target class (t \\in {1, \\dots,\nK}) with high confidence, so that (z_{t} = z_{\\max}) is\nmuch larger than all non-target logits (z_j) ((j \\neq t)).\nWe measure this confidence by the logit margin, which is how much larger\nthe target logit is than the log-sum-exp of all non-target logits:
\n[m = z_{t}
\n(2)
\nIn Equation 1, subtracting (z_{\\max} = z_{t}) sets\nthe shifted target logit to (\\tilde{z}_{t} = 0) (so\n(e^{\\tilde{z}_{t}} = 1) in the denominator), while the (K
\n1) non-target exponentials sum to a small positive number (e^{-m} =
\\sum_{j \\neq t} e^{z_j - z_{t}} \\ll 1).\nFloating-point numbers have much finer precision near (0) than near (1): in\nfp32, values near (0) can be as small as ({\\sim}10^{-38}), whereas\nvalues added to (1.0) are lost to rounding below ({\\sim}10^{-7}).\nBecause standard implementations form the denominator (1 + e^{-m}) by\nadding (e^{-m}) to (1.0), the (e^{-m}) term is lost to rounding\nonce the margin (m) exceeds a threshold (c_{d}) that we call the\nprecision cliff ((c_{16} \\approx 5.5) in bf16,\n(c_{32} \\approx 16.6) in fp32, and (c_{64} \\approx\n36.7) in fp64).\nWhen (m > c_{d}), the denominator (1 + e^{-m}) evaluates to\nexactly (1.0), which causes (q_{t} = 1/(1 + e^{-m})) to\nevaluate to (1.0) and zeroes out both the loss and the target gradient\n(\\tilde{g}{t} = q{t} - 1 = 0.0).\nHowever, the non-target gradients (\\tilde{g}_j = q_j = e^{z_j -\nz_{t}} / (1 + e^{-m})) ((j \\neq t)) are evaluated\nnear (0) without adding (1.0) to their numerators, so they remain positive all\nthe way down to exponent underflow ((m \\approx 88)).\nBecause the target gradient rounds to (0.0) while the non-target gradients do\nnot, the gradients no longer sum to zero ((\\sum_k \\tilde{g}_k > 0)).\nWhen training with adaptive optimizers such as Adam and AdamW\n(Kingma & Ba, 2015; Loshchilov & Hutter, 2019), which normalize each gradient\nby a moving average of its magnitude, this asymmetry causes the recurring\n“slingshot” loss spikes observed by Thilak et al. (2022) and the\ngrokking failures studied by Prieto et al. (2025).\nWhile prior works used surrogate losses (Prieto et al., 2025), gradient\nprojection (Liu et al., 2026), or fp64 casting\n(Nanda et al., 2023; Varma et al., 2023) to ameliorate these issues, we\ninstead fix log-softmax, softmax, and cross-entropy directly\n(Section 2).
The second bug, Softmax Cancellation (Figure 1b), occurs in\nknowledge distillation (Hinton et al., 2015) and KL regularization,\nwhere the target (\\mathbf{p} = \\operatorname{softmax}(\\mathbf{z}_{\\mathrm{T}})) comes from teacher logits\n(\\mathbf{z}_{\\mathrm{T}}).\nStandard implementations compute (\\mathbf{q} = \\operatorname{softmax}(\\mathbf{z}_{\\mathrm{S}})) and\n(\\mathbf{p} = \\operatorname{softmax}(\\mathbf{z}_{\\mathrm{T}})) separately and subtract them to get\n(\\tilde{\\mathbf{g}} = \\mathbf{q} - \\mathbf{p}).\nThis means that if (\\mathbf{z}{\\mathrm{S}}) is close to (\\mathbf{z}{\\mathrm{T}}), subtracting these\nnearly equal probabilities loses precision, and slight differences between the\nforward-pass softmax and backward-pass log-softmax formulas leave a\nnon-zero gradient even when (\\mathbf{z}{\\mathrm{S}} = \\mathbf{z}{\\mathrm{T}}).\nWe fix this by computing (\\mathbf{q} - \\mathbf{p}) directly from (\\mathbf{z}_{\\mathrm{S}} -\n\\mathbf{z}_{\\mathrm{T}}) in log space (Section 3).
\nWhen the target logit (z_{t} = z_{\\max}) is the largest logit,\nshifting by (z_{\\max}) in Equation 1 sets\n(\\tilde{z}_{t} = 0) and (\\tilde{z}_j = z_j -\nz_{t} < 0) for (j \\neq t), so the denominator is\n(\\sum_i e^{\\tilde{z}i} = 1 + \\sum{j \\neq t} e^{z_j -\nz_{t}} = 1 + e^{-m}).\nMathematically, the hard-target loss is (\\mathcal{L}(\\mathbf{z}, t) =\n-\\operatorname{logsoftmax}(\\mathbf{z})_{t} = \\log(1 + e^{-m})), with gradients
\n[g_{t}\n= \\underbrace{\\frac{1}{1 + e^{-m}} - 1}_{\\text{Standard}}\n= \\underbrace{-\\frac{e^{-m}}{1 + e^{-m}}}_{\\text{Fixed}}\\,,\n\\qquad\ng_j = \\frac{e^{z_j - z_{t}}}{1 + e^{-m}}\n\\quad \\text{for } j \\neq t\\,,]
\n(3)
\nwhich sum to zero ((\\sum_k g_k = 0)).\nIn floating point, however, the two expressions for (g_{t}) in\nEquation 3 behave very differently once (m > c_{d}),\nwhere (e^{-m}) is too small to change (1.0) when added to it\n(Figure 1a).\nStandard implementations add (1) and (e^{-m}) to form the denominator (1 +\ne^{-m}), which rounds to (1) when (m > c_{d}).\nIn the forward pass, (\\log(1 + e^{-m})) then evaluates to (\\log(1) = 0),\nzeroing out both (\\operatorname{logsoftmax}(\\mathbf{z})_{t}) and (\\mathcal{L}(\\mathbf{z},\nt)).\nIn the backward pass, standard cross-entropy implementations compute\n(q_{t} = 1/(1 + e^{-m}) = 1) and subtract (1) to get\n(\\tilde{g}{t} = q{t} - 1 = 0).\nIn contrast, our fixed expression (-e^{-m}/(1 + e^{-m})) keeps\n(e^{-m}) in the numerator: even when the denominator (1 + e^{-m})\nrounds to (1.0), dividing (-e^{-m}) by (1.0) preserves (g_{t}\n\\approx -e^{-m}) down to exponent underflow ((m \\approx 88)).
\nWhile (\\tilde{g}{t}) rounds to zero when (m > c{d}), the\nnon-target gradients do not.\nFor each non-target class (j \\neq t), standard cross-entropy computes\n(\\tilde{g}_j = q_j = e^{z_j - z_{t}} / (1 +\ne^{-m})).\nBecause the numerator (e^{z_j - z_{t}}) is evaluated directly\nfrom the negative difference (z_j - z_{t}) without adding\n(1), it remains accurate down to exponent underflow while the denominator rounds\nto (1).\nWhen (m > c_{d}), standard cross-entropy therefore zeroes out only the\ntarget entry (t) while all (K - 1) non-target gradients\nremain positive:
\n[\\tilde{g}_k = \\mathbf{1}{[k \\neq t]}\\,g_k\\,, \\qquad\n\\sum_k \\tilde{g}k = -g{t}\n= \\frac{e^{-m}}{1 + e^{-m}} > 0\n\\quad \\text{for } m > c_{d}\\,.]
\n(4)
\nThe same bug affects standalone (\\operatorname{softmax}(\\mathbf{z})) (Qiu & Yao, 2026): once\n(m > c_{d}), its backward pass rounds (q_{i_{\\max}}) ((i_{\\max} =\n\\arg\\max_i z_i)) to (1) and zeroes (\\tilde{g}{i{\\max}}) while the other\ngradients stay non-zero.
\nThis bug can be fixed by never adding (1) to (e^{-m}) directly.\nIn (\\operatorname{logsoftmax}(\\mathbf{z})) and (\\operatorname{softmax}(\\mathbf{z})), only the largest logit\n(z_{i_{\\max}}) is shifted to (0) (where (e^0 = 1) in the denominator).\nInstead of including (e^0 = 1) in the sum of exponentials, we set\n(\\tilde{z}{i{\\max}} = 0), compute the margin (m) by excluding (i_{\\max})\nfrom the log-sum-exp (Blanchard et al., 2021), and combine them with a\nnumerically stable (\\operatorname{softplus}(x) = \\log(1 + e^x)) (with\n(\\operatorname{softmax}(\\mathbf{z})_k = \\exp(\\operatorname{logsoftmax}(\\mathbf{z})_k))):
\n[\\tilde{z}_k\n= \\mathbf{1}{[k \\neq i_{\\max}]}{(z_k - z_{i_{\\max}})}\\,,\n\\;\\;\nm = -\\operatorname*{LSE}{j \\neq i{\\max}} \\tilde{z}_j\\,,\n\\;\\;\n\\operatorname{logsoftmax}(\\mathbf{z})_k = \\tilde{z}_k - \\operatorname{softplus}(-m)\\,.]
\n(5)
\nFor hard-target cross-entropy with target (t), we can use\n(t) directly in place of (i_{\\max}) without an (\\arg\\max):
\n[m = -\\operatorname*{LSE}_{j \\neq t}(z_j - z_{t})\\,,\n\\qquad\n\\mathcal{L}(\\mathbf{z}, t)\n= \\log(1 + e^{-m})\n= \\operatorname{softplus}(-m)\\,.]
\n(6)
\nBecause the (t)-th slot is replaced with the constant (0) rather than\n(z_{t} - z_{t}), (z_{t}) only appears in\n(\\mathcal{L}) through the (K - 1) differences (z_j -\nz_{t}) ((j \\neq t)), so autodiff correctly computes the\ntarget gradient as the negated sum of the non-target gradients:
\n[g_j = e^{z_j - z_{t} - \\mathcal{L}}\n= \\frac{e^{z_j - z_{t}}}{1 + e^{-m}}\n\\quad \\text{for } j \\neq t\\,, \\qquad\ng_{t} = -\\sum_{j \\neq t} g_j\n= -\\frac{e^{-m}}{1 + e^{-m}}\\,.]
\n(7)
\nHere is a JAX implementation of log-softmax, softmax, and hard-target\ncross-entropy, which remains accurate down to exponent underflow and\nsatisfies (\\sum_k g_k = 0):
\nimport jax, jax.numpy as jnp, optax\n\ndef log_softmax(z):\n i_max = jnp.argmax(z, axis=-1)\n z_max = jnp.take_along_axis(z, i_max[..., None], axis=-1) # Max logit.\n is_max = jax.nn.one_hot(i_max, z.shape[-1], dtype=bool) # Max mask.\n shifted = jnp.where(is_max, 0.0, z - z_max) # Exact 0 at argmax.\n others = jnp.where(is_max, -jnp.inf, shifted) # Exclude argmax from sum.\n m = -jax.nn.logsumexp(others, axis=-1, keepdims=True) # Logit margin.\n return shifted - jax.nn.softplus(-m) # log_softmax(z).\n\nsoftmax = lambda z: jnp.exp(log_softmax(z))\n\ndef cross_entropy(z, t):\n z_t = jnp.take_along_axis(z, t[..., None], axis=-1) # Target logit.\n is_t = jax.nn.one_hot(t, z.shape[-1], dtype=bool) # Target mask.\n others = jnp.where(is_t, -jnp.inf, z - z_t) # Exclude target from sum.\n m = -jax.nn.logsumexp(others, axis=-1) # Logit margin.\n return jax.nn.softplus(-m) # -log_softmax(z)_t.As an example, consider (\\mathbf{z} = [9.0, -9.0, -9.0]) and (t = 0).\nThe standard jax.nn and optax implementations return zero for\nthe target log-probability, loss, and target gradient, while our functions\nreturn accurate values whose gradients sum to zero:
z, t = jnp.array([9.0, -9.0, -9.0], dtype=jnp.float32), jnp.array(0)\nstd_ce = optax.softmax_cross_entropy_with_integer_labels\njax.nn.log_softmax(z) # -> [ 0.0, -18.0, -18.0 ]\njax.grad(lambda z: jax.nn.softmax(z)[0])(z) # -> [ 0.0, -1.523e-8, -1.523e-8]\njax.value_and_grad(std_ce)(z, t) # -> 0.0, [ 0.0, 1.523e-8, 1.523e-8]\nlog_softmax(z) # -> [-3.046e-8, -18.0, -18.0 ]\njax.grad(lambda z: softmax(z)[0])(z) # -> [ 3.046e-8, -1.523e-8, -1.523e-8]\njax.value_and_grad(cross_entropy)(z, t) # -> 3.046e-8, [-3.046e-8, 1.523e-8, 1.523e-8]Given student logits (\\mathbf{z}_{\\mathrm{S}} \\in \\mathbb{R}^{K}) and teacher logits\n(\\mathbf{z}_{\\mathrm{T}} \\in \\mathbb{R}^{K}), standard soft-target cross-entropy\ncomputes (\\mathbf{p} = \\operatorname{softmax}(\\mathbf{z}_{\\mathrm{T}})) and evaluates
\n[\\mathcal{L}(\\mathbf{z}{\\mathrm{S}}, \\mathbf{z}{\\mathrm{T}})\n= -\\sum_k p_k\\,\\operatorname{logsoftmax}(\\mathbf{z}_{\\mathrm{S}})_k\\,, \\quad\n\\mathbf{q} = \\operatorname{softmax}(\\mathbf{z}_{\\mathrm{S}})\\,, \\quad\n\\tilde{g}k = \\frac{\\partial \\mathcal{L}}{\\partial z{\\mathrm{S},k}}\n= q_k - p_k\\,.]
\n(8)
\nAs the student logits approach the teacher logits ((\\mathbf{z}{\\mathrm{S}} \\to \\mathbf{z}{\\mathrm{T}})),\nthe true gradient (\\mathbf{g} = \\mathbf{q} - \\mathbf{p}) shrinks toward (\\mathbf{0})\nand equals (\\mathbf{0}) when (\\mathbf{z}{\\mathrm{S}} = \\mathbf{z}{\\mathrm{T}}).\nIn floating point, however, computing (\\mathbf{q}) and (\\mathbf{p}) separately\nand subtracting them (Equation 8) loses precision when\n(\\mathbf{z}{\\mathrm{S}}) is close to (\\mathbf{z}{\\mathrm{T}}), because the small logit difference\n(\\boldsymbol{\\delta} = \\mathbf{z}{\\mathrm{S}} - \\mathbf{z}{\\mathrm{T}}) is lost when (\\mathbf{q}) and\n(\\mathbf{p}) are rounded to probabilities before being subtracted\n(Figure 1b).\nFurthermore, because (\\mathbf{p} = \\operatorname{softmax}(\\mathbf{z}_{\\mathrm{T}})) in the forward pass\nand (\\mathbf{q} = \\exp(\\operatorname{logsoftmax}(\\mathbf{z}_{\\mathrm{S}}))) in the backward pass are\nevaluated using slightly different sequences of operations, (\\mathbf{q}) and\n(\\mathbf{p}) do not match to machine precision even when (\\mathbf{z}_{\\mathrm{S}} =\n\\mathbf{z}_{\\mathrm{T}}), leaving a non-zero gradient (\\tilde{\\mathbf{g}} = \\mathbf{q} - \\mathbf{p}\n\\neq \\mathbf{0}).
\nTo avoid subtracting nearly equal probabilities, we pass (\\mathbf{z}_{\\mathrm{T}}) directly\nto our loss function and compute (\\mathbf{q} - \\mathbf{p}) in log space using\n(\\operatorname{expm1}(x) = e^x - 1) and (\\operatorname{log1p}(x) = \\log(1 + x)).\nGiven (\\boldsymbol{\\delta} = \\mathbf{z}{\\mathrm{S}} - \\mathbf{z}{\\mathrm{T}}), we center (\\boldsymbol{\\delta}) under\n(\\mathbf{p}) and write the log-ratio (s_k = \\log(q_k /\np_k)) and difference (u_k = q_k - p_k) as:
\n[d_k = \\delta_k - \\sum_i p_i \\delta_i\\,,\n\\qquad\ns_k = d_k
\n\\operatorname{log1p}\\left(\\sum_i p_i \\operatorname{expm1}(d_i)\\right)\\,,
\\qquad\nu_k = p_k \\operatorname{expm1}(s_k)\\,.]
(9)
\nBecause (\\sum_i p_i d_i = 0) and (\\operatorname{expm1}(x) \\ge x), the\ninput to (\\operatorname{log1p}) is always non-negative.\nWhen (|s_k| \\ge 1), we fall back to (u_k = q_k -\np_k) to avoid (0.0 \\times \\infty) if (p_k) rounds to (0.0)\nwhile (\\operatorname{expm1}) overflows, and we set (g_k = u_k - q_k\n\\sum_i u_i) so that (\\sum_k g_k = 0).\nBecause Equation 9 avoids adding or subtracting (1), (\\mathbf{g})\nremains accurate as (\\mathbf{z}{\\mathrm{S}} \\to \\mathbf{z}{\\mathrm{T}}) and is (\\mathbf{0}) when (\\mathbf{z}_{\\mathrm{S}} =\n\\mathbf{z}_{\\mathrm{T}}).\nHere is a JAX implementation of soft-target cross-entropy (using\nlog_softmax and softmax from Section 2), which\nuses jax.lax.stop_gradient (sg) to return the forward loss\nwhile substituting (\\mathbf{g}) in the backward pass:
def soft_target_cross_entropy(z_S, z_T):\n log_q = log_softmax(z_S)\n q, p = jnp.exp(log_q), softmax(z_T)\n delta = z_S - z_T # Logit difference.\n d = delta - jnp.sum(p * delta, axis=-1, keepdims=True) # Center under p.\n s = d - jnp.log1p(jnp.sum(p * jnp.expm1(d), axis=-1, keepdims=True)) # log(q / p).\n u = jnp.where(jnp.abs(s) < 1, p * jnp.expm1(s), q - p) # Difference q - p.\n g = u - q * jnp.sum(u, axis=-1, keepdims=True) # Ensure sum(g) = 0.\n sg = jax.lax.stop_gradient\n loss = -jnp.sum(p * sg(log_q), axis=-1) # Forward loss.\n 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\n(\\mathbf{z}_{\\mathrm{T}}).\nStandard soft-target cross-entropy shrinks the second gradient\nentry by nearly (3\\times), flips the sign of the third, and fails to sum to\nzero, while our implementation matches the exact fp64 gradient:
\nz_T = jnp.array([2.0, 1.0, 0.0], dtype=jnp.float32)\nz_S = z_T + jnp.array([-1e-7, 1e-7, 0.0], dtype=jnp.float32)\nstd_soft_ce = lambda z_S, z_T: -jnp.sum(jax.nn.softmax(z_T) * jax.nn.log_softmax(z_S))\njax.grad(std_soft_ce)(z_S, z_T) # -> [-5.960e-8, 1.452e-8, -1.123e-9]\njax.grad(soft_target_cross_entropy)(z_S, z_T) # -> [-4.596e-8, 4.144e-8, 4.513e-9]\n# Exact (float64) gradient: [-4.596e-8, 4.144e-8, 4.513e-9]In grokking (Power et al., 2022), a network quickly memorizes the training\nset ((q_{t} \\to 1)) and generalizes much later.\nWhile studying grokking on modular addition with Adam\n(Kingma & Ba, 2015; Loshchilov & Hutter, 2019), Thilak et al. (2022)\nobserved the slingshot mechanism, where training after memorization\nalternates between flat loss plateaus near zero and sudden loss spikes that\ncoincide with jumps in test accuracy.
\nThese slingshot spikes are a direct consequence of Softmax Collapse\n(Section 2).\nRecall that AdamW updates parameters (\\boldsymbol{\\theta}) by subtracting\n(\\eta\\big(\\hat{\\mathbf{m}} / (\\sqrt{\\hat{\\mathbf{v}}} + \\epsilon) +\n\\lambda\\boldsymbol{\\theta}\\big)), where (\\hat{\\mathbf{m}}) and (\\hat{\\mathbf{v}}) are\nbias-corrected moving averages of the parameter gradient (\\mathbf{g} =\n\\nabla_{\\boldsymbol{\\theta}}\\mathcal{L}(\\boldsymbol{\\theta})) and its square (with decay rates (\\beta_1,\n\\beta_2)).\nAfter memorization, the logit margins (m) grow as the network pushes its\ntarget probabilities toward (1).\nOnce an example's margin exceeds the precision cliff (c_{d}), standard\ncross-entropy rounds its target gradient (\\tilde{g}_{t}) to zero while its\nnon-target gradients (\\tilde{g}_j > 0) ((j \\neq t)) remain positive.\nWhen all examples in a batch sit above (c_{d}), the parameter gradients nearly\nvanish and Adam's second moment (\\hat{v}_i) decays exponentially toward zero\nwhile weight decay slowly shrinks the parameters.\nEventually, weight decay or the residual non-target updates pull a few examples\nback below (c_{d}), where their target gradients (\\tilde{g}_{t} \\approx\n-e^{-m}) abruptly turn back on.\nBecause (\\sqrt{\\hat{v}_i}) decayed while those examples were above\n(c_{d}), Adam divides the restored gradients by a very small denominator\n(\\sqrt{\\hat{v}_i} + \\epsilon), producing a large parameter step that spikes\nthe training loss, resets (\\hat{v}_i), and repeats the cycle.\nWith our fixed cross-entropy, the target-logit gradient (g_{t} =\n-e^{-m}/(1 + e^{-m})) does not round to zero at (c_{d}), so\n(\\hat{m}_i) and (\\sqrt{\\hat{v}_i}) decay smoothly together and slingshot\nspikes disappear (Figure 2).
\n![\"paper_fig_slingshot_legend.png\"](\"https://jonbarron.info/softmaxfix/img/paper_fig_slingshot_legend.png\")
![\"paper_fig_slingshot_a.png\"](\"https://jonbarron.info/softmaxfix/img/paper_fig_slingshot_a.png\")
(a) Mean logit margin (m)
\n![\"paper_fig_slingshot_b.png\"](\"https://jonbarron.info/softmaxfix/img/paper_fig_slingshot_b.png\")
(b) Training loss
\n![\"paper_fig_slingshot_c.png\"](\"https://jonbarron.info/softmaxfix/img/paper_fig_slingshot_c.png\")
(c) Test accuracy
\nFigure 2: The slingshot mechanism (Thilak et al., 2022) with and without\nweight decay (5 random seeds, fp32).\nTop row ((a + b \\bmod 31), (\\lambda = 1), (\\epsilon = 10^{-8}),\nVarma et al., 2023): Standard cross-entropy hits the\nfp32 cliff ((m = c_{32} \\approx 16.6), where loss\nrounds to (0)) and repeatedly drops and rebounds, slowing and degrading\ngrokking; fixed cross-entropy crosses (c_{32}) smoothly, groks\nfaster on all 5 seeds, and saturates just above (c_{32}) near\n(\\log({{}^{1}!/!{}_{\\epsilon}}) \\approx 18.4) (Section 5).\nBottom row ((a^2 + b^2 \\bmod 97), (\\lambda = 0), (\\epsilon =\n10^{-16}), Doshi et al., 2024): Without weight decay, hitting\n(c_{32}) around step (1250) triggers a spike that drops standard\ncross-entropy test accuracy from (61\\%) to (23\\%), whereas fixed\ncross-entropy continues smoothly to (85\\%).
\nWhether slingshot spikes appear to help or hurt generalization depends on how\nthe network's weight norm is controlled (Figure 2).\nWhen a network uses LayerNorm before the output layer and no weight decay\n(Thilak et al., 2022; Prieto et al., 2025), growing the pre-LayerNorm\nweights does not change the logits, but it does shrink the effective step size\n(\\eta / |\\boldsymbol{\\theta}|_2); each slingshot spike inflates (|\\boldsymbol{\\theta}|_2) and\nacts as an accidental learning-rate decay.\nWhen weight norms are already controlled by weight decay\n(Varma et al., 2023), slingshot spikes at (c_{32}) only slow\ndown and hurt grokking ((94 \\pm 5\\%) vs. (100 \\pm 0\\%) test accuracy on (a + b\n\\bmod 31); Figure 2, top).\nAnd when neither LayerNorm nor weight decay is used (Doshi et al., 2024), a\nsingle slingshot spike at (c_{32}) is destructive: on (a^2 + b^2 \\bmod\n97) ((\\lambda = 0), (\\epsilon = 10^{-16})), standard cross-entropy tracks fixed\ncross-entropy up to (61\\%) test accuracy until (m) hits (c_{32})\naround step (1250), where a spike drops test accuracy to (23\\%) (and from (60\\%)\nto (21\\%) on (a^3 + b \\bmod 97)), while fixed cross-entropy continues smoothly\nto (85\\%) (Figure 2, bottom).
\nIn Figure 2a, our fixed cross-entropy eliminates the\nprecision cliff at (c_{32} \\approx 16.6), yet with weight decay\n((\\lambda = 1), top) its margin still saturates just above (c_{32})\nnear (18), whereas without weight decay ((\\lambda = 0), bottom) its margin\ncontinues to climb.\nThis plateau is the equilibrium where weight decay balances AdamW's normalized\ngradient step.\nAt a stationary point where the parameters stop changing, the moving averages\nconverge to (\\hat{m}_i = g_i) and (\\hat{v}_i = g_i^2) (so\n(\\sqrt{\\hat{v}_i} = |g_i|)), and setting the update to zero gives:
\n[\\frac{\\hat{m}_i}{\\sqrt{\\hat{v}_i} + \\epsilon} + \\lambda\\theta_i = 0\n\\quad\\Longrightarrow\\quad\n\\frac{g_i}{|g_i| + \\epsilon} + \\lambda\\theta_i = 0\n\\quad \\forall i\\,.]
\n(10)
\nHere (\\eta), (\\beta_1), and (\\beta_2) cancel out, so RMSProp\n(Tieleman & Hinton, 2012) and LaProp (Ziyin et al., 2020) share the same\nstationary point.\nAfter memorization, the gradients shrink exponentially with the margin\n((|g_i| \\propto e^{-m})).\nIf (\\epsilon = 0), the normalized gradient (g_i / |g_i| =\n\\mathrm{sign}(g_i)) has a magnitude of (1) no matter how small (|g_i|)\nbecomes, which causes the margin to grow until the gradients underflow.\nWhen (\\epsilon > 0), however, the normalized gradient has magnitude (\\approx 1)\nwhile (|g_i| \\gg \\epsilon), and shrinks as (g_i / \\epsilon \\propto\ne^{-m}) once (|g_i| \\ll \\epsilon) (Figure 3).\nBalancing this shrinking step against weight decay (\\lambda\\theta_i) requires\n(|g_i| \\approx \\epsilon\\lambda|\\theta_i|)\n(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\n(\\epsilon\\lambda) (Rosset et al., 2004; Soudry et al., 2018).\nSince (|g_i| \\propto e^{-m}), the margin stops growing at:
\n[m^\\star(\\epsilon) \\approx \\log({{}^{1}!/!{}_{\\epsilon}})\\,.]
\n(11)
\nEquation 11 explains why the default (\\epsilon = 10^{-8})\ncauses slingshot spikes in fp32: its equilibrium\n(\\log({{}^{1}!/!{}{\\epsilon}}) \\approx 18.4) sits just above (c{32}\n\\approx 16.6) (Figure 3), so standard cross-entropy hits the cliff\nbefore equilibrium (whereas for (\\epsilon \\ge 10^{-7}),\n(\\log({{}^{1}!/!{}{\\epsilon}}) \\le 16.1 < c{32}) and standard\ncross-entropy settles without spiking).\nTo test Equation 11 across many orders of magnitude down to\n(\\epsilon = 0) in fp32 (Figure 4), we train with our\nfixed cross-entropy and LaProp on (a + b \\bmod 31), updating\n(\\sqrt{\\hat{\\mathbf{v}}}) directly without squaring (g_i) (which underflows in\nfp32 once (|g_i| < 2^{-63}), or (m \\approx 38)).\nWith weight decay ((\\lambda = 1), Figure 4a), the margin for each\n(\\epsilon > 0) climbs along the (\\epsilon = 0) trajectory and levels off at\n(m^\\star(\\epsilon) = \\log({{}^{1}!/!{}_{\\epsilon}})), while (\\epsilon = 0) grows\nuntil the fp32 gradients underflow at (m \\approx 81.2).\nWithout weight decay ((\\lambda = 0), Figure 4b), nothing balances\nthe (g_i / \\epsilon) step, so margins keep rising past\n(\\log({{}^{1}!/!{}_{\\epsilon}})).
\n![\"paper_fig1_legend.png\"](\"https://jonbarron.info/softmaxfix/img/paper_fig1_legend.png\")
Figure 3: Normalized Adam step\n(|g_{t}| / (|g_{t}| + \\epsilon)) vs. logit\nmargin (m) at (\\epsilon = 10^{-8}).\nWith our fixed cross-entropy, the step stays near (1) and decays\nsmoothly around (\\log({{}^{1}!/!{}_{\\epsilon}}) \\approx 18.4), whereas\nstandard cross-entropy zeroes out (\\tilde{g}_{t}) at\n(c_{16} \\approx 5.5) and (c_{32} \\approx 16.6).
\n![\"paper_fig1b_adam_eps.png\"](\"https://jonbarron.info/softmaxfix/img/paper_fig1b_adam_eps.png\")
![\"paper_fig6a.png\"](\"https://jonbarron.info/softmaxfix/img/paper_fig6a.png\")
(a) Weight decay ((\\lambda = 1))
\n![\"paper_fig6b.png\"](\"https://jonbarron.info/softmaxfix/img/paper_fig6b.png\")
(b) No weight decay ((\\lambda = 0))
\nFigure 4: Logit margins with our fixed cross-entropy and LaProp on (a + b \\bmod\n31) in fp32 ((\\beta_1 = 0.9), (\\beta_2 = 0.98); 3 seeds\nper setting as solid lines, analytical predictions\n(m^\\star(\\epsilon) = \\log({{}^{1}!/!{}_{\\epsilon}})) as dashed lines).\n(a) With weight decay ((\\lambda = 1)), margins level off at\n(\\log({{}^{1}!/!{}_{\\epsilon}})).\n(b) Without weight decay ((\\lambda = 0)), margins overshoot\n(\\log({{}^{1}!/!{}_{\\epsilon}})).
\nTo see how the precision cliff behaves across floating-point formats, we follow\nVarma et al. (2023) and Lyle et al. (2025) and train a 1-layer\nwidth-(64) ReLU Transformer on (a + b \\bmod 31) ((50\\%) training split) for\n(10{,}000) epochs of full-batch AdamW ((\\eta = 10^{-3}), (\\beta_1 = 0.9),\n(\\beta_2 = 0.98), (\\lambda = 1)), varying the loss (standard vs. fixed\ncross-entropy), precision (bf16, fp32, fp64), and\n(\\epsilon \\in {10^{-8}, 10^{-12}, 10^{-16}}) (Figure 5,\nTable 1).\nAt (\\epsilon = 10^{-8}) (Figure 5, top), standard cross-entropy\nstalls at (c_{16} \\approx 5.5) in bf16 and repeatedly spikes at\n(c_{32} \\approx 16.6) in fp32, whereas fixed cross-entropy\ncrosses both cliffs smoothly to (\\log({{}^{1}!/!{}_{\\epsilon}}) \\approx 18.4) and\nreaches (100\\%) test accuracy in all three precisions.\nReducing (\\epsilon) to (10^{-16}) (Figure 5, bottom) raises the\nequilibrium to (\\log({{}^{1}!/!{}{\\epsilon}}) \\approx 36.8 > c{64}),\nbreaking even fp64 standard cross-entropy at (c_{64} \\approx\n36.7) while fixed cross-entropy converges smoothly across all precisions.
\nIn Table 1 we summarize these results and compare against two\nrecent methods that modify the gradients or the loss to mitigate Softmax\nCollapse.\nLiu et al. (2026) subtract the mean logit gradient from each entry so the\ngradients sum to zero (“Zero-sum Proj.”).\nHowever, once (m > c_{32}) and (\\tilde{g}_{t} = 0),\nsubtracting the mean only gives the target logit (-{{}^{1}!/!{}_{K}}) of\nthe total non-target gradient, stalling (m) near (18) and dropping test\naccuracy below (50\\%) for (\\epsilon \\le 10^{-12}).\nPrieto et al. (2025) propose “StableMax,” which replaces (\\exp(x)) in\nsoftmax with a piecewise rational function ((x + 1) for (x \\ge 0) and\n(1/(1 - x)) for (x < 0)).\nBecause this function grows linearly instead of exponentially, logit margins\nblow up to ({\\sim}10^4) and test accuracy lags our fixed cross-entropy across\nall (\\epsilon).
\nWe observe the same behavior across four grokking benchmarks at (\\epsilon =\n10^{-12}) (Figure 6): 1-layer Transformers on (a + b \\bmod 113)\n(Nanda et al., 2023) and (a + b \\bmod 31) (Varma et al., 2023) with\n(\\lambda = 1), and 2-layer MLPs on (a^2 + b^2 \\bmod 97) and (a^3 + b \\bmod 97)\n(Doshi et al., 2024) with (\\lambda = 0.2).\nOn all four tasks, fp32 standard cross-entropy spikes once it hits\n(c_{32}), whereas fp32 fixed cross-entropy matches\nfp64 standard cross-entropy and reaches (100\\%) test accuracy.
\n![\"paper_fig2a.png\"](\"https://jonbarron.info/softmaxfix/img/paper_fig2a.png\")
(a) Mean logit margin (m)
\n![\"paper_fig2b.png\"](\"https://jonbarron.info/softmaxfix/img/paper_fig2b.png\")
(b) Training loss
\n![\"paper_fig2c.png\"](\"https://jonbarron.info/softmaxfix/img/paper_fig2c.png\")
(c) Test accuracy
\nFigure 5: Standard cross-entropy fails at the precision cliff (c_{d}) of its\nfloating-point format; fixed cross-entropy removes the cliff across all\nprecisions ((a + b \\bmod 31), (\\lambda = 1); 2 of 5 seeds shown).\nTop: (\\epsilon = 10^{-8}).\nBottom: (\\epsilon = 10^{-16}), where even fp64 standard\ncross-entropy fails at (c_{64}) while fixed cross-entropy\nconverges smoothly.
\n| Test Accuracy (%) | Final Margin (m) | ||||||||||||
|---|---|---|---|---|---|---|---|---|---|---|---|---|---|
| Variant(\\epsilon =) | (10^{-8}) | (10^{-12}) | (10^{-16}) | (10^{-8}) | (10^{-12}) | (10^{-16}) | |||||||
| fp32 | Fixed CE | 100.0 | (\\pm) | 0.0 | 99.6 | (\\pm) | 0.7 | 99.0 | (\\pm) | 2.3 | 17.9 | 27.4 | 36.8 |
| fp32 | Standard CE | 94.3 | (\\pm) | 5.4 | 96.4 | (\\pm) | 5.1 | 96.5 | (\\pm) | 4.3 | 13.0 | 15.1 | 12.7 |
| bf16 | Fixed CE | 100.0 | (\\pm) | 0.1 | 99.8 | (\\pm) | 0.2 | 98.9 | (\\pm) | 2.4 | 17.9 | 27.4 | 36.8 |
| bf16 | Standard CE | 93.3 | (\\pm) | 7.4 | 94.9 | (\\pm) | 4.9 | 94.2 | (\\pm) | 8.3 | 5.1 | 5.6 | 5.8 |
| fp64 | Fixed CE | 100.0 | (\\pm) | 0.0 | 99.5 | (\\pm) | 1.1 | 99.0 | (\\pm) | 2.0 | 17.9 | 27.4 | 36.8 |
| fp64 | Standard CE | 100.0 | (\\pm) | 0.0 | 100.0 | (\\pm) | 0.1 | 97.2 | (\\pm) | 3.2 | 17.9 | 27.4 | 30.1 |
| fp32 | Zero-sum Proj. (Liu et al., 2026) | 97.9 | (\\pm) | 3.2 | 48.9 | (\\pm) | 2.6 | 49.3 | (\\pm) | 2.5 | 17.0 | 18.9 | 18.9 |
| fp32 | StableMax (Prieto et al., 2025) | 97.3 | (\\pm) | 4.7 | 95.6 | (\\pm) | 2.0 | 97.3 | (\\pm) | 4.3 | 3200 | 5661 | 6313 |
| fp64 | StableMax (Prieto et al., 2025) | 97.4 | (\\pm) | 4.2 | 96.7 | (\\pm) | 5.3 | 95.6 | (\\pm) | 2.7 | 3170 | 21252 | 2322 |
Table 1: Comparison of loss functions on modular addition ((a + b \\bmod 31), (\\lambda\n= 1), (3)–(5) seeds) across (\\epsilon \\in {10^{-8}, 10^{-12}, 10^{-16}}).\nFixed CE reaches the predicted equilibrium margin ((\\log({{}^{1}!/!{}_{\\epsilon}})\n\\approx 18.4, 27.6, 36.8)) and ({\\ge}98.9\\%) test accuracy across all\nprecisions, whereas Standard CE stalls once (m^\\star(\\epsilon)) exceeds the\ncliff (c_{d}) ((c_{16} \\approx 5.5) in bf16,\n(c_{32} \\approx 16.6) in fp32, and (c_{64}\n\\approx 36.7) in fp64).
\n![\"paper_fig3_legend.png\"](\"https://jonbarron.info/softmaxfix/img/paper_fig3_legend.png\")
![\"paper_fig3a.png\"](\"https://jonbarron.info/softmaxfix/img/paper_fig3a.png\")
(a) (a{+}b \\bmod 113)
\n Nanda et al. (2023)
![\"paper_fig3b.png\"](\"https://jonbarron.info/softmaxfix/img/paper_fig3b.png\")
(b) (a{+}b \\bmod 31)
\n Varma et al. (2023)
![\"paper_fig3c.png\"](\"https://jonbarron.info/softmaxfix/img/paper_fig3c.png\")
(c) (a^2{+}b^2 \\bmod 97)
\n Doshi et al. (2024)
![\"paper_fig3d.png\"](\"https://jonbarron.info/softmaxfix/img/paper_fig3d.png\")
(d) (a^3{+}b \\bmod 97)
\n Doshi et al. (2024)
Figure 6: Results across four grokking tasks at (\\epsilon = 10^{-12}) (2 of 3–5\nseeds shown): 1-layer Transformers (a, b) and 2-layer MLPs (c, d).\nFixed cross-entropy in fp32 (blue) matches fp64\nstandard cross-entropy (red), while standard cross-entropy in\nfp32 (orange) breaks at (c_{32}).
\nTo see how Softmax Cancellation (Section 3) affects optimization,\nwe initialize a student network with the weights of a teacher trained to (100\\%)\naccuracy on (a + b \\bmod 53) ((\\boldsymbol{\\theta}{\\mathrm{S}}^{(0)} = \\boldsymbol{\\theta}{\\mathrm{T}})), and run\n(2{,}000) steps of Adam ((\\eta = 2\\times 10^{-3}), (\\lambda = 0),\n(\\epsilon = 10^{-8})) at distillation temperature (T = 4)\n(Hinton et al., 2015), which keeps the scaled margin (m/T)\nbelow (c_{16}) so (\\mathbf{p}) does not round to one-hot.\nBecause (\\mathbf{z}{\\mathrm{S}} = \\mathbf{z}{\\mathrm{T}}) initially, the true gradient is zero and the\nstudent should not move.\nHowever, standard soft-target cross-entropy has a non-zero residual gradient at\n(\\mathbf{z}{\\mathrm{S}} = \\mathbf{z}{\\mathrm{T}}) due to the mismatch between its forward softmax\nand backward log-softmax (Figure 1b).\nWhenever this residual gradient exceeds (\\epsilon), Adam normalizes it into\nsteps of size (\\eta) (Figure 7): in fp64 the residual\n(({\\sim}10^{-17})) lies well below (\\epsilon = 10^{-8}) so the student stays\nnear (\\mathrm{KL} = 0), in fp32 (({\\sim}10^{-8})) it causes an initial\n(\\mathrm{KL}) spike and a ({\\sim}10^{-6}) floor, and in bf16\n(({\\sim}10^{-3})) it drives the student away from the teacher and degrades test\naccuracy.\nOur fixed soft-target cross-entropy returns (\\mathbf{g} = \\mathbf{0}) at (\\mathbf{z}_{\\mathrm{S}} =\n\\mathbf{z}_{\\mathrm{T}}), so the student stays at (\\mathrm{KL} = 0) and (100\\%) accuracy\nacross all precisions.
\n![\"paper_fig4_legend.png\"](\"https://jonbarron.info/softmaxfix/img/paper_fig4_legend.png\")
![\"paper_fig4b_kl.png\"](\"https://jonbarron.info/softmaxfix/img/paper_fig4b_kl.png\")
(a) Held-out (\\mathrm{KL}(\\mathbf{p}{\\mathrm{T}} \\parallel \\mathbf{q}{\\mathrm{S}}))
\n![\"paper_fig4c_acc.png\"](\"https://jonbarron.info/softmaxfix/img/paper_fig4c_acc.png\")
(b) Test accuracy
\nFigure 7: Soft-target distillation when the student is initialized with the\nteacher's weights ((\\boldsymbol{\\theta}{\\mathrm{S}}^{(0)} = \\boldsymbol{\\theta}{\\mathrm{T}}), (T = 4),\n(\\epsilon = 10^{-8}), 2 of 3 seeds shown).\nFixed soft-target cross-entropy stays at (\\mathrm{KL} = 0) and (100\\%)\naccuracy across all precisions; standard soft-target cross-entropy\namplifies rounding error into parameter drift in bf16 and an\ninitial (\\mathrm{KL}) spike and ({\\sim}10^{-6}) floor in fp32.
\nWe have presented simple fixes for two numerical bugs in standard\nimplementations of softmax, log-softmax, and cross-entropy:\nSoftmax Collapse, which zeroes out the dominant logit's gradient at high\nconfidence while leaving non-dominant gradients non-zero, and *Softmax\nCancellation*, which loses precision when student logits approach teacher logits\nin soft-target distillation.\nRewriting these functions in log space using (\\operatorname{softplus}), (\\operatorname{expm1}), and\n(\\operatorname{log1p}) restores accurate, zero-sum gradients across bf16,\nfp32, and fp64 without internally casting to higher\nprecision.\nFixing these bugs eliminates slingshot loss spikes during grokking, prevents\nstudent networks from drifting away from the teacher during distillation, and\nreveals that logit margins under AdamW and weight decay naturally converge to\n(\\log({{}^{1}!/!{}_{\\epsilon}})).\nThough the extreme margins and near-zero student–teacher gaps probed by our\nexperiments are uncommon in standard large-scale training, our fixed functions\nare easy drop-in replacements that can be implemented without meaningful\nslowdown, and may prove useful whenever models make high-confidence predictions\nor distill from close teachers.
\nAcknowledgments:\nWe thank Jean-Baptiste Alayrac, Forrester Cole, Roy Frostig, Mark Hamilton, Miki\nRubinstein, Keith Rush, and Brandon Zhao for their helpful advice.
\n