Scaling Muon for Diffusion Transformers

• URL: http://arxiv.org/abs/2608.20818v3
• Date: Wed, 26 Aug 2026 04:08:20 GMT
• Title: Scaling Muon for Diffusion Transformers
• Abstract summary: matrix-aware Muon improves large model training by balancing updates across singular directions, yet its scaling behavior and end-to-end efficiency on large Transformers (DiTs) remain unclear.<n>We first establish Muon's scaling behavior on DiTs from 1.3B to 15B parameters, showing that its optimization and generative quality advantages over AdamW persist across model scales.<n>We introduce Periodic Row-wise Muon, which performs a full NS5 spectral update once every (K) steps and applies a low compute and communication cost row-wise constrained update based on the current momentum at the remaining steps
• Score: 32.20758643343064
• License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/
• Abstract: The matrix-aware optimizer Muon improves large model training by balancing updates across singular directions, yet its scaling behavior and end-to-end efficiency on large Diffusion Transformers (DiTs) remain unclear. We first establish Muon's scaling behavior on DiTs from 1.3B to 15B parameters, showing that its optimization and generative quality advantages over AdamW persist across model scales. However, at scale, the 5-step Newton--Schulz iteration (NS5) performed at every optimization step, together with full-momentum materialization, introduces substantial computation and communication overhead that can offset Muon's step-efficiency advantage. We introduce \emph{Periodic Row-wise Muon}, which performs a full NS5 spectral update once every \(K\) steps and applies a low compute and communication cost row-wise constrained update based on the current momentum at the remaining steps. We further co-design a distributed implementation that operates directly on sharded momentum during non-refresh steps and accelerates spectral refreshes through bucketed all-gather and communication--computation overlap. Across all scales, Muon improves the best observed generative quality over AdamW by 12.9--19.1\%. Compared with vanilla Muon, Periodic Row-wise Muon remains within 0.5\% in best generative quality on the 1.3B--4B models and improves it by 4.5\% at 9B. It reduces optimizer time by 46.9--54.3\%, end-to-end step time by 15.7--24.3\%, and logical communication volume by 66.7\%, while reaching its respective best generative quality with 33.7--64.8\% less active training time. These results show that Periodic Row-wise Muon preserves Muon's generative quality advantage while translating it into end-to-end training efficiency for large DiTs.

Related papers

• On the Design of Qwen3.8-Next Architecture: Evaluation, Efficiency, and Training Stability [57.03393882933515]
  We describe the architecture and ablations of Qwen3.8-Flash-Next, a sparse mixture-of-experts model with 125B parameters, 6B activated per token, and 51B parameters of n-gram embedding tables held off the accelerator.<n>On fourteen pre-training benchmarks the model leads the 397B-A17B predecessor on eight and trails it on the rest by at most 2.6 points, at 1/3 the activated parameters, 1/3 the training tokens, and roughly 1/9 the training FLOPs.
  arXiv  Detail & Related papers  (2026-08-31T06:35:07Z)
• SOAP, Muon, and Beyond: Pushing LLM Pretraining Scales [48.03325171473355]
  Higher-order LLMs such as Muon and SOAP offer faster convergence than AdamW, but their computational cost and numerical stability challenges have limited adoption at scale.<n>We adapt and enhance preconditioned methods to overcome the practical challenges of large-scale LLM pretraining regimes.
  arXiv  Detail & Related papers  (2026-07-13T21:47:08Z)
• DMuon: Efficient Distributed Muon Training with Near-Adam Overhead [13.192517400076746]
  We present DMuon, an open-source distributed Muon implementation that integrates into existing training pipelines as a drop-in module.<n>DMuon achieves a 1.48x-3.01x speedup in end-to-end step time and a 6.85x-163.00x speedup in-step time, bringing per-step latency to near-AdamW levels.
  arXiv  Detail & Related papers  (2026-06-25T15:23:03Z)
• Nemotron 3 Super: Open, Efficient Mixture-of-Experts Hybrid Mamba-Transformer Model for Agentic Reasoning [262.0414794728708]
  We describe the pre-training, post-training, and quantization of Nemotron 3 Super.<n>Nemotron 3 Super is a 120 billion (active 12 billion) parameter hybrid Mamba-Attention Mixture-of-Experts model.
  arXiv  Detail & Related papers  (2026-04-14T07:02:32Z)
• MuonRec: Shifting the Optimizer Paradigm Beyond Adam in Scalable Generative Recommendation [60.1890607252082]
  MuonRec is the first framework that brings the proposed Muon iteration to RecSys training.<n>We develop an open-source training recipe for recommendation models and evaluate it across both traditional sequential recommenders and modern generative recommenders.
  arXiv  Detail & Related papers  (2026-02-28T02:32:44Z)
• 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)
• 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)
• Muon: Training and Trade-offs with Latent Attention and MoE [4.500362688166346]
  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.
  arXiv  Detail & Related papers  (2025-09-29T07:51:06Z)
• 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.