Understanding and improving Shampoo and SOAP via Kullback-Leibler Minimization
Authors: Wu Lin
Scott C. Lowe
Felix Dangel
Runa Eschenhagen
Zikun Xu
Roger Baker Grosse
Venue: ICLR 2026 Poster (2026)
Keywords: Shampoo
SOAP
covariance estimation
Kullback–Leibler divergence
Gaussian
optimization
Primary area: optimization
TL;DR: A new Kullback–Leibler perspective to interpret and improve Shampoo and SOAP

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 \textbf{KL-Shampoo} and \textbf{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 promising foundation for designing structured methods in NN optimization.

Published as a conference paper at ICLR 2026
UNDERSTANDING AND IMPROVING SHAMPOO AND SOAP VIA KULLBACK-LEIBLER MINIMIZATION
Wu Lin', Scott C. Lowe!, Felix Dangel', Runa Eschenhagen", Zikun Xu', and Roger B. Grosse 1.4
Vector Institute, Toronto, Canada University of Cambridge, Cambridge, United Kingdom Miirosoff, United States
University of Toronto, Toronto, Canada
ABSTRACT
Shampoo and its efficient variant, SOAP, employ structured second-moment esti- mations 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 '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-Leibe (KL) divergence minimization, revealing previously over- looked 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- per-iteration runtime. Notably, KL-Shampoo does not rely Adam attain competitive performance, eliminating the memory overhead introduced by Adam. Across our experiments, KL-Shampoo consistently outper- forms SOAP, Shampoo, and even KL-SOAP, establishing the KL-based approach as a promising foundation for designing structured methods in NN optimization.
INTRODUCTION
Optimizer Shampoo (Gupta al., 2018) has recently rivaled and, in several benchmarks, surpassed Adam (Kingma & Ba, 2015) in training a wide range of neural network (NN) models (Dahl et al., 2023; Kasimbeg et al., 2025). Consequently, Shampoo and its efficient variant, SOAP (Vyas et al., 2025a), have drawn increasing attention. In practice, however, Shampoo typically requires step-size grafting with Adam to achieve competitive performance (Agarwal et al., 2020; Shi et al., 2023), SOAP mitigates this by applying Adam in Shampoo's eigenbasis and further reducing per- per-iteration runtime. Unfortunately, this 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-we which approximate the flattened gradient second moment (Duchi et al., 201 1)- 1)-priarrlly through the Frobenius norm. However, few studies have examined these schemes from the perspective of Kullback-Leibh (KL) divergence. Compared to the Frobenius norm, the KL divergence between aero-mean Gaussian covariance matrices is more appropriate for interpreting Shampoo's and SOAP's preconditioners as Gaussian covariance estimators, since the second moment they approximate can be viewed the covariance matrix of zero-mean Gaussian. Historically, the Frobenius norm (Greenstadt, 1968; Dennis & Moré, 1977) has been replaced by the KL divergence (Fletcher, 1991) or its second-order truncation (Nocedal & Wright, 2006) as the foundation for designing preconditioner estimation schemes in quasi- -easi- methods. The KL perspective has provided a unified framework for interpreting (Güler et al., 2009; Waldrip & Niven, 2016) and extending (Kanamori & Ohara, 201 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 (SPD) constraint (Amari, 2016;
"Corresponding author; yorker lin @ gmail.com; Code: gi thub.hu.b.. www.b.ccm/yyrkkeeee i Kl.-Mehhdss

---

Published as a conference paper at ICLR 2026
NanoCiP-T- watp-------------- PiM--iineWhh11
NanoFWWWYT 1-7-------P----- P.AM--PaWWWh
1.lma-134M- 13MM-C7-22
KanoMoT: 227-- 307M-OppeWWhaa
3.450
3.450
3.45
3.30
3:425
3.400
3.40
3.30
190 1X88
3.375
8.35
3.350
3.325
3.30
3,300
3.300
3.276
3:275
3. 0.35
3,250
3:20
250
3.12
50
100
130
50
100
150
0
50
100
100
0
100
200
2000
Minutes
Minutes
Minutes
Minutes
Shampoo
KI-SOAP
SOAP
Shampon (p- 3. grafting)
Shampoo (n- grahting)
Figure 1: Empirical results (random search with 150 runs per method) on language models using 6 1 demonstrate the advantages of KL-based methods over Shampoo and SOAP, while matching SOAP's per-iteration runtime. All methods take the same number of iterations in these experiments. Surprisingly, KL-Shampoo outperforms KL-SOAP. We include the best Shampoo run based on thet-- implementation from Meta (Shi et al., 2023). See Fig. 9 in Appx. H for evaluating KL-Shampoo larger model (Llama3 with 450M).
Minh & Murino, 2017) that preconditioners in Shampoo and SOAP must satisfy as 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 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, thereby connecting structured adaptive methods with both classical quasi-Newton ideas and structural Gaussian whitening within a single framework. Our framework naturally extends to tensor- valued settings, where some existing theoretical interpretations based on singular value decomposition (SVD) or the spectral norm may not apply. This perspective reveals key limitation (illustrated in Fig. 5): the Kronecker- structured estimators used by Shampoo and SOAP do not adequately solve the corresponding KL- minimization problem. This limitation, turn, opens new opportunities for improvement. Leveraging this insight, refine the estimation rules of Shampoo and SOAP and develop practical KL-based schemes- KL-SShmpoo and KL-SOAP that match or exceed the performance 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 (Table 1). We then generalize our framework as a divergence-based estimation approach and consider other Shampoo variants, including a trace-scaling variant- matrix version of Adafactor. In addition, the practical techniques we develop for KL-Shampoo (Sec. 4) can be adapted to strengthen Shampoo variants and make the trace-scaling variant competitive without step-size grafting (Fig. 10, Appx. H), while achieving SOAP-level per-iteration runtime (Fig. 11, Appx. H). Empirically, we show that KL-based methods are competitive for training range of NNs and remain as flexible as Shampoo and SOAP for tensor- valued weights. Surprisingly, KL-Shampoo consistently outperforms other Shampoo variants in experiments (Figs. 1, and 9). Overall, approach provides a principled and unifying way to design structured preconditioned methods for NN optimization (see Sec. 3.3).
BACKGROUND
Notation For presentation simplicity, we focus on matrix-valued weights and the optimization update for a single parameter matrix O E Rda rather than set of weight matrices for NN training. We use Mat(-) to unflatten its input vector into a matrix and vec(-) flatten its input matrix into vector. For example, 0 = vec( )) is the flattened weight vector and θ Mat(O) is the unflattened weight matrix. Vector g is (flattened) gradient vector for the weight matrix. We denote y, (8 and S to be step size, weight for moving average, and preconditioning matrix for adaptive method, respectively. Diag(-) returns a diagonal matrix whose diagonal entries are given by its input vector, whilst diag(.) extracts the diagonal entries of its input matrix as vector.
2

---

Published as a conference paper at ICLR 2026
Shampoo Given matrix gradient G and the flattened gradient g vec(G), the original Shampoo method (Gupta et al., 2018) considers a bronecker-factoree approximation, (S.) @ (S.) of the flattened gradient second moment, Bglgg J, where p denotes matrix power, S. g [GG
S, = Eg G G). and & denotes Kronecker product. In practice, we often approximate the expectation, Eglgg"1, with exponentially moving average (EMA) on the outer product (Morwani et al., 2025). The original Shampoo method uses the 1/4 power (i.e., = 1/a) and other works (Anil et al., 2020; Shi al., 2023; Morwani et al., 2025) suggest using the 1/2 power (i.e., p 1/2). At each iteration, Shampoo follows this update rule with EMA on S. and Sg:
S. + (1 - 82) + B2GG", S. + 1 (1-111100 -2) + B2G (Kronecker 2n moment est.),
0 - - + + t O S (Preconditioinn);
(1)
where S = 2 & Sha is Shampoo's 's preconditioning matrix, and we leverage the Kronecker structure of 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 cigendccomposition of Sk, such Qx Qx eigen(Sk)) for k E {a, b). every few steps and storing Qs and A (Anil et al., 2020; Shi et al., 2023). Therefore, the power of Sk is computed using an elementwise power in λ such as S " Q Diag X P where .Pp denotes elementwise p-th power. This compu- tation becomes an approximation if the decomposition is not performed at every step.
Using Adam for Shampoo's 's stabilization increases memory usage. If the cigendecomposition 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. Unfortunately, this increases its memory usage introduced by Adam (see Table I).
SOAP SOAP improves Shampoo with the p = 1/2 power by running Adam in the eigenbasis of Shampoo's 's preconditioner 2 @ (S, S. @ Sg. Notably, SOAP reuses Shampoo's Kronecker estimation rule for computing the eigenbasis Q == Qa @ Qs and incorporates Adam's 's 2nd moment, denoted by d, for preconditioning. where Qs is Shampoo's 's Kronecker eigenbasis Sş for k E {a. b} defined above. As a result, SOAP effectively employs an augmented preconditioner, S Q which cannot be expressed as Kronecker product of any matrices with the same shape as S. and S.. Because we omit momentum (i.e. let Adam's 31 = 0), SOAP takes the following step with the Adam update becoming an RMSProp update (Tieleman & Hinton, 2012):
S. + (1- 12) + 82GGT, S, + -1) + 32 G
(Shampoo's 2dd moment est.),
d (1 - (1-3)) + 329
(RMSProp's diagonal 2nd moment est. in the eigenhasis),
θ - θ - 70% Mat و
0 0- -0 7 ig
(Preconditioning),
where و Q veo( Q GQO) is "projected" gradient vector in eigenbasis Q and recall is SOAP's preconditioner. Here, we leverage the Kronecker structure and
that S = QDia@Diaa(@)) orthogonality of the eigenbasis to move from the left to the right in the last line of Eq. (2). Note that this EMA weight Sg is defined as S g(Adam) where @ffddmi is Adam's (RMSProp's 's 32. We use this definition rather than Adam's '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 is approximated via QR decomposition, whereas Shampoo typically employs eigendecom- position to compute both the eigenbasis and the cigenvalues. Vyas et al. (2025a) therefore suggest replacing the slower eigendecomposstion with the faster QR decomposition, such as Qx + (r kx) for & E {a, b). This makes SOAP more computationally efficient than Shampoo.
Running Adam in the eigenbasis increases memory usage. Introducing Adam's (RMSProp's) s) 2nd
moment estimation increases SOAP's memory consumption (see Table 1). This is because this estimation, d Rd uses extra memory and cannot be exactly expressed Kronecker product of any two vectors with compatible dimensions, such as d o r, @ m. where r. E Rd,x1 and ry E Rdsx1
3

---

Published as a conference paper at ICLR 2026
The original Shampoo's 's Kronecker estimation rule (p 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 = 1/2) (Morwani et al., 2025; Eschenhagen et al., 2025) focus bounds induced by the Frobenius norm. SOAP reuses Shampoo's s Kronecker estimation rule and additionally introduces Adam's 2nd -ommet estimation in the eigenbasis (Vyas et al., 2025a). None of these works interpret or motivate their estimations as covariance estimation, thereby missing the opportunity to introduce the KL perspective.
3 SECOND MOMENT ESTIMATION VIA KULLBACK-LEIBLAN MINIMIZATION
We focus on Shampoo with p 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 previously unrecognized limitation of Shampoo's 's estimation rule: the boonecker-structuree 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, propose a KL-bbased scheme for Shampoo, which we term the idealized KI.-Shampoo.
KL Minimization For simplicity, we begin by introducing a KL perspective a matrix- case and drop subscripts when referring to the flattened gradient 2nd moment, like Eigg" "1 := Pglgg" where g = vee( G) is flattened gradient vector of a matrix-valued gradient G E R xds The goal estimate a Kronecker-structured preconditioning matrix, S S. X Sb, that closely approximates the 2nd moment, where S. R x da and S₁ Ra are both symmetric positive-definil (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,
KL Perspective for Covariance Estimation KL(E S) DKL (0, kI) (O)
log det ((ggg (1 + const,
where Elgg" ] and S are considered as Gaussian's 's covariance, det(.) denotes the matrix determinant of its input, and x² 0 is damping ensure the desitive-definiteness of Esgg" 1l I if necessary.
+I
Natural extension to tensor-valued weights Our framework naturally extends to tensor cases, G E GR Rddadd xde, by considering a preconditioner S S₂ @ S, @ Se to approximate the flattened second moment, E[gg" ]. where matrix S, E R x is SPD for k E {a, b, c).
KL minimization as a divergence-
projection over SPD matrices. The KL divergence between zero-mean Gaussians coincides (up a constant factor) with the log-determinant divergence widely used in matrix optimization over SPD matrices (Dhillon & Tropp, 2008; Kulis et al., 2009; Sra, 2016), which does not require a zero-mean assumption when viewed purely as a matrix divergence. Thus, we can interpret our KL- AL-minimization problems as projecting target SPD matrix onto the Kronecker-structure This viewpoint naturally extends to settings where the target matrix is not second moment, such as the curvature matrices used in quasi-Newton methods (Fletcher, 1991; Waldrip & Niven, 2016), thereby reconnecting structured adaptive methods with quasi-Newton ideas.
SPD manifold (see Table 2).
3.1 INTERPRETING SHAMPOO'S '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 B2 1) for our descriptions and focus on Shampoo with power p- = 1/2, presenting a KL minimization perspective and interpreting its estimation rule from this perspective. We will obtain Shampoo's 's estimation rule by solving a KL minimization problem.
Claim 1. (Shampoo's 's Kronecker-based covariance estimation) The optimal solution of KL minimiza- tion mins, (T S) with one-sided preconditioner S (1),,S.) 39 I, is S EE where dy is the dimension of matrix S E dR xds for k E ba and G = Mat(g).
Likewise, we can obtain the estimation rule for So by considering S I, St).
4

---

Published as a conference paper at ICLR 2026
NanoGi 2-A.--------M-WWWeMM
W
1Amx-136M8-C---1
NanoMedi- Ope 21
10
4.0
3.5
3.33
3.4
3:20
3.30
3.3
3.3
3.15
3:25
1
10
10
Berrsions (thousands)
Berations (shonsands)
Heraation (thonords)
Iterations (thouseeds))
KL-Shamppo
Shampoc (p:s
Shampoo (p 2. gralting)
Shampon (p= 4))
Siaampon (p 4. grafting
Figure 2: Empirical results (random search with 150 runs per method) on language models using 6 6 show that KL-Shampoo (with QR decomposition) does not rely on step-size grafting with Adam to perform well. Shampoo without grafting does not perform well, even when using state-of-the-ard implementation (Shi et al., 2023). In particular, Shampoo with power p = 1/2 fails to train the RWKV7 model in all 150 runs when grafting is disabled.
Shampoo
SOAP
Shampoo
KL-SOOP
Kronecker factors (Sa, Sq)
da + df da + dh
da + df da + dB N/A
d² d
da + dp. d² + di
Kronecker factors' eigenbasis (Qu, Q6)
d dh
Kronecker factors' eigenvalues (Au, Ao)
da + db
da + do
do + db
Adam's 2dd moment in the eigenbasis (d)
N/A
da
N/A
d.
(interpreted as augmented eigenvalues, Sec. 5)
Momentum
dadb
da
dg
dadb
Step-size grafting with Adam
da
N/A
N/A
N/A
Table I: Memory usage of each method considered this work. We store and update 'lll 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, Q QDiag(d)Q can not be expressed as a Kronecker product due the augmented eigenvalues d, while Shampoo's and KL-Shampoo's preconditioners, QDaag X & A can, where Q == Qa & Qi.
Shampoo's ' estimation rule as Kronecker- (proof in Appx. A), Shampoo's estimation rule with power p =1/2 in Eq. (1) can be viewed as the optimal solution of a KI. minimization problem (up to a constant scalar) when one Kronecker factor is updated independently and the other fixed as the identity, which is known 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).
covariance estimation According to Claim 1
3.2 IMPROVING SHAMPOO'S S ESTIMATION: IDEALIZED KL-SHAMPOO
Our KL perspective reveals a key limitation-empireett demonstrated in Fig. 5-of Shampoo's Kronecker estimation with = 1/2 as 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 KI-Shhapoo'' covariance estimation for S. and S₁)) The optimal solution of KL minimization mins,,s, KL (E(gg" ] S) with two-sided precontioner S = -. & So satisfies
1 (S:) -G' d
1 E (S:) G). d
Stationarity Condition: S. =
(4)
'Sinee the current QR/eigen implementation in PyTorch does not support half precision (bfloat16), we perform QR/eigen to compute Qs for & E {a. ) in single precision (float32) and then cast and store them in half precision (bfloat16) 16). Other matrices and vectors can be computed and updated in half precision.
5

---

Published as a conference paper at ICLR 2026
Idealized KL-Shampoo's estimation Claim 2 (proof in Appx. B) establishes closed-form condition (see Eq. (4)) when simultancously learning both Kronecker factors to minimize the KL problem. In machine learning, Lin et al. (2019; 2024) treated the condition of multilinear exponential-famin (see Sec. 5 of Lin et al. (2019)) for Kronecker- optimization via natural gradient descent matrix Gaussian, while more recently, Vyas et al. (2025b) considered a similar condition motivated heuristically by gradient whitening. However, we cannot directly use this condition due the correlated update between St and S For example, solving St requires knowing S in Eq. (4) or vice versa. In practice, this condition is unachievable because the expectations in Eq. (4) must be approximated. Thus, we consider an estimated S to approximate S for k E {a, bj and propose an exponential moving average (EMA) scheme:
theoretical example
82 G S. (1 -1-
32 G
S. (1 -1-1-18 12) + db
(5)
da
Divergence
Preconditioner Structure
Estimation Scheme
KL-Shampoo
Kullback-Leibler von Neumann
Kronecker factors
MLE matrix diag. MME
Adafactor
diag. Kronecker factors
Matrix Adafactor
von Neumann
Kronecker factors
matrix MME
(Shampoo with trace scaling)
Table 2: Important divergences, structured SPD matrices, and equívalent Gaussian estimations in matrix cases, which can be extended tensor cases. Under the additional zero-mean assumption, minimizing KL divergence (for SPD matrices) is equivalent to maxximum-likelibee estimation (MLE) (for zero-mean Gaussian distributions), whereas minimizing VN divergence amounts to (normalized) matrix second moment matching estimation (MME) (see Sec. 3.3).
KL-Shampoo as an MLE scheme for zero-mean Gaussian whitening Statistically, the con- dition in Eq. (4) corresponds to the maximun-likelihod estimation (MLE) of a zero-mean ma- trix Gaussian (Dutilleul, 1999) when E[ggT] is considered as finite average 1
N E 9:9.
This because MLE is equivalent to minimizing the KI. divergence: KT, E 9.9, S)
sample generated from /O S). Thus, satisfying Eq. (4) for S = S. &S, implies that gradient matrix generated from a mean matrix Gaussian, G MA.( S St)), with row-wise covariance S and solumn-wise covariance S obtained by maximum likelihood. Under the condition, Shampoo-styll preconditioning naturally
- Ei log ((g: 0, S) + const, where gi is considered
(S:) (0t)-1/2G(S) G( -1/2 ~ _M Ia, Ip).
induces matrix-Gaussinn (row and column) whitening:
This also implies that the SOAP-like projection (i.c., gradient in the cigenbasis) amounts to covariance diagonalization under the optimal eigenbasis, as we will discuss in Sec. 5. Our KL-bbase approach further extends this tensor- tensor-Gaussian whitening for tensor-valued cost typically associated with SVD-based methods. Notably, this kind of whitening arises from minimizing KL divergence rather than the Frobenius norm.
coaddennsswwthhen the prohibitive
EMA scheme as a stochastic proximal gradient step for the KL minimization Our framework allows us to further justify our EMA scheme in Eq. (5) as stochastic proximal-gradient step (see Claim 3 and a proof in Appx. C) and establish a formal connection to the theoretical example of Lin et (2019; 2024). Notably, our approach uses S-1/2 1/2 for preconditioning (Eq. (1)), following Shampoo, whereas Lin et al. (2019; 2024) propose using S-1 A straightforward implementation of our scheme is computationally expensive, since it requires expensive matrix inversions (highlighted in red in Eq. (5)) and the slow eigendecomposition for Shampoo- type preconditioning (e.g.. S 1/2). However, these issues can be alleviated-in Sec. we propose an efficient implementation.
Claim 3. (KL-Shampoo's 's moving average scheme) The moving average scheme for Sy. (Eq. (5)) in idealized KL-Shampoo is a stochastic proximal- gradiett step with step-size 32 solve the KL. minimization problem in Eq. (3), for k E {a, b}. Recall that B2 in Eq. (5) is closely related to Adam's
B2 as 82 1-130 (fAddm)) where 3/ 2f.damm) is Adam's B2.

---

Published as a conference paper at ICLR 2026
3.3 COMPARISON WITH KRONECKER SCHEMES USING ALTERNATIVE DIVERGENCES
Unifying framework via divergence-basee projection Our perspective leads to a conceptual framework that allows us Gaussian assumption because we can view the minimization problem as projection problem over SPD matrices rather than Gaussian distributions. Moreover, the Gaussian assumption may not be satisfied when using other divergences.
the following divergences. Notably, we do not require the zero-mean
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 SVD of the second moment Elgg ] (Van Loan & Pitsianis, 1993). However, this solution is often unattainable in practice and is computationally expensive for two reasons: (1) the expectation Eggg" " must be approximated; and (2) performing the SVD is costly-yielding complexity (O(22 () in general even for rank- SVD-which is often higher than the eigen decompositions with complexity (O(d)) for k {a. 6} that we aim to avoid. In tensor cases, this optimal result based on SVD no longer holds. Instead, we analyze the stationarity conditions (Claim 6, Appx. F) and derive a new variant, idealized F-Shampoo (Fig. 6, Appx. F), that is structurally similar to KL-Shampoo. While a straightforward implementation of F-Shampoo performs poorly in practice, the techniques (Sec. 4) we develop for KL--Shapooo can be adapted to improve its performance (Fig. 7, Appx. H).
Von Neumann (VN) divergence (VN-Shampoo as matrix Adafactor) 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 trace-scaled Shampoo is equivalent to running Adafactor (Shazeer & Stern, 2018) in Shampoo's '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: trace-scaled Shampoo is exactly a matrix generalization of Adafactor obtained by minimizing the VN divergence (T'suda 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 Claim 7 (Appx. G). Our generalization of Adafactor is also applicable in tensor cases. By contrast, KL-Shampoo minimizes the KL divergence instead of the VN divergence. A straightforward implementation of Shampoo with trace scaling-referred poorly in practice. However, the techniques we develop for
referred to
as idealized KI--Shappo in See. 4 can be adapted (Fig. 8, Appx. G) to substantially improve its performance so that it matches Shampoo with step-size grafting and outperforms SOAP, as shown in Figs. 10 and 11 (Appx. H). When considering the additional Gaussian assumption, KL-Shampoo is derived from the maximum likelihood principle (Sec. 3.2) while VN-Shampoo can be interpreted as matrix S:
moment matching (Table 2). This is because row-wise matrix 2nd moment E[GG TH(S))
where we make use of the Gaussian assumption: G ~ MN((0; S St). Similarly, use column-wise matrix 2dd moment G G) to obtain the scheme for Sg: S Tr(S;) literature does not consider.
induces VN-Shampoo's scheme for S. (Eq. (14), Appx. G): S
E|GG T>(S))*
B G] This is new interpretation that existing
natural question then arises: which divergence 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) and in other applications (Kulis et al., 2009).
4 EFFICIENT IMPLEMENTATION
SLLSSOMPOO 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 eigende- composition. Vyas et al. (2025a) demonstrated that the cigendecomposition used in Shampoo's implementation (Shi et al., 2023) is more computationally expensive than QR decomposition. Moti- vated by this result, we aim to improve KI--Shapoo''s 's computational efficiency by replacing the
7

---

Published as a conference paper at ICLR 2026
cigendecomposition 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 significant overhead. Specifically, the eigenvalues are essential for a reduction in the computational cost of KI-Shampoo in two reasons: (1) they remove the need to compute the matrix -1/2 power, S- (Q. -2:-
(Q,Diag(A, 1/) )) ), used for KL-Shampoo's 's preconditioning; (2) they climinate expensive matrix inversions in its Kronecker estimation rule (Eq. (5)), such as S P. = QoDiag (Q in the update for Sa:
B2 B2 GS =(1 - -2) + de
B2 G.P.G, db
Sa - (1 S2)
(6)
where Qx and X are eigenbasis and eigenvalues of Sy. for k E {a, b). respectively.
KI.-based estimation rule for the eigenvalues X, and A using an outdated eigenbasis 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-none norm perspective, using an instantaneous scheme: 2 (inss) := diag (Q S₂Qu) ) for k E {a, 6). However, our empirical results (Fig. 4) indicate that this approach becomes suboptimal when an outdated eigenbasis Qs is reused to reduce the frequency and cost of QR decompositions. In contrast, our KI. perspective (see Claim 4 and its proof in Appx. D) provides 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) for eigenvalue estimation, which can be justified as a stochastic proximal-gradient step under our KL perspective, similar to Claim 3. This scheme updates the eigenvalues at every iteration while updating the eigenbasis less frequently through an efficient QR- 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. Since naturally scales the eigenvalues by the dimensions of the Kronecker factors, step-size grafting should not be necessary for KI.-Shampoo, as argued by Eschenhagen et al. (2025) and confirmed by our empirical results (Fig. 2). Furthermore, applying this scheme enables other Shampoo variants to be competitive and even outperform SOAP, as empirically demonstrated in Figs. 7, 10 and 11 of Appx. H. These results underscore the importance of our EMA eigenvalue scheme for Shampoo- based methods.
Claim 4. (Covariance estimation for eigenvalues Na and Xx) The optimal solution of KL minimiza- tion minx,.A, KL (B(gg' S) with preconditioner S Diag Xx, Q @ Qia Xx) satisfies this condition:
1 - @ B[G P.GIQ.). \E
diag Q J(.))
dg
where Pt Dig ((2))=1)) Q% is also defined in Eq. (6) and considered as approximation of ST for k {a, 63 when using an outdated eigenbasis = Qa @ Qu precomputed by QR.
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)) applies Shampoo's 's scheme estimate its Kronecker factors and then performs RMSProp (Adam) updates in the eigenbasis of these factors. Consequently, the interpretation of SOAP's Kronecker factor estimation is identical to that of Shampoo. RMSProp's 's second-moment moment estimation in the eigenbasis can itself be interpreted as the optimal solution to separate KL divergence minimization problem, as from the Frobenius-norm viewpoint
established Claim 5 in Appx. E. The KL perspective-disteeee (George et al., 2018; Eschenhagen et al., 2025) provides new lens for understanding RMSProp's estimation in the eigenbasis as augmented eigenvalue estimation for a covariance matrix.
Viewing KL-Shampoo and SOAP's eigenvalue estimations as corrections for outdated eigenbasis When an outdated eigenbasis is used, RMSProp's 's scheme (Step 4a of Fig. 3) for eigenvalue estimation analogous
can be viewed correction in an augmented (full-diagonal) eigen space, QDiag(d) in spirit to the Frobenius-norm interpretation (Eschenhagen et al., 2025) but derived under the KL framework. This perspective also highlights& close similarity KL-Shampoo's estimation scheme: recall that we introduced a comparable correction (Step 3a of Fig. 3) for KI-Shampoo, but in the original Kronecker-factored diagonal eigen space. QDiag Ao @ \²)
8

---

Published as a conference paper at ICLR 2026
(Original) Shampoo with power p=1/2
Replacing the slow eigendeecompositonn with more efficient QR updates (replace Step 3)
versus Our idealized KL-Shampoo
1: Gradient Computation == Vε(0)
3a: Eigenvalue Est. with EMA (each iter)
G = Mat(g) R Xd
Stronecker-fuctored Eigenvalue X_ @ As)
2: Covariance Estimation (each iter)
Ao A, (1-33)
Aa
+ 32 diag Q A. Q.:)
Sa
Sa
Aa
diag Q2 A Q6)
Sbb
- (1 - 82)
Sb
+
32
A
3b: Infrequent Eigenbasis Estimation using QR (every T > ittrs)
GG
(Original)
A.
Q: f- (x for & E {a.b)
GQ.Ding(A, -1--1) G' /db
(KL)
G
(Original)
SOAP (ues Shampoo's eigenbasis) versus
A =
G Q. Q Ding(A.- QQ.
(KI.)
Our KL-SOAP (uses KI.-Shampoo's eigenbasis): Using augmented preconditioners and QR up-
3: Figendecomppsitiion (every T 221 iters)
Ak. Q+ <- eig(S.) for k E b)
dates (replace Step 4)
4: Preconditioning using Q = Q & Q
Augmented Eigenvalue Est. with EMA (each iter) (interpret RMSProp's is 2 Moment d as Augmented Eigenvalue)
0 + 0 xQQ Diag(A A2 & x,)
Figure 3: Left: Simplified Shampoo-based schemes without momentum, damping, and weight decay. See Fig. 5 for an empiri- cal comparison and Fig. 12 for the prac- tical KL-Shampoo. Top Right: For com- putational 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 R using outdated eigenbasis Q₁ for k E {a, b), and use the QR procedure to estimate Qk. Bor- tom Right: Simplified SOAP-basee schemes without momentum. Notably, KL-SOAP needs estimation for A in Step 3a to com- pute the eigenbasis Q, whereas SOAP does not. Here, we view RMSProp's 2nd moment in the eigenbasis as augmented eigenvalues highlighted in blue.
d (+-31332e 9 = Q= g = = Q. GQ²)
3-g
4b: Preconditioning using Augmented Figenvalues
with Eigenbasis Q == Q ② Q
θ + 0 - (QDiag (d) -1/3 ')g
Equivalent to running RMSProp in Eigenbasis:
Mat(O) + Mat(O) QQ Mat
9
Q
RMSProp
Improving SOAP's estimation Similar to SOAP, we propose KL-SOAP, which utilizes KL- Shampoo's estimation to update Kronecker factors and additionally employs RMSProp (Adam) in KI.-Shampoo's 's eigenbasis. Our unified KI. perspective enables us to reuse Claim to justify the use of RMSProp's (Adam's) 's) 2nd moment estimation as augmented eigenvalue estimation in KL-SOAP. Notably, when using KL-Shampoo's eigenbasis Q = Q% @ Q obtained from the optimal condition in Eq. (4), we can see the (Gaussian) covariance of gradient in the eigenbasis Kronecker-diagonalized rather than fully diagonalized: ğ = vec( GQ;) = (Q*) g (O, Diag (>;) )))())))) & Diag Diag(>;)).
EXPERIMENTAL SETUP AND EMPIRICAL EVALUATIONS
We consider four sets of experiments to demonstrate the benefits of using the KI. divergence and its effectiveness. See Appx. H for the experimental setup and additional experiments.
In the first set of experiments, we demonstrate that our KL-based perspective enables principled redesign of Shampoo, resulting in KL-Shampoo, and achieves superior performance without step-size grafting. We evaluate Shampoo with matrix powers p 1/2 and 1/4, using tftte-of-the-art implementation (Shi et al., 2023). As shown in Fig. 2, 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.
9

---

Published as a conference paper at ICLR 2026
NanoGi B
NacolPP-- FiaeWisann
Liana 134M
5.0
0.0
6.0
4.5
4.5
3.5
4.0
5.0
4,0
3.5,
0.0
4.5
3.5
4.0
3.5
3.5
5.24
3.18
10
5
10
10
Itecatioons
(thoussadds)
Iserations (thousands)
Iterations (showandss
Iterations (thousands
KL. Sl-nlagga
RC (letaantareess)
Slauppo
Shainpon
VNN Shampoon
Figure 4: Empirical results (random search using 150 runs for each method) demonstrate that our EMA scheme for the eigenvalue estima- tion makes KL-Shampoo competitive when us- ing an outdated eigenbasis. Without this scheme, KL-Shampoo performs poorly under an out- dated eigenbasis Qs even when employing the instantaneous eigenvalue estimation X( diag(Q, (Q, Sx)u) at every iteration, as suggested by Eschenhagen et al. (2025) for k E {a, b). Adapting the EMA scheme also makes other variants of Shampoo competitive (Figs. 7 and 10, Appx. H) and allows the trace-scaling variant to outperform SOAP (Fig. 11, Appx. H).
Figure 5: Empirical saullss-based on random search with 150 runs per method- demonstrate the advantages of KL-Shampoo (two-sided) over other Shampoo variants un- der comparable settings, including Shampoo with
1/2 (one-sided, no grafting), F- Shampoo (two-sided, Frobenius- barses,dd)),,, and VN-Shampoo (trace scaling, two-sided von- Neumann-divergence-beee
We make these variants practical by incorporating QR step and an EMA scheme for eigenvalue estimation (Fig. 3). To ensure fair comparison and mini- mize implementation bias, we implement Sham- poo, F-Shampoo, and VN-Shampoo ourselves, aligning them closely with KL-Shampoo. See Fig. 11 (Appx. H) for detailed comparison be- tween KL-Shaapoo and VN-Shampoo.
In the second set of experiments, we demonstrate that our QR- scheme enables KL-Shampoo and KL-SOAP to achieve the same pre-iteration runtime as SOAP. We use the official SOAP im- plementation for comparison. As shown in Fig. 1, 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, underscore the importance of using EMA scheme for the eigenvalue estimation when working with an outdated eigenbasis. As shown in Fig. 4, the EMA scheme enables KL-Shampoo to perform well in practice, under stale eigenbases. Moreover, this scheme can be adapted to strengthen the trace scaling variant of Shampoo (Fig. 10, Appx. H), enabling it to outperform SOAP (Fig. 11, Appx. H).
In the last set of experiments, we evaluate the benefits of using the two-sided estimation scheme under KI. perspective. Specifically, compare the two-sided approach (KL-Shampoo) against the one- sided approach (Shampoo) 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- step and our EMA scheme for eigenvalue estimation, as described in Fig. 3. As shown in Fig. 5, KL-Shampoo consistently and significantly outperforms Shampoo, even when Shampoo employs similar basdd estimation rule.
CONCLUSION
We introduced 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 practical implementation of our approach, and extends naturally to tensor-valued estimation. Our empirical results demonstrate the effectiveness of approach for improving the estimation schemes for Shampoo and SOAP.
10

---

Published as a conference paper at ICLR 2026
ACKNOWLEDGEMENTS
The authors would like to thank Mubarak Shah (University of Central Florida), Frank Nielsen (Sony), Mark Schmidt (University of British Columbia), and Emtiyaz Khan (RIKEN) for their helpful feedback on the manuscript, and Cheng Wu (Microsoft) and Xiya Liu (Microsoft) for their insights on language models and their support. Resources used in preparing this research were provided, in part, by the Province of Ontario, the Government of Canada through CIFAR, the companies sponsoring the Vector Institute (see www. vectori nst tutee..ttte......... a #partners), and Microsoft Azure MI. compute.
REFERENCES
Naman Agarwal, Rohan Anil, Elad Hazan, Tomer Koren, and Cyril Zhang. Disentangling adaptive gradient methods from learning rates. arXiv preprint arXiv:2002.1 1803, 2020.
db::1044850//rrt.2002.1.3..
Shun-ichi Amari. Information geometry and applications, volume 194. Springer, 2016. ISBN
9784431559788 1tt:0001077988-444333
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: 0001100005500////..5553...3 207.2.
Rohan Anil, Vineet Gupta, Tomer Koren, Kevin Regan, and Yoram Singer. Scalable second order optimization for deep learning. arXiv preprint arXiv:2002.09011, 2020. 100:10448500/rrx:22000000000
Rajendra Bhatia, Positive definite matrices. Princeton University Press, 2007. ISBN 9780691 129181. URL http:://ww/ jstor. tsstor.or/sttteeeeeeeee.. wttgstaalle///o/ooooo/.......elt http:/..cctttttteeeoeeeeee 7xxv2
Peng Bo. RWKV-7: Surpassing GPT. https://gittubb.com comp:ttttt.c.......oooo/oe
modded-nanogpt rwkv, 2024. Accessed: 2025-06.
Nicolas Boumal, Bamdev Mishra, P.-A. Absil, and Rodolphe Sepulchre. for optimization on manifolds. Journal of Machine Learning Research, 15(42):1455-145) 2014 URL http:htt/// http:/mmrrrrrrrg wwwpor/pppppp////oooooooooeaarer vv1 /wwuoaoulaaaa...ooo........ 14a html.
Manopt, matlab toolbox
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, 733)2200-217, 1967. ISSN 0041-5553, doi: (0010166001155556669000000
George E Dahl, Frank Schneider, Zachary Nado, Naman Agarwal, Chandramouli Shama Sas- try, Philipp Hennig, Sourabh Medapati, Runa Eschenhagen, Priya Kasimbeg, Daniel Suo, al. Benchmarking neural network training algorithms. arXiv preprint arXiv:2306.07179, 2023. doi: 1001100.850///rv..330066.
Jan de Boer, Victor Godet, Jani Kastikainen, and Esko Keski-Vakee Vakkuri. Quantum informa- tion geometry of driven CFTs. Journal of High Energy Physics, 2023(9):1-89, 2023.
0001000100//HHEPP9222000888
John E Dennis, Jr and Jorge J Moré. Quasi-newton methods, motivation and theory. SIAM review, 19 (1):46-99) 1977.
Inderjit S Dhillon and Joel Tropp. Matrix nearness problems with bregman divergences. SIAM Journal on Matrix Analysis and Applications, 2944) 2008. 10.:10.111370066600 137/060649021.
John Duchi, Elad Hazan, and Yoram Singer. Adaptive subgradient methods for online learning and stochastic optimization. Journal of Machine Learning Research, 2261):2122-2155 2011. URL http:// /jtr///mmrr// or.g pwpppppprase a 2. duchi 11 html.
Pierre Dutilleul. The MLE algorithm for the matrix normal distribution. Journal of Staristical Computation and Simulation, 64(2):105-123, 1999. doi:10. 108000994465999888199777
11

---

Published as a conference paper at ICLR 2026
Sai Surya Duvvuri, Fnu Devvrit, Rohan Anil, Cho-Jui Hsich, and Inderjit S Dhillon. Combining preconditioners through Kronecker approximation for deep learning. In The Twelfth International Conference Learning Representations, 2024. URL https: id=8 jhzz B.
Runa Eschenhagen, Aaron Defazio, Tsung- Tsung-Hsien Lee, Richard E Turner, and Hao-Jun Michael Shi. Purifying Shampoo: Investigating Shampoo's 's heuristies by decomposing its preconditioner. arXiv preprint arXiv:2506.03595, 2025. doi: 01.48800kmii..25000035533 03595.
Roger Fletcher. A new variational result for quasi-newton formulae. SIAM Journal on Optimization, 1(1):18-1,, 1991. doi: 137/0000000
Thomas George, César Laurent, Xavier Bouthillier, Nicolas Ballas, and Pascal Vincent. Fast approximate natural gradient descent in kronecker factored eigenbasis. In S. Bengio, H. Wallach, H. Larochelle, K. Grauman, N. Cesa-Bianchi, and R. Garnett (eds.), Ad- vances in Neural Information Processing Systems, volume 31. Curran Associates, Inc.. 2018. URL https: http:::://ppooceeeiiee Fi1e/48000547h33 7 nwwrips http:///ws...uuu..c.ee//eeaaaeee httpss//ppppceeeennnne.e.essssss
cerrs_fiees// wwwes/iieeeeeeeee fiiess/pape://018// Pap.-epe--.......... paf.
Athanasios Glentis. A minimalist optimizer design for LLM pretraining. https: /ttp///gttb/////////ttt/ttt.hhtt com/ MaAII-LLaMMMMimmalaa LIM Ltt_LLLMePeaaainnnnggggggaaeeeee 2025, Accessed: 2025-06.
John Greenstadt. Variations variable-metric methods. In SIAM REVIEW, volume 10, pp. 474, 1968.
Osman Güler, Filiz Gürtuna, and Olena Shevehenko. Duality in quasi-newton methods and new variational characterizations of the DFP and BFGS updates. Optimization Methods Software, 24 (1):45 62, 2009.
Vincet Gupta, Tomer Koren, and Yoram Singer. Shampoo: Preconditioned stochastic tensor op- timization. In Jennifer Dy and Andreas Krause (eds.), Proceedings of the 35th International Conference Machine Learning, volume 80 of Proceedings of Machine Learning Research, pp. 1842-1850. PMLR, 10-15 Jul 2018. URL https: htt.ttss:///prrceeeee http///eerceee.ceee/eeee/oess.se wwt.mmr.rs.//..... gupta18 html.
Nicholas Higham and Theo Mary. Mixed precision algorithms in numerical linear algebra. Acta Numerica, 31:347-414, 2022.
Keller Jordan. NanoGPT (124M) in 3 minutes. httpas///sttubb..cm/eeelleerrra modded- nanoop,,, 2024. Accessed: 2025-06.
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. 10 10011.0000050778888888880 0011566788.011666000 1.61377..
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-118 2013b. doi: 10. http://.aam.011.4....
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: /ooprreeiiii ew net. forum? id=CtM5x(RS$m
Mohammad Emtiyaz Khan, Reza Babanezhad, Wu Lin, Mark Schmidt, and Masashi Sugiyama. Faster stochastic variational inference using groximal-gradien methods with general divergence functions. In Proceedings of the Thirty-Second Conference on Uncertainty in Artificial Intelligence, pp. 319-328. AUAI Press, 2016. https:/ papers/ 218 pdf.
: /www auai. weeaai.rrg/ae /uagguual 201 // prwpeedeeeeeeeeeeoooooooooo ings/
Diederik Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In International Conference Learning Representations, 2015. 1r.11044550//axxv11411222 1412.6900.
12

---

Published as a conference paper at ICLR 2026
Brian Kulis, Mátyás A Sustik, and Inderjit Dhillon. Low-rank kernel learning with Bregman matrix divergences. Journal of Machine Learning Research, 10(2), 2009.
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 Ma- chine Learning, volume 97 of Proceedings of Machine Learning Research, pp. 3992-4002 PMLR, 09-15 Jun 2019. URL https: / httppp/oooeeiing//ee/ee
mmrr. wwwppprsssaa.e... 1inll1111 9b html.
Wu Lin, Valentin Duruisseaux, Melvin Leok, Frank Nielsen, Mohammad Emtiyaz. Khan, and Mark Schmidt. Simplifying momentum-based positive-defintte 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: pttpp/eeeeeeegsssee/ee/ep www.mmlr.............. wal..press//ee S....press/v2222 / lin23c. html.
Wu Lin, Felix Dangel, Runa Eschenhagen, Juhan Bae, Richard Turner, and Alireza Makhzani. Can we remove the square-root in adaptive gradient methods? A second-orde orderr perspective. In Ruslan Salakhutdinov, Zico Kolter, Katherine Heller, Adrian Weller, Nuria Oliver, Jonathan Scarlett, and Felix Berkenkamp (eds.), Proceedings of the 4 International Conference on Machine Learning. volume 235 of Proceedings of Machine Learning Research, pp. 29949-29973. PMLR, 21-27 Jul 2024. URL https: /ttppppeeoogg..s.p//o/pppoo..ssp mlr prwss v 1ewee/innnn...eeee...m..e...e html.
Jingyuan Liu, Jianlin Su, Xingcheng Yao, Zhejun Jiang, Guokun Lai, Yulun Du, Yidao Qin, Weixin Xu, Enzhe Lu, Junjie Yan, et al. Muon is scalable for Ilm training, arXiv preprint 2rX9922502.16982. 2025.
Karl Löwner. Über monotone matrixfunktionen. Mathematische Zeitschrift. 38(1):177-216, 1934. doi: 10. 1007/BF01 170633.
Hà Quang Minh and Vittorio Murino. Covariances in computer vision and machine learning. Synthesis Lectures Computer Vision, 7(4):1-170, 2017. ISSN 2153-1056. 000:10.1007798-3-001-00111 6.
Depen Morwani, Itai Shapira, Nikhil Vyas, Eran Malach, Sham M Kakade, and Lucas Janson. new perspective on Shampoo's 's preconditioner. In The Thirteenth International Conference on Learning Representations, 2025. URL https: httppp:://pppreeeeeee //tp//pppeeeeeeeeeeee/ee/ee wet.iiewwweeeeeeee.eee.eee.eee forum wtt.:ooommmmmcce1100000
Yurii Nesterov et al. Lectures on convex optimization, volume 137. Springer Cham, 2018. ISBN 9783319915784. 1011100000099783111999000
Michael A Nielsen and Isaac I. Chuang. Quantum computation and quantum information. Cambridge university press, 2010.
Jorge Nocedal and Stephen J Wright. Numerical optimization. Springer, 2006. ISBN 978-0-387- 40065-5. doi: 10.10007977-00-88 40007-5555-774455555
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.
Neal Parikh and Stephen Boyd. Proximal algorithms. Foundations and trends in Optimization, 1(3): 127-239, 2014. doi: 152.1///4400000033.
Xavier Pennec, Pierre Fillard, and Nicholas Ayache. Riemannian framework for tensor com- puting- International Journal of Computer Vision, 66(1):41-66. Jan 2006. ISSN 1573-1405. Noi:00.000//1126666 1010077/1112---00-2222--
Noam Shazeer and Mitchell Ster.. Adafactor: Adaptive learning rates with sublinear memory cost. In Proceedings of the 35th International Conference on Machine Learning, pp. 4596-4604 4604. PMLR, 2018. URL https::tp//oo.. htttt:p://prooeee:iii ngs ml
wtt.ssawww........e..eem..mmeeee prwss/v800 sha zeer 8a html.
13

---

Published as a conference paper at ICLR 2026
Hao-Jun Michael Shi, Tsung-Hsier 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.066 06497, 2023. doi: 20044400aaxxiv.233.0.0o
Suvrit Sra. Positive definite matrices and the s-divergence. Proceedings of the American Mathematical Society, 144(7)2787-27977 2016.
Tijmen Tieleman and Geoffrey Hinton, RMSProp: Divide the gradient by running average of its recent magnitude. Coursera, 2012.
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-0018, 2005. URL https: : / jmlr org/ wwtpppapp////eeaaeeeeeeeeeeeesee / v tsuda05a. html.
C. Van Loan and N. Pitsianis. Approximation with Kronecker Products, pp. 293-314. Springer Netherlands, Dordrecht, 1993. ISBN 978-94-015-81 196-7. 1111:11100777774444-----111---11
Nikhil Vyas, Depen Morwani, Rosie Zhao, Itai Shapira, David Brandfonbrener, Lucas Janson, and Sham Kakade. SOAP: Improving and stabilizing Shampoo using Adam for language modeling, In The Thirteenth International Conference Learning Representations, 2025a. URL
net/ http://oommmddddxxppppppp
Nikhil Vyas, Rosie Zhao, Depen Morwani, Mujin Kwun, and Sham Kakade. Improving SOAP using iterative whitening and Muon. https://ni kh 1 vyass tt StAP/CAPPMMMMM .fff 2025b.
Steven H Waldrip and Robert K Niven. Maximum entropy derivation of quasi-Newton methods.
SIAM Journal on Optimization, 26(4):2495-2511, 2016. doi: 11770
Cameron R. Wolfe. An extension of the NanoGPT repository for training small MOE models. https://gitiie www.:ttw..w...oooooeeeoooooooo.e http://///...ccctoo/////oooooota 2025. Accessed: 2025-06.
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::///// /ttp::////ppeeeeewww// net httt:/ooommiiddd55555500
14

---

Published as a conference paper at ICLR 2026
APPENDIX
A Proof of Claim 1
16
Proof of Claim 2
16
C Proof of Claim 3
17
D Proof of Claim 4
19
Proof of Claim 5
20
Two-sided Shampoo Scheme based on Frobenius norm
G Key Distinction between Shampoo with trace scaling and KL-Shampoo
21
Experimental Setup and Additional Experiments
24
15

---

Published as a conference paper at ICLR 2026
A
PROOF OF CLAIM 1
We will show that the optimal solution KL minimization ming, (Rlg 1.S) with one-sided preconditioner S = (.) @ I, is S E.GG").
By definition in Eq. (3) and substituting (?/d.5a) Lb, we can simplify the objective function as
KL [Elgg]) 1. S)
log (S) + TrrsTTrSSiiii I]) + const.
(ds log d
- S.) + Tr Saits--Baa Bgg T)) + const.
(Kronecker identity for matrix det.)
2
dp
(ds logdet det (S. + Tr(s Egg ()) + const.
(identity log-dcterminant)
2
(db log del S. + BTr(TTTTTT100 gg T)) + const.
(linearity of the expectation)
2
1
(b log det S + E[Tr(d, GI, + const.
(identity for a Kronoeker vector product)
2
db log det S.) + TT
const.
db log Pa) E[Tr )]) + const.,
(8)
where = Mat(g) and P. = S
If we achieve the optimal solution, the stationarity condition must be satisfied regardless of the gradient with respect to S, or S Pa. such as
0 = 0g. KLL [ gg 1)
0p, KL (L 1,
do ( 1---1 +GG+
Eq. and matrix calculus identities)
2
db ( S. +
2
Notice that the KL divergence unbounded above. Thus, the optimal (minimal) solution exists. must be 5 E[GG | to satisfy this stationarity condition.
B PROOF CLAIM 2
We will show that the optimal solution of KL minimization mins..s, .s. KL(EEgg S) with two- sided preconditioner S Sa @ S, should satisfy this condition: S = ([G((;) G and S + E (5%) GG.
Similar to the proof of Claim in Appx. A, we can simplify the objective function
KL E gg 1, S)
llog (egdees))= ((T(+EETTSS111 gg T)) + const.
2
1
(ds log (eleet((,)) + d, log det S.) E[Tr(S- T))) const, (identity for log-determinant)
1
(ds log det ( S.)) de log det S.) E[Tr(S, GS," G )]) + const. (identity Kronecker-vectt-prrduaae)
2
2 do log det P.) - da log det + B. )]) + const.,
where P:
S for k ().
16

---

Published as a conference paper at ICLR 2026
The optimal solution must satisfy the stationarity condition with respect to {Sa: So}. Notice that the gradient with respect to { } can be expressed in terms of the gradient with respect to {Sa, S as 0g KL = -S (0s, KL)SS, and sg KL = (3s, K S.. Thus, the optimal solution must satisfy the following stationarity condition with respect to "
S, 0 - Ss. (B J, S). Using Eq. (9) and simplifying the left expression
0=
ds. (KL(Bggg) S)
Öp. (L 1.
] doP-t BGG...G
1
(10)
gives us this equation
0 - S + -[G()))""""))
that the optimal solution must satisfy.
This naturally leads to the following expression:
S = ds
B(G
Likewise, we can obtain the following expression by simplifying the right expression of the stationarity condition.
1 E|G (S:) -1 G).
S de
C PROOF OF CLAIM 3
To simplify the notation, we define H = Eigg" ]. and re-express the objective function in the KL minimization problem as L(S)= == (L S) = (H,H, We now introduce the proximal- gradient framework (Parikh & Boyd, 2014; Khan et al., 2016) to formally state and prove Claim 3. We assume that an estimated ( is given at iteration I. We use a non-negative function, J(S(1,, ( s(++1)), to measure the closeness between the current and the next iteration. Function f(--)) is known as proximal function. A (unconstrained) proximal -raddent step at iteration + 1 with given proximal function, ((. is defined as the optimal solution of another minimization problem,
1
((++1) arg min (V 5= X)
f(((t) .X),
B2
at every iteration with step-size 32 based on the lincarization of the objective function L. We consider a weighted quadratic function the proximal function.
f(S(t) x) ==
1 2 x - 1 vec X - -(()) T Wvec(X 5(0)
1
2
where W is a given weight matrix. For example, W is the Hessian of the KL divergence W V vee(Y) (L() Y) y 7l(
vec s= This matrix is also known as the Fisher- Rao Riemannian metric for 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, (L( X), at x S(t)
@ vec(S-))
When S S. X Sp admits Kronecker product, want to choose weight matrix W so that this proximal function can be separated into two terms:
1 X X, S(() 2
X 80 X. -S(1) & S. lP
2
Xx.
S 1l,
|x, -5(0) 12
2
17

---

Published as a conference paper at ICLR 2026
Here, consider the weight matrix as the block-diagonal Hessian of the KL divergence, such by setting the cross-block terms highlighted red to zero, where Wj =
as W = Wa 0
0
Wo
8 (Y) KI(S Ya @ exactly the block-diagonal approximation of the Fisher-Rao matrix for a zero-mean matrix Gaussian considered by Lin et al. (2019; 2024).
s for k (a, We can show that this weight matrix is
Now, we can formally state the claim and provide proof of it.
Claim 3. (formal version) The moving average scheme for S == Sa @ S, in idcalized KL-Shampoo a mroximal-gradion step at iteration f + 1,
((++)), ((++)) arg min (vs, Gl fls,fl X,) + (Vs, Xx) + Sb
1
Xxa 00 X, ( I|::
X.
2.32
- ((+1) = (1 - 32)5 + B,E
- GT, s s(++1) -12) B2 (5! - G
with step-size 82 to solve the KL minimization problem in Eq. (3), if we use proximal function using the weight matrix, W, defined above.
In mini-batch cases, approximate the expectations using a current batch gradient (Morwani et al., 2025) (see Eq. (5)), which leads to a stochastic proximal- -gradiett step.
Proof. Because the weight matrix W is block-diagonal, we can slice this objective function for the proximal step into two terms.
(Vs, s. (|s-su) X.) + (Vs, s=)()) X,) + X @ X, - 1
(Vs,fssssm)
2/2
-VVs.flsssv s Xa) 232
Xx.- - + + (|
s X + 2,3
X,,
S((²) 12w
(block X.)
Oblock Xs)
Importantly, W. and Wo are independent of X, and X.. Thus, solve this objective by indepen- dently for each X for E b).
We now show that solving this proximal problem gives rise the estimation rule for ((+++) at iteration t+ 1. We focus on the first term since the second term does not depend on X We can show that W.
can be expressed Wa 82 Pe (L() Y. Yo) x st
(8 voe(()))) vecc s==
This matrix W. is also considered in Lin al. (2024). Importantly, W. is invertible and W- -2 vec(S.)
do @ veee s=g With this result, the optimal solution of X, must satisfy this stationarity condition
1 1-11 |X, - -(((0) 1 (lw.) (note: ||aa S(1 == veo(Xa-S5)) -S(0) W. Wavee(X,
( 5=500 X.) +
233
- C S= +
1
Weveo(XX - -s( (note: (Vs. L s= C Xa) = (V (loc(s.))sssso)) vvo(Xx))
32
- vec(X.) = veo(
32 W- vec(S.) S=500
18

---

Published as a conference paper at ICLR 2026
It is easy to see that the optimal solution of the proximal step is
vec(S = veo veo(Xx) = veo(S( Wa- vec((S)) s=50)
=vec(S(.) 32 ( ( d vve((So)
di @vec(S,)) s=s( V vec(Ss)) s=g
=W.
2
=vec(((())
V o
(chain Jacobian W which known as Bregman duality (Lin al., 2019))
db
-vec(S()) + 2 vec( 1
(2(-d55h)) + )
(recall the definition of C and use Eq. (10))
db
- - s=801
B2 ve((B(Gs)))) GT))
- = vec( + ds
which is equivalent to the moving average scheme in Eq. (5) for updating S, at iteration + 1.
Likewise, we can obtain the moving average scheme for S..
D PROOF OF CLAIM 4
We will show that the optimal solution of KL minimization minx A, (L S. with a two-sided preconditioner S = (Qagiag( Aa ) Q & (QuDiag A Q should satisfy this con- de
dition: A I diag Q E|GP, [Qa) and X ã diag (Q, E[G G]G]),), where PE ==
Q&Diag ((A)) 0 Q and Qu known and precomputed by QR for l E b).
Let St = Q Diag( A QT for l b). Because Qs is orthogonal, is easy see that S.1 Qu Diag(( (As) @ QT
Similar to the proof of Claim 2 in Appx. B, we can simplify the following objective function by substituting S. and S.. Here, also utilize the orthogonality of Qs for k E {a.
KI. (E [g 1,
1
(ds log (S.) + do log dSt + E(TT GT))) const.
1
(ds log det (iagiie.DDaaa.aeee + d, log (e A Q + E[Tr(S, Gs, GT))) ' + const.
2
((do log(1 + (da Σ log log(1 + EIT GS G )]} + const, (use the orthogonality of Qa and Qs)
2
((d) Σ log() (dua Σ log(2 (1) + ET Diag A 0 - GQ, (A A0-1 10-1)00 GT))) + const.
2
=S,
=S,
(11)
19

---

Published as a conference paper at ICLR 2026
The optimal Xa and A should satisfy the stationarity condition.
0 = 02 KL
2 d 1 + 0 (i Q.
Q. GQoDiag(A, QQ G
(use Eq. (LI)
( 0 R [Tr Q. Q.)))
8x_ 1 - diag (Q. GP, Qa))
(utilize the trace the diagonal stracture)
1 1- 2
R diag Q GP, QP.G'
1 0 11 2 2 o diag Q B Q.))
= 2
0 = deAe diag Q G G.G. IQ.))
We obtain the optimal solution by solving this equation.
X ds diag G
Similarly, we can obtain the other expression.
E PROOF OF CLAIM 5
Claim 5. (SOAP and KL-SOAP's covariance estimation for augmented eigenvalues d) The optimal solution of KL minimization: mina (B 1, with preconditioner S (D)
d' EE [veec @Q_ GQ+)) 02
where d Rdadax1 is viewed augmented = [
=
eigenvalue vector, و = Q" g is defined the update of (KL-JSSAP (see Eq. (2)), and Q = Qa 3 Q can be an outdated eigenbasis of (KL-)Shammpo) 's preconditioner.
This proof is similar to the proof of Claim Appx. D. We will show that the optimal solution of KL minimization mina KL (Elgg" S) with an augmented preconditioner S = (QDiag(d)Q) TT)
is d' E (ve.)) 02) where d Rd,d,x1 is augmented eigenvalue vector, Q
GQ,))
Qa & Qs. and Qs is given and precomputed by QR for l E ba,
We can simplify the objective function by substituting Here, we also utilize the orthogonality of Qx for & E {a.
KL E Bigg 1,
(log (Bgadd((Di(()) + (riDgiag c IQQ" Elgg T)) + const.
log(d:)) TTQQDagg DiQQTEggg EEgg DD) const.
(Q Qa X Qi orthogonal)
1
log(d,) + (TDiiiaaag - gg })] const. (linearity of the expectation)
2
T log(d:)) T GQ. Diag vvec( GQ,)) const. (identity of Kronecker-vector product)
2
1
log(d] (d)) + E (usummdd @ (vec(Q2. GQ,)) 02 const. (leverage trace and diagonal struct.)
(12)
20

---

Published as a conference paper at ICLR 2026
The optimal d should satisfy the stationarity condition.
0 0aKL (g'gg' .g9 S,
2 E vec GQ.) [{) (use Eq. (12) and compute its derivative)
- (d [d-Elp [e[ G.) (Q.))2))]
Notice that the KL divergence is unbounded above. Thus, the optimal (minimal) solution exists and it must be d' = E [[veecQQ, Qu GQ.) to satisfy the condition.
F TWO-SIDED SHAMPOO SCHEME BASED ON FROBENIUS NORM
Claim 6. (Shampoo's 's estimation scheme based Frobenius norm) The optimal solution of the
Frobenius norm minimization mins..s, Frob (ggg' S) == ||E|gg S Frob with a two-sided precontioner S Sa 36 Sb should satisfy the following condition.
S T) E[GS
S = TT((S:)2) B GG
(13)
Remark: Although the solution be obtained via rank-1 singular value decomposition (SVD) (Van Loan & Pitsianis, 1993) on this outer product, Elgg } it can be computationally expensive to compute the solution due 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 tensor-valued cases.
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 S. Here, we utilize the definition of the norm and re-express the norm using the matrix trace.
[gg 1 S, 2 Sp
-Tr((Eggg] S, & S₁) (E gg ]
S, & S.))
equivalent definition the square the norm)
-Tr(S) @ S3 gB SS, 8 S,)) const.
(St symmetric for E {a, b))
TT(() (Tr (S. @ Sb)) + const.
(Property of Kronecker product)
TT((S))T((;))) [T (S. @ Sb)) + const. (linearity of the expectation)
T((3) 2E ve.) const.
(Property Kronecker product)
= Tr - ( + const. (Property of a trace)
We can simplify the stationarity condition with respect to S, as below.
0 0s. |ggg ] Sa @ S. 11
as. (5 (r( 2 [T S.G.S)) + const.)
2(Tr(S))) -FGG,,))
Thus, the optimal solution should satisfy this condition S
EGGS; Similarly, we
Tr ((S;)²)
can obtain the other condition. Morwani al. (2025) also consider a similar condition (see Eq. 4 of their paper).
G KEY DISTINCTION BETWEEN SHAMPOO WITH TRACE SCALING AND KL-SHAMPOO
We will show that Shampoo's estimation with trace scaling is generalization of Adafactor. Our inter- pretation of Shampoo's 's update is grounded in a generalization of the divergence used in Adafactor quantum relative entropy (Tsuda al., 2005) Bregman divergence (Bregman, 1967) defined on
21

---

Published as a conference paper at ICLR 2026
Idealized F-Shampoo: two-sided Shampoo based on Frobenius norm (p= 1/2)
Replacing the slow eigen step with a more efficient QR step (replace Step 3)
F-Shampoo:
1: Gradient Computation y == Vε(0)
3a:
Frequent Eigenvalue Estimation with EMA (each iter)
G = Mat(g) Rd.xx Xdb
2: Covariance Estimation (each iter)
Ao A,
+ diag Q A. diag Q A Qb)
- (1-32)
Sa
Sa
Aaa
(1 32) 5b
+ 82
Sb
As
3b: Infrequent Eigenbasis Estimation using QR (every T 2 1 iters)
G.S.G. (T)(S))
(vil
A=
Qu - (sx for E {a, b}
GQu Ao. Q GT /GT//O /EA)
(TT/TT()
Figure 6: Left: Simplified two-sided Shampoo schemes based on the Frobenius norm without mo- mentum. We consider two variants. Variant spired by Claim 6, while Variant 2 is similar to KL- Shampoo's update scheme, which utilizes eigenval- ues. Note that Variant 1 of the idealized Shampoo is known as the two-sided Shampoo in the litera- ture (Morwani al., 2025). Right: Adapting our exponential moving average (EMA) approach en- ables F-Shampoo to use the faster QR procedure and makes it more competitive, as empirically shown in Fig. 7.
A
G Q.Daie A. Q. G/ (
in-
3: Eigendecomposition (every T iters)
Au. Qs + eig(S.) for E {a.b]
4: Preconditioning using Q = Qa @ Q
0- (Q (i Aa @ Ax)
)9
Limma-7433--74-21
5.0
0.0
4.5
4.0-
3.6
3.38
Too:
3.26
8.30
3.24
3.22
3:25
3.20
Itoracions (thousands)
Iterations (thousands)
k-Bhampoo
1-Shampoo
Shampon.Proh
Etampoo
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, improves the performance of the two-sided Shampoo based on Frobenius norm (see Eq. of Morwani al. (2025) and Claim rerrrrree to Variant 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 idcal. 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 reference, we also include the best Shampoo run with power 1/2 and grafting based on the state-of-the-arr version from Meta (Shi et al., 2023).
the trace of the matrix logarithm. This new view of Shampoo's 's estimation is distinct from the existing Frobenius-norm perspective. By contrast, KL-Shampoo's 's update is based on the KL divergence (classical relative anothey 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(-), the Bregman divergence based on this function is defined as
3x(X,)) F(X) -
YY). -Y)).
22

---

Published as a conference paper at ICLR 2026
As example, the KL divergence (classical relative entropy) ( Y) is a Bregman divergence with convex function F(M) = log det(M).
Bp(X,Y) = F - F(Y) - (x(( (X-))
( - log det X) + log det (Y) +T(YY T--T-TT(X------- (XX-
(dofn. of function F F(-))
log det (Y) (egdet((X) + Ir(Y-1X x) ddim(x) ( Y)
where (F(M) = M The KL. divergence is also known as the log-determinet divergence because function F is defined the logarithm of the matrix determinant. Notably, the Hessian of this F()) gives rise to the Fisher-Rao metric, which is also known as the affine-invariant metric (up to constant scalar) (Lin et al., 2023).
Now, we introduce quantum relative entropy, which is also known as von Neumann (VN) divergence, show that Shampoo with trace scaling is generalization of Adafactor. The VN divergence VN(X, Y) is defined as a Bregman divergence with convex function F(M) == Tr M (M) M):
VN(X, Y) == (r(X,Y)
= -(()) F(Y) - (XX
= -Tg(XL X YLogM(Y) + Y Log - Y)
(defn. of function F(-)
-TT(XLoMM(X*) X LogM( + - LogM(Y) (X
(property of the
-Tr(XLogM X - LogM((Y)))
[Tx(X(LogM(X) LogM(YY)) Tr(X) + TY(Y).
where LogM(.) is the matrix logarithm function and Tsuda al. (2005) show that (F Log The Hessian of this F(). gives rise to the quantum physics (de Boer et al., 2023).
(BKM) metric in
Claim 7. (Shampoo's 's estimation scheme with trace scaling) The optimal solution of the von Neumann (VN) divergence (quantum relative entropy) minimization ming..s. (N 1, ==
Tr(S) Tr (B[gg [Log((S)) + const. with a two-sided precontioner S = S. @ S, should satisfy the following condition.
1 GG St S. TT(S:)
E GJ,
(14)
Tr(S%)
where LogM(.) is the matrix logarithm function.
The optimal solutions is Shampoo's estimation rule (power p (t) with trace scaling:
EG G
S: = E
S (E
force S. and So to be diagonal matrices and solve the minimization problem, we obtain Adafactor's update as shown below.
St = Diag (aang(EEGGG (E(GG))) = Diag (F((())(
B
Diag Diag(E(13) T G
Diag E
2G)
St = Diag
Tr Tr(EGGTT]
TT (
sum [E[1 (Go2) sum (Go2)111
If the expectations are not computed exactly, the resulting update scheme is not the optimal solution. For example, Adafactor's update scheme not optimal due to the EMA scheme on the diagonal Kronecker factors.
Remark:
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 $ S, 60 Sb. We will use this
23

---

Published as a conference paper at ICLR 2026
(Kronecker sum) identity, LSgM(S, & S.) = (og I. Ia LogM(S)). simplify the matrix logarithm.
VN(Elgg 1.
1, S) - -r((s) (rig ( + const.
-Tr(S.) Ir(S,) (Egggg [LogM(S)) + const.
Tr(S,) (r:):: Tr (Elgg (LogM( (s.) I, La (og) + const.
Tr((s) (r,)))) Tr (Blgg [LoggM @ Ib) + Tr 1 (I & LogM(S.))) + const.
-Tr(S.) Tr(S)) E Tr (r (TogM Ø [.] T(T (I. x LogM(S + const.
TT((S)) (T,)))) [Tr Log
I [Ir(G' GLo)M(S))) const.
Tr((S)) (r,).) Tr LogGG
[TT [TT(G GLogM(S.)) + const.
-Tr(S.) (5,) E Tr ((GGGLL E [T GLogM( + const.
Lo.M(
Tr(ExpM((P.) Ir(ExpM(PP)) E[ GG P.)] - B[Tr(G' GP.) + const. (15)
where Pa == LSx)(SA)) for k {a, and ExpM(-) is the matrix exponential function.
Notice that the optimal solution should satisfy the stationarity condition. We consider the gradient with respect to P& because this condition must be satisfied regardless of S₃ and Ps for & {a, The condition for the derivative of Eq. (15) with respect to Pa is
0 = 0p. VN(Elgg) J) PxpM((a) Tr(ExpM(P.) - E E[GG
=8,
=T:
where Tsuda et al. (2005) show that 0p, (r(E) Exp) Thus, we can see that the optimal solution must satisfy this condition
S GG
Tr(S5)
Similarly, we can obtain the second condition.
| S Tr(S;)
We can verify that the following solution satisfies these conditions.
G S = E GG ST T+EE
Notice that the optimal S, and S, are unique. However, their Kronecker, which is S S x St is unique. Prior studies (Morwani et al., 2025; Vyas et al., 2025a; Eschenhagen et al., 2025) have shown that this solution 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 Sş to be a diagonal matrix for k E {a, b}.
H EXPERIMENTAL SETUP AND ADDITIONAL. EXPERIMENTS
Experimental Setup 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) (134M), M), and NanoMoE (Wolfe, 2025) (227 M). We consider NanoMoE. it contains 3D weight tensors. This model provides natural testhed for evaluating tensor
24

---

Published as a conference paper at ICLR 2026
Idealized VN-Shampoo: Improving Shampoo (p=1/2) with trace scaling
VN-Shampoo: Replacing the slow eigen step with a more efficient QR step (replace Step 3)
1: Gradient Computation y == Vε(0)
3a: Frequent Eigenvalue Estimation with EMA (each iter)
G : Mat(g) R X
2: Covariance Estimation (each iter)
Aa
+ 32 diag Q. A. Qa) diag Qg A Qb)
A + (1-
2
Sa
Sa
Aa A
Sb
(1 32)
5b
+
82
3b: Infrequent Eigenbasis Estimation using QR (every T > 1 iters)
GG
(v1)
A
Qx gr (r, for k {a.
((A))(
(v2)
G G
(v1)
Figure 8: Left: Simplified VN-Shampoo schemes motivated by Claim 7 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 idcalized VN- IN-Shampoo is known as Shampoo with trace scaling in the literature. Right: Adapting our exponential moving average (EMA) approach en- ables VN-Shampoo to use the faster QR procedure and makes it competitive, as empirically shown in Fig. 10 and Fig. 11.
A
(v2)
3: Eigendecomposition (every
iters)
Au. Qs + eig(Sx) ) for E {a.b)
4: Preconditioning using Q = Qa @ Q
0 xQ Diag( a @ An) 1122"" gg
1/ 1/V II.
(v1)
extension of KL-Shampoo and KL-SOAP, derived directly from our KL perspective. In doing so, demonstrate that methods retain the same flexibility Shampoo and SOAP in handling tensor-valued weights without reshaping them into matrices. We train NanoGPT and NanoRWKV7 tokens), Llama using subset of C4 (2 tokens), and NanoMoE
using subset of Web using& subset of OpenWebTeex (2.5 B tokens). All models except NanoMoF. are trained using mini-batches with a batch size of 0.5 M. We use batch size of 0.25 M 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 NonoMoE: linear warmup cosine step-size. We tune all available hyperparameters for each method--nclludng step-size, moving average, weight decay, damping, and momentum using random search with 150 runs. Our hyperparameter search follows 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 the results from the first stage and focus 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 al. (2025a).
narrower range. In our experiments,
We conduct three additional sets of experiments, following the same experimental setup as described in the main text, further evaluate our approach. Due to limited computational resources, we focus on two language models-NanoGP? (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., 2025e-eeferred to as idcalized 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. 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., refered to as idealized VN--Shmpoo--an find that it performs poorly in practice even when using cigendecomposition. By contrast, incorporating our moving-average scheme enables to perform well and use the fast QR decomposition, as demonstrated in Fig. 10.
In the third additional experiment, we evaluate the suitability of KL versus VN divergence for refining Shampoo's 's estimation rule in comparable setting, where both variants outperform SOAP while matching SOAP-level pre-iteration runtime. As shown in Fig. 11, KL-Shampoo consistently outperforms VN-Shampoo, even when VN-Shampoo is made practical and competitive using similar
25

---

Published as a conference paper at ICLR 2026
techniques those employed in KL-Shampoo. These results underscore the advantages of the KL divergence over the VN divergence.
1lamad-45001 Fineein
3.30
3.30
3.28
3.28
3.26
3.20
Teet Los0
3.24
3.24
3.22
3.22
3.20
100
250
Minutes
Iterations hchossadds)
K1-Stampon
SDAP
Moon
Figure 9; Empirical results (random search with 120 rums per method) show that KL-Shampoo performs better on a larger model with 450 million parameters. We perform QR every 10 iterations in this experiment using engleeprecision (FP32) and do not tune this frequency to optimize KL- Shampoo's ss runtime. From these figures, we can see that both KL-Shampoo and Muon outperform SOAP. KL-Shampoo is better than Muon (Liu et al., 2025) per iteration, but not per time. Although our results show worse performance vs. wall-clock time against Muon, we note that it is possible to improve KL-Shampoo's runtime, for example, by performing the QR decompositions less frequently and using mixed-prees -eeecippreiiionn QR (Higham & Mary, 2022). This is because the QR decomposition is the main computational bottleneck.
26

---

Published as a conference paper at ICLR 2026
KamaGPT FineWiBB
Elema-1a 14---
7.0
0.0
6.5
3.8
6.0
9.30
3.35
3.35
3.30
3.25
(thoossndss)
(thormsandss
Figure 10: 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, makes Shampoo with trace scaling scaling-referred to as Variant 1 of idealized VN- coctmpoo-pracciiin
and enables it 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-arr version from Meta (Shi al., 2023). As reference, we also include the best Shampoo run with power 1/2 and grafting based on the implementation from Meta.
NamGGTT.125M.PP Pl
Llamo- 10M.n--3MM-C--0
3.40
3.400
3.38
3.375
3.350
3.34
3.32
3.30
3.250
3.225
3.20
3.200
3.24
20
40
60
90
300
120
20
40
60
100
120
140
Xt-Shampoo
VN.Shampoo (v2, QR, EMA)
VS-SSlampeo (vl, QR, EMA)
SOAP
Figure 11: Empirical results (random search using 150 runs for each method) demonstrate that the advantages of KL-Shampoo over VN-Shampoo under comparable settings. In particular, we strengthen VN-Shampoo (i.e., Shampoo with trace scaling) by incorporating the QR step and the EMA scheme for eigenvalue estimation, described in Fig. 8, to achieve SOAP- pre-iteration runtime. To ensure fair comparison and eliminate implementation bias, we use our own implementation of -N--hampoo, aligned closely with that of KL-Shampov For runtime comparison, we include the best SOAP run as reference. All methods take the same number of iterations in these experiments.
27

---

Published as a conference paper at ICLR 2026
Practical version of KL-Shampoo
la: Gradient Computation := Vε(0)
G = Mat(g) R x
lb: Use Gradient Momentum
M < (1 M + 3,G
2: Covariance Estimation (each iter)
S.
S
An
32) 5b
+ 8/
S
As
A. Gi.lD A
Q0 G
/d. d,
[GQ, Ding( 28
[/ Diag 1
A = G Q. GOQ. (A -Q G/d. 2. G Qa Diag(A. 1//////11"" Qu Diag(A ]
Eigenvalue Estimation with EMA
(each iter)
Aa
Aa
diag Q. Qa)
Ao
La
- (1 (1-32)
A6
+
82
diag Qg A Qb) = (1-
+ 82
16
la == 2, sum Q2 GQ: D
) 1) = mean( QQ GQ, Diag(A, A, 0 1/2) 102 1)
l. = sum (QQ aT Q. (ia -2 1) mean (Diag( AA.
Q. GQ.] 00 .0)
3b: Infrequent Eigenbasis Estimation using QR (every T > 1 iters)
Qs Qa) for & (a,
Add Weight Decay
<- 0 720
4b: Preconditioning using Q == Qu @ Q
+ 0 (Q (ia @ )=) 10)--0
Q (vec(M)
Figure 12: A practical version of KL-Shampoo with momentum 2,8, and weight decay X. In practice,
also use either damping clipping when computing X @-112 for {a, b). In a low-precision (6 setting, we store 2 @-112 rather than X for k E {a, b} for numerical stability. In the original Shampoo, S& initialized by a non-zero matrix to keep eigenvalues AR non-zero. In KL-Shampoo, we directly initialize A to be non-zero (e.g., 0.1) while keeping Sy. to be zero for b).
-KT--Shmpoo uses momentum in the original space while KI-SSOP,, like SOAP, uses momentum in the rotated space.
28


Reviews:
[{"review_id": "MOGnCYU3w1", "reviewer": "Reviewer_ENoC", "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.", "questions": "Page 3, Figure 1 Caption: \"pre-iteration runtime\",  \"per-iteration runtime\"?\n\n\nAppendix 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.\n\nCould the authors provide theoretical guarantees, showing the advantage of the proposed algorithms?\n\nCould the authors provide more numerical results to compare with other emerging large scale algorithms?", "limitations": "", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 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. \n\n- The proposed methods are not just theoretical constructs; they are practical algorithms that deliver comparable performance over Shampoo and SOAP.\n\n- The paper is generally well-written.", "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.\n\n- The authors do not provide any theoretical convergence guarantees.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "akpFpy1r19", "reviewer": "Reviewer_1qBq", "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.", "questions": "1. Can the authors add comparison to Muon as well - for both runtime and number of iterations?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 4, "strengths": "1. The paper provides a novel perspective of viewing the Kronecker factorization updates based on KL-divergence.\n\n2. 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.\n\n3. The empirical results of the method are pretty strong, outperforming recent second order methods such as Shampoo and SOAP.", "weaknesses": "1. The paper is missing comparisons to Muon, another popular second order optimizer proposed recently.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "RKCewdBGJ9", "reviewer": "Reviewer_TCse", "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$.", "questions": "Please address my concerns above", "limitations": "", "rating": 6, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 3, "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.\n\n\n[1] Morwani, Depen, et al. \"A New Perspective on Shampoo's Preconditioner.\" arXiv preprint arXiv:2406.17748 (2024).", "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.\n\n\n\nThe 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.\n\nThe experiments never report approximation error under the proposed KL objective.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "OSJS5XTzm0", "reviewer": "Reviewer_Mv1d", "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.", "questions": "1. In figure 3, why do the delta definitions have two lines each?\n2. In figure 3 bottom right, where is the modification that defines “augmented”?\n3. Are the output layers and/or layernorms treated differently in the empirical implementation (as is often the case with similar methods)?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The paper provides a new perspective on pre-conditioning for LLM optimization that yields principled algorithms.\n2. The presentation of related work and how the KL methods fit in is very nice.\n3. 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.", "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. \n2. 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)