Provable Scaling Laws of Feature Emergence from Learning Dynamics of Grokking

尽管“顿悟”（grokking）现象，即延迟泛化，已被广泛研究，但对于复杂结构化输入而言，是否存在着一种数学框架，能够刻画出哪些特征会涌现、这些特征如何以及在何种条件下出现，并且与训练过程中的梯度动力学密切相关，这一问题仍然悬而未决。我们提出了一种名为 $\mathbf{Li_2}$ 的新框架，用于描述两层非线性网络在发生顿悟行为时的三个关键阶段：（I）**懒惰学习**（Lazy learning）、（II）**独立特征学习**（Independent feature learning）和（III）**交互式特征学习**（Interactive feature learning）。在懒惰学习阶段，顶层网络会过拟合于随机的隐层表示，模型表现出记忆行为。得益于懒惰学习和权重衰减，从顶层反向传播的梯度 $G_F$ 现在携带了关于目标标签的信息，并具有特定结构，使得每个隐层节点能够**独立地**学习其自身的表示。有趣的是，这种独立学习的动力学过程恰好对应于某个能量函数 $E$ 的**梯度上升**过程，而该能量函数的局部极大值点正是那些正在涌现的特征。我们在群算术任务中研究了由这些局部最优诱导出的特征是否具备泛化能力、其表示能力如何，以及它们随样本量变化的趋势。当在学习后期隐层节点开始相互作用时，我们从理论上证明了 $G_F$ 如何发生变化，从而将学习重点转向那些尚未被掌握的缺失特征。本研究从梯度动力学的第一性原理出发，阐明了权重衰减、学习率和样本量等关键超参数在顿悟过程中的作用，推导出了特征涌现、记忆和泛化可证明的缩放律，并揭示了近期优化器（如 Muon）之所以有效的根本原因。我们的分析还可推广至多层网络结构。

论文试图解决深度学习中‘grokking’（延迟泛化）现象的理论解释问题，即模型在训练初期过拟合、后期突然泛化的机制。特别是针对复杂结构化输入，当前缺乏一个数学框架来刻画特征如何涌现、在何种条件下发生，以及其与梯度动力学的关系。这是一个新兴且尚未充分探索的问题。

提出名为$\mathbf{Li_2}$的新框架，将两层非线性网络的grokking过程分解为三个阶段：(I) 懒惰学习（Lazy learning），(II) 独立特征学习（Independent feature learning），(III) 交互式特征学习（Interactive feature learning）。关键在于揭示了顶层梯度$G_F$携带标签信息，驱动隐藏单元独立地通过梯度上升优化一个能量函数$E$，其局部极大值对应可泛化的特征；并在后期阶段分析了特征交互和缺失特征的补全机制。

从梯度动力学第一性原理出发，解释了weight decay、学习率、样本量等超参数在grokking中的作用，推导出特征出现、记忆化与泛化的可证明缩放律；在群算术任务上验证理论，并揭示了Muon等现代优化器有效性的内在原因。实验设计结合理论分析与可控任务，虽未明确提及开源代码，但框架可扩展至多层网络，值得进一步探索其在更复杂任务和架构中的应用。

1. 'Grokking: Generalization Beyond Overfitting on Small Algorithmic Datasets' (2022) 2. 'The Mechanics of Memorization and Generalization in Deep Networks' (2023) 3. 'Scaling Laws and Inductive Biases in Delayed Generalization' (2023) 4. 'Weight Decay Implies Implicit Regularization of Memorization' (2021) 5. 'On the Role of Gradients in Neural Network Generalization' (2022)