Dion: Distributed Orthonormalized Updates

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

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

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

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

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.Source: §5.1 · p. 8
• 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.Source: §5.1 · pp. 8–9
• 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.Source: §5.1 · p. 9

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.Source: §5.2 · p. 9

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.Source: §5.3 · p. 9

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.Source: §6 · p. 10

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.Source: §6.1 · p. 10

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.Source: §6.2 · p. 10

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.Source: §6.3 · pp. 10–11

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.Source: §6.4 · p. 11, Figure 4 · p. 11

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.Source: §6.5 · p. 11

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.Source: §7 · p. 11
• 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.Source: §7 · p. 11

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.Source: §7 · p. 11

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.Source: A Accelerating Power Iteration with… · p. 17
• 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.Source: A Accelerating Power Iteration with… · p. 17

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.Source: A Accelerating Power Iteration with… · p. 17
• 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.Source: A Accelerating Power Iteration with… · p. 17, §A.1 · p. 17

Implementation

• Less than 1% of CQR calls fail, making fallback to standard orthogonalization negligible in overhead.The fallback handles occasional CQR failures during training.Source: §A.2 · p. 18

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.Source: §B.1 · p. 18

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.Source: §B.2 · p. 19
• 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.Source: §B.2 · p. 19

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.Source: §B.3 · p. 20
• 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.Source: §B.3 · p. 20

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.Source: §B.4 · pp. 20–21
• 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.Source: §B.4 · p. 21

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.Source: §B.5 · p. 21

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.Source: §B.6 · p. 22
• 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.Source: §B.6 · p. 22

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.Source: §B.7 · p. 23
• 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.Source: §B.7 · p. 24
• 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.Source: §B.7 · p. 25

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.Source: §C.2 · p. 26

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.Source: §C.2 · p. 26
• 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.Source: §C.2 · p. 26

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.Source: §C.2 · p. 26

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.Source: D Update Rules for Non-Matrix Parameters · pp. 26–27, §D.1 · p. 27

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.Source: §D.2 · p. 27

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.Source: §D.2 · pp. 27–28

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.Source: §D.2 · p. 28

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.Source: §D.3 · p. 28

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.Source: §E.1 · p. 28
• 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.Source: §E.2 · p. 28