Post-Grokking Collapse at the Representation-Readout Interface in Muon-Trained Transformers
本文证明了经由 Muon 训练的 Transformer 会表现出一种“后顿悟崩溃”（post-grokking collapse），即在学习模运算后，由于表示-读取接口（representation-readout interface）的不稳定性，导致泛化失效，这一现象以发散的优化器动力学和掩蔽效应为特征，且可以通过冻结特定模型组件或重缩放任务对齐电路来解决。
原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。 阅读完整免责声明
想象一下，你正在教一个机器人解决一个非常特定的数学谜题：将两个数字相加，并求出除以一个大质数后的余数。起初，机器人看起来只是在死记硬背答案，就像一个在背闪卡的学生。但随后，神奇的事情发生了，这被称为“顿悟”（grokking）。突然间，在经历了一段看似毫无进展的时期后，机器人不再仅仅是记忆，而是真正理解了其中的规律。它能够完美地解决从未见过的数字。这对科学家来说意义重大，因为这有助于他们理解人工智能是如何学会“思考”而非仅仅是“记忆”的。
然而，这里有一个陷阱。有时，在机器人掌握了解法之后，它竟然开始再次遗忘，尽管它仍在接受针对同一个谜题的训练。这就像一个学生终于理解了代数，却又开始在同样的问题上犯些低级错误。科学家们使用特殊的工具来观察机器人的“大脑”内部，看看究竟发生了什么。他们观察机器人用来解决数学问题的特定模式（就像歌曲中的音符），以查看机器人是否仍在演奏正确的音符。核心问题在于：机器人为什么会遗忘？是因为它停止学习音符了，还是发生了其他情况？
消失的天才之谜
在这项研究中，研究人员观察了一种特殊的 AI 机器人——Transformer，它被训练用来解决模加法谜题。他们使用了两位不同的“老师”（优化器）来教授这个机器人。一位名叫 AdamW，它是标准且可靠的教员。另一位名叫 Muon，它是一位更新、更快的老师，似乎能帮助机器人更快地理解这个谜题。
研究人员发现，Muot 确实是一个速度达人。它帮助机器人达到“顿悟”这一“恍然大悟”时刻的速度，大约是 AdamW 的一半。但转折在于：Muon 找到的解法极其脆弱。一旦机器人解决了谜题，它往往会再次开始失败，准确率从 100% 跌至接近于零，然后又恢复过来。这种情况反复发生，尽管机器人仍在接受针对同一问题的训练。AdamW 也并非完全安全，但它失败的频率较低，且程度较轻。
破碎的握手
那么，问题出在哪里呢？研究人员发现，问题并不在于机器人忘记了数学。机器人的“大脑”里依然完美地保留着解法。问题在于机器人两个部分之间的破碎握手：负责“思考”的部分（表示层/representation）与负责“回答”的部分（输出头/readout）之间的断裂。
想象一下，机器人的大脑是一个装满书籍（思想）的图书馆，而机器人的嘴巴是朗读这些书的图书管理员（答案）。
- 图书馆（隐藏层）每秒钟都在对书进行重新编写。
- 图书管理员（输出头）则试图阅读这些书并给出答案。
在使用 Muon 老师的情况下，图书馆每次更新时都会用一种略微不同的语言来重写书籍，而图书管理员却一直试图用旧语言来阅读。起初，他们能够互相理解。但由于图书馆改变语言的速度比图书管理员适应的速度更快，他们最终无法再理解彼此。图书馆仍在书写正确的数学逻辑，但图书管理员读到的却是错误的词汇。
研究人员通过冻结其中一方证明了这一点。当他们停止图书馆改变语言时，机器人保持完美；当他们停止图书管理员改变阅读风格时，机器人也保持完美。但当两者都被允许自由移动时，它们就产生了偏差，机器人开始失败。
机器中的幽灵
最令人困惑的部分是：即使在机器人失败时，“图书馆”仍然包含着完美的解法。研究人员使用了一种特殊的过滤器，专门观察那个掌握了数学知识的部分。当他们这样做时，解法依然是 100% 正确的！
事实证明，机器人遭受的是电路掩蔽（circuit masking）。想象一下，图书馆正在大声喊出正确答案，但机器人的其他部分也在以完全相同的音量喊出错误的答案。图书管理员听到了嘈杂混乱的声音，从而给出了错误的答案。正确的答案依然在那里，只是被噪音淹没了。
研究人员展示了，如果你仅仅调高正确部分（数学家族）的音量而不改变其他任何东西，机器人会立即恢复到 100% 的准确率。数学并没有消失，它只是迷失在了噪音之中。
为什么传统工具失效了
通常情况下，科学家通过观察机器人使用的“音符”来检查其工作情况。他们检查正确的音符是否存在，以及是否足够响亮。在这项研究中，这些工具显示机器人运行正常。音符都在那里，而且听起来与失败前几乎完全一样。
但机器人却在失败。这是因为这些工具只观察了“音符”，却没有观察“音量”或图书管理员正在使用的“语言”。音符是完美的，但图书管理员已经不再正确地聆听它们了。这就像检查一台收音机，看到电台广播信号清晰，却没意识到收音机已经调到了错误的频率。
深度探索
研究人员还测试了增加机器人“深度”（添加更多思考层）会发生什么。他们发现，Muon 让学习变得更快，但“破碎握手”的问题也随之恶化。在更深的模型中，机器人的大脑在一个区域编写数学，而读取答案的部分在另一个区域，它们分离得更快。
他们还尝试通过移除 Muon 老师的特殊“归一化”（normalization）技巧（一种保持更新稳定的方法）来修复它。没有了这个技巧，机器人学到的数学非常狭窄且集中，但最终会彻底崩溃且无法恢复。这表明，虽然这个特殊技巧有助于机器人快速学习并扩展其知识，但也让它更容易产生这种特定类型的漂移式失败。
总结
本文的主要教训是：如果机器人的各个部分不同步，快速学习可能会带来危险。 机器人并没有忘记数学，它只是失去了“知道数学的部分”与“表达数学的部分”之间的连接。
研究人员发现，如果在机器人学会解法后冻结“表达”部分（输出和嵌入层），机器人就会保持稳定且永不遗忘。这表明，不稳定并非因为机器人数学不好，而是因为这两个部分移动的速度不同，从而失去了共同的语言。
简而言之，这个机器人是一个天才，但因为它的大脑和嘴巴在说不同的方言，它突然开始结巴了。数学在它内心依然完美，但如果没有一种方法来使这些方言保持一致，这位天才就无法展示他的才华。
您所在领域的论文太多了？
获取与您研究关键词匹配的最新论文每日摘要——附技术摘要，使用您的语言。