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.
翻译:诸如Muon等正交动量优化器通过短极分解迭代对矩阵值动量更新进行近似白化/正交化,以改进Transformer训练。然而,极因子近似通常需要多次大型矩阵乘法,由此产生的额外开销可能相当可观且依赖于硬件。我们提出MUD(动量去相关),一种互补的白化方法,它用受经典Gram-Schmidt和高斯-赛德尔思想启发的三角(类Cholesky)白化替代方法取代了Muon的极更新。我们证明了行正交矩阵是MUD映射的不动点,将内部步骤与Gram矩阵的对称高斯-赛德尔预条件相关联,并证明了不动点附近的二次局部收敛性。在困惑度达到时间方面,与经过调参的AdamW和Muon相比,MUD在困惑度达到时间上实现了一致的10-50%的挂钟时间改进——尽管通常每步收敛速度略慢于Muon,但优化器开销显著降低:相对于Muon,MUD在大多数设置下将峰值令牌/秒提高了约1.3-2.6倍,在A100上训练GPT-2 large时甚至接近3倍。我们还演示了训练ESM-2 150M蛋白质语言模型,在该模型上MUD以显著更短的挂钟时间达到了与Muon同级别的验证困惑度。