Grokking（顿悟式泛化）是指模型在长期过拟合后突然实现泛化的惊人现象，对我们关于深度学习的认知构成了挑战。尽管学界在理解Grokking现象方面已取得显著进展，但泛化延迟背后的原因及其对正则化的依赖性仍不明确。在本研究中，我们提出，在无正则化的条件下，Grokking任务会将模型推至数值稳定性的边缘，在Softmax函数中引入浮点误差，我们将该现象称为Softmax坍塌（Softmax Collapse, SC）。我们证明，SC会阻碍Grokking现象的发生，而缓解SC则可在无正则化的情况下实现Grokking。在探究SC的根本原因时，我们发现，在越过过拟合临界点后，梯度会与我们所称的朴素损失最小化（naïve loss minimization, NLM）方向高度对齐。梯度的这一分量不会改变模型的预测结果，但会通过缩放logits（通常是沿权重当前方向缩放权重）来降低损失。我们证明，logits的这种缩放解释了Grokking特有的泛化延迟现象，且最终会引发SC，使后续学习停滞。为验证我们的假设，我们提出了两项应对Grokking任务挑战的核心贡献：StableMax，一种可防止SC、实现无正则化Grokking的新型激活函数；以及⊥Grad，一种通过完全阻止NLM来促进Grokking任务快速泛化的训练算法。这些贡献为Grokking现象提供了新的见解，阐明了其泛化延迟、对正则化的依赖以及现有诱导Grokking方法的有效性。本文代码可在https://github.com/LucasPrietoAl/grokking-at-the-edge-of-numerical-stability获取。
英文摘要：
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ïve 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.
Raymond and Beverly Sackler School of Physics and Astronomy, Tel Aviv University, Tel Aviv 69978, Israel（雷蒙德与贝弗利·萨克勒物理与天文学学院，特拉维夫大学）；École Polytechnique Fédérale de Lausanne (EPFL), Lausanne, Switzerland（洛桑联邦理工学院（EPFL），洛桑，瑞士）
INRIA, CNRS, Département d’Informatique de l’Ecole Normale Supérieure / PSL, France（法国国家信息与自动化技术研究院（INRIA）、法国国家科学研究中心（CNRS）、巴黎综合理工学院计算机系 / PSL研究院）；Department of Computing, Imperial College London, United Kingdom（伦敦帝国理工学院计算机系，英国）
Carnegie Mellon University（卡内基梅隆大学）；Bosch AI（博世人工智能）；Determined AI
We empirically demonstrate that full-batch gradient descent on neural network training objectives typically operates in a regime we call the Edge of Stability. In this regime, the maximum eigenvalue of the training loss Hessian hovers just above the numerical value $2 / \text{(step size)}$, and the training loss behaves non-monotonically over short timescales, yet consistently decreases over long timescales. Since this behavior is inconsistent with several widespread presumptions in the field of optimization, our findings raise questions as to whether these presumptions are relevant to neural …
Recent studies of gradient descent with large step sizes have shown that there is often a regime with an initial increase in the largest eigenvalue of the loss Hessian (progressive sharpening), followed by a stabilization of the eigenvalue near the maximum value which allows convergence (edge of stability). These phenomena are intrinsically non-linear and do not happen for models in the constant Neural Tangent Kernel (NTK) regime, for which the predictive function is approximately linear in the parameters. As such, we consider the next simplest class of predictive models, namely those that are…
Conflicting Biases at the Edge of Stability: Norm versus Sharpness Regularization
Munich Center for Machine Learning (MCML)（慕尼黑机器学习中心）；Prusa Research（普拉萨研究公司）；Institute for Robotics and Mechatronics, DLR-German Aerospace Center（德国航空航天中心机器人与机电研究所）
Technion - Israel Institute of Technology（以色列理工学院）；Habana-Labs（哈瓦那实验室）
Background: Recent developments have made it possible to accelerate neural networks training significantly using large batch sizes and data parallelism. Training in an asynchronous fashion, where delay occurs, can make training even more scalable. However, asynchronous training has its pitfalls, mainly a degradation in generalization, even after convergence of the algorithm. This gap remains not well understood, as theoretical analysis so far mainly focused on the convergence rate of asynchronous methods. Contributions: We examine asynchronous training from the perspective of dynamical stabili…
Variational Learning Finds Flatter Solutions at the Edge of Stability
Michigan State University（密歇根州立大学）；Institute of Science Tokyo（东京科学研究所）；RIKEN Center for AI Project（理化学研究所人工智能项目中心）；Georgia Institute of Technology（佐治亚理工学院）
Muon Sublates the Edge of Stability in LLM Pretraining
School of Mathematical Sciences, Shanghai Jiao Tong University（上海交通大学数学科学学院）；Alibaba Inc.（阿里巴巴集团）；School of Mathematical Sciences, Institute of Natural Sciences and MOE-LSC, Shanghai Jiao Tong University（上海交通大学数学科学学院、自然科学研究院及教育部重点实验室）
Muon 优化器在 LLM 预训练中打破经典边缘稳定性耦合，实验表明其损失中性边界与时间对齐独立，支持分裂 EoS 图景。
Recent findings (e.g., arXiv:2103.00065) demonstrate that modern neural networks trained by full-batch gradient descent typically enter a regime called Edge of Stability (EOS). In this regime, the sharpness, i.e., the maximum Hessian eigenvalue, first increases to the value 2/(step size) (the progressive sharpening phase) and then oscillates around this value (the EOS phase). This paper aims to analyze the GD dynamics and the sharpness along the optimization trajectory. Our analysis naturally divides the GD trajectory into four phases depending on the change of the sharpness. We empirically id…
Deep learning experiments by Cohen et al. [2021] using deterministic Gradient Descent (GD) revealed an Edge of Stability (EoS) phase when learning rate (LR) and sharpness (i.e., the largest eigenvalue of Hessian) no longer behave as in traditional optimization. Sharpness stabilizes around $2/$LR and loss goes up and down across iterations, yet still with an overall downward trend. The current paper mathematically analyzes a new mechanism of implicit regularization in the EoS phase, whereby GD updates due to non-smooth loss landscape turn out to evolve along some deterministic flow on the manif…
Recently, the stability of graph filters has been studied as one of the key theoretical properties driving the highly successful graph convolutional neural networks (GCNs). The stability of a graph filter characterizes the effect of topology perturbation on the output of a graph filter, a fundamental building block for GCNs. Many existing results have focused on the regime of small perturbation with a small number of edge rewires. However, the number of edge rewires can be large in many applications. To study the latter case, this work departs from the previous analysis and proves a bound on t…
Recently, researchers observed that gradient descent for deep neural networks operates in an ``edge-of-stability'' (EoS) regime: the sharpness (maximum eigenvalue of the Hessian) is often larger than stability threshold $2/η$ (where $η$ is the step size). Despite this, the loss oscillates and converges in the long run, and the sharpness at the end is just slightly below $2/η$. While many other well-understood nonconvex objectives such as matrix factorization or two-layer networks can also converge despite large sharpness, there is often a larger gap between sharpness of the endpoint and $2/η$.…
Learning Dynamics of Deep Linear Networks Beyond the Edge of Stability
Computational Mathematics Science and Engineering, Michigan State University（密歇根州立大学计算数学科学与工程系）；Department of Electrical Engineering & Computer Science, University of Michigan（密歇根大学电气与计算机工程系）
University of Oxford（牛津大学）；Swiss Data Science Center（瑞士数据科学中心）
Graph neural networks are experiencing a surge of popularity within the machine learning community due to their ability to adapt to non-Euclidean domains and instil inductive biases. Despite this, their stability, i.e., their robustness to small perturbations in the input, is not yet well understood. Although there exists some results showing the stability of graph neural networks, most take the form of an upper bound on the magnitude of change due to a perturbation in the graph topology. However, the change in the graph topology captured in existing bounds tend not to be expressed in terms of…
