# Why Transformers Need Adam: A Hessian Perspective

Yushun Zhang  Email<EMAIL_ADDRESS>Our code is available at
https://github.com/zyushun/hessian-spectrum. The Chinese University of Hong
Kong, Shenzhen, China  Shenzhen Research Institute of Big Data  Congliang Chen
Email<EMAIL_ADDRESS>The Chinese University of Hong Kong,
Shenzhen, China  Shenzhen Research Institute of Big Data  Tian Ding Email:
<EMAIL_ADDRESS>Shenzhen Research Institute of Big Data  Ziniu Li Email:
<EMAIL_ADDRESS>The Chinese University of Hong Kong, Shenzhen, China
Shenzhen Research Institute of Big Data
Ruoyu Sun Correspondence author. Email<EMAIL_ADDRESS>The Chinese
University of Hong Kong, Shenzhen, China  Shenzhen Research Institute of Big
Data  Zhi-Quan Luo Email<EMAIL_ADDRESS>The Chinese University of Hong
Kong, Shenzhen, China  Shenzhen Research Institute of Big Data

###### Abstract

SGD performs worse than Adam by a significant margin on Transformers, but the
reason remains unclear. In this work, we provide an explanation of SGD’s
failure on Transformers 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 badly on problems with block heterogeneity. To
validate that heterogeneity hampers SGD, we check various Transformers, CNNs,
MLPs, and quadratic problems, and find that SGD works well on problems without
block heterogeneity but performs badly when the heterogeneity exists. Our
initial theoretical analysis indicates that SGD fails because it applies one
single learning rate for all blocks, which cannot handle the heterogeneity
among blocks. The failure could be rescued if we could assign different
learning rates across blocks, as designed in Adam.

##  1 Introduction

Transformers (Vaswani et al., 2017) have become a major workhorse behind much
remarkable progress in AI development (e.g., (Achiam et al., 2023)). However,
the understanding of Transformer training remains limited. For instance,
Transformer training largely relies on Adam(W) (Kingma and Ba, 2014;
Loshchilov and Hutter, 2017). In contrast, Stochastic Gradient Descent with
momentum (SGD) 111We introduce the update rules of Adam(W) and SGD in Appendix
B.1., which is the de-facto optimizer for convolution neural networks (LeCun
et al., 1998) (CNNs), performs poorly on Transformers (see Figure
LABEL:fig:cv_figure and LABEL:fig:nlp_figure in Appendix D as evidence). Yet,
the cause of failure remains unclear. It is an intriguing question why SGD
fails on Transformers. First, from a theoretical perspective, this topic helps
reveal more properties of the Transformer training loss landscape, which helps
boost understanding of training objectives and current optimizers. Second,
from a computational perspective, Adam requires heavier memory storage than
SGD. Compared with SGD, Adam needs to store an additional second-order
momentum, which takes at least the same memory of the model size. The memory
consumption of Adam has become a major overhead in large-scale transformer
training (Rasley et al., 2020). Identifying the cause of failure may help
researchers design more efficient optimizers.

In this work, we explore the failure mode of SGD through the lens of Hessian.
We start by investigating the full Hessian spectrum of Transformers, i.e., the
full eigenvalue density of Hessian (see Figure LABEL:fig_full_spectrum as an
example). By theory, the full Hessian spectrum largely determines the behavior
of gradient methods (Nesterov, 2013; Goh, 2017; Sun, 2019; Goujaud et al.,
2022), so we suspect it can also help uncover the reasons behind SGD failure
on Transformers. Using numerical linear algebra tools (Bai et al., 1996), we
empirically compare the full spectra of CNNs (the scenarios where SGD is on
par with Adam) and those of Transformers (where SGD largely lags behind Adam).
Unfortunately, in Figure LABEL:fig_full_spectrum, we find that the spectra on
CNNs and Transformers are often largely similar despite the different
optimizer behaviors on these tasks. As such, we have not identified critical
phenomena in the full Hessian spectra associated with SGD’s failure on
Transformers. To reveal the cause of failure, we need a more fine-grained
investigation into Hessian.

What else could make SGD fail on Transformers, but not on CNNs? By dissecting
the structures of CNNs and Transformers, we observe that CNN architectures are
constructed by the repetitive stacking of similar parameter blocks
(convolution layers), while Transformers involve the interleaved stacking of
disparate parameter blocks (e.g. Query, Key, Value blocks in attention and MLP
layers). We suspect that these design differences may lead to different
properties in optimization. Intuitively, disparate parameter blocks would
contribute differently to the total loss. So each block might prefer a
distinct treatment by optimizers, which can be provided by Adam but not by
SGD. This motivates us to investigate the Hessian spectrum of each parameter
block, abbreviated as the blockwise Hessian spectrum.

By inspecting the blockwise Hessian spectrum, we find a possible explanation
why SGD fails on Transformers. We find that Transformers are “heterogeneous":
for all Transformers we checked, Hessian spectra of different parameter blocks
are dramatically different. We then verify that such heterogeneity indeed
hampers SGD via extensive experiments on various Transformers, CNNs, MLPs ,
and quadratic problems. Our contributions are summarized as follows.

  * •

Exploring the failure mode of SGD on transformers. We provide an explanation
why SGD largely lags behind Adam on Transformers through the lens of blockwise
Hessian spectrum. This explanation consists of two parts. First, we identify a
phenomenon called “block heterogeneity", which states that Hessian spectra are
dramatically different across parameter blocks. We find this phenomenon
prevalent on all Transformers in our experiments, but not on CNNs. Second, we
verify that block heterogeneity hampers SGD. More specifically, on various
kinds of Transformers, CNNs, MLPs, we find that SGD always performs badly on
problems with block heterogeneity but works well otherwise.

  * •

Theoretical results on quadratic models. We build various convex quadratic
problems, some with block heterogeneity and others without block
heterogeneity. We find that GD is significantly worse than Adam on problems
with block heterogeneity, but performs similarly to Adam otherwise. We provide
a theoretical result which shows that GD indeed works badly on quadratic
problems with block heterogeneity. Our analysis suggests that one important
reason for (S)GD failure is that it uses one single learning rate for all
blocks, which cannot handle the heterogeneity among blocks. The failure could
be saved if we assign different learning rates across blocks like Adam.

##  2 Main Results

###  2.1 Problem Settings

#### Notations.

We denote the full-batch training loss as
$\mathcal{L}(w)\equiv\frac{1}{n}\sum_{i=1}^{n}\mathcal{L}_{i}(w)$, where $n$
is the number of minibatches, $\mathcal{L}_{i}(w)$ is the loss of $i$-th
minibatch and $w\in\mathbb{R}^{d}$ is the neural network parameters. We denote
the gradient and Hessian of the training loss w.r.t. neural network parameters
as $\nabla\mathcal{L}(w)\in\mathbb{R}^{d}$ and
$\nabla^{2}\mathcal{L}(w)\in\mathbb{R}^{d\times d}$, respectively. We use
$\nabla\mathcal{L}_{i}(w)\in\mathbb{R}^{d}$ and
$\nabla^{2}\mathcal{L}_{i}(w)\in\mathbb{R}^{d\times d}$ to denote the $i$-th
minibatch counterparts. We use $[d]$ to denote the index set
$\\{1,2,\cdots,d\\}$. Given an arbitrary partition
$\\{\mathcal{D}_{l}\\}^{L}_{l=1}$ over $[d]$ with
$d_{l}\triangleq|\mathcal{D}_{l}|$, we can split $w$ into $L$ parameter blocks
$\\{w_{l}\\}_{l=1}^{L}$, where $w_{l}=\mathbb{R}^{d_{l}}$ consists of
parameters with indexes in the $l$-th block $\mathcal{D}_{l}$. We denote
$[\nabla^{2}\mathcal{L}(w)]_{l}\in\mathbb{R}^{d_{l}\times d_{l}}$ as the
Hessian of $l$-th parameter-block $w_{l}$, where
$[\nabla^{2}\mathcal{L}(w)]_{l,i,j}=\frac{\partial^{2}}{\partial_{w_{l,i}}\partial_{w_{l,j}}}\mathcal{L}(w_{l})$.
Note that $[\nabla^{2}\mathcal{L}(w)]_{l}$ is the $l$-th principle block sub-
matrix of $\nabla^{2}\mathcal{L}(w)$.

#### Preliminaries.

Hessian of large-scale NNs are intractable to obtain due to the numerical
difficulties of computing and storing them. Fortunately, there are extensive
numerical linear algebra methods to approximate the Hessian spectrum without
calculating the Hessian explicitly. Among which, the Stochastic Lanczos
Quadrature method (SLQ) (Bai et al., 1996) is known to be one of the oldest
yet the most powerful methods. SLQ is a combination of the stochastic trace
estimator (Hutchinson, 1989), the Gaussian quadrature (Golub and Meurant,
2009; Epperson, 2013), and the Lanczos algorithm (Lanczos, 1950). In short,
SLQ takes a handful of Hessian-vector products as input, where the vectors are
sampled from the normalized Rademacher distribution. It returns a smooth curve
on $\mathbb{R}$ that approximates the histograms of eigenvalues (see Figure
LABEL:fig_full_spectrum as an example). Since the derivation of SLQ is quite
involved, we defer the detailed description of SLQ to Appendix B.2. The
algorithmic form of SLQ can be seen in Algorithm 5 in the Appendix. We provide
running time analysis and detailed setup of SLQ in Appendix C.1.

In this work, we apply SLQ to various CNNs and Transformers on different
datasets. As for CNNs, we study CNNs including ResNet18 (11M) and VGG16 (138M)
on ImageNet (He et al., 2016; Simonyan and Zisserman, 2014). As for
Transformers, we study ViT-base (86M) on ImageNet (Dosovitskiy et al., 2020),
BERT (40M) on Cornell Movie-Dialogs Corpus (Devlin et al., 2018), and GPT2
(125M) on Openwebtext (Radford et al., 2019). For a fair comparison with
ResNet18, we also consider smaller models like
GPT2-nano222https://github.com/karpathy/nanoGPT/ (11M) on English corpus. The
largest model we conduct is GPT2 with $125$ million parameters. Detailed
experimental setups are shown in Appendix C.1. For all these Transformer-based
tasks, SGD performs significantly worse than Adam; while for CNN-based takes,
Adam and SGD are similar. The performance comparison is shown in Figure
LABEL:fig:cv_figure and LABEL:fig:nlp_figure in the Appendix. We estimate the
spectrum of (1) the full Hessian $\nabla^{2}\mathcal{L}(w)$ and (2) the
Hessian of all parameter blocks $[\nabla^{2}\mathcal{L}(w)]_{l},l\in[L]$. For
(2), we split $w$ according to the default partition in PyTorch
implementation, e.g., Embedding layer, Query in 1st attention layer, Key in
1st attention layer, Value in 1st attention layer, etc..

Note that the term “block" is a bit different from the term “layer". For
instance, Query and Key could be in the same layer, but they are different
parameter blocks. In the following contexts, we refer to “full Hessian
spectrum" as the estimated spectrum of $\nabla^{2}\mathcal{L}(w)$ and
“blockwise Hessian spectrum" as the estimated spectrum of
$[\nabla^{2}\mathcal{L}(w)]_{l},l\in[L]$.

###  2.2 Full Hessian Spectrum Is Not Informative Enough

Motivated by classical theory, we start by investigating the full Hessian
spectrum of Transformers. We choose to study the full Hessian spectrum for two
reasons. First, as mentioned in Section 1, the full Hessian spectrum is known
to largely determine the behavior of gradient methods. For instance, the range
of the spectrum will affect the convergence rate (Nesterov, 2013). Second, the
full Hessian spectrum has been shown to be informative in explaining multiple
phenomena in neural nets (NNs), e.g., how BatchNorm accelerates the training
(Ghorbani et al., 2019). Due to these reasons, we hypothesize that the full
Hessian spectrum can also help uncover the reasons behind SGD failure on
Transformers.

We compare the full spectra of CNNs (the scenarios where SGD is on par with
Adam) and those of Transformers (where SGD largely lags behind Adam). The
results are shown in Figure LABEL:fig_full_spectrum. Unfortunately, our
experiments show that the full Hessian spectrum might not be enough to explain
the SGD failure on Transformers. We explain as follows. The primary
information of the spectrum lies in its (1) dispersion, (2) shape, and (3)
evolution along training. As for (1), we find that eigenvalues are dispersed
very similarly and we do not observe any notably large eigenvalues in the
spectra of Transformers. So (1) does not seem to be related to the failure of
SGD. We further inspect (2) and (3). For all CNNs and Transformers in Figure
LABEL:fig_full_spectrum, we find similar phenomena: the shape of the spectrum
is approximately symmetrical around 0 at initialization. As training proceeds,
the majority of large negative eigenvalues disappear and the shape evolves
into a combination of a “bulk" and some “outliers". Since the spectral shape
and evolution are quite similar on Transformers and CNNs, they cannot explain
the failure of SGD on Transformers, either. In summary, we have not identified
critical phenomena in the full Hessian spectra linked to SGD’s failure on
Transformers.

###  2.3 Main Findings Through Blockwise Hessian Spectra

We have shown that the full Hessian spectrum seems not sufficient for
uncovering the reason why SGD fails on Transformers. What else could make SGD
fail on Transformers, but not on CNNs? We find two important features that are
overlooked in the full Hessian spectrum analysis above.

  * •

(1)  The Hessian structure is overlooked. There are numerical results showing
that the Hessians are close to block-diagonal matrices on MLPs (Collobert,
2004; Roux et al., 2007; Martens and Grosse, 2015). We restate their findings
in Appendix C.2. Further, (Collobert, 2004, Section 7) theoretically shows
that the block diagonal structure comes from (i) the layer-by-layer design in
NNs and (ii) the Cross-Entropy loss. Following this line of work, we also
observe a near block-diagonal Hessian in small Transformers in Figure
LABEL:fig_block_diagnal, where the variables in each principle block
correspond to the parameters of each block in the Transformer. These results
suggest that near block-diagonal Hessian might be common in NNs.

  * •

(2) The build-up rules of Transformers is overlooked. Transformers are built
upon stacking of disparate parameter blocks, e.g. Query, Key, Value blocks in
attention, and MLP layers. This is largely different from CNNs, which are
constructed by the repetitive stacking of similar parameter blocks
(convolution layers). Since each parameter blocks in Transformers are designed
differently, they might have different properties for optimization, which
might further affect the optimizer behavior.

Combining (1) and (2), we hypothesize that the blockwise Hessian spectrum
might provide some insights on distinguishing CNNs and Transformers. What
extra information can be contained in the blockwise spectrum but not in the
full spectrum? By definition, the blockwise Hessians form the principal block
sub-matrix of the full Hessian. For block diagonal matrices, blockwise Hessian
encodes the location of eigenvalues, i.e., which block does an eigenvalue (of
the full Hessian) lie in, while the full spectrum ignores this. In the
following, we study blockwise Hessian spectra on various models. We will show
that blockwise spectra indeed carry more information than the full spectrum
for distinguishing CNNs and Transformers.

We here demonstrate the shape of blockwise spectra in CNNs and Transformers.
We sample four blocks per model and present the spectra in Figure
LABEL:fig_blockwise_spectrum. In Transformers, the Hessian spectra of
embedding, attention, and MLP blocks are largely different. In contrast, in
CNNs, the spectra of convolution layers are similar. We further verify this
observation for the rest of the parameter blocks. We calculate the Jensen-
Shannon (JS) distance between two eigenvalue densities of all possible block
pairs. We show the results in Figure LABEL:fig_blockwise_heatmap.

The results in Figure LABEL:fig_blockwise_heatmap showed an interesting
phenomenon: for all the Transformers we checked, the blockwise Hessian spectra
are largely different across blocks. In the following, we refer to this
phenomenon as “block heterogeneity". As a control group experiment, the
blockwise Hessian spectra of CNNs are similar and the block heterogeneity is
not observed. We refer to this phenomenon as “block homogeneity". These
results indicate that block heterogeneity is informative in distinguishing
CNNs and Transformers. Intuitively, the block-homogeneous Hessian in a CNN
comes from the repetitively similar convolution layers, while the block-
heterogeneous Hessian in a Transformer stems from the interleaved stacking
disparate layers such as Query, Value, and MLPs. In the following, we will
show that the block heterogeneity is strongly correlated with the failure of
SGD on Transformers.

###  2.4 SGD Fails on Simple MLPs with Block Heterogeneity

We now further link block heterogeneity to the failure of SGD. We construct a
simple example without attention layers but has block heterogeneity and find
that SGD still fails. We consider a 4-layer MLP on MNIST. Using standard MLP,
we find that the Hessian exhibits a homogeneous block structure and SGD
converges as fast as Adam. We then scale the output of each layer by a
constant $c$. We define $c$ as the degree of heterogeneity since larger $c$
would bring more heterogeneity to the Hessian (see Table 3 in the Appendix).
As shown in Figure LABEL:fig_mlp_gap, we find that the gap between Adam and
SGD increases with the degree of heterogeneity. This example shows that
attention might not be the root cause of SGD failure, but the block
heterogeneity (brought by attention) might be.

###  2.5 Implication on Choosing SGD or Adam

The above findings could bring up an interesting 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 (Gu and Dao,
2023). 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. We present two reasons.

  * •

In large-scale model training, memory is limited but the model size grows, so
Adam could become a major computational burden due to the memory of 2nd-order
momentum (as mentioned in Section 1 and (Rasley et al., 2020)).

  * •

SGD is a lightweight alternative to Adam, but the performance of SGD may be
bad. In large NN training, it might take days (or even weeks) to obtain
informative feedback from training logs, so it is expensive to conduct trial
and error to choose an appropriate optimizer. Specifically, if SGD performs
badly, it would be difficult to judge whether the bad performance is due to
improper learning rates or other more fundamental reasons. Then, it would be
hard to decide whether to switch the optimizer.

Motivated by the above reasons, we here propose a quantitative metric that
could predict the failure of SGD before the training is launched. 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 each heatmap of Figure
LABEL:fig_blockwise_heatmap. We denote it as $JS^{0}$. We present $JS^{0}$ of
various models in Table 1. We note that $JS^{0}$ establishes a quantitative
difference between the loss landscape of Transformers and CNNs. As argued in
the previous sections, this quantity can predict the failure of SGD. Note that
$JS^{0}$ is independent of optimizers and could be efficiently checked in
priori before training.

Model | ResNet18 | VGG16 | ViT-base | BERT | GPT2-nano | GPT2
---|---|---|---|---|---|---
$JS^{0}$ | 0.10 | 0.09 | 286.41 | 53.38 | 14.93 | 83.23
Table 1:  $JS^{0}$ denotes the average JS distance between the initial Hessian
spectra of each pair of parameter blocks. We find that the $JS^{0}$ on CNNs is
at least $140\times$ smaller than on Transformers. $JS^{0}$ establishes a
quantitative difference between the loss landscape of Transformers and CNNs.

To summarize this section, we provide an explanation for SGD’s failure. We
identify a phenomenon called block heterogeneity. This phenomenon naturally
exists in Transformers, but not in CNNs. This phenomenon can exist in non-
Transformer architectures as well. Our experiments suggest a strong connection
between the failure of SGD and block heterogeneity. Our findings could provide
empirical guidance to decide whether to use Adam or SGD.

##  3 Case Study of Quadratic Models and Preliminary Theory

Now we construct quadratic functions with block diagonal Hessian. We will
compare the performance of GD and Adam on Hessian with and without block
heterogeneity. Initial theoretical results on these quadratic models will be
provided. Note that insights on quadratic models could be important for
understanding realistic NNs, as mentioned by researchers such as LeCun et al.
(2002) and OpenAI team (Kaplan et al., 2020).

###  3.1 Settings and Additional Notations

We consider $\min_{w}\mathcal{L}(w)=\frac{1}{2}w^{T}Hw-h^{T}w$ where
$H\in\mathbb{R}^{d\times d}$ is positive definite and $h\in\mathbb{R}^{d}$. We
denote $\mathcal{L}^{*}$ as the minimum value of $\mathcal{L}(w)$. We set $H$
as a block diagonal matrix: $H=\operatorname{diag}(H_{1},\cdots,H_{L})$, where
$H_{l}\in\mathbb{R}^{d_{l}\times d_{l}}$ and $d=\sum_{l=1}^{L}d_{l}$. We use
$w_{l}\in\mathbb{R}^{d_{l}}$ to denote the variable in the $l$-th block and
$w=(w_{1}^{T},\cdots,w_{L}^{T})^{T}\in\mathbb{R}^{d}$. Similarly for
$h_{l}\in\mathbb{R}^{d_{l}}$. Similarly, we use $[\nabla
L(w)]_{l}\in\mathbb{R}^{d_{l}}$ to denote the gradient in the $l$-th block and
denote
$[\mathcal{L}(w)]_{l}=\frac{1}{2}(w_{l}^{t})^{T}H_{l}w_{l}^{t}-h_{l}^{T}w_{l}$
as the objective function w.r.t. the $l$-th block. Note that
$\mathcal{L}(w)=\sum_{l=1}^{L}[\mathcal{L}(w)]_{l}$. We denote
$\lambda_{1}\geq\lambda_{2}\cdots\geq\lambda_{d}$ as the eigenvalues of $H$.
Similarly for $\lambda_{l,1}\cdots\lambda_{l,d_{l}}$. We denote
$\kappa=\frac{\lambda_{1}}{\lambda_{d}}$ and
$\kappa_{l}=\frac{\lambda_{l,1}}{\lambda_{l,d_{l}}}$ as the condition number
of $H$ and $H_{l}$, respectively. We say an algorithm has complexity
$\tilde{\mathcal{O}}(C)$ if it takes $\mathcal{O}(C\log(1/\epsilon))$
iterations to achieve error
$\frac{\mathcal{L}(w)-\mathcal{L}^{*}}{\mathcal{L}\left(w^{0}\right)-\mathcal{L}^{*}}\leq\epsilon$,
where $w^{0}$ is the initial point.

###  3.2 Experimental Observations

For all experiments, we choose $h=0$. We consider four types of Hessian $H$ as
follows.

  * •

Case 1: Hessian with Transformer-type spectra. We choose $L=4$ and
$d_{l}=100$. For $l\in[L]$, we construct $H_{l}=Q_{l}\Lambda_{l}Q_{l}^{T}$
where $Q_{l}$ are independent standard Gassian random matrix and $\Lambda_{l}$
are diagonal matrices. For the diagonal elements in $\Lambda_{l}$, we sample
$d_{l}$ 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 $\Lambda_{l}$ lie in the interval
$[1,2000]$. This ensures strong convexity and controls the condition number of
$H$ equals $2000$. The spectra of $H_{l}$ are in Figure LABEL:fig:heter-block.

  * •

Case 2: Hessian with CNN-type spectra. We consider the same setup as Case 1.
For the diagonal elements in $\Lambda_{l}$, we sample $d_{l}$ numbers
according to the spectrum of the 1st to 4th convolution layers in ResNet18. We
shift and proportionally scale the elements $\Lambda_{l}$ to the interval
$[1,2000]$ to ensure the strong convexity and the condition number $H$ equals
2000. The spectra of $H_{l}$ are shown in Figure LABEL:fig:homo-block.

  * •

Case 3: Hessian with simplified heterogeneous spectra. We choose $L=3$ and
$d_{l}=3$. For $l\in[L]$, we construct $H_{l}=Q_{l}\Lambda_{l}Q_{l}^{T}$ where
$Q_{l}$ are independent standard Gassian random matrix and $\Lambda_{l}$ are
diagonal matrices. We set the diagonal elements of $\Lambda_{l}$ as
$\\{1,2,3\\},\\{99,100,101\\},\\{1998,1999,2000\\}$ for $l=1,2,3$,
respectively. The spectra of $H_{l}$ are different due to their different
supports. The condition number of Hessian $H$ is $2000$.

  * •

Case 4: Hessian with simplified homogeneous spectra. We consider the same
setup as Case 3. We set the diagonal elements of $\Lambda_{l}$ as
$\\{1,99,1998\\},\\{2,100,1999\\},\\{3,101,2000\\}$ for $l=1,2,3$,
respectively. The spectra of $H_{l}$ are similar. The condition number of
Hessian $H$ is 2000.

Note that in all four cases, the Hessian share the same condition number. We
study two types of optimizers: one that assigns a single learning rate for all
blocks, and ones that assign different learning rates across blocks. The
optimizers are introduced as follows.

  * •

Single-learning-rate optimizer. We study gradient descent (GD).

$w^{t+1}=w^{t}-\eta\nabla\mathcal{L}(w)=w^{t}-\eta(Hw^{t}-h)$ (1)

We use the optimal learning rate $\eta=\frac{2}{\mu+L}$ (Nesterov, 2013). We
use a standard Gaussian random vector as initialization.

  * •

Coordinate-wise-learning-rate optimizer. We study Adam with a constant
learning rate and with no bias correction for simplicity (Algorithm 3). We set
$\beta_{1}=0$ to erase the effect of momentum. This helps us to focus on the
effect of coordinate-wise learning rate (or the effect of diagonal
preconditioning) in Adam. We use $\epsilon=0$. We consider $\beta_{2}=1$ and
$\beta_{2}=0.999$, respectively. When $\beta_{2}=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.

$w^{t+1}=w^{t}-\eta(D_{Adam}^{0})^{-1}\nabla\mathcal{L}(w)=w^{t}-\eta(D_{Adam}^{0})^{-1}(Hw^{t}-h),$
(2)

where
$D_{Adam}^{0}=\operatorname{diag}(\nabla\mathcal{L}(w^{0})\circ\nabla\mathcal{L}(w^{0}))^{\frac{1}{2}}$
and $\nabla\mathcal{L}(w^{0})=Hw^{0}-h$. When $\beta_{2}<1$, the coordinate-
wise learning rates adaptively change along iteration. The update rule is as
follows.

$w^{t+1}=w^{t}-\eta(D_{Adam}^{t})^{-1}\nabla\mathcal{L}(w)=w^{t}-\eta(D_{Adam}^{t})^{-1}(Hw^{t}-h),$
(3)

where
$D_{Adam}^{t}=\operatorname{diag}\left((1-\beta_{2})\left(\sum_{k=1}^{t}\beta_{2}^{t-k}\nabla\mathcal{L}(w^{k})\circ\nabla\mathcal{L}(w^{k})\right)+\beta^{t}\operatorname{diag}(\nabla\mathcal{L}(w^{0})\circ\nabla\mathcal{L}(w^{0}))\right)^{\frac{1}{2}}$
and $\nabla\mathcal{L}(w^{k})=Hw^{k}-h$. We grid search $\eta$ and use a
standard Gaussian random vector as initialization.

We remark that when $\beta_{2}<1$, Adam would repeatedly bounce among non-
optimal points, causing non-convergence. This will be shown in Proposition 2.

#### Summary of experimental observations.

We now compare the performance of Adam and GD on these four types of Hessians.
The results are shown in Figure LABEL:fig_quadratic. For Hessian with
heterogeneous blocks like Case 1 and 3, GD largely lags behind Adam. For
Hessian with homogeneous blocks (Case 2 and 4), GD is on par or even better
than Adam. We emphasize that all Hessians have the same condition number in
these four cases. Further, Hessian of Case 3 and 4 share all the eigenvalues
(not just the extreme ones). The performance gap between Adam and GD is due to
the different blockwise spectra caused by the different locations of
eigenvalues. 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.

###  3.3 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 $\kappa$. However, we point out that $\kappa$ is not informative
enough to describe the performance gap in Figure LABEL:fig_quadratic since
$\kappa$ 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.

Unfortunately, the blockwise spectrum is rarely discussed in the optimization
area. A most related notion is “block Lipschitz constant" (Beck and
Tetruashvili, 2013) for studying block coordinate descent (BCD) type methods,
but it was not linked to the performance of SGD or Adam before. To our
knowledge, we are not aware of any theory of Adam or GD built on the block
diagonal structures or the blockwise spectra of Hessian. Further, there seems
to be no result distinguishing the complexity of single-learning-rate methods
and the coordinate-wise counterpart, either. 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_{w}\mathcal{L}(w)=\frac{1}{2}w^{T}Hw-h^{T}w$ where
$H\in\mathbb{R}^{d\times d}$ is positive definite and $h\in\mathbb{R}^{d}$.
Let $w_{GD}^{t}$ be the output of GD after $t$ steps. There exists a block
diagonal matrix $H$, $h$ and an initial point $w^{0}$, s.t., for any $\eta$,
we have:

$\mathcal{L}(w_{GD}^{t+1})-\mathcal{L}^{*}\geq\left(1-\frac{2}{\kappa+1}\right)\left(\mathcal{L}(w_{GD}^{t})-\mathcal{L}^{*}\right)$
(4)

where $\kappa$ is the condition number of $H$.

Proposition 1 shows that GD has complexity $\tilde{\mathcal{O}}(\kappa)$ and
such complexity is tight. Now we prove that: if we are allowed to choose a
customized learning rate for each block sub-matrix $H_{l}$, then the
complexity will be improved. We call this method blockwise GD (BGD). Note that
BGD is impractical to implement since it is hard to choose $L$ learning rates
in advance. BGD serves as an artificial method for theoretical understanding.

#### Blockwise GD (BGD).

Choose learning rates $\\{\eta_{l}\\}_{l=1}^{L}$ for all blocks in advance and
performs the following:

$w^{t+1}=w^{t}-(D_{BGD})^{-1}\nabla\mathcal{L}(w)=w^{t}-(D_{BGD})^{-1}(Hw^{t}-h),$

where
$D_{BGD}=\operatorname{diag}\left(\left(\underbrace{\eta_{1},\cdots,\eta_{1}}_{d_{1}\text{times}},\cdots,\underbrace{\eta_{L},\cdots,\eta_{L}}_{d_{L}\text{times}}\right)^{T}\right)\in\mathbb{R}^{d\times
d}.$

BGD can be also viewed as a diagonally preconditioned version of GD. We now
show that this diagonal preconditioner can significantly improve the
complexity of GD.

######  Theorem 1.

(Upper bound for GD with blockwise learning rate.) Consider
$\min_{w}\mathcal{L}(w)=\frac{1}{2}w^{T}Hw-h^{T}w$ where
$H\in\mathbb{R}^{d\times d}$ is positive definite and $h\in\mathbb{R}^{d}$. We
assume $H$ is a block diagonal matrix, i.e.,
$H=\operatorname{diag}(H_{1},\cdots,H_{L})$, where
$H_{l}\in\mathbb{R}^{d_{l}\times d_{l}}$ and $d=\sum_{l=1}^{L}d_{l}$. Let
$w_{BGD}^{t}$ be the output of BGD after $t$ steps. There exists a set of
$\\{\eta_{l}\\}_{l=1}^{L}$ such that:

$\mathcal{L}(w_{BGD}^{t+1})-\mathcal{L}^{*}\leq\max_{l\in[L]}\left(1-\frac{1}{\kappa_{l}}\right)\left(\mathcal{L}(w_{BGD}^{t})-\mathcal{L}^{*}\right)$
(5)

where $\kappa_{l}$ is the condition number of $H_{l}$.

Theorem 1 shows that: if we are allowed to choose a learning rate for each
$H_{l}$, the complexity could be $\tilde{O}(\max_{l\in[L]}\kappa_{l})$, in
lieu of $\tilde{O}(\kappa)$ as in GD. That is to say, the rate of BGD only
depends on the "slowest block". Note that in certain heterogeneous block
Hessian such Case 1 and ${\bf 3}$, $\kappa_{l}$ is about $20\times$ and
$2000\times$ smaller than $\kappa$, respectively. In these cases, blockwise-
learning-rate methods can be about $20\times$ and $2000\times$ faster than
single-learning-rate methods like GD, respectively.

Despite the advantages of BGD, it is impractical to find the proper learning
rates for all blocks, especially when the number of blocks $L$ is large.
Fortunately, these learning rates could be approximated by the preconditioner
in Adam, e.g., $D_{Adam}^{0}$ in (2). We now provide a complexity bound of
Adam.

######  Assumption 1.

(Bounded initialization.) The initialization is bounded and is far from the
optimal solution:

$C_{l,1}\lambda_{l,1}\leq\left|[\nabla\mathcal{L}(w^{0})]_{l,i}\right|\leq
C_{l,2}\lambda_{l,1},\text{ for }i\in[d_{l}],l\in[L].$

######  Theorem 2.

(Upper bound for Adam with $\beta_{2}=1$.) Consider the same setting as in
Theorem 1 and consider Adam with $\beta_{1}=0$ and $\beta_{2}=1$ as in (2).
Assume the initialization satisfies Assumption 1. Let $w_{Adam}^{t}$ be the
output of Adam after $t$ steps. Let $\eta=\min_{l\in[L]}\frac{1}{C_{l,1}}$. We
have

$\mathcal{L}(w_{Adam}^{t+1})-\mathcal{L}^{*}\leq\max_{l\in[L]}\left(1-\frac{1}{\kappa_{Adam,l}}\right)\left(\mathcal{L}(w_{Adam}^{t})-\mathcal{L}^{*}\right)$
(6)

where $\kappa_{Adam,l}=r\kappa_{l}$,
$r:=\max_{l\in[L]}\left(\frac{C_{l,2}}{C_{l,1}}\right)^{2}$ and $\kappa_{l}$
is the condition number of $H_{l}$.

The proofs of all the above theorems are shown in Appendix E. Theorem 2 states
that Adam (with $\beta_{2}=1$) has complexity
$\tilde{O}\left(r\cdot\max_{l\in[L]}\kappa_{l}\right)$. We remark that
condition $\beta_{2}=1$ is necessary because any $\beta_{2}<1$ causes non-
convergence issue (Da Silva and Gazeau, 2020; Bock and Weiß, 2019). We restate
their results in Proposition 2. The non-convergence is also observed in Figure
LABEL:fig_quadratic (c), where we find that the iterates of Adam quickly
converge to near-optimal solutions, and then bounce back. As such,
$\beta_{2}=1$ is necessary for asymptotic analysis. As shown in (Da Silva and
Gazeau, 2020), the non-convergence is due to the constant learning rate.

######  Proposition 2.

(Non-convergence of constant-learning-rate Adam with $\beta_{2}<1$.) (Da Silva
and Gazeau, 2020, Proposition 12) Consider
$\min_{w\in\mathbb{R}}\mathcal{L}(w)=\frac{1}{2}w^{2}$. Consider Adam with
$\beta_{1}=0$ and $\beta_{2}<1$ as in (3). Let $w_{Adam}^{t}$ be the output of
Adam after $t$ steps. There exists a discrete limit cycle for (3), i.e., Adam
produces oscillations and
${\small\liminf_{t\rightarrow\infty}\left(\mathcal{L}(w_{Adam}^{t})-\mathcal{L}^{*}\right)>0}$.

We now compare the complexity of Adam and that of GD. By Theorem 2, Adam is
faster than GD when $r\cdot\max_{l\in[L]}\kappa_{l}\leq\kappa$. In the
quadratic model with heterogeneous blocks (Case 3), our simulation over 1000
trials shows that $r\leq 100$ with probability $\geq\frac{2}{3}$ when using
standard Gaussian random initialization. Since
$\max_{l\in[L]}\kappa_{l}\approx 1$, we have
$r\cdot\max_{l\in[L]}\kappa_{l}\leq 100$, w.h.p., and is $20\times$ smaller
than $\kappa=2000$. So Adam could be $20\times$ faster than GD, w.h.p.. This
is indeed observed in Figure LABEL:fig_quadratic where Adam outperforms GD by
a significant margin. We summarize the complexity of GD, BGD, and Adam in
Table 2.

Optimizer | GD | BGD | Adam with | Adam with
---|---|---|---|---
|  |  |  $\beta_{1}=0$ and $\beta_{2}=1$ (2) |  $\beta_{1}=0$ and $\beta_{2}<1$ (3)
Complexity | $\tilde{O}(\kappa)$ | $\tilde{O}(\max_{l\in[L]}\kappa_{l})$ | $\tilde{O}\left(r\cdot\max_{l\in[L]}\kappa_{l}\right)$ | ✗
Table 2:  The complexity of GD, BGD and Adam for minimizing a strongly convex
quadratic function with block diagonal Hessian. The symbol ✗ means non-
convergence. $\kappa$ and $\kappa_{l}$ denote the condition number of the full
Hessian and the block submatrix, respectively.
$r:=\max_{l\in[L]}\left(\frac{C_{l,2}}{C_{l,1}}\right)^{2}$ and $C_{l,1}$,
$C_{l,2}$ are constants in Assumption 1.

#### How to obtain a tighter complexity bound of Adam?

It is valid to ask whether the complexity upper bound
$\kappa_{Adam,l}=r\kappa_{l}$ can be tightened, e.g., remove the factor of
$r$. We point out it would be difficult if there is no extra structure on
$H_{l}$. A key technical step is to bound the condition number of the
preconditioned matrix $\kappa\left((D_{Adam,l}^{0})^{-1}H_{l}\right)$.
Intuitively, a diagonal preconditioner of $H_{l}$ is expected to be powerful
when $H_{l}$ itself has a near-diagonal structure, e.g., pure diagonal,
tridiagonal or diagonal dominant. Unfortunately, it is unclear whether these
structures hold in real Transformers. Without any assumption on $H_{l}$, we
find that the diagonal preconditioner of $D_{Adam}^{0}$ could increase the
condition number. For instance, when using standard Gaussian initialization,
in case 3, we find $\kappa\left((D_{Adam,l}^{0})^{-1}H_{l}\right)$ equals
$3.9\kappa_{1}$, $14.1\kappa_{2}$, $4.2\kappa_{3}$ for the 3 blocks,
respectively (all averaged over 1000 trials). It would be interesting to
explore if there are special structures of $H_{l}$ in real Transformers such
that Adam preconditioner can reduce $\kappa_{l}$, rather than increase it. We
leave it as a future direction.

Although Adam preconditioner might not always reduce the “local" condition
number $\kappa_{l}$, the coefficient in the complexity is now independent of
the “global" condition number $\kappa$. As argued above, such changes in
coefficient could lead to considerable improvement over GD. Such improvement
in complexity is attributed to the block diagonal structure in Hessian as well
as its heterogeneous blockwise spectrum. To our knowledge, such improvement is
not shown in the existing literature. In summary, our theory indicates that:
for problems with block heterogeneity, the single-learning rate methods like
GD can largely lag behind coordinate-wise learning rate methods like Adam.

##  4 Related Works

#### On The failure of SGD on Transformers

There is extensive literature studying the difficulties of Transformer
training. We summarize these works in Appendix A. We here focus on the
literature that explores why SGD fails on Transformers. One representative
hypothesis is that SGD fails because it cannot handle the heavy-tailed
stochastic noise in language tasks (Zhang et al., 2020). However, Chen et al.
(2021); Kunstner et al. (2023) reported that the gap between Adam and SGD
maintains even in the full-batch case with no stochasticity, so there might be
other reasons. Further, SGD performs poorly on Vision Transformers on ImageNet
(Figure LABEL:fig:nlp_figure), so the data modality (e.g., language or image
tasks) might not be as crucial as the architecture. Zhang et al. (2019c)
showed that Transformers have the “unbounded smoothness" issue and SGD with
gradient clipping performs better than SGD in this case. Although clipping is
an effective trick, it does not save SGD from failure as we still observe a
huge gap between clipped SGD and Adam 333For all NLP tasks, clipping is
performed after backpropagation. So in Figure LABEL:fig:nlp_figure, SGD in NLP
tasks essentially refers to clipped SGD. . So there might be other reasons
that hamper SGD. Different from these works, we find SGD fails 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
(Reddi et al., 2018). The convergence for the unmodified versions is later
established in (Shi et al., 2020; Zhang et al., 2022b) for RMSprop and Adam.
More convergence analyses of general adaptive gradient methods are listed in
Appendix A. We here focus on the literature that explores the benefit of Adam.
Xie et al. (2022) show that Adam can help avoid saddle points, which is an
orthogonal direction to this work. Wang et al. (2022a); Crawshaw et al.
(2022); Li et al. (2023) 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 fails while Adam works well (Figure LABEL:fig_quadratic).
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).

Bernstein et al. (2018); Wu et al. (2020); Kunstner et al. (2023); Liu et al.
(2023) build a relation between Adam and the sign-based methods. Wu et al.
(2020) further showed that sign-based 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 (2023) numerically found that the Adam can reduce the
directional sharpness along trajectories, while its relation to fast
convergence remains mysterious. A recent work (Jiang et al., 2023) 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 (Liu et al., 2023). The theory in (Jiang et
al., 2023) implies that Adam is faster when the Hessian is diagonal. However,
as argued above, it is unclear whether the diagonal Hessian commonly holds in
real problems. In fact, we find the Hessian is closer to a block-diagonal
(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. Sagun et al. (2016, 2017); Chaudhari et al. (2019) found that the
Hessian spectra of MLPs and CNNs consist of a “bulk" together with a few
“outliers". Papyan (2020); Wu et al. (2020); Liao and Mahoney (2021) further
characterized the bulks and outliers in theory. Papyan (2018, 2019)
numerically built the relation between these "outliers" and the Gauss-Newton
matrix. Sankar et al. (2021) numerically explored the relation between Hessian
of CNNs and Gauss-Newton matrix in each layer. They further showed that 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. (2018) 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. (2018); Zhang et al. (2019b) explored the relation between
the Hessian spectrum of CNNs and some training phenomena such as the effect of
batch sizes. Ghorbani et al. (2019); Yao et al. (2020) focused on 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) as well as its impacts on different optimizers.

##  5 Conclusion

In this work, we investigate the failure mode of SGD on Transformers through
the lens of Hessian. By numerically explore various Transformers, CNNs, MLPs
and quadratic problems, we establish a phenomenon called block heterogeneity
in Hessian and link it to the failure of SGD. We point out that SGD fails
because it applies one single learning rate for all blocks, which cannot
handle the heterogeneity among blocks. Initial theory are provided to support
the claim.

## References

  * Achiam et al. [2023] J. Achiam, S. Adler, S. Agarwal, L. Ahmad, I. Akkaya, F. L. Aleman, D. Almeida, J. Altenschmidt, S. Altman, S. Anadkat, et al.  Gpt-4 technical report.  _arXiv preprint arXiv:2303.08774_ , 2023.
  * Adams et al. [2018] 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.03451_ , 2018.
  * Avron and Toledo [2011] H. Avron and S. Toledo.  Randomized algorithms for estimating the trace of an implicit symmetric positive semi-definite matrix.  _Journal of the ACM (JACM)_ , 58(2):1–34, 2011.
  * Bachlechner et al. [2021] T. Bachlechner, B. P. Majumder, H. Mao, G. Cottrell, and J. McAuley.  Rezero is all you need: Fast convergence at large depth.  In _Uncertainty in Artificial Intelligence_ , pages 1352–1361. PMLR, 2021.
  * Bai and Golub [1996] Z. 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.
  * Bai et al. [1996] 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.
  * Beck and Tetruashvili [2013] A. Beck and L. Tetruashvili.  On the convergence of block coordinate descent type methods.  _SIAM journal on Optimization_ , 23(4):2037–2060, 2013.
  * Bernstein et al. [2018] J. Bernstein, Y.-X. Wang, K. Azizzadenesheli, and A. Anandkumar.  signsgd: Compressed optimisation for non-convex problems.  In _International Conference on Machine Learning_ , pages 560–569. PMLR, 2018.
  * Bock and Weiß [2019] S. Bock and M. Weiß.  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, Proceedings, Part II 28_ , pages 232–243. Springer, 2019.
  * Brezinski [1990] C. Brezinski.  A direct proof of the christoffel-darboux identity and its equivalence to the recurrence relationship.  _Journal of Computational and Applied Mathematics_ , 32(1-2):17–25, 1990.
  * Chaudhari et al. [2019] P. Chaudhari, A. Choromanska, S. Soatto, Y. LeCun, C. Baldassi, C. Borgs, J. Chayes, L. Sagun, and R. Zecchina.  Entropy-sgd: Biasing gradient descent into wide valleys.  _Journal of Statistical Mechanics: Theory and Experiment_ , 2019(12):124018, 2019.
  * Chen et al. [2022] 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.
  * Chen et al. [2021] 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.
  * Chen et al. [2018] 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:1804.09849_ , 2018.
  * Chen et al. [2019] X. Chen, S. Liu, R. Sun, and M. Hong.  On the convergence of a class of adam-type algorithms for non-convex optimization.  In _7th International Conference on Learning Representations, ICLR 2019_ , 2019.
  * Chowdhery et al. [2023] A. Chowdhery, S. Narang, J. Devlin, M. Bosma, G. Mishra, A. Roberts, P. Barham, H. W. Chung, C. Sutton, S. Gehrmann, et al.  Palm: Scaling language modeling with pathways.  _Journal of Machine Learning Research_ , 24(240):1–113, 2023.
  * Collobert [2004] R. Collobert.  Large scale machine learning.  Technical report, Université de Paris VI, 2004.
  * Crawshaw et al. [2022] M. Crawshaw, M. Liu, F. Orabona, W. Zhang, and Z. Zhuang.  Robustness to unbounded smoothness of generalized signsgd.  _Advances in Neural Information Processing Systems_ , 35:9955–9968, 2022.
  * Cullum and Willoughby [2002] J. K. Cullum and R. A. Willoughby.  _Lanczos algorithms for large symmetric eigenvalue computations: Vol. I: Theory_.  SIAM, 2002.
  * Da Silva and Gazeau [2020] 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):5072–5113, 2020.
  * Dao et al. [2022] 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.
  * Défossez et al. [2022] A. Défossez, L. Bottou, F. Bach, and N. Usunier.  A simple convergence proof of adam and adagrad.  _Transactions on Machine Learning Research_ , 2022.
  * Dehghani et al. [2023] 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 Learning_ , pages 7480–7512. PMLR, 2023.
  * Devlin et al. [2018] J. Devlin, M.-W. Chang, K. Lee, and K. Toutanova.  Bert: Pre-training of deep bidirectional transformers for language understanding.  _arXiv preprint arXiv:1810.04805_ , 2018.
  * Dong et al. [2021] 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.
  * Dosovitskiy et al. [2020] A. Dosovitskiy, L. Beyer, A. Kolesnikov, D. Weissenborn, X. Zhai, T. Unterthiner, M. Dehghani, M. Minderer, G. Heigold, S. Gelly, et al.  An image is worth 16x16 words: Transformers for image recognition at scale.  _arXiv preprint arXiv:2010.11929_ , 2020.
  * Duchi et al. [2011] J. Duchi, E. Hazan, and Y. Singer.  Adaptive subgradient methods for online learning and stochastic optimization.  _Journal of machine learning research_ , 12(7), 2011.
  * Epperson [2013] J. F. Epperson.  An introduction to numerical methods and analysis.  2013\.
  * Gadat and Gavra [2022] S. Gadat and I. Gavra.  Asymptotic study of stochastic adaptive algorithms in non-convex landscape.  _The Journal of Machine Learning Research_ , 23(1):10357–10410, 2022.
  * Ghorbani et al. [2019] B. Ghorbani, S. Krishnan, and Y. Xiao.  An investigation into neural net optimization via hessian eigenvalue density.  In _International Conference on Machine Learning_ , pages 2232–2241. PMLR, 2019.
  * Goh [2017] G. Goh.  Why momentum really works.  _Distill_ , 2017.  doi: 10.23915/distill.00006.  URL http://distill.pub/2017/momentum.
  * Golub and Meurant [2009] G. H. Golub and G. Meurant.  _Matrices, moments and quadrature with applications_ , volume 30.  Princeton University Press, 2009.
  * Golub and Strakoš [1994] G. H. Golub and Z. Strakoš.  Estimates in quadratic formulas.  _Numerical Algorithms_ , 8:241–268, 1994.
  * Golub and Welsch [1969] G. H. Golub and J. H. Welsch.  Calculation of gauss quadrature rules.  _Mathematics of computation_ , 23(106):221–230, 1969.
  * Goujaud et al. [2022] B. Goujaud, D. Scieur, A. Dieuleveut, A. B. Taylor, and F. Pedregosa.  Super-acceleration with cyclical step-sizes.  In _International Conference on Artificial Intelligence and Statistics_ , pages 3028–3065. PMLR, 2022.
  * Gu and Dao [2023] A. Gu and T. Dao.  Mamba: Linear-time sequence modeling with selective state spaces.  _arXiv preprint arXiv:2312.00752_ , 2023.
  * Gur-Ari et al. [2018] G. Gur-Ari, D. A. Roberts, and E. Dyer.  Gradient descent happens in a tiny subspace.  _arXiv preprint arXiv:1812.04754_ , 2018.
  * He et al. [2016] 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.
  * Huang et al. [2020] X. 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.
  * Hutchinson [1989] M. F. Hutchinson.  A stochastic estimator of the trace of the influence matrix for laplacian smoothing splines.  _Communications in Statistics-Simulation and Computation_ , 18(3):1059–1076, 1989.
  * Jiang et al. [2023] K. Jiang, D. Malik, and Y. Li.  How does adaptive optimization impact local neural network geometry?  _Advances in Neural Information Processing Systems_ , 36, 2023.
  * Kaplan et al. [2020] J. Kaplan, S. McCandlish, T. Henighan, T. B. Brown, B. Chess, R. Child, S. Gray, A. Radford, J. Wu, and D. Amodei.  Scaling laws for neural language models.  _arXiv preprint arXiv:2001.08361_ , 2020.
  * Kingma and Ba [2014] D. P. Kingma and J. Ba.  Adam: A method for stochastic optimization.  _arXiv preprint arXiv:1412.6980_ , 2014.
  * Kunstner et al. [2023] F. Kunstner, J. 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.
  * Lanczos [1950] C. Lanczos.  An iteration method for the solution of the eigenvalue problem of linear differential and integral operators.  1950\.
  * LeCun et al. [1998] Y. LeCun, L. Bottou, Y. Bengio, and P. Haffner.  Gradient-based learning applied to document recognition.  _Proceedings of the IEEE_ , 86(11):2278–2324, 1998.
  * LeCun et al. [2002] Y. LeCun, L. Bottou, G. B. Orr, and K.-R. Müller.  Efficient backprop.  In _Neural networks: Tricks of the trade_ , pages 9–50. Springer, 2002.
  * Li et al. [2023] H. Li, A. Rakhlin, and A. Jadbabaie.  Convergence of adam under relaxed assumptions.  _Advances in Neural Information Processing Systems_ , 36, 2023.
  * Liao and Mahoney [2021] Z. Liao and M. W. Mahoney.  Hessian eigenspectra of more realistic nonlinear models.  _Advances in Neural Information Processing Systems_ , 34:20104–20117, 2021.
  * Lin et al. [2016] L. Lin, Y. Saad, and C. Yang.  Approximating spectral densities of large matrices.  _SIAM review_ , 58(1):34–65, 2016.
  * Liu et al. [2023] 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.
  * Liu et al. [2019] 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 arXiv:1908.03265_ , 2019.
  * Liu et al. [2020] L. Liu, X. Liu, J. Gao, W. Chen, and J. Han.  Understanding the difficulty of training transformers.  _arXiv preprint arXiv:2004.08249_ , 2020.
  * Loshchilov and Hutter [2017] I. Loshchilov and F. Hutter.  Decoupled weight decay regularization.  _arXiv preprint arXiv:1711.05101_ , 2017.
  * Luo et al. [2018] 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.
  * Martens and Grosse [2015] J. Martens and R. Grosse.  Optimizing neural networks with kronecker-factored approximate curvature.  In _International conference on machine learning_ , pages 2408–2417. PMLR, 2015.
  * Merrill et al. [2020] 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.09697_ , 2020.
  * Molybog et al. [2023] I. Molybog, P. Albert, M. Chen, Z. DeVito, 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 preprint arXiv:2304.09871_ , 2023.
  * Nesterov [2013] Y. Nesterov.  _Introductory lectures on convex optimization: A basic course_ , volume 87.  Springer Science & Business Media, 2013.
  * Nguyen and Salazar [2019] T. Q. Nguyen and J. Salazar.  Transformers without tears: Improving the normalization of self-attention.  _arXiv preprint arXiv:1910.05895_ , 2019.
  * Noci et al. [2022] 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_ , 35:27198–27211, 2022.
  * Pan and Li [2023] Y. Pan and Y. Li.  Toward understanding why adam converges faster than sgd for transformers.  _arXiv preprint arXiv:2306.00204_ , 2023.
  * Papyan [2018] V. Papyan.  The full spectrum of deepnet hessians at scale: Dynamics with sgd training and sample size.  _arXiv preprint arXiv:1811.07062_ , 2018.
  * Papyan [2019] V. Papyan.  Measurements of three-level hierarchical structure in the outliers in the spectrum of deepnet hessians.  _arXiv preprint arXiv:1901.08244_ , 2019.
  * Papyan [2020] V. Papyan.  Traces of class/cross-class structure pervade deep learning spectra.  _The Journal of Machine Learning Research_ , 21(1):10197–10260, 2020.
  * Pearlmutter [1994] B. A. Pearlmutter.  Fast exact multiplication by the hessian.  _Neural computation_ , 6(1):147–160, 1994.
  * Radford et al. [2019] A. Radford, J. Wu, R. Child, D. Luan, D. Amodei, I. Sutskever, et al.  Language models are unsupervised multitask learners.  _OpenAI blog_ , 1(8):9, 2019.
  * Rasley et al. [2020] J. Rasley, S. Rajbhandari, O. Ruwase, and Y. He.  Deepspeed: System optimizations enable training deep learning models with over 100 billion parameters.  In _Proceedings of the 26th ACM SIGKDD International Conference on Knowledge Discovery & Data Mining_, pages 3505–3506, 2020.
  * Reddi et al. [2018] S. J. Reddi, S. Kale, and S. Kumar.  On the convergence of adam and beyond.  In _International Conference on Learning Representations_ , 2018.
  * Roux et al. [2007] N. Roux, P.-A. Manzagol, and Y. Bengio.  Topmoumoute online natural gradient algorithm.  _Advances in neural information processing systems_ , 20, 2007.
  * Saad [2011] Y. Saad.  _Numerical methods for large eigenvalue problems: revised edition_.  SIAM, 2011.
  * Sagun et al. [2016] L. Sagun, L. Bottou, and Y. LeCun.  Eigenvalues of the hessian in deep learning: Singularity and beyond.  _arXiv preprint arXiv:1611.07476_ , 2016.
  * Sagun et al. [2017] L. Sagun, U. Evci, V. U. Guney, Y. Dauphin, and L. Bottou.  Empirical analysis of the hessian of over-parametrized neural networks.  _arXiv preprint arXiv:1706.04454_ , 2017.
  * Sankar et al. [2021] 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 to regularization.  In _Proceedings of the AAAI Conference on Artificial Intelligence_ , volume 35, pages 9481–9488, 2021.
  * Shi et al. [2020] N. Shi, D. Li, M. Hong, and R. Sun.  Rmsprop converges with proper hyper-parameter.  In _International Conference on Learning Representations_ , 2020.
  * Simonyan and Zisserman [2014] K. Simonyan and A. Zisserman.  Very deep convolutional networks for large-scale image recognition.  _arXiv preprint arXiv:1409.1556_ , 2014.
  * Sun [2019] R. Sun.  Optimization for deep learning: theory and algorithms.  _arXiv preprint arXiv:1912.08957_ , 2019.
  * Sun and Ye [2021] R. Sun and Y. Ye.  Worst-case complexity of cyclic coordinate descent: O (n^ 2) o (n 2) gap with randomized version.  _Mathematical Programming_ , 185:487–520, 2021.
  * Ubaru et al. [2017] S. Ubaru, J. Chen, and Y. Saad.  Fast estimation of tr(f(a)) via stochastic lanczos quadrature.  _SIAM Journal on Matrix Analysis and Applications_ , 38(4):1075–1099, 2017.
  * Vaswani et al. [2017] A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, Ł. Kaiser, and I. Polosukhin.  Attention is all you need.  _Advances in neural information processing systems_ , 30, 2017.
  * Wang et al. [2022a] B. Wang, Y. Zhang, H. Zhang, Q. Meng, Z.-M. Ma, T.-Y. Liu, and W. Chen.  Provable adaptivity in adam.  _arXiv preprint arXiv:2208.09900_ , 2022a.
  * Wang et al. [2023a] B. Wang, J. Fu, H. Zhang, N. Zheng, and W. Chen.  Closing the gap between the upper bound and lower bound of adam’s iteration complexity.  _Advances in Neural Information Processing Systems_ , 36, 2023a.
  * Wang et al. [2023b] 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, 2023b.
  * Wang et al. [2022b] H. Wang, S. Ma, L. Dong, S. Huang, D. Zhang, and F. Wei.  Deepnet: Scaling transformers to 1,000 layers.  _arXiv preprint arXiv:2203.00555_ , 2022b.
  * Wang et al. [2019] Q. Wang, B. Li, T. Xiao, J. Zhu, C. Li, D. F. Wong, and L. S. Chao.  Learning deep transformer models for machine translation.  _arXiv preprint arXiv:1906.01787_ , 2019.
  * Wikipedia [2023] Wikipedia.  Gaussian quadrature — Wikipedia, the free encyclopedia, 2023.  URL https://en.wikipedia.org/w/index.php?title=Gaussian_quadrature&oldid=1191539517.  [Online; accessed 20-January-2024].
  * Wortsman et al. [2023] M. Wortsman, P. J. Liu, L. Xiao, K. Everett, A. Alemi, B. Adlam, J. D. Co-Reyes, I. Gur, A. Kumar, R. Novak, et al.  Small-scale proxies for large-scale transformer training instabilities.  _arXiv preprint arXiv:2309.14322_ , 2023.
  * Wu et al. [2020] 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.04261_ , 2020.
  * Xie et al. [2022] 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 24430–24459. PMLR, 2022.
  * Xiong et al. [2020] 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.
  * Yang et al. [2023] A. Yang, B. Xiao, B. Wang, B. Zhang, C. Bian, C. Yin, C. Lv, D. Pan, D. Wang, D. Yan, et al.  Baichuan 2: Open large-scale language models.  _arXiv preprint arXiv:2309.10305_ , 2023.
  * Yang et al. [2022] G. Yang, E. J. Hu, I. Babuschkin, S. Sidor, X. Liu, D. Farhi, N. Ryder, J. Pachocki, W. Chen, and J. Gao.  Tensor programs v: Tuning large neural networks via zero-shot hyperparameter transfer.  _arXiv preprint arXiv:2203.03466_ , 2022.
  * Yao et al. [2018] 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.
  * Yao et al. [2020] 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.
  * Zaheer et al. [2018] M. Zaheer, S. Reddi, D. Sachan, S. Kale, and S. Kumar.  Adaptive methods for nonconvex optimization.  _Advances in neural information processing systems_ , 31, 2018.
  * Zeng et al. [2022] A. Zeng, X. Liu, Z. Du, Z. Wang, H. Lai, M. Ding, Z. Yang, Y. Xu, W. Zheng, X. Xia, et al.  Glm-130b: An open bilingual pre-trained model.  _arXiv preprint arXiv:2210.02414_ , 2022.
  * Zhai et al. [2023] 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.
  * Zhang et al. [2019a] B. Zhang, I. Titov, and R. Sennrich.  Improving deep transformer with depth-scaled initialization and merged attention.  _arXiv preprint arXiv:1908.11365_ , 2019a.
  * Zhang et al. [2019b] G. Zhang, L. Li, Z. Nado, J. Martens, S. Sachdeva, G. Dahl, C. Shallue, and R. B. Grosse.  Which algorithmic choices matter at which batch sizes? insights from a noisy quadratic model.  _Advances in neural information processing systems_ , 32, 2019b.
  * Zhang et al. [2019c] J. Zhang, T. He, S. Sra, and A. Jadbabaie.  Why gradient clipping accelerates training: A theoretical justification for adaptivity.  _arXiv preprint arXiv:1905.11881_ , 2019c.
  * Zhang et al. [2020] 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.
  * Zhang et al. [2022a] 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.01068_ , 2022a.
  * Zhang et al. [2022b] Y. Zhang, C. Chen, N. Shi, R. Sun, and Z.-Q. Luo.  Adam can converge without any modification on update rules.  _Advances in Neural Information Processing Systems_ , 35:28386–28399, 2022b.
  * Zhou et al. [2018] 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:1808.05671_ , 2018.
  * Zou et al. [2019] F. Zou, L. Shen, Z. Jie, W. Zhang, and W. Liu.  A sufficient condition for convergences of adam and rmsprop.  In _Proceedings of the IEEE/CVF Conference on computer vision and pattern recognition_ , pages 11127–11135, 2019.

##  Appendix A More Related Works

#### On the difficulties of Transformer training.

Transformers are known to be difficult to train. Researchers have attributed
the training difficulties to various phenomena in different components of
Transformers, including: the logits divergence or the rank degeneracy in the
outputs of attention layers [Dong et al., 2021, Noci et al., 2022, Wortsman et
al., 2023, Zhai et al., 2023, Dehghani et al., 2023, Chowdhery et al., 2023];
the growth of parameter norm in attention layers [Merrill et al., 2020]; over-
reliance on residue branches [Liu et al., 2020]; and some negative impact of
layer norm [Chen et al., 2018, Zhang et al., 2019a, Huang et al., 2020]. These
phenomena have a strong correlation with gradient vanishing or explosion in
Transformers [Zhang et al., 2019a, Liu et al., 2020, Huang et al., 2020, Xiong
et al., 2020, Noci et al., 2022, Wang et al., 2022b, Wortsman et al., 2023,
Molybog et al., 2023], which leads to training difficulties.

Several solutions have been proposed. Liu et al. [2020] 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 in principle. Researchers further
develop training tricks such as warmup learning rate [Liu et al., 2019, Xiong
et al., 2020], temperature scaling [Noci et al., 2022], better initialization
[Zhang et al., 2019a, Huang et al., 2020, Wang et al., 2022b, Bachlechner et
al., 2021, Yang et al., 2022], and variants of Layer Norm [Nguyen and Salazar,
2019, Wang et al., 2019, Xiong et al., 2020, Wang et al., 2022b, Dehghani et
al., 2023]. Recent researchers also suggest using z-loss regularization
[Chowdhery et al., 2023, Yang et al., 2023] and tuning hyperparameters of Adam
[Zhang et al., 2022b, Wortsman et al., 2023]. All these tricks can help
mitigate gradient explosion or vanishing. Nevertheless, training large-scale
Transformers remains challenging [Zhang et al., 2022a, Zeng et al., 2022,
Wortsman et al., 2023, Molybog et al., 2023, Chowdhery et al., 2023].
Different from all aforementioned works, we investigate the training
difficulties of Transformers through the eigenvalues of Hessian. We establish
a strong correlation between the blockwise Hessian spectra of Transformers and
the failure of SGD. We realize that our attempt is just a first step towards
understanding Transformer training, and we believe there is rich information
hidden in Hessian and we leave more fine-grained analysis as future works.

#### Convergence analysis of adaptive gradient methods

There is extensive convergence analysis for adaptive gradient methods. For
instance, researchers study the convergence of AMSGrad [Reddi et al., 2018,
Zhou et al., 2018], RMSprop [Zaheer et al., 2018], AdaFom [Chen et al., 2019],
AdaBound [Luo et al., 2018], and Adam with iterate-dependent hyperparameters
[Zou et al., 2019, Chen et al., 2022, Gadat and Gavra, 2022]. The convergence
of Adam is also explored in [Défossez et al., 2022, Wang et al., 2023a]. There
is also an active line of theoretical research on the convergence of AdaGrad
[Duchi et al., 2011], we recommend [Wang et al., 2023b] 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.

##  Appendix B More Preliminaries

###  B.1 Preliminaries on Optimizers

Here we provide a detailed description of the optimizers mentioned in the full
script. We consider the minimizing
$\mathcal{L}(w)\equiv\frac{1}{n}\sum_{i=1}^{n}\mathcal{L}_{i}(w)$, where $n$
is the number of minibatches, $\mathcal{L}_{i}(w)$ is the loss of $i$-th
minibatch and $w\in\mathbb{R}^{d}$ is the neural network parameters. We denote
the gradient of the training loss w.r.t. neural network parameters as
$\nabla\mathcal{L}(w)\in\mathbb{R}^{d}$. We use
$\nabla\mathcal{L}_{i}(w)\in\mathbb{R}^{d}$ to denote the $i$-th minibatch
counterparts. We use $w^{t}$ to denote the variable at the $t$-th step. In
Algorithm 2 and 3, $\circ$, division and square-root are elementwise
operations. In the line 7 and 8 of Algorithm 2, $(\beta_{1})^{t}$ and
$(\beta_{2})^{t}$ indicates the $t$-th power of $\beta_{1}$, $\beta_{2}$. In
the PyTorch default setting,
$(\beta_{1},\beta_{2},\epsilon)=(0.9,0.999,\text{1e-8})$ for Adam and
$\beta_{1}=0.9$ for SGD.

Algorithm 1 Stochastic Gradient Descent with Momentum (SGD)

1: Initialize $w^{0}$ and choose $0\leq\beta_{1}<1$ and $\eta_{0}>0$

2: for $t=1\to\infty$ do

3: Uniformly sample $\tau^{t}$ from the index set $\\{1,2,\cdots,n\\}$

4: $m^{t}=\beta_{1}m^{t}+\nabla f_{\tau^{t}}(x^{t})$

5: $x^{t+1}=x^{t}-\eta_{t}m^{t}$

6: end for

Algorithm 2 AdamW

1: Initialize $x^{0}$, $m^{0}=v^{0}=0$, $0\leq\beta_{1}<1$,
$0\leq\beta_{2}<1$, $\epsilon>0$, $\eta^{0}>0$, and weight decay coefficient
$\lambda$

2: for $t=1\to\infty$ do

3: Uniformly sample $\tau^{t}$ from the index set $\\{1,2,\cdots,n\\}$

4: $w^{t+1}=w^{t}-\eta^{t}\lambda w^{t}$

5: $m^{t}=\beta_{1}m^{t}+(1-\beta_{1})\nabla\mathcal{L}_{\tau^{t}}(w^{t})$

6:
$v^{t}=\beta_{2}v^{t}+(1-\beta_{2})\nabla\mathcal{L}_{\tau^{t}}(w^{t})\circ\nabla\mathcal{L}_{\tau^{t}}(w^{t})$

7: $\hat{m}^{t}=\frac{m^{t}}{1-(\beta_{1})^{t}}$

8: $\hat{v}^{t}=\frac{v^{t}}{1-(\beta_{2})^{t}}$

9: $w^{t+1}=w^{t+1}-\eta_{t}\frac{\hat{m}^{t}}{\sqrt{\hat{v}^{t}}+\epsilon}$

10: end for

Algorithm 3 Adam with no bias correction

1: Initialize $x^{0}$,
$m^{0}=\nabla\mathcal{L}_{\tau^{t}}(w^{0}),v^{0}=\nabla\mathcal{L}_{\tau^{t}}(w^{0})\circ\nabla\mathcal{L}_{\tau^{t}}(w^{0})$,
$0\leq\beta_{1}<1$, $0\leq\beta_{2}<1$, $\epsilon>0$, $\eta^{0}>0$

2: for $t=1\to\infty$ do

3: Uniformly sample $\tau^{t}$ from the index set $\\{1,2,\cdots,n\\}$

4: $m^{t}=\beta_{1}m^{t}+(1-\beta_{1})\nabla\mathcal{L}_{\tau^{t}}(w^{t})$

5:
$v^{t}=\beta_{2}v^{t}+(1-\beta_{2})\nabla\mathcal{L}_{\tau^{t}}(w^{t})\circ\nabla\mathcal{L}_{\tau^{t}}(w^{t})$

6: $w^{t+1}=w^{t+1}-\eta_{t}\frac{m^{t}}{\sqrt{v^{t}}+\epsilon}$

7: end for

###  B.2 Preliminaries on the Stochastic Lanczos Quadrature Method

#### Additional notations.

Given a real symmetric matrix $H\in\mathbb{R}^{d\times d}$, we denote $tr(H)$
as its trace and $Q^{T}\Lambda Q$ as its spectral decomposition, where
$Q=\left[q_{1},\ldots,q_{d}\right],\Lambda=\operatorname{diag}\left(\lambda_{1},\ldots,\lambda_{d}\right)$
and $\lambda_{1}\geq\lambda_{2}\cdots\geq\lambda_{d}$. We denote the condition
number of $H$ as $\kappa=\lambda_{1}/\lambda_{d}$. We define matrix function
as $f(H):=Q^{T}f(\Lambda)Q$, where
$f(\Lambda)=\operatorname{diag}\left(f\left(\lambda_{1}\right),\ldots
f\left(\lambda_{d}\right)\right)\in\mathbb{R}^{d\times d}$. We use
$\mathbb{N}$ to denote the set of positive integers. We use $\|\cdot\|_{2}$ to
denote the Euclidean norm.

Approximation of the Hessian spectrum can be formulated as a trace estimation
problem, as introduced in [Lin et al., 2016, Ubaru et al., 2017]. First, the
spectrum (eigenvalue density) of Hessian $H$ can written as:
$\phi(t)=\frac{1}{d}\sum_{i=1}^{d}\delta\left(t-\lambda_{i}\right)$, where
$\lambda_{i}$ are the eigenvalues of $H$ and $\delta$ is the Dirac
$\delta$-function. Then, we replace the delta functions by a Gaussian blurring
function: $\phi(t)\approx
g(t):=\frac{1}{d}\sum_{i=1}^{d}f\left(\lambda_{i}\right)$, where
$f(\lambda):=\frac{1}{\sigma\sqrt{2\pi}}\exp\left(-\frac{(t-\lambda)^{2}}{2\sigma^{2}}\right)$.
By definition of matrix function, it is easy to see that
$g(t)=\frac{1}{d}tr(f(H))$. As such, spectrum approximation could be
formulated as a trace estimation problem, i.e., estimating
$\frac{1}{d}\operatorname{tr}(f(H))$, where $H\in\mathbb{R}^{d\times d}$ is a
real symmetric matrix.

Trace estimation problems could be solved efficiently by the Stochastic
Lanczos Quadrature Method (SLQ) [Golub and Strakoš, 1994]. For the ease of
readers, we provide a detailed description of SLQ in our context. SLQ consists
of the following steps.

#### Step 1.

We Approximate the trace of matrix function as
$\frac{1}{d}tr(f(H))=\mathbb{E}(v^{T}f(H)v)\approx\frac{1}{n_{v}}\sum_{i}^{n_{v}}v_{i}^{T}f(H)v_{i}$,
where $v=u/\|u\|_{2}$ and $u$ is a Rademacher random vector (each entry of $u$
independently takes $\pm 1$ with probability $1/2$). This step is called
Huchinson’s estimation [Hutchinson, 1989].

Note that we can also replace the Rademacher random vector $u$ by a unit
Gaussian vector (i.e., $u\sim N(0,I_{d\times d})$) and the unbiasedness still
holds [Avron and Toledo, 2011]. In our implementation, we sample $u\sim
N(0,I_{d\times d})$ because there is an efficient built-in PyTorch function
for generating Gaussian vectors.

SLQ estimates $v_{i}^{T}f(H)v_{i}$ for $i\in[n_{v}]$ and then take the
average. To understand SLQ, we only need to understand how it estimates each
individual quadratic form. To simplify the notation regarding $i$, from now
on, we will discuss how to estimate $v^{T}f(H)v$, where $v=u/\|u\|_{2}$ and
$u$ is a unit Gaussian vector.

#### Step 2-1.

We rewrite $v^{T}f(H)v$ as a Riemann-Stieltjes integral [Golub and Meurant,
2009]:

$v^{T}f(A)v=\sum_{i=1}^{d}\left(v^{T}q_{i}\right)^{2}f\left(\lambda_{i}\right)=\int_{\lambda_{d}}^{\lambda_{1}}f(\lambda)d\mu(\lambda),$
(7)

where $\mu$ is a measure on $(\mathbb{R},\mathbb{B})$ defined as follows
($\mu(\lambda)$ denotes the measure of set $\\{x;x\leq\lambda\\}$):

$\mu(\lambda)=\begin{cases}0&\lambda<\lambda_{d}\\\
\sum_{i=1}^{k}\left(v^{T}q_{i}\right)^{2}&\lambda_{k}\leq\lambda<\lambda_{k+1}\\\
\sum_{i=1}^{d}\left(v^{T}q_{i}\right)^{2}&\lambda\geq\lambda_{1}\end{cases}.$
(8)

#### Step 2-2.

Unfortunately, this integral is difficult to compute. This is because the
measure $\mu$ are related to the eigen-pairs of $H$, which are unknown. It
seems unclear how to directly integrate over an unknown measure. As such, we
further approximate this integral by a computationally friendly quantity, such
as:

$\int_{\lambda_{d}}^{\lambda_{1}}f(\lambda)d\mu(\lambda)\approx\sum_{j=1}^{m}c_{j}f(x_{j}).$
(9)

We hope to design $\\{(c_{j},x_{j})\\}_{j=1}^{m}$ with a reasonable number of
$m$ such that the estimation error is small. Fortunately, the Gaussian
Quadrature method provides a generic design principle of
$\\{(c_{j},x_{j})\\}_{j=1}^{m}$ [Golub and Meurant, 2009, Epperson, 2013]. It
is proved that: when $f(\lambda)$ is not "too complicated" (e.g. $f(\lambda)$
is a polynomial), then there exists $\\{(c_{j},x_{j})\\}_{j=1}^{m}$ which
gives a high quality estimation of integral (7). The required number of $m$ is
related to "how complicated the $f(\lambda)$ is". Such
$\\{(c_{j},x_{j})\\}_{j=1}^{m}$ are called the Gaussian Quadrature rules.
$c_{j}$ and $x_{j}$ are called the "weights" and the "nodes" of the Gaussian
Quadrature rules. A representative theorem is as follows: when $f(\lambda)$ is
a polynomial with degree $<2m$, then the Gaussian Quadrature rules give the
exact approximation of integral (7).

######  Theorem 3.

[Rewrited based on [Wikipedia, 2023]] Suppose we have a sequence of orthogonal
polynomials $\left\\{p_{k}(x)\right\\}_{k=1}^{m}$ w.r.t. measure $\mu$, that
is: $\int_{\lambda_{d}}^{\lambda_{1}}p_{n}(x)p_{m}(x)d\mu(x)=\delta_{m,n}$,
where $\delta_{m,n}=1$ if $m=n$ and $\delta_{m,n}=0$, otherwise. Assume $f(x)$
is a polynomial with degree $<2m$, then there exists
$\left\\{\left(c_{j},x_{j}\right)\right\\}_{j=1}^{m}$ s.t.
$\int_{\lambda_{d}}^{\lambda_{1}}f(\lambda)d\mu(\lambda)=\sum_{i=j}^{m}c_{j}f\left(x_{j}\right)$.
The equality holds when $x_{j}$ are the roots of $p_{m}(x)$ and
$c_{j}=\int_{\lambda_{d}}^{\lambda_{1}}\prod_{j\neq
i}\frac{x-x_{i}}{x_{j}-x_{i}}d\mu$. Such choice of
$\left\\{\left(c_{j},x_{j}\right)\right\\}_{j=1}^{m}$ are called the Gaussian
Quadrature rules.

Theorem 3 shows the existence of good
$\left\\{\left(c_{j},x_{j}\right)\right\\}_{j=1}^{m}$ and their general form.
In fact, it is also shown that Gaussian Quadrature is optimal: no
other$\left\\{\left(c_{j},x_{j}\right)\right\\}_{j=1}^{m}$ can achieve zero
approximation error for higher degree polynomials $f(\lambda)$ [Golub and
Meurant, 2009]. However, it is often difficult to find these quadrature rules
[Golub and Welsch, 1969]. There are at least three questions in sequel:

  * •

1) how to find the orthogonal polynomials
$\left\\{p_{k}(x)\right\\}_{k=1}^{m}$ w.r.t. an unknown measure $\mu$?

  * •

2) how to efficiently find the roots of $p_{m}(x)$, which gives the nodes
$x_{j}$?

  * •

3) how to efficiently calculate the weights
$c_{j}=\int_{\lambda_{d}}^{\lambda_{1}}\prod_{j\neq
i}\frac{x-x_{i}}{x_{j}-x_{i}}d\mu$?

We first answer question 2) and 3) and leave question 1) for later discussion.

Now suppose that we have found the orthogonal polynomials
$\left\\{p_{k}(x)\right\\}_{k=1}^{m}$ w.r.t. $\mu$. Recall that any orthogonal
polynomial has the following "three-term" recursion [Golub and Meurant, 2009]:

$p_{k+1}(x)=\left(x-\alpha_{k+1}\right)p_{k}(x)-\beta_{k}p_{k-1}(x),k=0,1,\ldots,$

where $p_{-1}(x)\equiv 0,p_{0}(x)\equiv 1,\alpha_{k+1}=\frac{\left\langle
xp_{k},p_{k}\right\rangle}{\left\langle p_{k},p_{k}\right\rangle}$ and
$\beta_{k}=\frac{\left\langle p_{k},p_{k}\right\rangle}{\left\langle
p_{k-1},p_{k-1}\right\rangle}$. Define
$P_{m}(x)=\left(p_{0}(x),p_{1}(x),\ldots
p_{m-1}(x)\right)^{T}\in\mathbb{R}^{m}$, we can rewrite the recursion formula
in matrix form (given $x$): $xP_{m}=J_{m}P_{m}+\beta_{m}p_{m}(x)e^{m}$, where
$e^{m}$ is the last column of identity matrix $I_{m,m}$ and $J_{m}$ is called
Jacobi matrix of order $m$:

$J_{m}=\left(\begin{array}[]{ccccc}\alpha_{1}&\sqrt{\beta_{1}}&&&\\\
\sqrt{\beta_{1}}&\alpha_{2}&\sqrt{\beta_{2}}&&\\\
&\sqrt{\beta_{2}}&\alpha_{3}&\sqrt{\beta_{3}}&\\\
&&\ddots&\ddots&\ddots\end{array}\right)\in\mathbb{R}^{m\times m}$

It turns out that $J_{m}$ can help us find the Gaussian Quadrature rules
$\left\\{\left(c_{j},x_{j}\right)\right\\}_{j=1}^{m}$ and thus provide answers
for question 2) and 3). This is shown in the following theorem.

######  Theorem 4.

[Golub and Meurant, 2009] For the Gaussian Quadrature, $\\{x_{j}\\}_{j=1}^{m}$
are the eigenvalues of $J_{m}$ and $\\{c_{j}\\}_{j=1}^{m}$ are the squares of
the first elements of the normalized eigenvectors of $J_{m}$.

The proof of Theorem 4 is based on Christoffel-Darboux relation [Brezinski,
1990]. Now, the remaining question is: how to find the Jacobian matrix $J_{m}$
of a sequence of orthogonal polynomials w.r.t. an unknown measure $\mu$? Note
that we no longer need to answer question 1) if $J_{m}$ is found, since
$J_{m}$ is sufficient for us to find the Gaussian quadrature rules. However,
it seems impossible to find $J_{m}$ if no information of $\mu$ is provided.
The good news is: when the $\mu$ is specified as in (8), there exists an
efficient way to find $J_{m}$.

#### Step 3.

When $\mu$ is specified as in (8), $J_{m}$ can be exactly found in $m$ steps
using the Lanczos algorithm [Lanczos, 1950], 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 Lanczos
algorithm) is called the Lanczos quadrature [Golub and Strakoš, 1994, Bai and
Golub, 1996, Golub and Meurant, 2009]. An extremely elegant but highly
nontrivial result is as follows:

######  Theorem 5.

[Golub and Meurant, 2009] Given a real symmetric matrix
$H\in\mathbb{R}^{d\times d}$ and an arbitrary vector $v\in\mathbb{R}^{d}$ with
unit Euclidean norm, we define the measure $\mu$ as in (8) based on this $H$
and $v$. Then $m$ steps of the Lanzcos algorithm return the Jacobian matrix
$J_{m}$ of orthogonal polynomials w.r.t. to $\mu$.

After $J_{m}$ is found by the Lanczos algorithm, we perform spectral
decomposition of $J_{m}\in\mathbb{R}^{m\times m}$ to get its eigen-pairs.
Using Theorem 4, we successfully get the Gaussian quadrature rules and thus we
can approximate the quadratic form $v^{T}f(H)v$. By averaging over different
random vectors $v$ we can then approximate $\frac{1}{d}tr(f(H))$. This
concludes the derivation of SLQ for the trace estimation problem.

The full procedure of SLQ is shown in Algorithm 5. We note that SLQ is
efficient in theory. Ubaru et al. [2017] show that SLQ converges faster than
any other polynomial expansion method for spectrum estimation (e.g., Chebyshev
methods used in [Adams et al., 2018]). See [Ubaru et al., 2017, Theorem 4.1]
for a formal statement.

We remark that there are at least four versions of the Lanczos algorithm in
Step 3. Here, we adopt the version in Algorithm 4 since it is known to be the
most numerically stable version [Cullum and Willoughby, 2002, Saad, 2011,
Wikipedia, 2023]. Throughout this work, we choose $f\left(\cdot\right)$ as the
Gaussian blurring function
$f(\lambda):=\frac{1}{\sigma\sqrt{2\pi}}\exp\left(-\frac{(t-\lambda)^{2}}{2\sigma^{2}}\right)$
for spectrum approximation. We plot the spectrum by sweeping $t$ from the
minimal node to the maximal node in Gaussian Quadrature rules.

Algorithm 4 The Lanczos Algorithm

1: Input a matrix-vector product $Hv_{1}\in\mathbb{R}^{d}$, where $H$ is a
real symmetric matrix and $v_{1}$ is an arbitrary vector with Euclidean norm
1. Choose $m\in\mathbb{N}$

2: Initialization: Let $w_{1}^{\prime}=Hv_{1}$,
$\alpha_{1}=(w_{1}^{\prime})^{T}v_{1}$, $w_{1}=w_{1}^{\prime}-\alpha_{1}v_{1}$

3: for $j=2\to m$ do

4: Let $\beta_{j}=\left\|w_{j-1}\right\|_{2}$ (also Euclidean norm)

5: If $\beta_{j}\neq 0$, then let $v_{j}=w_{j-1}/\beta_{j}$, else pick as
$v_{j}$ an arbitrary vector with Euclidean norm 1 that is orthogonal to all of
$v_{1},\ldots,v_{j-1}$

6: Let $w_{j}^{\prime}=Av_{j}$

7: Let $\alpha_{j}=(w_{j}^{\prime})^{T}v_{j}$

8: Let $w_{j}=w_{j}^{\prime}-\alpha_{j}v_{j}-\beta_{j}v_{j-1}$

9: end for

10: Let $V$ be the matrix with columns $v_{1},\ldots,v_{m}$

11: Let $T=\left(\begin{array}[]{cccccc}\alpha_{1}&\beta_{2}&&&&0\\\
\beta_{2}&\alpha_{2}&\beta_{3}&&&\\\ &\beta_{3}&\alpha_{3}&\ddots&&\\\
&&\ddots&\ddots&\beta_{m-1}&\\\ &&&\beta_{m-1}&\alpha_{m-1}&\beta_{m}\\\
0&&&&\beta_{m}&\alpha_{m}\end{array}\right)$

12: Return $T$

Algorithm 5 The Stochastic Lanczos Quadrature Method

1: Choose $\texttt{num}_{v},m\in\mathbb{N}$. Sample $\texttt{num}_{v}$ i.i.d.
$v_{i}$ from normalized Rademacher distribution, $i\in[\texttt{num}_{v}]$

2: for $i=1\to\texttt{num}_{v}$ do

3: Run $m$ steps of the Lanczos Algorithm 4 with input $Hv_{i}$, returns
$T\in\mathbb{R}^{m\times m}$

4: Compute eigenvalue decomposition $T=Q\Lambda Q^{T}$

5: Compute the nodes $x_{i}=\left(\Lambda_{ii}\right)_{i=1}^{m}$ and weights
$c_{i}=\left(Q_{1,i}^{2}\right)_{i=1}^{m}$

6: Return $q_{i}(t)=\sum_{i=1}^{m}c_{i}f\left(x_{i};t,\sigma^{2}\right)$

7: end for

8: Return
$\frac{1}{\texttt{num}_{v}}\sum_{i=1}^{\texttt{num}_{v}}f\left(\ell_{i};t,\sigma^{2}\right)$

##  Appendix C More Eperimental Details

###  C.1 Implementation Details on SLQ

#### Implementation and Running Time Analysis.

We provide a simple PyTorch implementation of SLQ. The only query SLQ makes to
the neural network is the Hessian vector product, which is attained using the
auto-differentiation framework [Pearlmutter, 1994]. 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 [Dao et al., 2022] does not support the
calculation of Hessian-vector product, we implement all attention blocks in
the naive way. For the calculation of the blockwise Hessian spectrum
$\nabla^{2}\mathcal{L}(w_{l})$, we sample $u_{l}\sim N(0,I_{d_{l}\times
d_{l}})$ and set $v_{l}=u_{l}/\|u_{l}\|_{2}\in\mathbb{R}^{d_{l}}$. Then we run
Algorithm 5 by taking $\nabla^{2}\mathcal{L}(w_{l})$ and $v_{l}$ as inputs. We
choose the hyperparameters as $m=100$ and $n_{v}=10$ in all experiments.
$\sigma$ is tuned based on visual effects. These hyperparameters are reported
to reach highly accurate estimation with error $<10^{-14}$ [Ghorbani et al.,
2019].

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
$\nabla^{2}\mathcal{L}(w)d=\frac{1}{n}\sum_{i=1}^{n}\nabla^{2}\mathcal{L}_{i}(w)d$,
so each Hessian-vector product operation requires (i) calculating
$\nabla^{2}\mathcal{L}_{i}(w)d$; (ii) repeating (i) on all data. We point out
that (i) can be computed efficiently and precisely with just two
backpropagation passes [Pearlmutter, 1994]. The major computational bottleneck
lies in (ii) due to the large $n$. Our largest-scale experiment for Hesian
spectrum is GPT2 (125M) on Openwebtext, where the number of tokens $n=9$
Billon. To calculate $\nabla^{2}\mathcal{L}(w)d$ on all these 9B tokens, it
requires about 9 GPU days on eight A100-80GB GPUs. Since SLQ requires at least
1,000 times query of $\nabla^{2}\mathcal{L}(w)d$, a complete run of SLQ would
take at least 9,000 days on eight A100-80GB GPUs, which is unaffordable. In
this work, we use the largest possible batch size (with gradient accumulation
tricks) to approximate $\nabla^{2}\mathcal{L}(w)$ under the constraints of GPU
bandwidth and time limit. More detailed setup of SLQ are shown as follows.

  * •

ResNet18 (20M) and VGG16 (138M) on ImageNet. We use the code base of PyTorch
Examples 444https://github.com/pytorch/examples/blob/main/imagenet/main.py. 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
555https://github.com/huggingface/pytorch-image-models. We use batch size
$=1024$. 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, $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.

  * •

BERT(40M) on Cornell Movie-Dialogs Corpus. We use the code base from the blog
666https://medium.com/data-and-beyond/complete-guide-to-building-bert-model-
from-sratch-3e6562228891. We use batch size $=327,680$ tokens. For the
calculation of the blockwise Hessian spectra, we apply SLQ to 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
777https://github.com/karpathy/nanoGPT/. 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 MLP
blocks (note that the $12$-th attention and MLP blocks are the final ones). In
total, it takes about 7 days on one A100 GPU to estimate all the blockwise
Hessian spectra and the full Hessian spectrum.

We train all the models under the default configurations in all cases. For CNN
training, we follow the rule of thumb to choose the learning rate of AdamW: we
set the learning rate of AdamW to be $100\times$ smaller than the recommended
learning rate of SGD.

###  C.2 Implementation Details on Figure LABEL:fig_block_diagnal

We use the code base of NanoGPT to train a decoder-only Transformers on 4
consecutive tokens randomly selected from Openwebtext. We set the model
configuration as: context window = 2, number of heads = 2 and the embedding
dimension = 4. In MLP layers, all widths equal to 4. We change the number of
layers from 1 to 4. We remove all the Layer Norms in the model. The rest of
the model configurations are set to their default values in the code base. We
compute the Hessian on all the parameters blocks in attention and MLP layers.
In Figure LABEL:fig_block_diagnal (a), the variables in three blue boxes
corresponds to the parameters in Query and Key; Value and Projection; and MLP,
respectively. Similarly for the rest of the figures.

Due to the intensive overhead of computing and storing the whole Hessian, we
have yet to check the block-diagonal structure on larger models. Rigorously
speaking, so far we have not gotten sufficient evidence to claim this
structure commonly holds in larger Transformers. It requires new numerical
methods to efficiently check the block-diagonal Hessian structure in larger
models without explicitly calculating them. We leave it as an interesting
future direction.

For completeness, we remark that Collobert [2004], Roux et al. [2007], Martens
and Grosse [2015] also observed the block-diagonal structure in (approximated)
Hessian of small-scaled MLPs. We restate their findings in Figure
LABEL:fig_block_diagnal_existing. [Collobert, 2004, Section 7] further
theoretically proved that the block diagonal structure stems from (i) the
layer-by-layer structure and (ii) the Cross-Entropy loss. These results
suggest that block-diagonal structure might be common in NNs.

###  C.3 Implementation Details on the MLP experiments in Figure
LABEL:fig_mlp_gap

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
$c\in\mathbb{N}$. We scale $c$ from 1 to 15. We record the average JS distance
between the initial Hessian spectra of each pair of layers and find that the
average JS distance increases as $c$ increases. As such, we define $c$ as the
degree of block heterogeneity. For each $c$, we train SGD and Adam with
default hyperparameters by grid-searching the learning rate from 1e-4 to 1e-1
and report the best test accuracy after 1 epoch.

$c$ | 1 | 2 | 3 | 4 | 5 | 6 | 7 | 8 | 9-15
---|---|---|---|---|---|---|---|---|---
Blockwise distance | 1e-3 | 1.61 | 33.99 | 65.32 | 260.77 | 298.39 | 351.69 | 353.48 | $\operatorname{NaN}$
Table 3: Consider a 4-layer MLP on MNIST. $c$ denotes the scaler that is
multiplied to the output of each layer. Blockwise distance denotes the average
JS distance between the initial Hessian spectra of each pair of layers. We
find that the average JS distance increases as $c$ increases. For $c\geq 9$,
the calculation of JS distance encounters numerical instability due to the
division of near-zero values and returns $\operatorname{NaN}$.

##  Appendix D More Experimental Results

###  D.1 Performance comparison of AdamW and SGD on Different Architectures

Here, we show the performance comparison of AdamW and SGD on different models.

##  Appendix E Proofs

###  E.1 Proof of Proof of Proposition 1

Let $H=\begin{bmatrix}L&0\\\ 0&\mu\end{bmatrix}$, where $L>\mu>0$. We choose
the initial point as
$w^{0}=(w_{1}^{0},w_{2}^{0})=(\sqrt{\mu/L},\sqrt{L/\mu})$. By the update rule
of GD, we have

$\displaystyle\mathcal{L}(w^{t+1})$ $\displaystyle=$
$\displaystyle\mathcal{L}\left(w^{t}-\eta\nabla\mathcal{L}(w^{t})\right)$ (10)
$\displaystyle=$ $\displaystyle\frac{1}{2}(w^{t}-\eta Hw^{t})^{T}H(w^{t}-\eta
Hw^{t})$ $\displaystyle=$ $\displaystyle(w^{t}_{1})^{2}|1-\eta
L|L+(w^{t}_{2})^{2}|1-\eta\mu|\mu$ $\displaystyle=$ $\displaystyle|1-\eta
L|^{t}L\frac{\mu}{L}+|1-\eta\mu|^{t}\mu\frac{L}{\mu}$ $\displaystyle=$
$\displaystyle\mu|1-\eta L|^{t}+L|1-\eta\mu|^{t}$

To proceed, we discuss the following cases:

When $\eta\leq 1/L$, since $|1-\eta L|^{t}$ and $|1-\eta\mu|^{t}$ are
monotonically decreasing , the optimal solution is $\eta=1/L$.

When $\eta\geq 1/\mu$, since $|1-\eta L|^{t}$ and $|1-\eta\mu|^{t}$ are
monotonically increasing , the optimal solution is $\eta=1/\mu$.

When $1/L\leq\eta\leq 1/\mu$, (10) can be written as $g_{t}(\eta)=\mu(\eta
L-1)^{t}+L(1-\eta\mu)^{t}$. Take the first-order and the second-order
derivative of the g, we can obtain
$g_{t}^{\prime}(\eta)=tL\mu(\eta_{L}-1)^{t-1}-t\mu L(1-\eta\mu)^{t-1}$ and
$g_{t}^{\prime\prime}(\eta)=t(t-1)L^{2}\mu(\eta
L-1)^{t-2}+t(t-1)\mu^{2}(1-\eta\mu)$. Since $g^{\prime\prime}_{t}(\eta)\geq 0$
for all $\eta\in[1/L,1\mu]$, the function g is convex. By solving the equation
that $g_{t}^{\prime}(\eta)=0$, we can obtain $\eta=\frac{2}{L+\mu}$ is a
solution for all $t$. Plugging this result into (10) and rearranging the
terms, we conclude the proof of Proposition 1.

###  E.2 Proof of Theorem 1

Without loss of generality, we assume $h=0$. This is because minimizing
$\mathcal{L}(w)=\frac{1}{2}w^{T}Hw-h^{T}w$ is equivalent to minimizing
$\mathcal{L}(w)=\frac{1}{2}\left(w-w^{*}\right)^{T}H\left(w-w^{*}\right)$
where $w^{*}=H^{-1}h$. By a linear transformation $z=w-w^{*}$, BGD for
minimizing $\frac{1}{2}\left(w-w^{*}\right)^{T}H\left(w-w^{*}\right)$ starting
from $w^{0}$ is equivalent to BGD for minimizing $\frac{1}{2}z^{T}Hz$ starting
from $z^{0}=w^{0}-w^{*}$. Thus we can assume $w^{*}=0$, or equivalently,
$h=0$. When $h=0$, the update form of BGD becomes

$w^{t+1}=w^{t}-\eta D_{BGD}^{-1}Hw^{t},$

where
$D_{BGD}=\operatorname{diag}\left(\left(\underbrace{\eta_{1},\cdots,\eta_{1}}_{d_{1}\text{times}},\cdots,\underbrace{\eta_{L},\cdots,\eta_{L}}_{d_{L}\text{times}}\right)^{T}\right)\in\mathbb{R}^{d\times
d}.$

In the following, we will first prove the rate of BGD for each block and then
combine them together to get the total rate. For the $l$-th block, we choose
$\eta_{l}=\frac{1}{\lambda_{l,1}}$. The update rule of BGD for the $l$-th
block becomes:

$w_{l}^{t+1}=w_{l}^{t}-\frac{1}{\lambda_{l,1}}H_{l}w_{l}^{t}$

.

Now we apply the standard convergence proof of GD to the $l$-th block. By
descent lemma, we have

$\displaystyle[\mathcal{L}(w^{t+1})]_{l}$ $\displaystyle\leq$
$\displaystyle[\mathcal{L}(w^{t})]_{l}+\langle
w_{l}^{t+1}-w_{l}^{t},[\nabla\mathcal{L}(w^{t})]_{l}\rangle+\frac{\lambda_{l,1}}{2}\|w_{l}^{t+1}-w_{l}^{t}\|_{2}^{2}$
(11) $\displaystyle\leq$ $\displaystyle[\mathcal{L}(w^{t})]_{l}+\langle
w_{l}^{t+1}-w_{l}^{t},H_{l}w_{l}^{t}\rangle+\frac{\lambda_{l,1}}{2}\eta^{2}w_{l}^{t}H_{l}^{T}H_{l}w_{l}^{t}$
(12) $\displaystyle=$
$\displaystyle[\mathcal{L}(w^{t})]_{l}-\frac{1}{2\lambda_{l,1}}w_{l}^{t}H_{l}^{T}H_{l}w_{l}^{t}$
(13) $\displaystyle\leq$
$\displaystyle[\mathcal{L}(w^{t})]_{l}-\frac{\lambda_{l,d_{l}}}{2\lambda_{l,1}}w_{l}^{t}H_{l}w_{l}^{t}$
(14) $\displaystyle=$
$\displaystyle\left(1-\frac{1}{\kappa_{l}}\right)[\mathcal{L}(w^{t})]_{l}$
(15)

Substract both sides by $\mathcal{L}^{*}$ and we get
$[\mathcal{L}(w^{t+1})]_{l}-[\mathcal{L}^{*}]_{l}\leq\left(1-\frac{1}{\kappa_{l}}\right)\left([\mathcal{L}(w^{t})]_{l}-[\mathcal{L}^{*}]_{l}\right)$.
Summing up both sides over $l\in[L]$,we have

$\displaystyle\mathcal{L}(w^{t+1})-\mathcal{L}^{*}$ $\displaystyle=$
$\displaystyle\sum_{l=1}^{L}[\left(\mathcal{L}(w^{t+1})]_{l}-[\mathcal{L}^{*}]_{l}\right)$
$\displaystyle\leq$
$\displaystyle\sum_{l=1}^{L}\left(1-\frac{1}{\kappa_{l}}\right)\left([\mathcal{L}(w^{t})]_{l}-[\mathcal{L}^{*}]_{l}\right)$
$\displaystyle\leq$
$\displaystyle\max_{l\in[L]}\left(1-\frac{1}{\kappa_{l}}\right)\sum_{l=1}^{L}\left([\mathcal{L}(w^{t})]_{l}-[\mathcal{L}^{*}]_{l}\right)$
$\displaystyle=$
$\displaystyle\max_{l\in[L]}\left(1-\frac{1}{\kappa_{l}}\right)\left(\mathcal{L}(w^{t})-\mathcal{L}^{*}\right).$

This concludes the proof of Theorem 1.

###  E.3 Proof of Theorem 2

Similarly as the proof of Theorem 1, we assume $h=0$. The update rule of Adam
becomes

$w^{t+1}=w^{t}-\eta(D_{Adam}^{0})^{-1}Hw^{t},$

where
$D_{Adam}^{0}=\operatorname{diag}(\nabla\mathcal{L}(w^{0})\circ\nabla\mathcal{L}(w^{0}))^{\frac{1}{2}}=\operatorname{diag}(|Hw^{0}|)$.
We denote $d_{t}=\eta(D_{Adam}^{0})^{-1}Hw^{t}$ and thus we have
$w^{t}=\frac{1}{\eta}H^{-1}D_{Adam}^{0}d^{t}$ and $w^{t+1}=w^{t}-d_{t}$. These
relations also hold for each block by changing the notation to $H_{l}$
$w_{l}^{t}$, $D_{Adam}^{0}$, and $d_{l}^{t},$ etc.. Following the framework in
[Sun and Ye, 2021], 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 two
terms characterizes the rate of convergence. We now express both terms using
$d_{l}^{t}$. For the cost-to-go term for the $l$-th block, we have

$[\mathcal{L}(w^{t})]_{l}-[\mathcal{L}^{*}]_{l}=\frac{1}{2}(w_{l}^{t})^{T}H_{l}w_{l}^{t}=\frac{1}{2\eta^{2}}(d_{l}^{t})^{T}D_{Adam,l}^{0}H_{l}^{-1}D_{Adam,l}^{0}d_{l}^{t}.$
(16)

For the per-step improvement, we have

$\displaystyle[\mathcal{L}(w^{t})]_{l}-[\mathcal{L}(w^{t+1})]_{l}$
$\displaystyle=$
$\displaystyle\frac{1}{2}(w_{l}^{t})^{T}H_{l}w_{l}^{t}-\frac{1}{2}(w_{l}^{t+1})^{T}H_{l}w_{l}^{t+1}$
(17) $\displaystyle=$
$\displaystyle\frac{1}{2}(w_{l}^{t})^{T}H_{l}w_{l}^{t+1}-\frac{1}{2}(w_{l}^{t}-d^{t})^{T}H_{l}(w_{l}^{t}-d_{l}^{t})$
$\displaystyle=$
$\displaystyle(d_{l}^{t})^{T}H_{l}w_{l}^{t}-\frac{1}{2}(d_{l}^{t})^{T}H_{l}d_{l}^{t}$
$\displaystyle=$
$\displaystyle\frac{1}{2}(d_{l}^{t})^{T}\left(\frac{2}{\eta}D_{Adam,l}^{0}-H_{l}\right)d_{l}^{t}.$

To proceed, we denote $\hat{H}=(D_{Adam}^{0})^{-1}H$ and we denote its
eigenvalues as
$\hat{\lambda}_{1}\geq\hat{\lambda}_{2}\geq\cdots\hat{\lambda}_{d}$.
Similarly, we denote $\hat{H}_{l}=(D_{Adam,l}^{0})^{-1}H_{l}$ and its
eigenvalues
$\hat{\lambda}_{l,1}\geq\hat{\lambda}_{l,2}\geq\cdots\hat{\lambda}_{l,d_{l}}$.
Let $\eta=\min_{l\in[L]}C_{l,1}$, we have

$\displaystyle\frac{[\mathcal{L}(w^{t})]_{l}-[\mathcal{L}^{*}]_{l}}{[\mathcal{L}(w^{t})]_{l}-[\mathcal{L}(w^{t+1})]_{l}}$
$\displaystyle=$
$\displaystyle\frac{\frac{1}{\eta^{2}}(d_{l}^{t})^{T}D_{Adam,l}^{0}H_{l}^{-1}D_{Adam,l}^{0}d_{l}^{t}}{(d_{l}^{t})^{T}\left(\frac{2}{\eta}D_{Adam,l}^{0}-H_{l}\right)d_{l}^{t}}$
(18) $\displaystyle\leq$
$\displaystyle\left\|\frac{1}{\eta^{2}}\left(\frac{2}{\eta}D_{Adam,l}^{0}-H_{l}\right)^{-1}D_{Adam,l}^{0}H_{l}^{-1}D_{Adam,l}^{0}\right\|_{2}$
$\displaystyle\overset{(*)}{\leq}$
$\displaystyle\frac{C_{l,2}^{2}\lambda_{l,1}^{2}}{(\min_{l\in[L]}C_{l,1}^{2})\lambda_{l,1}\lambda_{l,d_{l}}}$
(19) $\displaystyle\leq$
$\displaystyle\left(\max_{l\in[L]}\frac{C_{l,2}^{2}}{C_{l,1}^{2}}\right)\kappa_{l},$
(20)

where $(*)$ is due to: by Assumption 1, $D_{Adam,l}^{0}\preccurlyeq
C_{l,2}\lambda_{l,1}I$,
$\frac{2}{\eta}D_{Adam,l}^{0}-H_{l}\succcurlyeq\left(\frac{2}{C_{l,1}}C_{1l,}\lambda_{l,1}-\lambda_{l,1}\right)I\succcurlyeq\lambda_{l,1}I$,
where $\preccurlyeq$ and $\succcurlyeq$ are matrix inequalities. By
rearranging both sides of (20), we have
$[\mathcal{L}(w^{t+1})]_{l}-[\mathcal{L}^{*}]_{l}\leq\left(1-\frac{1}{\left(\max_{l\in[L]}\frac{C_{l,2}^{2}}{C_{l,1}^{2}}\right)\kappa_{l}}\right)\left([\mathcal{L}(w^{t})]_{l}-[\mathcal{L}^{*}]_{l}\right)$.
Summing up both sides over $l\in[L]$ and we conclude the proof.

$\displaystyle\mathcal{L}(w^{t+1})-\mathcal{L}^{*}$ $\displaystyle=$
$\displaystyle\sum_{l=1}^{L}\left([\mathcal{L}(w^{t+1})]_{l}-[\mathcal{L}^{*}]_{l}\right)$
$\displaystyle\leq$
$\displaystyle\sum_{l=1}^{L}\left(1-\frac{1}{\left(\max_{l\in[L]}\frac{C_{l,2}^{2}}{C_{l,1}^{2}}\right)\kappa_{l}}\right)\left([\mathcal{L}(w^{t})]_{l}-[\mathcal{L}^{*}]_{l}\right)$
$\displaystyle\leq$
$\displaystyle\max_{l\in[L]}\left(1-\frac{1}{\left(\max_{l\in[L]}\frac{C_{l,2}^{2}}{C_{l,1}^{2}}\right)\kappa_{l}}\right)\sum_{l=1}^{L}\left([\mathcal{L}(w^{t})]_{l}-[\mathcal{L}^{*}]_{l}\right)$
$\displaystyle=$
$\displaystyle\max_{l\in[L]}\left(1-\frac{1}{\left(\max_{l\in[L]}\frac{C_{l,2}^{2}}{C_{l,1}^{2}}\right)\kappa_{l}}\right)\left(\mathcal{L}(w^{t})-\mathcal{L}^{*}\right).$