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

Whether Muon’s step-level advantages on large DiTs translate into distributed training efficiency remains unclear. This paper evaluates Muon across 1.3B–15B DiTs and introduces Periodic Row-wise Muon, which preserves broadly comparable generation quality to vanilla Muon while reducing optimizer time by 46.9–54.3%.

  • Problem

    Evidence on Muon’s scaling in DiTs is limited, and optimizer scalability at billion-parameter scale must account for realized compute and communication costs as well as optimization gains.

    Source: §1 · p. 1

  • Method

    Periodic Row-wise Muon combines periodic NS5 spectral refreshes with RowNorm updates between refreshes and a sharded distributed implementation.

    Source: §3 · p. 3, §6 · p. 10

  • Results

    46.9–54.3% lower optimizer time, while Periodic Row-wise Muon preserves generation quality broadly comparable to vanilla Muon across the evaluated scales.

    Source: §6 · p. 10

  • Takeaways & Limitations

    The results support Periodic Row-wise Muon as an efficiency-preserving alternative to vanilla Muon in the evaluated DiT settings.

    Source: §4.3 · p. 7

    Evaluation covers one DiT family, dataset, resolution, and 32-node H100 configuration, and does not cover other training regimes.

    Source: §6 · p. 10

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

The introduction argues that Muon’s value for large DiTs depends on translating its optimization advantage into lower distributed training cost, not merely fewer steps or lower loss. It establishes Muon’s quality advantage at scale and introduces Periodic Row-wise Muon to reduce the overhead of spectral updates and full-momentum communication.

Introduction

  • Across four DiT scales from 1.3B to 15B parameters, Muon achieves lower validation loss and improves best observed FD-DINO over AdamW by 12.9–19.1%.The results establish Muon’s algorithmic scalability while exposing its additional computation and communication as a per-step cost.Source: §1 · p. 2
  • Within 0.5% on 1.3B–4B, Periodic Row-wise Muon matches vanilla Muon’s best FD-DINO, then improves it by 4.5% at 9B and 2.7% at 15B.It reduces optimizer time by 46.9–54.3%, end-to-end step time by 15.7–24.3%, and logical optimizer-communication volume by 66.7% across all scales.Source: §1 · p. 2
  • Periodic Row-wise Muon applies a full NS5 spectral update once every K steps and uses RowNorm (Lau & Su, 2026) on current momentum between refreshes.Spectral steps couple rows and singular directions; intervening row-wise updates provide local scale control and maintain stable optimization behavior.Source: §1 · p. 2
  • Its distributed implementation applies RowNorm directly to sharded momentum between refreshes and overlaps bucketed all-gathers with NS computation during refreshes.For normalization rows spanning ranks, only norm statistics are all-reduced; this removes full-momentum materialization on non-refresh steps and reduces exposed refresh communication latency.Source: §1 · p. 2

2 MUON FOR DITS

Muon applies finite Newton–Schulz spectral updates to DiT hidden matrices, while non-matrix parameters remain on AdamW. Its repeated full-matrix computation and distributed momentum handling motivate reducing how often spectral updates are performed.

2.1 SCALING MUON FOR DITS TRAINING

  • Muon applies five Newton–Schulz iterations to hidden weight-matrix momentum, while biases, normalization parameters, and other non-matrix parameters remain on AdamW.The implemented finite map Φ5 produces a globally coupled spectral direction; it is not assumed to equal an exact SVD polar decomposition or produce an exactly orthogonal matrix.Source: §2.1 · pp. 2–3

2.2 COMPUTATION COMPLEXITY

  • For a matrix with r = min(m, n), NS5 costs an additional Θ(JNSr) arithmetic factor per matrix relative to AdamW’s Θ(mn) = Θ(rc).This ratio ignores implementation constants and hardware throughput.Source: §2.2 · p. 3

2.3 DISTRIBUTED COST UNDER LARGE SCALE DISTRIBUTED TRAINING

  • Under row-wise sharding, vanilla Muon all-gathers full momentum to every rank, runs NS5 on the full matrix, and retains each rank’s local output slice (Liu et al., 2025).This entails replicated spectral computation, full-momentum communication, and temporary full-matrix materialization; communication time also depends on factors such as bucket sizes, launch latency, topology, and overlap, motivating less frequent spectral transformations.Source: §2.3 · p. 3

3 PERIODIC ROW-WISE MUON

Periodic Row-wise Muon alternates periodic NS5 spectral updates with RowNorm steps, pairing globally coupled spectral structure with lower-cost row-wise updates and a sharded distributed path.

3 PERIODIC ROW-WISE MUON

  • Periodic Row-wise Muon alternates periodic NS5 spectral refreshes with RowNorm updates from current momentum, reducing communication and full-matrix materialization between refreshes.The method treats the branches as updates induced by distinct constraint geometries, not as approximations of one another.Source: §3 · pp. 3–4

3.1 COMPLEMENTARY SPECTRAL AND ROW-WISE GEOMETRIES

  • Spectral updates impose a globally coupled spectral constraint, whereas RowNorm independently preserves each row’s direction under an ℓ2,∞ constraint; periodic refreshes restore global structure.The stability motivation assumes matrices remain away from rank degeneracy, where moderate momentum changes imply moderate changes in the ideal polar direction.Source: §3.1 · p. 4

3.2 PERIODIC ROW-WISE MUON

  • K sets refresh frequency, while fixed γ scales the RowNorm branch relative to the spectral branch because equal Frobenius norms do not ensure equal stable step amplitudes.The global learning-rate schedule remains that of vanilla Muon, with γ maintaining a constant relative scale between branches.Source: §3.2 · pp. 4–5

3.3 DISTRIBUTED EXECUTION

  • Distributed execution overlaps bucketed momentum all-gathers with NS5 refreshes and computes non-refresh RowNorm on shards, using only n-scalar all-reduces for tall matrices.NS5 costs Θ(5r^2c) versus Θ(rc) for RowNorm; for K = 3, optimizer-specific payload approaches one-third of vanilla Muon when norm-statistics payload is small.Source: §3.3 · p. 5

4 EXPERIMENTS

Across DiTs from 1.3B to 15B parameters, Muon improves optimization and generative quality over AdamW, while its per-step overhead limits wall-clock gains. Periodic Row-wise Muon preserves competitive generation quality while reducing training time and communication costs.

4.1 EXPERIMENTAL SETUP

  • The experiments train 1.3B, 4B, 9B, and 15B MMDiTs from scratch for 60,000 steps on 256 H100 GPUs, using the same global batch size of 4,096 across optimizers.Models use 512 × 512 resolution; evaluation uses 50,000 fixed test prompts and a shared classifier-free guidance scale of 5.0.Source: §4.1 · pp. 5–6

4.2 SCALING MUON FOR DIFFUSION TRANSFORMERS

  • Muon achieves lower validation loss than AdamW across all four scales, and its best observed FD-DINO improvement over AdamW is 12.9–19.1%.Despite its per-step optimization advantage, higher compute and communication costs shift Muon’s loss trajectory later in normalized wall-clock time; AdamW reaches intermediate loss levels earlier across all four scales.Source: §4.2 · pp. 6–7, Figure 1 · p. 6

4.3 PERIODIC ROW-WISE MUON GENERATION QUALITY AND TRAINING EFFICIENCY

  • Periodic Row-wise Muon’s best FD-DINO remains within 0.5% of vanilla Muon at 1.3B–4B, and improves it by 4.5% at 9B and 2.7% at 15B.At 9B, its GenEval2 AM rises from 51.77 to 57.33 and GM from 11.35 to 15.97; at 15B, its primary fidelity metrics remain close while Coverage and Density improve.Source: Figure 2 · p. 7, §4.3 · p. 7, Table 1 · p. 8
  • 33.7%, 36.1%, 64.8%, and 57.4% less active training time lets Periodic Row-wise Muon reach its best FD-DINO at 1.3B, 4B, 9B, and 15B, respectively.The largest active-time gains occur at 9B and 15B.Source: Figure 3 · p. 7, §4.3 · p. 7

4.4 SYSTEM EFFICIENCY

  • Periodic Row-wise Muon reduces optimizer time by 46.9–54.3% and end-to-end step time by 15.7–24.3% versus vanilla Muon across four scales.At 15B, optimizer time falls by 54.3%, with a 23.4% reduction in total step time.Source: Table 2 · p. 8, §4.4 · p. 8
  • With K = 3, Periodic Row-wise Muon cuts logical communication volume by 66.7% at every scale, replacing two of every three NS5 phases with sharded RowNorm updates.Bucketed all-gather overlaps refresh communication with NS5 computation, while norm-statistics communication overlaps local computation.Source: §4.4 · p. 8

4.5 ABLATION STUDIES

  • On 1.3B, the complete Periodic Row-wise Muon configuration with K = 3 and γ = 0.15 provides the best generation quality among the ablated variants.In the 15B system ablation, sharded RowNorm cuts communication volume by 66.7%; adding bucketed all-gather and overlap reduces optimizer and end-to-end step time by 35.6% and 12.0%.Source: §4.5 · p. 8

5 RELATED WORK

Prior work motivates evaluating matrix-aware optimization for DiTs across optimization, compute, quality, and distributed efficiency; this work complements methods that speed up or reduce Muon’s spectral transformations.

5 RELATED WORK

  • Unlike AdamW’s element-wise updates, Shampoo and SOAP use structured preconditioners, while Muon applies a finite Newton–Schulz transformation to momentum matrices for global spectral structure (Loshchilov & Hutter, 2017; Gupta et al., 2018).Source: §5 · pp. 8–9
  • Prior Muon work demonstrates scaling to large language models, while DiT studies establish predictable scaling and examine matrix-aware optimization for diffusion training (Liu et al., 2025; Shah et al., 2025; Esser et al., 2024; Yin et al., 2025; Schaipp, 2025; Chen et al., 2026a).These directions motivate evaluating optimizer scalability through optimization progress, algorithmic compute, generative quality, and realized distributed efficiency.Source: §5 · p. 9
  • Periodic Row-wise Muon complements prior cost-reduction methods by refreshing spectra less often and applying RowNorm to current momentum between refreshes (Li et al., 2025; Zhang et al., 2026b).Unlike approaches that post-process every-step orthogonalization, RowNorm serves as a dense substitute between refreshes; its locality avoids full-momentum materialization on non-refresh steps (Lau & Su, 2026).Source: §5 · p. 9
  • Prior efforts reduce Muon’s spectral cost through faster iterations, less frequent or smaller-scope transforms, row-wise post-processing, and distributed runtimes (Amsel et al., 2026; Zhang et al., 2026a; Grishina et al., 2025; Khaled et al., 2026; Dev et al., 2026; Bolatov et al., 2026; Tang et al., 2026; Ahn et al., 2025; Li et al., 2025; Zhang et al., 2026b; Ahn & Xu, 2025; Chen et al., 2026b).These include GPU-friendly polynomial and Gram-matrix methods, block-periodic or reused updates, alternating spectral and sign-based updates, tiled or subsampled transforms, and runtimes that reduce communication, redundant computation, or exposed latency.Source: §5 · p. 9

6 CONCLUSION AND LIMITATIONS · APPENDIX

Muon’s optimization and generative-quality advantages over AdamW persist as DiTs scale from 1.3B to 15B parameters, but its per-step spectral computation and momentum communication create distributed overhead. Periodic Row-wise Muon targets that overhead while retaining broadly comparable generation quality; the evaluation and method retain stated limitations.

6 CONCLUSION AND LIMITATIONS

  • Muon’s optimization and generative-quality advantages over AdamW persist as DiTs scale from 1.3B to 15B parameters, despite overhead from every-step NS5 computation and full momentum communication.Source: §6 · p. 9
  • Across all scales, Periodic Row-wise Muon preserves generation quality broadly comparable to vanilla Muon while reducing optimizer time by 46.9–54.3%.The conclusion reports this comparison across the evaluated model scales.Source: §6 · p. 10
  • Periodic Row-wise Muon combines periodic spectral refreshes, low-cost RowNorm updates, and a sharded distributed implementation to address Muon’s overhead.Source: §6 · p. 10
  • Evaluation is limited to one DiT family, dataset, resolution, and 32-node H100 configuration, and does not cover other training regimes.Source: §6 · p. 10
  • The method uses fixed global K and γ; layer-wise and adaptive schedules remain unexplored.Source: §6 · p. 10
  • Refresh steps still require full momentum communication and materialization, and realized speedup may vary with hardware topology and parallelism strategy.Source: §6 · p. 10

A MATHEMATICAL ANALYSIS … B.1 PER-MATRIX ARITHMETIC

The appendix distinguishes the exact polar geometry from the finite NS5 implementation and gives conditional—not unconditional—optimization guarantees. It also explains how periodic spectral refresh reduces refresh events without proportionally reducing total optimizer cost.

A MATHEMATICAL ANALYSIS; A.1 NOTATION AND UPDATE CONVENTION

  • The appendix distinguishes the exact polar factor from finite NS5 and conditional descent guarantees from unconditional convergence; the effective direction includes all optimizer blocks and weight decay.The RowNorm multiplier scales only RowNorm blocks, not the complete optimizer direction.Source: A MATHEMATICAL ANALYSIS · p. 13, §A.1 · p. 13

A.2 ROWNORM AS A ROW-WISE VARIATIONAL DIRECTION

  • For matrices with nonzero rows, RowNorm uniquely maximizes alignment with its input over matrices whose row norms are at most one.The maximizing rows are the corresponding input rows normalized to unit length.Source: §A.2 · p. 13

A.3 POLAR AND ROW-WISE VARIATIONAL GEOMETRIES

  • Under full row rank, the polar factor uniquely solves a globally coupled spectral-norm-constrained problem and attains the nuclear norm, unlike RowNorm’s independent row-wise problems.Finite NS5 is motivated by the polar geometry but is not identified with the exact polar factor.Source: §A.3 · p. 14

A.4 CONDITIONAL STABILITY OF THE IDEAL POLAR DIRECTION

  • When the smallest singular value stays above σ0, the exact polar map is 1/σ0-Lipschitz, but this bound does not establish stability of finite NS5 between refreshes.The result motivates periodic correction only when matrices remain away from rank degeneracy.Source: §A.4 · pp. 14–15

A.5 CONDITIONAL FINITE-HORIZON DESCENT

  • Under conditional alignment assumptions, the finite-horizon stationary-point bound has a noise-dependent residual and covers deterministic schedules, but does not establish asymptotic convergence or verified training alignment.The result applies to the complete effective direction and does not require tracking vanilla Muon iterates.Source: §A.5 · p. 15

A.6 THE BLOCKWISE ROLE OF THE ROWNORM MULTIPLIER

  • The RowNorm multiplier’s admissible value depends on bounds that make the net descent term positive; the analysis motivates branch-specific calibration but does not select its experimental value.Scaling RowNorm blocks does not multiply the complete direction’s alignment and second-moment constants by γ and γ2, respectively.Source: §A.6 · p. 16

B COMPLEXITY AND DISTRIBUTED EXECUTION; B.1 PER-MATRIX ARITHMETIC

  • The period removes exactly 1−1/K of spectral refresh events, but not that share of total optimizer FLOPs or time; RowNorm arithmetic is Θ(rc).Total optimizer cost also includes momentum maintenance, non-Muon updates, and other work performed every step.Source: §B.1 · p. 17

B.2 CORRECTNESS OF SHARDED ROWNORM … D.2 MODEL CONFIGURATIONS

The appendices establish exact-arithmetic correctness for sharded RowNorm and describe communication accounting, distributed update procedures, time-to-quality conditions, and the experimental dataset and model setup.

B.2 CORRECTNESS OF SHARDED ROWNORM

  • Sharded RowNorm exactly matches dense RowNorm in real arithmetic for both complete-row and split-row sharding; floating-point differences are limited to reduction-order roundoff.For tall matrices, ranks combine column statistics before applying the shared dense denominator to their local entries.Source: §B.2 · p. 17

B.3 PER-RANK COMMUNICATION VOLUME

  • When reduced statistics are negligible relative to full matrices, period-averaged logical communication payload approaches 1/K of vanilla Muon’s refresh payload, not necessarily its collective or step time.The payload accounting excludes headers, padding, launch latency, and parameter communication already needed for forward and backward propagation.Source: §B.3 · p. 18

B.4 IDEALIZED OVERLAP MODEL

  • Idealized bucket overlap changes exposed latency rather than logical bytes, and actual end-to-end claims rely on profiler windows because traces can include synchronization, launch overhead, contention, and imperfect concurrency.The ideal dependency graph alone does not establish measured end-to-end speedup.Source: §B.4 · p. 18
  • Refresh steps apply NS5, while intervening steps use γ-scaled RowNorm; the distributed path bucket-all-gathers momentum on refreshes and uses sharded RowNorm with all-reduced statistics off-refresh.The implementation overlaps refresh-bucket communication with computation and overlaps off-refresh statistic communication with local RowNorm work.Source: §B.4 · p. 19

B.5 WALL-CLOCK BREAK-EVEN CONDITION

  • Lower per-step systems cost improves time-to-quality only if realized speedup exceeds any increase in steps to the same quality; differing GPU counts require a GPU-hour comparison.The break-even condition compares measured mean step times and the number of steps required by each method to reach the prespecified threshold.Source: §B.5 · p. 19

C ALGORITHMS

  • Section C presents the Periodic Row-wise Muon algorithm and its system-level distributed counterpart as Algorithms 1 and 2.Source: C ALGORITHMS · p. 19

D.1 GPIC DATASET

  • GPIC (Chandrasegaran et al., 2026) supplies training, validation, and held-out test splits of 100M, 200K, and 1M examples, respectively, after filtering and deduplication.Qwen3-VL-4B-Instruct generates captions at multiple detail levels; evaluation uses the held-out test set.Source: §D.1 · p. 20

D.2 MODEL CONFIGURATIONS

  • All four latent-space MMDiT models share components and connectivity, varying only in hidden dimension, depth, and attention-head count; Muon updates Transformer-block matrices while AdamW updates the remaining parameters.Source: §D.2 · p. 20

D.3 OPTIMIZER AND TRAINING HYPERPARAMETERS

The optimizer comparison matches training data and schedules across methods, while allowing wall-clock time and total compute to differ with optimizer costs.

D.3 OPTIMIZER AND TRAINING HYPERPARAMETERS

  • AdamW updates all parameters, while Muon variants use matrix-valued updates for hidden-layer weights and AdamW for remaining parameters; all Muon runs use five Newton–Schulz iterations with fixed coefficients (3.4445, −4.7750, 2.0315).At each model scale, optimizers share the model, training length, global batch size, and learning-rate schedule, and process the same number of examples.Source: §D.3 · p. 20

D.4 GENERATION EVALUATION · E ADDITIONAL EXPERIMENTAL RESULTS · E.1 SELECTION OF K AND γ

Periodic Row-wise Muon uses K = 3 and γ = 0.15 as a quality–cost trade-off, with that configuration selected on a smaller development run and transferred across model scales. The paper evaluates generated images using a fixed protocol and reports additional training, generation, and profiling analyses.

D.4 GENERATION EVALUATION

  • Generation evaluation uses 50,000 fixed GPIC test prompts, 512 × 512 images, classifier-free guidance scale 5.0, and checkpoints every 5,000 steps from 5,000 to 60,000.Unless otherwise specified, main comparisons and ablations use the same generation and evaluation protocol.Source: §D.4 · pp. 20–21
  • FD-DINOv2 (Stein et al., 2023), FID (Heusel et al., 2017), and MMD measure distances between generated and real-image or feature distributions.FD-DINOv2 uses DINOv2 feature space, FID uses Inception feature space, and MMD is non-parametric.Source: §D.4 · pp. 20–21
  • Precision and Density characterize sample fidelity, while Recall and Coverage characterize how well generated samples cover the real-data distribution and thus generation diversity.Source: §D.4 · p. 21
  • HPSv2.1 (Wu et al., 2023) measures text–image alignment and human-preference consistency, while GenEval2 (Kamath et al., 2025) assesses object, attribute, relation, and compositional correctness.Source: §D.4 · p. 21

E ADDITIONAL EXPERIMENTAL RESULTS

  • Additional experiments cover the K and γ selection, complete training-loss trajectories, checkpoint-level generation evaluations, and profiler traces.The selection results identify K = 3 and γ = 0.15 as the default configuration.Source: E ADDITIONAL EXPERIMENTAL RESULTS · p. 21, §E.1 · p. 22

E.1 SELECTION OF K AND γ

  • K = 3 and γ = 0.15 are selected: raising K from 2 to 3 cuts NS5 refreshes 33% for a 2.8% FD-DINO increase, while K = 3 to 4 cuts refreshes 25% but raises FD-DINO 8.2%.For K = 3, γ = 0.15 achieves the lowest late-stage mean FD-DINO among the tested values.Source: Figure 5 · p. 22, §E.1 · p. 22
  • The selected K and γ are transferred unchanged to all 1.3B–15B models, without scale-specific retuning.The paper describes this as avoiding separate hyperparameter tuning at each model scale.Source: §E.1 · p. 22
  • Hyperparameter selection used the 1.3B model for 30k steps, tested K ∈ {2, 3, 4} and γ ∈ {0.10, 0.15, 0.25, 0.35}, and chose by mean FD-DINO at 20k, 25k, and 30k.The run used 2 H100 nodes (16 GPUs total) and local batch size 16.Source: §E.1 · p. 22

E.2 TRAINING-LOSS TRAJECTORIES ACROSS MODEL SCALES · E.3 ACTIVE TIME EFFICIENCY OVER ADAMW

Periodic Row-wise Muon tracks Muon’s stable training-loss trajectory across 1.3B–15B models, while achieving lower FD-DINO than AdamW across comparable active training times and at training’s end.

E.2 TRAINING-LOSS TRAJECTORIES ACROSS MODEL SCALES

  • All three optimizers show stable, decreasing training loss across model scales, with no evidence of divergence.Figure 6 covers the complete 60k-step trajectories for 1.3B, 4B, 9B, and 15B models.Source: §E.2 · p. 22, Figure 6 · p. 23
  • Muon generally achieves lower training loss during the main training regime, while Periodic Row-wise Muon closely tracks its trajectory.The close tracking indicates that replacing two out of every three NS5 updates with the cheaper RowNorm update does not destabilize optimization.Source: §E.2 · p. 22
  • Overall training-loss levels decrease as model size grows from 1.3B to 15B, consistent with the main experiments’ validation-loss scaling behavior.Source: §E.2 · p. 22
  • Because minibatch training loss is substantially stochastic, individual values are not used as proxies for generation quality.Generation performance is evaluated using FD-DINO, FID, and additional generation metrics; Figure 6 instead verifies stability over the complete 60k-step trajectory, beyond sparse generation checkpoints.Source: §E.2 · pp. 22–23
  • Figure 6 provides a higher-resolution view of optimization dynamics than validation-loss comparisons, spanning 60k steps for the 1.3B, 4B, 9B, and 15B models.Source: §E.2 · p. 22, Figure 6 · p. 23

E.3 ACTIVE TIME EFFICIENCY OVER ADAMW

  • Across model sizes, Periodic Row-wise Muon achieves lower FD-DINO than AdamW over most overlapping active training times and at the end of training.This indicates an advantage in both generation quality for a comparable time budget and final generation quality.Source: Figure 7 · p. 23, §E.3 · p. 23

E.4 COMPLETE CHECKPOINT EVALUATION RESULTS · E.5 ADDITIONAL PROFILER TRACES ACROSS MODEL SCALES

Complete checkpoint trajectories show that Muon variants improve on AdamW while quality rankings and metric-optimal checkpoints vary over training. Across 1.3B–9B, profiler traces show lower three-step normalized time for Periodic Row-wise Muon, with the same execution pattern reported at 15B.

E.4 COMPLETE CHECKPOINT EVALUATION RESULTS

  • Tables 7–10 provide checkpoint evaluations every 5k steps from 5k to 60k across four model scales, covering fidelity, diversity, alignment, and compositionality metrics.The reported profile includes FD-DINO, FID, MMD-DINO, Precision, Recall, Coverage, Density, HPSv2, and arithmetic and geometric GenEval2 means.Source: §E.4 · p. 23, Table 7 · p. 25, Table 8 · p. 26
  • Across training trajectories, both Muon variants substantially improve over AdamW at multiple stages, while their relative ordering varies by checkpoint.The checkpoint optimizing one generation metric does not necessarily optimize others, so the paper reports the full final-checkpoint profile and best-observed FD-DINO separately rather than combining metrics from different checkpoints.Source: §E.4 · pp. 23–24
  • 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.Source: §E.4 · p. 24

E.5 ADDITIONAL PROFILER TRACES ACROSS MODEL SCALES

  • Across all three profiled scales, vanilla Muon runs a full NS5 update each step, whereas Periodic Row-wise Muon uses one refresh per K = 3 steps and RowNorm on the intervening steps.Refresh steps overlap bucketed all-gather for upcoming matrices with NS5 computation on available matrices; RowNorm steps overlap norm-statistic communication with local computation.Source: §E.5 · p. 24
  • Over three consecutive steps, normalized time falls from 5.59 to 4.20 (1.3B), 7.26 to 5.68 (4B), and 12.47 to 10.26 (9B) with Periodic Row-wise Muon versus vanilla Muon.The 15B trace shows the same qualitative behavior; the authors attribute the cross-scale gains to periodic spectral refresh, sharded RowNorm, and communication–computation overlap.Source: §E.5 · p. 24, Figure 8 · p. 29
  • Figure 8 extends the 15B systems analysis with 1.3B, 4B, and 9B profiler traces, using the same global time normalization and separating optimizer-specific CPU phases, GPU computation, communication, and non-optimizer work.The normalization defines the mean step time of the 1.3B AdamW run as one time unit.Source: §E.5 · p. 24, Figure 8 · p. 29

Related papers

Loading 2608.20818v1…