Dion2: A Simple Method to Shrink Matrix in Muon

AI summary

Overview

Research area: Machine learning optimization — specifically the Muon optimizer and its computational scalability in large-scale neural network training.

Technical level: Intermediate. The paper assumes familiarity with optimizers (Adam, Muon, momentum), matrix orthonormalization, and distributed training concepts such as data parallelism and sharding, but its core idea is simple enough to describe without deep linear algebra.

Scope: The paper introduces Dion2, a simple row-or-column sampling scheme that shrinks the matrix entering Muon's Newton-Schulz orthonormalization step, and evaluates it on 304M- and 1.01B-parameter language models trained on FineWeb.

What This Paper Is About

Muon is a fast-growing alternative to Adam, but its orthonormalization step has superlinear cost that becomes a compute and communication bottleneck at scale. Prior work (Dion) reduced this cost by orthonormalizing only a low-rank approximation of the momentum matrix, but that required a complicated power-iteration procedure. Dion2 asks whether an even simpler approach — just picking a fraction of rows or columns and orthonormalizing only those — can achieve the same reduction in cost without losing update quality.

Key Contributions

1. A drastically simplified shrinking rule. Instead of building a low-rank approximation with amortized power iteration, Dion2 selects an α-fraction of rows (or columns) of the momentum matrix at each iteration and orthonormalizes only that submatrix, producing a sparse weight update.

2. A selective decay (error-feedback) mechanism adapted to submatrix selection. Only the selected rows or columns of the momentum are decayed by a multiplicative factor μ, a simplification of Dion's error-feedback rule that the authors show is essential to the method working.

3. Empirical validation across two scales. Dion2 matches full Muon quality at 25% selection on a 1.01B-parameter model trained on 100B FineWeb tokens (final validation loss 2.635 vs. Muon's 2.623), and outperforms Dion on a 304M model despite being much simpler.

4. Demonstration that random selection works nearly as well as ℓ1-norm selection, which matters because random selection is easier to synchronize in distributed settings.

Main Findings

• Shrinking to 25% preserves quality at scale: On the 1.01B model / 100B-token run, 0.25-Dion2 achieved final validation loss 2.635 versus 2.623 for full Muon. The gap narrowed compared with the 300M model, which the authors interpret as favorable scaling.

• Random selection matches ℓ1-norm selection: At 300M scale, final losses for ℓ1 vs. random were 2.9154 vs. 2.9148 (α=0.5), 2.9262 vs. 2.9296 (α=0.25), and 2.9452 vs. 2.9469 (α=0.125). At 1B scale, the two selection methods also produced nearly identical final validation losses.

• Dion2 beats Dion despite being simpler: On the 300M model, Dion2 initially lagged behind Dion but eventually caught up and surpassed it, achieving a lower final loss.

• Error feedback is the critical ingredient: Ablating selective decay (decaying the entire momentum matrix, M ← μM, rather than only M[K,:]) led to significantly higher loss, showing the mechanism is necessary rather than incidental.

• Compute savings are measurable: In a micro benchmark using a 4-layer GPT initialized across 4 devices with varying model dimension, optimizer step time dropped clearly as the selection fraction decreased, averaged over the final 20 of 100 training steps. Communication costs were identical between Muon and Dion2 in that setup.

• Additional distributed benefit: Because only the selected submatrix M[K,:] needs to be synchronized across data-parallel workers, Dion2 supports a compressed DP-sync strategy, which matters when data parallelism spans pods over lower-bandwidth data-center networking.

Methodology in Plain English

At each training step, Dion2 accumulates the gradient into a momentum matrix as usual. It then picks a fraction α of that matrix's rows (or columns), runs Newton-Schulz orthonormalization only on that smaller submatrix, shrinks those same selected rows of the momentum by a factor μ, and applies the resulting orthogonalized update only to those rows of the weight. Everything outside the selection is untouched for that step, so the update is sparse.

Two selection strategies were tested: picking rows or columns with the largest ℓ1 norm, and picking them uniformly at random. The submatrix is always taken along the shorter dimension of the momentum matrix. Experiments reused the Dion codebase on the FineWeb dataset with a constant learning rate of 0.02 and a 25% decay period, at 304M parameters (1024 d_model / 16 layers / 8 heads, batch 2.0M tokens, 10K steps, 20B tokens) and 1.01B parameters (2048 d_model / 16 layers / 16 heads, batch 2.0M tokens, 50K steps, 100B tokens). The momentum decay factor μ was set to 0.95.

The intuition for why this works comes from Dion and Trion: Dion showed that error feedback — decaying only the component captured by the approximation — is what makes low-rank shrinking viable, and Trion showed that even a crude approximation works well, suggesting the quality of the approximation matters less than the error feedback. Dion2 takes this to the extreme, treating "select some rows and zero the rest" as the simplest possible approximation.

Why This Matters

Muon has already been adopted in frontier models, but its scalability has depended on a lucky alignment of architecture, parallelism strategy, and framework choices — for example, fine-grained MoE architectures that keep most matrices small, and pipeline-plus-expert parallelism that distributes the work naturally. Dion2 offers a lever that relaxes those constraints, making Muon's orthonormalization cost tunable rather than fixed. If it holds at production scale, it makes Muon a more plausible general replacement for Adam.

Real-world applications:

• Large language model pretraining, where optimizer overhead is a direct fraction of millions of GPU-hours.
• Training on distributed clusters spanning data centers, where shrinking data-parallel synchronization volume across low-bandwidth links is particularly valuable.
• Mixture-of-experts and other sparse architectures, where overall model FLOPs are lower and the Newton-Schulz step therefore becomes a relatively larger cost.
• Optimizer implementations in frameworks such as PyTorch FSDP2, where a preliminary implementation is already released in Microsoft's Dion repository.

Industry relevance: any organization training foundation models at scale pays for orthonormalization in wall-clock time and interconnect bandwidth. A one-line-style sampling change that cuts both, while preserving update quality, is directly deployable in existing Muon-based training stacks.

Future Directions

• Large-scale validation. The authors state that the evaluations are preliminary and that large-scale experiments on production training setups are needed to fully assess the compute and communication benefits.

• Broader selection strategies. Only ℓ1-norm and uniform random selection were tested. The authors note other strategies could perform equally well or better, and that approximating top-ℓ1 selection per datacenter rather than globally would make it parallelizable.

• Whether sparse updates bring additional benefits. The authors find the success of sparse, randomly selected submatrix updates surprising and suggest investigating whether sparsity itself offers advantages beyond cost reduction.

• System-level measurement of communication savings. The micro benchmark only captured compute savings because communication costs were identical in that small setup; the compressed DP-sync advantage remains to be measured empirically.

Target Audience

Machine learning researchers and engineers working on optimizers for large-scale model training, particularly those already using or evaluating Muon and concerned about its orthonormalization overhead. It is also relevant to distributed-training systems engineers interested in reducing synchronization costs across data-parallel replicas, and to practitioners who want a low-complexity alternative to Dion or Trion that can be dropped into an existing training codebase.

Authors’ abstract

The Muon optimizer enjoys strong empirical performance and theoretical grounding. However, the super-linear cost of its orthonormalization step introduces increasing overhead with scale. To alleviate this cost, several works have attempted to reduce the size of the matrix entering the orthonormalization step. We introduce Dion2, a much simpler method for shrinking the matrix involved in Muon's computation compared to prior approaches. At a high level, Dion2 selects a fraction of rows or columns at each iteration and orthonormalizes only those. This sampling procedure makes the update sparse, reducing both computation and communication costs which in turn improves the scalability of Muon.