Dion3: Full-Stack Orthogonal Updates

Abstract: The Muon optimizer incurs a significant overhead cost due to its cubic-time Newton-Schulz orthogonalization step. When weights are sharded, communication overhead compounds this computational cost, eroding the benefits of Muon in many settings. We present Dion3, a revision of Muon that targets this overhead at every level of the stack. Our Gram Newton-Schulz algorithm reduces the FLOP cost of orthogonalization, our CuteDSL kernels accelerate it by exploiting symmetry, and our megabatching strategy reduces communication overhead. Moreover, we propose a simple change to the update rule that cuts costs even further: selecting only a fraction of the momentum matrix's rows to orthogonalize at each step. This update rule improves on Dion (another "compressed" version of Muon), in both speed and performance. Overall, Dion3 matches or improves on the loss achieved by Muon but reduces optimizer step time by up to 6x. Dion3 is available via the dion package (https://github.com/microsoft/dion) as a drop-in replacement for Muon.

• The paper introduces a full-stack Muon optimization system that combines Gram Newton–Schulz, symmetry-aware GPU kernels, fractional updates, and megabatched communication to reduce optimizer-step time by up to 6.5×.
• Gram Newton–Schulz reformulates orthogonalization around smaller symmetric Gram matrices, reducing FLOPs by 55% versus symmetric Newton–Schulz and requiring restart strategies to remain stable in half precision.
• Fractional updates improve both efficiency and reported quality, with a 1/4 selection fraction outperforming NorMuon across 3B–14B models and delivering a 0.7-percentage-point accuracy gain at 14B parameters.

“Dion3: Full-stack orthogonal updates” (2608.11612) addresses the principal systems limitation of Muon-family optimizers: the cost of repeatedly orthogonalizing momentum matrices with Newton–Schulz iterations. The paper argues that Muon’s favorable optimization behavior is accompanied by a scaling burden that becomes increasingly significant for large rectangular weight matrices and distributed training. Dion3 responds with a coordinated algorithm–kernel–communication design comprising Gram Newton–Schulz, symmetry-aware GPU kernels, fractional row or column updates, and megabatched communication. The resulting optimizer preserves or improves validation loss relative to Muon and NorMuon while reducing optimizer-step time by as much as $6\times$.

Scaling limitations of Muon

Muon applies momentum followed by an approximate polar decomposition of each matrix-valued update. If a momentum matrix has dimensions $n \times m$, with $n \leq m$, standard Newton–Schulz performs repeated products involving both $n \times m$ and $n \times n$ matrices. With $T$ iterations, its dominant cost under symmetry-aware multiplication is proportional to $T(3\alpha+1)n^3$, where $\alpha=m/n$ is the aspect ratio. The cost is therefore superlinear in the smaller matrix dimension and grows strongly with rectangularity.

This complexity distinguishes Muon from SGD and AdamW, whose optimizer operations are predominantly elementwise and linear in parameter count. The distinction is particularly consequential for MoE architectures: sparse activation reduces forward and backward computation while dense expert matrices can remain large, causing orthogonalization to occupy a larger fraction of the training step. Distributed training introduces a second problem. Under FSDP, a momentum matrix must be assembled across shards before orthogonalization and redistributed afterward. Thus, Muon incurs both cubic local computation and collective communication that ordinary sharded AdamW does not require.

The paper frames Muon’s successful deployment in large systems such as Kimi K2 as dependent on a favorable alignment among architecture, framework state layout, and parallelism strategy. Dion3 instead seeks to make orthogonal optimization less sensitive to those deployment choices.

Gram Newton–Schulz

The central mathematical contribution is Gram Newton–Schulz (GNS), which reformulates the standard Newton–Schulz iteration around the smaller symmetric Gram matrix. For a matrix $X \in \mathbb{R}^{n \times m}$, the polar factor can be expressed as

$\operatorname{polar}(X)=(X X^\top)^{-1/2}X.$

Standard Newton–Schulz iterates directly on $n \times m$0. GNS first forms $n \times m$1, applies an inverse-square-root iteration to the resulting $n \times m$2 matrix, and finally multiplies the approximate inverse square root by $n \times m$3. The reformulation is not merely an approximation: in exact arithmetic, it is algebraically equivalent to standard Newton–Schulz for the same sequence of odd polynomial coefficients.

The paper derives this equivalence by writing each odd polynomial as $n \times m$4. The scalar iteration can then be expressed through a Gram variable $n \times m$5, an auxiliary factor $n \times m$6, and an accumulated multiplier $n \times m$7. Lifting these recurrences to matrices allows the algorithm to avoid constructing the large intermediate matrices generated by standard Newton–Schulz.

The computational consequence is substantial. GNS requires only two large rectangular matrix multiplications, one to form the Gram matrix and one to reconstruct the output. The iterative portion uses small symmetric products. For $n \times m$8 and a typical transformer aspect ratio $n \times m$9, the paper reports a 55% reduction in FLOPs relative to standard Newton–Schulz with symmetric GEMMs, and a 68% reduction relative to a conventional implementation without symmetry exploitation. The asymptotic dependence changes from a product-like $n \leq m$0 structure to an additive $n \leq m$1 structure.

The advantage increases for highly rectangular matrices, including MLP and expert weights. For square matrices, the implementation falls back to standard Newton–Schulz with symmetric kernels because it launches fewer operations and is faster in wall-clock measurements.

Numerical stability and restarting

Mathematical equivalence does not imply numerical equivalence. The naive GNS implementation is unstable in half precision because forming the Gram matrix introduces spurious negative eigenvalues. Although $n \leq m$2 is positive semidefinite in exact arithmetic, rounding error can produce small negative eigenvalues near the numerical nullspace. The inverse-square-root polynomial iteration is well behaved on nonnegative inputs but can amplify negative inputs exponentially.

Eigenvector drift provides a second instability mechanism. In exact arithmetic, the intermediate matrices share eigenvectors with the initial Gram matrix. Finite-precision matrix products perturb these eigenspaces, causing the spectral evolution to deviate from the scalar polynomial analysis. The resulting drift can increase intermediate spectral norms and lead to divergence even when negative eigenvalues are absent.

The proposed remedy is restarting. After a small number of GNS iterations, the algorithm reconstructs an intermediate approximation in the original rectangular space, recomputes its Gram matrix, and begins a new short iteration sequence. For the five-iteration Polar Express configuration used in the paper, restarting after iteration two provides the preferred stability–performance tradeoff. The stabilized implementation also uses float16 rather than bfloat16, distributes scalar terms to avoid numerically unfavorable fused operations, and applies a conservative safety factor to the polynomial coefficients.

The empirical importance of this intervention is clear: naive half-precision GNS causes Llama-430M training to develop loss spikes, infinities, and eventual failure, whereas the restarted version preserves validation behavior relative to standard Newton–Schulz.

Figure 1: Naive half-precision Gram Newton–Schulz destabilizes Llama-430M training.

Figure 2: A restart after two iterations stabilizes Gram Newton–Schulz with Polar Express coefficients.

The stability analysis is one of the paper’s more valuable theoretical components. It identifies a failure mode that is invisible in exact-arithmetic derivations and shows that implementation-level details—precision format, restart location, safety factors, and treatment of diagonal terms—are part of the algorithmic specification rather than incidental engineering choices.

Symmetry-aware CuteDSL kernels

GNS exposes a large number of symmetric matrix operations. Dion3 exploits this structure with custom CuteDSL kernels for symmetric GEMM and fused matrix-quadratic operations on NVIDIA Hopper and Blackwell GPUs. The kernels compute only the lower triangular output, then transpose and copy each off-diagonal tile into the upper triangle. This avoids redundant arithmetic while preserving a dense symmetric output layout.

The implementation uses a triangular tile scheduler and an epilogue that writes both the computed lower-triangular tile and its transposed counterpart. The paper reports approximately $n \leq m$3 speedups over cuBLAS GEMM for sufficiently large symmetric matrices, both with and without fused epilogues. The performance benefit is especially important for GNS because its iterative core consists primarily of symmetric multiplications.

Figure 3: Symmetry-aware CuteDSL kernels compute only one triangular half and approach a $n \leq m$4 speedup over cuBLAS for large matrices.

Across Llama, Qwen, Gemma, and MoE configurations, the combination of GNS and custom kernels accelerates the Newton–Schulz component by approximately $n \leq m$5–$n \leq m$6. The largest gains occur for matrices with high aspect ratios, such as Gemma-1B’s MLP projections. In an estimate based on Kimi K2’s exposed optimizer operations, the kernelized GNS implementation is reported to be $n \leq m$7 faster than standard Newton–Schulz.

The kernel contribution also illustrates the paper’s full-stack thesis. The algorithmic reformulation reduces the number and shape of GEMMs, but the resulting operations must be implemented in a way that realizes their arithmetic structure on modern tensor-core hardware. Conversely, symmetry kernels alone improve standard Newton–Schulz but cannot remove its dependence on rectangular multiplications.

Fractional orthogonal updates

Dion3’s third contribution changes the optimizer trajectory by orthogonalizing only a selected fraction of the momentum matrix. At each step, the optimizer selects the rows or columns with the largest $n \leq m$8 norms, orthogonalizes the resulting submatrix, updates only the selected coordinates, and applies error-feedback decay to the selected momentum components. Unselected components remain in the momentum buffer and can accumulate until they become eligible for future updates.

If the selected fraction is $n \leq m$9, the orthogonalized matrix has approximately $n \times m$0 rows. This reduces local orthogonalization cost by roughly a factor determined by $n \times m$1 for the dominant products and reduces communication volume by approximately $n \times m$2. Selection also increases the effective aspect ratio, making GNS comparatively more advantageous.

The error-feedback mechanism is essential. Without it, rows omitted repeatedly could be permanently under-trained. By leaving the residual momentum component undamped, Dion3 increases the likelihood that neglected rows will be selected in later steps. At $n \times m$3, the method reduces to Muon or NorMuon up to finite-precision and implementation details.

The paper reports a learning-rate transfer rule:

$n \times m$4

because selecting only $n \times m$5 rows reduces the Frobenius norm of an orthogonalized update by approximately $n \times m$6. Consequently, the learning rate should scale approximately as $n \times m$7.

Figure 4: Optimal Dion3 learning rates follow the transfer relation $n \times m$8 across selection fractions.

The result that fractional updates improve optimization quality is contrary to the paper’s initial expectation. On 1B-parameter models trained for 100B tokens, Dion3 variants with $n \times m$9 remain below the fully tuned NorMuon validation-loss curve throughout training and finish approximately $n \times n$0 cross-entropy points lower. The best result in the sweep occurs at $n \times n$1.

Figure 5: Fractional Dion3 updates produce lower validation loss than NorMuon across the reported training trajectories.

At larger scales, Dion3 with $n \times n$2 outperforms NorMuon in validation loss at every tested model size from 3B to 14B parameters. The improvements are $n \times n$3, $n \times n$4, $n \times n$5, and $n \times n$6 cross-entropy points for 3B, 4B, 7B, and 14B models, respectively. Downstream accuracy improves at three of the four scales, including a $n \times n$7 percentage-point gain at 14B, although the 4B result is lower by $n \times n$8 points.

These results suggest that structured update sparsification can act as more than a computational approximation. The selection mechanism changes the temporal allocation of optimization effort across rows, while error feedback preserves residual information. Whether the observed improvement is caused by implicit regularization, altered noise statistics, or a better allocation of update magnitude remains unresolved.

Megabatched distributed communication

Dion3 also addresses the latency structure of distributed orthogonalization. A naive FSDP implementation processes matrices in groups of approximately world_size, requiring a number of all-to-all rounds proportional to the number of matrices divided by the world size. This creates two inefficiencies: every round incurs fixed launch and synchronization latency, and small messages underutilize NVLink or other interconnects.

Megabatching groups all matrices with the same shape into a single collective. The local shards are packed, communicated in one operation, orthogonalized as a batch, and scattered back. Since transformer models contain relatively few distinct matrix shapes, the number of communication rounds becomes effectively independent of model depth.

The measured benefit is workload dependent. For a 1B model on eight FSDP shards, megabatching reduces Muon optimizer-step time from $n \times n$9 ms to $T$0 ms, a 35% reduction. The benefit falls to $T$1–$T$2 for larger or more compute-bound configurations. The paper’s all-to-all microbenchmarks show why: small transfers are latency dominated and achieve poor bandwidth utilization, whereas payloads in the tens of megabytes approach a substantial fraction of peak interconnect bandwidth.

Figure 6: NCCL all-to-all performance is latency dominated for small messages and reaches high bandwidth only after payload coalescing.

The authors further propose compressed data-parallel synchronization in settings where data-parallel communication crosses slower datacenter networks. If row selection can be performed without globally synchronizing the full momentum matrix, only the selected submatrix needs to participate in synchronization. This possibility is particularly relevant for hybrid ICI/DCN training systems.

End-to-end performance

Dion3’s contributions compound across the stack. On large models, kernelized GNS and the symmetry-aware implementation provide at least a $T$3 improvement over standard Muon. Fractional updates with $T$4 and $T$5 provide additional reductions of approximately $T$6 and $T$7, respectively. The resulting total speedups are reported as $T$8 for $T$9 and $T(3\alpha+1)n^3$0 for $T(3\alpha+1)n^3$1 relative to standard Muon in the largest benchmark regimes.

Figure 7: The combined Dion3 stack reduces optimizer-step time by up to approximately $T(3\alpha+1)n^3$2 across single-GPU and four-GPU FSDP configurations.

For a 7B model trained on four GH200s, the paper presents a more direct comparison with AdamW. Standard Muon costs approximately $T(3\alpha+1)n^3$3 as much optimizer time as AdamW in the reported setup. After adding the proposed improvements, Dion3 reduces that overhead to approximately $T(3\alpha+1)n^3$4 AdamW.

Figure 8: The full Dion3 stack reduces Muon’s optimizer cost from $T(3\alpha+1)n^3$5 AdamW to approximately $T(3\alpha+1)n^3$6 AdamW.

The optimizer remains slower than AdamW because orthogonalization is intrinsically more expensive than elementwise moment updates. Nevertheless, the practical gap is substantially narrowed. The paper estimates that standard Newton–Schulz can account for anywhere from roughly 2% to 17% of end-to-end training time, depending on model architecture, batch size, parallelism, and overlap. This range is important: orthogonalization may be almost hidden in a large, efficiently pipelined MoE pretraining run, yet become a dominant cost in small-batch dense-model fine-tuning.

Theoretical and practical implications

The paper’s theoretical contribution is not a new convergence theorem for orthogonal optimization. Its significance lies instead in showing that an apparently cubic matrix iteration can be algebraically reorganized to expose lower-dimensional symmetric structure without changing its exact-arithmetic output. This establishes a general design pattern for optimizer kernels: derive the matrix function at the level of invariants such as Gram matrices, then map those invariants to hardware-specialized primitives.

The practical implications are more immediate. Dion3 provides a drop-in implementation through the dion and gram-newton-schulz packages, supports Muon and NorMuon baselines, and integrates FSDP, DDP, mixed sharding, CUDA graphs, megabatching, and custom GPU kernels. The implementation emphasis is justified by the paper’s numerical findings: naïvely inserting row selection can introduce repeated casts and destroy convergence, while poorly arranged fused operations can reduce GNS stability. Production use therefore requires treating precision and kernel fusion as algorithmic concerns.

Several limitations remain. The quality experiments primarily use dense decoder-only transformers, ClimbMix or FineWeb-Edu, and selected model scales. The strongest optimization improvements are measured over 10B or 100B tokens rather than frontier-scale pretraining budgets. The observed benefit from fractional updates may depend on row-selection statistics, architecture, normalization, weight decay, or the error-feedback coefficient. The paper also cautions that forming Gram matrices squares condition numbers, making GNS unsuitable when high-accuracy polar factors are required or when matrices are poorly conditioned beyond the tolerance of Muon-style approximate updates.

Future research should examine whether fractional orthogonal updates generalize across modalities, optimizer schedules, multimodal architectures, and long-horizon reinforcement-learning workloads. A useful theoretical direction would characterize error-feedback selection as a stochastic or adaptive coordinate process and determine when it improves versus merely preserves the stationary behavior of full orthogonal updates. Systems research should extend symmetry-aware kernels to additional accelerator families, tensor layouts, and heterogeneous clusters, while algorithmic work could explore adaptive $T(3\alpha+1)n^3$7, randomized selection, block selection, or selection criteria based on curvature and optimizer-state uncertainty rather than row norms alone.

Conclusion

Dion3 presents a coherent full-stack response to Muon’s principal scalability bottlenecks. Gram Newton–Schulz reduces the algorithmic cost of orthogonalization, CuteDSL kernels exploit the resulting symmetric computation, fractional updates reduce both computation and communication, and megabatching addresses collective latency. The stabilized implementation is essential because naïve Gram reformulations are numerically divergent in half precision.

The empirical evidence supports the paper’s central claim: Dion3 can preserve Muon-family optimization quality while reducing optimizer-step time by up to approximately $T(3\alpha+1)n^3$8, with validation-loss improvements over NorMuon in the reported 3B–14B experiments. The unexpected gains from fractional updates make the method relevant not only as an engineering optimization but also as a candidate source of new optimization dynamics. Its broader importance will depend on whether those quality improvements persist across architectures and training regimes beyond the configurations studied.