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

1 Introduction

• Muon orthogonalizes each weight matrix’s momentum, and its implicit preconditioner shares the Hessian’s row-block structure when the momentum self-outer-product is close to diagonal.Prior work finds the layer-wise neural-network Hessian row-block diagonally dominant; the paper therefore focuses on why the momentum self-outer-product becomes diagonally dominant.Source: §1 · pp. 1–2
• Existing analyses use general optimization frameworks or specific models, but leave open why orthogonalization suits the neural-network Hessian.The Hessian is described as the optimal preconditioner under a quadratic approximation of the loss.Source: §1 · p. 2
• In four canonical models, the diagonal of E[GG⊤] outgrows the off-diagonal as hidden dimensions grow, making Muon’s implicit preconditioner asymptotically diagonally dominant.The paper gives explicit rates for symmetric matrix factorization, deep linear networks, two-layer ReLU networks, and single-layer linear attention.Source: §1 · pp. 2–3
• Simulations reproduce the predicted rates, while Qwen3.5, Llama-3.2, and Granite-4.0 at initialization show diagonal dominance strengthening with model width.Figure 1 juxtaposes a Transformer weight-matrix Hessian with the momentum self-outer-product; Figure 2 shows empirical gradient self-outer-products for attention projections.Source: Figure 1 · p. 2, §1 · p. 3, Figure 2 · p. 3

Mechanistic analyses of Muon

• Muon analyses follow curvature-based and update-based routes, but neither compares the implicit preconditioner’s structure with the network Hessian.The curvature route includes spectral flattening and loss-decrease analyses; the update route studies the polar factor and its spectral dynamics.Source: §2 · pp. 3–4

Block structure of the neural-network Hessian

• Neural-network Hessians concentrate mass in dense principal blocks indexed by output neurons, with heterogeneity across parameter blocks and, under anisotropic inputs, within blocks.Muon’s (V V⊤)^(1/2) ⊗ I has the row-block structure precisely when V V⊤ is diagonal.Source: §2 · p. 4

3 Preliminaries

• The object of study is E[GG⊤] for a weight matrix at Gaussian initialization, with G the loss gradient with respect to that matrix.The paper notes that this expectation equals the initialization momentum self-outer-product expectation.Source: §3 · p. 4

3.1 Setup

• The four setups are symmetric matrix factorization, deep linear networks, two-layer ReLU networks, and single-layer linear attention, each fitting a fixed target.The attention setup has query, key, value, and output weights; the deep linear and ReLU setups use Gaussian inputs and fixed linear targets.Source: §3.1 · pp. 4–5

3.2 Assumptions

• Each model uses standard Gaussian initialization with deterministic targets independent of the weights, while linear attention also assumes independent Gaussian tokens.Expectations are over the weights, and for attention over both weights and tokens.Source: §3.2 · p. 5

4 Main Results

• The section gives closed forms for E[GG⊤] in each model and analyzes their diagonal dominance, with proofs deferred to Appendices B–E.Source: §4 · p. 5

4.1 Symmetric matrix factorization

• For symmetric matrix factorization, the diagonal-to-off-diagonal ratio grows as Θ(k), and E[GG⊤] approaches an isotropic matrix whose leading term is target-independent.The target contributes only lower-order terms.Source: §4.1 · p. 6

4.2 Deep linear network

• In deep linear networks, hidden-layer E[GiG⊤] is exactly diagonal, while the output-layer diagonal-to-off-diagonal ratio grows with the product of hidden dimensions.With equal hidden dimensions d, the output-layer diagonal and off-diagonal orders are Θ(d^(2(L−1))) and Θ(d^(L−1)), respectively.Source: §4.2 · p. 6

4.3 Two-layer ReLU network

• In the two-layer ReLU network, both layers’ diagonal-to-off-diagonal ratios grow as Θ(d1), although the first layer’s off-diagonal is target-free and arises from ReLU’s lack of odd symmetry.For the second layer, the off-diagonal term comes only from the target.Source: §4.3 · p. 6

4.4 Single-layer linear attention

• For linear attention, the expected off-diagonals are exactly zero for Q, K, and V, while the output projection’s ratio grows as Θ(dv d0^3).The output ratio is unaffected by dk and grows cubically with model width when d0 = d2.Source: §4.4 · p. 7

4.5 Summary and proof sketch

• Across all four models, E[GG⊤] becomes diagonally dominant as hidden dimensions grow, making Muon’s implicit preconditioner asymptotically block-diagonal in the Hessian’s row-block structure.The section also sketches a shared proof strategy based on expected inner products between gradient rows.Source: §4.5 · p. 7
• The proof strategy computes expected row inner products: independent Gaussian rows cancel most off-diagonal terms, leaving lower-order contributions from the fixed target.The derivations use Wishart moments for matrix factorization, deep linear networks, and linear attention, and Stein’s lemma with the arc-cosine kernel for ReLU.Source: §4.5 · p. 7

Simulation setup

• The simulations vary dimensions governing each model’s theoretical scaling and estimate the expected gradient self outer-product over independent initializations.The analysis summarizes mean diagonal magnitude, mean off-diagonal magnitude, and their ratio.Source: §5 · pp. 7–8

Symmetric matrix factorization and two-layer ReLU network

• The matrix-factorization ratio bR increases with k, while both layers’ ratios increase with d1 in the two-layer ReLU network, matching the predicted Θ(k) and Θ(d1) scalings.For matrix factorization, bD and dOD follow Θ(k^3) and Θ(k^2), respectively.Source: Figure 3 · p. 8, §5 · pp. 8–9

Deep linear network

• In the depth-4 linear network, bR increases with each hidden dimension for all layers; hidden-layer expected off-diagonals are zero, and the output-layer ratio grows approximately linearly with each varied dimension.The output-layer orders are D = Θ((d1d2d3)^2) and OD = Θ(d1d2d3).Source: Figure 4 · p. 8, §5 · p. 8

Single-layer linear attention

• For attention, the output projection’s bR grows approximately linearly with dv, is nearly unchanged with dk, and grows superlinearly with d0; Q, K, and V have zero expected off-diagonals.Their finite estimated bR values reflect Monte Carlo variance and follow the corresponding diagonal trends.Source: Figure 5 · p. 9, §5 · p. 9

• In Qwen3.5-0.8B-Base, Granite-4.0-1B-Base, and Llama-3.2-1B, bR increases with model width for all four attention projections.The study uses 25 random initializations and 64 fixed FineWeb-Edu sequences per configuration.Source: §6 · p. 9

7 Conclusion and Limitations

• Across four models, diagonal dominance grows at explicit rates: query, key, and value weights are exactly diagonal in linear attention, while other cases have growing diagonal-to-off-diagonal ratios.The output-weight ratio in linear attention grows with the head dimension and as the cube of model width.Source: §7 · p. 9
• The theory covers initialization and expectation, not the training dynamics of V_tV_t^T; whether dominance reflects degenerate behavior at large scale remains open.The paper cites empirical evidence that dominance persists during training, but leaves its dynamics and possible degeneracy for future work.Source: §7 · p. 10

Ethics statement

• The experiments use synthetic simulations and publicly available language models and corpora, with no human subjects, personal or sensitive data, or newly collected datasets.The authors foresee no ethical concerns or direct negative societal impacts.Source: Ethics statement · p. 10

Reproducibility statement

• The appendices restate theorem assumptions and supporting lemmas, while the numerical checks and language-model experiments report their settings and the code is provided in an anonymous repository.The reproducibility statement names the dimensions, sample sizes, initialization, model sizes, datasets, optimizer settings, and training steps as reported details.Source: Reproducibility statement · p. 10

A Extended Related Work

• The analysis connects row normalization and orthogonalization: when VV^T is diagonally dominant, they are the same operation.The paper places this connection alongside prior row-wise normalization optimizers and says it explains their common thread.Source: A Extended Related Work · pp. 16–17

B.1 Setup; B.2 Standing assumption

• The matrix-factorization appendix sets up an over-parameterized symmetric loss with a deterministic positive-semidefinite target and assumes i.i.d. standard Gaussian weights.Under this assumption, WW^T is Wishart with k degrees of freedom and identity scale.Source: §B.1 · p. 17, §B.2 · p. 17

B.3 Main results restated

• The expected matrix is invariant under orthogonal basis changes that fix M⋆.The formula expresses E[GG^T] using the identity, M⋆, and (M⋆)^2.Source: §B.3 · p. 17
• As k → ∞, the reported matrix-factorization expectations scale as Θ(k^3) on diagonal entries, E[(GG^T)_ii], and Θ(k^2) on off-diagonal entries, E[(GG^T)_ij].The supplied matrix formula is E[GG^T] = α3(d, k) I_d − 2α2(d, k) M⋆ + k(M⋆)^2.Source: §B.3 · p. 17

B.4 Preliminary lemmas; B.5 Proofs

• The third-moment coefficient α3 is k^3 + 3(d + 1)k^2 + (d^2 + 3d + 4)k; triples (a, b, c) are partitioned by multiplicity pattern.Source: §B.5 · p. 19

B.5 Proofs

• The paper interprets the Θ(k) gap between diagonal and off-diagonal scales as the regime where orthogonalization suppresses the isotropic Θ(k^3)I_d part relative to the target-aligned signal.This interpretation is stated for the matrix-factorization calculation.Source: §B.5 · p. 19

C.1 Setup; C.2 Standing assumption

• The deep linear-network setup defines a depth-L population squared-error loss and forward, backward, and residual factors, under Gaussian weights, Gaussian inputs, and a deterministic target map.The weights are independent standard Gaussians; the input is distributed as N(0,I_d0), and Φ is deterministic and independent of the weights.Source: §C.1 · p. 19, §C.2 · p. 20

C.3 Main results restated

• For every layer i < L, E[GiGi⊤] is a scalar multiple of the identity, so its off-diagonal entries vanish.The proof states that the final layer's target-dependent term contributes entries from ΦΦ⊤.Source: §C.3 · p. 20, §C.5 · p. 22

C.5 Proofs

• The proof derives the layerwise results using expansions, recursive moments, and independence between Ai and Bi.The proof attributes the vanishing off-diagonal entries for i < L to Ti being a multiple of Idi.Source: §C.5 · pp. 21–22

D.1 Setup

• The two-layer ReLU setup uses a deterministic linear target and Gaussian inputs and weights, with expectations over the weights.The input is x ∼ N(0, Id0), and the entries of W1 and W2 are i.i.d. N(0, 1).Source: §D.1 · p. 23, §D.2 · p. 23

D.3 Main results restated

• E[G2G2⊤] = αId2 + (d1/4)ΦΦ⊤, where α = EW1[tr(H2)].Here H is the hidden-activation second-moment matrix and α = EW1[tr(H2)].Source: §D.3 · p. 23
• For W1, the expected gradient self outer-product is permutation-invariant in the hidden index, with stated diagonal and off-diagonal formulas.The supplied result passages also give high-dimensional limits for κ1 and κ2: 1/(4π²) + 1/16 and 1/(4π), respectively.Source: §D.3 · p. 23

D.7 Discussion

• For W2, the leading-order diagonal contribution comes entirely from tr(H2), while the target term is lower order in d1.The passage attributes the leading contribution to a term depending only on W1.Source: §D.7 · p. 27
• Unlike the deep linear case, the W1 off-diagonal dominance is asymptotic in d1 rather than exact, which the discussion attributes to ReLU’s lack of odd symmetry.The discussion contrasts this with a deep-linear symmetry that forces off-diagonal expectations to zero.Source: §D.7 · p. 27

E.1 Setup

• The attention setup defines query, key, value, and output weights, with mutually independent Gaussian weight matrices and Gaussian input tokens.The target Φ is deterministic under the stated assumption.Source: §E.1 · p. 27, §E.2 · p. 27

E.3 Main results restated

• The hidden dimensions d_v and d_0 drive diagonal dominance, while d_k enters the stated quantities only through bounded factors; nonzero off-diagonal terms have relative size O(1/d_k).The d_k dependence enters through the ratio A2/A1; the off-diagonal rate applies when a ≠ b and (ΦΦ⊤)_ab ≠ 0.Source: §E.3 · p. 28

E.4 Preliminary lemmas

• The Wishart moment bounds give γ = Θ(d_0^3), A1 and A2 = Θ(d_0^6), and A3 = A1 + A2.The displayed moment orders establish these scalings as d_0 grows.Source: §E.4 · p. 30

E.5 Proofs

• Conditional on X, the expected gradient self outer-products for Q, K, and V are isotropic, with equal diagonal entries and zero off-diagonal entries.The proof uses sign changes to eliminate off-diagonal entries and permutations to show the diagonal entries are equal.Source: §E.5 · pp. 31–32

D GV (W), and (DWV, WOD)

• The conditional calculations provide closed forms for the WO gradient self outer-product and isotropic constants, which are then averaged over X using Wishart moments.The proof averages the conditional expressions to obtain the stated order for the averaged WO constant.Source: D GV (W), and (DWV , WOD) · pp. 32, 34

D GV (W), and (DWV , WOD)

• The WO ratio’s d_k dependence is confined to factors bounded above and below, so it does not change the ratio’s scaling with d_k.Source: D GV (W), and (DWV , WOD) · p. 35

E.6 Discussion

• For query, key, and value weights, isotropy follows from weight-space symmetries for any dataset, differentiable loss, symmetric initialization, and multi-head attention, but RoPE is not covered.For these matrices, the expected preconditioner is proportional to the identity, so orthogonalization reduces to scalar normalization.Source: §E.6 · p. 35
• The exact zeros arise from weight-space symmetries per dataset, while Gaussian tokens only translate moment constants; data-dependent initialization can break the symmetry argument.Independence of the data from initialization is what the argument genuinely requires.Source: §E.6 · p. 35
• When min(T,d_0)=1, c_K=c_Q; for 2≤T,d_0≤512, c_K>c_Q for all d_2 and d_v.Source: §E.6 · p. 35
• Under fan-in scaling, the WO ratio is Θ(1) in d_k, d_v, and d_0, with its size set by sequence length through the data moments.Deng et al. (2026) observed near-diagonal GG⊤ along GPT-2 and LLaMA trajectories, consistent with sequence length rather than hidden dimension driving diagonalization under this scaling.Source: §E.6 · p. 36

F Asymptotic Approximation between Orthogonalization and Row Normalization

• At Gaussian initialization, Muon’s Θ(mn min(m,n)) orthogonalization collapses to Θ(mn) row-wise ℓ2 normalization, exactly in expectation for attention query, key, and value weights.Deng et al. (2026) report that the equivalence holds throughout GPT-2 and LLaMA training, with the dominance ratio growing with model size; extending the analysis beyond initialization and to deep nonlinear networks is left for future work.Source: F Asymptotic Approximation between… · p. 36

G Simulation of Theoretical Results

• The evaluation estimates expected gradient self outer-products and tracks diagonal magnitude, off-diagonal magnitude, and their ratio as relevant dimensions vary.The estimates average gradients from independent initializations.Source: G Simulation of Theoretical Results · p. 36
• The experiments examine both representative configurations and dimension sweeps for each model.The figure caption specifies a symmetric matrix factorization example with d = 25, k = 80, and a fixed rank-5 target, estimated over 10,000 initializations.Source: G Simulation of Theoretical Results · p. 36, Figure 7 · p. 37
• The Figure 7 heatmap is normalized by its empirical mean diagonal magnitude and uses a logarithmic color scale.It depicts the empirical expected gradient self outer-product for the stated symmetric matrix factorization configuration.Source: Figure 7 · p. 37

G.1 Symmetric Matrix Factorization

• For symmetric matrix factorization, D = Θ(k^3), OD = Θ(k^2), and R = Θ(k), consistent with simulations showing the diagonal magnitude outgrowing the off-diagonal magnitude.The empirical sweep varies k with d = 25 and a fixed target; a representative expected Gram is strongly diagonal.Source: §G.1 · pp. 37–38

G.2 Deep Linear Network

• In the depth-4 linear network, hidden-layer expected off-diagonal entries are zero, while the output-layer ratio scales as R = Θ(d_j) when hidden dimension d_j varies.For output-layer gradients, D = Θ(d_j^2) and OD = Θ(d_j); for other layers, the diagonal scales as Θ(d_j) for the corresponding gradient and Θ(d_j^2) for the others.Source: §G.2 · pp. 38–39

G.3 Two-Layer ReLU Network

• For both layers of the two-layer ReLU network, the diagonal-dominance ratio grows as R = Θ(d_1), while expected off-diagonal entries remain nonzero.As d_1 increases, G1 has D = Θ(d_1) and OD = Θ(1); G2 has quadratic diagonal and linear off-diagonal growth.Source: §G.3 · pp. 40–42

G.4 Single-Layer Attention

• In single-layer attention, expected off-diagonal entries are exactly zero for Q, K, and V, while the output-projection ratio grows with d_v and d_0.For G_O, R = Θ(d_v d_0^3); increasing d_k affects the ratio only through a bounded factor.Source: §G.4 · pp. 43–44

H Empirical Evaluation on Large Language Models

• The study tests both projection-Gram structure at released model widths and width scaling with head dimension held fixed.The experiments use the models’ native initialization schemes and fixed FineWeb-Edu sequences.Source: H Empirical Evaluation on Large… · p. 44, §H.2 · p. 46

H.1 Gradient Gram Structure at Initialization

• All four attention projections show a clear diagonal component in Qwen3.5-0.8B-Base, Granite-4.0-1B-Base, and Llama-3.2-1B at random initialization.The empirical Grams use 25 random initializations and 64 fixed FineWeb-Edu sequences.Source: §H.1 · p. 46

H.2 Model Width Scaling

• The empirical diagonal-dominance ratio increases with model width for Q, K, V, and O across all three models.The increase is strongest for Q and K in Qwen, and for V and O in Granite and Llama.Source: §H.2 · p. 46, Figure 16 · p. 48