The AdEMAMix Optimizer: Better, Faster, Older

## Abstract

Abstract Momentum based optimizers are central to a wide range of machine learning applications. These typically rely on an Exponential Moving Average (EMA) of gradients, which decays exponentially the present contribution of older gradients. This accounts for gradients being local linear approximations which lose their relevance as the iterate moves along the loss landscape. This work questions the use of a single EMA to accumulate past gradients and empirically demonstrates how this choice can be sub-optimal: a single EMA cannot simultaneously give a high weight to the immediate past, and a non-negligible weight to older gradients. Building on this observation, we propose AdEMAMix, a simple modification of the Adam optimizer with a mixture of two EMAs to better take advantage of past gradients. Our experiments on language modeling and image classification show—quite surprisingly—that gradients can stay relevant for tens of thousands of steps. They help to converge faster, and often to lower minima: e.g., a 1.3 1.3 1.3 1.3 B parameter AdEMAMix LLM trained on 101 101 101 101 B tokens performs comparably to an AdamW model trained on 197 197 197 197 B tokens ( + 95 % percent 95 +95\% + 95 % ). Moreover, our method significantly slows-down model forgetting during training. Our work motivates further exploration of different types of functions to leverage past gradients, beyond EMAs.

## 1 Introduction

With large neural networks, deep-learning has revolutionized numerous fields, such as computer vision and natural language processing. At the heart of this paradigm lies the challenge of optimizing complex, non-convex loss functions using noisy gradient estimates. This optimization process is typically carried out using variants of Stochastic Gradient Descent (SGD) (Robbins & Monro, ) or adaptive methods such as Adam and AdamW (Kingma & Ba, ; Loshchilov & Hutter, ), which have become ubiquitous in training state-of-the-art models (Devlin et al., ; Brown et al., ; Dosovitskiy et al., ; Radford et al., ; Zhai et al., ; Dehghani et al., ; Touvron et al., ; Dubey et al., ).

A key component in many of these iterative optimization algorithms is momentum, which has long been shown to accelerate convergence (Nemirovskii & Nesterov, ) and often leads to solutions with superior generalization properties (Sutskever et al., ). By accumulating gradient vectors over successive optimization steps, momentum helps overcome small local variations of the loss landscape, potentially escaping shallow local minima, and accelerate in plateau regions (Qian, ; Ruder, ; Goh, ). Both SGD with momentum (SGD+M) and Adam incorporate momentum under the form of Exponential Moving Averages (EMAs) of past gradients ${\mathcal{G}}^{T}=\{{\bm{g}}^{(0)},\ldots,{\bm{g}}^{(T)}\}$:

$$ $\displaystyle\text{EMA}(\beta,{\mathcal{G}}^{T})\triangleq\beta\cdot\text{EMA} (\beta,{\mathcal{G}}^{(T-1)})+(1-\beta){\bm{g}}^{(T)}=\sum_{i=0}^{T}\beta^{i}( 1-\beta){\bm{g}}^{(T-i)}.$ (EMA) $$

Two considerations support the use of EMAs. From a practical standpoint, the recursive formula of EMA allows for efficient implementations, which do not require maintaining a buffer of past gradients. From a theoretical standpoint, gradient descent with momentum leads to optimal convergence rates for quadratics (Polyak, ; Nesterov, ). However, those results do not guarantee any optimality for general non-quadratic cases (Goujaud et al., ).

The use of momentum in optimization is grounded in the varying nature of gradients. As local linear approximations of the loss landscape, their information can quickly become outdated as the optimization process progresses (Pascanu et al., ).
For this reason, practitioners typically employ moderate momentum values (i.e. $\beta\approx 0.8$ or $0.9$), effectively creating a moving average of recent gradients while discarding older information. Selecting larger $\beta$ values seems counter-intuitive, as it would suggest that older gradients maintain their relevance over extended periods of training. While it is tempting to see the use of small $\beta$s as a confirmation of the limited temporal relevance of gradients, our work reveals instead that older gradients can efficiently be used.
When we increase $\beta$, we decrease the relative importance of recent gradients, and the iterate now fails to respond to local changes in the loss landscape. We observe that a single EMA cannot both give a significant weight to recent gradients, and give a non-negligible weight to older gradients (see Fig. 3(a)). However, a linear combination between a “fast-changing” (e.g. $\beta=0.9$ or $\beta_{1}=0$) and a “slow-changing” (e.g. $\beta=0.9999$) EMA allows the iterate to beneficiate from (i) the great speedup provided by the larger (slow-changing) momentum, while (ii) still being reactive to small changes in the loss landscape (fast-changing). More precisely, we find the following statement to convey an important intuition behind this approach:
{adjustwidth}2em2em
*While changing the direction of the slow momentum is difficult, any adjustment orthogonal to that direction is easy—which favors fast progress in sinuous canyon-like landscapes.*
A toy illustration of this can be seen in Fig. 2. Based on this idea, we propose AdEMAMix (Adaptive EMA Mixture), a novel Adam based optimizer which successfully leverages very old gradients to reach better solutions.

Figure: (a) $110$M parameters.
Refer to caption: x1.png

Contributions. Our contributions can be summarized as follows:

- •
We propose AdEMAMix, a novel optimizer which better leverages past gradients by avoiding a common pitfall of EMA-based optimizers (see § 3).
- •
We empirically demonstrate the superiority of our method over Adam by training ViT and language models (Transformers and Mamba) of up to $1.3$B parameters (see § 4). In addition, we show gains from switching mid-training from Adam to AdEMAMix (see Fig. 5).
- •
We show AdEMAMix forgets the training data slower when compared to Adam (see Fig. 4).
- •
More broadly, our findings contribute to a deeper understanding of the optimal balance between using historical gradients and adapting to the rapidly changing loss landscape. Our work invites further research in methods combining old and recent gradients, beyond EMAs.

Figure: (a) $\|{\bm{x}}^{(t)}-{\bm{x}}^{\star}\|_{2}$.
Refer to caption: x4.png

## 2 Related Work

Works on understanding momentum. From the seminal work of Polyak (), many publications analyzed the effect of gradient descent + momentum (GD+M) in both convex and non-convex settings (Ghadimi et al., ; Flammarion & Bach, ; Kidambi et al., ; Defazio, ; Sebbouh et al., ). While the acceleration in the noise-free setting has been long theorized for convex functions, several publications indicate this effect might not necessarily extend to stochastic settings (Yuan et al., ; Kidambi et al., ; Leclerc & Madry, ), emphasizing instead a link between momentum and effective learning rate. Recent work have been seeking to understand the impact of momentum on generalization through studying the implicit bias of momentum methods (Ghosh et al., ; Papazov et al., ), exposing a preference of SGD+M for lower norm solutions. Those further exposed a link between higher momentum and higher effective learning rate and higher variance reduction. Despite the volume of prior work on the subject, our understanding of momentum methods in non-convex stochastic settings is still incomplete (Yuan et al., ). Oscillatory behaviours, and the sometimes ambiguous effect of variance on optimization render the analysis tedious. From a theoretical standpoint, our work raises several questions. First, given that we gain from averaging very old gradients, what can it reveal of the loss landscape and the consistency of one batch’s gradient during training? Second, would our approach not decrease the variance up to a point that is harming generalization (Ghosh et al., )? While no answer to those questions is given in this work, we provide a toy justification which indicates that large momentums can have a positive impact in noise-free non-convex settings (see Fig. 2)—indicating the improvement of our approach is at least partially explainable without considering variance-reduction effects. We moreover expose a link between momentum and forgetting the training data (see Fig. 4), which to our knowledge is novel.

Works on deep-learning optimizers. Despite the popularity of Adam and AdamW (Kingma & Ba, ; Loshchilov & Hutter, ) in training deep neural networks, optimizer design is a rich field of research and we focus on a few of the works most relevant to this study. Adafactor (Shazeer & Stern, ) improves Adam’s memory efficiency by factorizing the second moment estimate. Lamb (You et al., ) extends Adam by adding layerwise normalized updates. Chen et al. () use algorithm discovery to derive the Lion optimizer. Contrary to Adam, Lion uses a single momentum term and the sign function to produce updates with the same magnitude across dimensions. Interestingly, Chen et al. () also report better scores are obtained when using a slightly larger momentum term ($\beta=0.99$). In this work we show how increasing the momentum well beyond this value can still be beneficial. See App. C.3.2 for a detailed comparison between AdEMAMix and Lion. Recently, Liu et al. () introduced Sophia, a scalable second-order optimizer designed for LLM training. Sophia uses a Hessian-based pre-conditioner which better normalizes the step size, penalizing steps in high curvature direction and accelerating in low curvature directions. Understanding in which circumstances those novel optimizers bring improvements is still being investigated (Kaddour et al., ), and Adam’s dominance remains mostly unchallenged.

Work incorporating an additional momentum term. Lee et al. () introduce Grokfast, which uses a pre-filtering step on the gradient to amplify the low frequencies and help solve groking. When combined with Adam, it effectively applies the Adam’s EMAs on top of another gradient averaging method (e.g. EMA for Grokfast-EMA). Somewhat similarly, Chen et al. () refer to the Double EMA (DEMA) used in some finance applications (Mulloy, ) as one motivation for their AdMeta optimizer. Our motivation behind AdEMAMix is to combine both a high sensitivity to the recent gradients as well as incorporating very distant gradient, in this respect, using nested EMAs is not the right candidate as it reduces the influence of recent gradients. A more detailed review of AdMeta and DEMA is in App. C.3.1. Most relevant to us, Lucas et al. (, AggMo) also observe that using a combination of EMAs can enable the use of larger $\beta$s, and incorporates a sum of $K$ momentum terms into *GD*
They show their approach reaches similar performances as baseline optimizers, with a faster convergence. In contrast, we modify *Adam*, and introduce schedulers that are critical to reaching good performances at larger scales.
This allows us to outperform Adam by a significant margin.
In App. C.3.3, we notice no further improvement by adding more momentum terms. Finally, Szegedy et al. () propose a general framework to derive and study optimizers with linear combinations of memory vectors—which encompasses EMA mixtures.

Works on distributed optimization. Perhaps surprisingly, recent work on distributed optimization—DiLoCo and SlowMo (Douillard et al., ; Wang et al., )—are relevant to our work. $N$ workers ${\bm{\theta}}^{(t)}_{1},\ldots,{\bm{\theta}}^{(t)}_{N}$ are trained independently for $K$ steps (e.g. $K=500$). Every $K$ steps, the delta updates $\{\Delta{\bm{\theta}}_{i}\}_{i=1}^{N}\triangleq\{{\bm{\theta}}_{i}^{(t+K)}-{
\bm{\theta}}_{i}^{(t)}\}_{i=1}^{N}$ are averaged and applied to each worker using an outer optimizer *with momentum $\beta$*: ${\bm{\theta}}_{i}^{(t+K)}={\bm{\theta}}_{i}^{(t)}-\eta\cdot\text{Opt}(\frac{1}
{N}\sum_{i}\Delta{\bm{\theta}}_{i},\beta)$. The application of the outer momentum every $500$ steps increases the importance of old gradients in the optimization trajectory. We believe this observation might in parts explain the strong results provided by those methods, and further motivated our work.

## 3 Our method: AdEMAMix

Figure: (a) Limitation of EMAs.
Refer to caption: x7.png

Setup & notations. Let ${\mathcal{L}}_{\bm{\theta}}:{\mathcal{X}}\mapsto\mathbb{R}$ be a loss function parameterized by ${\bm{\theta}}$, and mapping inputs ${\bm{x}}\in{\mathcal{X}}$ to $\mathbb{R}$. Given a sampled batch ${\bm{x}}$, let ${\bm{g}}=\nabla_{{\bm{\theta}}}{\mathcal{L}}_{\bm{\theta}}({\bm{x}})$ be a stochastic gradient of the loss w.r.t. ${\bm{\theta}}$. To minimize the empirical loss, the Adam optimizer (Kingma & Ba, ) relies on first and second moments, resp. ${\bm{m}}$ and ${\bm{\nu}}$, estimated via two EMAs parametrized by $(\beta_{1},\beta_{2})\in[0,1[^{2}$. A weight-decay parameter $\lambda\in\mathbb{R}^{+}$ is often used as in Loshchilov & Hutter ():

$$ $\displaystyle\begin{cases}{\bm{m}}^{(t)}=\beta_{1}{\bm{m}}^{(t-1)}+(1-\beta_{1 }){\bm{g}}^{(t)},\quad\quad\hat{{\bm{m}}}^{(t)}=\frac{{\bm{m}}^{(t)}}{1-\beta_ {1}^{t}}\\ {\bm{\nu}}^{(t)}=\beta_{2}{\bm{\nu}}^{(t-1)}+(1-\beta_{2}){{\bm{g}}^{(t)}}^{2} ,\quad\quad\hat{{\bm{\nu}}}^{(t)}=\frac{{\bm{\nu}}^{(t)}}{1-\beta_{2}^{t}}\\ {\bm{\theta}}^{(t)}={\bm{\theta}}^{(t-1)}-\eta\big{(}\frac{\hat{{\bm{m}}}^{(t) }}{\sqrt{\hat{{\bm{\nu}}}^{(t)}}+{\epsilon}}+\lambda{\bm{\theta}}^{(t-1)}\big{ )}.\end{cases}$ (AdamW) $$

With $t>0$ being the timestep, $\eta$ being the learning rate, and ${\epsilon}$ a small constant. Initially ${\bm{m}}^{(t=0)}={\bm{\nu}}^{(t=0)}=\mathbf{0}$. To prevent the bias induced by the initial ${\bm{m}}^{(t=0)}$ and ${\bm{\nu}}^{(t=0)}$, the outputs of the two EMAs are corrected into $\hat{{\bm{m}}}^{(t)}$ and $\hat{{\bm{\nu}}}^{(t)}$. Those are used to compute the final weight update, scaled by the learning rate.

How far to look into the past? A typical value for $\beta_{1}$ is $0.9$. Fig. 3(a) shows the weights given to past gradients for different values of $\beta$. The larger the $\beta$, the more uniform the average is. To put this in perspective—observing that $\sum_{i=0}^{\infty}\beta^{i}(1-\beta)=1$ for $\beta\in[0,1[$—the number of successive previous steps receiving a cumulative weight of $0.5$, is $t_{half}=\frac{\ln(0.5)}{\ln(\beta)}-1$. For $\beta=0.9$, $t_{half}\approx 6$, meaning that half of the weight is given to the previous six gradients. This observation can also be extended to SGD with e.g. polyak or nesterov momentums (Polyak, ; Nesterov, ), which typically relies on similar $\beta$ values. The value of $\beta_{1}$ is rarely increased beyond $\sim 0.9$.
In our experiments with AdamW, increasing $\beta_{1}$ further degraded the performance (see App. C.1.6). Does this mean older gradients are outdated? We show that this is not the case, rather, increasing beta is reducing the sensitivity to recent gradients too much. We design AdEMAMix such that the sensitivity to recent gradients is kept, while also incorporating information from much older gradients using an additional momentum term. This allows for the use of much larger $\beta$ values e.g. $0.9999$. To compare, for $\beta=0.9999$, $t_{half}\approx 6{,}930$, spreading half of the mass over the previous $6{,}930$ past gradients.

AdEMAMix. To keep a high sensitivity to recent gradients, while also incorporating information from older gradients, we add a second EMA (changes compared to AdamW are in Blue):

$$ $\displaystyle\begin{cases}{\bm{m}}_{1}^{(t)}=\beta_{1}{\bm{m}}_{1}^{(t-1)}+(1- \beta_{1}){\bm{g}}^{(t)},\quad\quad\hat{{\bm{m}}}_{1}^{(t)}=\frac{{\bm{m}}_{1} ^{(t)}}{1-\beta_{1}^{t}}\\ {\color[rgb]{0,0,1}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,1}{\bm{m}}_{2} ^{(t)}}{\color[rgb]{0,0,1}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,1}= \beta_{3}{\bm{m}}_{2}^{(t-1)}+(1-\beta_{3}){\bm{g}}^{(t)}}\\ {\bm{\nu}}^{(t)}=\beta_{2}{\bm{\nu}}^{(t-1)}+(1-\beta_{2}){{\bm{g}}^{(t)}}^{2} ,\quad\quad\hat{{\bm{\nu}}}^{(t)}=\frac{{\bm{\nu}}^{(t)}}{1-\beta_{2}^{t}}\\ {\bm{\theta}}^{(t)}={\bm{\theta}}^{(t-1)}-\eta\big{(}\frac{\hat{{\bm{m}}}_{1}^ {(t)}{\color[rgb]{0,0,1}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,1}+\alpha {\bm{m}}_{2}^{(t)}}}{\sqrt{\hat{{\bm{\nu}}}^{(t)}}+{\epsilon}}+\lambda{\bm{ \theta}}^{(t-1)}\big{)}.\end{cases}$ (AdEMAMix) $$

In our experiments, while the values of $\beta_{1},\beta_{2}$ remain similar to those of equation AdamW, we often use $\beta_{3}=0.9999$. We find $\alpha\in[4,10]$ to work well in practice.

Tackling early training instabilities. Early training instabilities are commonplace when training deep models, and empirically motivated the introduction of methods such as learning rate warmup and gradient clipping. Gilmer et al. () show how the use of learning rate warmup can be justified from a curvature perspective; allowing the iterates to move to parts of the optimization landscape where larger learning rates are stable. While we use learning rate warmup in all our experiments, we still noticed AdEMAMix models using a large $\beta_{3}$ would diverge early. This, despite not using bias correction over ${\bm{m}}_{2}$, which lets the momentum buffer fill itself slowly. Those failed runs are characterized by updates of large magnitudes in the early phase of training (see App. C.1.9, Fig. 27).
For this reason, we progressively increase the values of $\beta_{3}$ and $\alpha$ using schedulers. For $\alpha$ we use a linear scheduler. A linear scheduler for $\beta_{3}$ would be ill-fitted as the same increment of $\beta_{3}$ has a different impact for different values of $\beta_{3}$. For instance, observe that an increase of $\beta$ of $\delta_{\beta}=0.0001$ barely increases the $t_{half}$ for $\beta=0.9$, while $0.999\rightarrow 0.999+\delta_{\beta}$ increases the $t_{half}$ of $77$. For this reason, we design the $\beta_{3}$ scheduler to *increase $t_{half}$ linearly* (see App. A.1 for a derivation of that scheduler). The two schedulers are summarized below:

$$ $\displaystyle\alpha^{(t)}$ $\displaystyle=f_{\alpha}(t,\alpha,T_{\alpha})=\min(\frac{t\alpha}{T_{\alpha}}, \alpha),$ ( $f_{\alpha}$ ) $\displaystyle\beta_{3}^{(t)}$ $\displaystyle=f_{\beta_{3}}(t,\beta_{3},\beta_{start},T_{\beta_{3}})=\min\big{ (}\exp\big{(}\frac{\ln(\beta_{start})\ln(\beta_{3})}{(1-\frac{t}{T_{\beta_{3}} })\ln(\beta_{3})+\frac{t}{T_{\beta_{3}}}\ln(\beta_{start})}\big{)},\beta_{3} \big{)}.$ ( $f_{\beta_{3}}$ ) $$

With $T_{\alpha}$ and $T_{\beta_{3}}$ are resp. the warmup times for $\alpha^{(t)}$ and $\beta^{(t)}_{3}$ to reach their final and maximal values. In practice we always set those two to the same value: $T_{\alpha}=T_{\beta_{3}}=T_{\alpha,\beta_{3}}$, and we typically use $T_{\alpha,\beta_{3}}=T$, with $T$ being the total number of iterations. $\beta_{start}$ is always set to $\beta_{1}$ in our experiments. The use of those schedulers is not always required, especially, we found those have no impact when AdEMAMix is activated later during training (see Fig. 5). The full AdEMAMix optimizer, including the schedulers, is shown in Alg. 1.

Figure: Algorithm 1 AdEMAMix optimizer. Differences with AdamW are in blue.

Figure: (a) $t_{B}=90k$.
Refer to caption: x10.png

Using $\beta_{1}=0$ to save memory. By setting $\beta_{1}=0$, we can save on memory by replacing $\hat{{\bm{m}}}_{1}^{(t)}$ by ${\bm{g}}^{(t)}$. In this case, the memory cost of AdEMAMix is the same as for AdamW. We show in App. C.1.4 (Fig. 15(b)) and App. C.1.7 that using $\beta_{1}=0$ often works, at the cost of potentially less stable training.

Computational overheads. Adding an additional EMA requires additional memory and compute. The added compute is negligible in comparison to what is required for the forward-backward, and has little impact over the total runtime as shown in Fig. 5(a). Moreover—when considering larger distributed setups—it is worth noting that AdEMAMix is not increasing communication (gradient reduction) over Adam. Therefore, we expect the runtime overhead of AdEMAMix to shrink in those settings, as data movements occupy a larger fraction of the total runtime. A more significant overhead is in terms of memory when $\beta_{1}\neq 0$, as AdEMAMix requires to allocate both ${\bm{m}}_{1}$ and ${\bm{m}}_{2}$, which are of the same size as the model parameters ${\bm{\theta}}$. We believe this issue is of lesser consequences as Fully-Sharded-Data-Paralellism (Zhao et al., , FSDP) can always be used to distribute the optimizer states across compute nodes. In the cases where memory remains an issue, one mitigation strategy could be to apply factorization strategies as in Shazeer & Stern (); Zhao et al. ().

Hyperparameter sensitivity. While we introduce up to four new hyperparameters: $\alpha$, $\beta_{3}$, $T_{\alpha}$, and $T_{\beta_{3}}$, in practice we always set $T_{\alpha}=T_{\beta_{3}}=T_{\alpha,\beta_{3}}$, and use $T_{\alpha,\beta_{3}}=T$ in most cases. We show in App. C.1.4 that $T_{\alpha}$ and $T_{\beta_{3}}$ should only be large enough to prevent instabilities early during training. While all of our experiments on language modeling use $\beta_{3}=0.9999$, other values such as $0.999$ or even $0.99999$ still can outperform the AdamW baseline (see App. C.1.4, Fig. 12). On vision tasks, the scatter plots in Fig. 6 show all the AdEMAMix models trained for those experiments, highlighting how easy it can be to find good $\alpha$ and $\beta_{3}$ values. Overall, we find the ranges of values of $\alpha,\beta_{3}$ and $T_{\alpha,\beta_{3}}$ providing improvements over AdamW to be wide. See App. C.1.4 for hyperparameter sweeps.

Limitations. AdEMAMix consists in leveraging very old gradients. Therefore, the method is best suited to settings where the number of iterations is important. We report on this effect in App. C.1.5, additionally showing how smaller values of $\beta_{3}$ (e.g. $\beta_{3}=0.999$) can be better for low iterations scenarios. Moreover, retaining gradient information over many thousands steps can pose a problem in domains requiring fast adaptation to a sudden distribution shift.

## 4 Results

In this section we use AdEMAMix to train language models (§ 4.1 and § 4.2) and vision transformers (§ 4.3) ranging from $24$M to $1.3$B parameters.

### 4.1 Transformer LLM Training

Experimental setup. We use a transformer architecture (Vaswani et al., ) with learnt positional encoding. All of our experiments are using sequences of $1{,}024$ tokens. We experiment with three sizes of transformers: $110$M, $335$M, and $1.3$B. For the learning rate, we use $3$k warmup steps followed by—unless specified—a cosine decay to $10^{-5}$. We extensively tuned the hyperparameters for both AdamW and AdEMAMix models (see App. B.1). We use the RedPajama v2 (Computer, ) dataset for all of our experiments. We use batch sizes of $64$, $96$ and $128$ for respectively our $110$M, $335$M, and $1.3$B parameter models. Depending on the model, we vary the number of iterations from $256k$ to $1.5$M. For AdEMAMix, we use $\beta_{3}=0.9999$ and $\alpha\in\{5,8,10\}$ depending on the model. A full description of the architecture and hyperparameters used is in App. B.1. We train on up to $8$ A100 NVIDIA GPUs using data-parallelism.

Figure: (a) Train time for $256k$ iters.
Refer to caption: x14.png

Why not simply increasing AdamW’s $\beta_{1}$? While our toy experiment in Fig. 2 already gives some intuition on why increasing Adam’s $\beta_{1}$ is likely not to have the same effect as having an additional EMA as in AdEMAMix, we verify this intuition by training multiple $110$M models using Adam with large $\beta_{1}\in\{0.99,0.999,0.9999,0.99999\}$. When we use a large $\beta_{1}$ from the beginning of training, we observe instabilities for larger $\beta_{1}$ values and no $\beta_{1}>0.9$ improves upon the AdamW baseline. One can imagine this to be due to increasing $\beta_{1}$ too early. Therefore, we also modify AdamW and add the same scheduler on $\beta_{1}$ as we use on AdEMAMix’s $\beta_{3}$. $\beta_{1}$ is now increased steadily over the entire training duration. While this mostly stabilizes the training, none of the experiments outperformed the baseline using $\beta_{1}=0.9$. Moreover, to rule out any effect that could be due to early training instabilities, we run the same experiments starting from a pre-trained AdamW checkpoint trained for $300k$ iterations. We simply resume training and either increase $\beta_{1}$ suddenly or using a scheduler. Here again—unlike when using AdEMAMix—none of those experiments outperform the baseline. The details of those experiments are available in App. C.1.6. Those experiments show that simply increasing the $\beta_{1}$ value in AdamW is not enough, which justifies our design of AdEMAMix.

Better perplexity for the same number of steps. For all model sizes, all the number of iterations used—ranging from $256$k to $1$M—, AdEMAMix always outperforms the AdamW baseline. In Fig. 1, we show the validation loss curves for AdamW and AdEMAMix models trained—for each size—on various numbers of tokens. For $110$M parameter models, training for $256k$ iterations gives similar results as training an AdamW model for $500k$ iterations. The gap between baseline and our method seems to be increasing as we increase the number of iterations. For $1.3$B parameter models, training using $770k$ steps is on par with training the baseline for $1.5$M iterations—reaching the same performance with $51\%$ fewer tokens (an economy of $96$M tokens). In Fig. 3(b), we observe similar improvements when using a constant and then linear decay learning rate scheduler.

Training speed comparison. We measure the time per iteration for all of our experiments. In Fig. 5(a) we plot the time required to train our $110$M and $1.3$B parameter models for $256k$ iterations. We observe that the impact of using an additional EMA on the training speed is negligible. If we were to train new models with a fix time budget, the extra iterations of the baseline would not be sufficient to close the gap. For instance, to match a $110$M parameter AdEMAMix model trained for $256k$ iterations, we need to train an AdamW model for $500k$ iterations, and the time advantage of AdamW would only allow us to do $2{,}379$ additional iterations. Moreover, as mentioned in § 3, we expect the time overhead to decrease when multi-node training is used, as IOs would become an important bottleneck, and AdEMAMix is not increasing IOs.

Consistency of the gain across token-budgets. In Fig. 1, given that enough tokens have been used w.r.t. the model size, we observe consistent gains accross token budgets. It seems our method is able to bring a constant improvement over the baseline. This can be seen more clearly in Fig. 3(b), showing results when using a constant learning rate scheduler—which removes the confounder of the cosine learning rate decay. We observe how, after an initial phase in which the gap grows, this gap becomes seemingly constant after a sufficient number of iterations.

Resume from AdamW vs. training AdEMAMix from scratch. So far we trained AdEMAMix models from scratch. We show it is also possible to switch from an AdamW to an AdEMAMix state in the middle of training. When switching to AdEMAMix at step $t_{switch}$, we initialize ${\bm{m}}_{2}^{(t_{switch})}=\mathbf{0}$ and replace $t$ by $t-t_{switch}$ in the scheduler equations—if those are used. However, we find that schedulers are not required when resuming training and report results without them in the main paper (see App. C.1.9 for more details). In Fig 5(b) and Fig.5(c), we show that (i) it is possible to improve upon the baseline when activating AdEMAMix later during training, and (ii) the earlier the switch, the larger the gain, with diminishing returns. This indicates that the improvement of AdEMAMix cannot be attributed solely to early training dynamics, but rather, late training dynamics seem to play an important role. This is further corroborated by the reverse experiment—which switches from AdEMAMix to AdamW mid-training and show a performance degradation (see results in App. C.1.3).

AdEMAMix models forget the training data slower. As an attempt to understand the reason for AdEMAMix’s improvements over AdamW, we study how fast a training batch is forgotten after being used during training. We focus on following one specific batch $B$. We start by removing $B$ from the RedPajama training data and train AdamW and AdEMAMix models. For those runs, $B$ is therefore akin to a batch from the validation set. We measure the loss on $B$ while training. Now we can resume training from various checkpoints, and inject $B$ into the training data at various times $t_{B}$. By comparing the two runs—one having trained on $B$, while the other never saw $B$—we can visualize how $B$ is learned, and how it is forgotten. After training on $B$, we expect the loss on $B$ to decrease suddenly and then increase as the model forgets the contribution of that batch. When comparing the forgetting curves for AdamW and AdEMAMix in Fig. 4, we see striking differences. AdamW models forget much faster—the loss over $B$ increases faster—than AdEMAMix models. Moreover, at the end of training, batches processed by AdEMAMix see their loss being improved over many thousands of iterations. Additional experiments on forgetting can be found in App. C.1.2.

### 4.2 Mamba LM Training

Experimental setup & results. Our experimental setup is similar to § 4.1, except that we now train $168$M parameter Mamba models (Gu & Dao, ) using the FineWeb dataset (Penedo et al., ). See App. B.2 for more details.
In Fig. 3(c) the improvement using AdEMAMix is consistent with our experiments on Transformer models. This shows how AdEMAMix’s gains can extend beyond Transformer models.

### 4.3 ViT Training

Figure: (a) $24$M params, $11$M images. (ImageNet-21k)
Refer to caption: x17.png

Experimental setup. In this section we consider a different setting in which the data is now a limited resource, and we do multiple epochs (e.g. $37$ or $320$). We use two subsets of the ImageNet (Russakovsky et al., ) dataset: (i) the widely used ImageNet-1k subset, consisting of $1.3$M images and $1{,}000$ possible classes; (ii) a filtered and pre-processed version of the ImageNet-21k (Ridnik et al., ) containing $11$M images corresponding to $10{,}450$ classes. For each, we measure the test loss on a held-out test set. We use the ViT architecture (Dosovitskiy et al., ) at two different scales: $24$M and $86$M parameters. Importantly, if it is common in the vision literature to pre-train large ViT models and finetune them on smaller datasets, this work focuses on pretraining optimization and we therefore train and test on the same distribution. The models’ hyperparameters are detailed in App. B.3, we use a batch size of $4096$ for all our experiments. We used training hyperparameters from Dosovitskiy et al. () as a starting point and did some additional tuning of the learning-rate, dropout, and weight decay for our AdamW baselines. We then use the hyperparameters of the best AdamW baseline and train AdEMAMix models with various values of $\alpha\in\{1,5,10,15,20\}$ and $\beta_{3}\in\{0.999,0.9999\}$. We train models for $320$ and $37$ epochs for resp. the ImageNet-1k and ImageNet-21k datasets, corresponding in both cases to $100k$ iterations. Data-augmentation techniques have been shown to be central to the efficient training of ViTs Touvron et al. (; ). We use simple data-augmentations, which includes mixup (Zhang et al., ). We train on $8$ A100 NVIDIA GPUs using Pytorch Fully Sharded Data-Parallelism (Zhao et al., , FSDP).

AdEMAMix for different capacity/data ratios. We consider three scenarios that differ by their capacity/data ratios. First, we trained $24$M parameter models on $11$M images (ImageNet-21k), for $37$ epochs. In this setting, as can be seen in Fig. 6(a), it is trivial to find AdEMAMix parameters outperforming the baseline in terms of both training and test accuracy. We now increase the model size to $86$M parameters. Fig. 6(b) shows it is still easy to find parameters outperforming the baseline. We now keep the model size of $86$M parameters and switch to the smaller ImageNet-1k dataset ($1.3$M images), which—given our $100k$ iterations—increases the number of epochs to $320$. In this setting, Fig. 6(c) shows that outperforming the baseline becomes difficult. These experiments show how AdEMAMix seems to perform best in scenarios with large volumes of data w.r.t. the capacity of the model. Overall, we found AdEMAMix to consistently reduce the training loss more efficiently than AdamW. When this decrease in training loss correlates with a decrease in test loss, AdEMAMix outperforms the AdamW baseline.

## 5 Conclusion

In this work, we find that old gradients can be leveraged to efficiently train large language models and ViTs. Our proposed optimizer combines two momentum terms. A slow (large $\beta$) momentum gathers information over many timestep, while a fast (low $\beta$) momentum can adapt the trajectory of the iterates to the rapidly changing loss landscape. We demonstrate the superiority of our optimizer over AdamW through a set of experiments on text modeling and image classification. We moreover reveal how our optimizer forgets the training data at a slower pace.

## 6 Acknowledgements

We thank Alaaeldin Mohamed Elnouby Ali for guiding us throughout the training of our ViT models, as well as Federico Danieli and Abhinav Moudgil for their help training Mamba models. We also thank Angelos Katharopoulos and Ronan Collobert for their insightful comments and feedback.

## Appendix A Implementation details

### A.1 Deriving the β 3 subscript 𝛽 3 \beta_{3} italic_β start_POSTSUBSCRIPT 3 end_POSTSUBSCRIPT scheduler

Let’s consider $S(t)$, the sum of the weights given to the last $t$ gradients by an EMA parameterized by $\beta\in[0,1[$:

$$ $S(t)=(1-\beta)\sum_{i=0}^{t}\beta^{i}$ $$

We want to know which timestep $t$ would correspond to a cumulative weight of $0.5$:

$$ $(1-\beta)\sum_{i=0}^{t}\beta^{i}=0.5\Leftrightarrow\beta^{t+1}=0.5 \Leftrightarrow t=\frac{\ln(0.5)}{\ln(\beta)}-1$ $$

Let $f(\beta)=\frac{\ln(0.5)}{\ln(\beta)}-1$. This function provides how many past consecutive gradients receive a cumulative weight of $0.5$.

Its inverse is:

$$ $f^{-1}(t)=0.5^{\frac{1}{t+1}}$ $$

We want a scheduler which increases $\beta$ from $\beta_{start}$ to $\beta_{end}$ such that $f(\beta)$ increases linearly. Given an interpolating parameter $\mu\in[0,1]$, this scheduler can be written as:

$$ $\beta(\mu)=f^{-1}((1-\mu)f(\beta_{start})+\mu f(\beta_{end}))$ $$

By replacing $f$ and $f^{-1}$ by their respective formula, one can arrive to:

$$ $\beta(\mu)=\exp\big{(}\frac{\ln(\beta_{start})\ln(\beta_{end})}{(1-\mu)\ln( \beta_{end})+\mu\ln(\beta_{start})}\big{)}$ $$

By setting $\beta_{end}=\beta_{3}$ and $\mu=\frac{t}{T_{\beta_{3}}}$, we arrive to the $\beta_{3}$-scheduler introduced in § 3. We show the shape of our scheduler and compare it to a linear scheduler in Fig. 7.

Figure: Figure 7: AdEMAMix’s $\beta_{3}$ scheduler. We compare our scheduler to a linear scheduler for $\beta_{start}=0.9$ and $\beta_{end}\in\{0.999,0.9999\}$. While our scheduler looks more aggressive at first glance, it increases fast for smaller values of $\beta$, and slowly for larger ones associated. This makes sense as the same increase of $\beta$ for larger $\beta$ values has a bigger impact than that same increase applied to a smaller value of $\beta$. The two linear schedulers look practically the same, despite values of $\beta_{end}$ differing by one order of magnitude. This is not the case with our scheduler.
Refer to caption: x20.png

### A.2 Pytorch implementation

The following is a code skeleton for our AdEMAMix optimizer in Pytorch (Paszke et al., ). The full implementation of AdEMAMix in Pytorch is provided in the following repository: [https://github.com/apple/ml-ademamix](https://github.com/apple/ml-ademamix)

Figure: Listing 1: AdEMAMix code skeleton using Pytorch

### A.3 Optax implementation

The following is a code skeleton for our AdEMAMix optimizer in Optax, an optimization library based on JAX (DeepMind et al., ). The full implementation of AdEMAMix in Jax is provided in the following repository: [https://github.com/apple/ml-ademamix](https://github.com/apple/ml-ademamix)

Figure: Listing 2: AdEMAMix code skeleton using JAX and Optax

## Appendix B Experimental details

### B.1 Transformer LLM experiments

Architecture details. Our architecture is based on the transformer decoder of Vaswani et al. (). We use learnt positional encoding. We use a SentencePiece (Kudo & Richardson, ) tokenizer with a vocabulary of $32000$ tokens. The model specific parameters for the different sizes are summarized in Table 1. Our implementation is using Jax (DeepMind et al., ), and we train using bfloat16, except for normalization modules and softmax which use float32. The optimizer states and operations are in float32.

**Table 1: Model parameters for our LLM experiments.**
| Model parameters | $110$M | $330$M | $1.3$B |
| --- | --- | --- | --- |
| Hidden size | $768$ | $1024$ | $2048$ |
| FFW expansion factor | $4$ | $4$ | $4$ |
| Attention heads | $12$ | $16$ | $16$ |
| Layers | $12$ | $24$ | $24$ |

How did we tune the hyperparameters? Starting from our smallest models ($110$M parameters), we first tuned hyperparameters for our AdamW models. We then use the best hyperparameters found as a starting point for AdEMAMix and tuned $\beta_{3}$ and $\alpha$. When we use schedulers for $\alpha$ and $\beta_{3}$, we set $T_{\alpha}=T_{\beta_{3}}=T$, with $T$ being the total number of iterations. Table 2 summarizes the hyperparameters we tried for this model size. The impact of AdEMAMix’s hyperparameters is discussed extensively in App. C.1.4. For our $330$M parameter models, we mostly kept the best hyperparameters found for the $110$M model and re-tuned the learning rate and gradient clipping. For AdEMAMix, we additionally tested multiple $\beta_{3}$ and $\alpha$ values. We summarize this process in Table 3. Finally, for our $1.3$B parameter models, we re-iterated the same process, re-tuning only the learning rate and gradient clipping parameters for our AdamW runs. When trying to transfer the best learning rate found to AdEMAMix, we found it to be too high, causing instabilities we couldn’t fix using gradient clipping. For this reason, we also tuned the learning rate for AdEMAMix for this model size. This process is described in Table 4.

**Table 2: Hyperparameter tuning for our $110$M parameter LLM models. In this table we report the hyperparameters we tuned for our $110$M parameter models. When multiple values are given, we bold the parameters we found to give the best results.**
| Hyperparameter | Value |
| --- | --- |
| Learning rate $\eta$ | $0.005$, $0.002$, $\mathbf{0.001}$, $0.0005$, $0.0001$ |
| Number of warmup steps | $2000$, $\mathbf{3000}$, $4000$, $5000$, $6000$ |
| Sequence length | $1024$ |
| Weight decay $\lambda$ | $\mathbf{0.1}$, $0.0$ |
| Learning rate decay scheduler | no-decay, cosine-decay |
| Batch size | $64$ |
| Gradient clipping | None, $5.0$, $1.0$, $\mathbf{0.5}$ |
| AdamW $\beta_{1}$ | $\mathbf{0.9}$, $0.99$, $0.999$, $0.9999$ |
| AdamW $\beta_{2}$ | $0.95$, $\mathbf{0.999}$ |
| AdEMAMix $\beta_{1}$ | $0.9$ |
| AdEMAMix $\beta_{2}$ | $0.999$ |
| AdEMAMix $\beta_{3}$ | $0.999$, $\mathbf{0.9999}$, $0.99999$ |
| AdEMAMix $\alpha$ | $2$, $4$, $6$, $\mathbf{8}$, $\mathbf{10}$, $15$, $20$ |

**Table 3: Hyperparameter tuning for our $330$M parameter LLM models. In this table we report the hyperparameters we tuned for our $330$M parameter models. When multiple values are given, we bold the parameters we found to give the best results.**
| Hyperparameter | Value |
| --- | --- |
| Learning rate $\eta$ | $0.001$, $\mathbf{0.0005}$, $0.0001$ |
| Number of warmup steps | $3000$ |
| Sequence length | $1024$ |
| Weight decay $\lambda$ | $0.1$ |
| Learning rate decay scheduler | cosine-decay |
| Batch size | $96$ |
| Gradient clipping | $1.0$, $0.5$, $\mathbf{0.1}$ |
| AdamW $\beta_{1}$ | $0.9$ |
| AdamW $\beta_{2}$ | $0.999$ |
| AdEMAMix $\beta_{1}$ | $0.9$ |
| AdEMAMix $\beta_{2}$ | $0.999$ |
| AdEMAMix $\beta_{3}$ | $0.999$, $\mathbf{0.9999}$, $0.99999$ |
| AdEMAMix $\alpha$ | $5$, $\mathbf{8}$, $10$, $15$ |

**Table 4: Hyperparameter tuning for our $1.3$B parameter LLM models. In this table we report the hyperparameters we tuned for our $1.3$B parameter models. When multiple values are given, we bold the parameters we found to give the best results.**
| Hyperparameter | Value |
| --- | --- |
| Number of warmup steps | $3000$ |
| Sequence length | $1024$ |
| Weight decay $\lambda$ | $0.1$ |
| Learning rate decay scheduler | cosine-decay |
| Batch size | $128$ |
| Gradient clipping | $1.0$, $0.5$, $\mathbf{0.1}$ |
| Learning rate $\eta$ for AdamW | $0.001$, $\mathbf{0.0005}$, $0.0003$, $0.0001$, $0.00005$ |
| AdamW $\beta_{1}$ | $0.9$ |
| AdamW $\beta_{2}$ | $0.999$ |
| Learning rate $\eta$ for AdEMAMix | $0.0005$, $\mathbf{0.0003}$ |
| AdEMAMix $\beta_{1}$ | $0.9$ |
| AdEMAMix $\beta_{2}$ | $0.999$ |
| AdEMAMix $\beta_{3}$ | $0.999$, $\mathbf{0.9999}$, $0.99999$ |
| AdEMAMix $\alpha$ | $1$, $3$, $\mathbf{5}$, $8$, $10$, $15$ |

Hyperparameters for experiments switching from AdamW and AdEMAMix. For experiments in Fig. 5(b) and Fig. 5(c), when we switch from AdamW to AdEMAMix during training, for our $110$M parameter models (Fig. 5(b)), we use $\alpha=2,\beta_{3}=0.9999$ and $T_{\alpha,\beta_{3}}=0$. For our $1.3$B parameter models (Fig. 5(c)), we use $\alpha=1,\beta_{3}=0.9999$ and $T_{\alpha,\beta_{3}}=0$. Other hyperparameters are inherited from the AdamW model we are switching from.

Hyperparameters used our constant learning rate scheduler experiments. For Fig. 3(b) we use $\eta=10^{-4}$, the remaining of the hyperparameters are identical as those used for our other $1.3$B experiments and provided in Table 4.

Hyperparameters used in our forgetting experiments. For the experiments in Fig. 4, we used $110$M parameter models and hyperparameters from Table 2.

### B.2 Mamba LM experiments

Architecture details. We use a Mamba architecture (Gu & Dao, ) with an embedding dimension of $768$, an expansion factor of $2$, a state size of $16$, and a depth of $24$. We use a batch size of $120$ sequences of $1024$ tokens. We use the EleutherAI/gpt-neox-20b tokenizer from Black et al. (), using the HuggingFace library (Wolf et al., ).

Hyperparameter details. We used parameters mostly taken from Gu & Dao ():

- •
Learning rate: $\eta=0.0006$,
- •
Warmup steps: $3k$,
- •
$\eta$-scheduler type: cosine decay to $10^{-5}$,
- •
Weight-decay: $\lambda=0.1$,
- •
Gradient clipping value: $1$,
- •
Total training steps $T\in\{64k,128k\}$.

For experiments with AdamW, we use $\beta_{1}=0.9$ and $\beta_{2}=0.999$. For AdEMAMix, we use $\beta_{1}=0.9,\beta_{2}=0.999,\beta_{3}=0.9999,T_{\alpha,\beta_{3}}=T$ and $\alpha=8$.

### B.3 ViT experiments

Architecture details. We use a ViT architecture following Dosovitskiy et al. (). The architecture details for both our $24$M and $86$M parameter models can be seen in Table 5. We do not use an EMA of the iterates as our final model. Our implementation is in Pytorch (Paszke et al., ), and use bfloat16.

**Table 5: Model parameters for our ViT experiments.**
| Model parameters | $24$M | $86$M |
| --- | --- | --- |
| Patch size | $16$ | $16$ |
| Number of patches | $14\times 14$ | $14\times 14$ |
| Image size | $224$ | $224$ |
| Embedding dim. | $384$ | $768$ |
| MLP dim. | $1536$ | $3072$ |
| Layers | $12$ | $12$ |
| Number of heads | $6$ | $12$ |

Data augmentation and Mixup. While more sophisticated augmentation methods exist (Touvron et al., ), we simply use random-resized cropping, random horizontal flip ($p=0.5$), and random color jitter (applied with a probability $p=0.8$).

How did we tune the hyperparameters? For each model size, we started by tuning the hyperparameters for the AdamW baseline. The hyperparameters we tuned and the values we retained for our three different settings are in Table 6. For our experiments on ImageNet-1k, given that all models overfit the dataset, we select the best model according to the minimum validation loss, akin to using early stopping. For AdEMAMix, for each setting, we use the hyperparameters of the best AdamW model and only tune $\alpha$ and $\beta_{3}$. For each setting, we train $8$ models using $\alpha\in\{1,5,10,15,20\}$ and $\beta_{3}\in\{0.999,0.9999\}$. All of the AdEMAMix models are shown in Fig. 6 and Fig. 30. Only the best AdamW model is shown in Fig. 30.

**Table 6: Hyperparameters tuning for our ViT AdamW experiments. We started our hyperparameter search from the values given in (Dosovitskiy et al. (), Table 3), which recommends using (learning-rate, weight decay, dropout)$=(0.003,0.3,0.1)$ when training on ImageNet-1k and $(0.001,0.03,0.1)$ when training on ImageNet-21k. The various hyperparameters we experimented with are in the following table, values which gave us the lowest validation loss are in bold. For each setting, AdEMAMix experiments use the same hyperparameters as the best AdamW baselines.**
| Setting | $24$M params, $11$M images | $86$M params, $11$M images | $86$M params, $1.3$M images |
| --- | --- | --- | --- |
|  | (ImageNet-21k) | (ImageNet-21k) | (ImageNet-1k) |
| Learning rate | $0.001,\mathbf{0.003},0.005$ | $\mathbf{0.001},0.003$ | $0.001,0.003,\mathbf{0.005}$ |
| Weight decay | $\mathbf{0.03},0.1,0.3$ | $0.03,\mathbf{0.1},0.3$ | $0.03,0.1,\mathbf{0.3},0.5,0.7$ |
| Dropout | $0.1$ | $0.1$ | $0.1$ |
| AdamW $\beta_{1}$ | $\mathbf{0.9},0.99,0.999$ | $\mathbf{0.9},0.99,0.999$ | $\mathbf{0.9},0.99,0.999$ |
| Batch size | $4096$ | $4096$ | $4096$ |

## Appendix C Additional results

### C.1 LLM experiments

#### C.1.1 In-Context Learning (ICL) results

In-Context Learning (ICL) results. We evaluate our largest ($1.3$B) LLM models on in-context learning tasks. We use the lm-eval package (Gao et al., ). We consider the following tasks:

HellaSwag (Zellers et al., ),
Winogrande (Sakaguchi et al., ),
ARC (Bhakthavatsalam et al., ),
BoolQ (Clark et al., ),
LogiQA (Liu et al., ),
MathQA (Amini et al., ),
MMLU (Son et al., ),
OpenbookQA (Mihaylov et al., ),
PIQA (Bisk et al., ),
PubmedQA (Jin et al., ),
RewardBench (Lambert et al., ),
Sciq (Welbl et al., ),
TruthfulQA (Lin et al., ).

The two AdEMAMis and AdamW models we benchmark have been trained for $1$M steps, with a batch size of $128$, corresponding to around $131$B tokens. The results of the evaluation are in Table 7. We observe how the model trained with AdEMAMix is outperforming the AdamW model on nearly all of the tasks, sometimes by a large margin for e.g. PubmedQA. In addition to Table 7, we track the evolution of some of those scores during training for MMLU, RewardBench, ARC-Challenge, and ARC-Easy, results can be seen in Fig. 8. For that figure, we modify the way MMLU is evaluated. Instead of appending all the multiple answers, we concatenate the question prompt and each answer separately, and pick the most likely answer. We found it allows a much better comparison of smaller models, which otherwise fluctuates around random guessing.

**Table 7: In-context learning results for our $1.3$B parameter models. We compare AdamW and AdEMAMix models trained on $131$B tokens. We use the lm-eval package.**
| Task | AdamW | AdEMAMix | Task | AdamW | AdEMAMix |
| --- | --- | --- | --- | --- | --- |
| ARC-Challenge | $0.262$ | $\mathbf{0.274}$ | OpenbookQA | $\mathbf{0.240}$ | $0.238$ |
| ARC-Easy | $0.612$ | $\mathbf{0.619}$ | PIQA | $\mathbf{0.715}$ | $\mathbf{0.715}$ |
| BoolQ | $0.569$ | $\mathbf{0.576}$ | PubmedQA | $0.556$ | $\mathbf{0.632}$ |
| HellaSwag | $0.426$ | $\mathbf{0.436}$ | RewardBench | $0.569$ | $\mathbf{0.573}$ |
| LogiQA | $\mathbf{0.235}$ | $0.225$ | RewardBench (reasoning) | $0.617$ | $\mathbf{0.630}$ |
| MathQA | $0.226$ | $\mathbf{0.236}$ | Sciq | $0.903$ | $\mathbf{0.907}$ |
| Winogrande | $0.563$ | $\mathbf{0.580}$ | TruthfulQA | $\mathbf{0.361}$ | $0.352$ |
| MMLU | $0.244$ | $\mathbf{0.248}$ |  |  |  |

Figure: (a) MMLU.
Refer to caption: x21.png

#### C.1.2 More results on forgetting

Evolution of forgetting during training. In Fig. 4 in the main paper, we follow the loss on one specific batch over the entire training. In the following experiment, we do a closeup on the forgetting of different batches as training progresses. The goal being to visualize how later batches are ultimately more remembered, and compare the forgetting profiles of AdamW and AdEMAMix. For this, in Fig. 9(a) and Fig. 9(b), we follow fixed batches at different stages of the training process. Every $10k$ iterations, we measure the loss over a specific batch $B$ before and after training on that batch.

Figure: (a) AdamW.
Refer to caption: x25.png

Why are training batches forgotten more at the end of training? Given the previous observation about how both AdamW and AdEMAMix models are forgetting the training batches at different paces given different stages of training, one can wonder what is causing this phenomenon? To answer this question, we contrast the results of Fig. 4 with results obtained using a different scheduler with a constant learning rate and a linear decay. We train $110$M parameter models for $300k$ iterations, using a max learning rate of $0.001$ and a batch size of $64$. Results for those experiments are in Fig. 10. Those results indicate that the decaying learning rate might be the main parameter controlling how much a batch is remembered during training. This has interesting implications when selecting which learning rate decay to use. A cosine decay, with a rather long period of decay, might remember more than a constant learning rate scheduler with only a small number of steps of linear decay at the end.

Figure: (a) $t_{B}=230k$.
Refer to caption: x27.png

#### C.1.3 Removing the second EMA mid-training

Removing the second EMA mid-training. In Fig. 5(b) and Fig. 5(c), we looked at what is happening when we switch from AdamW to AdEMAMix in the middle of training. In this section, we will look at the opposite conversion: removing the ${\bm{m}}_{2}$ parameter of AdEMAMix during training, effectively switching back to AdamW. Results for this experiment can be seen in Fig. 11. Right after the switch, we observe a drop of the loss, followed by an increase and finally back to convergence. Ultimately, the final loss value is in between the ones obtained training only using AdamW and only using AdEMAMix.

Figure: Figure 11: Switching from AdEMAMix to AdamW at $t_{switch}=128k$. Using our AdEMAMix optimizer, we train two $110$M parameters models with different $T_{\alpha}=T_{\beta_{3}}\in\{90k,256k\}$. For historical reasons, those experiments were done using a batch size of $32$ instead of $64$ (as we used for all the other $110$M parameter experiments in this paper). We switch from AdEMAMix to AdamW at $t_{switch}=128k$. At first, the removal of the $+\alpha{\bm{m}}_{2}$ contribution in the update is suddenly decreasing the effective learning rate, causing the loss to drop suddenly. The loss then increases slightly and ends up in between the AdamW only (blue curve) and AdEMAMix only (grey curves) losses.
Refer to caption: x31.png

#### C.1.4 Hyperparameter sensitivity

Hyperparameter sensitivity. Depending on whether the $\alpha$ and $\beta_{3}$ schedulers are used, AdEMAMix adds up to $4$ new hyperparameters: $\alpha$, $\beta_{3}$, $T_{\alpha}$ and $T_{\beta_{3}}$. In all our experiments we tied $T_{\alpha}=T_{\beta_{3}}=T_{\alpha,\beta_{3}}$. In this section we analyze the sensitivity of AdEMAMix to those hyperparameters. We study the impact of $\alpha$ and $\beta_{3}$ in Fig. 12, revealing wide ranges of values for which hyperparameters are outperforming the AdamW baseline. We study the sensitivity of $T_{\alpha,\beta_{3}}$ as well as justify the choice of using a scheduler on $\alpha$ in Fig. 13. When training from scratch $T_{\alpha,\beta_{3}}$ needs simply to be large enough to avoid early divergence. In Fig. 15 we study the sensitivity to the gradient clipping and AdEMAMix’s $\beta_{1}$ value. While gradient clipping can help stabilize training and smooth the loss curves, it has little impact over the final loss value. Reducing the value of $\beta_{1}$, we observe some loss spikes, yet the final loss value is relatively unaffected.

Figure: (a) Sensitivity to $\alpha$.
Refer to caption: x32.png

Figure: (a) Sensitivity to $T_{\alpha}=T_{\beta_{3}}=T_{\alpha,\beta_{3}}$.
Refer to caption: x34.png

Figure: (a) No $\alpha$-scheduler ($110$M).
Refer to caption: x36.png

Figure: (a) Sensitivity to gradient clipping.
Refer to caption: x38.png

#### C.1.5 Impact of training for fewer iterations

Sensitivity to the number of iterations. As we rely on old gradients, on question that arises is whether AdEMAMix would still perform well when reducing the number of iterations. In this section we compare the loss obtained by $110$M parameter models when halving the number of iterations and doubling the batch size, in such a way that the number of tokens consumed for training is always the same ($17$B). In Fig. 16, we observe that decreasing the number of iterations too much increases the final loss of the model. This effect is more pronounced for AdEMAMix. However, at both $32k$ and $64k$ iterations, AdEMAMix still outperforms AdamW. In Fig. 17 we observe that reducing $\beta_{3}$ can mostly alleviate the problem.

Figure: (a) AdamW with fewer steps.
Refer to caption: x40.png

Figure: (a) AdEMAMix with fewer steps $\beta_{3}=0.9999$ (same as Fig. 16(b)).
Refer to caption: x43.png

#### C.1.6 Limitations of a single EMA

Results on a 2D toy example. A natural question that arises from our method is whether it is possible to obtain the same results without the additional EMA. In this section we aim to convince the reader that the two EMAs are required. First, we propose to study a small toy 2D optimization problem:

$$ $(x^{\star},y^{\star})\in\operatorname*{arg\,min}_{(x,y)}f(x,y),\quad\text{ with }f(x,y)=8(x-1)^{2}\times(1.3x^{2}+2x+1)+0.5(y-4)^{2}$ $$

This function was introduced by Liu et al. () as a function with sharp curvature along the $x$-axis, and flatter curvature along the $y$-axis.
Initializing the first iterate ${\bm{x}}^{(0)}=(0.3,1.5)$, we start optimizing with (i) Adam with a large $\beta_{1}=0.999$, and (ii) AdEMAMix with $\beta_{1}=0.9$ and $\beta_{3}=0.999$. In both cases $\beta_{2}=0.999$. To make the experiment interesting, we initialize the EMA buffers for both methods to either $(-3,0)$ or $(-0.8,-3)$. This has for effect to give an initial "speed" to the first iterate. As a result of this speed, Adam with a large $\beta_{1}$ is unable to correct his trajectory. In contrast, using two EMAs, one with a small, and one with a large $\beta$—as in AdEMAMix—converges to the solution. The fast changing EMA can correct for the bias introduced by the slow changing EMA. See Fig. 18 for results.

Figure: (a) ${\bm{m}}_{1}^{(0)}={\bm{m}}_{2}^{(0)}=[-3,0]$.
Refer to caption: x46.png

Results training $110$M parameters LMs. To further demonstrate that simply increasing $\beta_{1}$ in AdamW does not provide nearly the same gains as AdEMAMix, we run several additional experiments. In Fig. 19 we show what we obtain when training from scratch using a single EMA with a $\beta_{1}\in\{0.99,0.999,0.9999,0.99999\}$. Naively increasing $\beta_{1}$ in AdamW does not work; adding a scheduler on $\beta_{1}$ to smoothly increase its value during the entirety of the training also fails. In Fig. 20, we show results when increasing $\beta_{1}$ in the middle of training, with and without scheduler on $\beta_{1}$. This differs from the previous setting as we bypass the initial training phase capable of causing instabilities (see § 3), as well as bypass the initial iterations during which the bias correction done by AdamW can have an impact. Here again we observe increasing the $\beta_{1}$ value does not provide any noticeable gain.

Figure: (a) Vanilla AdamW.
Refer to caption: x48.png

Figure: (a) Increasing $\beta_{1}$ without warmup.
Refer to caption: x50.png

#### C.1.7 Removing m 1 subscript 𝑚 1 m_{1} italic_m start_POSTSUBSCRIPT 1 end_POSTSUBSCRIPT by setting β 1 = 0 subscript 𝛽 1 0 \beta_{1}=0 italic_β start_POSTSUBSCRIPT 1 end_POSTSUBSCRIPT = 0

Effect of the batch size when $\beta_{1}=0$. In this section we investigate the impact of setting $\beta_{1}=0$, which allows to replace ${\bm{m}}_{1}^{(t)}$ by the current gradient ${\bm{g}}^{(t)}$, saving memory. In this scenario the memory complexity of AdEMAMix is the same as the one of AdamW. We noticed in Fig. 15(b) how reducing $\beta_{1}$ can make the training slightly more unstable, triggering some loss spikes. In Fig. 21 we study the impact of the batch size on those instability, revealing that $\beta_{1}=0$ is less stable when using larger batch sizes. However, despite the spikes, the final loss for $\beta_{1}=0$ and $\beta_{1}=0.9$ are very similar—$\beta_{1}=0.9$ yielding slightly better results for smaller batch sizes.

Figure: Figure 21: Effect of batch size with $\beta_{1}=0$. For a $110$M parameter model, we use $\beta_{1}\in\{0.0,0.9\}$ and vary the batch-size $\in\{4,8,16,32,64,128\}$. Lighter curves use $\beta_{1}=0$, darker curves are their $\beta_{1}=0.9$ counterpart. For small batch sizes, no instabilities are observed, and $\beta_{1}=0.9$ reaches slightly better performances. For larger batch sizes, instabilities are observed with using $\beta_{1}=0$, yet the final loss values are similar as the ones obtained with $\beta_{1}=0.9$. This indicates that one can potentially save memory without compromising results.
Refer to caption: x52.png

Setting $\beta_{1}=0$ for $1.3$B parameter models. In this section, using a $1.3$B parameter model, we re-do the same experiment as in Fig. 1, but we set $\beta_{1}=0$. In Fig. 22, we show the results of this experiment.

Figure: Figure 22: Using $\beta_{1}=0$ with a $1.3$B parameter model. We use $\beta_{2}=0.999$ for all those experiments. We observe how it is possible to converge despite using $\beta_{1}=0$. The final loss with $\beta_{1}=0$ is slightly better than the one with $\beta_{1}=0.9$.
Refer to caption: x53.png

#### C.1.8 Comparing m 1 + α ⁢ m 2 subscript 𝑚 1 𝛼 subscript 𝑚 2 m_{1}+\alpha m_{2} italic_m start_POSTSUBSCRIPT 1 end_POSTSUBSCRIPT + italic_α italic_m start_POSTSUBSCRIPT 2 end_POSTSUBSCRIPT versus ( 1 − α ) ⁢ m 1 + α ⁢ m 2 1 𝛼 subscript 𝑚 1 𝛼 subscript 𝑚 2 (1-\alpha)m_{1}+\alpha m_{2} ( 1 - italic_α ) italic_m start_POSTSUBSCRIPT 1 end_POSTSUBSCRIPT + italic_α italic_m start_POSTSUBSCRIPT 2 end_POSTSUBSCRIPT

In the equation AdEMAMix, we combine ${\bm{m}}_{1}$ and ${\bm{m}}_{2}$ using ${\bm{m}}_{1}+\alpha{\bm{m}}_{2}$. In this section we contrast this to using a convex combination $(1-\alpha){\bm{m}}_{1}+\alpha{\bm{m}}_{2}$, with $\alpha\in[0,1]$.

Are those two somewhat equivalent? Given the combination $\eta({\bm{m}}_{1}+\alpha{\bm{m}}_{2})$, it is possible to rewrite it as a convex combination as such:

$$ $\displaystyle\eta({\bm{m}}_{1}+\alpha{\bm{m}}_{2})=\hat{\eta}((1-\hat{\alpha}) {\bm{m}}_{1}+\hat{\alpha}{\bm{m}}_{2}),\quad\text{with $\hat{\eta}=\eta(\alpha +1)$ and $\hat{\alpha}=\frac{\alpha}{\alpha+1}$}.$ $$

Therefore, the two formulas are equivalent up-to a reparametrization. However, in practice, we use schedulers on $\eta$ and $\alpha$. While at any timestep $t$, it is possible to reparametrize one formula into the other, there is not one single reparametrization that holds for all timesteps. Assuming that we do not change the nature of the schedulers for $\alpha$ and $\eta$, those two formulations are therefore *not equivalent*. In Fig. 23 we show the evolution of the weights given to ${\bm{m}}_{1}$ and ${\bm{m}}_{2}$ for the two approaches. We assume a cosine scheduler for $\eta$ and a linear scheduler for $\alpha$. We can see how the two formulations give very different weights to ${\bm{m}}_{1}$ and ${\bm{m}}_{2}$.

Figure: (a) $\eta^{(t)}({\bm{m}}_{1}+\alpha^{(t)}{\bm{m}}_{2})$.
Refer to caption: x54.png

Is the convex combination better? Now that we established that the two formulations are different, which one is better? To answer this we run two grid searches, varying $\eta\in\{0.4,0.6,0.8,1,1.2,1.4,1.6\}$ and $\alpha\in\{2,4,6,8,10,12,14\}$ for the formulation in equation AdEMAMix, and $\alpha\in\{0.3,0.4,0.5,0.6,0.7,0.8,0.9\}$ for the convex combination formulation. Results in Fig. 24 show how the formulation used in AdEMAMix provides a better loss *for all the pairs $(\eta,\alpha)$ tested*, also outperforming AdamW for all of those. In contrast, for the convex combination, while most $(\eta,\alpha)$ pairs outperform AdamW, the improvement is smaller.

Figure: (a) ${\bm{m}}_{1}+\alpha{\bm{m}}_{2}$.
Refer to caption: x56.png

Getting the same results by changing the schedulers for the convex combination. We concluded before that—unless we tamper with the schedulers—it is not possible to find a reparametrization that would work for all timesteps and make the two formulations equivalent. What if we could change the schedulers? Which schedulers would we need to use with the convex formulation, to get the same results as with the other formulations. The answer is simple, let $\eta^{(t)}$ and $\alpha^{(t)}$ denote the cosine and linear schedulers for resp. $\eta$ and $\alpha$. The convex combination would need to use the following schedulers:

$$ $\displaystyle\hat{\eta}^{(t)}=\eta^{(t)}(\alpha^{(t)}+1),\quad\quad\hat{\alpha }^{(t)}=\frac{\alpha^{(t)}}{\alpha^{(t)}+1}.$ $$

Therefore, to get equivalent results as for the original AdEMAMix formulation, the learning rate scheduler would need to be $\hat{\eta}^{(t)}$, which is not a cosine scheduler. The $\alpha$ scheduler, $\hat{\alpha}^{(t)}$, would not be linear. We plot those schedulers in Fig. 25.

Figure: (a) $\eta$ schedulers.
Refer to caption: x58.png

#### C.1.9 Miscellaneous results

Impact of $\beta_{3}$ and $\alpha$ schedulers when starting AdEMAMix from AdamW. As mentioned in § 3, we do not necessarily need the $\alpha$ and $\beta_{3}$ schedulers when resuming training from a sufficiently late checkpoint. Fig. 5(b) and Fig. 5(c) do not use schedulers, unless we start using AdEMAMix from scratch. In Fig. 26, we vary the warmup periods for $\alpha$ and $\beta_{3}$ when starting AdEMAMix from an AdamW checkpoint at $t_{switch}=300k$. We observe how increasing $T_{\beta_{3}}=T_{\alpha}$ only slows down the convergence. This serves to illustrate that the schedulers for $\alpha$ and $\beta_{3}$ are only required to stabilise the optimization during the early training phase.

Figure: Figure 26: No need to use schedulers when switching from AdamW to AdEMAMix sufficiently late. For a $110$M parameter model, we switch from AdamW to AdEMAMix after $t_{switch}=300k$ iterations. We vary $T_{\beta_{3}}=T_{\alpha}\in\{0,25k,50k,100k\}$, the other prameters are fixed: $\beta_{3}=0.9999,\alpha=2$, the learning rate is inherited from AdamW and stays on his decaying course following a cosine decay. AdEMAMix’s ${\bm{m}}_{2}$ is initialized to $\mathbf{0}$. The larger $T_{\beta_{3}}=T_{\alpha}$, the slower the convergence is, the final loss value are very similar. This shows that the schedulers for $\alpha$ and $\beta_{3}$ are not important when starting AdEMAMix from a sufficiently late checkpoint.
Refer to caption: x60.png

AdEMAMix from scratch without schedulers. In this section we provide more justification for the use of schedulers on $\alpha$ and $\beta_{3}$. We reveal that, without the use of schedulers, the norms of the updates increases significantly, even for small $\alpha$ values. Unstable and large weight updates are known to occur in the early phases of training when learning rate warmup is not used (Gotmare et al., ). Using large momentum values too early seem to increase this phenomenon. See Fig. 27.

Figure: (a) Norm of $1000$ consecutive updates.
Refer to caption: x61.png

Impact of the linear decay duration when using a constant $\eta$-scheduler. In Fig. 3(b) we show results using AdEMAMix with a linear warmup $\rightarrow$ constant $\rightarrow$ linear decay learning rate scheduler. We used $100k$ of linear decay. In this section we experiment with $200k$ steps of linear decay, the rest of the parameters are the same: we use a $1.3$B parameter model with a max learning rate of $10^{-4}$ and remaining hyperparameters as in Table 4. Results can be seen in Fig. 28.

Figure: Figure 28: Testing a different decay duration compared to Fig. 3(b). Increasing the length of the learning rate decay improves the results for each methods. The gap between AdEMAMix using $\beta_{3}=0.99999$ and $\beta_{3}=0.9999$ is shrinking, but the gap between AdEMAMix models and AdamW remains.
Refer to caption: x63.png

Same figure as Fig. 1, including the AdamW trained on $197$B tokens. In Fig. 1 we represent the AdamW experiment trained on $197$B tokens by a blue horizontal line for aesthetic reasons. In Fig. 29 we include this missing curve to the same plot.

Figure: Figure 29: Same as Fig. 1 adding AdamW $197$B.
Refer to caption: x64.png

### C.2 ViT experiments

Top-k accuracies. In Fig. 6, we plot the test and train loss for the final iterates. In Fig. 30, we give a more detailed view by reporting the evolution of the training loss, test loss, and top-1 accuracy. Looking at the first row, the training loss for AdEMAMix is systematically better than the AdamW baseline. The second row shows that in cases where the test loss correlates well with the train loss, AdEMAMix works well. The top-1 accuracy carries the same message.

Figure: (a) $24$M params, $11$M images. (ImageNet-21k)
Refer to caption: x65.png

### C.3 Comparison with other methods

#### C.3.1 Comparison with AdMeta and DEMA

Double Exponential Moving Average (DEMA). Originally introduced by Mulloy (), a DEMA originally aimed to emphasize the weight of recent asset price fluctuations, making the DEMA indicator more reactive to changes compared to simple EMAs. Given notations introduced in the main paper, let $\text{EMA}_{\beta}^{(T-N:T)}\triangleq\text{EMA}(\beta,{\bm{g}}^{(T-N)},\ldots
,{\bm{g}}^{(T)})$, let $N$ be the window size representing how many consecutive values are considered in the average, the formula for DEMA can be written as follows:

$$ $\text{DEMA}(\beta,{\bm{g}}^{(T-2N)},\ldots,{\bm{g}}^{(T)})=2\cdot\text{EMA}_{ \beta}^{(T-N:T)}-\text{EMA}(\beta,\text{EMA}_{\beta}^{(T-2N:T-N)},\ldots,\text {EMA}_{\beta}^{(T-N:T)}).$ $$

If a simple EMA tends to give a significant weight to more recent observations, a DEMA emphasizes this behaviour by removing some of the weight given to older observations. This is not what we suggest doing in this work, we want both high sensitivity to recent observations and non-negligible weights given to older observations.

AdMeta. Chen et al. () take inspiration over DEMA and use nested EMAs in their AdMeta-S optimizer:

$$ $\displaystyle\begin{cases}{\bm{m}}_{1}^{(t)}&=\beta_{1}{\bm{m}}_{1}^{(t-1)}+{ \bm{g}}^{(t)}\\ {\bm{h}}^{(t)}&=\kappa{\bm{g}}^{(t)}+\mu{\bm{m}}_{1}^{(t)}\\ {\bm{m}}_{2}^{(t)}&=\beta_{2}{\bm{m}}_{2}^{(t-1)}+(1-\beta_{2}){\bm{h}}^{(t)} \\ {\bm{\theta}}^{(t)}&={\bm{\theta}}^{(t-1)}-\eta{\bm{m}}_{2}^{(t)}.\end{cases}$ (AdMeta-S) $$

With $\mu$ and $\kappa$ parameterized by $\beta_{1}\in[0,1[$ as such:

$$ $\displaystyle\mu$ $\displaystyle=25-10(\beta_{1}+\frac{1}{\beta_{1}})$ $\displaystyle\kappa$ $\displaystyle=\frac{10}{\beta_{1}}-9.$ $$

In their AdMeta-S experiments, they use $\beta_{1}=0.9$, corresponding to $(\mu,\kappa)=(4.88,2.11)$. $\beta_{2}$ takes values ranging from $0.1$ to $0.4$. As a results, unlike AdEMAMix, AdMeta is not leveraging very old gradients.
Analysing the AdMeta algorithm, we see that it consists in two nested EMAs. We show the shape of nested EMAs in Fig. 31. In sharp contrast with our approach we observe that (i) it reduces the weights given to recent gradients and (ii) it gives a small weight to old gradients.

Figure: (a) Linear scale.
Refer to caption: x68.png

#### C.3.2 Comparison with Lion

The Lion optimizer. Chen et al. () derived the following optimizer. We change notations to facilitate the comparison with AdEMAMix:

$$ $\displaystyle\begin{cases}{\bm{\theta}}^{(t)}&={\bm{\theta}}^{(t-1)}-\eta\cdot \big{(}\text{sign}(\alpha{\bm{m}}^{(t-1)}+(1-\alpha){\bm{g}}^{(t)})+\lambda{ \bm{\theta}}^{(t-1)}\big{)}\\ {\bm{m}}^{(t)}&=\beta{\bm{m}}^{(t-1)}+(1-\beta){\bm{g}}^{(t)}.\end{cases}$ (Lion) $$

The Lion optimizer uses a sign function and updates its EMA after updating the parameters. Moreover, it does not normalize the updates. While Lion and AdEMAMix are quite different from each others, we can draw one similarity. Indeed, the interpolation $\alpha{\bm{m}}^{(t-1)}+(1-\alpha){\bm{g}}^{(t)}$ is similar to combining two EMAs, one of them using $\beta=0$. We also explore the possibility of setting AdEMAMix’s $\beta_{1}$ to 0 0 in § 3, we show in App. C.1.7 and App. C.1.4 (Fig. 15(b)) that this often works despite a more unstable training. Interestingly, Chen et al. () find that larger $\beta=0.99$ values work best. Beside this similarity, the two optimizers behave very differently, the biggest difference being the use of the sign function, which Chen et al. () claim can help regularize the training. We test the Lion optimizer on language modeling. To tune the hyperparameters, we took values from Chen et al. () as a starting point, as well as recipes provided by the Optax Jax library. A summary of our hyperparameter tuning is in Table 8:

**Table 8: Hyperparameter tuning for Lion $110$M parameter models. When multiple values are given, we bold the parameters we found to give the best results.**
| Hyperparameter | Value |
| --- | --- |
| Learning rate $\eta$ | $0.00005,0.0001,0.0002,\mathbf{0.0004},0.0006,0.0008,0.001$ |
| Number of warmup steps | $3000$ |
| Sequence length | $1024$ |
| Weight decay $\lambda$ | $0.01,0.125,0.166,\mathbf{0.25},0.5,1.0$ |
| Learning rate decay scheduler | cosine-decay |
| Batch size | $64$ |
| Gradient clipping | $1.0,\mathbf{0.5}$ |
| Lion $\alpha$ | $0.5,0.7,\mathbf{0.9}$ |
| Lion $\beta$ | $\mathbf{0.99},0.9999$ |

The training curve associated to the best hyperparameters is in Fig. 32. We observe that Lion is not outperforming our carefully tuned AdamW baseline. Moreover, no attempt to increase $\beta$ beyond $0.99$ was successful, as those models mostly diverged. This emphasizes one of our main contributions: the introduction of schedulers enabling the use of large momentum values.

Figure: Figure 32: Comparing AdamW with the best Lion model found. Hyperparameters used are in Table 8. We train many Lion models and select the one with the lowest final validation loss. Lion is not outperforming our carefully tuned AdamW baseline.
Refer to caption: x70.png

#### C.3.3 Adding a third momentum term ( ∼ similar-to \sim ∼ AggMo)

Lucas et al. (, AggMo) propose to add an arbitrary number ($K$) of momentum terms to the gradient descent algorithm:

$$ $\displaystyle\begin{cases}{\bm{m}}_{i}^{(t)}&=\beta_{i}{\bm{m}}_{i}^{(t-1)}+{ \bm{g}}^{(t)}\quad\text{for $1\leq i\leq K$}\\ {\bm{\theta}}^{(t)}&={\bm{\theta}}^{(t-1)}-\frac{\eta}{K}\sum_{1\leq i\leq K}{ \bm{m}}_{i}^{(t)}.\end{cases}$ (AggMo) $$

In order to see whether more momentum terms can be beneficial, we modify our AdEMAMix optimizer to include a third momentum ${\bm{m}}_{3}$. We name the resulting optimizer Ad3EMAMix:

$$ $\displaystyle\begin{cases}{\bm{m}}_{1}^{(t)}=\beta_{1}{\bm{m}}_{1}^{(t-1)}+(1- \beta_{1}){\bm{g}}^{(t)},\quad\quad\hat{{\bm{m}}}_{1}^{(t)}=\frac{{\bm{m}}_{1} ^{(t)}}{1-\beta_{1}^{t}}\\ {\color[rgb]{0,0,1}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,1}{\bm{m}}_{2} ^{(t)}}{\color[rgb]{0,0,1}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,1}= \beta_{3}{\bm{m}}_{2}^{(t-1)}+(1-\beta_{3}){\bm{g}}^{(t)}}\\ {\color[rgb]{0,0,1}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,1}{\bm{m}}_{3} ^{(t)}}{\color[rgb]{0,0,1}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,1}= \beta_{4}{\bm{m}}_{3}^{(t-1)}+(1-\beta_{4}){\bm{g}}^{(t)}}\\ {\bm{\nu}}^{(t)}=\beta_{2}{\bm{\nu}}^{(t-1)}+(1-\beta_{2}){{\bm{g}}^{(t)}}^{2} ,\quad\quad\hat{{\bm{\nu}}}^{(t)}=\frac{{\bm{\nu}}^{(t)}}{1-\beta_{2}^{t}}\\ {\bm{\theta}}^{(t)}={\bm{\theta}}^{(t-1)}-\eta\big{(}\frac{\hat{{\bm{m}}}_{1}^ {(t)}{\color[rgb]{0,0,1}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,1}+\alpha ({\bm{m}}_{2}^{(t)}+{\bm{m}}_{3}^{(t)})}}{\sqrt{\hat{{\bm{\nu}}}^{(t)}}+{ \epsilon}}+\lambda{\bm{\theta}}^{(t-1)}\big{)}.\end{cases}$ (Ad3EMAMix) $$

We Train a $110$M parameter model with same hyperparameters as in Table 2, but we use $\alpha=4$ instead of $\alpha=8$. We apply the same scheduler to $\beta_{4}$ as for $\beta_{3}$. In Fig. 33 we show the resulting validation loss curves for various $(\beta_{1},\beta_{3},\beta_{4})$ triplets. Our experiments show no advantage of adding an extra $\beta_{4}$-EMA. Rather than carrying the message that we should use more momentum terms, our work shows how—in order to use a large momentum—a term should be added to stay sensitive to local fluctuations of the loss landscape. In that sense, we believe two terms should be enough.

Figure: Figure 33: Adding an extra momentum to AdEMAMix does not provide better results. We train multiple $110$M parameter Ad3EMAMix models with various $(\beta_{1},\beta_{3},\beta_{4})$ triplets (lighter curves). None of those models end up outperforming AdEMAMix, indicating that the additional ${\bm{m}}_{3}$ momentum term is seemingly not bringing any advantage.
Refer to caption: x71.png

## Peer Reviews

confidence: 4
contribution: 3
initial_rating: 6
presentation: 4
questions: 1. The motivation on the Rosenbrock function is rather weak.  Why do you consider the Rosenbrock function? How does it relate to LLM training?  Why do you think the implication on the Rosenbrock function can be transferred to LLM?  To me, there is a huge gap between these two.

2. Need more evidence that "beta1 = 0.9 is indeed suboptimal for LLMs". While I fully understand Figure 3, i still did not fully convinced that "LLM training really needs to use more historical gradients, in one way or another". More LLM-related evidence is needed to convince the that it is indeed a major bottleneck. Despite the impressive performance, the overall motivation still seems quite vague to me. I suggest the authors put more effort into designing more experiments for better motivation.

3. It is difficult to understand the highlighted sentence in the introduction: "While changing the direction of the slow momentum is difficult, any adjustment orthogonal to that direction is easy—which favors fast progress in sinuous canyon-like landscapes." I don't understand what you mean by "change the direction of the slow momentum" or "any adjustment orthogonal to that direction". Please explain more.  Also, the authors mentioned "canyon-like landscapes (in Rosenbork function)." but never connected it to LLM in the script. This also makes the motivation rather weak.

4. I don't quite understand how to read Figure 2. For instance: 

   -- in Figure 2 (a), is AdaEMAMix considered better than Adam? It is not clear whether we can draw such conclusion based on the figure. In my opinion, at least 10 x more iterations are needed to draw valid conclusions.

   -- in Figure 2 (c), did AdaEMAMix converge to the optimal solution? 

5. The proposed method boosts performance by using extra memory (to store an additional copy of momentum). Though many readers might regard it as a drawback,  I personally think such a trade-off is acceptable as long as the performance gain is significant.  Further, I think AdaEMAMix can be combined with some orthogonal methods to reduce memory. For instance, AdaEMAMix can be combined with the recent method Adam-mini [1] to reduce the memory for V. I suggest the authors try it out.



6. Some missing related works: [2] proves that vanilla Adam can converge under a wide range of beta1  = 0, 0.5, 0.9, 0.99, etc., as opposed to the divergence result in Reddi et al. 2018. This result lays down a preliminary foundation of this work.  Without the theoretical guarantee, it would be dangerous to play with beta1 of Adam.

[1] Zhang, Y., Chen, C., Li, Z., Ding, T., Wu, C., Ye, Y., ... & Sun, R. (2024). Adam-mini: Use fewer learning rates to gain more. *arXiv preprint arXiv:2406.16793*.

[2] Zhang, Y., Chen, C., Shi, N., Sun, R., & Luo, Z. Q. (2022). Adam can converge without any modification on update rules. *Advances in neural information processing systems*, *35*, 28386-28399.
soundness: 3
strengths: The prerformance of AdaEMAMix is quite impressive. The writing is excellent. The ablation studies are thorough.
summary: The paper propose a new optimizer called AdEMAMix, which add an additional EMA sequence to Adam. The intuition is to to keep a high sensitivity to recent gradients (using m1), while also incorporating information from older gradients (using m2).
weaknesses: The motivation is rather vague. See below.

---

confidence: 4
contribution: 4
initial_rating: 10
presentation: 4
questions: The authors state “While no answer to those questions is given in this work, we
provide a toy justification which indicates that large momentums can have a positive impact in
noise-free non-convex settings (see Fig. 2)—indicating the improvement of our approach is at least partially explainable without considering variance-reduction effects.” Is there empirical support for this in the LLM experiments? looking at Figure 17 the benefit seems to drop with increasing batch size. Note that the maximum batch size used here (512k) is smaller than that used for LLMs like Llama (4m), though I agree that the model size is also small here. Could the authors provide an experiment with 2m batch size to see the trend?
soundness: 4
strengths: These are strong results on a very important problem. They also provide many optimizer ablations in the Appendix showing the robustness of their proposed optimizer.
summary: This paper proposes AdEMAMix, a new optimizer which outperforms Adam on language model training and ViT training. Their empirical results show large benefits (~50% reduction) over Adam in the regime of noisy gradients i.e. small batch size or longer runs. The main idea behind the optimizer is to maintain two momentum terms and combine them to get the final movement direction. The coefficient of momentum and their combination are also dynamically adapted.
weaknesses: Since many of the experiments are with small batch size it would have been interesting to explore the effect of weight averaging. For example, is it the case that weight averaging helps AdamW and AdEMAMix equally? Or not?

---

confidence: 4
contribution: 3
initial_rating: 5
presentation: 3
questions: Overall, I think the idea introduced in this paper is interesting. The paper can be improved if the above  weaknesses can be addressed.
soundness: 3
strengths: - The paper is well written and easy to follow. The experiments on the Rosenbrock function are convincing. There are some similar (loss landscape) models proposed recently to analyze learning rate schemes [1]. It maybe interesting to draw some theory/experimental connections in the case of momentum.   

- The experiments seem to be comprehensive covering different settings and tasks. The improvement over baseline is shown.

[1] Understanding Warmup-Stable-Decay Learning Rates: A River Valley Loss Landscape Perspective.
summary: This paper proposes to use an additional momentum buffer for the Adam optimizer. The motivation behind is that a single momentum buffer may not be able to utilize both past and current gradient information efficiently. To this end, one buffer with large momentum parameter $\beta$ is for slow-changing directions, and the other one with small $\beta$ is for fast adaptation to current gradient. Then, the two momentum buffers are mixed according to some fixed ratio ($\alpha$), and the rest of updates (i.e., pre-conditioning and weight-decay) follow those of AdamW. However, a large $\beta$ (for slow-changing momentum buffer) can cause training instabilities in the initial stage. To this end, the paper proposes to gradually increase this parameter until the target value is reached.  Then, experiments are performed on various language modelling tasks to show that the proposed algorithm is faster than AdamW given a fixed computation budget.
weaknesses: - There are some additional hyperparameters introduced. Noticeably, it seems that $\alpha$ (that controls the mixing ratio)  is important for the algorithm. It would be important to study the sensitivity of the algorithm to this hyperparameter, given that it essentially controls the contribution of each momentum buffer to the current update. 

- It would be better if some convergence guarantees of the algorithm can be provided even in the convex setting. For example, what is the relationship between the two momentum parameters that would guarantee convergence?

---

confidence: 4
contribution: 3
initial_rating: 6
presentation: 4
questions: 1. Doesn’t the introduction of the additional beta_3 term potentially lead to an exploding step size if the incoming gradients are extremely small for many steps (the additional term in the numerator causes it to decay slower than the denominator in the update)? This isn’t unforeseeable in language model training, where the model could encounter multiple documents of rare tokens consecutively.
2. Do the authors have any experiments adding an EMA to other optimizers like Adafactor or Signum? Does the additional EMA always provide a gain, even with factored gradients?
3. Related to the previous question to address the memory overhead, does storing the slow-moving EMA in lower precision affect AdEMAMix performance?
4. Why was the optimal Adam learning rate being too high for AdEMAMix only manifesting at higher scales? Does this perhaps suggest AdEMAMix being more difficult to tune at larger scales? I’m wondering if there was there further investigation into this eg. Was there a certain part of the network that was destabilizing?
5. Did the authors use certain stabilization strategies, eg. QK-LayerNorm in your experiments? I’m wondering if this could be used for the author’s benefit, to stabilize training and offer greater ease in hyperparameter tuning. 
6. Why was alpha set to 2 and 1 for the switching Adam -> AdEMAMix experiments whereas the optimal values found in hyperparameter tuning were much higher?
soundness: 3
strengths: - The paper is written and organized well. The optimizer is a simple modification to Adam and is conceptually easy to understand, with the algorithmic differences presented clearly.
- The authors extensively benchmark their proposed optimizer against Adam across architectures, dataset domains and hyperparameters.
- The spirit of AdEMAMix (having a ‘fast’ and ‘slow’ tracking of gradients) could be useful for mitigating forgetting in continual learning regimes, and is also shown to be relevant for current pretraining regimes where we often train much longer past eg. Chinchilla-optimal tokens.
- The authors have experiments which address one’s immediate concerns of the proposed optimizer, like hyperparameter tuning and memory overhead.
summary: This paper introduces AdEMAMix, an optimizer modifying Adam by incorporating a mixture of two Exponential Moving Averages (EMAs) to better leverage past gradients. Traditional Adam or AdamW optimizers rely on a single EMA to smooth gradients, which prioritizes recent gradients over older ones. However, the authors argue that this approach is suboptimal because it cannot simultaneously emphasize both recent and very old gradients. They conduct experiments across language (LLMs and Mamba) and vision (ViTs), showing that AdEMAMix outperforms and tends to forget training data slower than Adam.
weaknesses: - The proposed optimizer is not that novel, with several existing approaches in maintaining a longer horizon at the cost of memory which have been mentioned by authors; in particular, AggMo is essentially the same as AdEMAMix but applied to gradient descent, and the only difference in the setting of this work is applying the same principle to Adam.
- Although the authors have performed experiments on the hyperparameter stability of AdEMAMix, they do introduce two new parameters beta_3 and alpha which require the use of schedulers for larger values to avoid divergence in early iterations, which may require more extensive tuning depending on the setting (eg. fast domain shifts).
- There is an additional memory overhead in the order of the model size to incorporate the second EMA, which the authors propose can be mitigated by setting beta_1 to be 0; however, these runs exhibit instabilities and sometimes large spikes in training especially at larger batch sizes (Figure 22). These spikes are generally undesirable even though the loss seems to ‘recover’, it is unclear to me that such spikes wouldn’t have more detrimental effects at larger scales.

---

confidence: 4
contribution: 3
initial_rating: 5
presentation: 2
questions: - In figure 2, authors claim the proposed algorithm does not exhibit oscilations when compared to Adam. However, in both 2 (a) and 2 (c) it appears the AdEMAMix does in fact overshoot the minimizer, further than Adam at that. It looks to me that Adam with “default” hypterparamters appears to be the most stable of all the options. In general this is a little concerning as modern loss surfaces will be much more complex than the Rosenbrock function and these large oscilations/overshooting minimizers may be difficult to deal with or even diagnose. A seperate point on this figure but potentially contributing to my confusion is that the colour scheme in the three subplots is not consistent, if the authors could use a shared legend or something along those lines the results may be more interpretable.

- Given this method requires an additional buffer for the second EMA, this authors discuss it potentially being slower and mention that it is also possible to set  $\beta_1 $ to zero to compensate for this. I’m curious if any of the main experiments or the wall time comparison are using this  $\beta_1= 0 $. The hyperparameter configuration for figure 1 uses  $\beta_1 =0.9 $ so I would just  want ensure the wall time comparison is too to ensure figure 1 doesnt change substantially if the x axis is changed from steps to hours.

- What does older gradients being outdated mean? Is there any way to quantify this? Any intuition? All the data is weighted equally as far as minimizing the loss is concerned despite being at different points in the loss surface when you evaluate it. In general this whole idea of older/newer gradients being less/more relevant seems vaugly okay, but it would be nice to see a more scientific defintion or explanation. This doesn’t need to be totally rigorus and have proof but the current explanations seem hand-wavey and leads to far more questions than answers.

- Why not just use a normal non exponential moving average? You can have a uniformly weighted moving average that should not diverge, has this been tried?

- Is there any intuition for why  $\alpha $ is so large (in particlar  $\ge1 $?) Does this lead to a smaller step size being necessary? Can you also set it to zero until some point to fight early instability?

- Why is there no bias correction on  $m_2 $? I don’t really think it matters much much and most analysis ignores it but I’m curious why that decision was made in comparison to the typical Adam bias correction step.

- Regarding figure 4, while it does appear that the proposed algorithm “forgets” slightly less quickly than AdamW, is the claim that this is why it works better? It appears to me that the biggest predictor of “forgetting” is the injection time  $t_B $ rather than the algorithm, although perhaps if this minor difference is compounded across batches this leads to the superior performance of the algorithm. 

- In many experiments, notably figure 5, the x axis is limited to not show the optimization performance in earlier parts of training. Given that instability is cited as a challenge for this algorithm, to a point where complex schedules need to be used, I’m curious how unstable the well tuned algorithm is. Are there instability issues in the 0-200k range in figure 5 (b) with AdEMAMix from 0 for example?

- Recent analysis in [1] has shown Adam(W) should work well if the gradient and Hessian become correlated, and [2] provided a mechanism where that shows up in language models. Can the authors reconcile this with their modifications to the algorithm? It seems to me those works would suggest this new variant is not a great idea so I’m wondering if the authors have any thoughts for why that isn’t the case.

- Any idea why there is less good performance on imagenet 1k? It seems AdamW is performing better there although all the plots are fairly close.


[1]
Michael Crawshaw, Mingrui Liu, Francesco Orabona, Wei Zhang, Zhenxun Zhuang

Robustness to Unbounded Smoothness of Generalized SignSGD

https://openreview.net/forum?id=8oj_2Ypp0j

[2]
Frederik Kunstner, Robin Yadav, Alan Milligan, Mark Schmidt, Alberto Bietti

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

https://arxiv.org/abs/2402.19449
soundness: 3
strengths: The primary strength of this paper apperas to be an algorithm with strong empiracle performance. The authors conduct a thorough set of experiments and ablation in the setting of language modelling and vision using multiple architectures and configuarations. The idea is interesting and combines other ideas from the optimization literature (the Adam update, multiple momentum terms). I am particularly interested in the authors study of how the loss with respect to different batches changes over time depending on the order in which they where trained on. I think this direction is a promising area of research for understanding why some optimizers work better than others. The paper also contains many interesting questions that are of interest to the optimization community, although struggles to answer them (as does the rest of the field).
summary: This work proposes a new optimization algorithm named AdEMAMix, a variant of Adam(W). They add an aditional moving average of gradients to the algorithm, on top of the existing first and second moment moment averages. In contrast to the traditional  $\beta_1 = 0.9 $ moving average typically used in Adam, the new moving average that updates much slower ( $\beta_3 = 0.9999 $). This additional moving average is then used with the traditional moving average to update the parameters, in a similar fashion to the Adam update. The authors argue this assists in preventing the “forgetting” of data and better uses past information than the traditional Adam update, and show strong empiracle performance on some vision and language settings.
weaknesses: Some specific concerns are addressed in the questions section, but I think the fundamental challenge this paper faces is explaining why the algorithm is better. As is typical with deep learning optimizers, it would be too hard for any provable reason why the algorithm may be superior, but I do think some more effort in the paper could be focused on explaining why this is a good idea, especially with the increased memory requirments. As said in the strengths, I do like the amount of questions being asked but it would be nice to see if the answer to any of those questions is leading to a better optimizer. I think it’s possible that the “forgetting slower” direction may lead to an answer, but the majority of the experiments and argumentation for that is relegated to the appendix. While practically showing plots with better performance is good, from a scientific perspective understanding why that performance is better is important.

I’m slightly concernerd by the number of additional hyperparameters being introduced, which appear to be somewhat sensitive given the use of schedules for most of them. This one can make the algorithm more difficult to use in practice, and two runs the risk of having some lucky configurations for a given problem that works great but in general may not be great. I also feel that the additional memory requirment is a non trivial drawback. While it does not increase communication costs, just running in FSDP is not an easy option for many people who do not have easy access to multi-gpu servers.

A more cosmetic note, the figures in this work can lead to confusion at times. In multiple of the figures, there are three subplots where some and perhaps all are unrelated. In some cases, Some subplots are meant to be related, but use different colours for the same thing. This is also likely the reason many of the figure captions are  $\approx $ half a page. I think it would be valuable for the authors to give the plots another pass before any camera ready version. 

Overall despite being a potentially effective algorithm, I do feel that this paper is somewhat lacking in the motivation and justification of that. I would be open to raising my score in the discussion period but in it’s current state I’m not sure this work is ready. I would be interested in why the authors think this is a good idea and if there are better ways to show that to be the case. This does not need to be large experiments or rigorous theory. Well designed small experiments or simple derivations can be just as effective give an intuition to explain the performance on minimizing the loss or solving some problem existing optimizers have.

More specific questions can be found below.

Meta-review: This paper proposes a new heuristic to improve the momentum-based methods. In particular, it introduces another moving average sequence of stochastic gradient with a larger momentum parameter and adds it to the standard EMA sequence with a large weight. The paper has demonstrated superior performance in various settings. All reviewers agree that the paper has done a great work in demonstrating the effectiveness of the paper. However, a concern is that the paper does not provide any convergence guarantee of the proposed method. Hence, I will recommend a weak acceptance. The authors should take reviewers' comments into account for improving their paper.

Decision: Accept (Poster)