Grokking or Glitching? How Low-Precision Drives Slingshot Loss Spikes

In plain words

When a neural network is trained for a very long time without any regularisation, its loss sometimes jumps up sharply and then recovers, again and again. This is called the Slingshot Mechanism. Earlier work blamed the way the optimiser behaves. This paper argues that the cause is the limited precision of the numbers the computer uses. Once the model is very confident about the right answer, the computer rounds the gradient for that answer to exactly zero. The gradients for the wrong answers are not zero, and they push the model's internal averages in one direction, so those averages grow. When the right-answer gradient finally comes back, the optimiser reacts with a very large step, and the loss jumps. The authors also show that the same drift can make parameters grow even when no loss spike appears.

TL;DR

• Slingshot loss spikes in unregularised training with cross-entropy and Adam are argued to come from float32 rounding in the loss, not from the optimiser itself §1 §3.1.
• The rounding breaks the rule that gradients across classes sum to zero. The classifier mean drifts, and this feeds back on the feature mean, so both grow §3.2.
• Computing the loss in float64, projecting the logit gradient to zero sum, or placing BatchNorm directly before the classifier removes the spikes in the authors' runs §4.2.
• All numbers are author-reported on small setups (modular division, CIFAR-10, one nanoGPT run). No official code was found.

Problem

Under unregularised training with cross-entropy and Adam, networks show periodic loss spikes after long training. Prior work ties this to Edge of Stability or to optimiser and gradient-scale effects §1 §2. The paper claims the trigger is numerical: the drift it describes is a precondition for the spike, not a side effect of curvature §1 app. A.2.

The authors also report that the same numerical drift can make parameter and logit norms grow in practical settings where no visible spike appears §1 §4.3.

Method

The method is an analysis rather than a training algorithm. The chain of argument is:

1. Absorption in the loss. The softmax cross-entropy denominator is computed by log-sum-exp, so Z = z_m + log(sum_k exp(z_k - z_m)) eq. 1 §3.1. If the margin between the largest logit and the others exceeds (p-1) ln 2, the second term is absorbed and Z = z_m. For float32 (p = 24) that margin is about 16 §3.1 §2.
2. Softmax Collapse (SC). Under absorption, the gradient for the correct class r is exactly zero. Gradients for wrong classes stay small but non-zero, with g_k = exp(z_k - z_r) eq. 2 §3.1.
3. Broken zero-sum. In exact arithmetic the class gradients on the rows of the classifier sum to zero, so the global classifier mean W_G gets no gradient app. C.1. Under SC the sum equals eps * h, where eps = E[sum_{k != r} y_hat_k] is the residual probability on wrong classes. For a class-balanced batch, E[delta W_G] = -(eta * eps / K) * mu_G eq. 3 [thm. 1, §3.2]. The authors call this a drift of W_G against the global feature mean mu_G.
4. Feature drift. Under SC and the orthogonality assumption in Proposition 1, the gradient on a feature h has a component eps * W_G along W_G app. C.2 §3.2.
5. Feedback loop: Numerical Feature Inflation (NFI). The two means form a positive feedback loop. Their norms grow as (1 + eta * eps / sqrt(K))^t, and W_G becomes anti-parallel to mu_G eq. 6 eq. 7 eq. 8 [thm. 2, §3.2] app. C.3.
6. Trigger. As the means grow, some samples' margins fall below the absorption threshold. The correct-class gradient returns. Adam's second moment is small from the long quiet period, so the step is very large. The authors report the update at the spike step is about 50 times the previous epoch's §3.3.

# Softmax cross-entropy logit gradient in float32  [§3.1] [eq. 1] [eq. 2]
# Input:  logits (K,) float32; r = index of the correct class
# Output: gradient g (K,) with respect to the logits

def ce_logit_grad(logits, r):
    # Log-sum-exp denominator; the other logits are absorbed once
    # logits[r] - max(other logits) > 23 * ln 2 (about 16)  [§3.1]
    denom = logsumexp(logits)
    # Probabilities; probs[r] rounds to exactly 1.0 when absorbed
    probs = exp(logits - denom)
    g = probs
    # Softmax Collapse: the correct-class gradient is exactly 0  [§3.1]
    g[r] = probs[r] - 1.0
    return g

# Coupled drift of the global classifier mean and feature mean  [app. C.3]
# W_G, mu_G: (d,) vectors; eps: residual mass on wrong classes  [app. C.1]
# This is the paper's linear model of the feedback, not the training loop.

def nfi_step(W_G, mu_G, lr, eps, K):
    # Both updates use the old state, as in the matrix form of [app. C.3]
    alpha = lr * eps / K    # weight-mean step  [eq. 3]
    beta = lr * eps         # feature-mean step; K-free in the appendix
    new_W_G = W_G - alpha * mu_G
    new_mu_G = mu_G - beta * W_G
    return new_W_G, new_mu_G

# Zero-sum projection of the logit gradient  [§4.1]
# Input:  g (B, K): per-sample gradient with respect to the logits
# Output: g (B, K) whose entries sum to zero over classes for each sample

def zero_sum_project(g):
    # Subtract the class mean from every class entry  [§4.1]
    class_mean = mean(g, axis=1, keepdims=True)   # (B, 1)
    return g - class_mean                         # (B, K)

Architecture dependence: a spike needs eps to fall slowly. If eps decays fast, for example like 1/t, the growth is polynomial or logarithmic and no spike appears. The authors say ResNet18 learns fast enough to avoid this regime §4.1 fig. 2.

Mini-batch: with mini-batches, not all samples collapse at once, so no visible spike may appear. Collapsed samples can still drive the drift, and late-stage growth of mu_G and the logits continues §4.3.

Language models: softmax collapse also occurs in nanoGPT training, with tokens whose loss is exactly zero. Here higher precision in the loss increases logit growth, which the authors attribute to Zipf-distributed token frequencies creating a large W_G with the same sign as the feature mean §4.3 app. A.4.

Hyperparameters

Values are as stated in the paper. "not stated" means the source does not give the value.

Results

Author-reported. Setups: modular division (Transformer, MLP), CIFAR-10 (MLP, VGG11 with and without BN, ViT, ResNet18), and one nanoGPT run on FineWeb.

The paper does not report accuracy or benchmark numbers for any mitigation.

Paper vs code

No official code. The paper's TeX source has no repository link, and the abstract and arXiv metadata have none either. Searches: the project page does not exist (none linked); GitHub repository search for "numerical feature inflation", "slingshot grokking" and "slingshot float" found no repository belonging to the authors. A GitHub user search for "hanqing liu tsinghua" found one account whose six public repositories do not include this work. The NanoDo link in the paper refers to a third-party project and is not this paper's code. The paper's comparison code for another work (Stollenwerk et al.) is not this paper's code either.

Pitfalls

• Adam step recomputed from the stated inputs. The paper gives m_t = 1.46e-8, sqrt(v_t) = 2.7e-8 and eta = 1e-3, and reports an update of about 4e-4 §4.1. Those inputs give 1e-3 * 1.46e-8 / 2.7e-8, about 5.4e-4. The re-emerged gradient is also set to the threshold value 1.19e-7, a bound, not a measurement. Treat the estimate as an order of magnitude.
• Two different growth ratios at the spike. §3.3 says about 50 times the previous epoch; §4.1 says more than 40-fold. Neither is a measured constant in the appendix; do not hard-code one.
• Absorption depends on the kernel, not only on the formula. The threshold is derived for log-sum-exp in float32 §3.1. Whether g_r is exactly zero depends on the framework's fused cross-entropy, autocast, and the dtype of the logits. The mixed-precision LLM runs use BF16 with the loss in float32 tab. 4. Compare your loss path against the claimed threshold before you expect spikes.
• The growth law is an asymptotic sketch. Theorem 2 writes lim_{t->inf} ||W_G|| propto (1 + eta eps / sqrt(K))^t. That limit is of an unbounded quantity, and it needs the assumption that eps changes slowly compared with ||W_G|| [thm. 2, §3.2]. The paper checks this only by measurement fig. 2b.
• The sqrt(K) comes from an asymmetric step. The classifier-mean step is eta eps / K eq. 3. The feature-mean step in the appendix model is eta eps, with no 1/K and no batch factor, and it assumes that the centred classifier rows cancel over a batch app. C.3. The rate 1 + eta eps / sqrt(K) depends on that choice. Derive it again before you use it for prediction.
• Orthogonality and balance are assumptions. Proposition 1 needs W_G orthogonal to the centred classifier subspace app. C.2. The orthogonality lemma needs class-balanced data and a ReLU network in the NC state, citing Dang et al. app. C.2. The authors state that the LLM case is not class-balanced §4.3.
• Many results are proof sketches. The abstract says "proves". The main text labels Theorem 1 and Proposition 1 as "proof sketch" §3.2. The appendix derivations use the stated assumptions, not the full training loop.
• Zero-sum projection is one implementation. The projection is applied to the logit gradient g = grad_z L per sample §4.1. The text does not say whether it is applied before or after the batch mean, or how it interacts with a mean-reduced loss. The pseudocode above follows the text's formula.
• Figure captions omit datasets. The Adam-epsilon, LayerNorm and BatchNorm mitigation panels fig. 3 do not name the dataset in their captions. Check the appendix or your own rerun before you copy the setting.
• Batch size for the VGG11 mini-batch run conflicts with the appendix. §4.3 uses batch 256 for VGG11 and about 10^6 steps. Appendix D.2 lists full batch for all but ViT and 10,000 epochs. The mini-batch run is not in the appendix table.
• Adam betas differ across tables. Modular uses (0.9, 0.999) tab. 2; CIFAR-10 and the toy model use (0.9, 0.95) tab. 3 §4.1. The text tests (0.9, 0.999) and (0.9, 0.9) for the toy case §4.1. Use the table for each experiment.
• Full batch versus mini-batch. Most spikes were observed in full-batch or large-batch runs. Small mini-batches reduce the spike, because noise slows convergence to the absorption threshold and biases toward flatter minima app. A.3.2. Expect a smaller effect at small batch sizes.
• Smaller learning rates spike more. The authors report that smaller learning rates trigger spikes more readily, and they offer a heuristic explanation app. A.1.2. This is observed, not proven.
• Bias term speeds up spikes. The authors report that a bias in the classifier speeds the spike, because the weight and bias updates differ in scale when the feature norm is about 10^2 app. A.1.3.
• Underflow was ruled out only by clamping. Clamping the logit margin to 100 in float32 still shows spikes, so the authors attribute the effect to mantissa absorption, not exponent underflow app. A.1.4.
• Logit divergence is not reproduced. The authors could not reproduce the logit divergence of Wortsman et al. with their setup app. A.4. Their LLM observation follows Stollenwerk et al. §4.3 app. D.3. Do not treat the direction of the logit effect as settled across setups.
• Label smoothing changes the mechanism. With label smoothing, the Hessian stays bounded, and the spikes that appear are EOS-type, not precision-driven §4.2 app. A.2.
• Commented-out material is not in the paper. The TeX has several commented blocks (for example a logit-divergence paragraph and a float64-masking appendix). They are not part of the compiled text, and this card does not use them. Equation and figure numbers were counted from the TeX source; check them against the PDF before citing precisely.
• Modular and CIFAR settings are small. The authors' claims are on modular division (p = 97), CIFAR-10 with a 1,000-sample test set, and one nanoGPT run. The Results table says so; the general claim is not tested on large vision or language training.

Lineage

• Builds on: Thilak et al., "The Slingshot Mechanism" (arXiv:2206.04817), which gives the spike phenomenon and the Adam epsilon fix this paper analyses §2 §4.2.
• Builds on: Prieto et al., "Grokking at the Edge of Numerical Stability", which names Softmax Collapse and the float32 absorption threshold §2 §3.1. No card for it here.
• Compared against: Edge of Stability (Cohen et al. 2021) and its adaptive variants app. A.2; Nanda et al. 2023 on numerical instability §2; Wortsman et al. on output logit divergence app. A.4.
• Superseded by: not known.