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

• URL: http://arxiv.org/abs/2608.02502v1
• Date: Mon, 03 Aug 2026 17:04:43 GMT
• Title: CMuon: Accelerating and Stabilizing Diffusion Transformer Training via Chunked Momentum Orthogonalization
• Abstract summary: Diffusion Transformers (DiTs) have achieved state-of-the-art (SOTA) performance in visual generative modeling, yet their training remains computationally prohibitive.<n>Momentum Orthogonalization (Muon) offers a promising alternative to AdamW, its direct application to DiTs yields suboptimal late-stage convergence.
• Score: 10.379982302944201
• License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/
• Abstract: 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.

Related papers

• DMuon: Efficient Distributed Muon Training with Near-Adam Overhead [13.192517400076746]
  We present DMuon, an open-source distributed Muon implementation that integrates into existing training pipelines as a drop-in module.<n>DMuon achieves a 1.48x-3.01x speedup in end-to-end step time and a 6.85x-163.00x speedup in-step time, bringing per-step latency to near-AdamW levels.
  arXiv  Detail & Related papers  (2026-06-25T15:23:03Z)
• N-vium: Mixture-of-Exits Transformer for Accelerated Exact Generation [68.47358899451255]
  N-vium is a mixture-of-exits transformer that partially parallelizes computation across depth on standard hardware.<n>N-vium attaches prediction heads at multiple depths and defines the next-token distribution as a learned mixture over these exits.
  arXiv  Detail & Related papers  (2026-05-13T08:46:17Z)
• TrasMuon: Trust-Region Adaptive Scaling for Orthogonalized Momentum Optimizers [24.534939825452884]
  TrasMuon preserves the near-isometric geometry of Muon while stabilizing magnitudes through (i) global RMS calibration and (ii) energy-based trust-region clipping.<n>We demonstrate that while reintroducing adaptive scaling improves optimization efficiency, it typically exacerbates instability due to high-energy outliers.<n> Empirical experiments on vision and language models demonstrate that TrasMuon converges faster than baselines.
  arXiv  Detail & Related papers  (2026-02-13T22:11:59Z)
• Variance-Adaptive Muon: Accelerating LLM Pretraining with NSR-Modulated and Variance-Scaled Momentum [19.385264518362472]
  Large Language Models (LLMs) achieve competitive performance across diverse natural language processing (NLP) tasks.<n>We propose two variants of Muon, Muon-NSR and Muon-VS, which apply variance-adaptive normalization to momentum.<n>Experiments on GPT-2 and LLaMA pretraining demonstrate that our proposed methods accelerate convergence and consistently achieve lower validation loss than both competitive, well-tuned AdamW and Muon baselines.
  arXiv  Detail & Related papers  (2026-01-21T02:41:56Z)
• TNT: Improving Chunkwise Training for Test-Time Memorization [62.78875147721906]
  Recurrent neural networks (RNNs) with deep test-time memorization modules, such as Titans and TTT, represent a promising, linearly-scaling paradigm distinct from Transformers.<n>We introduce TNT, a novel training paradigm that decouples training efficiency from inference performance through a two-stage process.<n>TNT achieves a substantial acceleration in training speed-up to 17 times faster than the most accurate baseline configuration.
  arXiv  Detail & Related papers  (2025-11-10T17:45:09Z)
• MuonBP: Faster Muon via Block-Periodic Orthogonalization [24.232069944820513]
  We show how to adjust the learning rate from the baseline to MuonBP and give guarantees for this algorithm.<n>When training an 8B model with eight-way tensor tensor and ZeRO statewiseing, MuonBP achieves 8% Muon with no degradation in performance.
  arXiv  Detail & Related papers  (2025-10-19T19:56:05Z)
• Octic Vision Transformers: Quicker ViTs Through Equivariance [29.044546222577804]
  We introduce Octic Vision Transformers (octic ViTs) to capture geometric symmetries.<n>Our octic linear layers achieve 5.33x reductions in FLOPs and up to 8x reductions in memory.<n>We train octic ViTs supervised (DeiT-III) and unsupervised (DINOv2) on ImageNet-1K.
  arXiv  Detail & Related papers  (2025-05-21T12:22:53Z)
• On Disentangled Training for Nonlinear Transform in Learned Image Compression [59.66885464492666]
  Learned image compression (LIC) has demonstrated superior rate-distortion (R-D) performance compared to traditional codecs.<n>Existing LIC methods overlook the slow convergence caused by compacting energy in learning nonlinear transforms.<n>We propose a linear auxiliary transform (AuxT) to disentangle energy compaction in training nonlinear transforms.
  arXiv  Detail & Related papers  (2025-01-23T15:32:06Z)
• Dynamic Diffusion Transformer [67.13876021157887]
  Diffusion Transformer (DiT) has demonstrated superior performance but suffers from substantial computational costs. We propose Dynamic Diffusion Transformer (DyDiT), an architecture that dynamically adjusts its computation along both timestep and spatial dimensions during generation. With 3% additional fine-tuning, our method reduces the FLOPs of DiT-XL by 51%, accelerates generation by 1.73, and achieves a competitive FID score of 2.07 on ImageNet.
  arXiv  Detail & Related papers  (2024-10-04T14:14:28Z)
• Inverse-Free Fast Natural Gradient Descent Method for Deep Learning [52.0693420699086]
  We present a fast natural gradient descent (FNGD) method that only requires inversion during the first epoch. FNGD exhibits similarities to the average sum in first-order methods, leading to the computational complexity of FNGD being comparable to that of first-order methods.
  arXiv  Detail & Related papers  (2024-03-06T05:13:28Z)
• CageViT: Convolutional Activation Guided Efficient Vision Transformer [90.69578999760206]
  This paper presents an efficient vision Transformer, called CageViT, that is guided by convolutional activation to reduce computation. Our CageViT, unlike current Transformers, utilizes a new encoder to handle the rearranged tokens. Experimental results demonstrate that the proposed CageViT outperforms the most recent state-of-the-art backbones by a large margin in terms of efficiency.
  arXiv  Detail & Related papers  (2023-05-17T03:19:18Z)
• Adan: Adaptive Nesterov Momentum Algorithm for Faster Optimizing Deep Models [134.83964935755964]
  In deep learning, different kinds of deep networks typically need different extrapolations, which have to be chosen after multiple trials.<n>To relieve this issue and consistently improve the model training speed deep networks, we propose the ADAtive Nesterov momentum Transformer.
  arXiv  Detail & Related papers  (2022-08-13T16:04:39Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.