Transformer在模运算任务中出现训练后性能骤降,根源在表示与输出层的接口不稳。
Post-Grokking Collapse at the Representation-Readout Interface in Muon-Trained Transformers
- 发现模型在解决模加法任务后,虽快速收敛但泛化能力崩溃。
- 九种配置下均出现性能下降,最低准确率降至27.59%。
- 适用于研究模型稳定性与表示-读出接口设计的科研人员。
在标准分割下,Muon模型包含隐藏矩阵和AdamW嵌入/输出头。其在模113加法任务上更快达到解,但解不稳定:所有九种配置在$(a+b) mod 113$任务上均经历沟壑(grokking)后失去泛化能力。五组随机种子中,仅一组参考点维持阈值以上,最低达27.59%。该不稳定性贯穿两种模数、两种宽度、两种训练比例、减法任务及不同深度。故障源于表示-读出接口,仅在联合可逆映射下可识别,损失函数无法选择。训练集解决后梯度降至$10^{-6}$量级,优化器响应差异显著:Muon步长弹性为-0.03,AdamW为+1.5,且Muon组每参数移动速度是后者的8.0倍。从比特相同状态出发,冻结任一组件可防止失败;冻结嵌入/输出头时,在451,400步内共五次运行均无低于阈值评估,而未冻结组记录137–321次。移除Muon归一化与正交化无效:表示维度从326个有效共轭对降至4,无循环性坍缩,最终终止失败。傅里叶滤波分离了电路失效与掩蔽效应。43个检查点、五组种子、三类模式中,仅任务对齐族能精确达到100%。电路失效时其任务解失效;掩蔽状态下仍保持完美,但全模型降至45.85%,在每个样本上产生正裕量,尽管含错误,却因近似相等的对抗余项被压倒。重缩放使其恢复至99.9%;沟壑本质为向上一致性的条件达成。任务选择该族,减法下$(k,k)$转为$(k,-k)$。突然坍缩期间,标准傅里叶支持不变,功率分布余弦仍为0.9899。
原文摘要 · Abstract (English)
Under the standard split, Muon gets hidden matrices and AdamW embeddings/output head. Muon groks modular addition faster, but its solutions do not hold. All nine configurations on $(a+b) \bmod 113$ grok and later lose generalization. Across five seeds the selected AdamW reference falls below threshold on four, reaching 27.59%. Instability persists across two moduli, two widths, two training fractions, subtraction, and depth. The failure arises at the representation-readout interface, identified only jointly up to an invertible map unselected by the loss. After solving the training set, the gradient falls to order $10^{-6}$ and the optimizers respond differently: step-size elasticity is -0.03 for Muon versus +1.5 for AdamW, and the Muon group moves 8.0 times faster per parameter. From bit-identical states, freezing either group prevents failure. Freezing embeddings/readout removes it in five runs over 451,400 post-grokking steps and five paired seeds: unfrozen arms record 137-321 sub-threshold evaluations, frozen arms none. Removing Muon's normalization and orthogonalization is no substitute: it collapses representation from 326 effective conjugate pairs to 4, shows no recurrent collapse, and fails terminally. Fourier filtering separates circuit failure from masking. Across 43 checkpoints over five seeds and three regimes, the task-aligned family reaches exactly 100% alone. In circuit failure it no longer solves the task; in masking it remains perfect while the full model reaches 45.85%, giving a positive margin on every example, including errors, but being outvoted by a near-equal adversarial remainder. Rescaling it restores 99.9%; grokking is the same condition resolving upward. The task selects the family, swapping $(k,k)$ for $(k,-k)$ under subtraction. Across an abrupt collapse, standard Fourier support is unchanged and the power-distribution cosine remains 0.9899.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。