KL-SOAP: Hybrid Shampoo Optimizer

• KL-SOAP is a hybrid optimizer that integrates KL-Shampoo’s Kronecker factor updates with SOAP’s second-moment vector preconditioning for enhanced scaling and coupling.
• Its reparametrization supports BFloat16 storage by storing combined factor matrices and using efficient subspace QR decompositions to cut computational overhead.
• KL-SOAP improves memory efficiency and runtime while mitigating performance degradation, making it robust for high-dimensional deep learning tasks.

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" (Milligan et al., 25 May 2026), 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 (Milligan et al., 25 May 2026). 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$0. The preconditioner at time $v$1 is

$v$2

and the parameter update is

$v$3

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 (Milligan et al., 25 May 2026). Let

$v$4

Then

$v$5

and the factor updates are

$v$6

Using $v$7, KL-SOAP updates the vector moment as

$v$8

then defines

$v$9

This formulation makes explicit that KL-SOAP uses two distinct state types. The first is the Kronecker-factored curvature state $A_1\in\mathbb R^{d_1\times d_1}$0, inherited from KL-Shampoo. The second is the coordinatewise second moment $A_1\in\mathbb R^{d_1\times d_1}$1 in the time-varying basis $A_1\in\mathbb R^{d_1\times d_1}$2, 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 (Milligan et al., 25 May 2026). Its proposed remedy is a reparametrization of the preconditioner.

Instead of storing $A_1\in\mathbb R^{d_1\times d_1}$3 and $A_1\in\mathbb R^{d_1\times d_1}$4 in FP32, the method stores the combined factor

$A_1\in\mathbb R^{d_1\times d_1}$5

together with the orthogonal basis $A_1\in\mathbb R^{d_1\times d_1}$6. Algebraically,

$A_1\in\mathbb R^{d_1\times d_1}$7

and then

$A_1\in\mathbb R^{d_1\times d_1}$8

The exposition states that storing $A_1\in\mathbb R^{d_1\times d_1}$9, described as nearly diagonal, and $A_2\in\mathbb R^{d_2\times d_2}$0 in BFloat16 is safe, while QR is performed in FP32 only on small subspaces. It also states that the vector $A_2\in\mathbb R^{d_2\times d_2}$1 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 (Milligan et al., 25 May 2026). Rather than performing a full $A_2\in\mathbb R^{d_2\times d_2}$2 QR, the method chooses an index set $A_2\in\mathbb R^{d_2\times d_2}$3 of size $A_2\in\mathbb R^{d_2\times d_2}$4, for example $A_2\in\mathbb R^{d_2\times d_2}$5 with $A_2\in\mathbb R^{d_2\times d_2}$6, and partitions

$A_2\in\mathbb R^{d_2\times d_2}$7

It then performs

$A_2\in\mathbb R^{d_2\times d_2}$8

and updates only those rows and columns of $A_2\in\mathbb R^{d_2\times d_2}$9 that touch $P_{\rm shampoo}=A_1\otimes A_2.$0: $P_{\rm shampoo}=A_1\otimes A_2.$1 while $P_{\rm shampoo}=A_1\otimes A_2.$2 and $P_{\rm shampoo}=A_1\otimes A_2.$3 are rotated similarly.

The cost is given explicitly as $P_{\rm shampoo}=A_1\otimes A_2.$4 for QR plus $P_{\rm shampoo}=A_1\otimes A_2.$5 for the two block updates, versus $P_{\rm shampoo}=A_1\otimes A_2.$6 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. $P_{\rm shampoo}=A_1\otimes A_2.$7, $P_{\rm shampoo}=A_1\otimes A_2.$8, and $P_{\rm shampoo}=A_1\otimes A_2.$9 are stored in BFloat16; gradients and small QR computations remain in FP32. The algorithm computes $\Phi=U_1\otimes U_2$0 in FP32, projects it into the current basis via $\Phi=U_1\otimes U_2$1, updates $\Phi=U_1\otimes U_2$2, periodically performs the subspace QR step every $\Phi=U_1\otimes U_2$3 iterations, applies the basis transforms implicitly through two matrix multiplications rather than materializing $\Phi=U_1\otimes U_2$4, updates $\Phi=U_1\otimes U_2$5, rescales by $\Phi=U_1\otimes U_2$6, and maps the preconditioned direction back to parameter space.

5. Computational and memory characteristics

The paper gives a cost analysis in terms of $\Phi=U_1\otimes U_2$7 (Milligan et al., 25 May 2026). A full QR costs $\Phi=U_1\otimes U_2$8 per factor, for a total of $\Phi=U_1\otimes U_2$9. With subspace QR and block size $P_{\rm shampoo}$0,

• QR on the $P_{\rm shampoo}$1 block costs $P_{\rm shampoo}$2.
• Rotating the two touched blocks costs $P_{\rm shampoo}$3.
• The combined per-factor cost is $P_{\rm shampoo}$4.
• The total across both factors is $P_{\rm shampoo}$5.

The exposition adds that with $P_{\rm shampoo}$6, for example $P_{\rm shampoo}$7 or $P_{\rm shampoo}$8, the subspace approach can be $P_{\rm shampoo}$9 or $v=\mathbb E[(\Phi^\top g)^{\odot2}],$0 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 $v=\mathbb E[(\Phi^\top g)^{\odot2}],$1 and $v=\mathbb E[(\Phi^\top g)^{\odot2}],$2, each with $v=\mathbb E[(\Phi^\top g)^{\odot2}],$3 numbers. In FP32 this is

$v=\mathbb E[(\Phi^\top g)^{\odot2}],$4

whereas in BFloat16 it is

$v=\mathbb E[(\Phi^\top g)^{\odot2}],$5

which yields 50% savings. The vector $v=\mathbb E[(\Phi^\top g)^{\odot2}],$6 is also stored in BFloat16 for $v=\mathbb E[(\Phi^\top g)^{\odot2}],$7 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 (Milligan et al., 25 May 2026). 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 $v=\mathbb E[(\Phi^\top g)^{\odot2}],$8 with at most 0.001 degradation in loss.

Method	Parametrization	FP32 loss $v=\mathbb E[(\Phi^\top g)^{\odot2}],$9 BFP16 loss ($$P_{\rm soap}=\Phi\;\mathrm{Diag}(v)\;\Phi^\top, \qquad P_{\rm soap}^{-\tfrac12}=\Phi\,\mathrm{Diag}(v^{-1/2})\,\Phi^\top.$$0)
KL-SOAP	old	3.351 $$P_{\rm soap}=\Phi\;\mathrm{Diag}(v)\;\Phi^\top, \qquad P_{\rm soap}^{-\tfrac12}=\Phi\,\mathrm{Diag}(v^{-1/2})\,\Phi^\top.$$1 3.362 (+0.011)
KL-SOAP	new	3.346 $$P_{\rm soap}=\Phi\;\mathrm{Diag}(v)\;\Phi^\top, \qquad P_{\rm soap}^{-\tfrac12}=\Phi\,\mathrm{Diag}(v^{-1/2})\,\Phi^\top.$$2 3.346 (+0.000) ★
KL-Shampoo	old	3.345 $$P_{\rm soap}=\Phi\;\mathrm{Diag}(v)\;\Phi^\top, \qquad P_{\rm soap}^{-\tfrac12}=\Phi\,\mathrm{Diag}(v^{-1/2})\,\Phi^\top.$$3 3.345 (+0.000)
KL-Shampoo	new	3.345 $$P_{\rm soap}=\Phi\;\mathrm{Diag}(v)\;\Phi^\top, \qquad P_{\rm soap}^{-\tfrac12}=\Phi\,\mathrm{Diag}(v^{-1/2})\,\Phi^\top.$$4 3.344 (−0.001) ★
SOAP	old	3.358 $$P_{\rm soap}=\Phi\;\mathrm{Diag}(v)\;\Phi^\top, \qquad P_{\rm soap}^{-\tfrac12}=\Phi\,\mathrm{Diag}(v^{-1/2})\,\Phi^\top.$$5 3.372 (+0.014)
SOAP	new	3.353 $$P_{\rm soap}=\Phi\;\mathrm{Diag}(v)\;\Phi^\top, \qquad P_{\rm soap}^{-\tfrac12}=\Phi\,\mathrm{Diag}(v^{-1/2})\,\Phi^\top.$$6 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 $$P_{\rm soap}=\Phi\;\mathrm{Diag}(v)\;\Phi^\top, \qquad P_{\rm soap}^{-\tfrac12}=\Phi\,\mathrm{Diag}(v^{-1/2})\,\Phi^\top.$$7 to $$P_{\rm soap}=\Phi\;\mathrm{Diag}(v)\;\Phi^\top, \qquad P_{\rm soap}^{-\tfrac12}=\Phi\,\mathrm{Diag}(v^{-1/2})\,\Phi^\top.$$8. This suggests that, in the reported setting, the BFloat16 penalty is removed for KL-SOAP.

The paper further reports that for $$P_{\rm soap}=\Phi\;\mathrm{Diag}(v)\;\Phi^\top, \qquad P_{\rm soap}^{-\tfrac12}=\Phi\,\mathrm{Diag}(v^{-1/2})\,\Phi^\top.$$9 and $(A_1,A_2)$0, subspace QR is $(A_1,A_2)$1 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 $(A_1,A_2)$2, the subspace frequency is 5–10 steps, the inner loop usually uses $(A_1,A_2)$3 subspace QR per update, the damping for vector $(A_1,A_2)$4 is $(A_1,A_2)$5, the Shampoo damping for $(A_1,A_2)$6 is $(A_1,A_2)$7, and the optimizer coefficients are $(A_1,A_2)$8 and $(A_1,A_2)$9. 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.