Adam Exploits $\ell_\infty$-geometry of Loss Landscape via Coordinate-wise Adaptivity
Authors: Shuo Xie
Mohamad Amin Mohamadi
Zhiyuan Li
Venue: ICLR 2025 Spotlight (2025)
Keywords: Adam
coordinate-wise adaptivity
adaptive algorithms
infinity norm
Primary area: optimization

Abstract: Adam outperforms SGD when training language models. Yet this advantage is not well-understood theoretically --  previous convergence analysis for Adam and SGD mainly focuses on the number of steps $T$ and is already minimax-optimal in non-convex cases, which are both $\widetilde{O}(T^{-1/4})$. In this work, we argue that the exploitation of nice $\ell_\infty$-geometry is the key advantage of Adam over SGD. More specifically, we give a new convergence analysis for Adam under novel assumptions that loss is smooth under $\ell_\infty$-geometry rather than the more common $\ell_2$-geometry, which yields a much better empirical smoothness constant for GPT-2 and ResNet models. Our experiments confirm that Adam performs much worse when the favorable $\ell_\infty$-geometry is changed while SGD provably remains unaffected. We also extend the convergence analysis to blockwise Adam under novel blockwise smoothness assumptions.

Published as a conference paper at ICLR 2025
ADAM EXPLOITS Loo GEOMMETY OF LOSS LANDSCAPE VIA COORDINATE-WIST ADAPTIVITY
Shuo Xie, Mohamad Amin Mohamadi & Zhiyuan Li Toyota Technological Institute at Chicago (shuox, mohamadamin, zhiyuanli
ABSTRACT
Adam outperforms SGD when training language models. Yet this advantage is not well- understood theoretically previous convergence analysis for Adam and SGD mainly focuses on the number of steps T and is already minimax-optimal in non-convex cases, which both O(T-1/4). In this work, we argue that the exploitation of nice geomeety is the key ad- vantage of Adam over SGD. More specifically, we give new convergence analysis for Adam under novel assumptions that loss is smooth under -geomery rather than the more common fz-geometry, which yields much better empirical smoothness constant for GPT-2 and ResNet models. Our experiments confirm that Adam performs much worse when the favorable foo geometry is changed while SGD provably remains unaffected. We also extend the convergence analysis to blockwise Adam under novel blockwise smoothness assumptions.
INTRODUCTION
Large language models (LLMs) have gained phenomenal capabilities as their scale grows (Kaplan al., 2020; Brown al., 2020; Touvron et al., 2023; OpenAl, 2023; Reid al., 2024). However, pre-training LLMs is incredibly time-consuming. Adam (Kingma & Ba, 2014) is the current to- go optimization algorithm for LLMs due to its fast convergence. In contrast, SGD, a popular and arguably the simplest optimizer, optimizes language model losses much more slowly than Adam.
However, the optimization benefit of Adam over SGD cannot be explained by existing theory. Cur- rent convergence analyses for Adam and SGD focus on the dependence the number of steps under assumptions on loss smoothness and gradient bounds (Défossez et al., 2022), and has been shown that both Adam and SGD achieve the minimax convergence rate (TT-1/4) in the non-convex set- tings (Arjevani et al., 2023). Thus according the theory, in the worst case, SGD would be more desirable than Adam because it has the same convergence rate, and yet Adam is less memory-efficient due to its coordinate- wise adaptivity, which needs to store the empirical moving average of second- order moments of past stochastic gradients. Therefore, we hypothesize that the coordinate- adaptivity in Adam is exploiting some unknown properties of L.LMs which SGD cannot make use of.
Towards this end, we identified big difference between Adam and SGD, which is ignored in previ- ous works. That is, SGD is auation-equivariann, while Adam is only permutation equivariant (Defini- tion 2.1). Intuitively, if we rotate the loss landscape, the trajectory of SGD would be the same (up to some rotation), while the trajectory of Adam can be completely different. If Adam optimizes much more slowly after rotation, it suggests Adam is exploiting some non-rotation- iirioonn--iviiantt properties of the loss, which is not captured by standard smoothness assumptions in the convergence analysis.
findings by comparing Adam on the original and rotated loss. The perfor- mance of Adam on the rotated loss does become much worse than Adam on the original loss. We also test memory-efficient and rotation-equivariin variant of SGD, AdaSGD (Wang & Wiens, 2020), de- fined in Algorithm 2. Surprisingly, the rotated Adam performs even worse than the SGD variant. The results suggest it impossible explain the superior optimization performance of Adam just using rotation-invariant assumptions on the loss function, which raises the natural question,
Figure summarizes
SGD
What cann-rottion-imaariinn properties of loss functions make Adam converge faster than SGD?
We hypothesize that the ég tg-lipschitzness loss gradient does not provide a tight-enough charac- terization of loss landscape of deep learning models in practice, such that we separate Adam and

---

Published as a conference paper at ICLR 2025
Training Loss
Evaluation Loss
9.20
9.20
Adam
Adam
6.00
AHSSD
6.00
AdsSC
Reiated
Rodated
SCID
4.40
SCDD
1,65
3.20
3.20
3.00
3.00
2.0
2,90
2.85
2.85
50000
80000 100000
40000
60000
Iterations
Iterations
Figure Training and evaluation losses of Adam, AdaSGD and SGD on GPT-2. rotated Adam means running Adam on a rotated loss. Adam on the original loss converges the fastest as expected. But convergence of Adam on rotated loss much slower, notably even worse than AdaSGD.
other rotation-equivarian algorithms. Inspired by the similarity between Adam and SignGD and the fact that SignGD is the normalized steepest descent with respect to fo-mrmm we propose to use 20 -norm related smoothness as a better tool to analyze Adam. In particular, our main results use the (1, nomm of the Hessian of the loss normalized by variable dimension d, as the smoothness measure, instead of its spectral norm. And we prove a convergence rate of O( (TT for Adam without
noise, or O(( (1/4) with noise. Our results have the same dependence on T as previous results, but much smaller smoothness constant when measured empirically. We empirically verify that (1, 1) norm of Hessian positively correlates with final training loss of Adam on both synthetic tasks like quadratic loss and real tasks like training GPT2 Open OpeenWe'lee WebText and ResNet CIFAR1O.
We summarize our contributions below:
1. We show by experiments that the empirical optimization advantage of Adam over SGD can not be explained solely under rotation rotation-invariant assumptions. (Figure 1)
2. We propose new complexity metric for the optimization problem, which is the (1, 1)- norm of the Hessian matrix of loss, || We present a novel convergence result for Adam depending on this metric in the case of 31 = O. (Theorem 3.5)
3. We further generalize the theoretical analysis (Theorem 3.11) for Adam to blockwise Adam (Algorithm 3) whose convergence rate can be characterized by a novel smoothness mea- sure (Definition 3.9). Adam and AdaSGD are two notable examples of blockwise Adam. In Adam, all blocks are of size 1. In AdaSGD, there is only one block.
4. We empirically verify that when Adam converges more slowly on the rotated loss, the (1, 1)- norm of Hessian also increases, which suggests that our new complexity metric for Adam's convergence is practically relevant. (Section 4).
2 PRELIMINARIES
Notations. For œ E IR4, we define the vector p-norm || ppp as (Ei 1/D for p E [1, oo]. For a matrix A E RRd xds, its (1, 1) norm || ||l1, is defined as Ei= Σ |At, and its operator norm induced by vector p-norm I|p as sup T| is the dual norm of || Ilp. For a square matrix A E Rdxd A is defined as the unique square root of A A. For deterministic loss function L(x), we consider optimization over L with access only to independent stochastic functions {z) {[ such that (x) L(x)) for any input x E Rd. Rotation. For an invertible function T R2 - Rd, T is rotating transformation if there exists an orthogonal matrix T e Rddd such that T(x) Tx. T a permutating transformation if there exists permutation mation is always rotating transformation. We will use R denote rotating transformation.
and - llg Amllly denoted by || where 1 + 1/-
[d] + [d] such that T(x:) = (1) ... [x(()) A permutating transfor-
Definition 2.1. For initialization
and stochastic losses [Lr) we can get when running algorithu A on (zo, (L))= For transformation T. we can also get 3; when running A with the same hyperparameters on (Øa. (L,)T with h. = =T T (xx) and L, = L, TT.
The code is httph:: com www.o///m//aaaaaoaaomaaooaaaaaaa ami h/adam-aaaeee coordi nate- adapt in ty.
2

---

Published as a conference paper at ICLR 2025
Algorithm Adam
Algorithm 2 AdaSGD (Wang & Wiens, 2020)
Hyperparam: B1, 32. 2 O, total steps T,
Hyperparam: 31, 32. 2 O, total steps T,
learning rate (no)e=1 initial mo,
learning rate {ne )= initial mo. Up Input: initial 20, stochastic losses {L+
Input: initial Vo. stochastic losses {La}E=
0o,4 < vo for t =
for 1,2,-.. ..2....TT)
Gt.i -- V.L.(xx-1)
gt.i V,Lo(EEr11
Ma. + 31mm- + -1-1 -1)g
ma,i t 31m--155 + Br)
U/ - -2 + (1 +11-B2)gi
- 82U=-1 + (1 (1- (1 /d)
mas
me,s
21,1 t Tt-1,3 1/t Vurte
Ti, t It-1,i 1)t Vos,it
return BT
return WT
An algorithm A is equivariant w.r.t. T if it always holds that ā. TT-(ax) for any hyperparam- eters, initialization and stochastic losses. A is rolation-equvaariant if it is equivariant w.r.t any rotating transformation R. A is permutation--qquivariiee if it is equivariant w.r.. any permutating transformation.
Theorem 2.2 shows the difference between Adam and AdaSGD, whose proof is in Appendix B.
Theorem 2.2. AdaSCD is rotation-equiiaarann. Adam and SignGD are only permulationneqquiarrian.
MAIN RESULTS: CONVERGENCE RATES OF Adam
In this section, we present our main theoretical results. Theorem 3.5 gives a convergence analysis of Adam for stochastic non-convex smooth losses with coordinate-wise gradient noise. The convergence measured by & norm of the gradient. For deterministic losses, our best rate (Theorem 3.2) is achieved by SignGD (Adam with 31 B2 0). For stochastic losses with bounded gradient noise variance, our best rate (Corollary 3.6) is achieved by RMSProp (Adam with 81 = 0, 82 E 1).
Then extend analysis of Adam to contains both Adam and AdaSCD as special cases. We also come up with novel smoothness mea- sures (Definition D.2) corresponding to the set of blocks used in blockwise Adam.
general blockwise Adam (Theorem 3.11), which
Similar to previous works (Défossez et al., 2022), our analysis can be extended to the general case of Adam, where both 31, 82 are non-zero. But the rate becomes strictly worse than the RMSProp (the case of 31 = 0), as there will be some extra polynomials of T-3, We decide not to include the result for the general case, on the one hand for ease of presentation, and on the other hand, because such result can not explain the optimization benefit of momentum (31 > 0) in practice and does not add any insight on the benefit of Adam. We hypothesize that the theoretical assumptions are missing some important features of loss landscape transformers and we leave this for future work.
3.1 WARMUP: SignGD (Adam WITH 31 32 = 0)
In this section, we present the convergence analysis for SignGD as warm-up and illustrate how Adam could benefit from non-rotation-invarinn property the loss, which in particular the fo smoothness. The key observation is that SignGD is normalized steepest descent w.r.t. too norm (Xie & Li, 2024), so it is more natural to analyze its convergence using foo norm geometry of the loss. Definition 3.1. Given norm ||-Il on Rd and I|. as its dual norm, say a function L is H -mooohhh w.r.t ||-|| for any x, y Rd we have that [(xx L < H || -ll.
Theorem 3.2. Let L be H-smooth w.r.t. loo norm and {x, IT be the iterates of SignGD (Adam with 3 = B2 = 0) on with initialization By and learning rate n, it holds that
[[V((x)) S L(wo) min L(x)
Hn
min
Tn
2
we choose 7)
2[L((mo) min M(a)) then mini VVI(x;)) S TH
217 (1ma)) - min f.(m))
Heee it is slightly different from Wang & Wiens (2020). We exponential average of the gradient for m. instead of momentum. Our definition makes AdaSGD the Adam in one-dimensional problem.
3

---

Published as a conference paper at ICLR 2025
Theorem 3.2 gives the convergence rate for SignGD and the proof is in Appendix C.
3.2 MAIN RESULT: RMSProp (Adam WITH 31 0, 32 E [0,
It is well-known that SignGD might not converge in the stochastic case as the expectation of descent direction for mini-batch loss may not be a descent direction for RMSProp is proposed to address this issue by using moving average of the squared gradient per coordinate to reduce the correla- tion between the denominator and the numerator, thus making the expected update direction less biased (Hinton et al., 2012). In this subsection a positive 32 in Adam helps convergence in the stochastic case. The main challenges here are from both lower bounding the first-order term and upper bounding the second-order term in the modified descent lemma (the RMSProp counterpart of Equation 7).
formalize the above intuition and show indeed
9: VoL +E
H
9: Vo +
L(xx) L(x S V((z.)
2 "
+E
We can only upper bound
veete by 1-
g:
1 without more fine-grained analysis on the relationship
between gradients in each step, which will greatly hurt the dependence of convergence rate on 1 1-32- as large VI- the average coordinate
However, though some large Jt,i make
Ver,s+
moving speed should be close Therefore, we introduce Definition 3.3, which is slightly stronger than Definition 3.1 but enables decomposing the second order term into each coordinate according Lemma D.3. also facilitates coordinate-wise analysis for the first order term. We note this definition also appears in Assumption 2.3 of the concurrent work (Maladkar et al., 2024).
Definition 3.3. For any H = (HI...., Ha) E IRd, say T. is H-smooth coordinate- wisely w.r.t Lo norm, if and only if V [(x) V,I(y) s
H, | yl1 for any E [d], æ. y Rdd
estimate for coordinate-wist smoothness. H; in Definition 3.3 is determined by sup Σラ= (x)) which is difficult compute because requires taking supreme over the entire domain. A computationally-rrattaii alternative is approximate Σ H, locally by the (1, 1)-ormm of Hessian loss along the training trajectory. We provide an efficient approximation algorithm with concentration guarantees in Appendix E.3, which uses dessian-vectee product against random Cauchy vectors.
(1,1)-norm
By definition, H-smoothness coordinate-wistt wisell w.r.t. loo norm implies Σ =i=1 H;-smoothness w.r.t. lo norm. We also need Assumption 3.4 to measure the influence of noise the stochastic setting. Assumption 3.4 (Coordinate-wirs
noise). There exist
constants
03 such that
( (=L(=)] S for any E [d], t E and a Rd.
We present the main result in Theorem 3.5. The sketch of the proof is presented in Section 3.4 and the complete proof for the generalized blockwise Adam algorithm is presented in Appendix D. The proof incorporates some key steps from Li Lin (2024), extending them accommodate the generalized algorithm and different smoothness assumptions.
Let (T)) 1 be independent stochastic losses satisfying Assumption 3.4 and that their expectation I. is H H-coordinate-wicn wiislyy smooth w.r.t. foo norm. For Adam with 31 O.
Theorem 3.5 (Main, Adam),
we have that
d
32 dup
min E V 5o E+ SSET
+
Σ 1=1
oi + dvee
T(1 32)
with E
w/ E ((z)) - L(xT)) +
B,8 1
17 Ed Hi VI BB Ea_ 01
and F
T11-331
Ed Ein 1 Einn 2T T-m
2 In
1+
In
The convergence rate of R.Sprop can be determined by choosing hyperparameters 7 and 32 in The- orem 3.5 minimize F. By assuming to of ( + +E, H?+2) H?3m2//oly((7 and
4

---

Published as a conference paper at ICLR 2025
1 poly(T), can simplify the term by considering F = Ollo T). The two terms involving
log 7
E oj have a lower bound Σ
which is achieved with - 822
log
Then the three terms involving yy has a lower bound θ (V
(L(een) min, [(x)) E H: reached by
7
-= (v L(wo)) mint L(w) These hyperparameter choices yield the optimal convergence rate for T-: H,
stochastic case Corollary 3.6. For convenience, we define R - (L(zz) min L(x)) Σ Hi: which will be the core term in Corollaries 3.6 and 3.7.
Corollary 3.6 (Stochastic Case, general o,). Let {Lo}ī=1 be independent stochastic losses sat- isfying Assumption 3.4 and that their expectation L is H #--cooddnaat--wise/ smooth w.r.f Loo
norm. For 31 = 0, - B2 = 0(=ogT) E = 0, r = θ
(V
M(mo) min. L(x) and Vo >
T Σ H,
E of [V + +2 H/311)//ool((((( we have that
( +
R 1 T
log 7 T
$
min <IST
E1lgi|1 1 =0
+ dr
i=1
with 8T V
den T(I- B2] exp
(7)
d Ei=
log
T
Here ōy can be smaller than any polynomial of T by manipulating the value of T(1-Ba) 0(1). log 7
Then Dd Ti Log 1
the leading term w.r.t. in the rate whose coefficient only involves
EL.O TGe It suggests that the rate can be much improved when noise is small. Below get the convergence rate with the same hyperparameters in deterministic case in Corollary 3.7.
Corollary 3.7 (Deterministic Case, o; 0). Let TT= be deterministic losses satisfying Assump-
tion 3.4 and that their expectation L is H T--corddnnte-wiien wistly smooth w.r.f. foo norm. For B1 = 0,
1 32 1 1o) 0, n θ (v Σ. HH/12)/polyl/ for any pol polynomial poly(T), we have that
[( ming L(x) and Uo EEi oī EE + [ +
T TEE H.
|Igell1 -0
R T
min <IST
8g
with 87 V
deg T[]-9.] cxp
(_TQ-$1) 9o)
Corollary 3.7 almost recovers Theorem 3.2, except for the smoothness constant. Specifically, it uses
sup, V V²I((=))((..) that is larger than supa V [x=))>> in Theorem 3.2 as ||- I1- always holds. This gap is due to technical difficulty of analyzing Adam as mentioned in Section 3.2.
Dependence on e, un and 32. While many previous works rely on the relatively large magnitude of e compared to vn and give a bound in the regime of SGD when the adaptive effect is dominated by the constant e (Zaheer et al., 2018: 8; De al., 2018), our result actually prefers to be while maintaining the value of Un + We also note the dependence of our bound in Theorem 3.5 on vo is very mild and logarithmic. Theorem 3.5 has similar convergence rates for all th of magnitude at most poly(T), while most previous result only addresses the case where Vo,; is at the scale of noise (Li & Lin, 2024) or O. The main reason for this adaptivity to wide range of Uo is our specific choice of 32 = 1 - O ), which allows the initial large up to decay fast and resume normal training. Other existing results using 82 = 1 O(1/T) (Défossez et al., 2022: Li & Lin, 2024) cannot allow large initial value ui because un only decays a constant fraction throughout the training and the effective learning rate will be too small.
3.3 A UNIFIED ANALYSIS FOR BLOCKWISE Adam
In this subsection, we present convergence analysis for a broader class adaptive algorithms de- fined in Algorithm 3. It can be viewed coarser version of Adam because it does pre-conditioning
5

---

Published as a conference paper at ICLR 2025
Algorithm 3 Blockwise Adam
Hyperparam: 32, 2 O, block partition 4, total steps T, learning rate {m 1n-3- initial me. vo-
Input: initial To. stochastic losses {1=)
0o.5 < Vo for t = 1,2,-..,5
gt.: + V,L,(1:-1)
ma, t -1m + 31)
E
(o oZ.c
Vt.b + 12 (1 (2)
db
3t,i f It-1. - 1/t oce
return DT
blockwisely instead coordinate-wisely. Since Adam and AdaSGD can be viewed as special cases of blockwise Adam with Addem
respectively, any convergence results for Algorithm would imply convergence of Adam and AdaSGD. Finally we also note that such block- wise Adam has been recently studied empirically by some concurrent work, where the algorithm is named by Adam-mini (Zhang et al., 2024b) and Adalayer (Zhao et al., 2024).
and Padssao
notations. For partition function T [d] [B] where B is the number of blocks, (b) is defined as ob (b) {i) = b) and d, = #(b), the number of parameters in block We define the vector (b) [Te =b)== and the submatrix (())(1)) as [AL)]a(()=))))))=)
We first introduce
Definition 3.8 (f-norm). We define the (oo, 2)-norm w.r.t partition & of x as the foo norm of the B
|
|| For convenience, we will denote by || or váb
vector
which is maxbe[[] b=1
Vdb
just call it 4 4-norm. We denote its dual norm by || which is equal Σ ddg (((b)
Definition 3.9 (#-smoothness). We say a diagonal matrix A follows partition & iff its diagonal elements are conslant within each block, i.e., there are u1," ,aB A;,; (.p(s) for any i E [d].
say a tvice-lifferentiabl I. is H-smooth under partition & if there exists a diagonal matrix A following partition & such that H Tr(A) and A dominates [ for all æ. We further define the - of loss L, denoted by H (L, D), as the smallest constant H such that L is H- -mooohh under parition that is.
H(L, 4)
min
Tr(A)
(1)
follows 4 A2 [an)
We note that @-ssmoothness is different from the smoothness under Ø-norm, where the latter is equal sue sup $1 [
a... For each æ, it holds that
2 L(x)uull || Wlll.s
sup sup V(L(E)u= sup |
au V2L(x)u
sup
| |Pl SI
When A N L we have that sup alla $1 u LLL(x)u S sup (ill Si u Au. When diag- onal follows partition Φ we have that supp u Au Tr(A). So d-smoothness defined in Definition 3.9 is always no smaller than the smoothness under h-norm.
Another advantage with the definition is that the (T, b) will always non-increase with a more fine-grained partition. For two partition functions 2 and 22. we say e includes 22 and only 2 = 2, (j) for any j.j E [d] such that Φ
2 (j). If diagonal matrix A follows partition follows o A Z [x=) for all a C (A] A follows 22 > L(x) for all æ: and HI,, & is no smaller than H (L, 22). Since 4 AdaSED includes any partition and any partition includes & Aden H (L, (ddase)) is the largest and H (I, P is the smallest among all the partitions Tb
Φ1 and partition D includes 22 then A also follows 42 So we have that {
With Assumption 3.10 on noise, we can prove the result for blockwise Adam Theorem 3.11. Assumption 3.10 (Generalized version of Assumption 3.4). There exists constant ob such that E V (x) ( < doo for any block b E [B), I N and x E Rd.
Theorem 3.11 (Main, Blockwise Adam). For specific partition Φ consider the updates defined in Algorithm 3. Under Assumption 3.10, we have that
B V do [o) 2V + V2F,
T 4,3
B doob + dvc
min
E
d d +
Σ
T ESEST
T(1 - 3z)
6=1
b=1

---

Published as a conference paper at ICLR 2025
with E 2. (L(0)) Lxxr) + +11+11111
F. T(1-82)
(H +2/1- 32 Σi= drob)
and F
1/
2 2 (1+ 1f (10))ll000))
Ebc(a) HC " 1- + 12.
No+
When is (HI..... Ha)- -mmoth soordinate-wisely w.r.t. foc norm, we show in Appendix D.2 that diag(H1..... Ha) == [()) for all æ E I2 and diag(H1,... ................................ Ha) follows & Paaan partition. So H (L, Padan) is at most Vi=1 H, and we can use Theorem 3.11 to derive Theorem 3.5. We can also estimate H[T., AAcam) by sUP=ER= (a=))l,,,,
discussed in Section 3.2. For Aaassa being the mapping + (L, PAaasoD) is the same as the smoothness under 4 Pa norm, whose value is d [
Different norms for smoothness. As an implication of Theorem 3.11, we can get analogs of Corollaries 3.6 3.7 for AdaSGD, with the corresponding noise and smoothness assumptions. When the optimization not noise-dominated and / R - (L(z)) L(w) (LL.))) becomes the lead- T
ing term, the choice of & now matters lot. The key difference between AdaSGD and Adam lies in the gap between H (L, AdasaD) and (L, 4 Adaa)) and comparing these coefficients can provide insight into which algorithm may be more effective under different conditions.
Previous analyses of Adam's convergence (Shi & Li, 2021; Défossez et al., 2022; Li & Lin, 2024) usually assume smoothness under the la norm and the rate of Adam ends up being identical to the rate AdaSCD, which fails to explain why Adam often performs better than AdaSGD practice. By adopting an fo norm smoothness assumption, the coefficient for Adam's convergence rate changes from d sup, | La))ll l to HL, Padan). where the latter is typically much smaller when Adam optimizes faster because sup,.. ||>]((>)]] ||,, 1,1 is much smaller than dsuu,.. || [
Finally, we note that the Dadan Smoothnss (L, Padan) is not rotation-invariant in the sense that HL, Padan) A (L o R, (Aaan) for typical rotation R. In practice, the (1, 1)-morm of Hessian matrix can vary lot when rotation is performed on the loss shown in Section 4.1. In contrast, P P smoothness (L, Padasao) is invariant under loss rotations.
3.4 PROOF SKETCH OF THEOREM 3.11
We will use ÿ. = Elgn x <c] ((x,,-)) to denote the full batch gradient and consider the decrease of L(x;)) in a single step L. We can upper bound the second order term in the Taylor expansion (Equation 2) with H (x)(x) that achieves Tr(H) H(L Φ). Then we can get
(r) 12
M(x.) T(x n
Σ
H
2
(1) +c
24b6 +
b=1
B
B
|19
////o
= -1
VeL +e +
n" Hads
>
(3)
De.b +
The proof contains two main parts: lower bounding the first order term using || 1la and upper bounding the second order term. We address the second-order term by employing Lemma 3.12 to bound the sum by T + 1-82 no,te
8a
In urate where we set Uj Up.6 and gr [9,(00)122 /Vd..
Lemma 3.12. Given any 0 < 82 < 1, for any scalar sequences (ve)T and (T satisfying No 2 0, 01 > and f2 - 82 2 (1 - B2)9 for 1. 2 1, holds that EI of. S T + I-3 In "r.
Now we turn to the first term. Ideally, for each block b E [B]. we would like to connect the first order term to [[ 112 taking expectation, i.e., E, (i E.g. (c.( Ver,te
vest
where we use [.[ as abbreviation for B|||x2]) However, this is not correct because both the numerator and denominator in (0) depend on the stochastic gradient gt. To circumvent this
difficulty, we lower bound each conditional expectation E, (2) by Voste 2 2.3.0,,tte
F.gT minus error

---

Published as a conference paper at ICLR 2025
terms related to noise magnitude obe We can further have 1sg
9e. Vo,tte
where 24,6
Vilertot
8201-1,6 1,0 (1- 32) (100, /do of) This leads Lemma 3.13 whose proof in Appendix D.
Lemma 3.13 (first-order approximation). With Assumption 3.10, it holds that for any block E
[B]
T (1 (1,(M) E
T
||
BoTdbor - dror E VI
VTD +
E
In
(4)
t=1 00D
11.6 tE
v1 B2
Vo,b
te
Combining Lemmas 12 and 3.13 and Equation gives an upper bound for E
E
Σb=1
Volte using
Finally, we employ Cauchy inequality (Equation 5) and Lemma 3. 14 to upper bound 1 a,-
G -t=
Ex=1
This completes the proof.
T
B
T
B
10,
Σ
@gel,o,- =
Σ
de
Σ
Σ
Σ
day Vibb vrb tc. (5)
v
/eis +
t=
+1
b=1
+
b=1
t= ++++11
Lemma 3.14. With Assumption 3. 3010, holds that for any block E [B]
Σ
E
VoL
28
/u, +
T
abb
T
Vc 2
E
gr,(m) /d/
(6)
1- 82
0rb +
4 EXPERIMENTS
In order to empirically investigate and confirm the implications of our proposed theory, we compare the training performance Adam with AdaSGD, SGD and rotated Adam multiple different tasks. The details of all experiments are in Appendix E.
4.1 QUADRATIC LOSS
We perform controlled experiments on quadratic loss study the relationship between optimiza- tion speed of Adam and the shape of Hessian in terms of O smoothnns,.ssss. More specifically, we consider = diag(1,
2T 3r, .... 3y0z) 11000x1000 and manually generate orthog- j T Ex with AdaSGD and Adam and optimize L.(x) = f R, ER,æ with Adam for 100 steps. Because AdaSGD is rotation-equivariann, the optimization process of AdaSGD is the same on all L, We tune learning rates for each setting with 10 random seeds and present their lowest average loss with standard deviation in Table 1.
onal matrices R4. Then we optimize Lo(x)
Optimizer |
(1, 1) (eor))nooo//a /d
Loss (3) = B2 = 0)
Loss (3) = 0.9, 0.99)
AdaSGD
0.01164
0.00887 + 0.00119
0.00405 +: 0.00021
Adam
0.01164
0.00022 0.00007
0.00002 + 0.00001
Adam (R₂)
0.08324
0.00314 + 0.00031
0.00066 + 0.00008
Adam (R²)
0.50729
0.00567 = 0.00053
0.00134 + 0.00007
Adam (Rx)
1.23731
0.00751 = 0.00086
0.00183 + 0,00009
Adam (R4)
2.59919
0.00978 + 0.00132
0.00254 + 0.00008
Table 1: The final loss values obtained by different optimizers and the (1, 1)) of Hessian matrix for the corresponding objectives. The spectral norm of the Hessian always 1. Adam optimizes worse when the (1, -nomm of Hessian matrix increases, as suggested by Corollary 3.7. Moreover, when (1, 1)-morm divided by d is much smaller than spectral norm, Adam tends to optimize faster than AdaSGD, supporting tmsmoothhesss as predictor of blockwise Adam's 's optimization speed.
We find clear pattern that Adam optimizes worse when the (1, -)-mmrm of Hessian matrix increases, as suggested by our Corollary 3.7. Moreover, when (1, -norm divided by dimension is smaller than spectral norm, Adam tends to optimize faster than AdaSGD, suggested by Theorem 3.11.
8

---

Published as a conference paper at ICLR 2025
Train and Evaluation Loss
Train and Evaluation Accuracy
2.5
100
5e-1
90
20-1
SGD
80
SGD
1c-1
AdaSGD
AdsSGD
Adam
73
Adam
Rotated AdamW
Adam
30-2
70
500
1000
1500
2000
2500
5000
3500
4000
500
1000
1500
2000
2500
5000
3500
2000
Berations
Iterations
Figure 2: Training (solid line) and test (dashed line) losses and accuracy of ResNet18 on CIFAR-10
with Adam, AdaSGD, rotated Adam, and SGD. Adam converges faster than other algorithms.
4.2 GPT-2 ON LANGUAGE MODELING TASK
We train GPT-2 on the WpenWeeText corpus. The training losses and evaluation losses of different optimizers are plotted in Figure As mentioned in Section Adam converges faster than AdaSGD while they both converge faster than rotated Adam. Since we propose the (1, n of Hessian as a non non-rotation ann-rrationninvariann metric that can affect the convergence rate of Adam, we also measure for the original loss function and rotated loss function on checkpoints trained with different losses. The results are presented in Table 2. The same correlation between norms and convergence rates holds here. The smaller the norm is, the faster the optimizer works.
Optimizer
Smoothness Metric
Upper Bound d V L(x) (V R L R
Estimated Value
AdaSGD Adam
(LL, PadssoD) H (L, Fadam) H o R. Pddan)
4.2446
2.3538
Rotated Adam
14.3745
Table 2:
Hessian norms for the last GPT-2 checkpoints trained with different optimizers.
We also explore how learning rate can affect the performance of different optimizers and find another advantage of Adam over AdaSGD: it can maintain stable training at a larger learning rate, which often beneficial to faster and more efficient convergence. The results and details are in Appendix E.4.
GPT-2 small models have more than 100 million parameters, and thus the size of its hessian of loss as well as the rotation matrix is more than 16 which is way more than the storage of the modern computers. We introduce how to rotate the loss efficiently in Appendix E.2 and how estimate Hessian norms in Appendix E.3.
4.3 RESNET18 ON CIFAR-10
To further test whether the correlation between 4--mootthess and the optimization performance holds for architectures other than transformers, we conduct an experiment on ResNet18 trained on CIFAR-10 (Krizhevsky & Hinton, 2009).
Figure depicts the loss and accuracy curves for the best performing hyperparameters chosen over the training set's final loss for batch size 256. The results for other batch sizes are in Table 4. When comes to optimization speed, even for ResNet18, 8, Adam is always better than rotated Adam and they are always better than AdaSGD and SGD across different batch sizes. Note that this does not contradict with common practice of training ResNet with SGD, where the goal is to get better gener- alization and the training budget is large so all optimizers can casily achieve full training accuracy. In our experiment, we study optimization speed and intentionally limit the number of steps.
We have intentionally limited the number of training iterations to emphasize the difference of optimizers terms of training speed over generalization.

---

Published as a conference paper at ICLR 2025
Optimizer
Smoothness Metric
Upper Bound
Estimated Value
AdaSGD Adam
(L, PAdeseo) H (L, 2 (den)
V2 L(x)) x
1.5355
V L(
0.0036
Rotated Adam
(L R, 2 Pddan)
R V
0.9868
Table 3:
Hessian norms for optimal ResNet checkpoints trained with different optimizers.
We also measure the Hessian for checkpoints obtained at batch size 256 and the results are in Table 3. The correlation between norms and convergence rates still holds here. When the (1, 1)-ormm is smaller than d times spectral norm, Adam optimizes faster than AdaSGD.
RELATED WORKS
Comparison between Adam and SGD Previous work tries to analyze the difference between Adam and SGD from different perspectives. Zhou et al. (2018) proves a faster convergence rate of Adam than SGD when the stochastic gradients are sparse. Zhang al. (2020) suggests that SGD suffers more from heavy-tailed noise than Adam. Pan & Li (2023) claims that Adam has lower directional sharpness because of the effect of coordinate-wise clipping. Other works also consider the coordinate-wer wisc normalization of Adam (Balles & Hennig, 2018; Kunstner et al., 2022). Kunstner al. (2024) shows that the heavy-tailed class imbalance in language modeling tasks will cause SGD to converge slower when it can only optimize majority class well. Zhang et al. (2024a) finds that Adam is better at handling the block heterogeneity of Hessian matrix, which a specific phenomenon in transformers. When viewing Adam an adaptive method, there are works showing that adaptive methods have an advantage of achieving optimal convergence rate without relying on problem- dependent constant (Ward et al., 2020; Levy et al., 2021).
Convergence rate Adam There are many works showing convergence rate for Adam (Zhou et al., 2018; Chen et al., 2018; Zou al., 2019; Shi & Li, 2021; Guo et al., 2021; Défossez al., 2022: Zhang et al., 2022). Most of them rely on the smoothness measured w.r.t. 82 norm. Zhang al. (2019) proposes (Ln, Li -smosmoossssss condition should be more reasonable than globally bounded smoothness. Li et al. (2024) further generalizes the (Lo, L1) smoothness condition. However, they still focus the default rotation-invariant norm. To the best of our knowledge, we are the first to assume smoothness under foo norm for analyzing Adam.
the loss function, which is
Comparison with Li & Lin (2024) Li & Lin (2024) employs the same & norm for gradient and improves the dependence on dimension d compared to previous results for &a norm. But they still assume & norm smoothness while we adapt their results under &oo norm smoothness potentially further improve dependence on d. Another drawback of Li & Lin (2024) is setting vo based on noise magnitude a. which impractical real experiments because a is unknown. Overestimating a will result in slow convergence because large vy causes Adam to behave similarly with SGD without adjusting the learning rate adaptively. contrast, we allow for general initialization for vo and convergence rate can work well in both noisy setting and deterministic setting. We also use
1 1 32 θ
T to obtain our convergence rate while Li & Lin (2024) requires 1 - 33 = θ (
CONCLUSION
We give new convergence analysis (Theorem 3.5) for Adam in the stochastic non-convex setting ing a novel smoothness assumption. We show the convergence rate for the E norm of the gradient is
O in the deterministic case (Corollary 3.7) and ( log 1 vT
in the stochastic case (Corol- lary 3.6). We also extend our analysis to blockwise Adam on loss L. with respect to arbitrary partition of the parameters & (Theorem 3.1 using the corresponding smoothness H (L, Q) (Defi- nition D.2). Our bound for Adam involves (1, 1)-morm of Hessian, rather than the spectral norm of Hessian, which relevant to the convergence speed of AdaSCD. This leads to significantly better smoothness conditions for deep learning models including ResNet-18 18 and GPT2 empirically. Our experiments also verify that the smoothness measure HLL Φ) positively correlates with the opti- mization speed of blockwise Adam with respect to the partition T.
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGEMENT
The authors would like to thank Khashayar Gatmiry and Sharut Gupta for helpful discussion and preliminary experiments exploring the idea of rotated Adam and Xiaoyu Chen for discussions on extending the analysis to blockwise Adam and steepest descent w.r.t. general e-norm. ZL is supported by OpenAl Superalignment Grant.
REFERENCES
Yossi Arjevani, Yair Carmon, John C Duchi, Dylan Foster, Nathan Srebro, and Blake Woodworth. Lower bounds for non-convex stochastic optimization. Mathematical Programming. 2023.
Lukas Balles and Philipp Hennig. Dissecting adam: The sign. magnitude and variance of stochastic gradients. In International Conference on Machine Learning, 2018. 10
Jeremy Bernstein, Yu-Xiang Wang. Kamyar Azizzadenesheli, and Animashree Anandkumar.
signSGD: Compressed optimisation for non-convex problems. In International Conference on Machine Learning, 2018. 14
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in neural information processing systems, 2020.
Xiangyi Chen, Sijia Liu, Ruoyu Sun, and Mingyi Hong. On the Convergence of A Class of Adam- Type Algorithms for Non-Convex Optimization. In International Conference on Learning Repre- sentations, 2018. 10
Soham De, Anirbit Mukherjee, and Enayat Ullah. Convergence guarantees for RMSProp and ADAM in non-convex optimization and an empirical comparison Nesterov acceleration. arXiv preprint (rXivv:807700666 2018. 5
Alexandre Défossez, Leon Bottou, Francis Bach, and Nicolas Usunier. A Simple Convergence Proof of Adam and Adagrad. Transactions on Machine Learning Research, 2022. 1, 5,
Alina Ene, Huy L Nguyen, and Adrian Vladu, Adaptive gradient methods for constrained convex optimization and variational inequalities. In Proceedings of the AAAI Conference on Artificial Intelligence, 2021. 14
Zhishuai Guo, Yi Xu, Wotao Yin, Rong Jin, and Tianbao Yang. A novel convergence analysis for algorithms of the adam family. arXiv preprint arXiv:2112.03459, 2021. 10
Geoffrey Hinton, Nitish Srivastava, and Kevin Swersky. Neural networks for machine learning lecture 6a overview of mini-batch gradient descent. Cited on, 2012.
Ruichen Jiang, Devyani Maladkar, and Aryan Mokhtari. Convergence analysis of adaptive gradient methods under refined smoothness and noise assumptions. arXiv preprint 04552, 2024. 14
Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei. Scaling laws for neural language models. arXiv preprint arXhv:2001.08361 2020.
Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint 11X1v:1412.6980, 2014.
Alex Krizhevsky and Geoffrey Hinton. Learning multiple layers of features from tiny images. 2009.
Frederik Kunstner, Jacques Chen, Jonathan Wilder Lavington, and Mark Schmidt. Noise Is Not the Main Factor Behind the Gap Between Sgd and Adam on Transformers, But Sign Descent Might Be. In The Eleventh International Conference Learning Representations, 2022. 10
11

---

Published as a conference paper at ICLR 2025
Frederik Kunstner, Robin Yadav, Alan Milligan, Mark Schmidt, and Alberto Bietti. Heavy-Tailed Class Imbalance and Why Adam Outperforms Gradient Descent on Language Models. CoRR, 2024. 10
Kfir Levy, Ali Kavis, and Volkan Cevher. Storm+: Fully adaptive sgd with recursive momentum for nonconvex optimization. Advances in Neural Information Processing Systems, 2021. 10
Haochuan Li, Alexander Rakhlin, and Ali Jadbabaie. Convergence of Adam under relaxed assump- tions. Advances in Neural Information Processing Systems, 2024. 10
Huan Li and Zhouchen Lin. On the O( vd Convergence Rate of RMSProp and Its Momentum TIT Extension Measured by & Norm. arXiv preprint arXiv:2402.00389 2024. 4, 5, 7, 10, 18
Selena Zihan Ling, Nicholas Sharp, and Alec Jacobson. Vectoradam for rotation equivariant geom-
etry optimization. Advances in Neural Information Processing Systems, 2022. 14
Yuxing Liu, Rui Pan, and Tong Zhang. AdaGrad under Anisotropic Smoothness. arXiv preprint arXhv:2406.1524, 2024. 14
Lucas Maes, Tianyue H Zhang, Alexia Dolicoeur-Martines Ioannis Mitliagkas, Damien Scieur, Understanding Adam Requires Better Rotation Dependent Assumptions. arXiv preprint arXiv:2410. 19964, 2024. 14, 25
mon Lacoste-Julien, and Charles
Devyani Maladkar, Ruichen Jiang, and Aryan Mokhtari. Convergence Analysis of Adap- tive Gradient Methods under Refined Smoothness and Noise Assumptions. arXiv preprint arXiv:2406.04592, 2024. 4
OpenAL. GPT-4 technical report. arXiv, 2023. 1
Yan Pan and Yuanzhi Li. Toward understanding why adam converges faster than sgd for transform- ers. arXiv preprint arXiv:2306.00204 2023. 10
Machel Reid, Nikolay Savinov, Denis Teplyashin, Dmitry Lepikhin, Timothy Lillicrap, Jean- baptiste Alayrac, Radu Soricut, Angeliki Lazaridou, Orhan Firat, Julian Schrittwieser, et al. Gem- ini 1.5: Unlocking multimodal understanding across millions of tokens of context. arXiv preprint arXiv:2403.05530. 2024. 1
Naichen Shi and Dawei Li. Rmsprop converges with proper hyperparameter. In International conference on learning representation, 2021. 7, 10
Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, et al. Llama: Open and efficient foundation language models. arXiv preprint arXiv:2302. 13977 2023
Jiaxuan Wang and Jenna Wiens. AdaSGD: Bridging the gap between SGD and Adam. arXiv preprint arXiv:2006. 16541, 2020. 1,3
Rachel Ward, Xiaoxia Wu, and Leon Bottou. Adagrad stepsizes: Sharp convergence over nonconvex landscapes. Journal of Machine Learning Research, 2020. 10
Shuo Xie and Zhiyuan Li. Implicit Bias of AdamW: too Norm Constrained Optimization. arXiv preprint arXiv:2404. 0444544 2024. 3
Manzil Zaheer, Sashank Reddi, Devendra Sachan, Satyen Kale, and Sanjiv Kumar. Adaptive meth- ods for nonconvex optimization. Advances in neural information processing systems, 2018.
Jingzhao Zhang, Tianxing He, Suvrit Sra, and Ali Jadbabaie. Why Gradient Clipping Accelerates Training: A Theoretical Justification for Adaptivity. In International Conference on Learning Representations, 2019. 10
Jingzhao Zhang, Sai Praneeth Karimireddy, Andreas Veit, Seungyeon Kim. Sashank Reddi, Sanjiv Kumar, and Suvrit Sra. Why are adaptive methods good for attention models? Advances in Neural Information Processing Systems, 2020. 10
12

---

Published as a conference paper at ICLR 2025
Yushun Zhang, Congliang Chen, Naichen Shi, Ruoyu Sun, and Zhi- Zhi-Quan Luo. Adam Can Converge Without Any Modification On Update Rules. In Advances in Neural Information Processing Systems, 2022. 10
Yushun Zhang, Congliang Chen, Tian Ding, Ziniu Li, Ruoyu Sun, and Zhi-Quan Luo. Why Trans- formers Need Adam: A Hessian Perspective. arXiv preprini arXiv:2402. 240.:28002,6788, 2024a. 10
Yushun Zhang, Congliang Chen, Ziniu Li, Tian Ding, Chenwei Wu, Yinyu Ye, Zhi-Quan Luo, and Ruoyu Sun. Adam-mini: Use Fewer Learning Rates To Gain More. arXiv preprint
arXiv:2406.16791, 2024b. 6
Rosie Zhao, Depen Morwani, David Brandfonbrener, Nikhil Vyas, and Sham Kakade. Deconstruct- ing what makes a good optimizer for language models. arXiv preprint arXiv:2407.07972, 2024.
Dongruo Zhou, Jinghui Chen, Yuan Cao, Yiqi Tang, Ziyan Yang, and Quanquan Gu. On the convergence of adaptive gradient methods for nonconvex optimization. arXiv preprint arXiv: 05671, 2018. 10
Fangyu Zou, Li Shen, Zequn Jie, Weizhong Zhang, and Wei Liu. sufficient condition for conver- gences of adam and rmsprop. In Proceedings of the IEEE/CVF Conference on computer vision and pattern recognition, 2019. 10
13

---

Published as a conference paper at ICLR 2025
CONTENTS
1 Introduction
2 Preliminaries
3 Main Results: Convergence Rates of Adam
3.1 Warmup: SignCD (Adam with 31 = B2 = 0)
3.2 Main result: RMSProp (Adam with 31 0, E [0. 1))
3.3 A unified analysis for blockwise Adam
3.4
Proof sketch of Theorem 3.11
4
Experiments
4.1
Quadratic loss
4.2
GPT-2 language modeling task
4.3
ResNet18 on CIFAR-10
5
Related Works
10
6
Conclusion
10
A More Related Work
14
Invariance Property of Adam and SGD
15
Convergence rate SignGD for deterministic loss
15
D Proof Details
15
D.I Proof for convergence rate of blockwise Adam
15
D.2 Proof for convergence rate of Adam
23
Experiment Details
24
E.1 Training details
24
E.2 Adam on a rotated loss
25
E.3 Computation of matrix norms
25
E,4 More results
28
A MORE RELATED WORK
We note the definition of P smoothness been mentioned in previous work. Bernstein et al. (2018) relies on to prove the convergence rate SignGD. Previous and concurrent work Ene al. (2021); Liu et al. (2024); Jiang al. (2024) uses this assumption for analyzing AdaGrad. We are the first to analyze Adam under such smoothness assumption. Moreover, we are the first to empirically measure it for loss functions in real task with theoretical guarantee. Concurrent work Maes et al. (2024) proposes another algorithm to measure (1, 1) norm of Hessian matrix. Ling et al. (2022) analyzes the rotation equivariance property of Adam in the geometry optimization setting.
14

---

Published as a conference paper at ICLR 2025
B INVARIANCE PROPERTY OF Adam AND SGD
Theorem 2.2. AdaSGD is rotation-equivariine Adam and SignGD are only permutationn-quiivrriant
Proof of Theorem 2.2. For SGD and AdaSGD, we will show they are rolation-cquivarian by induc- tion. For any rotating transformation R(x) Ræ, suppose F. R Rx,)
vg holds for R
R V= and m. = R m.- From the update rule of SGD, we have that ãt = ēt-1 Tj R x- n1 m R ( m,,) R ay. For the update rule of AdaSGD, we fur-
8 S I Then we have that 9: = Val. (m.,) = R V
ther have that |1 = g because R is an orthogonal matrix. Then ūN i and the derivation similar.
For Adam and SignCD, it is easy to show by induction they are equivariant w.r.t any permu- tating transformation because the operation gradient performed on each coordinate sepa- rately. We only need to show they are not equivariant w.r.. rotating transformation. We choose R = L ve'
A Lx(x) L(x) 2xf + 2.2. Due the update rule of SignGD, it can only update a and æ in the direction of [1, 1] and [1. - 1]. But when rotating the update direction on back the space of æ. The update direction can only be [L, 0] or [0. 1] that different from the update direction in the original space. Because the first step in Adam takes the same direction in SignGD, we simultaneously show that both SignGD and Adam are not rotation-equii eetiiiiiieiivaiia.....
C CONVERGENCE RATE OF SignGD FOR DETERMINISTLE LOSS
Theorem 3.2. Let L be H-smooth w.E.. too norm and 1T= be the iterates of SignCD (Adam with 31 B2 0) on with initialization By and learning rate 1), it holds that
[| S I(xo) - min L(x)
Hn
min
+ 2
18ST
Tn
we choose 17 2[L(mo) - min M())
then minigesT
VL(x;)]1 $ [2(((ma)) (((L((a))-inn) min M(r))
TH
Proof of Theorem 3.2. We will directly prove a more general verion of Theorem 3.2. Because H -ssoooth with respect to ||. |oxo, we have that
H æ x141
-L((z) T x141) L(x++1) -VL(Z))
$ - -n [VL[\\ + 1²2 r?²
(7)
This implies that
T
I((00) [((Tr)
Hn
min LEST
[V((x:)l. <
Σ
[L(xx))l. -
+ 2
T 1=1
Tn
which completes the proof.
D PROOF DETAILS
D. PROOF FOR CONVERGENCE RATE OF BLOCKWISE Adam
We will Lemma 3.12 better control the growth of the sum of second order term.
Lemma 3.12. Given any < B2 1, for any scalar sequences {vz) (v and {ge T satisfying
o?
to 0, U1 0 and 82 2 B2)9 for ' it holds that Σt=1 i2
< T + I-33
In wr to
15

---

Published as a conference paper at ICLR 2025
Proof of Lemma 3.12. Notice that 1
In for any positive We can have that
Ut
12
U1
(1
}2)
Σ
T
1
B2
1
Vt-1
1 - $2
T
82
ST
1 82
1/
t=1
82
VT
T +
(8)
82
00
when vo o O. When vo 0, we can still have that
T
Σ
91
1
gr
Ut
B2
1=2
1
B2
Or
<
+
(T
1) +
In
1
B2
B2
VT
T
In
1
32
U1/e
As mentioned in Section 3.4, we deal with the first order term by approximating it with a deter- ministic term. Recall the notation defined in Section 3.4. Ii denotes the gradient of mini-bater L2 (x)-1) at step t. And E [gt [ = VI(x;-1) because EL, = L. The full-batch gradient is g. VL(x,-)) Different kinds of second-order momentum are defined the following way
Ut 3 112 /db (1 32)
Σ 3y (1e0- /dn, j=0
dr (1 32) | /l/ /do + of) -200-1166
Lemma 3.13 (first-order approximation). With Assumption 3.10, it holds that for any block E
[B]
(b)
gr.(
dr0b82
In vr +
E
E
VI 32Td,0
E
(4)
Ut.b +
+
VI 32
Uo.b
+
t=
Proof of Lemma 3.13. The first order change in block b can decomposed into two terms.
T
T
T
It.
E
E
+
E
(9)
Ur,b
+e
Utbb
+e
Vur
+e
Vor
+e
[=1
(b)
(=1
(b)
[=1
(b)=b
16

---

Published as a conference paper at ICLR 2025
The first term E E E (i)
Suide Vioto equals to E EE Ee())=b Volote
when taking expecta-
tion conditional on 11-1. For the second term, it holds for each that
Σ
Veno+e
Vor. +E
(()=bb
91,19r, (ōr.b - v(,b)l
E +b))=b
Ur.b + +EV FEE/vu. Ur.b + Vv. t + / ++)
9r,1 - -B2) (0, /ds + og - ||9,,00112/ /db)
Σ
Vub. + +e ErL + (ver, + /ūx. +e)
=b))=b
91,19r,;(1 - 82) (Vlo 11 /oo + to + VIlen,es 1l/
( (l 112 do + to -00-vleeee
Vlgrnc 12/
Σ
Vod + F 00, + (VeLs +E+ v 00s +C)
=b)
D0,Diii v1 B2 (Ilgar1)) 12 /d + of
Vilen.ooll///
(1)
Vub. + t / 11 te
(b)==b
v VI 11 do + o - (g 12/ 2
(2) 1
Th
(10)
2
b(1)=b
Dr,b +e
E V Vller,cool2 do + to Vil9r.co /2/ [xc__1)
(
2 VIllares /// /d + of 1e.c /db |2:-1l
(1 (1- 3.
(11)
2 +()==b
(VL. + e) V
The first inequality (1) because 4r + E 2 (1 - B2) 19 12. /2//db and 0r,6 + E 2 (1
(2) (190..0112 11///d /dg + og) The second inequality (2) is obtained with AM-GM inequality. For
| ll
the term in Equation 10, it will be exactly 2 Volate after taking expectation conditional on Φ;-1
because only (Viee..o11 /de + oB Vilea.ool//ad /db depends 24. For the term in Equa-
tion 11, we have the following inequality
(Vil 11/ + +03
2 119.,w)11///b
E
xt-1
-E
(gr 11/ /db + of + Σ / - -2 (9.,0))1122 / ( 112 /db + o? 8t-1
$(j)=b
(1)
2 (l /2/0b + +o?) 2 (1 (0) 11. /db + of V /lgr, 1|2 /d xt-1
2
(1 /2/ + of
2/lo,,,0 /ds + o || 11: /db
22 ( 11: do + of
(I /do + oz
1 /db
<2/ (0, 11 /d + ofor.
The first inequality (1) replaces E E E)) Rz//db xx-1l with (9.,0))11/ /do + of based on As- sumption 3.10. The second inequality (2) is because lz norm is convex function. Then we can
17

---

Published as a conference paper at ICLR 2025
bound Equation 11 by
(1 B2) (Vionooll /d + oî
v Vllan.o) 12/0b/ |-1-1
(Ve. + /),00
=())=b
(1- 92) 2v (I.((6) /// + tgoo
(Vr.6 + E)
20,6 te
(b))=b
F
320b
V + e
(())
Then back to Equation 9, we have that
9t,i
gt.iht,i
gr,igh,
E
E
+E
14,6
+e
Dt,b +
Urb Fc
v V0 +
Am
(()
tm
1ml ((1)
L.tt
1 2 12V1-82003 R
(+ 12
E
-E.
01, +
12bb +
2
Ut +
1=1
(()==b
1=1
@()==b
R
V
B20
(b)==b Üz.b +
Ut,6 +
For the second term, we can apply Lemma 3.12 and get that
Σ (9. 1 //b
B2
UT& + VD +
<T +
In
U2D +
1 B2
Combining these two terms, we can get that
UT. + Vp.b
E
E
t=1 +(1==b 10 VI
E
In
VI
32
+ E
1=1
=b
Next we need Lemma 3.14 to deal with the denominator in the approximated first order term. The lemma is largely inspired by Lemma 6 in Li Lin (2024), where we further generalize to the case of biock-wise Adam.
Lemma 3.14. With Assumption 3.10, holds that for any block E [B]
T
28
T ob
T
T
( 11 1/2/db /db
E +
Σ
UrD +e
< 1 32
00,6 +
+
Ve
+ 2
(6)
2
v Ve.b +
t=
Proof of Lemma 3. 3.14. For each S T. we have that
E
Ur.b +e
E
120r-11 + (1 (1-))=++111llll ()) 11/ /do
oā) te
3201-1 + (1 - 32) 00² +e
1gr /1///b /db Ur,b +
(12)
E
+ (1 32
B2v:-18 + (1 32) (1)(1)
11 1//00--00/ + +of)) +e]
(0.0001122 /db
18
+ +
+ (1 - +11382
Ur.b +
18

---

Published as a conference paper at ICLR 2025
We use 120 + (1 B2)(E (()= R + oo) + 1204-1.6 + (1 - B2)0 c ei the last
step. And for each I t - 1, we have that
350 + (1 32)0z +
-E
35 04 11 + 88(1 - 82) (9)) /db + (1 13)og +
TE
18th U/ 1,b + B3 (2) (t) /db (1 - 88)00 + Br
(1)
32 1 1 + 13 32)1E
g g 112 /db F2 x2- + (1- -3500 +
32 +
11 1. + 13 - /22) gr-- 12 /dg + (1 -1- 132 102 +
82 U/ -- 1,b + (1 32 3 30 +
82 U 11 + 35(1 - B2) g (b)) 1|=² db + (1 - 32 35t1 )06 +E
85 - B2) g (b) (0)11² 112 /db
+ E
12 11 11- 18 - )22 (
1/ / + (1 ) + c
g 1 /db
<E
V 12 01---16 + (1 - 10 + / /33(1 32)
+
- b +
The first inequality (1) is because v a concave function. The second inequality (2) is based on the noise Assumption 3.10. By summing the above inequality over = 1,... 1, we have that
E
1201-11 32)03 +e
[-1
g -0- | 1/2//// Vil_o,b +
E3
8500,0 + (1- +e
+
3 (1 - 82)
s=1
1-1
||9- 12/d0
BBoob + Vo8+ +
35
B2)
V +
Back to Equation 12, we have that
t-1
g - (000 /2/b0
Urb t s / B3ro,s +
+c+
Σ
35
B2)
Ui_o,b +
==0
By summing the above inequality over
I + I,... ,T, we have that
Σ
(Vr.o +e x
T
Bynn +
T/
T
of
+e+ Σ
Σ
35
B2)EE
([ 112 /db
2
V vu +
+=
4 +1
t= ++
T
| 11= /d/
3
T
1 322 E
Vo.b
2 Vo8 +ct
12
1 3B2 1=1
V0.. +
l1 12/0d
8
B2)
T
Vub.b +
Vo + (1 + v (22
1- B2
2
Vu.b +
ge, 12
2,3
T
T
/ Vo,b +
Tb
Ve +
E
1- 82
2
2
UU.b +
1=1
19

---

Published as a conference paper at ICLR 2025
This following Lemma D.1 is to control the growth of by,6 so that the right hand side in Lemma 3.12 is indeed ④
T
+ log T I-P2
instead of O T- when all the constants are poly(T).
T
Lemma D.1. Suppose there exists H dominates L(x) and follows the partition & (Defini- tion 3.9). denote the diagonal element in block by Hb, i.e, Hig = Ha(sy. With Assump-
tion 3. 10, it holds that
2ln 1+ Ee of (V 1 + Ese( Ho + I-3,)
E ( Vr. +C
In
+ Im
+
+
Proof of Lemma D.I. From the definition of Ug,b and Assumption 3.10, we have that
E max VtDb = max bE[B]
200.0 + (1- 32)
8/
| 112 /db
bE[B]
s=1
< 1100 0c +11 - 33
8/ -
max BE(A
99. 112 /db
(13)
We can bound each
|9s.c0)112 dd as following
(1)
S 2/d. | g (b) a 1112 2/d, ( Blga, -. / 112
2/d | L 11 2/db g (lg x.-11 1|
(2)
< 4/d. V (b) ((x)) 12 4/d ( (xx--) V (o) 112 + 2/d, [[9 EEgg, (2.0))/..000 1||
((r) 12 /dg + 4 VI (1, - + Σ |190,(09 E ([g x 1]]| /dg
s max PEE
where we employ (a + b) < 2a² + 26 in (1) and (2). This bound holds for any specific b E [B] so it is also an upper bound for max ( 1|- /dg. We will further simplify the bound and first upper bound the distance between gradients with the distance between parameters
B Ho 10,-1.00) 1.(b) 2
[V -
S
(14)
b=1
The distance between parameters can be characterized by the total updates and use Lemma 3.12 in (1) below
do |
1
20,00111=
Ja.j
1
g
do
Va.b
do
8=1 Ug +
+b))=b
#=1
((j)=b
= 1/2
Eogi= 2zg /oo (1)
1/
1
32
In
1/.b + c
Ma.b
B2
40,b +
s=1
S 21 + 21
B2
maxye([] U1,
(15)
32
1o +
We can bound the distance between stochastic gradient and deterministic gradient with Assump- tion 3.10
E
(v) Elg.,(0) pc. 11/ /dv <
o
d=
20

---

Published as a conference paper at ICLR 2025
Combining these results, we can get that
(1-B2)LE
32
max
[s, 10)l1 dd
be B
2
B Hodhn - 1)² -2 2
B
<(1 B2)
max B
(x0))( /dy +
2
1=1
b=1
h=1
B
+ (1- - 82)
91- 2
21-81 4
Hnduns (s 1) 1- B2
82
E
EIn
max
Vt,0 +E
vo +
mm1
b=1
<4 max B
Vor 112 /dv + 2 Σ o +47 Σ
/
Hndpn 2
Σ
Hndon 2
32 E In
maxy E [B) V2,6/ +
1 B2
U0 + 6
We define = + Noll + +
b= of + 4 [
V (z) [xx0)]
[/ /dg for simplicity. We also 1 T such that
define G maxigisT E max 44.b + e. There exists
B Hzd, + 41
B2
do In G
G = E max U. t S C' 4
Σ
Hg
1 B2
10
+
We will upper bound the last term by the linear term G as following
B
B
82
G
4n 2
H db +
4n²TT
Σ
Hid
de
In
1- B2
20 +E
h=1
b=1
82
G(1 82)
82
1 TT
Σ
Hado
+ In4 Σ
H2ddo 2T
B2
4 bE[B] T
(1 - (10 +c)
bB(
bE[[]
G
32
82
+ TT
HRd 1n
Σ
Hn T
2
1 - B2
(1 (10 + cc)
bE(B]
bE[B]
2
G <
B2
+ (vo +6)
4 Σ
TT
2
(1 - 32) + e)
The last two inequalities come from I S _ and In(x) S x². Then get that
2
B2
G 2C +8 Σ 13d.n + 32 (E)
Σ
Hndon
2T
(1 -
(2) +
hE(B
and
E maXNe[B] UT,6 +
In
Up +
2C + 8 SiE (E(B)
A H² + 32(1v0 + +c) (EOELE) Hidon (1-32](00te)
< ln
te
2 + ΣB b= of + 2 V VI(xx)) 11
2 2 TT
EE HBdnn
In
+
Up +
Up +
(1 (10 +c)
Σm of
VL xn)
1/ E 2Td + 1-35
1
<2ln
1+
+ In 32.
U0 +
21

---

Published as a conference paper at ICLR 2025
Finally, we give the proof for Theorem 3.11. When ((i) i, i.e., each parameter forms a single block, it becomes the proof for Theorem 3.5.
Theorem 3.11 (Main, Blockwise Adam). For a specific partition O we consider the updates defined in Algorithm 3. Under Assumption 3.10, we have that
B Σ
4,2 T(1 32)
T
B doob + dve
min
E
de
S2V2E+ V2E
d
+
$ EEST
b=1
with E
E ((x0)) (xT) +
T(-
FF
HH ②) +2 82 Ex- door) and F
1+ E + + EnE(m H T 1 + In 32.
2 In
to
Proof of Theorem 3.11. From Definition 3.9, there exists a diagonal matrix H that follows and always dominates (1)) satisfying H[,, +) Tr(H) Evelo/ Hydg. In single step, we can have that
B
[(x1) - L(x1-1) & VL(x Txx; (xi--t (x(1)
H
Σ
(re,s 21 1.1)
2 hm1
(1)=bb
B (b)
B
gr.(
Hb
Ub.b F
2
4t.b
+
b=1
b=1
If we sum over i from to T and take expectation, can get
9c.
+ 1
|1
E [[(()T) L(xo)] S EE
17
H
VUrDD +E
Ur +
<
E
1]
T
B
9 (b)
+ 5
Σ
B3
Hada
T + 1 1 182
BB
In
2T,6 + E
Vebb +
10,6 +
b=1
1==1
The second inequality comes from applying Lemma 3.12. By Lemma 3.13, we have that
2
1
E
(t)
2 n/
E [(xx)
L(xr))
"
E
Σ
Hads
T
B2
In
UT +
T
D1 +e
T
1 82
20,6
+
b-
+
Σ
B
doop / 11- - 33
T
32
E In
UT +e
T
1
60.b + e B2
b=1
B
B
<
EE [(xx)) - L((TT)) +
Σ
Hods + +2 1-
B2
doo
1/T
b=
b=
32
VT.b +
T(1
-82)
Hids + 2221 1- 32
ab
max
be [
Eln
20,6 +
1=1
b=1
B Σ b=1
< E (xxo) - L(xr)) +
Hodo + +221 - B2
dbo
1/
b=1
32
B
E Caxxs
UTD + e
+ T(1 - (2)
Hndn + 22 1 32
dg
In
8o +
b=1
b=1
From Lemma D.1, we can define
2 E [l(())) L(xT)] + 11
32F ) T(1
Hads + 2 1 82 B
E
doop
T/
B2)
1=
b=1
22

---

Published as a conference paper at ICLR 2025
with
E of + [[L())))
+ +ZbElB]
+ 1
F 2In 1+
+ In 32.
V0 +
Then it holds that
B
Pi,s +
E
$
E.
VoL
b=
By Lemma 3.14 and Cauchy inequality, we have that
2
B db
Dt.
S
T
(b)
2
Σ
T
Σ
day Gz, +6
T
T
01, +
T
+1
b=1
b=1
+1
b=1
4,3
B
S /2E
4E +
dv 00
+
doob + dvv
T(1 - B2)
b=1
487
$ <2/2E + V2VE
d Σ
doos + +dVe.
T(1 - 32)
This completes the proof.
D.2 PROOF FOR CONVERGENCE RATE OF Adam
We rely on Definition D.2 that generalizes Definition 3.3 obtain Theorem 3.5 from Theorem 3.
Definition D.2 (Generalized version of Definition 3.3). For any partition function T [d] + [B] and H = (H....,, HR), we say a function L is H-blockwisely-smook
w.r.t. Φ-norm, if and only if V L(z) ( L())l, I Hd || y 4 any b [B]. x, y E Rd
The following Lemma D.3 will show that H(I, +) is upper bounded by Eb deHs when L is 77-blockwisely-smookk w.r.t. ø-norm.
Lemma D.3. For any twic differentiable loss which is H-blockwisell-smooll w.r.t. D-norm (Defi- nition D.2), we have for any x and E Rdd
A AL'L(T))
< Σ Ho
(16)
b=1
Then the diagonal matrix A defined by A.,; Hapy follows partition + and dominates L(x)
23

---

Published as a conference paper at ICLR 2025
Proof of Lemma D.3. From Definition D.2, we know that
(b) A
V (b) ((x))
H₂ 2 sup
A Anl ll
v dr max E|B]
ddy
V (b).: [(x)
sup
æ.
dh mAXKEE|
dey
Eb= (o),
L(x) (8))
sup
A woll
A
do maxve(B
day
(b) = V2
((x) (M)
(b),t0)
sup
Amil
..A.
do max
dy
1
V (i),00)) (1c)
sup
( S1 vdl
(b) y=1
(0) 11, S S du.
Vdg
(1))) (0)
sup a.A.A.
vdb
(1) (b)
M=1
2
Then for any x and A, we know that
B
dv
(A(0)) ((3))))))
Ho A
dp
A (b)
A
(b)
V2
(1))))))))
(b)
and
H₁
A
(0)112 1
b=1
Ho
Σ
Hy ||
(+) 11
(6)
6=1
==1
B
ds
(6)
(A()):
(0) (0)
B
Vdl
A (8))
K (A
(1))))
+
dg
A
(80
2
b=1
dd
(b)
B
22
V
(1))
2 A2/
(b)
hm
1
E EXPERIMENT DETAILS
E. TRAINING DETAILS
In Adam and its variants (including AdaSGD) we set (3,, 32) 0.99) for experiments the quadratic loss (Section 4.1) and ResNet18 (Section 4.3) and (3,, (2) = (0.9. 0.95) for experiments on GPT-2 (Section 4.2). Momentum in SGD is also set to 0.9. Weight decay is always deactivated.
For the quadratic loss experiments in Section 4.1, we generate orthogonal matrices R, in the fol- lowing way. We first sample M E IRdxd where Misi is i.i.d. sampled from N(0 1). Then
24

---

Published as a conference paper at ICLR 2025
A M MMM is skew-symmetric matrix and exp (t.A) represents continuous family of matri-
ces. We define R, = exp (t, for different 4.. When & = O, we know R; = I When t; + oo, R2 converges to random orthogonal matrix in distribution. We pick t 0.002, t2 0.008, 0.015, 14 = 0.1 for our experiments. The initial 2n is decided by sampling from Unif( [o. 1 1 when there is no rotation. When R, is not identity matrix, we will start training from R2 2% to ensure the initial loss values are the same across different rotations.
For the experiments in Section 4.2, we train GPT-2 small (124M parameters) on the Open Web corpus containing more than 9B tokens for 100k iterations with sequence length of 512 sequence length and 480 sentences per batch. We use cosine learning rate schedule of the same peak learning rate 6 x 10 for all the adaptive optimizers, which is also the default of nanoGPT codebase. We did a grid search to find the maximum possible peak learning rate for SGD*.
For the experiments in Section 4.3, we applied random crop and random horizontal flip augmen- tations over the training data to promote better generalization. We tuned each optimizer through searching over the same grid of learning rates The number of iterations is adjusted per batch size to result in 20 epochs for each training run (for instance, 4000 iterations were used for a batch size of 256, and 1000 iterations were used for batch size of 1024). For the training loss and training accuracy plotted in Figure 2, it is measured on a subset of augmented training data that is the same size of evaluation set. The evaluation loss and accuracy are measured on the entire evaluation set without the augmentation. Track running stats is set to false at initialization.
E.2 Adam ON A ROTATED LOSS
A key difficulty in implementing rotated Adam arises from applying orthogonal rotation on the parameters before calculating the loss. It is computationally infeasible to apply a 125M x 125M orthogonal matrix the 125M-sized parameter vector. To avoid such computation, design new orthogonal transformer to rotate the parameters of the network. In what follows, we elaborate on this rotation.
RandPerm. Given vector of size d, we can orthogonally rotate it by repeatedly applying these consecutive operations: 1. Permute the entries of the vector according to a randomly chosen permu- tation E S.. 2. Reshape the permuted vector into 3D tensor of size (#1, 82, 82,83], 83). apply fixed
orthogonal rotation of size 8 x 8 on each side of the tensor and then reshape it back to a vector of size d.
This operation performs an orthogonal transformation R on the input vector We can chain mul- tiple operations of this kind and construct RandPerm", where k is positive number indicating the number of consecutive RandPerm s applied. Building upon this rotation, we train GPT-2 125M with Adam on L RandPerm' 2 to analyze our hypothesis regarding the &o geometry of the loss landscape and to verify that Adam will indeed suffer from the induced orthogonal equivariance. Figure I con- firms our findings, as the performance of rotated Adam with RandPerm' is significantly worse than Adam. This suggests that Adam is highly sensitive to the rotation and adaptivity alone can't explain its advantage.
Concurrent work (Maes et al., 2024) also run Adam on rotated loss. They use different way to achieve global rotation efficiently. They also conduct extensive module-wise rotation experiments for more fine-grate fine-grained analysis.
E.3 COMPUTATION OF MATRIX NORMS
As mentioned in Section 3.3, it is computationally infeasible to get the full Hessian matrix and directly compute norms of it. Instead we leverage Hessian vector product function in Jax to probe the Hessian matrix. We use Lanczos algorithm to estimate spectral norm.
*Our codebase is built upon nanoGPT codehase https: https///t/.tttttt///aooorrrr
com
5 "We tried 0.00001, 0.00003, 0.0001, 0.0003, 0.001, 0.003, 0.01, 0.0,0, 0.1, We used the following values: 6.25 10-4, x,25 - 2.5x 5.0 3,.0 2.0 2,0x
4.0 x 10-2 8.0 10 1.6 x 10-1, 3.2 10 10 1.6.6 1. x 10".
25

---

Published as a conference paper at ICLR 2025
We propose Algorithm 4 to estimate the sum of absolute values for each and sum all the rows to get (1, -norm of Hessian matrix. We first subsample a fixed batch training data for estimating the (1, -ommm of Hessian matrix. The high-level idea is to compute the matrix vector products between Hessian of training loss on this batch and a sequence of random Cauchy vectors. Then take the & norm of the coordinate-wise median of the resulting sequence of Hessian vector products. Because the Cauchy distribution is I-stable, the resulting product is also vector of Cauchy random variables, and the magnitude of each element equals norm the corresponding row of the Hessian. Thus with infinitely many samples, the { norm of the coordinate median converges almost surely the (1,1)-norm) of the Hessian.
We choose = 200 for the measurement experiments on GPT-2 and n = 50 for the measurement experiments on ResNet18. We also prove a non-asymptotic high-probabillty multiplicative bound for the estimation error which depends mildly on the dimension d in Theorem E.2.
Algorithm 4 Estimation of (Lroom of Hessian,
(()
Input: Number of Cauchy vectors parameter E Rd loss L
1: for = 1 to n
2: Sample a independent Cauchy vector o(0) e Rd where (0)) itd. Cauchy(O, for = I,....d d.
3: H. L(x) y(()
(Using hessian-vector product)
4: return プラ median H
First we prove that median of random variables following uniform distribution is sub-Gaussian. Lemma F.1. Suppose Z1,".. Zm iid
Unif([O, 1)). Then
P( median ... Zn) 2
2 E) 2exp (-2ne²)
for any 2 0.
Proof of Lemma F.I. Define S = Ei=1 1 2,5--- Since z.5s follows i.i.d. Bernoulli distribu- tion with P1 1/2 c. S Bin(n,P1).
M, median [Z, Zn) S 1/2 - if and only at least
n+1 2
Z s are smaller than c. And
can apply Hoeffding's s inequality on S and get that
n
2 n
P(M, _ 2
- e)
P(S >
2
P(S
-np1
-mP1)
2
2(
mp1) )2
< exp
exp (-2ne?)
n
We can know P(M, 2 2 + < exp(-2ncc from the symmetry of distribution.
Theorem E.2. For the estimate in Algorithm 4 with Cauchy vectors, it holds that
median ( H,.:)
nd2
2ne cos ((1 A)
<2dexp(-
2 2exp
2
T²
for every A E (0.1) when 2
22 cast (1+A)7]
In other words, for any e, 8 (0, 1), loss L at parameter æ and nd extra computation time to get an estimation of (1,1)-nom L with at most multiplicative error and at least probability 8.
2(In d 1/2 In }/) hessian-vector product of the
26

---

Published as a conference paper at ICLR 2025
Proof of Theorem E.2. Define as Ei=n [ for E [d]. When 0j
iid. Cauchy(O, 1),
it holds that H5, L y. v0()) follows Cauchy(O, a,) because Cauchy distribution is 1-stable. And H5,,' Hnn are independent. Then it suffices to show that
d Σ
d
P
ag median ( (K,al, .M (l
Σ
a2
/=1
f=1
n
2nc² cos (11+
52dexp(- 2 -) + 2exp
n²
for any (G,E) such that Y,,,, .y, Cauchy(O, 1) for any j [d]. Furthermore,
(Dyyal,.. : (yml)) & (aan(X))),""" tan()))))) for X,,1.-' Xjn ~ Unif(0, }). So we only need show that
iid
d
d 6
Σ a, mediaa(taan(Xa)) ... (la(((jn))
Σ
aj
aj
5=1
(17)
n
2nc² cos cost((1 ((1+ 1)
expexp( 2
+ 2exp
x²
iid for any {Xj,k} such that X,,, Xjjn
~
Unif(0, 3) for any E [d].
Fix A E 10, We define
(a)
if - A f ((++)))) (x (1 + )) + tan((1 + ))
f(x)
- -) tan - ))
oos
(1 <ェ<
cos Then f()) is differentiable function on (0, 2) that equals to tan(x) in the middle and is linear on both ends.
We can decompose Equation 17 into
d
P
(mndianttan(X,,)) (an((
0j
j=1
j=l
d
d
Σ aj tan tanntamneet (median Mamedian(X,,,"..
Σ
é
0j
j=1
==1
<P
Σ aj lan Lan (median( X,,,' ".,...)),)))))) .x,n))
Σ
aj (median( X,1.... 1. Xzn)) >0
j=1
d
Σ a (em.mmaniaan.....mm_ Xj,n)) Σ
+
aj
>
1=1
j=1
For the first part, we have that
Σ aj (mediaan X,,,.". ,,,n)) Σ
aj (median (Ny,, ".. Xz,))) =0
j=1
P (eaaniddiin(X,... ,x,m) - (median X,,,. .x,n)) =0)
1=
P [mmdian(XX,,) ,,. ... Xz,n) π
A Z
4
J=1
<2dexp (- n 2
27

---

Published as a conference paper at ICLR 2025
where we apply Lemma E.1 in the last step.
For the second part, we first know that median ( ... (( - x. is Gub-Gaus variable with 4
a² =² from Lemma E.I. Then [(mediaa( 1 (,n)) E, (.(ediann Xi,, xzn)))
16n
is sub-Gaussian variable with o² cos (1++)) I6n because (*) S (( And
Σ/=1 az (median( (X,, xz,))) - (mediam ",,,, x,nn) sub-Gaussian vari- able with o² cos' ((1+))5 16n
[ a;) When 2 2 cos'oo((10 ((++)))) holds that
[M/(diiaa, (jm)) 1 s E [(mddian (X , .,,n)) - 1(
I max f((x)
median Xz,,, ,.n) π
1 cos ((1 + 4)
V2T S 2
Avn
and
d
a, (mediamd(Xn)) ... ,),n))
ag
aj
j=1
j=1
<P
o (median AonedinnXX,1...... ... .x,,1))
Σ
aj a (eedian ,x,m)))))
> =
aj
5=1
5=1
[ (median( (X,,1 (Xn)) 11> > )
+
j=1
2n cos ((1 cos +)
<2
exp
x²
Combining these two parts, we can get that
d
P
Σ a median
tta((Xj,))
Σ
aj
Σ aj
n
(- 2ne cos + A)
<dd <2dexp(
2exp
x²
for n 2 2e cos (1+A)))
E.4 MORE RESULTS
As mentioned in Section 4.2, we also explore how learning rate can affect the performance of differ- ent optimizers by using peak learning rate 3 x 10-4 and 1.8 10-3 for all the adaptive optimizers. The training losses are plotted in Figure 3. All the optimizers perform worse with smaller learning rate, which aligns with the common understanding that optimizers tend to work better with larger learning rate as long as the training is still stable. When larger learning rate is used, the per- formance of Adam is improved but the performance of rotated Adam becomes worse than with the default learning rate. The training with AdaSGD even completely failed. This suggests that another advantage of Adam over AdaSCD: it can maintain stable training at a larger learning rate, which is often beneficial to faster and more efficient convergence.
As mentioned in Section 4.3, we tried different batch sizes when training ResNet and the results are
in Table 4.
28

---

Published as a conference paper at ICLR 2025
610
610
610
Lan0
Tisisius Lani
2.0
300
3.0
280
290
2.0
28
2004
2000
50000
00000
1E000
2004
4000
50000
60000
1100
280
2108
2000
50000
6000
(0000)
Figure 3: Training losses of Adam, AdaSGD and rotated Adam on GPT-2 under different learning rates. All the optimizers perform worse with smaller learning rate x 10-4 Only Adam will perform better with larger learning rate 1.8 x 10 10-3
Batch Size
SGD
AdaSGD
Adam
Rotated Adam
16
0.0777
0.114
0.064
0.0905
64
0.0698
0.0854
0.0472
0.0574
256
0.0723
0.0787
0.0359
0.0485
1024
0.1115
0.0915
0.0735
0.0817
Table 4: Training losses of ResNet for different optimizers and different batch sizes within 20 epochs. For each setting, we choose the optimal performance over the learning rates. The perfor- mance of Adam is consistently the best among all four optimizers.
29


Reviews:
[{"review_id": "JanWTJnZhZ", "reviewer": "Reviewer_GxXu", "summary": "The paper studies the superior performance of Adam compared to SGD from a theoretical perspective. \n\nSpecifically, existing analysis on Adam is commonly performed under the $\\ell_2$ norm. This paper argues that Adam's advantages arise from its sensitivity to the $\\ell_\\infty$ geometry of the loss landscape, rather than the $\\ell_2$ geometry. By introducing $\\ell_\\infty$ smoothness measures, the authors develop a new theoretical framework to explain Adam's faster convergence rates. The authors further extend this framework to cover blockwise Adam variants such as Adalayer and Adam-mini.\n\nTo support the theoretical findings, the paper presents experimental evidence that:\n\n1. Adam performs poorly on rotated loss landscapes, demonstrating its dependence on non-rotation-invariant properties. \n\n2. Adam outperforms SGD in cases where the $\\ell_\\infty$ smoothness measure is significantly smaller than the $\\ell_2$ smoothness measure, demonstrating that $\\ell_\\infty$ measures are more adequate to characterize Adam's performance.", "questions": "1. Zhang et al. (2024a) also characterize the advantage of Adam to SGD. They have a theory on quadratic problems. Can the authors compare the results to that of Zhang et al. (2024a)?\n\n2. The authors point out that rotating the loss can easily hamper the performance of Adam. However, in practice, Adam performs on par or better than SGD for most neural networks. This means that deep neural networks are not ``randomly rotated''. Why? Can the authors discuss this?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The paper draws an insight that Adam is permutation-invariant, but not rotation-invariant is crucial, while SGD is rotation-invariant. I believe this property is highly related to the performance difference between Adam and SGD. \n\n2. Based on $\\ell_\\infty$ smoothness measures, the paper provides a general framework to analyze Adam and its blockwise variants such as Adam-mini and Adalayer. This contribution is timely, given the increasing interest in blockwise optimization approaches aimed at reducing memory overhead in training large language models. The framework of this paper can potentially guide the design of new Adam variants.\n\n3. Experiments show that different $\\ell_\\infty$ smoothness measures indeed lead to different performance of Adam. This provide a strong evidence of the theoretical findings.", "weaknesses": "While the paper provides good insights into Adam's sensitivity to $\\ell_\\infty$ geometry, the proposed theorems using the $ ||\\cdot ||_{1,1}$ norm may not fully capture this sensitivity, particularly in explaining the performance gap between SGD and Adam. Two specific concerns are as follows:\n\n- For convex problems with a positive semi-definite Hessian $B$, it holds that:\n\n$$   ||B||_{1,1} \\geq \\mathrm{trace}(B) \\geq || B ||_2 $$\n\nThus, for a wide class of problems, we have $ \\sup_x ||B||2 $ smaller than $\\sup_x ||B||_{1,1}$. However, the authors claim ``the latter is typically much smaller when Adam optimizes faster.''  This assertion seems counterintuitive given the inequality. It requires further justification in practical neural network training scenarios.\n\n\n- In the quadratic loss experiments (Section 4.1), Table 1 shows that Adam still outperforms both SGD and AdaSGD, even for quadratic cases $\\sup_x ||\\nabla^2 L(x)||_{1,1}$ is larger than $\\sup_x ||\\nabla^2 L(x)||_2$. This result appears inconsistent with the authors' theoretical argument, as pointed out in the last bullet point.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "7ETweAYg9f", "reviewer": "Reviewer_3yjk", "summary": "This paper presents a new theoretical analysis of the Adam optimizer, highlighting its advantages over vanilla SGD in training deep neural networks like GPT-2 and ResNets. By relying on loss smoothness under $l\\infty$-norm geometry instead of the more common $l2$-norm smoothness assumption, the authors argue that Adam's success over SGD is due to its coordinate-wise adaptivity which allows to exploit some properties of the presented models.The analysis is also extended to blockwise Adam which is a generalized form of Adam.\n\nThe proposed theory is empirically verified Adam on rotated and unrotated loss landscapes, verifying the claim that Adam uses non-rotation-invariant features.", "questions": "1. A thorough proofread is needed as some typos/redundancies exist even though the paper is quite polished (eg. l.49 \"If Adam optimizes much slower more slowly after\" -> \"If Adam optimizes much slower\" or \"If Adam optimizes more slowly\"; l.102: \"for $p \\in [1, infty]$ \" -> \"for $p \\in [1, \\infty]$\"; l. 459: \"OpenWebText corups\"->\"OpenWebText corpus\".\n\n2. A naive question one could ask is what would the results of rotated Adam be when compared to Adam on standard CNNs/MLPs and recurrent neural networks? What could be the intuition behind that?\n\n3. Are there any specific patterns in the rotations that worsen performance, or does any random rotation lead to degradation? Could the pattern be extracted from the data/model? A more detailed work on the rotation sensitivity would be very interesting.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. Even if the paper could be even more polished (see Question 1 for a comprehensive list of needed corrections), the paper is overall very-well written and interesting. I carefully read the main text and the appendices and found the proofs very clear and did not find mathematical errors.\n\n2. The authors introduce what seems to me is a novel framework to better capture Adam's coordinate-wise adaptivity, namely the $l\\infty$ geometry that allows them to get tighter convergence bounds than previous work (Défossez et al., 2022) that used $l2$ smoothness of the objective function.\n\n3. The authors demonstrate Adam's sensitivity to rotation. This is used to show that Adam overperforms SGD thanks to non-rotation-invariant properties.\n\n4. Blockwise-Adam is proposed allowing adaptive updates across parameter groups, which could be used for large models.\n\n5. An empirical validation on models like GPT-2 and ResNet-18 is provided, which supports the theoretical findings. The results on rotated Adam are in my opinion particularly interesting and the Appendix D.1 answers a natural question that arises when reading the paper as to how to apply an orthogonal rotation on the parameters on large models.\n\n6. The provided insights could be used to further improve adaptive methods in deep learning and maybe even design adaptive methods \"adapted\" to specific models as  the discussed coordinate-wise adaptivity is probably heavily linked to specific properties of the studied models.", "weaknesses": "1. The convergence rate improvements seem a bit incremental when compared to (Défossez et al., 2022) and might not translate into practical gains.\n\n2. The use of non-standard $l\\infty$ smoothness assumption may limit the generalizability of the proposed results.\n\n3. The empirical analysis of rotation sensitivity is very interesting but a bit limited in scope. The effort of including a ResNet-18 to explore a different kind of architecture is commendable but a more diverse set of architectures would be very interesting to study. It might provide deeper insight into why Adam is impacted by certain rotations and also in which cases do the non-rotation-invariant properties of Adam easily emerge.\n\n4. The blockwise Adam variant is interesting but under-exploited as there is no practical guidance on choosing parameter blocks, even though the related works of Adamini (Zhang et al., 2024b) and Adalayer (Zhao et al., 2024) are mentioned.\n\n5. The high memory consumption of Adam is only mentioned in the introduction. In practice, Adam's higher memory usage when compared to SGD (three times as much) is a real drawback, particularly for large models relying on the Transformer architecture which is studied in this work (GPT-2). One could ask if training a three times bigger model with SGD before reducing its size with common techniques (distillation, pruning) would lead to better results.\n\n6. The focus on the case $\\beta_1=0$ seems a bit far-fetched when looking at the results of Table 1 as the final losses obtained are far better with ($\\beta_1 = 0.9, \\beta_2 = 0.99$)", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "cRBIufmOPh", "reviewer": "Reviewer_chAM", "summary": "The authors give a unified proof for AdaSGD, Adam, and blockwise Adam under a new Lipschitz and noise assumption. The new assumption does not only consider the gradient Lipschitz of the full gradient but also the Lipschitz property for each block (for SGD, the full parameter, for Adam, each coordinate).  From the theoretical results, the authors find that (1,1)-norm is the critical value for convergence. Experimental results validate the conclusion for Adam-type algorithms.", "questions": "1.  Do the results in Section 3.3 cover results in Section 3.2? If so, why should we introduce section 3.2? If not, can you specify the major difference?\n\n2. A simple example: when we optimize a quadratic function with a positive diagonal matrix. In this case, Adam converges much faster than SGD while the (1,1)-norm of the matrix is always larger than the 2-norm of the matrix, because the (1,1)-norm is the sum of diagonal value while the 2-norm is the maximum value of the diagonal entries.  It seems that the result in the simple case contradicts the results provided in section 3.3. \n\nCan the authors give some concrete examples showing the correctness of the claim?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 3, "strengths": "1.  The authors propose a unified algorithm that contains Adam, AdaSGD, and blockwise Adam.\n\n2. The authors propose a more detailed assumption on gradient Lipschitz to characterize the underlying function carefully. Thus, they can give a tighter bound than giving the overall Lipschitz constant.\n\n3. The authors find that (1,1)-norm of hessian is positively related to the performance of Adam in both theoretical analysis and experimental validation.", "weaknesses": "1.  From my point of view the theorem in section 3.3 has already covered the results in section 3.2, making section 3.2 meaningless.\n\n2. The authors claim that in their proof, we can see the reason that Adam can be better than SGD, while the explanation of the results is only given by $\\sup_x ||\\nabla^2 L(x)||_{1,1} \\leq \\sup_x ||\\nabla^2 L(x)||_2$. It should have some reasonable examples.\n\n3. In Table 1,  since the convergence of AdaSGD is related to 2-norm instead of (1,1)-norm, why do the authors not report the 2-norm instead of reporting (1,1)-norm twice?\n\n4. There are some typos in the paper: \n\ne.g., The first line of notation $\\sum_{i=1}^d$ instead of $\\sum_{i=1^d}$, $\\infty$ instead of $infty$. The name of the algorithm is \"Adam-mini\" instead of $Adamini$. I do not carefully check every detail of the writing, but the authors should go through and correct the typos.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Spotlight)