Understanding and improving Shampoo and SOAP via Kullback-Leibler Minimization

## Abstract

Abstract Shampoo and its efficient variant, SOAP, employ structured second-moment estimations and have shown strong performance for training neural networks (NNs). In practice, however, Shampoo typically requires step-size grafting with Adam to be competitive, and SOAP mitigates this by applying Adam in Shampoo’s eigenbasis—at the cost of additional memory overhead from Adam in both methods. Prior analyses have largely relied on the Frobenius norm to motivate these estimation schemes. We instead recast their estimation procedures as covariance estimation under Kullback-Leibler (KL) divergence minimization, revealing a previously overlooked theoretical limitation and motivating principled redesigns. Building on this perspective, we develop KL-Shampoo and KL-SOAP , practical schemes that match or exceed the performance of Shampoo and SOAP in NN pre-training while achieving SOAP-level per-iteration runtime. Notably, KL-Shampoo does not rely on Adam to attain competitive performance, eliminating the memory overhead introduced by Adam. Across our experiments, KL-Shampoo consistently outperforms SOAP, Shampoo, and even KL-SOAP, establishing the KL-based approach as a compelling foundation for designing structured methods in NN optimization.

## 1 Introduction

Optimizer Shampoo (Gupta et al., 2018) has recently dethroned Adam (Kingma & Ba, 2015) as the winner of several competitions in training a wide range of neural network (NN) models (Dahl et al., 2023; Kasimbeg et al., 2025). Consequently, Shampoo and its variant, SOAP (Vyas et al., 2025a), have drawn increasing attention. In practice, Shampoo requires step-size grafting with Adam to achieve competitive performance (Agarwal et al., 2020; Shi et al., 2023). SOAP addresses this by applying Adam in Shampoo’s eigenbasis and further reducing per-iteration runtime.
However, reliance on Adam introduces additional memory overhead in both methods. Prior work (Morwani et al., 2025; Eschenhagen et al., 2025; An et al., 2025; Xie et al., 2025) has investigated their structural preconditioner schemes—which approximate the flattened gradient 2^nd moment (Duchi et al., 2011)—through the Frobenius norm.
However, few studies have examined these schemes from the perspective of Kullback–Leibler (KL) divergence.
Compared to the Frobenius norm, the KL divergence between zero-mean Gaussian covariance matrices is more appropriate for interpreting Shampoo’s and SOAP’s preconditioners as Gaussian covariance (Amari, 2016; Minh & Murino, 2017), since the second moment they approximate can be viewed as the covariance matrix of a zero-mean Gaussian. A similar KL perspective has provided a unified framework to interpret (Fletcher, 1991; Waldrip & Niven, 2016) and extend (Kanamori & Ohara, 2013a, b) structural preconditioner estimation in quasi-Newton methods such as BFGS and DFP—something the Frobenius norm does not. Moreover, the KL divergence intrinsically respects the symmetric positive-definite constraint (Amari, 2016; Minh & Murino, 2017) that preconditioners in Shampoo and SOAP must satisfy as adaptive (preconditioned) methods (Nesterov et al., 2018)—a property the Frobenius norm lacks. This constraint implies that the entries of the preconditioning matrix do not play equivalent roles and therefore should not be treated equally (Pennec et al., 2006; Bhatia, 2007)—a point the Frobenius norm ignores.

In this work, we introduce a KL perspective that interprets the estimation schemes of Shampoo and SOAP as solutions to KL-minimization problems for covariance estimation. Our approach naturally extends to tensor-valued settings, where some existing theoretical interpretations may not apply. This perspective reveals a key limitation (illustrated in [Fig.˜5](https://arxiv.org/html/2509.03378v4#S6.F5)): the Kronecker-structured estimators used by Shampoo and SOAP do not adequately solve the corresponding KL-minimization problem. This limitation, in turn, opens new opportunities for improvement.
Leveraging this insight, we refine the estimation rules of Shampoo and SOAP
and develop practical KL-based schemes—KL-Shampoo and KL-SOAP—that meet or exceed the performance of Shampoo and SOAP for NN (pre-)training while maintaining SOAP-level per-iteration runtime.
Notably, KL-Shampoo does not rely on Adam to achieve competitive performance, thereby avoiding Adam’s additional memory overhead ([Footnote˜1](https://arxiv.org/html/2509.03378v4#footnote1)).
In addition, our practical techniques developed for KL-Shampoo ([Sec.˜4](https://arxiv.org/html/2509.03378v4#S4)) can be adapted to strengthen existing Shampoo variants and make the trace-scaling variant competitive without step-size grafting ([Fig.˜10](https://arxiv.org/html/2509.03378v4#A8.F10), [Appx.˜H](https://arxiv.org/html/2509.03378v4#A8)), while also achieving SOAP-level per-iteration runtime ([Fig.˜10](https://arxiv.org/html/2509.03378v4#A8.F10), [Appx.˜H](https://arxiv.org/html/2509.03378v4#A8)).
Empirically, we show that KL-based methods are competitive for training a range of NNs and remain as flexible as Shampoo and SOAP for tensor-valued weights. Surprisingly, KL-Shampoo consistently outperforms the other methods in our experiments ([Figs.˜1](https://arxiv.org/html/2509.03378v4#S2.F1) and [5](https://arxiv.org/html/2509.03378v4#S6.F5)). Overall, our KL-based approach provides a principled way for designing structured methods in NN optimization.

## 2 Background

Figure: Figure 1: Empirical results (random search using 150 runs for each method) on language models demonstrate the advantages of KL-based methods over Shampoo and SOAP while matching SOAP’s pre-iteration runtime. All methods take the same number of iterations in these experiments. Surprisingly, KL-Shampoo also outperforms KL-SOAP. We include the best Shampoo run based on a state-of-the-art implementation from Meta (Shi et al., 2023) in the plots.
Refer to caption: x1.png

#### Notation

For presentation simplicity, we focus on matrix-valued weights and the optimization update for a single parameter matrix ${\bm{\Theta}}\in\mathcal{R}^{d_{a}\times d_{b}}$, rather than a set of weight matrices for NN training.
We use $\operatorname{Mat}(\cdot)$ to unflatten its input vector into a matrix and $\operatorname{vec}(\cdot)$ to flatten its input matrix into a vector. For example, ${\bm{\theta}}:=\operatorname{vec}({\bm{\Theta}})$ is the flattened weight vector and ${\bm{\Theta}}\equiv\mathrm{Mat}({\bm{\theta}})$ is the unflattened weight matrix.
Vector ${\bm{g}}$ is a (flattened) gradient vector for the weight matrix.
We denote $\gamma$, $\beta_{2}$ and ${\bm{S}}$ to be a step size, a weight for moving average, and a preconditioning matrix for an adaptive method, respectively.
$\mathrm{Diag}(\cdot)$ returns a diagonal matrix whose diagonal entries are given by its input vector, whilst $\mathrm{diag}(\cdot)$ extracts the diagonal entries of its input matrix as a vector.

#### Shampoo

Given a matrix gradient ${\bm{G}}$ and the flattened gradient ${\bm{g}}=\mathrm{vec}({\bm{G}})$, the original Shampoo method (Gupta et al., 2018) considers a *Kronecker-factored* approximation, $({\bm{S}}_{a})^{2p}\otimes({\bm{S}}_{b})^{2p}$, of the flattened gradient second moment, $\mathbb{E}_{{\bm{g}}}[{\bm{g}}{\bm{g}}^{\top}]$, where $p$ denotes a matrix power,
${\bm{S}}_{a}:=\mathbb{E}_{{\bm{g}}}[{\bm{G}}{\bm{G}}^{\top}]$, ${\bm{S}}_{b}:=\mathbb{E}_{{\bm{g}}}[{\bm{G}}^{\top}{\bm{G}}]$,
and $\otimes$ denotes a Kronecker product.
In practice, we often approximate the expectation, $\mathbb{E}_{{\bm{g}}}[{\bm{g}}{\bm{g}}^{\top}]$, with an exponentially moving average (EMA) on the outer product (Morwani et al., 2025).
The original Shampoo method uses the $\nicefrac{{1}}{{4}}$ power (i.e., $p=\nicefrac{{1}}{{4}}$) and other works
(Anil et al., 2020; Shi et al., 2023; Morwani et al., 2025) suggest using the $\nicefrac{{1}}{{2}}$ power (i.e., $p=\nicefrac{{1}}{{2}}$).
At each iteration, Shampoo follows this update rule with EMA on ${\bm{S}}_{a}$ and ${\bm{S}}_{b}$:

$$ $\displaystyle{\bm{S}}_{a}$ $\displaystyle\leftarrow(1-\beta_{2}){\bm{S}}_{a}+\beta_{2}{\bm{G}}{\bm{G}}^{\top},\quad{\bm{S}}_{b}\leftarrow(1-\beta_{2}){\bm{S}}_{b}+\beta_{2}{\bm{G}}^{\top}{\bm{G}}\quad\text{{\small\color[rgb]{.5,.5,.5}\definecolor[named]{pgfstrokecolor}{rgb}{.5,.5,.5}\pgfsys@color@gray@stroke{.5}\pgfsys@color@gray@fill{.5}(Kronecker 2\textsuperscript{nd} moment est.)}},$ $\displaystyle{\bm{\theta}}$ $\displaystyle\leftarrow{\bm{\theta}}-\gamma{\bm{S}}^{-\nicefrac{{1}}{{2}}}{\bm{g}}\iff{\bm{\Theta}}\leftarrow{\bm{\Theta}}-\gamma{\bm{S}}_{a}^{-p}{\bm{G}}{\bm{S}}_{b}^{-p}\quad\text{{\small\color[rgb]{.5,.5,.5}\definecolor[named]{pgfstrokecolor}{rgb}{.5,.5,.5}\pgfsys@color@gray@stroke{.5}\pgfsys@color@gray@fill{.5}(preconditioning)}},$ (1) $$

where
${\bm{S}}:={\bm{S}}_{a}^{2p}\otimes{\bm{S}}_{b}^{2p}$ is Shampoo’s preconditioning matrix, and we leverage the Kronecker structure of ${\bm{S}}$ to move from the left expression to the right expression in the second line.

Shampoo’s implementation employs eigendecomposition.
Shampoo is typically implemented by using the eigendecomposition of ${\bm{S}}_{k}$, such as ${\bm{Q}}_{k}\mathrm{Diag}({\bm{\lambda}}_{k}){\bm{Q}}_{k}^{\top}=\mathrm{eigen}({\bm{S}}_{k})$, for $k\!\in\!\{a,b\}$, every few steps and storing ${\bm{Q}}_{k}$ and ${\bm{\lambda}}_{k}$ (Anil et al., 2020; Shi et al., 2023).
Therefore, the power of ${\bm{S}}_{k}$ is computed using an elementwise power in ${\bm{\lambda}}_{k}$ such as ${\bm{S}}_{k}^{-p}={\bm{Q}}_{k}\mathrm{Diag}\big({\bm{\lambda}}_{k}^{\odot-p}\big){\bm{Q}}_{k}^{\top}$, where $\cdot^{\odot p}$ denotes elementwise $p$-th power. This computation becomes an approximation if the decomposition is not performed at every step.

Using Adam for Shampoo’s stabilization increases memory usage.
If the eigendecomposition is applied infrequently to reduce iteration cost, Shampoo has to apply step-size grafting with Adam to maintain performance (Agarwal et al., 2020; Shi et al., 2023) as empirically shown in [Fig.˜2](https://arxiv.org/html/2509.03378v4#S3.F2).
Unfortunately, this increases its memory usage introduced by Adam (see [Footnote˜1](https://arxiv.org/html/2509.03378v4#footnote1)).

#### SOAP

SOAP improves Shampoo with the $p=\nicefrac{{1}}{{2}}$ power by running Adam in the eigenbasis of Shampoo’s preconditioner $({\bm{S}}_{a})^{2p}\otimes({\bm{S}}_{b})^{2p}={\bm{S}}_{a}\otimes{\bm{S}}_{b}$.
Notably, SOAP reuses Shampoo’s Kronecker estimation rule for computing the eigenbasis ${\bm{Q}}:={\bm{Q}}_{a}\otimes{\bm{Q}}_{b}$ and incorporates Adam’s 2^nd moment, denoted by ${\bm{d}}$, for preconditioning, where ${\bm{Q}}_{k}$ is Shampoo’s Kronecker eigenbasis ${\bm{S}}_{k}$ for $k\in\{a,b\}$ defined above.
As a result, SOAP effectively employs an *augmented* preconditioner, ${\bm{S}}:={\bm{Q}}\mathrm{Diag}({\bm{d}}){\bm{Q}}^{\top}$, which cannot be expressed as a Kronecker product of any two matrices with the same shape as ${\bm{S}}_{a}$ and ${\bm{S}}_{b}$.
Because we omit momentum (i.e. let Adam’s $\beta_{1}=0$), SOAP takes the following step with the Adam update becoming an RMSProp update (Tieleman & Hinton, 2012):

$$ $\displaystyle{\bm{S}}_{a}$ $\displaystyle\leftarrow(1-\beta_{2}){\bm{S}}_{a}+\beta_{2}{\bm{G}}{\bm{G}}^{\top},\quad{\bm{S}}_{b}\leftarrow(1-\beta_{2}){\bm{S}}_{b}+\beta_{2}{\bm{G}}^{\top}{\bm{G}}\quad\text{{\small\color[rgb]{.5,.5,.5}\definecolor[named]{pgfstrokecolor}{rgb}{.5,.5,.5}\pgfsys@color@gray@stroke{.5}\pgfsys@color@gray@fill{.5} (Shampoo's 2\textsuperscript{nd} moment est.)}},$ $\displaystyle{\bm{d}}$ $\displaystyle\leftarrow(1-\beta_{2}){\bm{d}}+\beta_{2}\hat{{\bm{g}}}^{\odot 2}\quad\text{{\small\color[rgb]{.5,.5,.5}\definecolor[named]{pgfstrokecolor}{rgb}{.5,.5,.5}\pgfsys@color@gray@stroke{.5}\pgfsys@color@gray@fill{.5}(RMSProp's diagonal 2\textsuperscript{nd} moment est. in the eigenbasis)}},$ $\displaystyle{\bm{\theta}}$ $\displaystyle\leftarrow{\bm{\theta}}-\gamma{\bm{S}}^{-\frac{1}{2}}{\bm{g}}\iff{\bm{\Theta}}\leftarrow{\bm{\Theta}}-\gamma{\bm{Q}}_{a}^{\top}\operatorname{Mat}\left(\frac{\hat{{\bm{g}}}}{\sqrt{{\bm{d}}}}\right){\bm{Q}}_{b}\quad\text{{\small\color[rgb]{.5,.5,.5}\definecolor[named]{pgfstrokecolor}{rgb}{.5,.5,.5}\pgfsys@color@gray@stroke{.5}\pgfsys@color@gray@fill{.5}(preconditioning)}},$ (2) $$

where $\hat{{\bm{g}}}:={\bm{Q}}^{\top}{\bm{g}}=\operatorname{vec}({\bm{Q}}_{a}^{\top}{\bm{G}}{\bm{Q}}_{b})$ is a “projected” gradient vector in eigenbasis ${\bm{Q}}$ and recall that ${\bm{S}}:={\bm{Q}}\mathrm{Diag}({\bm{d}}){\bm{Q}}^{\top}$ is SOAP’s preconditioner.
Here, we leverage the Kronecker structure and orthogonality of the eigenbasis to move from the left to the right in the last line of [Eq.˜2](https://arxiv.org/html/2509.03378v4#S2.E2).
Note that this EMA weight $\beta_{2}$ is defined as $1-\beta_{2}^{\text{(Adam)}}$, where $\beta_{2}^{\text{(Adam)}}$ is Adam’s (RMSProp’s) $\beta_{2}$.
We use this definition rather than Adam’s because we want to further interpret this moving-average scheme through the lens of our KL perspective.

SOAP’s implementation utilizes QR decomposition. SOAP requires only the eigenbasis, which can be approximated via a QR decomposition, whereas Shampoo typically requires an eigendecomposition to compute both the eigenbasis and the eigenvalues. Vyas et al. (2025a) therefore suggest replacing the slower eigendecomposition with the faster QR decomposition, such as ${\bm{Q}}_{k}\leftarrow\mathrm{qr}({\bm{S}}_{k}{\bm{Q}}_{k})$ for $k\in\{a,b\}$.
This makes SOAP more computationally efficient than Shampoo.

Runing Adam in the eigenbasis increases memory usage. Introducing Adam’s (RMSProp’s) 2^nd moment estimation increases SOAP’s memory consumption (see [Footnote˜1](https://arxiv.org/html/2509.03378v4#footnote1)).
This is because this estimation, ${\bm{d}}\in\mathcal{R}^{d_{a}d_{b}\times 1}$, uses extra memory and cannot be expressed as a Kronecker product of any two vectors, such as ${\bm{d}}\neq{\bm{r}}_{a}\otimes{\bm{r}}_{b}$, where ${\bm{r}}_{a}\in\mathcal{R}^{d_{a}\times 1}$ and ${\bm{r}}_{b}\in\mathcal{R}^{d_{b}\times 1}$.

The original Shampoo’s Kronecker estimation rule ($p=\nicefrac{{1}}{{4}}$) (Gupta et al., 2018; Duvvuri et al., 2024) is proposed based on a matrix Loewner bound (Löwner, 1934), while recent estimation rules ($p=\nicefrac{{1}}{{2}}$) (Morwani et al., 2025; Eschenhagen et al., 2025) focus on bounds induced by the Frobenius norm.
SOAP reuses Shampoo’s Kronecker estimation rule and additionally introduces Adam’s (RMSProp’s) 2^nd-moment estimation rule in the eigenbasis (Vyas et al., 2025a).
None of these works interpret or motivate their estimation rules as covariance estimation, thereby missing the opportunity to introduce the KL perspective.

**Table 1: Memory usage of each method considered in this work. We store and update (^1^11Since the current QR/eigen implementation in PyTorch does not support half-precision (bfloat16), we perform QR/eigen to compute ${\bm{Q}}_{k}$ for $k\in\{a,b\}$ in single precision (float32) and then cast and store them in half precision (bfloat16). Other matrices and vectors can be computed and updated in half precision.)all of them in half precision (bfloat16). The memory overhead introduced by Adam is highlighted in red. Note that SOAP’s and KL-SOAP’s preconditioners, ${\bm{Q}}\mathrm{Diag}({\bm{d}}){\bm{Q}}^{\top}$, can not be expressed as a Kronecker product due to the augmented eigenvalues ${\bm{d}}$, while Shampoo’s and KL-Shampoo’s preconditioners, ${\bm{Q}}\mathrm{Diag}({\bm{\lambda}}_{a}\otimes{\bm{\lambda}}_{b}){\bm{Q}}^{\top}$, can, where ${\bm{Q}}:={\bm{Q}}_{a}\otimes{\bm{Q}}_{b}$.**
|  | Shampoo | SOAP | KL-Shampoo | KL-SOAP |
| --- | --- | --- | --- | --- |
| Kronecker factors (${\bm{S}}_{a}$, ${\bm{S}}_{b}$) | $d_{a}^{2}+d_{b}^{2}$ | $d_{a}^{2}+d_{b}^{2}$ | $d_{a}^{2}+d_{b}^{2}$ | $d_{a}^{2}+d_{b}^{2}$ |
| Kronecker factors’ eigenbasis (${\bm{Q}}_{a}$, ${\bm{Q}}_{b}$) | $d_{a}^{2}+d_{b}^{2}$ | $d_{a}^{2}+d_{b}^{2}$ | $d_{a}^{2}+d_{b}^{2}$ | $d_{a}^{2}+d_{b}^{2}$ |
| Kronecker factors’ eigenvalues (${\bm{\lambda}}_{a}$, ${\bm{\lambda}}_{b}$) | $d_{a}+d_{b}$ | N/A | $d_{a}+d_{b}$ | $d_{a}+d_{b}$ |
| Adam’s 2^nd moment in the eigenbasis (${\bm{d}}$)<br>(interpreted as augmented eigenvalues, [Sec.˜5](https://arxiv.org/html/2509.03378v4#S5)) | N/A | $d_{a}d_{b}$ | N/A | $d_{a}d_{b}$ |
| Momentum | $d_{a}d_{b}$ | $d_{a}d_{b}$ | $d_{a}d_{b}$ | $d_{a}d_{b}$ |
| Step-size grafting with Adam | $d_{a}d_{b}$ | N/A | N/A | N/A |

## 3 Second Moment Estimation via Kullback–Leibler Minimization

We first focus on Shampoo with $p=\nicefrac{{1}}{{2}}$ and show that its second-moment estimation can be viewed as a structured covariance estimation problem solved via Kullback–Leibler (KL) minimization. This perspective reflects the natural connection between the flattened gradient second moment (Duchi et al., 2011) that Shampoo approximates and a covariance matrix.
From the KL perspective, we reveal a previously unrecognized limitation of Shampoo’s estimation rule: the Kronecker-structured estimators used by Shampoo and SOAP do not adequately solve the corresponding KL-minimization problem.
This limitation, in turn, opens new opportunities for improvement.
Building on this insight, we propose a KL-based estimation scheme for Shampoo, which we term the idealized KL-Shampoo.

#### KL Minimization

For simplicity, we begin by introducing a KL perspective in a matrix-valued case and drop subscripts when referring to the flattened gradient 2^nd moment, like $\mathbb{E}[{\bm{g}}{\bm{g}}^{\top}]:=\mathbb{E}_{{\bm{g}}}[{\bm{g}}{\bm{g}}^{\top}]$, where ${\bm{g}}=\mathrm{vec}({\bm{G}})$ is a flattened gradient vector of
a matrix-valued gradient ${\bm{G}}\in\mathcal{R}^{d_{a}\times d_{b}}$.
The goal is to estimate a Kronecker-structured preconditioning matrix, ${\bm{S}}={\bm{S}}_{a}\otimes{\bm{S}}_{b}$, that closely approximates the 2^nd moment, where ${\bm{S}}_{a}\in\mathcal{R}^{d_{a}\times d_{a}}$ and ${\bm{S}}_{b}\in\mathcal{R}^{d_{b}\times d_{b}}$ are both symmetric positive-definite (SPD).
Motivated by the natural connection between the second moment and a covariance matrix, we treat these as covariance matrices of zero-mean Gaussian distributions and achieve this goal by minimizing the KL divergence between the two distributions,

where $\smash{\mathbb{E}[{\bm{g}}{\bm{g}}^{\top}]}$ and ${\bm{S}}$ are considered as Gaussian’s covariance, $\mathrm{det}(\cdot)$ denotes the matrix determinant of its input, and $\kappa\!\geq\!0$ is a damping weight to ensure the positive-definiteness of $\smash{\mathbb{E}[{\bm{g}}{\bm{g}}^{\top}]\!+\!\kappa{\bm{I}}}$ if necessary.
Mathematically, this KL divergence coincides (up to a factor of $\nicefrac{{1}}{{2}}$) with the log-determinant divergence widely used in matrix optimization (Dhillon & Tropp, 2008; Kulis et al., 2009; Sra, 2016) for *SPD matrices* rather than Gaussian distributions, which does not require a zero-mean assumption for SPD matrices.
This additional zero-mean Gaussian viewpoint provides a probabilistic interpretation of this SPD-aware “distance”, even when the target matrix is not itself a second moment, such as the curvature matrix used in quasi-Newton methods (Fletcher, 1991; Waldrip & Niven, 2016).
Moreover, the KL divergence naturally extends to tensor-valued cases, such as a 3D tensor gradient, ${\bm{G}}\in\mathcal{R}^{d_{a}\times d_{b}\times d_{c}}$, by considering a structured preconditioner ${\bm{S}}={\bm{S}}_{a}\otimes{\bm{S}}_{b}\otimes{\bm{S}}_{c}$ to approximate the flattened gradient second moment, $\smash{\mathbb{E}[{\bm{g}}{\bm{g}}^{\top}]}$, where matrix ${\bm{S}}_{k}\in\mathcal{R}^{d_{k}\times d_{k}}$ is SPD for $k\in\{a,b,c\}$.

#### Justification of using the KL divergence

Many works (Morwani et al., 2025; Eschenhagen et al., 2025; An et al., 2025; Xie et al., 2025) primarily focus on matrix-valued weights and interpret Shampoo’s and SOAP’s estimation rules from the Frobenius-norm perspective. However, this norm does not account for the SPD constraint implicitly imposed on Shampoo’s and SOAP’s preconditioners, which ensures that the preconditioned gradient direction is a descent direction (Nesterov et al., 2018). As emphasized in the literature (Pennec et al., 2006; Bhatia, 2007), it is more appropriate to consider a “distance” that respects this constraint. We adopt the KL divergence because it naturally incorporates the SPD constraint, is widely used for covariance estimation (Amari, 2016; Minh & Murino, 2017), and provides a unified framework for reinterpreting and improving Shampoo’s estimation—even in tensor-valued settings where some existing interpretations based on singular value decomposition (SVD) (Van Loan & Pitsianis, 1993) may not apply.
In numerical optimization, the KL divergence, known as a merit function (Byrd & Nocedal, 1989), offers a unifying interpretation (Fletcher, 1991; Waldrip & Niven, 2016) and extension (Kanamori & Ohara, 2013a, b) of the low-rank estimation schemes of BFGS and DFP.
In contrast, the standard Frobenius norm cannot recover these updates without additional weighting (see Sec. 6.1 of Nocedal & Wright (2006)). In statistical estimation and inference, this KL divergence is also preferred over the Frobenius norm (James & Stein, 1961; Kivinen & Warmuth, 1999; Davis & Dhillon, 2006; Khan & Lin, 2017; Lin et al., 2019; Kunstner et al., 2021).

### 3.1 Interpreting Shampoo’s estimation as covariance estimation

Similar to existing works (Morwani et al., 2025; Eschenhagen et al., 2025; Vyas et al., 2025a), we first disable the moving average (i.e., let $\beta_{2}=1$) for our descriptions and focus on Shampoo with power $p=\nicefrac{{1}}{{2}}$, presenting a KL minimization perspective and interpreting its
estimation rule
from this perspective.
We will show that Shampoo’s estimation rule can be obtained by solving a KL minimization problem.

###### Claim 1

(Shampoo’s Kronecker-based covariance estimation)
The optimal solution of KL minimization $\min_{{\bm{S}}_{a}}\mathrm{KL}\big(\mathbb{E}[{\bm{g}}{\bm{g}}^{\top}],{\bm{S}}\big)$ with a one-sided preconditioner ${\bm{S}}=(\nicefrac{{1}}{{d_{b}}}{\bm{S}}_{a})\otimes{\bm{I}}_{b}$ is ${\bm{S}}_{a}^{*}=\mathbb{E}[{\bm{G}}{\bm{G}}^{\top}]$, where $d_{k}$ is the dimension of matrix ${\bm{S}}_{k}\in{\mathbb{R}}^{d_{k}\times d_{k}}$ for $k\in\{a,b\}$ and ${\bm{G}}=\mathrm{Mat}({\bm{g}})$.

Likewise, we can obtain the estimation rule for ${\bm{S}}_{b}$ by considering ${\bm{S}}={\bm{I}}_{a}\otimes(\nicefrac{{1}}{{d_{a}}}{\bm{S}}_{b})$.

###### Claim 1

#### Shampoo’s estimation rule as Kronecker-based covariance estimation

According to [˜1](https://arxiv.org/html/2509.03378v4#Thmclaim1) (proof in [Appx.˜A](https://arxiv.org/html/2509.03378v4#A1)), Shampoo’s estimation rule with power $p=\nicefrac{{1}}{{2}}$ in [Eq.˜1](https://arxiv.org/html/2509.03378v4#S2.E1) can be viewed as the optimal solution of a KL minimization problem (up to a constant scalar) when one Kronecker factor is updated independently and the other is fixed as the identity, which is known as a one-sided preconditioner (An et al., 2025; Xie et al., 2025).
In practice, Shampoo further approximates the required expectations using the EMA scheme in [Eq.˜1](https://arxiv.org/html/2509.03378v4#S2.E1).

Figure: Figure 2: Empirical results (random search using 150 runs for each method) on language models demonstrate that KL-Shampoo does not rely on step-size grafting with Adam to perform well. Shampoo without grafting does not perform well, even when using the state-of-the-art implementation (Shi et al., 2023). In particular, Shampoo with power $p=\nicefrac{{1}}{{2}}$ fails to train the RWKV7 model in all 150 runs when grafting is disabled.
Refer to caption: x2.png

### 3.2 Improving Shampoo’s
estimation:
Idealized KL-Shampoo

Our KL perspective reveals a key limitation—empirically demonstrated in [Fig.˜5](https://arxiv.org/html/2509.03378v4#S6.F5)—of Shampoo’s Kronecker estimation with $p=\nicefrac{{1}}{{2}}$ as a one-sided approach: it does not adequately solve the KL-minimization problem when both factors are learned jointly. Motivated by this, we design an improved estimation rule that updates the two factors simultaneously. We term this scheme as idealized KL-Shampoo, which is a two-sided approach.

###### Claim 2

(Idealized KL-Shampoo’s covariance estimation for ${\bm{S}}_{a}$ and ${\bm{S}}_{b}$)
The optimal solution of KL minimization $\min_{{\bm{S}}_{a},{\bm{S}}_{b}}\mathrm{KL}\big(\mathbb{E}[{\bm{g}}{\bm{g}}^{\top}],{\bm{S}}\big)$ with a two-sided precontioner ${\bm{S}}={\bm{S}}_{a}\otimes{\bm{S}}_{b}$ should satisfy the following condition.

$$ $\displaystyle{\bm{S}}_{a}^{*}=\frac{1}{d_{b}}\,\mathbb{E}[{\bm{G}}\big({\bm{S}}_{b}^{*}\big)^{-1}{\bm{G}}^{\top}],\quad{\bm{S}}_{b}^{*}=\frac{1}{d_{a}}\,\mathbb{E}[{\bm{G}}^{\top}\big({\bm{S}}_{a}^{*}\big)^{-1}{\bm{G}}].$ (4) $$

###### Claim 2

#### Idealized KL-Shampoo’s estimation

[˜2](https://arxiv.org/html/2509.03378v4#Thmclaim2) (proof in [Appx.˜B](https://arxiv.org/html/2509.03378v4#A2)) establishes a closed-form condition (see [Eq.˜4](https://arxiv.org/html/2509.03378v4#S3.E4)) when simultaneously learning both Kronecker factors to minimize the KL problem.
Statistically, this condition corresponds to the maximum-likelihood covariance estimator of a zero-mean matrix Gaussian (Dutilleul, 1999) when $\mathbb{E}({\bm{g}}{\bm{g}}^{\top})$ is considered as a finite average $\frac{1}{N}\sum_{i=1}^{N}{\bm{g}}_{i}{\bm{g}}_{i}^{\top}$. This is because $\mathrm{KL}(\frac{1}{N}\sum_{i=1}^{N}{\bm{g}}_{i}{\bm{g}}_{i}^{\top},{\bm{S}})=-\frac{1}{N}\sum_{i=1}^{N}\log\mathcal{N}({\bm{g}}_{i};0,{\bm{S}})+\text{const}$, where ${\bm{g}}_{i}$ is considered as a sample generated from $\mathcal{N}(0,{\bm{S}})$.
Thus, satisfying this condition immediately implies that gradient matrix ${\bm{G}}$ is generated from a zero-mean matrix Gaussian, ${\bm{G}}\sim\mathcal{MN}(\mathbf{0},{\bm{S}}_{a}^{*},{\bm{S}}_{b}^{*})$, with row-wise covariance ${\bm{S}}_{a}^{*}$ and column-wise covariance ${\bm{S}}_{b}^{*}$ obtained by maximum likelihood.
Under this condition, Shampoo-style preconditioning naturally induces matrix-Gaussian (row and column) whitening: $\big({\bm{S}}_{a}^{*}\big)^{-1/2}{\bm{G}}\big({\bm{S}}_{b}^{*}\big)^{-1/2}\sim\mathcal{MN}(\mathbf{0},{\bm{I}}_{a},{\bm{I}}_{b})$.
As we will discuss in [Sec.˜5](https://arxiv.org/html/2509.03378v4#S5), the SOAP-like projection onto the optimal eigenbasis leads to covariance diagonalization.
Our KL-based approach further extends this to tensor-Gaussian whitening for tensor-valued gradients—without the prohibitive cost typically associated with SVD-based methods. Notably, this kind of whitening arises from minimizing KL divergence, not the Frobenius norm.
In machine learning, Lin et al. (2019, 2024) treated this condition as a theoretical example of a multilinear exponential-family (see Sec 5 of Lin et al. (2019)) for Kronecker-based optimization via natural gradient descent on a matrix Gaussian manifold, while more recently, Vyas et al. (2025b) considered a similar condition motivated heuristically by gradient whitening.
However, we cannot directly use this condition due to the *correlated* update between ${\bm{S}}_{a}^{*}$ and ${\bm{S}}_{b}^{*}$.
For example, solving ${\bm{S}}_{a}^{*}$ requires knowing ${\bm{S}}_{b}^{*}$ in [Eq.˜4](https://arxiv.org/html/2509.03378v4#S3.E4) or vice versa.
In practice, this condition is unachievable because the expectations in [Eq.˜4](https://arxiv.org/html/2509.03378v4#S3.E4) must be approximated.
Thus, we consider an estimated ${\bm{S}}_{k}$ to approximate ${\bm{S}}_{k}^{*}$ for $k\in\{a,b\}$ and propose an exponential moving average (EMA) scheme:

$$ $\displaystyle{\bm{S}}_{a}\leftarrow(1-\beta_{2}){\bm{S}}_{a}+\frac{\beta_{2}}{d_{b}}{\bm{G}}{\color[rgb]{0.75,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{0.75,0,0}{\bm{S}}_{b}^{-1}}{\bm{G}}^{\top},\quad{\bm{S}}_{b}\leftarrow(1-\beta_{2}){\bm{S}}_{b}+\frac{\beta_{2}}{d_{a}}\,{\bm{G}}^{\top}{\color[rgb]{0.75,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{0.75,0,0}{\bm{S}}_{a}^{-1}}{\bm{G}}.$ (5) $$

Our KL perspective allows us to further justify this EMA scheme as a stochastic proximal-gradient step (see [˜3](https://arxiv.org/html/2509.03378v4#Thmclaim3) and a proof in [Appx.˜C](https://arxiv.org/html/2509.03378v4#A3)) and establish a formal connection to the theoretical example of Lin et al. (2019, 2024).
Notably, our approach uses ${\bm{S}}^{-\nicefrac{{1}}{{2}}}$ for preconditioning ([Eq.˜1](https://arxiv.org/html/2509.03378v4#S2.E1)), following Shampoo, whereas Lin et al. (2019, 2024) propose using ${\bm{S}}^{-1}$.
A straightforward implementation of our scheme is computationally expensive, since it requires expensive matrix inversions (highlighted in red in [Eq.˜5](https://arxiv.org/html/2509.03378v4#S3.E5)) and the slow eigendecomposition for Shampoo-type preconditioning (e.g., ${\bm{S}}^{-\nicefrac{{1}}{{2}}}$).
However, these issues can be alleviated—in [Sec.˜4](https://arxiv.org/html/2509.03378v4#S4) we propose an efficient implementation.

###### Claim 3

(KL-Shampoo’s moving average scheme)
The moving average scheme for ${\bm{S}}_{k}$ ([Eq.˜5](https://arxiv.org/html/2509.03378v4#S3.E5)) in idealized KL-Shampoo is a stochastic proximal-gradient step with step-size $\beta_{2}$ to solve the KL minimization problem in [Eq.˜3](https://arxiv.org/html/2509.03378v4#S3.E3), for $k\in\{a,b\}$. Recall that $\beta_{2}$ in [Eq.˜5](https://arxiv.org/html/2509.03378v4#S3.E5) is closely related to Adam’s $\beta_{2}$ as $\beta_{2}=1-\beta_{2}^{\text{(Adam)}}$, where
$\beta_{2}^{\text{(Adam)}}$ is Adam’s $\beta_{2}$ .

###### Claim 3

### 3.3 Comparison with Kronecker Schemes Using Alternative Divergences

#### Frobenius norm (F-Shampoo)

Morwani et al. (2025) consider a two-sided Shampoo variant based on the Frobenius norm and derive the optimal solution via rank-1 singular value decomposition (SVD) of the second moment $\mathbb{E}[{\bm{g}}{\bm{g}}^{\top}]$ (Van Loan & Pitsianis, 1993).
However, this solution is often unattainable in practice and is computationally expensive for two reasons:
(1) the expectation
$\mathbb{E}[{\bm{g}}{\bm{g}}^{\top}]$
must be approximated; and
(2) performing the SVD is costly—yielding complexity ($O(d_{a}^{2}d_{b}^{2})$) in general even for rank-1 SVD—which is higher than the eigen decompositions with complexity
($O(d_{k}^{3})$) for $k\in\{a,b\}$ that we aim to avoid.
Instead, we analyze the stationarity conditions
([˜6](https://arxiv.org/html/2509.03378v4#Thmclaim6), [Appx.˜F](https://arxiv.org/html/2509.03378v4#A6)) and derive a new variant, idealized F-Shampoo ([Fig.˜6](https://arxiv.org/html/2509.03378v4#A6.F6.fig1), [Appx.˜F](https://arxiv.org/html/2509.03378v4#A6)), that is structurally similar to KL-Shampoo.
While a straightforward implementation of F-Shampoo performs poorly in practice, the techniques ([Sec.˜4](https://arxiv.org/html/2509.03378v4#S4)) we develop for KL-Shampoo can be adapted to improve its performance ([Fig.˜7](https://arxiv.org/html/2509.03378v4#A6.F7), LABEL:{app:extra_exp}).

#### Von Neumann divergence (VN-Shampoo)

Another variant, often discussed in the literature (Morwani et al., 2025; Vyas et al., 2025a; Eschenhagen et al., 2025), is Shampoo with trace scaling. Vyas et al. (2025a) established that Shampoo with trace scaling is equivalent to running Adafactor (Shazeer & Stern, 2018) in Shampoo’s eigenbasis. In contrast, KL-Shampoo is not equivalent to running Adafactor in its eigenbasis.
To clarify this distinction, we make the theoretical connection between Shampoo and Adafactor more explicit: Shampoo with trace scaling is exactly a matrix generalization of Adafactor obtained by minimizing the von Neumann (VN) divergence (Tsuda et al., 2005; Dhillon & Tropp, 2008; Nock et al., 2012) and recovers Adafactor when its Kronecker factors are restricted to be diagonal, as we establish in
[˜7](https://arxiv.org/html/2509.03378v4#Thmclaim7) ([Appx.˜G](https://arxiv.org/html/2509.03378v4#A7)). By contrast, KL-Shampoo minimizes the KL divergence instead of the VN divergence. A straightforward implementation of Shampoo with trace scaling—referred to as idealized VN-Shampoo—performs poorly in practice. Howeve, the techniques we develop for KL-Shampoo in [Sec.˜4](https://arxiv.org/html/2509.03378v4#S4) can be adapted ([Fig.˜8](https://arxiv.org/html/2509.03378v4#A7.F8.fig1), [Appx.˜G](https://arxiv.org/html/2509.03378v4#A7)) to substantially improve its performance so that it matches Shampoo with step-size grafting and outperforms SOAP, as shown in [Figs.˜10](https://arxiv.org/html/2509.03378v4#A8.F10) and [10](https://arxiv.org/html/2509.03378v4#A8.F10) (LABEL:{app:extra_exp}).

A natural question then arises: which divergence is more suitable? Theoretically, the KL divergence is broadly applicable to arbitrary SPD matrices (Bhatia, 2007; Boumal et al., 2014) and is widely used for covariance matrices (Minh & Murino, 2017). In contrast, the Frobenius norm does not respect the SPD constraint, and the VN divergence is primarily motivated, studied, and applied for unit-trace SPD matrices (Tsuda et al., 2005; Nielsen & Chuang, 2010). Empirically, adopting the KL divergence yields larger improvements than both the Frobenius norm and the VN divergence for designing Shampoo’s schemes (see [Fig.˜5](https://arxiv.org/html/2509.03378v4#S6.F5)) and in other applications (Kulis et al., 2009).

## 4 Efficient Implementation: KL-Shampoo with QR Decomposition

We develop techniques that enable KL-Shampoo to match SOAP-level per-iteration runtime and to achieve
competitive performance without step-size grafting, all without relying on eigendecomposition.
Vyas et al. (2025a) demonstrated that the eigendecomposition used in Shampoo’s implementation (Shi et al., 2023) is more computationally expensive than QR decomposition. Motivated by this result, we aim to improve KL-Shampoo’s computational efficiency by replacing the eigendecomposition with QR decomposition. However, incorporating QR decomposition into KL-Shampoo is non-trivial because the eigenvalues of the Kronecker factors are required, and QR does not directly provide them without a significant overhead. Specifically, the eigenvalues are essential for a reduction in the computational cost of KL-Shampoo
in two reasons: (1) they remove the need to compute the matrix $-\nicefrac{{1}}{{2}}$ power, $\smash{{\bm{S}}^{-\nicefrac{{1}}{{2}}}=({\bm{Q}}_{a}\mathrm{Diag}({\bm{\lambda}}_{a}^{\odot-\nicefrac{{1}}{{2}}}){\bm{Q}}_{a}^{\top})\otimes({\bm{Q}}_{b}\mathrm{Diag}({\bm{\lambda}}_{b}^{\odot-\nicefrac{{1}}{{2}}}){\bm{Q}}_{b}^{\top})}$, used for KL-Shampoo’s preconditioning;
(2) they eliminate expensive matrix inversions in its Kronecker estimation rule ([Eq.˜5](https://arxiv.org/html/2509.03378v4#S3.E5)), such as ${\bm{S}}_{b}^{-1}={\bm{P}}_{b}:={\bm{Q}}_{b}\mathrm{Diag}({\bm{\lambda}}_{b}^{\odot-1}){\bm{Q}}_{b}^{\top}$ in the update for ${\bm{S}}_{a}$:

$$ $\displaystyle{\bm{S}}_{a}$ $\displaystyle\leftarrow(1-\beta_{2}){\bm{S}}_{a}+\frac{\beta_{2}}{d_{b}}{\bm{G}}{\bm{S}}_{b}^{-1}{\bm{G}}^{\top}=(1-\beta_{2}){\bm{S}}_{a}+\frac{\beta_{2}}{d_{b}}{\bm{G}}{\bm{P}}_{b}{\bm{G}}^{\top},$ (6) $$

where ${\bm{Q}}_{k}$ and ${\bm{\lambda}}_{k}$ are eigenbasis and eigenvalues of ${\bm{S}}_{k}$ for $k\in\{a,b\}$, respectively.

#### KL-based estimation rule for the eigenvalues 𝝀 a {\bm{\lambda}}_{a} and 𝝀 b {\bm{\lambda}}_{b} using an outdated eigenbasis

We aim to estimate the eigenvalues using an outdated eigenbasis and replace the slow eigendecomposition with a fast QR decomposition in KL-Shampoo.
Eschenhagen et al. (2025)
propose estimating the eigenvalues from a Frobenius-norm perspective, using an instantaneous scheme: ${\bm{\lambda}}_{k}^{\text{(inst)}}:=\mathrm{diag}({\bm{Q}}_{k}^{\top}{\bm{S}}_{k}{\bm{Q}}_{k})$ for $k\in\{a,b\}$.
However, our empirical results ([Fig.˜5](https://arxiv.org/html/2509.03378v4#S6.F5)) indicate that this approach becomes suboptimal when an outdated eigenbasis ${\bm{Q}}_{k}$ is reused to reduce the frequency and cost of QR decompositions.
In contrast, our KL perspective (see [˜4](https://arxiv.org/html/2509.03378v4#Thmclaim4) and its proof in [Appx.˜D](https://arxiv.org/html/2509.03378v4#A4)) provides a principled alternative, allowing us to use an outdated eigenbasis.
Building on this claim, we introduce an exponential moving average (EMA) scheme (Step 3a of [Fig.˜3](https://arxiv.org/html/2509.03378v4#S4.F3.fig1)) for eigenvalue estimation, which can be justified as a stochastic proximal-gradient step under our KL perspective, similar to [˜3](https://arxiv.org/html/2509.03378v4#Thmclaim3). This scheme updates the eigenvalues *at every iteration* while updating the eigenbasis less frequently through an efficient QR-based procedure, similar to SOAP. We can view this estimation as an eigenvalue correction for using an outdated eigenbasis, as will be discussed in [Sec.˜5](https://arxiv.org/html/2509.03378v4#S5).
Since it naturally scales the eigenvalues by the dimensions of the Kronecker factors, step-size grafting should not be necessary for KL-Shampoo, as argued by Eschenhagen et al. (2025) and confirmed by our empirical results ([Fig.˜2](https://arxiv.org/html/2509.03378v4#S3.F2)).
Furthermore, applying this scheme enables other variants of Shampoo to be competitive and even outperform SOAP, as empirically demonstrated in [Figs.˜7](https://arxiv.org/html/2509.03378v4#A6.F7), [10](https://arxiv.org/html/2509.03378v4#A8.F10) and [10](https://arxiv.org/html/2509.03378v4#A8.F10) of LABEL:{app:extra_exp}.
These empirical results underscore the importance of our EMA scheme on eigenvalues.

Figure: Figure 3: *Left:* Simplified Shampoo-based schemes without momentum, damping, and weight decay. See [Fig.˜5](https://arxiv.org/html/2509.03378v4#S6.F5) for an empirical comparison. *Top Right:* For computational efficiency, we replace the eigen step with our exponential moving average (EMA) scheme to estimate eigenvalues and infrequent eigenbasis estimation using QR, where we estimate eigenvalues ${\bm{\lambda}}_{k}$ using an outdated eigenbasis ${\bm{Q}}_{k}$ for $k\in\{a,b\}$, and use the QR procedure to estimate ${\bm{Q}}_{k}$. *Bottom Right*: Simplified SOAP-based schemes without momentum. Notably, KL-SOAP needs estimation for ${\bm{\lambda}}_{k}$ in Step 3a to compute the eigenbasis ${\bm{Q}}$, whereas SOAP does not.

###### Claim 4

(Covariance estimation for
eigenvalues ${\bm{\lambda}}_{a}$ and ${\bm{\lambda}}_{b}$)
The optimal solution of KL minimization
$\smash{\min_{{\bm{\lambda}}_{a},{\bm{\lambda}}_{b}}\mathrm{KL}\big(\mathbb{E}[{\bm{g}}{\bm{g}}^{\top}],{\bm{S}}\big)}$ with preconditioner $\smash{{\bm{S}}\!=\!({\bm{Q}}_{a}\mathrm{Diag}({\bm{\lambda}}_{a}){\bm{Q}}_{a}^{\top})\!\otimes\!({\bm{Q}}_{b}\mathrm{Diag}({\bm{\lambda}}_{b}){\bm{Q}}_{b}^{\top})}$ satisfies this condition:

$$ $\displaystyle{\bm{\lambda}}_{a}^{*}=\frac{1}{d_{b}}\mathrm{diag}\big({\bm{Q}}_{a}^{\top}\mathbb{E}[{\bm{G}}{{\bm{P}}}_{b}^{*}{\bm{G}}^{\top}]{\bm{Q}}_{a}\big),\quad{\bm{\lambda}}_{b}^{*}=\frac{1}{d_{a}}\mathrm{diag}\big({\bm{Q}}_{b}^{\top}\mathbb{E}[{\bm{G}}^{\top}{{\bm{P}}}_{a}^{*}{\bm{G}}]{\bm{Q}}_{b}\big),$ (7) $$

where ${{\bm{P}}}_{k}^{*}:={\bm{Q}}_{k}\mathrm{Diag}\left(({\bm{\lambda}}_{k}^{*})^{\odot-1}\right){\bm{Q}}_{k}^{\top}$ is also defined in [Eq.˜6](https://arxiv.org/html/2509.03378v4#S4.E6) and considered as an approximation of ${\bm{S}}_{k}^{-1}$ for $k\in\{a,b\}$ when using an outdated eigenbasis ${\bm{Q}}={\bm{Q}}_{a}\otimes{\bm{Q}}_{b}$ precomputed by QR.

###### Claim 4

## 5 Interpreting and Improving SOAP via KL Minimization

We extend the KL perspective to better understand and improve the estimation scheme used in SOAP.

#### Interpreting SOAP’s estimation as covariance estimation

Recall that SOAP ([Eq.˜2](https://arxiv.org/html/2509.03378v4#S2.E2)) applies Shampoo’s scheme to estimate its Kronecker factors and then performs RMSProp updates in the eigenbasis of these factors. Consequently, the interpretation of SOAP’s Kronecker factor estimation is identical to that of Shampoo. RMSProp’s second-moment estimation in the eigenbasis can itself be interpreted as the optimal solution to a separate KL divergence minimization problem, as established in [˜5](https://arxiv.org/html/2509.03378v4#Thmclaim5) (see [Appx.˜E](https://arxiv.org/html/2509.03378v4#A5) for a proof).
The KL perspective—distinct from the Frobenius-norm viewpoint (George et al., 2018; Eschenhagen et al., 2025)—provides a new lens for understanding RMSProp’s estimation in the eigenbasis as the estimation of augmented eigenvalues of a covariance matrix under KL divergence.
When an outdated eigenbasis is used, RMSProp’s scheme (Step 4a of [Fig.˜3](https://arxiv.org/html/2509.03378v4#S4.F3.fig1)) for eigenvalue estimation can be viewed as a correction in an augmented (full-diagonal) eigen space, ${\bm{Q}}\mathrm{Diag}({\bm{d}}){\bm{Q}}^{\top}$, analogous in spirit to the Frobenius-norm interpretation but derived under the KL framework. This perspective also highlights a close similarity to KL-Shampoo’s estimation scheme: recall that we introduced a comparable correction (Step 3a of [Fig.˜3](https://arxiv.org/html/2509.03378v4#S4.F3.fig1)) for KL-Shampoo, but in the original Kronecker-factored diagonal eigen space, ${\bm{Q}}\mathrm{Diag}({\bm{\lambda}}_{a}\otimes{\bm{\lambda}}_{b}){\bm{Q}}^{\top}$.
Notably, when using KL-Shampoo’s eigenbasis ${\bm{Q}}^{*}={\bm{Q}}_{a}^{*}\otimes{\bm{Q}}_{b}^{*}$ obtained from the optimal condition in [Eq.˜4](https://arxiv.org/html/2509.03378v4#S3.E4), we can see the (Gaussian) covariance of gradient vector ${\bm{g}}$ in the eigenbasis is Kronecker-diagonalized rather than fully diagonalized: $\mathrm{vec}(({\bm{Q}}_{a}^{*})^{\top}{\bm{G}}{\bm{Q}}_{b}^{*})=({\bm{Q}}^{*})^{\top}{\bm{g}}\sim\mathcal{N}(0,\mathrm{Diag}({\bm{\lambda}}_{a}^{*})\otimes\mathrm{Diag}({\bm{\lambda}}_{b}^{*}))$.

###### Claim 5

(SOAP and KL-SOAP’s covariance estimation for augmented eigenvalues ${\bm{d}}$)
The optimal solution of KL minimization: $\min_{{\bm{d}}}\mathrm{KL}\big(\mathbb{E}[{\bm{g}}{\bm{g}}^{\top}],{\bm{S}}\big)$ with preconditioner ${\bm{S}}={\bm{Q}}\mathrm{Diag}({\bm{d}}){\bm{Q}}^{\top}$
is ${\bm{d}}^{*}=\mathbb{E}\left[\left(\mathrm{vec}({\bm{Q}}_{a}^{\top}{\bm{G}}{\bm{Q}}_{b})\right)^{\odot 2}\right]=\mathbb{E}\left[\hat{{\bm{g}}}^{\odot 2}\right]$, where ${\bm{d}}\in\mathcal{R}^{d_{a}d_{b}\times 1}$ is viewed as an augmented eigenvalue vector, $\hat{{\bm{g}}}={\bm{Q}}^{\top}{\bm{g}}$ is defined at the update of (KL-)SOAP (see [Eq.˜2](https://arxiv.org/html/2509.03378v4#S2.E2)), and ${\bm{Q}}={\bm{Q}}_{a}\otimes{\bm{Q}}_{b}$ can be an outdated eigenbasis of (KL-)Shampoo’s preconditioner.

###### Claim 5

#### Improving SOAP’s estimation

Similar to SOAP, we propose KL-SOAP, which utilizes KL-Shampoo’s estimation to update Kronecker factors and additionally employs Adam (RMSProp) in KL-Shampoo’s eigenbasis.
Our unified KL perspective enables us to reuse [˜5](https://arxiv.org/html/2509.03378v4#Thmclaim5) to justify the use of Adam’s (RMSProp’s) 2^nd moment estimation as augmented eigenvalue estimation in KL-SOAP.

## 6 Experimental Setup and Empirical Evaluations

Figure: Figure 4: Empirical results (random search using 150 runs for each method) demonstrate that our EMA scheme for the eigenvalue estimation makes KL-Shampoo competitive when using an outdated eigenbasis. Without this scheme, KL-Shampoo performs poorly under an outdated eigenbasis ${\bm{Q}}_{k}$ even when employing the instantaneous eigenvalue estimation ${\bm{\lambda}}_{k}^{\text{(inst)}}=\mathrm{diag}({\bm{Q}}_{k}^{\top}{\bm{S}}_{k}{\bm{Q}}_{k})$ at every iteration, as suggested by Eschenhagen et al. (2025) for $k\in\{a,b\}$. Adapting the EMA scheme also makes other variants of Shampoo competitive ([Figs.˜7](https://arxiv.org/html/2509.03378v4#A6.F7) and [10](https://arxiv.org/html/2509.03378v4#A8.F10), [Appx.˜H](https://arxiv.org/html/2509.03378v4#A8)) and allows the trace-scaling variant to outperform SOAP ([Fig.˜10](https://arxiv.org/html/2509.03378v4#A8.F10), [Appx.˜H](https://arxiv.org/html/2509.03378v4#A8)).
Refer to caption: x3.png

We consider four sets of experiments to demonstrate the benefits of using the KL divergence and the effectiveness of KL-based methods. See [Appx.˜H](https://arxiv.org/html/2509.03378v4#A8) for additional experiments.

#### Experimental Setup

In all the experiments, we consider training four language models based on existing implementations: NanoGPT (Jordan, 2024) (123 M), NanoRWKV7 (Bo, 2024) (162 M), Llama (Glentis, 2025) (134 M), and NanoMoE (Wolfe, 2025) (227 M).
We consider NanoMoE, as it contains 3D weight tensors. This model provides a natural testbed for evaluating a tensor extension of KL-Shampoo and KL-SOAP, derived directly from our KL perspective.
In doing so, we demonstrate that our methods retain the same flexibility as Shampoo and SOAP in handling tensor-valued weights without reshaping them into matrices.
We train NanoGPT and NanoRWKV7 using a subset of FineWeb (1 B tokens), Llama using a subset of C4 (2 B tokens), and
NanoMoE using a subset of OpenWebText (2.5 B tokens).
All models except NanoMoE are trained using mini-batches with a batch size of 0.5 M. We use a batch size of 0.25 M to train NanoMoE to reduce the run time.
We use the default step-size schedulers from the source implementations; NanoGPT and NanoRWKV7: linear warmup + constant step-size + linear cooldown; Llama and NanoMoE: linear warmup + cosine step-size.
We tune all available hyperparameters for each method—including step-size, moving average, weight decay, damping, and momentum—using random search with 150 runs. Our hyperparameter search follows a two-stage strategy, with 75 runs in each stage. In the first stage, we search over a wider range of hyperparameters. In the second stage, we refine the search space based on the results from the first stage and focus on a narrower range.
In our experiments, Shampoo by default performs eigendecomposition every 10 steps, while SOAP, KL-Shampoo, and KL-SOAP perform QR decomposition every 10 steps, as suggested by Vyas et al. (2025a).

In the first set of experiments, we demonstrate that our KL-based perspective enables a principled redesign of Shampoo, resulting in KL-Shampoo, and achieves superior performance without step-size grafting. We evaluate Shampoo with matrix powers
$p=\nicefrac{{1}}{{2}}$ and $p=\nicefrac{{1}}{{4}}$, using a state-of-the-art implementation (Shi et al., 2023). As shown in [Fig.˜2](https://arxiv.org/html/2509.03378v4#S3.F2), Shampoo requires step-size grafting to perform well, whereas KL-Shampoo performs robustly without it. Moreover, KL-Shampoo outperforms Shampoo with grafting—even in terms of step-wise progress—even when Shampoo is equipped with eigendecomposition and step-size grafting via Adam.

In the second set of experiments, we demonstrate that our QR-based scheme enables KL-Shampoo and KL-SOAP to achieve the same pre-iteration runtime as SOAP. We use the official SOAP implementation for comparison. As shown in [Fig.˜1](https://arxiv.org/html/2509.03378v4#S2.F1), KL-Shampoo and KL-SOAP outperform SOAP.
Remarkably, KL-Shampoo also consistently surpasses KL-SOAP while using less memory.

In the third set of experiments, we underscore the importance of using our EMA scheme for the eigenvalue estimation when working with an outdated eigenbasis. As shown in [Fig.˜5](https://arxiv.org/html/2509.03378v4#S6.F5), the EMA scheme enables KL-Shampoo to perform well in practice, even under stale eigenbases. Moreover, this scheme can be adapted to strengthen the trace scaling variant of Shampoo ([Fig.˜10](https://arxiv.org/html/2509.03378v4#A8.F10), LABEL:{app:extra_exp}), enabling it to outperform SOAP ([Fig.˜10](https://arxiv.org/html/2509.03378v4#A8.F10), LABEL:{app:extra_exp}).

In the last set of experiments, we evaluate the benefits of using the two-sided estimation scheme under our KL perspective. Specifically, we compare the two-sided approach (KL-Shampoo) against the one-sided approach (Shampoo) in a comparable setting. To ensure fairness and eliminate implementation bias, we use our own implementation of Shampoo aligned closely with that of KL-Shampoo. For this comparison, we extend Shampoo with a QR-based step and our EMA scheme for eigenvalue estimation, as described in [Fig.˜3](https://arxiv.org/html/2509.03378v4#S4.F3.fig1).
Similarly, we also consider two more Shampoo variants discussed in [Sec.˜3.3](https://arxiv.org/html/2509.03378v4#S3.SS3).
As shown in [Fig.˜5](https://arxiv.org/html/2509.03378v4#S6.F5), KL-Shampoo consistently outperforms other Shampoo variants, even when these variants employ a similar QR-based estimation rule.

## 7 Conclusion

We introduced a KL perspective for interpreting Shampoo’s and SOAP’s structured second-moment estimation schemes. This perspective uncovers a previously unrecognized limitation of Shampoo, motivates an alternative estimation strategy to overcome it, enables a practical implementation of our approach, and extends naturally to tensor-valued estimation. Our empirical results demonstrate the effectiveness of our approach for improving Shampoo’s and SOAP’s estimation schemes.

## References

- Agarwal et al. (2020)
Naman Agarwal, Rohan Anil, Elad Hazan, Tomer Koren, and Cyril Zhang.
Disentangling adaptive gradient methods from learning rates.
*arXiv preprint arXiv:2002.11803*, 2020.
doi:[10.48550/arxiv.2002.11803](https://doi.org/10.48550/arxiv.2002.11803).
- Amari (2016)
Shun-ichi Amari.
*Information geometry and its applications*, volume 194.
Springer, 2016.
ISBN 9784431559788.
doi:[10.1007/978-4-431-55978-8](https://doi.org/10.1007/978-4-431-55978-8).
- An et al. (2025)
Kang An, Yuxing Liu, Rui Pan, Yi Ren, Shiqian Ma, Donald Goldfarb, and Tong Zhang.
ASGO: Adaptive structured gradient optimization.
*arXiv preprint arXiv:2503.20762*, 2025.
doi:[10.48550/arxiv.2503.20762](https://doi.org/10.48550/arxiv.2503.20762).
- Anil et al. (2020)
Rohan Anil, Vineet Gupta, Tomer Koren, Kevin Regan, and Yoram Singer.
Scalable second order optimization for deep learning.
*arXiv preprint arXiv:2002.09018*, 2020.
doi:[10.48550/arxiv.2002.09018](https://doi.org/10.48550/arxiv.2002.09018).
- Bhatia (2007)
Rajendra Bhatia.
*Positive definite matrices*.
Princeton University Press, 2007.
ISBN 9780691129181.
URL [http://www.jstor.org/stable/j.ctt7rxv2](http://www.jstor.org/stable/j.ctt7rxv2).
- Bo (2024)
Peng Bo.
RWKV-7: Surpassing GPT.
[https://github.com/BlinkDL/modded-nanogpt-rwkv](https://github.com/BlinkDL/modded-nanogpt-rwkv), 2024.
Accessed: 2025-06.
- Boumal et al. (2014)
Nicolas Boumal, Bamdev Mishra, P.-A. Absil, and Rodolphe Sepulchre.
Manopt, a matlab toolbox for optimization on manifolds.
*Journal of Machine Learning Research*, 15(42):1455–1459, 2014.
URL [http://jmlr.org/papers/v15/boumal14a.html](http://jmlr.org/papers/v15/boumal14a.html).
- Bregman (1967)
Lev M Bregman.
The relaxation method of finding the common point of convex sets and its application to the solution of problems in convex programming.
*USSR computational mathematics and mathematical physics*, 7(3):200–217, 1967.
ISSN 0041-5553.
doi:[10.1016/0041-5553(67)90040-7](https://doi.org/10.1016/0041-5553(67)90040-7).
- Byrd & Nocedal (1989)
Richard H Byrd and Jorge Nocedal.
A tool for the analysis of quasi-Newton methods with application to unconstrained minimization.
*SIAM Journal on Numerical Analysis*, 26(3):727–739, 1989.
doi:[10.1137/0726042](https://doi.org/10.1137/0726042).
- Dahl et al. (2023)
George E Dahl, Frank Schneider, Zachary Nado, Naman Agarwal, Chandramouli Shama Sastry, Philipp Hennig, Sourabh Medapati, Runa Eschenhagen, Priya Kasimbeg, Daniel Suo, et al.
Benchmarking neural network training algorithms.
*arXiv preprint arXiv:2306.07179*, 2023.
doi:[10.48550/arxiv.2306.07179](https://doi.org/10.48550/arxiv.2306.07179).
- Davis & Dhillon (2006)
Jason Davis and Inderjit Dhillon.
Differential entropic clustering of multivariate gaussians.
*Advances in Neural Information Processing Systems*, 19, 2006.
- de Boer et al. (2023)
Jan de Boer, Victor Godet, Jani Kastikainen, and Esko Keski-Vakkuri.
Quantum information geometry of driven CFTs.
*Journal of High Energy Physics*, 2023(9):1–89, 2023.
doi:[10.1007/JHEP09(2023)087](https://doi.org/10.1007/JHEP09(2023)087).
- Dhillon & Tropp (2008)
Inderjit S Dhillon and Joel A Tropp.
Matrix nearness problems with Bregman divergences.
*SIAM Journal on Matrix Analysis and Applications*, 29(4):1120–1146, 2008.
doi:[10.1137/060649021](https://doi.org/10.1137/060649021).
- Duchi et al. (2011)
John Duchi, Elad Hazan, and Yoram Singer.
Adaptive subgradient methods for online learning and stochastic optimization.
*Journal of Machine Learning Research*, 12(61):2121–2159, 2011.
URL [http://jmlr.org/papers/v12/duchi11a.html](http://jmlr.org/papers/v12/duchi11a.html).
- Dutilleul (1999)
Pierre Dutilleul.
The MLE algorithm for the matrix normal distribution.
*Journal of Statistical Computation and Simulation*, 64(2):105–123, 1999.
doi:[10.1080/00949659908811970](https://doi.org/10.1080/00949659908811970).
- Duvvuri et al. (2024)
Sai Surya Duvvuri, Fnu Devvrit, Rohan Anil, Cho-Jui Hsieh, and Inderjit S Dhillon.
Combining axes preconditioners through Kronecker approximation for deep learning.
In *The Twelfth International Conference on Learning Representations*, 2024.
URL [https://openreview.net/forum?id=8j9hz8DVi8](https://openreview.net/forum?id=8j9hz8DVi8).
- Eschenhagen et al. (2025)
Runa Eschenhagen, Aaron Defazio, Tsung-Hsien Lee, Richard E Turner, and Hao-Jun Michael Shi.
Purifying Shampoo: Investigating Shampoo’s heuristics by decomposing its preconditioner.
*arXiv preprint arXiv:2506.03595*, 2025.
doi:[10.48550/arxiv.2506.03595](https://doi.org/10.48550/arxiv.2506.03595).
- Fletcher (1991)
Roger Fletcher.
A new variational result for quasi-Newton formulae.
*SIAM Journal on Optimization*, 1(1):18–21, 1991.
doi:[10.1137/0801002](https://doi.org/10.1137/0801002).
- George et al. (2018)
Thomas George, César Laurent, Xavier Bouthillier, Nicolas Ballas, and Pascal Vincent.
Fast approximate natural gradient descent in a kronecker factored eigenbasis.
In S. Bengio, H. Wallach, H. Larochelle, K. Grauman, N. Cesa-Bianchi, and R. Garnett (eds.), *Advances in Neural Information Processing Systems*, volume 31. Curran Associates, Inc., 2018.
URL [https://proceedings.neurips.cc/paper_files/paper/2018/file/48000647b315f6f00f913caa757a70b3-Paper.pdf](https://proceedings.neurips.cc/paper_files/paper/2018/file/48000647b315f6f00f913caa757a70b3-Paper.pdf).
- Glentis (2025)
Athanasios Glentis.
A minimalist optimizer design for LLM pretraining.
[https://github.com/OptimAI-Lab/Minimalist_LLM_Pretraining](https://github.com/OptimAI-Lab/Minimalist_LLM_Pretraining), 2025.
Accessed: 2025-06.
- Gupta et al. (2018)
Vineet Gupta, Tomer Koren, and Yoram Singer.
Shampoo: Preconditioned stochastic tensor optimization.
In Jennifer Dy and Andreas Krause (eds.), *Proceedings of the 35th International Conference on Machine Learning*, volume 80 of *Proceedings of Machine Learning Research*, pp. 1842–1850. PMLR, 10–15 Jul 2018.
URL [https://proceedings.mlr.press/v80/gupta18a.html](https://proceedings.mlr.press/v80/gupta18a.html).
- James & Stein (1961)
William James and Charles Stein.
Estimation with quadratic loss.
In *Proceedings of the fourth Berkeley symposium on mathematical statistics and probability*, volume 1, pp. 361–379. University of California Press, 1961.
- Jordan (2024)
Keller Jordan.
NanoGPT (124M) in 3 minutes.
[https://github.com/KellerJordan/modded-nanogpt](https://github.com/KellerJordan/modded-nanogpt), 2024.
Accessed: 2025-06.
- Kanamori & Ohara (2013a)
Takafumi Kanamori and Atsumi Ohara.
A Bregman extension of quasi-Newton updates I: an information geometrical framework.
*Optimization Methods and Software*, 28(1):96–123, 2013a.
doi:[10.1080/10556788.2011.613073](https://doi.org/10.1080/10556788.2011.613073).
- Kanamori & Ohara (2013b)
Takafumi Kanamori and Atsumi Ohara.
A Bregman extension of quasi-Newton updates II: Analysis of robustness properties.
*Journal of computational and applied mathematics*, 253:104–122, 2013b.
doi:[10.1016/j.cam.2013.04.005](https://doi.org/10.1016/j.cam.2013.04.005).
- Kasimbeg et al. (2025)
Priya Kasimbeg, Frank Schneider, Runa Eschenhagen, Juhan Bae, Chandramouli Shama Sastry, Mark Saroufim, Boyuan Fend, Less Wright, Edward Z Yang, Zachary Nado, et al.
Accelerating neural network training: An analysis of the AlgoPerf competition.
In *The Thirteenth International Conference on Learning Representations*, 2025.
URL [https://openreview.net/forum?id=CtM5xjRSfm](https://openreview.net/forum?id=CtM5xjRSfm).
- Khan & Lin (2017)
Mohammad Khan and Wu Lin.
Conjugate-computation variational inference: Converting variational inference in non-conjugate models to inferences in conjugate models.
In Aarti Singh and Jerry Zhu (eds.), *Proceedings of the 20th International Conference on Artificial Intelligence and Statistics*, volume 54 of *Proceedings of Machine Learning Research*, pp. 878–887. PMLR, 20–22 Apr 2017.
URL [https://proceedings.mlr.press/v54/khan17a.html](https://proceedings.mlr.press/v54/khan17a.html).
- Khan et al. (2016)
Mohammad Emtiyaz Khan, Reza Babanezhad, Wu Lin, Mark Schmidt, and Masashi Sugiyama.
Faster stochastic variational inference using proximal-gradient methods with general divergence functions.
In *Proceedings of the Thirty-Second Conference on Uncertainty in Artificial Intelligence*, pp. 319–328. AUAI Press, 2016.
URL [https://www.auai.org/uai2016/proceedings/papers/218.pdf](https://www.auai.org/uai2016/proceedings/papers/218.pdf).
- Kingma & Ba (2015)
Diederik P Kingma and Jimmy Ba.
Adam: A method for stochastic optimization.
In *International Conference on Learning Representations*, 2015.
doi:[10.48550/arxiv.1412.6980](https://doi.org/10.48550/arxiv.1412.6980).
- Kivinen & Warmuth (1999)
Jyrki Kivinen and Manfred K. Warmuth.
Boosting as entropy projection.
In *Proceedings of the Twelfth Annual Conference on Computational Learning Theory*, pp. 134–144, New York, NY, USA, 1999. Association for Computing Machinery.
ISBN 1581131674.
doi:[10.1145/307400.307424](https://doi.org/10.1145/307400.307424).
- Kulis et al. (2009)
Brian Kulis, Mátyás A Sustik, and Inderjit S Dhillon.
Low-rank kernel learning with bregman matrix divergences.
*Journal of Machine Learning Research*, 10(2), 2009.
- Kunstner et al. (2021)
Frederik Kunstner, Raunak Kumar, and Mark Schmidt.
Homeomorphic-invariance of EM: Non-asymptotic convergence in KL divergence for exponential families via mirror descent.
In Arindam Banerjee and Kenji Fukumizu (eds.), *Proceedings of The 24th International Conference on Artificial Intelligence and Statistics*, volume 130 of *Proceedings of Machine Learning Research*, pp. 3295–3303. PMLR, 13–15 Apr 2021.
URL [https://proceedings.mlr.press/v130/kunstner21a.html](https://proceedings.mlr.press/v130/kunstner21a.html).
- Lin et al. (2019)
Wu Lin, Mohammad Emtiyaz Khan, and Mark Schmidt.
Fast and simple natural-gradient variational inference with mixture of exponential-family approximations.
In *International Conference on Machine Learning*, volume 97 of *Proceedings of Machine Learning Research*, pp. 3992–4002. PMLR, 09–15 Jun 2019.
URL [https://proceedings.mlr.press/v97/lin19b.html](https://proceedings.mlr.press/v97/lin19b.html).
- Lin et al. (2023)
Wu Lin, Valentin Duruisseaux, Melvin Leok, Frank Nielsen, Mohammad Emtiyaz Khan, and Mark Schmidt.
Simplifying momentum-based positive-definite submanifold optimization with applications to deep learning.
In Andreas Krause, Emma Brunskill, Kyunghyun Cho, Barbara Engelhardt, Sivan Sabato, and Jonathan Scarlett (eds.), *Proceedings of the 40th International Conference on Machine Learning*, volume 202 of *Proceedings of Machine Learning Research*, pp. 21026–21050. PMLR, 23–29 Jul 2023.
URL [https://proceedings.mlr.press/v202/lin23c.html](https://proceedings.mlr.press/v202/lin23c.html).
- Lin et al. (2024)
Wu Lin, Felix Dangel, Runa Eschenhagen, Juhan Bae, Richard E Turner, and Alireza Makhzani.
Can we remove the square-root in adaptive gradient methods? A second-order perspective.
In Ruslan Salakhutdinov, Zico Kolter, Katherine Heller, Adrian Weller, Nuria Oliver, Jonathan Scarlett, and Felix Berkenkamp (eds.), *Proceedings of the 41st International Conference on Machine Learning*, volume 235 of *Proceedings of Machine Learning Research*, pp. 29949–29973. PMLR, 21–27 Jul 2024.
URL [https://proceedings.mlr.press/v235/lin24e.html](https://proceedings.mlr.press/v235/lin24e.html).
- Löwner (1934)
Karl Löwner.
Über monotone matrixfunktionen.
*Mathematische Zeitschrift*, 38(1):177–216, 1934.
doi:[10.1007/BF01170633](https://doi.org/10.1007/BF01170633).
- Minh & Murino (2017)
Hà Quang Minh and Vittorio Murino.
Covariances in computer vision and machine learning.
*Synthesis Lectures on Computer Vision*, 7(4):1–170, 2017.
ISSN 2153-1056.
doi:[10.1007/978-3-031-01820-6](https://doi.org/10.1007/978-3-031-01820-6).
- Morwani et al. (2025)
Depen Morwani, Itai Shapira, Nikhil Vyas, Eran Malach, Sham M Kakade, and Lucas Janson.
A new perspective on Shampoo’s preconditioner.
In *The Thirteenth International Conference on Learning Representations*, 2025.
URL [https://openreview.net/forum?id=c6zI3Cp8c6](https://openreview.net/forum?id=c6zI3Cp8c6).
- Nesterov et al. (2018)
Yurii Nesterov et al.
*Lectures on convex optimization*, volume 137.
Springer Cham, 2018.
ISBN 9783319915784.
doi:[10.1007/978-3-319-91578-4](https://doi.org/10.1007/978-3-319-91578-4).
- Nielsen & Chuang (2010)
Michael A Nielsen and Isaac L Chuang.
*Quantum computation and quantum information*.
Cambridge university press, 2010.
- Nocedal & Wright (2006)
Jorge Nocedal and Stephen J Wright.
*Numerical optimization*.
Springer, 2006.
ISBN 978-0-387-40065-5.
doi:[10.1007/978-0-387-40065-5](https://doi.org/10.1007/978-0-387-40065-5).
- Nock et al. (2012)
Richard Nock, Brice Magdalou, Eric Briys, and Frank Nielsen.
Mining matrix data with bregman matrix divergences for portfolio selection.
In *Matrix Information Geometry*, pp. 373–402. Springer, 2012.
- Parikh & Boyd (2014)
Neal Parikh and Stephen Boyd.
Proximal algorithms.
*Foundations and trends in Optimization*, 1(3):127–239, 2014.
doi:[10.1561/2400000003](https://doi.org/10.1561/2400000003).
- Pennec et al. (2006)
Xavier Pennec, Pierre Fillard, and Nicholas Ayache.
A Riemannian framework for tensor computing.
*International Journal of Computer Vision*, 66(1):41–66, Jan 2006.
ISSN 1573-1405.
doi:[10.1007/s11263-005-3222-z](https://doi.org/10.1007/s11263-005-3222-z).
- Shazeer & Stern (2018)
Noam Shazeer and Mitchell Stern.
Adafactor: Adaptive learning rates with sublinear memory cost.
In *Proceedings of the 35th International Conference on Machine Learning*, pp. 4596–4604. PMLR, 2018.
URL [https://proceedings.mlr.press/v80/shazeer18a.html](https://proceedings.mlr.press/v80/shazeer18a.html).
- Shi et al. (2023)
Hao-Jun Michael Shi, Tsung-Hsien Lee, Shintaro Iwasaki, Jose Gallego-Posada, Zhijing Li, Kaushik Rangadurai, Dheevatsa Mudigere, and Michael Rabbat.
A distributed data-parallel PyTorch implementation of the distributed Shampoo optimizer for training neural networks at-scale.
*arXiv preprint arXiv:2309.06497*, 2023.
doi:[10.48550/arxiv.2309.06497](https://doi.org/10.48550/arxiv.2309.06497).
- Sra (2016)
Suvrit Sra.
Positive definite matrices and the s-divergence.
*Proceedings of the American Mathematical Society*, 144(7):2787–2797, 2016.
- Tieleman & Hinton (2012)
Tijmen Tieleman and Geoffrey Hinton.
RMSProp: Divide the gradient by a running average of its recent magnitude.
*Coursera*, 2012.
- Tsuda et al. (2005)
Koji Tsuda, Gunnar Rätsch, and Manfred K Warmuth.
Matrix exponentiated gradient updates for on-line learning and Bregman projection.
*Journal of Machine Learning Research*, 6(34):995–1018, 2005.
URL [https://jmlr.org/papers/v6/tsuda05a.html](https://jmlr.org/papers/v6/tsuda05a.html).
- Van Loan & Pitsianis (1993)
C. F. Van Loan and N. Pitsianis.
*Approximation with Kronecker Products*, pp. 293–314.
Springer Netherlands, Dordrecht, 1993.
ISBN 978-94-015-8196-7.
doi:[10.1007/978-94-015-8196-7_17](https://doi.org/10.1007/978-94-015-8196-7_17).
- Vyas et al. (2025a)
Nikhil Vyas, Depen Morwani, Rosie Zhao, Itai Shapira, David Brandfonbrener, Lucas Janson, and Sham M Kakade.
SOAP: Improving and stabilizing Shampoo using Adam for language modeling.
In *The Thirteenth International Conference on Learning Representations*, 2025a.
URL [https://openreview.net/forum?id=IDxZhXrpNf](https://openreview.net/forum?id=IDxZhXrpNf).
- Vyas et al. (2025b)
Nikhil Vyas, Rosie Zhao, Depen Morwani, Mujin Kwun, and Sham Kakade.
Improving SOAP using iterative whitening and Muon.
[https://nikhilvyas.github.io/SOAP_Muon.pdf](https://nikhilvyas.github.io/SOAP_Muon.pdf), 2025b.
- Waldrip & Niven (2016)
Steven H Waldrip and Robert K Niven.
Maximum entropy derivation of quasi-Newton methods.
*SIAM Journal on Optimization*, 26(4):2495–2511, 2016.
doi:[10.1137/15M1027668](https://doi.org/10.1137/15M1027668).
- Wolfe (2025)
Cameron R. Wolfe.
An extension of the NanoGPT repository for training small MOE models.
[https://github.com/wolfecameron/nanoMoE](https://github.com/wolfecameron/nanoMoE), 2025.
Accessed: 2025-06.
- Xie et al. (2025)
Shuo Xie, Tianhao Wang, Sashank J Reddi, Sanjiv Kumar, and Zhiyuan Li.
Structured preconditioners in adaptive optimization: A unified analysis.
In *Forty-second International Conference on Machine Learning*, 2025.
URL [https://openreview.net/forum?id=GzS6b5Xvvu](https://openreview.net/forum?id=GzS6b5Xvvu).

## Appendices

## Appendix A Proof of ˜ 1

We will show that the optimal solution of KL minimization $\min_{{\bm{S}}_{a}}\mathrm{KL}\big(\mathbb{E}[{\bm{g}}{\bm{g}}^{\top}],{\bm{S}}\big)$ with a one-sided preconditioner ${\bm{S}}=(\nicefrac{{1}}{{d_{b}}}{\bm{S}}_{a})\otimes{\bm{I}}_{b}$ is ${\bm{S}}_{a}^{*}=\mathbb{E}[{\bm{G}}{\bm{G}}^{\top}]$.

By definition in [Eq.˜3](https://arxiv.org/html/2509.03378v4#S3.E3) and substituting ${\bm{S}}=(\nicefrac{{1}}{{d_{b}}}{\bm{S}}_{a})\otimes{\bm{I}}_{b}$, we can simplify the objective function as

$$ $\displaystyle\mathrm{KL}\big(\mathbb{E}[{\bm{g}}{\bm{g}}^{\top}],{\bm{S}}\big)$ $\displaystyle\quad=\frac{1}{2}\big(\log\mathrm{det}({\bm{S}})+\mathrm{Tr}({\bm{S}}^{-1}\mathbb{E}[{\bm{g}}{\bm{g}}^{\top}])\big)+\text{const.}$ $\displaystyle\quad=\frac{1}{2}\big(d_{b}\log\mathrm{det}(\frac{1}{d_{b}}{\bm{S}}_{a})+\mathrm{Tr}({\bm{S}}^{-1}\mathbb{E}[{\bm{g}}{\bm{g}}^{\top}])\big)+\text{const.}\quad\text{{\small\color[rgb]{.5,.5,.5}\definecolor[named]{pgfstrokecolor}{rgb}{.5,.5,.5}\pgfsys@color@gray@stroke{.5}\pgfsys@color@gray@fill{.5}(Kronecker identity for matrix det.)}}$ $\displaystyle\quad=\frac{1}{2}\big(d_{b}\log\mathrm{det}({\bm{S}}_{a})+\mathrm{Tr}({\bm{S}}^{-1}\mathbb{E}[{\bm{g}}{\bm{g}}^{\top}])\big)+\text{const.}\quad\text{{\small\color[rgb]{.5,.5,.5}\definecolor[named]{pgfstrokecolor}{rgb}{.5,.5,.5}\pgfsys@color@gray@stroke{.5}\pgfsys@color@gray@fill{.5}(identity for a log-determinant)}}$ $\displaystyle\quad=\frac{1}{2}\big(d_{b}\log\mathrm{det}({\bm{S}}_{a})+\mathbb{E}[\mathrm{Tr}({\bm{S}}^{-1}{\bm{g}}{\bm{g}}^{\top})]\big)+\text{const.}\quad\text{{\small\color[rgb]{.5,.5,.5}\definecolor[named]{pgfstrokecolor}{rgb}{.5,.5,.5}\pgfsys@color@gray@stroke{.5}\pgfsys@color@gray@fill{.5}(linearity of the expectation)}}$ $\displaystyle\quad=\frac{1}{2}\big(d_{b}\log\mathrm{det}({\bm{S}}_{a})+\mathbb{E}[\mathrm{Tr}(d_{b}{\bm{S}}_{a}^{-1}{\bm{G}}{\bm{I}}_{b}{\bm{G}}^{\top})]\big)+\text{const.}\quad\text{{\small\color[rgb]{.5,.5,.5}\definecolor[named]{pgfstrokecolor}{rgb}{.5,.5,.5}\pgfsys@color@gray@stroke{.5}\pgfsys@color@gray@fill{.5}(identity for a Kronecker vector product)}}$ $\displaystyle\quad=\frac{d_{b}}{2}\big(\log\mathrm{det}({\bm{S}}_{a})+\mathbb{E}[\mathrm{Tr}({\bm{S}}_{a}^{-1}{\bm{G}}{\bm{G}}^{\top})]\big)+\text{const.}$ $\displaystyle\quad=\frac{d_{b}}{2}\big(-\log\mathrm{det}({\bm{P}}_{a})+\mathbb{E}[\mathrm{Tr}({\bm{P}}_{a}{\bm{G}}{\bm{G}}^{\top})]\big)+\text{const.},$ (8) $$

where ${\bm{G}}=\mathrm{Mat}({\bm{g}})$ and ${\bm{P}}_{a}:={\bm{S}}_{a}^{-1}$.

If we achieve the optimal solution, the gradient stationary condition
must be satisfied regardless of the gradient with respect to ${\bm{S}}_{a}$ or ${\bm{S}}_{a}^{-1}\equiv{\bm{P}}_{a}$, such as

$$ $\displaystyle\mathbf{0}$ $\displaystyle=\partial_{{\bm{S}}_{a}^{-1}}\mathrm{KL}\big(\mathbb{E}[{\bm{g}}{\bm{g}}^{\top}],{\bm{S}}\big)$ $\displaystyle=\partial_{{\bm{P}}_{a}}\mathrm{KL}\big(\mathbb{E}[{\bm{g}}{\bm{g}}^{\top}],{\bm{S}}\big)$ $\displaystyle=\frac{d_{b}}{2}\big(-{\bm{P}}_{a}^{-1}+\mathbb{E}[{\bm{G}}{\bm{G}}^{\top}]\big)\quad\text{{\small\color[rgb]{.5,.5,.5}\definecolor[named]{pgfstrokecolor}{rgb}{.5,.5,.5}\pgfsys@color@gray@stroke{.5}\pgfsys@color@gray@fill{.5}(use \lx@cref{creftype~refnum}{eq:proof1_eq} and matrix calculus identities)}}$ $\displaystyle=\frac{d_{b}}{2}\big(-{\bm{S}}_{a}+\mathbb{E}[{\bm{G}}{\bm{G}}^{\top}]\big).$ $$

Notice that the KL divergence is unbounded above. Thus, the optimal (minimal) solution exists. It must be ${\bm{S}}_{a}^{*}=\mathbb{E}[{\bm{G}}{\bm{G}}^{\top}]$ to satisfy this stationary condition.

## Appendix B Proof of ˜ 2

We will show that the optimal solution of KL minimization $\min_{{\bm{S}}_{a},{\bm{S}}_{b}}\mathrm{KL}\big(\mathbb{E}[{\bm{g}}{\bm{g}}^{\top}],{\bm{S}}\big)$ with a two-sided preconditioner ${\bm{S}}={\bm{S}}_{a}\otimes{\bm{S}}_{b}$ should satisfy this condition: ${\bm{S}}_{a}^{*}=\frac{1}{d_{b}}\,\mathbb{E}[{\bm{G}}\big({\bm{S}}_{b}^{*}\big)^{-1}{\bm{G}}^{\top}]$ and ${\bm{S}}_{b}^{*}=\frac{1}{d_{a}}\,\mathbb{E}[{\bm{G}}^{\top}\big({\bm{S}}_{a}^{*}\big)^{-1}{\bm{G}}]$.

Similar to the proof of [˜1](https://arxiv.org/html/2509.03378v4#Thmclaim1) in [Appx.˜A](https://arxiv.org/html/2509.03378v4#A1), we can simplify the objective function as

$$ $\displaystyle\mathrm{KL}\big(\mathbb{E}[{\bm{g}}{\bm{g}}^{\top}],{\bm{S}}\big)$ $\displaystyle\quad=\frac{1}{2}\big(\log\mathrm{det}({\bm{S}})+\mathbb{E}[\mathrm{Tr}({\bm{S}}^{-1}{\bm{g}}{\bm{g}}^{\top})]\big)+\text{const.}$ $\displaystyle\quad=\frac{1}{2}\big(d_{b}\log\mathrm{det}({\bm{S}}_{a})+d_{a}\log\mathrm{det}({\bm{S}}_{b})+\mathbb{E}[\mathrm{Tr}({\bm{S}}^{-1}{\bm{g}}{\bm{g}}^{\top})]\big)+\text{const.}\,\text{{\small\color[rgb]{.5,.5,.5}\definecolor[named]{pgfstrokecolor}{rgb}{.5,.5,.5}\pgfsys@color@gray@stroke{.5}\pgfsys@color@gray@fill{.5}(identity for a log-determinant)}}$ $\displaystyle\quad=\frac{1}{2}\big(d_{b}\log\mathrm{det}({\bm{S}}_{a})+d_{a}\log\mathrm{det}({\bm{S}}_{b})+\mathbb{E}[\mathrm{Tr}({\bm{S}}_{a}^{-1}{\bm{G}}{\bm{S}}_{b}^{-1}{\bm{G}}^{\top})]\big)+\text{const.}\,\text{{\small\color[rgb]{.5,.5,.5}\definecolor[named]{pgfstrokecolor}{rgb}{.5,.5,.5}\pgfsys@color@gray@stroke{.5}\pgfsys@color@gray@fill{.5}(identity for a Kronecker-vector-product)}}$ $\displaystyle\quad=\frac{1}{2}\big(-d_{b}\log\mathrm{det}({\bm{P}}_{a})-d_{a}\log\mathrm{det}({\bm{P}}_{b})+\mathbb{E}[\mathrm{Tr}({\bm{P}}_{a}{\bm{G}}{\bm{P}}_{b}{\bm{G}}^{\top})]\big)+\text{const.},$ (9) $$

where ${\bm{P}}_{k}:={\bm{S}}_{k}^{-1}$ for $k\in\{a,b\}$.

The optimal solution must satisfy the gradient stationarity condition with respect to $\{{\bm{S}}_{a},{\bm{S}}_{b}\}$.
Notice that the gradient with respect to $\{{\bm{S}}_{a}^{-1},{\bm{S}}_{b}^{-1}\}$ can be expressed in terms of
the gradient with respect to $\{{\bm{S}}_{a},{\bm{S}}_{b}\}$ as
$\partial_{{\bm{S}}_{a}^{-1}}\mathrm{KL}=-{\bm{S}}_{a}\big(\partial_{{\bm{S}}_{a}}\mathrm{KL}\big){\bm{S}}_{a}$ and
$\partial_{{\bm{S}}_{b}^{-1}}\mathrm{KL}=-{\bm{S}}_{b}\big(\partial_{{\bm{S}}_{b}}\mathrm{KL}\big){\bm{S}}_{b}$.
Thus, the optimal solution must satisfy the following gradient stationary condition with respect to $\{{\bm{S}}_{a}^{-1},{\bm{S}}_{b}^{-1}\}$:

$$ $\displaystyle 0=\partial_{{\bm{S}}_{a}^{-1}}\mathrm{KL}\big(\mathbb{E}[{\bm{g}}{\bm{g}}^{\top}],{\bm{S}}\big),\quad 0=\partial_{{\bm{S}}_{b}^{-1}}\mathrm{KL}\big(\mathbb{E}[{\bm{g}}{\bm{g}}^{\top}],{\bm{S}}\big).$ $$

Using [Eq.˜9](https://arxiv.org/html/2509.03378v4#A2.E9) and simplifying the left expression

$$ $\displaystyle 0$ $\displaystyle=\partial_{{\bm{S}}_{a}^{-1}}\mathrm{KL}\big(\mathbb{E}[{\bm{g}}{\bm{g}}^{\top}],{\bm{S}}\big)$ $\displaystyle=\partial_{{\bm{P}}_{a}}\mathrm{KL}\big(\mathbb{E}[{\bm{g}}{\bm{g}}^{\top}],{\bm{S}}\big)$ $\displaystyle=\frac{1}{2}\big(-d_{b}{\bm{P}}_{a}^{-1}+\mathbb{E}[{\bm{G}}{\bm{P}}_{b}{\bm{G}}^{\top}]\big)$ (10) $$

gives us this equation

$$ $\displaystyle 0=\frac{1}{2}(-d_{b}{\bm{S}}_{a}^{*}+\mathbb{E}[{\bm{G}}\big({\bm{S}}_{b}^{*}\big)^{-1}{\bm{G}}^{\top}])$ $$

that the optimal solution must satisfy.

This naturally leads to the following expression:

$$ $\displaystyle{\bm{S}}_{a}^{*}=\frac{1}{d_{b}}\mathbb{E}[{\bm{G}}\big({\bm{S}}_{b}^{*}\big)^{-1}{\bm{G}}^{\top}].$ $$

Likewise, we can obtain the following expression by simplifying the right expression of the gradient stationary condition.

$$ $\displaystyle{\bm{S}}_{b}^{*}=\frac{1}{d_{a}}\mathbb{E}[{\bm{G}}^{\top}\big({\bm{S}}_{a}^{*}\big)^{-1}{\bm{G}}].$ $$

## Appendix C Proof of ˜ 3

To simplify the notation, we
define
${\bm{H}}:=\mathbf{E}[{\bm{g}}{\bm{g}}^{\top}]$, and re-express the objective function in the KL minimzation problem as $\mathcal{L}({\bm{S}}):=\mathrm{KL}(\mathbf{E}[{\bm{g}}{\bm{g}}^{\top}],{\bm{S}})=\mathrm{KL}({\bm{H}},{\bm{S}})$.
We now introduce the proximal-gradient framework (Parikh & Boyd, 2014; Khan et al., 2016) to formally state and prove [˜3](https://arxiv.org/html/2509.03378v4#Thmclaim3).
We assume that an estimated ${\bm{S}}^{(t)}$ is given at iteration $t$.
We use a non-negative function, $f({\bm{S}}^{(t)},{\bm{S}}^{(t+1)})$,
to measure the closeness between the current and the next iteration.
Function $f(\cdot,\cdot)$ is known as a proximal function.
A (unconstrained) proximal-gradient step at iteration $t+1$ with a given proximal function, $f(\cdot,\cdot)$, is defined as
the optimal solution of
another minimization problem,

$$ $\displaystyle{\bm{S}}^{(t+1)}:=\arg\min_{{\bm{X}}}\langle\nabla_{\bm{S}}\mathcal{L}\big|_{{\bm{S}}={\bm{S}}^{(t)}},\,{\bm{X}}\rangle+\frac{1}{\beta_{2}}f({\bm{S}}^{(t)},{\bm{X}}),$ $$

at every iteration with step-size $\beta_{2}$ based on the linearization of the objective function $\mathcal{L}$.

We consider a weighted quadratic
function as the proximal function.

$$ $\displaystyle f({\bm{S}}^{(t)},{\bm{X}}):=\frac{1}{2}\|{\bm{X}}-{\bm{S}}^{(t)}\|_{{\bm{W}}}^{2}=\frac{1}{2}\mathrm{vec}\big({\bm{X}}-{\bm{S}}^{(t)}\big)^{\top}{\bm{W}}\mathrm{vec}\big({\bm{X}}-{\bm{S}}^{(t)}\big)$ $$

where ${\bm{W}}$ is a given weight matrix.
For example, ${\bm{W}}$ can be the Hessian of the KL divergence ${\bm{W}}:=\nabla_{{\bm{Y}}}^{2}\mathrm{KL}({\bm{S}}^{(t)},{\bm{Y}})\big|_{{\bm{Y}}={\bm{S}}^{(t)}}=\frac{-1}{2}\big(\frac{\partial{\bm{S}}^{-1}}{\partial{\bm{S}}}\big)\big|_{{\bm{S}}={\bm{S}}^{(t)}}$. This matrix is also known as the Fisher-Rao Riemannian metric for a zero-mean Gaussian (Amari, 2016).
Note that this proximal function has been used in the quasi-Newton literature (Nocedal & Wright, 2006). Indeed, we can show that this proximal function is exactly a second-order Taylor approximation of the KL divergence, $\mathrm{KL}({\bm{S}}^{(t)},{\bm{X}})$, at ${\bm{X}}={\bm{S}}^{(t)}$.

When ${\bm{S}}={\bm{S}}_{a}\otimes{\bm{S}}_{b}$ admits a Kronecerk product, we can specify this weight matrix ${\bm{W}}$ so that this proximal function can be separated into two terms:

$$ $\displaystyle\frac{1}{2}\|{\bm{X}}_{a}\otimes{\bm{X}}_{b}-{\bm{S}}^{(t)}\|_{{\bm{W}}}^{2}$ $\displaystyle=\frac{1}{2}\|{\bm{X}}_{a}\otimes{\bm{X}}_{b}-{\bm{S}}_{a}^{(t)}\otimes{\bm{S}}_{b}^{(t)}\|_{{\bm{W}}}^{2}$ $\displaystyle=\frac{1}{2}\|{\bm{X}}_{a}-{\bm{S}}_{a}^{(t)}\|_{{\bm{W}}_{a}}^{2}+\frac{1}{2}\|{\bm{X}}_{b}-{\bm{S}}_{b}^{(t)}\|_{{\bm{W}}_{b}}^{2}$ $$

Here, we define the weight matrix as the block-diagonal Hessian of the KL divergence, such as ${\bm{W}}:=\begin{bmatrix}{\bm{W}}_{a}&{\color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}\mathbf{0}}\\
{\color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}\mathbf{0}}&{\bm{W}}_{b}\end{bmatrix}$ by setting the cross-block terms highlighted in red to zero, where ${\bm{W}}_{k}:=\partial_{{\bm{Y}}_{k}}^{2}\mathrm{KL}({\bm{S}}^{(t)},{\bm{Y}}_{a}\otimes{\bm{Y}}_{b})\big|_{{\bm{Y}}={\bm{S}}_{a}^{(t)}\otimes{\bm{S}}_{b}^{(t)}}$ for $k\in\{a,b\}$.
We can show that this weight matrix is exactly the block-diagonal approximation of the Fisher-Rao matrix for a zero-mean matrix Gaussian considered by Lin et al. (2019, 2024).

Now, we can formally state the claim and provide proof of it.

Claim 3. (formal version)
The moving average scheme for ${\bm{S}}:={\bm{S}}_{a}\otimes{\bm{S}}_{b}$ in idealized KL-Shampoo is a proximal-gradient step at iteration $t+1$,

$$ $\displaystyle{\bm{S}}_{a}^{(t+1)},{\bm{S}}_{b}^{(t+1)}:=\arg\min_{{\bm{X}}_{a},{\bm{X}}_{b}}\langle\nabla_{{\bm{S}}_{a}}\mathcal{L}\big|_{{\bm{S}}={\bm{S}}^{(t)}},\,{\bm{X}}_{a}\rangle+\langle\nabla_{{\bm{S}}_{b}}\mathcal{L}\big|_{{\bm{S}}={\bm{S}}^{(t)}},\,{\bm{X}}_{b}\rangle+\frac{1}{2\beta_{2}}\|{\bm{X}}_{a}\otimes{\bm{X}}_{b}-{\bm{S}}^{(t)}\|_{{\bm{W}}}^{2},$ $\displaystyle\iff{\bm{S}}_{a}^{(t+1)}=(1-\beta_{2}){\bm{S}}_{a}^{(t)}+\beta_{2}\mathbb{E}[{\bm{G}}\big({\bm{S}}_{b}^{(t)}\big)^{-1}{\bm{G}}^{\top}],\quad{\bm{S}}_{b}^{(t+1)}=(1-\beta_{2}){\bm{S}}_{b}^{(t)}+\beta_{2}\mathbb{E}[{\bm{G}}^{\top}\big({\bm{S}}_{a}^{(t)}\big)^{-1}{\bm{G}}]$ $$

with step-size $\beta_{2}$ to solve the KL minimization problem in [Eq.˜3](https://arxiv.org/html/2509.03378v4#S3.E3), if we use a proximal function using the weight matrix, ${\bm{W}}$, defined above.

In mini-batch cases, we approximate the expectations using a current batch gradient (Morwani et al., 2025) (see [Eq.˜5](https://arxiv.org/html/2509.03378v4#S3.E5)), which leads to a stochastic proximal-gradient step.

#### Proof:

Because the weight matrix ${\bm{W}}$ is block-diagonal, we can slice this objective function for the proximal step into two terms.

$$ $\displaystyle\langle\nabla_{{\bm{S}}_{a}}\mathcal{L}\big|_{{\bm{S}}={\bm{S}}^{(t)}},\,{\bm{X}}_{a}\rangle+\langle\nabla_{{\bm{S}}_{b}}\mathcal{L}\big|_{{\bm{S}}={\bm{S}}^{(t)}},\,{\bm{X}}_{b}\rangle+\frac{1}{2\beta_{2}}\|{\bm{X}}_{a}\otimes{\bm{X}}_{b}-{\bm{S}}^{(t)}\|_{{\bm{W}}}^{2}$ $\displaystyle\quad=\underbrace{\langle\nabla_{{\bm{S}}_{a}}\mathcal{L}\big|_{{\bm{S}}={\bm{S}}^{(t)}},\,{\bm{X}}_{a}\rangle+\frac{1}{2\beta_{2}}\|{\bm{X}}_{a}-{\bm{S}}_{a}^{(t)}\|_{{\bm{W}}_{a}}^{2}}_{\text{(block ${\bm{X}}_{a}$)}}+\underbrace{\langle\nabla_{{\bm{S}}_{b}}\mathcal{L}\big|_{{\bm{S}}={\bm{S}}^{(t)}},\,{\bm{X}}_{b}\rangle+\frac{1}{2\beta_{2}}\|{\bm{X}}_{b}-{\bm{S}}_{b}^{(t)}\|_{{\bm{W}}_{b}}^{2}}_{\text{(block ${\bm{X}}_{b}$)}}$ $$

Importantly, ${\bm{W}}_{a}$ and ${\bm{W}}_{b}$ are independent of ${\bm{X}}_{a}$ and ${\bm{X}}_{b}$. Thus, we solve this objective by independently for each ${\bm{X}}_{k}$ for $k\in\{a,b\}$.

We now show that solving this proximal problem gives rise to the estimation rule for ${\bm{S}}_{a}^{(t+1)}$ at iteration $t+1$. We focus on the first term since the second term does not depend on ${\bm{X}}_{a}$.
We can show that ${\bm{W}}_{a}$ can be expressed as ${\bm{W}}_{a}=\partial_{{\bm{Y}}_{a}}^{2}\mathrm{KL}({\bm{S}}^{(t)},{\bm{Y}}_{a}\otimes{\bm{Y}}_{b})\big|_{{\bm{Y}}={\bm{S}}_{a}^{(t)}\otimes{\bm{S}}_{b}^{(t)}}=-\frac{d_{b}}{2}\big(\frac{\partial{\bm{S}}_{a}^{-1}}{\partial{\bm{S}}_{a}}\big)\big|_{{\bm{S}}={\bm{S}}^{(t)}}$. This matrix ${\bm{W}}_{a}$ is also considered in Lin et al. (2024).
Importantly, ${\bm{W}}_{a}$ is invertible and ${\bm{W}}_{a}^{-1}=\frac{-2}{d_{b}}\big(\frac{\partial{\bm{S}}_{a}}{\partial{\bm{S}}_{a}^{-1}}\big)\big|_{{\bm{S}}={\bm{S}}^{(t)}}$
With this result, the optimal solution of ${\bm{X}}_{a}$ must satisfy this stationary condition

$$ $\displaystyle 0$ $\displaystyle=\partial_{{\bm{X}}_{a}}\big(\langle\nabla_{{\bm{S}}_{a}}\mathcal{L}\big|_{{\bm{S}}={\bm{S}}^{(t)}},\,{\bm{X}}_{a}\rangle+\frac{1}{2\beta_{2}}\|{\bm{X}}_{a}-{\bm{S}}_{a}^{(t)}\|_{{\bm{W}}_{a}}^{2}\big)$ $\displaystyle=\nabla_{{\bm{S}}_{a}}\mathcal{L}\big|_{{\bm{S}}={\bm{S}}^{(t)}}+\frac{1}{\beta_{2}}{\bm{W}}_{a}({\bm{X}}_{a}-{\bm{S}}_{a}^{(t)})$ $\displaystyle\iff$ $\displaystyle{\bm{X}}_{a}={\bm{S}}_{a}^{(t)}-\beta_{2}{\bm{W}}_{a}^{-1}\nabla_{{\bm{S}}_{a}}\mathcal{L}\big|_{{\bm{S}}={\bm{S}}^{(t)}}$ $$

It is easy to see that the optimal solution of the proximal step is

$$ $\displaystyle{\bm{S}}_{a}^{(t+1)}$ $\displaystyle:={\bm{X}}_{a}^{*}={\bm{S}}_{a}^{(t)}-\beta_{2}{\bm{W}}_{a}^{-1}\nabla_{{\bm{S}}_{a}}\mathcal{L}\big|_{{\bm{S}}={\bm{S}}^{(t)}}$ $\displaystyle={\bm{S}}_{a}^{(t)}-\beta_{2}\underbrace{\big(\frac{-2}{d_{b}}(\frac{\partial{\bm{S}}_{a}}{\partial{\bm{S}}_{a}^{-1}}\big|_{{\bm{S}}={\bm{S}}_{a}^{(t)}})\big)}_{={\bm{W}}_{a}^{-1}}\nabla_{{\bm{S}}_{a}}\mathcal{L}\big|_{{\bm{S}}={\bm{S}}^{(t)}}$ $\displaystyle={\bm{S}}_{a}^{(t)}+\frac{2\beta_{2}}{d_{b}}\nabla_{{\bm{S}}_{a}^{-1}}\mathcal{L}\big|_{{\bm{S}}={\bm{S}}^{(t)}}\,\text{{\small\color[rgb]{.5,.5,.5}\definecolor[named]{pgfstrokecolor}{rgb}{.5,.5,.5}\pgfsys@color@gray@stroke{.5}\pgfsys@color@gray@fill{.5}(use the chain rule and utlize the Jacobian matrix contained in ${\bm{W}}_{a}^{-1}$ )}}$ $\displaystyle={\bm{S}}_{a}^{(t)}+\frac{2\beta_{2}}{d_{b}}\underbrace{\big(\frac{1}{2}(-d_{b}{\bm{S}}_{a}^{(t)}+\mathbb{E}[{\bm{G}}\big({\bm{S}}_{b}^{(t)}\big)^{-1}{\bm{G}}^{\top}])\big)}_{=\nabla_{{\bm{S}}_{a}^{-1}}\mathcal{L}\big|_{{\bm{S}}={\bm{S}}^{(t)}}}\,\text{{\small\color[rgb]{.5,.5,.5}\definecolor[named]{pgfstrokecolor}{rgb}{.5,.5,.5}\pgfsys@color@gray@stroke{.5}\pgfsys@color@gray@fill{.5}(recall the definition of $\mathcal{L}$ and use \lx@cref{creftype~refnum}{eq:grad_S_a})}}$ $\displaystyle=(1-\beta_{2}){\bm{S}}_{a}^{(t)}+\frac{\beta_{2}}{d_{b}}\mathbb{E}[{\bm{G}}\big({\bm{S}}_{b}^{(t)}\big)^{-1}{\bm{G}}^{\top}],$ $$

which is exactly the moving average scheme in [Eq.˜5](https://arxiv.org/html/2509.03378v4#S3.E5) for updating ${\bm{S}}_{a}$ at iteration $t+1$.

Likewise, we can obtain the moving average scheme for ${\bm{S}}_{b}$.
$\square$

## Appendix D Proof of ˜ 4

We will show that the optimal solution of KL minimization $\min_{{\bm{\lambda}}_{a},{\bm{\lambda}}_{b}}\mathrm{KL}\big(\mathbb{E}[{\bm{g}}{\bm{g}}^{\top}],{\bm{S}}\big)$ with a two-sided preconditioner ${\bm{S}}=({\bm{Q}}_{a}\mathrm{Diag}({\bm{\lambda}}_{a}){\bm{Q}}_{a}^{\top})\otimes({\bm{Q}}_{b}\mathrm{Diag}({\bm{\lambda}}_{b}){\bm{Q}}_{b}^{\top})$ should satisfy this condition:
${\bm{\lambda}}_{a}^{*}=\frac{1}{d_{b}}\mathrm{diag}\big({\bm{Q}}_{a}^{\top}\mathbb{E}[{\bm{G}}{{\bm{P}}}_{b}^{*}{\bm{G}}^{\top}]{\bm{Q}}_{a}\big)$ and ${\bm{\lambda}}_{b}^{*}=\frac{1}{d_{a}}\mathrm{diag}\big({\bm{Q}}_{b}^{\top}\mathbb{E}[{\bm{G}}^{\top}{{\bm{P}}}_{a}^{*}{\bm{G}}]{\bm{Q}}_{b}\big)$,
where ${{\bm{P}}}_{k}^{*}:={\bm{Q}}_{k}\mathrm{Diag}\left(({\bm{\lambda}}_{k}^{*})^{\odot-1}\right){\bm{Q}}_{k}^{\top}$, and ${\bm{Q}}_{k}$ is known and precomputed by QR for $k\in\{a,b\}$.

Let ${\bm{S}}_{k}:={\bm{Q}}_{k}\mathrm{Diag}({\bm{\lambda}}_{k}){\bm{Q}}_{k}^{\top}$ for $k\in\{a,b\}$. Because ${\bm{Q}}_{k}$ is orthogonal, it is easy to see that
${\bm{S}}_{k}^{-1}:={\bm{Q}}_{k}\mathrm{Diag}(\big({\bm{\lambda}}_{k}\big)^{\odot-1}){\bm{Q}}_{k}^{\top}$.

Similar to the proof of [˜2](https://arxiv.org/html/2509.03378v4#Thmclaim2) in [Appx.˜B](https://arxiv.org/html/2509.03378v4#A2), we can simplify the following objective function by substituting ${\bm{S}}_{a}$ and ${\bm{S}}_{b}$. Here, we also utilize the orthogonality of ${\bm{Q}}_{k}$ for $k\in\{a,b\}$.

$$ $\displaystyle\mathrm{KL}\big(\mathbb{E}[{\bm{g}}{\bm{g}}^{\top}],{\bm{S}}\big)$ $\displaystyle\quad=\frac{1}{2}\big(d_{b}\log\mathrm{det}({\bm{S}}_{a})+d_{a}\log\mathrm{det}({\bm{S}}_{b})+\mathbb{E}[\mathrm{Tr}({\bm{S}}_{a}^{-1}{\bm{G}}{\bm{S}}_{b}^{-1}{\bm{G}}^{\top})]\big)+\text{const.}$ $\displaystyle\quad=\frac{1}{2}\big(d_{b}\log\mathrm{det}({\bm{Q}}_{a}\mathrm{Diag}({\bm{\lambda}}_{a}){\bm{Q}}_{a}^{\top})+d_{a}\log\mathrm{det}({\bm{Q}}_{b}\mathrm{Diag}({\bm{\lambda}}_{b}){\bm{Q}}_{b}^{\top})+\mathbb{E}[\mathrm{Tr}({\bm{S}}_{a}^{-1}{\bm{G}}{\bm{S}}_{b}^{-1}{\bm{G}}^{\top})]\big)+\text{const.}$ $\displaystyle\quad=\frac{1}{2}\big((d_{b}\sum_{i}\log(\lambda_{a}^{(i)}))+(d_{a}\sum_{j}\log(\lambda_{b}^{(j)}))+\mathbb{E}[\mathrm{Tr}({\bm{S}}_{a}^{-1}{\bm{G}}{\bm{S}}_{b}^{-1}{\bm{G}}^{\top})]\big)+\text{const.}\,\text{{\small\color[rgb]{.5,.5,.5}\definecolor[named]{pgfstrokecolor}{rgb}{.5,.5,.5}\pgfsys@color@gray@stroke{.5}\pgfsys@color@gray@fill{.5}(use the orthogonality of ${\bm{Q}}_{a}$ and ${\bm{Q}}_{b}$ )}}$ $\displaystyle\quad=\frac{1}{2}\big((d_{b}\sum_{i}\log(\lambda_{a}^{(i)}))+(d_{a}\sum_{j}\log(\lambda_{b}^{(j)}))+\mathbb{E}[\mathrm{Tr}(\underbrace{{\bm{Q}}_{a}\mathrm{Diag}({\bm{\lambda}}_{a}^{\odot-1}){\bm{Q}}_{a}^{\top}}_{={\bm{S}}_{a}^{-1}}{\bm{G}}\underbrace{{\bm{Q}}_{b}\mathrm{Diag}({\bm{\lambda}}_{b}^{\odot-1}){\bm{Q}}_{b}^{\top}}_{={\bm{S}}_{b}^{-1}}{\bm{G}}^{\top})]\big)+\text{const.}$ (11) $$

The optimal ${\bm{\lambda}}_{a}$ and ${\bm{\lambda}}_{b}$ should satisfy the gradient stationary condition.

$$ $\displaystyle 0$ $\displaystyle=\partial_{{\bm{\lambda}}_{a}}\mathrm{KL}\big(\mathbb{E}[{\bm{g}}{\bm{g}}^{\top}],{\bm{S}}\big)$ $\displaystyle=\frac{1}{2}\big(d_{b}{\bm{\lambda}}_{a}^{\odot-1}+\partial_{{\bm{\lambda}}_{a}}\mathbb{E}[\mathrm{Tr}({\bm{Q}}_{a}\mathrm{Diag}({\bm{\lambda}}_{a}^{\odot-1}){\bm{Q}}_{a}^{\top}{\bm{G}}\overbrace{{\bm{Q}}_{b}\mathrm{Diag}({\bm{\lambda}}_{b}^{\odot-1}){\bm{Q}}_{b}^{\top}}^{={\bm{P}}_{b}}{\bm{G}}^{\top})]\big)\,\,\,\text{{\small\color[rgb]{.5,.5,.5}\definecolor[named]{pgfstrokecolor}{rgb}{.5,.5,.5}\pgfsys@color@gray@stroke{.5}\pgfsys@color@gray@fill{.5}(use \lx@cref{creftype~refnum}{eq:proof_eq3})}}$ $\displaystyle=\frac{1}{2}\big(d_{b}{\bm{\lambda}}_{a}^{\odot-1}+\partial_{{\bm{\lambda}}_{a}}\mathbb{E}[\mathrm{Tr}(\mathrm{Diag}({\bm{\lambda}}_{a}^{\odot-1}){\bm{Q}}_{a}^{\top}{\bm{G}}{\bm{P}}_{b}{\bm{G}}^{\top}{\bm{Q}}_{a})]\big)$ $\displaystyle=\frac{1}{2}\big(d_{b}{\bm{\lambda}}_{a}^{\odot-1}+\partial_{{\bm{\lambda}}_{a}}\mathbb{E}[{\bm{\lambda}}_{a}^{\odot-1}\odot\mathrm{diag}\big({\bm{Q}}_{a}^{\top}{\bm{G}}{\bm{P}}_{b}{\bm{G}}^{\top}{\bm{Q}}_{a}\big)]\big)\,\,\,\text{{\small\color[rgb]{.5,.5,.5}\definecolor[named]{pgfstrokecolor}{rgb}{.5,.5,.5}\pgfsys@color@gray@stroke{.5}\pgfsys@color@gray@fill{.5}(utilize the trace and the diagonal structure)}}$ $\displaystyle=\frac{1}{2}\big(d_{b}{\bm{\lambda}}_{a}^{\odot-1}-\mathbb{E}[{\bm{\lambda}}_{a}^{\odot-2}\odot\mathrm{diag}\big({\bm{Q}}_{a}^{\top}{\bm{G}}{\bm{P}}_{b}{\bm{G}}^{\top}{\bm{Q}}_{a}\big)]\big)$ $\displaystyle=\frac{1}{2}\big(d_{b}{\bm{\lambda}}_{a}^{\odot-1}-{\bm{\lambda}}_{a}^{\odot-2}\odot\mathrm{diag}\big({\bm{Q}}_{a}^{\top}\mathbb{E}[{\bm{G}}{\bm{P}}_{b}{\bm{G}}^{\top}]{\bm{Q}}_{a}\big)\big)$ $\displaystyle\iff 0$ $\displaystyle=d_{b}{\bm{\lambda}}_{a}-\mathrm{diag}\big({\bm{Q}}_{a}^{\top}\mathbb{E}[{\bm{G}}{\bm{P}}_{b}{\bm{G}}^{\top}]{\bm{Q}}_{a}\big)\big)$ $$

We obtain the optimal solution by solving this equation.

$$ $\displaystyle{\bm{\lambda}}_{a}^{*}=\frac{1}{d_{b}}\mathrm{diag}\big({\bm{Q}}_{a}^{\top}\mathbb{E}[{\bm{G}}{\bm{P}}_{b}^{*}{\bm{G}}^{\top}]{\bm{Q}}_{a}\big)\big)$ $$

Similarly, we can obtain the other expression.

## Appendix E Proof of ˜ 5

This proof is similar to the proof of [˜4](https://arxiv.org/html/2509.03378v4#Thmclaim4) in [Appx.˜D](https://arxiv.org/html/2509.03378v4#A4).
We will show that the optimal solution of KL minimization $\min_{{\bm{d}}}\mathrm{KL}\big(\mathbb{E}[{\bm{g}}{\bm{g}}^{\top}],{\bm{S}}\big)$ with an augmented preconditioner ${\bm{S}}=({\bm{Q}}\mathrm{Diag}({\bm{d}}){\bm{Q}}^{\top})$ is
${\bm{d}}^{*}=\mathbb{E}\left[\left(\mathrm{vec}({\bm{Q}}_{a}^{\top}{\bm{G}}{\bm{Q}}_{b})\right)^{\odot 2}\right]$, where ${\bm{d}}\in\mathcal{R}^{d_{a}d_{b}\times 1}$ is an augmented eigenvalue vector, ${\bm{Q}}:={\bm{Q}}_{a}\otimes{\bm{Q}}_{b}$, and ${\bm{Q}}_{k}$ is given and precomputed by QR for $k\in\{a,b\}$.

We can simplify the objective function by substituting ${\bm{S}}$. Here, we also utilize the orthogonality of ${\bm{Q}}_{k}$ for $k\in\{a,b\}$.

$$ $\displaystyle\mathrm{KL}\big(\mathbb{E}[{\bm{g}}{\bm{g}}^{\top}],{\bm{S}}\big)$ $\displaystyle\quad=\frac{1}{2}\big(\log\mathrm{det}({\bm{Q}}\mathrm{Diag}({\bm{d}}){\bm{Q}}^{\top})+\mathrm{Tr}({\bm{Q}}\mathrm{Diag}({\bm{d}}^{\odot-1}){\bm{Q}}^{\top}\mathbb{E}[{\bm{g}}{\bm{g}}^{\top}])\big)+\text{const.}$ $\displaystyle\quad=\frac{1}{2}\big(\sum_{i}\log(d_{i}))+\mathrm{Tr}({\bm{Q}}\mathrm{Diag}({\bm{d}}^{\odot-1}){\bm{Q}}^{\top}\mathbb{E}[{\bm{g}}{\bm{g}}^{\top}])\big)+\text{const.}\quad{\small\color[rgb]{.5,.5,.5}\definecolor[named]{pgfstrokecolor}{rgb}{.5,.5,.5}\pgfsys@color@gray@stroke{.5}\pgfsys@color@gray@fill{.5}\text{(${\bm{Q}}={\bm{Q}}_{a}\otimes{\bm{Q}}_{b}$ is orthogonal)}}$ $\displaystyle\quad=\frac{1}{2}\big(\sum_{i}\log(d_{i}))+\mathbb{E}\big[\mathrm{Tr}({\bm{Q}}\mathrm{Diag}({\bm{d}}^{\odot-1}){\bm{Q}}^{\top}{\bm{g}}{\bm{g}}^{\top})\big)\big]+\text{const.}\quad{\small\color[rgb]{.5,.5,.5}\definecolor[named]{pgfstrokecolor}{rgb}{.5,.5,.5}\pgfsys@color@gray@stroke{.5}\pgfsys@color@gray@fill{.5}\text{(linearity of the expectation)}}$ $\displaystyle\quad=\frac{1}{2}\big(\sum_{i}\log(d_{i}))+\mathbb{E}\big[\mathrm{Tr}((\mathrm{vec}({\bm{Q}}_{a}^{\top}{\bm{G}}{\bm{Q}}_{b}))^{\top}\mathrm{Diag}({\bm{d}}^{\odot-1})\mathrm{vec}({\bm{Q}}_{a}^{\top}{\bm{G}}{\bm{Q}}_{b})\big]+\text{const.}\,{\small\color[rgb]{.5,.5,.5}\definecolor[named]{pgfstrokecolor}{rgb}{.5,.5,.5}\pgfsys@color@gray@stroke{.5}\pgfsys@color@gray@fill{.5}\text{(identity of Kronecker-vector product)}}$ $\displaystyle\quad=\frac{1}{2}\big(\sum_{i}\log(d_{i}))+\mathbb{E}\big[\mathrm{sum}({\bm{d}}^{\odot-1}\odot(\mathrm{vec}({\bm{Q}}_{a}^{\top}{\bm{G}}{\bm{Q}}_{b}))^{\odot 2}\big]+\text{const.}\,{\small\color[rgb]{.5,.5,.5}\definecolor[named]{pgfstrokecolor}{rgb}{.5,.5,.5}\pgfsys@color@gray@stroke{.5}\pgfsys@color@gray@fill{.5}\text{(leverage trace and diagonal struct.)}}$ (12) $$

The optimal ${\bm{d}}$ should satisfy the gradient stationary condition.

$$ $\displaystyle 0$ $\displaystyle=\partial_{{\bm{d}}}\mathrm{KL}\big(\mathbb{E}[{\bm{g}}{\bm{g}}^{\top}],{\bm{S}}\big)$ $\displaystyle=\frac{1}{2}\big({\bm{d}}^{\odot-1}-\mathbb{E}\big[{\bm{d}}^{\odot-2}\odot\mathrm{vec}({\bm{Q}}_{a}^{\top}{\bm{G}}{\bm{Q}}_{b})^{\odot 2}\big)\big]\big)\quad{\small\color[rgb]{.5,.5,.5}\definecolor[named]{pgfstrokecolor}{rgb}{.5,.5,.5}\pgfsys@color@gray@stroke{.5}\pgfsys@color@gray@fill{.5}\text{(use \lx@cref{creftype~refnum}{eq:obj_soap} and compute its derivative)}}$ $\displaystyle\iff 0$ $\displaystyle=\frac{1}{2}\big({\bm{d}}-\mathbb{E}\big[\mathrm{vec}({\bm{Q}}_{a}^{\top}{\bm{G}}{\bm{Q}}_{b})^{\odot 2}\big)\big]\big)$ $$

Notice that the KL divergence is unbounded above. Thus, the optimal (minimal) solution exists and it must be ${\bm{d}}^{*}=\mathbb{E}\big[\mathrm{vec}({\bm{Q}}_{a}^{\top}{\bm{G}}{\bm{Q}}_{b})^{\odot 2}\big]$ to satisfy the condition.

## Appendix F Two-sided Shampoo Scheme based on Frobenius norm

###### Claim 6

(Shampoo’s estimation scheme based on Frobenius norm)
The optimal solution of the Frobenius norm minimization $\min_{{\bm{S}}_{a},{\bm{S}}_{b}}\mathrm{Frob}\big(\mathbb{E}[{\bm{g}}{\bm{g}}^{\top}],{\bm{S}}\big):=\|\mathbb{E}[{\bm{g}}{\bm{g}}^{\top}]-{\bm{S}}\|_{\text{Frob}}$ with a two-sided precontioner ${\bm{S}}={\bm{S}}_{a}\otimes{\bm{S}}_{b}$ should satisfy the following condition.

$$ $\displaystyle{\bm{S}}_{a}^{*}=\frac{1}{\mathrm{Tr}(({\bm{S}}_{b}^{*})^{2})}\,\mathbb{E}[{\bm{G}}{\bm{S}}_{b}^{*}{\bm{G}}^{\top}],\quad{\bm{S}}_{b}^{*}=\frac{1}{\mathrm{Tr}(({\bm{S}}_{a}^{*})^{2})}\,\mathbb{E}[{\bm{G}}^{\top}{\bm{S}}_{a}^{*}{\bm{G}}],$ (13) $$

Remark:
Although the solution can be obtained via rank-1 singular value decomposition (SVD) (Van Loan & Pitsianis, 1993) on this outer product, $\mathbb{E}[{\bm{g}}{\bm{g}}^{\top}]$, it can be computationally expensive to compute the solution due to the high dimensionality of the product. Moreover, the optimal solution is only achievable when the expectation of the outer product is computed exactly.
Obtaining the optimal solution using SVD is even more expensive in tensor-valued cases.

###### Claim 6

#### Proof:

To simplify the proof, we will consider the square of the objective function, as the optimal solution remains unchanged.
We simplify the square of the objective function by substituting ${\bm{S}}$. Here, we utilize the definition of the norm and re-express the norm using the matrix trace.

$$ $\displaystyle\|\mathbb{E}[{\bm{g}}{\bm{g}}^{\top}]-{\bm{S}}_{a}\otimes{\bm{S}}_{b}\|_{\text{Frob}}^{2}$ $\displaystyle\quad=\mathrm{Tr}\big((\mathbb{E}[{\bm{g}}{\bm{g}}^{\top}]-{\bm{S}}_{a}\otimes{\bm{S}}_{b})^{\top}(\mathbb{E}[{\bm{g}}{\bm{g}}^{\top}]-{\bm{S}}_{a}\otimes{\bm{S}}_{b})\big)\quad\text{ {\small\color[rgb]{.5,.5,.5}\definecolor[named]{pgfstrokecolor}{rgb}{.5,.5,.5}\pgfsys@color@gray@stroke{.5}\pgfsys@color@gray@fill{.5}(an equivalent definition of the square of the norm)}}$ $\displaystyle\quad=\mathrm{Tr}\big({\bm{S}}_{a}^{2}\otimes{\bm{S}}_{b}^{2}-2\mathbb{E}[{\bm{g}}{\bm{g}}^{\top}]({\bm{S}}_{a}\otimes{\bm{S}}_{b})\big)+\text{const.}\quad\text{{\small\color[rgb]{.5,.5,.5}\definecolor[named]{pgfstrokecolor}{rgb}{.5,.5,.5}\pgfsys@color@gray@stroke{.5}\pgfsys@color@gray@fill{.5}(${\bm{S}}_{k}$ is symmetric for $k\in\{a,b\}$)}}$ $\displaystyle\quad=\mathrm{Tr}\big({\bm{S}}_{a}^{2}\big)\mathrm{Tr}\big({\bm{S}}_{b}^{2}\big)-2\mathrm{Tr}\big(\mathbb{E}[{\bm{g}}{\bm{g}}^{\top}]({\bm{S}}_{a}\otimes{\bm{S}}_{b})\big)+\text{const.}\quad\text{{\small\color[rgb]{.5,.5,.5}\definecolor[named]{pgfstrokecolor}{rgb}{.5,.5,.5}\pgfsys@color@gray@stroke{.5}\pgfsys@color@gray@fill{.5}(Property of a Kronecker product)}}$ $\displaystyle\quad=\mathrm{Tr}\big({\bm{S}}_{a}^{2}\big)\mathrm{Tr}\big({\bm{S}}_{b}^{2}\big)-2\mathbb{E}\big[\mathrm{Tr}\big(({\bm{g}}{\bm{g}}^{\top})({\bm{S}}_{a}\otimes{\bm{S}}_{b})\big)\big]+\text{const.}\quad\text{{\small\color[rgb]{.5,.5,.5}\definecolor[named]{pgfstrokecolor}{rgb}{.5,.5,.5}\pgfsys@color@gray@stroke{.5}\pgfsys@color@gray@fill{.5}(linearity of the expectation)}}$ $\displaystyle\quad=\mathrm{Tr}\big({\bm{S}}_{a}^{2}\big)\mathrm{Tr}\big({\bm{S}}_{b}^{2}\big)-2\mathbb{E}\big[\mathrm{Tr}\big({\bm{g}}^{\top}\mathrm{vec}({\bm{S}}_{a}{\bm{G}}{\bm{S}}_{b})\big)\big]+\text{const.}\quad\text{{\small\color[rgb]{.5,.5,.5}\definecolor[named]{pgfstrokecolor}{rgb}{.5,.5,.5}\pgfsys@color@gray@stroke{.5}\pgfsys@color@gray@fill{.5}(Property of a Kronecker product)}}$ $\displaystyle\quad=\mathrm{Tr}\big({\bm{S}}_{a}^{2}\big)\mathrm{Tr}\big({\bm{S}}_{b}^{2}\big)-2\mathbb{E}\big[\mathrm{Tr}\big({\bm{G}}^{\top}{\bm{S}}_{a}{\bm{G}}{\bm{S}}_{b}\big)\big]+\text{const.}\quad\text{{\small\color[rgb]{.5,.5,.5}\definecolor[named]{pgfstrokecolor}{rgb}{.5,.5,.5}\pgfsys@color@gray@stroke{.5}\pgfsys@color@gray@fill{.5}(Property of a trace)}}$ $$

We can simplify the stationary condition with respect to ${\bm{S}}_{a}$ as below.

$$ $\displaystyle 0$ $\displaystyle=\partial_{{\bm{S}}_{a}}\|\mathbb{E}[{\bm{g}}{\bm{g}}^{\top}]-{\bm{S}}_{a}\otimes{\bm{S}}_{b}\|_{\text{Frob}}^{2}$ $\displaystyle=\partial_{{\bm{S}}_{a}}\big(\mathrm{Tr}\big({\bm{S}}_{a}^{2}\big)\mathrm{Tr}\big({\bm{S}}_{b}^{2}\big)-2\mathbb{E}\big[\mathrm{Tr}\big({\bm{G}}^{\top}{\bm{S}}_{a}{\bm{G}}{\bm{S}}_{b}\big)\big]+\text{const.}\big)$ $\displaystyle=2\big(\mathrm{Tr}({\bm{S}}_{b}^{2}){\bm{S}}_{a}-\mathbb{E}[{\bm{G}}{\bm{S}}_{b}{\bm{G}}^{\top}]\big)$ $$

Thus, the optimal solution should satisfy this condition
${\bm{S}}_{a}^{*}=\frac{1}{\mathrm{Tr}\big(({\bm{S}}_{b}^{*})^{2}\big)}\mathbb{E}[{\bm{G}}{\bm{S}}_{b}^{*}{\bm{G}}^{\top}]$.
Similarly, we can obtain the other condition.
Morwani et al. (2025) also consider a similar condition (see Eq. 4 of their paper).
$\square$

Figure: Figure 6: *Left:* Simplified two-sided Shampoo schemes based on the Frobenius norm without momentum. We consider two variants. Variant 1 is inspired by [˜6](https://arxiv.org/html/2509.03378v4#Thmclaim6), while Variant 2 is similar to KL-Shampoo’s update scheme, which utilizes eigenvalues. Note that Variant 1 of the idealized F-Shampoo is known as the two-sided Shampoo in the literature (Morwani et al., 2025). *Right:* Adapting our exponential moving average (EMA) approach enables F-Shampoo to use the faster QR procedure and makes it more competitive, as empirically shown in [Fig.˜7](https://arxiv.org/html/2509.03378v4#A6.F7).

Figure: Figure 7: Empirical results from a random search with 150 runs per method on language models demonstrate that our exponential moving average (EMA) scheme for eigenvalue estimation, as described in [Fig.˜6](https://arxiv.org/html/2509.03378v4#A6.F6.fig1), improves the performance of the two-sided Shampoo based on Frobenius norm (see Eq. 4 of Morwani et al. (2025) and [˜6](https://arxiv.org/html/2509.03378v4#Thmclaim6))—referred to as Variant 1 of idealized F-Shampoo. All these methods perform QR or eigen decompostion at every 10 iterations. Note that F-shampoo cannot match the performance of Shampoo with step-size grafting. This also illustrates using the Frobenius norm for preconditioner estimation is not ideal. To ensure a fair comparison and eliminate implementation bias, we use our own implementation of F-Shampoo, aligned closely with that of KL-Shampoo. As a reference, we also include the best Shampoo run with power $p=\nicefrac{{1}}{{2}}$ and grafting based on the state-of-the-art version from Meta (Shi et al., 2023).
Refer to caption: x5.png

## Appendix G Key Distinction between Shampoo with trace scaling and KL-Shampoo

We will show that Shampoo’s estimation with trace scaling is a generalization of Adafactor. Our interpretation of Shampoo’s update is grounded in a generalization of the divergence used in Adafactor—quantum relative entropy (Tsuda et al., 2005)—a Bregman divergence (Bregman, 1967) defined on the trace of the matrix logarithm. This new view of Shampoo’s estimation is distinct from the existing Frobenius-norm perspective. By contrast, KL-Shampoo’s update is based on the KL divergence (classical relative entropy)—another Bregman divergence, but one defined on the (scalar) logarithm of the matrix determinant.

We now introduce the definition of a Bregman divergence to formally discuss the distinction between Shampoo with trace scaling and KL-Shampoo.
Given a strictly convex and differentiable (scalar) function $F(\cdot)$, the Bregman divergence based on this function is defined as

$$ $\displaystyle\mathcal{B}_{F}({\bm{X}},{\bm{Y}}):=F({\bm{X}})-F({\bm{Y}})-\mathrm{Tr}\big([\nabla F({\bm{Y}})]({\bm{X}}-{\bm{Y}})\big).$ $$

As an example, the KL divergence (classical relative entropy) $\mathrm{KL}({\bm{X}},{\bm{Y}})$ is a Bregman divergence with convex function $F({\bm{M}}):=-\frac{1}{2}\log\mathrm{det}({\bm{M}})$.

$$ $\displaystyle\mathcal{B}_{F}({\bm{X}},{\bm{Y}})$ $\displaystyle=F({\bm{X}})-F({\bm{Y}})-\mathrm{Tr}\big([\nabla F({\bm{Y}})]({\bm{X}}-{\bm{Y}})\big)$ $\displaystyle=\frac{1}{2}\big(-\log\mathrm{det}({\bm{X}})+\log\mathrm{det}({\bm{Y}})+\mathrm{Tr}({\bm{Y}}^{-1}({\bm{X}}-{\bm{Y}})\big)\quad\text{{\small\color[rgb]{.5,.5,.5}\definecolor[named]{pgfstrokecolor}{rgb}{.5,.5,.5}\pgfsys@color@gray@stroke{.5}\pgfsys@color@gray@fill{.5}(defn. of function $F(\cdot)$)}}$ $\displaystyle=\frac{1}{2}\big(\log\mathrm{det}({\bm{Y}})-\log\mathrm{det}({\bm{X}})+\mathrm{Tr}({\bm{Y}}^{-1}{\bm{X}})-\mathrm{dim}({\bm{X}})\big)=\mathrm{KL}({\bm{X}},{\bm{Y}})$ $$

where $\nabla F({\bm{M}})=-\frac{1}{2}{\bm{M}}^{-1}$.
The KL divergence is also known as the log-determinant divergence because function $F$ is defined as the logarithm of the matrix determinant.
Notably, the Hessian of this $F(\cdot)$ gives rise to the Fisher-Rao metric, which is also known as the affine-invariant metric (up to a constant scalar) (Lin et al., 2023).

Now, we introduce quantum relative entropy, which is also known as von Neumann (VN) divergence, to show that Shampoo with trace scaling is a generalization of Adafactor.
The VN divergence $\mathrm{VN}({\bm{X}},{\bm{Y}})$ is defined as a Bregman divergence with convex function $F({\bm{M}}):=\mathrm{Tr}\big({\bm{M}}\mathrm{LogM}({\bm{M}})-{\bm{M}}\big)$:

$$ $\displaystyle\mathrm{VN}({\bm{X}},{\bm{Y}}):=\mathcal{B}_{F}({\bm{X}},{\bm{Y}})$ $\displaystyle\quad=F({\bm{X}})-F({\bm{Y}})-\mathrm{Tr}\big([\nabla F({\bm{Y}})]({\bm{X}}-{\bm{Y}})\big)$ $\displaystyle\quad=\mathrm{Tr}\big({\bm{X}}\mathrm{LogM}({\bm{X}})-{\bm{X}}-{\bm{Y}}\mathrm{LogM}({\bm{Y}})+{\bm{Y}}-\mathrm{LogM}({\bm{Y}})({\bm{X}}-{\bm{Y}})\big)\quad\text{{\small\color[rgb]{.5,.5,.5}\definecolor[named]{pgfstrokecolor}{rgb}{.5,.5,.5}\pgfsys@color@gray@stroke{.5}\pgfsys@color@gray@fill{.5}(defn. of function $F(\cdot)$)}}$ $\displaystyle\quad=\mathrm{Tr}\big({\bm{X}}\mathrm{LogM}({\bm{X}})-{\bm{X}}-\mathrm{LogM}({\bm{Y}}){\bm{Y}}+{\bm{Y}}-\mathrm{LogM}({\bm{Y}})({\bm{X}}-{\bm{Y}})\big)\quad\text{{\small\color[rgb]{.5,.5,.5}\definecolor[named]{pgfstrokecolor}{rgb}{.5,.5,.5}\pgfsys@color@gray@stroke{.5}\pgfsys@color@gray@fill{.5}(property of the trace)}}$ $\displaystyle\quad=\mathrm{Tr}\big({\bm{X}}\mathrm{LogM}({\bm{X}})-{\bm{X}}+{\bm{Y}}-\mathrm{LogM}({\bm{Y}}){\bm{X}}\big)\,$ $\displaystyle\quad=\mathrm{Tr}\big({\bm{X}}[\mathrm{LogM}({\bm{X}})-\mathrm{LogM}({\bm{Y}})]\big)-\mathrm{Tr}({\bm{X}})+\mathrm{Tr}({\bm{Y}}),$ $$

where $\mathrm{LogM}(\cdot)$ is the matrix logarithm function and Tsuda et al. (2005) show that $\nabla F({\bm{M}})=\mathrm{LogM}({\bm{M}})$. The Hessian of this $F(\cdot)$ gives rise to the Bogoliubov-Kubo-Mori (BKM) metric in quantum physics (de Boer et al., 2023).

###### Claim 7

(Shampoo’s estimation scheme with trace scaling)
The optimal solution of the
von Neumann (VN) divergence (quantum relative entropy)
minimization $\min_{{\bm{S}}_{a},{\bm{S}}_{b}}\mathrm{VN}\big(\mathbb{E}[{\bm{g}}{\bm{g}}^{\top}],{\bm{S}}\big):=\mathrm{Tr}({\bm{S}})-\mathrm{Tr}\big(\mathbb{E}[{\bm{g}}{\bm{g}}^{\top}]\mathrm{LogM}({\bm{S}})\big)+\text{const.}$ with a two-sided precontioner ${\bm{S}}={\bm{S}}_{a}\otimes{\bm{S}}_{b}$ should satisfy the following condition.

$$ $\displaystyle{\bm{S}}_{a}^{*}=\frac{1}{\mathrm{Tr}({\bm{S}}_{b}^{*})}\,\mathbb{E}[{\bm{G}}{\bm{G}}^{\top}],\quad{\bm{S}}_{b}^{*}=\frac{1}{\mathrm{Tr}({\bm{S}}_{a}^{*})}\,\mathbb{E}[{\bm{G}}^{\top}{\bm{G}}],$ (14) $$

where $\mathrm{LogM}(\cdot)$ is the matrix logarithm function.

The optimal solutions is Shampoo’s estimation rule (power $p=\frac{1}{2}$) with trace scaling:

$$ $\displaystyle{\bm{S}}_{a}^{*}=\mathbb{E}[{\bm{G}}{\bm{G}}^{\top}],\quad{\bm{S}}_{b}^{*}=\frac{\mathbb{E}[{\bm{G}}^{\top}{\bm{G}}]}{\mathrm{Tr}(\mathbb{E}[{\bm{G}}{\bm{G}}^{\top}])}$ $$

If we force ${\bm{S}}_{a}$ and ${\bm{S}}_{b}$ to be diagonal matrices and solve
the minimization problem, we obtain Adafactor’s update as shown below.

$$ $\displaystyle{\bm{S}}_{a}^{*}$ $\displaystyle=\mathrm{Diag}\big(\mathbb{E}[{\bm{G}}{\bm{G}}^{\top}]\big)=\mathrm{Diag}\big(\mathbb{E}[\big({\bm{G}}^{\odot 2}\big)\mathbf{1}]\big)$ $\displaystyle{\bm{S}}_{b}^{*}$ $\displaystyle=\mathrm{Diag}\left(\frac{\mathbb{E}[{\bm{G}}^{\top}{\bm{G}}]}{\mathrm{Tr}\big(\mathbb{E}[{\bm{G}}{\bm{G}}^{\top}]\big)}\right)=\frac{\mathrm{Diag}\big(\mathbb{E}[\mathbf{1}^{\top}{\bm{G}}^{\odot 2}]\big)}{\mathrm{Tr}\big(\mathbb{E}[\mathbf{1}^{\top}\big({\bm{G}}^{\odot 2}\big)\mathbf{1}]\big)}=\frac{\mathrm{Diag}\big(\mathbb{E}[\mathbf{1}^{\top}\big({\bm{G}}^{\odot 2}\big)]\big)}{\sqrt{\mathrm{Tr}\big(\mathbb{E}[\mathbf{1}^{\top}\big({\bm{G}}^{\odot 2}\big)\big)\mathrm{Tr}\big(\mathbb{E}[\big({\bm{G}}^{\odot 2}\big)\mathbf{1}]\big)}}$ $$

###### Claim 7

#### Remark:

If the expectations are not computed exactly, the resulting update scheme is not the optimal solution. For example, Adafactor’s update scheme is not optimal due to the EMA scheme on the diagonal Kronecker factors.

#### Proof:

We will show that Shampoo’s update scheme with trace scaling is an optimal solution to this minimization problem.
We first simplify the objective function when ${\bm{S}}={\bm{S}}_{a}\otimes{\bm{S}}_{b}$.
We will use this (Kronecker sum) identity,
$\mathrm{LogM}({\bm{S}}_{a}\otimes{\bm{S}}_{b})=\mathrm{LogM}({\bm{S}}_{a})\otimes{\bm{I}}_{b}+{\bm{I}}_{a}\otimes\mathrm{LogM}({\bm{S}}_{b})$, to simplify the matrix logarithm.

$$ $\displaystyle\mathrm{VN}(\mathbb{E}[{\bm{g}}{\bm{g}}^{\top}],{\bm{S}})$ $\displaystyle=\mathrm{Tr}({\bm{S}})-\mathrm{Tr}\big(\mathbb{E}[{\bm{g}}{\bm{g}}^{\top}]\mathrm{LogM}({\bm{S}})\big)+\text{const.}$ $\displaystyle=\mathrm{Tr}({\bm{S}}_{a})\,\mathrm{Tr}({\bm{S}}_{b})-\mathrm{Tr}\big(\mathbb{E}[{\bm{g}}{\bm{g}}^{\top}]\mathrm{LogM}({\bm{S}})\big)+\text{const.}$ $\displaystyle=\mathrm{Tr}({\bm{S}}_{a})\,\mathrm{Tr}({\bm{S}}_{b})-\mathrm{Tr}\big(\mathbb{E}[{\bm{g}}{\bm{g}}^{\top}]\big(\mathrm{LogM}({\bm{S}}_{a})\otimes{\bm{I}}_{b}+{\bm{I}}_{a}\otimes\mathrm{LogM}({\bm{S}}_{b})\big)\big)+\text{const.}$ $\displaystyle=\mathrm{Tr}({\bm{S}}_{a})\,\mathrm{Tr}({\bm{S}}_{b})-\mathrm{Tr}\big(\mathbb{E}[{\bm{g}}{\bm{g}}^{\top}]\big(\mathrm{LogM}({\bm{S}}_{a})\otimes{\bm{I}}_{b}\big)+\mathrm{Tr}\big(\mathbb{E}[{\bm{g}}{\bm{g}}^{\top}]\big({\bm{I}}_{a}\otimes\mathrm{LogM}({\bm{S}}_{b})\big)\big)+\text{const.}$ $\displaystyle=\mathrm{Tr}({\bm{S}}_{a})\,\mathrm{Tr}({\bm{S}}_{b})-\mathbb{E}\big[\mathrm{Tr}\big({\bm{g}}{\bm{g}}^{\top}\big(\mathrm{LogM}({\bm{S}}_{a})\otimes{\bm{I}}_{b}\big)\big]-\mathbb{E}\big[\mathrm{Tr}\big({\bm{g}}{\bm{g}}^{\top}\big({\bm{I}}_{a}\otimes\mathrm{LogM}({\bm{S}}_{b})\big)\big)\big]+\text{const.}$ $\displaystyle=\mathrm{Tr}({\bm{S}}_{a})\,\mathrm{Tr}({\bm{S}}_{b})-\mathbb{E}\big[\mathrm{Tr}\big({\bm{G}}^{\top}\mathrm{LogM}({\bm{S}}_{a}){\bm{G}}{\bm{I}}_{b}\big)\big]-\mathbb{E}\big[\mathrm{Tr}\big({\bm{G}}^{\top}{\bm{I}}_{a}{\bm{G}}\mathrm{LogM}({\bm{S}}_{b})\big)\big]+\text{const.}$ $\displaystyle=\mathrm{Tr}({\bm{S}}_{a})\,\mathrm{Tr}({\bm{S}}_{b})-\mathbb{E}\big[\mathrm{Tr}\big({\bm{G}}^{\top}\mathrm{LogM}({\bm{S}}_{a}){\bm{G}}\big)\big]-\mathbb{E}\big[\mathrm{Tr}\big({\bm{G}}^{\top}{\bm{G}}\mathrm{LogM}({\bm{S}}_{b})\big)\big]+\text{const.}$ $\displaystyle=\mathrm{Tr}({\bm{S}}_{a})\,\mathrm{Tr}({\bm{S}}_{b})-\mathbb{E}\big[\mathrm{Tr}\big({\bm{G}}{\bm{G}}^{\top}\mathrm{LogM}({\bm{S}}_{a})\big)\big]-\mathbb{E}\big[\mathrm{Tr}\big({\bm{G}}^{\top}{\bm{G}}\mathrm{LogM}({\bm{S}}_{b})\big)\big]+\text{const.}$ $\displaystyle=\mathrm{Tr}(\mathrm{ExpM}({\bm{P}}_{a}))\,\mathrm{Tr}(\mathrm{ExpM}({\bm{P}}_{b}))-\mathbb{E}\big[\mathrm{Tr}\big({\bm{G}}{\bm{G}}^{\top}{\bm{P}}_{a}\big)\big]-\mathbb{E}\big[\mathrm{Tr}\big({\bm{G}}^{\top}{\bm{G}}{\bm{P}}_{b}\big)\big]+\text{const.}$ (15) $$

where ${\bm{P}}_{k}:=\mathrm{LogM}({\bm{S}}_{k})$ for $k\in\{a,b\}$ and $\mathrm{ExpM}(\cdot)$ is the matrix exponential function.

Notice that the optimal solution should satisfy the gradient stationary condition.
We consider the gradient with respect to ${\bm{P}}_{k}$ because this condition must be satisfied regardless of ${\bm{S}}_{k}$ and ${\bm{P}}_{k}$ for $k\in\{a,b\}$.
The condition for the derivative of [Eq.˜15](https://arxiv.org/html/2509.03378v4#A7.E15) with respect to ${\bm{P}}_{a}$ is

$$ $\displaystyle 0=$ $\displaystyle\;\partial_{{\bm{P}}_{a}}\mathrm{VN}(\mathbb{E}[{\bm{g}}{\bm{g}}^{\top}],{\bm{S}})=\underbrace{\mathrm{ExpM}({\bm{P}}_{a})}_{={\bm{S}}_{a}}\underbrace{\mathrm{Tr}(\mathrm{ExpM}({\bm{P}}_{b}))}_{=\mathrm{Tr}\big({\bm{S}}_{b}\big)}-\mathbb{E}\big[{\bm{G}}{\bm{G}}^{\top}\big]$ $$

where Tsuda et al. (2005) show that $\partial_{{\bm{P}}_{k}}\mathrm{Tr}(\mathrm{ExpM}({\bm{P}}_{k}))=\mathrm{ExpM}({\bm{P}}_{k})$.

Thus, we can see that the optimal solution must satisfy this condition

$$ $\displaystyle{\bm{S}}_{a}^{*}=\frac{\mathbb{E}\big[{\bm{G}}{\bm{G}}^{\top}\big]}{\mathrm{Tr}({\bm{S}}_{b}^{*})}$ $$

Similarly, we can obtain the second condition.

$$ $\displaystyle{\bm{S}}_{b}^{*}=\frac{\mathbb{E}\big[{\bm{G}}{\bm{G}}^{\top}\big]}{\mathrm{Tr}({\bm{S}}_{a}^{*})}$ $$

We can verify that the following solution satisfies these conditions.

$$ $\displaystyle{\bm{S}}_{a}^{*}=\mathbb{E}\big[{\bm{G}}{\bm{G}}^{\top}\big],\quad{\bm{S}}_{b}^{*}=\frac{\mathbb{E}\big[{\bm{G}}^{\top}{\bm{G}}\big]}{\mathrm{Tr}(\mathbb{E}\big[{\bm{G}}{\bm{G}}^{\top}\big])}$ $$

Notice that the optimal ${\bm{S}}_{a}$ and ${\bm{S}}_{b}$ are not unique. However, their Kronecker, which is ${\bm{S}}^{*}={\bm{S}}_{a}^{*}\otimes{\bm{S}}_{b}^{*}$, is unique.
Prior studies (Morwani et al., 2025; Vyas et al., 2025a; Eschenhagen et al., 2025) have shown that this solution is an optimal Kronecker approximation of the flattened gradient second moment under the Frobenius norm.

In the Adafactor case, the result can be similarly derived when considering ${\bm{S}}_{k}$ to be a diagonal matrix for $k\in\{a,b\}$.

$\square$

Figure: Figure 8: *Left:* Simplified VN-Shampoo schemes motivated by [˜7](https://arxiv.org/html/2509.03378v4#Thmclaim7) to incorporate trace scaling. We consider two variants to incorporate trace scaling into the original Shampoo. Variant 1 is inspired by Adafactor’s update scheme, while Variant 2 is similar to KL-Shampoo’s update scheme. Note that Variant 1 of the idealized VN-Shampoo is known as Shampoo with trace scaling in the literature. *Right:* Adapting our exponential moving average (EMA) approach enables VN-Shampoo to use the faster QR procedure and makes it competitive, as empirically shown in [Fig.˜10](https://arxiv.org/html/2509.03378v4#A8.F10) and [Fig.˜10](https://arxiv.org/html/2509.03378v4#A8.F10).

## Appendix H Additional Experiments

We conduct three additional sets of experiments, following the same experimental setup as described in the main text, to further evaluate our approach. Due to limited computational resources, we focus on two language models—NanoGPT (123M) and Llama (134M)—in these additional experiments.

In the first additional experiment, we evaluate the two-sided Shampoo based on Frobenius norm (Morwani et al., 2025; Eschenhagen et al., 2025)—referred to as idealized F-Shampoo—and find that it performs poorly in practice even when we improve its performance using QR and EMA on the eigenvalues, as shown in [Fig.˜7](https://arxiv.org/html/2509.03378v4#A6.F7). This indicates using the Frobenius norm for preconditioner estimation is not ideal.

In the second additional experiment, we evaluate Shampoo with trace scaling (Morwani et al., 2025; Vyas et al., 2025a; Eschenhagen et al., 2025)—referred to as idealized VN-Shampoo—and find that it performs poorly in practice even when using eigendecomposition.
By contrast, incorporating our moving-average scheme enables it to perform well and use the fast QR decomposition, as demonstrated in [Fig.˜10](https://arxiv.org/html/2509.03378v4#A8.F10).

In the third additional experiment, we evaluate the suitability of KL versus VN divergence for refining Shampoo’s estimation rule in a comparable setting, where both variants outperform SOAP while matching SOAP-level pre-iteration runtime. As shown in [Fig.˜10](https://arxiv.org/html/2509.03378v4#A8.F10), KL-Shampoo consistently outperforms VN-Shampoo, even when VN-Shampoo is made practical and competitive using similar techniques to those employed in KL-Shampoo.
These results underscore the advantages of the KL divergence over the VN divergence.

Figure: Figure 9: Empirical results from a random search with 150 runs per method on language models demonstrate that our exponential moving average (EMA) scheme for eigenvalue estimation, as described in [Fig.˜8](https://arxiv.org/html/2509.03378v4#A7.F8.fig1), makes Shampoo with trace scaling—referred to as Variant 1 of idealized VN-Shampoo—practical and enables it to match or exceed the performance of Shampoo with step-size grafting. All these methods perform QR or eigen decompostion at every 10 iterations. Without this scheme, Shampoo with trace scaling performs poorly in practice, as shown in the figure. We implement VN-Shampoo (i.e., Shampoo with trace scaling) ourselves, as it is not available in existing implementations, including the state-of-the-art version from Meta (Shi et al., 2023). As a reference, we also include the best Shampoo run with power $p=\nicefrac{{1}}{{2}}$ and grafting based on the implementation from Meta.
Refer to caption: x6.png

## Peer Reviews

confidence: 2
contribution: 3
initial_rating: 6
presentation: 3
questions: Page 3, Figure 1 Caption: "pre-iteration runtime",  "per-iteration runtime"?


Appendix G: It might be helpful to briefly summarize the conclusions of Fig. 7 and Fig. 8 in the main text, as they provide strong support for the choice of KL divergence over VN divergence.

Could the authors provide theoretical guarantees, showing the advantage of the proposed algorithms?

Could the authors provide more numerical results to compare with other emerging large scale algorithms?
soundness: 3
strengths: - The KL divergence view is a well-motivated approach. The connection to covariance estimation and the justification via proximal-gradient steps are interesting. 

- The proposed methods are not just theoretical constructs; they are practical algorithms that deliver comparable performance over Shampoo and SOAP.

- The paper is generally well-written.
summary: The paper reinterprets the preconditioner estimation in Shampoo and SOAP as a covariance estimation problem solved via Kullback-Leibler (KL) divergence minimization.  This perspective may reveal a previously overlooked limitation in the one-sided estimation approach of Shampoo. Building on this insight, the authors propose new estimation rules, leading to KL-Shampoo and KL-SOAP. Through  experiments on multiple language models, the authors demonstrate that KL-Shampoo and KL-SOAP outperform Shampoo and SOAP.
weaknesses: - The authors only demonstrate that KL-Shampoo and KL-SOAP outperform Shampoo and SOAP. They do not compare with other state-of-the-art large scale algorithms.

- The authors do not provide any theoretical convergence guarantees.

---

confidence: 4
contribution: 4
initial_rating: 8
presentation: 3
questions: 1. Can the authors add comparison to Muon as well - for both runtime and number of iterations?
soundness: 4
strengths: 1. The paper provides a novel perspective of viewing the Kronecker factorization updates based on KL-divergence.

2. It provides a new update rule for the kronecker factor updates based on this perspective, and also provides a computationally efficient way for implementing it.

3. The empirical results of the method are pretty strong, outperforming recent second order methods such as Shampoo and SOAP.
summary: The paper proposes a KL-divergence perspective of looking at the updates provided by Shampoo and SOAP. It provides a new update rule for the kronecker factored preconditioner based on the KL divergence perspective and provides empirical results for the same. Empirically, KL-Shampoo and KL-SOAP outperform Shampoo and SOAP on various language modeling benchmarks.
weaknesses: 1. The paper is missing comparisons to Muon, another popular second order optimizer proposed recently.

---

confidence: 3
contribution: 3
initial_rating: 6
presentation: 3
questions: Please address my concerns above
soundness: 4
strengths: The paper tackles a well motivated problem and offers a clear explanation. The literature typically justifies SOAP with Frobenius norm approximations [1]. The authors replace that with a covariance estimation view that minimizes KL divergence between zero mean Gaussians. Under this view, the usual Kronecker estimators in Shampoo and SOAP do not solve the right objective when both factors are learned jointly, which motivates a two sided estimator. This reveals a limitation in the standard Kronecker update and leads to principled fixes, presented as KL Shampoo and KL SOAP. The result is better performance at roughly SOAP level iteration cost.


[1] Morwani, Depen, et al. "A New Perspective on Shampoo's Preconditioner." arXiv preprint arXiv:2406.17748 (2024).
summary: The authors recast diagonal estimation as covariance estimation via KL minimization: they treat the true second moment $E[gg^\top]$ as a Gaussian covariance and seek the vector $d$ that minimizes the KL divergence between $\mathcal{N}(0, E[gg^\top])$ and $\mathcal{N}(0, Q,\mathrm{Diag}(d),Q^\top)$. They prove that the KL-optimal $d$ equals the expected squared gradient in the eigenbasis, which provides a clean justification for SOAP’s “RMSProp in an eigenbasis” update as the solution to a KL problem. From this view, the classic Kronecker estimators in Shampoo and SOAP are suboptimal for the KL objective, which motivates revised estimators. They propose KL-SOAP, which keeps SOAP’s basic pattern but replaces Shampoo’s factor estimation with KL-Shampoo’s estimation to obtain a better eigenbasis $Q$.
weaknesses: The paper argues that the Frobenius-norm view misses the right geometry and proposes a KL objective instead. Since methods like SOAP ultimately aim to approximate a preconditioner, it is natural to assess estimation quality by a distance such as the Frobenius norm, and it should not be surprising that different distance functions produce different optimal approximations. I appreciate the point that reconstruction error is not the final goal, yet the justification still feels incomplete. The paper does not fully persuade me that the KL objective, rather than Frobenius or another principled choice, is the uniquely appropriate metric for evaluating structured preconditioners in this setting.



The paper presents KL-SOAP as a principled improvement to SOAP, but the mechanical change is small. In practice the algorithmic delta is limited to swapping SOAP’s Shampoo-based eigenbasis for a KL-Shampoo eigenbasis and adding an EMA rule for factor eigenvalues under basis staleness. The step form, the augmented diagonal, and the memory profile remain essentially the same. The narrative should reflect that KL-SOAP is best understood as SOAP with a stronger basis construction and more careful eigenvalue tracking, not as a distinct optimizer.

The experiments never report approximation error under the proposed KL objective.

---

confidence: 4
contribution: 3
initial_rating: 6
presentation: 3
questions: 1. In figure 3, why do the delta definitions have two lines each?
2. In figure 3 bottom right, where is the modification that defines “augmented”?
3. Are the output layers and/or layernorms treated differently in the empirical implementation (as is often the case with similar methods)?
soundness: 3
strengths: 1. The paper provides a new perspective on pre-conditioning for LLM optimization that yields principled algorithms.
2. The presentation of related work and how the KL methods fit in is very nice.
3. The method seems to work better than baselines in the experiments. And hyper parameter tuning seems relatively fair across baselines and runtime/memory is considered.
summary: This paper proposes KL-Shampoo and KL-SOAP, which modify the Shampoo and Soap algorithms to try to find the best approximation to the gradient second moment in the KL norm rather than the Frobenius norm. In practice this yields an update rule where current estimates of the second moment of the RHS covariance are used to update the LHS covariance and visa-versa. The paper also proposes to use QR updates, but combined with frequent estimates of eigenvalues via EMA. Empirical results show gains over the baselines.
weaknesses: 1. It’s a lot to cover, but the presentation of the algorithm itself could be improved. It is too nonlinear with options here and there. Some clear, linear algorithm boxes for each separate algo in the appendix could provide clarity. 
2. I understand there are computational constraints, but the experiments are small-scale, so it is not clear whether the gains with scale. Perhaps even more than scaling up model size (which can get expensive), I would be interested to see substantially longer runs at the same size. The current experiments are effectively *very* early in training (less than chinchilla and far below the >>100 token-to-param rations that are more standard now), so it is not clear whether the gains matter later on as well.


Decision: Accept (Poster)