CMuon: Accelerating and Stabilizing Diffusion Trans...

CMuon: Accelerating and Stabilizing Diffusion Trans... | AI Research

Key Takeaways

CMuon is a strategy designed to improve the training efficiency and stability of Diffusion Transformers (DiTs).
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.

Paper AbstractExpand

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.

CMuon is a strategy designed to improve the training efficiency and stability of Diffusion Transformers (DiTs). By modifying how the Momentum Orthogonalization (Muon) optimizer processes weight matrices, the authors address a specific bottleneck where standard training methods cause "subspace interference," leading to poor convergence in the later stages of model training.

The Problem: Subspace Interference

Diffusion Transformers often combine functionally distinct weights—such as those used in AdaLN (Adaptive LayerNorm) and QKV (Query, Key, Value) layers—into single, unified tensors to improve computational efficiency. While the Muon optimizer is effective at accelerating training by orthogonalizing momentum matrices, applying it to these fused tensors creates a problem.
Because Muon treats the entire fused matrix as one unit, it forces a shared preconditioner onto different types of data. This creates implicit coupling between unrelated parameter subspaces, which distorts the update directions and prevents the model from reaching optimal performance.

The Solution: Chunked Momentum Orthogonalization

To resolve this, the researchers introduced Chunked Muon (CMuon). Instead of applying orthogonalization to the entire fused matrix, CMuon partitions these matrices into their original, independent sub-components before performing the calculation.
By treating each functional block separately, the optimizer can apply a preconditioner tailored to the specific statistics of that block. This prevents cross-subspace interference and ensures that the update geometry remains accurate for each part of the network. The method requires only minor code changes to the optimizer and introduces negligible computational overhead.

Performance and Results

In experiments using a 675M-parameter DiT model on the ImageNet 256x256 dataset, CMuon achieved a Fréchet Inception Distance (FID) of 1.18 in 200 epochs. The authors report that this approach provides a 2x training speedup compared to AdamW, while also maintaining a consistent convergence trajectory throughout the entire training process, effectively overcoming the plateaus observed with vanilla Muon.

Franklin Analysis

The effectiveness of CMuon is supported by the authors' comparative analysis of "coupled" versus "decoupled" optimization. By mathematically demonstrating that a shared preconditioner mixes covariance structures across unrelated blocks, the paper provides a clear justification for why chunking improves stability. The results suggest that for architectures relying on fused tensors, the way an optimizer handles parameter grouping is as critical as the optimization algorithm itself.