Abstract
Grokking is a puzzling phenomenon in neural networks where full generalization occurs only after a substantial delay following the complete memorization of the training data. Previous research has linked this delayed generalization to representation learning driven by weight decay, but the precise underlying dynamics remain elusive. In this paper, we argue that post-memorization learning can be understood through the lens of constrained optimization: gradient descent effectively minimizes the weight norm on the zero-loss manifold. We formally prove this in the limit of infinitesimally small learning rates and weight decay coefficients. To further dissect this regime, we introduce an approximation that decouples the learning dynamics of a subset of parameters from the rest of the network. Applying this framework, we derive a closed-form expression for the post-memorization dynamics of the first layer in a two-layer network. Experiments confirm that simulating the training process using our predicted gradients reproduces both the delayed generalization and representation learning characteristic of grokking.
Abstract (translated)
Grokking 是神经网络中的一种令人困惑的现象，即在完全记住训练数据之后，会出现一段延迟期，在此期间不会发生泛化，直到某个时刻才突然实现全面的泛化。先前的研究已经将这种延迟泛化与受权重衰减驱动的表现学习联系起来，但其背后的确切动态机制仍然不清楚。在这篇论文中，我们主张可以通过约束优化的视角来理解记忆后学习：在零损失流形上，梯度下降有效最小化了权重范数。我们在极限条件下（即学习率和权重衰减系数无限小的情况下）正式证明了这一点。为了进一步剖析这一机制，我们引入了一种近似方法，将一部分参数的学习动态与网络其余部分解耦。通过应用这个框架，我们推导出了双层网络中第一层在记忆后阶段的封闭形式表达式。实验结果证实，使用我们的预测梯度模拟训练过程可以再现 grokking 的延迟泛化和表示学习特性。 简单来说，这篇论文提出了一种理解神经网络中 Grokking 现象的新方法，即通过约束优化来解析权重衰减对模型泛化能力的影响，并且理论推导与实验结果吻合良好。
URL
https://arxiv.org/abs/2511.01938
