Why Transformers Need Adam: A Hessian Perspective
Authors: Yushun Zhang
Congliang Chen
Tian Ding
Ziniu Li
Ruoyu Sun
Zhi-Quan Luo
Venue: NeurIPS 2024 poster (2024)
Keywords: Transformers
Adam
Optimization
Primary area: optimization_for_deep_networks

Abstract: SGD performs worse than Adam by a significant margin on Transformers, but the reason remains unclear. In this work, we provide an explanation through the lens of Hessian: (i) Transformers are "heterogeneous'': the Hessian spectrum across parameter blocks vary dramatically, a phenomenon we call "block heterogeneity"; (ii) Heterogeneity hampers SGD: SGD performs worse than Adam on problems with block heterogeneity.  To validate (i) and (ii), we check various Transformers, CNNs, MLPs, and quadratic problems, and find that SGD can perform on par with Adam on problems without block heterogeneity, but performs worse than Adam when the heterogeneity exists. Our initial theoretical analysis indicates that SGD performs worse because it applies one single learning rate to all blocks, which cannot handle the heterogeneity among blocks. This limitation could be ameliorated if we use coordinate-wise learning rates, as designed in Adam.

Why Transformers Need Adam: A Hessian Perspective
Yushun Zhang" Congliang Chen Tian Ding Ziniu Li Ruoyu Sun 123. Z2i-Qua Luo
'The Chinese University of Hong Kong, Shenzhen, China
Shenzhen International Center For Industrial And Applied Mathematics, Shenzhen, China SShenzhen Research Institute of Big Data, Shenzhen, China
Cyushunzhang, congliangghenn,ziiiiu ziniuli)elink cuhk. edu cn dingtiandsribd.cc cn, sunruoyu@cuhk. edu. en, luozq@cuhk. edu.
Abstract
SGD performs worse than Adam by significant margin Transformers, but the reason remains unclear. In this work, we provide an explanation through the lens of Hessian: (i) Transformers
"heterogeneous". the Hessian spectrum across parameter blocks vary dramatically, a phenomenon we call "block heterogeneity", (ii) Heterogeneity hampers SGD: SGD performs
than Adam on problems with block heterogeneity. To validate (i) and (ii), we check various Transformers, CNNs. MLPs, and quadratic problems, and find that SGD can perform on par with Adam on problems without block heterogeneity, but performs worse than Adam when the heterogeneity exists. Our initial theoretical analysis indicates that SGD performs worse because it applies one single learning rate to all blocks, which cannot handle the heterogeneity among blocks. This limitation could be ameliorated if we use coordinate-ve learning rates, as designed in Adam.!
Introduction
Transformers [83] have become major workhorse behind Al development (e.g., [1]). However, the understanding of Transformer training remains limited. For instance, Transformer training largely relies on the Adam optimizer [45, 57]. In contrast, stochastic gradient descent with momentum (SGD)?, the de-facto optimizer for convolution neural networks (CNNs) [49], performs significantly worse than Adam on Transformers (e.g., Figure 3). Yet, the reasons behind this performance gap remain unclear. Understanding why SGD performs worse than Adam Transformers question. First, from a theoretical perspective, this can help us better understand the training of Transformers and more generally, neural networks. Second, from computational perspective, the understanding may inspire the design of better algorithms for training neural networks.
intriguing
In this work, we explore why SGD largely underperforms Adam on Transformers through the lens of Hessian. start by investigating full Hessian spectrum of Transformers, i.e., the full eigenvalue density of Hessian (see Figure 1). By theory, the full Hessian spectrum largely determines the behavior of gradient-based methods [63, 33, 79, 37], so we suspect it may also help explain SGD's unsatisfactory performance. Using tools from numerical linear algebra [7]. empirically compare the full spectra of CNNs (where SGD is on par with Adam) and those of Transformers (where SGD largely lags behind Adam). Unfortunately, Transformers are often largely similar despite the different optimizer behaviors. identified critical features in the full Hessian spectra associated with the gap between Adam and SGD on Transformers. To reveal the cause, a more fine- mneggraieed investigation into the Hessian needed.
shown in Figure the spectra for CNNs and
such, we have not
*: Correspondence author. !Our code is available setps://githon www//yushun/heeeeeeeeeo aarrusshunhesiiann spectrum. We introduce update nules of Adam(W) SGD in Appendix C.1.
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

VODTA LONDAR
VODLS 123010
VCGLO 123040
VCGSS (2300))
Ko-Bee CLOPE
Wordeict CLIM
Autbiiie (LIM)
WAT (1.900
(422
(272
CTT (23M)
(a) Initialization
(b) 25% steps
(C) 50% steps
(d) 100% steps
Figure I: The full Hessian spectra of CNNs (VGG16 and ResNet18) and Transformers (GPT2, GPT2-nano, and ViT-base) at different training stages. The x-axis records the eigenvalues and the y-axis records the frequency in the log scale. To allow comparison in the same figure, the plotted spectra are normalized by their 10th largest eigenvalues. We find that the spectra on CNNs and Transformers are largely similar.
What would cause SGD to perform significantly worse than Adam Transformers, but not on CNNs? By dissecting the structures of CNNs and Transformers, we notice that CNNs are constructed by the repetitive stacking of similar parameter blocks (convolution layers), while Transformers involve the non-sequential stacking of disparate parameter blocks (e.g. Query, Key, Value, Output projection blocks in attention and MLP layers). We hypothesize that these architectural differences might lead to different optimization properties. Intuitively, disparate parameter blocks contribute differently to the overall loss. So each block might benefit from specialized treatment by optimizers, a flexibility offered by Adam but not by SGD. This observation motivates us to investigate the Hessian spectrum of each parameter block, which we refer to as the blockwise Hessian spectrum.
By inspecting the blockwise Hessian spectrum, we discover a possible explanation for why SGD worse: the "heterogeneity" inherent in Transformers. We provide both empirical and theoretical evidence to support this explanation. Our contributions can be summarized as follows:
Why SGD underperforms Adam on Transformers. We explain why SGD is worse than Adam Transformers by examining the blockwise Hessian spectrum. First, we identify phenomenon called "block heterogeneity", which refers to the large differences in the Hessian spectra across parameter blocks. This block heterogeneity is observed in all examined Transformers but not in CNNs. Second, we verify that block heterogeneity hinders SGD. Across various Transformers, CNNs, and MLPs, we show that SGD consistently performs worse than Adam problems with block heterogeneity but can perform similarly to Adam, otherwise.
Theoretical results on quadratic models. We construct convex quadratic problems with and without block heterogeneity and find that gradient descent (GD) largely underperforms Adam on problems with block heterogeneity, but can perform comparably otherwise. Our theoretical analysis shows that GD can be slower than Adam on quadratic problems with block heterogeneity. We point out GD is slower than Adam because it uses a single learning rate for all blocks. The deficiency be mitigated by assigning different learning rates across blocks, Adam does.
We emphasize that we do not claim block heterogeneity is the only cause for the performance gap between Adam and SGD, but just that it is at least one important cause. We verify, both empirically and theoretically, that SGD underperforms Adam when block heterogeneity is present.
2 Problem Settings and Initial Attempts
2.1 Problem Settings
Notations. We denote the training loss L(w), where w E Rd is the neural network parameters. We denote the gradient and Hessian of the training loss w.r.t. neural network parameters V V((w) Rd and V²((w) E Rdxdd respectively. We use [d] to denote the index set 2,, .d.. Given an arbitrary partition (Di) over [d] with di [Dr), we can split w into L parameter blocks (wi)f where wy d consists of parameters with indexes in the I-th block D.. We denote [(w]]
,d}.

---

R4xx as the Hessian of 1-th parameter-block wr, where [VL[(w)]sss [V][(w)] is the 1-th principal block sub-matrix of V (tc).
8² [(wi). Note that S0r,
Setups. Hessian of large-scale NNs are intractable to compute and store. In this work, we apply a numerical tool called Stochastic Lanczos Quadrature method (SLQ) [7] to approximate the Hessian spectrum. SLQ uses a smooth curve on I to approximate the histograms of eigenvalues (see Figure I as example). A detailed introduction to SLQ is provided in Appendix C.2. All experimental setups are shown in Appendix D. We focus primarily on the following models/tasks.
CNNs. We study ResNet18 (1IM) and VGGI6 (138M) on ImageNet [40, 78]. On these tasks, SGD performs par with Adam. See Figure 9 in Appendix B for the evidence.
Transformers. We study Transformer with various scales and modalities, including GPT2 (125M) on OpenWeb Web [71]; ViT-base (86M) on ImageNet [27]; BERT (40M) on Cornell Movie-Dialog, Corpus [25]; GPT2-nano (11M) on English corpus. On these tasks, SGD performs significantly worse than Adam. See Figure 10 in Appendix B for the evidence.
For each model. we estimated (1) the full Hessian spectrum ([w(), and (2) the blockwise Hessian spectrum [V>[(w)),, E [L]. For the latter, we split w according to the default partition in Py TyTrch implementation, e.g.. Embedding layer, Query in each attention layer, Key in each attention layer, Value in each attention layer, etc. Note that the term "block" differs from the term "layer". For instance, Query and Key can reside in the same layer but are different parameter blocks.
2.2 Full Hessian Spectrum Is Not Informative Enough
We study the full Hessian spectrum of Transformers for two reasons. First, as stated in Section 1, the Hessian spectrum significantly influences the behavior of gradient methods [63]. Second, previous research shows that the Hessian spectrum provides insights into neural network phenomena, like BatchNorm's 's effect on training speed [32]. Therefore, we hypothesize that the Hessian spectrum may also explain why SGD largely lags behind Adam on Transformers.
We compare the full Hessian spectra of CNNs (where SGD performs similarly to Adam) and those of Transformers (where SGD underperforms Adam), as shown in Figure 1. Unfortunately, the results suggest that the full Hessian spectrum alone may not suffice explain the gap between Adam and SGD on Transformers. We elaborate as follows. The primary information in the spectrum lies in its (A) dispersion, (B) shape, and (C) evolution during training. Regarding (A), we observe that the eigenvalues are dispersed similarly across different models, with no notably large outlier for Transformers. Thus, dispersion does not seem to be related to why SGD is worse than Adam. We further investigate (B) and (C). For all CNNs and Transformers in Figure 1, we observe similar phenomena: the spectrum's shape is approximately symmetrical around O at initialization. As training proceeds, the majority of negative eigenvalues disappear, and the shape evolves into combination of a "bulk" and some "outliers"" Since the spectral shape and evolution are quite similar for both Transformers and CNNs, they cannot explain why SGD is worse than Adam on Transformers, either. In summary, we have not identified any critical phenomena in the full Hessian spectra that can be linked to the performance gap between Adam and SGD Transformers.
2.3 Motivations of Investigating Blockwise Hessian Spectra
What other factors could cause SGD to perform significantly worse than Adam on Transformers but not on CNNs? We identify one critical feature that has been overlooked in the full Hessian spectrum analysis above: the building-up rules of Transformers. As shown in Figure 3, CNNs are constructed by the repetitive stacking of similar parameter blocks (convolution layers). contrast, Transformers consist of disparate parameter blocks, e.g. Query, Key, Value in attention, and MLP layers. Further, these blocks are stacked in a non- non-sequettial manner. We hypothesize that the "different designs among parameter blocks" can be reflected in the Hessian of these parameter blocks, which might affect al algorithmic behavior. This inspires us to investigate the blockwise Hessian spectra, i.e., the spectrum of principal blocks of Hessian [VZ((w)),, E [Z].
In parallel the motivation above, we further provide another evidence that blockwise spectra might be helpful. Classical literature showed that the Hessians of neural nets are near-block bloockkkkkddooaal matrices [18], i..., the magnitudes in the Hessian principle blocks are much larger than those in the off-diagonal blocks. We restate their findings in Figure 2 (a). This implies that the majority of
phttps://gitta hwt.htbbb...okkpppttyyaaan////
3

---

(a) Hessian an MI.P
(h) Hessian of an
(c) Hessian ML.P
(d) Hessian of an
[18] after step
ML. at 1% step
at 50% step
MI. at 100% step
Figure 2: (a): The Hessian of an MLP after 1 training step reported in [18]. (b,c,d): We calculate the
Hessians of an MLP (with neurons) different training stages. We find the sear-block-diagome
structure maintains along the training.
Hessian information indeed lies in its principle blocks, and the blockwise Hessian of neural nets might contain valuable information.
To summarize, the "heterogeneous" building-up rules of Transformers inspire us to check the block- wise Hessian, i.e., the principle blocks of the Hessian. The classical results of neural nets [18] further support to explore this direction since they find that the majority of Hessian information indeed lies in its principle blocks. In the following, we study the blockwise Hessian spectra of various neural networks. For case of implementation, we define parameter blocks under the Py PyTorch partition. We show that the blockwise spectra indeed carry more information than the full spectrum for distinguishing CNNs and Transformers.
Remark: why near-block- ?iaigonll We briefly restate the analysis in [18, Section 7] to explain the near-block-e block-diagonal Hessian structure of neural nets Consider minimizing {(f x). v) where .(-,-) is the Cross-Entropy (CE) loss, [(0, x) = Σi= vad(ws x)) is an 1-hidden-layer neural network with input x E Rd, weight w; E Rdd 04 E IR, and label y E {0, 1), then the off-diagonal-bloog Hessian elements will contain
@ef((0 x)., po(ul>) (1 po-oooolii po(y) (wTTx) x)
T for i + j.
(1)
Ow,,wy
where (o(lx) 1/(1 exp(-yf((.. x))) and d'(-) is the derivative of 6(-). Note that the term (o(y)z) ()) (1 - po x)) will vanish rapidly since the training objective is to maximize po(y|z)). Con- sequently, this drives the Hessian towards near-block-diagonal configuration, with each block representing an output neuron. This result is validated in Figure 2: we find that the near-block- diagonal structure appears at 1 % step and it maintains along the training.
3
Main Results
3.1
Transformers Exhibit Block Heterogeneity in Hessian, while CNNs Do Not
We now compare the shape of blockwise spectra in VGG16 [40] (CNN) and BERT [25] (Transformer). We sample four blocks for each model and present the spectra in Figure 3. In BERT, the spectra of embedding, attention, and MLP blocks are largely different. contrast, in ResNet, the spectra of convolution layers are similar. We further verify this observation for the rest of the parameter blocks. We calculate the Consen-Shanne (JS) distance between two eigenvalue densities of all possible block pairs and show the results in Figure 4. We summarize our findings in Observation 1.
TAAL
(a) VGG16
(b) VGG16
(c) BERT
(d) BERT
Figure 3: (a) (c): The blockwise Hessian spectra of VGG16 (CNN) and BERT (Transformer) at initialization. The z-axis records the eigenvalues and the y-axis records the frequency the log scale. To allow comparison in the same figure, we sample 4 blocks in each model. The plotted spectra are normalized by their 10th largest eigenvalues. The spectra are similar among blocks for VGG and differ significantly across blocks for BERT. (b) (d) Adam v.s. SGD for training VGG16 and BERT.
'We note that this analysis is rather informal: rigorous theoretical study is left for future work.

---

(a) ResNet18
(b) BERT
(c) GPT2-nano
(d) VGG16
(e) ViT-base
( GPT2
Figure 4: The JS distance among blockwise Hessian spectra at initialization. We find that the JS distance of blockwise spectra in CNNs is significantly smaller than that in Transformers.
Observation 1: For all Transformers we checked, the blockwise Hessian spectra are largely different from each other. In contrast, the blockwise Hessian spectra of CNNs are similar.
In the following, we refer to the phenomenon Transformers as "block heterogeneity". and refer to that of CNN as "block homogeneity". The observations in Figure 3 and 4 indicate that block heterogeneity is informative distinguishing CNNs and Transformers. In the following, we will show that the block heterogeneity is strongly correlated with the performance gap between SGD and Adam on Transformers.
3.2 SGD Performs Worse than Adam on Various Tasks with Block Heterogeneity
Figure 3 and 4 have shown that (1) SGD is worse than Adam on Transformers. (2) Transformers have block heterogeneity. Now we further link block heterogeneity to SGD's s unsatisfactory performance on non-Transformer models. This would directly establish connection between "block heterogeneity" and "why SGD is worse than Adam", without going through Transformers or attention blocks as an intermediary. We consider man-made example and one real-world example.
Example 1: man-made MLP. We consider a 4-layer MLP on MNIST and change the degree of heterogeneity by scaling each layer by constant Figure 5 (a) shows SGD gradually performs worse than Adam as heterogeneity grows.
(a) Final performance of Adam and SGD on man-made ML.P
(b) MLP-mixer
SGD V.S. Adam on MLP.
mixer
Figure 5: (a) SGD v.s. Adam man-made MLP with different degrees of heterogeneity Each point records the best-converged test accuracy under learning rate grid search. SGD performs worse as heterogeneity grows. (b) The IS distance among blockwise Hessian spectra for MLP-mixer [81] at initialization. We observe heterogeneity. (c) SGD performs worse than Adam on MLP-mixer.
Example 2: MLP-mixer. We consider MLP-mixer [81], famous alI-MLP architecture that outperforms CNNs and ViT' on some vision tasks. Figure (b) (c) show that the initial Hessian of MLP-mixer has block heterogeneity and SGD lags behind Adam this architecture
5

---

We summarize the findings so far in Observation 2.
Observation 2: For all tasks that we checked, SGD is worse than Adam when block heterogeneity exists, regardless of whether Transformers or attention mechanisms are utilized.
3.3 Reduced Block Heterogeneity in Pre-trained Transformers
We remark that different Transformers exhibit different levels of block heterogeneity. Although all examined Transformers show strong block heterogeneity, we find that this heterogeneity can be mitigated, resulting in less performance deterioration for SGD. As illustrated in Figure pre-trained GPT2 on SFT tasks can exhibit less block heterogeneity compared to pre-training GPT2 from scratch (Figure 4 (f)). In this case, although SGD is still slower than Adam, it achieves a similar loss at convergence. Compared with training GPT2 from scratch (Figure 10 (d) in Appendix B), the performance gap between SGD and Adam is significantly narrowed down. These findings suggest that the heterogeneity induced by architectural design can be alleviated by selecting "good" weights. This partly explains why simpler methods like SGD and even its zeroth-order version can still be effective for fine-tuning language models, albeit with slower convergence [39, 60].
24-
44
120
(a) JS Distance
(b) Training Loss
Figure 6: We finettune GPT2 (pre-trained) on Alpaca Eval, and plot (a) the JS distance among blockwise Hessian spectra: (b) the training loss of SGD and Adam.
In Figure 13 in Appendix B, we further report the evolution of the block heterogeneity of ViT-base along the training. Similarly to GPT2 in Figure 6, we find that the block heterogeneity of ViT-base tends to reduce after the training. In addition, we find that SGD can perform better when initializing at the weight with less heterogeneity, e.g., initializing at 50% total training steps. We hypothesize that "the attenuation of Hessian heterogeneity" is common phenomenon after training. and we leave detailed investigation as a future direction.
Observation 3: Block heterogeneity in Hessian tends to reduce after (pre)-training.
3.4 Implication on Choosing SGD or Adam
We have shown that SGD can largely underperform Adam on various architectures. This leads to an intriguing question: Can we predict the incompetence of SGD before the training begins?
Our findings can bring up an empirical guidance: we can compute the blockwise spectrum of initial Hessian, and then decide whether to use Adam or SGD. Such a method could be useful in scenarios in training large models that are not mainstream Transformers or CNNs, e-g., Mamba [38]. In these cases, there is not much prior experience in choosing optimizers. It would be intriguing to decide whether SGD is suitable for the task before the training is launched. One might argue that simple trial is enough: try both SGD and Adam; if Adam is remarkably better, then pick Adam; if Adam and SGD are similar, then pick SGD. Nevertheless, this simple approach may not be casy for large models. First, for large models, it may take days to know run of an algorthm is good not. Second, it requires tuning hyperparameters at least a few times to get a reasonably good judgement, making the cost of trial even higher.
We here propose a quantitative metric that could predict the incompetence of SGD before the training. With the help of this metric, we could save much expense on the trial and error for SGD. The metric is simply the averaged JS distance among blockwise Hessian spectra at initialization, i.e., the averaged value in the heatmap of Figure 4. We denote it as J We present J of various models in Table 1. Note that J establishes a quantitative difference between the loss landscape of Transformers and CNNs. Further, J is independent of optimizers and could be checked before training.

---

Table 1: JS denotes the average distance between the initial Hessian spectra of each pair of parameter blocks. A larger JS suggests that the task is more difficult for SGD.
Model
ResNer18
VGGI6
GPT2 (pretrained) MLP-mixer
BERT
GPT2
ViT-base
JS
0.10
0.09
18.84
34.90
53.38
83.23
286.41
To validate the effectiveness of the quantitative metric JS we summarize J of different models and the corresponding SGD performance in Figure 7. We find that the performance gap between SGD and Adam becomes greater JS increases. Thus, J.5 can serve as potential indicator to predict whether SGD may underperform Adam.
Easy SGD
Difficult for SGD
Homogeneity in Hessian
Heterogeneity Hessian
VGG16 JS 0.09
GPT2 (pretrained)
MLP-mixer JS 34.90
GPT2 JS 83.23
J5 18.84
Figure Comparison of JS and the performance of SGD on different models. We find the performance gap between SGD and Adam becomes greater as JS increases.
Finally, we remark J is rather expensive to compute due to the overhead of SLQ: it requires comparable time to one training run. Fortunately, find the original SLQ is rather redundant for measuring hessian heterogeneity. We propose some simple tricks to significantly reduce the computation time, while still effectively detecting the Hessian heterogeneity, We call simplified SLQ and we present it in Table in Appendix B. As result, the simplified SLQ can obtain the same message in Table while only taking negligible time (e.g., 0.001s for ResNet18).
4 Case Study of Quadratic Models and Preliminary Theory
Now we study quadratic functions with block diagonal Hessian, with or without block heterogeneity. Note that insights on quadratic models could be important for understanding realistic NNs, as mentioned by researchers such as LeCun et [50] and OpenAl team [44].
Setups and additional notations. We consider the following quadratic minimization.
min L(w) = w² Hw hTT w, wERd
where H Idddd is positive definite and h Rdd. We denote L the minimum value of L(w). We as a block diagonal matrix: H diag(H H2), where H; E Ridi and d Σl_ di- We use wi Id denote the variable in the /-th block and (wf'.... wF)T Rd Similarly for hi E RRd, Similarly, we [VW( E Rd: to denote the gradient in the (-th block and denote [C(w)] (ww))" Hww hf wy as the objective function w.r.t. the 1-th block. Note that L(w) [t](((w), We denote 2 A.... X as the eigenvalues of H. Similarly for Ara At,1 Au,d We denote As and ky Xz.d as the condition number of H and H,, respectively.
Xo We say algorithm has complexity Ô(C) if it takes Ooolog(11/ (/)) iterations achieve error L(w)-L-
I-C c, where n is the initial point.
4.1 Experimental Observations
We consider four types of Hessian H follows. For all cases, set condition number 5000.
7

---

(a) Hessian with GPT2 (b) Hessian with ResNet18 (c) Hessian with simplified Hessian with simplified block-wise spectrum biockwise spectrum
heterogeneous blocks
homogenenus blacks
Figure 8: The performance Adam and GD on homo/heterogeneous quadratic problems. The condition numbers of Hessian equal to 5000 for all four cases. When blocks are heterogeneous, GD largely lags behind Adam, and GD performs similarly to Adam if otherwise.
Case 1: Hessian with Transformer-type spectra. We choose = 4 and de = 25. For E [L], we construct H, Qf where Qi are matrices with i..d. standard Gassian entries and A are diagonal matrices. For the diagonal elements in A₁, we sample di numbers according to the spectrum of the embedding layer; 3rd Query, 3rd Value, 3rd MLP (fc layer) in GPT2. Shifting and proportional scaling are performed to ensure all elements in A, lie in the interval [1. 5000]. This ensures strong convexity and controls the condition number of H equals 5000. The spectra of Hi are in Figure 14 in Appendix B. We choose = O for all cases.
Case 2: Hessian with CNN-type spectra. We the same setup as in Case 1. For the diagonal elements in A, we sample di numbers according to the spectrum of the 1st to 4th convolution layers in ResNet18. We then shift and scale A; the interval [1, 5000] to ensure strong convexity and a condition number of 5000. The spectra of H2 are shown in Figure 15 Appendix B.
and di 3. For E [L], we construct Hi Qf where QI are independent standard Gas- sian random matrix and A are diagonal matrices. We the diagonal elements of A as 21, 3), (99, 100, 101}, {4998, 4999, 5000} for / = 1,2, 3, respectively. The spectra of H1 are different due their different supports. The condition number of Hessian H is 5000.
Case 3: Hessian with simplified heterogencous spectra. We choose L
Case 4: Hessian with simplified homogeneous spectra. We consider the same setup as Case 3. We set the diagonal elements of A {1, 99, 4998). {2, 100, 4999), {3, 101, 5000} for 1, 2, respectively. The spectra of Hi are similar. The condition number is 5000.
Now we study two types of optimizers: one that assigns a single learning rate for all blocks, and one that assign different learning rates across blocks.
Single-learning-ratt optimizer. We study gradient descent (GD).
u r+1 w² v/W)) w² (HH -))
We use the optimal learning rate 7 = W+L [63]. We use standard Gaussian initialization.
optimizer. We study Adam with a constant learning rate and with no bias correction for simplicity (Algorithm 3). We set 3 = to erase the effect of momentum. This helps to focus on the effeet of coordinate-wis learning rate (or the effect of diagonal preconditioning) in Adam. We
O. We consider 32 = 1 and 32 0.99, respectively. When /3 = 1, Adam assigns coordinate- wise learning rates according to the initial gradient, but these learning rates are fixed along iteration. The update rule is as follows.
w0 " w n(Daatmm)
h).
(3)
where DO Adam (V(()) V((w")) à and VL(w") Hw h. When B2 1. the coordinate- wise learning rates adaptively change along iteration. The update rule is as follows (note that v((w)) =HHHu Hk h.).
w = w¹ n(D' VVL(w) = =w' ( -n(D'eeam
'HH I -h),
where
(4)
1
Σ 3 13-'"(((")>> VL(w") ( B'VL(w") VV(
D' diag (1-32)
We grid search 7 and the standard Gaussian initialization. We remark that when 82 < 1, Adam would bounce among non-optimal points. This will be shown in Proposition 2.
Summary of experimental observations. Figure presents two phenomena. For Hessian with heterogeneous blocks (Case and 3), GD largely lags behind Adam. For Hessian with homogencous
8

---

blocks (Case 2 and 4), GD is on par with Adam. We emphasize that all Hessians have the same condition number. Further, Hessian in Case 3 and 4 share all the eigenvalues (not just the extreme ones). The gap between Adam and GD is purely due to the different blockwise spectra caused by the different locations of eigenvalues. Case 3 and 4 help reveal the causal relation between "block heterogeneity in Hessian" and "GD is worse than Adam". We hypothesize that GD performs badly because it uses one single learning rate for all blocks, which cannot handle the heterogeneity among blocks. Such heterogeneity can be better handled using different learning rates across blocks, as designed in Adam.
4.2 Initial Theoretical Results
We now provide initial theoretical results to characterize how GD lags behind Adam in problems with heterogenous Hessian. Note that classical optimization theory depicts the rate of first-order methods by the condition number of the full Hessian K. However, we point out that is not informative enough to describe the performance gap in Figure since k is the same in all four cases. To distinguish Adam and GD, we need to utilize more fine-grained quantities like blockwise spectra of sub-matrices.
Note that the blockwise spectrum is not common in the optimization area. The most related notion is perhaps "block Lipschitz constant" [8] for studying block coordinate descent (BCD) type methods, but was not linked to the performance of SGD or Adam before. To our knowledge, we are not aware of any theory of Adam GD built on the block diagonal structures or the blockwise spectra of Hessian. We now make an initial attempt in this direction. We first present the lower bound for GD.
Proposition 1. (Lower bound for GD.) Consider min, L(w) 1 w Hw - h² w where H E Rdxd is positive definite and h R.d. Let wODD be the output of GD after t steps. There exists block diagonal matrix H, h and an initial point wº, s.t., for any 17, have;
(1-x2 x (C(wGD) L²)
L(wED)
L.
(5)
where K is the condition number of H.
Proposition shows that GD has complexity (Kx) and such complexity is tight. Now we prove that Adam can achieves better complexity. This is because it chooses different learning rates for different block sub-matrix Hi via its diagonal preconditinoner Du Adama We consider generic random initialization that covers commonly used distributions such as Gaussian, Uniform, etc.
Assumption 1. (Random initialization.) ) Assume the initialization w" is sampled from continuous distribution, i.e., the probability measure (induced by w²) of any sero-Lebesgue-mett -eesure set is 0. Theorem 1. (Upper bound for Adam with 32 = 1.) Consider the same setting as Proposition I and consider Adam with 31 = O and B2 = as in (3). Assume the initialization satisfies Assumption 1. Adam
Let A'da be the output of Adam after t steps. Let 17 = minee]]] JT,"" Then w.p.1., we have
(1-
L(w 1+1 - - S max max
(E(w'Atam) -
(6)
[E[L]
KAdom.1
where E Adam. Tlit, Kt the condition number of H2, constant relates to w defined as: r= maxig[L) C2
[[V[(
Maal
[VC(u"))as)
minug(r] C7,1"
where Ci,1 min
ie|dy
A
CC.2
max iE(d)
A11
(7)
The proofs of the above theorems are shown in Appendix E. Theorem states that Adam (with 32 = 1) has complexity Ô (r- max ki). We note that coefficient r depends on the ratio between initial gradient and the principal eigenvalue for each block, and smaller ratio would give faster convergence. We further remark that condition 3r = 1 is necessary because any /32 < 1 causes non-convergence issue [10, 21J. We restate their results in Proposition The non-convergence also observed in Figure 8 (c), where we find that the iterates of Adam quickly converge to near-optimal solutions, and then bounce back. As such, 82 1 is necessary for asymptotic analysis. The analysis for 32 1 is still meaningful since it still shows the effect of Adam's 's preconditioner.
As shown in [21], the coneeenveeee is due to the constant learning rate. Reducing the learning rate reduces the gap between L(w Adam) and L", but does not remove it.
Proposition 2. (Non-convergence of constant-learning- rate Adam with 32 < 1.) [21, Proposition 12, Figure 1] Consider minweR L(w) = tw².. Consider Adam with 31 = and 32 < 1 as in (4). Let w Adam be the output of Adam after steps. There exists a discrete limit cycle for (4) and lim infi-ooo (C(w'ddmm) - L*) > 0.

---

We now compare the complexity of Adam and that of GD. By Theorem 1, Adam is faster than GD when r- [ Ni I K. In the quadratic model with heterogeneous blocks (Case 3), our simulation over 1000 trials shows that r Σ 1000 with probability 2 3 when using standard Gaussian random initialization. Since [ Ki ~ 1, we have [ Rit I 1000, w.h.p., and is about 5x smaller than y = 5000. So Adam could be 5x faster than GD, w.h.p.. This is indeed observed in Figure 8 where Adam outperforms GD by significant margin. We summarize the complexity of GD and Adam in Table 2.
some common misconceptions. During the review process, we find that readers might conclude that "Theorem implics Adam under homogeneity has worse complexity than Adam under heterogeneity". We now clarify that this claim is not correct, and there is conclusion "whether Adam under homogeneity faster or slower than Adam under heterogeneity". Similarly, Theorem I does not imply "Adam always converges similarly as GD under homogeneity". Though it is observed on CNNs, there is no general conclusion of this kind. For interested readers, we provide a detailed explanation in Appendix B.
Remark:
Table 2:
The complexity of GD and Adam for minimizing strongly convex quadratic function with block diagonal Hessian. The symbol X means non-convergence. к and Ką denote the condition number of the full Hessian and the block submatrix, respectively. r is defined in (7).
Optimizer
GD
Adam with
Adam with
3 (r- maxie[L] Ki) ((x)
0 and B2 (3) 31 = 0 and B2 1 (4)
Complexity
X
How to obtain a tighter complexity bound of Adam? It valid ask whether the complexity upper bound in Theorem 1 can be tightened, e.g., improve the factor of r. We point out it would be difficult if there is no extra structure on H,. A key technical step to bound the condition
number of the preconditioned matrix ((D) (D'
Intuitively, diagonal preconditioner
of Hi is powerful when Hi itself has near-diagonal structure, e.g., pure diagonal, tridiagonal or diagonal dominant [30]. Unfortunately, it is unclear whether these structures hold in Transformers. Without any assumption on Hi, we find that the diagonal preconditioner of Do Adam could increase the condition number. For instance, when using standard Gaussian initialization. in case 3, we find
(Dadam,)) Hi) equals 7.09k1, 18.9 18.98k2. 18.76kg for the 3 blocks, respectively (all averaged over 1000 trials). would be interesting to explore there are special structures of Hj in Transformers such that Adam preconditioner can reduce Kr, rather than increase it. We leave it a future direction.
More discussions on the theoretical advantage of Adam. Although Adam preconditioner might not always reduce the "local" condition number Ki, the coefficient in the complexity is now independent of the "global" condition number K. As argued above, such changes in coefficient could lead to considcrable improvement over GD. Such improvement in complexity is attributed to the block diagonal structure in Hessian as well as its hetcrogeneous blockwise spectrum. To our knowledge, such improvement is not shown in the existing literature. One possible reason is that: for the optimization community, it is very rare to analyze (near-) block- diak--da Hessian structure since typical problems do not have such structure. For instance, in the classical non-linear programming dataset [48], all problems have non-block-die block--iiidi Hessian. We suggest a different perspective to characterize modern optimization problems. We believe our perspective is new because it is built upon multiple non-trivial findings.
In summary, our theory indicates that: for problems with block heterogeneity, the single-learning rate methods like GD can largely lag behind coordinate-wist learning rate methods like Adam.
5 Conclusion
In this work, we explore why SGD largely underperforms Adam on Transformers. we establish a phenomenon called block heterogeneity in Hessian and link it to the performance gap between Adam and SGD. We numerically verify our claim on various Transformers, CNNs, MLPs, and quadratic problems. Initial theory is also provided to support the claim.
10

---

Acknowledgements
Yushun Zhang would like to thank Yinyu Ye, Wentao Ding, Guiyu Hong. Yingru Li, and Bohan Wang for the valuable discussions. The work of Ruoyu Sun was supported by NSFC (No. 12326608); Hetao Shenzhen- Hong Kong Science and Technology Innovation Cooperation Zone Project (No. HZQSWS- CCCYB-2024016) together with Tian Ding: University Development Fund UDF0100149) the Chinese University of Hong Kong, Shenzhen; Guangdong Provincial Key Laboratory of Mathematical Foundations for Artificial Intelligence (2023B1 21 2212010001) The work of Z..Q Luo supported by the Guangdong Major Project of Basic and Applied Basic Research N0.200338033000000 the Guangdong Provincial Key Laboratory of Big Data Computing, and the National Key Research and Development Project under grant 0022YFA1003900
Broader Impacts
We explore why SGD performs worse than Adam for training Transformers. Our work can help the community better understand large Al model training. However, it would be a potential threat the Al models are used for illegal usage.
References
[1] J. Achiam, S. Adler, S. Agarwal, L. Ahmad, Akkaya, L. Aleman, Almeida, J. Altenschmidt, S. Altman, S. Anadkat, et al. Gpt-4 technical report. arXiv preprint arXiv:2303.08774 2023.
[2] R. P. Adams, J. Pennington, M. J. Johnson, J. Smith, Y. Ovadia, B. Patton, and J. Saunderson. Estimating the spectral density of large implicit matrices. arXiv preprint arXiv: 1802.0345, 2018.
[3] K. Ahn, X. Cheng, M. Song, C. Yun, A. Jadbabaie, and S. Sra. Linear attention is (maybe) all you need (to understand transformer optimization). arXiv preprint arXiv:2310.01082, 2023.
[4] H. Avron and S. Toledo. Randomized algorithms for estimating the trace of implicit symmetric positive semi-definite matrix. Journal of the ACM (JACM), 58(2):1-34, 2011.
[5] T. Bachlechner, B. Majumder, H. Mao, G. Cottrell, and J. McAuley, Rezero all you need: Fast convergence at large depth. In Uncertainty in Artificial Intelligence, pages 1352-1361. PMLR, 2021.
[6] 7. Bai and G. H. Golub. Bounds for the trace of the inverse and the determinant of symmetric positive definite matrices. Annals of Numerical Mathematics, 4:29-38, 1996.
[7] Z. Bai, G. Fahey, and G. Golub. Some large-scale matrix computation problems. Journal of Computational and Applied Mathematics, 74(1-2):71-89, 1996.
[8] A. Beck and L.. Tetruashvili. On the convergence of block coordinate descent type methods. SIAM journal on 0p 2013.
[9] J. Bernstein, Y.-X. Wang, K. Avizzadenesheli, and A. Anandkumar. signsgd: Compressed optimisation for non-convex problems. In International Conference on Machine Learning. pages 560-569, PMLR, 2018.
[10] S. Bock and M. WeiB. Non-convergence and limit cycles in the adam optimizer. In Artificial Neural Networks and Machine Learning-ICANN 2019: Deep Learning: 28th International Conference on Artificial Neural Networks, Munich, Germany, September 17-19, 2019, Pro- ceedings, Part Il 28, pages 232-243. Springer, 2019.
[11] C. Brezinski. A direct proof of the christoffel-daroou identity and its equivalence to the recurrence relationship. Journal of Computational and Applied Mathematics, 32(1-2):17-25. 1990.
11

---

[12] P. Chaudhari, A. Choromanska, S. Soatto, Y. LeCun, C. Baldassi, C. Borgs, J. Chayes, L. Sagun, and R. Zecchina. Egt: Biasing gradient descent into wide valleys. Journal of Statistical Mechanics: Theory and Experiment, 2019(12):124016 2019.
[13] C. Chen, L. Shen, F. Zou, and W. Liu. Towards practical adam: Non-convexity, convergence theory, and mini-batch acceleration. The Journal of Machine Learning Research, 23(1): 10411-10457, 2022.
[14] J. Chen, F. Kunstner, and M. Schmidt. Heavy-tailed noise does not explain the gap between sgd and adam on transformers. In 13th Annual Workshop on Optimization for Machine Learning. 2021.
[15] M. X. Chen, O. Firat, A. Bapna, M. Johnson, W. Macherey, G. Foster, L. Jones, N. Parmar, M. Schuster, Z. Chen, et al. The best of both worlds: Combining recent advances in neural machine translation. arXiv preprint arXiv: 09849, 2018.
[16] X. Chen, S. Liu, R. Sun, and M. Hong. On the convergence of a class of adam-type algorithms for non-convex optimization. 7th International Conference on Learning Representations, ICLR 2019, 2019.
[17] A. Chowdhery, S. Narang, J. Devlin, M. Bosma, G. Mishra, A. Roberts, P. Barham, H. W. Chung, C. Sutton, Gehrmann, et al. Palm: Scaling language modeling with pathways. Journal of Machine Learning Research, 24(240): 2023.
[18] R. Collobert. Large scale machine learning, Technical report, Université de Paris VI, 2004.
[19] M. Crawshaw, M. Liu, F. Orabona, W. Zhang, and 7. Zhuang. Robustness to unbounded smoothness of generalized signsgd. Advances in Neural Information Processing Systems, 35: 9955-9968, 2022.
[20] J. K. Cullum and R. A. Willoughby. Lanczos algorithms for large symmetric eigenvalue computations: Vol. I: Theory. SIAM, 2002.
[21] A. B. Da Silva and M. Gazeau. A general system of differential equations to model first-order adaptive algorithms. The Journal of Machine Learning Research, 21(1):5022-5113, 2020.
[22] T. Dao, D. Fu, S. Ermon, A. Rudra, and C. Ré. Flashattention: Fast and memory-efficient exact attention with io-awareness. Advances in Neural Information Processing Systems, 35: 16344-16359, 2022.
[23] A. Défossez, I.. Bottou, F. Bach, and N. Usunier. A simple convergence proof of adam and adagrad. Transactions on Machine Learning Research, 2022.
[24] M. Dehghani, J. Djolonga, B. Mustafa, P. Padlewski, J. Heek, J. Gilmer, A. P. Steiner, M. Caron, R. Geirhos, I. Alabdulmohsin, et al. Scaling vision transformers to 22 billion parameters. In International Conference on Machine Leurning, pages 7480-7512, PMLR, 2023.
[25] J. Devlin, M.-W. Chang, K. Lee, and K. Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint
2018.
[26] Y. Dong, J.-B. Cordonnier, and A. Loukas. Attention is not all you need: Pure attention loses rank doubly exponentially with depth. In International Conference on Machine Learning. pages 2793-2803. PMLR, 2021.
[27] A. Dosovitskiy, L. Beyer, A. Kolesnikov, D. Weissenborn, X. Zhai, T. Unterthiner, M. De- hghann, M. Minderer, G. Heigold, S. Gelly, et al. An image is worth 16x16 words: Transformers for image recognition at scale. arXiv preprint arXiv:2010.119299 2020.
[28] J. Duchi, E. Hazan, and Y. Singer. Adaptive subgradient methods for online learning and stochastic optimization. Journal of machine learning research, 12(7), 2011.
[29] J. F. Epperson. An introduction to numerical methods and analysis. 2013.
[30] G. F. Forsythe and F. G. Straus. On best conditioned matrices. Proceedings of the American Mathematical Society, 6(3):340-345, 1955.
12

---

Mathematics of [31] S. Gadat and I. Gavra. Asymptotic study of stochastic adaptive algorithms in non-convex landscape. The Journal of Machine Learning Research, 23(1):1035--10410 2022.
[32] B. Ghorbani, S. Krishnan, and Xiao. An investigation into neural net optimization via hessian eigenvalue density. In International Conference Machine Learning, pages 2232-2241. PMLR, 2019.
[33] G. Goh. Why momentum really works. Distill, 2017. doi: http://distill.puts pub/2011ppp 2017/momentum
URL
[34] Golub and G. Meurant. Matrices, moments and quadrature with applications, volume 30.
Princeton University Press, 2009.
[35] G. H. Golub and 7., Strakoš. Estimates in quadratic formulas. Numerical Algorithms, 8:
241-268, 1994.
[36] G. H. Golub and J. H. Welsch. Calculation of gauss quadrature rules. computation, 23(106):221-238 1969.
[37] B. Goujaud, D. Scieur, A. Dieuleveut, A. B. Taylor, and Pedregosa. Super-acceleration with cyclical step-sizes. In International Conference on Artificial Intelligence and Statistics, pages 3028-3000 2128-3065. PMLR, 2022.
[38] A. Gu and T. Dao. Mamba: Linear-time sequence modeling with selective state spaces. arXiv preprint arXiv:2312.0075: 2023.
[39] Gur-Ari, preprint (xXiv:8812.0477. 2018.
Roberts, and Dyer. Gradient descent happens in a tiny subspace. arXiv
[40] K. He, X. Zhang, S. Ren, and J. Sun. Deep residual learning for image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 770- 778, 2016.
[41] S. Huang, F. Perez, J. Ba, and M. Volkovs. Improving transformer optimization through better initialization. In International Conference on Machine Learning, pages 4475-4483, PMLR, 2020.
[42] M. F. Hutchinson. A stochastic estimator of the trace of the influence matrix for laplacian smoothing splines. Communications in 1076, 1989.
and Computation, 18(3):0599-
[43] K. Jiang, D. Malik, and Y. Li. How does adaptive optimization impact local neural network geometry? Advances in Neural Information Processing Systems, 36, 2023.
[44] J. Kaplan, McCandlish, T. Henighan, T. Brown, B. Chess, R. Child, S. Gray, A. Rad- ford, J. Wu, and D. Amodei. Scaling laws for neural language models. arXiv preprint arXiv:2001.0s8 1.0300,.0036,..... 2020.
[45] D. P. Kingma and J. Ba. Adam: A method for stochastic optimization. arXiv preprint 1rXiv:1412.6998 2014.
[46] F. Kunstner, Chen, J. W. Lavington, and M. Schmidt. Noise is not the main factor behind the gap between sgd and adam on transformers, but sign descent might be. arXiv preprint arXiv:2304. 13960, 2023.
[47] C. Lanczos. An iteration method for the solution the eigenvalue problem of linear differential and integral operators. 1950.
[48] G. Lavezzi, K. Guye, and M. Ciarcià. Nonlinear programming solvers for unconstrained and constrained optimization problems: a benchmark analysis. arXiv preprint arXiv:2204.05297, 2022.
[49] L.eCun, I.. Bottou, Y. Bengio, and Haffner. Gradient-based learning applied to document recognition. Proceedings of the IEEE, 86(11):2278- -2224, 1998.
13

---

[50] Y. LeCun, L. Bottou, G. B. OII, and K.-R. Müller. Efficient backprop. In Neural networks: Tricks of the trade, pages 9-50. Springer, 2002.
[51] H. Li, A. Rakhlin, and A. Jadbabaie. Convergence of adam under relaxed assumptions. Advances in Neural Information Processing Systems, 36, 2023.
[52] Z. Liao and M. W. Mahoney. Hessian eigenspectra of more realistic nonlinear models. Advances in Neural Information Processing Systems, 34:20104-20117, 2021.
[53] L. Lin, Y. Saad, and C. Yang. Approximating spectral densities of large matrices. SIAM review, 58(1):34 65, 2016.
[54] H. Liu, Z. Li, D. Hall, P. Liang, and T. Ma. Sophia: A scalable stochastic second-order optimizer for language model pre-training. arXiv preprint arXiv:2305.14342, 2023.
[55] L. Liu, H. Jiang, P. He, W. Chen, X. Liu, J. Gao, and J. Han. On the variance of the adaptive learning rate and beyond. arxiv 2019, arXiv preprint 4rXiv:1908.03265 2019.
[56] L. Liu, X. Liu, J. Gao, W. Chen, and J. Han. Understanding the difficulty of training trans- formers. arXiv preprint arXiv:2004.08249 2020.
[57] I. Loshchilov and F. Hutter. Decoupled weight decay regularization. arXiv preprint arXiv:1711.0510/ 2017.
[58] L. Luo, Y. Xiong, Y. Liu, and X. Sun. Adaptive gradient methods with dynamic bound of learning rate. In International Conference on Learning Representations, 2018.
[59] K. Lv, Y. Yang, T. Liu, Q. Gao, Q. Guo, and X. Qiu. Full parameter fine-tuning for large language models with limited resources. arXiv preprint arXiv:2306.09782 2023.
[60] S. Malladi, T. Gao. E. Nichani, A. Damian, J. D. Lee, D. Chen, and S. Arora. Fine-tuning language models with just forward passes. Advances in Neural Information Processing Systems, 36:53038-53011 2023.
[61] W. Merrill, V. Ramanujan, Y. Goldberg, R. Schwartz, and N. Smith. Effects of parameter norm growth during transformer training: Inductive bias from gradient descent. arXiv preprint arXiv:2010.09699 2020.
[62] I. Molybog, P. Albert, M. Chen, Z. Vetito, D. Esiobu, N. Goyal, P. S. Koura, S. Narang, A. Poulton, R. Silva, et al. A theory on adam instability in large-scale machine learning. arXiv 2023.
preprint arXiv:2304.098:
[63] Y. Nesterov. Introductory lectures on convex optimization: A basic course, volume 87. Springer Science & Business Media, 2013.
[64] T. Q. Nguyen and J. Salazar. Transformers without tears: Improving the normalization of self-attention. arXiv preprint 1910059910 2019.
[65] L. Noci, S. Anagnostidis, L. Biggio. A. Orvieto, S. P. Singh, and A. Lucchi. Signal propagation in transformers: Theoretical perspectives and the role of rank collapse. Advances in Neural Information Processing Systems, 55:27198-27211, 2022.
[66] Y. Pan and Y. Li. Toward understanding why adam converges faster than sgd for transformers. arxiv preprint arXiv:2306. 00204, 2023.
[67] V. Papyan. The full spectrum of deepnet hessians at scale: Dynamics with sgd training and sample size. arXiv preprint arXiv:/8/1.07062. 2018.
[68] V. Papyan. Measurements of three-level hicrarchical structure in the outliers in the spectrum of deepnet hessians. arXiv preprint arXiv:1901.08244 2019.
[69] V. Papyan. Traces of class/cross-class structure pervade deep learning spectra. The Journal of Machine Learning Research, 21()):01177-0226, 2020.
14

---

[70] B. A. Pearlmutter. Fast exact multiplication by the hessian. Neural computation, 6(1):147-160, 1994.
[71] Radford, Wu, R. Child, D. Luan, D. Amodei, Sutskever, et Language models are unsupervised multitask learners, OpenAl blog, 188)9, 2019.
[72] S. Reddi, S. Kale, and S. Kumar. On the convergence of adam and beyond. In International Conference on Learning Representations, 2018.
[73] Y. Saad. Numerical methods for large eigenvalue problems: revised edition. SIAM, 2011.
[74] L. Sagun, L. Bottou, and Y. LeCun. Eigenvalues of the hessian deep learning: Singularity and beyond. arXiv preprint 1 2016.
[75] L. Sagun, U. Evci, V. U. Guney, Y. Dauphin, and L. Bottou. Empirical analysis of the hessian over-parametrized neural networks. arXiv preprint 100.v:1706.0445, 2017.
[76] A. R. Sankar, Y. Khasbage, R. Vigneswaran, and V. N. Balasubramanian. A deeper look at the hessian eigenspectrum of deep neural networks and its applications regularization. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 35, pages 9481-9488, 2021.
[77] N. Shi, D. Li, M. Hong, and R. Sun. Rmsprop converges with proper hyper-parameter. In International Conference on Learning Representations, 2020.
[78] K. Simonyan and A. Zisserman. Very deep convolutional networks for large-scale image recognition. arXiv preprint arXiv:1409.155: 2014.
[79] Sun. Optimization for deep learning: theory and algorithms. arXiv preprint arXiv:/912.088 2019.
[80] R. Sun and Y. Ye. Worst-case complexity of cyclic coordinate descent: O (n" 2) (n gap with randomized version. Mathematical Programming, 185:487-520, 2021.
gap
[81] I. O. Tolstikhin, Houlsby, A. Kolesnikov, L. Beyer, X. Zhai, T. Unterthiner, Yung, A. Steiner, D. Keysers, J. Uszkoreit, et al. Mip-mixer: An all-mlp architecture for vision. Advances in neural information processing systems, 34:24261-24277 2021.
[82] S. Ubaru, J. Chen, and Saad. Fast estimation of tr(f(a)) via stochastic lanczos quadrature. SIAM Journal on Matrix Analysis and Applications, 38 28(4):1075-109)) 2017.
[83] A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, L. Kaiser, and I. Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017.
[84] B. Wang, Y. Zhangg H. Zhang, Q Meng. 7.M Ma, T.-Y, Liu, and W. Chen. Provable adaptivity in adam. arXiv preprint arXiv:2208. 09900, 2022.
[85] B. Wang, J. Fu, H. Zhang, N. Zheng, and W. Chen. Closing the gap between the upper bound and lower bound of adam's 's iteration complexity. Advances Neural Information Processing Systems, 36, 2023.
[86] B. Wang, H. Zhang, Z. Ma, and W. Chen. Convergence of adagrad for non-convex objectives: Simple proofs and relaxed assumptions. In The Thirty Sixth Annual Conference on Learning Theory, pages 161-190. PMLR, 2023.
[87] H. Wang, S. Ma, L. Dong, S. Huang, D. Zhang, and F. Wei. Deepnet: Scaling transformers to 1,000 layers. arXiv preprint arXiv:2203.00 005502033.0055. 2022.
[88] Q. Wang, B. Li, T. Xiao, Zhu, C. Li, D. F. Wong, and I.. S. Chao. Learning deep transformer models for machine translation. arXiv preprint arXiv: /906.0/87, 2019.
[89] Wikipedia. Gaussian quadrature Wikipedia, the free encyclopedia, 2023. URL
https:/ www.eewwwwwiaaa..eeeeee aeaediii.org//// www.///ginnioeooooooooeoooooooee htwp//nx....ee wttpiidnewpppppttteeeee www.ttpeeetaasssiieiieoemmmmmeee httpssssssa.uuurrruurrelllllaa.. 1191539517. [Online; accessed 20-January-2024].
15

---

[90] M. Wortsman, P. J. Liu, L. Xiao, K. Everett, A. Alemi, B. Adlam, J. D. Co-Reyes, I. Gur, A. Kumar, R. Novak, al. Small-scale proxies for large-scale transformer training instabilities. arXiv preprint arXiv:2309. 14322, 2023.
[91] Y. Wu, X. Zhu, C. Wu, A. Wang, and R. Ge. Dissecting hessian: Understanding common structure of hessian in neural networks. arXiv preprint arXiv:2010.0426. 2020.
[92] T. Xiao, M. Singh, E. Mintun, T. Darrell, P. Dollár, and R. Girshick. Early convolutions help transformers see better. Advances in neural information processing systems, 34:30392-30400, 2021.
[93] Z. Xie, X. Wang, H. Zhang, I. Sato, and M. Sugiyama. Adaptive inertia: Disentangling the effects of adaptive learning rate and momentum. In International conference on machine learning, pages 11430-24455 PMLR, 2022.
[94] R. Xiong, Y. Yang. D. He, K. Zheng, S. Zheng, C. Xing, H. Zhang, Y. Lan, L. Wang, and T. Liu. On layer normalization in the transformer architecture. In International Conference on Machine Learning, pages 10524-10533, PMLR, 2020.
[95] A. Yang, B. Xiao, B. Wang, B. Zhang, C. Bian, C. Yin, C. D. Pan, D. Wang, D. Yan, et al. Baichuan 2: Open large-scale language models. arXiv preprint arXiv:2309. 0ot::230900033 10305, 2023.
[96] G. Yang, E. J. Hu, I. Babuschkin, S. Sidor, X. Liu, Farhi, N. Ryder, J. Pachocki, W. Chen, and Gao. Tensor programs v: Tuning large neural networks via zero-shot hyperparameter transfer. arXiv preprint arXiv:2203.034000 2022.
[97] Z. Yao, A. Gholami, Q. Lei, K. Keutzer, and M. W. Mahoney. Hessian-based analysis of large batch training and robustness to adversaries. Advances in Neural Information Processing Systems, 31, 2018.
[98] Z. Yao, A. Gholami, K. Keutzer, and M. W. Mahoney. Pyhessian: Neural networks through the lens of the hessian. In 2020 IEEE international conference on big data (Big data), pages 581-590. IEEE, 2020.
[99] M. Zaheer, S. Reddi, D. Sachan, S. Kale, and S. Kumar. Adaptive methods for nonconvex optimization. Advances in neural information processing systems, 31, 2018.
[100] A. Zeng, X. Liu, Z. Du, Z. Wang, H. Lai, M. Ding, Z. Yang, Y. Xu, W. Zheng, X. Xia, et al. Glm-1306s An open bilingual pre-trained model. arXiv preprint arXiv:22 2022.
[101] S. Zhai, T. Likhomanenko, E. Littwin, D. Busbridge, J. Ramapuram. Y. Zhang, J. Gu, and
J. M. Susskind. Stabilizing transformer training by preventing attention entropy collapse. In International Conference on Machine Learning, pages 40770-40803. PMLR, 2023.
[102] B. Zhang, I. Titov, and R. Sennrich. Improving deep transformer with depth-scaled initialization and merged attention. arXiv preprint arXiv: /365, 2019.
[103] G. Zhang, L. Li. Z. Nado, Martens, S. Sachdeva, G. Dahl, C. Shallue, and Grosse. Which algorithmic choices matter at which batch sizes? insights from noisy quadratic model. Advances in neural information processing systems, 32, 2019.
[104] J. Zhang, T. He, S. Sra, and A. Jadbabaie. Why gradient clipping accelerates training: A theoretical justification for adaptivity. arXiv preprint arXiv:1905. 1881, 2019.
[105] J. Zhang, S. P. Karimireddy, A. Veit, S. Kim, S. Reddi, S. Kumar, and S. Sra. Why are adaptive methods good for attention models? Advances in Neural Information Processing Systems, 33: 15383 15393, 2020.
[106] S. Zhang, S. Roller, N. Goyal, M. Artetxe, M. Chen, S. Chen, C. Dewan, M. Diab. X. Li. x. V. Lin, et al. Opt: Open pre-trained transformer language models. arXiv preprint arXiv:2205.0006
2022.
[107] Y. Zhang, C. Chen, N. Shi, R. Sun, and Z.. Luo. Adam can converge without any modification update rules. Advances in Neural Information Processing Systems, 35:28386-28399 2022.
16

---

[108] D. Zhou, J. Chen, Y. Cao, Y. Tang, Z. Yang, and Q. Gu. On the convergence of adaptive gradient methods for nonconvex optimization. arXiv preprint arXiv:/808.0560 05671, 2018.
[109] F. Zou, L. Shen, Z. Jie, W. Zhang, and W. Liu. A sufficient condition for convergences adam and rmsprop. In Proceedings of the IEEE/CVF Conference on computer vision and pattern recognition, pages 11127-11135, 2019.
17

---

A Related Works
On the unsatisfactory performance of SGD on Transformers There is an active line of works that explores why SGD performs significantly worse than Adam on Transformers. One representative hypothesis is that SGD cannot handle the heavy-tailed stochastic noise in language tasks [105]. However, Chen et al. [14], Kunstner et al. [46] reported that the gup between Adam and SGD maintains even in the full-batch case with no stochasticity, so there might be other reasons. Further, SGD performs worse than Adam on Vision Transformers on ImageNet (See Figure 10, Also see [92] for more evidence), so the data modality (e.g., language or vision tasks) might not be as crucial as the architecture. [104] showed that NLP tasks have "unbounded smoothness" issue and SGD with gradient clipping performs better than SGD in this case. Although clipping is an effective trick, we still observe a huge gap between clipped SGD and Adam 5, so there might be other reasons. Different from these works, we find SGD underperforms Adam because it uses one single learning rate for all blocks, which cannot handle the Hessian heterogeneity among blocks.
Understanding of Adam. There was once a long-standing debate on the possible divergence
of Adam [72]. The convergence for the unmodified versions is later established in [77, 107] for RMSprop and Adam. More convergence analyses of general adaptive gradient methods are listed later in this section. We here focus on the literature that explores the benefit of Adam. Xie et al. [93] show that Adam can help avoid saddle points, which is an orthogonal direction to this work. Wang et al. [84], Crawshaw et al. [19], Li et al. [51] show that Adam and its variant outperform SGD under relaxed smoothness conditions, based on the intuition that Adam can adaptively change its learning rate along iteration (over time). We pointed out that the theory is not complete: even for quadratic functions where the smoothness is fixed, SGD sometimes performs largely worse than Adam (Figure 8). This indicates that the benefit of Adam is not merely due to its ability to adaptively change the learning rate (over time), and there are other reasons for Adam's success. We show that an important benefit of Adam is its ability to handle the heterogeneity across blocks (over space).
Recent works [9, 91, 46, 54, 3] build a relation between Adam and the sign-based methods. Wu et al. [91] further showed that sign-hased methods can be effective when the Hessian is diagonal and satisfies several other properties. However, as put by the authors, it seems "unclear to what extent these properties hold for real problems". Pan and Li [66] numerically found that the Adam can reduce the directional sharpness along trajectories, while its relation to fast convergence remains mysterious. A recent work [43] point out that Adam biases the trajectories towards regions where Hessian has "uniform diagonal entries" while SGD cannot. The distribution of Hessian diagonal entries is also investigated in [54]. The theory in [43] implies that Adam is faster when the Hessian is diagonal. However, as argued above, it is unclear whether the diagonal Hessian structure commonly holds in real problems. In fact, we find the Hessian is closer to a block- boock-diaoonal (instead of pure diagonal) structure on some small Transformers. In these cases, blockwise eigenvalues carry more information than diagonal entries, providing extra details such as the location of eigenvalues. We find that these extra details are important for distinguishing Adam and SGD.
Hessian Spectrum Analysis. There are several important attempts to explore the Hessian spectrum of MLPs and CNNs. Early works [74, 75, 12] found that the Hessian spectra of MLPs and CNNs consist of a "bulk" together with few "outliers". Papyan [69], Wu et al. [91], Liao and Mahoney [52] further characterized the bulks and outliers in theory. Papyan [67, 68) numerically built the relation between these "outliers" and the Gauss-Newton matrix. Sankar et al. [76] numerically explored the relation between Hessian of CNNs and most CNN layers contribute similarly to the overall loss surface. We find that this result is restricted to CNNs and does not hold on Transformers due to the heterogeneity. Gur- Ari et al. [39] showed that for MLPs and CNNs, gradient descent converges to a small subspace spanned by a few top eigenvectors of the Hessian. Yao et al. [97]. Zhang et al. [103] explored the relation between the Hessian spectrum of CNNS and some training phenomena such as the effect of batch sizes. Ghorbani et al. [32], Yao et al. [98] focused explaining the effectiveness of techniques such as BatchNorm. Note that all these works are restricted to MLPs and CNNs, while we study the Hessian of Transformers (in addition to CNNs and MLPs) well as its impacts on different optimizers.
matrix in each layer. They further found that
For all NLP tasks, clipping is performed immediately after backpropagation. So in Figure 10, SGD in NLP tasks essentially refers to clipped SGD.
18

---

On the difficulties of Transformer training. Transformers are known to be difficult to train. Researchers have attributed the training difficulties to various phenomena different components of Transformers, including: the logits divergence or the rank degeneracy in the outputs of attention layers [26, 65, 90, 101, 24, 17]; the growth of parameter norm in attention layers [61]; over-reliance on residue branches [56]; and some negative impact of layer norm [15, 102, 41J. These phenomena have a strong correlation with gradient vanishing or explosion Transformers [102, 56, 41, 94, 65, 87, 90, 62], which leads to training difficulties.
Several solutions have been proposed. Liu al. [56] numerically observed that adaptive gradient methods can (partly) overcome gradient vanishing by giving "consistent update magnitude", while it seems unclear how consistent update magnitude would help optimization principle. Researchers further develop training tricks such as warmup learning rate [55, 94], temperature scaling [65], better initialization [102, 41, 87, 5, 96], and variants of Layer Norm [64, 88, 94, 87, 24). Recent researchers also suggest using 7-loss regularization [17, 95] and tuning hyperparameters of Adam [107, 90]. All these tricks can help mitigate gradient explosion or vanishing. Nevertheless, training large-scale Transformers remains challenging [106, 100, 90, 62, 17J. Different from all aforementioned works, we investigate the training difficulties of Transformers through the eigenvalues Hessian. We establish strong correlation between "the blockwise Hessian spectra of Transformers" and "why SGD largely underperforms Adam on Transformers". We realize that our attempt is just a first step towards understanding Transformer training, and believe there is rich information hidden in Hessian and we leave more fine- mee-granned analysis as future works.
Convergence analysis of general adaptive gradient methods There is extensive convergence analysis for adaptive gradient methods. For instance, researchers study the convergence of AMSGrad [72, 108], RMSprop [99], AdaFom [16], AdaBound [58], and Adam with sverate-depender hyper- parameters [109, 13, 31]. The convergence of Adam is also explored in [23, 85]. There is also an active line of theoretical research the convergence of AdaGrad [28], we recommend [86] for more detailed introduction. In this work, we do not focus on the convergence analysis. Rather, we explore the quantitative difference between the loss landscape of CNNs and Transformers and how it impact the behaviors of SGD and Adam.
19

---

B More Results and Discussions
Performance comparison of AdamW W and SGD on different Architectures. Here, we show the performance comparison of AdamW and SGD on different models. All the vision models are trained on ImageNet. Language models are trained on different English corpus. We grid-search the learning rates for SGD and Adam under the same budget and report the best result for each optimizer. See Appendix D.1 for more implementation details.
(a) ResNet18
(b) VGG16
Figure 9: Performance of AdamW and SGD on CNNs including ResNet18 and VGGI6. SGD and Adam perform similarly on these tasks.
(a) VIT
(b) BERT
(c) Gano
(d) GPT2
Figure 10: Performance of AdamW and SGD on Transformers including ViT, BERT, GPT2-nano,
and GPT2. SGD performs significantly worse than Adam on these tasks.
More results for SGD under careful tuning. For ViT-base, GPT2-nano, nanno, and BERT, we further present the performance of SGD under learning rate grid search. For ViT-base training on ImageNet, we report the results after 30 epochs (or equivalently, about 30k iterations). We cannot afford further training ViT-base due to the limited hardware resources (a complete run of 90 epochs would take weeks for each curve). As shown Figure 11, SGD consistently performs worse than Adam on ViT-base, GPT2, and BERT. The best results for SGD are picked out and presented in Figure 10.
ViT-base
(b) GPT2-nano nano
BERT
Figure 11: Performance of SGD under careful tuning. On baseb GPT2-nano, and BERT, carefully tune the learning rate SGD and show all the results here. For all these Transformer tasks, SGD is still significantly worse than AdamW even after careful tuning.
Note that we are not the first ones to report that "SGD performs worse than Adam on VIT". influential work [92] also reports that SGD is worse than Adam vanilla VIT. The authors report that "SGD yields significantly worse results than AdamW (on ViT)", and "ViT often fails to converge with SGD (their Figure 3). These results align with our findings in Figure 11.
20

---

Detailed training curves of Figure 5. In Figure 12, we present the detailed training curves for the man-made MLP Figure We find that SGD performs worse as heterogeneity grows, while Adam still performs well.
(a) Training curves of Adam
(b) Training curves of SGD
The training curves of SGD and Adam on MNIST with 4-layer MLPs under different degrees of block heterogeneity c. We observe that SGD performs worse as heterogeneity grows. while Adam remains unaffected.
Figure 12:
The evolution of Hessian heterogeneity along training. For ViT-base, we further investigate the evolution of block heterogeneity of Hessian along the training. As shown in Figure 13, find that heterogeneity attenuates along the training. We further take the checkpoint of ViT-base at 50% training step and switch AdamW SGD, observe that now SGD performs better than training from scratch as in Figure 10 (a). SGD performs better here because there is less heterogeneity when initializing the middle of training.
(a) 25% training step (b) 50% training step (c) 100% training step (d) SGD v.s. Adam when resumed 50% step
Figure 13: For ViT-base, we plot the evolution of heterogeneity of Hessian along training. We find that heterogeneity attenuates along training. (d): we take the checkpoint of ViT-base at 50% training step and switch AdamW to SGD, we find that now SGD performs better than training from scratch as in Figure 10 (a).
Simplifed SLQ for calculating J in Section 3.4. We note that JS" in Table is rather expensive compute due the computational overhead of SLQQ it requires comparable time one training run. Fortunately, we find the original SLQ is redundant for measuring hessian heterogeneity. We propose the following simple tweaks significantly reduce the computation time, while still effectively detecting the Hessian heterogeneity. We call it simplified SLQ.
Change the hyperparemters of SLQ, including:
We change num, 10 to num, 1. SLQ, num,, decides the number of random Gaussian vectors to approximate the expected quadrature. It is reasonable to reduce num, because in high dimensional space, random vectors tend concentrate around their mean, so one random sample can already be informative enough.
We change the Lanzcos step 100 to m 10. The reduction on Lanzcos step will have coarse estimation on the middle eigenvalue, but won't affect much the heterogeneity, which is more dependent the extreme eigenvalues.
Randomly sample a subset of blocks and reduce batch size for estimating the spectrum. We uniformly sample 50% blocks and choose batch size 32 (previously batch size = 1024).
21

---

We report the result and runtime in Table 3. As result, the simplified SLQ can obtain the same message the original SLQ: JS of ResNet is about 100x smaller than that of BERT. Further, the simplified SI.Q is highly efficient to compute. With this simplified SIQQ, we believe our method can efficiently scale larger models. The result is tested on a single V100.
Table 3: JS computed by simplified SLQ. We find that the simplified SLQ can obtain the same message as the original SLQQ J5 of ResNet is about 100x smaller than that of BERT. Further, the simplified SLQ is efficient compute.
Model JSO
Time for JSO
Time for Training
Time for JSO / Time for training
BERT
98.8344
2
4h
0.0014
ResNet18
0.3569
65%
87.5 h
0.0002
Blockwise spectra for quadratic models in Section 4.1. We here visualize the blockwise spectrum for the quadratic models in Case and Case These spectra are collected from GPT2 and ResNet18, 88, respectively.
(a) Spectrum of Hi
(b) Spectrum of Ha
(c) Spectrum of Hs
(d) Spectrum of H4
Figure 14: Histogram of eigenvalues of each block in Case 1 (the heterogeneous case). The cigenvalues in the four blocks are sampled from the spectrum of the embedding layer; 3rd Query, 3rd Value, 3rd MLP (fc layer) in GPT2, respectively. All the eigenvalues scaled such that: the objective function is strong convex; the condition number Hessian equals 5000; their relative ranges preserved; and the block heterogeneity is preserved.
shifted and proportionally
(a) Spectrum of Hi
(b) Spectrum of Ha
(c) Spectrum of Hs
(d) Spectrum of H4
Figure 15: Histogram cigenvalues of each block in Case 2 (the homogeneous case). The eigenvalues in the four blocks sampled from the spectrum of 1st to 4th convolution layers in ResNet18. respectively. All the eigenvalues are shifted and proportionally scaled such that: the objective function is strong convex; the condition number of Hessian equals 5000; their relative ranges preserved; and the block homogeneity preserved.
Does our theory imply that Adam under homogeneity is slower than Adam under heterogeneity? During the review process, we find that readers might conclude that "Theorem implies Adam under homogeneity has worse complexity than Adam under heterogeneity". We would like to clarify that this conclusion is not correct, and there is conclusion on "whether Adam under homogeneity is faster or slower than Adam under heterogeneity We explain as follows.
Our theoretical result states that Adam has complexity O (max; K1). If our result implies the above conclusion, one needs the following argument: when changing heterogeneity to homogeneity, max, K
We would like thank the anonymous Reviewer pTmmy for raising this insightful question.
22

---

increases, and thus Adam is slower. However, "changing heterogeneity to homogeneity" does not necessarily mean "max, Kı increases". Actually, max; K2 can change an arbitrary way (can increase, decrease, keep the same) when changing the heterogeneity. We provide three examples below.
We will use Adam (homo) to denote the convergence rate of Adam homogeneous Hessian, similarly for Adam (hetero).
Example I: Adam (homo) is same as Adam (hetero).
Case 1-1 (homogeneous): eigenvalues are (1,2), (1,2)
Case (heterogeneous): eigenvalues are (1,2), (11,12)
Since max; Kj are the same for both Case 1-1 and 1-2, Adam (homo) is the same as Adam (hetero).
Example 2: Adam (homo) is faster than Adam (hetero).
Case 2-1 (homogencous): eigenvalues are (1,1.5), (1,1.5)
Case 2-2 (heterogeneous); eigenvalues (1,2), (11,12)
Since Case 2-1 has smaller max; Ką than Case 2-2, Adam (homo) is faster than Adam (hetero).
Example Adam (homo) is slower than Adam (hetero).
Case 3-1 (homogeneous): cigenvalues are (1,12), [1,12]
Case 3-2 (heterogeneous): eigenvalues are [1,2), (11,12)
Since Case 3-1 has larger max; Ky than Case 3-2, Adam (homo) is slower than Adam (hetero).
To sum up, there is no conclusion "whether Adam under homogeneity is faster or slower than Adam under heterogeneity ". Either case can happen.
One possible source of confusion may come from the numerical examples (Case and 4 in Section Section 1). If comparing two figures, Adam (homo) in Case is slower than Adam (hetero) in Case 4. But as argued above, this is just one example, and it does not show Adam (homo) is always slower than Adam (hetero).
23

---

C More Preliminaries
C.1 Preliminaries Optimizers
Here we provide a detailed description of the optimizers mentioned the full script. We consider the minimizing (ww) Σ Li(w), where the number of minibatches, Li(w) is the loss of i-th minibatch and w E R the neural network parameters. We denote the gradient of the training loss w.r.t neural network parameters as V ([w) E IR We use V (w:))) E I to denote the i-th minibatch counterparts. We use w' denote the variable at the t-th step. In Algorithm and o, division and square-root are elementwise operations. In the line and 8 of Algorithm 2, (5,)) and (32)* indicates the 1-th power 31, 32. In the PyTorch default setting, (3,,82, 6) (0.9, 0.999, 1e-8) for Adam and B1 = 0.9 for SGD.
Algorithm 1 Stochastic Gradient Descent with Momentum (SGD)
1: Initialize w and choose S 31 < and !/o >
2: for ( 1=
do
3: Uniformly sample 7' from the index set {1.2, ... ,n}
4: m' = 8mm' VL, (x')
5: x x=+1 x' mm
6: end for
Algorithm 2 AdamW
Initialize 20, m 00 = 0, < 1.0 s B2 < e 0, 170
0, and weight decay
coefficient X
2: fort=
do
3: Uniformly sample 7' from the index set f1,2, ..,,.,.,,,
4: w'+1 = 17" Aw
5: m' Bim' + (1 31) C,, (w*)
6: of Bzut +(1 32) V.) (w') (C+((w")
7:
m' m I-(3,)
8:
6 I-(023
9: 111 w0t1
m
1]t Vor4c
10: end for
Algorithm 3 Adam with no bias correction
1: Initialize x0, m = VL, (w²),, v VL, (L,) & < <1,0g & <1,c>0 0,
17" >>
2: 1rr=1-1 xc do
3: Uniformly sample 7 from the index set 21,
4: m² = 8im + (1 V (w²)
5: of = Bav (1 32) V2) (wo" (Ly(
6: w++1 wi+ Nt Vor4e
7: end for
C.2 Preliminaries the Stochastic Lanczos Quadrature Method
Additional notations. Given a real symmetric matrix H E Rdxdd we denote tr(H) as its trace and Q" AQ as its spectral decomposition, where Q I41..... qa] A diag (Ai, Xa) and A > X2 > Ad. We denote the condition number of /l as A Aq- We define matrix function as (H) == Q" fAAQQ where f(A) = diag (f (Ai)....f fAa)) E Rdxd We use to denote the set of positive integers. We use || ||2 denote the Euclidean norm.
Approximation of the Hessian spectrum can be formulated as trace estimation problem, as introduced in [53, 82]. First, the spectrum (cigenvalue density) of Hessian H can written as:

---

o(t) - ( (t- Xc), where A are the eigenvalues of H and ő is the Dirac 8-function. Then, we replace the delta functions by a Gaussian blurring function: <(t) g(t) == 2 f (As), where f(A) = oV exp
20t (-12>"
By definition of matrix function, is easy to see that s(() - (r(((H) As such, spectrum approximation could be formulated as trace estimation problem, i.e., estimating tr(f(H)) where H E R is real symmetric matrix.
Trace estimation problems could be solved efficiently by the Stochastic Lanczos Quadrature Method (SLQ) [35]. For the ease of readers, we re-organize and summarize the existing literature (135, 82, 32)) and provide a detailed description of SLQ in our context. SLQ consists of the following steps.
Step We Approximate the trace of matrix function as ((r(J(H)) = E( ((H)u) s
Σ. ((H),, where = 2/|lul12 and u is a Rademacher random vector (each entry of independently takes #1 with probability 1/2). This step called Huchinson's 's estimation [42].
Note that we can also replace the Rademacher random vector u by a unit Gaussian vector (i.e., u (0, Iaxa)) and the unbiasedness still holds [4]. In our implementation, we sample u ~ (0,, Taxa) because there an efficient built-in PyTorch function for generating Gaussian vectors.
SLQ estimates of f(H)c for i E [ne] and then take the average. To understand SLQ, we only need understand how it estimates each individual quadratic form. To simplify the notation regarding i, from now on, we will discuss how to estimate v² f(H)e, where = w/||ul]2 and is a unit Gaussian vector.
Step 2-1. We rewrite f(H)) as a Riemann-Stieltjes integral [34]:
d of f(A)U (F ((a)= = S/ f(A)d((A),
(8)
where is measure on (R, IB) defined follows (p(A) denotes the measure of set {æ, S X)):
A
((A)=
(oT
qe)
2
A S < AR+1
d -=1
(oT
q:)
A Z A
Step 2-2. Unfortunately, this integral is difficult to compute. This is because the measure are related to the eigen-pairs of H, which are unknown. seems unclear how to directly integrate over an unknown measure. As such, we further approximate this integral by computationally friendly quantity, such as:
A1 J
m
f(A)du(A) Σ Gf(a)))
(10)
We hope to design {(c,, 23) with a reasonable number of m such that the estimation error is small. Fortunately, the Gaussian Quadrature method provides generic design principle of (y);,)))? } [34, 29]. It is proved that: when f(A) is not "too complicated" (e.g. f(X) is a polynomial), then there exists ((c,,,,)) 2,)] which gives high quality estimation of integral (8). The required number of m related to "how complicated the f(() is". Such {(c,, ), are called the Gaussian Quadrature rules. Cj and zj are called the "weights" and the "nodes" of the Gaussian Quadrature rules. representative theorem is as follows: when (A) is a polynomial with degree < 2m, then the Gaussian Quadrature rules give the exact approximation of integral (8).
Theorem 2. (Rewrited based on [89]] Suppose we have a sequence of orthogonal polynomials {p<())) w.r.t. measure p. that is: f pa(x))mm(a)) (x))p(((()) ôm,no where 8m,n if m and 8m,, 0, otherwise. Assume fxx is a polynomial with degree < 2m, then there exists {lo 1,)1m s.t. S f(A)dp(A) = Em off (æz). The equality holds when xj the roots
25

---

of Pm(()) and Oj Si ls du. Such choice of [1),= are called the Gaussian Quadrature rules.
Theorem 2 shows the existence of good {(e,, 1n) and their general form. In fact, it is also shown that Gaussian Quadrature is optimal: other () can achieve zero approximation error for
higher degree polynomials f()) [34]. However, often difficult to find these quadrature rules [36].
There are at least three questions in sequel:
1) how to find the orthogonal polynomials (x)) w.r.t. unknown measure μ!?
2) how to efficiently find the roots of pm (2)), which gives the nodes x,²
3) how to efficiently caleulate the weights Gj S 1i - du?
We first answer question 2) and 3) and leave question 1) for later discussion.
Now suppose that we have found the orthogonal polynomials {px(xx) (x)) w.r.t Recall that any orthogonal polynomial has the following "three-term" recursion [34]:
Pk+1 (I)) (x- .a++1) pk(x) PkPk-1
0, 1....,
where p-1(x) = O, (o = 1,04+1 =
(ape.Pa) and 3g [pe.Pe)
(Pt-1-PC-1) Define Pm (x) =
(pc.pe)
(pa(2), (r).... Pm (x))²
E RRm, we can rewrite the recursion formula in matrix form (given rt):
xPm= = Jm Pm + 3mPm Pm (2) em where e the last column of identity matrix I Im and Jm is called
Jacobi matrix of order
an
T
1
02
2
Jm
33
R m
It turns out that Jm can help us find the Gaussian Quadrature rules {(cj. 21)]" and thus provide answers for question 2) and 3). This is shown in the following theorem.
Theorem 3. (34] For the Gaussian Quadrature, { are the eigenvalues of Jm and { are the squares of the first elements of the normalized eigenvectors of Jm-
The proof of Theorem is based on Prristoffel-Darbbone relation [11]. Now, the remaining question is: how to find the Jacobian matrix Jm of sequence of orthogonal polynomials w.r.t. an unknown measure !u? Note that we no longer need to answer question 1) if Jm is found, since Jm is sufficient for us find the Gaussian quadrature rules. However, seems impossible find J,m if no information of μ is provided. The good news is: when the is specified as in (9), there exists an efficient way to find Jm.
Step 3. When is specified as in (9), Jm can be exactly found in m steps using the Lanczos algorithm [47], as shown in Algorithm 4. This method takes a real symmetric matrix as input and returns a tridiagonal matrix. It was originally proposed to solve eigenvalue problems. Later, researchers found a deep connection between the Lanczos algorithm and orthogonal polynomials, which further connects this method to the Gaussian quadrature. The method (of finding the Gaussian quadrature by the Lanezos algorithm) is called the Lanezos quadrature [35, 6, 34]. An extremely elegant but highly nontrivial result is as follows:
Theorem 4. [34] Given a real symmetric matrix H E Rddd and an arbitrary vector v E IR with unit Euclidean norm, we define the measure I as (9) based on this and Then m steps of the Lanzcos algorithm return the Jacobian matrix Jm of orthogonal polynomials w.r.t. to μ.
After Jm is found by the Lanezos algorithm, we perform spectral decomposition of Jm E mmmx x to get its eigen-pairs. Using Theorem 3, we successfully get the Gaussian quadrature rules and thus we can approximate the quadratic form 8² [H)).. u. By averaging over different random vectors v we can then approximate (H)). This concludes the derivation of SLQ for the trace estimation problem.
26

---

The full procedure of SLQ is shown in Algorithm 5. We note that SLQ is efficient in theory. Ubaru al. [82] show that SLQ converges faster than any other polynomial expansion method for spectrum estimation (e.g., Chebyshev methods used in [2]). See [82, Theorem 4.1] for formal statement.
We remark that there are least four versions of the Lanczos algorithm in Step 3. Here, adopt the version in Algorithm 4 since it is known to be the most numerically stable version [20, 73, 89]. Throughout this work, we choose f (-) as the Gaussian blurring function f(A)
0/22 exp ((-A)² for spectrum approximation. We plot the spectrum by sweeping from the minimal node to the maximal node Gaussian Quadrature rules.
Algorithm 4 The Lanczos Algorithm
1: Input matrix-vector product Hvj E Rd, where is a real symmetric matrix and ei is an arbitrary vector with Euclidean norm 1. Choose E N
2: Initialization: Let w/ = HU1, a1 = (wi)"er 21. W1 = w/ 0111
3: for j
4: Let B3 ||wy-1 (also Euclidean norm)
5: ;3 * 0, then let By wjj1/
else pick as by an arbitrary vector with Euclidean norm 1 that is orthogonal to all of
U1..... 11-1
6:
Let w/j = Auj
7:
Let a; (w') 2 Uj
8:
Let W; w/j Ug B1
9:
end for
Let V be the matrix with columns .1..... Um
Or
82
0
82
022
B3
11:
T
Bg
03
8m-1
3m-1
am-1
Pm
0
Bm
O'm
12: Return
Algorithm 5 The Stochastic Lanezos Quadrature Method
1: Choose num,, E N. Sample num, i.id. v; from normalized Rademacher distribution,
[num,]
2: for = num, do
3: Run m steps of the Lanczos Algorithm 4 with input He,, returns E Rmxx
4: Compute eigenvalue decomposition T QAQ"
5: Compute the nodes 3; (Aa) and weights c; (Q
6: Return (t) = cf (x)
7: end for
8: Return Σn / (lstt,²²)
27

---

D More Eperimental Details
D.1 Implementation Details on SLQ and Training Configurations
Implementation and Running Time Analysis. We provide simple PyTorch implementation of SI.Q. The only query SLQ makes to the neural network is the Hessian vector product, which is attained using the auto-differentiation framework [70]. To assure the accuracy of the Lanczos algorithm, we remove all the randomness in the forward and backward passes, including: data shuffling order, data augmentation, and dropout, etc.. Since Flash Attention [22] does not support the calculation of Hessian vecttr vector product, we implement all attention blocks in the naive way. For the calculation of the blockwise Hessian spectrum {²L((u), we sample in (O, Idxx;) and set v1 u/ / 2|2 E R4: Then we run Algorithm 5 by taking () and as inputs. We choose the hyperparameters as m = 100 and ni, in all experiments. a is tuned based on visual effects. These hyperparameters are reported to reach highly accurate estimation with error < 10 [32].
We now briefly discuss the computational cost of SLQ. The major computational expense of SLQ is the repeated Hessian-vector product operations in Lanczos algorithm in Step 3. Recall V = (w) d, so each Hessian-vector product operation requires (i) calculating V V2 (w)d; (ii) repeating (i) on all data. We point out that (i) can be computed efficiently and precisely with just two backpropagation passes [70]. The major computational bottleneck lies in (ii) due to the large n. Our largest-scale experiment for Hesian spectrum is GPT2 (125M) Openwebtext, where the number of tokens n = Billon. To calculate V FL((w) d on all these 9B tokens, requires about 9 GPU days on eight A100-80GB GPUs. Since SLQ requires at least 1, 1,000 times query of V²L(ww) (w)d, a complete run of SL.Q would take at least 9,000 days on eight A A100-80GB GPUs, which is unaffordable. In this work, we use the largest possible batch size (with gradient accumulation tricks) to approximate V² [(w) under the constraints of GPU bandwidth and time limit. More detailed setup of SLQ are shown as follows.
" Ei= V
R8sNet18 (18M) and VGG16 (138M) on ImageNet. We use the code base of PyTorrh Examples 7, We use batch size- = 1024. For the calculation of the blockwise Hessian spectra, we apply SLQ to all parameter blocks except for the BatchNorm layers. In total, it takes about 3 days on one V100 GPU to estimate all the blockwise Hessian spectra and the full Hessian spectrum.
ViT-base (86M) on ImageNet. We use the code base of PyTorch Image Models 8, We use batch size 1024. Due the large number of parameters, Hessian spectra for all parameter blocks. Instead, we apply SLQ to: the embedding layer; the output layer; the 1-st, 6-th, 12-th attention blocks; and the 1-st, 6-th, 12-th MLP blocks (note that the 12-th attention and MLP blocks are the final ones). In total, it takes about 3 days on one V100 GPU to estimate all the blockwise Hessian spectra and the full Hessian spectrum.
are not able calculate the blockwise
(ERT(40M) on Cornell Movie-Dialogs Corpus. We use the code base from the blog 9 We use batch size = 327, 680 tokens. For the calculation of the blockwise Hessian spectra, we apply SLQ all parameter blocks except for the LayerNorm layers. In total, it takes about 12 hours on one V100 GPU to estimate all the blockwise Hessian spectra and the full Hessian spectrum.
GPT2-nano (11M) on Shakespeare. We use the code base of NanoGPT 10 We use batch size 163, 840 tokens. For the calculation of the blockwise Hessian spectra, we apply SLQ to all parameter blocks with even indices, except for the LayerNorm layers. In total, it takes about 12 hours on one V100 GPU to estimate all the blockwise Hessian spectra and the full Hessian spectrum.
GPT2 (125M) on Openwebtext "! We use the code base of NanoGPT. We use batch size 245, 760 tokens. Due to the large number of parameters, we are not able to calculate the blockwise Hessian spectra for all parameter blocks. Instead, we apply SLQ to: the embedding layer; the output layer; the 1-st, 4-th, 8-th, 12-th attention blocks; and the 1-st, 4-th, 8-th, 12-th MI.P blocks (note that the 12-th attention and MLP blocks are the final ones). In total, it takes
wwwperchhiiaaeemmm-mee modees
bww.od/bbbbbooooooeetoeeeeeeeee. geeeee-gideetoooo guide-to-building-- seildinng-ert-modell -frome---ooo--oo aaron--racch--- www.acrraeeeeeee0e 366562228891
"I httt:p///uggggggee.o../attttt/tp
28

---

about 7 days on one A100 GPU to estimate all the blockwise Hessian spectra and the full Hessian spectrum.
Training configuration. In all cases, we train all the models under the default configurations in the above codebase. We grid-search the learning rates for SGD and Adam under the same budget and report the best result for each optimizer. We use the cosine-decay learning rate schedule for vision tasks. For SFT task, we use nanoGPT codebase. We first pre-train GPT2 on OpenwebText for 25B tokens and then fine-tune it on a subset of Alpaca Eval 12.
D.2 Ablation Studies of SLQ on Small Tranformer
On a small GPT model with 75k parameters, we compare (1) the true Hessian spectrum and (2) the estimated Hessian spectrum curve by SLQ method. As shown in Figure 16, we find that the SLQ method can produce accurate estimation.
embd embd_layer layer
(b) attn.c__aeee c_attn
attn.c.proj
(d) mlp __ppc_proo
Figure 16: (a, b, c, d): The comparison of (1) the true Hessian spectrum: and (2) the estimated Hessian spectrum curve by SLQ method. The experiments are conducted on a small GPT model with 75k parameters. We find that the SL.Q method can produce accurate estimation.
D.3 Implementation Details on Figure
We employ a synthetic dataset designed for binary classification, with 100 data points generated through the process outlined below. Our model is 1-hidden-layer neural network, featuring an input size of 64 and a layer width of 8, utilizing the hyperbolic tangent (Tanh) as the activation function. We train this model over 1000 iterations using the Adam optimizer with learning rate set to 1 10-4 achieving the classification accuracy of 100%.
def generate_
classes, input_dim): Generate synthetic data for specified dimensions
x = +O y= = []
for i in range (n_classes):
center np. randon rand (input_dim) class_samples
Random class center
np. random handom.randee randn (n_ (tw.samplespper.eemom class, input_ dim
0.5 center # Add some noise
x. ap (class_saaae sclaseesamples)
y extend [[i] n. s_aamles.per..co class)
10
11
up. vstack (X)
Combine all class samples
mp array (y) return X, y
Convert labels to NumPy array
D.4 Implementation Details on the MLP experiments in Figure 5
We train a 4-layer MLP on MNIST. We use batch size 128 and width 300, 128, and 64 for the hidden layers. We use ReLU activation. We change the degree of heterogeneity by scaling the output of each layer with constant e N. We scale from 15. For each c, we train SGD and Adam with default hyperparameters by grid-searching the learning rate from le-4 to le-1 and report the best test accuracy after epoch.
12
29

---

E Proofs
E.1 Proof of Proof of Proposition
L
Let H 0
where L > 0. We choose the initial point w" (wf. wch)
(Vn/L, v L//)). By the update rule of GD, we have
[(w(t1)
L (w" V ((ww))
5(00 - "HHW'' H(w' -1
(w1) - 17b66h + (wx)) - nulu
11 1- L 1 I - 1-pp p L
μ11 nLI + L|1 mp
(11)
To proceed, we discuss the following cases:
When 7) $ I/L, since |1 rL and |1 nul' are monotonically decreasing, the optimal solution is L.LL.
When 1/14, since [1 nLI1 and |1 null monotonically increasing. the optimal solution is = 1//1
When 1/L. < 17 - 1/μ, (11) can be written as 9. (n) = - 1 np)) Take the first-order
and the second-order derivative the g. we can obtain g/(n) = tLu(M 1) - tu 7744)=11
and (n) f(t - ()LL"((nL~- 1)t-2 ttt-t(tli 1)122111111111 - nu). Since g (n) 0 for all n E L1 1p]. function g is convex. By solving the equation that g((n) = 0, we can obtain = Σ+ッ is
solution for all t. Plugging this result into (11) and rearranging the terms, we conclude the proof of Proposition 1.
E.2 Proof of Theorem 1
first show that Ct22 and Cr,1 are non-zero w.p.1. under the random initialization Assumption 1. We define set S; (w;; hi w = 0} where h4 E Rd is the i-th row of H. Since positive definite, there is at least one non-zero entry in h2 E [d]. As such, S; is a (d 1)-dimensional subspace of
dd and thus has zero Lebesgue measure in Rd, Since wo follows continuous distribution, we have Pr (fw"; hi "h = 07) 0, for [d]. Then have
Pr (V((w") has at least one zero entry)
= Pr (Hw" has at least one zero entry)
(12)
Pr (U (wo h2 w 0})
(13)
<
Pr (fwo h "0 0})
(14)
0.
(15)
Therefore, V((w)) is elementwise non-zero w.p.1.., so Ct.1 and Ct,2 are non-zero for all / E [L], w.p.1.. In the following analysis, We will assume Ci and Ct,2 are non-zero.
Without loss of generality, assume h 0. This is because minimizing L(w) 1 w Hw RTT m is equivalent to minimizing L(w) 1 (w "= (w w²) where wt = 1/. By a linear transformation
w w², Adam for minimizing (w w H (w w") starting from w is equivalent to Adam for minimizing } Hz starting from 20 w w². Thus we can assume "" 0, or equivalently, h 0. The update rule of Adam becomes
T
30

---

= w² Adam) ) Hu"
where Da = diag( (L((g)))()))) V)( = diag( (Hw)) We denote d, = (dam)) th and thus we have w = - Adam d' and w = w¹ d,. These relations also hold for each block by changing the notation to H1 wf. D Ademm and df, etc.. Following the framework in [80], we try to bound the error yet to be optimized (a.k.a., cost-to-go) and the per-step improvement, respectively. The ratio of these terms characterizes the of convergence. We now express both terms using df. For the cost-to-go term for the 1-th block, we have
[C(w")] - [C*] = (uff)² Hawi = 27/
1 (dt) D'dam, Hi D
D' A
(16)
For the per-step improvement, we have
[C(w")] ))- - [C(w++1) (w]) Howf
2 (w( (w Hiu
1
2( (wt)² Hmw!
1 (wi Hy)hi -di)
(d)) Haws 1 (di) Hidj
2
(d))²² (20a Adom.1 H1 di.
D%
(17)
To proceed, we denote H (Du Adam) H 11 and we denote its eigenvalues as ò i2 2 ... da.
Similarly, we denote H, (Do Adam.) and its eigenvalues dµι d1.2 2,. Let
" = mimg Ct we have
[[(w")], (C") [C(w*)]; [C(w(+)]]
# (d))z Do Adaam, H Do D'
D Hi) dli (df)²
-1
D D'ddam. - H1)
D dl Do Adam,
(18)
172
IAC
C2
(19)
(min C2,1 Mr,1
maXie(t) C222
Ni,
(20)
minge[L] Cf,1
where (*) is due by Assumption 1. D Adam..
C1,AA,14
3 Dddam.
Hi >
1f
C Cu.AAA1---u > \r where & and are matrix inequalities.
By rearranging both
sides of (20), we have [[("" ]] [C*] < 1
([C(w")] - [C*)). Summing
ma Cf min CP.
up both sides over 1 E [L] and we conclude the proof.
L(w - L- Σ ([((w(+))]
((C(m)) [c
maxie C222 minre([]] Cf
31

---

Σ ([C(w")) I1 - [c7]]
max
[e
maxiel C
mince(M] C1
mi
l=1
1
max (E(4)
(C(w') -L").
max]c[4] 0222 minrc([] Cf,3
ni
32

---

NeurIPS Paper Checklist
1. Claims
Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Question:
Answer: [Yes]
The main results in the experimental sections match the claims in the abstract and introduction.
Justification:
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contribu- tions made in the paper and important assumptions and limitations. A No NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected generalize to other settings.
It is fine include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.
Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification:
The limitation is discussed in Section 5.
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are discussed in the paper.
The authors are encouraged create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to vi- olations of these assumptions (e-g., independence assumptions, noiseless settings, model selll-pecification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might be used by review- grounds for rejection, worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
Justification: The full set of assumptions are Section 4.2 and the complete proof E. Guidelines:
33

---

The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross-referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
Experimental Result Reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main experi- mental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided not)?
Answer: [Yes]
Justification: In Appendix D, we carefully describe all the needed information to reproduce the experimental results.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution a specific model and empirical evaluation, may be necessary to either make it possible for others replicate the model with the same dataset, provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case a large language model), releasing of a model checkpoint, other means that appropriate the research performed.
While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.
(b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the particular way they provide for reproducibility. In the case closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, described in supplemental material?
Answer: [Yes]
Justification:
We provide code in the supplementary materials.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
34

---

Please see the NeurIPS code and data submission guidelines quides/CodeSubbe osoliccy for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" is acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https:// nips. cc/cc///// cccppblic//eccccccccccccccccccct hwwpul//guides/////eeeeee ceniddes/CoeSuumiiis. wwwisssssssssiiieeee cnsionnooccy) for more details.
The authors should provide instructions data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only a subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs data and code is permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: The experimental details are described in Appendix D.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, or as supplemental material. Experiment Statistical Significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about statistical significance of the experiments?
Answer: [No]
Justification: We did not provide the error bar since the experiments on large Transformers are too expensive to repeat for multiple runs.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
It should be clear whether the error bar is the standard deviation or the standard error of the mean.
It is OK to report I-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are of range (e.g. negative error rates). If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
35

---

8. Experiments Compute Resources
Question: For each experiment, does the paper provide sufficient information the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experi- ments?
Answer: [Yes]
Justification:
The computation resource is described in Section D.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experi- mental runs as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary failed experiments that didn't make into the paper).
Code Of Ethies
Does the research conducted in the paper conform, in every respect, with the NeurlPS Code of Ethics httts::///eurrr/e httppppc.ccpppppppppeeeeeeeeeeee httpppbbll//thiicccutteeeeee ines?
Question:
Answer: [Yes]
Justification:
We conform with the NeurIPS code of Ethics.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No. they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., there a special consideration due to laws or regulations in their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification: Positive and negative social impacts are discussed in Section 5.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious unintended uses (e.g.. disinformation, generating fake profiles, surveillance). fairness considerations (e.g., deploy- ment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications. let alone deployments. However, if there is a direct path any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).
36

---

11. Safeguards
Does the paper describe safeguards that have been put place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: Our paper poses no such risks.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines or restrictions access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards challenging, and many papers not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Question: Are the creators or original owners of assets (e.g., code, data, models). used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification: All assets are properly cited.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset used and, if possible, include a URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thcode. com/ datasess has curated licenses for some datasets. Their licensing guide can help determine the license of dataset. For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information is not available online, the authors are encouraged to reach out to the asset's creators.
13. New Assets
Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Question:
Answer: [NA]
Justification:
We do not release new assets.
Guidelines:
The answer NA means that the paper does not release assets. Researchers should communicate the details of the dataset/code/moddl part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create anonymized URL or include anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, applicable, as well as details about compensation (if any)?
37

---

Answer: [NA]
Justification: Our paper does involve crowdsourcing nor research with human subjects. Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with
human subjects.
Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should included in the main paper.
According the NeurIPS Code Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals Equivalent for Research with Human Sub- jects
Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Question:
Answer: [NA]
Justification: Our paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with
human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors adhere to the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
38


Reviews:
[{"review_id": "7Jk2uOhHnx", "reviewer": "Reviewer_5ncB", "summary": "In this work the authors assert that the Adam optimizer works well on Transformers while Stochastic Gradient Descent (SGD) does not, and they attempt to explain this phenomenon by inspecting the spectrum of Transformers (i.e., the eigenvalues of the model's Hessian matrix) and other models, such as convolutional networks, where SGD is competitive with Adam.  They show that the spectrum of Transformers and convolutional networks are similar, providing little insight.  The main innovation of the authors is to investigate the spectrum of individual components/layers of these networks - the so-called block spectrum - and the main contribution of the work is to show that the block spectrum is empirically correlated with the ineffectiveness of SGD relative to Adam.", "questions": "See \"Weaknesses\" for my questions.", "limitations": "yes.", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1) The paper is generally clear and well-written - this is appreciated\n2) The idea of inspecting the block-wise spectrum is potentially innovative\n3) The empirical results presented by the authors are fairly convincing.", "weaknesses": "1) The premise of this paper is that SGD performs worse for attention-based models, and it is unclear to me if this is true.  I say this for two reasons:\na) The authors cite [45,56] on Line 16 after making this claim, presumably to support the claim.  However, these are quite old publications and neither of these papers specifically make this claim.  Are there other publications that make these claims? \nb) The authors also provide their own empirical results (e.g., Fig. 3) to support this claim, however, these results seem highly insufficient to me to support this claim.  How did the authors choose/optimize the learning rate for SGD, and/or the learning rate schedule? \n\nIf SGD does not generally perform worse than attention models on Transformers, then it seems to me that the work has little scientific value, since the differences in the block spectrum are just inconsequential differences, or at least that their implications are different than suggested by the authors.  e.g., perhaps the authors could argue that the heterogeneous block spectra implies that SGD needs a lot more tuning to work well, for example.    I'd ask the authors to please address this point, and I will consider raising my rating.       \n\n2) Novelty of using block spectrum (potentially).  To me this is the main value of this work.  I'm not as familiar with this topic, and it is unclear to me how non-obvious the use of a block spectrum would be for analysis, or how surprising the results of using this approach are: perhaps these findings would be somewhat obvious to researchers familiar with this topic?    I will defer to other reviewers for this assessment.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks for the careful evaluation of our paper. We respectfully provide our response as follows.\n**1-1. Are there other publications that make these claims (SGD worse than Adam on attention-based model)?**\n\nThanks for raising this comment. Yes, it is widely reported that SGD largely underperforms Adam on attention-based models, **even after careful tuning**. Here are some references:\n\n1. [1] provides empirical evidence that SGD is worse than Adam on BERT, but not on SGD. On Bert, the authors comment that \"Although hyperparameters for SGD are finetuned, a large performance gap is still observed\" (their Figure 1)\n\n2. [2] provides empirical evidence that SGD is worse than Adam on ViT. The authors reported that \" SGD yields significantly worse results than AdamW (on ViT)\", \"ViT often fails to converge with SGD \" ( their Figure 3)\n\n3. [3] carefully tuned SGD on ViT and find that SGD cannot work (their Figure 4). They reported that \" the SGD algorithm did not perform well for this task”\"  even after careful tuning: “ The optimizer hyperparameters were set according to the values obtained during hyperparameter tuning. ”\n\n4. [4] carefully tuned SGD on modern GPT models and find that SGD cannot work (their Figure 1).\n\n5. Besides, main-stream large language models use Adam, including GPT3 [5], Llama series [6], Gemini series [7], etc. This serves as additional evidence that SGD is worse than Adam for Transformer training. \n\n\n**1-2. How did the authors choose/optimize the learning rate for SGD, and/or the learning rate schedule?**\n\nFor the learning rate schedule, we use the same default cosine decay with warmup for both SGD and Adam. As for the learning rate, we grid search the learning rate of SGD and report the best performance. We grid-search the learning rate as follows.\n\n- For BERT: we use lr  = 1e-4 for Adam.  For SGD, we grid search lr  over [1e-4, 1e-3, 1e-2, 1e-1]\n\n- For ViT: we use lr  = 1e-4 for Adam.  For SGD, we grid search lr  over [1e-4, 5e-4, 1e-3, 5e-3, 1e-2, 5e-2, 1e-1]\n\n- For GPT2-nano: we use lr = 1e-2 for Adam, For SGD, we grid search lr  over [1e-5, 5e-5, 1e-4, 5e-4, 1e-3, 5e-3, 1e-2, 5e-2, 1e-1, 5e-1, 1]\n\nWe also visualize the performance of SGD under different learning rates.  This can be seen in **Figure 1 in the attached PDF**.  For all these Transformers we investigate, we find  SGD (even after fine-tuning carefully) performs significantly worse than Adam.\n\n**2. Novelty of using block spectrum (potentially).**\n\nThanks for the question. We believe our idea is non-trivial and non-obvious. The novelty of our perspective is also confirmed by other three reviewers, for instance:\n\n- Reviewer PTmy: \"The discovery of block heterogeneous for the reason of the bad performance of SGD on Transformers is interesting and novel.\"\n\n- Review tJ3K: \"The connection between Adam's performance and Hessian block-wise heterogeneity is novel, as are the related empirical observations.\"\n\n- Reviewer A8Vc: \" I find the ideas introduced throughout this work interesting and potentially useful to the community.\"\n\n**Our comment on why blockwise spectrum is non-trivial.** For completeness, we further comment on why the perspective of the blockwise spectrum is new and non-trivial.\n\n**--The natural idea does not work.** To explore \"why Transformers need Adam\", a natural idea is to investigate the full Hessian eigenvalue (spectrum). This is because: by optimization theory, full Hessian eigenvalues largely decide the behavior of gradient methods, and there is an active line of work trying to understand NN training via full Hessian spectrum (e.g., reference [8] here, also see the reference [12,32,39,67,68,69,75,76,97,98,103] in our paper ).  Unfortunately, we find there is no noticeable difference between the full spectrum of CNNs and that of Transformers (Figure 1 in the paper).  As such, the natural idea of full spectrum does not work.\n\n**--Why our perspective is non-trivial.**  Our major conceptual innovation is that we connect the following findings to \"why Transformers need Adam\". \n\n1. Transformers contains various kinds of building blocks, while CNN consists of similar convolutional layers. Based on this, we conjecture that blockwise differences might be crucial. \n\n2. \"Hessian of NNs have near-block-diagonal structure\" (a non-trivial but highly overlooked finding in this field). Based on this, we realize that the blockwise spectrum has rich information and could be used to quantify the blockwise difference. Note that for models with dense Hessian, blockwise spectrum can still be computed, but loses substantial information. \n\nAs such, our blockwise spectrum perspective is novel to the deep learning community because it is based on multiple non-trivial findings. \n\nWe believe this perspective is also new to optimization community: for optimization community,  it is very rare to analyze (near-) block-diagonal Hessian structure since typical problems do not have such structure.  For instance, in the classical non-linear programming dataset [9], all problems have non-block-diagonal Hessian.  We point out a new perspective to characterize modern optimization problems.\n\nAs such, we believe our perspective is new, non-trivial, and potentially useful to a wide range of audiences.\n\n[1] Why are adaptive methods good for attention models? \n\n[2] Early convolutions help transformers see better.\n\n[3] AdaGC: A Novel Adaptive Optimization Algorithm with Gradient Bias Correction\n\n[4] Deconstructing What Makes a Good Optimizer for Language Models\n\n[5] Language models are few-shot learners\n\n[6] Llama 2: Open foundation and fine-tuned chat models\n\n[7] Gemini: A natural language system for spoken-language understanding\n\n[8] An investigation into neural net optimization via hessian eigenvalue density\n\n[9] Nonlinear Programming Solvers for Unconstrained and Constrained Optimization Problems: a Benchmark Analysis", "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": "Y66pPEv7IQ", "reviewer": "Reviewer_pTmy", "summary": "The paper investigates why SGD performs worse than Adam in Transformers from the lens of Hessian. They first find that Transformers are \"heterogeneous\", that is, the Hessian spectrum across parameter blocks varies dramatically. Then they conduct various tasks on Transformers, CNNs, MLPs, and quadratic problems, to verify that block heterogeneity hampers SGD. Finally, they derive some initial theoretical analysis to indicate that SGD fails because it applies one single learning rate for all blocks, which cannot handle the heterogeneity among blocks.", "questions": "See weakness", "limitations": "The authors have addressed the limitations.", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 4, "strengths": "1. The paper is well written, clear to read, and the story is interesting.\n\n2. The discovery of block heterogeneous for the reason of the bad performance of SGD on Transformers is interesting and novel.\n\n3. The experiments are completed to verify the points described in the paper.", "weaknesses": "1. The paper demonstrates that for different models when there the block heterogeneous at initialization increases, SGD becomes much worse than Adam. However, these comparisons are between different models. It is natural to ask for the same Transformers model, when the initialization changes to induce different block heterogenous levels, how would the performance of SGD compared with Adam change?\n\n2. The theoretical results of the convergence of Adam for the quadratic models show that homogeneous spectra have worse complexity than heterogeneous spectra, is this the truth in practice?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We are grateful for the careful review of our paper and the great questions. Please find our respectful reply below.\n\n**1. For the same Transformers model, when the initialization changes to induce different block heterogenous levels, how would the performance of SGD compared with Adam change?**\n\nThanks for the interesting question. Here is our finding: For the same model, when heterogeneity is reduced, the gap between Adam and SGD is also reduced. We provide evidence as below.\n\nThere are typically two approaches to change initialization (please let us know if neither of these aligns with your thoughts).\n\n- Approach 1: Change the variance of random initialization. \n\n- Approach 2: Choose a checkpoint in the middle of training as the initialization.\n\nUnfortunately, we find Approach 1 is not suitable for controlling the heterogeneity: it seems unclear how the variance of random initialization is explicitly related to the heterogeneity in Hessian. As such, we use Approach 2 to change the initialization and track the performance of SGD.  We present two experiments below.  (**Exp 1** is already presented in the paper; **Exp 2** is new.)\n\n**Exp 1:**  On finetuning tasks of GPT2, when we change the initialization to the pre-trained weights, GPT2 (pre-trained) exhibits much less hetereogeneity compared with GPT2 with random initialization. This can be seen in **Figure 4 (f) and Figure 6 (a) in the paper**.  As a result, SGD can reach a similar loss as Adam here. This is shown in **Figure 6 (b)** in the paper.  \n\n**Exp 2:** We realize that **Exp 1** might not fully answer your question: although **Exp 1** does not change the architecture,  it changes the training dataset (from pre-training dataset to SFT dataset), and thus changes the loss function. To rigorously address your concern, we further take the checkpoint of ViT at the 46-th epoch (which is about 50% training steps) and calculate Hessian heterogeneity. We make two findings: (1) For ViT checkpoint at the 46-th epoch, Hessian heterogeneity is largely reduced compared with that of random initialization. (2 ) If we continue to rain it with SGD, we find SGD now can perform much better than training from scratch. This result can be seen in **Figure 2 in the attached PDF**.\n\nThis phenomenon above also provides one interesting practical implication: if we can find a good initialization with less heterogeneity, it is possible to train transformers with SGD. Yet, as discussed above, designing such initialization is non-trivial,  as its explicit relation with heterogeneity is still unclear.  We leave this topic as an intriguing future direction.\n\n**2. The theoretical results of the convergence of Adam for the quadratic models show that homogeneous spectra have worse complexity than heterogeneous spectra, is this the truth in practice?**\n\nThanks for the great question. Let us clarify: our theory does **not** imply that Adam has worse complexity on homogeneous spectra. For simplicity of discussion, let's state your comment as a conjecture.\n\n**Conjecture 1:** Adam on quadratic problems with homogeneous spectra has worse complexity than Adam on heterogeneous spectra. \n\nIf we understand your question correctly  (please tell us if not so), your question states that our theory implies Conjecture 1. We kindly point out that Conjucture 1 is **not correct**. We'll clarify below.\n\n(1) **What we proved and why it does NOT imply Conjecture 1**.\n**Our theoretical result:** Adam has complexity $O(\\max_l \\kappa_l)$. GD has complexity $\\Omega(\\kappa)$\n\nIf our result implies Conjecture 1, one needs the following argument: when changing heterogeneity to homogeneity, $\\max_l \\kappa_l$ increases, and thus Adam is slower. \n\nHowever, \"changing heterogeneity to homogeneity\" does **not** necessarily mean \"$\\max_l \\kappa_l$  increases\". Actually, \"$\\max_l \\kappa_l$\" **can change in an arbitrary way** (can increase, decrease, or keep the same) when changing the heterogeneity. See detailed examples below.\n\n(2) **Detailed Examples on Comparing Homo and Hetero Cases**.\nWe provide three examples below.\n\n**Notation:** We use Adam (homo) to denote the rate of Adam on homogeneous Hessian. Similarly for Adam (hetero)\n\n**Example 1:  Adam (homo) is same as Adam (hetero)**\n\ncase 1-1: homo: eigenvalue {1,2}, {1,2}\n\ncase 1-2:  hetero: eigenvalue {1,2}, {11,12}\n\nSince $\\max_l \\kappa_l$ are the same for both case 1-1 and 1-2, Adam (homo) is same as Adam (hetero)\n\n\n**Example 2: Adam (homo) is faster than Adam (hetero)**\n\ncase 2-1: homo: eigenvalue {1,1.5}, {1,1.5}\n\ncase 2-2:  hetero: eigenvalue {1,2}, {11,12}\n\nSince case 2-1 has smaller  $\\max_l \\kappa_l$ than case 2-2, Adam (homo) is faster than Adam (hetero)\n\n\n\n**Example 3: Adam (homo) is slower than Adam (hetero)**\n\ncase 3-1: homo: eigenvalue {1,11}, {2,12}\n\ncase 3-2:  hetero: eigenvalue {1,2}, {11,12}\n\nSince case 3-1 has larger $\\max_l \\kappa_l$ than case 3-2, Adam (homo) is slower than Adam (hetero)\n\nTo sum up, there is no conclusion on \"whether Adam under homogeneity is faster or slower than Adam under heterogeneity \".  Either case can happen.\n\n\n(3) **possible source of confusion:** We thought (correct us if not so) the confusion may come from the numerical examples (Case 3 and 4 in Section 3.1) in the paper. If comparing two figures, Adam (homo) in Case 3 is slower than Adam (hetero) in Case 4. But as argued above, this is just one example, and it does NOT show Adam (homo) is ALWAYS slower than Adam (hetero). \n\n\nAs a result, all the above three examples can happen, and \"Changing hetero to homo\" does not necessarily mean \"Adam becomes slower\".", "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": "lEljzgg0PV", "reviewer": "Reviewer_tJ3K", "summary": "This paper investigates why Adam outperforms SGD in Transformers by examining Hessian estimations. While the overall Hessian spectrum appears similar in both Transformers and other neural networks, the block-wise Hessian spectrum is heterogeneous across blocks in Transformers and homogeneous in other neural networks. The authors empirically confirm that this heterogeneity hampers SGD's performance through experiments involving Transformers, CNNs, MLPs, and quadratic problems. They attribute Adam's advantage to its coordinate-wise learning rate and provide initial theoretical analysis in a simplified setting.", "questions": "-  Including a more detailed experimental setting will help with reproducibility. The hyperparameters for experiments such as learning rate are not included. Since SGD is sensitive to hyperparameters, it's unclear whether the gap is due to the claimed reason or if hyperparams are insufficiently tuned. For example, the training curve of SGD in Figure 9(a) in ViT-ImageNet is much lower than Adam and that is not the case in [Figure 12, Kunstner et.al. Heavy-Tailed Class Imbalance and Why Adam Outperforms Gradient Descent on Language Models]. Can the author comment on this?\n- Figure 1 shows the full spectrum across different time points to illustrate its evolution during training, yet all other plots in the paper related to the block-wise spectrum are presented at initialization. Is there a reason why the evolution of the block-wise spectrum is not shown?\n- Are there any ablation studies on how well the Hessian spectrum is estimated through the SQL method, perhaps on reasonably small networks?\n- Typo in Figure 7 y-axis label", "limitations": "Limitations are well-addressed", "rating": 6, "confidence": 2, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "- The connection between Adam's performance and Hessian block-wise heterogeneity is novel, as are the related empirical observations.\n- Related work is comprehensively cited in the Appendix.\n- Although the finding itself is not groundbreaking, it suggests potential for future research in theoretical analysis and practical improvements of Adam, particularly regarding its advantage in adaptively updating different blocks.", "weaknesses": "- The experiments support the claim that Adam outperforms SGD when Hessian block-wise heterogeneity structure is present. However, it is unclear whether this correlation implies causation beyond intuitive understanding. In general, the conceptual contribution is not very strong.\n- The definition of \"block\" also seems arbitrarily determined. For example, Fig. 3 shows that all convolution layers have a similar spectrum, while the spectrum differs between Query, Value, and MLP blocks in Transformers. Given that one set is the same type of \"block\" (conv layer) and the other is not, this does not seem surprising. Since how the block is defined entirely determines Hessian heterogeneity versus homogeneity, this observation might be overfitting the transformer architecture and could be difficult to generalize to non-mainstream architectures as claimed. \n- The case study on quadratics helps understand and attempt initial theoretical analysis, but it only partially represents behaviours in larger networks. The noise in SGD does not play a significant role here, and Adam with $\\beta_2 = 0.99$ performs very differently than in larger networks.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks for  the careful review of our work. Here is our respectful response.\n\n**1. It is unclear whether this correlation implies causation beyond intuitive understanding.** \n\nWe provide the following evidence on causation. \n\n**I: Theory:** On quadratic models, our theory suggests that: when heterogeneity happens, the complexity of Adam is better than GD. \n\n**II: Control-variable experiments:**. We design the experiment as follows: \n1) We take two Hessian matrices with an identical set of eigenvalues (cases 3 and 4 in Section 3.1). This ensures the performance of GD is fixed. \n\n2) We re-arrange these eigenvalues in two different ways,  which lead to heterogeneity and homogeneity (case 3 and 4).  \n\n3)  We find Adam is faster when heterogeneity exists,  resulting in a huge gap with GD\n\nThis experiment implies \"heterogeneity\" can cause the \"gap between Adam and SGD.\"\n\n\n**2. Block seems arbitrarily determined.**  \n\nWe kindly point out that block is not arbitrarily determined. Here is our principle: we define \"blocks\" as the irreducible dense principal sub-blocks in Hessian. \n\nSince it requires heavy calculation on Hessian to precisely find blocks by the principle above,  we find a computationally-friendly way: We use the PyTorch partition to determine the blocks.  We find that partition in PyTorch can usually match our principle. For instance, in Figure 2, each block corresponds to the parameters in Q, K, V; and MLP, respectively.  Under this common partition strategy, we observe heterogeneity in Transformers but not in CNNs\n\n**3. This observation might be overfitting the transformer architecture.** \n\nWe kindly point out that the observation is not overfitting to Transformers since we also observed heterogeneity on non-Transformer models such as MLP mixer (a real-world non-attention-based architecture, **Figure 5 in paper**), and SGD also performs worse than Adam here.  \n\n\n**4. The case study on quadratics helps understand and attempt initial theoretical analysis, but it only partially represents behaviors in larger networks.** \n\nWe agree that quadratic models cannot capture all behaviors of real NNs. To cover more behaviors in larger NN,  we need analysis on more fine-grained settings such as noisy quadratic models or 1-layer attention models.  It is an intriguing direction to pursue in the future. \n\n**5. It's unclear whether the gap is due to the claimed reason or if hyperparams are insufficiently tuned.** \n\nWe believe the gap is due to Hessian heterogeneity, instead of improper hyperparam for SGD. For all Transformers,  we grid search the learning rate of SGD and report the best performance.  We grid-search lr of SGD as follows.\n\n- For BERT: we use lr  = 1e-4 for Adam.  For SGD, we search lr  over [1e-4, 1e-3, 1e-2, 1e-1]\n\n- For ViT: we use  lr  = 1e-4 for Adam.  For SGD, we search lr  over [1e-4, 5e-4, 1e-3, 5e-3, 1e-2, 5e-2, 1e-1]\n\n- For GPT2-nano: we use lr = 1e-2 for Adam, For SGD, we search lr  over [1e-5, 5e-5, 1e-4, 5e-4, 1e-3, 5e-3, 1e-2, 5e-2, 1e-1, 5e-1, 1]\n\nWe visualize the performance of SGD under different lr.  This can be seen in **Figure 1 in the attached PDF**. We find that SGD (even after careful tuning) is still worse than Adam, including ViT on ImageNet. \n\n**Additional evidence for ViT:**  We here provide more evidence that \"SGD is worse than Adam on ViT\" from the literature. \n\n- On Cifar-10, [1] carefully tuned SGD on ViT and find that SGD cannot work (Figure 4). They reported that \"the SGD algorithm did not perform well for this task” even after careful tuning: “ The optimizer hyperparam were set according to the values obtained during hyperparam tuning. ”\n\n- On ImageNet, [2] provides empirical evidence that SGD is worse than Adam on vanilla ViT. The authors reported that \" SGD yields significantly worse results than AdamW (on ViT)\", \"ViT often fails to converge with SGD \" (their Figure 3).  \n\nAs for Figure 12 in Kunstner et.al., there are two possible reasons for the mismatch: (1) they used SimpleViT, which is a simplified version of the vanilla ViT as we conducted. It is possible that SimpleViT exhibits less heterogeneity and is more friendly to SGD. (2) perhaps they did not tune Adam carefully enough and Adam can perform better if well-tuned.\n\n[1] AdaGC: A Novel Adaptive Optimization Algorithm with Gradient Bias Correction\n\n[2] Early convolutions help transformers see better\n\n**6. Block-wise spectrum along training.** \n\nFollowing your suggestions, we plot the block-wise spectrum at 25%, 50%, and 100% training steps for GPT2 and ViT.  We find an interesting phenomenon: Hessian heterogeneity tends to reduce along training. This can be seen in **Figure 2 and 3 in the attached PDF.**\n\nWe also take the checkpoint of ViT at the 46th epoch (which is about 50% steps) and continue to train with SGD. We find SGD now can perform much better. We provide a similar explanation here:  SGD now performs better because there is less Hessian heterogeneity in ViT at the 46-th epoch (50\\% training steps). This can be seen in **Figure 2 in the attached PDF.** \n\nThis phenomenon above also provides one interesting practical implication: if we can find a good initialization with less heterogeneity, it is possible to train transformers with SGD. However, designing such initialization is non-trivial, as its explicit relation with heterogeneity is still unclear. We leave this topic as an intriguing future direction.\n\nTo sum up, we make two new findings when investigating the blockwise spectrum along training:  1. Heterogeneity tends to reduce along training. 2. as heterogeneity is reduced, we can switch Adam to SGD. We will include these results in the revised script.\n\n\n**7. Ablation studies on how well the Hessian spectrum is estimated through the SQL method, perhaps on reasonably small networks?**\n\nWe provide an ablation study on a 4-layer Transformer. The results are shown in **Figure 4 in the attached PDF.**  We find SQL provide accurate approximation over the true eigenvalue density.", "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": "kxF96ht3fk", "reviewer": "Reviewer_A8Vc", "summary": "This paper explores the performance gap between SGD and Adam when training transformers. The authors analyze the eigenspectrum of the Hessian in different neural network architectures and attribute the failure of SGD to the different Hessian spectra in transformers compared to CNNs and MLPs. They empirically show that while the full Hessian spectrum is similar at initialization across architectures, transformers exhibit block heterogeneity, meaning different parameter blocks (e.g., key/query layer versus MLP layer) have distinct spectra. Conversely, other architectures such as CNNs and MLPs are block homogeneous. By artificially inducing block heterogeneity in MLPs, they demonstrate through numerical experiments that SGD's performance declines as the block spectrum variation increases, while Adam's performance remains stable. They attribute Adam's success in block heterogeneous cases to its ability to assign different learning rates to different parameters. The authors also propose using the Jensen-Shannon distance of Hessian blocks as a quantitive measure of the performance gap. They also provide a preliminary theoretical analysis comparing the convergence rates of SGD and Adam on a quadratic function with a block-form Hessian, highlighting Adam's dependence on the fine-grained condition number of individual Hessian blocks.", "questions": "1. Can you comment on the computational complexity of calculating the JS distance and the Hessian approximation? Specifically, can you compare the runtime of computing this quantitative measure to the training time of the models in your paper? For significantly larger models, do you expect this approach to scale efficiently and be more practical than training the model?\n\n2. Is the architecture the only factor affecting the Hessian structure? I am curious if the training distribution, particularly label imbalance, also impacts the Hessian structure, since in some previous works previous works, such as [46], the performance gap has been attributed to the heavy-tailed distribution of labels in language data.", "limitations": "Yes", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "This paper tackles a key question: understanding the optimization challenges of transformers with SGD and Adam is crucial for developing more efficient techniques for training large language models.\n\nThe paper is well-organized and easy to read. I find the ideas introduced throughout this work interesting and potentially useful to the community.", "weaknesses": "Quantifying the performance gap between SGD and Adam using the average JS distance between block-wise Hessian spectra is an interesting idea. However, based on the runtime reported in the appendix, the Hessian approximation itself seems computationally expensive, which can raise concerns about how practical it is for large models, especially since the experiments in this paper are only on small models.\n\nAdditionally, the main theoretical analysis is conducted in a pretty simplified setup where the momentum part of Adam is fully ignored ($\\beta_1=0$ in Algorithm 3).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks for the insightful question and careful evaluation of our paper. We provide response as follows.\n\n**1. In theoretical analysis, the momentum part of Adam is fully ignored  ( $\\beta_1=0$ in Algorithm 3).**\n\nThanks for the comment. The current analysis ignores momentum because we want to focus on the benefit of assigning coordinate-wise lr, so we remove the effect of all other components to control variables. Nevertheless,  we agree that choosing $\\beta_1=0$ simplifies the problem. However, exploring the role of momentum is an independent and equally challenging problem. We consider it as an important future direction.\n\n\n\n**2. On the computational complexity of calculating the JS distance and the Hessian approximation; compare runtime to the training time; scale efficiently to larger models**;\n\nThanks for the great question. Currently, our JS distance is indeed rather expensive to compute: it requires comparable to one training run.  Fortunately,  we find the current deployment of SQL is redundant for measuring hessian heterogeneity. We propose the following simple tweaks to significantly reduce the computation time, while still effectively detecting the Hessian heterogeneity. We call it **simplified SQL**:\n\n1. Change the hyperparemters of SQL, including: \n- 1-1: we change num_v = 10 to num_v = 1. In SQL, num_v decides the number of random Gaussian vectors to approximate the expected quadrature. It is reasonable to reduce num_v because in high dimensional space, random vectors tend to concentrate around their mean, so one random sample can already be informative enough.\n- 1-2: we change the Lanzcos step m = 100 to m = 10. The reduction on Lanzcos step will have a coarse estimation on the middle eigenvalue, but won't affect much the heterogeneity, which is more dependent on the extreme eigenvalues\n2. Randomly sample a subset of blocks and reduce batch size for estimating the spectrum. We uniformly sample 50% blocks and choose batch size = 32 (previously batch size = 1024). \n\nWe report the result and runtime in the Table below.  As a result, the simplified SQL can obtain the same message as the original SQL: **JS0 of ResNet is about 100x smaller than that of BERT.**  Further, the simplified SQL is highly efficient to compute. With this simplified SQL, we believe our method can efficiently scale to larger models.\n\n|  Model   |        JS0         | Time for JS0 | Time for Training | Time for JS0 / Time for training |\n| :------: | :----------------: | :----------: | :---------------: | :------------------------------: |\n|   BERT   | 98.83440094176784  |     20s      |        4h         |           0.001388889            |\n| ResNet18 | 0.3568875006145019 |     65s      |       87.5h       |           0.000206349            |\n\nTested on: single V100\n\n**3. Is the architecture the only factor affecting the Hessian structure? I am curious if the training distribution, particularly label imbalance, also impacts the Hessian structure**\n\nWe agree that there exist other factors (such as data distribution and label imbalance) that affect the Hessian structure.  However,  for ViT on ImageNet,  we find SGD (even after careful tuning) is still largely worse than Adam. As such, we believe architecture plays a more crucial rule.", "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:
Dear reviewers and AC:

We attached a PDF with the following four figures. Please check.

**Figure 1:** On ViT, BERT, and GPT2-nano, we carefully grid search the learning rate for SGD and report all the results. We find that on all these tasks, SGD (even after careful tuning) is significantly worse than Adam. 

**Figure 2 and 3:**  For ViT and GPT2, we plot the evolution of heterogeneity of Hessian along training. We find that heterogeneity is reduced along with training. Further, when switching Adam to SGD in the middle of training where the heterogeneity is reduced (e.g., the 46th epoch of ViT training), SGD can perform better. 

**Figure 4:**  We conduct ablation study on small GPT to show that our SQL method can accurately calculate the Hessian spectrum.

Meta-review: Great paper!  Review scores: 6 6 7 5

The paper explores and partially answers the question why SGD is worse than Adam when training transformer models.  The reviews did identify potential weaknesses which the rebuttal of the authors did successfully clear out.  I didn't find the weaknesses pointed out by review with the lowest score (the 5, reviewer 5ncB) very convincing.  Also the review got a detailed rebuttal  as well, but that wasn't acknowledged by the reviewer.

So overall this is a really nice paper that should be accepted!
Decision: Accept (poster)