Heavy-Tailed Class Imbalance and Why Adam Outperforms Gradient Descent on Language Models
Authors: Frederik Kunstner
Alan Milligan
Robin Yadav
Mark Schmidt
Alberto Bietti
Venue: NeurIPS 2024 spotlight (2024)
Keywords: language model
optimization
transformers
heavy-tailed
class imbalance
gradient descent
Adam
adaptive methods
sign descent
Primary area: optimization_for_deep_networks
TL;DR: We provide evidence that gradient descent struggles to fit low-frequency classes and that this leads to poor performance in the presence of heavy-tailed class imbalance as found in language modelling tasks while Adam does not suffer from this problem

Abstract: Adam has been shown to outperform gradient descent on large language models 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.

Heavy-Tailed Class Imbalance and Why Adam Outperforms Gradient Descent on Language Models
Frederik Kunstner kunstner&cs wbc.
Alan Milligan'
Robin Yaday
alanmil@cs. ubc.
mobiny12@stuce ubc. ca.
Mark Schmidt schmidtm@cs ubc.ca
Alberto Bietti
abietti@flatironnnsiii org
University of British Columbia
Canada CIFAR AI Chair
Flatiron Institute
Abstract
Adam has been shown to outperform gradient descent large language models 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 slow decrease the average loss as most samples come from infrequent words. On the other hand, Adam and sign-based methods 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 linear model with cross-entropy loss, we show that class imbalance leads 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- frwfrreuency classes while sign descent does not.
less sensitive this problem. To establish that this behavior
Introduction
The recent success of large language models such as GPT-3 (Brown et al., 2020) and its successors has relied on costly training procedures at unprecedented scale. A key ingredient their training is the Adam optimizer (Kingma and Ba, 2015), which outperforms stochastic gradient descent (SGD) on language problems by a large margin. Despite this large performance gap, we have poor understanding of why Adam works better and it has been difficult to find new optimizers that consistently improve over Adam (Schmidt et al., 2021). Not only is it computationally difficult to validate new optimizers large models, but also lack theoretical guidance; we do not know what "problem" Adam solves to outperform SGD.
The success of Adam language transformers has been well documented. Multiple works have found metrics or statistics that correlate with the improved performance of Adam. showing that yields uniform updates across parameters despite imbalanced gradients (Liu et al., 2020), gives better descent direction than the gradient (Pan and Li, 2023), and takes a path over which a robust variant of the condition number is smaller (Jiang al., 2022). But these observations do not provide a mechanism explaining what property of the problem leads to the improved performance of Adam.
Plausible mechanisms have been put forward, but they do not provide complete explanation. Zhang et al. (2020b) show that Adam-like methods are more resilient to heavy-tailed noise, which seems more prominent in language than in vision tasks. But noise is not the primary cause of the gap, it already appears in deterministic training (Kunstner et al., 2023). An alternative hypothesis is that the magnitude of the gradient and Hessian are correlated, which justifies clipping (Zhang al., 2020a). But to justify methods that normalize element-wise, like Adam and sign-like methods, we additionally need the gradient and Hessian be correlated across parameters (Crawshaw et al., 2022). While there empirical evidence for this behavior in neural networks, we do not have a good understanding of why this occurs, nor why this would be more pronounced on language rather than vision tasks.
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

a)
Samples/class
b)
Overall loss
C)
SGD
d)
Adam
# samples
10
Train loss
10
10
10-
10
10
0-
102
10
5k
6k
10k
15k
5k
10k
15k
Class index (sorted)
Step
Step
Step
SGD (with momentum)
2 10% samples, least freq. classes 25 10% samples, most freq. classes
Adam (with momentum)
Figure 1: Gradient descent does not make progress on low-frequency classes, while Adam does. Training GPT2- 2-lllll on WikiTextt-13.. (a) Distribution of the classes sorted by class frequency, split into groups corresponding = =10% of the data. (b) Overall training loss. (c, d) Training loss for each group using SGD and Adam. SGD makes little to no progress on low-frequency classes while Adam makes progress on all groups. (b) is the average of (c, d) for the respective optimizer.
1.1
Contributions
Our goal is to answer the following question: what is the "problem" that makes SGD slow on language tasks, that Adam "fixes" to perform better?
We argue the problem is what we call heavy-tailed class imbalance, where rare classes account for a large fraction of the data. Language data is imbalanced as some words are much more frequent than others, typically following a power-law. common modeling assumption is Zipf's law, where the kth most frequent word has frequency x 1 1/k (Piantadosi, 2014). For language tasks framed as next-token prediction, this property is reflected in the tokens and leads to heavy-tailed class imbalance. This contrasts with typical vision datasets such as MNIST, CIFAR, and ImageNet, which are curated to have uniform classes, but also with imbalanced problems with small number of classes. For example, in binary classification, extreme imbalance implies the minority class has limited impact on the loss; with an imbalance of 99:1, only 1% of the data comes from the minority class.
The performance gap arises because SGD makes slow progress on rare classes, see Figure 1. On binary problem, slow performance on 1% of the data need not have large impact on the average loss if we make fast progress on the remaining 99% of the samples. In contrast, the heavy-tailed class imbalance found in language tasks makes possible for low-frequency classes account for most of the data and significantly contribute to the loss, leading to slow performance overall.
We show that heavy-tailed class imbalance makes SGD slow across tasks in Section 2. We show that modifying vision datasets to exhibit heavy-tailed imbalance leads to slow progress with SGD on architectures where the performance gap with Adam is typically smaller. The impact of heavy-tailed imbalance can even be seen on linear models. Additionally, the performance of SGD improves with techniques that address imbalance such as upweighting rare classes.
Our findings provide simple model where Adam outperforms SGD, a softmax linear model under heavy-tailed class imbalance, which we analyze in Section 3. We show empirically that correlation between the magnitude the gradient and Hessian across coordinates, used to justify the benefits of Adam, appears naturally even on a linear model with class imbalance. We provide intuition as to how this pattern emerges through an assignment mechanism that leads to a correlation between class frequencies and the magnitude of the gradient and Hessian across parameters. We additionally prove that, on a simple dataset and in continuous time, GD is slow on low-frequency classes while sign descent is insensitive to the class frequencies.
We do not claim that class imbalance the only reason Adam outperforms SGD, as other properties of the data or architectures likely also contribute to this gap. Instead, we show that Adam consistently outperforms SGD under heavy-tailed class imbalance. The difficulty of minimizing the loss of minority classes has been explored for binary problems or problems few classes (Anand et al., 1993; Francazi et al., 2023), but the recent scaling of large language models to predictions over more than 100 000 classes puts the problem on new scale. Our findings indicate that heavy-tailed class imbalance has a significant impact on training performance and should consideration for future optimizers to perform well on language and other tasks exhibiting heavy- class imbalance.

---

Imbalanced MNIST
a)
MNIST
b)
GD
d)
Adam
Train Loss
10
10
10-
5-
5-
100
200
0
100 Epoch
200
100
200
0
100
200
Epoch
Epoch
Epoch
GD (with momentum)
2 50% samples, least freq. classes 2 50% samples, most freq. classes
Adam (with momentum)
Figure 2: Adam outperforms GD for training a CNN under heavy-tailed class labels. (a) Perfor- mance on the MNIST dataset. (b) Performance on a modified MNIST with two groups of classes. The first group consists of the 10 original classes with 5k samples each, while the second consists of #10k added classes with 5 examples each. (c, d) Performance of GD and Adam on the two groups. The initial loss is higher for imbalanced MNIST as there are =10 classes instead of 10, leading to a loss of - 1og(1/10*) = 9.2 for uniform prediction instead of - log(1/10 = 2.3.
2 Experimental results and ablation studies
Figure suggests a correlation between class frequencies and optimization performance that impacts SGD more than Adam. The goal of this section is to verify that (i) class imbalance is root cause for the performance gap between SGD and Adam, and (ii) whether this gap can be reproduced with simpler algorithms, such as deterministic optimizers, or using sign descent proxy for Adam.
To test these hypotheses, perform experiments focusing the training loss understand what makes optimization difficult. We use a simple training procedure, with a constant step-size tuned by grid search. For visualization, we split the data into groups of classes with similar frequencies, as in Figure 1. For instance, for 10 groups, the first group corresponds to RJ x10% of the samples from the most frequent classes. This grouping is only used for visualization and does not affect training. The models, datasets and training procedures are described in Appendix A.
objective is to
In Appendix B, we give additional information and additional ablation experiments on language models. We show that the heavy-tailed class distribution appears across datasets and tokenizers, and that the separation across class frequencies observed the training loss in Figure also affects the validation loss. We show that similar dynamics appear on smaller language models, including when training only the last layer while keeping the embedding and attention modules frozen initialization. Finally, we show that stochasticity is not necessary to reproduce the impact of heavy-tailed class imbalance, and that it also appears when using deterministic updates (i.e., GD instead of SGD), As result, we use deterministic updates whenever possible, denoted by GD in the figures.
2.1 Reproducing the frequency gap with vision models
Language transformers are often contrasted with vision CNNs, where we do not see a large perfor- mance gap between SGD and Adam. Our hypothesis is that key differentiation between the two settings is the heavy-tailed class imbalance present in language data. In this section, we show that making heavy-tailed vision datasets leads to slower performance with SGD and larger performance gap with Adam. These experiments show that heavy-tailed imbalance has a significant impact on performance and can make otherwise "easy" problem into "hard" one for SGD.
CNN. We first use a CNN on a variant of MNIST with heavy-tailed class imbalance. We augment the dataset to have two equally-sized groups of classes with relative frequency difference of 1000. The first group consists of the original 10 classes with #55k samples/class. For the second, we create ss 10k new classes with 5 samples/class. We create "barcode" in a corner of the image, see Appendix A. The performance of GD and Adam is shown in Figure 2. On the original MNIST dataset, both optimizers drive the loss to O, and Adam still makes progress on both groups in the imbalanced case. But on the imbalanced variant, GD makes almost no progress half of the data corresponding to the low-frequency classes and progress stalls. However, eventually converge run for much longer (see Appendix D.2), indicating that the problem is one of slow optimization rather than getting stuck in local minima.
classes by copying existing images and adding
ResNet. We replicate this effect with a ResNet18 on an imbalanced variant ImageNet. We subsample classes with frequencies Ty 1/k and compare against uniform subset with a similar
3

---

a)
Small ImageNet
b)
Imbalanced ImageNet
c)
SGD
10- d)
Adam
10-
10
Train loss
5-
0
500
1000
1500
500
1000
1500
500
1000
1500
0
600
1000
1500
Step
Step
Step
Step
SGD (with momentum) Adam (with momentum)
x 10% samples, least freq. classes
2 10% samples, most freq. classes
Figure 3: Adam outperforms SGD for training ResNet under heavy-tailed class labels. (a) Per- formance on a subset of ImageNet and (b) an imbalanced subset of ImageNet with class frequencies Tk & 1/k. (c, d) Performance of GD and Adam on groups corresponding to =10% of the data.
number of samples. In Figure 3, we see that SGD and Adam perform similarly on uniform data but performance gap appears across class frequencies the heavy-tailed imbalanced dataset. As in Figures and 2, SGD is slower on imbalanced data, especially on low-frequeccy classes.
Vision Transformers. This performance gap also appears with vision transformers (ViTs). In Ap- pendix C, we see that SGD and Adam both perform well on ImageNet, but exhibit a similar perfor- mance gap in Figure on the imbalanced variant. While ViT's may require more data, data augmentations, or regularization to generalize as well as ResNets (Steiner et al., 2022), there does not seem to be large gap between SGD and Adam without class imbalance.
2.2 Reproducing the frequency gap with a linear model on uniform data
To highlight that heavy-tailed imbalance alone can lead the observed difficulties, we reproduce this behavior in simple setting: softmax linear model with cross- cross-entropy loss. We create a dataset where the class frequencies approximate MT < 1/k and draw samples uniformly from [0. in d dimensions, independently of the label. While there is no relationship to learn, the optimization problem is still well posed and linear model can separate the data if n & d. As on the transformer of Figure 1, GD makes less progress on low-frequency classes than Adam, shown in Figure 4.
This example illustrates that a problem that might look innocuous at first is hard to optimize with GD due to heavy-tailed imbalance, while the performance of Adam is less negatively impacted. Nonetheless, imbalance alone not sufficient to make GD slow. possible to generate pathological datasets with heavy-tailed imbalance where GD fits all classes fast, by making all the samples (close to) orthogonal. In this case, each sample is learned independently of the others, and there is no difference across classes. However, perfectly orthogonal data is unlikely, especially as we expect samples from similar classes to be assigned a similar (correlated) representation. We discuss this issue and give additional examples the linear model in Appendix D.
2.3 Interactions between optimizer and imbalance
We have shown that heavy-tailed class imbalance lead different performance across class frequencies, but it is not clear which component of the training process has the highest impact on this behavior. We next experiment with simple algorithms to answer the following questions. (i) Is the impact of class imbalance due to stochasticity, or does it happen with deterministic training? (ii) Which component of Adam leads to an improved performance? and (iii) If imbalance is the problem, can we improve the performance of SGD by reweighting the losses?
Class imbalance already impacts deterministic optimization. A natural hypothesis to explain the impact of class imbalance is that it may be due to small batch sizes in SGD; rare classes could be sampled less often, and thus learned more slowly. have little impact on the gap between SGD and Adam (Kunstner et al., 2023). Our experiments in Figures 2, 4 and 5 and further examples in Appendix B.3 reproduce the dynamics of Figure I with full batch GD and Adam, indicating the problem already arises in the deterministic setting.
the other hand, stochasticity has been found to
Adam and sign descent both perform well under imbalance. Following Kunstner et al. (2023), we check whether the benefit of Adam is due to change in the magnitude of the update or its direction. Changing the magnitude as in normalized GD known perform better on separable problems (Nacson et al., 2019), while the benefits of Adam have been attributed to the change of direction close

---

a) Samples/class
b) Overall loss
c)
GD
d)
Adam
10
# samples
Train loss
10º
5-
5-
0
0-
10
102
100
150
50
100
150
150
Class index (sorted)
Epoch
Epoch
Epoch
GD (with momentum)
= 9% samples, least freq. classes
Adam (with momentum)
x 9% samples, most freq. classes
Figure The impact of heavy-tailed class imbalance is reproducible with linear models. Softmax regression on synthetic data. The inputs are drawn from uniform distribution on [0, 1] The target classes are heavy-tailed (a) and independent of the inputs, but the model can still fit the data overparamettrized. (b, d) Overall training loss and performance of GD and Adam each subset.
to sign descent (Tieleman and Hinton, 2012; Balles and Hennig. 2018). We compare the performance of GD, Adam, normalized GD and sign descent, with and without momentum, for training the last layer of a small transformer in Figure 5 and on additional problems in Appendix E. Normalization and momentum helps across problems, but they have less impact the performance gap across class frequencies than changing the update direction. Sign descent and Adam have a similar performance.
Upweighting low-frequency classes can help. Given our hypothesis that the performance gap between (SJGD and Adam is due class imbalance, expect interventions directly targeting imbalance to improve performance. In Appendix E. we show that upweighting the loss of low- frequency classes can improve the performance of SGD. While reweighting is not complete solution as it changes the objective function, this experiment supports the hypothesis that the optimization problem is due heavy-tailed class imbalance.
An investigation on linear models
Heavy-tailed imbalance already leads to slow performance the linear softmax model of Figure but we do not have a good understanding of why GD becomes slow while Adam is less affected. In this section, we explore the effect of heavy-tailed class imbalance the special case of softmax linear models, showing that it leads to correlated, imbalanced gradients and Hessians. In Section 3.1, we give an example a quadratic where imbalanced Hessians lead performance gap between GD and Adam. In Section 3.2, we show that class imbalance leads to imbalanced gradients and Hessians that are correlated with class frequencies through an assignment mechanism, showing that this pattern emerges naturally. Finally, we prove that on simple imbalanced problem and in continuous time, GD is slow on frweurey classes while sign descent fast all classes in Section 3.3.
3.1 Intuition on a weighted quadratic problem
Consider the following toy problem which is purposefully oversimplified provide high-level intuition about the optimization dynamics. Suppose we have functions fi. corresponding to the losses for each class, that are on the same scale in the sense that gradient descent with step-size a makes fast progress on any fs. For concreteness, take fi(w) = : ///l where GD with step-size of converges one step. Instead of running GD on each function independently. suppose we run GD on the weighted average / (Wi,, ..., We) Ei= / (w:) with positive weights 2 ... 2 Res Σ; ग. 1, corresponding to the class frequencies. If these weights span multiple orders of magnitude, we expect similar behavior as in Figures to 5, as illustrated in Figure 6. GD makes slow progress on functions with low weights as the gradient w.r.t. Wy scaled by Tk.
w = w (t-1) a. fr (w( (t-1)) = (1 a
This slow convergence on functions with low weights cannot be fixed by increasing the step-size, as increasing it beyond 1 I/m: would cause instabilities on the highest-frequency "class" f1- The problem is that use the same step size for all functions, which have different scales. Adam and sign descent
5

---

Overall loss
GD (-m)
Adam (-m)
NormGD (-m)
Sign (-m)
Train Loss
10-
10-
10
10-
104
500
1000
500
1000
500
1000
500
1000
500
1000
GD (+m)
Adam (+m)
NormGD (+m)
Sign (+m)
Train Loss
10-
10-
10
10
10-
500
1000
500
1000
500
1000
500
1000
500
1000
Epoch
Epoch
Epoch
Epach
Epoch
GD
(-m)
Adam (-m)
NomGD (-m) Sign (-m)
GD (+m)
Adam (+m)
NormG (+m)
Sign (+m)
Figure 5: Sign descent, as simplified form of Adam, performs well on low-frequency classes. Training the last layer of a simplified one-layer transformer with GD, Adam, normalized GD, and sign descent, with and without momentum m+) Momentum and normalizing the magnitude help but have smaller effects than using sign descent, which recovers similar dynamics Adam.
are less sensitive to this problem as their updates are independent of Tk.
(t-1) Th (W (1 1) wi
(t) 00gg wi
(t-1)
(ignn(/(( (w))
J' )-
While sign descent or Adam with a fixed step-size need not converge and can oscillate around the minimum, they perform much better in early iterations, independently of mi.
Another perspective is that the imbalance in the weights Ti, Te makes the problem ill-conditioned. The weights not only affect the gradient of f but also its Hessian, which is Diag( [ Te.). A common intuition for Adam is that using the magnitude of the coordinates of the gradient preconditioner is a good proxy for the Hessian diagonal (Duchi et al., 2011; I; Kingma and Ba, 2015), which would also lead to larger step-sizes for coordinates with small πy. While this does not hold in general (Kunstner et al., 2019), the gradient can be a reasonable approximation to the Hessian on this problem. The gradient is (xt (1,w1,-..
If the weights ...... Te vary by orders of magnitude more than the parameters |w]] | [Wel. the gradient and Hessian will be correlated, and preconditioning by the gradient magnitude Hessian diagonal will yield similar directions.
3.2 Correlations between the magnitude of the gradient and Hessian across coordinates
What is lacking to explain Adam's improved performance is an understanding of how a correlation between the gradient and Hessian arises in realistic problems. This feature has been observed on neural networks, but we do not yet know why it appears, even on the softmax linear problem. The caricature of the diagonal quadratic problem of the previous section provides some intuition, but does not directly apply to the softmax linear model of Figure 4 as that problem is neither quadratic nor separable. Nonetheless, a similar pattern emerges in the rows W1. We of its parameter matrix W E gexd, the magnitude of the gradient and Hessian across rows and the class frequencies become correlated during training due to class imbalance. In this section, we establish this observation empirically and provide mechanism for how it emerges.
In Figure 7, we show the gradient norm against the Hessian trace with respect to each row wa throughout the trajectory of Adam on the softmax linear model of Figure 4. While there is no corelaatio at initialization, the gradient and Hessian blocks become correlated with class frequencies during training and become imbalanced. This imbalance in the diagonal blocks is the main feature of the Hessian as the than off-diagonal blocks orders of magnitude smaller, shown in Figure 9, Similar dynamics occur with GD, although only on high-frequene classes as GD makes little progress on freeueeuyncy classes, Appendix F. This correlation also appears in the last layer of large models such as GPT2-Sman small used in Figure 1, as shown in Figure 8.
To explain this behavior, we show that the impact of samples on the Hessian follows an assignment mechanism: if the model assigns samples to their correct class, the Hessian with respect to wy is

---

Overall loss
Gradient descent
Adam
Sign descent
GD
Sign Adam
Loss
0
0
0
Iteration
Iteration
Iteration
Iteration
10% of the weights, smallest weights s 1 10% the weights, largest weights
Figure 6: Class-separation on the quadratic problem of Section 3.1 with weights T x 1/k. GD fits functions with low weights more slowly, while Adam and sign descent have the same dynamics across all functions and all the lines overlap every parameter W; is initialized wt;
primarily influenced by samples from class k, leading to a correlation between the magnitude of the gradient, Hessian, and class frequencies. To capture this effect, we introduce some notation and simplifying assumption. Suppose we have n samples with inputs X; E Idd and labels yi E [c]. where class lc has frequency Tą "s/n. The parameters of the linear model are W E Rexd We write p(x) = =(Wx for the predicted probabilities where a is the softmax, and summarize the data as
7 Ei_ x; x²
nh E.EEE
Eing. x X;: H = |
H ni Ei.. X,,!
Assumption (correct assignment). The model correctly assigns samples to class if it predicts k with non- neg-negliiib probability p on samples from that class (p(x;)s & = = = (1/c)) for x, from class k), and predicts k with near-random chance otherwise (p(x;)ı = (1/)) for where yi o k).
Proposition If initialized at Wo the gradient and Hessian of the loss L w.r.f. WE are
Vw ((Wo) = πタ 1x
V ((W)) = [1 -1 H),
(1)
During training, the model correctly assigns samples to class ie with probability p (Assumption 1),
Vw L p1 x² + O())
|
and 1| C1l
Tr(V a) as > (2)
22 C= p H O()),
p Tr(H*)
for classes where the frequency does not vanish too quickly, w2//).
The assumption that c + oo is used to obtain simple and interpretable equation in the correlation.
In practice, > 10x appears sufficient to make the dependence ग& appear, in Figures 7 and 8.
At initialization, Equation (1) shows that the Hessian blocks are uniform across classes while the gradients depend on Tg. If the data is uniform across classes (|x**|| ~ xx( while the frequencies differ by orders of magnitude, the the gradient blocks will mirror the class frequencies for high- frequency classes where ny > 1/c. This confirms the pattern observed initialization in Figures 7 and 8. During training, Equation (2) indicates a correlation between gradient norm and Hessian trace if classes have similar values of [|x*]| Tr(H+) and predicted probabilities p, confirming the behavior observed during training in Figures and 8 for the high frequency classes. As Adam fits low-frequency classes faster Figure 4, they have value of p closer (shown in Appendix F) and deviate slightly from the trend in Figure 7, as expected from Equation (2).
We now give the main intuition and defer the derivation of the asymptotics to Appendix G. We ignore off-diagonal blocks here, as they are orders of magnitude smaller than diagonal blocks (Figure 9), and show in Appendix G. 1.1.1.11 that they are expected to small.
Proof idea. Our loss is L( L(W)= _ Σ ((W, y:), where l is a softmax linear model,
(WW, x, y) =
exp(ze) with (2)) Σ, exp(z,)
Writing pxx) = or (Wx) for the vector predicted probabilities, the gradient and Hessian blocks are Vw ( x.
(1ly = k] p(x)x)x,
2 ((W, x, v/) p1x)x (1=()())))
(4)
The contribution of a sample (x, to the gradient w.r.t w& primarily depends whether the sample belongs to class & through the 1 = term, while the contribution the Hessian block depends on whether the model assigns that sample to class & through p(x)2- At initialization, p(x) 1/c
7

---

Correlation at
t= 10
t= 50
t= 100
t= 0
10-1-
Hessian
Tr(V f)
10
10 6
10-1
Grad. ,,
10
10-3
Grad.
10'
10-1
Grad. WV
10'
10-1
Grad.
10'
V. fll
IV All
least freq. classes (=s 10% samples)
most freq. classes (s 10% samples)
The gradient norm and Hessian trace across blocks become correlated during training, over the path taken by Adam in training the linear model of Figure 4. The blocks correspond to the rows WI;.., W. of the parameter matrix W. The color indicates the class frequency, showing that lower (higher) frequency classes have smaller (larger) gradient norm and Hessian trace.
Figure 7:
for all samples, and averaging the terms in Equation (4) yields Equation (1). Highlighting this effect during training is more challenging due to the dependency on the predictions. However, if W start assign samples to their correct classes (Assumption 1). we can obtain similar decomposition as Equation (1). For a given class k; the probabilities for correct labels are all p while the probabilities for incorrect ones are bounded by (1//c), which vanishes in the limit of + co.
This assignment mechanism explains why the gradient, Hessian, and class probabilities become correlated on the linear model. While the gradient does not directly approximate the Hessian, the main feature of the imbalance in the Hessian comes from the weighting by the class frequencies क yyy.. Tes which is present in both the gradient and the Hessian, as shown in Figures 7 and 9. This correlation is not a global property of the problem, as there are parameters for which the opposite pattern holds, see Appendix F, but it appears during training if the optimization algorithm makes progress. While the per-coordinate normalization of Adam or sign descent was not designed specifically address class imbalance, they appear to benefit from this property to make faster progress.
Our results complement prior work on optimization with class imbalance on problems with two or few classes, which argued that the gradient is dominated by the majority class, and result is biased towards making progress on the majority class at the expense of the minority class (Anand et al., 1993; Ye al., 2021; Francazi al., 2023). While this explains why GD might not make fast progress on rare c classes, it was not clear why this would lead to slow performance on average, especially under heavy-tailed imbalance where there no "majority". Our results show that, in addition to imbalance in the gradients, class imbalance leads to optimization difficulties through imbalanced Hessians.
3.3 Improvement of sign-based approaches over gradient descent
While the above arguments provide high-level intuition as to why the gradient might reasonable proxy for the Hessian, it remains difficult to formally describe this effect and prove the benefits of Adam over GD without strong assumptions. Doing so would require a fine fine-grained analysis of the dynamics, as the correlation only appears during training. To obtain a provable a guarantee highlighting the benefit of baged methods, we consider a stripped- problem where the only difficulty lies in the class imbalance:
Simple imbalanced setting. Consider e. classes with frequencies Ti,... Te where all samples from a class are the same, x; ek if k, where eş the kth standard basis vector in IRC
This setting is trivial as a possible solution is W al with - or taking one step of gradient descent with an arbitrarily large step-size. However, we will see that the dynamics with small step- sizes already exhibit the separation by class frequencies observed experimentally. In this simplified setting, we show that the continuous time variant of gradient descent, gradient flow, and sign descent proxy for Adam, obtain qualitatively different convergence rates (proof in Appendix H).
Theorem 3. On the simple imbalanced setting, gradient flow and continuous time sign descent initialized W minimize the loss of class k. (t()) (og(a(W(()))) (W)))))) at the rate
Gradient flow: (t) = ((1/mxt),
Continuous time sign descent: (x(() (ae
8

---

Correlation at
t= 10
t= 100
t= 1000
t= 10000
t=0
20
Hessian
Tr(V2
10-1
10.5.
10-5
Grad. 10-1
10-5 Grad. 10
10-3 Grad. 10-1
10-5 Grad. 10-1
10-5
Grad. 10-1 w/ fll
|, fll
[V All
least freq. classes (=s 10% samples)
most freq. classes (=110 10% samples)
Figure 8: The Headienn--Hessian blocks also become correlated in the last layer of large models. Reproducing Figure on the transformer of Figure 1. Evolution of the gradient norm and Hessian trace for each row w. of the last layer through optimization with Adam. Colors indicates class frequency. Lower (higher) frequency classes have smaller (larger) gradient norm and Hessian trace.
The difference between the sublinear rate of gradient flow (1/t) and linear rate of sign descent (e ') similar to existing results for separable logistic regression, where normalized updates converge faster as they keep increasing the margin despite small gradients (Nacson et al., 2019). While the setting studied here is separable, we still observe the separation across class frequencies on problems that are not separable, either because the problem has examples with different output for the same inputs, as in Figure 1, or when adding regularization, as in Appendix D.3. The novel element is that the convergence of gradient flow strongly depends on the class frequencies π, while the convergence of sign descent independent the class frequencies.
This setting admittedly oversimplified and does not capture some of the features observed in our experiments. For example, in Theorem 3, the loss is monotonically decreasing for all classes. This no longer holds once we introduce bias term and the loss from low-frequency classes will instead first increase, can be seen for example in Figure 4. This setting is also biased towards sign descent, as the inputs are aligned with the basis vectors. Finally, the problem is inadequate step-sizes, as it can be solved in one large step. On data with non-orthogonal classes, large step-sizes would lead to training instabilities and oscillations in the loss of frequent classes, as can be seen in Figures 2 to 5. Nevertheless, this result formally establishes the benefit of sign-based updates and we believe captures the key difficulty encountered by GD under heavy-tailed class imbalance.
study large
Discussion and limitations
Interaction with stochasticity. Our experiments include both stochastic and deterministic training regimes and show that stochasticity is not the cause of the slow performance of SGD on low-frequency classes, as already appears between full batch GD and Adam. This observation is consistent with prior work showing that the performance gap between SGD and Adam language transformers already appears with deterministic training (Kunstner et al., 2023). However, we do not attempt to quantify the interaction between stochasticity and class imbalance and leave it for future work.
Training performance vs. generalization. Our main focus is on optimization performance. Our observations need not generalize to the validation loss, especially settings prone to overfitting, as good training performance may lead to overfitting classes with few samples (Sagawa et al., 2020). However, some form of memorization might be needed in long-tailed settings (Feldman, 2020), and SGD cannot even fit the training data, generalization cannot be good. On the transformer of Figure I, we observe similar dynamics across frequencies on the validation loss, shown Appendix B.2. Training dynamics on the empirical and population loss are also often similar, particularly early in training (sce, e.g., Nakkiran et al., 2021: 1: Ghosh et al., 2022), and the one-pass training regime commonly used in large language models might mitigate those issues by blurring the line between train and test loss.
Additional difficulties due to text data. We study the effect of the distribution of the classes, the nex/ token to predicted, but other optimization difficulties might arise from the heavy-tailedness of text data. For example, the sequence of tokens used as inputs to the embedding layer are also heavy-tailed. This imbalance might lead to slow progress for rare tokens with GD, giving another potential cause for a performance gap. With stochastic training, this imbalance leads to sparse updates

---

Hessian at t= 0
t= 10
t= 50
t= 100
0
2
4
6
The diagonal Hessian blocks are orders of magnitude larger than off-diagonal blocks. Showing the magnitude of subset of the Hessian blocks (logıo( [Tr(V) ()) for a [160 x 160] subset of the Hessian, sampling 40 classes log-uniformly and 40 input dimensions uniformly.
Figure 9:
to the embedding layer, close to the setting that initially motivated AdaGrad (Duchi et al., 2011) and follow-up works attempting to correct this frequency imbalance in the input tokens (Défossez and Bach, 2017; Li et al., 2022). Beyond the inputs, full sentences (Williams et al., 2015) and latent rules or mechanisms required to understand a paragraph (Michaud al., 2023) may also display heavy tails, and Adam could be beneficial if those are captured by intermediate layers (e.g., Meng et al., 2022; Wang al., 2023; Bietti et al., 2023). The choice of tokenization has also been shown to impact downstream performance, which has been attributed to the lack of samples on rare tokens (Gowda and May, 2020) and the improved efficiency of more uniform tokenizers (Zouhar al., 2023). Our results indicate that tokenization also has a large impact on optimization performance.
Difficulties due to architectures. Beyond the class distribution, additional optimization difficulties may arise from the architectures, due depth, signal propagation (Noci et al., 2022: He et al., 2023), vanishing gradients and higher order derivatives (Liu et al., 2020; Orvieto et al., 2022). The simplified transformer of Ahn et al. (2023) also exhibits many of the difficulties observed in the literature on regression instead of a classification problem. However, a phenomenon similar to the assignment mechanism could still explain the benefit of Adam. The oscillations in the loss observed at the feature level by Rosenfeld and Risteski (2023) suggests a link between subsets of the samples and subsets of the parameters. For example, if convolution filter detects specific background color and captures specific feature of the data, the magnitude of the gradient and Hessian at intermediate layers could be influenced by the relative frequency of the feature in the data, leading to another form of imbalance.
Recent ablations on the benefit of Adam for language transformer. Parallel to our work, recent investigations have looked into the benefits of Adam on language transformers. Zhang et al. (2024a) argue that the Hessian has a block-diagonal structure, with similar magnitude within blocks but very different magnitudes across blocks, and that Adam may improve performance by using a different step-size for different blocks. This hypothesis is supported by recent ablations studies. Zhang et al. (2024b) show that the element-wiss preconditioning in Adam is not necessary and can be replaced by single parameter across such blocks while maintaining performance, which they coin Adam-mini. Similarly, Zhao et al. (2024) show that the performance of Adam can be recovered by training most of the network with GS) except for the last layer and LayerNorm parameters. Both approaches still need to treat the last layer separately, either using a step-size for each row w. of the last layer in the case of Zhang et al. (2024b) or by using Adam to train the last layer in the case of Zhao et al. (2024). These observations complement our approach, which focuses on the impact of heavy- class imbalance on the last layer, and are consistent with our conclusion that one of the main benefit of Adam is to counteract the slow progress on rare classes by preconditioning the last layer.
Conclusion
We have shown that heavy-tailed class imbalance leads to a performance gap between (S)GD and Adam. This effect is reproducible across architectures and data types, but is most salient on language tasks which naturally exhibit heavy-tailed imbalance. As vision tasks are typically more uniform, imbalance is a key differentiating feature of the training difficulties in language tasks. The correlation between entries of the gradient and Hessian that occurs due to class imbalance provides justification for the intuition that Adam-like algorithms "adapt to curvature". We provide an explanation for how this correlation arises during training through the assignment mechanism and prove on a simplified setting that gradient descent performs poorly wrw-urege classes while sign descent does not.
10

---

Acknowledgements
We thank Greg d' d'Eon, Aaron Mishkin, Victor Sanches Portella, and Danica Sutherland for useful discussions and comments the manuscript. This research supported by the Canada CIFAR AI Chair Program, the Natural Sciences and Engineering Research Council of Canada (NSERC) through the Discovery Grants Group and the Digital Research Alliance of Canada (all anaaaanccccoo ca).
and was enabled by the support provided by the BC DRI
References
Kwangjun Ahn, Xiang Cheng, Minhak Song, Chulhee Yun, Ali Jadbabaie, and Suvrit Sra (2023). "Linear attention is (maybe) all you need (to understand transformer optimization)". In: arXiv preprint arXiv:2310.01082
Rangachari Anand, Kishan G. Mehrotra, Chilukuri K. Mohan, and Sanjay Ranka (1993). "An improved algorithm for neural network classification of imbalanced training sets". In: IEEE Transactions on Neural Networks 4.6, pp. 962-969.
Jimmy Ba, Jamie Ryan Kiros, and Geoffrey E. Hinton (2016). "Layer Normalization". In: Neural Information Processing Systems (NeurIPS), Deep Learning Symposium.
Lukas Balles and Philipp Hennig (2018). "Dissecting Adam: The Sign. Magnitude and Variance of Stochastic Gradients". In: International Conference Machine Learning (ICML). Vol. 80, pp. 413-422.
Lucas Beyer, Xiaohua Zhai, and Alexander Kolesnikov (2022). "Better plain ViT baselines for ImageNet-Ik". In: CoRR abs/2205.01588
Alberto Bietti, Vivien Cabannes, Diane Bouchacourt, Herve Jegou, and Leon Bottou (2023). "Birth of a Transformer: A Memory Viewpoint". In: Neural Information Processing Systems (NeurIPS).
Tom B. Brown et al. (2020). "Language Models are Few-Shot Learners". In: Neural Information Processing Systems (NeurIPS).
Vivien Cabannes, Berfin Simsek, and Alberto Bietti (2024). "Learning associative memories with gradient descent". In: International Conference on Machine Learning (ICML).
Michael Crawshaw, Mingrui Liu, Francesco Orabona, Wei Zhang, and Zhenxun Zhuang (2022). "Ro- bustness to Unbounded Smoothness of Generalized SignSGD". In: Neural Information Processing Systems (NeurIPS).
Alexandre Défossez and Francis R. Bach (2017). AdaBatch: Efficient Gradient Aggregation Rules for Sequential and Parallel Stochastic Gradient Methods.
Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei (2009). "ImageNet: large- scale hierarchical image database". In: Conference on Computer Vision and Pattern Recognition (CVPR).
John C. Duchi, Elad Hazan, and Yoram Singer (2011). "Adaptive Subgradient Methods for Online Learning and Stochastic Optimization". In: Journal of Machine Learning Research (JMLR) 12, pp. 2121-2159.
Vitaly Feldman (2020). "Does learning require memorization? a short tale about a long tail". In: Symposium on Theory of Computing (STOC), pp. 954-959.
Emanuele Francazi, Marco Baity-Jesi, and Aurélien Lucchi (2023). "A Theoretical Analysis of the Learning Dynamics under Class Imbalance". In: International Conference on Machine Learning (ICML). Vol. 202, pp. 10285-10322.
Philip Gage (1994). "A new algorithm for data compression". In: C Users Journal 12.2, pp. 23-38.
Nikhil Ghosh, Song Mei, and Bin Yu (2022). "The Three Stages of Learning Dynamics in High- dimensional Kernel Methods". In: International Conference on Learning Representations (ICLR).
Thamme Gowda and Jonathan May (2020). "Finding the Optimal Vocabulary Size for Neural Machine Translation". In: Findings of the Association for Computational Linguistics (EMNLP), pp. 3955- 3964.
11

---

Bobby He, James Martens, Guodong Zhang, Aleksandar Botev, Andrew Brock. Samuel L. Smith, and Yee Whye Teh (2023). "Deep Transformers without Shortcuts: Modifying Self-attention for Faithful Signal Propagation". In: International Conference Learning Representations (ICLR).
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun (2016). "Deep Residual Learning for Image Recognition". In: Conference on Computer Vision and Pattern Recognition (CVPR), pp. 770-778.
Abdolhossein Hoorfar and Mehdi Hassani (2008). "Inequalities on the Lambert function and hyper- power function". In: Journal of Inequalities in Pure and Applied Mathematics 9.2.
Sergey Ioffe and Christian Szegedy (2015). "Batch Normalization: Accelerating Deep Network Training by Reducing Internal Covariate Shift". In: International Conference on Machine Learning (ICML). Vol. 37, pp. 448-456.
Kaiqi Jiang, Dhruv Malik, and Yuanzhi Li (2022). "How Does Adaptive Optimization Impact Local Neural Network Geometry?" In: arXiv preprint arXiv:2211.0223
Diederik P. Kingma and Jimmy Ba (2015). "Adam: A Method for Stochastic Optimization". In: International Conference on Learning Representations (ICLR).
Taku Kudo (2018). "Subword Regularization: Improving Neural Network Translation Models with Multiple Subword Candidates". In: Annual Meeting of the Association for Computational Linguis- tics (ACL), pp. 66-75.
Frederik Kunstner, Jacques Chen, Jonathan Wilder Lavington, and Mark Schmidt (2023). "Noise is not the main factor behind the gap between SGD and Adam on transformers, but sign descent might be". In: International Conference Learning Representations (ICLR).
Frederik Kunstner, Philipp Hennig, and Lukas Balles (2019). "Limitations of the empirical Fisher approximation for natural gradient descent". In: Neural Information Processing Systems (NeurIPS), pp. 4158-4169.
Yann LeCun, Léon Bottou, Yoshua Bengio, and Patrick Haffner (1998). "Gradient Based Learning Applied to Document Recognition". In: Proceedings of the IEEE. Vol. 86. 11, pp. 2324..324.
Yan Li, Dhruv Choudhary, Xiaohan Wei, Baichuan Yuan. Bhargav Bhushanam, Tuo Zhao, and Guanghui Lan (2022). "Frequency-awarc SGD for Efficient Embedding Learning with Provable Benefits". In: International Conference on Learning Representations, ICLR.
Liyuan Liu, Xiaodong Liu, Jianfeng Gao, Weizhu Chen, and Jiawei Han (2020). "Understanding the Difficulty of Training Transformers". In: Conference on Empirical Methods in Natural Language Processing (EMNLP), pp. 5747-5763.
Mitchell P. Marcus, Beatrice Santorini, and Mary Ann Marcinkiewicz (1993). "Building a Large Annotated Corpus of English: The Penn Treebank". In: Computational Linguistics 19.2. pp. 313- 330.
Kevin Meng, David Bau, Alex Andonian, and Yonatan Belinkov (2022). "Locating and editing factual associations in GPT". In: Neural Information Processing Systems (NeurIPS).
Stephen Merity, Caiming Xiong, James Bradbury, and Richard Socher (2017). "Pointer Sentinel Mixture Models". In: International Conference on Learning Representations (ICLR).
Eric J. Michaud, Ziming Liu, Uzay Girit, and Max Tegmark (2023). "The quantization model of neural scaling". In: Neural Information Processing Systems (NeurIPS).
Mor Shpigel Nacson, Jason D. Lee, Suriya Gunasckar, Pedro Henrique Pamplona Savarese, Nathan Srebro, and Daniel Soudry (2019). "Convergence of Gradient Descent on Separable Data", In: International Conference Artificial Intelligence and Statistics (AISTATS). Vol. 89, pp. 3420- 3428.
Preetum Nakkiran, Behnam Neyshabur, and Hanie Sedghi (2021). "The deep bootstrap framework: Good online learners are good offline generalizers". In: International Conference on Learning Representations (ICLR).
Lorenzo Noci, Sotiris Anagnostidis, Luca Biggio, Antonio Orvieto, Sidak Pal Singh, and Aurélien Lucchi (2022). "Signal Propagation in Transformers: Theoretical Perspectives and the Role of Rank Collapse". In: Neural Information Processing Systems (NeurIPS).
12

---

Antonio Orvieto, Jonas Kohler, Dario Pavllo, Thomas Hofmann, and Aurélien Lucchi (2022). "Van- ishing Curvature in Randomly Initialized Deep ReLU Networks". In: International Conference on Artificial Intelligence and Statistics (AISTATS). Vol. 151, pp. 7942-7975.
Yan Pan and Yuanzhi Li (2023). Toward Understanding Why Adam Converges Faster Than SGD for Transformers. NeurlPS 2022 Workshop on Optimization for Machine Learning. arXiv/2306.0022)
Vardan Papyan, XY Han, and David L Donoho (2020). "Prevalence of neural collapse during the terminal phase of deep learning training". In: Proceedings of the National Academy of Sciences (PNAS) 117.40, pp. 24652-246666 24663.
Adam Paszke et al. (2019). "Py Toych An Imperative Style, High- Bigh-Performance Deep Learning Library". In: Neural Information Processing Systems (NeurIPS), pp. 8024-8035.
Steven T. Piantadosi (2014). "Zipf"s word frequency law in natural language: A critical review and future directions". In: Psychonomic bulletin & review 21, pp. 1112-1130.
Alec Radford, Jeff Wu, Rewon Child, David Luan, Dario Amodei, and Ilya Sutskever (2019). Language Models are Unsupervised Multitask Learners. Tech. Report.
Elan Rosenfeld and Andrej Risteski (2023). "Outliers with Opposing Signals Have an Outsized Effect on Neural Network Optimization", In: arXiv preprint arXiv/2311.044
Shiori Sagawa, Aditi Raghunathan, Pang Wei Koh, and Percy Liang (2020). "An investigation of why overparameterization exacerbates spurious correlations" In: International Conference on Machine Learning (ICML).
Robin M. Schmidt, Frank Schneider, and Philipp Hennig (2021). "Descending through a Crowded Valley Benchmarking Deep Learning Optimizers". In: International Conference on Machine Learning (ICML). Vol. 139, pp. 9367-9376.
Rico Sennrich, Barry Haddow, and Alexandra Birch (2016). "Neural Machine Translation of Rare
Words with Subword Units". In: Annual Meeting of the Association for Computational Linguistics (ACL).
Nitish Srivastava, Geoffrey E. Hinton, Alex Krizhevsky, lla Sutskever, and Ruslan Salakhutdinov (2014). "Dropout: simple way prevent neural networks from overfitting". In: Journal of Machine Learning Research (JMLR) 15.1, pp. 1929-1958.
Andreas Steiner, Alexander Kolesnikov, Xiaohua Zhai, Ross Wightman, Jakob Uszkoreit, and Lucas Beyer (2022). "How to train your ViT? Data, Augmentation, and Regularization in Vision Transformers" In: Transactions of Machine Learning Research (TMLR) 2022.
Christos Thrampoulidis, Ganesh Ramachandra Kini, Vala Vakilian, and Tina Behnia (2022). "Imbal- ance Trouble: Revisiting Neural- Geometry". In: Neural Information Processing Systems (NeurIPS).
Tijmen Ticleman and Geoffrey Hinton (2012). RMSPROP: Divide the gradient by a running average of its recent magnitude. Lecture notes http: / / /wwp///wwccoooo cs. toronto. /et/ htth:/rriiiieeeeeec/////////eele
pdf. Hugo Touvron, Matthieu Cord, Matthijs Douze, Francisco Massa, Alexandre Sablayrolles, and Hervé Jégou (2021). "Training data-efficient image transformers & distillation through attention". In:
International Conference Machine Learning (ICML). Vol. 139, pp. 10347-10357.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin (2017). "Attention is All you Need". In: Neural Information Processing Systems (NeurIPS). pp. 1998-6008.
Kevin Ro Wang, Alexandre Variengien, Arthur Conmy, Buck Shlegeris, and Jacob Steinhardt (2023). "Interpretability in the Wild: Circuit for Indirect Object Identification in GPT-2 Small". In: International Conference on Learning Representations (ICLR).
Jake Ryland Williams, Paul R. Lessard, Suma Desu, Eric M. Clark, James P. Bagrow, Christopher M. Danforth, and Peter Sheridan Dodds (2015). "Zipf"s law holds for phrases, not words". In: Scientific 1 P. 12209.
reports 5.1,
13

---

Han-Jia Ye, De-Chuan Zhan, and Wei-Lun Chao (2021). "Procrustean Training for Imbalanced Deep Learning". In: International Conference on Computer Vision (ICCV), pp. 92-102.
Jingzhao Zhang, Tianxing He, Suvrit Sra, and Ali Jadbabaie (2020a). "Why Gradient Clipping
Accelerates Training: A Theoretical Justification for Adaptivity". In: International Conference Learning Representations (ICLR).
Jingzhao Zhang, Sai Praneeth Karimireddy, Andreas Veit, Seungyeon Kim, Sashank J. Reddi, Sanjiv Kumar, and Suvrit Sra (2020b). "Why are Adaptive Methods Good for Attention Models?" In: Neural Information Processing Systems (NeurIPS), pp. 15383-15333
Yushun Zhang, Congliang Chen, Tian Ding, Ziniu Li, Ruoyu Sun, and Zhi-Quas Luo (2024a). Why Transformers Need Adam: A Hessian Perspective.
Yushun Zhang, Congliang Chen, Ziniu Li, Tian Ding, Chenwei Wu, Yinyu Ye, Zhi-Quan Luo, and Ruoyu Sun (2024b). Adam-mini: Use Fewer Learning Rates To Gain More. Preprint.
A.Kiv/2406.16799
Rosie Zhao, Depen Morwani, David Brandfonbrener, Nikhil Vyas, and Sham M. Kakade (2024). Deconstructing What Makes Good Optimizer for Language Models.
Vilém Zouhar. Clara Meister. Juan Luis Gastaldi, Li Du, Mrinmaya Sachan, and Ryan Cotterell (2023). "Tokenization and the Noiseless Channel". In: Annual Meeting of the Association for Computational Linguistics (ACL), pp. 5184-5207.
14

---

Supplementary Material
Table of Contents
A Experimental details
15
B Language problems
Vision problems
20
D Linear models
22
E Alternative optimizers
25
F Dynamics of the gradient and Hessian
27
G Corrclation between the gradient and Hessian across blocks
31
H Continuous time GD and sign descent simple imbalanced problem
33
A Experimental details
This section documents the datasets, models, software, and experimental setup. The code is available at wttpg/tttttt.ooootttttonoee class- imwalance-eaeeeeeeeemeeaaaiaa sgd- adam.
A.1 Datasets
WikiText-103 (Merity al., 2017), using sequences of 024 tokens and the BPE tokenizer (Sennrich et al., 2016), with vocabulary of size 50 608.
WikiText-2 (Merity et al., 2017) is used in Appendix B.1 illustrate that other combinations of datasets and tokenizers lead to heavy-tailed distributions.
PTB (Marcus et al., 1993), using sequences of 35 tokens built from word-based tokenizer (basic_englist provided by torchtext), for vocabulary of size 9220. For deterministic runs, the validation set reduced training set, labeled TinyPTB.
MNIST (LeCun al., 1998).
ImageNet (Deng et al., 2009).
A.2 Custom datasets
The Random Heavy-Tailed Labels dataset is synthetic dataset exhibiting heavy-tailed class imbalance. The number of samples per class and the number of classes are picked to approximate a power-law distribution. We create "groups" of classes, where each class within group has the same relative frequency;
1 class with n samples,
2 classes with 2 samples,
2 classes with 2 samples.
group
group 2
The inputs drawn from uniform distribution [0. 1], independently of the class label. The inputs are d = (m + 1m dimensions, the number of samples is n = m 2 and the number of classes is 2m+1 1. We use two variants of the datasets; large in Figure Appendix E (m = 11, = 22 528, d = 24576, e = 4 4995) and a small one in Appendix D (m 8, n 248, d 22304 511).
The Barcoded MNIST dataset modified variant of MNIST. We start with 50k examples from the original MNIST dataset across 10 classes, and create 51 150 (5 x (10 x 210 1)) new
15

---

images. The new examples are copies of existing image with an added "barcode", a 10-bir number encoded in a corner of the image, as in the examples below. The class label a combination of the original class and this barcode.
2 1
4
5
6
7
9
The Barcoded-only dataset contains 10 x 21 classes with samples each. To obtain an im- balanced dataset, we combine the barcoded images with the original samples from the MNIST dataset to get 101 200 examples spread across 10 250 (10 x 2 + 10) classes classes; 10 240 with 5 examples per class and 10 classes with 5k examples per class, labeled MNIST+Bareoded
The Heavy Tailed ImageNet dataset subset of ImageNet (Deng et al., 2009), subsampled to exhibit heavy-tailed class imbalance. We sort the original 1000 classes by frequency and sample [1300/k] / images from the kth class, leading to n = 107 samples.
The Small ImageNet dataset is a uniform subset of ImageNet to contrast the with the heavy 000 samples.
tailed variant. We sample 10 images per class to get
A.3 Models
The 2-layer transformer used in Appendix B.3 transformer Vaswani et (2017), based on the PyTorch implementation of Transf Transf ormerrncoderrLay (Paszke et al., 2019).
Embedding + 2x [Attention - Linear ReLU + Linear] Classifier.
The model includes LayerNorm, dropout, and skip connections (He et al., 2016; Ba et al., 2016; Srivastava et al., 2014). The embedding dimension and width of the linear layers is 1000 and the attention modules use 4 heads.
The simplified transformer used in Figure 5 and Appendix B.3 does not use encoder blocks, and only uses attention:
Embedding + Attention + Classifier.
We remove LayerNorm, dropout, and the block [Linear ReLU + Linear] containing the non-linearity. In Figure 5, we freeze the embedding and attention layers at initialization, and only the last classification layer is trained. The model is then linear model on fixed feature transformation.
The GPT2-Small model (Radford al., 2019) is used in Figure 1. The blocks includes Layer- Norm, residual connections, and dropout on the embedding and dense layers. We use sinusoid positional encodings in the transformer architecture (Vaswani et al., 2017). The embedding dimension is 768, the width of the intermediate layers is 3072, and we use 12 encoder blocks with 12 self attention heads.
The convolutional network used in Figure 2 and Appendix a 2-layer convolution
Conv + Relu + MaxPool - Conv + Relu - MaxPool + Linear
The linear model used in Figures 4 and and Appendix E uses a bias vector.
The ResNet18 model (He al., 2016) is used Figure 3. Additionally, variant replacing the BatchNorm layers with LayerNorm is used in Appendix C.
The ViT model (Beyer et al., 2022) used Appendix follows the architecture of ViT- S/16 (Touvron et al., 2021). based on the vit vit-pytorch implementation saccidaannsvii-pyttrre v1. 6.56.5).
com/
A.4 Training procedures
Our primary focus is on the performance of the optimizers on the training error, using the simplest training procedure possible. We use a constant step-size throughout training, set by grid search. We start with a sparse grid powers of 10 [10-6 6 10-2 10a and increase the density to half-powers around the best step-size. The step-size is selected to minimize the maximum over 3 seeds of the training loss at the end of training. For some settings, this selection still produces runs that are unstable; the training loss is the smallest at the end but oscillates a lot during training, reaching
16

---

loss values that are orders of magnitude worse than at initialization. For those runs, we use the next smaller step-size, which has similar performance at the end but is more stable. We use the following batch sizes with gradient accumulation (computing the gradient through multiple passes)
The large transformer experiment in Figure uses mini-batches of 512 sequences of 1024 tokens.
The stochastic experiments with smaller transformer Appendix B.3 uses sequences of 35 tokens.
of 512
Both ResNet18 variants and the Simple Vision Transformer were trained using mini-batcher of 1024. The training images were normalized and randomly cropped to 224 x 224 pixels as is standard for ImageNet training.
Other experiments use the entire dataset to compute updates
Our experiments ran on a cluster using a mix A100, PI00, V100, and H100 GPUs. The large scale experiment in Figure took 3 days H100, while all other experiments ran in 2-8 hours. The total amount of compute used for this project is #3 GPU-years, including preliminary experiments.
A.5 Optimization algorithms
Given momentum buffers my initialized at mo and a (possibly) stochastic gradient 9t, we implement the update of GD, normalized GD and sign descent with heavy-ball momentum as
gt.
for gradient descent,
m = fm,-1 + de.
with de = 9r/ Gell2 for normalized GD, sign(Ot)
Ft+1 (3774
for sign descent.
For GD and Adam, we use the standard implementation PyTorch (Paszke et al., 2019). For algorithms,
either momentum with 0.9 (31 = 0.9 for Adam) or momentum (3 0, 31 = 0), indicated by solid lines and the legend (+m) for runs with momentum, and dashed lines and the legend (-m) for runs without momentum.
A.6 Summary of settings used
Table 1: Summary of models, datasets and batch-size used
Model
Dataset
Batch size
Used in
GPT2-Small
WT103
512
Figure and Figure 11
2-layer transformer
PTB
512
Figures 12, 25 and 31
1-layer transformer
TinyPTB
Full
Figures 13 and 23
l-layer transformer
TinyPTB
Full
Figure 5 (last layer only)
CNN
Barcoded MNIST
Full
Figure 18
CNN
MNIST
Full
Figures and 18
CNN
MNIST+Barcoded
Full
Figures 2, 18, 24, 25 and 29
Linear
Random HT labels, m=11
Full
Figures 4, 22, 25, 26, 32 and 33
Linear
Random HT labels, m=7
Full
Figures 19 and 20
Simple ViT
ImageNet
1024
Figure 16
ResNet18
Small and HT ImageNet
1024
Figures 3, 25 and 30
ResNet18+th 1+LN
Small and HT ImageNet
1024
Figure 15
Simple ViT
Small and HT ImageNet
1024
Figure 17
17

---

B Language problems
This section provides additional ablations on language models, showing that the impact of class imbalance holds across models of different sizes and using deterministic updates.
B.1 shows that the heavy-tailed distribution in text data occurs across datasets and tokenizers.
B.2 shows that the imbalanced training speed across frequencies translates to the validation loss.
B.3 shows that the imbalance training speed across frequencies and the gap between SGD and Adam can be reproduced with smaller transformers. This effect also appears when training only the last layer, and in the deterministic setting, comparing GD and Adam.
B.1 Class distribution for common datasets and tokenizers
Figure 10 provides additional examples of the heavy-tailed distribution of tokens using the basic english tokenizer in torchtezt (Paszke et al., 2019), Byte-Pair Encoding (BPE, Sennrich al., 2016; Gage, 1994) and Unigram (Kudo, 2018) on the PTB and WikiText-2 datasets. The relationship between the relative frequency rank i and and the relative frequency me is roughly we x 1/k:
Word
Unigram
BPE
#samples
PT8
109
WHRT ext2 esamples
107
10a
10
10'
10"
10a
101
10"
10
10"
10"
10'
10
10"
Token rank
Taken rank
Token rank
Figure 10: Different tokenizers and datasets lead to heavy-tailed token distributions. Comparison of word and subword tokenization (BPE, Unigram) on the PTB and WikiText2 datasets.
B.2 Effect of class imbalance on validation loss
In Figure 11, we show the validation error on same problem as Figure 1, training GPT2-Small on 103. The validation loss exhibits the same separation across class frequencies, and the faster progress Adam on low-frequency classes is also visible. While this trend does not hold for all the settings investigate, isolate the source of the training difficulties, the benefit of Adam on frequency classes does not immediately lead to overfitting.
settings use smaller datasets and deterministic training to
Overall loss
b)
SGD
C)
Adam
Validation loss
30-
10
10-
5-
5-
10k
15k
Step
Step
Step
Figure 11: The class-separation behavior of Figure holds on the validation loss. Same experiment
as Figure 1, training GPT2-Small on WikiText- 103, but showing the validation loss. (a) Distribution of the classes sorted by class frequency, split into groups corresponding R 110% of the data. (b) Overall validation loss. (c. d) Validation loss for each group using SGD and Adam. SGD makes little no progress on low-frequency classes while Adam makes progress on all groups. (b) is the average of (c, d) for the respective optimizer.
18

---

B.3 Smaller transformers and deterministic training
In Section 2.3, we argued that the qualitatively different behavior on low-frequency classes between SGD and Adam in Figure is not due stochasticity. In this section, we provide additional results showing that this behavior appears across multiple batch sizes on language transformers of different sizes and that it can be reproduced in the deterministic setting.
In Figure 12, we show that a similar qualitative behavior appears when training a smaller model (2-layer transformer) on a smaller dataset (PTB). In Figure 13, we repeat the experiment with a 1-layer transformer, trained in full batch on TinyPTB (the validation set of PTB). The separation between GD and Adam low-frequency classes in the deterministic settings is also visible in Figures 2, 4, 5 and 7 the main paper. These results indicate that stochasticity is not necessary to reproduce the behavior observed in Figure 1. Finally, we repeat the experiment but freeze all the layers except the last, and still observe this behavior in Figure 14.
# samples/class
Overall loss
SGD (+m)
Adam (+m)
samples
10
Train Loss
10
10-
10-
10
5-
100
102
10
Class index (sorted)
Epoch
Epoch
Epoch
SGD (with momentum)
x 1 10% samples, least freq. classes
Adam (with momentum)
z 10% samples, most freq. classes
Figure 12: Similar behavior as Figure 1 on smaller problem. Training 2-layer transformer on PTB with Adam and SGD using larger batch-sizes. As Figure 1, SGD makes little to no progress on low-frequency classes while Adam makes progress all subsets. Subplots: (1) Distribution of the classes and subsets of the data sorted by class frequency, each corresponding to ze 10% of the samples. (2) Overall training loss. (3, 4) Training loss for each subset for SGD and Adam.
Overall loss
GD (+m)
Adam (+m)
# samples/class
samples
10
Train Loss
10
10
10-
10a
5-
100
0-
10"
10?
101
200
100
Class index (sorted)
Epoch
Epoch
Epoch
GD (with momentum)
2 1 10% samples, least freq, classes
Adam (with momentum)
2 1 10% samples, most freq. classes
Figure 13: Similar behavior as Figure 1 on a one-layer transformer with deterministic updates. Trained on TinyPTB. Subplots: (1) Distribution of the classes and subsets of the data sorted by class frequency. (2) Overall training loss. (3, 4) Training loss for each subset for GD and Adam.
a)
Overall loss
b)
GD
c)
Adam
# samples/class
samples
10
Train loss
10-
10-
10-
10a
5
5-
5-
10
0
10P
10
10
50
100
150
100
150
0
50
100
150
Class index (sorted)
Epach
Epoch
Epoch
GD (with momentum) Adam (with momentum)
R 1 10% samples, least freq. classes
2 1 10% samples, most freq. classes
Figure 14: Similar behavior as Figure 1 when training only the last layer. Training the last
layer of a 1-layer transformer on PTB with Adam and GD with deterministic updates. Subplots: (1) Distribution of the classes and subsets of the data sorted by class frequency. (2) Overall training loss. (3, 4) Training loss for each subset for GD and Adam.
19

---

C
Vision problems
This section gives additional results on vision tasks to complement Section 2.1.
Figure 15 shows similar behavior on ResNet18 with LayerNorm instead of BatchNorm.
Figure 16 shows similar behavior with a vision transformer.
Figure 18 confirms that GD can solve the Barcoded MNIST variant without imbalance.
C.1 ResNet18 with LayerNorm
In Figure 15, we use the same settings Figure 3. training a ResNet18 on a uniform and unbalanced subset ImageNet, but replace the normalization layers with LayerNorm (Ba al., 2016) instead of
BatchNorm (loffe and Szegedy, 2015). We observe similar pattern as in Figure 3. Although Adam slightly outperforms SGD on the uniform dataset, the performance gap grows on the imbalanced one.
a)
Small ImageNet b) Imbalanced ImageNet
SGD
d)
Adam
Train loss
0-
2000 Step
4000
2000
4000
2000
4000
2000
4000
Step
Step
Step
SGD (with momentum) Adam (with momentum)
10% samples, least freq. classes ~ 10% samples, most freq. classes
Figure 15: Adam outperforms SGD on ResNet with LayerNorm under heavy-tailed imbalance. (a) Performance on uniform subset of ImageNet (b) and on an imbalanced subset with class frequencies n 1/k. l.. (c, d) Performance of GD and Adam across frequencies.
C.2 Vision Transformers
In Figure 16, we train vision transformer on the ImageNet dataset, without subsampling, to confirm that the training behavior is similar. While vision transformers might require more data or regularization than their ResNet counterparts to achieve comparable generalization performance, the optimization problem does not appear to be more difficult for SGD than for Adam.
ImageNet1K
Figure 16: Adam and SGD perform similarly training Vision Transformer with balanced Classes. Training loss on the full ImageNet dataset (without subsampling). There is little performance in training performance.
Train Loss
5
0
50000 Steps
100000
In Figure 17, we train the same vision transformer on the uniform and imbalanced subsets of ImageNet. As in prior experiments with vision data, the performance of Adam appears unaffected by the change in class frequencies while the performance of SGD degrades.
Small ImageNet
b) Imbalanced ImageNet
SGD
d)
Adam
Train loss
0
500
1000
1500
500
1000
1500
500
1000
1500
500
1000
1500
Step
Step
Step
Step
SGD (with momentum) Adam (with momentum)
x 10% samples, least freq. classes ~ 10% samples, most freq. classes
Figure 17: Adam outperforms SGD vision transformer under heavy-tailed imbalance. (a) Per- formance on a uniform subset of ImageNet (b) and on an imbalanced subset with class frequencies Th & 1 1/k. (c, d) Performance of GD and Adam across frequencies.
20

---

C.3
Sanity check on Barcoded MNIST
Figure in Section 2.1 showed that the performance gap between GD and Adam on the imbalanced variant of MNIST with barcoded images is larger than on plain MNIST. In this section, verify that the training difficulties encountered on the CNN on the imbalanced MNIST dataset Figure are indeed due to class imbalance. As create new images and new classes by adding barcode in the corner of existing images, could be that the dataset becomes harder to fit.
In Figure 18, we run Adam and GD to train the same network on the MNIST dataset only, the barcoded-only subset of the imbalanced MNIST and the combination of the two, leading to an imbalanced dataset. While Adam is faster GD the barcoded-only dataset, both algorithms reach negligible error within 200 steps. In contrast, on the combined imbalanced dataset MNiST+Barcoded, GD fails to make progress the low-frequency classes and stalls.
MNIST
BarcodedMNIST
MNIST + Barcoded
Train Loss
Train Loss
10
Train Loss
5-
5
100
200
300
0
100
200
300
100
200
300
Epochs
Epochs
Epochs
GD (with momentum)
Adam (with momentum)
GD optimizes on balanced barcoded data. Training a CNN on only the barcoded portion of the data, which has balanced classes. While Adam is slightly faster, both optimizers reach negligible error within 200 steps. As the level of imbalance is increased, GD performs increasingly worse than Adam.
Figure 18:
21

---

D Linear models
Section 2.2 showed that GD is already slow on linear models. We give additional details here.
D.1 discusses the impact of the distribution of the inputs, as it is possible to construct problems exhibiting class imbalance without negatively impacting training.
D.2 shows that while (S)GD appears stuck in some experiments, it is not due to being stuck in a local minima. eventually converges, although very slowly, if run long enough.
D.3 shows that while some of our datasets are separable, leading to weights going to oo, class imbal- ance also impacts optimization when the weights remain small, e.g. when using 12 regularization.
D.1 Impact of input distribution
Imbalance alone is not sufficient to induce slow performance of GD on low-frequency classes. It is possible to generate dataset with heavy-tailed class imbalance where GD fits all classes fast, by making all inputs X, (close to) orthogonal, (X, Xx) ~ 0 for i * j. If all samples are orthogonal, (x,, x,) 0 Vi * j. decomposition similar that used in the proof of Theorem shows that each sample is learned independently of the other, and class frequency has no impact. However, completely orthogonal data is rare. the last layer of neural networks, we expect samples from the same class to be mapped to similar representation (Papyan et al., 2020), phenomenon also observed under class imbalance (Thrampoulidis et al., 2022). Using a bias term also increases alignment between samples, as it is equivalent to adding dimension where each sample has the same value.
In the setting of Theorem 3, class imbalance has impact because samples from the same class collinear, even though samples from separate classes are orthogonal. A more realistic mixture model where samples from the same class are aligned ((xx, x,) > 6 y: = yjy but independent otherwise (1xx,, x,)l < if yss # Yj). as the setting of Feldman (2020) would also exhibit class separation. The class imbalance appears in Figure because we draw the inputs from high-dimensional uniform distribution on [0, 1] ensuring that for any two samples x;, xj. (x;, xj) >> 0. If the data was sampled from (0, 1) in sufficiently high dimension, the samples would be independent enough to avoid the slowdown due to class imbalance. We illustrate this in Figure 19, where we use smaller synthetic data with inputs drawn from ((1, 1) and (00, 1). The cero-mean data is be approximately orthogonal and does not exhibit a slow progress low- uwwfrequency classes.
(a) Aligned data sampled from (/1, 1)d
Overall loss
GD (-m)
Adam (-m)
Train Loss
5.0
5.0
2.5
2.5-
2.5
0.0
0.0
0.0
10
10
20
Epoch
Epoch
Epoch
(b) Independent data sampled from (00, 1)d
Overall loss
GD (-m)
Adam (-m)
Train Loss
5.0
5.0
5.0
2.5
2.5
2.5
0.0
10
20
10
20
10
20
Epoch
Epoch
Epoch
Figure 19: The distribution of the inputs can have a large impact on optimization. Linear model on the Random Heavy-Tailed Labels dataset, with Inputs sampled from (1,,, (a) and A/ 1) (b).
22

---

The behavior of GD on aligned data appears to be a better representation of the behavior of GD on language transformers, as we observe a performance separation per class frequency on GD, even when tuning only the last layer of a language transformer in Figure 5. Although the embedding weights are initialized to be zero-mean Gaussian noise, the representation of the tokens in a transformer are aligned, and this alignment increases with depth (Noci et al., 2022, e.g.).
D.2 An early iteration problem
As GD is slower than Adam at fitting the low-frequency classes, it might seem that GD does not fit the low-frequency classes at all. But when run for longer. GD converges and fits all classes. We show this behavior on the linear model and the CNN on imbalanced MNIST in Figure 20. This highlight that the difference between the algorithms primarily a difference the start of training. However, this "start" can be quite long. In the transformer of Figure 1, the average loss on 10% of the data corresponding to the least frequent classes is still higher than at initialization after 15k steps.
(a) Linear model on synthetic data Medium (1k steps)
Short (100 steps)
Long (10k steps)
Train Loss
Train Loss
Train Loss
2-
2-
2-
0-
0-
25
50
75
100
250
500
750
1000
2500
5000
7500
10000
Epoch
Epoch
Epoch
(b) CNN on MNIST
CNN on MNIST+Barcoded
CNN on MNIST+Barcoded
300 steps
10k steps
15
15
Train Loss
10
10
5-
5-
01
0-
100
200
300
0
5000
10000
Epochs
Epochs
Figure 20: Training with GD eventually drives the loss down for all classes. Using the same step-size for different horizons (100, Ik, 10k). GD eventually drives the loss down for all classes, but the loss for the least-frequen classes only decreases below its value at initialization after Ik steps. (a) Linear model on synthetic data, (b) CNN on MNIST.
23

---

D.3 Impact of regularization
The data used with the linear model of Figure separable, meaning the predicted probabilities for the correct class will converge to I while the magnitude of weights go to oc. This might lead to concerns that the observed behavior is tied to the weights growing without bounds. In Figure 21, we show that the gap between GD and Adam still appears with regularization limiting the magnitude of the weights. However, as regularization increased, the 1.2 penalty makes it difficult to fit low- frequency classes, the problem looks more like X |1-11 1 and the gap between the methods disappears.
A=10
10
== 10
A= 10
10-
10
NIL
5-
5-
0-
0-
100
200
160
200
100
200
100
200
NLL per group
10-
GD (+m)
5-
6-
5-
100
200
100
200
100
200
200
NLL per group
10-
10
10
10
Adam (+m)
0-
0-
0-
100
200
100
200
100
200
100
200
Epoch
Epoch
Epoch
Epoch
Figure 21: The separation between GD and Adam still appears when using Le regularization. Using varying levels of regularization À the linear model of Figure 4. The plots show the negative log-likelihood and do not include the Le penalty.
24

---

E Alternative optimizers
Figure 5 in Section 2.3 we compared GD and Adam to normalized GD and sign descent on the last layer of one-module transformer on TinyPTB, showing that Adam and sign descent perform similarly. We repeat this experiment on other settings here to confirm that sign descent leads to similar benefits as Adam on low-frequency classes, and that changing the direction, as in sign descent, has more impact than just changing the magnitude, as in normalized GD.
We also observe this behavior on the following problems:
Figure 22: A linear model Random Heavy-Taile Labels, in Figure 4.
Figure 23; A one-module transformer TinyPTB, in Figure 13, training all layers.
Figure 24: A CNN on MNIST+Barcoded, as Figure 2.
Overall loss
GD
m)
Adam (-m)
NormGD ( m)
Sign m)
Train Loss
10-
10-
10-
10-
10-
0
100
200
100
200
100
200
100
200
100 Sign (+m)
200
GD (+m)
Adam (+m)
NormGD (+m)
Train Loss
10
10
10
10
10
0-1
0-
0-
100 Epoch
200
0
100
200
100
200
0
100
200
0
100 Epoch
200
Epoch
Epoch
Epoch
Figure 22: All optimizers on the linear model of Figure 4.
Overall loss
GD (-m)
Adam (-m)
NormGD (-m)
Sign (-m)
Train Loss
10
10
10-
10-
10-
0-
100
200
100
200
100
200
100
200
100
200
GD (+m)
Adam (+m)
NormGD (+m)
Sign (+m)
Train Loss
10
10
10-
10-
100
200
100
200
100
200
100
200
100 Epoch
200
Epoch
Epoch
Epoch
Epoch
Figure 23: All optimizers on the transformer of Figure 13.
Overall loss
GD (-m)
Adam (-m)
Sign (-m)
NormGD (-m)
Train Loss
10-
10-
10-
10
10-
200
400
200
400
200
400
200
400
200
400
GD (+m)
Adam (+m)
Sign (+m)
NormGD (+m)
Train Loss
10-
10-1
10-
200
Epoch
Epoch
Epoch
Epoch
Epoch
GD (-m)
Adam
(-m)
NormGD (-m)
Sign (-m) Sign (+m)
GD (+m)
Adam (+m) NomGD (+m)
Figure 24: All optimizers on the CNN of Figure 2. First column: Overall training loss. Remaining:
Loss by frequency groups for each optimizer, with and without momentum (+m, bottom/ m,, top).
25

---

E.1
Up-weighting low-frequency classes can improve the performance of SGD
To support Section 2.3, we show show that upweighting low-frequency classes helps reduce the performance gap between SGD and Adam problems with heavy-tailed class imbalance, providing evidence that the optimization difficulties are associated with class imbalance.
While reweighting the loss of samples from class k by 1/Tu to address the class imbalance seems intuitive, optimizing the reweighted loss no longer guaranteed to lead progress on the original loss, especially if the weights are large. Indeed, we find that on some problems this reweighting does not improve performance (although SGD and Adam perform similarly on the reweighted loss, not shown). However, the less extreme reweighting of 1//Te appears to consistently outperform SGD.
In Figure 25, run SGD on the reweighted loss with the two weighting schemes, 1/TH and 1/VTTI and plot its performance on the original, unweighted loss. We compare the performance of the two reweighting schemes with SGD and Adam, all with momentum, on the following problems.
The small transformer on PTB in Figure 12 (stochastic training)
The Linear model on synthetic data in Figure 4 (deterministic training)
The CNN on
dataset in Figure 2 (deterministic training)
The & 8 the Heavy-Taiie Tailed ImageNet dataset in Figure 3 (stochastic training)
We found that the combination of both Adam and reweighting did not improve over running Adam on the original loss and do no include it in Figure 25.
PTB
Linear
HT MNIST
HT ImageNet
10
SGD
rSGD 12 rSGD
Train Loss
Adam
5
5-
5-
0
50
100
500
1000
200
500
1000
1500
Epochs
Epochs
Epochs
Steps
Figure 25: Reweighting the loss improves the performance of (S)GD on low-frequency classes.
The plots show the unweighted loss, while (S)GD and Adam optimize reweighted loss. Reweighted (S)GD (r(S)DD) with weights 1/ consistently outperforms plain SGD, although it can lead to spikes, on the CNN on the MNIST dataset. Reweighting with weights 1/T. is sometimes better (Linear, MNIST) but can be worse (PTB, ImageNet) as it optimizes different objective. We use deterministic updates for the first problems, labeled Epoch, and stochastic updates for the R8sNer18 on heavy-tailed ImageNet.
26

---

F Dynamics of the gradient and Hessian
This section provides additional details on the dynamics of (S)GD and Adam discussed in Section 3.2.
Figure 26 shows the dynamics of GD and Adam on the linear model synthetic data in Figure (deterministic training). This figure complements Figure which shows the dynamics over the path taken by Adam.
Figure 27 and additionally shows the average predicted probabilities p for each frequency group, showing that the deviation from the linear relationship for rare classes coincides with the predicted probabilities p for those classes going 1.
The following figures show the correlation on additional problems, on
Figure The GPT2-Small model on Wiki'ext- 103 Figure (stochastic training). This figure complements Figure 8 which shows the dynamics over the path taken by Adam.
Figure 29 The CNN on the MNIST+Barcodde dataset in Figure 2 (deterministic training)
Figure 31 The small transformer on PTB Figure (stochastic training)
Figure The ResNet18 on the Heavy-TTailed ImageNet dataset Figure (stochastic training)
Figure 32 illustrates that this correlation does not hold globally and only emerges throughout training by showing that negative correlation instead be found by looking at the oppositve path of the path taken by Adam, W, (when W, are the iterates generated by Adam).
F1 Linear model on synthetic data
(a) Dynamics the path of GD
Correlation at
10
t=
t= 100
t=0 0
10-1_
Hessian
Tr(V i
10
10
10-1
Grad, [|V fll
10'
1
Grad.
10'
-
Grad.
10'
10
Grad,
10'
|V fll
||V fl
IV.,
(b) Dynamies the path of Adam
Correlation at
t=10
t=
t= 100
t= 0
10-1
Hessian
Triv O
10-3
10.5
10-3
Grad,
10
10
Grad. [V, fll
10'
10
Grad.
10'
10
Grad,
10'
[IV
I|V fll
Figure 26: Evolution of the gradient norm and Hessian trace through optimization. Taken the path GD (a) and Adam (b) on the linear problem of Figure 4. The blocks correspond to the rows Wi,..., w. of the parameter matrix W. The color indicates the class frequency, showing that lower (higher) frequency classes have smaller (larger) gradient norm and Hessian trace. Figure 26b is replication of Figure given here for convenience. The deviation from the correlation is explainable by the fact that difference classes are learned at difference speed, leading to a different value of in Proposition 2, shown Figure 27. For GD, frequent classes are learned faster than infrequent ones, while for Adam, is similar among the most frequent groups of classes while - 1 for the least frequent classes.
27

---

Adam
GD
100
10
Probability (split by group)
10
10-4
10
10.
10.
109
10
10a 109
10' Epoch
102
Epoch
Figure 27: Evolution the predicted probabilities for the correct class. Complement to Figure 26, taken over the path of GD and Adam on the linear problem of Figure 4. For GD, frequent classes are learned faster than infrequent ones, Adam has similar behavior on the most frequent groups of classes but also increases the predicted probability for the correct class on infrequent groups. The color indicates the class frequency, showing that lower (higher) frequency classes have smaller (larger) gradient norm and Hessian trace.
F.2
GPT2-Small WikiText-103
Dynamics the path of SGD
Correlation at
100
1000
10000
t=0
10-
Hessian
Tr(V
10
10-9
10
Grad.
10 fll
10 Grad. 10
10 Grad. 10
10
Grad. 10
10 Grad, 10
||V
IV fll
[[V Al
||
IIV
(b) Dynamics the path of Adam
Correlation at
t= 10
t= 100
1000
t= 10000
t=0
-
Hessian
Tr(V2 f
14-4
10-8
1
Grad. 10 [|V., fll
10-5 Grad. 10 [V fll
10-1 Grad. 10 ||V fll
10-5 Grad. 10
Grad, 10 10
||V fll
||V fll
Figure 28: The grasient-Hessian blocks also become correlated in the last layer of large models. Reproducing Figure 7 on the GPT2- 1arr2-mmlllWiiiiextttt 103 problem of Figure 1. Evolution of the gradient norm and Hessian trace for each row w. of the last layer throughout optimization, over the path taken by SGD (a) and Adam (b). The color indicates the class frequency. showing that lower (higher) frequency classes have smaller (larger) gradient norm and Hessian trace.
28

---

F33 CNN on BNNIST
(a) Dynamics over the path of GD
Correlation at
t= 10 t= 50
t= 100
10-1-
Hessian
TrVV f)
10-3
10
10-
Grad. ,| fll
10'
10-1
Grad. V fll
10'
10-a
Grad.
10-1
Grad. v.
101
|| fll
(b) Dynamics over the path of Adam
Correlation at t=
t=10
t= 50
t= 100
10-1-
Hessian
Tr(V (
10-3
10-5_
10-3
Grad. v]
10-3
Grad. [[V fll
10'
10-3
Grad.
10-3
Grad.
10'
|| Al
Figure 29: Evolution of the gradient norm and Hessian trace through optimization. Taken over the path of GD and Adam on the CNN on imbalanced MNIST in Figure 2. Note that this problem only has two groups of classes with different frequencies; 10 classes have m55k samples while 10k classes have 5 samples.
F.4 ResNet18 on Heavy-Tailed ImageNet
(a) Dynamics over the path of SGD
Correlation at
t= 10
50
t= 100
t=0
10-2
Hessian
Tr(V2 f)
10
10-5
10-3
Grad. ||V
10
1
Grad.
10'
10
Grad.
10'
10
Grad,
10'
All
|V All
|| All
IV Al
(b) Dynamics over the path of Adam
Correlation at
t=
t= t=
t= 100
t=0
10
Hessian
Tr(va f)
10
10-4
10
Grad.
10
1
Grad. I fll
10'
10
Grad.
10
10
Grad. VV fl
10'
||V fl
||V fl
Figure 30: Evolution of the gradient norm and Hessian trace through optimization. Taken over the path of SGD and Adam on the R on Heavy-Tailed ImageNet in Figure 3.
29

---

E5 Small Transformer on PTB
(a) Dynamics over the path of SGD
Correlation at
t=
t= t=
t= 100
t= 0
10
Hessian
Tr(V D
10-3
10
10
Grad. I|V fll
10'
10
Grad. IV fll
10'
10->
Grad.
10'
10-2
Grad.
10'
[| fll
(b) Dynamics over the path of Adam
Correlation at t=
t= 10
5=
t= 100
100
Hessian
Tr f)
10-2
10-4
10-3
Grad.
10-1
Grad. |V
10'
10->
Grad.
10'
10-1
Grad,
10'
||V n. Al
All
||V
vV Al
Figure 31: Evolution of the gradient norm and Hessian trace through optimization. Taken over the path of SGD and Adam on the small Transformer on PTB in Figure 12.
The correlation depends on the path
Proposition 2 requires that the optimizer make progress and assign samples their correct classes. Indeed, the positive correlation observed the previous figures is not a global property of the loss function. Not only does it not hold initialization, where the Hessian is uniform, the correlation can even be reversed in some areas of the parameter space, as shown in Figure 32.
Opposite correlation at W=
t= 10
t=
t= 100
Wa
2
t=
Hessian
Tr(V
1
10.6.
10-3
Grad.
10'
10-1
Grad. ||V w.
10'
10-3
Grad. IV.,fi fll
10'
10-3
Grad. I|
10'
[|V.. fll
Figure 32: The correlation only holds while training. Correlation between the gradient and Hessian blocks through the path [-W.). where W, are the iterates of Adam on the linear model of Figure 4.
30

---

Correlation between the gradient and Hessian across blocks This section gives the proof of Proposition 2 in Section 3.2
G
Proposition 2. If initialized at Wo 0, the gradient and Hessian of the loss L w.E.f. W. are
During training, if the model correctly assigns samples to class with probability p (Assumption 1), ((W)) (1 -
Vwe ((W)) x L
(1)
Vw (1- p x + O()),
|
and Vl C1
Tr(V as - oc, (2)
W/ f 10==p(1-ppm H + O();
p TT.HH)
for classes where the frequency does not vanish too quickly, Tx w(/c).
The requirement that the class frequencies do not vanish, mh = w(+/o), necessary to make possible discuss class frequencies c co, unless the class frequencies do not depend on While the frequencies Tæ and the number classes c can be independent, for example if TT follows an exponential decay, Th x 2- it does not hold for all distributions. While it may seem that this result only holds for relatively frequent classes, as it requires T&C + co, we can see that nearly all the data comes from classes where this correlation holds when the classes тę x 1/k. Denote by H (c) Ex=1 1/k ((log c). After normalization, we have T& = 1/kH(c). The correlation result holds as long as T&C - oo, and so it at least holds for the first k v c/log(c)2 classes as T&C 2 log(c)
distributed as
oo. While this only cover 1/ log(c)² fraction of the classes, those classes account for nearly all the data as
H(Te/ ) = H(c)
( log(c) 2 log log(c) log(c)
Tk
+ 1.
1=1
Proof of Proposition 2. We first recall the gradient and Hessian for each block Wly Woi
Vw, ((W, y) [1 = k] - xxx)))x,,
V (WW,,y) = - (xx(xx))
and the definitions of the moments of the data, per class and overall.
x = = Σ X H² ET x
H = 4/ Ei x.x,
Our first step is to rewrite the sums for gradient and Hessian to separate the influence of samples of the correct class A: and the other samples.
1 Σ [1)
Vw L(W)
k] p(xx)))x,
[1(y = k] - x(x))+)xx,
(Split by class)
(1/y: k] - p(x:)a)xi,
(Use class frequencies т; -nj/n)
==1
TTk nk
(1-pt pxx)xx))
Σ
J Σ (-p(xx)))xx.
blitk ny (:
i= 1:ys
1 Σ (1-)))(1
n
22 L(W)
n
Σ (1x))(11p(x)aal
x(x))
Σ T (1-x))((1-((x))))
ng
ny #::
j= boliik
We simplify the first terms using the assumption that p(x:)x= ) = p for samples of the correct class,
mT
Σ
(1-
xx*
T Σ p(x)) - p(x))x)xxx - = =p(1 - - H*
Tik
nkk d.
31

---

We introduce the following shorthands for the second terms.
d2 c Σ Σ
(-p(x:)a)x,
DA >
- (p()))))
jz fl sty:
j=1, A
Using those simplifications, we obtain that
Vw. L(W)= = (1 p) dc.
V ((W) p(1 - p)T H D
The terms dg. D₁ are averages of terms weighted by cp(x;)t, which by assumption is 0(1), and such both |d and Tr(D;) are O(1). The ratio between the two will be dominated by the contribution of their first term as long as Ty dominates 1/c. in the sense that lim - 0, as
Vw, CII lim
[|(1 x) d
lim
01000 Tr V2 C)
6-+00
Tr(p(1 - TT H + DR)
(1 p)** + cTp
de
x
lim
Tr(p(1 p) H +
D2)
p (T()
CT
G.1 Off-diagonal blocks orders of magnitude smaller than diagonal blocks
Our discussion Section 3.2 ignored the impact of off-diagonal blocks. In this section, we show that they are small. The diagonal and off off-diagonal blocks of the matrix for k + k'.
HEA
V &(W. X.W. y)
= xxx) =(x)&{
and for j k.
Hiy
Vw. V. x,,
p xx )xx
From this, that of the diagonal blocks, as
that, average, the magnitude of the off-diagonal blocks will be smaller than
HRG
Haj
because E 1 f p(x)&p(x)x> = (1 - p())x), This means that the matrix T [e x c] formed by taking the trace the blocks, Tyx = ,x,).. is diagonally dominant.
Figures 9 and 33 show that the magnitude of the entries of the Hessian in off off--digonall blocks is orders of magnitude smaller than those of the diagonal blocks. Instead of plotting the [ed ed) Hessian, we subsample 40 classes and 40 input dimensions and plot the resulting [160 160] entries at different points throughout trajectory of Adam on the problem of Figure Figure 9 shows the matrices with classes sampled uniformly and Figure 33 with classes sampled log-uniformly
Hessian at t=0
t= 10
t= 50
t= 100
00
4
6
8
Figure 33: The off-diagonal blocks are much smaller than the diagonal blocks. Showing the magnitude (ogio(((V'L)), for [160 160] subset of the Hessian, sampling 40 classes and 40 input dimensions uniformly.
32

---

H Continuous time GD and sign descent on a simple imbalanced problem
We give the proof of Theorem 3 on the simple imbalanced setting, restated here for convenience.
Simple imbalanced setting. Consider classes with frequencies T1.-.. Te where all samples from class are the same, X; C₁ y = k, where ej is the kth standard basis vector in IR
Theorem 3. On the simple imbalanced setting, gradient flow and continuous time sign descent initialized at W minimize the loss of class k, (x) - log(a
(oggla(W) (W ([)ek)))) at the rate
Gradient flow: (t(t) = ((1/Tat),
Continuous time sign descent: (x) = (
We separate the proof for gradient flow into parts. Lemma 4 simplifies the dynamics into smaller, independent differential equations, Lemma 5 solves the differential equation and Lemma bounds the loss. The proof uses similar tools as for the gradient flow dynamics studied by Cabannes et al. (2024), but we focus instead the loss per class. We treat continuous time sign descent separately in Lemma 7.
Notation. If W is a (a b] matrix, then Wi...., Wa are the rows and w .... w are the vectors, and Wij is the entry at the ith column, jthh row. For brevity, we use
1 as the term appears often.
Lemma 4 (Separation of the dynamics). The dynamics of the parameter matrix W separate into e 2- 2-dimensional differential equations, wis (t) = (+()) and wik (t) = for j & k, where
d
exp(ax)
ax (0) 0,
dl
0k T&
exp(ax) + (c 1) exp(bx)
exp(bg)
by (0) = 0,
bk TT
df
exp(ax) + (c- exp(bx)
Proof. Our goal is to simplify the dynamics starting ( 0 and following the gradient flow d W = CWW),, where W lc d). For the simplified setting, we have that d = are the inputs de
are the standard basis vectors in R. The derivative of w.r.t. a single element Wky is
Owny L(W)
1 = -} + (w""))))
As Own, only depends w for all k, The dynamics are independent across the columns of W, giving e independent equations in R
d w = - o
w (0) 0,
dt
To further simplify the dynamics, we use the fact that the weights that are not associated with the correct class have the same dynamics. For any indices i, different from k. wa(() (tx(t). They have the same derivatives if they have the same value, as
exp(win) TLo(w")) = dt Wyk: dt Wit kLo(W*)) = Th Σn exp(wwk) Th Exx exp(ww<)
d
exp(wax
so they will have the same dynamics and the equation can be reduced system of variables. Wka = and Wi = by for any + k, with
d 0k Th
exp(a<)
ag (0) = 0,
dt
exp(ax) + (c- 1) exp(bx)
exp(bx)
by (0) = 0,
bk T&
di
exp(ax) + (c exp(bx)
Lemma 5 (Solution of the dynamics). For a given class with frequency the dynamics of the parameters and in Lemma 4 evolve as follows, using the shortcuts f(t) 1 + eut and c
1 [(t) =w ( exp (x (7)(1)
b(t) - a(t).
(t)) =
33

---

Proof. We want the solution to the differential equation
exp(a)
a(0)
exp(a) + (c- 1) exp(b)
d
cxp(b)
((0)
dt
exp(a) + (c- 1) exp(b)
The general solution, ignoring the initial conditions, uses the Lambert W function and constants K1, K2. For brevity, introduce the shorteut c
a(t) =
1
ce Ki K2 + cant - - exp c (2Tt + t K2) K3
K2)
2c
x2
b(1) = K1
ce-K Ki
ce
K2 + cant - 2 expl
c
(2Tl + 6 K: K2) K3
z²c
22
We need set K1, K₂ satisfy the initial conditions a(O) b(0) 0. As b(l) K₁ a(t)/z,
must have that K = O, giving the simplification
1 ze
cK2 + cant - 22W ( (exp( (2 K2) K1) ki))); b(1) -a())
a(t)
To set K2, we need to have
0 = =ca((0) = cK₂ - 22 ( exp K2 (xx)(ks
W ( exp K2 (Ks
= 2
K2
Since W(xe") = x for >> O, the equation is satisfied for K2 = g. as we get W (fe')
1 giving
a(t) = ((1 1+ cet eut - 2W ( (1 + ext) )))
1
exp
b(t) 2 a(t).
Lemma (Bound for the loss). For sufficiently large such that + cry. 2 = log 2
(t = θ T&t
Using the simplification derived in Lemma 4 and the solution of the differential equation in Lemma 5, we can rewrite the loss for a specific class as a function of time as
exp(wix)
(W) log(a (Wee))) log
Σj= exp(wx)))
exp(ax((t)
(x :0))== Lx = log exp((x(()) (c (x)x))xp(ba(())
log(1 (c 1) exp(cbb (f))),
where the equality uses that (t) = (c - 1)bx For brevity, we will drop the index k in ag. bys lx and Ty. and use the shorteut
1, bounding the quantity
f(t) log(1 exp(cb(t))).
Expanding the definition of b(t) using Lemma 5, we have
((() (10) =w (= exp (()
z =xx((cb(()) zexp -
where
S(1) end.
To simplify the W function, we use the fact that for x > e (Hoorfar and Hassani, 2008, Theorem 2.7)
log(z) E
w(x) log(æ)) log(log(2)) + 8(x)
where
2 S 8(x)
log(log(z))
6
WWolframAlpha solution for 1: httt:///wwwooaaaaaaooc////////tt
34

---

To use this bound on W (t))) we need 1 f(t) 2 e, which is satisfied for t sufficiently large, once f(t) 2 z(log: 2 + 1).
Using that log(
= (t) log(z), and writing h(t) = 6( exp( ( f()))) we have ((1) =W { exp (0))
((t) - ( -lo[==))))))))=>)>))) log (1)))) - -log()) -((()) ((log(f(t) =log(2)) -t(t),
giving the simplification
1 ( f(t) == zW(= ( exp (1
exp(cc(tt) = 2 exp -
((t))
= =log(2)) + h(1)) = 2 eexp(h (h(t)) f(t) 2 log 2
=
This gives the average loss
E(I) = log(1 + eex((c(())) log (1+ z eex((((())
f(t) - <log
To bound this expression, we can use that eep((h()) assumed to apply the bound on the W function, and use the bounds I+ S log(1 + x) s x to get
- 2 0 after f(t) 2 z log 2, which we have already
ex)(h((t)
2 eex
ft) log 2 + 2 exp(((t)) S e(t) -
f(t) z log
As h(t) is upper bounded by a constant and lim h(()) = 0. lim exp(t(t)) = 1, we have
{(t) o 0 f(t) log
θ πt
Lemma 7.
The loss at time t for continuous time sign descent la (1) = log(1 + (c 1) exexp(-c))
Proof. The same decomposition as in Lemma 4 hold, with the dynamics
ax (0) = 0. dt ax = 1, (x(t) = t, (0) = 0, dt
by
1, bx(t) -t,
leading to the following loss
(t) = log(1 + (c 1) exp(-ct)) = O(2 exp(-ctt
35

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: The main claims of the abstract and introduction refer the following sections.
Section 2 and Figures 1 to 5 support the claim that heavy-tailed class imbalance leads performance gap between SGD and Adam, and that this gap can be made to appear by taking typically uniform datasets and making the class distribution heavy-tailed.
Section 3 and Figures 6 to 9 supports the analysis on a softmax linear model under heavy-tailed class imbalance, where Adam outperforms SGD. We provide toy ample where a correlation between the magnitude of the gradient and Hessian across coordinates can be argued to benefit Adam (Section 3.1), show that the class imbalance on a linear model leads to such a correlation and explain how through an assignment mechanism in (Section 3.2), and prove that, simplified problem and in continuous time, gradient descent performs poorly on low-frequency classes while sign descent is unaffected (Section 3.3).
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
is fine to include aspirational goals as motivation as long as is clear that these goals not attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
We discuss the main limitations of our results, regarding the interaction between class imbalance and stochasticity, the validity of our findings on generalization error, and additional optimization difficulties not captured by class imbalance in Section Throughout the paper, point out subtleties that further discussed in the appendix, such that is possible to create pathological imbalance datasets that are still easy to optimize with GD Section 2.2, that the correlation between class frequencies, gradients and Hessian due the assignment mechanism is not a global property and requires the model to fit the data in Section 3.2, and point out what properties of the optimization dynamics are not capture by the simple imbalanced setting in Section 3.3.
Justification:
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well sellsspecificatonn, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
36

---

The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-text system might not be used reliably provide closed captions for online lectures because it fails handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
Question:
For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
Justification: The proof Proposition 2 relies on Assumption and is given in Appendix G. The proof of Theorem 3 the simple imbalanced setting is given in Appendix H.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should properly referenced.
4. Experimental Result Reproducibility
Question: Does the paper fully disclose all the information needed reproduce the main ex- perimental results of the paper to the extent that affects the main claims and/or conclusions of the paper (regardless whether the code and data are provided or not)?
Answer: [Yes]
The experimental details needed to reproduce the main experimental results are given in Appendix A. We use standard architectures and datasets where possible and simple training procedure for reproducibility. and the main claims of section Section 2.2 and Section can easily be reproduced on linear models with synthetic data.
Justification:
Guidelines:
The answer NA means that the paper does not include experiments.
37

---

If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution a specific model and empirical evaluation, may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of large language model), releasing of a model checkpoint, other means that appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how reproduce that algorithm.
(b) the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is a new model (e.g., a large language model), then there should either be way access this model for reproducing the results way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors welcome describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification: The code to reproduce our experiments is uploaded to the openreview submis- sion and will be made publicly available.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips.ca
for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https: / wtips..c/ppblic/guiiii wttp://ddddeeeeeemiimmceeeeeee ssi /ttes//Cooesmmmiiiooooo........
for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
38

---

The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing much information possible in supplemental material (appended the paper) recommended, but including URLs to data and code permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification:
The experimental setting presented at high-level in Section 2 and detailled in Appendix A. The accompagnying code provides a full specification the experiments.
Guidelines:
The answer NA means that the paper does include experiments.
The experimental setting should be presented in the core of the paper level of detail that is necessary to appreciate the results and make sense of them.
The full details can provided either with the code, appendix, or as supplemental material.
7. Experiment Statistical Significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [NA]
The main figures do not report error bars, as the figures show detailled trajecto- ries for specific runs that do not lend themself to show the behavior averaged over mutliple runs. Instead, we account for factors of variability by reproducing the observed behavior, the performance gap across class frequencies, across multiple datasets, architectures, and training procedures.
Justification:
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
should be clear whether the error bar is the standard deviation or the standard error of the mean.
It is OK to report 1-sigma enoo bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
39

---

For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are of range (e.g. negative enror rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments Compute Resources
Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Answer: [Yes]
Justification: Appendix A.4 lists the type of compute and the estimated overall total compute budget used, including preliminary experiments. The details of per-experiment compute type and budget is available in the code, where each experiment file specifies the hardware configuration and runtime.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make into the paper).
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurlPS Code of Ethics http://rpppcccccpplaa/cccccooooa
Answer: [Yes]
Justification: The authors confirm that the research was conducted conforming to the Code of Ethics.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g.. there special consid- eration due to laws or regulations in their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification:
The paper focuses foundational research understand the behavior of generic algorithms used to optimize neural networks. The paper not tied to a particular applications and do not direct path to social impact impact.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
40

---

Examples of negative societal impacts include potential malicious unintended uses (e.g., disinformation, generating fake profiles, surveillance), faimesss considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: This paper does not release data or models.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
recognize that providing effective safeguards challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Question:
Are the creators or original owners of assets (e.g., code, data, models), used in the paper. properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Appendix A gives references for the datasets, models, and code used this project, along with citations to the original papers and URLs where available.
Justification:
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset used and, possible, include URL.
41

---

The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (c.g., website). the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswith waperswithcode.comm come datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information is available online, the authors are encouraged to reach out to the asset's s creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification:
The paper does not release new assets. Existing publicly available datasets used to create variants with more classes, using the procedures are detailled in Appendix A.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/mmodel as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset is used.
At submission time, remember anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification: the paper does not involve crowdsourcing nor research with human subjects
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included the main paper.
According to the NeurlPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
42

---

Answer:
[NA]
Justification: the paper does not involve research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
43


Reviews:
[{"review_id": "1VAsqab88F", "reviewer": "Reviewer_ujWn", "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.", "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?", "limitations": "There are no potential negative impacts of this work. The authors discuss the limitations of their work at length in Section 4.", "rating": 7, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 3, "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.\n\n- 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.\n\n- 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.\n\n- 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).\n\n- It would be nice to elaborate on lines 331-333 and include more discussion on this aspect in Section 3.\n\n- Some corrections/suggestions:\n\n    - Missing ‘the’ in line 152. Extra ‘the’ in line 225. Missing ‘than’ in line 607. Extra ‘be’ in line 631.\n\n    - It would be good to add some description in lines 496-497.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "**Clarification on barcoded MNIST**\n> 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?\n\nGood 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.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "0K6kxz40gA", "reviewer": "Reviewer_HWpb", "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.", "questions": "When the parameters are not separable, does the conclusion still hold?", "limitations": "N/A", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The authors explain that class imbalance is a reason that Adam can outperform SGD.\n\n2. For linear models, the authors establish the relationship between hessian and gradient showing the \"correctness\" of Adam who approximates Lipschitz with gradients.\n\n3. 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.\n\n2. 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?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "**Does the conclusion still hold when the parameters are not separable**\n> When the parameters are not separable, does the conclusion still hold?\n> 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?\n\nWe 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. \n\nIn [**the supplementary pdf for the response**](https://openreview.net/attachment?id=xsoeNVdrpL&name=pdf), 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)\n\nThose 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$. \n\nIf 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.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Q2eGbRrZFy", "reviewer": "Reviewer_LhRX", "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 \n1. reproduce the performance gap by training different language models on some standard datasets, like WikiText-103 and PTB;\n2. 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;\n3. introduce a simple linear model trained on synthetic uniform data, whose class frequencies follow a power law, and again reproduce the performance gap;\n4. 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\n5. demonstrate that the performance gap can be reduced by upweighting the loss of low-frequency classes.\n\nHaving 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.", "questions": "1. Do you know why SGD eventually finds the minimum in the imbalanced MNIST case but does not in the imbalanced ImageNet?", "limitations": "Yes, they authors addressed the limitations of their study.", "rating": 7, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 3, "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.\n- The problem and hypothesis are very clearly stated.\n- 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. \n- 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.\n- The paper also provides some theoreticall understanding as to why the performance gaps actually takes place.\n\nIn 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.\n\nAlthough the paper is very well written and structured, the following suggestions might improve readability:\n- 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. \n- 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.\n- Perhaps section 2.2 could be moved inside section 3? \n\nTypos: lines 152, 160 and 225.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "**Clarification**  \n> Do you know why SGD eventually finds the minimum in the imbalanced MNIST case but does not in the imbalanced ImageNet?\n\nWe 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.\n\n> why SGD [does not] find the minimum in the imbalanced MNIST case but [eventually does] in the imbalanced ImageNet\n\n**Does GD get stuck on Imbalanced MNIST (Fig. 2)**  \nGD 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**](https://openreview.net/attachment?id=xsoeNVdrpL&name=pdf).\n\nThe 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).\n\nWe 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.\n\n**Inconsistency in references to SGD or GD between caption/figures**  \nThose are indeed typos, apologies for the confusion. We will double-check those references.\n\n**Remaining comments**  \nWe will fix the typos and add references to the appendix to highlight the results on other language tasks.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "TIhvS0BgAV", "reviewer": "Reviewer_CdGj", "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).\n\nThe 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.", "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?\n- 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.", "limitations": "The paper discusses its limitations and future directions in detail. For additional potential limitations, please refer to the 'Weakness' section.", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 4, "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.\n- The experiments are well-designed and robust, including thorough ablation studies that strongly support the hypothesis.\n- 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "**Weight blocks for the gradient-Hessian correlation plots**\n> 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?\n\nThis 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.\n\n**Effect beyond last layer**\n> If so, how does class imbalance affect other layers beyond the last layer?\n\nThis 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.\n\nHowever, 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.\n\n**Correlation appears on realistic models**\n> 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.\n\nWhile 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. \n\nWe include an additional plot in [**the supplementary pdf for the response**](https://openreview.net/attachment?id=xsoeNVdrpL&name=pdf) 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.\n\n**Related work**\n> 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?\n\nWe will include a discussion of those works. Follow-up works to [Zhang et al. (2024a)](https://arxiv.org/abs/2402.16788) from [Zhang et al. (2024b)](https://arxiv.org/abs/2406.16793) and [Zhao et al (2024)](https://arxiv.org/abs/2407.07972), (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.\n\nIn addition to their work showing that the Hessian of transformers has a block-diagonal structure ([Zhang et al., 2024a](https://arxiv.org/abs/2402.16788)), [Zhang et al. (2024b)](https://arxiv.org/abs/2406.16793) 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)](https://arxiv.org/abs/2407.07972) 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. \n\nBoth the experiments of [Zhang et al. (2024b)](https://arxiv.org/abs/2406.16793) and [Zhao et al (2024)](https://arxiv.org/abs/2407.07972)  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.\n\n---\n\n**References**\n- [**Zhang et al. (2024a)**](https://arxiv.org/abs/2402.16788)  \n  Y. Zhang, C. Chen, T. Ding, Z. Li, R. Sun, Z. Luo  \n  Why Transformers Need Adam: A Hessian Perspective  \n  https://arxiv.org/abs/2402.16788\n- [**Zhang et al (2024b)**](https://arxiv.org/abs/2406.16793)  \n  Y. Zhang, C. Chen, Z. Li, T. Ding, C. Wu, Y. Ye, Z. Luo, R. Sun  \n  Adam-mini: Use Fewer Learning Rates To Gain More  \n  https://arxiv.org/abs/2406.16793\n- [**Zhao et al (2024)**](https://arxiv.org/abs/2407.07972)  \n  R. Zhao, D. Morwani, D. Brandfonbrener, N. Vyas, S. Kakade  \n  Deconstructing What Makes a Good Optimizer for Language Models  \n  https://arxiv.org/abs/2407.07972", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
We thank the reviewers for their careful readings and thoughtful comments. 

We answer specific questions in individual replies:
- [**ujWn:** Clarification on barcoded MNIST](https://openreview.net/forum?id=T56j6aV8Oc&noteId=2QGkcNTqaR)
- [**HWpb:** Does the conclusion still hold when the parameters are not separable](https://openreview.net/forum?id=T56j6aV8Oc&noteId=ASSfu6V1bw)
- [**LhRX:** Does GD get stuck on Imbalanced MNIST (Fig. 2)](https://openreview.net/forum?id=T56j6aV8Oc&noteId=cKpTYeF9u4)
- [**CdGj:** Impact of class imbalance beyond the last layer](https://openreview.net/forum?id=T56j6aV8Oc&noteId=vrl5iFlM2d)

We will address the minor comments such as typos in a revision.

Meta-review: This paper provides new insights into why Adam consistently outperforms SGD by a significant margin in language models. While this phenomenon has been observed by both practitioners and researchers for some time, previous explanations have largely focused on the architectures of models like Transformers. In contrast, this paper introduces a new hypothesis, proposing that the heavy-tailed class imbalance in training datasets is the key factor driving this performance gap. The authors conducted a series of rigorous experiments that solidly support their hypothesis, demonstrating that this gap is reproducible across various tasks and architectures. Additionally, the paper offers theoretical analysis in simplified settings, which deepens our fundamental understanding of this phenomenon.

All reviewers unanimously agreed on the acceptance and gave high ratings. I recommend accepting it as a spotlight paper, given its solid foundation and the potential for significant impact on the field.
Decision: Accept (spotlight)