Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

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

Ben S. Southworth, Stephen Thomas

Abstract

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.

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

Abstract

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 across most settings and up to nearly 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.
Paper Structure (19 sections, 4 theorems, 37 equations, 8 figures, 7 tables, 2 algorithms)

This paper contains 19 sections, 4 theorems, 37 equations, 8 figures, 7 tables, 2 algorithms.

Table of Contents

  1. Introduction
  2. MomentUm Orthogonalization based on Newton--Schulz (Muon)
  3. Descent methods
  4. Muon
  5. Approximate MomentUm Decorrelation (MUD)
  6. Orthogonalization and decorrelation
  7. MUD: a fast triangular approximation to whitening
  8. Analysis
  9. Algorithmic discussion
  10. Numerical experiments
  11. Transformer LLMs
  12. Throughput
  13. NanoGPT-style language modeling
  14. WikiText-103
  15. FineWeb-Edu
  16. ...and 4 more sections

Key Result

Proposition 1

The Gram matrix induced by any orthonormal row matrix $Q$, where $\mathcal{G} = QQ^T = I_k$, is a fixed point of $\mathcal{M}$.

Figures (8)

  • Figure 1: Nvidia GH200 (bs=12, T=2048, GPT2-Medium): validation perplexity vs. step and vs. wall-clock time for AdamW, Muon, and MUD.
  • Figure 2: Nvidia GH200 (bs=12, T=2048, GPT2-Medium): (left) speedup relative to AdamW vs. achieved validation perplexity; (right) relative perplexity improvement rate vs. perplexity, showing diminishing relative gains at low perplexity.
  • Figure 3: Nvidia GH200 (bs=12, T=2048, GPT2-Small): validation perplexity vs. step and vs. wall-clock time for AdamW, Muon, and MUD.
  • Figure 4: Nvidia A100 (bs=12, T=2048, GPT2-Medium): validation perplexity vs. step and vs. wall-clock time for AdamW, Muon, and MUD.
  • Figure 5: Nvidia A100 (bs=12, T=2048, GPT2-Medium): (left) speedup relative to AdamW vs. achieved validation perplexity; (right) relative perplexity improvement rate vs. perplexity, showing diminishing relative gains at low perplexity.
  • ...and 3 more figures

Theorems & Definitions (4)

  • Proposition 1: Orthonormal fixed point
  • Proposition 2: Inner step as symmetric Gauss Seidel preconditioning
  • Theorem 3: Local quadratic convergence
  • Corollary 4: Local eigenvalue clustering and condition number improvement