Dion: The distributed orthonormal update revolution is here
Kwangjun Ahn from Microsoft Research introduces Dion, a next-generation optimizer that improves upon Muon by using amortized power iteration. Dion enables efficient, scalable training for massive models by orthonormalizing a low-rank subspace, reducing compute and communication overhead in distributed settings.
While Adam and its variants have long been the standard for training AI models, the “orthonormal updates revolution,” led by optimizers like Muon, has demonstrated the potential for faster convergence, more stable training, and better performance with large batch sizes. This new class of optimizers has already been used in production-level models such as Kimi-K2 and GLM-4.5.
The Principle of Orthonormal Updates
Unlike a standard Stochastic Gradient Descent (SGD) step, which updates weights Xₜ directly along the negative gradient Gₜ (Xₜ = Xₜ₋₁ - ηGₜ), an orthonormal update first decomposes the gradient matrix Gₜ via Singular Value Decomposition (SVD) into UΣVᵀ. The optimizer then uses only the orthonormal component, UVᵀ, as the update direction.
The theoretical justification for this approach is that an orthonormal update transforms any input activation by an equal amount, which in practice leads to significant performance improvements. To avoid the prohibitive cost of a full SVD at every step, the Muon optimizer implements this orthonormalization using an iterative matrix multiplication method known as the Newton-Schulz iteration.
Muon’s Scalability Wall
Despite its benefits, Muon encounters a significant bottleneck when scaling to dense models with over 100 billion parameters. The core issue lies in the Newton-Schulz iteration, which requires dense matrix multiplications on the full weight matrices. This requirement clashes directly with common distributed training strategies like weight sharding (e.g., FSDP), where each GPU only holds a partial slice of the weights. To perform the full matrix multiplication, the system must either engage in heavy cross-shard communication to reconstruct the full matrix or perform redundant computations, both of which severely limit scalability.
Dion: Orthonormal Updates for Distributed Training
Dion is a next-generation optimizer designed to provide the benefits of orthonormal updates without the scalability limitations of Muon. The central design question it answers is: “Can we design orthonormal updates without full-matrix materialization?”
To achieve this, Dion replaces Muon’s Newton-Schulz iteration with amortized power iteration. This technique is far more compatible with sharded weights, as it does not require the materialization of the full matrix. While preserving the convergence and stability benefits of Muon, Dion introduces several key features that enhance scalability.
Key Features and Innovations
- Low-Rank Fraction as a Scalability Lever: Instead of orthonormalizing the entire gradient, Dion operates on a top-
rsubspace. This “low-rank fraction” acts as a new scalability lever, allowing practitioners to trade off computational and communication costs. A lower rank means a cheaper optimizer step. - Error-Feedback Mechanism: To maintain the quality of the updates at lower ranks, Dion incorporates an error-feedback mechanism, ensuring that performance is not sacrificed for speed.
- Efficiency at Scale: Empirical studies show that as model size increases, the performance gap between high-rank and low-rank Dion narrows. This suggests that for very large models, one can “get away with a smaller-rank fraction,” making Dion increasingly efficient at scale.
- Performance Benchmarks: Microbenchmarks demonstrate that Dion’s time-per-step is significantly lower than Muon’s, especially for large matrices. For a Llama 3 405-billion parameter dense model configuration, Dion is shown to be “a lot more tractable” and feasible, unlike Muon.
- Compatibility and Flexibility: Dion is designed for the realities of modern large-scale training. The open-source implementation includes efficient support for both one-way and two-way weight sharding. The underlying algorithm also allows for greater flexibility, leading to variants like “Lazy-Dion” for further speedups.
- Compressed Communication: Dion also integrates a PowerSGD-style compressed data-parallel synchronization mechanism, which further reduces communication overhead and is particularly beneficial in hybrid sharding strategies.
In summary, Dion makes orthonormal updates practical and efficient for training state-of-the-art, massive-scale AI models. By rethinking the underlying linear algebra, it provides a scalable, communication-efficient solution that is fully compatible with modern distributed training paradigms.