# How Does Orthogonalization Adapt to the Neural Network Hessian Structure? A Gradient Self Outer-Product Analysis at Initialization

> Muon's orthogonalization preconditioner matches neural-network Hessian structure because the gradient self outer-product becomes diagonally dominant as hidden dimensions grow, with explicit asymptotic rates across canonical models.

- **Source:** [arXiv](https://arxiv.org/abs/2610.28975)
- **Published:** 2026-10-24
- **Permalink:** https://picx.dev/p/i7OVTv
- **Whiteboard:** https://picx.dev/p/i7OVTv/image

## Summary

# 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)

| Model | Diagonal | Off-diagonal | Asymptotic dims |
|---|---|---|---|
| Matrix factorization | $\Theta(k^3)$ | $\Theta(k^2)$ | $k$ |
| Deep linear, hidden layer $i$ | $\Theta(d_i\prod_{j=1,j\neq i}^{L-1} d_j^2)$ | 0 | $d_1, \ldots, d_{L-1}$ |
| Deep linear, output layer | $\Theta(\prod_{j=1}^{L-1} d_j^2)$ | $\Theta(\prod_{j=1}^{L-1} d_j)$ | $d_1, \ldots, d_{L-1}$ |
| Two-layer ReLU, layer $W_1$ | $\Theta(d_1)$ | $\Theta(1)$ | $d_1$ |
| Two-layer ReLU, layer $W_2$ | $\Theta(d_1^2)$ | $\Theta(d_1)$ | $d_1$ |
| Linear attention, $W_Q, W_K$ | $\Theta(d_2^2 d_v^2 d_0^6/d_k)$ | 0 | $d_v, d_0, d_2$ |
| Linear attention, $W_V$ | $\Theta(d_2(d_2+d_v)d_0^6)$ | 0 | $d_v, d_0, d_2$ |
| Linear attention, $W_O$ | $\Theta(d_v^2 d_0^6)$ | $\Theta(d_v d_0^3)$ | $d_v, d_0$ |

### 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).

---

_Markdown view of https://picx.dev/p/i7OVTv, served by PicX — AI-generated visual whiteboard summaries of research papers._