Heavy-Tailed Class Imbalance and Why Adam Outperforms Gradient Descent on Language Models (2402.19449v2)

Abstract: Adam has been shown to outperform gradient descent on LLMs by a larger margin than on other tasks, but it is unclear why. We show that a key factor in this performance gap is the heavy-tailed class imbalance found in language tasks. When trained with gradient descent, the loss of infrequent words decreases more slowly than the loss of frequent ones. This leads to a slow decrease on the average loss as most samples come from infrequent words. On the other hand, Adam and sign-based methods are less sensitive to this problem. To establish that this behavior is caused by class imbalance, we show empirically that it can be reproduced across architectures and data types, on language transformers, vision CNNs, and linear models. On a linear model with cross-entropy loss, we show that class imbalance leads to imbalanced, correlated gradients and Hessians that have been hypothesized to benefit Adam. We also prove that, in continuous time, gradient descent converges slowly on low-frequency classes while sign descent does not.

• The paper shows that Adam’s adaptive preconditioning mitigates heavy-tailed class imbalance, ensuring uniform learning across rare classes.
• Experimental results reveal that SGD struggles with infrequent classes while Adam achieves consistent training loss reduction across model types.
• The study's insights offer practical modifications for SGD and guide future research on optimizer designs for imbalanced datasets.

Heavy-Tailed Class Imbalance: Exploring Adam's Superiority over Gradient Descent in LLMs

Introduction

The optimization of LLMs is crucial for advancing the field of NLP. An interesting observation made in recent times is the distinct advantage that the Adam optimizer holds over traditional stochastic gradient descent (SGD) when training these models. The paper discussed here explores understanding this phenomenon, attributing the performance disparity to the heavy-tailed class imbalance inherent in LLMling tasks.

Heavy-Tailed Class Imbalance

Language data characteristically displays a heavy-tailed class distribution, where a significant number of classes (or words) are relatively infrequent. Traditional gradient descent methods tend to make slow progress on these low-frequency classes, negatively impacting overall training efficiency. Contrarily, Adam and similar sign-based methods do not exhibit this limitation, thereby facilitating uniform class learning speeds. The researchers empirically substantiate their argument through experiments across various models—including language transformers and vision CNNs—highlighting the generalizability of their findings beyond language data.

Experimental Insights

The distinction between Adam and SGD becomes particularly pronounced when observing training performance disaggregated by class frequency. Experiments demonstrate that while SGD struggles with low-frequency classes—barely making progress—the training loss for these classes reduces much more uniformly under Adam. This behavior persists across different architectures and data types, reinforcing the core thesis that heavy-tailed class imbalance significantly contributes to the optimization gap between Adam and SGD. Intriguingly, the implementation of simpler optimizers, such as sign descent, reveals that altering the update direction rather than magnitude (as done by Adam) is essential for mitigating class imbalance effects.

Theoretical Perspectives

On a linear model exhibiting heavy-tailed class imbalance, it was evidenced that both the scale of gradient and Hessian reflect class frequencies, which leads to ill-conditioning—a situation where gradient descent performance degrades due to vastly different convergence speeds across classes. Adam's efficiency, in this context, could be partially attributed to its preconditioning capability, which approximately counteracts the ill-conditioning by normalizing gradient magnitudes. This finding suggests that, at least for softmax classification on linear models, Adam indirectly caters to the differential scaling induced by class frequencies, facilitating a more balanced training dynamic.

Broader Implications

This paper not only elucidates why Adam outperforms SGD in the context of LLMs but also sheds light on potential improvements across various fields where class imbalance is prevalent. The insights provided could lead to the development of new optimization algorithms or adjustments to existing ones—especially in tasks beyond LLMing. Moreover, the demonstrated effectiveness of simple modifications, such as loss reweighting, provides practical avenues for enhancing SGD's performance, narrowing the gap with Adam.

Future Directions

The comprehensive analysis presented sparks a plethora of questions for future research. Specifically, understanding the full ramifications of heavy-tailed class imbalance on model generalization and exploring other model architectures where similar optimization dynamics might be at play are compelling directions. The observed correlation between gradient scale and Hessian in the context of class frequencies also opens up theoretical avenues for developing novel optimizers or enhancing existing ones to leverage this relationship more explicitly.

In summary, the paper provides a thorough examination of the challenges posed by heavy-tailed class imbalance in optimizing LLMs, revealing the underlying reasons for Adam's superiority over SGD. Acting on these insights can not only improve the training efficiency of LLMs but also inform optimization strategies in other domains facing similar issues.