Combining Axes Preconditioners through Kronecker Approximation for Deep Learning
Authors: Sai Surya Duvvuri
Fnu Devvrit
Rohan Anil
Cho-Jui Hsieh
Inderjit S Dhillon
Venue: ICLR 2024 poster (2024)
Keywords: Optimization
Primary area: optimization

Abstract: Adaptive regularization based optimization methods such as full-matrix Adagrad which use gradient second-moment information hold significant potential for fast convergence in deep neural network (DNN) training, but are memory intensive and computationally demanding for large neural nets. We develop a technique called Combining AxeS PReconditioners (CASPR), which optimizes matrix-shaped DNN parameters by finding different preconditioners for each mode/axis of the parameter and combining them using a Kronecker-sum based approximation. We show tighter convergence guarantees in stochastic optimization compared to a Kronecker product based preconditioner, Shampoo, which arises as a special case of CASPR. Furthermore, our experiments demonstrates that CASPR approximates the gradient second-moment matrix in full-matrix Adagrad more accurately, and shows significant improvement in training and generalization performance compared to existing practical adaptive regularization based methods such as Shampoo and Adam in a variety of tasks including graph neural network on OGBG-molpcba, Transformer on a universal dependencies dataset and auto-regressive large language modeling on C4 dataset.

Published as a conference paper at ICLR 2024
CASPR: COMBINING AXES PRECONDITIONERS THROUGH KRONECKER APPROXIMATION FOR DEEP
LEARNING
Sai Surya Duvvuri'
Fnu Devvrit
Department of Computer Science The University of Texas at Austin saisurya@cs ut edu
Department of Computer Science
The University of Texas at Austin devvrit@cs ute.as.edd
Rohan Anil
Cho-Jui Hsieh
Inderjit S. Dhillon
Google DeepMind
CS Department, UCLA & Google chohsiehêcs uela. edu
Google i isdegoogle.o com
rohanani @ohananil@google.. com
ABSTRACT
Adaptive regularization based optimization methods such as full-matrix Adagrad which use gradient second-moment information hold significant potential for fast convergence in deep neural network (DNN) training, but are memory intensive and computationally demanding for large neural nets. We develop technique called Combining AxeS PReconditioners (CASPR), which optimizes matrix-shaped DNN parameters by finding different preconditioners for each mode/axis of the parameter and combining them using Kronecker-sum based approximation. The Kronecker-sum based combination allows us to show that CASPR ordered between a well-known Kronecker product based combination, Shampoo, and full-matrix Adagrad preconditioners Loewner order, as a result, is nearer to full-matrix Adagrad than Shampoo. We also show tighter convergence guarantees in stochastic optimization compared to Shampoo. Furthermore, our experiments demonstrates that CASPR approximates the gradient second-momen matrix in full- matrix Adagrad more accurately, and shows significant improvement in training and generalization performance compared existing practical adaptive regularization based methods such as Shampoo and Adam in a variety of tasks including graph neural network on OGBG- GGBG-molppba, Transformer on a universal dependencies dataset and ruto-regre large language modeling on C4 dataset.
INTRODUCTION
Adaptive methods, including Adagrad (Duchi et al., 2011) and its variants such as Adam (Kingma & Ba 2014) and RMSProp (Tieleman & Hinton, 2012), are widely employed for training neural networks. These approaches calculate updates by scaling each coordinate of the gradient direction (or momentum in case of Adam) (9:); with an adaptive learning rate, which is tuned using the second moments (g)) the gradients before updating the weights. This process is equivalent to meet-lutiplying the gradient by diagonal matrix (preconditioner) to perform the weight update. On the other hand, full-matrix Adagrad (Duchi et al. 201 stores cross-moments, i.c., the (-th) element of the (full-matrix statistic) Σ. 9t9. where Jt E IRd. However, this can be infeasible for modern neural networks with large number of parameters due the (dd)) memory requirement to store the mecontsmoments and (d3) computational complexity to perform inverse square root operations.
We address this issue by first proposing a set of preconditioners- axes preconditioners, one for each axis/dimension of matrix shaped parameters in deep neural network, which is derived by finding a block-diagonal approximation with identical blocks such that it is nearest to full-matrix
Work done student researcher at Google

---

Published as a conference paper at ICLR 2024
statistic. Each approximation offers the advantages of a) fast inverse square root operation (which then gives the axes preconditioner). and b) low memory. However, using only one axes preconditioner throughout the training can be suboptimal, statistic. To this end, we first propose a sequence of combinations of axes preconditioners using Kronecker-sums, and subsequently show that Kronecker product based preconditioner (an existing preconditioner- - Shampoo (Anil et al. 2020; Gupta et al. (0188) to be at one end of this sequence, we find our proposed novel combination CASPR to be at the other end of the sequence (Lemma 3.2p. CASPR preconditioner is carefully picked from the proposed set of combinations such that is ordered (sandwiched) between Shampoo and full-matrin Adagrad in Loewner order (Löwner) 1934). result, it is nearer to the more powerful full-matrix Adagrad preconditioner than Shampoo in Loewner order. This is acheived by utilizing the simple inequality in Loewner order: (A + B) = 4AB for commutable and positive definite A and B. Surprisingly, we show that the established Loewner order also helps us derive tighter regret bounds in the online convex optimization framework (Hazan al., 2016; Shalev-Shwarts et al. 2012) than Shampoo, which translates to convergence guarantees in stochastic convex optimization via online-to-bateis conversions (Cesa) Bianchi et al. 20044). Furthermore, we conduct numerical experiments to measure nearness of the CASPR approximation to full-matrix statistics encountered during deep neural network training and show that our CASPR approximates full-matrix statisties better than individual block-diagonal approximations (with identical blocks) corresponding to each axis preconditioner and Shampoo.
it only stores partial information about full-matrix
Contributions:
We develop a practical adaptive regularization method which uses Kronecker- sum based combination of axes preconditioners. We prove a tighter regret bound than a well- known practical adaptive method Shampoo (Gupta al. 2018 Anil et al. 2020) in online convex optimization framework. Note that regret bounds in online convex optimization framework have reductions to convergence guarantee in stochastic convex optimization non-smooth non-convex optimization (Cutkosky et al.] 2023) and smooth non-convex optimization guarantees (Agarwal et al., 2019). We conduct numerical experiments to compare the approximation error with respect to full matrixx statistic and demonstrate better approximation error than Shampoo, diagonal preconditioners, and individual axes preconditioners. We train graph neural network (Battaglia et al. 2018) on OGBG-molpcha dataset (Hu al. 2020), where CASPR gives relative improvement of 1.4% in test average precision compared with Shampoo and 7% improvement with respect to AdamW. We also train Transformer network (Vaswani et al., 2017) on universal dependencies dataset (Nivre et al. 2020), where we observe 1% relative improvement in validation error compared to Shampoo and 3% relative improvement compared AdamW. To demonstrate scalability of our method we train a 14 million parameter on ~ 167 billion tokens and 234 million parameter model with ~ 42 billion tokens on C4 dataset (Raffel et al., (2200. where we noticed upto I % relative improvement in log perplexity. Furthermore, it is worth noting that the simplicity of CASPR would only require small modification to existing implementations of Shampoo. As result, we notice that Shampoo and CASPR take almost similar training time in experiments.
Desa-Bianchi al. 2004).
Notation and Standard Matrix Identities. We use A @ B denote Kronecker product (Petersen et al.) 2008), and vec( A) to denote row-major order vectorization of matrix A. Properties of Kronecker product: producct: a) ( [A @ B²' (G) = vec(AGT) b) (A @ B) & D) = (AC) @ (BD). We use Loewner order A > B to denote that A - B is positive semi-definite (PSD) matrix. We use boock-diagAA,......... Am to denote a block-diagons matrix with blocks {AI,..., A. On the other hand, diag(A)), Ar, if i else O. preserving only the diagonal entries of A. We use (r( to denote trace of matrix and use ||A|| / (t((A) A) to denote Frobenius norm, and (u, v) u
to denote dot-product between vectors u and v. For x 0, A denotes the matrix pth root of A.
RELATED WORK
Hessian- methods. Newton's method is a second-order optimization technique that pre- multiplies the gradient by the inverse of Hessian, which constitutes second-order partial derivatives to update the model parameters. Although it can have fast convergence properties, it is computationally expensive and inapplicable for large-scale neural networks due to the high computational and memory complexity involved in calculating the Hessian matrix. BFGS Seeydden--eeecer---ee
-hannno) (Broyden 1967 Fletcher 1970] SShnnoo 1970; Goldfarb 1970) and LBFGS (Limited-mennl memory BFGS) (Liu & Nocedal| 1989) are quasi-Newton methods that implicitly approximate the Hessian matrix, without computing second-order partial derivatives. However, directly applying LBFGS on
2

---

Published as a conference paper at ICLR 2024
neural networks can be infeasible due high memory requirement in maintaining the gradients of previous iterations in memory.
Adaptive regularization methods. Adaptive regularization introduced in Duchi et al. (201 includes a diagonal preconditioner (diagonal Adagrad) that adapts the learning rate for each parameter individually based the historical sum of squared gradients. Although full-matrix version of this preconditioner is outlined in this work it has been disregarded for its high memory and computational requirements. There exist low-rank approximations of the full-matrix statistic such as AAgawal et al. (2019), which demonstrate better performance in anisotropic loss landscapes compared to diagonal preconditioners. However, this method requires storing hundreds of gradients from previous iterations. which is meeorry-itennsvve Alternatively, Kronecker product based approximation such as Shampoo (Gupta et al. 2018) reduces the time and memory complexity by leveraging the matrix structure of the gradient. Shampoo, which admits regret bound in the online convex optimization (OCO) framework (Hazan et al., 2016)), has demonstrated superior performance than diagonal preconditioners such as AdamW, RMSProp in large-batch settings (Anil et al. 2020) where the time to conduct large batch forward-backwand propagation is high compared to the costly optimizer update rule and also has shown adoption in industry-scaan scale clickthrough rate prediction models (Anil et al.) 2022). It is worth noting that regret upper bound guarantees can be converted to stochastic optimization convergence rate guarantees reach stationary point for smooth non-convex objectives in (Agarwal et al. 2019) and more recently, conversion to non-smooth non-convex objectives (Cutkosky et al.) 2023).
Fisher approximation to Hessian. Natural Gradient Descent (Amari 1998) is optimization method that uses inverse of Fisher information matrix of the model's predictive distribution as preconditioner. This has shown to be generalized Gauss-Newton approximation of the Hessian matrix (Martens 2020). Natural Gradient Descent (NGD) suffers from high complexity due to the need to compute the Fisher matrix and its inverse, which can be infeasible to compute for modern neural networks. Kronecker product approximations such as K-FAC (Martens & Grosse (0155 Ren & Goldfarb) 202T), have appeared to address the complexity issue. However, the gradients used to approximate the Fisher matrix require an additional backpropagation step through different loss function, which can further increase the computational complexity.
3 KRONECKER SUM BASED PRECONDITIONER
Full-matrix Adagrad Duchi et al. (201) Agarwal et 2019) maintains H, E R.dxd and updates the parameters Wi Red using the following steps:
H, = H + 9t9:
Wh+1 = we m7 H,
(2)
where fe: Radxd + IR is the loss induced by the mini-batch at iteration 1 E (1...., ,,,. and It Ve(we)). Here H, stores all gradient cross-moments (H:).. (91 Fi:t (())), ...,, (9+).). Storing the entire full-matrix statistic H, Bddd is infeasible for deep neural- network (DNN) as can be in the millions. Furthermore, the inverse square root takes (dd)) floating point operations (flops).
(5), ) where 9i:t
To address these issues, we aim to find an approximation il, of H, that: a) offers fast inverse operation, b) has low memory representation, and c) approximates H, well. To enforce these constraints formally, we introduce an optimization sub-problem to determine a sparse approximate statistic A constrained to a set of sparse matrices S, before conducting the inverse root find the preconditioner H, 1/2 and updating Wy:
A. arg min min||// Helle, W141 We -ncH, 1/2 112 9t.
(3)
Hes
Constraining Å to the set of all positive semi-definite (PSD) matrices S = (A E Rdddd A > 0} in subproblem 3) trivially gives the full-matrix statistic H, = H,, since H, ł O. Note that inverse square root H, 1/2 cannot be conducted on an indefinite i1, But, if we choose S to be the set of diagonal PSD matrices, = {A E ddddd A 1 O, Ar, i * j E [d]]. the optimal solution is A, = Biag(H,), which leads to the commonly used diagonal Adagrad update:
A, Å diag(o:g.)),
W w. 7/ H, 1/2 9t-
3

---

Published as a conference paper at ICLR 2024
where, diag(A) A diagonal elements of He in memory, and (dd) computations for the inverse square root operation. Nonetheless, the approximation error || - H, ll can be substantial when using the diagonal sparsity structure since off-diagonal elements of H1, are neglected.
if j. and 0 if + j. Maintaining the above H. requires only keeping d
This calls for more accurate approximations H, that also capture cross cross-moments of gradients. Note that gradients of individual layers in most DNNs assume a matrix structure, for example, weight matrices of fully-connected layers in DNNs. The gradient corresponding to weight matrix W E Rm*" can be written as G2 E I R with 9: vec(G:), where vee(A) E IRn denotes vector with elements of matrix in row-major order and total number elements = mn,, For such matrix-structured gradients, we now generalize the diagonal sparsity structure mentioned earlier to the following block -digoonal structure, Sy {block-diag/ R1, Rg Rgm) R, ł 0 E R=** Vi e [m]]. where individual blocks R, capture gradient second-moment within row i of G,. This is related to unit-wisc preconditioning in Osawa et al. (2023). Setting = Sp in ( results in:
min H HESR
H2 (Ri....
min
Σ RR
H1
11
Σ sc|m
min R4 H R.
R H
(9) [m]
where 9 denotes the i-th row of G., and H E R=x" is the i-th diagonal block of H, = Ee=1 9=9, However, storing all the R, requires mn? memory, which can still be impractical when m and n are high. Furthermore, the inverse root operation would require O(mn²) flops, as the inverse needs to be computed for each block individually. To alleviate this issue, we propose a simpler block-diagonal sparsity constraint Siøn, where all the RG are identical, in the following lemma.
Lemma 3.1. Consider the sparsity constraint SISR = {block- dinck--ii((ae R. R..., R) : R 2 0 E R"X"} in 3 then the optimal solution R for (3) (rewritten below) is follows:
R = arg min Rz Im & R
Σ 9-9.
1
m Ed.
where gs = (.).
Proof. Let H E R2 be the i-th diagonal block of H, = E.=1
9=9.
Then:
min feSteR
| - Hel
min R20
Σ
||
H. 113
min m m
Σ
tr(H! R).
R25
4E/m]
VE ml
R
m Σ H
1 m
1 m
1 G
m
Similarly, for column preconditioner corresponding to g. = G defining the sparsity constraint
SLØI = (L & J, L 2 0 E Rmxm) and solving the subproblem gives L = E. D.E.i G.G.
(see Appendix A Lemma A.I].
Individually A = L @ In and B = Im & R' are optimal approximations for the sparsity constraints St and SræR respectively for the subproblem (3. However, using the corresponding precondition- 1-1 -1/2 &) I, and B-1/2 T. -11) (via (3)) individually can be inefficient as they do not approximate gradient mooss- outside the block-diagonals of the full-matric statistic Hd, To this end, several works have used Kronecker product structure (Martens & Grosse 2015, et al., 202T Gupta et al. 2018) for preconditioner approximation. However, we take an alternative approach by the following general combination of axes preconditioners.
Lemma 3.2 (A general combination of axes preconditioners). Consider a sequence of preconditioners for p 2p11 1:
X (p) = (ii,- @ In + Jm @ R 1/P) /2)2p,
where, L, = Ei=i G.G.T + el R = Et= G G, + eIn then
lim X caspr (p) = 1- @ R/ x 7000
4

---

Published as a conference paper at ICLR 2024
where, Xi is the Shampoo preconditioner. which is existing Kronecker product based precondi- tioner (Cupta et al. 2018 Amil et al. 2020
Lemma 3.2 shows that the Shampoo preconditioner is a limiting case of the general CASPR precon-
ditioner X X.aepe (p). The proof of Lemma 3.2 is given in Appendix A.1 To compare with full-matrix preconditioner, tioner than the Shampoo preconditioner in Loewner order (öwwner 1934 Bhatia 2009).
show in the following lemma, that X, (1) is nearer full-matrix precondi-
We will use
this fact to prove tighter convergence guarantees than Shampoo in Section
Lemma 3.3 (CASPR preconditioner). Let Let I = =E... G,G. cl R - G. CI, and
max, rank (G+). then the CASPR preconditioner ((I, L,, + Im @ R4 R, )/2) (Algorithm
and Shampoo preconditioner L," @ R, 1/4 follow the Loewner order:
1/2 = ((t, s In Im 60 R //2) = 1, s R vi(reTa + Σ g )-12
We give the proof of this lemma in Appendix A.I Shampoo preconditioner is the geometric mean of X = 1/2 @ In and X = Lm @ R 1/2 X = (X x//1//2 L, @ R, 1/4 where last equality is due to (A & B)1/2 1/ @ 11/2 (Petersen et al. 2008) for positive semidefinite A and B. This geometric mean interpretation explains the occurrence of exponent -1/4 /4 for L, and R, in Shampoo preconditioner. The proposed preconditioner Xeaspr (1) can be rewritten as follows:
( (f.-1/4 80 In + Im 3 /2) (i-1/2 @ In + Im & /21122/22 L, i,1 @ R1 1/4
2
(x/ x2 (X/X2)1/2
2
Using the last equality, Xm (1) can be interpreted as the average of arithmetic mean and geometric mean (which is Shampoo preconditioner) of X and X
Using the CASPR preconditioner to update the parameter wy = VCC(W)) E Rd. where d mn, will give the following update rule:
X, (I, 1/4 2 In 1 & R 1/)
W/ Wh-1 ne
(vector update),
where Vee (et(G,) E Rdd Here we dropped the /r scale factor from Lemma 3.3 it can be absorbed in the learning rate parameter 17. Forming the entire CASPR preconditioner X, RRdd infeasible, since d can be very large and computing the update X.9 would cost (dd)). However, expanding X, and using the identity (A (A&C) vec( = (eC(A update rule:
will give the following feasible
W4+1 == W, - (1, (I, -., 1/4 + G.R, /4
(matrix update).
This update rule can be further simplified into two preconditioning steps which only requires comput- ing f////4 1/4
and R 1/4 follows:
U. = i,-1/4 + +G
U. = U. + U.R," 1/4
(CASPR update)
W. = W - n/U:
which performs the above every iteration by maintaining I., and R, in online fashion (see line which require storing total of (m² + n²²) elements of matrices L, and R, in memory. Given the simplicity of the update, it is easy change the existing implementations of Shampoo implement CASPR, since both CASPR and Shampoo compute the same axes preconditioners i, ř,-1/4 however Shampoo preconditions the gradient G₂ differently:
1/4 and R R,11/1
U. = - l/G.
U- U. 1/4
(Shampoo update)
5

---

Published as a conference paper at ICLR 2024
Algorithm 1 CASPR Algorithm
1: W1 E Rmxx Lg 0, Ro
2: for t E I..... do
3: Compute gradient G, V)((W)) E mx
4: Update preconditioners:
4-1 = R
5: Precondition gradient and update parameters:
Compute 1/4
-, = (Lr + +elm)-114 1,,, R = (Re + elm)
U² = E, L,2/AG + R. 1/4
U. = i,- + U,R
W: = -1j
6: end for
In line the preconditioning of the columns and rows of G. is achieved sequentially using L, 1,71/4 and R2 1/4 It worth noting that this step takes O(m + n²²) flops due to the inverse operation, which is the same cost incurred by Kronecker product based approaches.
3.1 COMPARISON OF CASPR AND SHAMPOP APPROXIMATIONS
1/4 Let X,aspr -(L & In + Im @ R. )/2) and X²'h L, @ R/ denote the precon- ditioners of CASPR and Shampoo respectively as in Lemma 3.3 (here e is dropped for brevity), where I,2 = - G,G. and R2 = E. GJ G,, then the corresponding induced approximations figee and H1.h for H, E, 999, will satisfy X.aspr respectively. This relationship is evidenced in the update rule for Adagrad, as in (2). Then the explicit expressions for the approximations freeer and Hsh and their eigendecompositions is as follows:
(H/ and X² (
Lemma 3.4 (eigendecomposition of CASPR preconditioner). Let the eigenpairs of Lu be {(11, uz),................. (Am, }m)} and R2 be {[o1, U1). (on: vn then the eigenpairs of iequsr (Xceepr) -(()) 1/4 8 In + Im R2 1/4 1/2) and 11.th (X²h)-2 L12/ @ R, are 1/2
{(((A, 1/4
1/4 + +0, )/2) 1/2 us 2 vz) i.j E [m], [n]] and {((),o,) un 2 v2) i.j E [m].
respectively.
The proof is given in Appendix A.1 which uses the Kronecker product structure of H1h and Kronecker sum structure of Aresspr to derive the eigenvalues and eigenvectors. Lemma 3.4 shows that both H.hh and freaspr possess the same set of eigenvectors, but the corresponding eigenvalues (A, 1/4 + o. 1/4 4/4 and (A/ are significantly different.
While Lemma 3.4 conveys discrepancies between each eigenvalues of Shampoo and CASPR, we measure overall approximation errors of Shampoo and CASPR with respect to H, in Figure I where CASPR demonstrates better approximation quality than Shampoo.
4 CONVERGENCE GUARANTEES
We set up our convergence analysis in the online convex optimization (OCO) framework, where the OCO learner makes a prediction We and receives a loss fr(W)), where fi is convex, and gradient 9t Vw (WW), which is used make the next prediction such that the regret minimized: (W Wr) E f((W)) - f(W)). Ry measures the closeness of the objective value of the predictions with the optimal comparator W* arg miny ΣT, f:WW.. There are online- Bianchi et al., (004) which can serve as convergence rate
to-batch conversion techniques guarantees in stochastic convex optimization. For instance, regret R in OCO is equivalent to a convergence rate R T in stochastic convex optimization. The following theorem establishes a regret upper bound CASPR algorithm.

---

Published as a conference paper at ICLR 2024
Query Key
Value
Dense
1.00
Shampoo- 4,4 o
Layer:1 Error
0.75
(1, "1, 112")))))))""))"
0.50
diagonal-Adagaad o
column stat stat 1
0.00
row stat 0
1.00
Layer4 Error
0.50
0.25
0.00
Kerations(x1000)
(terations(11000)
kerations(x1000)
Iterations(x1000)
Figure I: Error plots in Transformer: mine |cff, H| (which gives the best approxima- tion enror up to a scale c at iteration-t). Here L = E, G.G, R E, G, Gt, where G. E R mxx denotes the downsampled gradient (m = n = 128) of the fully-connected layer, and H, = E, n.9 where 9r vec(G.) IR mmn We note that diagonal-Adage tions have high relative error ~ 1, while the combination of axes statistics (I., R) such as Shampoo, CASPR, approximate full-matrix statistic H, comparably well during training of 6 layer Trans- former. Here, each row corresponds to a layer/encoder block which constitutes a self-attention sub-layer followed by MLP sub-layer. We picked the Query, Key, Value transformations from self-attention sub-layer, while the Dense transformation is picked from the MLP sub-layer. More details of this experiment are in Appendix A.2
column-stat, and row-stat approxima-
Theorem 4.1 (Regret upper bound of CASPR (Algorithm D.. Given that the loss functions ft. Rmxn
R. Vt E [T] are convex and G- Gipichiitzz in 22-norm i.e. ([V S G, W E Rmmk.
Algorithm 7 iincurs the following regret
Σ ((W))-f(WW))) S Dr
((I7
@ In In @ Rr 1/111//22 1 -2
)
S Drr
(14/4 @ R1 R1/4) O(T))
when 17 D/v2r, where r max, rank(G.)) D= | WF
The proof is given in Appendix A.3
The above theorem establishes a tighter regret bound for CASPR (first inequality) than the regret upper bound of Shampoo (second inequality). Furthermore, the upperbound (VVT) conveys that Algorithm is asymptotically optimal in the online convex optimization setting. We also discuss existing reductions from regret bound guarantees to non-convex optimization in Appendix A.6
EXPERIMENTAL RESULTS
We conduct several experiments to evaluate the performance of CASPR on graph neural networks (Battaglia et al. 2018)am Transformers (Vaswani et al.) 2017). We compare CASPR with adaptive regularization methods such as Adam (Kingma & Ba 2014) and Shampoo (Gupta et al. 2018)) Anil et al. 2020). Specifically, we employ blocking Shampoo Anil et al. 2020) and similarly in CASPR, where large parameters are partitioned into blocks to precondition them separately. CASPR has a computational and memory complexity of O(m n² and O(m² n?) respectively, which is the same as Shampoo, while diagonal preconditioners have time and memory complexity of ((m)).
For Shampoo and CASPR, inverse operations are performed once every 20 iterations, and the inverted statistics 1,41/4 and R 1/4 are reused to amortize the cost of the inverse operations. The inverse fourth root operations computed via iterative routine in Anil al. (2020) which uses coupled Newton algorithm (Guo & Higham 2006); [annazzo] (006). Given the wide applicability of Shampoo in several benchmarks and existing comparisons with Fisher based methods in Anil et al.
7

---

Published as a conference paper at ICLR 2024
(020)). make more thorough comparison with Shampoo. However, the possibility of adapting our method to approximate the Fisher matrix could be a fruitful avenue for future research.
To ensure fair comparison with the aforementioned methods, averages into Line 4 of Algorithm I aas follows: L == L + (1 G ft == 3 R + (1 G2 G. For more detailed insight into the Shampoo and CASPR settings used in our experiments, Appendix [A.4] We conduct experiments using the JAX (Bradbury et al.) (0118) and Flax (Heek et al., 2023) frameworks. The CASPR code is adapted from the Optax (Babuschkin et al.)
incorporate exponential moving
2020) implementation of Shampoo (Anil et al.),, which is JAX implementation.
5.1
GRAPH NEURAL NETWORK TRAINING ON OGBG-MOLPCBA
Our focus now shifts to evaluating the generalization performance of CASPR on a graph neural network (GNN) (Battaglia et al.] 2018). The GNN is trained on the OGBG-molpcha dataset (Hu et al., 2020 using init2winit (Gilmer et al. 2023). The GNN model, implemented using the Jraph library @ et al. 2020), constitutes I8 fully fully-connected layers. These layers collectively account for 3.5M parameters with the largest layer having dimensions of 1024 x 256. The training data consists of 350,343 graphs and the test set has 43,793 graphs. The nodes are represented 9-dimensional feature vector and edges are represented as 3-dimensional feature vector and the task is to predict 128 binary labels each corresponding
biological activity.
Our training process uses binary cross entropy loss for each class among the 128 with a batch size of 512. We utilize specific learning rate schedule, involving linear warmup followed by cosine decay. In this benchmark, we compare CASPR against Shampoo and Adam W. We use random search with upto 300 hyperparameters, where we search over weight decay, learning rate and momentum parameter for all the algorithms. We fix the dropout to 0.1. The test average precision are depicted in F which, demonstrates that CASPR shows relatively 1.4% better test average precision than Shampoo and 7% better test average precision than AdamW. W. While CASPR and Shampoo are run for 60,000 iterations and take almost the same amount of time (Figure Adam is calibrated to run for 720000 iterations such that training finishes in the same amount of time as CASPR or Shampoo. Our data- mata-parallel training runs required 4 TPUv2s, however, the walltime in Figure is computed on TPU v4s to utilize more latest optimizations that TPU v2s don't offer.
CASTA d9 75% 4h 11m 44s
0.30
Shampoo 60.42% di
Best Best Validation Accuracy
Adsmmi 201e ch 21
Toa Average Drondao
0.69
0.26
0.69
0.20
0.67
0.20
CASPR 0,250 2h35m
0,286 2132m
0.60
AdauW-0 20.3m
0.10
10
20
40
0.65
[terations (e1 (ik)
20
(teration(x1100)
Figure 3: GNN on OGBG-molpeba dataset: CASPR demonstrates better test average preci- sion (0.29) compared with Shampoo (0.286) and AdamW (0.271). Here Shampoo and CASPR are run for 60,000 steps, however, Adam is run for 72,000 steps, result, AdamW takes the same amount of time Shampoo and CASPR. Here, Shampoo and CASPR take approximately same amount of time, while CASPR gives a better test AP.
Figure/ 2: Transformer on Parts of Speech dataset: best validation accuracies reached by CASPR (69.76%), Shampoo (69.45%) AdamW (68.84%). Shampoo and CASPR is run for 18,750 iterations, but AdamW is run for 66,180 iterations so that it takes the same amount of time as CASPR and Shampoo. Shampoo and CASPR take almost the same amount of time, while CASPR gives a much better accuracy.
5.2 TRANSFORMER ON A UNIVERSAL DEPENDENCIES DATASET
To further verify our generalization performance by using a transformer encoder model (Vaswani et al, 2017) with around 16.6M parameters on a
tagging dataset from universal dependencies (Nivre et al. 2020) with about 202,989 tokens. The largest layer of our model is of size 256 x 2048 and a depth of 6 transformer encoder blocks where each block comprises a self-attention mechanism with 8 attention heads and MIP with hidden dimension 2048. We train our model for
8

---

Published as a conference paper at ICLR 2024
18750 iterations with batch size of 256, and 48 randomly selected hyperparameter configurations, where we search over weight decay, learning rate and momentum hyperparameters, while fixing the dropout to 0.3, and pick the which gave best validation error. We use inverse square root learning rate schedule. In this benchmark, Shampoo reports a best validation accuracy percentage of 69.45%, and AdamW gives validation accuracy of 68.84%, whereas CASPR reports 69. 69.76% giving relative improvement of 1 1% in eтrоr percentage (100-Accuracy) over Shampoo and 3% improvement over AdamW. Figure shows validation performance with similar observations. We run our Adam W baseline for 3.53 times more iterations so that it spends the same training time as Shampoo and CASPR. However, Shampoo and CASPR take about the same training time while they were run for the same number of iterations 18,750. We use one Nvidia A100 gpu for this benchmark.
5.3 AUTO-REGRESSIVE LANGUAGE MODELING
We train dncoder-only GLU based Transformer model (Saazeer 2020: Vaswani al. 2017) with 234 million parameters on C4 Dataset (Raffel et al., 2020) with varying batch size (number of sequences per iteration) and sequence length 1024. In Figure we showcase two experiments a) 14 million parameter model with model dimension 256, hidden dimension 682 and transformer layers trained with batch size of 8192 sequences with sequence length 1024 for 20,000 iterations totalling 167 billion tokens, and b) 234 million parameter model with model dimension - 1024, hidden dimension 2730 and 16 transformer layers trained with batch size of 256 sequences with sequences length 1024 for 160,000 iterations totalling ~ 42 billion tokens. Training involved 16 TPU v3s for 234M model with 256 batch size and 64 TPU v3s for 14M model with 8192 batch size, using paxml software (Google). In Figure g we compare Shampoo and CASPR's performance where CASPR demonstrates better loss compared to Shampoo. We used linear warmup with cosine decay to zero learning rate for both the optimizers. Given the superior performance of Shampoo and CASPR against Adam in OGBG-molpeba and universal dependencies benchmarks, and due to expensive nature of the language modeling experiments, we omitted Adam from our comparison.
3,90
2.90
3.05
CASPR 3,518
Slampgo
2.50
GASPR 2 784
Shatppoo 2.8IZ
3.00
log persteatt
3755
2.80
3/0
2.84
3.85
2.82
3.55
2.80
3.5000
1950
2000
30000
12005 10000 160CO 18500 20020
90000
100000
110000 120000
L30000 540000 150000 100000
(a)
(b)
Figure 4: GLU based decoder-only Transformer model trained on C4 dataset with varying batch size (number of sequences), number of parameters and total number of steps a) 8192 batch size and 14M parameters for 20,000 steps b) 256 batch size and 234M parameters for 160,000 steps. In both the training runs we see CASPR outperforming Shampoo in log perplexity.
6 CONCLUSIONS AND FUTURE WORK
In this paper, we have introduced the CASPR preconditioner for deep neural network optimization that considers the inherent matrix structure of fully-connected layer parameters to construct dedicated preconditioners for each dimension or axis of the parameter. CASPR uses a Kronecker-sum inspired combination to approximate the full-matrix statistic more accurately than existing Kronecker-product based method Shampoo. We establish tighter regret bound guarantees within the online convex optimization framework compared to Shampoo, and demonstrate better training and generalization performance in our deep learning experiments. As future work, CASPR can be adapted to approximate the Fisher matrix in Natural Gradient Descent (Amari 1998).
REFERENCES
Naman Agarwal, Brian Bullins, Xinyi Chen, Elad Hazan, Karan Singh, Cyril Zhang, and Yi Zhang. Efficient full-matrix adaptive regularization. In International Conference on Machine Learning,

---

Published as a conference paper at ICLR 2024
pp. 102-110. PMLR, 2019.
Shun-Ichi Amari. Natural gradient works efficiently in learning. Neural computation, 10(2)251-277. 1998.
Rohan Anil, Vineet Gupta, and Anudhyan Lottes,
James Boral. optax
shampoo.
scalable scalable_shampeeeeeeebeeeeeeeee_
Rohan Anil, Vineet Gupta, Tomer Koren, Kevin Regan, and Yoram Singer. Scalable second order optimization for deep learning, arXiv preprint arXiv:2002.090/1 2020.
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 arXiv:220..05310. 2022.
Igor Babuschkin, Kate Baumli, Alison Bell, Surya Bhupatiraju, Jake Bruce, Peter Buchlovsky, David Budden, Trevor Cai, Aidan Clark, Ivo Danihelka, Antoine Dedieu, Claudio Fantacci, Jonathan Godwin, Chris Jones, Ross Hemsley, Tom Hennigan, Matteo Hessel, Shaobo Hou, Steven Kapturowski, Thomas Keck, lurii Kemaev, Michael King, Markus Kunesch, Lena Martens, Hamza Merzic, Vladimir Mikulik, Tamara Norman, George Papamakarios, John Quan, Roman Ring, Francisco Ruiz, Alvaro Sanchez, Rosalia Schneider, Eren Sezener, Stephen Spencer, Srivatsan Srinivasan, Wojciech Stokowiec, Luyu Wang, Guangyao Zhou, and Fabio Viola. The DeepMind JAX Ecosystem, 2020. URL /ttt:://githb.... com deepmind
Peter Battaglia, Jessica B Hamrick, Victor Bapst, Alvaro Sanchez-Gonzalez, Vinicius Zambaldi, Mateusz Malinowski, Andrea Tacchetti, David Raposo, Adam Santoro, Ryan Faulkner, al. Relational inductive biases, deep learning, and graph networks. arXiv preprint 1rXiv:1806.01261, 2018.
Rajendra Bhatia. Positive definite matrices. In Positive Definite Matrices. Princeton university press, 2009.
James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman-Milne, and Qiao Zhang. JAX: composable transformations of Python+NumPy programs, 2018. URL
http: / /githbb...... com//oogle/jaax
Charles G Broyden. Quasi-Newton methods and their application to function minimisation. Mathe- matics of Computation, 2 21(99):368--3) -383668888111113 1967.
Nicolo Desa-Bianchi, Alex Conconi, and Claudio Gentile. On the generalization ability of on-line learning algorithms. IEEE Transactions on Information Theory, 50(9):2050-205 2004.
Ashok Cutkosky, Harsh Mehta, and Francesco Orabona. Optimal stochastic non-smooth non-convex optimization through online-to-non-convex conversion. arXiv preprint 07X11:2302.03777, 2023.
John Duchi, Elad Hazan, and Yoram Singer. Adaptive subgradient methods for online learning and stochastic optimization. Journal of machine learning research, 12(7), 2011.
Roger Fletcher. A approach variable metric algorithms. The computer journal, 13(3):3173322, 1970.
Justin M. Gilmer, George F. Dahl, Zachary Nado, Priya Kasimbeg, and Sourabh Medapati. init2winit: a jax codebase for initialization, optimization, and tuning research, 2023. URL http:://ttth//t Com/ www.goooooingiiiee
Jonathan Godwin Thomas Keck*, Peter Battaglia, Victor Bapst. Thomas Kipf, Yujia Li, Kimberly Stachenfeld, Petar Veličković, and Alvaro Sanchez-Gonzalezz Jraph: A library for graph neural networks jax., 2020. URL http:/ttt::/gttthaatt wttp:/tww.cocccccoooomoacoeecc com deepmi nd/ jraph
Donald Goldfarb. A family of variable-metric methods derived by variational means. Mathematics of computation, 24(109):23-26, 1970.
10

---

Published as a conference paper at ICLR 2024
Google. paxml. httpes://githbb..co///o/////////
Chun-Hua Guo and Nicholas J Higham. A schur-newton method for the matrix boldmath th root and its inverse. SIAM Journal on Matrix Analysis and Applications, 28(3):788 804, 2006.
Vineet Gupta, Tomer Koren, and Yoram Singer. A unified approach to adaptive regularization in online and stochastic optimization. arXiv preprint arXiv: 1i0.006666.6999 2017.
Vineet Gupta, Tomer Koren, and Yoram Singer. Shampoo: Preconditioned stochastic tensor optimiza- tion. In International Conference on Machine Learning, pp. 1842-1850. PMLR, 2018.
Elad Hazan et al. Introduction to online convex optimization. Foundations and Trends@ in Optimiza- tion, 2(3-4):157-325, 2016.
Jonathan Heek, Anselm Levskaya, Avital Oliver, Marvin Ritter, Bertrand Rondepierre, Andreas Steiner, and Mare van Zee. Flax: A neural network library and ecosystem for JAX, 2023. URL http://github. com www.oooogo/e.....cooooooo
Weihua Hu, Matthias Fey, Marinka Zitnik, Yuxiao Dong. Hongyu Ren, Bowen Liu, Michele Catasta, and Jure Leskovec. Open graph benchmark: Datasets for machine learning on graphs. In Hugo Larochelle, Marc Aurelio Ranzato, Raia Hadsell, Maria-Florina Balcan, and Hsuan-Tiees Lin (eds.), Advances in Neural Information Processing Systems 33: Annual Conference on Neural Information Processing Systems 2020, NeurIPS 2020, December 6-12, 2020, virtual, 2020, URL
neuriis..urp.pp.. cC wttppppppppp//////aaaaoopppp wtw.://fcccc...o--aaataaaae .htll
Bruno lannazzo. On the newton method for the matrix th root. SIAM journal on matrix analysis and applications, 28(2):503-522 2006.
Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv: 14/2.6980, 2014.
Dong C Liu and Jorge Nocedal. On the limited memory bfgs method for large scale optimization. Mathematical programming, 45(1-3):503-528, 1989.
Karl Löwner. Über monotone matrixfunktionen. Mathematische Zeitschrift, 38(1):177-216, 1934.
James Martens. New insights and perspectives the natural gradient method. The Journal of Machine Learning Research, 21(1):5766-581, 2020. 2020.
James Martens and Roger Grosse. Optimizing neural networks with kronecker-factored approximate curvature. In International conference on machine learning, pp. 2408-2417. PMLR, 2015.
Joakim Nivre,
de Marneffe, Filip Ginter, Jan Hajič, Christopher D Manning, Sampo Pyysalo, Sebastian Schuster, Francis Tyers, and Daniel Zeman. Universal dependencies v2: An evergrowing multilingual treebank collection. arXiv preprint arXiv:2004. 106443 2020.
Kazuki Osawa, Satoki Ishikawa, Rio Yokota, Shigang Li, and Torsten Hoefler. Asdl: A unified interface for gradient preconditioning in pytorch. arXiv preprint arXiv:2305.0440 2023.
Kaare Brandt Petersen, Michael Syskind Pedersen, et al. The matrix cookbook. Technical University of Denmark, 7(15):510, 2008
Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J Liu. Exploring the limits of transfer learning with a unified text-to-text transformer. The Journal of Machine Learning Research, 21(1):5485-555 -5551, 2020.
Yi Ren and Donald Goldfarb. Tensor normal training for deep learning models. Advances Neural Information Processing Systems, 0--20000-26055 2021.
Yi Ren, Achraf Bahamou, and Donald Goldfarb. Kronecker-factored quasi-newton methods for deep learning. arXiv preprint arXiv:2102.067337 2021.
Shai Shalev-Shwartz et al. Online learning and online convex optimization. Foundations and Trends@ in Machine Learning, 1(2):107-194, 2012.
11

---

Published as a conference paper at ICLR 2024
David F Shanno. Conditioning of quasi-Newton methods for function minimization. Mathematics of computation, 24(111)x477666, 1970.
Noam Shazeer. Glu variants improve transformer. arXiv preprint arXiv:2002.052022 2020.
T. Tieleman and G. Hinton. Lecture 6.5 4.5-rmsprop: Divide the gradient by running average of its recent magnitude. coursera: Neural networks for machine learning. 2012.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017.
12

---

Published as a conference paper at ICLR 2024
A APPENDIX
A.1 PROOFS
A.1.1 DERIVATION OF COLUMN PRECONDITIONER
Lemma A.1. Consider the sparsity constraint SLØI [L @ In L 0 E [ in then the optimal solution L for the subproblem (
(rewritten below) as follows:
1 EG,G..
= m @ I,
9.9s |E
8=1
8=1
Proof. Let P is permutation matrix such that /2 (ve (ec((") for any G E |R then:
(L 2 I,) vec(G) vcc(LG)
vec((G" ))
P vec((G" L)
P (In @ L) Ve
= PP @
Thus L 1= & L)P is block diagonaa matrix upto a permutation, substituting this the objective function in the lemma gives:
L arg min min & In
Σ 9.9. l|F
arg miin @ In) 1T) Σ
Pg.g. P ll
= arg min @ Σ 9.0. l|F
arg min min bllckk/i///okk- L) Σ g.9 ll
(G)) [ G.
(Using Lemma 3.1 and 9. = ved (ee(GG) ))
G.G.T
A.1.2 PROPERTIES OF KRONECKER PRODUCT, KRONECKER SUM AND LOEWNER ORDER
Lemma A.2 (Eigendecomposition of Kronecker-Sumi Let A E mm B Raxn be two positive definite matrices with eigenpairs {(),, u:) [m]) and ((o,, v;): j [n]], where Asi X and Ejj = oj. then the eigenpairs of the matrix A I, + I,, T..&B. Bre [(1+to,, (&0;) :i ie [m].j E [n]]
13

---

Published as a conference paper at ICLR 2024
Proof. While this is standard property, give its proof here for completeness. First, show that 24 @ Uy is an eigenvector of A @ In + Im @ B:
(A 2 In + Im @ B)u & uy = (A @ In + Im @ B) ved(uee))
(vec(Au)u)) + (ed B))
= (ou() + to,uoo)) (since Au, = X,u; and Bug = oz0z)
(A: + σz) vec(u.e,)
(A, + oj )u, 8 U/
Thus all the mn eigenpairs of A In + Im & B are {(>; + oj 8 U2) : E [m] E [n]]
Lemma A.3 (Cigendecomppsstionn of Kronecker-Prodde
). Lel E Rmxm, B E Raxe be two
positive definite matrices with eigenpairs {(),, ",,) i E [m]] and ((a), uy) :j [m]], where Ag, = A and Ey = J. then the eigenpairs of the matrix @ B are {(,,o u, @ 02) : [m].] E [n]}
Lemma A.4 (Geometric mean of commutable matrices (Gupta et al. 2018)). Ler 0 3 A 3 Bı
and O A2 B2 and further assume that A₁ A2A1, B B, B2B then A 1/2 A B\/2 /2 B²
Lemma A.5 (standard properties of Loewner order). Let Ai A 0ad Bi 2 B2 > then
(a) A @ Bi > A2 @ B3
(b) A コ Aズ²² and B 1 B,
(c) A = / and 11 = 1//2
A.1.3 CONNECTIONS MADE AMONG CASPR, SHIMMPO AND FULL-MATRIX ADAGRAD
Proof of Lemma 3.4 Using Lemma A.2 and A.3 for L, 1--1/4 @ In + Im & R2 1/4 and L,/2 @ R2/2
respectively, followed by the fact that eigenvalues of A" are p-th powers of eigenvalues of A for any integer p.
Proof of Lemma 3.2 In this proof, we drop the subscripts for L, and RS in the lemma statement.
Let L = UAU and R = VΣVT, given that A > 0 and Σ > 0 (since L > eln > 0). Using Lemma A.2 the eigenpairs of the general combination X.essr (p) in the lemma are {((((2-1/44)
1//4///222// 1/4p) //)/ap, us 8 02) i [m]..] E [n]]. Note that the eigenvectors are independent of p, so
eigenpairs of limp Xrespr (p) are ((1 1/4p + to, - /2) u4 & U2) E [m],
[n]}. We now analyze the individual limits for each eigenvalue as follows:
lim ((2- + +o //2 p0000
lim 20 1/40 1/ 1/2)
p0000
lim Jog((( 1/4p to, 1/4p /2) p-000
limp- log((A,
As numerator and denominator in the exponent of the last equality evaluate to zero in the limit, L' 'ôpptal''s 's rule can be used, which gives the following:
2(log(>,
)2,
114r
1/4p
+og((;
1/4 10, 1/40)
10
lim log((A, p-000
1/4p + t0 + ///44//22//11//
= lim
9-000
A, -1/4p -o,
1/4p
= logg(A, 1/4) + log(o, (Since lim X 1/44 -1/4p = 1)
:0
p000
applying exponential function both sides gives:
lim ((,-1) 1/4p + +0, 1/4) -A,10)) -1/1
p-030
By Lemma A.3 we know that the ((),a,) 114 (1,4,)_1/44_ un 8 v2)) : i [m].j E [n] are eigenpairs of the Shampoo preconditioner Xehh T.q 1--1/4 R 111
14

---

Published as a conference paper at ICLR 2024
Proof of Lemma 3.3
Showing ela E =9 1 (Lt) @ Ln. We first upperbound individual rank-1 terms gg" in Loewner order by proceeding in the Gupta et al. (proff of Lemma 9). Using the SVD of g vec(G) = vecl Vec(E )i,;)l)) Ei=1 (21 @ 04), where r' denotes the rank of the matrix G
gives:
9g" 0164 x)) 0.64 (1,)) 3r Σ (u,u, 0,44)
-1 Σ (ueuf @ (.) r(G @ ln)
(4)
rmm & G)
(S)
where, the first 3 follows from Ei wi w²)² 3 rr w:w?) and 4 use yada x In.
u.u. mm respectively.
Summing and 5 for gradients gr. across iterations E [T] gives, eld + E._1 9
( G.G. + +m 8 In = La @ In and ela + Σ._ Egag
Im @ (E, Gf + ten) = Im 0 Rre
respectively.
Showing (cla +
19.9.) (i/ @ In + Im 2 R 1//44
1/4
Expanding the RHS gives
(I, & In + Im & R2 1/22 1/2
In Im @ R, + L, 1/4 @ R 1/4
(6)
Now we use 1/2 @ In 3
ela
Σe +l++E=m [go9%
112
since
X1/2 is a monotone operator
and A B + +
-1/2
(t, 1/4 s Jn + mm 8 R 1/422-202 32
cTa Σ
Is9
+ 2. I,
1,4 e R 1/4
s=1
Note that A = 1/2 8 In and B = Im 8
1/2
1/2 commute. Using Lemma A.4 gives AB =
1/2
L, R 1/4 1 Ell Σ'= 1
Using this in 66)
gives the first inequality.
Showing (I, 1/4 @ In + Im & R 1/4 )2 /3 L, 1/4 Rt,"
1/4,
Using (A 2² = (A + B)² 4AB = 0, since AB BA, gives the second inequality.
A.2 SETUP OF APPROXIMATION EXPERIMENT
In Figure T we used the ritentiio-bbsed transformer (Vaswani et al. 2017) with 16.6M parameters from Section 5.2 where the details of the transformer are mentioned. Due to memory constraint we eached the gradients once ever 50 iterations when using Adam optimizer, and the cached gradients are used to construct the plot for each approximation in Figure [ The Query, Key, Value parameters were of size 256 8 x 32 (since 8 attention heads), which are flattened to 256 256 (putting together all the attention heads' parameters), Dense parameter from MLP sub-layer is of size 256 x 2048. Since constructing Σ, 9t99. can require large amounts of memory, we downsampled Query, Value, Key and Dense parameters to 128 128 parameter sizes. While Figure row stat and diagonal-Adage to be nearer to 1, close comparison among them is given in Figure 5
[ mentions errors of column stat,
A.3 REGRET BOUND ANALYSIS OF CASPR
We restate the theorem upper bounding regret here:
15

---

Published as a conference paper at ICLR 2024
Query
Key
Value
Dense
Layer-4 log(1-Erro)) Layer-1 log(1-Erro)
Shampoo oR
-2
CASPR ((1-1M + in "2)*
-4
diagonal-Adagrad daal)
-6
row stat In R
-2
-4
-6
-8
10 (terationsx11000)
kerations(x1000)
kerations(x1000)
Rerations(x1000)
Figure 5: Plot of log(1 minelleHt H.ll /|l in contrast to Figure 1 the column stat and row stat approximate the full-matrix statistic hetter than the diagonal-Adagrad approximation.
Here notice that
Theorem A.6 (Regret upper bound of CASPR (Algorithm Vt [T] are convex, Algorithm 7 gives the following regret
DT. Given that the loss functions fü.
S.W,) - S(W")) S Vr ((177)
@ In + Im @ RT ) -2)
$ Dr [1 R2(44)
when 1j D/v2r, where max, rank D max I[, W'
Proof for Theorem 4.1 Theorem A.6 We first give a general upperbound to regret Rr by proceeding as in (Hazan al. 2016
Lemma A.7 (A general regret upperbound for adaptive regularization Hazan et al. 2016 Let the
regret RT 1 f(W,) f(W*), ) then
RT < 211
1 W1 w0 11x
1
We w²) Txx ---1-x--1 x, - w²) + " Σ
9t X19:
T
2n 1=2
Proof. By the parameter update rule from A [we know that wx+1 = w) 7,g, where
X, =(L, 1/4 & I Im @ R,
2 1/2)2 Subtracting w" from both sides gives:
Wi+1 w² = wt w² - .1.t
|10211 w/ w0 |x, = [1(w; - w²) - 17 N1
|1(w; - w²)) w²) x, + n²99. Xigi 2ng. (wi - -w")
2n (I( w² w1 || w w=(lxx-1) + "1 20 X.9:
1
De (we w²)
16

---

Published as a conference paper at ICLR 2024
where, | VuT Au is generalized norm defined for positive definite A. Summing the above equation overt E {l,..., ,T gives the following:
T
T
7 T
Dt (we - w²) = 2n 1=1
(w - w² )xx) - 110+11 - 101l= xxx-,
+
Σ
Xt9t
2 t=1
1 (llw= w')ll,, l1(w2 w²) 001x2
21
+ l1(u2 " "='" >|(x;1 1(w3 w²* )|x + 11 [[(r-="")l? w -1((u++1 w²))l>>"
+ 2 gt T Xz9:
Σ
1
-x-1) (ll(1 - w²)) x- + (w2 - - w²)+
21
(wT -w") )² '(XT' -Xi-1 =w)
T
- -I[[T++ wz)"l ()(x" + 17
Σ
9i X1gt
(combining 2nd and 3rd terms, 4th and 5th, ...)
2
1 (Ilw1 w= "1x
21
+ (w2 w²) " T(X²! 1Xxx1_X,- -X-') w²) ... + (wT w²) w') Xy-1)(wr- w"))
1 T X19:
2 9t
The last incquality is because -I[(W++1- - following:
S 0. Note that convexity of fo gives the
F((W)- S tr( ( W*)) = dt (w+ - w
where (ee(W;) and g: = vec(G.). Using the upperbound on E, g. (wi w²) derived earlier and the above gives the following:
Rr Σ ((w.)-S(W'))
<
1 (1|w0
w0 w
+
1
X w²)
2n
20
+ シタ Σ 9t Xigr
T
1 |w 2m
x +
2t=1 (wi - - xx - -w"):
To upperbound this term first establish Loewner order inverse of CASPR preconditioners for consecutive iterations as follows:
Lemma A.8. Preconditioner inverses follow the Loewner order: 0 X .... 1 3X7
Proof. We can expand the CASPR preconditioner as follows:
X, (L, In + Im 2 R 2/22222
(I, 1/2 Ø Im 3 R 1/2 2L,
1/4 @ R. 1/44///
Note that ř, = i + G.G.T > أد and R = R
+ G. G, ł L1-1s then by Lemma A.5 n11/2 Furthermore, since i,-1/2 T
(a,b,c), / 8) Jn x 1 T, and Im 8 n 1 Im
17

---

Published as a conference paper at ICLR 2024
and m. ft,11/2 commutable, by Lemma A.4 we get L, /
R 1/4
- R
Putting together the equations, - In 3 L 1/2 & In. Im @ R2 1/2
3 Im @ Rī and
L, 1/4 @ R 1/4 1/4 @ R 1/4 with 7 gives the following
X₁ 3 X
X, m x (By Lemma A.3 b))
The above lemma implies that difference between inverses of consecutive preconditioners X, x positive semidefinite for all {l..... ,. thus the individual term
(w:- w -X," (ww, w²) I I1( = tr -X,'
Note that trace operator a norm for positive semidefinite matrices. Summing over all E ( gives
T
Ti = 2n
W1 - w0 uIll²
(w w²) "x)T(x,' - xw - w²)
2n
(D² tr(xTt) + D - + D2 1
X1
-x7=1)
2n
D2 tx((7')
2m
D² ir(( ur(((LT])) 1/4 & In + Im @ )/2)
where D max (1 77 lwu, w" 112 m1. Ty Wr W
To =2 Xege:
To upperbound this term, the Loewner order X7 (1I, 2 In + Im @ //232 =
(ela -. ) which gives the following bound:
T2
7
9r T
(cJa
Σ
Dags T -1/
)
2
Vrn
9t (erla + Σ In 11/2
(8)
2
= Lemma A.9 (Gupta et al.] (2017)) Let 91, 92.............................. gr be a sequence of vectors and let M, funcctffunciion function & over positive definite matrices, let
Ei E for < T. Given a real-valued
He arg m (tr(M, 6(H))
for S T, then,
T
T 1ge l9cll + ((H)) o(Ho)
H.
[=1
We proceed as in Gupta et al. (2018) by setting ((H) = tr(H) + retr(HH and solving the following reparameterization of the subproblem in the above lemma in variable X H-1:
H. = (arg minx= - (tr(M,X) 6(x 6(x-1)))-
Using that (x(tt(A 1)) - 2 X A 1/2 for all A > 0 and the convexity
of the above problem, gives:
1/2
H rela Σ 9-9
18

---

Published as a conference paper at ICLR 2024
By Lemma A.9 and (8)
Vrn
T2 S
9i (arla +
T
Σ
Io99 T
( 1129
2
t=1
#=1
vrn
Σ 1=1
St (erla +
Σ
959s T ) ) gt + tr(((eeaa Σ 9:9: ) 1/2) +reca((rela Σ 9o92 ) 1/2)
==1
Vrm
T
tr
919t (erls + Σ gogs )-1/2 1/2 + tr((rela Σ gigt )1/)) (rela Σ
)
(adding terms 1&3)
= -Vintr((rela
Σ gig. 1/31)
<rntr((F7) (by Lemma 3.3
Now setting 1| = 2/ gives
Pr $ Ti - T2 v2rD =V (( @ In + mm @ Ro //2 -2) )
We also know by Lemma 3.3 and monotonicity tr()) operator, that
1/44 tr ((i In Im & RT 122 >) S (1 1241)
thus proving second equality of Theorem 4.1 Given that Ss are G-Lipschitz functions, G.G.T 3 I and G, Ge 3 G2I,, thus LT TG Im Rr I TGG I, and tr((t//t) I m/GT1/4, 11/4 tr( tr(Re/) S /// Substituting in the above gives tr(L @ O(VT).
A.4 HYPERPARAMETER SEARCH SPACES AND OPTIMIZER CONFIGURATIONS
Here we mention the search spaces used for hyperparameter search for graph neural network bench- mark and transformer benchmark
Hyperparameter
Search Space
Hyperparameter Search Space
10-5 0.1] (0.9, 0.999 0.999 [0.001, 1]
[10-5 0.1
B1
B
0.9. 0.999) 0.999 [10-4,1]
82 weight decay
B2
weight decay
Table I: Search space for graph neural net- work benchmark in Section 5.1 é set to 10-6 for Shampoo and CASPR (which are the default values)
Table 2: Search space for Transformer bench-
mark in Section 5.2 is set to 10-8 for AdamW and 10-6 for Shampoo and CASPR (which are the default values)
A.5 POTENTIAL NEGATIVE SOCIENTAL IMPACT
CASPR tries reduce the number of iterations required to reach desirable accuracy, and con- sequently conserves GPU time. To the best of our knowledge, this poses no negative impact over society.
19

---

Published as a conference paper at ICLR 2024
Hyperparameter
Search Space
Hyperparameter
Search Space
?
(10 10 10-2 10-1)
{10- 10-3 10-2 10-1) {0.99, 0.999]
82 weight decay epsilon
{0.99, 0.999]
82
{0.0, 0.001, 10 100
Table 4: Search space of 234M parameter model 553) (1 is set to 0.9, weight decay is 0.0.
Table 3: Search space for 14M parameter lan-
guage model in Section 5.3 $1 is set to 0.9.
A.6 REGRET BOUND MINIMIZATION FOR NON-CONVEX OPTIMIZATION
The problem minimizing smooth non-convex functions f can be reduced online convex opti- mization (Agarwal et al. 2019) with sequence of objectives of the form
s<(w) = [(w) + cllw wee
where, > and L- Lipschitz smoothness constant. This method introduced in Agarwal al. (2019) can be used to obtain convergence guarantees to reach stationary point of the non-convex objective f.
Recently, more direct reduction regret minimization of linear functions f. was established in (Cutkosky et al., 2023), to obtain stationary point guarantees in non-smooth non-convex optimization. Thus using regret upper bound guarantees, one can establish stationary point guarantees in the non-convex regime, and hence focus the former in the paper.
A.7 COUPLED NEWTON ITERATION FOR INVERSE P-TH RoOT COMPUTATION
For computing inverses p-th roots we use Algorithm in Anil et al. 2020) in Appendix D, originally formulated in Guo & Higham (2066;; lannazzo (206)), which we repeat Algorithm 2 We limit the number of iterations in while loop to be 100, as done in the JAX implementation of Shampoo. Coupled Newton iteration for inverse 4-th roots has a time complexity of (( and memory complexity of O(n²). For our experiments, the inverse pth root operations were performed using Moating-poine format (float32) We determined the damping term cf in Algorithm [ by scaling the largest eigenvalue of (L) with as outlined in Algorithm 2 This scaling ensures that the modified matrix L L + =dmax = maintains é2-condition number not exceeding 1/s. We set 10-6 across all our experiments, which caps the condition number of L 10 which is less than the inverse of machine epsilon for float32. In the case of the 14M parameter language model showcased in Figure 4a. In that instance, experimented with values in the (10-6 10-8), *), as
detailed in Table 3
A.8 STANDARD DEVIATION IN EXPERIMENTS
We mention here the standard deviations averaged across seeds for the Transformer on universal dependencies and OGBG-molpcba benchmarks:
Optimizer
Transformer Universal Dependencies Accuracy
OGBG-molpcca Test mAP
Shampoo
69.45 0.26% 69.76 + 0.14%
0.2843 0.0028
CASPR
0.2873 + 0.0017
Adam W
68.28 - 0.39%
0.2701 + 0.0011
Table 5: Standard deviations for Transformer various benchmarks.
A.9 PRECONDITIONING MULTIDIMENSIONAL PARAMETERS
Our method (as well Shampoo) be applied to any parameter tensor. For example, in experiments with transformers we handle 3D parameters, such as attention layers by flattening the parameter (in the form of higher order tensors) into 2D tensor by merging the dimensions. The
20

---

Published as a conference paper at ICLR 2024
Algorithm 2 A coupled Newton iteration procedure for computing inverse p-th roots of a PSD matrix, with warm start and singular value projection
procedure MaxSV(G)
2: Parameters: > 0, Ratep
3: IR", where G R
4: = 0, enror xo, À = 0
5: while Nstep and error > 8
6: û e/|lu|l
7: v = Go
8: Aodd X; X 6
9: error = | Aotal: = i+1
return
11: end procedure
13: procedure Coupledlteration(G.. E N, X (optional))
Parameters:
15: Outputs: G-1/p
Amaax MaxSV(G)
17: G G+E Amax -1
18: a = 11pp
19: if X is provided then
20: M XPG
21: else
22:
ItP
G
X II
24:
M 2G
25:
while || - 11/00 > e
26:
Mi [1-a)I+a aM
27:
= XM,
28:
M M
29: return X
30: end procedure
21

---

Published as a conference paper at ICLR 2024
attention layers have 3-dimensional parameter (m, which precondition by merging the second and third dimensions to a 2-dimensional parameter (m, n * k)) and then applying CASPR. For higher dimensional parameters, one can apply similar transformation by merging more dimensions.
22


Reviews:
[{"review_id": "7OtTnasfb9", "reviewer": "Reviewer_a8ib", "summary": "In this work, the authors propose a new second-order method inspired by the Shampoo optimizer. The authors show that the proposed method is theoretically grounded and empirically tested.  The authors demonstrate the performance of the proposed method on several large models.", "questions": "See the weakness section", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 4, "strengths": "This is a solid work. Strong theoretical results and empirical results are provided. I do not check the proofs in the appendices.", "weaknesses": "The authors should discuss the time and space complexity of the coupled Newton algorithm used for computing the matrix inverse p-root. \nMoreover, numerically stability of the proposed method should be discussed. \nFor example, how to choose the damping term $\\epsilon$? Note that the coupled Newton algorithm could be sensitive to the choice of the damping term.   \nThe proposed method may not work well in a low numerical precision setting such as float-32 or bfloat-16. The authors should explicitly discuss this point.", "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": "mG6iAmwvCt", "reviewer": "Reviewer_gGzL", "summary": "The paper proposes an alternative to the recently successful Shampoo optimizer that is derived by minimizing the error of the approximation. They show that their method obtains a faster theoretical convergence guarantee in an online learning setting than the original Shampoo optimizer, and also in practice performs better.", "questions": "- On page 3 it says the second-moments are computed per row of the gradient. Is it correct to say that this means it is a neuron-level conditioner? This is something that seems to exist in the context. See for example, Fig 1 \"unit-wise\" preconditioning in https://arxiv.org/pdf/2305.04684.pdf.\n- Are the 300 hyperparameters optimized for each method or only for caspr? What is the target for this hyperparameter tuning? Validation accuracy is the performance reported so is the grid search optimizing that value?\n- Why are there no error bars included? Could the experiments be run on at least 3 seeds? Without a single rerun, the results could be random.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- clear development of the method and comparison to Shampoo\n- theoretical development that exists of Shampoo extended to caspr \n- experiments suggest that caspr improves over shampoo and potentially also other baselines (only compared to Adam(W))", "weaknesses": "- abstract and introduction could be written more clearly, it is only mentioned what is achieved in the abstract and introduction but not how this will be done\n- relatively limited benchmarking to other methods andmissing error bars on the experiments, it seems like a single training run is reported in the figures. This is the major potential weakness I see in the paper and I am worried that the results are not reproducible. See my questions below (Are the figures created from multiple runs averaged? What is the target for tuning hyperparameters if the plots show validation performance?). \n- seems unclear how to extend this method to layers that are not fully-connected\n\n### Minor\n- first page sentence ... \"..to be at one end of this sequence\" confusing", "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": "rKtxmAeiHr", "reviewer": "Reviewer_Vz7y", "summary": "This work proposes an optimization algorithm called CASPR that yields tighter convergence guarantees for stochastic optimization. It is inspired from the commonly known Shampoo optimizer (which is a Kronecker product based optimizer) and uses the Kronecker sum approximation to obtain better convergence rate with similar complexity as in the Shampoo case. The technique is well motivated and provides theoretical/experimental proofs/explanations for all the claims in the paper (convergence guarantees, basic properties of kronecker product/sum and Loewner order).", "questions": "In the evaluation you introduce a momentum for $L$ and $R$ matrices for a fair comparison against the other optimizers, can you please say what is the impact of momentum terms in this case? Have you experimented without momentum?\n\nI would like to point out some typos in the manuscript:\n- at page 4, Lemma 3.1, the Frobenius norm should be squared (the 2nd power is missing from the norm)\n- at page 5, when you rewrite the preconditioner $X_t^{caspr}(1)$ in terms of $X_t^L$ and $X_t^R$, I belive the matrices $L$ and $R$ are missing a tilde, based on the definitions of the $X_t^{L/R}$ terms 4 lines above", "limitations": "", "rating": 8, "confidence": 5, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "1. clearly explains the motivation of a Kronecker sum based precondition, the differences in the updates of CASPR and Shampoo and comparison with diagonal Adagrad statistics in Figure 1 (all in section 3)\n2. provides convergence guarantees\n3. experiments are performed on transformers and graph neural network tasks (commonly known as difficult tasks) against Shampoo and AdamW, the main competitor optimizers in the literature\n4. comprehensive appendix", "weaknesses": "1. the paper does not have experiments on computer vision tasks, such as ResNets, ViTs or other important architectures used for benchmarking in the literature", "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)