How Does Orthogonalization Adapt to the Neural Network Hessian Structure?

Summary (Overview)

• Core question: Why does Muon's orthogonalization preconditioner, which uses $(V_t V_t^\top)^{1/2} \otimes I$ as an implicit metric, align with the Hessian structure of neural networks?

• Main finding: At initialization, $\mathbb{E}[GG^\top]$ (which equals $\mathbb{E}[VV^\top]$) becomes diagonally dominant as hidden dimensions grow, causing Muon's preconditioner to match the row-block diagonally dominant structure of the neural-network Hessian.

• Key contribution: The paper provides quantitative asymptotic rates for diagonal dominance in four canonical models: symmetric matrix factorization, deep linear networks, two-layer ReLU networks, and single-layer linear attention (Theorems 1–4).

• Empirical validation: Simulations reproduce predicted scaling rates, and experiments on Qwen3.5, Llama-3.2, and Granite-4.0 language models confirm that diagonal dominance strengthens with model width.

• Intuition: Overparameterization of weight matrices drives the diagonal dominance—diagonal entries of $GG^\top$ accumulate over hidden coordinates while off-diagonal entries cancel in expectation.

Introduction and Theoretical Foundation

Background: Muon's Orthogonalization

Muon optimizes a matrix-shaped weight $W_t \in \mathbb{R}^{m \times n}$ by orthogonalizing the momentum buffer:

$$V_t = \beta V_{t-1} + (1-\beta) G_t, \qquad D_t = \mathrm{NS}_5(V_t) \approx (V_t V_t^\top)^{-1/2} V_t, \qquad W_{t+1} = W_t - \eta_t D_t$$

where $\mathrm{NS}_5$ is a fixed-point iteration for the orthogonal polar factor: if $V_t = P\Sigma Q^\top$ is a thin SVD, then $D_t = PQ^\top = (V_t V_t^\top)^{-1/2}V_t$.

The Equivalence to a Preconditioner

The orthogonalized update can be rewritten as:

$$\mathrm{vec}(D_t) = \left((V_t V_t^\top)^{-1/2} \otimes I_n\right) \mathrm{vec}(V_t)$$

so Muon descends in the metric:

$$H_{\mathrm{Muon}} = P_t^{1/2} \otimes I_n, \qquad P_t := V_t V_t^\top \in \mathbb{R}^{m \times m}$$

At initialization, $V_0 = G_0$, hence $P_0 = G_0 G_0^\top$.

The Hessian Connection

Prior work (Zhang et al., 2024a,b) showed that the neural-network Hessian is row-block diagonally dominant: for a weight $W \in \mathbb{R}^{m \times n}$, the $m \times m$ blocks indexed by rows (output neurons) are dominated by diagonal entries. When $P_t$ is close to diagonal, $H_{\mathrm{Muon}}$ shares this row-block structure (Figure 1).

The central question: How do neural networks produce such a diagonally dominant $V_t V_t^\top$ (equivalently, $G_t G_t^\top$ at initialization)?

Methodology

Four Canonical Models

1. Symmetric Matrix Factorization: For $W \in \mathbb{R}^{d \times k}$ fitting target $M^\star \in \mathbb{R}^{d \times d}$ (symmetric PSD, rank $r$):

$$\mathcal{L}(W) := \frac{1}{4}\|WW^\top - M^\star\|_F^2, \quad G := \frac{\partial \mathcal{L}}{\partial W} = (WW^\top - M^\star)W$$

2. Deep Linear Network: For $L \geq 1$ layers with weights $W_i \in \mathbb{R}^{d_i \times d_{i-1}}$ fitting linear target $\Phi$:

$$f(x) = W_L \dots W_1 x, \quad G_i := \frac{\partial \mathcal{L}}{\partial W_i} = A_i^\top R B_i^\top$$

where $B_i := W_{i-1} \cdots W_1$, $A_i := W_L \cdots W_{i+1}$, and $R := W_L \cdots W_1 - \Phi$.

3. Two-layer ReLU Network: With $h(x) := \sigma(W_1 x)$, residual $r(x) := f(x) - \Phi x$, and $D(x) := \mathrm{diag}(\sigma'(W_1 x))$:

$$G_1 = \mathbb{E}_x[D(x)W_2^\top r(x)x^\top], \quad G_2 = \mathbb{E}_x[r(x)h(x)^\top]$$

4. Single-layer Linear Attention: With attention matrix $A := \frac{1}{\sqrt{d_k}}XW_Q^\top W_K X^\top$:

$$f(X) = AXW_V^\top W_O^\top, \quad \mathcal{L} := \frac{1}{2}\|f(X) - X\Phi^\top\|_F^2$$

Analysis Approach

Each row of $G$ is a random vector; the $(a,b)$ entry of $GG^\top$ is the inner product of rows $a$ and $b$. Diagonal entries (squared norms) accumulate over hidden coordinates, while off-diagonal entries involve independent Gaussian vectors whose terms cancel in expectation. The proofs use:

• Wishart distribution moments for matrix factorization, deep linear networks, and linear attention
• Stein's lemma with the arc-cosine kernel (Cho & Saul, 2009) for ReLU networks

Empirical Validation / Results

Theoretical Results (Table 1)

Key Theorems

Theorem 1 (Matrix factorization): With $S := WW^\top \sim \mathcal{W}_d(k, I_d)$ Wishart:

$$\mathbb{E}[GG^\top] = \alpha_3 I_d - 2\alpha_2 M^\star + k(M^\star)^2$$

where $\alpha_2 = k(k+d+1)$ and $\alpha_3 = k^3 + 3(d+1)k^2 + (d^2+3d+4)k$. Diagonal entries are $\Theta(k^3)$; off-diagonal entries are $\Theta(k^2)$.

Theorem 3 (Two-layer ReLU): For $G_2$ with $\alpha := \mathbb{E}_{W_1}[\mathrm{tr}(H^2)] = \Theta(d_1^2)$:

$$\mathbb{E}[G_2 G_2^\top] = \alpha I_{d_2} + \frac{d_1}{4}\Phi\Phi^\top$$

For $G_1$: diagonal entries are $\Theta(d_1)$ and off-diagonal entries are $\Theta(1)$.

Theorem 4 (Linear attention): For $W_Q, W_K, W_V$, off-diagonal entries are exactly zero. For $W_O$:

$$\mathbb{E}[G_O G_O^\top] = \bar{\alpha}_O I_{d_2} + d_v\gamma\Phi\Phi^\top$$

with ratio $\Theta(d_v d_0^3)$ diverging with head dimension and as the cube of model width.

Simulation Results

• Matrix factorization: $\widehat{R}$ increases with $k$ following $R = \Theta(k)$
• Deep linear networks: $\widehat{R}$ increases with every hidden dimension for all layers
• Two-layer ReLU: $\widehat{R}$ increases with $d_1$ for both layers, matching $R = \Theta(d_1)$
• Linear attention: $\widehat{R}$ for $G_O$ grows linearly with $d_v$, is nearly unchanged with $d_k$, and grows superlinearly with $d_0$

Language Model Experiments

Figure 6 shows $\widehat{R}$ increasing with model width for all four attention projections (query, key, value, output) in Qwen3.5-0.8B-Base, Granite-4.0-1B-Base, and Llama-3.2-1B at random initialization, consistent with the $\Theta(d_v d_0^3)$ prediction.

Theoretical and Practical Implications

Theoretical Significance

1. Bridges optimization theory and neural network structure: The paper connects Muon's success to a concrete structural property of neural networks—the diagonal dominance of the gradient self outer-product—rather than abstract smoothness or convexity assumptions.

2. Identifies overparameterization as the driver: The explicit rates show that the diagonal dominance emerges from the hidden dimensions of weight matrices, providing a mechanistic explanation for why larger models benefit more from Muon.

3. Exact results for attention: For $W_Q, W_K, W_V$ in linear attention, $\mathbb{E}[GG^\top]$ is exactly diagonal, not merely asymptotically so—a stronger structural guarantee.

Practical Implications

1. Preconditioner design: The results justify the row-block structure of Muon's preconditioner, suggesting that orthogonalization is asymptotically equivalent to row-wise normalization (Appendix F).

2. Scaling behavior: The $\Theta(d_v d_0^3)$ ratio for the output projection means diagonal dominance strengthens as models get wider, explaining why Muon's advantage grows with model scale.

3. Model-specific predictions: The theory predicts which weight matrices benefit most from orthogonalization (e.g., query/key/value weights have exact diagonal dominance in attention).

Conclusion

Main Takeaways

This work provides the first quantitative asymptotic theory explaining why Muon's orthogonalization preconditioner matches the Hessian structure of neural networks: $\mathbb{E}[GG^\top]$ becomes diagonally dominant as weight matrices are overparameterized, with explicit rates across four canonical models. The diagonal entries accumulate over hidden coordinates while off-diagonal entries cancel in expectation, driving the dominance.

Future Directions

• Training dynamics: The theory is stated at initialization and in expectation; characterizing how $V_t V_t^\top$ evolves during training remains open (though empirical evidence suggests dominance persists).
• Degeneracy concerns: Whether this dominance reflects a degenerate behavior of large-scale models warrants investigation.
• Broader architectures: Extending the analysis beyond the four canonical models to more complex architectures (e.g., multi-layer attention, convolutional networks).