REVIEW 3 major objections 4 minor 6 cited by
Grokking at the Edge of Numerical Stability
T0 review · 3 major / 4 minor · reviewed 2026-08-10 · deepseek-v4-flash
Pith's one-line read Grokking without regularization fails because floating-point errors in Softmax zero out gradients, and two targeted fixes restore it.
desk verdict A genuinely new mechanistic account of grokking, with clean interventions and an honest limitation that keeps the causal loop from fully closing for biased networks. read the letter →
The pith
A machine-rendered reading of the paper's core claim, the machinery that carries it, and where it could break.
The reading
What carries the argument
The central object is the decomposition of the gradient after overfitting into two components: the NLM component, which for a positively homogeneous network (a network whose outputs scale by a constant when all weights are scaled by a constant) points along the current weight vector $\theta$ and only rescales the logits (so $f(\theta+d_{\mathrm{NLM}}(\theta);x)=c f(\theta;x)$ for some $c>1$), and the orthogonal component that actually changes predictions. The second piece of machinery is the absorption-error condition that defines SC: when the exponent gap between the true-class exponential and every other term in the Softmax sum exceeds the significand precision, the sum collapses to the true-class term, the cross-entropy loss is exactly zero, and the gradients from correctly classified samples vanish. StableMax replaces the exponential in Softmax with the piecewise function $s(x)=x+1$ for $x\ge 0$ and $s(x)=1/(1-x)$ for $x<0$, which grows linearly rather than exponentially and avoids the extreme summands that cause absorption. $\perp$Grad projects the gradient onto the hyperplane orthogonal to the weight vector, $\nabla_\perp L(\theta_t)=\nabla L(\theta_t)-\frac{\theta_t^\top\nabla L(\theta_t)}{\theta_t^\top\theta_t}\theta_t$, removing the NLM component so only the prediction-changing part drives updates. Together these two mechanisms carry the argument: SC explains why training stops, NLM explains the delay before it stops, and each intervention isolates one link of the chain.
What would settle it
Measure, after 100% training accuracy on modular addition, whether the radial gradient component that ⊥Grad removes ever changes predicted labels on the test set; if it does, that component is not purely logit scaling and the NLM account of the delay needs revision. A second check: train the same biased MLP with ⊥Grad on a loss, such as mean-squared error with bounded targets, where scaling logits does not reduce the loss; if it still generalizes faster than vanilla SGD, the speedup is not explained by blocking NLM.
Extended reading notes
Core claim
The paper's central claim is that grokking fails without regularization because training runs into Softmax Collapse (SC), a floating-point absorption error in the Softmax sum. For a correctly classified sample, once the true-class logit is so large that $\sum_k e^{z_k} \doteq e^{z_y}$ under floating-point arithmetic, the cross-entropy loss evaluates to exactly zero and the gradient contribution from that sample vanishes, so learning stops even though the model has not generalized. SC is the last step of a longer chain: after reaching 100% training accuracy, cross-entropy gradients align with the naive loss minimization (NLM) direction, $d_{\mathrm{NLM}}(\theta)=\alpha\theta$ for positively homogeneous networks, which scales all logits by a constant and lowers the loss without altering predictions; the same alignment is observed empirically in the biased MLPs and transformers actually trained. Two interventions support the chain: StableMax, a modified Softmax whose unbounded 'exponential' grows only linearly, prevents SC and produces grokking without regularization, while ⊥Grad, which updates only the gradient component orthogonal to the weight vector, removes the delay in generalization. In this account, weight decay works, MSE loss works on shallow networks, and label smoothing behaves differently for a common reason: whether the loss can be reduced indefinitely by scaling logits, and whether floating-point collapse can be triggered.
Load-bearing premise
The load-bearing premise is that, for the biased MLPs and transformers actually trained, the gradient component pointing along the current weights is exactly the part that only scales the logits without changing predictions; the formal proof covers only bias-free networks in which scaling all weights scales all outputs, and for the biased models the paper relies on empirical alignment rather than a theorem.
Editorial extensions
If this is right
- Weight decay induces grokking by counteracting NLM: it pulls weights back along the same radial direction, so logit scaling stops once the loss reduction no longer outweighs the regularization penalty.
- MSE loss groks on shallow networks without regularization because scaling logits cannot reduce an MSE loss indefinitely, so the NLM mechanism is absent.
- Temperature scaling and float64 only delay SC rather than prevent it; only sub-exponential replacements like StableMax avoid the collapse entirely.
- ⊥Grad reaches 100% test accuracy with no preceding overfitting phase and, in the tested settings, outperforms the best-tuned weight decay with no extra hyperparameters.
- Slingshot spikes may be a mechanism that prevents full SC, explaining why adaptive optimizers can occasionally grok without weight decay.
Reading between the lines
- The paper's diagnosis would predict that grokking curves shift with floating-point precision in a quantifiable way: identical runs in float16, float32, and float64 should stall at precisely the epochs where the SC fraction crosses a threshold; running the same sweep in bfloat16 would be a direct check.
- Because NLM is a property of losses that are monotonically decreased by logit scaling, any training scheme that bounds or penalizes logit norms should eliminate the delay; the paper's appendix logit-regularization experiment supports this, and the same logic suggests fixed-norm or normalized-gradient variants would grok without either StableMax or ⊥Grad.
- The formal gap for biased networks matters: if a quasi-homogeneous analysis shows that the radial gradient component is not purely logit-scaling in the presence of biases, the causal story would need to be revised, so a theoretical characterization of quasi-homogeneous models is the natural next step.
- In mixed-precision training beyond grokking benchmarks, SC may be an unexamined regularizer that silently zeroes gradients for well-classified examples; replacing Softmax with a sub-exponential variant could change training dynamics in settings where losses are driven near zero.
Editorial analysis
A structured set of objections, weighed in public.
Referee Report
Summary. The paper proposes that grokking failures without regularization are caused by floating-point absorption errors in the Softmax, termed Softmax Collapse (SC): once the correct-class logit is much larger than the others, the softmax denominator collapses to that logit, the cross-entropy loss becomes exactly zero, and gradients from correctly classified samples vanish, halting learning. The paper further introduces Naive Loss Minimization (NLM), a gradient component aligned with scaling the weights/logits that reduces CE loss after 100% training accuracy without changing predictions, and argues that NLM both delays generalization and eventually causes SC. Two interventions are proposed: StableMax, a numerically stable replacement for Softmax that enables grokking without regularization, and ⊥Grad, an optimizer that projects out the radial gradient component and removes the overfitting delay. The claims are supported by experiments on modular arithmetic, sparse parity, and MNIST, with additional validation on GPT-2 and ResNet. The paper also offers post-hoc explanations for why weight decay and MSE loss induce grokking.
Significance. If the central claims hold, the paper reframes a widely studied phenomenon as a numerical-stability problem rather than purely an implicit-bias or regularization effect, and it provides two simple, actionable interventions. The strengths are the explicit formal definitions (SC and NLM), the direct causal interventions (artificial SC, StableMax, ⊥Grad), the breadth of experimental settings, and the released code. The SC portion is well supported: the onset of SC coincides with stalled generalization, artificial SC stops generalization, and StableMax restores grokking. However, the NLM causal story has a load-bearing formal gap for the biased architectures actually used in the main experiments, as the paper itself acknowledges in its Limitations paragraph. The single-seed main figures also leave the quantitative robustness of the claims under-addressed.
major comments (3)
- [Sec. 4.2, Def. 5; Sec. 5.1, Def. 7; Fig. 5; Limitations] The NLM direction is defined by Eq. (9), which requires f(θ+d_NLM(θ);x)=c f(θ;x). For positively homogeneous networks this is satisfied by d_NLM(θ)=αθ, as shown in Sec. 4.2. However, the MLPs and transformers actually trained include bias terms (Fig. 5b, 5c), and for biased networks f(cθ;x) is not generally equal to c f(θ;x); scaling all parameters scales hidden pre-activations by different powers of c and leaves bias terms with their own scaling, so the defining logit-scaling property is not satisfied. The paper acknowledges this in the Limitations paragraph but still uses ⊥Grad, which removes the projection of the gradient onto the full parameter vector θ (Eq. 12), and interprets the resulting speedup as specifically preventing NLM. Fig. 5 measures only cosine similarity between weights and gradients, which does not verify the logit-scaling property. A direct test is needed: for trained biased models, measure how f((1+α)θ;x) compares to c f(θ;x), or measure the change in logits when perturbing along θ; alternatively, restrict the NLM claims to homogeneous models and modify ⊥Grad to project onto the last-layer parameters only. Without this, the claim that ⊥Grad removes NLM on the architectures studied is not established.
- [Figs. 2, 4, 6, 7] The main empirical claims—SC onset coinciding with stalled generalization, StableMax inducing grokking, and ⊥Grad removing the delay—are each illustrated with single training runs and no error bars or multiple seeds. Grokking dynamics are known to be seed-sensitive, and the quantitative timing of SC onset and generalization can vary across runs. Reporting mean curves with standard deviations or at least a few seeds would make the load-bearing qualitative claims robust to stochasticity. Table 1 reports seeds only for the realistic settings, not for the central grokking experiments.
- [App. B.1, Fig. 8] The description of the artificial SC intervention is internally inconsistent. The text says the goal is to 'set the gradients from the correct classes to zero', but the implementation is described as 'multiplying the logits for the right classes by 0'. Multiplying the correct logit by zero changes the softmax probability and produces a large negative gradient for the correct class, not a zero gradient. If the intervention actually zeroes the gradients, the sentence should say so; if it zeroes the logits, the experiment does not implement SC as defined in Def. 3 and its negative result would not support the SC explanation.
minor comments (4)
- [Fig. 5 caption] The caption sentence 'MLPs with (a) and without (b) bias terms' is reversed relative to the subcaptions, which correctly label (a) as without bias and (b) as with bias; the text in Sec. 4.2 refers to Fig. 5b for biased models, so the caption should be corrected.
- [App. A, proof of Prop. 2] The proof of Prop. 2 is more convoluted than necessary and the normalization step is not clearly justified. The result follows directly from Cauchy-Schwarz: ⟨−∇⊥L,∇L⟩ = −(‖∇L‖² − (θ·∇L)²/‖θ‖²) ≤ 0, with equality iff ∇L is parallel to θ. Consider replacing the current argument with this two-line derivation.
- [Fig. 4, right panel] The right panel of Fig. 4 is labeled '2-hot input' and 'random binary input', while Sec. 4.1 describes a random binary vector of dimension 14. Please clarify whether the inputs are 2-hot in 14 dimensions or dense random binary vectors; the two descriptions are not equivalent.
- [Sec. 5.1, Def. 7] The definition of ⊥Grad assumes θ_t ≠ 0 in the projection formula (Eq. 12). Since training is initialized with small random weights and the projection is well defined after the first step, this is not a practical issue, but the zero-parameter case should be noted for completeness.
Circularity Check
No significant circularity: the SC and NLM claims are tested by direct interventions (StableMax, ⊥Grad, artificial SC) and rest mostly on external results; the quasi-homogeneous gap is a rigor limitation, not a circular reduction.
full rationale
The paper's derivation chain is not circular. Softmax Collapse is operationally defined (Def. 3) and independently intervened upon: increasing FP precision, replacing Softmax by StableMax (Def. 4), and artificially zeroing correct-class gradients (App. B.1) all produce the predicted effects, so SC's causal role is not assumed by construction. NLM is formally defined (Def. 5), and the result that d_NLM = αθ for homogeneous networks follows from positive homogeneity (Def. 6) and standard CE-loss properties, citing external work (Lyu & Li 2020) for directional convergence; the paper's own ⊥Grad intervention (Def. 7) directly removes the radial component and shows faster generalization, which is a falsifiable test rather than a renamed fit. Explaining weight decay and MSE success via NLM is post-hoc interpretation, not a fitted input called a prediction. The one substantive weakness is flagged by the authors themselves in Sec. 7: the formal argument covers homogeneous/approximately homogeneous networks, and the extension to biased MLPs and transformers is supported only by last-layer homogeneity and empirical gradient-weight alignment (Fig. 5), so the claim that ⊥Grad specifically removes the exact NLM direction on biased models is not rigorously closed. This is a correctness/rigor gap, not circularity. The only self-citations, e.g. "the link between training trajectories and generalization is already established in prior art (Birdal et al., 2021; Andreeva et al., 2024)", are used as general context and are not load-bearing for the paper's central claims.
Assumptions & free parameters
assumptions (3)
- domain assumption The studied ReLU MLPs (without bias) are positively homogeneous and transformers are approximately homogeneous, so scaling all weights by c > 0 scales all logits by c^L.
- domain assumption Floating point addition follows the absorption error model of Def. 1, and standard Softmax implementations compute the sum in a way that can absorb small exponentials, yielding SC as in Eq. (2).
- domain assumption The gradient component orthogonal to the weights (after removing NLM) is a sufficient direction for generalization in grokking tasks.
Cite this review
Pith. "Pith review of Grokking at the Edge of Numerical Stability." pith.science (2026). https://pith.science/paper/RHDJK25A
@misc{pith2026250104697,
author = {Pith},
title = {Pith review of: Grokking at the Edge of Numerical Stability},
year = {2026},
howpublished = {\url{https://pith.science/paper/RHDJK25A}},
note = {Machine review of arXiv:2501.04697}
}
abstract
Grokking, the sudden generalization that occurs after prolonged overfitting, is a surprising phenomenon challenging our understanding of deep learning. Although significant progress has been made in understanding grokking, the reasons behind the delayed generalization and its dependence on regularization remain unclear. In this work, we argue that without regularization, grokking tasks push models to the edge of numerical stability, introducing floating point errors in the Softmax function, which we refer to as Softmax Collapse (SC). We demonstrate that SC prevents grokking and that mitigating SC enables grokking without regularization. Investigating the root cause of SC, we find that beyond the point of overfitting, the gradients strongly align with what we call the na\"ive loss minimization (NLM) direction. This component of the gradient does not alter the model's predictions but decreases the loss by scaling the logits, typically by scaling the weights along their current direction. We show that this scaling of the logits explains the delay in generalization characteristic of grokking and eventually leads to SC, halting further learning. To validate our hypotheses, we introduce two key contributions that address the challenges in grokking tasks: StableMax, a new activation function that prevents SC and enables grokking without regularization, and $\perp$Grad, a training algorithm that promotes quick generalization in grokking tasks by preventing NLM altogether. These contributions provide new insights into grokking, elucidating its delayed generalization, reliance on regularization, and the effectiveness of existing grokking-inducing methods. Code for this paper is available at https://github.com/LucasPrietoAl/grokking-at-the-edge-of-numerical-stability.
Figures
Figures from the paper (14 more)
Forward citations
Cited by 6 Pith papers
-
Decomposing Prediction Mechanisms for In-Context Recall
In a toy in-context recall task, label-based task initiation and observation-based continuation are distinct mechanisms with separate emergence times, and the same first-token versus second-token gap appears in an OLM...
-
Post-Grokking Collapse at the Representation-Readout Interface in Muon-Trained Transformers
Muon-trained transformers grok modular addition and then lose generalization because the hidden representation and the output readout drift apart; freezing the readout and embeddings after grokking removes the collapse.
-
Grokking Is Conditional and Fragile: A Fully-Tractable, Multi-Seed Study at 12K Parameters
In a fully tractable 12K Llama-style model, grokking is a conditional fragile phase transition gated by coverage (tracking modulus more than structure), weight decay, and floating-point reduction order, so evidence mu...
-
Grokking vs. Learning: Same Features, Different Encodings
Grokked and steadily trained models learn the same features, but steady training can produce much more compressible models in a parameter regime that grokking does not reach.
-
Mechanistic Insights into Grokking from the Embedding Layer
Trainable embeddings in a simple MLP cause delayed generalization (grokking) on modular arithmetic, and a higher embedding learning rate plus balanced sampling accelerates it.
-
Not All Explanations for Deep Learning Phenomena Are Equally Valuable
A position paper arguing that narrow, puzzle-solving explanations of deep learning edge case phenomena are low-value, and that these phenomena should instead be used to stress-test broad explanatory theories.
Reference graph
Works this paper leans on
-
[4]
Deep networks always grok and here is why
Ahmed Imtiaz Humayun, Randall Balestriero, and Richard Baraniuk. Deep networks always grok and here is why. arXiv preprint arXiv:2402.15555,
-
[6]
Imagenet classification with deep convolutional neural networks
11 Published as a conference paper at ICLR 2025 Alex Krizhevsky, Ilya Sutskever, and Geoffrey E Hinton. Imagenet classification with deep convolutional neural networks. In Advances in neural information processing systems, volume 25, pp. 1097–1105,
work page 2025
-
[7]
Language Models "Grok" to Copy
Ziming Liu, Eric J Michaud, and Max Tegmark. Omnigrok: Grokking beyond algorithmic data. InThe Eleventh International Conference on Learning Representations, 2023a. Ziming Liu, Ziqian Zhong, and Max Tegmark. Grokking as simplification: A nonlinear complexity perspective. In UniReps: the First Workshop on Unifying Representations in Neural Models, 2023b. A...
-
[8]
Emergence in non-neural models: grokking modular arithmetic via average gradient outer product
Neil Mallinar, Daniel Beaglehole, Libin Zhu, Adityanarayanan Radhakrishnan, Parthe Pandit, and Mikhail Belkin. Emergence in non-neural models: grokking modular arithmetic via average gradient outer product. arXiv preprint arXiv:2407.20199,
-
[9]
William Merrill, Vivek Ramanujan, Yoav Goldberg, Roy Schwartz, and Noah A. Smith. Parameter norm growth during training of transformers. CoRR, abs/2010.09697,
work page Pith review arXiv 2010
-
[10]
Grokking: Generalization beyond overfitting on small algorithmic datasets
Alethea Power, Yuri Burda, Harri Edwards, Igor Babuschkin, and Vedant Misra. Grokking: Generalization beyond overfitting on small algorithmic datasets. arXiv preprint arXiv:2201.02177,
-
[12]
Explaining grokking through circuit efficiency
Vikrant Varma, Rohin Shah, Zachary Kenton, J´anos Kram´ar, and Ramana Kumar. Explaining grokking through circuit efficiency. arXiv preprint arXiv:2309.02390,
-
[13]
Achieving Margin Maximization Exponentially Fast via Progressive Norm Rescaling
Mingze Wang, Zeping Min, and Lei Wu. Achieving margin maximization exponentially fast via progressive norm rescaling. arXiv preprint arXiv:2311.14387,
Show all 18 references
-
[14]
A presents the proofs for the propositions in the paper, App
12 Published as a conference paper at ICLR 2025 APPENDIX In support of the main paper, App. A presents the proofs for the propositions in the paper, App. B includes additional findings that support our main results, and App. D provides further discussion on conditions that lea...
2025
-
[15]
∇⊥L(θt) = 0 can also be the case if ∇L(θt) =0, which corresponds to the loss function being at a local optimum
If ∇L(θt) ̸= 0, this corresponds to the condition where the gradient is in the same direction with the parameter vector. ∇⊥L(θt) = 0 can also be the case if ∇L(θt) =0, which corresponds to the loss function being at a local optimum. B A DDITIONAL FINDINGS B.1 F URTHER EVIDENCE...
2025
-
[18]
clean up
Grokking is induced via StableMax and without weight decay. Taylor approximation of the Softmax. We have introduced StableMax as a change to the Softmax that leads to grokking without regularization. The motivation behind this is to prevent values in the sum of the Softmax tha...
2023
-
[113]
For a gradient update using regu- lar gradient descent θt+1 = θt − η∇L(θt) it is easy to see that ||θt+1 − θt|| →0 as ||∇L(θt)|| →0
Unexplored in the main paper, NLM also has the effect of reduc- ing the effective learning rate. For a gradient update using regu- lar gradient descent θt+1 = θt − η∇L(θt) it is easy to see that ||θt+1 − θt|| →0 as ||∇L(θt)|| →0. This problem has been observed before when trai...
2020
-
[2015]
Dashiell Stander, Qinan Yu, Honglu Fan, and Stella Biderman
doi: 10.1007/s11263-015-0816-y. Dashiell Stander, Qinan Yu, Honglu Fan, and Stella Biderman. Grokking group multiplication with cosets. In Forty-first International Conference on Machine Learning,
-
[2021]
Andrey Gromov
https://transformer-circuits.pub/2021/framework/index.html. Andrey Gromov. Grokking modular arithmetic. arXiv preprint arXiv:2301.02679,
2021 arXiv
-
[2022]
Grokking at the edge of linear separability
Alon Beck, Noam Levi, and Yohai Bar-Sinai. Grokking at the edge of linear separability. arXiv preprint arXiv:2410.04489,
-
[2023]
Why do we need weight decay in modern deep learning? arXiv preprint arXiv:2310.04415,
Francesco D’Angelo, Maksym Andriushchenko, Aditya Varre, and Nicolas Flammarion. Why do we need weight decay in modern deep learning? arXiv preprint arXiv:2310.04415,
-
[2024]
Risk and parameter convergence of logistic regression
Ziwei Ji and Matus Telgarsky. Risk and parameter convergence of logistic regression. arXiv preprint arXiv:1803.07300,
-
[6000]
Our scheduler is simi- lar to the one in Lyu & Li (2020) except at each step we divide the learning rate by the norm of the full gradient, instead of the loss
B.2 SGD WITH LEARNING RATE SCHEDULING To show that our results are not due to the inductive bias of adaptive moments in optimizers like AdamW, we replicate some of the AdamW results using SGD with a learning rate scheduler. Our scheduler is simi- lar to the one in Lyu & Li (20...
2020
Reviewed August 10, 2026 · model on record in the stance chip above.
Discussion (0). Continue with ORCID to comment.