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

Weight blocks for the gradient-Hessian correlation plots

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?

This is correct. The plots show the last layer for each problem. The last layer is a $[D \times C]$ matrix, where $D$ is the dimension of the last hidden layer and $C$ is the number of classes, and each $w_c$ correspond to a column of this matrix. We will make it clear in Appendix G.

Effect beyond last layer

If so, how does class imbalance affect other layers beyond the last layer?

This question is difficult to answer as there is no direct "mapping" between specific weights and classes beyond the last layer. It is likely that the slow performance of GD on the last layer leads to slower performance on earlier layers. For example, it might be that some feature relevant to a particular class can only really be learned once this class is sufficiently separated from the others. But formalizing this idea is far from trivial.

However, other ablation studies that appeared after our submission also provide evidence that one of the primary benefit of Adam over GD is to fix the effect of class imbalance in the last layer. We discuss those works below.

Correlation appears on realistic models

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.

While we agree that the toy model of §3.3 is not a realistic depiction of modern complex models, we do believe that the main take-away message, that the performance of GD suffers in heavy-tailed class imbalance, holds across models. Proposition 2 only describes a correlation between gradients and Hessians for rows of the last layer, but the performance gap between GD and Adam already appears when tuning only the last layer, keeping everything else frozen (see Figure 5). The effect described by Proposition 2 should still hold for complex models (assuming the inputs to the last layer do not change drastically to undo the effect of the scaling by the assignment property) and we observe this correlation beyond the linear model. Appendix G shows this effect on a small transformer on PTB, a CNN on the imbalanced MNIST dataset and a ResNet18 on the imbalanced ImageNet dataset.

We include an additional plot in the supplementary pdf for the response which shows that the correlation between gradients and Hessian observed in Fig 7 also holds on GPT2/WikiText, using the same training procedure as Fig. 1.

Related work

Could the authors please cite the relevant works by Zhang et al. and Xie and Li that also study the benefits of Adam?

We will include a discussion of those works. Follow-up works to Zhang et al. (2024a) from Zhang et al. (2024b) and Zhao et al (2024), (both of which where arXived after our submission) also give some insights on the benefits of Adam beyond the last layer, and we will include a discussion of those works.

In addition to their work showing that the Hessian of transformers has a block-diagonal structure (Zhang et al., 2024a), Zhang et al. (2024b) show in a follow-up ablation studies. Instead of looking for properties of the model, they study changes to the training procedure and show that the diagonal/sign-descent behavior of Adam is mostly redundant. Revisiting earlier layer-wise normalization approaches, they show empirically that most of the benefit of Adam can be recovered by using a learning rate per weight matrix rather than per coordinate (dividing by the average of the gradient magnitudes). They show that per-layer normalization achieves the same performance as Adam, as long as this scheme is not applied to the last layer. Zhao et al (2024) provide a similar ablation study, but further shows that most layers can be trained using unnormalized SGD updates, except for LayerNorm layers and the last layer.

Both the experiments of Zhang et al. (2024b) and Zhao et al (2024) indicate that the performance suffers unless the last layer is normalized independently for each class. Although neither paper specifically shows that this effect is due to class imbalance, their observations support our result that the primary benefit of Adam over SGD comes from the per-class normalization of the last layer to address class imbalance, and show that the impact of choosing an SGD or Adam-style update on the rest of the network has a smaller impact.

---

References

• Zhang et al. (2024a)
  Y. Zhang, C. Chen, T. Ding, Z. Li, R. Sun, Z. Luo
  Why Transformers Need Adam: A Hessian Perspective
  https://arxiv.org/abs/2402.16788
• Zhang et al (2024b)
  Y. Zhang, C. Chen, Z. Li, T. Ding, C. Wu, Y. Ye, Z. Luo, R. Sun
  Adam-mini: Use Fewer Learning Rates To Gain More
  https://arxiv.org/abs/2406.16793
• Zhao et al (2024)
  R. Zhao, D. Morwani, D. Brandfonbrener, N. Vyas, S. Kakade
  Deconstructing What Makes a Good Optimizer for Language Models
  https://arxiv.org/abs/2407.07972

Clarification

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

We think MNIST and ImageNet might have been flipped in this comment. GD and Adam are closer at the end of the given budget on imbalanced ImageNet (Fig. 3) than on on imbalanced MNIST (Fig. 2), where GD appears to stall. We assume the intended phrasing was as given below and will comment on this. Please let us know if we misunderstood your question.

why SGD [does not] find the minimum in the imbalanced MNIST case but [eventually does] in the imbalanced ImageNet

Does GD get stuck on Imbalanced MNIST (Fig. 2)
GD does eventually find a good solution for MNIST, if run for much longer than Adam. See the longer training run in the supplementary pdf for the response.

The apparent "plateau" reached by GD is not due to being stuck in a local minimum, but rather due to the optimization progress being slow compared to Adam. This issue is less apparent on ImageNet, possibly because the problem is more complex and Adam cannot find a good solution in 100 steps (note that Fig. 2 (MNIST) shows 300 iterations while Fig. 3 (ImageNet) shows 1500).

We comment on this behaviour in Appendix C.2, showing an example on a linear model where GD appears to stall on a short horizon (100 steps) but converges if run for longer (10k steps). As this question is likely to come up for more readers, we will mention this point explicitly in the MNIST paragraph (L96-103), add a forward reference to §C.2, and add the longer MNIST run to the appendix.

Inconsistency in references to SGD or GD between caption/figures
Those are indeed typos, apologies for the confusion. We will double-check those references.

Remaining comments
We will fix the typos and add references to the appendix to highlight the results on other language tasks.

Does the conclusion still hold when the parameters are not separable

When the parameters are not separable, does the conclusion still hold? 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?

We assume you have in mind phenomena such as the one identified by Nacson et al. (2019), where GD can be slower than normalized methods to converge (in direction) as the gradients and Hessian goes to 0. While our experiments use expressive models that can reach good classification accuracy on all classes, the observed behaviour does not require separability nor that the weights go to $\infty$. The behaviour we identify is different, and appears at the start of training, well before the gradients and Hessian go to 0.

In the supplementary pdf for the response, we show that the distinction between Adam and GD still appears when the model is regularized. The plots show the dynamics of the softmax loss (without the regularization) on a linear softmax model on the small Random Heavy-Tailed Labels dataset with varying levels of $L_2$ regularization ($\lambda \in [10^{-4}, 10^{-3}, 10^{-2}, 10^{-1}]$). For small values of $\lambda$, the weights of the model do not diverge to $\infty$ but the training dynamics are very similar to the unregularized case. Of course, if regularization is so large that the model cannot fit low-frequency classes, the loss per class after training reflects the class frequencies and the gap between GD and Adam is lessened. ($L_2$ penalty impacts performance on low-frequency more. Take a simple example with one class corresponding to $50\%$ of the data and $100$ classes corresponding to $.5\%$ of the data (each). Increasing $w_1$ can decrease the loss on ${\approx}50\%$ of the data. To achieve the same reduction on the remaining ${\approx}50\%$ of samples requires increasing $w_2, ..., w_{100}$, which is much more costly in $L_2$ penalty)

Those examples illustrate that while the observed behaviour does require the model to be expressive enough so that the model can fit low-frequency classes (otherwise, not fitting low-frequency classes will not be a problem), but it does not require separability nor that the weights diverge to $\infty$.

If the question regarding separability was aimed at the theoretical analysis, then we of course agree that the toy model of §3.3 relies on separability to obtain the result in Theorem 3. We believe that the conclusion that sign or per-class normalized methods can outperform GD holds more broadly. But finding a reasonable model where the computations can be carried out in closed form is far from trivial. We do not think that the gradient flow equations have a closed-form if we add an $L_2$ regularization term, and an analysis for discrete time gradient descent is even more complex.

Clarification on barcoded MNIST

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?

Good catch, thanks! You are correct, the all-0 10-bit pattern would be indistinguishable from the background. This was a typo in the appendix, which should read "$10 \times (2^{10}-1)$ classes". The code did generate $2^{10}-1$ barcodes, excluding the all-0 string (generate_combinations in code/src/optexp/datasets/barcoded_mnist.py:l27-33). We will fix the appendix.