---
title: 'KL-SOAP: Hybrid Shampoo Optimizer'
url: https://www.emergentmind.com/topics/kl-soap
type: topic
---

# KL-SOAP: Hybrid Shampoo Optimizer

KL-SOAP is a Shampoo-based optimizer that combines Kronecker-factored metric learning, as in KL-Shampoo, with the “second-moment-vector” preconditioning of SOAP. In the formulation reported in "Reparametrizing Shampoo and SOAP for Subspace Basis Updates and BFloat16 Storage" [2605.26327], KL-SOAP tracks Kronecker factors \((A_1,A_2)\) exactly as in KL-Shampoo while also maintaining a vector second moment \(v\) in the Kronecker eigenbasis, and it uses a reparametrization that supports BFloat16 storage and forms a complete basis by combining updated basis vectors with unchanged ones. The same work places KL-SOAP within a broader family of Shampoo-based methods that employ QR decomposition and argues that reparametrized, subspace-based basis updates reduce computational overhead while mitigating the performance degradation caused by BFloat16 storage.

## 1. Position within the Shampoo family

The paper situates KL-SOAP among three closely related Shampoo-based methods [2605.26327]. Shampoo and KL-Shampoo maintain two Kronecker factors \(A_1\in\mathbb R^{d_1\times d_1}\) and \(A_2\in\mathbb R^{d_2\times d_2}\), and form the preconditioner
\[
P_{\rm shampoo}=A_1\otimes A_2.
\]

SOAP instead works in the eigenbasis \(\Phi=U_1\otimes U_2\) of \(P_{\rm shampoo}\), tracks the elementwise “Adam” second moment
\[
v=\mathbb E[(\Phi^\top g)^{\odot2}],
\]
and forms
\[
P_{\rm soap}=\Phi\;\mathrm{Diag}(v)\;\Phi^\top,
\qquad
P_{\rm soap}^{-\tfrac12}=\Phi\,\mathrm{Diag}(v^{-1/2})\,\Phi^\top.
\]

KL-SOAP merges these two ideas. It keeps KL-Shampoo’s update mechanism for \((A_1,A_2)\), but also tracks SOAP’s vector moment \(v\). The preconditioner at time \(t\) is
\[
P_t
=
\underbrace{(U_{1,t}\otimes U_{2,t})}_{\Phi_t}\,
\mathrm{Diag}(v_t)\,
\underbrace{(U_{1,t}\otimes U_{2,t})^\top}_{\Phi_t^\top},
\]
and the parameter update is
\[
W_{t+1}=W_t-\gamma\,\mathrm{mat}\bigl(P_t^{-1/2}g_t\bigr).
\]

The paper’s stated motivation is that KL-SOAP often improves on pure-Shampoo through better scaling in “long-tail” coordinates and on pure-SOAP through better Kronecker-coupling. A plausible implication is that KL-SOAP is intended as a hybrid preconditioner that preserves the structural coupling of Kronecker factors while retaining elementwise adaptation in the learned basis.

## 2. Full mathematical formulation

In the full-basis formulation, KL-SOAP tracks two matrix factors and a vector moment [2605.26327]. Let
\[
g_t=\mathrm{vec}\bigl(\nabla\ell(W_t)\bigr).
\]
Then
\[
\Delta_{i,t}
=\begin{cases}
A_{2,t-1}^{-1/2}\,(\mathrm{mat}(g_t)\,A_{2,t-1}^{-1/2})
(\mathrm{mat}(g_t)\,A_{2,t-1}^{-1/2})^\top/d_2
& i=1,\\[6pt]
\bigl(\mathrm{mat}(g_t)^\top A_{1,t-1}^{-1/2}\bigr)
\bigl(\mathrm{mat}(g_t)^\top A_{1,t-1}^{-1/2}\bigr)^\top/d_1
& i=2,
\end{cases}
\]
and the factor updates are
\[
A_{i,t}=(1-\beta_2)\,A_{i,t-1}+\beta_2\,\Delta_{i,t},
\qquad
[U_{i,t},\Sigma_{i,t}]=\mathrm{eig}(A_{i,t}).
\]

Using \(\Phi_t=U_{1,t}\otimes U_{2,t}\), KL-SOAP updates the vector moment as
\[
v_t=(1-\beta_2)\,v_{t-1}+\beta_2\,(\Phi_t^\top g_t)^{\odot2},
\]
then defines
\[
P_t=\Phi_t\,\mathrm{Diag}(v_t)\,\Phi_t^\top,
\qquad
g_t^{\rm pre}=P_t^{-1/2}\,g_t,
\qquad
W_{t+1}=W_t-\gamma\,\mathrm{mat}(g_t^{\rm pre}).
\]

This formulation makes explicit that KL-SOAP uses two distinct state types. The first is the Kronecker-factored curvature state \((A_{1,t},A_{2,t})\), inherited from KL-Shampoo. The second is the coordinatewise second moment \(v_t\) in the time-varying basis \(\Phi_t\), inherited from SOAP. The optimizer therefore combines curvature adaptation induced by the Kronecker factors with diagonal rescaling in the learned eigenspace.

## 3. Reparametrization for BFloat16 storage

The paper identifies two practical issues for Shampoo-based methods: existing QR implementations require single-precision arithmetic and remain computationally expensive, and using BFloat16 storage to reduce memory usage can degrade performance [2605.26327]. Its proposed remedy is a reparametrization of the preconditioner.

Instead of storing \(U_{i,t}\) and \(\Sigma_{i,t}\) in FP32, the method stores the combined factor
\[
M_{i,t}:=U_{i,t}^\top\,A_{i,t}\,U_{i,t}\quad\in\mathbb R^{d_i\times d_i},
\]
together with the orthogonal basis \(U_{i,t}\). Algebraically,
\[
A_{i,t}^\text{old}=U_{i,t-1}\,M_{i,t-1}\,U_{i,t-1}^\top,
\qquad
[Q_i,R_i]=\mathrm{qr}\bigl(M_{i,t-1}\bigr),
\qquad
U_{i,t}=U_{i,t-1}Q_i,
\]
and then
\[
M_{i,t}
=Q_i^\top\,M_{i,t-1}\,Q_i,
\qquad
M_{i,t}\leftarrow (1-\beta_2)\,M_{i,t}+\beta_2\,\widetilde\Delta_{i,t},
\qquad
\widetilde\Delta_{i,t}=Q_i^\top\,\Delta_{i,t}\,Q_i.
\]

The exposition states that storing \(M_{i,t}\), described as nearly diagonal, and \(U_{i,t}\) in BFloat16 is safe, while QR is performed in FP32 only on small subspaces. It also states that the vector \(v_t\) is updated as before but can also be stored in BFloat16 without measurable loss. This suggests that the reparametrization is not merely a memory layout change; it is designed to preserve the numerical role of the basis while relocating the precision-sensitive computation to a reduced-dimensional QR step.

## 4. Subspace basis updates

A central feature of the reparametrized method is the subspace basis update [2605.26327]. Rather than performing a full \(d_i\times d_i\) QR, the method chooses an index set \(\mathcal I\) of size \(k\), for example \(k=B\,d_i\) with \(B\in(0,1)\), and partitions
\[
M_{i,t-1}=
\begin{bmatrix}
M_{XX}&M_{XY}\\
M_{YX}&M_{YY}
\end{bmatrix},
\qquad
M_{XX}\in\mathbb R^{k\times k}.
\]

It then performs
\[
[Q_{XX},R]=\mathrm{qr}(M_{XX}),
\qquad
U_{i,t}\bigl[:,\mathcal I\bigr]
=
U_{i,t-1}\bigl[:,\mathcal I\bigr]\,Q_{XX},
\]
and updates only those rows and columns of \(M_{i,t}\) that touch \(\mathcal I\):
\[
M_{i,t}\bigl[\mathcal I,\mathcal I\bigr]
=
Q_{XX}^\top\,M_{XX}\,Q_{XX},
\]
while \(M_{i,t}[\mathcal I,:]\) and \(M_{i,t}[:,\mathcal I]\) are rotated similarly.

The cost is given explicitly as \(O(k^3)\) for QR plus \(O(d_i\,k^2)\) for the two block updates, versus \(O(d_i^3)\) for a full QR. The paper also describes the resulting basis as complete because updated basis vectors are combined with unchanged ones. In operational terms, the algorithm updates only a subspace of the basis at each scheduled QR step while leaving the complement intact.

A typical loop includes the following elements. \(M_i\), \(U_i\), and \(v\) are stored in BFloat16; gradients and small QR computations remain in FP32. The algorithm computes \(\Delta_i\) in FP32, projects it into the current basis via \(M̄\leftarrow U_i^\top\Delta_iU_i\), updates \(M_i\), periodically performs the subspace QR step every \(T\) iterations, applies the basis transforms implicitly through two matrix multiplications rather than materializing \(\Phi=U_1\otimes U_2\), updates \(v\), rescales by \((v+\epsilon)^{-1/2}\), and maps the preconditioned direction back to parameter space.

## 5. Computational and memory characteristics

The paper gives a cost analysis in terms of \(d=\max(d_1,d_2)\) [2605.26327]. A full QR costs \(O(d^3)\) per factor, for a total of \(2\,O(d^3)\). With subspace QR and block size \(k=B\,d\),

- QR on the \(k\times k\) block costs \(O(k^3)=O(B^3d^3)\).
- Rotating the two touched blocks costs \(O(d\,k^2)=O(B^2d^3)\).
- The combined per-factor cost is \(O(B^3d^3+B^2d^3)=O(B^2d^3)\).
- The total across both factors is \(O(2B^2d^3)\).

The exposition adds that with \(B\ll1\), for example \(B=1/2\) or \(1/3\), the subspace approach can be \(4\times\) or \(9\times\) faster than full QR. This is a direct consequence of replacing cubic dependence on the full dimension by cubic and quadratic dependence on a smaller block.

The memory accounting is also explicit. Per factor, the method stores \(U_i\) and \(M_i\), each with \(d^2\) numbers. In FP32 this is
\[
2\times4d^2=8d^2 \text{ bytes},
\]
whereas in BFloat16 it is
\[
2\times2d^2=4d^2 \text{ bytes},
\]
which yields 50% savings. The vector \(v\in\mathbb R^{d_1d_2}\) is also stored in BFloat16 for \(2\times\) savings versus FP32. The paper’s broader claim is that the reparametrization makes Shampoo-based methods more memory- and time-efficient, and KL-SOAP is presented as a concrete beneficiary of that redesign.

## 6. Empirical behavior and implementation parameters

The reported empirical results focus on BFloat16 storage and subspace-QR runtime behavior [2605.26327]. Under the new reparametrization, KL-SOAP, KL-Shampoo, and SOAP can be stored entirely in BFloat16 without measurable loss in test loss; the same section also states that subspace-QR steps cut wall-clock runtime by approximately \(1.5\!-\!2\times\) with at most 0.001 degradation in loss.

| Method | Parametrization | FP32 loss \(\rightarrow\) BFP16 loss (\(\Delta\)) |
|---|---|---|
| KL-SOAP | old | 3.351 \(\rightarrow\) 3.362 (+0.011) |
| KL-SOAP | new | 3.346 \(\rightarrow\) 3.346 (+0.000) ★ |
| KL-Shampoo | old | 3.345 \(\rightarrow\) 3.345 (+0.000) |
| KL-Shampoo | new | 3.345 \(\rightarrow\) 3.344 (−0.001) ★ |
| SOAP | old | 3.358 \(\rightarrow\) 3.372 (+0.014) |
| SOAP | new | 3.353 \(\rightarrow\) 3.359 (+0.006) |

The mark ★ is identified in the source as best under BFloat16. For KL-SOAP specifically, the new parametrization changes the reported transition from \(3.351\rightarrow3.362\) to \(3.346\rightarrow3.346\). This suggests that, in the reported setting, the BFloat16 penalty is removed for KL-SOAP.

The paper further reports that for \(d_1=d_2=6{,}000\) and \(B=1/2\), subspace QR is \(>1.5\times\) faster than full-basis QR on an NVIDIA H100, yielding overall optimizer speedups of approximately 30%. In the implementation notes, the recommended block fraction is \(B\in\{1/2,1/3\}\), the subspace frequency is 5–10 steps, the inner loop usually uses \(K=1\) subspace QR per update, the damping for vector \(v\) is \(1\mathrm{e}{-6}\), the Shampoo damping for \(A_i\) is \(1\mathrm{e}{-6}\), and the optimizer coefficients are \(\beta_1=0.9\) and \(\beta_2=0.999\). For block selection, the recommended strategy is the two-phase “greedy Jacobi,” described as picking the largest off-diagonal and growing a block, implementable via top-k in PyTorch; random sampling also works but is slightly less robust. The work states that the method can be implemented in modern deep-learning frameworks such as JAX and PyTorch using standard operations including matmul, top-k, and small-scale QR.

In this presentation, KL-SOAP is therefore not only a hybrid preconditioner but also a precision-aware and basis-update-aware optimizer design. Its defining technical characteristics are the combination of KL-Shampoo factor updates with SOAP-style vector moments, the reparametrized storage of basis and factor information, and the restriction of QR operations to selected subspaces rather than the full basis.

Source: https://www.emergentmind.com/topics/kl-soap