# Grokking at the Edge of Numerical Stability

arXiv: https://arxiv.org/abs/2501.04697 · HTML: https://mlsanity.com/paper/2501.04697/

## Key facts

- Cited by 46 papers in the map (data through 14 Oct 2026).
- Builds on 33 of its 41 references that are in the map.
- First posted on arXiv 8 Jan 2025.

## Builds on (most cited first, 30 of 33)

- [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://mlsanity.com/paper/1810.04805/) (2018)
- [ImageNet Large Scale Visual Recognition Challenge](https://mlsanity.com/paper/1409.0575/) (2014)
- [PointNet: Deep Learning on Point Sets for 3D Classification and Segmentation](https://mlsanity.com/paper/1612.00593/) (2016)
- [Progress measures for grokking via mechanistic interpretability](https://mlsanity.com/paper/2301.05217/) (2023)
- [On Lazy Training in Differentiable Programming](https://mlsanity.com/paper/1812.07956/) (2018)
- [Grokking: Generalization Beyond Overfitting on Small Algorithmic Datasets](https://mlsanity.com/paper/2201.02177/) (2022)
- [Gradient Descent Maximizes the Margin of Homogeneous Neural Networks](https://mlsanity.com/paper/1906.05890/) (2019)
- [Hidden Progress in Deep Learning: SGD Learns Parities Near the Computational Limit](https://mlsanity.com/paper/2207.08799/) (2022)
- [Gradient descent aligns the layers of deep linear networks](https://mlsanity.com/paper/1810.02032/) (2018)
- [PPFNet: Global Context Aware Local Features for Robust 3D Point Matching](https://mlsanity.com/paper/1802.02669/) (2018)
- [Directional convergence and alignment in deep learning](https://mlsanity.com/paper/2006.06657/) (2020)
- [Omnigrok: Grokking Beyond Algorithmic Data](https://mlsanity.com/paper/2210.01117/) (2022)
- [A Toy Model of Universality: Reverse Engineering How Networks Learn Group Operations](https://mlsanity.com/paper/2302.03025/) (2023)
- [Explaining grokking through circuit efficiency](https://mlsanity.com/paper/2309.02390/) (2023)
- [Grokking as the Transition from Lazy to Rich Training Dynamics](https://mlsanity.com/paper/2310.06110/) (2023)
- [The Slingshot Mechanism: An Empirical Study of Adaptive Optimizers and the Grokking Phenomenon](https://mlsanity.com/paper/2206.04817/) (2022)
- [Risk and parameter convergence of logistic regression](https://mlsanity.com/paper/1803.07300/) (2018)
- [Grokking modular arithmetic](https://mlsanity.com/paper/2301.02679/) (2023)
- [Dichotomy of Early and Late Phase Implicit Biases Can Provably Induce Grokking](https://mlsanity.com/paper/2311.18817/) (2023)
- [Why Do We Need Weight Decay in Modern Deep Learning?](https://mlsanity.com/paper/2310.04415/) (2023)
- [Intrinsic Dimension, Persistent Homology and Generalization in Neural Networks](https://mlsanity.com/paper/2111.13171/) (2021)
- [Rotational Equilibrium: How Weight Decay Balances Learning Across Neural Networks](https://mlsanity.com/paper/2305.17212/) (2023)
- [Deep Networks Always Grok and Here is Why](https://mlsanity.com/paper/2402.15555/) (2024)
- [Grokking as a First Order Phase Transition in Two Layer Networks](https://mlsanity.com/paper/2310.03789/) (2023)
- [Grokking phase transitions in learning local rules with gradient descent](https://mlsanity.com/paper/2210.15435/) (2022)
- [Emergence in non-neural models: grokking modular arithmetic via average gradient outer product](https://mlsanity.com/paper/2407.20199/) (2024)
- [The Asymmetric Maximum Margin Bias of Quasi-Homogeneous Neural Networks](https://mlsanity.com/paper/2210.03820/) (2022)
- [Grokking Group Multiplication with Cosets](https://mlsanity.com/paper/2312.06581/) (2023)
- [Topological Generalization Bounds for Discrete-Time Stochastic Optimization Algorithms](https://mlsanity.com/paper/2407.08723/) (2024)
- [Grokking at the Edge of Linear Separability](https://mlsanity.com/paper/2410.04489/) (2024)

## Built on by (most cited first, 30 of 46)

- [Reinforcement Learning for Reasoning in Large Language Models with One Training Example](https://mlsanity.com/paper/2504.20571/) (2025)
- [Hierarchical Reasoning Model](https://mlsanity.com/paper/2506.21734/) (2025)
- [Less is More: Recursive Reasoning with Tiny Networks](https://mlsanity.com/paper/2510.04871/) (2025)
- [Provable Scaling Laws of Feature Emergence from Learning Dynamics of Grokking](https://mlsanity.com/paper/2509.21519/) (2025)
- [Muon Optimizer Accelerates Grokking](https://mlsanity.com/paper/2504.16041/) (2025)
- [Grokking in LLM Pretraining? Monitor Memorization-to-Generalization without Test](https://mlsanity.com/paper/2506.21551/) (2025)
- [Thinking with Looped Flows](https://mlsanity.com/paper/2609.11801/) (2026)
- [Adversarial Attacks Leverage Interference Between Features in Superposition](https://mlsanity.com/paper/2510.11709/) (2025)
- [To Grok Grokking: Provable Grokking in Ridge Regression](https://mlsanity.com/paper/2601.19791/) (2026)
- [What Happens During the Loss Plateau? Understanding Abrupt Learning in Transformers](https://mlsanity.com/paper/2506.13688/) (2025)
- [Mechanistic Insights into Grokking from the Embedding Layer](https://mlsanity.com/paper/2505.15624/) (2025)
- [The Geometric Inductive Bias of Grokking: Bypassing Phase Transitions via Architectural Topology](https://mlsanity.com/paper/2603.05228/) (2026)
- [Grokking or Glitching? How Low-Precision Drives Slingshot Loss Spikes](https://mlsanity.com/paper/2605.06152/) (2026)
- [Not All Explanations for Deep Learning Phenomena Are Equally Valuable](https://mlsanity.com/paper/2506.23286/) (2025)
- [Grokking as a Falsifiable Finite-Size Transition](https://mlsanity.com/paper/2603.24746/) (2026)
- [Grokking vs. Learning: Same Features, Different Encodings](https://mlsanity.com/paper/2502.01739/) (2025)
- [NeuralGrok: Accelerate Grokking by Neural Gradient Transformation](https://mlsanity.com/paper/2504.17243/) (2025)
- [OrthoGrad Improves Neural Calibration](https://mlsanity.com/paper/2506.04487/) (2025)
- [Egalitarian Gradient Descent: A Simple Approach to Accelerated Grokking](https://mlsanity.com/paper/2510.04930/) (2025)
- [Generalization at the Edge of Stability](https://mlsanity.com/paper/2604.19740/) (2026)
- [Distributional Spectral Diagnostics for Localizing Grokking Transitions](https://mlsanity.com/paper/2605.08237/) (2026)
- [Dead-Direction Conditioners: Gauge-Equivariant Preconditioning for Deep Networks](https://mlsanity.com/paper/2606.29176/) (2026)
- [Post-Grokking Collapse at the Representation-Readout Interface in Muon-Trained Transformers](https://mlsanity.com/paper/2608.07436/) (2026)
- [When Data Falls Short: Grokking Below the Critical Threshold](https://mlsanity.com/paper/2511.04760/) (2025)
- [Tuning the Implicit Regularizer of Masked Diffusion Language Models: Enhancing Generalization via Insights from k-Parity](https://mlsanity.com/paper/2601.22450/) (2026)
- [The Implicit Bias of Logit Regularization](https://mlsanity.com/paper/2602.12039/) (2026)
- [Dynamical Systems Theory Behind a Hierarchical Reasoning Model](https://mlsanity.com/paper/2603.22871/) (2026)
- [First-Passage Prediction of Grokking Delay: ACalibrated Law under AdamW with Causal Validation](https://mlsanity.com/paper/2605.18845/) (2026)
- [Dead Directions: Geometric Singular Learning](https://mlsanity.com/paper/2606.05957/) (2026)
- [What Does the Weight Norm Control in Grokking? Logit-Scale Mediation under Cross-Entropy](https://mlsanity.com/paper/2606.18465/) (2026)

## Abstract (arXiv)

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