Dion3: Full-Stack Orthogonal Updates — Summary

Summary (Overview)

• Dion3 is a comprehensive revision of the Muon optimizer that reduces the computational and communication overhead of its cubic-time Newton-Schulz orthogonalization step, achieving up to 6× speedup in optimizer step time while matching or improving training quality.
• The paper presents four compounding contributions: (1) Gram Newton-Schulz, a mathematically equivalent reformulation that reduces FLOP cost; (2) custom symmetric GEMM kernels in CuteDSL that exploit matrix symmetry; (3) a fractional row-selection update rule that orthogonalizes only a subset of momentum rows; and (4) megabatched communication to reduce distributed training overhead.
• The Gram Newton-Schulz algorithm reduces FLOPs from $O(T\alpha n^3)$ to $O((T + \alpha)n^3)$, a significant saving for large aspect ratios $\alpha$; for typical transformer MLP blocks ($\alpha = 4, T = 5$), this saves 55% of FLOPs versus standard Newton-Schulz with symmetric GEMMs.
• Dion3 with $f = 1/4$ (selecting only 25% of momentum rows) not only matches but improves validation loss across model scales from 3B to 14B parameters, with the largest gain of -0.027 loss and +0.7% downstream accuracy at 14B.
• The method is implemented in two open-source, pip-installable packages (dion and gram-newton-schulz), serving as drop-in replacements for Muon.

---

Introduction and Theoretical Foundation

Background and Motivation

Muon has become the optimizer of choice for frontier LLMs (e.g., Kimi K2, GLM-5) due to its ability to reach a given loss in fewer steps than AdamW. However, each Muon step is more expensive due to the Newton-Schulz orthogonalization — a cubic-time matrix operation. As model sizes grow, this overhead scales super-linearly, and distributed training adds communication costs that further erode Muon's benefits.

Muon Update Rule

The Muon optimizer is best described as steepest-descent with respect to the spectral norm. The update rule is:

$$\begin{array}{l} \boldsymbol{M} \leftarrow \mu \boldsymbol{M} + \boldsymbol{G} \\ \boldsymbol{W} \leftarrow \boldsymbol{W} - \eta \operatorname{polar}(\boldsymbol{M}) \end{array}\tag{1}$$

where $\mu$ is the momentum coefficient, $\eta$ is the learning rate, $\boldsymbol{M}$ is the momentum matrix, and the polar decomposition is:

Definition 1 (Polar Decomposition). If $\boldsymbol{X} = \boldsymbol{U}\boldsymbol{\Sigma}\boldsymbol{V}^{\top}$ is the SVD of a matrix, then $\operatorname{polar}(\boldsymbol{X}) = \boldsymbol{U}\boldsymbol{V}^{\top}$.

NorMuon Variant

NorMuon adds Adam-style per-neuron adaptive normalization:

$$\begin{array}{l} \boldsymbol{M} \leftarrow \mu \boldsymbol{M} + \boldsymbol{G}, \qquad \boldsymbol{O} \leftarrow \operatorname{polar}(\boldsymbol{M}) \\ \boldsymbol{v}_i \leftarrow \beta_2 \boldsymbol{v}_i + (1 - \beta_2) \cdot \frac{1}{m} \sum_j \boldsymbol{O}_{ij}^2, \qquad \widehat{\boldsymbol{O}}_{ij} \leftarrow \frac{\boldsymbol{O}_{ij}}{\sqrt{\boldsymbol{v}_i} + \epsilon} \\ \boldsymbol{W} \leftarrow \boldsymbol{W} - \eta \frac{\|\boldsymbol{O}\|_F}{\|\widehat{\boldsymbol{O}}\|_F} \widehat{\boldsymbol{O}} \end{array}\tag{2}$$

Standard Newton-Schulz

The standard Newton-Schulz iteration applies degree-5 odd polynomials to approximate the polar decomposition:

$$\boldsymbol{X}_{t+1} = a_t \boldsymbol{X}_t + b_t \boldsymbol{X}_t \boldsymbol{X}_t^{\top} \boldsymbol{X}_t + c_t (\boldsymbol{X}_t \boldsymbol{X}_t^{\top})^2 \boldsymbol{X}_t$$

Each iteration preserves singular vectors and transforms singular values via polynomial composition. With normalization $\boldsymbol{X}_0 = \boldsymbol{X}/\|\boldsymbol{X}\|_F$, all singular values lie in $[0,1]$, and the iterates converge to $\operatorname{polar}(\boldsymbol{X}_0)$.

FLOP Analysis of Standard Newton-Schulz: For an $n \times m$ matrix with aspect ratio $\alpha = m/n \geq 1$ and $T$ iterations:

$$\text{Cost} = T(4mn^2 + 2n^3) = 2T(2\alpha + 1)n^3 \text{ FLOPs}$$

With $T = 5$, this is $(20\alpha + 10)n^3$ FLOPs spread across 15 GEMMs.

Key Challenges Identified

1. Super-linear complexity: Orthogonalization requires $O(n^3)$ time versus linear scaling for Adam.
2. Distributed training: Weight sharding requires all-to-all communication to assemble matrices before orthogonalization.
3. Symmetric structure ignored: The matrices $A = XX^{\top}$ and $A^2$ are symmetric but standard implementations don't exploit this.
4. Aspect ratio dependence: Rectangular matrix multiplications dominate the cost, and modern MoE architectures have increasing aspect ratios.

---

Methodology

1. Gram Newton-Schulz

Core Theorem — The key insight is that odd polynomials can be rewritten as $p_t(x) = x h_t(x^2)$, enabling iteration on the small symmetric Gram matrix instead of the full rectangular matrix:

Theorem 2. If $p_t(x) = x h_t(x^2)$ for all $t \in \{1, \ldots, T\}$, then $(p_T \circ \cdots \circ p_1)(x) = q_T x$, where $q_T$ is defined by the iteration $r_0 = x^2$, $q_0 = 1$, and

$z_t = h_t(r_{t-1}), \qquad r_t = r_{t-1} z_t^2, \qquad q_t = q_{t-1} z_t$

Algorithm Structure (Naive Version):

1. Compute the $n \times n$ Gram matrix $R_0 = XX^{\top}$
2. Iterate to approximate $Q_T \approx (XX^{\top})^{-1/2}$
3. Output $Q_T X$

FLOP Analysis: With symmetric GEMMs costing $n^3$ FLOPs each:

$$\text{Cost} = T \cdot 4n^3 + 3mn^2 - 3n^3 = (4T + 3\alpha - 3)n^3 \text{ FLOPs}$$

This compares favorably to standard Newton-Schulz's $T(3\alpha + 1)n^3$ FLOPs with symmetric GEMMs.

Stabilization via Restarting: The naive version suffers from spurious negative eigenvalues in the Gram matrix due to half-precision rounding errors. The fix: run only the first two iterations, compute $X_2 = Q_2 X$, then restart with $X_2$ as the new input (recomputing the Gram matrix). This resets spurious eigenvalues at the cost of $3(\alpha - 1)n^3$ FLOPs.

Algorithm 3 (Stabilized Gram Newton-Schulz) uses:

• A restart at iteration $t = 3$
• float16 instead of bfloat16 for casting
• Reformulated intermediate polynomials for stability

2. Symmetric GEMM Kernels in CuteDSL

Custom GPU kernels for operations $AB$ and $\alpha AB + \beta C$ where $AB$ and $C$ are symmetric:

• Triangular Scheduler: Only lower-triangle tiles (including diagonal) are computed and assigned to thread block clusters.
• Transposed Tile Epilogue: Computed lower-triangle results are copied to their transposed locations in the upper triangle.

These kernels target NVIDIA Hopper and Blackwell architectures and achieve ~2× speedup over cuBLAS for large enough $n$.

3. Dion3 Update Rule (Fractional Row Selection)

Algorithm 4 (Dion3 update rule, single weight matrix):

1. Selection: Pick the $k = \lceil fn \rceil$ rows of $\boldsymbol{M}$ with largest $\ell_1$ norm (where $f \in (0,1]$ is the compression factor)
2. Orthogonalization: Compute $O = \operatorname{polar}(M[S, :])$ using Gram Newton-Schulz
3. Weight update: Update only selected rows of $\boldsymbol{W}$
4. Error Feedback: Decay only selected rows of $\boldsymbol{M}$ by factor $\mu$

The error feedback mechanism differs from standard momentum: instead of $M \gets \mu M$, it uses $M \gets \mu \widehat{M} + (M - \widehat{M})$, where $\widehat{M}$ matches $M$ at selected rows and is zero elsewhere. This boosts the residual component, encouraging future iterations to select previously ignored rows.

Key implementation details:

• A custom Triton kernel restores numerical fusion for the weight update (avoiding precision loss from multiple upcast/downcast rounds)
• NorMuon normalization steps run in float32
• CUDA graph capture minimizes kernel launch overhead at small scales

4. Megabatched Communication

Instead of batching matrices in groups of world_size, megabatching groups all matrices of the same shape into a single batch:

• Pack local momentum shards into one all-to-all
• Assemble, orthogonalize as a batch, scatter back together
• Reduces communication rounds from $O(N/\mathrm{world\_size})$ to $O(1)$ per optimizer step

This is especially effective when the optimizer is communication-bound (small models, few shards).

---

Empirical Validation / Results

Model Quality

Learning-Rate Transfer Rule: The optimal learning rate for fraction $f$ scales as $\eta \sqrt{f} = \text{const}$. This is derived from matching Frobenius norms: $\|\eta \cdot \operatorname{polar}(M)\|_F = \eta\sqrt{n}$ for Muon versus $\eta'\sqrt{fn}$ for Dion3, giving $\eta' = \eta/\sqrt{f}$.

Key Results (1B-parameter models, 100B tokens of ClimbMix):

• Dion3 with $f < 1$ outperforms fully-tuned NorMuon when learning rate is properly scaled
• Lowest loss achieved at $f = 1/8$ (approximately 0.01 lower than NorMuon)
• All Dion3 variants track below the NorMuon baseline throughout training

Scaling Results (10B tokens, 3B–14B parameters):

Dion3 improves validation loss at every scale (largest gain at 14B: -0.027) and wins downstream accuracy at 3 of 4 scales.

Optimizer Speedup

Figure 6 results (optimizer step time relative to standard Muon):

• Symmetric kernels + Gram Newton-Schulz: 1.5× combined speedup
• Adding $f = 1/2$: additional 2× reduction
• Adding $f = 1/4$: additional 3.7× reduction
• Overall: 3.6× ($f=1/2$) and 6.5× ($f=1/4$) speedup over standard Muon for larger models
• For MoE architectures (higher aspect ratios $\alpha = 8$), Gram Newton-Schulz + symmetric kernels alone achieve 2× speedup

Megabatching impact (Muon, per-GPU optimizer step time):

Megabatching has the largest impact when the optimizer is communication-bound (small models, few shards).

---

Theoretical and Practical Implications

Theoretical Contributions

1. Gram Newton-Schulz provides a mathematically elegant reformulation showing that Newton-Schulz iterations implicitly compute inverse square roots of the Gram matrix. The complexity improvement from $O(T\alpha n^3)$ to $O((T+\alpha)n^3)$ is a fundamental algorithmic advance.

2. Stability analysis identifies spurious negative eigenvalues from half-precision arithmetic as the key failure mode, with a practical restart strategy that fully mitigates the issue.

3. The learning-rate transfer rule ($\eta' = \eta/\sqrt{f}$) provides a principled way to adjust hyperparameters when compressing the update.

Practical Contributions

1. Full-stack optimization: The four contributions operate at different levels (kernels → algorithm → update rule → communication), compounding for maximum benefit.

2. Drop-in replacement: The dion and gram-newton-schulz packages make Muon-family optimizers practical across a wide range of settings without architectural or parallelism constraints.

3. Unexpected quality improvement: The finding that subsampling rows improves training quality (rather than merely approximating Muon) is surprising and suggests that partial updates may act as a form of regularization or implicit momentum diversity.

4. Kimi's success demystified: The paper explains that Muon's successful scaling at Moonshot AI relied on a fragile alignment of deprecated PyTorch features, fine-grained MoE architecture, and specific parallelism strategies — all of which Dion3 makes unnecessary.

---

Conclusion

Dion3 addresses the scalability challenges of Muon at every level of the stack:

• Gram Newton-Schulz reduces FLOP cost by iterating on the small symmetric Gram matrix
• CuteDSL symmetric kernels exploit matrix symmetry for ~2× speedup over cuBLAS
• Fractional row selection ($f = 1/4$) provides an additional 3.7× speedup while improving training quality
• Megabatching reduces communication rounds to $O(1)$ per step

The combined system achieves up to 6× faster optimizer steps than standard Muon while matching or improving loss. The unexpected quality improvement from subsampling warrants further investigation — the authors note support from Joo et al. [18], who showed that randomly masking update blocks improves SGD with momentum.

Future directions include:

• Understanding how widely the quality improvement generalizes across architectures and datasets
• Exploring alternative selection strategies beyond top-$\ell_1$-norm rows
• Extending the approach to other orthogonalization-based optimizers

The packages are available as open-source, making orthogonal optimizers practical and accessible for general-purpose LLM training.