Pro-KLShampoo: Enhancing KL-Shampoo Optimizers

• Pro-KLShampoo is a family of optimizers that reformulate KL-Shampoo by minimizing the KL divergence to achieve structured whitening using projection methods.
• It reduces computational and memory overhead by employing low-rank approximations, QR-based orthogonalization, and BFloat16 storage reparameterization.
• Empirical evaluations indicate that Pro-KLShampoo variants enhance training stability and efficiency in neural network pretraining compared to traditional methods.

Searching arXiv for the specified Pro-KLShampoo papers and closely related KL-Shampoo work. Pro-KLShampoo denotes a family of optimizer formulations built on KL-Shampoo, a Shampoo variant derived from minimizing the Kullback–Leibler divergence between a true gradient covariance and a Kronecker-structured approximation. In the literature supplied here, the term refers specifically to two 2026 developments: a projected, orthogonalization-based optimizer titled “Pro-KLShampoo: Projected KL-Shampoo with Whitening Recovered by Orthogonalization” (Sun et al., 7 May 2026), and a reparametrization framework for Shampoo-based methods titled “Reparametrizing Shampoo and SOAP for Subspace Basis Updates and BFloat16 Storage” (Milligan et al., 25 May 2026). Both works take KL-Shampoo as their point of departure, whose original formulation reframes Shampoo as covariance estimation under Gaussian assumptions and motivates a two-sided update through KL minimization rather than Frobenius-norm analysis (Lin et al., 3 Sep 2025). Across these papers, Pro-KLShampoo is characterized by attempts to preserve KL-Shampoo’s structured whitening behavior while reducing instability, memory cost, QR overhead, or precision sensitivity.

1. KL-Shampoo as the immediate precursor

KL-Shampoo begins from the observation that Shampoo maintains two Kronecker factors $A\in\mathbb R^{d_a\times d_a}$ and $B\in\mathbb R^{d_b\times d_b}$ to approximate the full second moment $E[g\,g^\top]$ of the flattened gradient $g\in\mathbb R^{d_a d_b}$. Rather than analyzing these factors through Frobenius-norm bounds, the method treats $g$ as zero-mean Gaussian with covariance $\Sigma = E[g\,g^\top]+\kappa I$, and seeks a structured estimate $C=A\otimes B$ by minimizing

$$D_{\mathrm{KL}}(\mathcal N(0,\Sigma)\,\|\,\mathcal N(0,C)) = \tfrac12[\log\det C + \mathrm{Tr}(\Sigma C^{-1})] + \mathrm{const}.$$

This viewpoint yields two central consequences. First, if one factor is fixed, the KL minimizer recovers a Shampoo-style update; specifically, under a one-sided restriction with $C=(1/d_b\,A)\otimes B$ and $B=I$, the minimizer satisfies

$B\in\mathbb R^{d_b\times d_b}$0

which recovers Shampoo’s update with power $B\in\mathbb R^{d_b\times d_b}$1. Second, the joint minimization produces coupled “ideal” equations,

$B\in\mathbb R^{d_b\times d_b}$2

which motivate KL-Shampoo’s practical two-sided moving-average estimator (Lin et al., 3 Sep 2025).

The practical update uses the gradient matrix $B\in\mathbb R^{d_b\times d_b}$3, forms

$B\in\mathbb R^{d_b\times d_b}$4

updates

$B\in\mathbb R^{d_b\times d_b}$5

and preconditions by

$B\in\mathbb R^{d_b\times d_b}$6

The paper describes this as a principled proximal-gradient step on the KL objective. It further introduces a QR-based “fast” KL-Shampoo that maintains orthonormal bases $B\in\mathbb R^{d_b\times d_b}$7 via QR every $B\in\mathbb R^{d_b\times d_b}$8 steps and updates only diagonal eigenvalue vectors $B\in\mathbb R^{d_b\times d_b}$9 between QR steps (Lin et al., 3 Sep 2025).

This KL perspective is the conceptual substrate from which later Pro-KLShampoo variants emerge. A plausible implication is that once Shampoo is interpreted as structured covariance estimation rather than merely matrix accumulation, low-rank restrictions, projected parametrizations, and orthogonalization-based approximations become natural design axes.

2. Core limitation identified in Shampoo and the role of KL minimization

The original KL-Shampoo analysis argues that prior Shampoo analyses obscured a limitation by centering the estimation problem around Frobenius-norm considerations. In empirical terms, the reported limitation appears as instability when Shampoo is used without Adam-style grafting. The paper states that across NanoGPT (123M), NanoRWKV7 (162M), Llama (134M), and NanoMoE (227M), Shampoo with $E[g\,g^\top]$0 fails to train reliably when Adam-grafting is disabled, and that all 120 random-search runs diverged on RWKV7 under this condition. Under the same hyperparameter-search budget, KL-Shampoo trains stably without Adam (Lin et al., 3 Sep 2025).

The same work contrasts three optimizer families. Shampoo updates factors through $E[g\,g^\top]$1 and its transpose and needs Adam grafting if QR is infrequent. SOAP augments Shampoo with an extra RMSProp-style diagonal update in the eigenbasis of $E[g\,g^\top]$2, which incurs an additional $E[g\,g^\top]$3 vector. KL-Shampoo removes this heuristic addition by using the KL-derived two-sided factor updates and requiring only the Kronecker eigenvalues $E[g\,g^\top]$4 and $E[g\,g^\top]$5, rather than an extra $E[g\,g^\top]$6-sized moment vector (Lin et al., 3 Sep 2025).

The reported empirical outcomes position KL-Shampoo as both stabilizing and competitive. On validation loss and perplexity, it matches or outperforms SOAP while using the same QR frequency, and in NanoGPT pretraining it achieves approximately $E[g\,g^\top]$7–$E[g\,g^\top]$8 lower loss after $E[g\,g^\top]$9 B tokens than SOAP, and approximately $g\in\mathbb R^{d_a d_b}$0 better than Shampoo with grafting. The same summary also notes a tensor-valued extension on NanoMoE that outperforms baseline Shampoo variants (Lin et al., 3 Sep 2025).

Within the Pro-KLShampoo lineage, these claims are significant because later variants do not abandon KL-Shampoo’s two-sided covariance logic; instead, they attempt to compress or reparametrize it. This suggests that the primary object preserved across the family is not the exact state representation but the algebraic form of KL-derived whitening.

3. Projected KL-Shampoo with whitening recovered by orthogonalization

“Pro-KLShampoo: Projected KL-Shampoo with Whitening Recovered by Orthogonalization” (Sun et al., 7 May 2026) introduces a distinct optimizer under the Pro-KLShampoo name. Its central structural observation is that the eigenvalue spectra of KL-Shampoo’s Kronecker factors exhibit a “spike-and-flat” shape across layers, depths, and training stages in GPT-2 and LLaMA: a small number of dominant eigenvalues followed by an approximately uniform tail.

The paper provides an exact explanation under a rank-$g\in\mathbb R^{d_a d_b}$1 signal-plus-noise model

$g\in\mathbb R^{d_a d_b}$2

with $g\in\mathbb R^{d_a d_b}$3, $g\in\mathbb R^{d_a d_b}$4 deterministic and $g\in\mathbb R^{d_a d_b}$5 i.i.d. zero-mean noise with variance $g\in\mathbb R^{d_a d_b}$6. For stationary points $g\in\mathbb R^{d_a d_b}$7 of

$g\in\mathbb R^{d_a d_b}$8

the right factor satisfies

$g\in\mathbb R^{d_a d_b}$9

so that $g$0 has exactly $g$1 identical bottom eigenvalues, yielding a perfect spike-and-flat spectrum (Sun et al., 7 May 2026).

To exploit this, the method restricts the larger Kronecker factor to the parametric family

$g$2

Here $g$3 tracks the top-$g$4 eigenspace of the whitened second moment $g$5, $g$6 stores the full spectrum on that subspace, and $g$7 collapses the remaining $g$8 directions to a single scalar. The associated stationarity conditions are

$g$9

with $\Sigma = E[g\,g^\top]+\kappa I$0. The paper also gives an upper bound on the approximation gap to full KL optimization, which vanishes when the tail eigenvalues are uniform (Sun et al., 7 May 2026).

The defining move of Pro-KLShampoo is then to recover per-direction whitening on the complement by orthogonalization rather than by explicitly estimating a full complement covariance. The preconditioned gradient decomposition is

$\Sigma = E[g\,g^\top]+\kappa I$1

The second term uses only scalar right-scaling on the complement. Pro-KLShampoo replaces it by its polar factor: $\Sigma = E[g\,g^\top]+\kappa I$2 The paper states that this algebraically exactly matches the full KL-Shampoo complement update in its eigenbasis, but without forming a large eigendecomposition (Sun et al., 7 May 2026).

Algorithmically, Pro-KLShampoo alternates over minibatches between projecting $\Sigma = E[g\,g^\top]+\kappa I$3 into the tracked subspace and its complement, updating exponential moving averages of $\Sigma = E[g\,g^\top]+\kappa I$4, $\Sigma = E[g\,g^\top]+\kappa I$5, and $\Sigma = E[g\,g^\top]+\kappa I$6, tracking the top-$\Sigma = E[g\,g^\top]+\kappa I$7 eigenspace of $\Sigma = E[g\,g^\top]+\kappa I$8 by a single QR step, and computing

$\Sigma = E[g\,g^\top]+\kappa I$9

The paper reports that for $C=A\otimes B$0 with $C=A\otimes B$1 and $C=A\otimes B$2, memory drops from $C=A\otimes B$3 floats in KL-Shampoo to $C=A\otimes B$4, while the second factor’s QR cost is reduced from $C=A\otimes B$5 to $C=A\otimes B$6 and the right-factor state from size $C=A\otimes B$7 to $C=A\otimes B$8 (Sun et al., 7 May 2026).

4. Reparametrization, subspace QR, and BFloat16 storage

A second 2026 line, “Reparametrizing Shampoo and SOAP for Subspace Basis Updates and BFloat16 Storage” (Milligan et al., 25 May 2026), also addresses KL-Shampoo and related methods, but under a different notion of Pro-KLShampoo. Here the key idea is not a spike-and-flat model with orthogonalization; instead, it is a reparametrization of the preconditioner that supports BFloat16 storage and subspace basis updates.

For a weight matrix $C=A\otimes B$9, the paper writes Shampoo’s factors as $$D_{\mathrm{KL}}(\mathcal N(0,\Sigma)\,\|\,\mathcal N(0,C)) = \tfrac12[\log\det C + \mathrm{Tr}(\Sigma C^{-1})] + \mathrm{const}.$$0 and $$D_{\mathrm{KL}}(\mathcal N(0,\Sigma)\,\|\,\mathcal N(0,C)) = \tfrac12[\log\det C + \mathrm{Tr}(\Sigma C^{-1})] + \mathrm{const}.$$1, and recalls that KL-Shampoo and related QR-based variants replace expensive eigendecomposition with QR updates every $$D_{\mathrm{KL}}(\mathcal N(0,\Sigma)\,\|\,\mathcal N(0,C)) = \tfrac12[\log\det C + \mathrm{Tr}(\Sigma C^{-1})] + \mathrm{const}.$$2 steps. The reparametrization introduces the projected factor

$$D_{\mathrm{KL}}(\mathcal N(0,\Sigma)\,\|\,\mathcal N(0,C)) = \tfrac12[\log\det C + \mathrm{Tr}(\Sigma C^{-1})] + \mathrm{const}.$$3

where $$D_{\mathrm{KL}}(\mathcal N(0,\Sigma)\,\|\,\mathcal N(0,C)) = \tfrac12[\log\det C + \mathrm{Tr}(\Sigma C^{-1})] + \mathrm{const}.$$4. Storing $$D_{\mathrm{KL}}(\mathcal N(0,\Sigma)\,\|\,\mathcal N(0,C)) = \tfrac12[\log\det C + \mathrm{Tr}(\Sigma C^{-1})] + \mathrm{const}.$$5 is therefore equivalent to storing $$D_{\mathrm{KL}}(\mathcal N(0,\Sigma)\,\|\,\mathcal N(0,C)) = \tfrac12[\log\det C + \mathrm{Tr}(\Sigma C^{-1})] + \mathrm{const}.$$6. The paper argues that under this parametrization one does not need to materialize $$D_{\mathrm{KL}}(\mathcal N(0,\Sigma)\,\|\,\mathcal N(0,C)) = \tfrac12[\log\det C + \mathrm{Tr}(\Sigma C^{-1})] + \mathrm{const}.$$7 or $$D_{\mathrm{KL}}(\mathcal N(0,\Sigma)\,\|\,\mathcal N(0,C)) = \tfrac12[\log\det C + \mathrm{Tr}(\Sigma C^{-1})] + \mathrm{const}.$$8 in full precision; instead, one updates $$D_{\mathrm{KL}}(\mathcal N(0,\Sigma)\,\|\,\mathcal N(0,C)) = \tfrac12[\log\det C + \mathrm{Tr}(\Sigma C^{-1})] + \mathrm{const}.$$9, described as remaining nearly diagonal, in BFloat16 and rotates it in low-rank subspaces (Milligan et al., 25 May 2026).

The work proves algebraic equivalence between updates in the original and reparametrized systems. Given

$C=(1/d_b\,A)\otimes B$0

one computes a QR factorization

$C=(1/d_b\,A)\otimes B$1

and deduces

$C=(1/d_b\,A)\otimes B$2

The rotated projected factor then becomes

$C=(1/d_b\,A)\otimes B$3

According to the paper, these identities establish that updating $C=(1/d_b\,A)\otimes B$4 is algebraically equivalent to the original $C=(1/d_b\,A)\otimes B$5 QR update (Milligan et al., 25 May 2026).

The method then introduces subspace QR decomposition. Instead of factorizing the full $C=(1/d_b\,A)\otimes B$6 matrix $C=(1/d_b\,A)\otimes B$7, it updates only a block of size $C=(1/d_b\,A)\otimes B$8, for example $C=(1/d_b\,A)\otimes B$9 or $B=I$0. Writing

$B=I$1

it performs

$B=I$2

forms the block-diagonal rotation $B=I$3, and applies the same rotation formulas to $B=I$4 and $B=I$5. The paper gives the resulting cost as $B=I$6 plus two block matrix multiplications of cost $B=I$7, rather than $B=I$8 (Milligan et al., 25 May 2026).

This version of Pro-KLShampoo is therefore best understood as a systems-oriented refinement of QR-based KL-Shampoo and SOAP: it preserves the optimizer algebra while reducing precision and basis-update overhead.

5. Computational profile and memory characteristics

The three papers jointly provide a detailed picture of the computational trade space surrounding KL-Shampoo and Pro-KLShampoo variants.

The original KL-Shampoo paper describes each iteration as requiring two small matrix–matrix multiplications, $B=I$9 and $B\in\mathbb R^{d_b\times d_b}$00, plus $B\in\mathbb R^{d_b\times d_b}$01 elementwise operations for preconditioning, while QR incurs $B\in\mathbb R^{d_b\times d_b}$02 every $B\in\mathbb R^{d_b\times d_b}$03 steps, with $B\in\mathbb R^{d_b\times d_b}$04 reported as working well. Its state consists of $B\in\mathbb R^{d_b\times d_b}$05, $B\in\mathbb R^{d_b\times d_b}$06, orthogonal bases $B\in\mathbb R^{d_b\times d_b}$07, and eigenvalues $B\in\mathbb R^{d_b\times d_b}$08, and it avoids SOAP’s extra $B\in\mathbb R^{d_b\times d_b}$09 vector (Lin et al., 3 Sep 2025).

The projected-orthogonalization Pro-KLShampoo of (Sun et al., 7 May 2026) reduces the second Kronecker factor from $B\in\mathbb R^{d_b\times d_b}$10 parameters to $B\in\mathbb R^{d_b\times d_b}$11 by restricting it to a rank-$B\in\mathbb R^{d_b\times d_b}$12 subspace plus flat tail, and it changes the leading compute from KL-Shampoo’s $B\in\mathbb R^{d_b\times d_b}$13 matrix multiplication with $B\in\mathbb R^{d_b\times d_b}$14 amortized QR to $B\in\mathbb R^{d_b\times d_b}$15 with $B\in\mathbb R^{d_b\times d_b}$16 QR. The explicit memory comparison given is $B\in\mathbb R^{d_b\times d_b}$17 floats for KL-Shampoo versus $B\in\mathbb R^{d_b\times d_b}$18 for Pro-KLShampoo (Sun et al., 7 May 2026).

The reparametrized Pro-KLShampoo of (Milligan et al., 25 May 2026) instead compares “vanilla KL-Shampoo” with “Pro-KLShampoo” in terms of matrix-multiplication-equivalent cost. Vanilla KL-Shampoo with full-basis QR every $B\in\mathbb R^{d_b\times d_b}$19 steps is summarized as approximately $B\in\mathbb R^{d_b\times d_b}$20 matrix-multiplication units plus preconditioning, whereas Pro-KLShampoo with projected updates and subspace QR is summarized as approximately $B\in\mathbb R^{d_b\times d_b}$21 units plus the same preconditioning term. On memory, vanilla KL-Shampoo stores two $B\in\mathbb R^{d_b\times d_b}$22 factors in FP32, while Pro-KLShampoo stores two $B\in\mathbb R^{d_b\times d_b}$23 projected factors in BFP16 (Milligan et al., 25 May 2026).

The following table organizes these reported state reductions.

Method	State description	Reported reduction
KL-Shampoo	$B\in\mathbb R^{d_b\times d_b}$24	Avoids SOAP’s extra $B\in\mathbb R^{d_b\times d_b}$25 vector
Pro-KLShampoo (Sun et al., 7 May 2026)	$B\in\mathbb R^{d_b\times d_b}$26	Right factor from $B\in\mathbb R^{d_b\times d_b}$27 to $B\in\mathbb R^{d_b\times d_b}$28
Pro-KLShampoo (Milligan et al., 25 May 2026)	Projected factors $B\in\mathbb R^{d_b\times d_b}$29 in BFP16	FP32 factors replaced by BFP16 projected factors

A plausible implication is that “Pro-KLShampoo” has become a label for multiple optimization pathways that target different bottlenecks in KL-Shampoo: one targets statistical structure in the spectrum, while the other targets numerical representation and basis-update mechanics.

6. Empirical behavior in pretraining workloads

The empirical record reported in these sources is concentrated on neural network pretraining, especially LLMs.

KL-Shampoo is described as stable without Adam grafting on NanoGPT (123M), NanoRWKV7 (162M), Llama (134M), and NanoMoE (227M), under the same 120-run hyperparameter search in which Shampoo without grafting fails to train reliably. On NanoGPT pretraining, KL-Shampoo is reported to achieve approximately $B\in\mathbb R^{d_b\times d_b}$30–$B\in\mathbb R^{d_b\times d_b}$31 lower loss after $B\in\mathbb R^{d_b\times d_b}$32 B tokens than SOAP, and approximately $B\in\mathbb R^{d_b\times d_b}$33 better than Shampoo with grafting; on NanoMoE, its tensor-Kronecker extension outperforms baseline Shampoo variants (Lin et al., 3 Sep 2025).

The projected-orthogonalization Pro-KLShampoo reports extensive pretraining results at four scales: GPT-2 124M and 350M on FineWeb-10B, and LLaMA 134M and 450M on C4. The paper states that Pro-KLShampoo consistently outperforms KL-Shampoo at every tested subspace rank $B\in\mathbb R^{d_b\times d_b}$34 in validation loss, peak per-GPU memory, and wallclock time to reach each loss level. Concrete examples include GPT-2 124M, where KL-Shampoo reaches $B\in\mathbb R^{d_b\times d_b}$35 versus Pro-KLShampoo $B\in\mathbb R^{d_b\times d_b}$36, and LLaMA 450M, where the gap is $B\in\mathbb R^{d_b\times d_b}$37. Reported peak per-GPU memory examples include $B\in\mathbb R^{d_b\times d_b}$38 GiB on GPT-2 350M and $B\in\mathbb R^{d_b\times d_b}$39 GiB on LLaMA 450M. Reported time-to-loss improvements relative to KL-Shampoo at matched loss are approximately $B\in\mathbb R^{d_b\times d_b}$40 for GPT-2 124M, $B\in\mathbb R^{d_b\times d_b}$41 for GPT-2 350M, $B\in\mathbb R^{d_b\times d_b}$42 for LLaMA 134M, and $B\in\mathbb R^{d_b\times d_b}$43 for LLaMA 450M (Sun et al., 7 May 2026).

The reparametrized Pro-KLShampoo emphasizes BFloat16 robustness and time-cost trade-offs. On nanoGPT (123 M) and Llama3 (119 M), switching from FP32 to BFP16 storage causes the original KL-Shampoo to degrade by up to $B\in\mathbb R^{d_b\times d_b}$44 in test loss, whereas Pro-KLShampoo remains within $B\in\mathbb R^{d_b\times d_b}$45 and slightly improves in one setting. For $B\in\mathbb R^{d_b\times d_b}$46, full-basis QR in FP32 is reported as approximately $B\in\mathbb R^{d_b\times d_b}$47 runtime, subspace QR with $B\in\mathbb R^{d_b\times d_b}$48 as approximately $B\in\mathbb R^{d_b\times d_b}$49, and matrix-multiply only as approximately $B\in\mathbb R^{d_b\times d_b}$50. Greedy block selection is reported to outperform random block selection; a single subspace update with $B\in\mathbb R^{d_b\times d_b}$51 and $B\in\mathbb R^{d_b\times d_b}$52 matches full-basis test loss while saving approximately $B\in\mathbb R^{d_b\times d_b}$53–$B\in\mathbb R^{d_b\times d_b}$54 runtime. On Llama3 (313 M), $B\in\mathbb R^{d_b\times d_b}$55 loses only $B\in\mathbb R^{d_b\times d_b}$56 in loss while reducing wall-clock by approximately $B\in\mathbb R^{d_b\times d_b}$57, and more frequent small-block updates recover most accuracy at still $B\in\mathbb R^{d_b\times d_b}$58 lower runtime versus full basis (Milligan et al., 25 May 2026).

Because these evaluations are reported in different experimental programs, direct cross-paper ranking is not warranted from the supplied evidence alone. What can be stated is that both variants report gains relative to KL-Shampoo under the workloads they test.

7. Relationship to SOAP, Muon, and open questions

The Pro-KLShampoo literature is notable for placing KL-Shampoo in relation to two different neighboring optimizer families.

First, KL-Shampoo is repeatedly compared with SOAP. In the original KL-Shampoo work, SOAP is described as running an extra RMSProp-style diagonal update in the eigenbasis of $B\in\mathbb R^{d_b\times d_b}$59, which requires an additional $B\in\mathbb R^{d_b\times d_b}$60 vector. KL-Shampoo removes this heuristic component through KL-derived factor estimation and no extra $B\in\mathbb R^{d_b\times d_b}$61 memories (Lin et al., 3 Sep 2025). The reparametrization paper then generalizes its projected-factor and subspace-QR construction to Shampoo-based methods employing QR, including KL-Shampoo, SOAP, and KL-SOAP, and states that it improves SOAP and KL-SOAP under BFP16 storage, enabling KL-SOAP to match or exceed KL-Shampoo (Milligan et al., 25 May 2026).

Second, the projected-orthogonalization Pro-KLShampoo directly links KL-Shampoo to Muon-style orthogonalization. The paper describes explicit Kronecker-factored preconditioning and orthogonalization of the gradient momentum as “two distinct frontiers,” typically developed in isolation, and proposes a hybrid in which a low-rank spectral restriction handles the spiked subspace while orthogonalization recovers whitening on the flat complement (Sun et al., 7 May 2026). This positioning matters conceptually because it reframes orthogonalization not as a separate optimizer family but as an exact algebraic surrogate for full KL-Shampoo whitening on a restricted complement.

Several limitations and unresolved issues are explicitly recorded. The original KL-Shampoo paper describes its results as preliminary and limited to language-model pretraining; vision and RL benchmarks remain to be tried. It also raises adaptive selection of QR frequency $B\in\mathbb R^{d_b\times d_b}$62, and possible trade-offs introduced by mixed-precision or distributed variants (Lin et al., 3 Sep 2025). The reparametrization paper addresses mixed precision in part through BFloat16 storage, but it still depends on QR decomposition in FP32, since existing QR implementations require single-precision arithmetic and remain computationally expensive when preconditioning matrices are large (Milligan et al., 25 May 2026).

A common misconception would be to treat Pro-KLShampoo as a single canonical algorithm. The supplied literature shows instead that the label covers at least two technically distinct developments: one based on projected low-rank-plus-flat structure with orthogonalization (Sun et al., 7 May 2026), and another based on projected-factor reparametrization, subspace QR, and BFloat16 storage (Milligan et al., 25 May 2026). What unifies them is their dependence on KL-Shampoo’s KL-minimization perspective and their effort to preserve whitening quality while reducing practical bottlenecks.