Structured Preconditioners in Adaptive Optimization: A Unified Analysis
Authors: Shuo Xie
Tianhao Wang
Sashank J. Reddi
Sanjiv Kumar
Zhiyuan Li
Venue: ICML 2025 poster (2025)
Keywords: Shampoo
adaptive optimization
layerwise
Primary area: optimization

Abstract: We present a novel unified analysis for a broad class of adaptive optimization algorithms with structured (e.g., layerwise, diagonal, and kronecker-factored) preconditioners for both online regret minimization and offline convex optimization. Our analysis not only provides matching rate to several important structured preconditioned algorithms including diagonal AdaGrad, full-matrix AdaGrad, and AdaGrad-Norm, but also gives an improved convergence rate for a one-sided variant of Shampoo over that of original Shampoo. Interestingly, more structured preconditioners (e.g., diagonal Adagrad, AdaGrad-Norm which use less space and compute) are often presented as computationally efficient approximations to full-matrix Adagrad, aiming for improved optimization performance through better approximations. Our unified analysis challenges this prevailing view and reveals, perhaps surprisingly, that more structured preconditioners, despite using less space and computation per step, can outperform their less structured counterparts. To demonstrate this, we show that one-sided Shampoo, which is relatively much cheaper than full-matrix AdaGrad could outperform it both theoretically and experimentally.

Structured Preconditioners in Adaptive Optimization: A Unified Analysis
Shuo Xie Tianhao Wang'
Sashank Reddi 2 Sanjiv Kumar Zhiyuan Li
Abstract
root of d x d matrix where d is the number of parame- ters. Thus, for large-scale settings, the huge computation and memory cost makes it prohibitively expensive. This has inspired works on designing more efficient adaptive optimizers by using a structured preconditioner, such as coordinate-wise adaptivity employed by AdaGrad (Strecter & McMahan, 2010; Duchi al., 2011). Kronecker prod- uct based preconditioner employed by Shampoo (Gupta et al., 2018; Anil et al., 2020) and layerwise adaptivity em- ployed by LARS (You et al., 2017) and LAMB (You et al., 2019)), which all aim to provide computational and mem- ory efficient approximation of full-matrix AdaGrad. These algorithms have been shown to be very effective for gen- eral deep learning settings. is often assumed that a better approximation of full-matrix preconditioner usually results in better optimization convergence, seen with methods like Shampoo (Gupta et al., 201 8). In contrast, precondi- tioners with more structure such as diagonal preconditioner have inferior performance. In this paper, we challenge these prevailing notions by providing theoretical and empirical evidence against them.
We present novel unified analysis for a broad class of adaptive optimization algorithms with structured (e.g.. layerwise, diagonal, and kronecker-Fie (rontecee-factoredd preconditioners for both on- line regret minimization and offline convex opti- mization. Our analysis not only provides match- ing rate to several important structured precondi- tioned algorithms including diagonal AdaGrad, full-matrix AdaGrad, and AdaGrad-Norm, Norm, but also gives an improved convergence rate for a one-sided variant of Shampoo over that of original Shampoo. Interestingly, more structured precon- ditioners (e.g., diagonal Adagrad, AdaGrad-Noos Norm which use less space and compute) are often pre- sented as computationally efficient approxima- tions to full-matrix Adagrad, aiming for improved optimization performance through better approx- imations. Our unified analysis challenges this prevailing view and reveals, perhaps surprisingly, that more structured preconditioners, despite us- ing less space and computation per step, can out- perform their less structured counterparts. To demonstrate this, we show that one-sided Sham- poo, which is relatively much cheaper than full- matrix AdaGrad could outperform it both theoret- ically and experimentally.
Conceptually, one could equate the degree of structure in preconditioner to the ease of its computation and storage. In this view, full-matrix AdaGrad can be considered as the least structured and most expensive preconditioner while MtaGrad-Norm (Ward et al., 2020), which only maintains a scalar and uses the same preconditioning for every direction, is the most structured and least expensive. The conventional wisdom here is that using a less structured preconditioner, which requires more space and compute per step, reduces the number of steps needed for training. Thus, choosing the right structure balances this orade-off between convergence speed and training step cost.
1. Introduction
Adaptive optimization algorithms (Streeter & McMahan, 2010; Duchi et al., 2011; Kingma & Ba, 2014) play a pivotal role in modern machine learning, especially in the expensive training of large foundation models. Within the machine learning community, full-matrix AdaGrad considered as an ideal adaptive preconditioner for fast convergence in terms of number of steps. The computation of full-matrix AdaGrad precondtioner typically involves inverse square
Among these preconditioning methods, Shampoo (Gupta et al., 2018) has gained notable attention for its Kronecker- factored preconditioning approach, which promises im- proved convergence in large-scale optimization tasks (Dahl et al., 2023). Shampoo was originally proposed as compu- tationally efficient surrogate for full-matrix AdaGrad (Duchi et al., 2011). Despite its popularity, existing analyses of Shampoo (Gupta et al., 2018) are limited and do not pro- vide a full justification for its effectiveness. In particular, we argue that the best-known regret bounds for both Sham- poo as well as full-matris AdaGrad are consistently worse
'Toyota Technological Institute Chicago "Google Research. Correspondence to: Shuo Xie <shuox @ edic,,du>, Zhiyuan Li <zhiyuanli @ edu>.
Proceedings of the 42nd International Conference Machine Learning, Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

Structured Preconditioners in Adaptive Optimization: A Unified Analysis
Algorithm
Subalgebra K.
|| 1grrlll
Regret Bound ||1l Ilgg.r
AdaGrad-Norm
c. la for E
[|ella va
vā
Σ 1
T
|
lggll (Strecter & McMahan, 2010)
d
AdaGrad
Diagonal Matrices
||
Σ
gh
| Σ
(Streeter & McMahan, 2010)
i=1
111
[
+
Full-matrix AdaGrad
All Matrices
||1||
Tr( Σ g.gi }+]
| Tr
Σ
g.9.
(Duchi et al., 2011)
One-sided Shampoo
1d xdz @ Lan
a Van
Tr((dR Σ G.G,)*
| Tr
(EC.G))
(Theorem 4.2)
Table 1. Regret bound obtained by Theorem 3.4 for different algorithms. Theorem 4.2 proves the rate specifically for one-sided shampoo and other rows match existing results in literature. H is chosen K. n S In the third column, æ denotes the parameter in vector form and X denotes the parameter in matrix form with æ = vee (x). In the fourth column, gi denotes the gradient in vector form and G, denotes the gradient matrix form with g. vee G..). We omit O(-) in complexity measures and regret bound for convenience.
than a more memory- and computationally efficient vari- ant of these methods like diagonal AdaGrad and AdaGrad- Norm (Streeter & McMahan, 2010; Duchi et al., 2011; Ward et al., 2020). This demonstrates that more structure on the preconditioner may not necessarily restrict its optimization performance. In fact, our unified analysis uncovers an inter- esting finding: we show simpler, more structured variant of Shampoo, called one-sided Shampoo, actually has sub- stantially better regret bound compared to original Shampoo and full- full-matrix Adagrad (Section 3.3).
Our unified analysis provides new regret bound (The- orem 4.2) for a one-sided variant of Shampoo, which is always better than the existing bound for two-sided Shampoo (Gupta et al., 2018) and could be smaller by a multiplicative factor of d, where we assume the matrix- shaped parameter is of size Vd-by- dd This also leads novel convergence rate of one-sided shampoo for stochastic convex optimization (Theorem 4.4) via standard offline-to-online reduction (Levy et al., 2018).
Conceptually, our findings challenge the conventional wisdom that using larger set of preconditioners which require more memory and compute leads to better opti- mization performance in terms of number of steps. In particular, while using larger set of preconditioners re- duces the gradient term I|l in our regret bound, it increases the other term involving the magnitude of the optimal solution by increasing the norm metric; thus, leading to a worse regret bound (Theorem 2.1 and Ta- ble 1). For instance, we demonstrate that one-sided Shampoo can outperform full-matric AdaGrad both theoretically (Section 4.3) and experimentally (Sec- tion 5). This suggest that mned Shampoo is not just computational-efficen surrogate of full-matrix AdaGrad, but could also be fundamentally better, de- pending on the optimization problem.
1.1. Main contributions
In light of the above discussion, we highlight the main contributions of the paper.
We present comprehensive and unified theoretical framework (Theorem 2.1) for adaptive optimization with structured preconditioners, encompassing several popular methods including diagonal AdaGrad, full- matrix AdaGrad, AdaGrad-Norm, and layerwise adap- tive methods for both online convex optimization and stochastic smooth convex optimization. In particular, our analysis integrates and extends key insights from existing work (Gupta et al., 2017). which presents a unified and elegant way (Algorithm 1) to derive the affrrementooeed structured preconditioners, but only allows analysis on a case-by-case basis. To our knowl- edge, this is the first truly unified analysis of a large family of adaptive optimization algorithms.
1.2. Notations
Let Md be the set of all d-by-d matrices, and Sd C Md be the subset of all symmetric matrices. We use Sd to denote the set of positive semii-efimite matrices, and S to denote the set of positive definite matrices. Dd is the set of all d- dimensional diagonal matrices, and D = Dd n n We denote by La the d-by-d identity matrix. For matrices A, B. denote their inner product by (A, B) Tr(A B).
To enable a unified analysis, we identify novel suffi- cient condition, sell-structured preconditioners (Defi- nition 3.1), which overcomes key technical barrier in Gupta et al. (2017); thereby, allowing a unified analy- sis for several important adaptive algorithms. A more detailed discussion is provided in Section 2 and a sum- mary of our results is presented in Table 1.
For any H E S such that H * 0, we denote H H H/Tr HHT)(H.. For H E Sd | = Vx Hr is the
2

---

Structured Preconditioners in Adaptive Optimization: A Unified Analysis
(semi-)nong of R. with respect H. For convex set H C $4, we define
Algorithm 1 Adaptive Regularization Meta-Algorithm
AdaReg (Gupta et al., 2017)
|x=l
sup | H
(L)
Hyperparam: e >> 0, convex set X' S CR², learning rate 7). preconditioners H. C S
(1
For a convex set X' E Rd and any norm |1-11, we define ||al|| = SuPEER ||ж|| and || == suPwex x| For any H > O, the projection of x onto X' with respect to ||. || is defined as 11 (z) == arg ein || a' H
Input: initialization T1, loss functions {L+)--1 -
Mo < ela
for t ,. do
g: V(((),))
Throughout the paper, we consider the factorization d dr n, and denote the corresponding matrix form of a: E Rd by X E [ddxxdx, We denote x (e) and X voe (xx) for conversion between the vector form and the matrix form. Then for a function (ax)) defined on æ E IRd, we extend its definition to matrices by letting ( denote L( L(veo(X) and we will use L(x) and L(X) interchange- ably when the context is clear. We define the gradient as g: V VI(x)) in vector form and G2 (T,,,, in matrix form, so gr = vec Gi) and G. vec (g:).
M, M + 919
H, <- arg arg me M, +
Fi+1 + = (Ie H, gi)
Return #1y..., TT
The above bound is obtained by first applying a standard bound for online mirror descent to get bound in the form of Σt |1ge H. plus the second term on the RHS of (2). Then the choice of He in Algorithm 1 enables the application of FTL. lemma (Kalai & Vempala, 2005) to further bound Σ E/=1 || by the first term on the RHS of (2).
For a matrix X E Rdd we denote its operator norm as X and its Frobenious norm as | For vector
Rdd we denotes its foc norm by ||ae max [1s]] and l2 norm by ||
To proceed from Theorem 2.1, Gupta et al. (2017) relies on crucial assumption that H 3 H, for each Or more generally, for any M > 0, define
2. Background: Unified Adaptive Regularization with Non-Unified Analysis
(M == arg minmmi A.M,HH") H TT (3)
MEH
Then we hope it holds that
Seminal work Gupta al. (2017) presented an adaptive reg- ularization meta-algorithm, AdaReg (Algorithm 1), which can be used to derive various adaptive optimization al- gorithms known at that time in a unified approach. For example, AdaReg becomes full-matrix AdaGrad (Duchi al.. 2011), diagonal AdaGrad (Duchi et al., 2011), and AdaGrad-Norm (Streeter & McMahan, 2010; Ward al., 2020) by choosing the set of preconditioners all PSD matrices, diagonal PSD matrices, and mutipliers of identity matrix respectively (see Table 1). The original AdaReg also allows other choices of potential function t e.g., (-)) = log det(-) for Online Newton Step (Hazan al., 2007), while we are only interested in the case of ((.) , Tr(-) in this work.
(M( = Px(M') for any 0 M 3 M'
(4)
In other words, we need Ps S + S to be operator monotone under the semi-definite ordering. With this as- sumption, a critical step in the derivation of the regret bound in Gupta et al. (2017) is to further rewrite and upper bound the second term on the RHS of Equation (2) by
the set of
| - 1h1-211if 117 + > 1=
x H..
H.
41|
HH +
11 -
(5)
Note that we need H I H, to ensure that || 112 H. is indeed a (pseudo) norm and that the last inequality holds. Such analysis has been done for a few notable variants of AdaGrad in Gupta et al. (2017), where the condition H 3 H, is verified in a case-by-case way for specific choice of H. However, the following question remains unclear for optimizers described by general H::
In addition to the unified approach to deriving various adap- tive optimization algorithms, Gupta et al. (2017) also at- tempts to give a unified analysis for the convergence rate or regret of these adaptive algorithms, which can be summa- rized by the following theorem.
Theorem 2.1 (Gupta et al. (2017)). Let {xt be the iterates of Algorithm T Then for any x= e X.
Question 1. For cone H C S does PH(M) 3 (M(M)) hold whenever 0 x M = M'?
(x(x))
Σ r=1 Lx(x")
Indeed, the answer is no for "ill-structured" H. and we mention two negative examples here for illustration.
[Mr, HT Tr(HT) 2 Tr Tr(H)))
Example 2.2. Let 7. = (A @ B 2 0 A E Mdd, B e MdR] with dr = d, i.e., the set of preconditioners for
ETn + 2 -- (I xx-1 x1122, | x 211
3

---

Structured Preconditioners in Adaptive Optimization: A Unified Analysis
two-sided Shampoo (Algorithm 3). We show that in the special case where d, dæ= 2, for M = diag(1, c, and M' diag(L, e, e; 1), (MM 3 (M does not hold for sufficiently small e > 0, although we have M = M'. See Appendix 1.4.1 for detailed proof.
in (3). Furthermore, for any M > PM(M) satisfies the
following properties: T((())))))))
(a) (M, Px(M)) Tr(PH(M))
(b) (s((M) = arg HEH, (M,H-1), HH-),, where we recall that (M = PH(M)/ (T(()((TT)))) Moreover, for any 0 M 3 M', ( E H holds, In particular, it implies that (M( 3 (M').
Example 2.3. The second example of ill-structured H in- volves tridiagonal PSD matrices, i.c., matrices that only have nonzero elements on the main diagonal, the first diag- onal above the main diagonal, and the first diagonal below the main diagonal. Specifically, for H containing all 3- dimensional PSD matrices, we provide numerical evidence demonstrating that the desired condition Equation (4) breaks for very simple instances, See details in Appendix A.4.2.
The closure under both matrix addition and matrix multipli- cation is crucial in the proof of Proposition 3.2. Violation of any of the two properties could lead to problems: for the preconditioner set of two-sided Shampoo in Example 2.2, it is not closed under matrix addition; while for Example 2.3 regarding tridiagonal matrices, we note that the of tridi- agonal matrices is not closed under matrix multiplication.
These failure modes naturally lead to the second question:
Question 2. Is there a sufficient yet general condition (which covers all existing examples) on H for the inequality hold?
3.2. Unified analysis for adaptive optimization
To proceed with the regret bound, recall from Proposi- tion 3.2 that H, (T()) is a solution to the fol- lowing constrained optimization problem
As one of the main contributions of our work, we give an affirmative answer Question 2 in Section 3 by proposing notion of well-structured preconditioners (Definition 3.1) and deriving a unified analysis correspondingly.
H. = arg min (M, H HEH,T((H)S1
(6)
3. Unified Analysis for Well-Structured Preconditioners
Also recall that M, = E. 9.9." (assuming c = 0 for illustration). We interpret the optimal value of this con- strained optimization problem as the magnitude of the se- quence of gradients gı:t (g) (g1.-..-90 with respect the best preconditioner in hindsight. Motivated by this, for any sequence of gradients gia, we define their adaptive gradient norm with respect to H. to be
We establish a unified framework for adaptive optimization with structured preconditioners. In Section 1, we propose the notion of well-structured preconditioners and show that they satisfy the desired condition Equation (4). Then in Section 3.2, we present a unified analysis for adaptive opti- mization with well-structured preconditioners. We discuss several prominent examples in Section 3.3.
1 H
||g ==
inf HCH,
9gg. (Σ9. E
H-1)
(7)
Indeed, this definition of adaptive gradient norm corre- sponds to the dual norm of the norm ||- I denoted
3.1. Well-structured preconditioners
For set of d-by-d matrices IC C Mdd we say that I. is subalgebra if it is closed under scalar multiplication, matrix addition, and matrix multiplication. More concretely, we require that for any E E and A, B E K. it holds that AB, + B E K. Based on this, we propose the following core concept of our paper.
as || ll Specifically, we define
ved ree(gg.)llie.. =
sup WER'd ww ho $1
vec vec (g(:) T w.
Then we have |Ig - lowing lemma, which is proved in Appendix A.3.
/vt /t the fol-
Definition 3.1 stell-structured preconditioner sets). H C St is said be a well-structured preconditioner set if H. Sq (nK IC for some matrix subalgebra I. S MM with Ia K.
Lemma 3.3. Let H be well-structured preconditioner sel under Definition 3.1. Then any t 2 1 and 91. .I.....9t ,: E Rd, it holds that
As positive response to Question 2, the following proposi- tion shows that our notion of well-structured preconditioner sets provides sufficient condition for (=)(-) to be operator monotone. See Appendix A for proof.
{ 9.g. H
1
inf HEX)
ved ved rec(gg))l
Proposition 3.2. Let H be a well-structured preconditioner set under Definition 3.1. For any M > 0, there exists a unique solution PH (M) > 0 to the optimization problem
Given this definition of adaptive gradient norm, combining (6) and the fact that Tr H₂) = 17 "M, H by Proposi-
tion 3.2, we obtain Tr (H,) = 17 -'llgg::

---

Structured Preconditioners in Adaptive Optimization: A Unified Analysis
Now recall the upper bound in (5) for the second term in the regret bound (2) from Theorem 2.1. Note that for each Proposition 3.2 guarantees that H H H. so we can further bound || HH I || Tr(H, - H1-1).
full-matrix AdaGrad in terms of the number of steps. See a more detailed analysis in Section 4.3.
Next, we present the convergence rate of Algorithm 1 for stochastic convex smooth loss functions. We use Defini- tion 3.6 to characterize the smoothness of a loss function. It is an extension of d-smoothness in (Xie et al., 2025), which only applies to block boeek-daggoal preconditioners (correspond- ing to blockwise Adam).
This allows us to telescope the sum in Equation (5) to get
| 7
1 H, + Σ 1
(8)
H
Observe that the upper bound is factored into two parts: 1) [|x'|,,, the diameter of the domain under norm || llw: and 2) lIg.r 11 the adaptive gradient norm with respect to H. We pause here for an important remark. Note that || -.-H.,, is defined by taking the supremum over æ E X, which precludes telescoping the sum over & E [T] at first glance. We address this issue by proposing the norm || || which allows us to extract the factor (HH-He-HH this unified analysis is possible thanks to Proposition 3.2 for well-structured preconditioner sets, in contrast to the case- by-case analysis done by Gupta et al. (2017). Furthermore, we remark that the above factored bound is crucial for us to identify the correct norm metric for Shampoo, leading to an improved analysis. See Section 4.1 for details.
Definition 3.6 (H-smoothnesss). For set H. C S and any loss function L : Rd + R, we define the н- smoothness of L, denoted by H H), as the smallest number H 2 0 such that there exists matrix H' e H satisfying H = Tr(H") and for any x E R it holds that H 3 L() 3 H In the case of convex L, this requirement becomes V (L((x) I H Furthermore, we ex- tend the notation to matrices E Mad by defining H(A, H) as H(x + 2 Az. H).
Again,
We need the following assumption on the stochastic noise. Assumption 3.7. For any t E [T] and any a E X. [[),(x)] = L(x) and there exists some Σ E S such that ([ V=(x)) (VL((x)
Finally, combining (8) and the original regret bound in (2) yields the final regret bound for Algorithm 1. This is sum- marized in the following Theorem 3.4. The complete proof can be found in Appendix D.
Now we are ready to state main results on the conver- gence rate of Algorithm I for stochastic convex functions.
Theorem 3.8. Let H be well-structured preconditioner set under Definition 3.1. Consider any independent stochastic convex loss functions Lı...., Lr satisfying Assumption 3.7, and let H (L, H) be the H-smoothness of their expectation L. Suppose the global minimizer of I,, denoted by x", is in X. Then for the iterates D1...., IT of Algorithm 1, denoting $1.7 = ". Dt, it holds that
Theorem 3.4. Let H. be a well-structured preconditioner set under Definition 3.1. Then for any convex loss functions LI....,Ly Lr. the regret of Algorithm I compared to any
x² E X can be bounded as
T Ele(x)) - En(()))
T
D2 +1) (G dVE) 2n
1=1
E [L(TTT)) - [(x')]
where G = 11ll1:11ll 11 D maxie(T) || - x: 11
16 T
| 2 H(L H.) + 4v2
| 1
+
4V2dVE
|
VT
T
Corollary 3.5. Under the setting of Theorem 2.1, further suppose that X is bounded set in Rdd. Then choosing n 12 the regret bound for Algorithm I becomes
where inf inf /EE
Our analysis naturally extends to Algorithm 4, which re- places the direct sum of past gradient outer products in Algorithm with an exponential moving average (EMA). This modification is widely used in adaptive optimizers, in- cluding Adam (Kingma & Ba, 2014) and AdaSGD (Wang & Wiens, 2020), which is an ema version of AdaGrad. The detailed discussion is in Appendix C.
T Σ (m.) - Σ (2(2")2222(1111 (G (G+DV)). [=1
1=1
Ignoring c, the bound reveals an intrinsic trade-off between || and !ell as H gets larger, ||xx | increases while ||l decreases. The previous common belief that more adaptivity (larger H) helps optimization could be largely due to the loose upper bound on ||x''|x, i.e., always measuring the size of domain X' by Frobenius norm instead of potentially much smaller I1- The fact that comeeckrred-fcctoree subalgebra induces smaller I1-11, than the entire matrix subalgebra is the key reason behind our surprising finding that one-sided Shampoo can outperform
3.3. Examples of well-structured preconditioner sets
Next, we demonstrate that our Definition 3.1 is general enough to cover existing examples, by discussing several important matrix subalgebra K. For each associated well- structured H = KC n St recall the optimization problem
5

---

Structured Preconditioners in Adaptive Optimization: A Unified Analysis
over H defined in (3). We will show that every minimizer PH (M) corresponds to the preconditioner used in specific adaptive optimization algorithm. We list the correspondence relationship below and the results are summarized in Table 1. The detailed derivation and calculation of the norms || ||a|| and I 2 can be found Appendix B.
for d parameters of an N-layer neural network, decompose Rd Rdi x x Rd where d >m= dns and each Rd corresponds to the dn parameters in the n-th layer. For each E [M], let Kn, S MM be matrix subalge- bra. Then we define K. = I = fo An An E Kn n E [N]], and is easy to verify that is also a subalgebra. Then the corresponding well-structured precon- ditioner set H. I n S4*) = 0; Hn contains preconditioners that apply individual types of transforms to gradient of parameters in different layers?. Such H made up by direct sum of smaller cones also has very composi- tional property in its induced complexity metrics, namely, |l = gaXigngN Il-ll,, and IlI-Il = E Il... This operation provides a useful tool for designing layerwise preconditioning methods (Bernstein & Newhouse, 2024a).
Example 3.9 (AdaGrad-Norma scalar matrices). For the scalar matrix subalgebra KC = fc Ia c E R}, we have H. = (c. La 2 0}. Then solving (3) for M, yields
1 Tr Ia H; T)
1 / la.
E
#=1
This is the preconditioner used in AdaGrad-Norm.
Example 3.10 (Diagonal AdaGrad: diagonal matrices). For the diagonal matrix subalgebra K = Ddd we have H. = D4. Correspondingly.
Other possible operations that can generate new matrix sub- algebra K.' from the original subalgebra K. include taking Kronecker product with the identity matrix, i.e. K' {A' A & Is A E K), and rotation by an orthogo- nal matrix, i.e. K' = (A' U AU A E XC} where U is an orthogonal matrix.
H, diag (1
c+
11
[d]
U
This is the preconditioner used in diagonal AdaGrad.
4. Improved Convergence Analysis for One-sided Shampoo
Example 3.11 (Full-matrix AdaGrad: all matrices). For K. = .Md we have H = St. In this case, solving (3) for Mr yields that H,- = 1 M which corresponds to the update rule of full-matrix AdaGrad.
We now turn to one-sided Shampoo (Algorithm 2), a special example of Algorithm 1. In Section 4.1, we present our main results on its regret bound and convergence rate, and then compare to the previous results for two-sided Shampoo in Section 4.2. In Section 4.3, we present comprehensive comparison between the regret bound of one-sided Shampoo and those of the AdaGrad variants, which suggests why Shampoo can outperform other adaptive algorithms on some real tasks.
Example 3.12 (One-sided Shampoo: factored matrices). Let be factored as dr Then for the factored matrix algebra K. Rdx xd2 @ Iag we have S @ Lnn- Now writing G. E Rduxda as the matricized version of g. E Rd, solving the corresponding problem in (3) leads to
1 H. - 17 Ia + dn s=1
1 E..C.
3 IdR:
4.1. Our results for one-sided Shampoo
which updates o, the same as one-sided Shampoo' dis- played in Algorithm 2. where we write the algorithm in the matrix form for convenience. More specifically, note
We first characterize the norm || 1| for one-sided Shampoo.
Lemma 4.1 (lI - |l for one-sided Shampoo). Recall from
that gt vec( G:) and H, L @ Ldg. Moreover, M, corresponds to L, by the fact that (M, (H2 & Idn) (L,, H for any Hr, E Sd4.
Example 3.12 that for one-sided Shampoo (Algorithm 2), H= 69 IdR where d = d. Then for any x E RRd, it
holds that || Vaz
|| with X = vec vec-v (x)), and
thus || Vāg ||
The detailed calculations of the norms ||- Ilx and II for the above examples can be found in Appendix B.
See Appendix B.4 for the proof.
Generate new well-structured preconditioner sets. Be- yond the previous examples, it is also possible to generate new well-structured preconditioner sets based on existing ones. We discuss in particular an example of layerwise com- bination for parameters in a neural network. Specifically,
With this, apply Theorem 3.4 to get the regret bound for one-sided Shampoo, as summarized below in Theo- rem 4.2. See Appendix E.2 for its proof.
Theorem 4.2 (Regret bound for one-sided Shampoo). For convex functions Lı,..., , Lr, the regret of one-sided Sham-
We add a normalized factor zn. compared to the L, in Algo- rithm so that it can exactly derived from Algorithm 1.
Thhis be applied to any partition of the parameters, not only for partition hased on layers.
6

---

Structured Preconditioners in Adaptive Optimization: A Unified Analysis
Algorithm One-sided Shampoo
Algorithm 3 Two-sided Shampoo (Gupta et al., 2018)
Hyperparam: learning rate 17 > O, convex set X C R xdR, Lo = cla, for e 2 0
Hyperparam: learning rate 1, convex set X' C TRdd xdx,
Lo eLdz, Ro eldg for e 0 Input: initialization zu. stochastic loss functions (I,) IR xd2 - RR , _T do
Input: initialization To. stochastic loss functions
{1x)/=1 ::::::dxxxx.. xdR + R
for t 1,2. .T
G: VL (X(-1)
G. +- (4((X-1)
L, 5 Li-1 + dR
G.G.G
L, - レ1 + G.G.
L @lan (X1-1- Mi Gi)
R f- R + G. G2
X, - 11x
X, f- 11 sR
(X,-1 - neL, G,R,
Return DT
Return TT
poo (Algorithm 2) compared to any X E Rddxda satisfies
ing rate 1) - /21l//lopp Then 31.T for = + +Et=1 2t.
B[L( L(x')]
T
T Σ L((x))
D2 +n) (G+dVC),
Σ (X,)
=
2d
Tdæ 11 HiGh +
16
4v20 vTdR
|| +
4v2dv/
1/2'lop
1=1
T
where Dop = maxielT X,
lolpop and G
where a = in v (,, and Hieh (L) is the left smoothness of the expected loss L identified in Lemma 4.3.
(Ei |
Vd Tr bounded in operator norm, ie.,
When the domain X is 12 < co, further choos-
ing 17- 22 2/ 21 it holds
4.2. Comparison with previous results on Shampoo
T
T
We compare our main results for one-sided Shampoo with the original results in Gupta et al. (2018). Here, we restate their original regret bound for easier comparison.
Σ L.(X))
Σ Lx(x')
T G.G.
$
Theorem 4.5 (Regret bound of two-sided Shampoo (Gupta et al., 201 ()). For convex functions {L.) suppose their gradients (G, = VL,(X ())) are matrices of rank at most r. Then the regret of two-sided Shampoo (Algorithm 3²) compared to any X E Rd xdR is bounded as
1 [==
d
<2v2||X
Tr
Ve
vdn
Next, before presenting the convergence rate of one-sided Shampoo, we first provide more interpretable formulation of the H-smoothness for one-sided Shampoo, which we call left smoothness. See Appendix 1 for proof of Lemma 4.3.
T
D
Σ (((i)) - Σ L(X") I
Tr(LT) Tr(R)),
+ m
2m
1=1
where Dy max.e(T] || X LT eldz +
Lemma 4.3 (Left smoothness for one-sided Shampoo). Let
Σt G.G., and RT elaa + - G. When
H Sth @ Idg be the well-structured preconditioner set for oie-sided Shampoo. Then the H-smoothness H (L, H.) defined in Definition 3.6 is equal to the smallest number H 2 such that there exists Hig Rdi. xdi, satisfying that H- = dn Tr(H2,) ) and that for any X,A RR XdR,
||''| < co, we further choose 17 v /2/r X then
T
Σ (X,) Σ L (v2]|\l1 Tr(L)) TT(R
We now compare our regret bound in Theorem 4.2 and the original regret bound in Theorem 4.5 by Gupta et al. (2018) when 0, i.e., LT EE=G, and Rr I G, G. For matrix M E Sd. it always holds that (T(M) I Tr(M)). Therefore,
V'L(X)IA (H,, AAT.
In this case, the H-smoothness is denoted by Hiett (T.).
Then the convergence rate of one-sided Shampoo can be ob- tained by specializing Theorem 3.8 to || 1 || from Lemma 4.1 and the left smoothness from Lemma 4.3.
T- Tr(Lr) TrLL) (( 2 Tr(L Tr -
Tr IL 2 Tr(D)).
Theorem 4.4 (Convergence rate of one-sided Shampoo). Let LI....,,A. Lr be stochastic convex loss functions satisfy- ing Assumption 3.7, and let æI...., DT be the correspond- ing iterates of one-sided Shampoo (Algorithm 2) with learn-
'The original two-sided Shampoo analysis in (Gupta et al., 2017) without per step projection to the bounded domain. We adapt their theorem into the projected version a standard way.
7

---

Structured Preconditioners in Adaptive Optimization: A Unified Analysis
This implies that regret bound in Theorem 4.5 is always smaller than the regret bound in Theorem 4.2 because 2 and || F 2 || as || X || > |opp for any X E X.
it holds that H 3 Idz. Then we have
max min Tr(D) S dz max H(H, (4 ding TH())
Moreover, in the worst case, the Frobenius norm can be /min(dr, dı)) times larger than the operator norm. As a concrete example, suppose G. satisfies that G.[i,j] = only for (j 1 i) = (mod min(dz, dR)) and all other elements are zero, Then each G, has rank r = min(dr, dR). At step T drdR, L1 = du min(d,, dn) Ia,, and R dg min(dr,, dR) Ldy; and thus Tr(L)) Tr(Rt) d while Tr(L)) dzd In this case, the regret bound of
On the other hand, if we choose H = 1s, 1 /d.. for any diagonal D = H. it holds that Tr(D) dr. So we prove that
de DId 2 17, Hid, ( H(H,7 D4 = dr.
We also know that maxix Ilap vied llx S 1 since | V || lop If the only nonzero entry of X is Xi, = 1, then || = | = 1. Overall, the
worst case of convergence rate is Lnda..
two-sided Shampoo is min(dz, du I d times larger than the our regret bound for one-sided Shampoo.
Full-matrix AdaGrad. We have seen in AdaGrad-Norm that maxij --
Il.p $1 vved (X*) / {dz, dR}. so the
Duvvuri et al. (2024) introduced CASPR as an alternative to Shampoo for approximating full-matrix AdaGrad and achieved the same regret bound as Theorem 4.5. Our previ- ous comparison thus applies to both Shampoo and CASPR.
worst case of convergence rate is min (dr,dn)da
To summarize, we have identified class of optimization problems, namely loss functions like Equation (9). with Hessian of bounded trace and optimizer of bounded spectral norm, for which one-sided Shampoo has much better worst case convergence rate than any other adaptive algorithms that could be derived from our unified analysis.
4.3. Comparison with AdaGrad variants
Next we show one-sided Shampoo can achieve the best the- oretical upper bound of the suboptimality gap for a specific class of functions, where each loss function has the form
L(X)= = (H, (X - Xx -
(9)
5. Experiments
where X E RRdZxdR, H E tt with Tr(H) < 1, and |X 1x S 1. For this function class, we compare the largest possible value of convergence rate of different al- gorithms given by Theorem 3.8, which is summarized in Table 2.
In this section we empirically demonstrate the superior performance of 1-sided shampoo over other variants of AdaReg (Algorithm 1) on simple but natural setting. More- over, such superior performance is predicted by our theo- retical analysis in Section 4.3, which in turn validates the practical utility of our theory in guiding optimizer selection.
For each algorithm defined by Algorithm with specific H. we will pick X = {x || S ||0*||x]] so that the global minimizer x² is reachable. To get the convergence rate, it boils down to calculate l = -|<="|1 and H HL, H. associated with each algorithm. We have already derived the explicit form of ||>x'| for each algorithm in Section 3.3, as shown in the third column of Table For the H-smoothness, note that L(x then it is straightforward to calculate H (L. H) according to Definition 3.6, as shown in the fourth column of Table 2.
Setup. We consider a linear regression problem |A y where A is the data matrix and y A is the label vector generated by ground-truth X Thus, the loss function can be equivalently written as
f(X) = (H, -X-) -X)))
H @ Lax for the loss L in (9), and
which is the same function that we studied in Section 4.3 and show that 1-sided shampoo outperforms other adaptive algorithms. We consider X dxdd with d 103. We set the eigenvalues of H by 01
Below we present the worst case of the convergence rate on this problem class. We can see that one-sided Shampoo is strictly better the other three adaptive algorithms.
= - = 1 and a = (i-10]= for 11 < i < 103. Each element of the solution X* is independently sampled from (o 4). We run MlaGrad-Norma AdaGrad, one-sided Shampoo and full- matrix AdaGrad for 100 steps from initialization Xo- O. We also the original Shampoo algorithm for comparison. Full-matrix AdaGrad is run in a memory memory-efficient way and the detail is in Appendix F.I. We will compare the last iterate loss and the average iterate loss separately. The learning rate is tuned over five seeds for last iterate loss and average iterate loss respectfully, selecting the one with the
One-sided Shampoo. The worst case of convergence rate for one-sided Shampoo is --
Since maxT(H)şi Amax H) = and maxixx llop |vec{ x=) = max x S1 [1 (X*) I |p =
AdaGrad-Norm.
/ /min [di, dRJ. the worst case of convergence rate is min (dz, dT)/T.
AdaGrad. For psd matrix H E RRd2xds xdz with Tr(H)= S 1,
8

---

Structured Preconditioners in Adaptive Optimization: A Unified Analysis
Algorithm
Subalgebra K.
|
(I,, H)
Convergence Rate
H
AdaGrad-Norm
{cla C
200
1l2
d Xmaoe (H)
2 2 (aa)(H)TT
AdaGrad
Diagonal matrices, (R) All matrices, Rdxd
x 1loc
da HH, ((H,H% "4")
|
2
Dc dg ((H, /
Full-Matrix AdaGrad
x
l|2
dR Tr(H)
| dR H | Tr Tr(H)/T
One-Sided Shampoo
Rd..xd. xd1, @ Ldæ
Ix
du Tr(H)
vdg
Table 2. Convergence rate for the loss function L(x) = (H. (x X- -x))) ). The results can be obtained by Theorem 3.8 with σ= 0 and omitting e. For each algorithm we pick the smallest domain which still ensures the minimizer lies in the domain. Here recall that 2 is the set of all d-dimensional diagonal matrices, and D D ns..
lowest average loss. We use precision float32 and for all the experiments.
tioners (Ba et al., 2017: George al., 2018; Martens et al., 2018; Yao et al., 2021; Jahani et al., 2021; Zhang al., 2022; Duvvuri et al., 2024). Indeed, our results suggest that such compromises might not harm the performance of adap- tive optimizers in practice, because more adaptivity is not always helpful, as discussed in Section 4.3 and Section
We also run the EMA version of the adaptive algorithms, whose results are consistent as shown in Appendix F2.
The overall results are shown in Section 5. We can see that one-sided Shampoo greatly outperforms other algorithms, corroborating the theoretical analysis in Sec- tion 4.3, Moreover, full-matrix AdaGrad is apparently the worst, suggesting more or even full adaptivity does not al- ways help optimization for fixed budget of training steps. One-sided Shampoo is also better than the original Shampoo algorithm.
Results.
Understanding Shampoo. There are also recent efforts to understand the Shampoo optimizer from various perspec- tives. Bernstein & Newhouse (2024b) interpret Shampoo as the steepest descent with respect to the spectral norm of the layerwise matrix-form parameters of the neural network. Recognizing such structures of matrix-form parameters and role of spectral-norm geometry in deep learning has led to development of new optimizer such as Muon (Jordan et al., 2024). In addition, the second-order perspective on Shampoo (Anil et al., 2020) has also led to fruitful results: Morwani et al. (2024) connect the preconditioner Sham- poo to the optimal Kronecker product approximation of the Gauss-Newton component of the Hessian, and Vyas et al. (2024) propose to view Shampoo Adafactor in the eigenbasis of the preconditioner of Shampoo.
Loss
Avarage Itarata
Step
Step
AdaGrad
AdaGrad
Full-Matrix AdaGrad
One-sided Shampoo
Shampao
Figure 1. For both the last iterate loss f(X,) and the average iterate training f(} E X.). one-sided Shampoo performs the best and Full mall-matrix AdaGrad performs the worst. Here
7. Conclusion and Future Works
We present a unified analysis for a broad class of adaptive optimization algorithms with well-structured precondition- (Definition 3.1) for both online regret minimization and smooth convex optimization. Our analysis not only provides matching rate to several important algorithms including di- agonal AdaGrad, full-matrix AdaGrad, and AdaGradNorm, but also gives an improved convergence rate for a one-sided variant of Shampoo over that of the original Shampoo. We reveal novel trade-off in final convergence rate between domain metric and adaptive gradient norm (Equation (7)) for regret minimization or adaptive smoothness (Definition 3.6) for smooth convex optimization. We hope this insight could be useful towards design of future adaptive optimizers. One important future direction is to identify more subalgebras or other structures that are useful for improving the perfor- mance by better adapting to the domain or loss smoothness.
f(X) (H, (X- X'XX- T)*)).
6. Related Work
Adaptive optimizers and structured preconditioners. The extensive costs for training large-scale deep learning models have motivated the development of efficient opti- mization algorithms, among which adaptive optimizers have been widely studied because of their ability to exploit the rich geometry of the loss landscape (Pascanu & Bengio, 2013; Martens & Grosse, 2015; Dozat, 2016; Loshchilov & Hutter, 2018; Shazeer & Stern, 2018; Reddi et al., 2019; You et al., 2019; Zhuang et al., 2020; Liu et al., 2023; Yuan et al., 2024). For the sake of memory and computational efficiency, many works involve approximations to full-matre mallimatrix precondi-

---

Structured Preconditioners in Adaptive Optimization: A Unified Analysis
Impact Statement
George, T., Laurent, C., Bouthillier, X., Ballas, N., and
Vincent, P. Fast approximate natural gradient descent in a kronecker factored eigenbasis. Advances in Neural Information Processing Systems, 2018. URL https: / //arxiv.org//ooo.. pdd 800. 03884.
The goal this paper is to advance the field of Machine Learning by providing theoretical and experimental insights for adaptive optimization algorithms. There are many po- tential societal consequences of our work, none of which we feel must be specifically highlighted here.
Gupta, V., Koren, T., and Singer, Y. A unified approach to adaptive regularization in online and stochastic opti- mization. arXiv preprint
Acknowledgment
2017. URL https: /tte://arxiie araaa.iiiiv..... wwtp 06556.
The authors sincerely thank Matt Streeter for his valuable discussions and insightful comments throughout this work.
Gupta, V., Koren, T., and Singer, Y. Shampoo: Pre- conditioned stochastic tensor optimization. In Interna- tional Conference on Machine Learning, 2018. URL 111./11822....53
References
fttps://arxiv.orgg
Anil, R., Gupta, V., Koren, T., Regan, K., and Singer, Y. Scalable second order optimization for deep learning- arXiv preprint 20Xiv:2002.09011 2020. URL https: // //arxiv. org/ pptp://df//2/22//////////// /2002.2022222 0 09018.
Hazan, E., Agarwal, A., and Kale, S. Logarithmic regret al- gorithms for online convex optimization. Machine Learn- ing. 2007. URL a article 10. www111//////09-0000
com/
Ba, J., Grosse, R., and Martens, J. Distributed second- order optimization using kronecker-factoonn approxima- tions. In International Conference on Learning Represen- tations, 2017. URL https:/ / jimmy (ba.ylbb.. gittub........ io/papers/naa htwperrr///nncccccc
Jahani, M.. Rusakov, S., Shi, Z., Richtárik, P., Mahoney, M. W., and Takáč, M. Doubly adaptive scaled algo- rithm for machine learning using second-order informa- tion. arXiv preprint vrXiv:2109.05/990 2021. URL https: bttpss//arxiv.coo orgg pdf /2109 051 98
Bernstein, J. and Newhouse, L. Modular duality in deep learning. arXiv preprint arXivv241002/265, 2024a. URL https: petps://arxiv...r wtt. http://df////////
Jordan, K., Jin, Y., Boza, V., You, J., Cesista, F., Newhouse, L., and Bernstein, Muon: An optimizer for hidden layers in neural networks, 2024. URL https: /web. ar archive htcpa/ee...rree/////222222200000 / //kellerjorda: githugthhhhggge io hww.:/oosss/ssssoo
Bernstein, and Newhouse, L. Old optimizer, new norm: An anthology. arXiv preprint arXiv:2409 2024b.
URL https:///axiii wwtp/////pp/ee// 20325..
Kalai. A. and Vempala, S.
Efficient algo-
rithms for online decision problems. Jour- nal of Computer and System Sciences,
Dahl, G. E., Schneider, F., Nado, Z., Agarwal, N., Sastry, C. S., Hennig, P., Medapati, S., Eschenhagen, R., Kasim- beg, P., Suo, D., et al. Benchmarking neural network train- ing algorithms. arXiv preprint arXiv:2306.07717 2023. URL httpt::tp///s.. //arxivv///xvv.v./// http//aaww..oogp//////3000 0 07179.
2005.
URL https::/www.wt http:://wwww.sceeeeddeee. com/ cctenne//attiaeeaa///22000000000
Kingma, D. P. and Ba. J. Adam: A method for stochastic op-
timization. arXiv preprint 1410:mm11122:0110000000 2014. URL
Dozat, T. Incorporating nesterov momentum into adam. 2016. URL https://openreviiet http///ewwwee//r
Levy, K. Y., Yurtsever, A., and Cevher, V. Online adap- tive methods, universality and acceleration. Advances in neural information processing systems, 2018. URL https: http:://aaxxv.oooo www.eiiiw...ooooooooooeooeee http://acxxi..oorpp/////ee/ wttp:///////109..22222 02864.
Duchi, J., Hazan, E., and Singer, Y. Adaptive sub- gradient methods for online learning and stochastic optimization. Journal of machine learning research, 2011. URL https::/wwet Smlr.oo org/ volume12/ www.umee///cchhllaae
Liu, H., Li, Z., Hall, D. L. W., Liang, P., and Ma. T. Sophia: A scalable stochastic second-order optimizer for lan- guage model pre-training. In The Twelfth International Conference on Learning Representations, 2023. URL https: /anss://araae //tp///aavvvv//aa///a//eeeeaaaev wttpoo/gg/p////o/ 143.2/3335...1..ppco .dff.
papers/
Duvvuri, S. S., Devvrit, F., Anil, R., Hsieh, C.-.,, and Dhillon, I. S. Combining axes preconditioners through kronecker approximation for deep learning. In The Twelfth International Conference on Learning Represen- tations, 2024. URL https: htttp:://openrevew..ct net./ pdffid=8 9ah 8DV18. (t
Loshchilov, I. and Hutter, F. Decoupled weight decay reg- ularization. In International Conference on Learning Representations, 2018. URL https: //t////aaivv/////////////////aar .org/ pdf/1711.0511 .ddf.
10

---

Structured Preconditioners in Adaptive Optimization: A Unified Analysis
Martens, J. and Grosse, R. Optimizing neural networks with kronecker-factored approximate curvature. In Inter- national conference machine learning, 2015. URL http:/ /ttp:///aaaiv.////a wwtpootg/////oooo/..///oott.. podt/1503. 05671.
Xie, S., Mohamadi, M. A., and Li, Z. Adam exploits $l Lifty&--eomeets of loss landscape via coordinate- wise adaptivity. In The Thirteenth International Confer- ence on Learning Representations, 2025. URL https: / / //opperevi ew net/ forum? d=PUnDD 6UEK5.
Martens, J., Ba, J., and Johnson, M. Kronecker-factored curvature approximations for recurrent neural networks. In International Conference on Learning Representa- tions, 2018. URL
Yao, Z.,, Gholami, A., Shen, S., Mustafa, M., Keutzer, K., and Mahoney, M. Adahessian: An adaptive second order optimizer for machine learning. In proceedings of the AAAI conference on artificial intelligence, 2021. URL https: /ttss///axxii org/ httg:/dp//2000/ 00719.
pdfRid=HMMKQQZA..
Morwani, D., Shapira, I., Vyas, N., Malach, E., Kakade, S., and Janson, L. A new perspective on shampoo's s preconditioner. arXiv preprint arXiv:2406. 2024.
You, Y.. Gitman, I., and Ginsburg, B. Large batch training of convolutional networks. arXiv preprint arXiv:1708.03880 2017. URL https://arxis: org/pdf/1708 03888.
URL httppp:://aa.....rrta/////......
Pascanu, R. and Bengio, Revisiting natural gradient for deep networks. arXiv preprint arXiv: 2013. URL hwtthhts////atttee orgg ptrg/pdf/13011
You, Y., L., Reddi, S., Hseu, J., Kumar, S., Bhojanapalli, S., Song, X., Demmel, J., Keutzer, K., and Hsieh, C.-J. Large batch optimization for deep learning: Training bert in 76 minutes. arXiv preprint arXiv: 19040006.. 2019. URL https: arriiv..vv org ptg/pdf////eeoooooooooggg//ggpp/ 1tpff/150 /1004 00962.
Petersen, K. B., Pedersen, M. S., et al. The matrix cookbook. Technical University of Denmark, 2008.
Yuan, H., Liu, Y., Wu, S., Zhou, X., and Gu, Q. Mars: Unleashing the power of variance reduction for training large models. arXiv preprint arXiv:241 2024. URL httph::p////.. /awpiv www.xii....g/pppoo htrg/pdf/2411 10438.
URL https: /ttce///cee./eee https::/ecc.wwatrrooo. wwwterloo MISS/matriixcoobbokk.
ca/ee00502/
Reddi, S. J., Kale, S., and Kumar, S. On the convergence of adam and beyond. arXiv preprint arXiv: 1904.092477777 2019. URL https:/ http///axivv..rr/ ort/ 1904. 09237.
Zhang, L., Shi, S., and Li, B. Eva: Practical second-order optimization with kronecker-vectorized approximation. The Eleventh International Conference on Learning Representations, 2022. URL https: http://axxivvvooo pdf 22008 02123.
Shazeer, N. and Stern, M. Adafactor: Adaptive learn- ing rates with sublinear memory cost. In International Conference on Machine Learning, 2018. URL https: l/arxiv.org/ wwtpo/////////o////////oo/////// /adf//1/44 04235.
Zhuang. J., Tang, T., Ding, Y., Tatikonda, S. C., Dvomek, N., Papademetris, X., and Duncan, J. Adabelief optimizer: Adapting stepsizes by the belief in observed gradients. Advances in neural information processing systems, 2020. URL https: / /arxivv.. org/paf/2010. 07468.
Streeter, M. and McMahan, H. B. Less regret via online con- ditioning. arXiv preprint ar.iv: 2010. URL https:/ /tttps://arivv..o otg/ 002.00022200000 4862.
Vyas, N., Morwani, D., Zhao, R., Shapira, I., Brandfon- brener, D., Janson, I., and Kakade, S. Soap: Im- proving and stabilizing shampoo using adam. arXiv preprint 2rXiv:2409.11321, 2024. URL https:// wrxiv.org/pdff//e
Wang, and Wiens, Adasgd: Bridging the gap between sgd and adam. arXiv preprint arXiv:2006. 16541, 2020. URL https: http///axivv///////////eee/...aa or..
wwtpoo/gg/////oooooooooooooo/tt/ http:dff/2000 6.00.0000222 16541.
Ward, R., Wu, X., and Bottou, L. Adagrad step- sizes: Sharp convergence over nonconvex land- scapes. Journal of Machine Learning Research, 2020. URL https: / /www.
colume21/18-- 11111181-52/111 182118--55510000 3df.

---

Structured Preconditioners in Adaptive Optimization: A Unified Analysis
A. Proof for Well-Structured Preconditioner Sets
A.1. Definition H. and its properties
Recall that IC is a subalgebra of d-by-d real-valued matrices, and the corresponding well-structured preconditioner set H K. S
Lemma A.1. For any A E K and any polynomial P(A) E K. Furthermore, for any invertible A E K, its inverse A-1 & K. Also, for any symmetric A E K. its pseudo inverse A E K.
Proof of Lemma A.I. The first statement follows from the fact that K. is a subalgebra, and the second and third statement consequences of the Cayley-Hamilee Hamilton theorem.
Lemma A.2. For H = K S where KC is a subalgebra of d-by-d real-valued matrices, define
H (AEK: (H, A) 2 0, E H).
(10)
Then H" = H. Consequently, for any H1, H2 E H. if (HI H H > 0 for all H E H. then Hi 2 H2.
Proof of Lemma A.2. Suppose there exists A E H such that A has a negative eigenvalue. Let A1 (A)..... (A) be the eigenvalues of in the decreasing order. Consider the matrix = (A 2 max(Ai (A), 1) 1) 0. The leading eigenspace of B is the same as the eigenspace of A corresponding its smallest eigenvalue, which is negative. Consequently, for large enough integer n, B" ||||| B" llop is approximately the projection matrix onto the eigenspace of the smallest negative eigenvalue of A. Therefore, for large enough integer n, (B"/I|B* /"|pB" A) < 0. However, since B" B"/|| is polynomial of A, we know that B" | |lop E H. This is contradiction to the definition of H. Hence, we conclude that for any A E H', it holds that A 2 0, and thus H C H. Moreover, for any A E H, it holds that (H, A) 2 0 for any H E H. because both and H are positive semi-definite. This shows that H. C H." and hence H = H. This completes the proof.
A.2. Properties of Px((-)
For any M > 0, recall the regularized optimization problem in (3). Note that we can assume 7) = without loss of generality, because the original problem equivalent to solve for M /112 with regularizer Tr(H) place of 1/2 Tr(H), Therefore, in the rest of this section, we focus on the following optimization problem:
(x(M) == arg (M, H Tr(H).
(11)
ME#
The results proved below applies to the original problem (3) after simple rescaling, and Proposition 3.2 follows from Proposition A.3 below.
For notational convenience, given any M > O, define
(M = (M
(12)
Proposition A.3. Let H be well-structured preconditioner set under Definition 3.1. For any M > there exists unique
solution Pxe (M) > 0 to the optimization problem in (11). Furthermore, for any M > 0, Ps (M) satisfies the following
properties:
(a) (M, -)) = -
(b) (M) = arg minnnennnnnngi (M, H where recall that Px(M) -MP)
(c) For any H E H. (-Px(M) MMx Ia. Px((M) =
Moreover, for any Mj Z M2 > O, it holds that (M (M2) H. and particular, PH (M1) 2 Px(M2).
12

---

Structured Preconditioners in Adaptive Optimization: A Unified Analysis
Proof of Proposition A.3. We first show that (M( exists and (M 0. Note that for H. Kn S4, since K. is linear subspace of Md and $ is closed subset of M we know that H is also closed subset of MM. Moreover, for any sequence { (Hn) 2221 such that H > 0 and either the smallest eigenvalue of Hn converges to 0 the largest eigenvalue of Hn converges to oc, the objective value far (H,) goes to co because M > O. Therefore, there exists (M(M) > 0 that attains the minimum objective value. For the uniqueness of (M(M), it suffices to note that the objective function fas (H) is strictly convex in H > 0 for M 0.
Proof for property (a). Suppose otherwise that (M PH((M) * (M)), and consider the following matrix:
H = PM((M) (M, (M // E H.
For this H. we have thus contradicting the optimality of Pk (M). Therefore, it must be true that (M, AM (+ (M))
>+ + (r) = 2/ (M, PM(MM Tr(PH(M)) < ()
Proof for property (b) Note that we can rewrite the original optimization problem as follows:
(M(
arg (M, H + Tr(H)
HER
H
arg min
M,
Tr(H) Tr(H).
(T(H))
HeN
Note that solving the above optimization problem is equivalent to first solving (M( = arg ( (M, H and then setting (M = (M)) (M) where the value of Tr (M)) ensures the previous property (a). Hence, we see that Pxx (M) solves the constrained version of the original optimization problem.
Proof for property (c). Since VH Tr(H) Ia and HH Tr(M - (see e.g. Equation (124) in Petersen al. (2008)), have
VH|M(H)= -H
la-
(13)
Then as H. is cone, by the optimality of Pr (M), it holds for any H E H. that
0 (VHIM(Px(MM)) H - PM) (-P (M) + Ia, H Px(M))
Proof for the operator monotonicity of Px((-) By property (c) of (+((). we know that for any M > 0, (M (M)) (M(M)++MMM)1 (MM)
(M) = O. Note that H + Px (M) is a bijection from H to H. by Lemma A.I, so we have (M - PM(M) H - (M(MM) -1) = 0 for all H E H. Applying this to both Mj and M2.
follows that for any H H
(Mi Px(M)) H (P-PP(MM)))) -MM-PAM PM H- -PM((M))
Rearranging the above equation, we obtain
(Pn(M,) - Pke (M)) H) (M1 M2, H) - Px(Mi) (T())((M)))
+ (M2, PH (x( - (M())
(Mi M2, H)
where the second equality follows from the first property of Px (M1 and Px from Proposition A.3. Since M1 > M2,
this implies that (Px( M Px) H) 2 0 for all H E H. By Lemma A. we know that P* (M, PH(M, E H. so it further follows from Lemma A.2 that (M > (x2 Since matrix square root is operator monotone, it holds that ( ž (2) We also know that (M) - (M E K. because IC is a subalgebra. Then we conclude that (M1) (M2) e H. from the definition of H. This completes the proof.
We can further extend the definition of Px(-)) to all positive semi-definin matrices. Specifically, for any M > 0, we have M + cla > 0 for any 6 > 0, so PH (M + eIa) well-defined. Also, by the operator monotonicity of Pxx (x) from
13

---

Structured Preconditioners in Adaptive Optimization: A Unified Analysis
Proposition A.3, (M + +ea) PM(M e'Ta) for e' 2 0. This implies that PR(M + +e has limit as e - 0.
Therefore, for any M = 0, we define
PH(M) = lim (M(M + +ea). e'U
(14)
Note that the above equality is also true for M > 0. To this, we apply the optimality of every PM(M + cTa) to get
(M + ela PH(M) (T(P)(M)) > cM +ela) (M cla)).
Also note that (M + cIa) > (M ōIa) for sufficiently small 6 > such that M ōIa > 0. Therefore, letting 0 both sides of the previous inequality, we can exchange the order of taking the limit and taking the inverse of (M + eTaa to obtain
(M, PM()) (T(P)(M)) 2 (.. (.iip lim Pu(M + +Iaa) =la)) Tr lim (M + eIaa).
(e
Then by the optimality of (M and its uniqueness, we conclude that (14) is also valid for any M > 0.
Indeed, the definition of (M( in (14) provides continuous extension of Py to ST, as summarized in the following proposition.
Proposition A.4. Let H be a well-structured preconditioner set under Definition 3.1. As function on S4 Ps can be continuously extended to be a function on S4. Moreover, for any M > 0 such that M + 0, Pi (M) 2 satisfies the following properties:
(a) span(M) S (Span(Px(MM) and (M, Px(M)) + Tr(Px(M)) = inf inf H TT(H).
(b) (M())) (r(P)(M)).
(c) (M, (M))) = inf inf (M, H where we recall that PH(M) =
(d) For any H E H. (PP(M))MPx(MM) II PH((M) 0, where M is the projection matrix onto (P))():(M))..
Moreover, for any Mi M₂ 0, it holds that (M) (M2) H. and in particular, (H 2 Px
Proof of Proposition A.4. We divide the proof into different parts for different properties of PH-
Proof for continuous extension of Ps. We first show that Px can be continuously extended to S and we consider the extension of Px. as given in (14). We first show that for any M > and any sequence (M, ] such that each M, > PH M). Note that for any (0, 1), there exist in 2 S. > such that 0 < (1- (M 3 M, = M M for all large enough n and moreover, lim Cn2 lim Sn Then by the operator monotonicity of (-).).
and lim M,, M, it holds that 1im PH (M,)
En Ia). Also have (=((1 (M S. La)) PH PH((M &
note that Px ((1 (M L)) vI 8P (M + La). Letting + co, both Px(M S,, La) and Px + te converge to PH (M), which then implies that -I- 5Px(P) = lim,- (M(M,) = (M(M). Since 5 is arbitrary,
we conclude that lim, (M,, = (M Next, consider any general sequence (M, ) such that M, 2 0 and lim,--oc M, = M. For any > 0, there exists > O such that for all M' > 0 with | - M' S 6, holds that (M - (M < e. For this 6, there exists
0 such that for all n > N, || - M F < 8/2. Then since (M,,)) limm-+20 PM,M, + 2dm Ia). where for every m 2 we have || (M, + 2mva Ia) S
&
IM Mn 1/F + 2mvd LalFF < 6, so (M) PxeMM, I Ta)lFF S for all > 1. which implies that (Px (M,)) < for all n > N. Therefore, it follows that lim,-+00 Px = Px(M). In conclusion, Px(()) can be extended to be a continuous function on S
2m/d
Below, we any M > such that M + 0.
14

---

Structured Preconditioners in Adaptive Optimization: A Unified Analysis
Proof for span MM) C span( (M))). Let I Pxx(MM)Px(MM be the projection matrix onto span (M))),
then M E H by Lemma A.1. It suffices to show that (Ia M: MM) 0. Using the fact that Tr(AB) & (T(A) (r(B)
for any A, B = 0, we can get the following inequality for any > 0:
(Ia HM. M) Tr((Ia- -MMMM S Tr((Ia (M) + +ela)) (r(P + +ela)-
= -Ia -
+ ela)) (M, PH(M +ela)
(15)
(M (ela) e ela)). which is bounded by an absolute constant for all E (0, 1). Therefore, letting e - O on both sided of (15), since (Ia II, PM(M + +La)) + (Ia - II, PH (M)) = 0 by the definition of IIM. we conclude that (Ia III M) 0. This implies that span S (M)))
By Proposition A.3, we know that (M, P (M++) + =
S Tr(P
Proof for the optimality of (M) We still use MM to denote the projection matrix onto span span(Px(MM) (M)) For any e > O, by the optimality Px(M + ELa), holds for any E H. that
(M + cla. (M + elaa Tr(P + cIJ) S (M + t H Tr(H)
(16)
Since span (M) S span(Px(ML (M)) we know that (M. +La) - HM). Then since lim Px(M + ela) = (M) it holds that M P
(aPPP+MM (M +ela) 11 ( M + e for sufficiently
small e >> 0, and also that Px(MM lim,-:o(II (M + +cI) This implies
This
lim M, PH(M c (M, (M)) e <
(17)
Therefore, assuming the existence of the limit of (cIa, (M + cla)= as c + 0, taking the limit of - 0 on both sides of (16) yields
(M, + Tr(Px(M)) + lim Px + + - +
(18)
c-0
Hence, suffices show that (ela, PH(M -eaa) + 0 as 0. To prove this, note that for any > 0, (x(M) +
vē(Ia- II)) O, and its inverse is given by PM(M) 1/2 II Also, (M + Vella IM) H by Lemma A.1. Therefore, by the optimality of Px M + cIa), we have
(M + ela, PM(M + +ela) Tr(Px (M + +ela))
+ Tr(PP(M) + Ve(la- IM)) S (M + cla, (PH(M) + VE(Ia - II)
= =(M ela PH(M) -6/2/ Il)) Tr + Ve(d rank((P
-(M, Px())( + (M))) + (MM) + +2/e(d rank(Px(MM))
where we apply Tr(La- - IIM = d rank (M) and the last equality follows from the fact that (M, La - IIM) 0.
Rearranging the above inequality, we obtain
cTa,, Px + <1)) (M, PM(M) (M, Px(M + -II)) TT(xx((M)) - Tr(P + +eTa))
+ (E-Th(P= + 2vE(d raakk ran(())(M))).
Now applying (17) and noting that c(La, (M +Ia) 2 taking the limit e - we obtain
lim (M eTa = -00
(19)
Then combining (18) and (19), we conclude that for any H E H.
(M, Px()))) (M)) & (M. +T(H)
This confirms the optimality of (M). Further applying the same argument as in the proof of Proposition A.3 yields the optimality of (M) = Tr(Px(M)) (M))
15

---

Structured Preconditioners in Adaptive Optimization: A Unified Analysis
Proof for property (d) of (M) By the optimality of (r((M), we have
(M, PM())) Tr(P)(M)) = inf HER (M, H ") TT(H
(M)(H)
= inf MEH
(M, H ') + Tr(II HITM).
(20)
[M((H)
To why the last equality holds, first note that inf meH SM (H) 2 iif J' (H) because Tr(H) MM M) Then for the other direction, note that can always approximate fM using fM (Hs) where H, (1-8)II HI H + òla. for which fM (Hs) & - (M, HM, (1 8) H
ôd + fM (H) as d + 0. This implies that HEH J' (H) 2 HeH JM[ H) and thus the optimal objective values of the two optimization problems are the same. Now define H = (M (La IIM) E whose inverse is HM (M i (La HIM) then HM > 0 solution to the optimization problem in (20) because fM (HM) = (M, (M + Tr(Px(MM) = inf Hnf (M) +IM. Since H is cone, the optimality of H implies
The gradient of Í' is given by (H)(H)
that for any H H.
0 0= (V H HM) = (- H- + M,, H HM)-
By the definition of HM. we have HM MHM PR(M) (M)² because span(M) S (pan(Pu(M)), There- fore, it follows that
= (Ps(M))MPP(M)) IM, H Px((M - - (Ia-IM))
(PM(M)))M(MMM) IIM H Px((M))
Proof for the operator monotonicity of the extended Pv. For any M1 2 M2 2 0, it follows from Proposition A.3 that for any > O, PM(M] + +ca) PH(M, + cIa). Taking the limit as + 0 on both sides yields (M) ł Px Similarly, as (M + +ela) Px (M2 + ela) E H for every e 0 by Proposition A.3, we have (+(M) (M lim,-+0 Px + eIa) - Pr (M2 + e) E H. because H. is closed set. This completes the proof.
A.3. Adaptive gradient norm corresponds to the dual norm
Proof of Lemma 3.3. We begin with the case of Cauchy-Schwarc inequality, it holds for all H H with Tr(H) S 1 that
Fix any g E Rd. First for any w E Rdd with |/l] S 1, by
g TTw w S Vg H gV w Hw S (99 .. .H- -/w S (gg HH-)
H
where the second inequality follows from the definition of ||w|| н. Now, further taking infimum over H E H with Tr(H)= < 1 and then taking supremum over Rd with ||l]|x < 1, obtain that
|ل
inf HER
(gg .. -H1).
(21)
Next, we need to show that the above inequality holds also for the other direction.
By Proposition A.4, we define H. Tr(Px(gg T) PH = inf SE H-1), We choose correspond-
H:g which satisfies that || w. 2l = 1. Then
ingly w. |H]a] xx
lg 1 2 g w. 9 H.g
(22)
|H 1l
Therefore, suffices to show that T w. 2 / gg Hig. which is equivalent
g Hig 2 High
g T T..H...
sup HEH
16

---

Structured Preconditioners in Adaptive Optimization: A Unified Analysis
By the property (d) of (g( ) from Proposition A.4, have (Pul99 T) ) Pxlgg") D IT gg H. Pn( [) for any H E H. Rearranging this equality, we obtain
g T Pulgg P HH ''g - g Px ') gg+ + (ggg Px(gg ))
= g Pn(gg ')g + (II 99 H Tr(P T)
where the second equality is because gg is exactly the projection matrix onto span(PH((e )). Applying the above equality with Tr(P:(gg))) )) HH in place of H and plugging in the definition of H., we further have
H g H. (99 () 88 H (99 )) 1 g H.g + (rg )) Tr H (gg ))²
g H.g.
This implies that for any H H with Tr(H) z 1, we have g H. S g Hig. Therefore, it follows from (22) that
g H
Igli 2 g Hig
inf H1 9 H. H.H.
2 g H'g
[gg HH- HH
inf
(23)
where the equality follows from the definition of H.. Finally, combining (21) and (23), we conclude that
inf HERI
V [ggT,H"* HH-1)
(24)
For the general case gızt (g1.....g)) Raxt note that for any E
Σ 9-9. = (vec(gia)) ved(gia)) T H @ b.) (vec(gi:)) vec(gn.)) (H @ 3.))
Then applying (24) with vec(g1.)) in place of and H. @ I, in place of H. obtain
inf HOLEHOL. / (veo(gi:)) voo(gna) (H @ I.)-1)
veo(gn:)ll
inf (H),
V ((veo(g1:)) voc(g1:) (H & I.)-1)
{ 9.9. 9 H-1
=V. inf HEH,T((H)S)
where the second equality is because Tr(H & I, Tr(), = Tr (H). This completes the proof.
A.4. Examples of ill-structured preconditioner sets
1.4.1. PRECONDITIONER SET OF TWO-SIDED SHAMPOO
Consider H {H E St H U & V for some U E Rduxxd, V E Rddxdda. In particular, we consider the special
case of d, = dg = 2, and the following two matrices
(e) diag(L,
G2(e) e,,
17

---

Structured Preconditioners in Adaptive Optimization: A Unified Analysis
where 0 is small constant to be determined later. Clearly Gi(e) < G2(() when e S 1, but will show that
())()) G.(e) does not hold for sufficiently small > For any H U V H, holds that both U and V
PSD, and we explicitly parametrize them
u1
U2
U1 1/
V
U2
U3
U2
where u1, us 2 4143 2 užz and similarly, v1, U3 2
2 Dg
Correspondingly,
01
U2
02
H-U
U2
Uq U3
U2
Ug
Ug
U2
U3
Ug
U2t2
0g02
1g00
Ug03
We first analyze H.a(e). For convenience, consider
Ha arg min (G((),, min( H) TT(H-
HER
arg min (G2()) & V) + Tr((U V)-
U.V.St
=1oo( (U.V)
Recall the properties of the Kronecker product that (U & -)-1= U V-1 and that Tr(U-1 V-1) Tr(U-)).
Tr(V-l). Then further using the definition (2(c), we have
(U, U1 + Uz + Tr(U
( e ugU1)
U1 U1 + U333 + U1 +
U3
U1 +
U/
+ cc(u10 + Uzv1)
Ujtl3 u2
0123
where we apply the explicit expression of U-1 and V-1, First observe that to minimize [G.(e) (U, V), we must have
U2 U2 0 because otherwise tit
witus and similarly mnny
gi
Therefore, it suffices further
0193
0109
minimize the following:
U1 us
01 +
Us
fazl (U. = 01 + Ugpy +
+c(u10s+ + ug0))
43
F1U3
(U.()
For the first term fo, (U, V), we can apply the AM-GM inequality to get fa, (U,V) > (4 (l/utgeee)))// where the equality is achieved when 11 = Ug, 01 = v3. and 11 = v2. Moreover, for sufficiently small e > 0, the contribution of the second term c(u) 2/ + uz ) is negligible, and thus we conclude that lim, Hoz() x/ 44
Similarly, for Gi(c), we have
U1 + u3
01 + 1/
(U, V) = U1
+ <(u1u3 + U3th + uytz))
U133
U1U3
2/
2 Ug
+ Ug
03
<(u U3
Ug 2zU3)
U1
U3
0/
1
1
1
1
UgU1
c(u103 43 + U303)
11
U133
U3
U313
(U.V)
where the equality holds when U2 = O. The first term fa (U, V) attains the minimum value when wi = 1, and the remainder can be made small by choosing large ua when is sufficiently small. Therefore, we conclude that lim HG.(e) diag(1, 0).
Now comparing the limits HG, (c) and HG. (c) as e + 0, we can see that for sufficiently small e > 0, it holds that HG > [1, 1]. Hence, for sufficiently small > O, HG, (c) N HGa(e) does not hold.
18

---

Structured Preconditioners in Adaptive Optimization: A Unified Analysis
A.4.2. TRIDIAGONAL MATRICES
Consider H {H E S H is tridiagonal]. Here, for H to be tridiagonal, it has nonzero elements on the main diagonal, the first diagonal above the main diagonal, and the first diagonal below the main diagonal. We consider the specific example of 3-by-3 tridiagonal matrices, and examine Px(-) for the following two matrices
(10000
1
M =
M'
1
2
It is clear that 0 < M 3 M'. We numerically solve for Px(M) and (x(MM) to get
1.382548
0.297594
0
100.000004
0.007229
(x(M)
0.297594
1.318491
0.297594
P2
0.007229
1.365999
0.366002 1.366032
0.297594
1.382548
0
0.366002
Note that (M) I (M') does not hold because the last diagonal entry of (MM is larger than that of (+(M))
B. Calculations for Examples of Well-Structured Preconditioner Sets
As mentioned in Section 3.3, we provide calculations for how derive each specific algorithm from Algorithm with specific choice of H and explain each entry in Table 1. Recall from Algorithm I that
M, = -Ia+ + Σ 9.9.-
8=1
B.1. AdaGrad-Norm For AdaGrad-Norm, we have H = {c. Ia c 0).
Calculation for H., Then for any H = c Id E H.
(M, (M. Tr H) = TY(M) + ofod 20 / d. Tr(M.)
where the equality is achieved by choosing
Th (M))
1 1 d n
Σ 11g. 12
This corresponds exactly the update rule of AdaGrad-Norm. which adjusts the global learning rate based the accumulated & norm of past gradients.
Calculation for 111x-- For the associated || || we have
|
|
H
sup USes1/d
Vc.z x L5x
Vd
Calculation for
Igi
We can caleulate the adaptive gradient norm similarly:
t
inf 0/o
Σ g.g. (cla)-1) = vāl
Σ 9 Vd,
11g.11
B.2. Diagonal AdaGrad
For diagonal AdaGrad, Dd (diag((),..., CI...., ca 2 0}.
19

---

Structured Preconditioners in Adaptive Optimization: A Unified Analysis
Calculation for H. For any H =diag(1,..... ..... Ca) E H. have
MG
(M, +
Tr
2 Σ 20y / M + +
where the equality is achieved by choosing
M
+
for = 1, I,...,d. d.
This corresponds to the update rule diagonal AdaGrad, which computes the historical sum of squared gradients for each individual coordinate.
Calculation for
||
1/
For the associated
|| |1 we have
|
sup ΣL asl
x: T diag(ee,........ ca)a
sup Ei e51
Σ Gaf
max iE [d]
[
|
|
Calculation for 191.e1l
We can calculate the adaptive gradient norm similarly:
|1g1: = inf
gog.
Ei= e.s
d
inf Σim 455
g
C2 s=1
where the last equality is because of the Cauchy inequality: for CI, .I....y Ca 2 0 such that Dd $1
d
G
2
gz..
i=1
i=1
==1
==1
B.3. Full-matrix AdaGrad For full-matrix AdaGrad, H = S
Calculation for H., Note that (M, H-1) + Tr(H) is convex function of H > O, we can get H, by first
calculating the gradient of the objective function and then setting it to zero. Specifically, for any H > O, we have
VH (M,, H 1) + (T(H) = - "MHH" + 2
Setting it to zero yields
H. - = elu + Σ 9.g.
This corresponds the update rule of full-matrix AdaGrad.
20

---

Structured Preconditioners in Adaptive Optimization: A Unified Analysis
Calculation for llx--
For the associated
||
1/,,, have
| =
sup HE)S
Vx H
sup H2)
(xx) T, H.
(1xx)) = =||0|12-
Calculation for 2l of Il: l to get
We can adapt the caleulation for H, to the constrained optimization problem in the definition
g1 l1
inf HEHIT(H)
Σ 9=g. H-1
[g00. # -11
)
1
Σ 9-9.
Tri(Ex= 9
Tr (Σ 9.9s
$
s=1
B.4. One-sided Shampoo For one-sided Shampoo, H. - @ Idz-
Calculation for H.. For any H Hr @ Iag H. H --1 =H²² 8 Idg. so we have H)) dn Tr(H,, and
M2 H (M, [HL & Idn)
(ela, H, & Iaa)
vec (e.) H ldg)
=dd Tr(H
Tr(G, H
=Tr (edRTa +
ΣG.G. H
Hy & Lag) (xx H2) for any X E Rdixda, Again, since the objective function is convex in H2 0, we can derive H2 by first calculating the gradient and then setting it zero. Taking derivative with respect to Hi, i. we obtain
where we use the fact that (vec (X) vec (X)
H (M, H + Ir(H)) -HH (edRts. + EG.. H + +
8=1
Setting we obtain
1( edRtd, + +.G.G.) s=l i.C. a=1
1
orgmin(M,, (Hy & Iaa) + 1 Tr(H, s Ida) HrES
n/dx
1 ela, + 1
dR
This is exactly the L, in Algorithm 2. Therefore, the preconditioner H, in one-sided Shampoo is given by
17 (ca. + dR s=1
1
E..c. & Ian
H. As result, Algorithm with H Sdh 8 Idp recovers Algorithm 2.
21

---

Structured Preconditioners in Adaptive Optimization: A Unified Analysis
Calculation for 1lx-
For the associated
||
1/,,, have
|
sup Hlag H. 200,T(HH)))
H
V /(vec (X) vee (X) H. & Ian)
sup HIS
((XTT Ht)
sup
X
VdR
up
where the last equality is achieved at H uu provides a proof for Lemma 4.1.
/dR for u being the leading eigenvector of X.X The derivation above
Calculation for +| Again, for the adaptive gradient norm, we can adapt the calculation for H, to the constrained optimization problem in the definition of ||lgi: x to get
Σ 9-9 9 H-1
inf HEH.T(H)S
H-
inf
Σ G.G.,HH.. H.
Σ E..0 )}
Σ |EGGG G.G.
1
dR dR G.G.Y
= < ((
=T> | $ - ||
[
< G.G
dR
s=1
C. Analysis for EMA adaptive optimization
As mentioned in Section our analysis can be generalized Algorithm 4, which is an EMA variant of Algorithm 1. To simplify the theoretical analysis, we instead focus Algorithm in which the current gradient is not scaled by 1 B3 when computing M, It serves as a bridge between Algorithms and 4. When B3 1, Algorithm 5 exactly recovers Algorithm I.
When the learning rate 7j is rescaled by VI- 332 and is rescaled by 1/(1 32). Algorithm 5 is equivalent Algorithm 4. Theorem C.1 extends Theorem 3.4 to get online regret bound for Algorithm 5. The proof is in Appendix D.3.
Theorem C.1. Let H be a well-structured preconditioner set under Definition 3.1. Then for any corvex loss functions In,..., Tor, the regret of Algorithm 5 compared any x² X can be bounded as
T
D2
Σ32
B
[i((ar) (((x)))L((x')) <
+1
inf (ER)
(MT,H-)
2îj
t=1
where D = maxte(M] |
x= 22
Theorem C.2 provides the convergence rate for Algorithm 5, which generalizes Theorem 3.8. The proof is in Appendix D.3. Theorem C.2. Let be a well-structured preconditioner set under Definition 3. I. Consider any independent stochastic convex loss functions Li...., LT satisfying Assumption 3.7, and let H (L, H) be the H-smoothness of their expectation L. Suppose the global minimizer of L, denoted by x", is in X. Then for the iterates æI,..., Er of Algorithm 5, denoting T-t F1. D1
Tis
Dt. it holds that
16
4
4/2
Bi(+1T)) - 1x(x")
||
H(L H) +
X
2 o +
|| dve
T
Σ
82
2
ΣT_ 35
Σ
82
22

---

Structured Preconditioners in Adaptive Optimization: A Unified Analysis
Algorithm EMA
Adaptive
Regularization
Meta-
Algorithm 5 Weighted Adaptive Regularization Meta- Algorithm
Algorithm
Hyperparam: > 0, convex set X' _ RR, learning rate
Hyperparam: e > 0, convex set X' C 12d, learning rate
7, preconditioners H C S 32 E (0, 1)
ft. preconditioners H. C S 32 E (0,1)
Input: initialization æ1, loss functions (I+)/=1 :Rd -R
Input: initialization loss functions {T,) :R
Mo ELe
M. + eld
for i 1,2....,7 do
for / =
T..,,aer do
9: <- VL((x,)
9: t- VL,(Tt)
M. < 82M-11 + (1 32)9:91
M, +- B2 9.9,
H, t arg min ( Mr, H H + Tr(H)
H, + arg minh M,, 1t-1) 17 TP()((
1t+1 +- : (o) H g.)
3+41 II (Ār H g1)
Return 2 FT
Return ,,,..., WT
where o = inf HEH, V H
When 32 Theorem C.2 recovers the result in Theorem 3.8 and provides O(T convergence rate. When 82
the optimality gap is upper bounded by 0((1 32) || (L, H. /1 B2 || σ) when T θ 1-
D. Proof for the Unified Analysis
D.1. Regret bound
We first present the proof for the main result on the regret bound for Algorithm with well-structured preconditioner set H.
Theorem 3.4. Let H be well- structutrd preconditioner under Definition 3.1. Then for any convex loss functions Li...., Lt, the regret of Algorithm I compared any x X can be bounded
D2 + (G+
Σ Lx(xx)) - EI(x*)
2n
t=1
where = ellgg.r 1 D max(e(T) || - x² 11
Proof of Theorem 3.4. First we will analyze the property of each H,. Recall from Proposition 3.2 that He satisfies (M,, H,") = = (H) and that H. Tr(H,) -1 H. arg min (M,, H H" Therefore.
Tr(,)) M, H = 12 = / (T.HH) n/ (M,, H.
-
(25)
Now recall the regret bound from Theorem 2.1:
T
((x))
Σ
L((x') 1 (Mr, H + + Tr(Hr) n7² Tr (HH))
1=1
1=1
+ Σ
|1 x x011 -1 - x 11212
1=1
27 V (Mr, H 17² Tr Ho)
|| 20112, || x² x²112r,
(26)
23

---

Structured Preconditioners in Adaptive Optimization: A Unified Analysis
where the equality follows from the facts in (25). Next, for the second term the right-hand side of (26), we rearrange the summation to obtain
||= x²IH.. - 1114411 x²1HA H.
- 121
x
+
ax x=11... - -112: x: 2211)2...)
1=2
= =|111 x"Il
2 x" 11 H. HH--
Notice that M, -M -9 2 0, and thus H, Hi-1 E H by Proposition 3.2. This implies
x x=1'= [| H.-H.., (z, x x²) (H, (z, x²) Tr(H. H || x=1
where the inequality follows from the definition of || || in (1). It then follows that
T
T
Σ | - x²11... - 11 - x²IlH.) - Tr H1) || - x=11
+
Tr(He - Hi-1) [ - xx ==1
1=2
< (T(HT) max IstST
|
x011
x
(Mr. H
max ISeST
||
x= x11
where the equality again follows from (25). Plugging this back into (26), we obtain
T
T
Σ L((x)) - - L((x)) & nV (Mr. HT') + 21y V
(Mr. H
||00 - x= 12
max ISAST
1=1
1=1
maxi<I<T ||0, x 112
+11
inf HEH,THHHSS
v VMMr H-
2n
maxigeST
||ar x
x0112 + +11
/Mr Mo. H11) + (Mo, H-1)
2n
for any H E H with HH) = In particular, we choose H aH, ي La for some a E (0, where H₃- arg HEw. / (Mr MM. H 17, (Mr Mo. H I|r according to the definition of the adaptive gradient norm (7), Since H- ( and H -1 Ia. we further have
T
maxiseST |
d
L (
Σ 1=1
[x(x)) x
+
V
(Mr Mo, (HT) +
(Mo, Ia)
2n
1
maxi<isT ||,,
x=
x=1
+ +1) m)
d' I1l1.r 11
2m
1 1-a a
Finally choosing a [ ll n +dve
1 ll
we obtain
maxISIST ||0, x 10
(xx))) Σ lx(
+ +1|
(1lg) 11 dve)
21
This completes the proof.
The proof for Corollary 3.5 is straightforward. Corollary D.1. Under the setting of Theorem 2.1, further suppose that X is bounded set in Rd Then choosing
V || the regret bound for Algorithm becomes
T
T
Σ
Lx(x)) Σ
5222)1222)111 (G +dve)
mm1
24

---

Structured Preconditioners in Adaptive Optimization: A Unified Analysis
Proof of Corollary 3.5. Note that D maX
a
x""
m8x(e[[]] | + || 2 21 because
" ar and x² are all in X. The proof completed by setting r = /2 || sc to minimize 2 211x111 + n
D.2. Convergence rate
Next, we present the proof for the convergence rate of Algorithm with a well-structured preconditioner set H.
Theorem 3.8. Let H be a well-structured preconditioner set under Definition 3.1. Consider any independent stochastic convex loss functions LI...., Lr satisfying Assumption 3.7, and let H (L, H) be the H-smoothness of their expectation L.. Suppose the global minimizer of I,, denoted by x², X. Then for the iterates ;1,..., of Algorithm I, denoting -T
T:T + +2 2t, it holds that
IE (( L(x')]
16
|2 (L, + 4v2 T
4v2dve
VT || 0+ +
||
T
where a= = HEH, V (2, H-1).
Proof of Theorem 3.8. Let H E H be the matrix given in Definition 3.6 for the loss function L, such that V2I(2)) 3 H Therefore, for any x, x' E X.
1 (a' x.) H (a' - x) I(x') S I(x) + V[(x) (x' xx) +
1 (x' x x (H") T H' (x' x + +
= L(x) +
2 (Hr)
Then we have that
[(x") = min L(x')
(("")"VVViil T H' (a' - x ((+)) 1 (HT) - (x'-
S min ((x) (x'
2
= L(x) 1 (L(x) -V()))
Applying the above inequality to Ti..... xy and denoting 9: = (L(x)). we then have
T
T
Σ (L(x)) L(xx))
9:
(H*)--
gt
E
9(9.
(H (H')
2
1=1
E
Σ 9:9.
(H'
HH (T((H)) H H 1)
(Tr(H
1=1
inf (L, H) THH
E
999
H
(27)
For any 8 > 0. we choose Hg E H such that (B J.g. g < 18 inf (B]_ ~.G. ], . H and
Tr(H,) S 1. Similarly, we choose Hy H such that ( H S 6 te ( H 1) and Tr(Hs) S 1. Correspondingly, we define H' H + (1 a Hy, which satisfies that H' E H and Tr(H') ) 1. Then by Jensen's inequality, we have
E inf HEH.THHH) V /MT Mo, H 1 < EM Mo. (H') S (M Mo. (H)
(M
25

---

Structured Preconditioners in Adaptive Optimization: A Unified Analysis
Plugging in Mr- Mo Σi= 9:9i' since E E
3 E
EE
9:9 !
TE, further have
T
inf (E))
(MT- Mo, 1-
9:91
(H')-1
T (2, (H")-)
T
9:
Hg
+ 1-a
T (Σ, H5'
T
9:9
.gg
+ T (T, H
1=1
where in the last step we choose a to be
[8] E i-= 9:96 J. Hgg 2
a
(EL D 9:9. ]. Hg (E, H
Recall that infHENII(H))) ( H 1). Then it follows from (27) and the definitions of Hg. Hy that
T
E
inf ())
V( Mo. H-1) <
2H( H.
Σ (L,(z)) L((x-)) -,,()))))xxx)))) + +
vT²² T
Since > 0 is arbitrary, this implies that
T > ( + VTo2
inf (H)
V /Mr
Mo. H
2H, H) E
Now recall the regret bound from Corollary 3.5, and follows from the above inequality that
Σ (L((m:))(((x= Le(x=')
22
inf HE))
(M. Mo, H-1) dve
T
(v
Σ (Ta(x:) - 1e(x*)) + VTo² 1702+av4 ave)
[H, H) E
[=1
Solving the above inequality yields
1 T T
16 $ T
4v2
4v2
Σ
(L((x)) - Te(x'))
|
(L,
+
||
0 +
|| dvc.
VT
T
This completes the proof.
D.3. Analysis for EMA style optimizers
Here we present the proof for the theorems in Appendix C.
Theorem C.I. Let H be a well-structured preconditioner set under Definition 3.1. Then for any convex loss functions Li...., Lr, the regret of Algorithm 5 compared any x X can be bounded as
Σ
T
3
-ame))Lx(xx) <
D
+
inf HEH V (Mr,
27
1=1
where maxte[T] |a, x 201
26

---

Structured Preconditioners in Adaptive Optimization: A Unified Analysis
Proof of Theorem C.I. We can choose L, = 35 L,, = and 1j. Then we claim Algorithm 5 is equivalent to Algorithm with hyperparameter 17 and loss functions {Le) which will be shown by induction.
Assume á. x, for Σ t. Then we know 9. (L,( 322 $ VL, for S t.
We consider the update in the step L of Algorithm 5.
M
Σ
82 9:9, + +, = B/2
Σ 1=
9:9: + eld =3yMe.
==1
H. arg min HEH
MG H 1/2 Tr(H)
mig (5rM,, H 1) + Tr (H) = =3 H, HER
H.'g. 9: -32 H. 33 gi H gD.
き.+1 = # (a. ( H 9.) = 11 (x. - H g.) 11 (xt - Hi 9i) = 2c+1
Therefore, obtain the regret bound for Algorithm 4 with Theorem 3.4.
Σ 32 [x,())) 1x(x)))
(L((r))) [e(x=)]
maxIsIST
|
x011/
+1
inf HEH
V(Mr, H H-1)
2n
Bg &
maxI || x=
x=112/
17)
inf HEH,T(HHS / (Mr. H-)
2îj
and
82
1 322
1/2
(zx(2x)) Tx(
1 1
/B2
maxi SIST
||x,
x=11 11
+1
inf HEH.TRMS) ( H-1)
$7 4
$3
27j
Theorem C.2. Let H. be well-structured preconditioner under Definition 3.1. Consider any independent stochastic convex loss functions Li...., LT satisfying Assumption 3.7, and let H (L, H) be the H-smoothness of their expectation Dr of Algorithm 5, denoting L. Suppose the global minimizer of L, denoted by x", is in X. Then for the iterates I1....,Zy
ē1. -E 32 D B2
Et=1
it holds that
EL(B1T)- ((xx') 5
16
|
H(L,H) H) +
4
|
4
| de/e
1=
B2
1==1
3/
ΣT 1= $2
where a infHEH,THHHSI H-1).
Proof of Theorem C.2. Similar the proof of Theorem 3.8, know
L(xx) - L(x") 2 (gg.Eg, (H1))1)
and
T Σ 3 L(ã,) I(#")² >
T
B Eg. (H))-1
1=1
1=1
> 2H. H) (HH)SLHEN
inf
Σ 82
Eg.Eg. HH!
27

---

Structured Preconditioners in Adaptive Optimization: A Unified Analysis
For any 8 > O, we choose Hg H such that Tr(
Tr(H., Hg) S and
Σ
B2
T-e Eg. Hg
s inf
Σ
B2
T-t
Eg.Eg .H
HEH,TAH)
We also choose Hy E H. such that H ) I Hnf H' aHg + Hy. Then H' E H and Tr (H(HH) S 1. We have that
Σ, 1 and Tr(Hz) S 1. We define
E
inf HEH V (Mr- 87 Mo, H-1 1)
B
M 8I 88 Mu. (H1))1)
E
MT 32 Mo. (H')-1)
T
Eg, B0. (H')-1
+ Σ
3 32- ( (H)-4)
t=1
Now plugging in the definition of H', we obtain
E
inf HEH,
/ (M. My,
Mu, H H-1)
3/ Eg. (
3T- { a -1)
1=1
1=1
T
1
82 Eg. (Hg
+
22
T-t E HH -1)
t=l
T Σ
Σ
82
Eg. Eg. (Hg)-
3
T-t
E
(Hx) -1).
1=1
where in the last step choose to be
> 21=1 822"
Eğ. (H
Et=
1==1 B2
Eg, Bo (H₂) + Vi= 1--- (Hz)-1)
Then we have that
T
T
inf HEHI
/ (MAr Mo, H H-1) s
(H,
E
82
T [Tx(x))) -a(o'))
+
Σ
8
1=1
t=1
When taking to have that
E
32
T (1)))) - [x(00)]
5 522DD
2H((.H) H)
E
32
T
[L((x))
Ix(
Σ
- B3
dve
1=1
1=1
and
T
T
E
3 1=1
T-t
[c((xs) - Lx(x'))
< ((,
H) +
2
2 Do + 4V2D33 d.e. dve.
1=1
28

---

Structured Preconditioners in Adaptive Optimization: A Unified Analysis
T-t
If we choose #1.-
32
2t. then from the convexity of L we know that
mm
Σt Ets1
82
B[((EIT)) - L(x*)| S
T-t 32 2
L((xx)) - Le(x)))
Σ 21=1 22
T
16D 16D² (L, H. +4, 2
2
32
3T Do + 4V2DD dve
1=
T
$
1
t=1
16
4v2
Do + 4
(L)
Ddve
Σ
7
32"
P-t
Σ t=1
-T
Σ 1=
322
1=
This completes the proof.
E. Proof for One-Sided Shampoo
E.1. Proof for left smoothness
Lemma E.1 (Left smoothness for one-sided Shampoo). Let H She @ Idn be the well-structured preconditioner set for one-sided Shampoo. Then the H-smoothhess HLL, H) defined in Definition 3.6 is equal to the smallest number H 2 0 such that there exists H dz E R& xd2 satisfying that H = dn Tr(Hā,) and that for any X. A E R& xdg,
|V2 A] S (H AA
In this case, the H-smoothhess is denoted by Hiett (L.)
Proof of Lemma 4.3. First, for any X, A E Rd xdR and x vec( X), we have
vec vec(A) L(x) vec(A) = (2I(X)(A, A].
Therefore, given the Kronecker product form of all H H. find H E with the smallest trace such that H = 4 H, it is equivalent find H Hag Lig H with the smallest trace such that for any A E Rduxda,
L V A] S vec (HHH & LdR) vec A)
TT T "HH A)
H AAT) T).
Further note that (H = dR Tr,Ha,) and thus we conclude that is equivalent to find Ha, E S with the smallest trace such that the above inequality holds for all E BRdx This completes the proof.
E.2. Proof for regret bound
Theorem 4.2 (Regret bound for one-sided Shampoo). For convex functions Li...., Lr, the regret of one-sided Shampoo (Algorithm 2) compared any X x E RR xdn satisfies
T
D2
(X,))
Σ [I(X'
<
op
+1)
(G (G+dVC),
2d
t=1
where Dop -maX,E(T)) ||, - X llop and G = vdR Tr (EIn
When the domain X' is bounded in operator
norm, i.e. 1|ol -0, further choosing n /2/d (1lpp it holds
(Xi)
Σ L(XX'
1=1
1=1
1
22/211X811o Tr
Σ G.G.T
ve
d
Vdn
29

---

Structured Preconditioners in Adaptive Optimization: A Unified Analysis
Proof of Theorem 4.2. We will apply Theorem 3.4 one-sided Shampoo. According to the analysis in Appendix B.4, Algorithm 1 with H. = (R22 xd2
and @ Idn) n n recovers one-sided Shampoo. We can plug in ||~|| || op
vdR
11g1 /d Ir ( (Σ G.G.)
into Theorem 3.4 and get that
Do
+ +n) n)
c+ min (dve, de 2G ))
Σ (x)))) Lx') V2
2d
with Dop = max | X' and G = Vd Tr C.G.G.T.
F. Additional Results for Experiments
E1. Efficient implementation of full-matrix AdaGrad.
Directly applying full-matrix AdaGrad to this 10 of--dimessional problem is impractical. Instead, we consider the eigende- composition of H to be U EU and define the transformation T(X) UU. We further define orthogonal matrices Vi...., Va E Raxd such that the first row of V, is in the same direction of T(X* (X+-Xo);; and define V diag(V,. (V..... Va). We can know that vee(T(XX Xo)=
v X e1 where v; [| Xll for i [d].
Then it holds that
= = (XXXXXXXXX X1) -X-XXX-XXXXXXXXXX UX-)
- -(X
(E la. vec VeC T(X - X veO -XX
(E Ia. V vec (T( X*) -x))
Further denoting a = V (x)) and ÿ = ved (Tc() J), then we obtain
d
f(X) = Σ
(i-1)a44 - (v @ ()) ?
1=1
Σ 0i (#(1-1)d+1 Ū(--))&+1 u;)2
Σ (a 1)=-)))+)2
Running full-matrix AdaGrad on f(X) starting from Xo can be implemented equivalently by using क as variable starting from y. Only ((111)411 will receive non-zero gradient so full-matrix AdaGrad actually only cares these d coordinates, which reduces the original problem to a problem only with d variables.
F.2. Results for EMA algorithms
As mentioned in Section 5, we compare AdaSGD, Adam, one-sided EMA Shampoo and full-matrix AdaSGD, which are EMA version of AdaGrad-Norm diagonal AdaGrad, one-sided Shampoo and full-matrix Adagrad. The results are plotted in Figure 2. We set 32- 0.95 and disable first-order momentum, i... 31 in Adam.
We tried 60 learning rates between 1 x 14 and 1 x 10 The relationship between loss and learning rate is shown in Figure 3.
30

---

Structured Preconditioners in Adaptive Optimization: A Unified Analysis
100
Average Iterate Loss
10"
Last Iterate Loss
10
10
10
Ada5GD
AdaSGD
Adam
Adam
10
Full-Matrix AdaSGD
Full-Matrix AdaSGD
One-sided EMA Shampoo
One-sided EMA Shampoo
20
40
60
80
100
20
40
60
100
Step
Step
Figure We plot the last iterate training loss f(X,)
(H. X*)T) and the average iterate training loss
Σ
82
X.)) over steps for optimizers obtained from Algorithm
10a
10a
10a
Average Iterate Loss
102
Last Iterate Loss
101
10a
10a
10-1
10-1
10-2
10-2
1
10-3
10
10
10
10
10a
102
10
10
10
10
101
10a
Learning Rate
Learning Rate
AdaGrad-Norm
AdaGrad
Full-Matrix AdaGrad
One-sided Shampoo
10a
10a
10a
Average Iterate Loss
Last Iterate Loss
10a
10a
10a
10-1
10-1
10-2
10-2
10-3
10-3
10:
10
100
10
10
10
10a
10a
Learning Rate
Learning Rate
AdaSGD
Adam
Full-Matrix AdaSGD
One-sided EMA Shampoo
Figure plot the last iterate loss and average iterate loss versus learning rate, For each learning rate the plotted value is the average last iterate loss and average of average iterate loss across five random soeds.
31


Reviews:
[{"review_id": "bqPTcXhVtr", "reviewer": "Reviewer_N1Ca", "summary": "The paper provides regret bounds for a family of adaptive algorithms with structured preconditioner matrices. The analysis generalizes the technique introduced in the original Shampoo work and applies to Adagrad, Adagrad-norm, Adagrad-diag and one-sided Shampoo. Intriguingly, the paper shows that, for certain loss functions, the bound provided by one-sided Shampoo could be tighter than the bound for full-matrix Adagrad and also demonstrates this empirically on a simplified loss surface.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for acknowledging the value of our unified analysis. Below we will address your concerns.  \n\n**Comparison with Shampoo and Shampoo^2:** We appreciate the suggestion to include two-sided Shampoo and Shampoo^2 in the experiments. As reviewer tCyu pointed out there is lack of motivation for our current experiment, we conduct experiments on a better setting which is more practical and better motivated by our theory analysis. We will add this experiment and focus more on it in the future revision.\n\nWe consider a synthetic linear regression problem $||\\mathbf{A} \\mathbf{X} - \\mathbf{y}||\\_2^2$ where $\\mathbf{A}$ is the data matrix and $\\mathbf{y} =\\mathbf{A} \\mathbf{X}^*$ is the label vector generated by ground-truth $\\mathbf{X}^*$. Thus, the loss function can be equivalently written as $f(\\mathbf{X}) = \\langle \\mathbf{H}, (\\mathbf{X}-\\mathbf{X}^*)(\\mathbf{X}-\\mathbf{X}^*)^\\top \\rangle$ for $\\mathbf{H} = \\mathbf{A}^\\top  \\mathbf{A}$, which is the same function that we studied in section 4.3 and show that 1-sided shampoo outperforms other adaptive algorithms. We consider $\\mathbf{X} \\in \\mathbb{R}^{d \\times d}$ with $d=10^3$. We set the eigenvalues of $\\mathbf{H}$ by $\\sigma_1 = \\cdots = \\sigma_{10}=1$ and $\\sigma_i = \\frac{1}{(i-10)^2}$ for $11 \\leq i \\leq 10^3$. Each element of the solution $\\mathbf{X}^*$ is independently sampled from $\\mathcal{N}(0, \\frac{1}{d})$. \n\nWe compared all the optimization algorithms by sweeping learning rate from 1e-4 to 1e2. The plots are [here](https://docs.google.com/document/d/e/2PACX-1vR37RVh5tPyZFbelExgmMfoXk_y7Egv0TvxXw4WqzVJ9zYHECzzRYjv-2zE_QGwcbkBYvXFCSSrweVl/pub). The results are consistent with the original experiment: (1). **one-sided Shampoo outperforms other variants of AdaReg;** **(2). Shampoo is slightly worse than one-sided Shampoo**; **(3). Shampoo^2 fails to make progress in optimizing this loss function.**  We are unsure why Shampoo^2 underperforms on this loss function but note that Morwani et al. (2024) also does not provide empirical evidence of the effectiveness of Shampoo^2. Investigating this behavior is beyond the scope of our current work, but may be a worthwhile direction for future study. The implementation of Shampoo and Shampoo^2 is [here](https://limewire.com/d/4fwuX#gtS4Jg1Wpi).\n\n**One-sided Shampoo as a new, improved algorithm:** We disagree with the reviewer’s opinion that this work doesn’t lead to an improved algorithm. To the best of our knowledge, we are the first to explicitly define and analyze the one-sided Shampoo algorithm even though it is informally mentioned on social media. And our section 4.2 shows that it improves original Shampoo theoretically and section 4.3 shows it can achieve the best rate on some specific loss functions, which is empirically verified by our added experiment above. \n\n**Plausible directions for designing better optimizer:** Thanks for this insightful question. We believe the definition of well-structured preconditioners can help designing novel optimization algorithms as discussed in line 261-308 though they may not necessarily be better. \n\nNotably, our framework already encompasses recent algorithms such as Adam-mini (Zhang et al., 2024) and Adalayer (Zhao et al., 2024) via specific choices of the subalgebra $\\mathcal{K}$. Inspired by these two examples, we identify three basic operations that can construct new matrix subalgebras that can be used for defining well-structured preconditioner. \n\n1. **Direct sum** of matrix subalgebras. This is defined in line 261-308. \n2. **Kronecker product** between a matrix subalgebra and identity matrix. For a matrix subalgebra $\\mathcal{K}$, we can define $\\mathcal{K}’$ by $\\mathcal{K} \\otimes \\mathbf{I}_d = \\\\{ \\mathbf{A} \\otimes \\mathbf{I}_d | \\mathbf{A} \\in \\mathcal{K}\\\\}$. \n3. **Rotations** via orthogonal matrices. For a matrix subalgebra $\\mathcal{K}$ and an orthogonal matrix $\\mathbf{U}$, we can define $\\mathcal{K’}$ by $\\mathbf{U}^\\top \\mathcal{K} \\mathbf{U} = \\\\{ \\mathbf{U}^\\top \\mathbf{A} \\mathbf{U} | \\mathbf{A} \\in \\mathcal{K} \\\\}$. Such $\\mathcal{K’}$ is still a matrix subalgebra. \n\nThese operations offer a principled way to design new preconditioners and optimizers. As noted in our conclusion (lines 430–435), a theory-driven strategy would be to evaluate terms in the regret bound or convergence rate under different $\\mathcal{K}$ and choose $\\mathcal{K}$ to optimize the trade-off. While this remains challenging to evaluate empirically on large models, we view this as an interesting direction for future work.", "claims_and_evidence": "The paper mostly has theoretical contributions and provides a simple empirical experiment supporting the theory.", "methods_and_evaluation_criteria": "N/A", "theoretical_claims": "I took a brief look and verified (not thoroughly) all the proofs in the Appendix", "experimental_designs_or_analyses": "N/A", "relation_to_broader_scientific_literature": "This work is well placed as a generalization of previous analysis of structured preconditioner matrices.", "essential_references_not_discussed": "N/A", "other_strengths_and_weaknesses": "The main strength of the work is proposing a general unified analysis for various preconditioner matrices, and showing that the bounds for a more general class could be worse than a structured one. Although this should not be considered a weakness, the analysis does not lead to an improved algorithm. This could be considered as a part of future work.", "other_comments_or_suggestions": "I would suggest adding two sided Shampoo and Shampoo^2 (based on the recent connections shown by Morwani et al. 2024) to the experimental plot.", "questions_for_authors": "1. What do the authors think are plausible directions for using these bounds in designing better optimizers?", "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": "T8TFwanxgX", "reviewer": "Reviewer_25wV", "summary": "The paper studies preconditioned adaptive methods for online convex optimization in a unified manner. They define a particular class of “well-structured” preconditioners that recovers all 3 variants of the AdaGrad and the Shampoo algorithm as special instances. The main take away message of the paper is that by selecting the set of preconditioners in some “well-structured” sense, the metric for the norm of the sum of gradient outer products and the distance to solution could be determined in a principled sense for achieving better regret bounds (under certain scenarios).", "questions": "", "limitations": "", "rating": 2, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely thank the reviewer for recognizing the value of our unified framework. We will address your concern below.\n\n**Clarification of our results:** We never claim that the regret bound for the full-matrix AdaGrad is worse than AdaGrad-norm or AdaGrad with diagonal preconditioner and we apologize if any words cause such misunderstanding.  We only argue that it is a misleading belief that **adding more structures to the preconditioner will always lead to worse performance**. We use the specific comparison between one-sided Shampoo and full-matrix AdaGrad to illustrate this clearly. And we only focus on optimization speed, specifically characterized by regret bound or convergence rate. Other properties such rotation-equivariance are beyond the scope of this paper. \n\n**Common belief that full-matrix preconditioner is “ideal”:** First we would like to clarify we agree with the reviewer on the statement that full-matrix preconditioner is not necessarily ideal as there is no theoretical justification on that. However, this is a common belief held by many researchers in the community. For example, Agarwal et al., 2019 explicitly argues that full-matrix preconditioning allows for better exploitation of anisotropic curvature in loss landscapes and their experiments show that full-matrix version can work better than Adam in some cases. The belief that full-matrix preconditioners are “ideal” is also implicit in the design of many optimizers, such as Shampoo, and Shampoo^2, which all attempt to approximate the full-matrix version more efficiently.\n\n**Vector vs Matrix Problems:** The reviewer argues our comparison is unfair when using a matrix-structured problem to highlight the benefits of one-sided Shampoo. We disagree with this interpretation. First, vector and matrix problems are mathematically **equivalent** up to reshaping and vectorization. Second, the models in practice are often built by complicated matrix-based modules and thus our formulation is relevant to practice.\n\nMost importantly, we do **not** claim that one-sided Shampoo is better than the AdaGrad family in general. We only demonstrate that it can be better on certain losses since it already sufficiently supports our claim. We agree that AdaGrad family may perform better on simple vector problems, which doesn’t contradict our argument.\n\n**Comparison within AdaGrad family:** Our goal is **not** to provide a comprehensive comparison among all AdaGrad variants. As mentioned above, we never make any claims within the AdaGrad family. That said, we found that AdaGrad-Norm and diagonal AdaGrad can outperform full-matrix AdaGrad empirically from figure 1. We are also happy to clarify that AdaGrad-Norm can outperform full-matrix AdaGrad in some setting in response to the next question. \n\n**Can AdaGrad-Norm be the best:** Yes. We can show that the regret bound of AdaGrad-Norm is the best for a specific setting when $\\mathcal{X}$ is the $\\ell_2$-norm constrained ball. Due to space constraints, we are happy to share the complete proof in follow-up discussion.\n\nEven in practice, Xie et al., 2025 shows that AdaSGD (EMA version of AdaGrad-Norm) can outperform rotated Adam on a GPT-2 model. While rotated Adam differs from the original, it shows an important insight that supports our claim: **coordinate-wise adaptivity is not always better the global adaptivity**. We think this is even a better example than a regression problem with real data for which AdaGrad-Norm can beat diagonal AdaGrad. \n\n**Definition of Structure:** We agree that the notion of “structure” can be more clearly explained even though we have discussed this in line 29-40. In our usage, a structured preconditioner imposes constraints on the preconditioning, formalized via a subalgebra $\\mathcal{K}$. As mentioned in the abstract and introduction, some examples of structured preconditioner includes layerwise, diagonal and kronecker-factored. In contrary, full-matrix preconditoner is the least structured preconditioner because there is no constraint. \n\n**Experiments on Shampoo:** Please see the response to reviewer N1Ca. \n\n**Insights for Developing new algorithms:** Please see the response to reviewer N1Ca. \n\nReferences\n\n1. Efficient full-matrix adaptive regularization. (Agarwal et al., ICML 2019)\n2. Adam Exploits $\\ell_\\infty$-geometry of Loss Landscape via Coordinate-wise Adaptivity. (Xie et al., ICLR 2025)", "claims_and_evidence": "The main claim is that more structured preconditioners, i.e., the set of preconditioners which are relatively small (e.g., AdaGrad with scalar steo size or diagonalmatric preconditioner) could outperform less structured (e.g., AdaGrad with full-matrix preconditioner) versions. Their claim is supported for the Shampoo algorithm; they propose a new variant called one-sided Shampoo which essentially keeps track of only the “left-matrix”. They prove a regret bound which is better than the original shampoo by improving the dependence on the norm of the optimal solution (Frobenius to Spectral).\n\nThis results shows that a more principled idea of preconditioner design might help design new adaptive algorithm.\n\nHowever, it is not clear how would such a design process work. For instance, the authors claim that regret bound for the full-matrix AdaGrad is worse than AdaGrad-norm or AdaGrad with diagonal preconditioner. However, there is no rigorous axplanation as to how this is the case. In fact, comparing the regret bounds from the original paper for the composite mirror descent update version, full-matrix AdaGrad had better regret; (i) it is robust against rotations unlike diagonal preconditioner and (ii) measures the variance of gradients across all coordinates as well as jointly between coordinates to adapt the step size per coordinate accordingly unlike the fixed step size rule for AdaGrad-Norm. The authors must be precise with their claim and provide an analytical verification.", "methods_and_evaluation_criteria": "Yes. They compare methods using their regret bounds and dependence of those bounds to dimension, problem structures etc.", "theoretical_claims": "I have checked the proof in details until Section 4 and also went over the proofs for Section 4. I haven’t spotted any mistakes.", "experimental_designs_or_analyses": "Yes. The experiments show some interesting results but unfortunately they are based on synthetic data (and for a matrix problem). Therefore, it is not enough to make any meaningful conclusion for the algorithms. I would be very surprised to see AdaGrad-norm beating diagonal version for a larger vector regression problem with real data (full-matrix would be too expensive).", "relation_to_broader_scientific_literature": "Adaptive methods are the main workhorse in many large scale problems. For instance, adaptive methods outperform SGD for training transformer-based models. \n\n1.\tThe paper attempts to unite the preconditioning of adaptive methods and understand their regret/convergence bounds based on how the preconditioner class is constructed. \n\n2.\tThey explain the possible effects of structure (i.e., size of the set of preconditioners) under certain scenarios.\n\n3.\tThey give a concrete example for Shampoo algorithm and propose a provably better version in terms of regret bounds.", "essential_references_not_discussed": "-", "other_strengths_and_weaknesses": "Strengths:\n\n1.\tI think the authors’ attempt at unifying adaptive preconditioning using matrix sub-algebra on top of positive semi-definiteness is an interesting perspective into understanding the set structure of different preconditioned algorithm classes.\n\nWeaknesses:\n\n1.\tThe claims about AdaGrad-class of algorithms is not verified mathematically, as opposed to the proposed variant of Shampoo that has better regret than the original algorithm. Similarly, the experiments are designed for matrix problems and might not capture the advantages of preconditioning for AdaGrad family.\n\n2.\tAlthough I find the definition of well-structuredness genuine, it is not clear how we could use this framework to improve existing methods and develop new algorithm. The recipe is not clear, which the authors could explain more in the main text.", "other_comments_or_suggestions": "1.\tI don’t agree that full-matrix AdaGrad is considered as the ideal preconditioning method and there is no mathematical evidence that supports the argument. Please update the first paragraph accordingly or verify the claim with references, examples or derivations.\n\n2.\tPlease clarify the idea of structure for preconditioners clearly, before using the term for the first time in the introduction. In the current version, it is confusing as \"structure\" has different connotations in the context of preconditioned methods.", "questions_for_authors": "1.\tHave you tried running Shampoo and one-sided Shampoo? It would make more sense to see them side by side on a matrix problem, as it is not directly comparable to AdaGrad family in an apples to apples sense.\n\n2.\tHave you compared AdaGrad family of algorithm for a vector problem? It is more informative to observe their performance for a regression problem, for instance, with a real-world dataset in terms of the comparison of regret bounds.\n\n3.\tCan you prove that AdaGrad-norm has better regret bounds than other variants using your framework?\n\n4.\tDo you propose that there is a variation of AdaGrad family (similar to what you do for Shampoo) that will yield better regret guarantee?", "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": "IjE0eCB129", "reviewer": "Reviewer_tCyu", "summary": "This paper presents a unified analysis of adaptive optimization algorithms with structured preconditioners, challenging the assumption that better approximations of full-matrix Adagrad or less structured preconditioners always yield superior performance. The authors introduce \"well-structured preconditioners\" enabling a general regret bound, revealing a trade-off between domain metric and adaptive gradient norm. They demonstrate that a one-sided Shampoo variant achieves improved regret bounds and outperforms full-matrix Adagrad theoretically and empirically, suggesting simpler, more structured preconditioners can be more effective. The work provides a new perspective on adaptive optimizer design and insights for efficient large-scale training.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for acknowledging the theoretical insight of our work. We will address your concerns below. \n\n**A larger set of preconditioner is not necessarily better:** We appreciate the reviewer’s concern regarding our claim. Due to computational limitations, it is indeed difficult to compare full-matrix AdaGrad with those structured preconditioners on real-world tasks. However, our intention is not to assert that simpler preconditioners always perform better, but rather to challenge the assumption that **a** **larger set is always superior**. We successfully show it is wrong via the theoretical analysis and synthetic task. \n\n**Practical advantage of one-sided Shampoo:** Thanks for the clarification. We realize that our paper might have unintentionally suggested that one-sided Shampoo has practical benefits on real-world tasks, which we do not claim. Our work is mainly theoretical, and our experiments are limited to a simplified setting that aligns with our analysis. We agree that demonstrating practical benefits would require larger-scale experiments, which is beyond the scope of this submission. We leave this for future work and welcome pointers to such work if it already exists. However, we want to clarify that the theoretical analysis (Theorem 3.3 and 3.7) holds for any convex functions so the analysis of one-sided Shampoo is not restricted to a specific type of function (equation 7). \n\n**Motivation for specific quadratic problem:** We propose a simple quadratic function so that we can easily compute those convergence rates in table 2 and verify the effectiveness of our theory. We agree it seems unclear why to choose such specific quadratic problem so we decide to add a new quadratic loss function on which the empirical results are consistent. We only introduce the loss function below and more details can be found in our response to reviewer N1Ca. \n\nWe consider a synthetic linear regression problem $||\\mathbf{A} \\mathbf{X} - \\mathbf{y}||\\_2^2$ where $\\mathbf{A}$ is the data matrix and $\\mathbf{y} =\\mathbf{A} \\mathbf{X}^*$ is the label vector generated by ground-truth $\\mathbf{X}^*$. Thus, the loss function can be equivalently written as $f(\\mathbf{X}) = \\langle \\mathbf{H}, (\\mathbf{X}-\\mathbf{X}^*)(\\mathbf{X}-\\mathbf{X}^*)^\\top \\rangle$ for $\\mathbf{H} = \\mathbf{A}^\\top  \\mathbf{A}$, which is the same function that we studied in section 4.3 and show that 1-sided shampoo outperforms other adaptive algorithms.\n\n**Detailed comparison between optimizers:** We would like to elaborate on the comparative insights our theory can offer. As mentioned in line 74-80, line 222-231 and line 430-434, we should really consider the tradeoff between different terms in the regret bound or convergence rate. A less structured preconditioner such as full-matrix AdaGrad may have strength in having smaller $|||\\mathbf{g}\\_{1:t}|||\\_\\mathcal{H}$ and its weakness lies in inducing larger $||\\mathcal{X}||_\\mathcal{H}$. \n\n**Property of subalgebra:** The conditions for the well-structured $\\mathcal{H}$ are required to ensure that $\\mathbf{H}\\_t\\succeq \\mathbf{H}\\_{t-1}$. In particular, the condition for being closed under scalar multiplication and matrix addition (i.e., $\\mathcal{K}$ is a linear subspace) and the condition for being closed under matrix multiplication are important in the proof of the desired property. Indeed, without the linear subspace condition, we have the failure example of two-sided Shampoo as discussed in Example 2.2 and Appendix A.3.1; moreover, without the matrix multiplication condition, we have the failure example of tridiagonal matrices as discussed in Example 2.3 and Appendix A.3.2. Also, as can be seen from the proofs in Appendix A.2, for any $\\mathbf{H}\\in\\mathcal{H}$, we require $p(\\mathbf{H}) \\in \\mathcal{H}$ for any polynomial $p$, which guarantees especially that $\\mathbf{H}^{-1}\\in\\mathcal{H}$. All these requirements naturally inspire us to formulate the well-structured preconditioner sets using the notion of matrix subalgebra.", "claims_and_evidence": "Most of the claims are well-evidenced, except for the claim that \"Conceptually, our findings challenge the conventional wisdom that using a larger set of preconditioners which require more memory and compute leads to better optimization performance in terms of number of steps.\". This is because of two major reasons: (a) The analysis of one-sided Shampoo focuses heavily on a particular type of loss function (Equation 7) and a simplified experimental setting. It's unclear if these findings generalize to other types of loss functions or more complex problems. (b) The experimental setup is highly specific and doesn't reflect the complexities of typical deep learning scenarios.", "methods_and_evaluation_criteria": "The proposed methods have theoretical merit, but the evaluation criteria are not sufficiently comprehensive to support strong claims about the practical benefits of one-sided Shampoo in real-world deep learning scenarios. More extensive and realistic experiments are needed.", "theoretical_claims": "I haven't conducted a detailed verification of the proofs. However, the theoretical claims appear to be generally sensible and don't immediately contradict known principles in adaptive optimization. A rigorous validation of the proofs would be necessary for complete confirmation.", "experimental_designs_or_analyses": "While the experimental design appears logically sound for the specific problem under consideration (quadratic optimization), the lack of realism raises concerns about the validity of extrapolating these results to more complex, real-world scenarios. The absence of experiments on standard deep learning datasets and models is a significant limitation.", "relation_to_broader_scientific_literature": "The paper's key contributions relate to the broader literature in the following ways:\n\n- Adaptive Optimization: Builds upon and refines existing adaptive optimization methods like AdaGrad, Adam, and Shampoo (Duchi et al., 2011; Kingma & Ba, 2014; Gupta et al., 2, particularly Gupta et al. (2017), but overcomes limitations by introducing the concept of \"well-structured preconditioners.\"\n\n- Shampoo Optimizer: Contributes to the understanding and improvement of the Shampoo optimizer (Gupta et al., 2018; Anil et al., 2020) by proposing and analyzing a one-sided variant with improved theoretical properties.", "essential_references_not_discussed": "As far as I know, no other reference needs to be discussed.", "other_strengths_and_weaknesses": "# Strengths\n- The theoretical analysis is generally well-presented and easy to follow (assuming sufficient background knowledge).\n# Weaknesses\n- The lack of experiments on standard deep learning benchmarks limits the practical significance of the findings.\n- The motivation for focusing on the specific quadratic problem in the experiments could be strengthened.", "other_comments_or_suggestions": "While Section 3.3 demonstrates that the proposed well-structured preconditioner framework encompasses a range of existing adaptive optimizers, the paper misses an opportunity to provide comparative insights within this unified framework. A more detailed analysis of the relative strengths and weaknesses of different optimizers, as revealed by the framework itself, would have further strengthened the contribution.", "questions_for_authors": "Could you elaborate on the intuition behind the development of the well-structured preconditioner sets? What key insights led you to believe that these specific algebraic properties (closure under scalar multiplication, addition, and multiplication) would be sufficient to address the limitations in previous analyses?", "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:

Meta-review: The paper provides a unified analysis for establishing regret bounds for adaptive gradient methods with structured preconditioners. The analysis is a generalization of the analysis in the original Shampoo work and it applies to widely studied variants of Adagrad and Shampoo.  The analysis shows improved guarantees for a one-sided version of Shampoo. Additionally, the analysis shows that the regret bounds for a more general class of preconditioners can be worse than a more structured class.

The reviewers appreciated the contribution and the new insights derived from the proposed analysis. Following the discussion, all reviewers expressed support for the paper but there remained some concerns. In particular, one of the reviewers noted during the discussion that the arguments on the regret rates and particular scenarios of comparison are straightforward, and that the empirical results are not extensive enough to draw a clear conclusion on the effectiveness of structure in pre-conditioner design.
Decision: Accept (poster)