arxiv: 2605.06316 · v1 · submitted 2026-05-07 · 💻 cs.LG · cs.AI

Recognition: unknown

Pro-KLShampoo: Projected KL-Shampoo with Whitening Recovered by Orthogonalization

Ruotong Sun , Ermin Wei

Authors on Pith no claims yet

Pith reviewed 2026-05-08 12:59 UTC · model grok-4.3

classification 💻 cs.LG cs.AI

keywords KL-ShampooKronecker preconditioningorthogonalizationLLM pre-trainingsecond-order optimizationprojected optimizerspike-and-flat spectrumgradient covariance

0 comments

The pith

KL-Shampoo preconditioners recover their full algebraic form after projection onto a low-rank spike-and-flat subspace plus orthogonalization of the tail.

A machine-rendered reading of the paper's core claim, the machinery that carries it, and where it could break.

The paper establishes that the Kronecker factors inside KL-Shampoo display eigenvalue spectra with a few large values followed by a nearly uniform tail, a pattern exact under a rank-rho signal-plus-noise model of the gradients. By keeping the full spectral decomposition only on an r-dimensional tracked subspace and assigning one shared eigenvalue to the remaining directions, then applying orthogonalization to the momentum in those directions, the resulting preconditioner matches the original KL-Shampoo expression exactly. This projection lowers the cost of maintaining and applying the second-order information. Experiments across GPT-2 and LLaMA models from 124M to 450M parameters show the projected method reaching better validation loss, lower peak memory, and shorter wall-clock time to each loss target than full KL-Shampoo at every tested subspace rank. The approach therefore offers a concrete route to scale explicit matrix preconditioning to larger pre-training runs.

Core claim

KL-Shampoo forms its preconditioner by minimizing KL divergence to a Kronecker product of two factors estimated from gradient statistics. These factors exhibit a spike-and-flat eigenvalue spectrum across layers and training stages. The method replaces one factor with a parametric family that stores the complete spectral structure on an r-dimensional subspace and a single shared eigenvalue on the orthogonal complement. Orthogonalization of the momentum vector on the complement then restores the exact algebraic form of the original full-rank KL-Shampoo preconditioner. The resulting algorithm therefore operates at reduced rank while preserving the preconditioning effect of the unprojected rule.

What carries the argument

The orthogonalization identity that recovers the full KL-Shampoo preconditioner from its projected spike-and-flat Kronecker factor.

If this is right

Pro-KLShampoo produces lower validation loss than KL-Shampoo at every subspace rank on GPT-2 124M, 350M, LLaMA 134M, and 450M pre-training runs.
Peak per-GPU memory drops because only the r-dimensional spectral structure is stored and applied explicitly.
Wall-clock time to reach each validation-loss milestone decreases relative to the full method.
The performance edge remains consistent across the four model scales examined.

Where Pith is reading between the lines

These are editorial extensions of the paper, not claims the author makes directly.

The same spike-and-flat projection plus orthogonalization step could be applied to other Kronecker-factored second-order methods that display comparable spectral structure.
Adaptive choice of the subspace dimension r based on running estimates of gradient rank could further reduce overhead without manual tuning.
The recovery identity supplies a template for mixing explicit low-rank preconditioning with momentum orthogonalization in other optimization settings that involve matrix-valued statistics.
Similar spectral assumptions may hold for gradient covariances in domains beyond language-model pre-training, opening the projection technique to vision or reinforcement-learning models.

Load-bearing premise

The Kronecker factors inside KL-Shampoo exhibit eigenvalue spectra with a few dominant values followed by an approximately flat tail.

What would settle it

If the measured eigenvalues of a KL-Shampoo Kronecker factor during training of any tested model deviate from the spike-and-flat pattern, the algebraic recovery would fail and the reported gains in loss, memory, and time would not appear.

Figures

Figures reproduced from arXiv: 2605.06316 by Ermin Wei, Ruotong Sun.

**Figure 1.** Figure 1: Eigenvalue spectra of practical version of KL-Shampoo’s Kronecker preconditioners on GPT view at source ↗

**Figure 2.** Figure 2: Validation loss versus time for KL-Shampoo and Pro-KLShampoo with view at source ↗

**Figure 3.** Figure 3: Validation loss versus training step on GPT-2 124M (left) and LLaMA 134M (right) for Pro view at source ↗

**Figure 4.** Figure 4: Training loss versus training step on GPT-2 (124M, left; 350M, right) for KL-Shampoo and Pro view at source ↗

**Figure 5.** Figure 5: Training loss versus training step on LLaMA (134M, left; 450M, right) for KL-Shampoo and view at source ↗

**Figure 6.** Figure 6: Eigenvalue spectra of ΦLt on GPT-2 (124M), normalized by the tail mean (vertical dashed line at r = 128). The spike-and-flat shape is present across all layer types, depths, and training stages, supporting the conjecture in §3.2, though less pronounced than the corresponding spectra of full KL-Shampoo’s preconditioner ( view at source ↗

**Figure 7.** Figure 7: Eigenvalue spectra of KL-Shampoo’s Kronecker preconditioners on LLaMA, normalized by the view at source ↗

**Figure 8.** Figure 8: Validation loss versus wallclock time for Muon and Pro-KLShampoo at view at source ↗

read the original abstract

Optimizers that exploit the matrix structure of gradients are central to modern LLM pre-training, with two distinct frontiers: explicit Kronecker-factored preconditioning -- most recently KL-Shampoo, which estimates the preconditioner via KL divergence minimization -- and orthogonalization of the gradient momentum, exemplified by Muon and analyzed as steepest descent under the spectral norm. The two routes are typically developed in isolation. We make a structural observation about KL-Shampoo's Kronecker preconditioners: their eigenvalue spectra exhibit a \emph{spike-and-flat} shape -- a few dominant eigenvalues followed by an approximately uniform tail -- across layers and training stages, holding exactly under a rank-$\rho$ signal-plus-noise gradient model. We exploit this structure by restricting one of KL-Shampoo's Kronecker factors to a parametric family aligned with the spike-and-flat shape: full spectral structure on a tracked $r$-dimensional subspace, single shared eigenvalue across the remaining $n-r$ directions. On these directions, we apply orthogonalization. An identity shows that this orthogonalization recovers the algebraic form of full KL-Shampoo's preconditioner. On four pre-training scales (GPT-2 124M / 350M, LLaMA 134M / 450M), Pro-KLShampoo consistently outperforms KL-Shampoo at every subspace rank we test in validation loss, peak per-GPU memory, and wallclock time to reach each loss level.

Editorial analysis

A structured set of objections, weighed in public.

Desk editor's note, referee report, simulated authors' rebuttal, and a circularity audit. Tearing a paper down is the easy half of reading it; the pith above is the substance, this is the friction.

Desk Editor's Note private letter to a colleague

Pro-KLShampoo projects KL-Shampoo factors onto a spike-and-flat parametric form and recovers the preconditioner via orthogonalization, with consistent empirical gains on small LLMs, but the gains rest on an unverified spectral assumption.

read the letter

The paper's main move is to notice that KL-Shampoo's Kronecker factors often show a few large eigenvalues followed by a roughly flat tail. They restrict one factor to full detail on an r-dimensional subspace and a single shared eigenvalue on the complement, then apply orthogonalization to the tail. An algebraic identity is supposed to make the result match the original KL-Shampoo preconditioner exactly. This is the concrete link they draw between Kronecker preconditioning and Muon-style orthogonalization methods. On GPT-2 124M/350M and LLaMA 134M/450M they report lower validation loss, lower peak memory, and faster time to target loss than plain KL-Shampoo at every tested rank. That is a practical result worth noting for anyone running pre-training at these scales. The structural observation and the projection that exploits it are the actual new pieces. The experiments are run at multiple model sizes and show the same pattern, which gives the claim some weight. The soft spot is the spike-and-flat assumption itself. If the real eigenvalue tails are not close to uniform or if the number of dominant directions shifts across layers and training steps, the restriction stops being a structure-preserving projection and becomes an ordinary low-rank approximation. The paper ties the shape to a rank-ρ signal-plus-noise model, but without direct checks on the actual gradient statistics the justification stays thin. The identity is stated in the abstract; the full derivation would need to be checked for exactness rather than approximation. This is aimed at people working on memory-efficient second-order methods for language model training. It has enough empirical consistency and a clear technical angle to go to a serious referee, though the review would likely press on validating the spectral claim and tightening the algebra. I would send it to review.

Referee Report

3 major / 1 minor

Summary. The manuscript proposes Pro-KLShampoo, a projected variant of KL-Shampoo that exploits an observed spike-and-flat eigenvalue structure in the Kronecker factors. One factor is restricted to a parametric family with full spectral structure on a tracked r-dimensional subspace and a single shared eigenvalue on the orthogonal complement; orthogonalization is then applied to the tail. An algebraic identity is claimed to recover the exact preconditioner of full KL-Shampoo. Experiments on GPT-2 (124M/350M) and LLaMA (134M/450M) report consistent gains in validation loss, peak per-GPU memory, and wall-clock time to target loss across tested subspace ranks.

Significance. If the spike-and-flat observation and the recovery identity are general and exact, the work would usefully bridge explicit Kronecker preconditioning with orthogonalization-based methods, potentially improving both theoretical understanding and practical efficiency for large-scale LLM training. The multi-scale empirical results are a positive signal, but the absence of the central derivation limits the assessed significance.

major comments (3)

[Abstract] Abstract: the claim that the eigenvalue spectra exhibit a spike-and-flat shape 'holding exactly under a rank-ρ signal-plus-noise gradient model' is the load-bearing justification for the parametric restriction; no derivation, proof, or empirical verification of this model appears in the provided text, so the restriction risks being an arbitrary low-rank approximation rather than a structure-exploiting projection.
[Abstract] Abstract: the statement that 'an identity shows that this orthogonalization recovers the algebraic form of full KL-Shampoo's preconditioner' is central to attributing any gains to the claimed mechanism rather than generic projection; the identity is asserted without derivation, equation, or section reference, preventing verification that it is exact and independent of r.
[Experiments] Experiments section: consistent outperformance is reported across four model scales and every tested subspace rank, yet the exact definition of the parametric family, the procedure for maintaining the tracked r-dimensional subspace, and the choice of r relative to layer dimensions are not specified, rendering the results non-reproducible and the contribution of the identity unisolated.

minor comments (1)

[Abstract] Abstract: the range of subspace ranks r tested and their relation to the dimensions of the Kronecker factors are not stated, which would aid interpretation of the memory and time gains.

Simulated Author's Rebuttal

3 responses · 0 unresolved

We thank the referee for the careful reading and constructive major comments. The points raised highlight omissions in the submitted version that limit verifiability. We address each comment below and will incorporate the requested derivations, equations, and implementation details into the revised manuscript.

read point-by-point responses

Referee: [Abstract] Abstract: the claim that the eigenvalue spectra exhibit a spike-and-flat shape 'holding exactly under a rank-ρ signal-plus-noise gradient model' is the load-bearing justification for the parametric restriction; no derivation, proof, or empirical verification of this model appears in the provided text, so the restriction risks being an arbitrary low-rank approximation rather than a structure-exploiting projection.

Authors: We agree that the derivation was omitted from the initial submission. Under the rank-ρ signal-plus-noise model, the gradient matrix is G = S + N with rank(S) = ρ and N isotropic; the resulting Kronecker factors of the second-moment matrix then possess exactly one dominant eigenvalue (spike) aligned with the signal subspace and identical eigenvalues on the orthogonal complement (flat tail). We will add a self-contained derivation (with all intermediate steps) as a new subsection in the revised Theory section, together with empirical spectral plots from the GPT-2 and LLaMA runs that confirm the structure persists across layers and training stages. This addition will make explicit that the parametric restriction is structure-exploiting rather than arbitrary. revision: yes
Referee: [Abstract] Abstract: the statement that 'an identity shows that this orthogonalization recovers the algebraic form of full KL-Shampoo's preconditioner' is central to attributing any gains to the claimed mechanism rather than generic projection; the identity is asserted without derivation, equation, or section reference, preventing verification that it is exact and independent of r.

Authors: The referee correctly notes that the identity was stated without supporting material. The identity follows because, once the tail eigenvalues are forced equal by the parametric form, orthogonalization (i.e., replacing the tail block with its spectral-norm or whitening factor) algebraically reproduces the full KL-Shampoo preconditioner for any choice of r. We will insert the complete derivation, including the key matrix identities and the proof that the result is independent of r, into the revised manuscript (new subsection immediately following the parametric-form definition). This will allow direct verification that performance gains are attributable to the claimed recovery mechanism. revision: yes
Referee: [Experiments] Experiments section: consistent outperformance is reported across four model scales and every tested subspace rank, yet the exact definition of the parametric family, the procedure for maintaining the tracked r-dimensional subspace, and the choice of r relative to layer dimensions are not specified, rendering the results non-reproducible and the contribution of the identity unisolated.

Authors: We acknowledge that the submitted Experiments section lacked the necessary implementation details. In the revision we will expand this section and add an appendix containing: (i) the precise parametric family (full diagonal spectrum on the tracked r-dimensional subspace, single shared scalar eigenvalue on the orthogonal complement); (ii) the subspace-tracking procedure (periodic top-r SVD or power-iteration updates on the accumulated second-moment matrix, performed every 100 steps); and (iii) the concrete choices of r (tested values r = 8, 16, 32, … together with the rule r = 0.25 × min(d1, d2) used for each layer). Pseudocode and hyper-parameter tables will also be supplied. These additions will render the experiments fully reproducible and isolate the contribution of the recovery identity. revision: yes

Circularity Check

0 steps flagged

No significant circularity; derivation is self-contained via empirical observation and algebraic identity.

full rationale

The paper grounds its construction in an empirical structural observation (spike-and-flat spectra across layers and stages) that is stated to hold exactly under an explicit rank-ρ signal-plus-noise gradient model, then applies a parametric restriction to one Kronecker factor and invokes an algebraic identity to show equivalence of the resulting preconditioner to the full KL-Shampoo form. This identity is presented as a mathematical equivalence under the stated restriction rather than a tautology derived from fitted outputs, and all performance claims are supported by direct empirical comparisons on GPT-2 and LLaMA scales at multiple ranks. No load-bearing self-citations, fitted inputs renamed as predictions, or ansatzes smuggled via prior work appear in the derivation chain; the central result remains independently verifiable against external benchmarks.

Axiom & Free-Parameter Ledger

1 free parameters · 1 axioms · 0 invented entities

The central claim rests on one domain assumption about gradient spectra and one tunable hyperparameter for subspace dimension; no new entities are postulated.

free parameters (1)

subspace rank r
Dimension of the tracked subspace kept with full spectral structure; selected per experiment and model scale.

axioms (1)

domain assumption Eigenvalue spectra of KL-Shampoo's Kronecker preconditioners exhibit a spike-and-flat shape across layers and training stages, holding exactly under a rank-ρ signal-plus-noise gradient model.
Invoked to justify restricting one Kronecker factor to the parametric family.

pith-pipeline@v0.9.0 · 5559 in / 1424 out tokens · 73949 ms · 2026-05-08T12:59:38.357782+00:00 · methodology

discussion (0)

Reference graph

Works this paper leans on

27 extracted references · 20 canonical work pages · 4 internal anchors

[1]

Asgo: Adaptive structured gradient optimization.arXiv preprint arXiv:2503.20762, 2025

Kang An, Yuxing Liu, Rui Pan, Yi Ren, Shiqian Ma, Donald Goldfarb, and Tong Zhang. Asgo: Adaptive structured gradient optimization.arXiv preprint arXiv:2503.20762,

work page arXiv
[2]

arXiv preprint arXiv:2002.09018 , year=

Rohan Anil, Vineet Gupta, Tomer Koren, Kevin Regan, and Yoram Singer. Scalable second order optimization for deep learning.arXiv preprint arXiv:2002.09018,

work page arXiv 2002
[3]

arXiv:2410.21265 , year=

Jeremy Bernstein and Laker Newhouse. Modular duality in deep learning.arXiv preprint arXiv:2410.21265, 2024a. Jeremy Bernstein and Laker Newhouse. Old optimizer, new norm: An anthology.arXiv preprint arXiv:2409.20325, 2024b. Jeremy Bernstein, Yu-Xiang Wang, Kamyar Azizzadenesheli, and Animashree Anandkumar. signsgd: Compressed optimisation for non-convex ...

work page arXiv
[4]

Purifying shampoo: Investigating shampoo’s heuristics by decomposing its preconditioner.arXiv preprint arXiv:2506.03595,

Runa Eschenhagen, Aaron Defazio, Tsung-Hsien Lee, Richard E Turner, and Hao-Jun Michael Shi. Purifying shampoo: Investigating shampoo’s heuristics by decomposing its preconditioner.arXiv preprint arXiv:2506.03595,

work page arXiv
[5]

Clarifying shampoo: Adapting spectral descent to stochasticity and the parameter trajectory.arXiv preprint arXiv:2602.09314,

Runa Eschenhagen, Anna Cai, Tsung-Hsien Lee, and Hao-Jun Michael Shi. Clarifying shampoo: Adapting spectral descent to stochasticity and the parameter trajectory.arXiv preprint arXiv:2602.09314,

work page arXiv
[6]

Gradient Descent Happens in a Tiny Subspace

10 Guy Gur-Ari, Daniel A Roberts, and Ethan Dyer. Gradient descent happens in a tiny subspace.arXiv preprint arXiv:1812.04754,

work page Pith review arXiv
[7]

Subspace optimization for large language models with convergence guarantees.arXiv preprint arXiv:2410.11289, 2024

Yutong He, Pengrui Li, Yipeng Hu, Chuyan Chen, and Kun Yuan. Subspace optimization for large language models with convergence guarantees.arXiv preprint arXiv:2410.11289,

work page arXiv
[8]

arXiv preprint arXiv:2407.11239 , year=

Ajay Jaiswal, Yifan Wang, Lu Yin, Shiwei Liu, Runjin Chen, Jiawei Zhao, Ananth Grama, Yuandong Tian, and Zhangyang Wang. From low rank gradient subspace stabilization to low-rank weights: Observations, theories, and applications. arXiv preprint arXiv:2407.11239,

work page arXiv
[9]

Adam: A Method for Stochastic Optimization

URL https: //github.com/KellerJordan/modded-nanogpt. Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization.arXiv preprint arXiv:1412.6980,

work page internal anchor Pith review arXiv
[10]

Understanding and improving shampoo and soap via kullback-leibler minimization.arXiv preprint arXiv:2509.03378,

Wu Lin, Scott C Lowe, Felix Dangel, Runa Eschenhagen, Zikun Xu, and Roger B Grosse. Understanding and improving shampoo and soap via kullback-leibler minimization.arXiv preprint arXiv:2509.03378,

work page arXiv
[11]

COSMOS: A hybrid adaptive optimizer for memory-efficient training of LLMs.arXiv preprint arXiv:2502.17410, 2025

Liming Liu, Zhenghao Xu, Zixuan Zhang, Hao Kang, Zichong Li, Chen , Weizhu Chen, and Tuo Zhao. Cosmos: A hybrid adaptive optimizer for memory-efficient training of llms.arXiv preprint arXiv:2502.17410,

work page arXiv
[12]

Decoupled Weight Decay Regularization

Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization.arXiv preprint arXiv:1711.05101,

work page internal anchor Pith review arXiv
[13]

A new perspective on shampoo’s preconditioner.arXiv preprint arXiv:2406.17748,

Depen Morwani, Itai Shapira, Nikhil Vyas, Eran Malach, Sham Kakade, and Lucas Janson. A new perspective on shampoo’s preconditioner.arXiv preprint arXiv:2406.17748,

work page arXiv
[14]

Subtrack++: Gradient subspace tracking for scalable llm training

Sahar Rajabi, Nayeema Nonta, and Sirisha Rambhatla. Subtrack++: Gradient subspace tracking for scalable llm training. arXiv preprint arXiv:2502.01586,

work page arXiv
[15]

A distributed data-parallel PyTorch imple- mentation of the distributed Shampoo optimizer for training neural networks at-scale.arXiv preprint arXiv:2309.06497, 2023

Hao-Jun Michael Shi, Tsung-Hsien Lee, Shintaro Iwasaki, Jose Gallego-Posada, Zhijing Li, Kaushik Rangadurai, Dheevatsa Mudigere, and Michael Rabbat. A distributed data-parallel pytorch implementation of the distributed shampoo optimizer for training neural networks at-scale.arXiv preprint arXiv:2309.06497,

work page arXiv
[16]

LLaMA: Open and Efficient Foundation Language Models

Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, et al. Llama: Open and efficient foundation language models.arXiv preprint arXiv:2302.13971,

work page internal anchor Pith review arXiv
[17]

SOAP: Improving and Stabilizing Shampoo using Adam

Nikhil Vyas, Depen Morwani, Rosie Zhao, Mujin Kwun, Itai Shapira, David Brandfonbrener, Lucas Janson, and Sham Kakade. Soap: Improving and stabilizing shampoo using adam.arXiv preprint arXiv:2409.11321,

work page internal anchor Pith review arXiv
[18]

Structured pre- conditioners in adaptive optimization: A unified analysis.arXiv preprint arXiv:2503.10537, 2025

Shuo Xie, Tianhao Wang, Sashank Reddi, Sanjiv Kumar, and Zhiyuan Li. Structured preconditioners in adaptive optimization: A unified analysis.arXiv preprint arXiv:2503.10537,

work page arXiv
[19]

Mousse: Rectifying the geometry of muon with curvature-aware precondition- ing, 2026

Yechen Zhang, Shuhao Xing, Junhao Huang, Kai Lv, Yunhua Zhou, Xipeng Qiu, Qipeng Guo, and Kai Chen. Mousse: Rectifying the geometry of muon with curvature-aware preconditioning.arXiv preprint arXiv:2603.09697,

work page arXiv
[20]

Zhao et al

Jiawei Zhao, Zhenyu Zhang, Beidi Chen, Zhangyang Wang, Anima Anandkumar, and Yuandong Tian. Galore: Memory- efficient llm training by gradient low-rank projection.arXiv preprint arXiv:2403.03507,

work page arXiv
[21]

SOAP [Vyas et al., 2024] improved its practical efficiency by running Adam in Shampoo’s eigenbasis

12 A Related work Kronecker-factored preconditioning.Shampoo [Gupta et al., 2018] introduced Kronecker-factored precon- ditioning as a tractable approximation to full-matrix Adagrad [Duchi et al., 2011]; distributed implementa- tions [Anil et al., 2020, Shi et al., 2023] scaled it to large models, with the Shampoo submission winning the AlgoPerf training-...

2018
[22]

0 5 10 15 20 25Right λi / ¯λtail self attn.q proj (768×768) self attn.k proj (768×768) self attn.v proj (768×768) self attn.o proj (768×768) mlp.gate proj (2048×768) mlp.down proj (768×2048) 0 5 10 15 20 25Left λi / ¯λtail Eigenvalue percentile, 0% – 100% (left to right) Layer 2 Layer 6 Layer 11 Step 1500 Step 3500 Tail mean (=1) r = 128 Figure 7: Eigenva...

2048
[23]

Recall that f(U) = log det(U ⊤ΦL U) + (n−r) log Tr(P⊥ΦL) is the reduced objective from (32)

This holds since U ∗⊤u⊥ = 0 (because u⊥ ⊥range(U ∗)). Recall that f(U) = log det(U ⊤ΦL U) + (n−r) log Tr(P⊥ΦL) is the reduced objective from (32). Along the pathU(t), define A(t) :=U(t) ⊤ ΦL U(t), s(t) := Tr(Φ L)−TrA(t), so that f(U(t)) = log detA(t) + (n−r) logs(t) . At t= 0 we have A(0) =A=U ∗⊤ΦL U ∗ and s(0) = (n−r)/β. We will show that d2 dt2 f(U(t)) ...

2006
[24]

For the noise contribution, write the(i, i ′)entry explicitly: ξ(R∗)−1ξ⊤ i,i′ = X j,j′ ξi,j (R∗)−1 j,j′ ξi′,j′

Substitute G=AB ⊤ +ξ into the L stationarity (2) and use E[ξ] = 0 to eliminate the cross-product terms: E[G(R ∗)−1 G⊤] =AB ⊤(R∗)−1BA⊤ +E[ξ(R ∗)−1ξ⊤]. For the noise contribution, write the(i, i ′)entry explicitly: ξ(R∗)−1ξ⊤ i,i′ = X j,j′ ξi,j (R∗)−1 j,j′ ξi′,j′. 21 Taking expectation, the uncorrelated-noise assumption E[ξijξi′j′] =σ 21(i,j)=(i′,j′) collaps...

2001
[25]

Similarly, ∥G⊥∥op =∥G(I n −U U ⊤)∥op ≤G max

Since U has orthonormal columns, ∥eG∥op =∥GU∥ op ≤ ∥G∥op ∥U∥op ≤G max (Assumption (iv)). Similarly, ∥G⊥∥op =∥G(I n −U U ⊤)∥op ≤G max. The clip also gives µ−1 ⊥ ≤C 2 and λ−1 L,i ≤C 2 for all i. Combine all the bounds: ∥∆L∥op ≤ 1 n G2 max C2 +C 2 G2 max = 2C2G2 max n ≤C 2G2 max ≤Θ, ∥∆S∥op ≤ 1 m G2 max C2 ≤C 2G2 max ≤Θ, δ⊥ ≤ C2 G2 max (n−r) m(n−r) = C2G2 max...

2024
[26]

Theorem 2 below, which has no floor after step-size balancing)

The irreducible noise floor 2ca √ k σF is linear in σF , a feature inherited from the sign-SGD–style argument [Bernstein et al., 2018] underlying the orthogonalization identity; this differs from the quadratic-in-σF floor in analyses where 27 the update is a positive-definite linear function of the gradient (cf. Theorem 2 below, which has no floor after s...

2018
[27]

Pro-KLShampoo reaches Muon’s final validation loss in approximately 10%less wallclock time at 134M and8%less at 450M

0 50 100 150 200 250 300 350 Minutes 3.05 3.10 3.15 3.20 3.25 3.30 Validation loss ∼10% wallclock saving Llama-134M 0 200 400 600 800 1000 1200 Minutes 2.75 2.80 2.85 2.90 2.95 3.00 ∼8% wallclock saving Llama-450M Muon Pro-KLShampoo ( r=32) Pro-KLShampoo ( r=64) Pro-KLShampoo ( r=128) Figure 8: Validation loss versus wallclock time for Muon and Pro-KLSham...

2025