---
title: 'Pro-KLShampoo: Enhancing KL-Shampoo Optimizers'
url: https://www.emergentmind.com/topics/pro-klshampoo
type: topic
---

# Pro-KLShampoo: Enhancing KL-Shampoo Optimizers

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” [2605.06316], and a reparametrization framework for Shampoo-based methods titled “Reparametrizing Shampoo and SOAP for Subspace Basis Updates and BFloat16 Storage” [2605.26327]. 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 [2509.03378]. 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
\[
A^* = E[\mathrm{Mat}(g)\,\mathrm{Mat}(g)^\top],
\]
which recovers Shampoo’s update with power \(p=\tfrac12\). Second, the joint minimization produces coupled “ideal” equations,
\[
A^* = (1/d_b)\,E[\mathrm{Mat}(g)\,B^{*-1}\,\mathrm{Mat}(g)^\top],\qquad
B^* = (1/d_a)\,E[\mathrm{Mat}(g)^\top\,A^{*-1}\,\mathrm{Mat}(g)],
\]
which motivate KL-Shampoo’s practical two-sided moving-average estimator [2509.03378].

The practical update uses the gradient matrix \(G=\mathrm{Mat}(\nabla \ell(w))\), forms
\[
\Delta_a = (1/d_b)\,G\,B^{-1}G^\top,\qquad
\Delta_b = (1/d_a)\,G^\top A^{-1}G,
\]
updates
\[
A\leftarrow (1-\beta_2)A+\beta_2\Delta_a,\qquad
B\leftarrow (1-\beta_2)B+\beta_2\Delta_b,
\]
and preconditions by
\[
w\leftarrow w-\gamma\,(A^{-1/2}GB^{-1/2}).
\]
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 \(U_a,U_b\) via QR every \(T\) steps and updates only diagonal eigenvalue vectors \(\Lambda_a,\Lambda_b\) between QR steps [2509.03378].

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 \(p=\tfrac12\) 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 [2509.03378].

The same work contrasts three optimizer families. Shampoo updates factors through \(G G^\top\) 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 \(A\otimes B\), which incurs an additional \(O(d_a d_b)\) vector. KL-Shampoo removes this heuristic addition by using the KL-derived two-sided factor updates and requiring only the Kronecker eigenvalues \(\Lambda_a\in\mathbb R^{d_a}\) and \(\Lambda_b\in\mathbb R^{d_b}\), rather than an extra \(d_a d_b\)-sized moment vector [2509.03378].

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 \(3\)–\(5\%\) lower loss after \(1\) B tokens than SOAP, and approximately \(10\%\) better than Shampoo with grafting. The same summary also notes a tensor-valued extension on NanoMoE that outperforms baseline Shampoo variants [2509.03378].

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” [2605.06316] 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-\(\rho\) signal-plus-noise model
\[
G = A B^\top + \xi,
\]
with \(A\in\mathbb R^{m\times\rho}\), \(B\in\mathbb R^{n\times\rho}\) deterministic and \(\xi\) i.i.d. zero-mean noise with variance \(\sigma^2\). For stationary points \((L^*,R^*)\) of
\[
\min_{L\succ0,\;R\succ0}\;
D_{\mathrm{KL}}(\mathcal N(0,\Sigma)\,\|\,\mathcal N(0,L\otimes R)),
\]
the right factor satisfies
\[
R^* = S_R^* + f_R^* I_n,\qquad
\rank(S_R^*)\le \rho,\qquad
f_R^* = \tfrac{\sigma^2\,\mathrm{tr}((L^*)^{-1})}{m},
\]
so that \(R^*\) has exactly \(n-\rho\) identical bottom eigenvalues, yielding a perfect spike-and-flat spectrum [2605.06316].

To exploit this, the method restricts the larger Kronecker factor to the parametric family
\[
\hat R = U S U^\top + \mu_\perp (I_n-UU^\top),
\qquad U\in \mathrm{St}(n,r),\; S\in S^r,\; \mu_\perp>0.
\]
Here \(U\) tracks the top-\(r\) eigenspace of the whitened second moment \(\Phi_L = E[G^\top L^{-1}G]\), \(S\) stores the full spectrum on that subspace, and \(\mu_\perp\) collapses the remaining \(n-r\) directions to a single scalar. The associated stationarity conditions are
\[
S^*=\tfrac1m\,E[G^\top(L^*)^{-1}G]_U,\qquad
\mu_\perp^*=\tfrac1{m(n-r)}\,\mathrm{tr}\bigl(E[G_\perp^\top(L^*)^{-1}G_\perp]\bigr),\qquad
L^*=\tfrac1n\,E[G(\hat R^*)^{-1}G^\top],
\]
with \(G_\perp = G(I-UU^\top)\). The paper also gives an upper bound on the approximation gap to full KL optimization, which vanishes when the tail eigenvalues are uniform [2605.06316].

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
\[
L^{-1/2}G\hat R^{-1/2}
=
L^{-1/2}GUS^{-1/2}U^\top
+
\mu_\perp^{-1/2}L^{-1/2}GP_\perp.
\]
The second term uses only scalar right-scaling on the complement. Pro-KLShampoo replaces it by its polar factor:
\[
\polar\bigl(\mu_\perp^{-1/2}L^{-1/2}G P_\perp\bigr)
=
L^{-1/2}G\,U_\perp
\bigl(U_\perp^\top E[G^\top L^{-1}G]U_\perp\bigr)^{-1/2}
U_\perp^\top.
\]
The paper states that this algebraically exactly matches the full KL-Shampoo complement update in its eigenbasis, but without forming a large eigendecomposition [2605.06316].

Algorithmically, Pro-KLShampoo alternates over minibatches between projecting \(G\) into the tracked subspace and its complement, updating exponential moving averages of \(L\), \(S\), and \(\mu_\perp\), tracking the top-\(r\) eigenspace of \(\Phi_L\) by a single QR step, and computing
\[
\Delta W =
-\alpha_{kl}(L^{-1/2}GUS^{-1/2}U^\top)
-
c_a\,\polar(\mu_\perp^{-1/2}L^{-1/2}GP_\perp).
\]
The paper reports that for \(W\in\mathbb R^{m\times n}\) with \(m\le n\) and \(r\ll n\), memory drops from \(2m^2+2n^2\) floats in KL-Shampoo to \(2m^2+2r^2+nr+1\), while the second factor’s QR cost is reduced from \(O(n^3)\) to \(O(r^3)\) and the right-factor state from size \(n^2\) to \(nr\) [2605.06316].

## 4. Reparametrization, subspace QR, and BFloat16 storage

A second 2026 line, “Reparametrizing Shampoo and SOAP for Subspace Basis Updates and BFloat16 Storage” [2605.26327], 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 \(W\in\mathbb R^{d_1\times d_2}\), the paper writes Shampoo’s factors as \(G_1\) and \(G_2\), and recalls that KL-Shampoo and related QR-based variants replace expensive eigendecomposition with QR updates every \(T\) steps. The reparametrization introduces the projected factor
\[
P_i := Q_i^\top G_i Q_i \in \mathbb R^{d_i\times d_i},
\]
where \(G_i = Q_i P_i Q_i^\top\). Storing \((P_i,Q_i)\) is therefore equivalent to storing \((G_i,Q_i)\). The paper argues that under this parametrization one does not need to materialize \(G_i\) or \(Q_i\) in full precision; instead, one updates \(P_i\), described as remaining nearly diagonal, in BFloat16 and rotates it in low-rank subspaces [2605.26327].

The work proves algebraic equivalence between updates in the original and reparametrized systems. Given
\[
P_i^{\rm old}=Q_i^{{\rm old}\,\top}G_iQ_i^{\rm old},
\]
one computes a QR factorization
\[
P_i^{\rm old}=Q_i^{\rm new}R_i,
\]
and deduces
\[
Q_i = Q_i^{{\rm old}\,\top}Q_i^{\rm new},\qquad
Q_i^{\rm new}=Q_i^{\rm old}Q_i.
\]
The rotated projected factor then becomes
\[
P_i^{\rm new}=Q_i^\top P_i^{\rm old}Q_i
=Q_i^{{\rm new}\,\top}G_iQ_i^{\rm new}.
\]
According to the paper, these identities establish that updating \((P_i,Q_i)\) is algebraically equivalent to the original \((G_i,Q_i)\) QR update [2605.26327].

The method then introduces subspace QR decomposition. Instead of factorizing the full \(d_i\times d_i\) matrix \(P_i\), it updates only a block of size \(d_{\rm sub}=B\,d_i\), for example \(B=1/4\) or \(1/2\). Writing
\[
P_i
=
\begin{pmatrix}
P_i^{(XX)} & P_i^{(XY)}\\
P_i^{(YX)} & P_i^{(YY)}
\end{pmatrix},
\qquad
P_i^{(XX)}\in\mathbb R^{d_{\rm sub}\times d_{\rm sub}},
\]
it performs
\[
P_i^{(XX)}=\widehat Q_i\widehat R_i,\qquad \widehat Q_i\in O(d_{\rm sub}),
\]
forms the block-diagonal rotation \(Q_i^{(\rm block)}=\widehat Q_i\oplus I_{d_i-d_{\rm sub}}\), and applies the same rotation formulas to \(Q_i\) and \(P_i\). The paper gives the resulting cost as \(O(d_{\rm sub}^3)\) plus two block matrix multiplications of cost \(O(d_i d_{\rm sub}^2)\), rather than \(O(d_i^3)\) [2605.26327].

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, \(G B^{-1}\) and \(G^\top A^{-1}\), plus \(O(d_a d_b)\) elementwise operations for preconditioning, while QR incurs \(O((d_a^2+d_b^2)\cdot \min(d_a,d_b))\) every \(T\) steps, with \(T\approx 10\) reported as working well. Its state consists of \(A\), \(B\), orthogonal bases \(U_a,U_b\), and eigenvalues \(\Lambda_a,\Lambda_b\), and it avoids SOAP’s extra \(O(d_a d_b)\) vector [2509.03378].

The projected-orthogonalization Pro-KLShampoo of [2605.06316] reduces the second Kronecker factor from \(n^2\) parameters to \(nr\) by restricting it to a rank-\(r\) subspace plus flat tail, and it changes the leading compute from KL-Shampoo’s \(O(mn^2)\) matrix multiplication with \(O((m^3+n^3)/\tau)\) amortized QR to \(O(m^2n+mnr+T_{\mathrm{NS}}mn)\) with \(O((m^3+r^3)/\tau)\) QR. The explicit memory comparison given is \(2m^2+2n^2\) floats for KL-Shampoo versus \(2m^2+2r^2+nr+1\) for Pro-KLShampoo [2605.06316].

The reparametrized Pro-KLShampoo of [2605.26327] instead compares “vanilla KL-Shampoo” with “Pro-KLShampoo” in terms of matrix-multiplication-equivalent cost. Vanilla KL-Shampoo with full-basis QR every \(T\) steps is summarized as approximately \(2+1+\tfrac{2}{T}\) matrix-multiplication units plus preconditioning, whereas Pro-KLShampoo with projected updates and subspace QR is summarized as approximately \(2+\tfrac{(B\,d_i)^3}{T}\) units plus the same preconditioning term. On memory, vanilla KL-Shampoo stores two \(d_i\times d_i\) factors in FP32, while Pro-KLShampoo stores two \(d_i\times d_i\) projected factors in BFP16 [2605.26327].

The following table organizes these reported state reductions.

| Method | State description | Reported reduction |
|---|---|---|
| KL-Shampoo | \(A,B,U_a,U_b,\Lambda_a,\Lambda_b\) | Avoids SOAP’s extra \(O(d_a d_b)\) vector |
| Pro-KLShampoo [2605.06316] | \(L,S,Q_L,Q_S,U,\mu_\perp\) | Right factor from \(n^2\) to \(nr\) |
| Pro-KLShampoo [2605.26327] | Projected factors \(P_i\) 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 language models.

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 \(3\)–\(5\%\) lower loss after \(1\) B tokens than SOAP, and approximately \(10\%\) better than Shampoo with grafting; on NanoMoE, its tensor-Kronecker extension outperforms baseline Shampoo variants [2509.03378].

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 \(r\in\{32,64,128\}\) 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 \(3.2796\) versus Pro-KLShampoo \(3.2745\), and LLaMA 450M, where the gap is \(2.7650\rightarrow 2.7453\). Reported peak per-GPU memory examples include \(73.07\rightarrow 72.25\) GiB on GPT-2 350M and \(73.42\rightarrow 71.89\) GiB on LLaMA 450M. Reported time-to-loss improvements relative to KL-Shampoo at matched loss are approximately \(2.3\%\) for GPT-2 124M, \(4.1\%\) for GPT-2 350M, \(11\%\) for LLaMA 134M, and \(13\%\) for LLaMA 450M [2605.06316].

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 \(+0.013\) in test loss, whereas Pro-KLShampoo remains within \(\pm 0.006\) and slightly improves in one setting. For \(d_1=d_2=6000\), full-basis QR in FP32 is reported as approximately \(1.0\times\) runtime, subspace QR with \(B=0.2\) as approximately \(0.85\times\), and matrix-multiply only as approximately \(0.4\times\). Greedy block selection is reported to outperform random block selection; a single subspace update with \(K=1\) and \(B=\tfrac12\) matches full-basis test loss while saving approximately \(6\)–\(7\%\) runtime. On Llama3 (313 M), \(T=10,B=\tfrac12,K=1\) loses only \(0.004\) in loss while reducing wall-clock by approximately \(9\%\), and more frequent small-block updates recover most accuracy at still \(5\%\) lower runtime versus full basis [2605.26327].

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 \(A\otimes B\), which requires an additional \(O(d_a d_b)\) vector. KL-Shampoo removes this heuristic component through KL-derived factor estimation and no extra \(d_a d_b\) memories [2509.03378]. 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 [2605.26327].

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 [2605.06316]. 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 \(T\), and possible trade-offs introduced by mixed-precision or distributed variants [2509.03378]. 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 [2605.26327].

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 [2605.06316], and another based on projected-factor reparametrization, subspace QR, and BFloat16 storage [2605.26327]. What unifies them is their dependence on KL-Shampoo’s KL-minimization perspective and their effort to preserve whitening quality while reducing practical bottlenecks.

Source: https://www.emergentmind.com/topics/pro-klshampoo