# Dion3: Full-stack orthogonal updates > Dion3 speeds up Muon optimization up to 6x by orthogonalizing only a fraction of momentum rows, which unexpectedly improves training quality across model scales. - **Source:** [arXiv](https://arxiv.org/abs/2608.11612) - **Published:** 2026-08-17 - **Permalink:** https://picx.dev/p/0QVKjG - **Whiteboard:** https://picx.dev/p/0QVKjG/image ## Summary # 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/\text{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):** | Model Size | NorMuon Loss | Dion3 ($f=1/4$) Loss | Δ Loss | NorMuon Acc (%) | Dion3 Acc (%) | Δ Acc | |---|---|---|---|---|---|---| | 3B | 2.269 | 2.257 | **-0.012** | 53.9 | **54.9** | +1.0 | | 4B | 2.243 | **2.232** | -0.011 | **55.2** | 54.9 | -0.3 | | 7B | 2.220 | **2.206** | -0.014 | 56.0 | **56.1** | +0.1 | | 14B | 2.189 | **2.162** | **-0.027** | 57.4 | **58.1** | **+0.7** | 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): | Model Size | Nodes | FSDP Shards | Batching (ms) | Megabatching (ms) | Change | |---|---|---|---|---|---| | 1B | 1 | 8 | 80.7 | 52.1 | **-35%** | | 1B | 4 | 32 | 61.9 | 59.3 | -4% | | 14B | 1 | 8 | 144.0 | 140.8 | -2% | | 14B | 4 | 32 | 94.7 | 89.1 | -6% | 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. --- _Markdown view of https://picx.dev/p/0QVKjG, served by PicX — AI-generated visual whiteboard summaries of research papers._