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

arXiv: https://arxiv.org/abs/2605.06316 · HTML: https://mlsanity.com/paper/2605.06316/

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Builds on 30 of its 44 references that are in the map.
- First posted on arXiv 7 May 2026.

## Builds on (most cited first, 30 of 30)

- [Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer](https://mlsanity.com/paper/1910.10683/) (2019)
- [LLaMA: Open and Efficient Foundation Language Models](https://mlsanity.com/paper/2302.13971/) (2023)
- [The FineWeb Datasets: Decanting the Web for the Finest Text Data at Scale](https://mlsanity.com/paper/2406.17557/) (2024)
- [Optimizing Neural Networks with Kronecker-factored Approximate Curvature](https://mlsanity.com/paper/1503.05671/) (2015)
- [Stochastic First- and Zeroth-order Methods for Nonconvex Stochastic Programming](https://mlsanity.com/paper/1309.5549/) (2013)
- [signSGD: Compressed Optimisation for Non-Convex Problems](https://mlsanity.com/paper/1802.04434/) (2018)
- [GaLore: Memory-Efficient LLM Training by Gradient Low-Rank Projection](https://mlsanity.com/paper/2403.03507/) (2024)
- [Shampoo: Preconditioned Stochastic Tensor Optimization](https://mlsanity.com/paper/1802.09568/) (2018)
- [An Investigation into Neural Net Optimization via Hessian Eigenvalue Density](https://mlsanity.com/paper/1901.10159/) (2019)
- [Gradient Descent Happens in a Tiny Subspace](https://mlsanity.com/paper/1812.04754/) (2018)
- [Old Optimizer, New Norm: An Anthology](https://mlsanity.com/paper/2409.20325/) (2024)
- [SOAP: Improving and Stabilizing Shampoo using Adam](https://mlsanity.com/paper/2409.11321/) (2024)
- [Fast Approximate Natural Gradient Descent in a Kronecker-factored Eigenbasis](https://mlsanity.com/paper/1806.03884/) (2018)
- [A Distributed Data-Parallel PyTorch Implementation of the Distributed Shampoo Optimizer for Training Neural Networks At-Scale](https://mlsanity.com/paper/2309.06497/) (2023)
- [Modular Duality in Deep Learning](https://mlsanity.com/paper/2410.21265/) (2024)
- [Scalable Optimization in the Modular Norm](https://mlsanity.com/paper/2405.14813/) (2024)
- [A New Perspective on Shampoo's Preconditioner](https://mlsanity.com/paper/2406.17748/) (2024)
- [ASGO: Adaptive Structured Gradient Optimization](https://mlsanity.com/paper/2503.20762/) (2025)
- [Memory-Efficient LLM Training with Online Subspace Descent](https://mlsanity.com/paper/2408.12857/) (2024)
- [Subspace Optimization for Large Language Models with Convergence Guarantees](https://mlsanity.com/paper/2410.11289/) (2024)
- [Structured Preconditioners in Adaptive Optimization: A Unified Analysis](https://mlsanity.com/paper/2503.10537/) (2025)
- [COSMOS: A Hybrid Adaptive Optimizer for Memory-Efficient Training of LLMs](https://mlsanity.com/paper/2502.17410/) (2025)
- [Streaming PCA: Matching Matrix Bernstein and Near-Optimal Finite Sample Guarantees for Oja's Algorithm](https://mlsanity.com/paper/1602.06929/) (2016)
- [From Low Rank Gradient Subspace Stabilization to Low-Rank Weights: Observations, Theories, and Applications](https://mlsanity.com/paper/2407.11239/) (2024)
- [Purifying Shampoo: Investigating Shampoo's Heuristics by Decomposing its Preconditioner](https://mlsanity.com/paper/2506.03595/) (2025)
- [Streaming k-PCA: Efficient guarantees for Oja's algorithm, beyond rank-one updates](https://mlsanity.com/paper/2102.03646/) (2021)
- [Mousse: Rectifying the Geometry of Muon with Curvature-Aware Preconditioning](https://mlsanity.com/paper/2603.09697/) (2026)
- [SubTrack++ : Gradient Subspace Tracking for Scalable LLM Training](https://mlsanity.com/paper/2502.01586/) (2025)
- [Understanding and Improving Shampoo and SOAP via Kullback-Leibler Minimization](https://mlsanity.com/paper/2509.03378/) (2025)
- [Clarifying Shampoo: Adapting Spectral Descent to Stochasticity and the Parameter Trajectory](https://mlsanity.com/paper/2602.09314/) (2026)

## Built on by (most cited first, 1 of 1)

- [ShamAN-Q: Shampoo Augmented NanoQuant for Sub-1-bit LLM Weights](https://mlsanity.com/paper/2609.38521/) (2026)

## Abstract (arXiv)

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 spike-and-flat shape -- a few dominant eigenvalues followed by an approximately uniform tail -- across layers and training stages, holding exactly under a rank-ρ 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.