{"slug": "a-fix-for-softmax-and-cross-entropy", "title": "A Fix for Softmax and Cross-Entropy", "summary": "A proposed fix for standard softmax, log-softmax and cross-entropy implementations addresses two numerical bugs the authors call Softmax Collapse and Softmax Cancellation. Softmax Collapse occurs once the logit margin exceeds a precision cliff of roughly 5.5 in bf16, 16.6 in fp32 and 36.7 in fp64, where the target gradient rounds to 0.0 while non-target gradients stay positive, so gradients no longer sum to zero; Softmax Cancellation produces non-zero gradients when student and teacher logits are equal. The authors report the fix eliminates the \"slingshot\" phenomenon, speeds up grokking with weight decay, prevents accuracy collapse without weight decay, stops student networks drifting from the teacher during distillation, and shows logit growth after memorization is bounded by an equilibrium between weight decay and Adam's epsilon hyperparameter.", "body_md": "# A Fix for Softmax and Cross-Entropy\n\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\n## 1. Introduction[#](#sec:intro)\n\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](#bib-higham2002accuracy); [Blanchard et al., 2021](#bib-blanchard2021accurately)):\n\nFor 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.\n\nThe first bug, *Softmax Collapse* ([Figure 1a](#fig:fig1a),\n[Prieto et al., 2025](#bib-prieto2025grokking)), 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\nIn [Equation 1](#eq:std_softmax), 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\n- 1\\) non-target exponentials sum to a small positive number \\(e^{-m} =\n\\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\n*precision 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](#bib-kingma2015adam); [Loshchilov & Hutter, 2019](#bib-loshchilov2019decoupled)), 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)](#bib-thilak2022slingshot) and the\ngrokking failures studied by [Prieto et al. (2025)](#bib-prieto2025grokking).\nWhile prior works used surrogate losses ([Prieto et al., 2025](#bib-prieto2025grokking)), gradient\nprojection ([Liu et al., 2026](#bib-liu2026grokking)), or fp64 casting\n([Nanda et al., 2023](#bib-nanda2023progress); [Varma et al., 2023](#bib-varma2023explaining)) to ameliorate these issues, we\ninstead fix log-softmax, softmax, and cross-entropy directly\n([Section 2](#sec:bug1)).\n\nThe second bug, *Softmax Cancellation* ([Figure 1b](#fig:fig4a)), occurs in\nknowledge distillation ([Hinton et al., 2015](#bib-hinton2015distilling)) 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](#sec:bug2)).\n\n## 2. Softmax Collapse[#](#sec:bug1)\n\nWhen the target logit $z_{t} = z_{\\max}$ is the largest logit,\nshifting by $z_{\\max}$ in [Equation 1](#eq:std_softmax) 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\nwhich sum to zero ($\\sum_k g_k = 0$).\nIn floating point, however, the two expressions for $g_{t}$ in\n[Equation 3](#eq:hard_exact) behave very differently once $m > c_{d}$,\nwhere $e^{-m}$ is too small to change $1.0$ when added to it\n([Figure 1a](#fig:fig1a)).\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$).\n\nWhile $\\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:\n\nThe same bug affects standalone $\\operatorname{softmax}(\\mathbf{z})$ ([Qiu & Yao, 2026](#bib-qiu2026why)): 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.\n\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](#bib-blanchard2021accurately)), 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\nFor hard-target cross-entropy with target $t$, we can use $t$ directly in place of $i_{\\max}$ without an $\\arg\\max$:\n\nBecause 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:\n\nHere 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$:\n\n``` python\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.\n```\n\nAs 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:\n\n```\nz, 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]\n```\n\n## 3. Softmax Cancellation[#](#sec:bug2)\n\nGiven 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\n\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](#eq:soft_buggy)) 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](#fig:fig4a)).\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}\\).\n\nTo 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:\n\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](#eq:soft_fix) 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\n`log_softmax` and `softmax` from [Section 2](#sec:bug1)), which\nuses `jax.lax.stop_gradient` (`sg`) to return the forward loss\nwhile substituting $\\mathbf{g}$ in the backward pass:\n\n``` python\ndef 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.\n```\n\nAs 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:\n\n```\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]\n```\n\n## 4. The Slingshot Mechanism[#](#sec:slingshot)\n\nIn grokking ([Power et al., 2022](#bib-power2022grokking)), 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](#bib-kingma2015adam); [Loshchilov & Hutter, 2019](#bib-loshchilov2019decoupled)), [Thilak et al. (2022)](#bib-thilak2022slingshot)\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.\n\nThese slingshot spikes are a direct consequence of Softmax Collapse\n([Section 2](#sec:bug1)).\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](#fig:fig_slingshot)).\n\nWhether slingshot spikes appear to help or hurt generalization depends on how\nthe network's weight norm is controlled ([Figure 2](#fig:fig_slingshot)).\nWhen a network uses LayerNorm before the output layer and no weight decay\n([Thilak et al., 2022](#bib-thilak2022slingshot); [Prieto et al., 2025](#bib-prieto2025grokking)), 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](#bib-varma2023explaining)), 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](#fig:fig_slingshot), top).\nAnd when neither LayerNorm nor weight decay is used ([Doshi et al., 2024](#bib-doshi2024grokking)), 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](#fig:fig_slingshot), bottom).\n\n## 5. Margin Equilibria[#](#sec:equilibria)\n\nIn [Figure 2a](#fig:fig_slingshot), 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\nHere $\\eta$, $\\beta_1$, and $\\beta_2$ cancel out, so RMSProp\n([Tieleman & Hinton, 2012](#bib-tieleman2012rmsprop)) and LaProp ([Ziyin et al., 2020](#bib-ziyin2020laprop)) 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](#fig:fig1b)).\nBalancing this shrinking step against weight decay $\\lambda\\theta_i$ requires\n$|g_i| \\approx \\epsilon\\lambda|\\theta_i|$\n([Wortsman et al., 2024](#bib-wortsman2023small); [Molybog et al., 2023](#bib-molybog2023theory); [Everett et al., 2024](#bib-everett2024scaling); [Kosson et al., 2024](#bib-kosson2024rotational)), which acts like an $L_2$ regularizer with weight\n$\\epsilon\\lambda$ ([Rosset et al., 2004](#bib-rosset2004margin); [Soudry et al., 2018](#bib-soudry2018implicit)).\nSince $|g_i| \\propto e^{-m}$, the margin stops growing at:\n\n[Equation 11](#eq:margin_law) 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](#fig:fig1b)), 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](#eq:margin_law) across many orders of magnitude down to\n$\\epsilon = 0$ in fp32 ([Figure 4](#fig:fig6)), 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](#fig:fig6a)), 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](#fig:fig6b)), nothing balances\nthe $g_i / \\epsilon$ step, so margins keep rising past\n$\\log({{}^{1}\\!/\\!{}_{\\epsilon}})$.\n\n## 6. Grokking Across Precisions and Benchmarks[#](#sec:experiments)\n\nTo see how the precision cliff behaves across floating-point formats, we follow\n[Varma et al. (2023)](#bib-varma2023explaining) and [Lyle et al. (2025)](#bib-lyle2025grokking) 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](#fig:fig2),\n[Table 1](#tab:baselines)).\nAt $\\epsilon = 10^{-8}$ ([Figure 5](#fig:fig2), 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](#fig:fig2), 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.\n\nIn [Table 1](#tab:baselines) we summarize these results and compare against two\nrecent methods that modify the gradients or the loss to mitigate Softmax\nCollapse.\n[Liu et al. (2026)](#bib-liu2026grokking) 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}$.\n[Prieto et al. (2025)](#bib-prieto2025grokking) 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$.\n\nWe observe the same behavior across four grokking benchmarks at \\(\\epsilon =\n10^{-12}\\) ([Figure 6](#fig:fig3)): 1-layer Transformers on $a + b \\bmod 113$\n([Nanda et al., 2023](#bib-nanda2023progress)) and $a + b \\bmod 31$ ([Varma et al., 2023](#bib-varma2023explaining)) 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](#bib-doshi2024grokking)) 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\n|  |  | Test Accuracy (%) |  |  |  |  |  |  |  |  | Final Margin $m$ |  |  | \n|---|---|---|---|---|---|---|---|---|---|---|---|---|---|\n|  | Variant$\\epsilon =$ | $10^{-8}$ |  |  | $10^{-12}$ |  |  | $10^{-16}$ |  |  | $10^{-8}$ | $10^{-12}$ | $10^{-16}$ | \n| fp32 | Fixed CE | 100.0 | $\\pm$ | 0.0 | 99.6 | $\\pm$ | 0.7 | 99.0 | $\\pm$ | 2.3 | 17.9 | 27.4 | 36.8 | \n| fp32 | Standard CE | 94.3 | $\\pm$ | 5.4 | 96.4 | $\\pm$ | 5.1 | 96.5 | $\\pm$ | 4.3 | 13.0 | 15.1 | 12.7 | \n| bf16 | Fixed CE | 100.0 | $\\pm$ | 0.1 | 99.8 | $\\pm$ | 0.2 | 98.9 | $\\pm$ | 2.4 | 17.9 | 27.4 | 36.8 | \n| bf16 | Standard CE | 93.3 | $\\pm$ | 7.4 | 94.9 | $\\pm$ | 4.9 | 94.2 | $\\pm$ | 8.3 | 5.1 | 5.6 | 5.8 | \n| fp64 | Fixed CE | 100.0 | $\\pm$ | 0.0 | 99.5 | $\\pm$ | 1.1 | 99.0 | $\\pm$ | 2.0 | 17.9 | 27.4 | 36.8 | \n| fp64 | Standard CE | 100.0 | $\\pm$ | 0.0 | 100.0 | $\\pm$ | 0.1 | 97.2 | $\\pm$ | 3.2 | 17.9 | 27.4 | 30.1 | \n| fp32 | Zero-sum Proj. ( [Liu et al., 2026](#bib-liu2026grokking) ) | 97.9 | $\\pm$ | 3.2 | 48.9 | $\\pm$ | 2.6 | 49.3 | $\\pm$ | 2.5 | 17.0 | 18.9 | 18.9 | \n| fp32 | StableMax ( [Prieto et al., 2025](#bib-prieto2025grokking) ) | 97.3 | $\\pm$ | 4.7 | 95.6 | $\\pm$ | 2.0 | 97.3 | $\\pm$ | 4.3 | 3200 | 5661 | 6313 | \n| fp64 | StableMax ( [Prieto et al., 2025](#bib-prieto2025grokking) ) | 97.4 | $\\pm$ | 4.2 | 96.7 | $\\pm$ | 5.3 | 95.6 | $\\pm$ | 2.7 | 3170 | 21252 | 2322 | \n\n## 7. Soft-Target Distillation[#](#sec:distillation)\n\nTo see how Softmax Cancellation ([Section 3](#sec:bug2)) 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](#bib-hinton2015distilling)), 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](#fig:fig4a)).\nWhenever this residual gradient exceeds $\\epsilon$, Adam normalizes it into\nsteps of size $\\eta$ ([Figure 7](#fig:fig4)): 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\n## 8. Conclusion[#](#sec:recommendations)\n\nWe have presented simple fixes for two numerical bugs in standard\nimplementations of softmax, log-softmax, and cross-entropy:\n*Softmax 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.\n\n**Acknowledgments:**\nWe thank Jean-Baptiste Alayrac, Forrester Cole, Roy Frostig, Mark Hamilton, Miki\nRubinstein, Keith Rush, and Brandon Zhao for their helpful advice.\n\n## References[#](#references)\n\n1. [Blanchard et al., 2021]  Accurately computing the log-sum-exp and softmax functions. *IMA Journal of Numerical Analysis* , 2021.\n2. [Doshi et al., 2024]  Grokking modular polynomials. *[arXiv:2406.03495](https://arxiv.org/abs/2406.03495)* , 2024.\n3. [Everett et al., 2024]  Scaling exponents across parameterizations and optimizers. *ICML* , 2024.\n4. [Higham, 2002]  *Accuracy and stability of numerical algorithms* . SIAM, 2nd edition, 2002.\n5. [Hinton et al., 2015]  Distilling the knowledge in a neural network. *[arXiv:1503.02531](https://arxiv.org/abs/1503.02531)* , 2015.\n6. [Kingma & Ba, 2015]  Adam: A method for stochastic optimization. *ICLR* , 2015.\n7. [Kosson et al., 2024]  Rotational equilibrium: How weight decay balances learning across neural networks. *ICML* , 2024.\n8. [Liu et al., 2026]  Grokking or glitching? How low-precision drives Slingshot loss spikes. *[arXiv:2605.06152](https://arxiv.org/abs/2605.06152)* , 2026.\n9. [Loshchilov & Hutter, 2019]  Decoupled weight decay regularization. *ICLR* , 2019.\n10. [Lyle et al., 2025]  What can grokking teach us about learning under non-stationarity? *Conference on Lifelong Learning Agents* , 2025.\n11. [Molybog et al., 2023]  A theory on Adam instability in large-scale machine learning. *[arXiv:2304.09871](https://arxiv.org/abs/2304.09871)* , 2023.\n12. [Nanda et al., 2023]  Progress measures for grokking via mechanistic interpretability. *ICLR* , 2023.\n13. [Power et al., 2022]  Grokking: Generalization beyond overfitting on small algorithmic datasets. *[arXiv:2201.02177](https://arxiv.org/abs/2201.02177)* , 2022.\n14. [Prieto et al., 2025]  Grokking at the edge of numerical stability. *ICLR* , 2025.\n15. [Qiu & Yao, 2026]  Why low-precision Transformer training fails: An analysis on Flash Attention. *ICLR* , 2026.\n16. [Rosset et al., 2004]  Margin maximizing loss functions. *NeurIPS* , 2004.\n17. [Soudry et al., 2018]  The implicit bias of gradient descent on separable data. *JMLR* , 2018.\n18. [Thilak et al., 2022]  The Slingshot mechanism: An empirical study of adaptive optimizers and the grokking phenomenon. *[arXiv:2206.04817](https://arxiv.org/abs/2206.04817)* , 2022.\n19. [Tieleman & Hinton, 2012]  Lecture 6.5—RMSProp: Divide the gradient by a running average of its recent magnitude. *COURSERA: Neural Networks for Machine Learning* , 2012.\n20. [Varma et al., 2023]  Explaining grokking through circuit efficiency. *[arXiv:2309.02390](https://arxiv.org/abs/2309.02390)* , 2023.\n21. [Wortsman et al., 2024]  Small-scale proxies for large-scale Transformer training instabilities. *ICLR* , 2024.\n22. [Ziyin et al., 2020]  LaProp: Separating momentum and adaptivity in Adam. *[arXiv:2002.04839](https://arxiv.org/abs/2002.04839)* , 2020.", "url": "https://wpnews.pro/news/a-fix-for-softmax-and-cross-entropy", "canonical_source": "https://jonbarron.info/softmaxfix/", "published_at": "2026-10-23 22:07:57+00:00", "updated_at": "2026-10-23 22:42:53.405300+00:00", "lang": "en", "topics": ["machine-learning", "ai-research", "neural-networks"], "entities": ["JAX", "PyTorch", "TensorFlow", "Adam", "AdamW", "Prieto et al.", "Higham", "Blanchard et al."], "also_reported_by": [], "alternates": {"html": "https://wpnews.pro/news/a-fix-for-softmax-and-cross-entropy", "markdown": "https://wpnews.pro/news/a-fix-for-softmax-and-cross-entropy.md", "text": "https://wpnews.pro/news/a-fix-for-softmax-and-cross-entropy.txt", "jsonld": "https://wpnews.pro/news/a-fix-for-softmax-and-cross-entropy.jsonld"}}