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

Mar 30, 2026

MuonEq: Balancing Before Orthogonalization with Lightweight Equilibration

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.

Da Chang, Qiankun Shi, Lvgang Zhang +5

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

Ask PDF

May 8, 2026

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

Muon fixes the direction of every matrix-valued update at the polar factor of its momentum, while each layer's step magnitude is addressed only by a static shape correction. We derive a dynamic per-layer scalar by adapting the LARS/LAMB trust-ratio principle to the orthogonalized setting, where the standard denominator candidates---the raw momentum norm or the polar-factor norm---either live in the wrong unit space or carry no update-scale information. The resulting method, OrScale, uses the norm of the parameter-space direction actually applied and anchors each layer's ratio at one via a per-layer calibration, so that the Moonlight recipe (tuned for AdamW, shared with Muon via RMS matching) transfers with no additional sweep; a component ablation confirms each design choice is individually load-bearing. Theoretically, OrScale retains a nuclear-norm

O(1/\sqrt{T})

convergence rate for any clipped multiplier and achieves a strict layer-adaptive descent gain

κ_{\mathrm{eff}}>1

under two conditions estimable from standard training diagnostics---a bound that predicts the gain should grow with architectural heterogeneity. Experiments confirm the prediction: with every hyperparameter inherited verbatim from the Moonlight recipe, OrScale matches or beats Muon+Moonlight across dense 125M--1.1B FineWeb-Edu pre-training, and on a 16B-A3B mixture-of-experts model---where the logged trust ratios separate cleanly by layer class---the gap widens by an order of magnitude to

0.130

nats (

3.8\%

relative) at parity wall-clock cost.

Yuxuan Lou, Yang You

National University of Singapore

Orthogonal Optimization Neural Network Optimization Muon Optimizer Matrix Optimization Layer-Wise Learning Rate Adaptation

Ask PDF

Sep 15, 2025

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

Neural network (NN) training is inherently a large-scale matrix optimization problem, yet the matrix structure of NN parameters has long been overlooked. Recently, the optimizer Muon \citep{jordanmuon}, which explicitly exploits this structure, has gained significant attention for its strong performance in foundation model training. A key component contributing to Muon's success is matrix orthogonalization. In this paper, we propose low-rank orthogonalization, which performs orthogonalization by leveraging the low-rank nature of gradients during NN training. Building on this, we introduce low-rank matrix-signed gradient descent (MSGD) and a low-rank variant of Muon. %Numerical experiments demonstrate the superior performance of low-rank orthogonalization, with low-rank Muon achieving promising results in GPT-2 and LLaMA pretraining---surpassing the carefully tuned vanilla Muon on tasks with large model sizes. {Numerical experiments demonstrate the advantages of low-rank orthogonalization: low-rank Muon generally matches or improves upon vanilla Muon on the GPT-2 and LLaMA pretraining tasks, with clearer improvements observed for relatively larger models.} Theoretically, we establish the iteration complexity of low-rank MSGD for finding an approximate stationary solution, and the iteration complexity of low-rank Muon for finding an approximate stochastic stationary solution under heavy-tailed noise. The code to reproduce our numerical experiments is available at https://github.com/dengzhanwang/Low-rank-Muon.

Chuan He, Zhanwang Deng, Zhaosong Lu

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

Neural Network Optimization Low-Rank Approximation Matrix Optimization

Ask PDF

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