A New Perspective on Shampoo's Preconditioner
Authors: Depen Morwani
Itai Shapira
Nikhil Vyas
eran malach
Sham M. Kakade
Lucas Janson
Venue: ICLR 2025 Poster (2025)
Keywords: optimization
hessian
shampoo
adagrad
second order optimization
Primary area: optimization
TL;DR: We give a new perspective on the Shampoo optimizer showing that it is close to the optimal kronecker factored approximation of Newton's method/Adagrad.

Abstract: Shampoo, a second-order optimization algorithm that uses a Kronecker product preconditioner, has recently received increasing attention from the machine learning community. Despite the increasing popularity of Shampoo, the theoretical foundations of its effectiveness are not well understood. The preconditioner used by Shampoo can be viewed as either an approximation of the Gauss--Newton component of the Hessian or the covariance matrix of the gradients maintained by Adagrad. Our key contribution is providing an explicit and novel connection between the optimal Kronecker product approximation of these matrices and the approximation
made by Shampoo. Our connection highlights a subtle but common misconception about Shampoo’s approximation. In particular, the square of the approximation used by the Shampoo optimizer is equivalent to a single step of the power
iteration algorithm for computing the aforementioned optimal Kronecker product approximation. Across a variety of datasets and architectures we empirically
demonstrate that this is close to the optimal Kronecker product approximation. We also study the impact of batch gradients and empirical Fisher on the quality of Hessian approximation. Our findings not only advance the theoretical understanding of Shampoo but also illuminate potential pathways for enhancing its practical performance.

Published as a conference paper at ICLR 2025
A NEW PERSPECTIVE ON SHAMPOO'S S PRECONDI- TIONER
Depen Morwani*
Itai Shapira"
Kempner Institute
SEAS
Harvard University
Harvard University
dmorwani 8g hawaardd.ddddd
i tagishaiiragg. hww.aaarvadda....
Nikhil Vyas*
Eran Malach
SEAS
Kempner Institute
Harvard University
Harvard University
wikhille 10g. harvard. edu
ema lach@a.c............ hww.ard, edu
Sham Kakade
Lucas Janson
Kempner Institute
Department of Statisties
Harvard University
Harvard University
sham@seas hadvard.edu
jansonég harvard. edu
ABSTRACT
Shampoo, a second-order optimization algorithm that uses a Kronecker product preconditioner, has recently received increasing attention from the machine learn- ing community. Despite the increasing popularity of Shampoo, the theoretical foundations of its effectiveness are not well understood. The preconditioner used by Shampoo can be viewed as either an approximation of the Gauss Newtoo com- ponent of the Hessian or the covariance matrix of the gradients maintained by Adagrad. Our key contribution is providing an explicit and novel connection be- tween the optimal Kronecker product approximation of these matrices and the approximation made by Shampoo. Our connection highlights subtle but com- mon misconception about Shampoo's 's approximation. In particular, the square the approximation used by the Shampoo optimizer is equivalent to single step of the power iteration algorithm for computing the aforementioned optimal Kro- necker product approximation. Across variety of datasets and architectures empirically demonstrate that this is close to the optimal Kronecker product ap- proximation. We also study the impact of batch gradients and empirical Fisher on the quality of Hessian approximation. Our findings not only advance the theoret- ical understanding of Shampoo but also illuminate potential pathways to enhance its practical performance.
INTRODUCTION
Second-order optimization methods offer significant theoretical advantages over first-order ap- proaches, promising faster convergence rates by incorporating curvature information. Recently, these methods have seen success practical large-scale training of neural networks such as Gemini 1.5 Flash (Gemini Team, 2024) and in the Algoperf benchmark (Dahl et al., 2023; MI.Commons, 2024). One of the primary challenges in this field arises from the substantial memory and compu- tational demands of traditional second-order methods, such as Adagrad (with full matrix) (Duchi et al., 2011) 1) and Newton's method. When applying classical techniques, they require storing and inverting P |P] dimensional matrix H (either covariance of the gradients for Adagrad or the Gauss-Newton component of the Hessian for Newton's 's method), where [P| denotes the model pa- rameters. With modern architectures often comprising billions of parameters, this leads quadratic memory and cubic computational requirements, rendering direct application practically infeasible.
$ Equal contribution. Randomized Author Ordering.

---

Published as a conference paper at ICLR 2025
MNIST-2
CIFAR-SM 5M
ImageNet
Cosine Similarity
Gauss -Newton
0.2
0.2
10
15
20
25
0.0
0
2020
2000
1006
1900
1002
0.0
0
1002 20156
aate asto: web
Training Steps
Training Steps
Training Steps
2.0
1.0
Cosine Similarity
0.0
0.0
0.5
Adagrad
6.6
0.6
5.4
3.4
0.2
5.2
3.2
0.0
0.0
2300
5050
8020
2000
0.0g
13000
30500 42000
32650
Training Steps
Training Steps
Training Steps
Optimal Kronecker
Shampoo?
Shampoo
Figure 1: Comparison of Kronecker product approximations for the Gauss-Newton Hessian (top) and Adagrad preconditioner (bottom) across datasets and architectures. Plots show the cosine simi- larity between the true matrix and three approximations: optimal Kronecker, Shampoo and Sham- poo. Shampoo? closely tracks the optimal Kronecker approximation, outperforming the original Shampoo method, consistent with Proposition 1. For MNIST-2 (binary subsampled), Shampoo? perfectly correlates with the optimal Kronecker approximation as proved in Corollary 2 for bino- mial logistic regression. See Appendix C for dataset and architecture details.
To address this issue, one line of work focuses on efficient approximations of the matrix H (Gupta al., 2018; Martens Grosse, 2015). These methods typically employ either a diagonal approxi- mation (e.g., Adam (Kingma, 2014)) or a layer-wise Kronecker product approximation of H. Such approaches are motivated by the significant memory and computational efficiency gains they of- fer compared to maintaining and inverting the full matrix H. Among the most prominent methods utilizing layer-wise wise Kronecker product approximations are K-FAC (Martens & Grosse, 2015) and Shampoo (Gupta al., 2018).
In this work, we primarily focus on the Shampoo optimizer (Gupta et al., 2018), which has recently gained increasing attention from the research community. Notably, in the Algoperf benchmark of optimization algorithms proposed for practical neural network training workloads (Dahl et al., 2023), Shampoo appears outperform all other existing methods. Another recent study, elucidating the Google Ads recommendation search pipeline, revealed that the Google Ads CTR model is trained using the Shampoo optimizer (Anil et al., 2022). Additionally, a recent work (Shi al., 2023) implemented a distributed data parallel version of Shampoo, demonstrating its superior speed in training ImageNet compared other methods.
Previous research has introduced the concept of optimal Kronecker product approximation (in Frobenius norm) for a matrix M (Loan Pitsianis, 1993). This involves finding the projection of M onto the set of matrices expressible as a Kronecker product of two smaller matrices. Loan & Pitsianis (1993) demonstrated that this optimal approximation can be computed numerically using a power iteration scheme. Our work presents novel result (Proposition 1) that establishes precise connection between Shampoo's 's approximation and this optimal Kronecker- factored approximation. Specifically, we show that the square of Shampoo's 's approximation is equivalent single step of the power iteration algorithm used to compute the optimal approximation.
Despite its empirical success, the theoretical foundations of Shampoo Our results bridge this gap by connecting Shampoo to optimal Kronecker approximation of Newton and Adagrad methods. Moreover, our theoretical insights were utilized in the design of SOAP Vyas
not fully understood.
2

---

Published as a conference paper at ICLR 2025
et al. (2024), recently proposed optimizer that improves upon AdamW and Shampoo language modeling tasks.
The main contributions of the work are summarized below:
We theoretically show (Proposition 1) that the square of the Shampoo's approximation of H is precisely equal to one round of the power iteration scheme for obtaining the optimal Kronecker factored approximation of the matrix H. Informally, for any covariance matrix H = E [gg of H is E G J G]. while the original Shampoo work (Gupta et al., 2018) proposes TR=x" representing a reshaped g into an m x matrix.
where g Rm we argue that the right Kronecker product approximation
E E[GGT] (EEGTT"2EE OE G1 with
We empirically verify that the result of one round of power iteration (i.e. square of the Shampoo's 's approximation) is very close to the optimal Kronecker factored approxima- tion (Figure I), and provide theoretical justification for the same (Section 3.2.1). Our ex- perimental results also show that the approximation proposed by the original Shampoo work (Gupta et al., 2018), which does not use the square, is significantly worse.
For the Hessian based viewpoint of Shampoo (Section 2.2.2), we empirically demon- strate the impact on the Hessian approximation of various practical tricks implemented to make Shampoo more computationally efficient such as averaging gradients over batch (Section 4. and using empirical Fisher instead of the actual Fisher (Section 4.2).
Remark. It is worth noting that previous works (Balles et al., 2020; Lin et al., 2024) have inves- tigated why Adagrad- based methods such as Adam and Shampoo incorporate an additional square root in their updates compared to the Hessian inverse. While this an important question, it lies outside the scope of our work. For more details on this topic, we refer readers to Appendix G.
BACKGROUND
In this section, provide the technical background, including basic definitions (Section 2.1), two perspectives on Shampoo's optimization (Adagrad in Section 2.2. Hessian in Section 2.2.2), and the theory of the optimal Kronecker product approximation (Section 2.3).
2.1 NOTATION AND BASIC DEFINITIONS
We use lowercase letters denote scalars and vectors, and uppercase letters denote matrices. For symmetric matrix A, A >> 0 (resp. A > 0) denotes that A is positive semi-definite (PSD) (resp. positive definite). Similarly, for symmetric matrices A and B, A B (resp. A > B) means A-B B > 0 (resp. A- B > 0). The identity matrix of size n is denoted by I, We use Míi, to refer to the (i.j) entry of the matrix M. The Kronecker product of two matrices A R2x and B Rexs is denoted by A @ B Berxqs It is defined such that (A @ [) sj + j'] A[i BBF J'] where 0 & < P,0 & j < 0,0 & 1' < r, 0 & j' < 8. The vectorization of a matrix A E Rmxx.. denoted by vec(A), is an cin-dimensional column vector obtained by stacking the columns of A on top of one another. We will usually denote vec(A) by a. The Frobenius norm of a matrix mm
denoted by ||,|, Il, is defined as ||||| F
T 1 Ati.j]
The following is basic lemma about Kronecker products that will be used later: Lemma (Henderson & Searle (1981)). (A @ B) vec(G) veo (GAT).
2.2
SHAMPOO
The Shampoo optimizer, introduced by Gupta et al. (2018), builds on the principles of the Adagrad algorithm (Duchi et al., 2011). 1). Shampoo can be understood through two primary perspectives: as component of the Hessian. We explore both perspectives below in Section 2.2.1 I and Section 2.2.2 respectively.
an extension of Adagrad, and as an approximation of the
'The Gauss-Newton component of the Hessian can also he expressed as a covariance matrix. For details, refer Section 2.2.2.
3

---

Published as a conference paper at ICLR 2025
2.2.1 ADAGRAD BASED PERSPECTIVE OF SHAMPOO
Adagrad is a preconditioned online learning algorithm that uses the accumulated covariance of the gradients as preconditioner. Given the model parameters 0, E Rº at iteration t, and the gradient 9t I of the loss function with respect to 8,, Adagrad maintains preconditioner 1- 9:g. The parameter update, for a learning rate 7, is given by:
Hada >7
0r+1 0r H gr.
Shampoo extends Adagrad by maintaining layer-wise Kronecker product approximation of the full Adagrad preconditioner HAda Let G2 Rmxx be the gradient for weight matrix W, E RR at iteration t. The following lemma forms the basis for Shampoo's 's approximation:
Lemma 2 (Gupta al. (2018)), Assume that Gt...., GT are matrices of rank at most Let gr. = vec(G,) denote the vectorization of G. all t. Then, for any & > 0, we have:
(elm (clm= M G.G. 1-1
1/2
(chn
1/2
em 1 - g69
cla C...)
Building on the above lemma, Shampoo maintains two preconditioners, L₁ E Rmxm and R, R=*x, which are initialized elm. and e/n, respectively. The updates for the preconditioners and the Shampoo parameter update, with learning rate 1), are given by:
LT Lr-1 + GT Rr = RT-1 + GTGT; WT+1 WT Gr
In Lemma 2, Shampoo approximates the Adagrad preconditioner H
Σi 9.9. using the
Kronecker product (2i-, G.C.)
1/2
G. 1/2
Our main focus is to study the optimal Kronecker product approximation of Had and how relates to Shampoo's approximation.
2.2.2 HESSIAN BASED PERSPECTIVE OF SHAMPOO
In this section, we describe the Hessian approximation viewpoint of Shampoo, explored by previous works (Anil et al., 2021; Osawa et al., 2023), as alternative to the Adagrad-based perspective. Our theoretical and empirical results apply both viewpoints.
Gauss-Newton (GN) component of the Hessian. For a datapoint (x. v). let f'(:: denote the output of neural network and ((f()). y) represent the training loss. Let E Rmxx be a weight matrix in the network and let D denote the training distribution. For cross-entropy (CE) loss, the Gauss- Newton component of the Hessian of the loss with respect to given by (see Appendix for details):
Os d os aw 6f aw
HaN E(aN))D
D> s~f(m)
[
Here, f(x) refers to the network's output, and Dr represents the training distribution of x (Pas- canu & Bengio, 2014). The right-hand side is commonly referred to as the Fisher matrix, while its counterpart using real labels, B(x,y)> [92.v9.. is called the empirical Fisher. Going forward, for simplicity, we will denote the Fisher matrix as (x,,- [e asumming x drawn from Dz and similarly for both r and y drawn from D.
Optimizers such as K-FAC and Shampoo, when viewed from the Hessian perspective, perform layerwise Kronecker product approximation the Fisher matrix HGN. The following lemma estab- lishes Shampoo's 's approximation:
Lemma (Adapted from Gupta et al. (2018); Anil et al. (2021)). Assume that Ga,e are matrices of rank most Let Ix vec(Ga..)) Then, for any >
E, Bx,)~((a) IT. < (Ez,s~ GT. 1/2 & (Ez,s-((w) Gr.)))
(1)
In Lemma 2 the matrix the left hand side is equal to HGN and the right hand side represents the Kronecker product approximation made by Shampoo. However, directly computing this ap- proximation at each step is computationally expensive. In practice, Shampoo applies two additional
4

---

Published as a conference paper at ICLR 2025
approximations make the process more efficient. First, it replaces the per-input gradient by batch gradient, ie, replacing (G,nf with EB.aGGg.sCBp.. where denotes batch of data points, is the concatenation of f(x) for all (z, v) E B and GB TET (m.y)e RR,--[[]] Gz,s the sampled batch gradient, with s|æ] representing the sampled label corresponding Second, Shampoo replaces sampled labels with real labels, i.e., it replaces EB (., with Eg BB] where GB [ (z,))) Gz,y is the batch gradient.
B.
Thus, if G2 and W; represent the batch gradient and weight matrix at iteration t, and is an expo- nential weighting parameter, then the Shampoo update is given by:
",,,,/,, I = M (1 A)G,GT R₂ = AB (1 - A)G Ga Wi+1 W W W, R,
where L, and Rg represent the left and right preconditioners maintained by Shampoo, respectively.
When viewed from the Hessian perspective, our focus is on studying:
The optimal Kronecker product approximation of the matrix HaN and its connection to Shampoo's approximation (detailed in Section 3).
The effect of the two aforementioned approximations (batch gradients and real labels) the quality of the approximation (detailed in Section 4).
2.3 OPTIMAL KRONECKER PRODUCT APPROXIMATION
In this subsection describe how to find the optimal Kronecker product approximation of matrix H Rmanmm under the Frobenius norm. This problem can be reduced to finding the best rank- one approximation of rearranged version of Н. We define the rearrangement operator reshape(). applied to a matrix H, as follows:
reshapo H) + 1l, nj j' H[mj mj' 17),
where i, i' [0, 10..,,.,,. 1] and j. E [0, .. 1], and reshape( H) E m=*xx" property the rearrangement operator is:
xn? One useful
H B
reshape (H) ab
where A R = vec(A) E I=T" B E R=x" and b vec( E R²² This property can be used prove the following result on optimal Kronecker product approximation:
Lemma 4 (Van Loan & Pitsianis (1993)). Let H e Rmaxxm and let L R"x" R E R=xm Then, the Kronecker product approximation of H is equivalent to the rank-one approximation of reshape( H)) under the Frobenius norm:
H L Ø R
|| reshape(Hf) veo(L) vec( R) L
Since the best rank-one approximation a matrix is given by its singular value decomposition (SVD), conctude the following:
Corollary 1. Let H mmxmm If the top singular vectors and singular value of reshapel H) are represented by u1, U1 and 01. respectively, then the matrices L pmxm and R E gaxn defined by
vec( L) 01 vec( (R) U1,
minimize the Frobenius norm H X R.
Obtaining SVD by power iteration. Power iteration a well-known method for estimating the top eigenvalue of matrix M. also be adapted compute the top singular vectors of matrix. The iterative procedure for the left singular vector and the right singular vector is given by
lx - Mr
Tk
M M E
(3)
where & denotes the iteration number.
Cosine similarity. We will use cosine similarity between matrices measure of approximation quality. For two matrices M, and M2, is defined as
Tr(M, ||
A cosine similarity value of 1 indicates perfect alignment, while value of 0 indicates orthogonality.
5

---

Published as a conference paper at ICLR 2025
3 OPTIMAL KRONECKER PRODUCT APPROXIMATION AND SHAMPOO
Loan & Pitsianis (1993) describe an approach to find the optimal Kronecker product approximation of matrix. Koroko al. (2023) extend this work derive layer-wise Kronecker product approxi- mations of the Hessian matrix for networks without weight sharing. In particular, their Proposition 3.1 relies on the rank-1 structure of gradients for single sample to efficiently compute Hox in son-weiht--shring networks. Our analysis does not rely on any assumptions on the gradients and hence applies to both weight-sharing and non-weight-sharee on these works, this restricted Section 3.1. Our primary contribution (Section 3.2) lies in estab- lishing a novel connection between the square of the Shampoo estimate and the optimal Kronecker approximation.
networks. While our analysis builds
3.1 OPTIMAL KRONECKER PRODUCT APPROXIMATION
This section applies the theory from Section 2.3 to find the optimal Kronecker product approxi- mation of a covariance matrix H = FEg-D,[ge '1 for g E Rmn. Both perspectives of Shampoo, described in Section 2.2, focus Kronecker product approximations of H in the form L @ R, where L E R and R E R=**, but for different distributions Dg- For the Adagrad viewpoint, Dg is the uniform distribution over gt. where refers the gradient at iteration L, giving H/ H For the Hessian viewpoint, Dg is the distribution over gradients with batch size 1 and sampled la- bels, leading to H HGN. To simplify notation, we use B[gg J to represent Dg Lgg J. as our results apply to any distribution p This section explores the optimal Kronecker product approxi- mation for such a generic matrix H, examines its connection to Shampoo, and presents experimental validations for
HAd and
HGN.
Since E Rm", each entry of g can be described by a tuple (i,j) E [m] x [n]. Consequently, each entry of H can be represented by the tuple ((1,5), (V, (')).). We now introduce the matrix xn² which is rearrangement of the entries of H (see Section 2). Using Equation (2), we have: H F[G@G]. Furthermore, by Lemma 4, if T.0 R is the optimal Kronecker product approximation of HI, then er is the optimal rank-1 approximation of f1, where & = vee L) and r vee( R). Thus, the problem reduces to finding the optimal rank-1 approximation of H. Applying the power iteration scheme from Equation (3) estimate the top singular vectors of H, and using Lemma 1, gives the following updates for the k-th step of power iteration:
A == reshape(H)) R==xxn,
l Hr E[G @ G]ra-1
Fk - 4--1 E[G & G1 4 vec(E[G' LG-1 Reshaping the vectors into matrices gives the updates:
La B[GR&-IGT]: Rs SIGTLx-1
(4)
3.2 ONE ROUND OF POWER ITERATION
Our primary approximation replaces the full power iteration scheme (Equation (4)) with just single iteration. This leads to the main contribution of our work:
Proposition 1. A single step of power iteration, starting from the identity matrix, for obtaining the optimal Kronecker product approximation of H exactly equivalent to the square of Shampoo 's approximation of H.
Proof. The initialization for this single iteration uses the identity matrix, i.e., Im and I, for L and R, respectively. This reduces the iterative update equations:
La E[GRG-IGT R - E Lx-
to the simplified one-step updates:
L E|GT]]
R E-G GJ.
With these expressions for L and R, L corresponds exactly to the square of Shampoo's 's approx- imation of H given by the right-hand side of Equation (1).
n

---

Published as a conference paper at ICLR 2025
As shown in Figure 1 this single step of power iteration closely approximates the optimal Kronecker product approximation for both H = HON (top) and H Hada (bottom). In contrast, the upper bound proposed by the original Shampoo work (Gupta et al., 2018) performs significantly worse.
3.2.1
WHY INITIALIZE WITH THE IDENTITY MATRIX?
MNIST-2
CIFAR-SM
ImageNet
Cosine Similarity
0..
Gauss Newton
0.5
5.4
2.4
0.2
622
0.2
0.0
20
5.0
2000
4205
50:5
2000
2000
3.0
20005
20020
30300 40000
55000
Training Steps
Training Steps
Training Steps
1.0
18
Cosine Simitarity
0.
5.0
3.5
Adagrad
5.6
3.5
0.4
0.4
00
0.5g
Training Steps
Training Steps
Training Steps
Figure 2: Effectiveness identity matrix initialization power iteration for Kronecker product ap- proximation. The plots compare oI, VE, a?o?
0101 for left (L) and 01 for right (R) singular VI, afo?
vectors across different datasets and architectures. Top row: Gauss-Newton component of the Hes- consistently sian (Hax). Bottom row: Adagrad preconditioner matrix AHw). Note that 0i
VI, año?
approaches 1 more closely than E. of
supporting our theoretical argument for the effectiveness
of one-step power iteration with identity initialization. See Appendix C.1 for experimental details.
Suppose the SVD of Á is given by म Σ, auu,, or equivalently, H Σg a,U, @ V, The convergence of the power iteration in one step depends on the inner product of the initialization vector with the top singular vector. Let us focus on the left side,2 i.e. the update - G. which as described earlier is equivalent to starting with the initialization In. Let vec (In) Σ, a, ie., In Σ, ar Vi. After one iteration, we obtain E := Σ, a,o,4, and correspondingly. := E, U, We interested in assessing how closely apaproxmmatess the leading eigenvector u1. The cosine similarity between and un is given by a1d, "f.
One reason why the cosine similarity might be large is Á is nearly rank-1 (i.e., o, is large), meaning H is closely approximated by a Kronecker product. However, as shown in Figure this assumption does not universally hold. Instead, we propose an alternative explanation for why a single step of power iteration typically sufficient: the coefficient ai usually larger than a, for all i > 2. We provide both a theoretical justification and empirical evidence for this.
We start by noting that good choice because it maximizes the dot product with possible top components, i.e., PSD matrices (Proposition 2), and is expected to have smaller dot product with the later components.
vec (Tw)
Tr(V)). Using the identity matrix initialization is
Lemma Loaan & Pitsianis (1993)). V1 is Positive Semi-Definite (PSD) matrix.
The discussion for the right side analogous.

---

Published as a conference paper at ICLR 2025
Since V1 is PSD matrix, want to initialize power iteration with matrix that is close to all PSD matrices. We now show that the identity matrix achieves this, specifically by maximizing the minimum dot product across the set of PSD matrices with unit Frobenius norm.
Proposition 2. Consider the set of PSD matrices of unit Frobenius norm of dimension mn denoted by Sm- Then
1 m Im = arg max min (vec(M)), vec(M'))
MeS,. M'ES.
This proposition argues that Im maximizes the worst-case dot product with possible top singular vectors. Now, we argue that its dot product with other singular vectors should be smaller: Lemma 6. V is positive-defiiite then V; for i are not PSD.
Therefore, the diagonal elements of V, for i > 2 need not be positive, potentially leading to cancel- lations (for > 2) in the trace Vi, which equals 0zz. Hence, we expect 0g for i > 2 to smaller than 01.
(for both left and right singular vectors) and Σ, of The latter can be interpreted as the cosine similarity if all a's
To quantify the benefit oj being larger than a, for i 2, we compare VE, o?
were equal, or as a measure of how close f is to being rank-1, since it equals the cosine similarity between d and H. Thus, VΣ. corresponds to the "Optimal Kronecker" cosine similarity shown in Figure In Figure we track both quantities during training and observe that v2, afo
consistently closer to than for both H = HGN (top) and H HAda (bottom).
3.2.2 EXACT KRONECKER PRODUCT STRUCTURE IN H
Our analysis shows that B @ E G] closely approximates the optimal Kronecker product approximation of Н. We now show that this approximation becomes exact when H itself Kronecker product. recover Å. While our earlier discussion focused on the direction of the top singular vectors of H, the rank-1 assumption allows us to derive an explicit expression for A, and consequently for H. Corollary 2. Under the assumption that H is rank-1,
this case, is rank-1, and a single round of power iteration will perfectly
H
EE GG T
2 E B[GTG])
Tr (E [GG
Proof. Let H = our i.. H = oU & V. Let Lm Tr(U)U Rm and I, Tr( (((VVV + R,, where Ru and R., are the residual matrices. After one round of power iteration, the left and right estimates provided by Shampoo are given by E [GG o oTT(V)) and E G G] oTr(U)V.. Hence, can see that Tr (E [GG
Tr(U) Tr(V). Thus,
H oU V (E GG
x E G GJ) T (G|
Since H HaN an m² matrix for binomial logistic regression, rank-1, so the equality in the corollary holds. In other words, the square of Shampoo's 's HaN estimate perfectly correlates with HON for binomial logistic regression. This is demonstrated in the first plot of Figure 1.
We note that (E G & G Tr (E [GGT]) I) as an estimate of Hessian (not Adagrad's 's co- variance matrix) was also derived by Ren & Goldfarb (2021). However, their assumptions were much stronger than ours. Specifically, they assume that the gradients follow a tensor-normal distri- bution, which implies that / is rank-1 and that is mean-zero (thus not applicable to the Adagrad viewpoint). contrast, our approach only requires a second moment assumption on the gradients: Á is rank-1. This weaker assumption allows our results to be applicable to a broader range of sce- narios, including binomial logistic regression. More importantly, derivation and experiments show that the direction [GG] & E G G] closely approximates the optimal Kronecker product, even if A is not rank-1.
8

---

Published as a conference paper at ICLR 2025
3.2.3 DISCUSSION ABOUT OPTIMIZATION
Our primary goal in this work poo's effectiveness, rather than proposing new algorithm. However, the connection we uncovered between Shampoo? and the optimal Kronecker product approximation has implications for opti- mization which we discuss below.
provide theoretical foundation for understanding Sham-
Let us refer E[GG"] & E[G G] as H1. As mentioned in Equation (1). Gupta et al. (2018) used the approximation 1/2 E[GG @ E[GTG]
practice, the gradient step in Shampoo is taken in the direction of | L., where p is tuned as hyperparameter (Anil et al., 2021; Shi al., 2023). Since H H1 p/2 searching over p in H7/2 yields the search space as HJ* meaning that this distinction does not affect optimization speed but deepens our understanding of Shampoo's 's mechanism.
In fact, empirical work has demonstrated the practical benefits of Shampoo?, where it has shown improved optimization performance compared to Shampoo (Anil et al., 2021; Shi et al., 2023). Moreover, Vyas et al. (2024) introduced the SOAP algorithm, incorporating our theoretical insights, including Shampoo" 2 and the trace correction in Section 3.2.2, and showed improvements compared to AdamW and Shampoo.
HESSIAN APPROXIMATION OF SHAMPOO
In this section, we investigate the impact of different practical considerations the Hessian ap- proximation in Shampoo, building on the insights from Section 2.2.2. Specifically, we examine the effects of averaging gradients across batch and using real labels instead of sampled labels. These factors influence how well the Shampoo optimizer approximates the Gauss-Newton matrix HGN, which we previously evaluated with batch size I and sampled labels.
4.1 AVERAGING GRADIENTS ACROSS THE BATCH
The first factor we analyze is averaging gradients across the batch. We transition from computing the gradient on a per-sample basis:
L En (2)) ( Gr G. R - [G) Gz,s]
to averaging across a batch B:
L - BB R- BBIBBBB.CGB..oo GB..l.
where s denotes the concatenation of f() for all e B and GB. |n] ExeB, Gr,s is the batch gradient, with s[x] representing the sampled label for to x.
As demonstrated in prior works, this change has no effect in expectation because Gz is mean-zero for all x when taking the expectation over s ~ f(x) (Bartlett, 1953), i.e. E., [G,. = 0. Lemma (Implicitly in Liu al. (2024); Osawa al. (2023)).
This averaging significantly improves running time, providing a multiplicative speed-up propor- tional to the batch size.
4.2 USING REAL LABELS INSTEAD OF SAMPLED LABELS
Next, we consider the impact of replacing sampled labels f(x) with real labels y, as often done in practice. This shift leads to the empirical Fisher approximation when batch size is 1. Prior work has extensively discussed this approximation and shown that, under certain conditions, the two quantities converge we move towards optima, in the presence of label noise (Grosse, 2021; Osawa et al., 2023; Kunstner et al., 2019). In Figure (top), we evaluate the approximation of HON with batch size finding that the quality remains high throughout training. Yet, as batch size increases, the quality degrades because gradients with real labels are not mean-zero. The following shows how the estimator changes with batch size:

---

Published as a conference paper at ICLR 2025
MNIST-2
CIFAR-5M
ImageNet
Similarity
Batch Size
2.0
10
15
20
25
2012
2013
0003
1300
12000
20033
Tost: 0030 wares
Training Steps
Training Steps
Training Steps
Cosine Similarity
Batch Size 256
0.6
0.5
0.4
0.4
3.2
0.2
02-
Training Steps
0.0
4020
Steps
0.0
Steps
Optimal Kronecker
Shampoo?, Sampled Labels
Shampoo?, Real Labels
Figure 3: Cosine similarity between approximations of HGN and the true Hessian. The top row shows results for batch size 1, where the empirical Fisher closely tracks the optimal Kronecker product approximation throughout training. In the bottom row (batch size 256), the approximation quality degrades as batch size increases (see Lemma 8). The batch size refers to that used in the Hessian approximation, not for optimization.
Lemma 8 (Grosse (2021)), Let B denote the batch and GB gradient. Then
1 [ E(ew)e Ga.y denote the batch
E. G..pl (1
1
En[GnG]]
Ex Er.wy Ger. Enny
[
This lemma shows that, depending on the batch size, the estimator interpolates between E..,, GT, (the empirical Fisher) and Fz., B..,,, [G..y] As shown Figure (top), at batch size 1, when G is equal to En.y GT 2. closely tracks the optimal Kro- necker product approximation. However, with increasing batch sizes (Figure 3, bottom row), the approximation quality begins to degrade.
We that this adjustment has the computational advantage of not requiring an additional back- propagation with sampled labels; instead, these computations can be performed alongside standard training.
5 CONCLUSION
Our primary contribution is establishing a precise connection between Shampoo's approximation and the optimal Kronecker-factored approximation of matrix H. We prove that the square of Sham- poo's approximation is equivalent one round of the power iteration scheme for obtaining this optimal approximation. Empirically, we verify that this single round closely tracks the optimal Kronecker-FF factereer-factredd approximation, significantly outperforming the original Shampoo method. Fi- nally, insights from our work have implications for optimization, with recent research showing im- provements over AdamW and Shampoo by incorporating our theoretical findings.
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGEMENTS
NV and DM are supported by a Simons Investigator Fellowship, NSF grant DMS-2134157, DARPA grant I91INF201001, and DOE grant DE-SC0022199, This work has been made possible in part by gift from the Chan Zuckerberg Initiative Foundation to establish the Kempner Institute for the Study of Natural and Artificial Intelligence. SK and DM acknowledge funding from the Office of Naval Research under award N00014-22-1-2377 and the National Science Foundation Grant under award #IIS 2229881. LJ acknowledges funding from the National Science Foundation DMS-2134157.

---

Published as a conference paper at ICLR 2025
REFERENCES
Rohan Anil, Vineet Gupta, Tomer Koren, Kevin Regan, and Yoram Singer. Towards practical second order optimization for deep learning, 2021. URL httph:tppps........tt /pppreveeeeeveeeeeeeeeeeeeeeeeee net/ / forum? ii Sc8cY4Jpi3s.
Rohan Anil, Sandra Gadanho, Da Huang, Nijith Jacob, Zhuoshu Li, Dong Lin, Todd Phillips, Cristina Pop, Kevin Regan, Gil Shamir, et al. On the factory floor: MI engineering for industrial- scale ads recommendation models. arXiv preprint "Xiv:2209.05310 2022.
Lukas Balles, Fabian Pedregosa, and Nicolas Le Roux. The geometry of sign gradient descent, 2020.
M. S. Bartlett. Approximate confidence intervals. Biometrika, 40(1/2):12-19, 1953. ISSN 00063444. URL httt:://www/ jst.or.og//sooooo..oo.oo wtt.ioo.rrg//taal/e/eeeeee http///eeeeee/33......
Minhyung Cho, Chandra Dhir, and Jachyung Lee. Hessian- optimization for learning deep multidimensional recurrent neural networks. In C. Cortes, N. Lawrence, D. Lee, M. Sugiyama, and R. Garnett (eds.). Advances in Neural Information Processing Systems, volume 28. Cur- ran Associates, Inc., 2015. URL https:tpp///cc htttpp::/pprceeediggge nwwrips......... wwtpppprppppllieeeeeeeee paper/2015/ http:///c6000000000aaaaaaaaaaaaa p.ppeaaa Paper.
George E. Dahl, Frank Schneider, Zachary Nado, Naman Agarwal, Chandramouli Shama Sastry, Philipp Hennig, Sourabh Medapati, Runa Eschenhagen, Priya Kasimbeg, Daniel Suo, Juhan Bae, Justin Gilmer, Abel Peirson, Bilal Khan, Rohan Anil, Mike Rabbat, Shankar Krishnan, Daniel Snider, Ehsan Amid, Kongtao Chen, Chris J. Maddison, Rakshith Vasudev, Michal Badura, Ankush Garg, and Peter Mattson. Benchmarking neural network training algorithms, 2023.
Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large-scale hier- archical image database. 2009 IEEE Conference on Computer Vision and Puttern Recognition, pp. 248-255, Icee, 2009.
John Duchi, Elad Hazan, and Yoram Singer. Adaptive subgradient methods for online learning and stochastic optimization. Journal of Machine Learning Research, 12(61):2121-2153) 2011. URL http: fmir.. wwtppapers//papaaaatpotaaerearra vv12 duchilla. html.
Sai Surya Duvvuri, Fnu Devvrit, Rohan Anil, Cho-Jui Hsieh, and Inderjit S Dhillon. Combining axes preconditioners through kronecker approximation for deep learning. The Twelfth International Conference Learning Representations, 2024.
Runa Eschenhagen, Alexander Immer, Richard E. Turner, Frank Schneider, and Philipp Hen- nig. Kronecket-factored approximate curvature for modern neural network architectures. In Thirty-seventh Conference on Neural Information Processing Systems, 2023. URL https:
/ //openzeview... net. fores//foommaie htt.:/orrroo.ddxxxx333.... d=
Kai-Xin Gao, Xiao-Lei Liu, Zheng-Hat Huang, Min Wang, Shuangling Wang, Zidong Wang, Dachuan Xu, and Fan Yu. Eigenvalue-corrected natural gradient based on new approximation, 2020.
Kaixin Gao, Xiaolei Liu, Zhenghai Huang, Min Wang. Zidong Wang, Dachuan Xu, and Fan Yu. A trace-restricted kronecker-factored approximation natural gradient. Proceedings of the AAAI Conference Artificial Intelligence, 35(9):7519-755 May 2021. doi: 10. 10.1609/aaii.ve v35i9. 16921. URL httpt:::////// aaai org/ aeggindde ii///dddexx.///// /pwp/AAAI/aa /tt.A////tttiteeeeeeeee/////////
Jezabel R Garcia, Federica Freddi, Stathi Fotiadis, Maolin Li, Sattar Vakili, Alberto Bernacchia, and Guillaume Hennequin. Fisher-legendre (fishleg) optimization deep neural networks. In The Eleventh International Conference Learning Representations, 2023. URL https://
openreview wwt. fortm//frmmiii d=c91 A
Google Gemini Team. Gemini 1.5: Unlocking multimodal understanding across millions of to-
kens of context. /ttts://storee googleapis. com/ detppiiepiiiee///////e wweppmind-aae wadd-mmeia/aeee gemini// gemi ni v1 Sa_________e_r_________________ _____reporttpdff 2024. [Online; accessed 19-May-2024J.
12

---

Published as a conference paper at ICLR 2025
Thomas George, César Laurent, Xavier Bouthillier, Nicolas Ballas, and Pascal Vincent. Fast approximate natural gradient descent in a kronecker factored eigenbasis. In S. Bengio, H. Wallach, H. Larochelle, K. Grauman, N. Cesa-Bianchi, and R. Gamett (eds.), Ad- vances in Neural Information Processing Systems, volume 31. Curran Associates, Inc., 2018. URL https: httttps:///proceednng..... neurips wwtparrsss.cpppppeeeeeeeeeee hcc/pappr_fflessee staerr/fieesspppe/20/111 Pap......
Roger Grosse. Adaptive gradient methods, normalization, and weight decay. https:
/ www..w... cs. wwwotto.ooooo.....o http:///otooooeeeeosss////e//ee/ 20rp:224i11222//aaaagg..... L05 1o5__ormalization.pp
2021.
Vineet Gupta, Tomer Koren, and Yoram Singer. Shampoo: Preconditioned stochastic tensor opti- mization. In International Conference on Machine Learning, pp. 1842-1850. PMLR, 2018.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recog- nition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp- 770-778, 2016.
Harold V Henderson and Shayle R Searle. The vec-permutation matrix, the vec operator and kro- necker products: A review. Linear and maltilinear algebra, 9(4):271-288, 1981.
Diederik P Kingma. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6980, 2014.
Abdoulaye Koroko, Ani Anciaux-Sedrakinn Ibtihel Ben Gharbia. Valérie Garès, Mounir Haddou, and Quang Huy Tran. Efficient approximations of the fisher matrix in neural networks using kronecker product singular value decomposition. ESAIM: Proceedings and Surveys, 73:218-237, 2023.
Frederik Kunstner, Philipp Hennig, and Lukas Balles. Limitations of the empirical fisher approxima- tion for natural gradient descent. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d'Alché-Buc, E. Fox, and R. Garnett (eds.), Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019, URL https: https::/pproceddiingg www.iss..eu..p....m cc/paper_ files/ https//pppeeeee////epp
Paperr pdf.
Yann LeCun, Léon Bottou, Yoshua Bengio, and Patrick Haffner. Gradient- based learning applied to document recognition. Proceedings of the IEEE, 86(11))2278-232) 1998.
Xi-Lin Li. Preconditioned stochastic gradient descent. IEEE Transactions on Neural Networks and Learning Systems, 2955):1454-1466 2018. doi: 10.1009/NNNSS:227222222999
Xi-Lin Li. Stochastic hessian fittings with lie groups 2024. URL WWRLLtttp:://exxi..orrg 2:422.1855
Wu Lin, Felix Dangel, Runa Eschenhagen, Juhan Bae, Richard E. Turner, and Alireza Makhzani. Can we remove the square-root in adaptive gradient methods? second-order perspective. arXiv 2402.03496, 2024.
Hong Liu, Zhiyuan Li, David Leo Wright Hall, Percy Liang, and Tengyu Ma. Sophia: A scalable stochastic second-order optimizer for language model pre-training. In The Twelftth International Conference on Learning Representations, 2024. httpttps:/ /twpr//eeeee./eeeeeee/e/ net forum? Add33xHeeANoo.
Zhuang Liu, Hanzi Mao, Chao-Yur Wu, Christoph Feichtenhofer, Trevor Darrell, and Saining Xie. A convnet for the 2020s. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pp. 11976-11986. June 2022.
C. F. Van Loan and N. Pitsianis. Approximation with kronecker products. In Bart L. R. Moor Mare S. Moonen, Gene H. Golub (ed.), Linear Algebra for Large Scale and Real-Time Applications, pp. 293-314. Springer, 1993. URL https: : llikk. springer. com/ chapter 1001007/ 978 978-94-00 -1---196----8-81-1
13

---

Published as a conference paper at ICLR 2025
James Martens. Deep learning via hessian-free optimization. In Johannes Fürnkranz and Thorsten Joachims (eds.), Proceedings of the 27th International Conference on Machine Learning (ICML- 10), June 21-24, 2010, Haifa, Israel, pp. 735-742. Omnipress, 2010. URL https: ptattps://mmm... /icml. cc/Conferente
201 0/papers/aaaaeeee 458 458.858. .ddf.
James Martens and Roger Grosse. Optimizing neural networks with kronecker- approximate curvature. In International conference on machine learning, pp. 2408-2417. PMLR, 2015.
James Martens and Ilya Sutskever. Learning recurrent neural networks with hessian-free free opti- mization. In Proceedings of the 28th International Conference on International Conference on Machine Learning, ICMI'11, pp. 1033-1040, Madison, WI, USA, 2011. Omnipress. ISBN
9781450306195.
James Martens. Jimmy Ba, and Matt Johnson. Kronecker-factored curvature approximations for recurrent neural networks. In International Conference on Learning Representations, 2018. URL
https: http:///ppeeeeeeeeeeeeeeeeee net forum? i d-HyMTKQZZb.
MI.Commons. MIc algoperf benchmark competition. https: / / ml commons. 2t20////22222 /08/ m1 aalc-algoperr wwwpcchhmrrkccooppeeeee on/ 2024. Accessed: 2024-10-01.
Preetum Nakkiran, Behnam Neyshabur, and Hanie Sedghi. The deep bootstrap framework: Good online learners are good offline generalizers. arXiv preprint arXiv:2010.08127,, 2020.
Kazuki Osawa, Yohei Tsuji. Yuichiro Ueno, Akira Naruse, Rio Yokota, and Satoshi Matsuoka. Large-scale distributed second-order optimization using croneckerffactore. approximate curva- ture for deep convolutional neural networks. In 2019 IEEE/CVF Conference Computer Vision and Pattern Recognition (CVPR), pp. 12351-12359, 2019. doi: 10.1109/CVPR.209..0025.
Kazuki Osawa, Satoki Ishikawa, Rio Yokota, Shigang Li, and Torsten Hoefler ASDL: A unified interface for gradient preconditioning pytorch. CoRR, abs/2305 04664,, 2023. doi: 10.48550/ ARXIV.2305.04668 URL httph:tt///// //t::///////////ooo www.dorr....oo........o.o..o.... /100 88500/ arxiv. 2305 684.
Vardan Papyan. Measurements of three-level hierarchical structure in the outliers in the spectrum of deepnet hessians. In Kamalika Chaudhuri and Ruslan Salakhutdinov (eds.), Proceedings of the 36th International Conference on Machine Learning, volume 97 of Proceedings of Machine Learning Research, pp. 5012-5021 PMLR, 09-15 Jun 2019. htLhttps: / httttps:/pproceeiii httpppodocdeeeeeeeooo ngs.
mlr. www..press/vee / wttpeess//////paaaaa.m papyan html.
Razvan Pascanu and Yoshua Bengio. Revisiting natural gradient for deep networks. In Yoshua Bengio and Yann LeCun (eds.), 2nd International Conference Learning Representations, ICLR 2014, Banff. AB, Canada, April 14-16, 2014, Conference Track Proceedings, 2014. URL http: // //arxiv. www//////aaa...eeeee.eee.....eee 111111111111 3584.
Omead Pooladzandi and Xi-Lin Li. Curvature-informed SGD via general purpose lie-group precon- ditioners, 2024. URL htt.s: /ttp:////peeeeeeeeeeeeeee nwt forum? wenerummid-saaa sddsswjxRnVp..
Yi Ren and Donald Goldfarb. Tensor normal training for deep learning models. In M. Ran- zato, A. Beygelzimer, Y. Dauphin, P.S. Liang, and J. Wortman Vaughan (eds.), Advances in Neural Information Processing Systems, volume 34, pp. 26040-26052. Curran Associates, Inc., 2021. URL https::/prroeett
neurips cc waaeeeeffieepppppr//eeee Swweebo0e-PPape... Paper.pdff
Adepu Ravi Sankar, Yash Khasbage, Rahul Vigneswaran, and Vineeth N Balasubramanian. A deeper look at the hessian eigenspectrum of deep neural networks and its applications to reg- ularization. Proceedings of the AAAI Conference Artificial Intelligence, 35(11):9481-944 May 2021. doi: 10. 1111166/9ai.i3/.. 135111.17142 URL https: ojjs.................. wea.saaaa...... org. hwtp//idex.../m wwtp..ppp/p//////p// AAAT a Maaaiiicl
Hao-Jun Michael Shi, Tsung-Hsien Lee, Shintaro Iwasaki, Jose Gallego-Posada, Zhijing Li, Kaushik Rangadurai, Dheevatsa Mudigere, and Michael Rabbat. A distributed data- aatallarallel py- torch implementation of the distributed shampoo optimizer for training neural networks at-scale, 2023.
14

---

Published as a conference paper at ICLR 2025
Charles F Van Loan and Nikos Pitsianis. Approximation with Kronecker products. Springer, 1993.
Nikhil Vyas, Depen Morwani, Rosie Zhao, Itai Shapira, David Brandfonbrener, Lucas Janson, and Sham Kakade. Soap: Improving and stabilizing shampoo using adam. arXiv preprint vrXiv:2409.11322 2024.
Zhewei Yao, Amir Gholami, Sheng Shen, Mustafa Mustafa, Kurt Keutzer, and Michael Mahoney. Adahessian: An adaptive second order optimizer for machine learning. In proceedings of the AAAA conference on artificial intelligence, volume 35, pp. 10665-10673, 2021.
15

---

Published as a conference paper at ICLR 2025
A LIMITATIONS
The main contribution of our work is to demonstrate that the square of Shampoo's approximation of H (whether H refers to Add or HoN) nearly equivalent to the optimal Kronecker approximation of H. While we empirically verify this across various datasets and provide theoretical arguments, the gap between the two depends on the problem structure. In some experiments with the ViT architec- ture (see Appendix B), observe that the gap relatively larger compared other architectures. Furthermore, it remains an open question to understand the conditions beyond those described in K-FAC (Martens & Grosse, 2015)-under which H is expected be close to Kronecker product. Again, in some of the ViT experiments (Appendix B), we find that the optimal Kronecker product approximation // performs much worse compared other architectures.
B ADDITIONAL EXPERIMENTAL. RESULTS
MNIST-2
CIFAR-5M (ResNet)
ImageNet
Casine Similarity
Gauss- Newttn
0.4
0.4
3.2
0.2
0.2
3.0
0.3
2000
0.0
25007
Training Steps
Training Steps
Training Steps
Optimal Kronecker
Shampoo
Shampoo
K-FAC
Figure 4: Comparison of different approximations to the Gauss-Newton Newton component of the Hessian. The results show that Shampoo achieves higher accuracy capturing the structure of the true Hessian compared to both Shampoo and K-FAC. K-FAC shows competitive performance in some cases, but Shampoo? generally offers better approximation.
B. VIT ARCHITECTURE
In this subsection, present the results for Vision Transformer (ViT) architecture trained on the CIFAR-5m dataset. This architecture features patch size of 4, a hidden dimension 512, an MLP dimension of 512, 6 layers, and attention heads.
For these experiments, we utilize three layers from the fourth transformer block: two layers from the MLP (referred to as 'FFN Linear Layer 1' and 'FFN Linear Layer 2') and the QK layer* (referred to as 'Q-K Projection Layer').
C EXPERIMENTS
Datasets and Architectures. We conducted experiments on three datasets: MNIST (LeCun al., 1998), CIFAR-5M (Nakkiran et al., 2020), and ImageNet (Deng et al., 2009), using logistic regres- sion, ResNet18 (He et al., 2016), and ConvNeXt-T (Liu et al., 2022) architectures, respectively. For MNIST, we subsampled two digits (f0, 13) and trained a binary classifier.
For MNIST, we used the only layer, i.e., the first layer the linear classifier, for computing the co- sine similarities. For Resnet18 18 and Imagenet, we picked arbitrary layers. In particular, for Resnet 8, we used one of the convolution layers within the first block ('layerl.1 1.covv1 in Resnet18 For Ima-
The QK layer is separated from the V part of the layer, following similar decomposition method described by (Duvvuri al., 2024)
'ttpps: : http://yttttttcctttttttttttttcct http:////oooooette//////////e//r sreenttt..
html #resnet18
16

---

Published as a conference paper at ICLR 2025
FFN Linear Layer
FFN Linear Layer
Q-K K Projection Layer
Cosine Simitarity
Gauss -Newton
5.2
0.2
2000
2200
2015
1030
0.0
0
2020
2000
1026
1000
1000
0.0
2010
2000
1500
Training Steps
Training Steps
Training Steps
2.0
1.0
Cosine Similarity
0.0
0.0
0.5
Adagrad
6.6
0.6
5.4
3.4
0.2
52
5.2
0.0
0020 Training Steps
6.0
0.0
Training Steps
Training Steps
Optimal Kronecker
Shampoo
Shampoo
K-FAC
Figure 5: Analogue of Figure for ViT architecture and the CIFAR-5m dataset for layers of the network. For some of the figures we observe relatively larger gaps between Shampoo² and optimal Kronecker approximation.
Table 1: Summary of Experimental Configurations. À denotes weight decay and Bn indicates mo- mentum.
Dataset
Architecture
Optimizer
Batch Size
Steps
Ir
81
MNIST
Lincar Classifier
GD
Full Batch
25
0.01
None
CIFAR-5M
ResNet18
SGD
128
10000
.02
None
.9
ImageNet
ConvNeXt-T
Adam W
2048
50000
3e-3
5e-3
genet, used the 1x1 convolutional layer within the 2nd block of convnext-T ('stages.2) ppwcov1'' p in Convnext-T").
Cosine similarity estimation for HGN. For estimating the Frobenius norm of HaN, used the identity:
E H2x u| E.
Haxll
Hessian-vector products with the Gauss-Newton component were performed using the Deep- NetHessian library provided by Papyan (2019).
For estimating the cosine similarity between HaN and its estimator HoN, we used the following procedure:
1. Estimate H, and calculate HaNlIE.
2. Define scaled HaN SaN |Hav| HoN. HHxll
3. Cos-sim (GN, HGN) = 1 Hessian-vector products.
Max S where the numerator is again estimated via 2|
htwp:////mmmmmmmmm////////m/ http:////ennnnnnnee/e//eeeeeaaaa itaaaa////eeeeee/eeiiiiiii.oooo. mode convnex__ting html #torchvi sion. models. convnext_tinnggggg
17

---

Similarly, the label Published as a conference paper at ICLR 2025
Note that in the above procedure, product form with both terms size m x m or n x n, where m x n is the size of weight matrix.
exactly calculate |HaN|| as it is generally of Kronecker
Cosine similarity estimation for HAda- We follow similar recipe as before, but using a difference method for computing the product U. For given time T. Hade Σn 9:92 Thus, H Ei=1 (ge e)gt. We maintain this by keeping a running estimate of the quantity for multiple random vectors during a training run, and use for estimating the product Hdda
C.I. FIGURE DETAILS
Optimal Kronecker method, wherever used was computed with five rounds of power iteration, start- ing from the identity. For H = HaN, the Hessian approximations Shampoo" Shampoo, and K-FAC were done using sampled labels and a batch size of 1. For -Adda and step t, we used gradient enocoutered during the training run in steps I.
K-FAC was computed with the "reduce" variant from Eschenhagen et al. (2023).
In Figure 2, the Optimal Kronecker legend represents the cosine similarity between the optimal Kronecker approximation of HaN and Hox- This is precisely equal to a1
VE,
1. (resp. R) represents the cosine similarity between the top left (resp. right) singular vector of HON and the estimate obtained after one round of power iteration starting from In (resp. Im). This is precisely equal to andy VI. afo?"
In Figure (top), the Hessian approximation is calculated with batch size i.e, |B] in Sec- tion 4.2. Similarly, in Figure 3 (bottom), [B] = 256.
D DEFERRED PROOFS
Lemma 6. If is positive-definite then V4 for i l are not PSD.
Proof. Consider PSD matrices Mi and M2 having the eigenvalue decomposition Mj Enarii and M2 = = A2:22:42 Then
M2)
Σ Ausdas (oT.92)²²
Thus, if M1 and M2 have unit frobenius norm and Mi is positive definite, then Tr( Tr) 0. Thus, if Vi is positive definite, then by orthogonality of successive singular vectors, V: for i >> cannot be positive semi-definite.
Proposition Consider the set of PSD matrices of unit Frobenius norm of dimension m denoted by Sm- Then
1
Vm Im arg max min (vec(M), veo() (M))
MeS.. M'ESm
Proof. Consider the eigendecomposition of any M e S, given by Ei= Auue Denote L = {i: d 11. As Σ X = 1, therefore, |A] > Consider any jE E A. Then
Vedo(y) /q
As vy orthogonal the other eigenvectors. Thus,
1
max min (vec(M), veo( MES, M'S,,
18

---

Published as a conference paper at ICLR 2025
Moreover, for the matrix /g g. for any matrix M',
ur(M)
V9 (1 My
V9
where tr(M") denotes the trace of the matrix M'. M Thus
However, we know tr(M²) -''' 3 as
(I. = tr(M))
/q
/9
44
Note that this is the only matrix with this property as any other matrix will at least have one eigen- value less than
Thus
V4"
I arg max min (vec(M), vec(M')
Mes, M'eS,
Lemma (Implicitly in Liu al. (2024); Osawa al. (2023)).
BBn G.n..GGn.GBa.. Rx))>
(w)) Ga, G
Proof. Evaluating GB..G
B, we get
GB
Σ
Gr
B
J.leR
Taking the expectation over for given B, and by using [x[Gr = 0 we get
1 Ee
E.GGn,.GB..
Een B.aw Ba Gr,
B
Now taking an expectation over batches, we get
Lemma 8 (Grosse (2021)), Let B denote the batch and GB
[&T Ei Gay denote the batch
1
gradient. Then
1 [B]
(1-
1
Ea[GaG]]
En Go,,
E
B
Proof.
Evaluating GDG., we get
GEGI
Σ
Gr
B
(2.2) w)eB
Taking the expectation over B on both the sides, get
EB
[
B
Er. [Gz. GT 2nd + (1312 - |B)) I.r,
= Eg
[GG
Eg Gz Gryy 1+ + (1- [B]
Ex.yl (Z..y E Gr. T
B
19

---

Published as a conference paper at ICLR 2025
E TECHNICAL BACKGROUND ON HESSIAN
Gauss-Newton (GN) component of the Hessian. For a datapoint (x, y), let f(æ) denote the output of neural network and [()(), y) represent the training loss. Let W Rmxx represent a weight matrix in the neural network and 'D denote the training distribution. Then, the Hessian of the loss with respect to W given by
O2 OW? = -(aay)>
o d of T aw 0f² ow
TC o² + (maa)=D Of aw?
(a.y)-
The first component, for standard losses like cross-entropy (CE) and mean squared error (MSE), positive semi-definite and is generally known as the Gaass-Newwt (GN) component (HaN). Previous works have shown that this part closely tracks the overall Hessian during neural network training (Sankar et al., 2021), and thus most second-order methods approximate the GN component. Denoting (C(((r).)) by G E R and Jz.9 = veo(Gz..)) for CE loss, it can also be shown that oW
af L of T O 8 OW
HON (ay))
LE DDDD s~f(w)
[
F RELATED WORK
We have already discussed two closely related works Koroko et al. (2023); Ren & Goldfarb (2021) in Sections 3.1 and 3.2.2 respectively. We discuss them again below for completeness.
Ren & Goldfarb (2021) study the Hessian perspective of Shampoo and show that, under the assump- tion that sampled gradients follow tensor-normal distribution, the square of the Hessian estimate of Shampoo is perfectly correlated with HGN- We also show the same result under much weaker conditions in Corollary 2. Moreover, in Proposition we show that, in general, the square of the Hessian estimate of Shampoo is closely related to the optimal Kronecker product approximation of HGN- We additionally also study the approximations used by Shampoo to make it computationally efficient (Section 4) and the Adagrad perspective of Shampoo's 's preconditioner.
Loan & Pitsianis (1993) develop the theory of optimal Kronecker product approximation of matrix (in Frobenius norm). Koroko et al. (2023) use it for finding layer-wise optimal Kronecker product approximation of HUN for a network without weight sharing. We extend their technique to net- works with shaght-hharig,, and show that the square of the Hessian estimate of Shampoo is nearly equivalent to the optimal Kronecker product approximation of HGN-
Another relevant work is Yao et al. (2021), which introduces AdaHessian, adaptive second-order optimizer that combines stochastic Hessian diagonal approximations with Adam-style momentum and weighted averaging.
F.1 OTHER RELATED WORKS
The literature related to second order optimization within deep learning is very rich, with methods that can be broadly classified as Hessian-frree and methods based on estimating the preconditioner H (which could refer to either H or HGx). Hessian-free methods (Martens, 2010) generally tend to approximate the preconditioned step (for Newton's 's method) using Hessian vector products, but do not maintain an explicit form of the Hessian. Estimating H (Martens & Grosse, 2015; Gupta et al., 2018) methods maintain an explicit form of the preconditioner that could be efficiently stored as well estimated.
F.2 HESSIAN-FREE
One of the seminal works related to second order optimization within deep learning was the intro- duction of Hessian-fre optimization (Martens, 2010). The work demonstrated the effectiveness of using conjugate gradient (CG) for approximately solving the Newton step multiple auto-encoder
20

---

Published as a conference paper at ICLR 2025
and classifications tasks. Multiple works (Martens & Sutskever, 2011: Cho et al., 2015) have tended this algorithm to other architectures such as recurrent networks and multidimensional neural nets. One of the recent works (Garcia et al., 2023) also takes motivation from this line of work, by approximately using single step CG for every update, along with maintaining a closed form for the inverse of the Hessian, for the single step be effective. Other recent works (Li, 2018; 2024; Pooladzandi & Li, 2024) have focused on designing iterative preconditioners to improve the conver- gence specifically for stochastic optimization algorithms.
F.3 ESTIMATING PRECONDITIONER
Given that it is costly to store the entire matrix H, various works have tried to estimate layer- wise Н. KFAC (Martens & Grosse, 2015) was one of the first work, that went beyond diagonal approximation and made Kronecker product approximation to layer-wise HaN- It showed that this structure approximately captures the per layer Hessian for MLPs. This approximation was extended to convolutional (Osawa et al., 2019) and recurrent (Martens et al., 2018) architectures. Subsequent works also improved the Hessian approximation, by further fixing the trace (Gao et al., 2021) as well the diagonal estimates (George al., 2018; Gao et al., 2020) of the approximation. A recent work (Eschenhagen et al., 2023) also demonstrated that K-FAC can be extended to large-scale training.
From the viewpoint of approximating Adagrad (Duchi et al., 2011), Gupta et al. (2018) introduced Shampoo, that also makes a Kronecker product approximation to Ada One of the subsequent work (Ren & Goldfarb, 2021) introduced a modification of Shampoo, that was precisely estimating the layer-wise HaN under certain distributional assumptions. Other works (Anil et al., 2021) intro- duced a distributed implementation of Shampoo, that has recently shown impressive performance for training large scale networks (Shi et al., 2023). Recently, another paper (Duvvuri et al., 2024) proposed modification of Shampoo, empirically and theoretically demonstrating that the new esti- mator approximates HAda better than Shampoo's 's approximation. Our work shows that the square of Shampoo's 's approximation Ads is nearly equivalent to the optimal Kronecker approximation.
G COMPARISON WITH EXTRA SQUARE ROOT IN ADAGRAD BASED APPROACHES
Multiple previous works (Balles et al., 2020; Lin et al., 2024) have tried to address the question of why Adagrad- based approaches like Adam and Shampoo, have an extra square root in their update compared Hessian inverse in their updates. This question is primarily concerned with the final update to the weights being used in the optimization procedure, once we have approximated the Hessian.
The primary contribution of this work is completely orthogonal to this question. We are addressing the question of optimal Kronecker approximation of the Hessian, and its connection to Shampoo's Hessian approximation. This is orthogonal to the Hessian power used in the final update.
21


Reviews:
[{"review_id": "2gU900gPMC", "reviewer": "Reviewer_sBYH", "summary": "The paper presents a theoretical examination of the Shampoo optimizer, a second-order optimization method widely used in machine learning for its efficiency in handling large-scale problems. Shampoo’s preconditioner, based on a Kronecker product, is analyzed in this work to address a gap in understanding its effectiveness. The authors establish that Shampoo's preconditioner approximates the Gauss-Newton matrix or gradient covariance matrix via a novel connection to the power iteration method. Specifically, they show that Shampoo’s approach can be viewed as a single iteration of power iteration aimed at producing an optimal Kronecker product approximation, which they demonstrate empirically to be close to the ideal. Their study further explores the role of batch gradients and the empirical Fisher matrix in improving Hessian approximation, suggesting potential enhancements to Shampoo’s performance in various applications.", "questions": "See above", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "The authors theoretically and empirically demonstrate that the square of Shampoo’s approximation of $H$ is exactly equivalent to one round of the power iteration method, aimed at obtaining the optimal Kronecker-factored approximation of the matrix $H$. This finding offers a clear theoretical basis for Shampoo’s preconditioner, revealing that its approximation is more rigorous and methodical than previously understood, with direct ties to the power iteration process. This insight deepens our understanding of Shampoo's structure and effectiveness, and it sets the groundwork for further enhancements based on this relationship.", "weaknesses": "Despite the fact that I find the results really interesting, I still have my concerns about the relevance of this article to the conference. Below I will try to explain what I mean (if I am wrong, I am open to discussion):\n\n- This paper can hardly be called theoretical, since the authors cite previous work for all lemmas in the paper;\n\n- The whole paper is built on one result (see Proposition 1), which seems to be a mynor result.  \n\nI also noticed some typos:\n\n- Line 104: Kronecker product ($\\otimes$) is used before definition;\n- Line 176 and following: $I_n$ is nowhere defined;\n- Line 235: Frobenius norm ($|| \\cdot ||_F$) undefined;\n- Line 374: PSD is not clear what it means.", "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": "R2Rwv0LMOK", "reviewer": "Reviewer_QDot", "summary": "The authors analyze the similarity between the optimal Kronecker product approximation and the approximation made by Shampoo and show that the square of the Shampoo preconditioner is equivalent to a single step of the power iteration algorithm for computing the optimal Kronecker product approximation. The authors show that the cosine similarity with H_{Ada} (and Gauss-Newton) of preconditioners over training iterations is very similar for Shampoo^2 and the optimal Kronecker product approximation on MNIST-2, CIFAR-5M and ImageNet. However, this similarity can decay with real labels using a large batch size.", "questions": "Why was the decision made to compare matrices using cosine similarity? For example, what do the plots look like when the Frobenius norm \\| A-B\\|_F is used instead? In Figure 3, do you know how Shampoo would perform in the same setting?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "Novel insight interpreting a popular optimization algorithm, with extensions to batches and a real data regime. Insights have already been used to create new optimization algorithms which perform well. A good number of numerical experiments clearly verify the claims made about the similarity with the optimal Kronecker product approximation and the reason for choosing L and R as multiples of the identity. The work is succinct with a clear plotting style.", "weaknesses": "Experimental details are contained in the Appendix, and we see that cosine similarity changes over the training steps, an investigation into whether this is optimization trajectory independent (or at least using the Shampoo or Shampoo^2 algorithm) could be performed.", "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": "JEiadMf0jB", "reviewer": "Reviewer_VKMZ", "summary": "Training modern deep neural networks is challenging due to their size and the inherent non-convexity of the training objective. \nThe most popular optimizers for training these models are first-order methods such as SGD, Adam, and AdamW. \nIn classical convex optimization, it is well-established empirically and theoretically that methods that incorporate second-order curvature information outperform first-order methods. \nHowever, directly extending these methods to the Deep Learning setting leads to prohibitive computational and memory costs. \nFortunately, more computationally friendly approximate second-order optimizers have emerged in recent years, and they have shown promising performance relative to methods like Adam and AdamW. \nForemost among these methods is Shampoo, which recently won the Algoperf benchmark for training neural networks commonly deployed in practice.\n\nDespite its excellent empirical performance, the reasons underlying Shampoo's success have remained a mystery. \nThis paper helps explain the method's success by providing a better theoretical understanding of what Shampoo's preconditioner approximates.\nThe paper provides experiments that corroborate their theory and show it can explain the observed practical performance of variants of Shampoo.\nFinally, the paper provides empirical studies investigating how common tricks employed in practice for computational efficiency impact the quality of the Hessian approximation.", "questions": "1. In the revision, could you also include plots showing training loss, as per my comment in Weaknesses?\n2. Do the authors have any thoughts on why $Shampoo^2$-Real Labels provides a much better approximation quality in the minibatch setting on CIFAR-10 than on ImageNet and Binary MNIST? This is not an essential question; I'm merely curious.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "**Clarity:**\n\nOverall, I found the paper easy and enjoyable to read. \nThe paper is very well-organized, and the authors make all their points clearly.\nFor instance, usually, each theoretical result comes with a plot demonstrating the result empirically. This was very nice, made the presentation more concrete, and helped to immediately validate the theory. \n\n**Contribution:**\n\nThe theoretical insights into the Shampoo preconditioner are valuable and timely, as Shampoo is arguably the most practically effective approximate second-order method currently available.\nThis paper's results can possibly lead to improved convergence analyses of Shampoo-like optimizers and inspire the design of new optimizers in this area.\nThis paper will be of interest to anyone in the ICLR community who is interested in more efficient model training and deep learning optimization.\n\nThe experiments in Section 4 give useful practical insights into how computational tricks commonly employed in practice impact the approximation quality of $H_{GN}$. \nI think these results can help practitioners to develop better heuristics that balance the tradeoff in the preconditioner delivering better approximation quality vs. computation time.\n\n**Correctness:**\n\nI read through all the proofs of the various results in the paper. They are clear and correct.\n\n**Overall:**\n\nI enjoyed the paper and think it makes a nice contribution to the deep learning optimization literature, so I recommend acceptance.", "weaknesses": "**Introduction**\n\nThe first paragraph's discussion of the cost of 2nd-order methods should be conditioned a bit; otherwise, it is a bit misleading. \nA quadratic space requirement and a cubic computational complexity only arise if you naively try to apply classical techniques like Newton's method to Deep Learning.\nBy leveraging automatic differentiation, we can compute hvps (which is usually all we need) without forming the Hessian at a cost of $\\mathcal O(np)$, where $n$ is the number of samples and $p$ is the number of parameters.\nThe same holds for the Fisher and Gauss-Newton matrices.\nOf course, if no minibatching is done, this too is prohibitively expensive as n and p are both very large. \nHowever, in practice, good performance is often obtained using subsampling. \nMany methods in the literature exploit this trick and are called Hessian-free. \nThe authors are aware of such methods, as they discuss them in the appendix.\nGiven this, in the final version, I would like the statements of the first paragraph to be qualified.\n\n**Experiments**\n\nOne area where the paper could improve is if, in addition to plotting the approximation quality against the iteration count, it also gave plots showing how well the optimizer with that approximation performed on the actual objective loss. \nPreferably, these would appear side-by-side. \nThis would give better insight into the impact of the preconditioner's approximation quality on the optimizer's convergence speed. \nFor instance, in the setting of Fig 3. it would be interesting to see if $Shampoo^2$-Real Labels exhibited much worse performance than $Shampoo^2$-Sampled Labels, as $Shampoo^2$-Real Labels provides a much worse approximation to $H_{GN}$ in the minibatch setting.   \n\n**Related Work**\n\nThe authors have been thorough in their list of related work in the appendix, but I feel they should add the citation:\n    \nYao, Z., Gholami, A., Shen, S., Mustafa, M., Keutzer, K. and Mahoney, M., 2021, May. AdaHessian: An adaptive second order optimizer for machine learning. In Proceedings of the AAAI Conference on Artificial Intelligence (Vol. 35, No. 12, pp. 10665-10673).\n\nAdaHessian is another popular stochastic second-order optimizer based on a stochastic approximation to the Hessian diagonal combined with Adam-style momentum and weighted averaging.", "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": "6U4CevdUN5", "reviewer": "Reviewer_kT8h", "summary": "This work provides a theoretical foundation for understanding Shampoo from both Adagrad and Hessian perspective. It also establishes a novel connection between Shampoo’s approximation and the optimal Kronecker product approximation, clarifying misconceptions and highlighting potential improvements.", "questions": "-", "limitations": "", "rating": 3, "confidence": 1, "soundness": 2, "presentation": 2, "contribution": 1, "strengths": "This work analyzes and understands Shampoo from the perspective of Adagrad and Hessian.\\\nA well-structured work.", "weaknesses": "-", "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)