论文精选

谱理论解释 grokking 现象:weight decay 驱动特征学习

A Spectral Theory of Grokking: Weight Decay induces Feature Learning

精选理由

研究 grokking 为什么突然泛化的论文,理论算出 grokking 时间由学习率乘 weight decay 决定,还用上千个网络验证了,做可解释性的可以看看。

arXiv 论文为 grokking(训练后泛化延迟出现)给出定量理论。作者证明在平方损失加 L2 weight decay 训练的齐次网络中,记忆化完成后残余误差会反馈进 NTK 动力学,残余驱动的核增长与 weight decay 相互竞争。理论预测 grokking 时间尺度由学习率与 weight decay 的乘积控制,在模加法任务上用 84×90 网格的 MLP 和 42×45 网格的单块 Transformer 验证了相图结构与反比缩放规律。

原文 · arXiv cs.LG

A Spectral Theory of Grokking: Weight Decay induces Feature Learning

In grokking an early fit to the training data separates from a much later improvement in generalization. During this delay, training can move from a fixed neural tangent kernel (NTK) regime to one in which task-relevant kernel eigendirections continue to evolve. We provide a quantitative theory for how this transition from lazy to rich learning can produce delayed generalization. For homogeneous networks trained with squared loss and $L_2$ weight decay, we show that a finite residual remains after memorization, with larger residual fractions in target components associated with smaller NTK eigenvalues. These residuals feed back into the dynamics of the NTK itself, and projecting the resulting dynamics onto task-relevant spectral directions yields a reduced system in which residual-driven kernel growth competes with weight decay. This system predicts that the grokking timescale is controlled by the product of learning rate and weight decay, that feature learning slows logarithmically near a critical decay above which task-aligned NTK structure can no longer support generalization, and that stronger decay can prevent fitting altogether. We test these predictions in modular addition. In a homogeneous MLP, task-aligned Fourier structure continues to emerge in the NTK after training accuracy has saturated, and an 84$\times$90-grid of trained networks across varying learning rate and weight decay recovers the predicted phase geometry and inverse-product scaling of the generalization time with learning rate and weight decay. A one-block Transformer shows similar macroscopic phase structure in a 42$\times$45-grid, as well as the same transition-time scaling despite violating exact homogeneity. Together, these results provide a mechanistic derivation connecting post-fit feature learning to both the onset of generalization and its phase structure in the learning rate and weight decay plane.