Muon: Training and Trade-offs with Latent Attention and MoE

• URL: http://arxiv.org/abs/2509.24406v1
• Date: Mon, 29 Sep 2025 07:51:06 GMT
• Title: Muon: Training and Trade-offs with Latent Attention and MoE
• Authors: Sushant Mehta, Raj Dandekar, Rajat Dandekar, Sreedath Panat,
• Abstract summary: We present a comprehensive theoretical and empirical study of the Muon for training transformers only with a small to medium decoder (30M - 200M parameters)<n>We provide rigorous theoretical analysis including: (i)showing the convergence rate under standard assumptions, (ii) spectral regularization properties that prevent gradient explosion, (iii) connection to natural gradient descent on the Stiefel manifold, and (iv) equivalence to steepest gradient descent under the spectral norm.
• Score: 4.500362688166346
• License: http://creativecommons.org/licenses/by/4.0/
• Abstract: We present a comprehensive theoretical and empirical study of the Muon optimizer for training transformers only with a small to medium decoder (30M - 200M parameters), with an emphasis on its mathematical foundations, convergence properties and synergistic interactions with modern architectural optimizations. Building on recent work showing Muon's scalability, we provide rigorous theoretical analysis including: (i)showing the convergence rate under standard assumptions, (ii) spectral regularization properties that prevent gradient explosion, (iii) connection to natural gradient descent on the Stiefel manifold, and (iv) equivalence to steepest gradient descent under the spectral norm. Crucially, we demonstrate that Muon expands the Pareto frontier in the compute-time trade-off by maintaining superior data efficiency at large batch sizes, a key finding of~\cite{essentialai2025muon} that we validate across our model scales. Empirically, Muon reaches the target loss with 48-52\% of the training calculated by AdamW while maintaining or improving the final perplexity, consistent with larger-scale results. When combined with Multi-Head Latent Attention (MLA) and Mixture-of-Experts (MoE), we observe multiplicative efficiency gains: MLA+MoE+Muon achieves 68\% memory reduction and 3.2$\times$ inference speedup, while improving perplexity by 8-12\%. We provide detailed procedures on 15 architectural and optimizer components, stability analyzes across 100+ training runs, and practical implementation guidelines including Newton-Schulz coefficients $(3.4445, -4.7750, 2.0315)$ optimized by~\cite{su2024muonblog}. Our theoretical analysis and comprehensive experiments establish Muon as a principled, robust alternative to AdamW that particularly excels when combined with modern efficiency techniques and large-batch training regimes.

Related papers

• Data Distribution as a Lever for Guiding Optimizers Toward Superior Generalization in LLMs [60.68927774057402]
  We show, for the first time, that a lower simplicity bias induces a better generalization.<n>Motivated by this insight, we demonstrate that the training data distribution by upsampling or augmenting examples learned later in training similarly reduces SB and leads to improved generalization.<n>Our strategy improves the performance of multiple language models including Phi2-2.7B, Llama3.2-1B, Gemma3-1B-PT, Qwen3-0.6B-Base-achieving relative accuracy gains up to 18% when fine-tuned with AdamW and Muon.
  arXiv  Detail & Related papers  (2026-01-31T07:40:36Z)
• TEON: Tensorized Orthonormalization Beyond Layer-Wise Muon for Large Language Model Pre-Training [30.25761898761426]
  The Muon generalization has demonstrated strong empirical performance in pre-training large language models.<n>In this work, we propose TEON, a principled generalization of Muon that extends orthogonalization beyond individual layers.<n>We present TEON's improved convergence guarantee over layer-wise Muon, and further develop a practical instantiation of TEON.
  arXiv  Detail & Related papers  (2026-01-30T18:30:12Z)
• ECO: Quantized Training without Full-Precision Master Weights [58.97082407934466]
  Error-Compensating (ECO) eliminates master weights by applying updates directly to quantized parameters.<n>We show that ECO converges to a constant-radius neighborhood of the optimum, while naive master-weight removal can incur an error that is inversely proportional to the learning rate.
  arXiv  Detail & Related papers  (2026-01-29T18:35:01Z)
• Beyond the Ideal: Analyzing the Inexact Muon Update [54.70108543057578]
  We show first analysis of the inexactized update at Muon's core.<n>We reveal a fundamental coupling between this inexactness and the optimal step size and momentum.
  arXiv  Detail & Related papers  (2025-10-22T18:01:07Z)
• NorMuon: Making Muon more efficient and scalable [71.49702449498085]
  We propose NorMuon (Neuron-wise Normalized Muon) as a successor to Adam.<n>We show NorMuon consistently outperforms both Adam and Muon, achieving 21.74% better training efficiency than Adam and 11.31% improvement over Muon on 1.1 B pretraining setting.
  arXiv  Detail & Related papers  (2025-10-07T01:13:41Z)
• On the Convergence of Muon and Beyond [31.900178928104648]
  The Muon has demonstrated remarkable success in matrix-structured parameters for neural networks.<n>A significant understanding gap persists between its theoretical and practical rate variants.<n>This work provides the first proof of optimality for a Muon-style and corroborates our findings on per-it convergence.
  arXiv  Detail & Related papers  (2025-09-19T09:43:37Z)
• DualSparse-MoE: Coordinating Tensor/Neuron-Level Sparsity with Expert Partition and Reconstruction [15.261077484922616]
  Mixture of Experts (MoE) has become a mainstream architecture for building Large Language Models (LLMs)<n>We identify dual sparsity at the tensor and neuron levels in pre-trained MoE modules as a key factor for both accuracy and efficiency.<n>We propose DualSparse-MoE, an inference system that integrates dynamic tensor-level dropping with static neuron-level reconstruction.
  arXiv  Detail & Related papers  (2025-08-25T18:08:32Z)
• Convergence Bound and Critical Batch Size of Muon Optimizer [1.2289361708127877]
  We provide convergence proofs for Muon across four practical settings.<n>We show that the addition of weight decay yields strictly tighter theoretical bounds.<n>We derive the critical batch size for Muon that minimizes the computational cost of training.
  arXiv  Detail & Related papers  (2025-07-02T11:03:13Z)
• Mind the Gap: Bridging Thought Leap for Improved Chain-of-Thought Tuning [54.65050470296886]
  We propose the CoT Thought Leap Bridge Task, which aims to automatically detect leaps and generate missing intermediate reasoning steps.<n>We demonstrate that models fine-tuned on bridged datasets consistently outperform those trained on original datasets.<n>Our approach effectively enhances distilled data and provides better starting points for reinforcement learning.
  arXiv  Detail & Related papers  (2025-05-20T17:59:31Z)
• 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)
• Mixture Compressor for Mixture-of-Experts LLMs Gains More [71.0473038084673]
  We propose a training-free Mixture-Compressor for Mixture-of-Experts large language models (MoE-LLMs)<n>Our MC integrates static quantization and dynamic pruning to collaboratively achieve extreme compression for MoE-LLMs with less accuracy loss.<n>For instance, at 2.54 bits, MC compresses 76.6% of the model, with only a 3.8% average accuracy loss.
  arXiv  Detail & Related papers  (2024-10-08T18:09:38Z)
• Sparse MoEs meet Efficient Ensembles [49.313497379189315]
  We study the interplay of two popular classes of such models: ensembles of neural networks and sparse mixture of experts (sparse MoEs) We present Efficient Ensemble of Experts (E$3$), a scalable and simple ensemble of sparse MoEs that takes the best of both classes of models, while using up to 45% fewer FLOPs than a deep ensemble.
  arXiv  Detail & Related papers  (2021-10-07T11:58:35Z)

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.