The AdEMAMix Optimizer: Better, Faster, Older
Matteo Pagliardini Pierre Ablin David Grangier

2


Geometrical Nonlinear Hall Effect Induced by Lorentz Force
    Wenhui Duan
    September 9, 2024
==========================================================



§ 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.3B parameter AdEMAMix LLM trained on 101B tokens performs comparably to an AdamW model trained on 197B tokens (+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.  




§ 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) <cit.> or adaptive methods such as Adam and AdamW <cit.>, which have become ubiquitous in training state-of-the-art models <cit.>.

A key component in many of these iterative optimization algorithms is momentum, which has long been shown to accelerate convergence <cit.> and often leads to solutions with superior generalization properties <cit.>. 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 <cit.>. Both SGD with momentum (SGD+M) and Adam incorporate momentum under the form of Exponential Moving Averages (EMAs) of past gradients ^T = {^(0),…,^(T)}: 

    EMAEMA(β, ^T) ≜β·EMA(β, ^(T-1)) + (1-β) ^(T) = ∑_i=0^T β^i (1-β) ^(T-i) .

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 <cit.>. However, those results do not guarantee any optimality for general non-quadratic cases <cit.>. 

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 <cit.>. 
For this reason, practitioners typically employ moderate momentum values (i.e. β≈ 0.8 or 0.9), effectively creating a moving average of recent gradients while discarding older information. Selecting larger β 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 β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 β, 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. <ref>). However, a linear combination between a “fast-changing” (e.g. β=0.9) and a “slow-changing” (e.g. β=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: 
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. <ref>. 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.  




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  <ref>).
    
  * We empirically demonstrate the superiority of our method over Adam by training ViT and language models (Transformers and Mamba) of up to 1.3B parameters (see  <ref>). In addition, we show gains from switching mid-training from Adam to AdEMAMix (see Fig. <ref>).
    
  * We show AdEMAMix forgets the training data at a slower pace when compared to Adam (see Fig. <ref>).
    
  * 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. 







§ RELATED WORK


Works on understanding momentum. From the seminal work of <cit.>, many publications analyzed the effect of gradient descent + momentum (GD+M) in both convex and non-convex settings <cit.>. 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 <cit.>, 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 <cit.>, 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 <cit.>. 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 <cit.>? 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. <ref>)—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. <ref>), which to our knowledge is novel.  

Works on deep-learning optimizers. Despite the popularity of Adam and AdamW <cit.> 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 <cit.> improves Adam's memory efficiency by factorizing the second moment estimate. Lamb <cit.> extends Adam by adding layerwise normalized updates. <cit.> 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, <cit.> also report better scores are obtained when using a slightly larger momentum term (β=0.99). In this work we show how increasing the momentum well beyond this value can still be beneficial. See App. <ref> for a detailed comparison between AdEMAMix and Lion. Recently, <cit.> 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 <cit.>, and Adam's dominance remains mostly unchallenged.      


Works incorporating an additional momentum term. <cit.> 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, <cit.> refer to the Double EMA (DEMA) used in some finance applications <cit.> 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. <ref>. 


Works on distributed optimization. Perhaps surprisingly, recent work on distributed optimization—DiLoCo and SlowMo <cit.>—are relevant to our work. N workers ^(t)_1,…,^(t)_N are trained independently for K steps (e.g. K=500). Every K steps, the delta updates {Δ_i}_i=1^N ≜{_i^(t+K)-_i^(t)}_i=1^N are averaged and applied to each worker using an outer optimizer with momentum β: _i^(t+K) = _i^(t) - η·Opt(1/N∑_i Δ_i, β). 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.    



§ OUR METHOD: ADEMAMIX





Setup & notations. Let _: ↦ be a loss function parameterized by , and mapping inputs ∈ to . Given a sampled batch , let = ∇__() be a stochastic gradient of the loss w.r.t. . To minimize the empirical loss, the Adam optimizer <cit.> relies on first and second moments, resp.  and , estimated via two EMAs parametrized by (β_1,β_2) ∈ [0,1[^2. A weight-decay parameter λ∈^+ is often used as in <cit.>:

    ^(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)) .
      AdamW

With t>0 being the timestep, η being the learning rate, and  a small constant. Initially ^(t=0)=^(t=0)=0. To prevent the bias induced by the initial ^(t=0) and ^(t=0), the outputs of the two EMAs are corrected into ^(t) and ^(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 β_1 is 0.9. Fig. <ref> shows the weights given to past gradients for different values of β. The larger the β, the more uniform the average is. To put this in perspective—observing that ∑_i=0^∞β^i(1-β)=1 for β∈ [0,1[—the number of successive previous steps receiving a cumulative weight of 0.5, is t_half=ln(0.5)/ln(β)-1. For β=0.9, t_half≈ 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 <cit.>, which typically relies on similar β values. The value of β_1 is rarely increased beyond ∼ 0.9.

In our experiments with AdamW, increasing β_1 further degraded the performance (see App. <ref>). 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 β values e.g. 0.9999. To compare, for β=0.9999, t_half≈ 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):

    _1^(t) = β_1 _1^(t-1) + (1-β_1) ^(t) ,     _1^(t) = _1^(t)/1-β_1^t
    _2^(t)= β_3 _2^(t-1) + (1-β_3) ^(t)
    ^(t) = β_2 ^(t-1) + (1-β_2) ^(t)^2 ,     ^(t) = ^(t)/1-β_2^t
    ^(t) = ^(t-1) - η(_1^(t)+ α_2^(t)/√(^(t))+ + λ^(t-1)) .
      AdEMAMix

In our experiments, while the values of β_1, β_2 remain similar to those of (<ref>), we often use β_3=0.9999. We find α∈[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. <cit.> 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 β_3 would diverge early. This, despite not using bias correction over _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. <ref>, Fig. <ref>).  

For this reason, we progressively increase the values of β_3 and α using schedulers. For α we use a linear scheduler. A linear scheduler for β_3 would be ill-fitted as the same increment of β_3 have a different impact for different values of β_3. For instance, observe that an increase of β of δ_β=0.0001 barely increases the t_half for β=0.9, while 0.999→ 0.999+δ_β increases the t_half of 77. For this reason, we design the β_3 scheduler to increase t_half linearly (see App. <ref> for a derivation of that scheduler). The two schedulers are summarized below:

    α^(t)   = f_α(t,α,T_α) = min(tα/T_α, α), f_α
    β_3^(t)   = f_β_3(t,β_3,β_start,T_β_3) = min(exp(ln(β_start)ln(β_3)/(1-t/T_β_3)ln(β_3)+ t/T_β_3ln(β_start)), β_3) . f_β_3

With T_α and T_β_3 are resp. the warmup times for α^(t) and β^(t)_3 to reach their final and maximal values. In practice we always set those two to the same value: T_α=T_β_3=T_α,β_3, and we typically use T_α,β_3=T, with T being the total number of iterations. β_start is always set to β_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. <ref>). The full AdEMAMix optimizer, including the schedulers, is shown in Alg. <ref>.






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. <ref>. 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, as AdEMAMix requires to allocate _2, which is of the same size as the model parameters . We believe this issue is of lesser consequences as 

Fully-Sharded-Data-Paralellism <cit.> can always be used to distribute the optimizer states across compute nodes. In the cases where memory remains an issue, we see two mitigation strategies. The first one consists in setting β_1=0, which removes entirely the need for _1. We show in App. <ref> (Fig. <ref>) that this strategy might work, but at the cost of less stable training. The second mitigation strategy is to apply factorization strategies as in <cit.>.   


Hyperparameter sensitivity. While we introduce up to four new hyperparameters:  α, β_3, T_α, and T_β_3. In practice we always set T_α=T_β_3=T_α,β_3, and use T_α,β_3=T in most cases. We show in App. <ref> that T_α and T_β_3 should only be large enough to prevent instabilities early during training. While all of our experiments on language modeling use β_3=0.9999, other values such as 0.999 or even 0.99999 still can outperform the AdamW baseline (see App. <ref>, Fig. <ref>). On vision tasks, the scatter plots in Fig. <ref> show all the AdEMAMix models trained for those experiments, highlighting how easy it can be to find good α and β_3 values. Overall, we find the ranges of values of α,β_3 and T_α,β_3 providing improvements over AdamW to be wide. See App. <ref> 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. <ref>, additionally showing how smaller values of β_3 (e.g. β_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, or general cases in which the distribution is non-stationary.




§ RESULTS


In this section we use AdEMAMix to train language models ( <ref> and  <ref>) and vision transformers ( <ref>) ranging from 24M to 1.3B parameters. 





 §.§ Transformer LLM Training


Experimental setup. We use a transformer architecture <cit.> with learnt positional encoding. All of our experiments are using sequences of 1,024 tokens. We experiment with three sizes of transformers: 110M, 335M, and 1.3B. For the learning rate, we use 3k 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. <ref>). We use the RedPajama v2 <cit.> dataset for all of our experiments. We use batch sizes of 64, 96 and 128 for respectively our 110M, 335M, and 1.3B parameter models. Depending on the model, we vary the number of iterations from 256k to 1.5M. For AdEMAMix, we use β_3=0.9999 and α∈{5,8,10} depending on the model. A full description of the architecture and hyperparameters used is in App. <ref>. We train on up to 8 A100 NVIDIA GPUs using data-parallelism. 



Why not simply increasing AdamW's β_1? While our toy experiment in Fig. <ref> already gives some intuition on why increasing Adam's β_1 is likely not to have the same effect as having an additional EMA as in AdEMAMix, we verify this intuition by training multiple 110M models using Adam with large β_1∈{0.99,0.999,0.9999,0.99999}. When we use a large β_1 from the beginning of training, we observe instabilities for larger β_1 values and no β_1>0.9 improves upon the AdamW baseline. One can imagine this to be due to increasing β_1 too early. Therefore, we also modify AdamW and add the same scheduler on β_1 as we use on AdEMAMix's β_3. β_1 is now increased steadily over the entire training duration. While this mostly stabilizes the training, none of the experiments outperformed the baseline using β_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 β_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. <ref>. Those experiments show that simply increasing the β_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 256k to 1M—, AdEMAMix always outperforms the AdamW baseline. In Fig. <ref>, we show the validation loss curves for AdamW and AdEMAMix models trained—for each size—on various numbers of tokens. For 110M 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.3B parameter models, training using 770k steps is on par with training the baseline for 1.5M iterations—reaching the same performance with 51% fewer tokens (an economy of 96M tokens). In Fig. <ref>, 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. <ref> we plot the time required to train our 110M and 1.3B 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 110M 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  <ref>, 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. <ref>, 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. <ref>, 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 _2^(t_switch)=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. <ref> for more details). In Fig <ref> and Fig.<ref>, 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. <ref>). 

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




 §.§ Mamba LM Training


Experimental setup & results. Our experimental setup is similar to  <ref>, except that we now train 168M parameter Mamba models <cit.> using the FineWeb dataset <cit.>. See App. <ref> for more details.
In Fig. <ref> the improvement using AdEMAMix is consistent with our experiments on Transformer models. This shows how AdEMAMix's gains can extend beyond Transformer models. 



 §.§ ViT Training




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 <cit.> dataset: (i) the widely used ImageNet-1k subset, consisting of 1.3M images and 1,000 possible classes; (ii) a filtered and pre-processed version of the ImageNet-21k <cit.> containing 11M images corresponding to 10,450 classes. For each, we measure the test loss on a held-out test set. We use the ViT architecture <cit.> at two different scales: 24M and 86M 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. <ref>, we use a batch size of 4096 for all our experiments. We used training hyperparameters from <cit.> 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 α∈{1,5,10,15,20} and β_3∈{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 <cit.>. We use simple data-augmentations, which includes mixup <cit.>. We train on 8 A100 NVIDIA GPUs using Pytorch Fully Sharded Data-Parallelism <cit.>.   

AdEMAMix for different capacity/data ratios. We consider three scenarios that differ by their capacity/data ratios. First, we trained 24M parameter models on 11M images (ImageNet-21k), for 37 epochs. In this setting, as can be seen in Fig. <ref>, 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 86M parameters. Fig. <ref> shows it is still easy to find parameters outperforming the baseline. We now keep the model size of 86M parameters and switch to the smaller ImageNet-1k dataset (1.3M images), which—given our 100k iterations—increases the number of epochs to 320. In this setting, Fig. <ref> 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.   
 




§ 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 β) momentum gathers information over many timestep, while a fast (low β) 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.



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


iclr2025_conference









§ IMPLEMENTATION DETAILS




 §.§ Deriving the β_3 scheduler


Let's consider S(t), the sum of the weights given to the last t gradients by an EMA parameterized by β∈ [0,1[:

    S(t) = (1-β) ∑_i=0^t β^i

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

    (1-β) ∑_i=0^t β^i = 0.5 ⇔β^t+1 = 0.5 ⇔ t = ln(0.5)/ln(β)-1

Let f(β) = ln(0.5)/ln(β)-1. This function provides how many past consecutive gradients receive a cumulative weight of 0.5. 

Its inverse is:

    f^-1(t) = 0.5^1/t+1


We want a scheduler which increases β from β_start to β_end such that f(β) increases linearly. Given an interpolating parameter μ∈[0,1], this scheduler can be written as:

    β(μ) = f^-1((1-μ) f(β_start) + μ f(β_end))

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

    β(μ) = exp(ln(β_start)ln(β_end)/(1-μ)ln(β_end)+ μln(β_start))

By setting β_end=β_3 and μ=t/T_β_3, we arrive to the β_3-scheduler introduced in  <ref>. We show the shape of our scheduler and compare it to a linear scheduler in Fig. <ref>.








 §.§ Pytorch implementation


The following is a code skeleton for our AdEMAMix optimizer in Pytorch <cit.>. 

[language=Python,breaklines=true,showstringspaces=false,caption=AdEMAMix code skeleton using Pytorch]
import math
import torch
from torch.optim import Optimizer


class AdEMAMix(Optimizer):

    def __init__(self, 
                 params, 
                 lr=1e-3, 
                 betas=(0.9,0.999,0.9999), 
                 alpha=5.0,
                 T_beta3=0, 
                 T_alpha=0, 
                 eps=1e-8, 
                 weight_decay=0.0):
        # init the optimizer 
    
    @torch.no_grad()
    def step(self):

        for group in self.param_groups:
        
            lr = group["lr"]
            lmbda = group["weight_decay"]
            eps = group["eps"]
            beta1, beta2, beta3_final = group["betas"]
            T_beta3 = group["T_beta3"]
            T_alpha = group["T_alpha"]
            alpha_final = group["alpha"]

            for p in group["params"]:

                grad = p.grad
                state = self.state[p]

                # State initialization
                if len(state) == 0:
                    state["step"] = 0 # step counter used for bias correction
                    state["m1"] = torch.zeros_like(p) # fast EMA
                    state["m2"] = torch.zeros_like(p) # slow EMA
                    state["nu"] = torch.zeros_like(p) # second moment estimate

                m1, m2, nu = state["m1"], state["m2"], state["nu"]

                # Bias correction: no correction for beta3's EMA
                state["step"] += 1
                bias_correction1 = 1 - beta1 ** state["step"]
                bias_correction2 = 1 - beta2 ** state["step"]

                # Calling the schedulers for alpha and beta3
                alpha = alpha_scheduler(state["step"], start=0, end=alpha_final, T=T_alpha)
                beta3 = beta3_scheduler(state["step"], start=beta1, end=beta3_final, T=T_beta3)

                # Update the EMAs
                m1.mul_(beta1).add_(grad, alpha=1 - beta1)
                m2.mul_(beta3).add_(grad, alpha=1 - beta3)
                nu.mul_(beta2).addcmul_(grad, grad, value=1 - beta2)

                # Compute step
                denom = (nu.sqrt() / math.sqrt(bias_correction2)).add_(eps)
                update = (m1.div(bias_correction1) + alpha * m2) / denom
                
                # Add weight decay
                update.add_(p, alpha=lmbda) 

                # Apply the update scaled by the learning rate 
                p.add_(-lr * update)

        return loss






 §.§ Optax implementation


The following is a code skeleton for our AdEMAMix optimizer in Optax, an optimization library based on JAX <cit.>. 

[language=Python,breaklines=true,showstringspaces=false,caption=AdEMAMix code skeleton using JAX and Optax]
from typing import NamedTuple
import chex
from optax._src import transform, combine, base, numerics
from jax import tree_util as jtu
import jax.numpy as jnp


class ScaleByAdemamixState(NamedTuple):
  count: chex.Array  
  count_m2: chex.Array  
  m1: base.Updates
  m2: base.Updates
  nu: base.Updates


def ademamix(lr, 
             b1=0.9, 
             b2=0.999, 
             b3=0.9999, 
             alpha=5.0, 
             b3_scheduler=None, 
             alpha_scheduler=None,
             eps=1e-8, 
             weight_decay=0.0):
  return combine.chain(
    scale_by_ademamix(b1, b2, b3, alpha, b3_scheduler, alpha_scheduler, eps),
    transform.add_decayed_weights(weight_decay),
    transform.scale_by_learning_rate(lr),
  )


def scale_by_ademamix(b1, 
                      b2, 
                      b3, 
                      alpha, 
                      b3_scheduler, 
                      alpha_scheduler, 
                      eps):

  def init_fn(params):
    m1 = tree_zeros_like(params)  # fast EMA
    m2 = tree_zeros_like(params)  # slow EMA
    nu = tree_zeros_like(params)  # second moment estimate
    return ScaleByAdemamixState(
        count=jnp.zeros([], jnp.int32), 
        count_mu2=jnp.zeros([], jnp.int32), 
        m1=m1, 
        m2=m2, 
        nu=nu
    )

  def update_fn(updates, state, params=None):
    del params
    c_b3 = b3_scheduler(state.count_m2) if b3_scheduler is not None else b3
    c_alpha = alpha_scheduler(state.count_m2) if alpha_scheduler is not None else alpha
    m1 = tree_update_moment(updates, state.m1, b1, 1) # m1 = b1 * m1 + (1-b1) * updates
    m2 = tree_update_moment(updates, state.m2, c_b3, 1)
    nu = tree_update_moment_per_elem_norm(updates, state.nu, b2, 2)
    count_inc = numerics.safe_int32_increment(state.count)
    count_m2_inc = numerics.safe_int32_increment(state.count_m2)
    m1_hat = tree_bias_correction(m1, b1, count_inc)
    nu_hat = tree_bias_correction(nu, b2, count_inc)
    updates = jtu.tree_map(
        lambda m1_, m2_, v_: (m1_+c_alpha*m2_)/(jnp.sqrt(v_)+eps), 
        m1_hat, 
        m2, 
        nu_hat
    )
    return updates, ScaleByAdemamixState(
        count=count_inc, 
        count_m2=count_m2_inc, 
        m1=m1, 
        m2=m2, 
        nu=nu
    )

  return base.GradientTransformation(init_fn, update_fn)





§ EXPERIMENTAL DETAILS




 §.§ Transformer LLM experiments


Architecture details. Our architecture is based on the transformer decoder of <cit.>. We use learnt positional encoding. We use a SentencePiece <cit.> tokenizer with a vocabulary of 32000 tokens. The model specific parameters for the different sizes are summarized in Table <ref>. Our implementation is using Jax <cit.>, and we train using , except for normalization modules and softmax which use . The optimizer states and operations are in .



How did we tune the hyperparameters? Starting from our smallest models (110M parameters), we first tuned hyperparameters for our AdamW models. We then use the best hyperparameters found as a starting point for AdEMAMix and tuned β_3 and α. When we use schedulers for α and β_3, we set T_α=T_β_3=T, with T being the total number of iterations. Table <ref> summarizes the hyperparameters we tried for this model size. The impact of AdEMAMix's hyperparameters is discussed extensively in App. <ref>. For our 330M parameter models, we mostly kept the best hyperparameters found for the 110M model and re-tuned the learning rate and gradient clipping. For AdEMAMix, we additionally tested multiple β_3 and α values. We summarize this process in Table <ref>. Finally, for our 1.3B 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 <ref>.     










Hyperparameters for experiments switching from AdamW and AdEMAMix. For experiments in Fig. <ref> and Fig. <ref>, when we switch from AdamW to AdEMAMix during training, for our 110M parameter models (Fig. <ref>), we use α=2, β_3=0.9999 and T_α,β_3=0. For our 1.3B parameter models (Fig. <ref>), we use α=1, β_3=0.9999 and T_α,β_3=0. Other hyperparameters are inherited from the AdamW model we are switching from.

Hyperparameters used our constant learning rate scheduler experiments. For Fig. <ref> we use η=10^-4, the remaining of the hyperparameters are identical as those used for our other 1.3B experiments and provided in Table <ref>.

Hyperparameters used in our forgetting experiments. For the experiments in Fig. <ref>, we used 110M parameter models and hyperparameters from Table <ref>. 



 §.§ Mamba LM experiments


Architecture details. We use a Mamba architecture <cit.> 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  tokenizer from <cit.>, using the HuggingFace library <cit.>. 

Hyperparameter details. We used parameters mostly taken from <cit.>:

    
  * Learning rate: η=0.0006,
    
  * Warmup steps: 3k,
    
  * η-scheduler type: cosine decay to 10^-5,
    
  * Weight-decay: λ=0.1,
    
  * Gradient clipping value: 1,
    
  * Total training steps T∈{64k,128k}.

For experiments with AdamW, we use β_1=0.9 and β_2=0.999. For AdEMAMix, we use β_1=0.9, β_2=0.999, β_3=0.9999, T_α,β_3=T and α=8.



 §.§ ViT experiments


Architecture details. We use a ViT architecture following <cit.>. The architecture details for both our 24M and 86M parameter models can be seen in Table <ref>. We do not use an EMA of the iterates as our final model. Our implementation is in Pytorch <cit.>, and use . 



Data augmentation and Mixup. While more sophisticated augmentation methods exist <cit.>, 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 <ref>. 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 α and β_3. For each setting, we train 8 models using α∈{1,5,10,15,20} and β_3∈{0.999,0.9999}. All of the AdEMAMix models are shown in Fig. <ref> and Fig. <ref>. Only the best AdamW model is shown in Fig. <ref>.    






§ ADDITIONAL RESULTS




 §.§ LLM experiments





  §.§.§ In-Context Learning (ICL) results

In-Context Learning (ICL) results. We evaluate our largest (1.3B) LLM models on in-context learning tasks. We use the  package <cit.>. We consider the following tasks: 

HellaSwag <cit.>,
Winogrande <cit.>,
ARC <cit.>,
BoolQ <cit.>,
LogiQA <cit.>,
MathQA <cit.>,
MMLU <cit.>,
OpenbookQA <cit.>,
PIQA <cit.>,
PubmedQA <cit.>,
RewardBench <cit.>,
Sciq <cit.>,
TruthfulQA <cit.>.


The two AdEMAMis and AdamW models we benchmark have been trained for 1M steps, with a batch size of 128, corresponding to around 131B tokens. The results of the evaluation are in Table <ref>. 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 <ref>, 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. <ref>. 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.









  §.§.§ More results on forgetting

Evolution of forgetting during training. In Fig. <ref> 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. <ref> and Fig. <ref>, 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. 




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. <ref> with results obtained using a different scheduler with a constant learning rate and a linear decay. We train 110M 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. <ref>. 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.







  §.§.§ Removing the second EMA mid-training

Removing the second EMA mid-training. In Fig. <ref> and Fig. <ref>, 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 _2 parameter of AdEMAMix during training, effectively switching back to AdamW. Results for this experiment can be seen in Fig. <ref>. 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.





  §.§.§ Hyperparameter sensitivity

Hyperparameter sensitivity. Depending on whether the α and β_3 schedulers are used, AdEMAMix adds up to 4 new hyperparameters: α, β_3, T_α and T_β_3. In all our experiments we tied T_α=T_β_3=T_α,β_3. In this section we analyze the sensitivity of AdEMAMix to those hyperparameters. We study the impact of α and β_3 in Fig. <ref>, revealing wide ranges of values for which hyperparameters are outperforming the AdamW baseline. We study the sensitivity of T_α,β_3 as well as justify the choice of using a scheduler on α in Fig. <ref>. When training from scratch T_α,β_3 needs simply to be large enough to avoid early divergence. In Fig. <ref> we study the sensitivity to the gradient clipping and AdEMAMix's β_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 β_1, we observe some loss spikes, yet the final loss value is relatively unaffected.       













  §.§.§ 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 110M 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 (17B). In Fig. <ref>, 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. <ref> we observe that reducing β_3 can mostly alleviate the problem. 








  §.§.§ 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^⋆,y^⋆)  ∈_(x,y) f(x,y),    with  f(x,y) = 8(x-1)^2 × (1.3 x^2+2x+1) + 0.5(y-4)^2

This function was introduced by <cit.> as a function with sharp curvature along the x-axis, and flatter curvature along the y-axis.
Initializing the first iterate ^(0)=(0.3, 1.5), we start optimizing with (i) Adam with a large β_1=0.999, and (ii) AdEMAMix with β_1=0.9 and β_3=0.999. In both cases β_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 β_1 is unable to correct his trajectory. In contrast, using two EMAs, one with a small, and one with a large β—as in AdEMAMix—converges to the solution. The fast changing EMA can correct for the bias introduced by the slow changing EMA. See Fig. <ref> for results.




Results training 110M parameters LMs. To further demonstrate that simply increasing β_1 in AdamW does not provide nearly the same gains as AdEMAMix, we run several additional experiments. In Fig. <ref> we show what we obtain when training from scratch using a single EMA with a β_1∈{0.99, 0.999, 0.9999, 0.99999}. Naively increasing β_1 in AdamW does not work; adding a scheduler on β_1 to smoothly increase its value during the entirety of the training also fails. In Fig. <ref>, we show results when increasing β_1 in the middle of training, with and without scheduler on β_1. This differs from the previous setting as we bypass the initial training phase capable of causing instabilities (see  <ref>), 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 β_1 value does not provide any noticeable gain.    










  §.§.§ Miscellaneous results


Impact of β_3 and α schedulers when starting AdEMAMix from AdamW. As mentioned in  <ref>, we do not necessarily need the α and β_3 schedulers when resuming training from a sufficiently late checkpoint. Fig. <ref> and Fig. <ref> do not use schedulers, unless we start using AdEMAMix from scratch. In Fig. <ref>, we vary the warmup periods for α and β_3 when starting AdEMAMix from an AdamW checkpoint at t_switch=300k. We observe how increasing T_β_3=T_α only slows down the convergence. This serves to illustrate that the schedulers for α and β_3 are only required to stabilise the optimization during the early training phase.  



AdEMAMix from scratch without schedulers. In this section we provide more justification for the use of schedulers on α and β_3. We reveal that, without the use of schedulers, the norms of the updates increases significantly, even for small α values. Unstable and large weight updates are known to occur in the early phases of training when learning rate warmup is not used <cit.>. Using large momentum values too early seem to increase this phenomenon. See Fig. <ref>.    



Impact of the linear decay duration when using a constant η-scheduler. In Fig. <ref> we show results using AdEMAMix with a linear warmup → constant  → 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.3B parameter model with a max learning rate of 10^-4 and remaining hyperparameters as in Table <ref>. Results can be seen in Fig. <ref>. 



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





 §.§ ViT experiments


Top-k accuracies. In Fig. <ref>, we plot the test and train loss for the final iterates. In Fig. <ref>, 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. 








 §.§ Comparison with other methods




  §.§.§ Comparison with AdMeta and DEMA


Double Exponential Moving Average (DEMA). Originally introduced by <cit.>, 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 EMA_β^(T-N:T)≜EMA(β, ^(T-N), …, ^(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:

    DEMA(β, ^(T-2N), …, ^(T)) = 2 ·EMA_β^(T-N:T) - EMA(β, EMA_β^(T-2N:T-N), …, EMA_β^(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. <cit.> take inspiration over DEMA and use nested EMAs in their AdMeta-S optimizer:

    _1^(t)   = β_1 _1^(t-1)+^(t)
    ^(t)   = κ^(t) + μ_1^(t)
    _2^(t)   = β_2 _2^(t-1)+(1-β_2)^(t)
    ^(t)   = ^(t-1)-η_2^(t) .
        AdMeta-S

With μ and κ parameterized by β_1 ∈ [0,1[ as such:

    μ   = 25 - 10 (β_1 + 1/β_1) 
    κ   = 10/β_1 - 9 .

In their AdMeta-S experiments, they use β_1=0.9, corresponding to (μ,κ)=(4.88,2.11). β_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. <ref>. 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. 





  §.§.§ Comparison with Lion


The Lion optimizer. <cit.> derived the following optimizer. We change notations to facilitate the comparison with AdEMAMix:

    ^(t)   = ^(t-1) - η·(sign(α^(t-1)+(1-α)^(t)) +λ^(t-1)) 
    ^(t)   = β^(t-1) + (1-β) ^(t) .
        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 α^(t-1)+(1-α)^(t) is similar to combining two EMA, one of them using β=0. While we mention the possibility of setting AdEMAMix's β_1 to 0 in  <ref>, we show in App. <ref> (Fig. <ref>) that this can cause instabilities. Interestingly, <cit.> find that larger β=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 <cit.> claim can help regularize the training. We test the Lion optimizer on language modeling. To tune the hyperparameters, we took values from <cit.> as a starting point, as well as recipes provided by the Optax Jax library. A summary of our hyperparameter tuning is in Table <ref>:



The training curve associated to the best hyperparameters is in Fig. <ref>. We observe that Lion is not outperforming our carefully tuned AdamW baseline. Moreover, no attempt to increase β beyond 0.99 was successful, as those models mostly diverged.



