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

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

How far to look into the past? A typical value for β 1 \beta_{1} is 0.9 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 ∑ i = 0 ∞ β i ​ ( 1 − β ) = 1 \sum_{i=0}^{\infty}\beta^{i}(1-\beta)=1 for β ∈ [ 0 , 1 [ \beta\in[0,1[ —the number of successive previous steps receiving a cumulative weight of 0.5 0.5 , is t h ​ a ​ l ​ f = ln ⁡ ( 0.5 ) ln ⁡ ( β ) − 1 t_{half}=\frac{\ln(0.5)}{\ln(\beta)}-1 . For β = 0.9 \beta=0.9 , t h ​ a ​ l ​ f ≈ 6 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, 1964 ; Nesterov, 1983 ) , which typically relies on similar β \beta values. The value of β 1 \beta_{1} is rarely increased beyond ∼ 0.9 \sim 0.9 . In our experiments with AdamW, increasing β 1 \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 0.9999 . To compare, for β = 0.9999 \beta=0.9999 , t h ​ a ​ l ​ f ≈ 6,930 t_{half}\approx 6{,}930 , spreading half of the mass over the previous 6,930 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 ) ) . \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}{\bm{m}}_{2}^{(t)}}{\color[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}+\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 β 1 , β 2 \beta_{1},\beta_{2} remain similar to those of equation AdamW , we often use β 3 = 0.9999 \beta_{3}=0.9999 . We find α ∈ [ 4 , 10 ] \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. (2022) 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 \beta_{3} would diverge early. This, despite not using bias correction over 𝒎 2 {\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 β 3 \beta_{3} and α \alpha using schedulers. For α \alpha we use a linear scheduler. A linear scheduler for β 3 \beta_{3} would be ill-fitted as the same increment of β 3 \beta_{3} has a different impact for different values of β 3 \beta_{3} . For instance, observe that an increase of β \beta of δ β = 0.0001 \delta_{\beta}=0.0001 barely increases the t h ​ a ​ l ​ f t_{half} for β = 0.9 \beta=0.9 , while 0.999 → 0.999 + δ β 0.999\rightarrow 0.999+\delta_{\beta} increases the t h ​ a ​ l ​ f t_{half} of 77 77 . For this reason, we design the β 3 \beta_{3} scheduler to increase t h ​ a ​ l ​ f t_{half} linearly (see App. A.1 for a derivation of that scheduler). The two schedulers are summarized below: α ( t ) \displaystyle\alpha^{(t)} = f α ​ ( t , α , T α ) = min ⁡ ( t ​ α T α , α ) , \displaystyle=f_{\alpha}(t,\alpha,T_{\alpha})=\min(\frac{t\alpha}{T_{\alpha}},\alpha), ( f α f_{\alpha} ) β 3 ( t ) \displaystyle\beta_{3}^{(t)} = f β 3 ​ ( t , β 3 , β s ​ t ​ a ​ r ​ t , T β 3 ) = min ⁡ ( exp ⁡ ( ln ⁡ ( β s ​ t ​ a ​ r ​ t ) ​ ln ⁡ ( β 3 ) ( 1 − t T β 3 ) ​ ln ⁡ ( β 3 ) + t T β 3 ​ ln ⁡ ( β s ​ t ​ a ​ r ​ t ) ) , β 3 ) . \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 β 3 f_{\beta_{3}} ) With T α T_{\alpha} and T β 3 T_{\beta_{3}} are resp. the warmup times for α ( t ) \alpha^{(t)} and β 3 ( t ) \beta^{(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 T_{\alpha}=T_{\beta_{3}}=T_{\alpha,\beta_{3}} , and we typically use T α , β 3 = T T_{\alpha,\beta_{3}}=T , with T T being the total number of iterations. β s ​ t ​ a ​ r ​ t \beta_{start} is always set to β 1 \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 .

Using β 1 = 0 \beta_{1}=0 to save memory. By setting β 1 = 0 \beta_{1}=0 , we can save on memory by replacing 𝒎 ^ 1 ( t ) \hat{{\bm{m}}}_{1}^{(t)} by 𝒈 ( t ) {\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 β 1 = 0 \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 β 1 ≠ 0 \beta_{1}\neq 0 , as AdEMAMix requires to allocate both 𝒎 1 {\bm{m}}_{1} and 𝒎 2 {\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., 2023 , 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 (2018) ; Zhao et al. (2024) .

Hyperparameter sensitivity. While we introduce up to four new hyperparameters: α \alpha , β 3 \beta_{3} , T α T_{\alpha} , and T β 3 T_{\beta_{3}} , in practice we always set T α = T β 3 = T α , β 3 T_{\alpha}=T_{\beta_{3}}=T_{\alpha,\beta_{3}} , and use T α , β 3 = T T_{\alpha,\beta_{3}}=T in most cases. We show in App. C.1.4 that T α T_{\alpha} and T β 3 T_{\beta_{3}} should only be large enough to prevent instabilities early during training. While all of our experiments on language modeling use β 3 = 0.9999 \beta_{3}=0.9999 , other values such as 0.999 0.999 or even 0.99999 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 β 3 \beta_{3} values. Overall, we find the ranges of values of α , β 3 \alpha,\beta_{3} and T α , β 3 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 β 3 \beta_{3} (e.g. β 3 = 0.999 \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 24 M to 1.3 1.3 B parameters.

### 4.1 Transformer LLM Training

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

Why not simply increasing AdamW’s β 1 \beta_{1} ? While our toy experiment in Fig. 2 already gives some intuition on why increasing Adam’s β 1 \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 110 M models using Adam with large β 1 ∈ { 0.99 , 0.999 , 0.9999 , 0.99999 } \beta_{1}\in\{0.99,0.999,0.9999,0.99999\} . When we use a large β 1 \beta_{1} from the beginning of training, we observe instabilities for larger β 1 \beta_{1} values and no β 1 > 0.9 \beta_{1}>0.9 improves upon the AdamW baseline. One can imagine this to be due to increasing β 1 \beta_{1} too early. Therefore, we also modify AdamW and add the same scheduler on β 1 \beta_{1} as we use on AdEMAMix’s β 3 \beta_{3} . β 1 \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 β 1 = 0.9 \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 300 ​ k 300k iterations. We simply resume training and either increase β 1 \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 β 1 \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 256 k to 1 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 110 M parameter models, training for 256 ​ k 256k iterations gives similar results as training an AdamW model for 500 ​ k 500k iterations. The gap between baseline and our method seems to be increasing as we increase the number of iterations. For 1.3 1.3 B parameter models, training using 770 ​ k 770k steps is on par with training the baseline for 1.5 1.5 M iterations—reaching the same performance with 51 % 51\% fewer tokens (an economy of 96 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 110 M and 1.3 1.3 B parameter models for 256 ​ k 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 110 M parameter AdEMAMix model trained for 256 ​ k 256k iterations, we need to train an AdamW model for 500 ​ k 500k iterations, and the time advantage of AdamW would only allow us to do 2,379 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 s ​ w ​ i ​ t ​ c ​ h t_{switch} , we initialize 𝒎 2 ( t s ​ w ​ i ​ t ​ c ​ h ) = 𝟎 {\bm{m}}_{2}^{(t_{switch})}=\mathbf{0} and replace t t by t − t s ​ w ​ i ​ t ​ c ​ h 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 ).