AF-Muon：用于捆绑嵌入模型的无 AdamW Muon 优化器

Muon 通过对矩阵参数应用谱范数最速下降更新来改进大规模训练，但实际模型还包含不适合稠密矩阵几何的参数块。一个重要的情况是绑定词汇表，它出现在语言模型和其他标记生成器中，并且可以接收多个结构不同的梯度源，从稀疏输入查找到密集输出分类器更新。在参考配方中，这些块被交给辅助 AdamW 优化器，该优化器恢复二阶矩状态并将别名表更新为通用张量。我们提出了 AF-Muon，它是 Muon 的无 AdamW 扩展，它保持隐藏权重矩阵的 Muon 矩阵更新，同时使用支持感知的有限上限线性最小化预言机来处理绑定词汇表，并使用 RMS 归一化更新来处理一维辅助参数。因此，AF-Muon 使用单个第一时刻缓冲区来训练每个参数类，并且没有第二时刻状态，在我们的基准测试中相对于混合 Muon 节省了大约 20% 的优化器状态内存。跨九个绑定词元设置 - 从 124M 到 1B 参数的仅解码器语言模型、完全共享的 T5 式编码器-解码器以及 ImageGPT 式图像词元、蛋白质和稀疏 MoE 变体，涵盖文本、图像和蛋白质序列数据 - AF-Muon 改善了混合 Muon 和 SCION 式 Sign 端点的平均验证损失和困惑度。长程运行和超参数敏感性研究证实增益是稳健的，同动量诊断将其归因于有限上限，它比 Sign 保留更多的行内幅度，同时限制行 RMS 的坐标浓度。这些结果将绑定词汇表识别为独特的优化器几何结构，并在模型、模态和架构中产生鲁棒的无 AdamW Muon 变体，在匹配训练中大约有 1% 的步时开销。