模型把训练集背熟后又练了上万步，忽然“顿悟”：Grokking 是怎么回事？

▧ 论文主题概念插画 · 暂无可用预览

配图：AI 生成的概念示意图，非论文原图或实验数据。

TL;DR：OpenAI 的研究者在小型算法数据集上发现，神经网络可以先把训练集彻底背下来、验证集却几乎答不对，继续训练很久之后，验证准确率才突然飙升到接近完美，他们称之为 grokking。论文只在特定的小型算法任务上观察并描述了这个现象，没有给出完整的理论解释，也不能直接推论到大语言模型。

这篇论文发表于 2022 年 1 月（arXiv 提交日为 1 月 6 日），早期版本曾在 2021 年的 ICLR 数学推理研讨会上亮相，并不是什么新闻，但它改变了人们对“过拟合”的直觉。作者是 Alethea Power、Yuri Burda、Harri Edwards、Igor Babuschkin 和 Vedant Misra。

实验怎么做的？ 他们造了一批“小型算法数据集”：给出一张二元运算表，比如模 97 的加法、减法、除法，或者 S5 置换群的乘法，把表里一部分条目拿来训练，剩下的留作验证，让模型补全空缺。模型是一个很小的两层 Transformer（宽度 128，4 个注意力头，约 40 万非嵌入参数），用 AdamW 训练，权重衰减设为 1。

数字很戏剧化。 以模 97 除法、一半数据用于训练为例，训练准确率在不到 1000 步时就接近满分，验证准确率却要接近 100 万步才达到同样水平，在 10 万步之前几乎没有像样的泛化。也就是说，模型先靠死记硬背“交了卷”，很久以后才真正学会规律。论文还发现，训练数据越少，需要的优化步数增长得越快：在约 25% 到 30% 的数据量附近，训练数据每减少 1%，达到泛化所需的中位步数就大约增加 40% 到 50%。

什么在帮忙？ 消融实验里，权重衰减的效果最明显，能让所需样本量减少一半以上；加入一定噪声也有帮助。另外，学习率必须落在大约一个数量级的窄区间内才会出现泛化。对输出层权重做 t-SNE 可视化时，模加法的嵌入会排成环形，群运算的嵌入则按子群的陪集聚成簇，说明模型最后学到的是带有数学结构的表示，而不只是查表。

为什么重要，又为什么别夸大。 传统直觉是：训练误差降到零、验证误差还很高，就该早停。Grokking 说明在这类任务里，“看起来已经过拟合”的模型可能还在缓慢地转向更泛化的解。但论文本身只把这些数据集当作研究过参数化网络泛化的试验场，作者并未声称大模型的能力都是这样“顿悟”出来的。

后续。 之后不少研究试图解释这一现象。例如 Neel Nanda 等人 2023 年的工作用机制可解释性方法拆解了模加法网络，发现它学到了基于傅里叶变换的算法（arXiv:2301.05217）。

论文地址：arXiv:2201.02177