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

TL;DR

CMuon (Chunked Muon) is a Muon-optimizer variant tailored for Diffusion Transformer training. The authors diagnose why vanilla Muon plateaus late in DiT training: standard DiT layers fuse functionally distinct weights (AdaLN modulation vectors and Q/K/V projections) into single tensors for compute efficiency, and running the momentum-orthogonalization step over those fused tensors induces “implicit subspace coupling” that distorts update directions. The fix is a one-line change — chunk the fused matrices into their functional sub-components before the Newton–Schulz msign step, then orthogonalize each chunk independently. On ImageNet 256², a 675M DiT trained with CMuon reaches FID 1.18 in 200 epochs, reported as >2× training speedup over AdamW and cleanly past the late-stage plateau of vanilla Muon.

Key claims

Directly applying vanilla Muon to standard DiTs yields suboptimal late-stage convergence — a specific failure mode of Muon on this architecture, not the AdamW baseline [Abstract].
The root cause is architectural: AdaLN and QKV layers fuse functionally distinct weights into single tensors for efficiency, and Muon’s per-tensor momentum orthogonalization treats those fused tensors as one operator, inducing implicit subspace coupling that distorts update directions [Abstract].
Partitioning the fused matrices into their functional sub-components before orthogonalization (Chunked Muon) removes the coupling and restores convergence [Abstract].
A 675M-parameter DiT trained with CMuon reaches FID 1.18 on ImageNet 256² in 200 epochs [Abstract].
Reported >2× training speedup over AdamW at matched FID, while overcoming the late-stage plateaus that vanilla Muon exhibits on the same architecture [Abstract].

Method

Muon (Momentum Orthogonalization) replaces AdamW’s per-parameter adaptive scaling with a matrix-level operation: take the momentum buffer for each 2-D weight and apply a Newton–Schulz iteration to approximate its polar factor (equivalently, msign(M) = UVᵀ from the SVD M = UΣVᵀ). This gives a spectrally-normalized update direction that has become the de-facto default for LLM pretraining.

The DiT-specific problem CMuon identifies is that this operation is per-tensor, and DiT codebases (following DiT / Peebles-Xie) concatenate multiple functionally distinct projections into a single tensor for GEMM efficiency: (1) AdaLN produces (scale, shift, gate) triples for the attention and MLP branches from one linear layer; (2) QKV projection produces Q, K, and V from one linear layer. Orthogonalizing the fused tensor mixes the singular subspaces of the sub-projections, so the update to (say) the Q block is no longer independent of the update to K or V, and the update to the AdaLN scale block is coupled to the shift and gate blocks.

CMuon splits each fused tensor along its output-feature axis into its functional chunks, applies the Newton–Schulz orthogonalization independently to each chunk, and re-concatenates. Everything else in the optimizer state and update rule is unchanged. This is a pure decoupling fix — no new hyperparameters, no additional passes over the tensor.

Results

675M DiT / ImageNet 256² / 200 epochs → FID 1.18 [Abstract].
Reported >2× speedup over AdamW at matched FID [Abstract].
Vanilla Muon exhibits a late-stage plateau on DiTs that CMuon does not — the chunking is what unlocks continued FID improvement past the plateau [Abstract].

Why it’s interesting

CMuon is the first filed paper on the wiki that makes a diffusion-transformer-specific correction to Muon, distinct from the LLM-side lineage tracked in Training stability at scale. Where SSO (Controlled LLM Training on Spectral Sphere) argues Muon is “half-aligned with µP” and fixes it by constraining weight spectra, and Aurora (Aurora: A Leverage-Aware Optimizer for Rectangular Matrices) fixes row-norm non-uniformity on rectangular matrices, CMuon locates a third pathology that only surfaces on architectures where multiple functionally distinct projections are fused for GEMM efficiency — a pattern near-universal in DiTs and largely absent from LLM decoder stacks. The Krea 2 report (Krea 2 Technical Report) already noted that Muon on a production MMDiT is stable “only after excluding the first and last linear layers and adding Nesterov momentum,” which reads in retrospect like a lower-friction hack for the same underlying coupling problem CMuon isolates.

On the Diffusion training efficiency side, this adds a new axis: rather than changing supervision (REPA / SRA / Self-Flow), latent design (UL, RAE), spatial tokens (Foveated), or step assignment (DiffusionBlocks), CMuon changes the optimizer’s interaction with the architecture. The claimed >2× AdamW speedup at 675M on ImageNet 256² is competitive with the multi-lever recipes in that concept (RAEv2’s 10× speedup) but from a single-file optimizer change — worth cross-checking against MMDiT / video-DiT stacks where the AdaLN and QKV fusions are the same.