🤖 machine learning

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

Pro-KLShampoo is a novel optimizer that enhances KL-Shampoo by exploiting the observed "spike-and-flat" eigenvalue structure of its preconditioners to combine full spectral tracking on a low-dimensional subspace with orthogonalization on the remaining directions, achieving superior performance in validation loss, memory efficiency, and training speed across multiple LLM scales.

Original authors: Ruotong Sun, Ermin Wei

Published 2026-05-08

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Ruotong Sun, Ermin Wei

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to teach a giant, complex robot (a Large Language Model) to speak human language. To do this, you show it millions of examples and let it make mistakes. Every time it makes a mistake, you give it a "nudge" (a gradient) to correct its path.

The problem is that these nudges are messy. Sometimes the robot needs a tiny, precise nudge in one direction, and a massive shove in another. If you just push it randomly, it learns slowly. To learn fast, you need a smart "preconditioner"—a tool that reshapes these nudges so they are perfectly balanced and efficient.

Two popular tools for this job are KL-Shampoo and Muon.

KL-Shampoo is like a master sculptor. It tries to carve the perfect shape for every single nudge. It's very accurate but requires a lot of heavy lifting (computational power and memory) to calculate the shape of every single piece of the puzzle.
Muon is like a gymnast. It doesn't try to reshape every piece; instead, it just straightens out the momentum of the nudges, making them move in a clean, orthogonal (right-angled) way. It's fast and light, but it might miss some of the subtle details that the sculptor catches.

The Big Discovery: The "Spike-and-Flat" Pattern
The authors of this paper looked closely at the "sculptor's" work (the KL-Shampoo preconditioner) and noticed a strange, beautiful pattern. They found that the "nudge-shapes" aren't random. Instead, they look like a spike-and-flat landscape:

The Spike: A few directions have huge, dominant values (like a few tall mountains).
The Flat: The rest of the directions are all roughly the same height (like a vast, flat plain).

This happens across all layers of the robot's brain, from the first layer to the last. It's as if the robot only really cares about a few specific directions, and everywhere else, it's just "flat" noise.

The Solution: Pro-KLShampoo
The authors built a new optimizer called Pro-KLShampoo (Projected KL-Shampoo). Think of it as a hybrid tool that gets the best of both worlds by using the "spike-and-flat" discovery.

Here is how it works, using a simple analogy:

The "VIP Lounge" (The Subspace):
The algorithm identifies the "spike" directions—the few important mountains. It puts these in a special "VIP Lounge" (a tracked subspace). Inside this lounge, it uses the heavy, precise sculpting tool (KL-Shampoo) to get the perfect shape for these few critical directions. It doesn't waste time on the rest.
The "Open Field" (The Complement):
For the "flat" directions (the rest of the landscape), it stops trying to sculpt every single pebble. Instead, it uses a clever trick called Orthogonalization (borrowed from the Muon tool).
- The Magic Trick: The authors proved mathematically that if you just "straighten out" (orthogonalize) the nudges in this flat area, it magically recovers the exact same algebraic result as if you had done the heavy sculpting there. It's like realizing that for a flat field, you don't need a map; you just need to walk in a straight line, and you'll end up in the exact same spot as if you had calculated every coordinate.

Why is this better?

Speed: By ignoring the heavy sculpting for the "flat" parts and just straightening them out, the algorithm runs much faster. It saves a lot of time on the computer's "wallclock" (real-world time).
Memory: It doesn't need to store the massive, complex shapes for the flat parts. It only stores the small "VIP" shapes and a single number for the rest. This saves a lot of computer memory.
Performance: In tests on different robot sizes (GPT-2 and LLaMA), Pro-KLShampoo learned faster and reached a lower error rate than the original KL-Shampoo, while using less memory.

In Summary
The paper says: "We found that the complex math behind modern AI training has a simple pattern: a few big spikes and a flat tail. We built a new tool that treats the spikes with extreme care and the flat tail with a simple, fast trick. This trick works just as well as the hard math but is much faster and cheaper to run."

The result is a smarter, faster way to train AI models that saves time and computer resources without sacrificing quality.

Technical Summary: Pro-KLShampoo

Problem Statement

Modern Large Language Model (LLM) pre-training relies heavily on optimizers that exploit the matrix structure of gradients. Two distinct frontiers have emerged:

Explicit Kronecker-factored preconditioning: Represented by KL-Shampoo, which estimates the preconditioner by minimizing the Kullback-Leibler (KL) divergence between the true gradient covariance and a Kronecker-structured approximation.
Orthogonalization of gradient momentum: Exemplified by Muon, which orthogonalizes the momentum to realize steepest descent under the spectral norm.

These approaches are typically developed in isolation. While KL-Shampoo offers strong theoretical grounding via KL minimization, it incurs high computational and memory costs due to maintaining full Kronecker factors. Conversely, orthogonalization methods like Muon are efficient but operate under a different geometric framework (spectral norm steepest descent) without the KL estimation perspective. The paper addresses the need to unify these approaches to achieve the statistical efficiency of KL-Shampoo with the computational efficiency of orthogonalization methods.

Methodology: Pro-KLShampoo

The authors propose Pro-KLShampoo (Projected KL-Shampoo with Whitening Recovered by Orthogonalization). The core innovation stems from a structural observation of KL-Shampoo's preconditioners: their eigenvalue spectra exhibit a "spike-and-flat" shape. This consists of a few dominant eigenvalues (the "spike") followed by an approximately uniform tail (the "flat"). This structure is observed empirically across layers and training stages and is theoretically exact under a rank- $\rho$ signal-plus-noise gradient model.

1. Structural Restriction

Pro-KLShampoo restricts one of KL-Shampoo's Kronecker factors (specifically the right-side factor $R$ for rectangular layers where $m \le n$ ) to a parametric family aligned with the spike-and-flat structure:

Projected Subspace ( $r$ -dimensional): The preconditioner retains full spectral structure on a tracked subspace spanned by $U \in \text{St}(n, r)$ .
Complement Subspace ( $n-r$ dimensions): The remaining directions are collapsed into a single shared eigenvalue $\mu_\perp$ .

This restriction reduces the storage of the right-side factor from $O(n^2)$ to $O(nr)$ and the QR decomposition cost from $O(n^3)$ to $O(r^3)$ .

2. Recovery via Orthogonalization

A naive restriction to a scalar tail loses information if the tail is not perfectly uniform. To recover the algebraic form of the full KL-Shampoo preconditioner on the complement subspace, Pro-KLShampoo applies orthogonalization (polar decomposition) to the complement gradient component.

The update rule is:
$\Delta W = -\alpha_{kl} L^{-1/2} G U S^{-1/2} U^\top - c_a \text{polar}\left(\mu_\perp^{-1/2} L^{-1/2} G P_\perp\right)$
Where:

The first term represents the full KL-Shampoo update within the tracked subspace.
The second term applies orthogonalization to the complement.
The authors prove an identity (Eq. 12) showing that this orthogonalization recovers the algebraic form of the full KL-Shampoo's per-direction whitening on the complement, effectively connecting Muon-style orthogonalization to the KL estimation framework.

3. Subspace Tracking

The optimal subspace $U$ is identified as the eigenspace of the matrix $\Phi_L = E[G^\top L^{-1} G]$ . Under the spike-and-flat conjecture, the top- $r$ eigenspace of $\Phi_L$ is the optimal choice. The algorithm tracks this subspace using a single step of power iteration (via QR decomposition) on a running estimate of $\Phi_L$ .

Key Contributions

Structural Observation & Parametric Restriction: The paper identifies the "spike-and-flat" eigenvalue structure in KL-Shampoo's preconditioners (empirically robust and theoretically exact under low-rank gradient models) and exploits it by restricting the KL objective to a parametric family.
Algebraic Recovery via Orthogonalization: It demonstrates that applying orthogonalization to the complement subspace recovers the algebraic form of the full KL-Shampoo preconditioner. This provides a theoretical bridge between Muon-style orthogonalization and KL estimation.
Convergence Guarantee: The authors provide a non-convex convergence guarantee at a rate of $O(1/\sqrt{T})$ for the idealized algorithm under operator-norm smoothness.
Empirical Performance: Extensive evaluation on GPT-2 (124M/350M) and LLaMA (134M/450M) models shows Pro-KLShampoo consistently outperforms KL-Shampoo in validation loss, peak per-GPU memory, and wallclock time to reach matched loss levels across all tested subspace ranks ( $r \in \{32, 64, 128\}$ ).

Experimental Results

The authors evaluated Pro-KLShampoo against AdamW, Muon, and KL-Shampoo on four pre-training configurations.

Validation Loss: Pro-KLShampoo achieved the lowest validation loss at every tested rank across all model scales.
Memory Efficiency: Peak per-GPU memory usage was lower than KL-Shampoo and comparable to Muon.
Wallclock Time: Pro-KLShampoo reached specific validation loss levels faster than KL-Shampoo.
- On GPT-2, the speedup was primarily driven by reduced per-step time (replacing $O(n^3)$ QR with $O(r^3)$ ).
- On LLaMA, the speedup was driven by faster convergence (lower training loss curves), in addition to per-step efficiency.
Ablation Studies:
- Removing orthogonalization (Smok-Hop variant) resulted in performance degradation, particularly at larger scales (LLaMA 450M), confirming the necessity of orthogonalization to recover per-direction whitening.
- Neither the subspace-only nor complement-only variants matched the performance of the full Pro-KLShampoo, indicating both components are essential.

Significance and Claims

The paper claims that Pro-KLShampoo successfully unifies two previously isolated optimization frontiers. By leveraging the intrinsic "spike-and-flat" structure of gradient covariances, it achieves the statistical efficiency of KL-Shampoo while utilizing orthogonalization to maintain computational efficiency.

The authors emphasize that the replacement of the scalar tail update with orthogonalization is not a heuristic but an algebraic recovery of the full KL-Shampoo form. This allows the method to retain the benefits of KL-based preconditioning (such as step-size robustness) while significantly reducing the computational overhead associated with full Kronecker factor maintenance. The work suggests that low-dimensional spectral structure is pervasive in deep network optimization and can be effectively exploited to design more efficient optimizers for LLM pre-training.

The paper remains modest regarding limitations, noting that the rank $r$ is set empirically, and future work is needed for per-layer adaptive calibration of the mixing weight $\alpha_{kl}$ and a convergence analysis of the practical variant involving Newton-Schulz iterations and Nesterov momentum.