[论文解读] Towards Understanding Grokking: An Effective Theory of Representation Learning

Ziming Liu, O. Kitouni+4 more|arXiv (Cornell University)|May 20, 2022

Neural dynamics and brain function被引用 24

一句话总结

这篇论文提出了一个有效理论和相图来解释grokking，将泛化与结构化表示相关联，并在变换器和玩具模型中识别出四个学习阶段（理解、grokking、记忆、困惑）

ABSTRACT

We aim to understand grokking, a phenomenon where models generalize long after overfitting their training set. We present both a microscopic analysis anchored by an effective theory and a macroscopic analysis of phase diagrams describing learning performance across hyperparameters. We find that generalization originates from structured representations whose training dynamics and dependence on training set size can be predicted by our effective theory in a toy setting. We observe empirically the presence of four learning phases: comprehension, grokking, memorization, and confusion. We find representation learning to occur only in a "Goldilocks zone" (including comprehension and grokking) between memorization and confusion. We find on transformers the grokking phase stays closer to the memorization phase (compared to the comprehension phase), leading to delayed generalization. The Goldilocks phase is reminiscent of "intelligence from starvation" in Darwinian evolution, where resource limitations drive discovery of more efficient solutions. This study not only provides intuitive explanations of the origin of grokking, but also highlights the usefulness of physics-inspired tools, e.g., effective theories and phase diagrams, for understanding deep learning.

研究动机与目标

• 研究为什么神经网络在对小型算法数据集过拟合后会泛化。
• 将泛化与结构化输入表示的出现联系起来。
• 发展一个受物理学启发的有效理论，用以预测关键训练数据大小和学习轨迹。
• 表征学习阶段及超参数如何影响延迟泛化。

提出的方法

• 构建一个简化的玩具模型，其中输入映射到可训练的嵌入，再由解码器处理它们的和。
• 基于嵌入中的平行四边形结构定义 Representation Quality Index (RQI)。
• 提出有效损失 ell_eff，支配嵌入动态并推导 grokking 速率来自 Hessian 谱。
• 通过改变表示学习率和解码器学习率以及权重衰减，在玩具和类变换器设置中推导相图。
• 将分析扩展到模组加法和非阿贝尔群，以测试跨任务的一般性。

实验结果

研究问题

• RQ1 在模型对训练数据过拟合时，grokking 的泛化原因是什么？
• RQ2 训练数据规模如何影响结构化表示的出现及泛化？
• RQ3 在什么条件下会出现延迟泛化，超参数如何延迟或加速它？
• RQ4 相图和有效理论预测是否在玩具模型之外对变换器架构和 MNIST 也成立？
• RQ5 表示结构在使跨不同运算的泛化能力中的作用是什么？

主要发现

• 泛化与结构化嵌入的出现相关，这些嵌入在表示空间中创建平行四边形。
• 临界训练数据分数决定何时出现线性且唯一的表示，从而实现泛化。
• Grokking 是理解与记忆之间的一个阶段，适当的超参数调整可以改变或消除它。
• 相图揭示四个学习阶段，显示更快的解码器或过度的解码器容量会促进记忆，而平衡学习有利于理解和 grokking。
• 在变换器设置中，降低解码器容量或应用正则化会加速泛化并可以抑制 grokking，有来自 MNIST 实验的证据。
• 有效理论预测 grokking 时间和相变点，与神经网络训练轨迹在定性上对齐。