揭示模型突然泛化背后的机制，发现权重衰减与数据分布共同决定性能飞跃。
Beyond Progress Measures: Theoretical Insights into the Mechanism of Grokking
- 通过理论分析发现，权重衰减促使嵌入空间令牌均匀化。
- 实验表明，性能突变需嵌入均匀性与数据分布协同作用。
- 提出新评估指标，可更准确追踪测试损失变化，适合研究泛化机制者。
Grokking 指模型在长期过拟合后测试准确率突然提升的现象，为理解模型泛化机制提供重要线索。现有研究认为 grokking 依赖于损失函数仅由权重衰减主导时的优化动态，但我们发现此时仅导致令牌均匀化，不足以引发 grokking。本文以 Transformer 在质数运算任务中的 grokking 为例，通过理论分析与实验证明：(i) 权重衰减最小化时，会促使嵌入空间中所有令牌趋于均匀；(ii) grokking 的发生由嵌入空间均匀性与训练数据分布共同决定。基于此，我们统一解释了多种已有进展度量，并提出一种新颖、简洁且有效的进展度量，能更准确追踪测试损失变化。最后，为验证理论框架的普适性，我们设计专用数据集，在 ResNet-18 上成功实现 grokking 现象。代码已开源。
原文摘要 · Abstract (English)
Grokking, referring to the abrupt improvement in test accuracy after extended overfitting, offers valuable insights into the mechanisms of model generalization. Existing researches based on progress measures imply that grokking relies on understanding the optimization dynamics when the loss function is dominated solely by the weight decay term. However, we find that this optimization merely leads to token uniformity, which is not a sufficient condition for grokking. In this work, we investigate the grokking mechanism underlying the Transformer in the task of prime number operations. Based on theoretical analysis and experimental validation, we present the following insights: (i) The weight decay term encourages uniformity across all tokens in the embedding space when it is minimized. (ii) The occurrence of grokking is jointly determined by the uniformity of the embedding space and the distribution of the training dataset. Building on these insights, we provide a unified perspective for understanding various previously proposed progress measures and introduce a novel, concise, and effective progress measure that could trace the changes in test loss more accurately. Finally, to demonstrate the versatility of our theoretical framework, we design a dedicated dataset to validate our theory on ResNet-18, successfully showcasing the occurrence of grokking. The code is released at https://github.com/Qihuai27/Grokking-Insight.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。