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

Conference: ECCV2026
Paper: Official Paper Page · PDF
Area: Training Efficiency / Diffusion Models
Keywords: chunked momentum orthogonalization, subspace interference, diffusion Transformer, functional partitioning, learning-rate scaling

TL;DR

CMuon independently orthogonalizes momentum for the functional QKV, FFN gate/up, and AdaLN components stored in fused DiT tensors, enabling a 675M DiT-XL to reach FID 1.18 on ImageNet-1K 256x256 in 200 epochs, better than AdamW's 1.21 after 400 epochs, with gains extending beyond early training.

Background & Motivation

Muon's appeal is that it does not merely rescale gradients element by element: it orthogonalizes momentum matrices for two-dimensional weights, reshaping the update's singular-value spectrum so that relatively weak directions are not continually overwhelmed by strong ones. However, faster initial loss reduction in a diffusion Transformer does not automatically imply better final generation quality. In this paper's DiT-B experiment, Muon and AdamW both reach FID 2.78 after 400 epochs. The problem is therefore not simply how to start training faster, but why the initial advantage disappears later.

The authors trace the issue to something often treated as an implementation detail: concatenating functionally distinct projections into one matrix for efficient computation. Q, K, and V can be produced by a single linear operation, as can AdaLN's scaling, shifting, and gating outputs. Yet when an optimizer orthogonalizes the entire tensor, the storage boundary also becomes a statistical coupling boundary. Projections sharing an input dimension need not share dominant gradient directions, and a common preconditioner can let one branch's statistics redirect another branch's update.

Core idea: retain efficient fused forward computation, but restore functional parameter boundaries inside the optimizer by partitioning first, orthogonalizing independently, and explicitly calibrating the resulting update scale, so implementation-level concatenation no longer forces distinct functions to share optimization geometry.

Method

Overall Architecture

CMuon changes the optimization step, not the generator's forward architecture, and introduces no representation-alignment loss. Each iteration first backpropagates through the flow-matching objective and forms a Nesterov-style momentum update. Selected two-dimensional weights then undergo functional partitioning, independent orthogonalization, and scale calibration before their updates are concatenated back into the original shape. AdamW handles the remaining parameters.

Here, "independent" means that orthogonalization statistics are not mixed across functional blocks. It does not mean that the network's components learn independently: all blocks still train jointly within the same model and objective, and backpropagation continues to transmit interactions through the network.

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    input["Flow Matching<br/>backpropagation"] --> momentum["Nesterov momentum<br/>for 2D projections"]
    input --> adam["Other parameters<br/>AdamW"]
    momentum --> split["Functional partitioning"]
    split --> orth["Independent orthogonalization"]
    orth --> scale["Scale calibration"]
    scale --> update["Concatenate and update weights"]
    adam --> update

Key Designs

1. Functional partitioning: separate computational fusion from optimizer grouping

Arbitrarily cutting a matrix into equal-sized pieces is not the central proposal: chunks should correspond to projections with distinct original functions. In DiT-XL with hidden dimension 1152, the QKV weight has shape [3456, 1152] and is split along dimension 0 into 3 matrices of shape [1152, 1152]. The FFN gate/up weight [6144, 1152] becomes 2 chunks of [3072, 1152]. AdaLN's [6912, 1152] becomes 6 chunks of [1152, 1152], covering modulation and gating for the attention and FFN branches. All three matrix types are split along their longer dimension, without changing the parameter count or forward semantics.

This matters because the row ranges within a fused tensor are not interchangeable feature groups. Q, K, and V respectively participate in attention matching and content aggregation, while FFN gate and up projections also serve different computational roles. CMuon allows those roles to use their own statistics during orthogonalization, rather than letting one strong-gradient functional block influence every other block through a shared preconditioner. An implementation must identify boundaries from the actual model's parameter layout; a matrix should not be split mechanically just because it is tall.

2. Independent orthogonalization: remove directional interference from a shared preconditioner

Muon is easiest to understand through the idealized exact polar factor: if the momentum matrix has singular-value decomposition \(M=U\Sigma V^\top\), then \(\mathrm{Orth}(M)=UV^\top\). This largely retains directional structure while reshaping singular-value magnitudes. In practice, Newton-Schulz iterations approximate the operation instead of computing an explicit SVD at every step. For rectangular matrices, semi-orthogonality is the more precise description: all rows and columns cannot necessarily be orthogonal simultaneously.

The authors illustrate the source of coupling using vertically stacked gradient blocks. The following is a normalized transcription of Section 3.3, assuming invertible Gram matrices. In the actual optimizer, the input should be understood as the momentum-adjusted update matrix, not as raw gradients bypassing momentum:

\[ U_i^{\mathrm{fused}}=G_i\left(\sum_{j=1}^{N}G_j^\top G_j\right)^{-1/2}, \qquad U_i^{\mathrm{chunk}}=G_i(G_i^\top G_i)^{-1/2}. \]

The difference lies in which statistics enter the inverse square root. The fused version uses the sum of all blocks' Gram matrices for every block. When their principal directions disagree, one block changes the directional scaling applied to another. The chunked version avoids mixing these statistics, then concatenates the updates in their original order. This derivation explains a possible interference mechanism, but is not a convergence theorem establishing that chunking wins for every dataset and architecture.

Algorithm 1 first accumulates the momentum buffer as \(M\leftarrow\mu M+G\), then constructs \(G+\mu M\) as the input to orthogonalization, and only then splits it into chunks. Orthogonalizing raw gradient blocks before accumulating momentum would be a different algorithm, because orthogonalization is nonlinear. The cache does not include the referenced appendix Algorithm 3, so its precise Newton-Schulz coefficients and iteration count cannot be supplied from this source.

3. Scale calibration: distinguish directional improvements from larger steps

Chunking changes matrix shapes and the norms of the resulting semi-orthogonal matrices. If the fused matrix's scaling is reused without adjustment, part of the apparent acceleration may simply come from a larger step. CMuon defaults to Moonlight scaling, computing \(\alpha_c=0.2\sqrt{\max(d_{\mathrm{out}},d_{\mathrm{in}})}\) from each chunk's shape and applying the corresponding update with weight decay. This makes the AdamW base learning rate easier to reuse, but does not imply that all scaling rules work without retuning.

For \(N\) vertically stacked chunks, each satisfying \(d_{\mathrm{out}}\ge d_{\mathrm{in}}\), the fused and chunked updates obey the following relationship under exact semi-orthogonality. Here \(G\) is the full stacked matrix and \(\alpha=0.2\sqrt{Nd_{\mathrm{out}}}\):

\[ \|\alpha\,\mathrm{Orth}(G)\|_F =\left\|\alpha_c \begin{bmatrix}\mathrm{Orth}(G_1)\\ \vdots\\ \mathrm{Orth}(G_N)\end{bmatrix}\right\|_F =0.2\sqrt{Nd_{\mathrm{out}}d_{\mathrm{in}}}. \]

Thus, default chunking does not simply enlarge the total update: it redistributes update norms and directions across chunks. A separate optional rescale switch multiplies the scale by \(\sqrt{N_{\mathrm{chunk}}}\) to accelerate early training. Once enabled, the equal-global-norm interpretation above no longer applies. The authors also apply the same rescaling to corresponding layers in unchunked Muon, helping separate the effects of larger steps from statistical decoupling.

A Worked Example

Consider the DiT-XL QKV layer. Backpropagation produces a [3456, 1152] gradient, the optimizer updates momentum in that original shape, and the Nesterov input is split along its rows into Q, K, and V chunks. Each chunk independently approximates its polar factor, receives the per-chunk Moonlight multiplier \(0.2\sqrt{1152}\), and is concatenated back into [3456, 1152] to update the same fused weight tensor.

Without chunking, the corresponding multiplier is \(0.2\sqrt{3456}\). Although each chunk receives a smaller multiplier, there are now 3 chunks contributing to the total norm, so its default value is preserved. Enabling rescale adds a factor of \(\sqrt{3}\) to each chunk's multiplier. This example only expands the paper's tensor shapes and scaling rules; it is not an additional numerical experiment, and does not require replacing the fused forward QKV layer with three modules.

Loss & Training

Training retains Flow Matching: data occupy time 0, noise occupies time 1, and the model predicts the velocity along a linear interpolation between them. The cached equations contain extraction damage. The following transcription follows the clear textual definitions in Section 3.1 without inventing a time-sampling distribution:

\[ x_t=(1-t)x+tz,\qquad \mathcal L(\theta)=\mathbb E_{x,z,c,t} \left[\|F_\theta(x_t,t,c)-(z-x)\|_2^2\right]. \]

Inference integrates from noise at time 1 toward time 0, and \(c\) represents the class condition in the paper's class-conditional experiments. The optimizer change neither modifies this generative objective nor introduces auxiliary representation-alignment techniques such as REPA.

Experiments use batch size 1024, bf16, a constant learning rate, gradient clipping at maximum norm 1.0, and EMA evaluation with decay 0.9999. AdamW uses \((\beta_1,\beta_2)=(0.9,0.95)\) and weight decay 0. Only two-dimensional attention, FFN, and AdaLN weights use Muon/CMuon; one-dimensional parameters, embeddings, and the final layer remain under AdamW. The DiT-XL learning-rate analysis tests \(\{1,2,3\}\times10^{-4}\), with \(2\times10^{-4}\) performing well near the default setting. The cache does not provide the complete appendix-level reproduction configuration.

Key Experimental Results

Main Results

The following selection comes from Table 2. The task is class-conditional ImageNet-1K 256x256 generation, using 30 NFE, EMA, and FID-50K throughout; lower FID is better. NFE counts function evaluations during sampling, not training steps. "Not reported" must not be interpreted as an unsuccessful run.

Model / VAE	Optimizer	FID@80ep	FID@200ep	FID@400ep
DiT-B 130M / VA-VAE	AdamW	5.87	3.49	2.78
DiT-B 130M / VA-VAE	Muon	5.50	3.31	2.78
DiT-B 130M / VA-VAE	CMuon	5.14	3.03	2.57
DiT-XL 675M / VA-VAE	AdamW	1.66	1.30	1.21
DiT-XL 675M / VA-VAE	Muon	1.65	1.29	Not reported
DiT-XL 675M / VA-VAE	CMuon	1.46	1.18	Not reported

At the same 200-epoch XL budget, CMuon improves FID over Muon by 0.11 and over AdamW by 0.12. Its 1.18 after 200 epochs also beats AdamW's 1.21 after 400 epochs, supporting the statement that it achieves better quality with half as many training epochs. The paper's "over 2x speedup" wording should not be substituted for a measured wall-clock or GPU-hour reduction.

Ablation Study

The following results are from Table 3, fixing the 130M DiT-B and changing only which projections are chunked. None corresponds to standard Muon.

Chunked projections	FID@80ep	FID@200ep	Late-stage change versus None
None	5.50	3.31	Baseline
FFN	5.43	3.28	Lower by 0.03
QKV	5.32	3.35	Higher by 0.04
AdaLN	6.00	3.23	Lower by 0.08, but worse early
FFN + QKV + AdaLN	5.14	3.02	Lower by 0.29

Table 3 reports 3.02 for the full configuration, while Table 2 reports 3.03 for the corresponding setting. The cache does not explain this difference, so both are retained as reported. More importantly, the prose claim that every individual block improves late-stage performance should not be repeated uncritically: QKV-only reaches 3.35, worse than 3.31. The evidence supports the combined configuration, not a monotonic benefit from each individual block.

Key Findings

• Table 5 controls chunking and rescaling separately: Muon yields FID@40ep/@80ep of 5.67/1.65, Muon + rescale 4.26/1.55, CMuon 5.32/1.50, and CMuon + rescale 3.78/1.46. These results support complementarity, but 80 epochs alone cannot establish the ultimate benefit for arbitrarily long training.
• In Table 4, Vanilla, KellerJordan, and Moonlight scaling produce Muon FIDs of 8.73, 7.40, and 3.31, versus CMuon FIDs of 6.94, 6.00, and 3.03. The authors reuse an AdamW learning rate without retuning the first two rules, so these are not fully tuned rankings of scaling strategies.
• In Table 6, AdamW reaches 1.26 at \(3\times10^{-4}\) after 200 epochs, while CMuon reaches 1.27 at \(2\times10^{-4}\) after 140 epochs: close, not better. That table also gives AdamW 1.29 at \(2\times10^{-4}\) and 200 epochs, another unexplained small difference from Table 2's 1.30.
• With SD-VAE, the DiT-B comparison at 200 epochs gives AdamW 3.60, Muon 3.30, and CMuon 3.06. Gains are therefore not confined to the default VA-VAE, although this remains evidence within the same dataset and resolution.

Highlights & Insights

• Tensor layout is not optimization-neutral. For a matrix-level optimizer, fusing projections changes the statistics seen by the preconditioner. Efficient computational layouts and meaningful optimizer groups should be designed together.
• Norm control makes the mechanism easier to distinguish. Default norm-preserving scaling and the separate rescale ablation help distinguish directional correction from larger steps, avoiding the assumption that all acceleration comes from decoupling.
• Inspect the full quality trajectory, not just early loss. On the B model, Muon ends at AdamW's 2.78 after 400 epochs, while CMuon reaches 2.57. This is more relevant to generative training objectives than an early-convergence plot alone.

Limitations & Future Work

• Limited evaluation scope. Evidence centers on ImageNet-1K 256x256 and 130M/675M DiTs. It does not directly validate large-scale text-to-image, video, or other architectures, nor establish that functional partitioning helps every matrix optimizer.
• Efficiency is primarily measured through training progress. The authors describe the extra overhead as negligible, but the supplied cache lacks complete wall-clock, memory, throughput, or distributed-communication tables. Whether multiple smaller matrix operations preserve hardware efficiency needs independent measurement.
• Incomplete ablations and statistical reporting. Single-block and all-block comparisons do not isolate every interaction. Pairwise combinations, random equal-sized partitions, seed variance, and confidence intervals are missing. Comparing functional and arbitrary partitions would better identify the source of the gains.
• Source and reporting boundaries must remain explicit. The cache ends with references and omits Appendices A/B and Algorithm 3. It also contains the cross-table differences 3.02/3.03 and 1.29/1.30. Exact Newton-Schulz settings, theoretical guarantees, or a supposedly corrected single result cannot be filled in without evidence.

Related Work & Insights

• Versus AdamW: AdamW uses elementwise adaptive statistics, whereas CMuon applies matrix-level momentum orthogonalization to selected two-dimensional weights while retaining AdamW elsewhere. This is a hybrid optimization scheme, not a complete replacement for AdamW.
• Versus Muon / Moonlight: Muon supplies the polar-factor update and Moonlight supplies shape-dependent scaling. CMuon's main addition is functional partitioning before orthogonalization with corresponding scale handling, not an entirely new orthogonalization operator.
• Versus REPA / VA-VAE: REPA improves training through representation alignment, VA-VAE changes latent-space quality, and CMuon changes optimization geometry. Results without REPA demonstrate that chunking can help on its own, but these experiments do not establish additional gains when combined with REPA.

Rating

• Novelty: 4/5. A simple change identifies how fused tensor boundaries affect matrix-level optimization and gives a concrete mechanistic explanation.
• Experimental Thoroughness: 4/5. Covers model scale, VAE, chunked layers, scaling, and learning rate, but lacks cross-task tests, repeated seeds, and hardware-cost measurements.
• Writing Quality: 3/5. The problem and algorithm are clear, but cross-table values and the prose interpretation of single-block ablations are inconsistent; cached equations and appendices are also incomplete.
• Value: 4/5. Directly useful for DiT training with fused projections, provided parameter layouts and actual throughput are checked before transferring the approach.