PaperlayerView exact revision on arXiv

Source-linked AI summary

Scaling Muon for Diffusion Transformers

Chenghao Li, Xiao Han, Xinxin Huang, Wei Liu, Boyang Li, Bing Xiao, Heran Zhang, Juanma Perez Rua, Ke Xu, Kangning Liu, Linjun Kuang, Na Li, Tan Wang, Tian Xie, Wei Peng, Yang Pei, Yifan Xu, Yuanhao Zhai, Yuwei Lin, Zhe Wang, Zihao He, Daniel Li, Junbiao Tang, Ziyang Jiang, Dake Chen

arXiv:2608.20818v1cs.LGcs.AIcs.CV

TL;DR

Large-scale DiT training leaves unclear whether Muon’s step-level advantages translate into efficient distributed training because NS5 and full-momentum handling add overhead. The paper introduces Periodic Row-wise Muon and a sharded implementation that reduce this overhead while retaining comparable generation quality. Across scales, Muon improves best observed FD-DINO over AdamW by 12.9–19.1%, while the periodic variant reduces optimizer time by 46.9–54.3%.

  • Problem

    Muon’s scaling behavior and end-to-end efficiency on large DiTs remain unclear because optimizer evaluation must jointly consider optimization, computation, and communication costs.

  • Method

    Periodic Row-wise Muon performs NS5 spectral refreshes every K steps, uses RowNorm between refreshes, and applies sharded communication and computation overlap.

  • Results

    12.9–19.1%: Muon improves best observed FD-DINO over AdamW across 1.3B–15B DiTs, while Periodic Row-wise Muon reduces optimizer time by 46.9–54.3%.

  • Takeaways & Limitations

    Periodic Row-wise Muon preserves generation quality broadly comparable to vanilla Muon while reducing the distributed overhead of spectral optimization.

  • Takeaways & Limitations

    Evaluation is limited to one DiT family, dataset, resolution, and 32-node H100 configuration, with fixed global K and γ.

Abstract

from arXiv · show

The matrix-aware optimizer Muon improves large model training by balancing updates across singular directions, yet its scaling behavior and end-to-end efficiency on large Diffusion Transformers (DiTs) remain unclear. We first establish Muon's scaling behavior on DiTs from 1.3B to 15B parameters, showing that its optimization and generative quality advantages over AdamW persist across model scales. However, at scale, the 5-step Newton--Schulz iteration (NS5) performed at every optimization step, together with full-momentum materialization, introduces substantial computation and communication overhead that can offset Muon's step-efficiency advantage. We introduce \emph{Periodic Row-wise Muon}, which performs a full NS5 spectral update once every \(K\) steps and applies a low compute and communication cost row-wise constrained update based on the current momentum at the remaining steps. We further co-design a distributed implementation that operates directly on sharded momentum during non-refresh steps and accelerates spectral refreshes through bucketed all-gather and communication--computation overlap. Across all scales, Muon improves the best observed generative quality over AdamW by 12.9--19.1\%. Compared with vanilla Muon, Periodic Row-wise Muon remains within 0.5\% in best generative quality on the 1.3B--4B models and improves it by 4.5\% at 9B. It reduces optimizer time by 46.9--54.3\%, end-to-end step time by 15.7--24.3\%, and logical communication volume by 66.7\%, while reaching its respective best generative quality with 33.7--64.8\% less active training time. These results show that Periodic Row-wise Muon preserves Muon's generative quality advantage while translating it into end-to-end training efficiency for large DiTs.

1 INTRODUCTION

Muon retains optimization and generative-quality advantages over AdamW as DiTs scale, but its distributed overhead motivates a joint evaluation of optimization progress and realized training cost.

  • At billion-parameter scale, optimizer evaluation must include GPU-hours, wall-clock time, and distributed compute and communication costs, not only optimization steps.
  • 12.9–19.1%: Muon improves the best observed FD-DINO over AdamW across DiTs from 1.3B to 15B parameters.
  • Muon achieves lower validation loss than AdamW across all four evaluated model scales.
  • Every-step NS5 computation and full-momentum communication make vanilla Muon substantially more expensive per step.

2 MUON FOR DITS

Muon applies finite spectral transformations to matrix-valued momentum, yielding globally coupled updates whose arithmetic and sharded execution costs exceed those of AdamW.

  • Muon applies five Newton–Schulz iterations to matrix momentum, while AdamW continues updating non-matrix parameters.
  • NS5 reshapes momentum singular values through a finite spectral map inspired by the matrix polar factor, rather than computing an exact SVD polar decomposition.
  • Θ(JNSr): NS5 has this arithmetic cost ratio relative to AdamW for one matrix, ignoring implementation constants and hardware throughput.
  • Under row-sharded execution, Newton–Schulz couples rows and singular directions, so Muon cannot generally compute spectral updates from local momentum shards alone.
  • Full-momentum communication, synchronization, temporary materialization, and repeated spectral computation can offset Muon’s optimization advantage.

3 PERIODIC ROW-WISE MUON

Periodic Row-wise Muon alternates periodic global NS5 refreshes with local RowNorm updates, using distributed communication strategies to reduce average computation and optimizer payload.

  • 3 PERIODIC ROW-WISE MUON: Periodic Row-wise Muon performs an NS5 spectral update periodically and applies RowNorm to current momentum on remaining steps.
  • 3.1 COMPLEMENTARY SPECTRAL AND ROW-WISE GEOMETRIES: RowNorm normalizes each row independently, unlike the globally coupled spectral constraint used by the ideal polar direction.
  • 3.1 COMPLEMENTARY SPECTRAL AND ROW-WISE GEOMETRIES: Moderate momentum changes imply moderate changes in the ideal polar direction away from rank degeneracy, motivating periodic spectral correction.
  • 3.2 PERIODIC ROW-WISE MUON: The period K controls spectral-refresh frequency, while γ calibrates the RowNorm branch’s effective step size relative to the spectral branch.
  • 3.3 DISTRIBUTED EXECUTION: Refresh steps materialize full momentum for NS5, while bucketed all-gather overlaps communication with computation to reduce exposed latency.
  • 3.3 DISTRIBUTED EXECUTION: For K = 3, optimizer-specific logical payload approaches one-third of vanilla Muon when the norm-statistics payload is small.

4 EXPERIMENTS

Experiments compare AdamW, vanilla Muon, and Periodic Row-wise Muon on DiTs from approximately 1.3B to 15B parameters. Muon preserves its optimization and generation-quality advantages, while Periodic Row-wise Muon reduces optimizer, step-time, and communication overhead.

  • Scaling Muon: Muon achieves lower validation loss than AdamW across all four model scales, and its optimization advantage remains stable as DiTs grow from 1.3B to 15B parameters.
  • Scaling Muon: 12.9–19.1%: Muon improves best-observed FD-DINO over AdamW across model scales.
  • Scaling Muon: Muon’s lower validation loss per step is partly or fully offset in wall-clock time by higher per-step computation and communication costs.
  • Periodic Row-wise Muon: Within 0.5%: Periodic Row-wise Muon remains near vanilla Muon in best-observed FD-DINO at 1.3B and 4B, while improving it by approximately 4.5% at 9B and 2.7% at 15B.
  • Periodic Row-wise Muon: 33.7–64.8% less active training time: Periodic Row-wise Muon reaches its respective best generative quality faster across the 1.3B, 4B, 9B, and 15B models.
  • System Efficiency: 46.9–54.3%: Periodic Row-wise Muon reduces optimizer time, while end-to-end step time falls by 15.7–24.3% relative to vanilla Muon.
  • System Efficiency: 66.7%: Logical communication volume decreases at every scale with K = 3, while sharded RowNorm adds approximately 0.7 small all-reduces per step.

5 RELATED WORK

Matrix-aware optimizers exploit parameter-matrix structure, while recent work targets Muon’s spectral-transformation cost through faster or less frequent updates.

  • Shampoo and SOAP construct structured preconditioners, whereas Muon uses a finite Newton–Schulz transformation.
  • Muon applies a finite Newton–Schulz transformation to momentum matrices, producing updates with global spectral structure.
  • Recent approaches reduce spectral-transformation cost through faster polynomial or Gram-matrix iterations.
  • Other approaches reduce transformation frequency, scope, or input size using periodic, alternating, tiled, or row-wise updates.

6 CONCLUSION AND LIMITATIONS

Muon retains its optimization and generative-quality advantages over AdamW as DiTs scale, while Periodic Row-wise Muon reduces the overhead of full spectral updates and communication. Its evaluation and scheduling choices leave several scope boundaries and open implementation questions.

  • Muon’s optimization and generative-quality advantages over AdamW persist from 1.3B to 15B parameters.
  • Periodic Row-wise Muon combines periodic spectral refreshes, RowNorm updates, and sharded execution to address Muon’s distributed overhead.
  • 46.9–54.3% lower optimizer time is achieved while preserving generation quality broadly comparable to vanilla Muon.
  • The evaluation covers one DiT family, dataset, resolution, and 32-node H100 configuration, excluding other training regimes.
  • Fixed global K and γ leave layer-wise and adaptive schedules unexplored, while refresh steps retain full-momentum communication and materialization.

A MATHEMATICAL ANALYSIS

The analysis distinguishes RowNorm’s independent row-wise geometry from the globally coupled spectral geometry underlying the polar direction and finite NS5. It provides conditional stability and finite-horizon descent results while explicitly separating these guarantees from exact polar-factor behavior and unconditional convergence.

  • Row-wise geometry: RowNorm uniquely maximizes the matrix inner product under independent row-norm constraints when every row is nonzero.Its solution is obtained by normalizing each row of the input independently.
  • Spectral geometry: The exact row-polar factor uniquely solves the spectral-norm constrained variational problem for full-row-rank matrices.The optimal value is the nuclear norm of the input matrix.
  • Relation between geometries: The polar direction is globally coupled through spectral-norm geometry, whereas RowNorm solves independent row-wise problems; finite NS5 is motivated by, but not identical to, the exact polar factor.The implementation uses the finite Newton–Schulz map rather than assuming convergence to an exact SVD-based polar decomposition.
  • Conditional stability: A 1/σ0-Lipschitz bound holds for the polar map on regions where the minimum singular value remains at least σ0 > 0.This conditional result motivates periodic spectral correction when momentum matrices remain away from rank degeneracy.
  • Optimization feasibility: Periodic switching with one NS5 step followed by K −1 RowNorm steps is covered by the finite-horizon descent result under conditional alignment and bounded-moment assumptions.The bound is finite-horizon and noise-dependent rather than an unconditional asymptotic convergence theorem.
  • Optimization feasibility: The RowNorm multiplier γ can be calibrated branch-specifically, but the theory does not determine its experimental numerical value.The admissible range depends on measured or assumed bounds making the net descent term positive; the experimental value comes from ablation.
  • Arithmetic analysis: Periodic scheduling removes exactly the fraction 1 −1/K of spectral refresh events, but not the same fraction of total optimizer FLOPs or wall-clock time.Fixed per-step work and period-averaged matrix-processing costs remain in the total cost model.

B.2 CORRECTNESS OF SHARDED ROWNORM

The sharded RowNorm construction preserves dense RowNorm exactly in real arithmetic, while periodic refreshes and bucketed execution reduce full-matrix communication and expose overlap opportunities.

  • Sharding along rows leaves complete oriented rows local when m ≤ n, so local RowNorm matches the corresponding dense slice.
  • For m > n, orientation transposes the matrix, splitting oriented rows across ranks and requiring globally reduced row statistics.
  • Proposition 3 establishes exact-arithmetic equivalence between concatenated local outputs and dense RowNorm, including split-row matrices.
  • Floating-point implementations remain subject only to usual reduction-order roundoff differences.
  • Periodic execution performs NS5 refreshes on gathered momentum but uses sharded RowNorm updates off-refresh, with bucketed asynchronous communication.
  • The ideal overlap model is not sufficient for end-to-end claims, which therefore rely on measured profiler windows spanning complete periods.

D.3 OPTIMIZER AND TRAINING HYPERPARAMETERS

Muon-based runs apply matrix-valued updates only to two-dimensional hidden-layer weights and use an auxiliary AdamW group for all remaining parameters.

  • Muon and Periodic Row-wise Muon apply matrix-valued updates to two-dimensional hidden-layer weights, while remaining parameters use AdamW.

D.4 GENERATION EVALUATION

Generation evaluation uses a fixed, shared protocol across checkpoints and model scales, while the refresh-period study selects K = 3 and γ = 0.15 as a quality–systems-cost trade-off.

  • Generation is evaluated every 5,000 steps from 5,000 to 60,000 using 50,000 fixed GPIC test prompts at 512 × 512 resolution.
  • The protocol reports distributional, fidelity, diversity, text–image alignment, and compositional-semantic metrics including FD-DINOv2, FID, MMD, Precision, Recall, Coverage, Density, HPSv2.1, and GenEval2.
  • All main comparisons and ablations use the same generation and evaluation protocol unless otherwise specified.
  • K controls full NS5 refresh frequency, while γ controls RowNorm update magnitude on non-refresh steps.
  • The selected default is K = 3 and γ = 0.15, transferred directly from the 1.3B development setting to all 1.3B–15B models without scale-specific retuning.
  • 33% fewer NS5 refreshes with K = 3 versus K = 2 increases best mean FD-DINO by only 2.8%, whereas K = 4 adds 25% savings but increases it by 8.2%.

E.2 TRAINING-LOSS TRAJECTORIES ACROSS MODEL SCALES

Across 1.3B, 4B, 9B, and 15B DiTs, Periodic Row-wise Muon remains stable, closely tracks Muon’s training-loss trajectory, and improves generation quality over AdamW at comparable active training time.

  • Figure 6 covers complete 60k-step training-loss trajectories for the 1.3B, 4B, 9B, and 15B models.
  • All three optimizers show stable decreasing training loss across every model scale without evidence of divergence.
  • Muon generally attains lower training loss, while Periodic Row-wise Muon closely tracks it despite replacing two of every three NS5 updates with RowNorm.
  • Training loss is not used as a proxy for generation quality because minibatch loss is substantially stochastic; generation metrics are evaluated separately.
  • Across model sizes, Periodic Row-wise Muon achieves lower FD-DINO through most overlapping active-time ranges and at the end of training than AdamW.

E.4 COMPLETE CHECKPOINT EVALUATION RESULTS

Complete checkpoint evaluations show that Muon variants substantially improve generation quality over AdamW across training, while their relative ordering varies by checkpoint and metric. Profiler traces across model scales show that Periodic Row-wise Muon reduces normalized execution time through periodic refreshes, RowNorm steps, and overlapped communication.

  • Checkpoint evaluation: Muon and Periodic Row-wise Muon achieve substantial improvements over AdamW at multiple checkpoints, with some checkpoint-level variation in their relative ordering.The complete evaluations cover all four model scales and checkpoints from 5k to 60k steps.
  • Checkpoint evaluation: Different generation metrics can peak at different checkpoints, so the evaluation reports final-checkpoint profiles separately from best-observed FD-DINO.This avoids constructing an artificial model by selecting different checkpoints for different metrics.
  • Checkpoint evaluation: Periodic Row-wise Muon typically reaches a generation-quality regime comparable to vanilla Muon with substantially less normalized active training time.Different optimizers enter their best-quality regimes at different points in training.
  • Profiler methodology: Profiler traces visualize non-optimizer phases, optimizer CPU phases, GPU computation, and communication using a common normalization where the 1.3B AdamW mean step time equals one unit.The traces compare vanilla Muon and Periodic Row-wise Muon across the 1.3B, 4B, and 9B models; the 15B result appears in Figure 4.
  • Profiler execution pattern: Periodic Row-wise Muon retains one full NS5 refresh per K = 3 period, uses RowNorm on intervening steps, and overlaps bucketed all-gather with NS5 computation.Norm-statistics communication is likewise overlapped with local computation on RowNorm steps.
  • Profiler results: 5.59 to 4.20, 7.26 to 5.68, and 12.47 to 10.26 are the normalized-time decreases for the 1.3B, 4B, and 9B models, respectively.The 15B trace shows the same qualitative behavior, indicating the systems gains persist across model scales.
Loading 2608.20818v1…