CMuon: Accelerating and Stabilizing Diffusion Transformer Training via Chunked Momentum Orthogonalization

AI summary

Overview

• Research area: Deep learning optimization for generative models — specifically, optimizers for training Diffusion Transformers (DiTs) used in visual image synthesis.
• Technical level: Intermediate. The core idea is intuitive (split a fused matrix before orthogonalizing), but the paper assumes familiarity with diffusion/flow-matching training, momentum-based optimizers, singular value decomposition, and Newton–Schulz iteration.
• Scope in one sentence: The paper diagnoses why the Muon optimizer's advantage over AdamW disappears late in DiT training, proposes a low-overhead fix called Chunked Muon (CMuon), and validates it on class-conditional ImageNet-1K at 256×256 for 130M- and 675M-parameter DiT models.

What This Paper Is About

Diffusion Transformers achieve state-of-the-art visual generation quality, but training them is computationally prohibitive. The Muon optimizer speeds up early training relative to the standard AdamW optimizer, yet in DiTs it plateaus later and ends up at roughly the same Fréchet Inception Distance (FID) as AdamW. This paper traces that plateau to a specific architectural detail — standard DiTs concatenate functionally unrelated weight blocks (AdaLN, QKV, FFN) into single tensors, and orthogonalizing the fused tensor implicitly couples those unrelated subspaces. The goal is to remove that coupling while keeping Muon's speed.

Key Contributions

1. Diagnosis of "subspace interference." The authors show that naively applying Muon to DiTs yields suboptimal convergence, and trace the bottleneck to unintended orthogonalization coupling among functionally disjoint parameters concatenated within AdaLN, QKV, and FFN projections.
2. The CMuon algorithm. A simple modification that partitions these fused matrices back into their original functional sub-components and orthogonalizes each chunk's momentum independently, eliminating implicit cross-subspace coupling. It requires only minor code changes in the optimizer and introduces negligible computational overhead.
3. Empirical speedup and final quality. For a 675M-parameter DiT (DiT-XL), CMuon reaches FID 1.18 on ImageNet 256×256 in 200 epochs with batch size 1024, versus AdamW's 400-epoch FID of 1.21 — reported as more than a 2× training speedup that persists through the whole training run, achieved without auxiliary acceleration techniques such as Representation Alignment (REPA).
4. Systematic ablations. Studies on which architectural blocks to chunk, on learning-rate scaling strategy (Vanilla, Keller–Jordan, MoonLight), on learning-rate rescaling by √N_chunk, and on learning rate / training length.

Main Findings

• Muon's early gains do not persist. On ImageNet-1K 256×256, vanilla Muon reduces FID faster than AdamW in early training but its improvement diminishes; by 200 epochs it becomes comparable to AdamW. On DiT-Base with VA-VAE at 400 epochs, AdamW and Muon both reach FID 2.78.
• CMuon keeps improving late. For DiT-XL (675M, VA-VAE, 30 NFE), CMuon reports FID 1.46 at 80 epochs and 1.18 at 200 epochs, versus AdamW's 1.66 / 1.30 / 1.21 at 80 / 200 / 400 epochs and Muon's 1.65 / 1.29 at 80 / 200 epochs.
• Gains hold at smaller scale and with a different VAE. For DiT-B (130M) with VA-VAE at 30 NFE: AdamW 5.87 / 3.49 / 2.78 (80 / 200 / 400 epochs), Muon 5.50 / 3.31 / 2.78, CMuon 5.14 / 3.03 / 2.57. With SD-VAE at 80 / 200 epochs: AdamW 5.87 / 3.60, Muon 5.39 / 3.30, CMuon 4.94 / 3.06.
• Chunking blocks is complementary. On DiT-B, chunking only one block type gives at most a 0.1 FID reduction at 200 epochs (None 3.31; FFN 3.28; QKV 3.35; AdaLN 3.23). Chunking FFN, QKV, and AdaLN jointly drops final FID from 3.31 to 3.02, with FID@80ep also improving from 5.50 to 5.14.
• CMuon helps under every scaling rule tested. At 200 epochs on DiT-B, Muon scored 8.73 / 7.40 / 3.31 and CMuon scored 6.94 / 6.00 / 3.03 for the Vanilla, Keller–Jordan, and MoonLight scaling strategies respectively. MoonLight was adopted as the default.
• Rescaling and chunking play different roles. Rescaling chunked blocks by √N_chunk mainly accelerates early training: at 40 epochs, Muon goes from 5.67 to 4.26 and CMuon from 5.32 to 3.78. By 80 epochs, the CMuon configurations dominate (1.50 and 1.46 versus 1.65 and 1.55 for the Muon configurations).
• Gains are not an artifact of poor AdamW tuning. Across learning rates {1, 2, 3}×10⁻⁴ on DiT-XL, 2×10⁻⁴ is near-optimal for both optimizers. CMuon reaches FID 1.27 at 140 epochs with 2×10⁻⁴, closely matching AdamW's 200-epoch result.
• Qualitative improvement. With the same DiT-XL setup, 200-epoch budget, and 30 NFE, CMuon produces visually cleaner and more coherent samples than vanilla Muon, reducing texture artifacts and structural distortions.
• Theoretical analysis location. The paper states that further theoretical discussion and analysis of CMuon is provided in Appendix 0.B; the specific theoretical results are not reproduced in the provided content.

Methodology in Plain English

1. Set up training as flow matching. The authors train DiTs with a continuous-time Flow Matching objective: a neural network F_θ is trained to regress the constant-velocity target (z − x) along a linear interpolation path x_t = (1−t)x + tz between data and Gaussian noise. Sampling integrates the learned vector field backward from t=1 to t=0 as an ODE initial value problem.
2. Explain how Muon works. Muon treats each 2D weight matrix as a matrix rather than a bag of scalars. It accumulates momentum, then orthogonalizes the momentum matrix — replacing its singular values with ones (the polar factor UVᵀ, approximated by Newton–Schulz iterations) — and applies a scaled update with a dimensional scaling factor of 0.2·√(max(m,n)).
3. Isolate the failure mode. The authors construct a toy analysis with N stacked gradient submatrices. Applying Muon to the stacked matrix gives each block a shared preconditioner (Σ_j G_jᵀG_j)^(−1/2), which mixes covariance structure across blocks with unrelated roles. Applying Muon per block gives (G_iᵀG_i)^(−1/2), which depends only on that block's own statistics. When the dominant principal directions of different blocks are misaligned, the shared preconditioner distorts each block's preferred descent geometry.
4. Fix it by chunking. For fused DiT layers, they split along the longer dimension so each chunk is a semantically independent sub-matrix, run Newton–Schulz on each chunk, then concatenate the results. The chunking applies to AdaLN modulation, attention QKV projections, and FFN gate+up projections.
5. Keep the update magnitude unchanged. Under the Moonlight scaling formulation, they prove the Frobenius norm of the update before and after chunking is identical (both 0.2·√(N·d_out·d_in)), so chunking only redistributes the norm across chunks rather than changing the global step size.
6. Add an optional early-training boost. Rescaling the chunked blocks' learning rate by √N_chunk accelerates early convergence without hurting final performance, exposed as an optional switch.
7. Hybrid optimizer assignment. CMuon/Muon optimize only the 2D weights in attention, FFN, and AdaLN projections; AdamW handles 1D parameters, embeddings, and the final layer, with β₁=0.9, β₂=0.95, and zero weight decay.
8. Evaluate. Training uses a global batch size of 1024, bf16 precision, gradient clipping at max norm 1.0, EMA decay 0.9999, and FID-50K on EMA models, with 30 NFEs unless otherwise specified.

Concrete chunking scheme for DiT-XL (d = 1152): QKV [3456, 1152] → 3 chunks of [1152, 1152]; MLP gate+up [6144, 1152] → 2 chunks of [3072, 1152]; AdaLN modulation [6912, 1152] → 6 chunks of [1152, 1152].

Why This Matters

• Impact on research: The paper reframes a widely observed but unexplained phenomenon — Muon's fading late-stage advantage — as an architectural artifact rather than an inherent optimizer limitation. This is a small, cheap fix that could change how practitioners configure Muon on any transformer that uses fused projection tensors, and it feeds into the broader line of linear-minimization-oracle optimization research (Scion, Gluon, PolarExpress).
• Real-world applications:
  • Training text-to-image diffusion models at scale, where the paper notes models such as an 8.3-billion-parameter HunyuanVideo 1.5 already use Muon across multi-stage text-to-image, text-to-video, and image-to-video regimens.
  • Few-step and real-time image generation pipelines (the paper cites pixel MeanFlow, which uses Muon for pre-training few-step pixel-space DiTs).
  • Reducing compute budgets for latent-space image synthesis using modern VAEs (VA-VAE, SD-VAE) in production generative systems.
  • Lowering the cost of training open-source SOTA image generation models built on DiT backbones.
• Industry relevance: The claimed speedup is a direct reduction in GPU-hours for one of the most expensive training workloads in visual AI, achieved without auxiliary techniques such as REPA and with negligible computational overhead in the optimizer itself.

Future Directions

• Theoretical grounding. The paper defers further theoretical analysis of CMuon to Appendix 0.B; formalizing precisely when subspace interference is severe enough to warrant chunking — and how to detect it automatically — remains open.
• Automatic chunking discovery. The current scheme hard-codes chunk counts for DiT-XL (3 for QKV, 2 for FFN gate+up, 6 for AdaLN). Whether these splits can be inferred from gradient statistics rather than fixed per architecture is not addressed.
• Generalization beyond DiT-XL and DiT-B. The chunk configuration table is specified for DiT-XL with d = 1152; how the scheme transfers to other hidden dimensions, other transformer families, and non-visual modalities is not reported.
• Interaction with other optimizers and acceleration methods. The paper reports results without REPA and compares against AdamW and vanilla Muon. How CMuon interacts with other Muon variants (Scion, Gluon, PolarExpress) or with representation-alignment methods is not reported in the provided content.

Target Audience

Machine learning researchers and engineers working on optimizers, large-scale training efficiency, and diffusion-based generative models. It is most useful to practitioners already running Muon or considering switching from AdamW for DiT training, and to optimizer researchers interested in how weight fusion in modern architectures interacts with matrix-wise preconditioning. Readers need working familiarity with diffusion/flow matching, momentum optimizers, and orthogonalization to follow the full derivation.

Authors’ abstract

Diffusion Transformers (DiTs) have achieved state-of-the-art (SOTA) performance in visual generative modeling, yet their training remains computationally prohibitive. While the recently proposed Momentum Orthogonalization (Muon) optimizer offers a promising alternative to AdamW, its direct application to DiTs yields suboptimal late-stage convergence. In this paper, we identify the root cause of this bottleneck: standard DiT architectures fuse functionally distinct weights (e.g., within AdaLN and QKV layers) into unified tensors for computational efficiency. Applying Muon to these fused tensors inadvertently induces implicit subspace coupling, which distorts update directions and degrades global optimization. To address this, we introduce Chunked Muon (CMuon), a simple yet highly effective strategy that partitions these matrices into independent sub-components prior to orthogonalization. Extensive experiments demonstrate that a 675M-parameter DiT trained with CMuon achieves a FID of 1.18 on ImageNet 256 in just 200 epochs. This represents more than a 2x training speedup over AdamW, while effectively overcoming the late-stage convergence plateaus of vanilla Muon.