MUD (MomentUm Decorrelation): Extends Muon Optimizer to Non-Square Matrices for Faster Full-Transformer Training

arXiv 2603.17970·medium signal

Muon improves transformer training via gradient orthogonalization but is limited to square weight matrices, leaving embedding layers, rectangular attention projections, and feed-forward layers untouched. MUD (Momentum Decorrelation) extends the orthogonalization principle to arbitrary-shaped gradient matrices, achieving whitening across the full transformer architecture. Benchmark results show faster convergence than both Muon and Adam on language model training at comparable compute budgets.