模型訓練與最佳化
CMuon Splits Fused DiT Weights Before Orthogonalization, Training a 675M Model Twice as Fast as AdamW
CMuon finds that applying Muon orthogonalization to fused QKV, AdaLN, and FFN weights as a whole causes interference between functionally distinct subspaces. The study improves late-stage convergence through blockwise momentum orthogonalization, although the claimed 2× speedup has so far been validated only on a single 675M DiT configuration.
Diffusion Transformers (DiTs) are beginning to replace AdamW with Muon, which improves training efficiency by applying Newton–Schulz orthogonalization to matrix-valued momentum. A new study, CMuon, identifies a hidden issue that arises when Muon is applied directly to common model implementations. For computational efficiency, DiTs typically fuse the Q, K, and V projections, as well as AdaLN scale, shift, and gating parameters, into larger weight tensors. Standard Muon applies a single shared orthogonalization operation to the entire tensor. This effectively constructs one preconditioner from the gradient statistics of all the constituent submatrices, coupling updates to parameters that originally serve independent functions.
Before orthogonalization, CMuon splits each fused parameter tensor back into blocks according to its original functional boundaries, then processes each block’s momentum matrix separately. This modification does not change the DiT forward architecture or require an additional representation-alignment loss; the paper says its computational overhead is negligible. The researchers trained a 675M-parameter model on ImageNet at 256×256 resolution with a batch size of 1,024. After 200 epochs, it achieved an FID of 1.18, outperforming the original Muon baseline in the same experiment, and the authors claim roughly 2× faster convergence than AdamW. The key distinction is that the advantage persists into the later stages of training instead of plateauing after a rapid initial improvement.
From an engineering perspective, the work highlights that weight fusion in a framework is not merely a memory-layout detail when matrix-based optimizers are used; it also changes the optimization geometry. Implementations must preserve the semantic boundaries of tensors such as QKV and AdaLN and ensure that the split dimension matches the model version. For now, the evidence is limited to a specific DiT, dataset, and training recipe, and the researchers have not yet released the complete code. FID alone also cannot capture text alignment or high-resolution generation quality. The next step is to determine whether the results can be reproduced in text-to-image and video models with billions of parameters, and whether the communication and kernel overhead introduced by blockwise processing remains negligible.