Heavy-Tailed Class Imbalance and Why Adam Outperforms Gradient Descent on Language Models

Summary

This paper addresses the important problem of why Adam outperforms SGD for language tasks, proposing that heavy-tailed class imbalance in the training dataset is the key factor in this performance gap. A series of experiments demonstrate that Adam consistently outperforms SGD under heavy-tailed class imbalance because SGD shows slow or no progress on low-frequency classes, whereas Adam progresses independently of class frequencies. This finding holds true across different data types (language and image data), architectures (CNNs and Transformers), and the stochasticity of the training algorithm (mini-batch and full-batch). The paper also theoretically investigates a linear softmax classification problem, proving that the convergence of gradient flow (as a proxy for GD) depends heavily on class frequencies, whereas the convergence of a continuous-time variant of sign descent (as a proxy for Adam) is independent of class frequencies. This theoretical result establishes the benefits of sign-based methods under heavy-tailed class imbalance for a linear model.

Strengths

- The finding that heavy-tailed class imbalance leads to the performance gap between Adam and SGD is a solid and novel contribution. This offers new insights into understanding the important question of why and when Adam outperforms SGD. - The experiments are well-designed and robust, including thorough ablation studies that strongly support the hypothesis. - I think it's great that the authors can rigorously establish the effect of class frequencies on the training speed of GD and sign descent in a simplified setting. Additionally, the investigation into the correlation between gradient and Hessian across coordinates provides valuable insights into the underlying mechanisms by which sign-based methods benefit from heavy-tailed class imbalance.

Weaknesses

- The linear model studied is an oversimplified setting and is designed to be biased towards sign descent. Therefore, it is unclear whether the insights gained from this model extend to practical settings, particularly in understanding why Adam benefits from a heavy-tailed class imbalance in real-world scenarios. Specifically, it remains uncertain whether the correlation between gradient, Hessian, and class frequencies observed in the linear model holds true in more complex, practical settings.

Questions

- In the gradient norm and Hessian trace plots in Figures 24-26 in Appendix G, how are the weight blocks $w_c$ defined? Do they correspond to the last layer of each architecture? If so, how does class imbalance affect other layers beyond the last layer? - Could the authors please cite the relevant works by Zhang et al. (https://arxiv.org/abs/2402.16788) and Xie and Li (https://arxiv.org/abs/2404.04454) that also study the benefits of Adam? In particular, the block heterogeneity discussed by Zhang et al. seems closely related to the intuition in this paper. Additionally, Xie and Li observe that Adam outperforms GD when the loss function has better properties under $\ell_\infty$ geometry, which could relate to heavy-tailed class imbalance. Therefore, I wonder if the authors can discuss and reconcile these works with the findings on heavy-tailed class imbalance.

Rating

8

Confidence

4

Soundness

4

Presentation

4

Contribution

4

Limitations

The paper discusses its limitations and future directions in detail. For additional potential limitations, please refer to the 'Weakness' section.

Summary

This paper argues that heavy-tailed class imbalances in natural language datasets, which follow because some words are much more frequent than others, causes (or significantly underlies) the performance gap typically observed between Adam and SGD. To make this argument, the authors 1. reproduce the performance gap by training different language models on some standard datasets, like WikiText-103 and PTB; 2. empirically demonstrate that the gap between Adam and SGD can also be reproduced in vision tasks, in which it’s normally absent, by training common vision models on artificially created image datasets with heavy-tailed class imbalances; 3. introduce a simple linear model trained on synthetic uniform data, whose class frequencies follow a power law, and again reproduce the performance gap; 4. rule out batch noise as the reason underlying the performance gap, by demonstrating that the latter remains even when they train their models with full (i.e. deterministic) gradient descent; and 5. demonstrate that the performance gap can be reduced by upweighting the loss of low-frequency classes. Having done these experimental investigations, the authors then attempt to provide some understanding and intuition as to why these class imbalances affect SGD. The authors study a collection of simplified models and training problem and argue that the class imbalance leads to correlations between gradients, Hessians and class probabilities, which appear to help Adam during training.

Strengths

- This paper studies an important problem in modern machine learning, viz. the performance gap of common training algorithms, specially when dealing with language models. - The problem and hypothesis are very clearly stated. - The paper is very well written and structured. Indeed, the authors first present their empirical evidence and only later attempt to provide some intuitive understanding of the phenomenon under investigation. - As summarised above, the paper provides significant empirical evidence that supports their main hypothesis, namely that heavy-tailed class imbalances in the training datasets cause the performance gap between Adam and SGD. - The paper also provides some theoreticall understanding as to why the performance gaps actually takes place. In sum, this paper is very sound, tackles a relevant problem and, in my view, represents an important contribution to the machine learning community.

Weaknesses

The main weaknesses or limitations of the paper pertain to the simplified models and assumptions in the theoretical investigations. However, the authors explicitly acknowledge and discuss them. Although the paper is very well written and structured, the following suggestions might improve readability: - The authors could explicitly mention that they reproduce the performance gap between Adam and SGD in language tasks, for different datasets and models. This information is hidden in Appendix A and could escape the casual reader. - Some captions in the figures are misleading. The authors perform experiments both with SGD and full GD and the way these are labelled, or referred to, in some captions is confusing. For example, the caption in figure 2 refers to SGD while the figure contains results for GD. Similarly the caption in figure 1 refers to GD but the figure contains results for SGD. - Perhaps section 2.2 could be moved inside section 3? Typos: lines 152, 160 and 225.

Questions

1. Do you know why SGD eventually finds the minimum in the imbalanced MNIST case but does not in the imbalanced ImageNet?

Rating

7

Confidence

4

Soundness

4

Presentation

3

Contribution

3

Limitations

Yes, they authors addressed the limitations of their study.

Summary

The paper shows that under the class imbalance setting which is natural in language tasks, Adam can be faster than SGD. Meanwhile, the authors investigate the linear model deeply showing the relationship between gradient and Hessian and the convergence speed of sign-gd and gd algorithm.

Strengths

1. The authors explain that class imbalance is a reason that Adam can outperform SGD. 2. For linear models, the authors establish the relationship between hessian and gradient showing the "correctness" of Adam who approximates Lipschitz with gradients. 3. Further, the authors point out that sign-gd can converge faster than gd for linear models.

Weaknesses

1. All of the results are based on that the parameters of different classes are separable. 2. Since the optimal solution of NN can not be infinity due to some generalization constraints (e.g. adding weight decay), will the sign algorithm still be faster than gd?

Questions

When the parameters are not separable, does the conclusion still hold?

Rating

7

Confidence

4

Soundness

3

Presentation

3

Contribution

3

Limitations

N/A

Summary

This paper investigates the reason why Adam outperforms (S)GD by a large margin on language models when the performance gap is much smaller in other settings. The authors argue that language data often has a heavy-tailed class imbalance, where a large fraction of the data consists of several classes with a small number of words, which leads to slow improvement in the training loss of these classes for (S)GD, whereas Adam (or sign descent) is more robust to this type of imbalance. Since the fraction of classes with fewer samples is relatively large, this contributes to a difference in the rate of decrease of the average loss as well, in contrast to conventionally considered settings for binary classification with class imbalance. The authors show that the performance gap can be reproduced in other settings, such as when training CNNs on MNIST, ResNet on ImageNet, and even linear models on synthetic high-dimensional data, when the heavy-tailed class imbalance is introduced in these datasets. The authors also show that similar effects are observed without stochasticity and with a simpler algorithm (sign gradient descent), and present some theoretical results in simple settings showing a difference in the rate of decrease in the training loss between GD and sign GD.

Strengths

- The paper presentation is exceptional: it is very well-written with aesthetic plots that illustrate the results really clearly. The contribution, as well as the motivation for each of the experiments conducted in the paper, is discussed in detail; I especially like the last paragraph of Section 1.1. - The experiments are designed well and thoroughly support the hypothesis that heavy-tailed class imbalance is a key reason for the performance gap between SGD and Adam on language models. Results showing that the gap can be reproduced in other settings when the heavy-tailed class imbalance is introduced are very convincing. - The paper contributes to our understanding of when and why Adam performs better than GD, which is fundamental to ultimately improving optimization.

Weaknesses

There are no major weaknesses. The authors adequately acknowledge and discuss the limitations of the work. Some minor concerns are as follows. - The theoretical results are in oversimplified settings. However, this seems justified since theory is not one of the primary contributions of the paper. It might be better to allocate less space to this part (e.g., Section 3.3). - It would be nice to elaborate on lines 331-333 and include more discussion on this aspect in Section 3. - Some corrections/suggestions: - Missing ‘the’ in line 152. Extra ‘the’ in line 225. Missing ‘than’ in line 607. Extra ‘be’ in line 631. - It would be good to add some description in lines 496-497.

Questions

For the Barcoded MNIST dataset, I think the number of new images should be $5\times 10\times (2^{10}-1)$ since one of the 10-bit patterns would be the same as the background. Can the authors clarify this?

Rating

7

Confidence

4

Soundness

4

Presentation

4

Contribution

3

Limitations

There are no potential negative impacts of this work. The authors discuss the limitations of their work at length in Section 4.