PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Chuyan Chen, Peng Sun, Kun Yuan

arXiv:2608.02502v1cs.AI

TL;DR

Muon can accelerate early DiT training, but its fully trained generative-quality advantage over AdamW remains uncertain. CMuon chunks fused DiT projections before independently orthogonalizing them; on ImageNet-1K 256 × 256, DiT-XL reaches FID 1.18 at 200 epochs, while AdamW requires 400 epochs to reach 1.21.

  • Problem

    Muon’s early acceleration on DiTs has not resolved whether it can surpass AdamW’s generative quality after full training, as late-stage convergence often plateaus.

    Source: §1 · p. 2

  • Method

    CMuon chunks fused DiT matrices into their functional sub-components and orthogonalizes each component’s momentum independently.

    Source: §1 · pp. 2–3

  • Results

    On ImageNet-1K 256 × 256, CMuon achieves FID 1.18 on DiT-XL at 200 epochs, while AdamW requires 400 epochs to reach FID 1.21.

    Source: §6 · p. 15

  • Takeaways & Limitations

    CMuon is presented as a highly efficient and stable optimization approach for large-scale DiT training.

    Source: §6 · p. 15

Abstract

from arXiv · show

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.

1 Introduction

Muon offered an alternative to AdamW, but its early acceleration on DiTs did not establish a sustained late-stage quality advantage. The paper attributes this plateau to orthogonalizing fused, functionally distinct weights together and introduces CMuon to orthogonalize their sub-components independently.

1 Introduction

  • Muon orthogonalizes momentum matrices and has accelerated large-scale pretraining, motivating its use as an alternative to AdamW [10] [16].The introduction also notes its adoption in LLM and deep-learning communities [14] [32] [37].Source: §1 · p. 1
  • Muon can accelerate early DiT training, but its final FID often only matches AdamW, leaving its fully trained quality advantage unclear.Earlier work reported approximately 2× acceleration in convergence during early epochs.Source: §1 · p. 2
  • Fusing independently initialized AdaLN, QKV, and FFN weights makes vanilla Muon’s joint orthogonalization couple unrelated sub-matrices and distort update directions.CMuon chunks fused matrices into their original functional sub-components before independently orthogonalizing their momentums.Source: §1 · pp. 2–3
  • FID 1.18 on ImageNet 256 × 256 at 200 epochs is reported for a 675M-parameter DiT; training speedup over AdamW is 2× and persists throughout training.The reported score uses CMuon without auxiliary acceleration techniques such as REPA.Source: §1 · p. 3
  • Ablations report consistent improvements over vanilla Muon across hyperparameter configurations, including matrix scaling, AdamW-reserved layers, and learning-rate scaling.The reported studies examine matrix scaling factors, layers such as embeddings and final projections reserved for AdamW, and learning-rate scaling.Source: §1 · p. 3

2 Related Works

The related work describes latent-space DiTs as a way to reduce the computational demands of image generation and reviews momentum orthogonalization’s growing role in model training. It then notes that Muon accelerates DiT training but may plateau late, motivating CMuon’s focus on architectural parameter coupling.

2.1 Diffusion Transformers for Image Generation

  • Latent diffusion uses VAE-compressed, lower-dimensional image representations to reduce the computational overhead of high-resolution generation while preserving reconstruction fidelity.The passage describes VAE-coupled diffusion as a standard framework for open-domain image synthesis.Source: §2.1 · pp. 3–4

2.2 Momentum Orthogonalization for Optimization

  • Muon and its extensions refine momentum orthogonalization, with cited studies reporting advantages across architectures and large-scale industrial pretraining [21] [22] [1] [26] [36] [14] [32].The passage frames extending Muon’s LLM pretraining success to other foundation models as a promising direction.Source: §2.2 · p. 4

2.3 Muon for DiT Training

  • Muon’s dense-matrix orthogonalization suits DiT projections, and prior work reports faster convergence, but naive application can yield suboptimal late-stage convergence from architectural parameter coupling [25] [37] [17].The cited examples include HunyuanVideo and pixel MeanFlow, where Muon accelerates training relative to AdamW.Source: §2.3 · p. 4

3 Methodology

The methodology pairs flow-matching training with matrix-wise momentum orthogonalization, then introduces chunking to prevent fused DiT parameter blocks from sharing an orthogonalization preconditioner.

3.1 Preliminaries on Diffusion Models

  • Flow Matching trains a network to regress the constant-velocity target z−x along a linear interpolation from data to Gaussian noise.The objective is squared error between the predicted vector field and z−x.Source: §3.1 · p. 5

3.2 The Muon Optimizer

  • Muon forms matrix-wise update directions by orthogonalizing momentum, using the polar factor to produce a semi-orthogonal matrix and scaling the parameter update.In practice, the polar factor is approximated with Newton–Schulz iterations rather than an explicit SVD.Source: §3.2 · p. 6

3.3 Subspace Interference in Muon

  • Chunking separates functionally distinct DiT parameter blocks that otherwise share Muon's preconditioner, which can mix their gradient covariance structures and distort descent geometry.The paper identifies AdaLN scale and shift projections as an example of independently initialized parameters that may have substantially different projection structures and gradient statistics.Source: §3.3 · pp. 6–7

3.4 Algorithm Design

  • For DiT-XL, fused matrices are split along their longer dimension into semantically independent submatrices before orthogonalization.The chunked scaling recomputes the factor from per-chunk dimensions so the stacked update preserves its overall Frobenius norm; optional √N_chunk learning-rate scaling can accelerate early convergence without compromising final performance.Source: §3.4 · p. 7, Learning-rate adjustment under chunking · pp. 8–9
  • CMuon applies momentum to gradients, splits selected fused matrices into chunks, orthogonalizes each chunk independently with Newton–Schulz iterations, and applies scaled updates.The targeted matrices are AdaLN modulation, attention QKV projections, and FFN gate+up projections.Source: CMuon update · p. 7, §3.4 · p. 7

4 Experiments

Experiments on ImageNet-1K at 256 × 256 compare AdamW, Muon, and CMuon across DiT scales. CMuon sustains later convergence and reaches a lower FID than AdamW with half the training epochs; the section also reports qualitative comparisons and chunking ablations.

4.1 Experimental Setup

4.2 Performance of CMuon in DiT Training

  • 1.18 FID in 200 epochs for CMuon outperforms AdamW’s 1.21 FID at 400 epochs.Across settings, vanilla Muon improves faster early but plateaus later, while CMuon sustains convergence.Source: §4.2 · pp. 10–11, Figure 2 · p. 10, Table 2 · p. 11
  • With the same DiT-XL setup, 200-epoch budget, and 30 NFE, CMuon samples are cleaner and more coherent than vanilla Muon samples.CMuon better captures fine details and mitigates texture artifacts and structural distortions.Source: Figure 3 · p. 11, §4.2 · p. 11
  • Chunking FFN, QKV, and AdaLN together improves FID at 80 and 200 epochs more than individual chunked-block configurations.The ablation evaluates 130M DiT-B models.Source: Table 3 · p. 12

5 Further Ablation Studies

The ablations find that joint chunking of FFN, QKV, and AdaLN is more effective than chunking any one block alone, and that CMuon outperforms Muon across tested learning-rate scaling strategies. They also distinguish early-stage gains from learning-rate rescaling and later performance gains from chunking, while showing CMuon remains effective under searched learning rates.

5.1 Ablation on Blocks for Chunking

  • At 200 epochs, joint FFN, QKV, and AdaLN chunking lowers FID from 3.31 to 3.02, while single-block chunking reduces final FID by at most 0.1.The joint configuration also shows a clear gain at 80 epochs.Source: §5.1 · p. 12

5.2 Effect of Learning Rate Scaling Strategies

  • CMuon achieves lower FID than Muon under Vanilla, Keller–Jordan, and MoonLight scaling; MoonLight performs best among these strategies for both optimizers.The authors use MoonLight as the default because it aligns better with the AdamW-tuned learning rate used in the comparison.Source: §5.2 · p. 13, Table 4 · p. 13

5.3 Ablation Study on Learning Rate Rescaling

  • At 40 epochs, the two rescaled configurations have the lowest FID; at 80 epochs, CMuon configurations have the best FID.The authors report complementary effects: rescaling primarily accelerates early training, while chunking primarily improves later performance.Source: §5.3 · pp. 13–14

5.4 Ablation Study on Learning Rate

  • With a learning rate of 2 × 10−4, CMuon reaches FID 1.27 at 140 epochs, closely matching AdamW’s FID 1.26 at 200 epochs with 3 × 10−4.The ablation evaluates learning rates {1, 2, 3} × 10−4 and reports results at 80, 140, and 200 epochs; 2 × 10−4 is near-optimal for both optimizers.Source: §5.4 · pp. 14–15

6 Conclusion

CMuon addresses Muon’s late-stage convergence plateau in Diffusion Transformers by chunking fused matrices into independent components before orthogonalization. On ImageNet-1K 256 × 256, DiT-XL reaches FID 1.18 at 200 epochs, while AdamW requires 400 epochs to reach 1.21.

6 Conclusion

  • At 200 epochs, CMuon achieves FID 1.18 on ImageNet-1K 256 × 256, surpassing AdamW’s FID 1.21 at 400 epochs.The result is reported for DiT-XL.Source: §6 · p. 15
  • CMuon partitions fused DiT matrices, including QKV and AdaLN projections, into independent sub-components before orthogonalization.The conclusion attributes Muon’s late-stage plateau to cross-subspace coupling from jointly orthogonalizing distinct projections.Source: §6 · p. 15
  • The conclusion characterizes this zero-overhead modification as a stable and efficient optimization approach for large-scale DiT training.Source: §6 · p. 15

A.1 Detailed Settings of Training and Sampling

The appendix specifies DiT-B and DiT-XL architectures, sampling settings, and the division of optimization between Muon-style updates and AdamW.

A.1 Detailed Settings of Training and Sampling

  • DiT-B has 12 layers, hidden size 768, 12 attention heads, and head dimension 64; DiT-XL has 28 layers, hidden size 1152, 16 heads, and head dimension 72.Both use an MLP ratio of 4.0 and patch size 2; DiT-XL scales primarily through depth and width.Source: §A.1 · p. 18
  • Sampling uses 30 steps for VA-VAE and 40 steps for SD-VAE and DC-AE, with training objectives and sampling hyperparameters following UCGM.Source: §A.1 · p. 18

A.2 Hyperparameter Settings for Optimization

  • Muon-style updates apply to 2D attention, FFN, and AdaLN weights; 1D parameters, embeddings, and the final layer use AdamW.The ImageNet-1K defaults include batch size 1024, EMA decay 0.9999, bf16 precision, and gradient clipping at maximum norm 1.0.Source: §A.2 · p. 18

A.3 Further Ablation Studies on Chunked Blocks

The further block ablation evaluates chunked-block configurations on a 130M DiT-B model using FID at 80 and 200 epochs. Its reported results identify QKV and FFN projections as consistently improving performance.

A.3 Further Ablation Studies on Chunked Blocks

  • QKV and FFN chunking consistently improve FID at both 80 and 200 epochs in the reported block-configuration results.The study examines chunked-block components in a 130M DiT-B model.Source: §A.3 · pp. 18–19
  • The results are summarized in Table 9, which compares FID across different block configurations.Source: §A.3 · p. 19

A.4 Image Synthesis on ImageNet 512 × 512

On 512×512 ImageNet, CMuon gives the best reported 400-epoch FID among AdamW, Muon, and CMuon. In Pixel-Space DiT, performance decreases with Muon, while CMuon is more stable; architecture-search modifications may favor AdamW in that configuration.

A.4 Image Synthesis on ImageNet 512 × 512

  • The 675M DiT-XL evaluation runs for 400 epochs and tracks FID at 80, 200, and 400 epochs; Table 8 lists default DiT-B and DiT-XL hyperparameters.The model uses a DC-AE encoder; the table describes ImageNet-1K settings.Source: §A.4 · p. 19, Table 8 · p. 19
  • At 400 epochs, CMuon reaches FID 1.60, versus FID 1.67 for AdamW and 1.65 for Muon on 512×512 ImageNet.The evaluation uses a 675M DiT-XL with a DC-AE encoder.Source: §A.4 · p. 19
  • The 130M DiT-B ablation compares QKV, FFN, and AdaLN chunking combinations using FID at 80 and 200 epochs.Table 9 describes the ablation scope but the supplied passage gives no cell values.Source: Table 9 · p. 20

A.5 Further Exploration on Pixel-Space DiT

  • In Pixel-Space DiT, performance decreases with Muon, but CMuon degrades less and retains more stable performance than Muon.Source: §A.5 · p. 20
  • PixelDiT’s AdamW advantage may be configuration-specific: bottleneck and final-layer search modifications could slightly overfit the model to Adam.Source: §A.5 · p. 20

A.6 Training Efficiency with System-Level Optimization

System-level optimizations address the throughput overhead of Muon-style Newton–Schulz iterations. With these optimizations, CMuon with Flash achieves throughput nearly identical to Adam for DiT-B and DiT-XL; the additional analyses describe the orthogonalization procedure.

A.6 Training Efficiency with System-Level Optimization

  • Parallel Newton–Schulz iterations overlap computation with inter-GPU communication, while fused or batched kernels reduce computation and launch overhead.The implementation parallelizes parameter blocks across GPUs and batches chunked parameters on the GPU.Source: §A.6 · p. 21, Optimization 2: Kernel-level… · p. 21
  • With system-level optimizations, CMuon with Flash achieves nearly identical throughput to Adam for both DiT-B and DiT-XL.Without these optimizations, Muon and CMuon have noticeable overhead relative to Adam; Table 12 reports maximum iterations per second on 8× H100 GPUs.Source: Maximum training throughput · pp. 21–22, Table 12 · p. 21

B Additional Analyses and Details

B.1 Implementation of Newton–Schulz Iteration

The implementation uses normalized Newton–Schulz iterations to obtain an approximately orthogonalized gradient matrix. Moonlight Muon scaling redistributes update norm across chunks while matching the approximately 0.2 RMS norm of AdamW updates.

B.1 Implementation of Newton–Schulz Iteration

  • Newton–Schulz implementation normalizes the gradient by its Frobenius norm, transposes tall matrices before iteration, and transposes the output back when needed.Source: §B.1 · p. 22

B.2 Moonlight Muon Scaling for RMS Norm Preservation

  • Moonlight Muon scaling redistributes the update across chunks while preserving its global Frobenius norm.The scaling factor is defined for both the full update and chunked components.Source: §B.2 · p. 23
  • Both AdamW and Moonlight Muon updates have an RMS norm of approximately 0.2 under this scaling.Source: §B.2 · p. 23
Loading 2608.02502v1…