The AdEMAMix Optimizer: Better, Faster, Older — part 1/5

Cleaned-text code points [0, 14717). Bibliography omitted here; source.md includes it.

# The AdEMAMix Optimizer: Better, Faster, Older

###### 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 B parameter AdEMAMix LLM trained on 101 101 B tokens performs comparably to an AdamW model trained on 197 197 B tokens ( + 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, 1951 ) or adaptive methods such as Adam and AdamW ( Kingma & Ba, 2015 ; Loshchilov & Hutter, 2019 ) , which have become ubiquitous in training state-of-the-art models ( Devlin et al., 2019 ; Brown et al., 2020 ; Dosovitskiy et al., 2021a ; Radford et al., 2021 ; Zhai et al., 2022 ; Dehghani et al., 2023 ; Touvron et al., 2023 ; Dubey et al., 2024 ) .

A key component in many of these iterative optimization algorithms is momentum, which has long been shown to accelerate convergence ( Nemirovskii & Nesterov, 1985 ) and often leads to solutions with superior generalization properties ( Sutskever et al., 2013 ) . 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, 1999 ; Ruder, 2016 ; Goh, 2017 ) . Both SGD with momentum (SGD+M) and Adam incorporate momentum under the form of Exponential Moving Averages (EMAs) of past gradients 𝒢 T = { 𝒈 ( 0 ) , … , 𝒈 ( T ) } {\mathcal{G}}^{T}=\{{\bm{g}}^{(0)},\ldots,{\bm{g}}^{(T)}\} : EMA ​ ( β , 𝒢 T ) ≜ β ⋅ EMA ​ ( β , 𝒢 ( T − 1 ) ) + ( 1 − β ) ​ 𝒈 ( T ) = ∑ i = 0 T β i ​ ( 1 − β ) ​ 𝒈 ( T − i ) . \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, 1964 ; Nesterov, 1983 ) . However, those results do not guarantee any optimality for general non-quadratic cases ( Goujaud et al., 2023 ) .

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., 2013 ) . For this reason, practitioners typically employ moderate momentum values (i.e. β ≈ 0.8 \beta\approx 0.8 or 0.9 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. β = 0.9 \beta=0.9 or β 1 = 0 \beta_{1}=0 ) and a “slow-changing” (e.g. β = 0.9999 \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: 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 ( Ad aptive EMA Mix ture), a novel Adam based optimizer which successfully leverages very old gradients to reach better solutions.

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 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.

## 2 Related Work

Works on understanding momentum. From the seminal work of Polyak (1964) , many publications analyzed the effect of gradient descent + momentum (GD+M) in both convex and non-convex settings ( Ghadimi et al., 2015 ; Flammarion & Bach, 2015 ; Kidambi et al., 2018 ; Defazio, 2020 ; Sebbouh et al., 2021 ) . 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., 2016 ; Kidambi et al., 2018 ; Leclerc & Madry, 2020 ) , 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., 2023 ; Papazov et al., 2024 ) , 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., 2016 ) . 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., 2023 ) ? 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, 2015 ; Loshchilov & Hutter, 2019 ) 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, 2018 ) improves Adam’s memory efficiency by factorizing the second moment estimate. Lamb ( You et al., 2020 ) extends Adam by adding layerwise normalized updates. Chen et al. (2023a) 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. (2023a) also report better scores are obtained when using a slightly larger momentum term ( β = 0.99 \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. (2023) 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., 2023 ) , and Adam’s dominance remains mostly unchallenged.

Work incorporating an additional momentum term. Lee et al. (2024) 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. (2023b) refer to the Double EMA (DEMA) used in some finance applications ( Mulloy, 1994 ) 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. (2019, AggMo) also observe that using a combination of EMAs can enable the use of larger β \beta s, and incorporates a sum of K 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. (2024) 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., 2023 ; Wang et al., 2020 ) —are relevant to our work. N N workers 𝜽 1 ( t ) , … , 𝜽 N ( t ) {\bm{\theta}}^{(t)}_{1},\ldots,{\bm{\theta}}^{(t)}_{N} are trained independently for K K steps (e.g. K = 500 K=500 ). Every K K steps, the delta updates { Δ ​ 𝜽 i } i = 1 N ≜ { 𝜽 i ( t + K ) − 𝜽 i ( t ) } i = 1 N \{\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 : 𝜽 i ( t + K ) = 𝜽 i ( t ) − η ⋅ Opt ​ ( 1 N ​ ∑ i Δ ​ 𝜽 i , β ) {\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 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

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, 2015 ) relies on first and second moments, resp. 𝒎 {\bm{m}} and 𝝂 {\bm{\nu}} , estimated via two EMAs parametrized by ( β 1 , β 2 ) ∈ [ 0 , 1 [ 2 (\beta_{1},\beta_{2})\in[0,1[^{2} . A weight-decay parameter λ ∈ ℝ + \lambda\in\mathbb{R}^{+} is often used as in Loshchilov & Hutter (2019) : { 𝒎 ( t ) = β 1 ​ 𝒎 ( t − 1 ) + ( 1 − β 1 ) ​ 𝒈 ( t ) , 𝒎 ^ ( t ) = 𝒎 ( t ) 1 − β 1 t 𝝂 ( t ) = β 2 ​ 𝝂 ( t − 1 ) + ( 1 − β 2 ) ​ 𝒈 ( t ) 2 , 𝝂 ^ ( t ) = 𝝂 ( t ) 1 − β 2 t 𝜽 ( t ) = 𝜽 ( t − 1 ) − η ⁡ ( 𝒎 ^ ( t ) 𝝂 ^ ( t ) + ϵ + λ ​ 𝜽 ( t − 1 ) ) . \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 t>0 being the timestep, η \eta being the learning rate, and ϵ {\epsilon} a small constant. Initially 𝒎 ( t = 0 ) = 𝝂 ( t = 0 ) = 𝟎 {\bm{m}}^{(t=0)}={\bm{\nu}}^{(t=0)}=\mathbf{0} . To prevent the bias induced by the initial 𝒎 ( t = 0 ) {\bm{m}}^{(t=0)} and 𝝂 ( t = 0 ) {\bm{\nu}}^{(t=0)} , the outputs of the two EMAs are corrected into 𝒎 ^ ( t ) \hat{{\bm{m}}}^{(t)} and 𝝂 ^ ( t ) \hat{{\bm{\nu}}}^{(t)} . Those are used to compute the final weight update, scaled by the learning rate.