PaperlayerView exact revision on arXiv

Source-linked AI summary

Dion: Distributed Orthonormalized Updates

Kwangjun Ahn, Byron Xu, Natalie Abreu, Ying Fan, Gagik Magakyan, Pratyusha Sharma, Zheng Zhan, John Langford

arXiv:2504.05295v3cs.LGcs.AImath.OC

TL;DR

Large-scale training makes Muon’s dense orthonormalization costly under weight sharding. Dion replaces it with amortized power iteration and low-rank error-feedback updates, retaining orthonormalized-update benefits while reducing wall-clock time at scale.

  • Problem

    Muon’s dense Newton–Schulz matrix operations conflict with sharded weights, creating compute and communication bottlenecks in large-scale training.

  • Method

    Dion applies amortized power iteration to a momentum buffer, uses low-rank updates with error feedback, and supports sharded synchronization.

  • Results

    Dion preserves orthonormalized updates’ practical benefits and delivers substantially lower wall-clock time at the largest evaluated scales.

  • Takeaways & Limitations

    Dion is presented as a practical optimizer for training foundation models with distributed weight sharding.

  • Takeaways & Limitations

    Experiments focus on GPT-style models, while benefits beyond LLMs are stated as a possibility for architectures with matrix-shaped parameters and dense activation vectors.

Abstract

from arXiv · show

Orthonormalized updates accelerate training, improve stability, and enable robust hyperparameter transfer, but existing methods like Muon rely on dense matrix operations that clash with sharded weights in large-scale LLM training, causing high compute and communication cost. We introduce Dion (Distributed Orthonormalization), a scalable and efficient update rule that replaces Newton-Schulz iteration with amortized power iteration on a momentum buffer, avoiding full-matrix reconstruction and integrating cleanly with weight sharding. The rank-fraction parameter with error feedback enables low-rank updates that balance quality with significant cost savings. On language models from 160M to 3B parameters, Dion retains the benefits of orthonormalized updates, while markedly reducing wall-clock time at scale, making it a practical optimizer for next-generation foundation models. Code is available at: https://github.com/microsoft/dion/

1 Introduction

Dion addresses the compute and communication bottlenecks of orthonormalized updates under large-scale weight sharding. It combines amortized power iteration, low-rank updates, and error feedback while retaining reported optimization benefits and reducing wall-clock time at scale.

  • Millions of GPU-hours motivate optimizer updates that reduce training costs for matrix-valued parameters.
  • Muon’s Newton–Schulz iterations require dense matrix multiplications that conflict with sharded parameters and can become compute- or communication-bound.
  • Dion replaces full-matrix Newton–Schulz processing with amortized power iteration on a slowly evolving momentum buffer, avoiding reconstruction and fitting weight sharding.
  • The rank-fraction parameter enables low-rank updates, while error feedback accumulates approximation errors in the momentum buffer for future steps.
  • At the largest evaluated scales, Dion delivers substantially lower wall-clock time while retaining Muon’s optimization benefits and addressing its system bottlenecks.
  • Dion retains orthonormalized updates’ fast convergence, large-batch tolerance, and hyperparameter-transfer robustness across evaluated language-model settings.
  • Dion supports 1D and 2D weight sharding, and lazy updates with CPU offloading can provide orthonormalized updates with minimal wall-clock overhead.

2 Notation Conventions

The paper represents sharding and parallelism directly in tensor indices, distinguishing model-sharding axes from replicated data parallelism and explicit collective transformations.

  • A device mesh is written Mesh({X : a, Y : b, Z : c}), with symbols denoting named mesh axes.
  • An array annotation such as A[IX, JY] indicates that dimensions I and J are sharded across the corresponding mesh axes.
  • Dion selects unsharded, 1D-sharded, 2D-sharded, or faster Lazy-Dion and CPU-Dion variants according to the parameter layout.
  • Low-rank factorizations introduce a rank axis R with size |R| = r.
  • AllGather removes a shard subscript, ReduceScatter reduces and shards across a logical axis, and AllToAll moves a subscript between axes.
  • The notation uses X and Y for two sharding axes and Z for replicated data parallelism.

3 Dion without Weight Sharding

For unsharded matrix parameters, Dion uses amortized low-rank power iteration and orthonormalization on momentum, with error feedback and compressed synchronization reducing the cost of Muon-like updates.

  • Muon orthonormalizes a momentum matrix through its approximate zeroth power, but Newton–Schulz requires full matrix-matrix products that challenge distributed training.
  • Dion warm-starts power iteration from the previous optimizer step to compute a rank-r approximation, then column-normalizes the update.
  • The update rule accumulates gradients into momentum, computes U and W with PowerIter1, subtracts the approximation for error feedback, and returns an orthonormal update.
  • Error feedback reuses the momentum buffer, requires no additional optimizer-state memory, and lets uncaptured information propagate across iterations.
  • Dion synchronizes low-rank matrices instead of complete gradients, reducing communication volume while producing identical results.
  • For an m × n matrix, compressed data-parallel synchronization reduces communication I/O from O(mn) to O((m + n)r).
  • Dion requires 8mnr + 6.5mr2 + 2.17r3 + O(mn) FLOPs, versus Muon’s five-iteration cost of 20mn2 + 10n3 + O(mn).

4 Dion with Weight Sharding

Dion extends its low-rank orthonormalization and synchronization scheme to 1D- and 2D-sharded weights, preserving locality while making communication depend on rank and matrix dimensions rather than full parameter size.

  • Dion provides 1D- and 2D-sharded implementations whose optimizer states and matrix dimensions remain distributed across device axes.
  • In 1D sharding, local momentum-factor products are reduced and scattered, orthonormalized, then gathered before the remaining update steps.
  • Column normalization requires synchronizing only O(r) scalars, making that operation lightweight.
  • Hybrid sharding combines intra-node FSDP with inter-node data parallelism, where Dion communicates compressed low-rank states instead of full gradients.
  • For 2D sharding, P reductions occur across X and W reductions across Y, preserving locality of both matrix multiplications.
  • Small- and large-batch regimes produce the same U, trading memory for bandwidth.
  • Distributed randomized Cholesky-QR computes an orthonormal basis while preserving Y-sharding and communicating only r × r matrices.
  • 2D-sharded communication scales as O(mr) along X and O(nr) along Y, with distributed orthonormalization adding only O(r2) micro-collectives.

5 Algorithmic Flexibility of Dion

Dion’s amortized power-iteration design supports several efficiency variants for sharded or bandwidth-constrained training. Lazy refreshes, CPU overlap, staged synchronization, and rank adaptation reduce overhead while preserving useful update quality.

  • Faster Dion variants: Lazy-Dion refreshes the right subspace only every k steps, avoiding QR between recomputations while using symmetric power iteration for stability.It reuses cached V to form low-rank updates between refreshes.
  • Faster Dion variants: CPU-Dion asynchronously runs symmetric power iteration on a CPU copy of the momentum matrix while the GPU executes training.Each iteration uses the latest available subspace estimate, which may be lagged.
  • Faster Dion variants: Three GEMMs plus column normalization are the incremental cost between recomputations, while CPU offloading can hide recomputation behind forward and backward execution.Host-memory transfers can also overlap with training execution.
  • Extreme DP communication constraints: Double Dion uses a small DP-synchronized rank r1 and a larger local rank r2, reducing the DP payload to (m+n)r1 under severe bandwidth constraints.Its two stages use distinct error-feedback rules, and delayed sketching can overlap communication with training at a slight convergence cost.
  • Effective Ranks Filtering: Effective rank typically lies between 400–800 for hidden dimension 1024 and varies by layer and training phase, motivating nonstationary rank selection.Filtering matches full-rank accuracy at lower cost, while effective-rank adaptation incurs only a minor gap.

6 Experimental Results

Experiments compare Dion with AdamW and Muon across model scales, batch sizes, learning rates, and faster variants. Dion preserves orthonormalized-update benefits, with low-rank and asynchronous variants maintaining competitive convergence and transfer behavior.

  • Experimental setup: Dion experiments compare against AdamW and Muon using GPT-style decoder-only Transformers trained on FineWeb or FineWeb-Edu.The study follows the modded-nanoGPT setup across multiple scales.
  • Larger Models Tolerate Lower Ranks: At 3B parameters, Dion with rank fractions r/d = 1/2 and 1/4 performs on par with full-rank Muon, while d/16 approaches d/2.Larger models are more tolerant of lower ranks.
  • Large-Batch Training: Dion preserves Muon’s robustness to large batches while remaining competitive with AdamW in the 160M-parameter experiments.The evaluation follows a critical batch size protocol.
  • Hyperparameter Transfer: Optimal learning rates are approximately identical across four model sizes for Muon and Dion at rank fractions 1/4 and 1/16.Each model was trained for a Chinchilla-optimal number of tokens.
  • Faster Dion Variants: Lazy-Dion and CPU-Dion show only minor convergence degradation compared with standard Dion, and Lazy-Dion with k = 32 remains competitive at r/d = 0.5.CPU-Dion is represented with a one-step delay in updating V.
  • Supplementary studies: The study also reports fine-tuning, rank-fraction, DeMo comparison, component ablations, and a preliminary 350M speedrun.These supplementary studies extend evaluation beyond the primary scaling experiments.

7 Related Work and Conclusion

Dion relates to preconditioned, communication-efficient, and low-rank optimizers while positioning error feedback as important for successful low-rank training. Its orthonormalization implementation trades QR’s stability for faster Cholesky QR within a condition-number-dependent scope.

  • Related work: Dion complements gradient sparsification and federated averaging, which can be combined with the optimizer to further reduce communication.Dion is presented as an efficient optimizer in itself.
  • Related work: Dion’s results suggest low-rank training can succeed when paired with an effective error-feedback rule, contrasting with observations from prior work.This connection distinguishes Dion’s low-rank approach from the cited findings.
  • Conclusion and future work: Future work includes quantizing optimizer states, refining error feedback at lower ranks, and testing architectures beyond GPT-style models.The paper specifically identifies matrix-shaped parameters with dense activation vectors as a broader target.
  • Cholesky QR: QR and RCQR provide numerical stability for Dion’s orthonormalization, but QR is computationally expensive.The routine targets accurate recovery of the top-r singular vectors.
  • Cholesky QR: CQR is faster than QR and RCQR but remains numerically stable only below condition number 5×10^3, with higher values risking instability and precision loss.Direct CQR initially produced occasional training instabilities and loss spikes.
  • Condition number dynamics: The condition number of P decreases and stabilizes after a few hundred iterations, enabling a warm-up switch from QR to CQR.Lower rank fractions reduce equilibrium condition numbers at fixed model size, whereas larger models increase them at fixed rank fraction.
  • Condition number dynamics: Larger models can improve CQR stability by using smaller rank fractions, while rare outlier-layer failures are handled by a fallback mechanism.The analysis excludes the top 10% of condition numbers as outlier layers.

A.2 Implementation details

Dion’s implementation and experiments show that adjustable rank fractions reduce communication while preserving competitive optimization and fine-tuning behavior. Its performance remains robust across batch sizes, adaptation strategies, and moderate compression.

  • Implementation: Less than 1% of CQR calls fail, making fallback to standard orthogonalization negligible in overhead.The fallback handles occasional CQR failures during training.
  • Batch-size robustness: Dion maintains a critical batch size at least as large as Muon’s and consistently larger than AdamW’s across tested ranks.Performance degrades gracefully as the update rank decreases.
  • Fine-tuning: Dion reduces fine-tuning wall-clock time by 1.5× for full fine-tuning and 2× for LoRA compared with Adam.Step-count speed-ups are 1.8× and 2.2×, respectively, while Muon and Dion show similar effectiveness.
  • Fine-tuning: Dion achieves performance comparable to full fine-tuning with LoRA rank r = 128 and higher, whereas Adam requires r = 256 and higher.This demonstrates effective lower-rank adaptation under tighter parameterization.
  • Rank fractions: Dion surpasses DeMo once δ > 1/32, exceeds AdamW from δ ≥ 1/8, and slightly outperforms Muon at full rank.The comparison uses update density δ = r/d for Dion and δ = k/s^2 for DeMo.
  • Rank fractions: Dion’s accuracy declines smoothly as r/d decreases, while DeMo peaks at δ = 1/32 and can outperform Dion at sufficiently low δ.Dion achieves its strongest performance at full rank.

B.4 Ablation Studies

The ablations validate Dion’s low-rank approximation and error-feedback design, while additional variants target high-rank computation, extreme communication constraints, and adaptive rank selection. Results show useful quality–cost trade-offs, but some variants remain preliminary or add overhead.

  • B.4 Ablation Studies: Single-step power iteration shows negligible convergence differences from per-step truncated SVD while costing substantially less.The power iteration is initialized from the previous right orthonormal basis.
  • B.4 Ablation Studies: Without error feedback, performance degrades steeply as rank decreases from r = d/2 to r = d/4, while Dion remains stable.Error feedback preserves optimization quality under low-rank approximation.
  • B.5 Results for the 350M Speedrun Configuration: Dion is 14.5% relatively slower in the 350M speedrun because QR is more expensive than Newton–Schulz at that scale.The 75% rank provides modestly better update quality, indicating full rank is not always best.
  • B.6 Dion Variants for Extreme DP Communication Constraints: Double Dion matches the other methods’ 1/128 DP communication volume while achieving the lowest validation loss.Its trade-off is roughly doubled optimizer compute and memory cost, and the authors characterize the result as preliminary.
  • B.6 Dion Variants for Extreme DP Communication Constraints: The delayed Double Dion variant converges more slowly but still outperforms Dion at 1/128 and roughly matches DeMo without delay.A one-step delay severely degrades DeMo’s performance by contrast.
  • B.7 Empricial Studies on the Update Rank: Effective rank is higher in middle layers, rises early in training, and decreases near the end of learning-rate decay.The observed distribution has peaks around ranks 400, 600, and 800.
  • B.7 Empricial Studies on the Update Rank: Full-rank training and rank filtering have similar final performance, while effective-rank and 75%-rank variants can be slightly worse by 0.01 validation loss.Filtering uses a 5e−6 threshold and may slightly improve performance.
  • B.7 Empricial Studies on the Update Rank: Adaptive filtering does not directly improve training speed because filtering and QR recomputation occur after normal QR computation.The result motivates future work on adaptive rank selection.

C.2 Detailed Communication Volumes for 2D-Sharded Dion

The communication analysis decomposes 2D-sharded Dion’s costs by mesh axis and batch regime. Its dominant traffic scales linearly with matrix dimensions and rank, while orthonormalization adds only quadratic-in-rank collectives.

  • Cost model: Dion’s cost model assumes a 1D bidirectional ring, reports payloads in elements, and omits DP synchronization along Z.Bandwidth-bound times ignore latency; modeled AllReduce combines ReduceScatter and AllGather.
  • Communication by axis: Across X, communication is dominated by P/U exchange traffic of O(mr), plus O(|L|r) for column-norm reductions.The latter is |L|r, equal to br in the small-batch regime and abr in the large-batch regime.
  • Communication by axis: Across Y, communication is dominated by W traffic of O(nr), while distributed orthonormalization adds only O(r^2) micro-collectives.The Y-axis collectives include AllReduce(W) and several ReduceScatter or AllGather operations involving rank-sized quantities.
  • Scaling implication: No collective scales with mn; overall communication is linear in mr or nr with quadratic r^2 terms, favorable when r ≪ m,n.This scaling supports low-rank distributed orthonormalization without full-matrix communication.
  • Learning-rate transfer: The section situates these costs within a single-learning-rate design spanning Dion, scalar optimizers, and parameter types.The supplied passages state that natural-norm normalization and dimension-dependent scaling target consistent transfer across model sizes.

D.2 Scaling Factor for LM Head Parameters

The LM head requires a dimension-dependent update scale because its logits depend on the interaction between unembedding vectors and activations. The reported experiments favor the 1/√d_in scaling used with Lion over alternatives.

  • Motivation: For an unembedding vector v and activation h ∈ R^d, the logit is s = v · h, so update magnitude affects logit changes.The analysis assumes i.i.d. zero-mean updates for ∆v when relating update RMS norm to logit changes.
  • Scaling rule: The proposed 1/√d scaling differs from Yang et al.’s 1/d factor, which the paper characterizes as a conservative worst-case bound.The paper reports stable training and best validation loss for 1/√d.
  • Empirical comparison: The unscaled factor 1 causes gradient explosions and instability, whereas 1/√d consistently outperforms both alternatives and a manually tuned Adam baseline.Figure 14 compares the three scale factors with Lion as the scalar optimizer.
  • Scalar optimizer: Dion+Lion outperforms Dion+Adam and enables a shared base learning rate across parameter types without extra tuning.The paper contrasts Lion’s constant RMS update with Adam’s need for a separately tuned scalar-parameter learning rate.

E.1 Results in Section 6.1

The experiments use a fixed learning rate for Muon and Dion across model sizes, alongside fixed training and scalar-optimizer settings. Hyperparameter-transfer results support the learning-rate choice.

  • Training setup: Muon and Dion both use a fixed learning rate of 0.01 across all model sizes.The passage attributes this choice to hyperparameter-transfer results in Figure 3.
  • Training setup: Each model size is trained on its Chinchilla-optimal token count with a 1M-token batch, constant learning rate, no warmup, and 10% linear cooldown.Adam uses learning rate 0.002 with β1 = 0.9 and β2 = 0.95 for non-matrix parameters.
Loading 2504.05295v3…