arXiv:2603.17970cs.LGcs.NA2026-03被引 3

MUD优化器通过低开销方法加速Transformer训练,比现有方法更快达到相同精度。

Beyond Muon: MUD (MomentUm Decorrelation) for Faster Transformer Training

  • 用三角分解替代极分解,降低动量更新的计算开销。
  • 在相同训练时间内,比AdamW快10%-50%,每秒处理令牌数提升1.3-3倍。
  • 适合追求训练效率的模型开发者,尤其在大模型和蛋白质语言模型中表现优异。

正交化动量优化器如Muon通过短极分解迭代对矩阵动量更新进行白化/正交化,从而提升Transformer训练效率。然而,极分解近似通常需要多次大型矩阵乘法,导致显著且依赖硬件的开销。本文提出MUD(MomentUm Decorrelation),一种互补的白化方法,将Muon的极更新替换为受经典Gram-Schmidt与Gauss-Seidel思想启发的三角(类似Cholesky)白化代理。我们证明行正交矩阵是MUD映射的不动点,内部步骤等价于对格拉姆矩阵的对称高斯-赛德尔预处理,并证明在不动点附近具有二次局部收敛性。在时间到困惑度(time-to-perplexity)方面,MUD相比调优后的AdamW和Muon实现一致10%-50%的时钟时间改进;虽然每步收敛略慢于Muon,但优化器开销大幅降低——相对于Muon,MUD在多数设置下峰值令牌/秒提升约1.3-2.6倍,GPT-2 large在A100上最高达近3倍。此外,我们展示了训练ESM-2 150M蛋白语言模型,MUD以显著更少的时钟时间达到与Muon相当的验证困惑度。

原文摘要 · Abstract (English)

Orthogonalized-momentum optimizers such as Muon improve transformer training by approximately whitening/orthogonalizing matrix-valued momentum updates via a short polar-decomposition iteration. However, polar-factor approximations typically require multiple large matrix multiplications, and the resulting overhead can be substantial and hardware-dependent. We introduce MUD (MomentUm Decorrelation), a complementary whitening approach that replaces Muon's polar update with a triangular (Cholesky-like) whitening surrogate inspired by classical Gram--Schmidt and Gauss-Seidel ideas. We show that row-orthonormal matrices are fixed points of the MUD map, relate the inner step to symmetric Gauss-Seidel preconditioning of the Gram matrix, and prove quadratic local convergence near the fixed point. In terms of time-to-perplexity, MUD yields consistent 10-50\% wall-clock improvements over tuned AdamW and Muon in time-to-perplexity, typically converging slightly slower per step than Muon but with substantially lower optimizer overhead -- relative to Muon, MUD improves peak tokens/s by roughly $1.3-2.6\times$ across most settings and up to nearly $3\times$ on GPT-2 large on an A100. We also demonstrate training a ESM-2 150M protein language model, where MUD matches Muon-level validation perplexity in significantly less wall-clock time.

优化器Transformer训练加速

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。