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

TL;DR

Pro-KLShampoo introduces a structured approach combining KL-Shampoo and orthogonalization, exploiting spectral properties to improve large language model pre-training efficiency and effectiveness.

Contribution

It reveals the eigenvalue spectrum structure of KL-Shampoo's preconditioners and leverages this to develop a hybrid optimizer with better performance.

Findings

Pro-KLShampoo outperforms KL-Shampoo in validation loss across multiple models.

It reduces memory usage and training time compared to standard KL-Shampoo.

The method is effective across different model scales and training stages.

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.