Why Transformers Need Adam: A Hessian Perspective

Thanks for the insightful question and careful evaluation of our paper. We provide response as follows.

1. In theoretical analysis, the momentum part of Adam is fully ignored ( $\beta_1=0$ in Algorithm 3).

Thanks for the comment. The current analysis ignores momentum because we want to focus on the benefit of assigning coordinate-wise lr, so we remove the effect of all other components to control variables. Nevertheless, we agree that choosing $\beta_1=0$ simplifies the problem. However, exploring the role of momentum is an independent and equally challenging problem. We consider it as an important future direction.

2. On the computational complexity of calculating the JS distance and the Hessian approximation; compare runtime to the training time; scale efficiently to larger models;

Thanks for the great question. Currently, our JS distance is indeed rather expensive to compute: it requires comparable to one training run. Fortunately, we find the current deployment of SQL is redundant for measuring hessian heterogeneity. We propose the following simple tweaks to significantly reduce the computation time, while still effectively detecting the Hessian heterogeneity. We call it simplified SQL:

1. Change the hyperparemters of SQL, including:

• 1-1: we change num_v = 10 to num_v = 1. In SQL, num_v decides the number of random Gaussian vectors to approximate the expected quadrature. It is reasonable to reduce num_v because in high dimensional space, random vectors tend to concentrate around their mean, so one random sample can already be informative enough.
• 1-2: we change the Lanzcos step m = 100 to m = 10. The reduction on Lanzcos step will have a coarse estimation on the middle eigenvalue, but won't affect much the heterogeneity, which is more dependent on the extreme eigenvalues

2. Randomly sample a subset of blocks and reduce batch size for estimating the spectrum. We uniformly sample 50% blocks and choose batch size = 32 (previously batch size = 1024).

We report the result and runtime in the Table below. As a result, the simplified SQL can obtain the same message as the original SQL: JS0 of ResNet is about 100x smaller than that of BERT. Further, the simplified SQL is highly efficient to compute. With this simplified SQL, we believe our method can efficiently scale to larger models.

Tested on: single V100

3. Is the architecture the only factor affecting the Hessian structure? I am curious if the training distribution, particularly label imbalance, also impacts the Hessian structure

We agree that there exist other factors (such as data distribution and label imbalance) that affect the Hessian structure. However, for ViT on ImageNet, we find SGD (even after careful tuning) is still largely worse than Adam. As such, we believe architecture plays a more crucial rule.

Thanks for the careful review of our work. Here is our respectful response.

1. It is unclear whether this correlation implies causation beyond intuitive understanding.

We provide the following evidence on causation.

I: Theory: On quadratic models, our theory suggests that: when heterogeneity happens, the complexity of Adam is better than GD.

II: Control-variable experiments:. We design the experiment as follows:

1. We take two Hessian matrices with an identical set of eigenvalues (cases 3 and 4 in Section 3.1). This ensures the performance of GD is fixed.

2. We re-arrange these eigenvalues in two different ways, which lead to heterogeneity and homogeneity (case 3 and 4).

3. We find Adam is faster when heterogeneity exists, resulting in a huge gap with GD

This experiment implies "heterogeneity" can cause the "gap between Adam and SGD."

2. Block seems arbitrarily determined.

We kindly point out that block is not arbitrarily determined. Here is our principle: we define "blocks" as the irreducible dense principal sub-blocks in Hessian.

Since it requires heavy calculation on Hessian to precisely find blocks by the principle above, we find a computationally-friendly way: We use the PyTorch partition to determine the blocks. We find that partition in PyTorch can usually match our principle. For instance, in Figure 2, each block corresponds to the parameters in Q, K, V; and MLP, respectively. Under this common partition strategy, we observe heterogeneity in Transformers but not in CNNs

3. This observation might be overfitting the transformer architecture.

We kindly point out that the observation is not overfitting to Transformers since we also observed heterogeneity on non-Transformer models such as MLP mixer (a real-world non-attention-based architecture, Figure 5 in paper), and SGD also performs worse than Adam here.

4. The case study on quadratics helps understand and attempt initial theoretical analysis, but it only partially represents behaviors in larger networks.

We agree that quadratic models cannot capture all behaviors of real NNs. To cover more behaviors in larger NN, we need analysis on more fine-grained settings such as noisy quadratic models or 1-layer attention models. It is an intriguing direction to pursue in the future.

5. It's unclear whether the gap is due to the claimed reason or if hyperparams are insufficiently tuned.

We believe the gap is due to Hessian heterogeneity, instead of improper hyperparam for SGD. For all Transformers, we grid search the learning rate of SGD and report the best performance. We grid-search lr of SGD as follows.

• For BERT: we use lr = 1e-4 for Adam. For SGD, we search lr over [1e-4, 1e-3, 1e-2, 1e-1]

• For ViT: we use lr = 1e-4 for Adam. For SGD, we search lr over [1e-4, 5e-4, 1e-3, 5e-3, 1e-2, 5e-2, 1e-1]

• For GPT2-nano: we use lr = 1e-2 for Adam, For SGD, we search lr over [1e-5, 5e-5, 1e-4, 5e-4, 1e-3, 5e-3, 1e-2, 5e-2, 1e-1, 5e-1, 1]

We visualize the performance of SGD under different lr. This can be seen in Figure 1 in the attached PDF. We find that SGD (even after careful tuning) is still worse than Adam, including ViT on ImageNet.

Additional evidence for ViT: We here provide more evidence that "SGD is worse than Adam on ViT" from the literature.

• On Cifar-10, [1] carefully tuned SGD on ViT and find that SGD cannot work (Figure 4). They reported that "the SGD algorithm did not perform well for this task” even after careful tuning: “ The optimizer hyperparam were set according to the values obtained during hyperparam tuning. ”

• On ImageNet, [2] provides empirical evidence that SGD is worse than Adam on vanilla ViT. The authors reported that " SGD yields significantly worse results than AdamW (on ViT)", "ViT often fails to converge with SGD " (their Figure 3).

As for Figure 12 in Kunstner et.al., there are two possible reasons for the mismatch: (1) they used SimpleViT, which is a simplified version of the vanilla ViT as we conducted. It is possible that SimpleViT exhibits less heterogeneity and is more friendly to SGD. (2) perhaps they did not tune Adam carefully enough and Adam can perform better if well-tuned.

[1] AdaGC: A Novel Adaptive Optimization Algorithm with Gradient Bias Correction

[2] Early convolutions help transformers see better

6. Block-wise spectrum along training.

Following your suggestions, we plot the block-wise spectrum at 25%, 50%, and 100% training steps for GPT2 and ViT. We find an interesting phenomenon: Hessian heterogeneity tends to reduce along training. This can be seen in Figure 2 and 3 in the attached PDF.

We also take the checkpoint of ViT at the 46th epoch (which is about 50% steps) and continue to train with SGD. We find SGD now can perform much better. We provide a similar explanation here: SGD now performs better because there is less Hessian heterogeneity in ViT at the 46-th epoch (50% training steps). This can be seen in Figure 2 in the attached PDF.

This phenomenon above also provides one interesting practical implication: if we can find a good initialization with less heterogeneity, it is possible to train transformers with SGD. However, designing such initialization is non-trivial, as its explicit relation with heterogeneity is still unclear. We leave this topic as an intriguing future direction.

To sum up, we make two new findings when investigating the blockwise spectrum along training: 1. Heterogeneity tends to reduce along training. 2. as heterogeneity is reduced, we can switch Adam to SGD. We will include these results in the revised script.

7. Ablation studies on how well the Hessian spectrum is estimated through the SQL method, perhaps on reasonably small networks?

We provide an ablation study on a 4-layer Transformer. The results are shown in Figure 4 in the attached PDF. We find SQL provide accurate approximation over the true eigenvalue density.

Thank you for addressing my concerns; I will raise my score accordingly.

We are grateful for the careful review of our paper and the great questions. Please find our respectful reply below.

1. For the same Transformers model, when the initialization changes to induce different block heterogenous levels, how would the performance of SGD compared with Adam change?

Thanks for the interesting question. Here is our finding: For the same model, when heterogeneity is reduced, the gap between Adam and SGD is also reduced. We provide evidence as below.

There are typically two approaches to change initialization (please let us know if neither of these aligns with your thoughts).

• Approach 1: Change the variance of random initialization.

• Approach 2: Choose a checkpoint in the middle of training as the initialization.

Unfortunately, we find Approach 1 is not suitable for controlling the heterogeneity: it seems unclear how the variance of random initialization is explicitly related to the heterogeneity in Hessian. As such, we use Approach 2 to change the initialization and track the performance of SGD. We present two experiments below. (Exp 1 is already presented in the paper; Exp 2 is new.)

Exp 1: On finetuning tasks of GPT2, when we change the initialization to the pre-trained weights, GPT2 (pre-trained) exhibits much less hetereogeneity compared with GPT2 with random initialization. This can be seen in Figure 4 (f) and Figure 6 (a) in the paper. As a result, SGD can reach a similar loss as Adam here. This is shown in Figure 6 (b) in the paper.

Exp 2: We realize that Exp 1 might not fully answer your question: although Exp 1 does not change the architecture, it changes the training dataset (from pre-training dataset to SFT dataset), and thus changes the loss function. To rigorously address your concern, we further take the checkpoint of ViT at the 46-th epoch (which is about 50% training steps) and calculate Hessian heterogeneity. We make two findings: (1) For ViT checkpoint at the 46-th epoch, Hessian heterogeneity is largely reduced compared with that of random initialization. (2 ) If we continue to rain it with SGD, we find SGD now can perform much better than training from scratch. This result can be seen in Figure 2 in the attached PDF.

This phenomenon above also provides one interesting practical implication: if we can find a good initialization with less heterogeneity, it is possible to train transformers with SGD. Yet, as discussed above, designing such initialization is non-trivial, as its explicit relation with heterogeneity is still unclear. We leave this topic as an intriguing future direction.

2. The theoretical results of the convergence of Adam for the quadratic models show that homogeneous spectra have worse complexity than heterogeneous spectra, is this the truth in practice?

Thanks for the great question. Let us clarify: our theory does not imply that Adam has worse complexity on homogeneous spectra. For simplicity of discussion, let's state your comment as a conjecture.

Conjecture 1: Adam on quadratic problems with homogeneous spectra has worse complexity than Adam on heterogeneous spectra.

If we understand your question correctly (please tell us if not so), your question states that our theory implies Conjecture 1. We kindly point out that Conjucture 1 is not correct. We'll clarify below.

(1) What we proved and why it does NOT imply Conjecture 1. Our theoretical result: Adam has complexity $O(\max_l \kappa_l)$. GD has complexity $\Omega(\kappa)$

If our result implies Conjecture 1, one needs the following argument: when changing heterogeneity to homogeneity, $\max_l \kappa_l$ increases, and thus Adam is slower.

However, "changing heterogeneity to homogeneity" does not necessarily mean "$\max_l \kappa_l$ increases". Actually, "$\max_l \kappa_l$" can change in an arbitrary way (can increase, decrease, or keep the same) when changing the heterogeneity. See detailed examples below.

(2) Detailed Examples on Comparing Homo and Hetero Cases. We provide three examples below.

Notation: We use Adam (homo) to denote the rate of Adam on homogeneous Hessian. Similarly for Adam (hetero)

Example 1: Adam (homo) is same as Adam (hetero)

case 1-1: homo: eigenvalue {1,2}, {1,2}

case 1-2: hetero: eigenvalue {1,2}, {11,12}

Since $\max_l \kappa_l$ are the same for both case 1-1 and 1-2, Adam (homo) is same as Adam (hetero)

Example 2: Adam (homo) is faster than Adam (hetero)

case 2-1: homo: eigenvalue {1,1.5}, {1,1.5}

case 2-2: hetero: eigenvalue {1,2}, {11,12}

Since case 2-1 has smaller $\max_l \kappa_l$ than case 2-2, Adam (homo) is faster than Adam (hetero)

Example 3: Adam (homo) is slower than Adam (hetero)

case 3-1: homo: eigenvalue {1,11}, {2,12}

case 3-2: hetero: eigenvalue {1,2}, {11,12}

Since case 3-1 has larger $\max_l \kappa_l$ than case 3-2, Adam (homo) is slower than Adam (hetero)

To sum up, there is no conclusion on "whether Adam under homogeneity is faster or slower than Adam under heterogeneity ". Either case can happen.

(3) possible source of confusion: We thought (correct us if not so) the confusion may come from the numerical examples (Case 3 and 4 in Section 3.1) in the paper. If comparing two figures, Adam (homo) in Case 3 is slower than Adam (hetero) in Case 4. But as argued above, this is just one example, and it does NOT show Adam (homo) is ALWAYS slower than Adam (hetero).

As a result, all the above three examples can happen, and "Changing hetero to homo" does not necessarily mean "Adam becomes slower".

Thanks for the careful evaluation of our paper. We respectfully provide our response as follows. 1-1. Are there other publications that make these claims (SGD worse than Adam on attention-based model)?

Thanks for raising this comment. Yes, it is widely reported that SGD largely underperforms Adam on attention-based models, even after careful tuning. Here are some references:

1. [1] provides empirical evidence that SGD is worse than Adam on BERT, but not on SGD. On Bert, the authors comment that "Although hyperparameters for SGD are finetuned, a large performance gap is still observed" (their Figure 1)

2. [2] provides empirical evidence that SGD is worse than Adam on ViT. The authors reported that " SGD yields significantly worse results than AdamW (on ViT)", "ViT often fails to converge with SGD " ( their Figure 3)

3. [3] carefully tuned SGD on ViT and find that SGD cannot work (their Figure 4). They reported that " the SGD algorithm did not perform well for this task”" even after careful tuning: “ The optimizer hyperparameters were set according to the values obtained during hyperparameter tuning. ”

4. [4] carefully tuned SGD on modern GPT models and find that SGD cannot work (their Figure 1).

5. Besides, main-stream large language models use Adam, including GPT3 [5], Llama series [6], Gemini series [7], etc. This serves as additional evidence that SGD is worse than Adam for Transformer training.

1-2. How did the authors choose/optimize the learning rate for SGD, and/or the learning rate schedule?

For the learning rate schedule, we use the same default cosine decay with warmup for both SGD and Adam. As for the learning rate, we grid search the learning rate of SGD and report the best performance. We grid-search the learning rate as follows.

• For BERT: we use lr = 1e-4 for Adam. For SGD, we grid search lr over [1e-4, 1e-3, 1e-2, 1e-1]

• For ViT: we use lr = 1e-4 for Adam. For SGD, we grid search lr over [1e-4, 5e-4, 1e-3, 5e-3, 1e-2, 5e-2, 1e-1]

• For GPT2-nano: we use lr = 1e-2 for Adam, For SGD, we grid search lr over [1e-5, 5e-5, 1e-4, 5e-4, 1e-3, 5e-3, 1e-2, 5e-2, 1e-1, 5e-1, 1]

We also visualize the performance of SGD under different learning rates. This can be seen in Figure 1 in the attached PDF. For all these Transformers we investigate, we find SGD (even after fine-tuning carefully) performs significantly worse than Adam.

2. Novelty of using block spectrum (potentially).

Thanks for the question. We believe our idea is non-trivial and non-obvious. The novelty of our perspective is also confirmed by other three reviewers, for instance:

• Reviewer PTmy: "The discovery of block heterogeneous for the reason of the bad performance of SGD on Transformers is interesting and novel."

• Review tJ3K: "The connection between Adam's performance and Hessian block-wise heterogeneity is novel, as are the related empirical observations."

• Reviewer A8Vc: " I find the ideas introduced throughout this work interesting and potentially useful to the community."

Our comment on why blockwise spectrum is non-trivial. For completeness, we further comment on why the perspective of the blockwise spectrum is new and non-trivial.

--The natural idea does not work. To explore "why Transformers need Adam", a natural idea is to investigate the full Hessian eigenvalue (spectrum). This is because: by optimization theory, full Hessian eigenvalues largely decide the behavior of gradient methods, and there is an active line of work trying to understand NN training via full Hessian spectrum (e.g., reference [8] here, also see the reference [12,32,39,67,68,69,75,76,97,98,103] in our paper ). Unfortunately, we find there is no noticeable difference between the full spectrum of CNNs and that of Transformers (Figure 1 in the paper). As such, the natural idea of full spectrum does not work.

--Why our perspective is non-trivial. Our major conceptual innovation is that we connect the following findings to "why Transformers need Adam".

1. Transformers contains various kinds of building blocks, while CNN consists of similar convolutional layers. Based on this, we conjecture that blockwise differences might be crucial.

2. "Hessian of NNs have near-block-diagonal structure" (a non-trivial but highly overlooked finding in this field). Based on this, we realize that the blockwise spectrum has rich information and could be used to quantify the blockwise difference. Note that for models with dense Hessian, blockwise spectrum can still be computed, but loses substantial information.

As such, our blockwise spectrum perspective is novel to the deep learning community because it is based on multiple non-trivial findings.

We believe this perspective is also new to optimization community: for optimization community, it is very rare to analyze (near-) block-diagonal Hessian structure since typical problems do not have such structure. For instance, in the classical non-linear programming dataset [9], all problems have non-block-diagonal Hessian. We point out a new perspective to characterize modern optimization problems.

As such, we believe our perspective is new, non-trivial, and potentially useful to a wide range of audiences.

[1] Why are adaptive methods good for attention models?

[2] Early convolutions help transformers see better.

[3] AdaGC: A Novel Adaptive Optimization Algorithm with Gradient Bias Correction

[4] Deconstructing What Makes a Good Optimizer for Language Models

[5] Language models are few-shot learners

[6] Llama 2: Open foundation and fine-tuned chat models

[7] Gemini: A natural language system for spoken-language understanding

[8] An investigation into neural net optimization via hessian eigenvalue density

[9] Nonlinear Programming Solvers for Unconstrained and Constrained Optimization Problems: a Benchmark Analysis