The AdEMAMix Optimizer: Better, Faster, Older — part 4/5
Cleaned-text code points [45607, 60460). Bibliography omitted here; source.md includes it.
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 110 M parameter models for 300 k 300k iterations, using a max learning rate of 0.001 0.001 and a batch size of 64 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.
#### 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 𝒎 2 {\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.
#### C.1.4 Hyperparameter sensitivity
Hyperparameter sensitivity. Depending on whether the α \alpha and β 3 \beta_{3} schedulers are used, AdEMAMix adds up to 4 4 new hyperparameters: α \alpha , β 3 \beta_{3} , T α T_{\alpha} and T β 3 T_{\beta_{3}} . In all our experiments we tied T α = T β 3 = T α , β 3 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 β 3 \beta_{3} in Fig. 12 , revealing wide ranges of values for which hyperparameters are outperforming the AdamW baseline. We study the sensitivity of T α , β 3 T_{\alpha,\beta_{3}} as well as justify the choice of using a scheduler on α \alpha in Fig. 13 . When training from scratch T α , β 3 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 β 1 \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 β 1 \beta_{1} , we observe some loss spikes, yet the final loss value is relatively unaffected.
#### 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 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 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 32 k 32k and 64 k 64k iterations, AdEMAMix still outperforms AdamW. In Fig. 17 we observe that reducing β 3 \beta_{3} can mostly alleviate the problem.
#### 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 ⋆ , y ⋆ ) ∈ arg min ( x , y ) f ( x , y ) , with f ( x , y ) = 8 ( x − 1 ) 2 × ( 1.3 x 2 + 2 x + 1 ) + 0.5 ( y − 4 ) 2 (x^{\star},y^{\star})\in\argmin_{(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. (2023) as a function with sharp curvature along the x x -axis, and flatter curvature along the y y -axis. Initializing the first iterate 𝒙 ( 0 ) = ( 0.3 , 1.5 ) {\bm{x}}^{(0)}=(0.3,1.5) , we start optimizing with (i) Adam with a large β 1 = 0.999 \beta_{1}=0.999 , and (ii) AdEMAMix with β 1 = 0.9 \beta_{1}=0.9 and β 3 = 0.999 \beta_{3}=0.999 . In both cases β 2 = 0.999 \beta_{2}=0.999 . To make the experiment interesting, we initialize the EMA buffers for both methods to either ( − 3 , 0 ) (-3,0) or ( − 0.8 , − 3 ) (-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 \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.
Results training 110 110 M parameters LMs. To further demonstrate that simply increasing β 1 \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 β 1 ∈ { 0.99 , 0.999 , 0.9999 , 0.99999 } \beta_{1}\in\{0.99,0.999,0.9999,0.99999\} . Naively increasing β 1 \beta_{1} in AdamW does not work; adding a scheduler on β 1 \beta_{1} to smoothly increase its value during the entirety of the training also fails. In Fig. 20 , we show results when increasing β 1 \beta_{1} in the middle of training, with and without scheduler on β 1 \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 β 1 \beta_{1} value does not provide any noticeable gain.
#### C.1.7 Removing m 1 m_{1} by setting β 1 = 0 \beta_{1}=0
Effect of the batch size when β 1 = 0 \beta_{1}=0 . In this section we investigate the impact of setting β 1 = 0 \beta_{1}=0 , which allows to replace 𝒎 1 ( t ) {\bm{m}}_{1}^{(t)} by the current gradient 𝒈 ( t ) {\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 β 1 \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 β 1 = 0 \beta_{1}=0 is less stable when using larger batch sizes. However, despite the spikes, the final loss for β 1 = 0 \beta_{1}=0 and β 1 = 0.9 \beta_{1}=0.9 are very similar— β 1 = 0.9 \beta_{1}=0.9 yielding slightly better results for smaller batch sizes.
Setting β 1 = 0 \beta_{1}=0 for 1.3 1.3 B parameter models. In this section, using a 1.3 1.3 B parameter model, we re-do the same experiment as in Fig. 1 , but we set β 1 = 0 \beta_{1}=0 . In Fig. 22 , we show the results of this experiment.
#### C.1.8 Comparing m 1 + α m 2 m_{1}+\alpha m_{2} versus ( 1 − α ) m 1 + α m 2 (1-\alpha)m_{1}+\alpha m_{2}
In the equation AdEMAMix , we combine 𝒎 1 {\bm{m}}_{1} and 𝒎 2 {\bm{m}}_{2} using 𝒎 1 + α 𝒎 2 {\bm{m}}_{1}+\alpha{\bm{m}}_{2} . In this section we contrast this to using a convex combination ( 1 − α ) 𝒎 1 + α 𝒎 2 (1-\alpha){\bm{m}}_{1}+\alpha{\bm{m}}_{2} , with α ∈ [ 0 , 1 ] \alpha\in[0,1] .
Are those two somewhat equivalent? Given the combination η ( 𝒎 1 + α 𝒎 2 ) \eta({\bm{m}}_{1}+\alpha{\bm{m}}_{2}) , it is possible to rewrite it as a convex combination as such: η ( 𝒎 1 + α 𝒎 2 ) = η ^ ( ( 1 − α ^ ) 𝒎 1 + α ^ 𝒎 2 ) , with η ^ = η ( α + 1 ) and α ^ = α α + 1 . \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 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 𝒎 1 {\bm{m}}_{1} and 𝒎 2 {\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 𝒎 1 {\bm{m}}_{1} and 𝒎 2 {\bm{m}}_{2} .
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 η ∈ { 0.4 , 0.6 , 0.8 , 1 , 1.2 , 1.4 , 1.6 } \eta\in\{0.4,0.6,0.8,1,1.2,1.4,1.6\} and α ∈ { 2 , 4 , 6 , 8 , 10 , 12 , 14 } \alpha\in\{2,4,6,8,10,12,14\} for the formulation in equation AdEMAMix , and α ∈ { 0.3 , 0.4 , 0.5 , 0.6 , 0.7 , 0.8 , 0.9 } \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.
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 η ( t ) \eta^{(t)} and α ( t ) \alpha^{(t)} denote the cosine and linear schedulers for resp. η \eta and α \alpha . The convex combination would need to use the following schedulers: η ^ ( t ) = η ( t ) ( α ( t ) + 1 ) , α ^ ( t ) = α ( t ) α ( t ) + 1 . \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 η ^ ( t ) \hat{\eta}^{(t)} , which is not a cosine scheduler. The α \alpha scheduler, α ^ ( t ) \hat{\alpha}^{(t)} , would not be linear. We plot those schedulers in Fig. 25 .
#### C.1.9 Miscellaneous results
Impact of β 3 \beta_{3} and α \alpha schedulers when starting AdEMAMix from AdamW. As mentioned in § 3 , we do not necessarily need the α \alpha and β 3 \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 β 3 \beta_{3} when starting AdEMAMix from an AdamW checkpoint at t s w i t c h = 300 k t_{switch}=300k . We observe how increasing T β 3 = T α T_{\beta_{3}}=T_{\alpha} only slows down the convergence. This serves to illustrate that the schedulers for α \alpha and β 3 \beta_{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 α \alpha and β 3 \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., 2019 ) . Using large momentum values too early seem to increase this phenomenon. See Fig. 27 .
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 100 k 100k of linear decay. In this section we experiment with 200 k 200k steps of linear decay, the rest of the parameters are the same: we use a 1.3 1.3 B parameter model with a max learning rate of 10 − 4 10^{-4} and remaining hyperparameters as in Table 4 . Results can be seen in Fig. 28 .
Same figure as Fig. 1 , including the AdamW trained on 197 197 B tokens. In Fig. 1 we represent the AdamW experiment trained on 197 197 B tokens by a blue horizontal line for aesthetic reasons. In Fig. 29 we include this missing curve to the same plot.
### 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.
### C.3 Comparison with other methods
#### C.3.1 Comparison with AdMeta and DEMA
Double Exponential Moving Average (DEMA). Originally introduced by Mulloy (1994) , 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 ) ) \text{EMA}_{\beta}^{(T-N:T)}\triangleq\text{EMA}(\beta,{\bm{g}}^{(T-N)},\ldots,{\bm{g}}^{(T)}) , let N 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 − 2 N ) , … , 𝒈 ( T ) ) = 2 ⋅ EMA β ( T − N : T ) − EMA ( β , EMA β ( T − 2 N : T − N ) , … , EMA β ( T − N : T ) ) . \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.