AF-Muon：面向绑定嵌入模型的无 AdamW Muon 优化器

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