Dion3: Full-Stack Orthogonal Updates

Dion3 is a revision of the Muon optimizer that targets overhead at every level of the stack. The Muon optimizer incurs significant overhead due to its cubic-time Newton-Schulz orthogonalization step, and communication overhead compounds this when weights are sharded. Dion3 presents four main contributions: (1) Gram Newton-Schulz, a mathematically equivalent reformulation of Newton-Schulz that iterates on the small symmetric Gram matrix, cutting FLOP costs; (2) custom GPU kernels for symmetric matrix multiplication written in CuteDSL; (3) a new update rule that subsamples rows or columns of the momentum matrix before orthogonalizing; and (4) megabatched communication to reduce the number of communication rounds. Overall, Dion3 matches or improves on the loss achieved by Muon but reduces optimizer step time by up to 6×.

The paper states: "Dion3 allows practitioners to easily realize the benefits of Muon without paying the high cost of its orthogonalization step—even for large-scale models in highly distributed settings. Our Gram Newton-Schulz algorithm and CuteDSL kernels speed up Muon by 1.5× for dense models and 2× for MoEs, a rare case of free lunch performance. Megabatching has an equally large effect in certain distributed settings. Dion3's fractional update rule (f = 1/4) provides an additional 3.7× speedup."

The Gram Newton-Schulz algorithm works by computing the Gram matrix XX T, using an iterative method to approximate (XX T)(-1/2), and then outputting Q T X. This uses just two rectangular matrix multiplications instead of many, and the paper notes: "When α = 1, they are equal. When α > 1, Gram Newton-Schulz is cheaper, often significantly so. For a typical Muon application (T = 5, α = 4/3), it saves 55% of the FLOPs used by standard Newton-Schulz with symmetric GEMMs, or 68% compared to a typical implementation without symmetric GEMMs."

The symmetric GEMM kernels compute the lower triangle of the output matrix and copy results to the upper triangle, saving about half the floating point operations. The paper reports: For large enough n, our kernels achieve a ∼ 2× speedup over cuBLAS, both with and without an epilogue addition of C.

The Dion3 update rule selects a fraction of rows from the momentum matrix, orthogonalizes only that submatrix, updates only those rows of the weight matrix, and applies error feedback by decaying only the selected rows. The paper explains: "Our update rule introduces a new hyperparameter f ∈ (0, 1] that controls the fraction of rows (or columns) to select. We think of f as a compression factor; f = 1 corresponds to Muon, and decreasing f speeds up the algorithm. We recommend f = 1/4 or f = 1/8."

Megabatching groups all matrices of the same shape into a single batch, reducing communication rounds to O(1). The paper states: "For a given shape, the local momentum shards are packed into one all-to-all, assembled, orthogonalized as a batch, and scattered back together. Transformers contain only a handful of distinct weight shapes, so megabatching reduces the number of communication rounds to O(1), independent of model depth."

Experiments show that Dion3 with f < 1 actually outperforms NorMuon when tuned correctly. The paper notes: "Dion3 with f < 1 is a genuinely new optimizer. In this section, we demonstrate that Dion3 achieves a slightly better loss than the baseline... the lowest loss is achieved at f = 1/8. At larger scales (3B to 14B parameters), Dion3 outperforms NorMuon at every scale, achieving the largest improvement (−0.027) at the largest scale (14B)."

The paper concludes: Dion3 allows practitioners to easily realize the benefits of Muon without paying the high cost of its orthogonalization step—even for large-scale models in highly distributed settings.

Improvements to AI Systems:

1. Optimizer-Level Speedup for Training Infrastructure: Integrate Dion3’s Gram Newton-Schulz algorithm and CuteDSL symmetric GEMM kernels into existing training frameworks (e.g., PyTorch, JAX). This reduces per-step optimizer overhead by up to 6×, enabling faster iteration on large dense models (1.5× speedup) and Mixture-of-Experts models (2× speedup) without altering model architecture or final loss.

2. Scalable Distributed Training with Reduced Communication: Adopt Dion3’s megabatched communication to collapse all-to-all rounds for weight shards into O(1) rounds per step. This allows training of very deep transformers (e.g., 14B+ parameters) across many GPUs/TPUs with significantly lower synchronization latency, making large-scale pretraining more feasible on limited bandwidth clusters.

3. Memory-Efficient Orthogonalization for High-Dimensional Layers: Use Dion3’s fractional update rule (f = 1/4 or 1/8) to orthogonalize only a subset of momentum rows/columns. This cuts compute by up to 3.7× and reduces memory footprint for the orthogonalization step, enabling Muon-style optimization for models with extremely wide layers (e.g., embedding tables, large hidden sizes) that previously exceeded GPU memory.

4. Improved Convergence with New Optimizer Variant: Replace standard Muon or AdamW with Dion3 (f < 1) as a drop-in optimizer. Empirically, it achieves lower loss than NorMuon at scales from 3B to 14B parameters (e.g., −0.027 loss improvement at 14B), meaning better final model quality for the same training compute budget.

5. Automatic Hyperparameter Tuning for Compression Factor: Build a system that dynamically selects f (e.g., 1/4 vs 1/8) based on layer shape and hardware, since the paper shows optimal f varies with scale. This allows AI systems to self-configure for maximum speed without sacrificing quality.

What the Improved AI System Can Do:

• Train a 14B-parameter transformer up to 6× faster per optimizer step while matching or exceeding Muon’s loss, enabling rapid experimentation cycles.

• Run distributed training across hundreds of nodes with near-constant communication rounds, regardless of model depth, making 100B+ parameter training practical on moderate clusters.

• Optimize models with very large hidden dimensions (e.g., 50k+ features) that previously choked on cubic-time orthogonalization, now using fractional updates to fit in memory.

• Achieve state-of-the-art loss on large-scale language modeling benchmarks with the same compute as AdamW, but with better final perplexity.

• Deploy as a plug-and-play optimizer in existing pipelines, requiring only a change of optimizer class and one new hyperparameter (f), with no architectural modifications.