Late-Stage Generalization Collapse in Grokking: Detecting anti-grokking with Weightwatcher
Hari K Prakash, Charles H Martin
cs.LG
2026-02-03
把两套经典grokking实验训到10^7步，泛化峰值后测试准确率崩回约50%，训练集仍满分。WeightWatcher从打乱权重谱里的Correlation Trap检出这一阶段，不需要任何数据。
Grokking 已经是固定剧情：网络很快把训练集背到满分，测试准确率却长时间停在接近瞎猜，再过很久才突然泛化。Power 等人 2022 年在算法任务上把它钉成一个现象之后，后续解释大多围着权重衰减、ℓ2 范数、记忆电路和泛化电路的竞争打转。
这些故事有一个共同盲区：实验通常在测试准确率冲上去之后就停。Hari K. Prakash（UC San Diego）和 Charles H. Martin（Calculation Consulting）把两套经典设置一直跑到 10^7 步，看见了第三段。训练准确率仍是满分，测试准确率从峰值掉回去。他们把这段叫 anti-grokking。它跟 Varma 等人的 ungrokking 不是一回事：ungrokking 是拿已经 grok 过的网络，换更小的数据集再训；anti-grokking 发生在原数据集上，主结果还不依赖权重衰减。
诊断工具是开源的 WeightWatcher（文中用 v0.7.5.5），理论来自 HTSR（Heavy-Tailed Self-Regularization，重尾自正则）和更新的 SETOL（Semi-Empirical Theory of Learning，半经验学习理论）。做法不碰训练数据也不碰测试数据，只看每一层权重矩阵 W。
对每一层构造相关矩阵 X = (1/N) W^T W，再看它的经验谱密度（ESD，特征值的经验分布）。训练得好的层，谱的右尾会变成幂律 ρ(λ) λ^{-α}。HTSR 把 α 当成层质量：
更硬的主信号是 Correlation Trap。把 W 的元素打乱得到 Wrand，条目之间的相关被拆掉，谱应该回到 Marchenko-Pastur（MP）体分布。如果打乱之后仍有远超 MP 体右边缘的异常大特征值，这些尖峰就是 Trap：权重里存在靠幅度撑起来的病态成分。SETOL 预测这种 Trap 会损害泛化。
对照实验沿用两套标准 grokking 设置，只是把训练拉长：
对照指标包括 ℓ2 范数，以及 Golechha 提出的 Activation Sparsity、Absolute Weight Entropy、Approximate Local Circuit Complexity。
MLP、WD=0 时，训练曲线分成三段：约 10^2 步训练集准确率拉升，10^4–10^5 步训练集饱和、测试仍低；约 10^5 步后测试迅速上升，在 10^6 步附近到峰值；再往后测试掉到约 0.5，训练集仍满分。图注把 0.5 写成崩溃终点，摘要则说崩回 chance。10 类 MNIST 的随机基线是 0.1，这两个说法对不齐，正文也没有另给一张峰值测试准确率的表。
层均 α 和 Trap 计数对得上这三段：
| 阶段 | 步数量级 | 层均 α | FC1 Trap 数 |
| pre-grokking | 10^5 | 4.0 ± 0.6 | 0 |
| grokking | 10^6 | 2.9 ± 0.2 | 0 |
| anti-grokking | 10^7 | 1.1 ± 0.3 | 7.5 ± 5.6 |
FC2 在 anti-grokking 阶段 α 掉到 1.3 ± 0.3，Trap 约 1 个。加 WD=0.01 之后，测试准确率 grok 完只轻微回落，然后长时间停在峰值附近，α 也在 2 一带稳住成平台。Trap 仍会出现，但 FC1 只有 2.0 ± 0.0 个，崩溃幅度小得多。Activation Sparsity 一类进度指标能跟上前两段，到第三段继续沿原趋势走，分不清是还在 grok 还是已经开始崩。
模加法上，训练准确率三个阶段都是 1.0 ± 0.0；测试从 pre-grok 的 0.40 ± 0.28 升到 grok 的 0.97 ± 0.02，再掉到 anti-grok 的 0.68 ± 0.11。层均 α 从 4.10 ± 0.83 收到 2.02 ± 0.44，再弹回 3.89 ± 0.68。Trap 在前两段是 0，第三段层均 4.88 ± 1.55。
两套模型的记忆不是同一种。MLP 在 pre-grok 是弱相关记忆，部分层 α≥5，还没训透；崩的时候第一层最大右奇异向量从平滑的全局模板收成像手写「8」的局部图案，权重行直接长出数字原型，论文叫 prototype overfitting。模加法从一开始整段 ESD 就是极重尾，grok 峰值时常只有一两个大特征值，被叫成 rule-based memorization；崩的时候 α 不降反升，更接近灾难性遗忘。
文末还扫了 gpt-oss-20b 和 gpt-oss-120b：两套生产级开源模型都有「异常多」的层 α<2，并声称也有 Trap（图上没画出数量）。没有干预实验，也没有这些层跟下游任务的对照。
对还在小模型上复现 grokking 的人，这篇把「训够了就会泛化」改成「泛化峰值之后还可能再崩」。早停和权重衰减在这里的作用很具体：WD=0.01 没有消灭 Trap，但把崩溃压成了平台。想用 grokking 当泛化机制的故事，至少要把时钟拉到千万步再下结论。
对训练监控更有用的是那句「不看数据」。α 和 Trap 只依赖权重矩阵，理论上可以在不持有验证集的场景里当过拟合探针。代价也清楚：主信号绑在 WeightWatcher / SETOL 这一套谱方法上，Charles Martin 本人是工具作者，这篇是 ICML workshop 论文的大幅扩展。把它当成可插即用的生产告警还早。
gpt-oss 那段是彩蛋，不是证据。看见 α<2 的层偏多，最多说明大模型权重谱里有类似病态，不能推出这些模型已经 anti-grok。
论文自己写了：α 是现象学分档，α≈2 经常对着最优点，但不保证。有的模型层 α 低于 2 仍能泛化，有的训得很好的层 α 会远大于 2。模加法 anti-grokking 阶段很多层谱已经 atypical，伴随 rank collapse 和多峰幂律，这时 α 拟合本身不可靠。
实验面更窄。主结果是 1000 张 MNIST 的 3 层 MLP，外加一个无 LayerNorm 的玩具 Transformer 做模加法。10^7 步、float64、权重初始化还乘了 8.0，这套超参离现代大模型训练很远。Trap 的因果链（单个超大权重元素导致打乱后仍越出 MP 边缘，再变成原型记忆）在附录里用 BBP 相变给了充分条件，但没有做「挖掉 Trap 对应的奇异方向后测试准确率是否回来」这种干预。gpt-oss 只有直方图，没有层计数，也没有和同系列更干净模型的对照。
Figure 1 把崩溃终点标成 0.5，摘要写成 chance。10 类 MNIST 的 chance 是 0.1。主实验也没给出 grokking 峰值测试准确率的点估计，只在扰动实验附录里看到 grok 段测试 0.8879 ± 0.0051、anti-grok 段 0.6184 ± 0.1086，而那段的相位切分还改过。读数字时要把主图、摘要、附录三套口径分开。