Why Transformers Need Adam: A Hessian Perspective

Summary

This paper explores the performance gap between SGD and Adam when training transformers. The authors analyze the eigenspectrum of the Hessian in different neural network architectures and attribute the failure of SGD to the different Hessian spectra in transformers compared to CNNs and MLPs. They empirically show that while the full Hessian spectrum is similar at initialization across architectures, transformers exhibit block heterogeneity, meaning different parameter blocks (e.g., key/query layer versus MLP layer) have distinct spectra. Conversely, other architectures such as CNNs and MLPs are block homogeneous. By artificially inducing block heterogeneity in MLPs, they demonstrate through numerical experiments that SGD's performance declines as the block spectrum variation increases, while Adam's performance remains stable. They attribute Adam's success in block heterogeneous cases to its ability to assign different learning rates to different parameters. The authors also propose using the Jensen-Shannon distance of Hessian blocks as a quantitive measure of the performance gap. They also provide a preliminary theoretical analysis comparing the convergence rates of SGD and Adam on a quadratic function with a block-form Hessian, highlighting Adam's dependence on the fine-grained condition number of individual Hessian blocks.

Strengths

This paper tackles a key question: understanding the optimization challenges of transformers with SGD and Adam is crucial for developing more efficient techniques for training large language models. The paper is well-organized and easy to read. I find the ideas introduced throughout this work interesting and potentially useful to the community.

Weaknesses

Quantifying the performance gap between SGD and Adam using the average JS distance between block-wise Hessian spectra is an interesting idea. However, based on the runtime reported in the appendix, the Hessian approximation itself seems computationally expensive, which can raise concerns about how practical it is for large models, especially since the experiments in this paper are only on small models. Additionally, the main theoretical analysis is conducted in a pretty simplified setup where the momentum part of Adam is fully ignored ($\beta_1=0$ in Algorithm 3).

Questions

1. Can you comment on the computational complexity of calculating the JS distance and the Hessian approximation? Specifically, can you compare the runtime of computing this quantitative measure to the training time of the models in your paper? For significantly larger models, do you expect this approach to scale efficiently and be more practical than training the model? 2. Is the architecture the only factor affecting the Hessian structure? I am curious if the training distribution, particularly label imbalance, also impacts the Hessian structure, since in some previous works previous works, such as [46], the performance gap has been attributed to the heavy-tailed distribution of labels in language data.

Rating

6

Confidence

3

Soundness

3

Presentation

4

Contribution

3

Limitations

Yes

Summary

This paper investigates why Adam outperforms SGD in Transformers by examining Hessian estimations. While the overall Hessian spectrum appears similar in both Transformers and other neural networks, the block-wise Hessian spectrum is heterogeneous across blocks in Transformers and homogeneous in other neural networks. The authors empirically confirm that this heterogeneity hampers SGD's performance through experiments involving Transformers, CNNs, MLPs, and quadratic problems. They attribute Adam's advantage to its coordinate-wise learning rate and provide initial theoretical analysis in a simplified setting.

Strengths

- The connection between Adam's performance and Hessian block-wise heterogeneity is novel, as are the related empirical observations. - Related work is comprehensively cited in the Appendix. - Although the finding itself is not groundbreaking, it suggests potential for future research in theoretical analysis and practical improvements of Adam, particularly regarding its advantage in adaptively updating different blocks.

Weaknesses

- The experiments support the claim that Adam outperforms SGD when Hessian block-wise heterogeneity structure is present. However, it is unclear whether this correlation implies causation beyond intuitive understanding. In general, the conceptual contribution is not very strong. - The definition of "block" also seems arbitrarily determined. For example, Fig. 3 shows that all convolution layers have a similar spectrum, while the spectrum differs between Query, Value, and MLP blocks in Transformers. Given that one set is the same type of "block" (conv layer) and the other is not, this does not seem surprising. Since how the block is defined entirely determines Hessian heterogeneity versus homogeneity, this observation might be overfitting the transformer architecture and could be difficult to generalize to non-mainstream architectures as claimed. - The case study on quadratics helps understand and attempt initial theoretical analysis, but it only partially represents behaviours in larger networks. The noise in SGD does not play a significant role here, and Adam with $\beta_2 = 0.99$ performs very differently than in larger networks.

Questions

- Including a more detailed experimental setting will help with reproducibility. The hyperparameters for experiments such as learning rate are not included. Since SGD is sensitive to hyperparameters, it's unclear whether the gap is due to the claimed reason or if hyperparams are insufficiently tuned. For example, the training curve of SGD in Figure 9(a) in ViT-ImageNet is much lower than Adam and that is not the case in [Figure 12, Kunstner et.al. Heavy-Tailed Class Imbalance and Why Adam Outperforms Gradient Descent on Language Models]. Can the author comment on this? - Figure 1 shows the full spectrum across different time points to illustrate its evolution during training, yet all other plots in the paper related to the block-wise spectrum are presented at initialization. Is there a reason why the evolution of the block-wise spectrum is not shown? - Are there any ablation studies on how well the Hessian spectrum is estimated through the SQL method, perhaps on reasonably small networks? - Typo in Figure 7 y-axis label

Rating

6

Confidence

2

Soundness

2

Presentation

3

Contribution

2

Limitations

Limitations are well-addressed

Summary

The paper investigates why SGD performs worse than Adam in Transformers from the lens of Hessian. They first find that Transformers are "heterogeneous", that is, the Hessian spectrum across parameter blocks varies dramatically. Then they conduct various tasks on Transformers, CNNs, MLPs, and quadratic problems, to verify that block heterogeneity hampers SGD. Finally, they derive some initial theoretical analysis to indicate that SGD fails because it applies one single learning rate for all blocks, which cannot handle the heterogeneity among blocks.

Strengths

1. The paper is well written, clear to read, and the story is interesting. 2. The discovery of block heterogeneous for the reason of the bad performance of SGD on Transformers is interesting and novel. 3. The experiments are completed to verify the points described in the paper.

Weaknesses

1. The paper demonstrates that for different models when there the block heterogeneous at initialization increases, SGD becomes much worse than Adam. However, these comparisons are between different models. It is natural to ask for the same Transformers model, when the initialization changes to induce different block heterogenous levels, how would the performance of SGD compared with Adam change? 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?

Questions

See weakness

Rating

7

Confidence

3

Soundness

3

Presentation

3

Contribution

4

Limitations

The authors have addressed the limitations.

Summary

In this work the authors assert that the Adam optimizer works well on Transformers while Stochastic Gradient Descent (SGD) does not, and they attempt to explain this phenomenon by inspecting the spectrum of Transformers (i.e., the eigenvalues of the model's Hessian matrix) and other models, such as convolutional networks, where SGD is competitive with Adam. They show that the spectrum of Transformers and convolutional networks are similar, providing little insight. The main innovation of the authors is to investigate the spectrum of individual components/layers of these networks - the so-called block spectrum - and the main contribution of the work is to show that the block spectrum is empirically correlated with the ineffectiveness of SGD relative to Adam.

Strengths

1) The paper is generally clear and well-written - this is appreciated 2) The idea of inspecting the block-wise spectrum is potentially innovative 3) The empirical results presented by the authors are fairly convincing.

Weaknesses

1) The premise of this paper is that SGD performs worse for attention-based models, and it is unclear to me if this is true. I say this for two reasons: a) The authors cite [45,56] on Line 16 after making this claim, presumably to support the claim. However, these are quite old publications and neither of these papers specifically make this claim. Are there other publications that make these claims? b) The authors also provide their own empirical results (e.g., Fig. 3) to support this claim, however, these results seem highly insufficient to me to support this claim. How did the authors choose/optimize the learning rate for SGD, and/or the learning rate schedule? If SGD does not generally perform worse than attention models on Transformers, then it seems to me that the work has little scientific value, since the differences in the block spectrum are just inconsequential differences, or at least that their implications are different than suggested by the authors. e.g., perhaps the authors could argue that the heterogeneous block spectra implies that SGD needs a lot more tuning to work well, for example. I'd ask the authors to please address this point, and I will consider raising my rating. 2) Novelty of using block spectrum (potentially). To me this is the main value of this work. I'm not as familiar with this topic, and it is unclear to me how non-obvious the use of a block spectrum would be for analysis, or how surprising the results of using this approach are: perhaps these findings would be somewhat obvious to researchers familiar with this topic? I will defer to other reviewers for this assessment.

Questions

See "Weaknesses" for my questions.

Rating

5

Confidence

3

Soundness

3

Presentation

3

Contribution

2

Limitations

yes.

TLDR: Reviewers please do acknowledge the rebuttals and react to them.

Dear reviewers, thanks for your reviews. Please do look at the rebuttals of the author and give the authors some feedback whether they could address your concerns. This is really important for the authors. Best regards Your AC.