Beyond Muon: MUD (MomentUm Decorrelation) for Faster Transformer Training
- URL: http://arxiv.org/abs/2603.17970v1
- Date: Wed, 18 Mar 2026 17:37:31 GMT
- Title: Beyond Muon: MUD (MomentUm Decorrelation) for Faster Transformer Training
- Abstract summary: We introduce MUD (MomentUm Decorrelation), a complementary whitening approach that replaces Muon's polar update with a triangular (Cholesky-like) whitening surrogate.<n>In terms of time-to-perplexity, MUD yields consistent 10-50% wall-clock improvements over tuned AdamW and Muon.<n>We also demonstrate training a ESM-2 150M protein language model, where MUD matches Muon-level validation perplexity in significantly less wall-clock time.
- Score: 0.0
- License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/
- 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.
Related papers
- Muon-C: Operator-Aligned Muon for Convolutional Kernels [5.886986014593717]
We introduce Muon, an operator-aligned that represents kernel momentum as frequency-wise channel-transfer matrices.<n>We show that the new geometry arises from combining the block partition and Fourier coordinates.<n>Muon-C reaches 9.87 FID at 40k iterations, compared with 22.26 for unfolded Muon and 51.31 for Adam.
arXiv Detail & Related papers (2026-09-09T03:45:44Z) - Beyond the Matrix Sign: Quadratic Spectral Descent [6.589542614722917]
We show that curvature can change both the singular values and the singular directions of the optimal update.<n>Experiments on GPT pre-training show that QSD consistently improves validation loss over Muon and recent Muon variants.
arXiv Detail & Related papers (2026-09-07T15:07:24Z) - Improving Neural Network Training by Decoupling the Magnitude and Direction of Weight Vectors [110.10775872372047]
We propose emphMagnitude-Direction-Direction (MD) Decoupling.<n>It factorizes each weight into a fixed-norm direction on a hypersphere and learnable per-row and per-column magnitude gains, updated at separate learning rates.<n>Across both Adam and Muon, MD Decoupling improves on well-tuned baselines, transfers the optimal LR across model width without retuning, and continues to help at scale on large Mixture-of-Experts models.
arXiv Detail & Related papers (2026-06-24T15:40:26Z) - Zeta: Dual Whitening for Matrix Optimization via Coordinate-Adaptive Preconditioning [56.24532075189964]
We propose Zeta, a dual whitening that applies coordinate whitening and spectral whitening in a strictly ordered pipeline.<n>We prove that Zeta matches or surpasses strong baselines across language modeling (0.6B to 8B parameters), mixture-of-experts architectures, and vision tasks.
arXiv Detail & Related papers (2026-06-12T07:10:17Z) - MiMuon: Mixed Muon Optimizer with Improved Generalization for Large Models [45.11415579822849]
We study generalization error of the Muon based on algorithmic stability and mathematical induction.<n>We propose an effective mixed Muon (MiMuon) by cautiously using gradient, which is a hybrid of Muon and momentum-based SGDs.<n>Our MiMuon algorithm has the same convergence rate of $O(frac1NTbig)$ as the Muon algorithm.
arXiv Detail & Related papers (2026-05-19T09:56:27Z) - When and Why Grouping Attention Heads Accelerates Muon Optimization [51.89517622356298]
We study whether Muon should be applied to the full attention projection, to individual heads, or to intermediate head groups.<n>Our analysis reveals a trade-off between the textbf group-wise whitening gain from group-wise updates and an additional update-norm cost caused by replacing full-matrix whitening with group-wise whitening.
arXiv Detail & Related papers (2026-05-09T13:06:31Z) - Muon$^2$: Boosting Muon via Adaptive Second-Moment Preconditioning [18.570226339282296]
Muon has emerged as a promising computation for large-scale foundation model pre-training.<n>We propose Muon$2$, a memory-efficient factorized variant that preserves most of the gains of Muon$2$ with negligible memory overhead.
arXiv Detail & Related papers (2026-04-11T00:27:40Z) - Mousse: Rectifying the Geometry of Muon with Curvature-Aware Preconditioning [72.62839712454196]
Recent advances in spectral optimization, notably Muon, have demonstrated that constraining update steps to the Stiefel manifold can significantly accelerate training and improveimation.<n>We argue that this "egalitarian" constraint is suboptimal for Deep Neural Networks, where the curvature spectrum is known to be highly heavy-tailed and ill-conditioned.<n>We propose textbfMousse (textbfMuon textbfOptimization textbfUtilizing textbfShampoo's textbfStruct
arXiv Detail & Related papers (2026-03-10T14:03:49Z) - Decoupling Variance and Scale-Invariant Updates in Adaptive Gradient Descent for Unified Vector and Matrix Optimization [14.136955342888987]
We reform the AdaGrad update and decompose it into a variance adaptation term and a scale-invariant term.<n>This produces $textbfDeVA$ ($textbfV$ariance $textbfA$daptation), a framework that bridges between vector-based variance adaptation and matrix spectral optimization.
arXiv Detail & Related papers (2026-02-06T17:06:42Z) - Muon in Associative Memory Learning: Training Dynamics and Scaling Laws [23.350512542598803]
We study Muon in a linear associative memory model with softmax retrieval and a hierarchical frequency spectrum over query-answer pairs.<n>We show that Muon mitigates this imbalance, leading to faster and more uniform progress.
arXiv Detail & Related papers (2026-02-05T14:49:40Z) - Delving into Muon and Beyond: Deep Analysis and Extensions [8.297062899157664]
We study Muon as the p = 0 endpoint of a family of spectral transformations of the form U boldsymbolp V'.<n>We find that RMS-normalized updates yield more stable optimization than first-moment updates.<n>Our results suggest that Muon is best understood as an effective form of spectral normalization, but not a universally superior optimization method.
arXiv Detail & Related papers (2026-02-04T15:40:47Z) - OLion: Approaching the Hadamard Ideal by Intersecting Spectral and $\ell_{\infty}$ Implicit Biases [29.60546958677364]
nameA combines spectral control from update directions with coordinate control from sign updates.<n>We prove convergence under a mild, empirically verified diagonal-isotropy assumption.<n>nameA matches or outperforms AdamW and Muon under comparable tuning while using only momentum-level state.
arXiv Detail & Related papers (2026-02-01T08:59:45Z) - Muon is Provably Faster with Momentum Variance Reduction [55.388203260208485]
Recent empirical research has demonstrated that deep learnings based on the linear linear oracle (LMO) over specifically chosen Non-Eudean.<n>Adam-type training methods outperform the minimization of large language models.
arXiv Detail & Related papers (2025-12-18T14:38:39Z) - MARS-M: When Variance Reduction Meets Matrices [47.405031764674014]
Matrix-based preconditioneds have been shown to be more efficient than scalar-based preconditioneds for large-scale neural networks.<n>We introduce MARS-M, a new technique that integrates the variance reduction technique in MARS with Muon.<n>Our empirical results on language modeling and computer vision tasks demonstrate that MARS-M consistently yields lower losses and improved performance across various downstream benchmarks.
arXiv Detail & Related papers (2025-10-20T16:49:22Z) - MuonBP: Faster Muon via Block-Periodic Orthogonalization [24.232069944820513]
We show how to adjust the learning rate from the baseline to MuonBP and give guarantees for this algorithm.<n>When training an 8B model with eight-way tensor tensor and ZeRO statewiseing, MuonBP achieves 8% Muon with no degradation in performance.
arXiv Detail & Related papers (2025-10-19T19:56:05Z) - Scale-Invariant Regret Matching and Online Learning with Optimal Convergence: Bridging Theory and Practice in Zero-Sum Games [60.871651115241406]
A considerable chasm has been looming for decades between theory and practice in zero-sum game solving through first-order methods.<n>We propose a new scale-invariance and parameter-free variant of PRM$+$, which we call IREG-PRM$+$.<n>We show that it achieves $T-1/2$ best-iterate and $T-1$ optimal convergence guarantees, while also being on par with PRM$+$ on benchmark games.
arXiv Detail & Related papers (2025-10-06T00:33:20Z) - Muon is Scalable for LLM Training [50.68746986439438]
We introduce Moonlight, a Mixture-of-Expert (MoE) model trained with 5.7T tokens using Muon.<n>Our model improves the current frontier, achieving better performance with much fewer training FLOPs compared to prior models.<n>We open-source our distributed Muon implementation that is memory optimal and communication efficient.
arXiv Detail & Related papers (2025-02-24T09:12:29Z)
This list is automatically generated from the titles and abstracts of the papers in this site.
This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.