训练越久，模型反而越会泛化吗？理解 Grokking 现象

在Grokking: Generalization Beyond Overfitting on Small Algorithmic Datasets这篇论文中神经网络在训练集上早已达到接近 100% 的准确率但继续训练很长时间后验证集准确率才突然从接近随机水平上升到接近完美作者把这种现象称为grokking。一、这篇文章到底研究了什么我们通常希望神经网络学习数据背后的规律而不是简单地记住训练样本。例如假设我们给模型一些这样的题目输入输出a ∘ b a \circ ba∘bc cca ∘ d a \circ da∘de eeb ∘ c b \circ cb∘cf ff模型需要根据已经看到的部分答案推断没有见过的组合。它可以采用两种方式记忆训练样本见过的输入直接查表没见过的输入不会做学习底层规则发现a ∘ b a \circ ba∘b背后存在一种统一的运算规律然后推断新的输入。这篇论文研究的正是模型什么时候会从第一种方式转向第二种方式。作者构造了许多小型的算法数据集包括模97 9797的加法、减法、除法以及抽象群S 5 S_5S5​中的排列运算。每条数据都可以写成a ∘ b c a \circ b ca∘bc其中a aa、b bb和c cc都被表示成没有内部结构的离散符号。模型不能直接看到“这是数字 3”或“这是某个排列”只能通过符号之间的关系自己发现规律。Datasets二、什么是 GrokkingGrokking 的典型过程可以分成两个阶段。第一阶段模型先学会“背答案”训练刚开始时模型很快就能把训练集做对。因为模型参数很多它可以给每个训练样本单独存储一个答案。于是会出现这样的情况训练准确率接近 100%验证准确率接近随机水平模型已经记住训练数据但还不会处理新数据在论文的模97 9797除法实验中训练准确率不到10 3 10^3103步就接近完美但验证准确率可能要到接近10 6 10^6106步才达到同样水平。也就是说模型在很长时间内都处于“训练集全会验证集不会”的状态。Grokking Example第二阶段模型逐渐找到统一规则如果继续训练模型的内部参数还会持续变化。它可能逐渐从大量零散的记忆转向一种更加统一、简单、可泛化的表示。此时验证准确率可能长期没有明显变化然后突然快速上升记忆训练样本 ⟶ 形成内部规则 ⟶ 泛化到未见样本 \text{记忆训练样本} \;\longrightarrow\; \text{形成内部规则} \;\longrightarrow\; \text{泛化到未见样本}记忆训练样本⟶形成内部规则⟶泛化到未见样本从外部观察这就像模型突然“顿悟”了。但这里的“顿悟”不是模型产生了人类式的意识也不是某一步突然写入了一个“理解模块”。更准确地说它是模型参数经过长时间优化后逐渐从一个复杂的记忆解转移到了一个更简单的规则解。三、模型为什么会从记忆转向规则关键在于能够拟合训练集的模型解并不只有一种。假设训练集中有100 100100个样本那么模型至少可以有两种解一个复杂的解专门记住这100 100100个答案一个简单的解学习背后的规则并因此正确回答更多未见样本。这两种解在训练集上的表现可能完全一样都是 100% 准确率。但它们的泛化能力不同。可以把它类比成学生做数学题。学生拿到一组题目后可能先把题目和答案背下来看到题目 A回答 7看到题目 B回答 3。这样他可以把练习册上的题全部答对但遇到新题就不会。如果他继续思考可能最终发现原来这些题都遵循同一个公式。这时他不再需要逐题查答案而是可以通过公式解决新的题目。Grokking 在模型身上呈现出来的现象与这个过程有些相似。四、为什么训练更久有时反而有利于泛化这和“模型偏好什么样的解”有关。在训练集上复杂的记忆解和简单的规则解都可以得到零训练误差。但在某些训练条件下优化过程可能逐渐偏向参数更简单、结构更规则的解。尤其是权重衰减可能发挥重要作用。权重衰减会惩罚过大的参数使模型不太容易维持一个极其复杂的记忆方案。论文发现在这些算法任务中权重衰减对数据效率的提升非常明显所需样本量相比多数其他干预甚至减少了一半以上。Weight Decay此外优化过程中的噪声也可能帮助模型找到泛化更好的解。小批量随机梯度下降或者对梯度和参数加入噪声有时会促使模型进入更平坦的损失区域而较平坦的解通常对参数的小变化不那么敏感。Optimization Noise因此模型可能经历下面的过程开始训练快速拟合训练集训练集接近100%继续优化与参数调整形成更简单的内部表示验证集性能突然提升五、这是否意味着训练越久越好不是。这正是我在理解这篇论文时产生的疑问按照通常的机器学习经验训练太久容易过拟合。既然如此为什么这篇论文里训练更久反而提升了泛化答案是传统过拟合和 Grokking 是两种可能出现的训练动态。普通过拟合普通过拟合通常表现为训练误差持续下降 ⟶ 验证误差持续上升 \text{训练误差持续下降} \;\longrightarrow\; \text{验证误差持续上升}训练误差持续下降⟶验证误差持续上升模型逐渐学会训练数据中的噪声和偶然细节因此在训练集上更好在验证集上更差。GrokkingGrokking 则可能表现为训练误差下降 ⟶ 验证性能暂时变差 ⟶ 验证性能再次提升 \text{训练误差下降} \;\longrightarrow\; \text{验证性能暂时变差} \;\longrightarrow\; \text{验证性能再次提升}训练误差下降⟶验证性能暂时变差⟶验证性能再次提升论文观察到验证损失在某些实验中会出现第二次下降。作者认为这种现象可能和经典的 double descent 有关联但也强调 Grokking 发生在训练集已经被拟合很久之后因此可能是不同的现象。Double Descent所以“训练越久越好”并不是一般规律。更准确的说法是在某些任务上模型可能需要很长时间才能从记忆解转向规则解但在另一些任务上继续训练只会让模型越来越严重地过拟合。六、什么类型的任务更容易出现 Grokking我的进一步理解是任务类型确实非常重要。如果一个任务存在明确、统一、可压缩的规则那么模型才有可能从训练样本中发现这个规则。比如模运算排列和群运算简单程序符号推理数学公式某些人工构造的算法任务这篇论文中的任务都具有明确的生成规则而且训练数据量很小、标签没有明显噪声模型又足够大可以同时表达“记忆解”和“规则解”。在这种情况下长时间训练有可能让模型逐渐找到更具泛化能力的规律。可以粗略地表示为明确规则 少量干净数据 足够大的模型 合适的优化 ⟶ 更可能出现 Grokking \text{明确规则} \text{少量干净数据} \text{足够大的模型} \text{合适的优化} \longrightarrow \text{更可能出现 Grokking}明确规则少量干净数据足够大的模型合适的优化⟶更可能出现Grokking不过有规则并不代表一定会 Grokking。论文中也有一些看似有规律的任务例如x 3 x y 2 y ( m o d 97 ) x^3 xy^2 y \pmod{97}x3xy2y(mod97)在允许的训练预算内模型并没有成功泛化基本上只是记住了训练数据。Failed Generalization这说明还需要满足几个条件规则必须足够简单或至少对当前模型来说足够容易发现模型必须能够表达这种规则优化过程必须有机会找到这个规则训练数据不能包含太多随机噪声训练步数和超参数要合适七、噪声会不会阻止 Grokking通常会。作者还专门研究了训练集中加入随机标签的情况。结果显示少量异常样本对泛化的影响可能不明显但大量异常样本会显著降低模型最终成功泛化的范围。Outliers这也很好理解。如果训练集中的每个样本都遵循同一个规则模型有机会发现统一规律。但如果训练集中混入大量错误答案模型可能不得不学习简单规则 许多例外 \text{简单规则} \text{许多例外}简单规则许多例外此时纯粹学习规则已经无法完全拟合训练集模型更容易被迫保留记忆和例外。因此Grokking 更容易出现在标签干净、规律明确的数据集上而不是噪声很多的现实数据中。八、模型真的“理解”了吗这里需要谨慎使用“理解”这个词。从行为上看模型确实从“只能回答见过的问题”变成了“可以回答没见过的问题”。这说明它学到了一些能够支持泛化的内部结构。作者还对模型的 embedding 进行了可视化发现模型有时会形成与底层数学对象对应的结构。例如模加法中的 embedding 可能呈现类似圆环的拓扑排列运算中的 embedding 则可能形成与群结构有关的簇。Embedding Structure但这并不等于模型像人一样理解了数学概念。我们更稳妥的说法是模型形成了一个能够复现底层规则、并支持新样本预测的内部表示。至于模型为什么一定会形成这个表示作者在这篇文章中并没有给出完整答案。论文提供了一个重要线索在S 5 S_5S5​实验中验证准确率和 sharpness 的 Spearman 相关系数为− 0.79548 -0.79548−0.79548。这意味着验证集表现更好的模型往往位于更平坦的损失区域。但这仍然是相关性并不能证明“平坦性”就是 Grokking 的唯一原因。Sharpness九、实际训练时要不要刻意训练久一点结论不要无条件地延长训练但对于规则型、小数据、低噪声任务值得专门测试更长的训练周期。对于一般的图像分类、文本分类或业务预测任务建议仍然观察验证集训练集性能上升、验证集性能也上升可以继续训练训练集性能上升、验证集性能下降可能正在普通过拟合训练集已经很好、验证集长期停滞可以尝试更长训练但要同时调整正则化和学习率验证集性能在长时间停滞后重新上升这可能是 Grokking 的迹象。实际操作中可以保存不同训练阶段的 checkpoint训练 10,000 步 → 保存模型并评估验证集 训练 100,000 步 → 保存模型并评估验证集 训练 1,000,000 步 → 保存模型并评估验证集 选择验证集表现最好的 checkpoint不要简单地认为“最后一个模型一定最好”也不要因为训练准确率已经达到 100% 就立刻停止。十、最佳实践该论文复现较为容易对数据集、模型、算力要求不高。详细复现的pipeline见grokking_reproduction十一、总结我对这篇论文的理解可以总结为模型训练过程不一定是从“不会”逐渐变成“会”而可能先从“记住训练样本”开始再经过很长时间的优化转向“发现数据背后的规则”。因此训练集准确率达到 100% 并不总意味着模型已经学会了任务。它可能只是找到了一个能够记忆训练数据的复杂解。如果任务背后存在明确、简单、可表达的规律那么继续训练可能让模型逐渐偏向更简单、更统一的规则解从而出现 Grokking。论文中的权重衰减和优化噪声实验也说明正则化和优化路径可能帮助模型摆脱单纯的记忆。Main Findings但这并不推翻“训练过久会过拟合”的常识。更准确的经验应该是普通任务要依靠验证集早停规则明确的小数据任务则要警惕过早早停并尝试检查验证性能是否会在长时间停滞后重新上升。这也是 Grokking 最有价值的地方它提醒我们过拟合并不总是学习过程的终点有时只是模型从记忆走向泛化之前的中间阶段。