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

Shenyang Deng; Garvit Goyal; Tianyu Pang; Shuhua Yu; Yaoqing Yang

cs.LGOct 20, 2026

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

Authors: Shenyang Deng, Garvit Goyal, Tianyu Pang, Shuhua Yu, Yaoqing Yang

Organizations: Dartmouth College · Meta

Abstract

Muon orthogonalizes a weight matrix's momentum before each step, and on neural networks this simple preconditioner beats AdamW in a wide range of tasks. Existing analyses explain Muon within general optimization frameworks or on a few specific models, and many of them view orthogonalization as a preconditioner. Since the Hessian is the optimal preconditioner under a quadratic approximation of the loss, these works leave a basic question open: what makes orthogonalization suited to the Hessian of a neural network? Prior work shows that the layer-wise Hessian of a neural network is row-block diagonally dominant, and Muon's implicit preconditioner, the Kronecker product $(\mathbf{V}\mathbf{V}^\top)^{1/2}\otimes\mathbf{I}$ with $\mathbf{V}$ the momentum, has the same row-block structure exactly when $\mathbf{V}\mathbf{V}^\top$ is diagonally dominant. Diagonal dominance of $\mathbf{V}\mathbf{V}^\top$ has been observed in language models, but why a neural network produces it has not been explained. This work provides a quantitative asymptotic theory for neural networks at initialization, and points out the intuition that the overparameterization of the weight matrices drives the diagonal dominance. We analyze $\mathbb{E}[\mathbf{G}\mathbf{G}^\top]$ , which equals $\mathbb{E}[\mathbf{V}\mathbf{V}^\top]$ at initialization, in four canonical models: symmetric matrix factorization, deep linear networks, two-layer ReLU networks, and single-layer linear attention. In each case the diagonal entries outgrow the off-diagonal ones at explicit rates as the hidden dimensions grow; for the query, key and value weights of attention the off-diagonal vanishes altogether. Simulations of the four models and experiments on language models at initialization confirm the predicted scaling.

Figures & tables

Figure 1: Left: Hessian of a Transformer weight matrix at initialization, whose mass concentrates on the diagonal blocks. Right: the momentum self outer-product $\bm{V}_{t}\bm{V}_{t}^{\top}$ of a Transformer weight matrix at initialization (log scale), whose mass concentrates on the diagonal.

Figure 2: Empirical $\widehat{\mathbb{E}}[\bm{G}\bm{G}^{\top}]$ (log scale) of the query, key, value and output projections at initialization, for single-layer linear attention (left, details in Section G.4 ) and the first attention block of Qwen3.5 (right, details in Section H.1 ). In both cases the mass concentrates on the diagonal. Each heatmap is normalized by its empirical mean diagonal magnitude.

Figure 3: Diagonal-dominance ratio

\widehat{R}

for symmetric matrix factorization against

k

(left;

d=25

, rank-5 target) and for the two-layer ReLU network against

d_{1}

(right;

d_{0}=512

d_{2}=12

Figure 4: Diagonal-dominance ratio

\widehat{R}

\bm{G}_{1},\dots,\bm{G}_{4}

for the depth-

4

linear network when

d_{1}

d_{2}

d_{3}

is varied (left to right). Each curve is normalized by its value at the baseline (

d_{1}=d_{2}=d_{3}=80

) to facilitate comparison across layers with different absolute scales.

Figure 5: Diagonal-dominance ratio

\widehat{R}

of the four attention projections when

d_{k}=d_{v}

d_{k}

d_{v}

, and

d_{0}=d_{2}

are varied (left to right). Each curve is normalized by its value at the baseline (first three:

T=32

d_{0}=512

d_{k}=d_{v}=80

d_{2}=12

; fourth:

T=16,d_{0}=d_{2}=32,d_{k}=d_{v}=20

) to facilitate comparison across projections with different absolute scales.

Figure 6: Model width scaling of

\widehat{R}

for the query, key, value, and output projections in the first head of the first eligible attention block of the three language models at random initialization.

Appendix figures & tables10 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 7: Empirical expected gradient self outer-product for symmetric matrix factorization with $d=25$ , $k=80$ , and a fixed rank-5 target, estimated over $10{,}000$ independent initializations. Each heatmap is normalized by its empirical mean diagonal magnitude. The panel uses the logarithmic color scale as shown.

Figure 8: Hidden-dimension scaling for symmetric matrix factorization with

d=25

and a fixed rank-5 target. From left to right, the panels show

\widehat{D}

\widehat{\mathrm{OD}}

, and

\widehat{R}=\widehat{D}/\widehat{\mathrm{OD}}

as the factorization hidden dimension

k

is varied. Empirical estimates use

10{,}000

independent initializations at each hidden dimension;

\widehat{\mathrm{OD}}

uses the paired estimator in ( 103 ).

Figure 9: Empirical expected gradient self outer-products for the depth- $4$ linear network with baseline dimensions $[512,80,80,80,512]$ , estimated from $10{,}000$ independent initializations. The four panels show $\widehat{\mathbb{E}}[\bm{G}_{i}\bm{G}_{i}^{\top}]$ for $i=1,\ldots,4$ . Each heatmap is normalized by its empirical mean diagonal magnitude. A common logarithmic color scale is used across all panels.

Figure 10: Empirical dimension scaling for the depth-

4

linear network. Columns vary

d_{1},d_{2},

and

d_{3}

individually while holding the remaining dimensions fixed at their baseline values. Rows show

\widehat{D}

\widehat{\mathrm{OD}}

, and

\widehat{R}=\widehat{D}/\widehat{\mathrm{OD}}

for

\bm{G}_{1},\ldots,\bm{G}_{4}

, each normalized by its value at the baseline configuration

[d_{0},d_{1},d_{2},d_{3},d_{4}]=[512,80,80,80,512]

to facilitate comparison across layers with different absolute scales. Empirical estimates use

10{,}000

independent initializations per configuration. The off-diagonal statistic is estimated using ( 106 ).

Figure 11: Empirical expected gradient self outer-products for the two-layer ReLU network with dimensions $[d_{0},d_{1},d_{2}]=[512,80,12]$ , estimated from $10{,}000$ independent weight initializations. The panels show $\widehat{\mathbb{E}}[\bm{G}_{1}\bm{G}_{1}^{\top}]$ and $\widehat{\mathbb{E}}[\bm{G}_{2}\bm{G}_{2}^{\top}]$ . Each heatmap is normalized by its empirical mean diagonal magnitude. A common logarithmic color scale is used across all panels.

Figure 12: Hidden-dimension scaling for the two-layer ReLU network with

d_{0}=512

and

d_{2}=12

. The top and bottom rows correspond to

\bm{G}_{1}

and

\bm{G}_{2}

, respectively, and show

\widehat{D}

\widehat{\mathrm{OD}}

, and

\widehat{R}=\widehat{D}/\widehat{\mathrm{OD}}

d_{1}

is varied. For

\bm{G}_{1}

\widehat{\mathrm{OD}}

is estimated from the zero-target Gram; for

\bm{G}_{2}

, it is estimated using ( 119 ).

Figure 13: Empirical expected gradient self outer-products for single-layer attention with $T=16$ , $d_{0}=32$ , $d_{k}=d_{v}=20$ , and $d_{2}=12$ , estimated from $10{,}000$ independent initializations. The four panels show $\widehat{\mathbb{E}}[\bm{G}_{p}\bm{G}_{p}^{\top}]$ for $p\in\{Q,K,V,O\}$ . Each heatmap is normalized by its empirical mean diagonal magnitude. A common logarithmic color scale is used across all panels.

Figure 14: Empirical dimension scaling for single-layer attention. Rows vary

d_{k}=d_{v}

jointly,

d_{k}

d_{v}

, and

d_{0}=d_{2}

(top to bottom). Columns show

\widehat{D}

\widehat{\mathrm{OD}}

, and

\widehat{R}=\widehat{D}/\widehat{\mathrm{OD}}

for

\bm{G}_{Q},\bm{G}_{K},\bm{G}_{V},

and

\bm{G}_{O}

, each normalized by its value at the baseline configuration (first three rows:

T=32

d_{0}=512

d_{k}=d_{v}=80

d_{2}=12

; fourth row:

T=16,d_{0}=d_{2}=32,d_{k}=d_{v}=20

), to facilitate comparison across projections with different absolute scales. Empirical estimates use

10{,}000

independent initializations per configuration; the off-diagonal statistics are estimated using ( 125 ).

Figure 15: Empirical expected gradient self outer-products for Qwen3.5-0.8B-Base, Granite-4.0-1B-Base, and Llama-3.2-1B (columns), estimated from $25$ independent random initializations and $64$ fixed FineWeb-Edu sequences of length $T=256$ . The rows show $\widehat{\mathbb{E}}[\bm{G}_{Q}\bm{G}_{Q}^{\top}]$ , $\widehat{\mathbb{E}}[\bm{G}_{K}\bm{G}_{K}^{\top}]$ , $\widehat{\mathbb{E}}[\bm{G}_{V}\bm{G}_{V}^{\top}]$ , and $\widehat{\mathbb{E}}[\bm{G}_{O}\bm{G}_{O}^{\top}]$ . We use head $0$ of block $3$ for Qwen and head $0$ of block $0$ for Granite and Llama. Each heatmap is normalized by its empirical mean diagonal magnitude. A common logarithmic color scale is used across all panels.

Figure 16: Model width scaling of the empirical diagonal-dominance ratio

\widehat{R}

for

\bm{G}_{Q},\bm{G}_{K},\bm{G}_{V},

and

\bm{G}_{O}

in Qwen3.5-0.8B-Base, Granite-4.0-1B-Base, and Llama-3.2-1B. Model width ( hidden_size ) is varied while the head dimension ( head_dim ) is held fixed at its released value. We analyze head

0

in block

3

for Qwen and in block

0

for Granite and Llama. Each point is estimated from

25

independent random initializations and

64

fixed FineWeb-Edu sequences of length

T=256

Neural Network Optimization

Neural network optimization studies algorithms and training dynamics that help models learn efficiently, reliably, and generalize well.

See all 332 papers

MuonEq: Balancing Before Orthogonalization with Lightweight Equilibration
Mar 30, 2026Da Chang, Qiankun Shi, Lvgang Zhang +5Neural Network Optimization Muon Optimizer
Shenzhen Institute of Advanced Technology, Chinese Academy of Sciences · Pengcheng Laboratory · University of Chinese Academy of Sciences +4
OrScale: Orthogonalised Optimization with Layer-Wise Trust-Ratio Scaling
May 8, 2026Yuxuan Lou, Yang YouOrthogonal Optimization Neural Network Optimization
National University of Singapore
Low-rank Orthogonalization for Large-scale Matrix Optimization with Applications to Foundation Model Training
Sep 15, 2025Chuan He, Zhanwang Deng, Zhaosong LuNeural Network Optimization Low-Rank Approximation
Department of Mathematics Linköping University, Sweden · Academy for Advanced Interdisciplinary Studies Peking University, Beijing, People’s Republic of China · Department of Industrial and Systems Engineering University of Minnesota, USA

Mar 30, 2026J/K move · Enter open · S save

MuonEq: Balancing Before Orthogonalization with Lightweight Equilibration

Da Chang, Qiankun Shi, Lvgang Zhang, Yu Li+4

Shenzhen Institute of Advanced Technology, Chinese Academy of Sciences · Pengcheng Laboratory · University of Chinese Academy of Sciences+4

Neural Network Optimization Muon Optimizer Orthogonal Optimization

Open paper AskPDF arXiv

Orthogonalized-update optimizers such as Muon improve training of matrix-valued parameters, but existing extensions typically either rescale updates after orthogonalization or use heavier whitening-based preconditioners before it. We introduce {\method}, a lightweight family of pre-orthogonalization equilibration schemes for Muon with three forms: two-sided row/column normalization (RC), row normalization (R), and column normalization (C). By rebalancing the momentum matrix before finite-step Newton--Schulz orthogonalization, {\method} improves the geometry seen by orthogonalization. We show that finite-step orthogonalization is governed by the input spectrum, especially stable rank and condition number, and that row/column normalization acts as a zeroth-order surrogate for whitening. For hidden matrix weights, R is the default variant. Theoretically, {\method} (R) retains the standard

\widetilde{\mathcal O}(T^{-1/4})

Muon-type nonconvex stationarity guarantee with decoupled weight decay and a horizon-free diminishing learning-rate schedule, and extends it to finite-step NS5 up to an explicit inexactness constant. In LLaMA2 pretraining on C4, {\method} (R) consistently outperforms Muon on 130M, 350M, and 1B models, with faster convergence and lower validation perplexity. The code is available at the MuonEq codebase.

Abstract

Neural Network Optimization

Explore similar work

MuonEq: Balancing Before Orthogonalization with Lightweight Equilibration

OrScale: Orthogonalised Optimization with Layer-Wise Trust-Ratio Scaling

Low-rank Orthogonalization for Large-scale Matrix Optimization with Applications to Foundation Model Training