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

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

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 B . We start by removing B B from the RedPajama training data and train AdamW and AdEMAMix models. For those runs, B B is therefore akin to a batch from the validation set. We measure the loss on B B while training. Now we can resume training from various checkpoints, and inject B B into the training data at various times t B t_{B} . By comparing the two runs—one having trained on B B , while the other never saw B B —we can visualize how B B is learned, and how it is forgotten. After training on B B , we expect the loss on B B to decrease suddenly and then increase as the model forgets the contribution of that batch. When comparing the forgetting curves for AdamW and AdEMAMix in Fig. 4 , we see striking differences. AdamW models forget much faster—the loss over B B increases faster—than AdEMAMix models. Moreover, at the end of training, batches processed by AdEMAMix see their loss being improved over many thousands of iterations. Additional experiments on forgetting can be found in App. C.1.2 .

### 4.2 Mamba LM Training

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

### 4.3 ViT Training

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 37 or 320 320 ). We use two subsets of the ImageNet ( Russakovsky et al., 2015 ) dataset: (i) the widely used ImageNet-1k subset, consisting of 1.3 1.3 M images and 1,000 1{,}000 possible classes; (ii) a filtered and pre-processed version of the ImageNet-21k ( Ridnik et al., 2021 ) containing 11 11 M images corresponding to 10,450 10{,}450 classes. For each, we measure the test loss on a held-out test set. We use the ViT architecture ( Dosovitskiy et al., 2021b ) at two different scales: 24 24 M and 86 86 M parameters. Importantly, if it is common in the vision literature to pre-train large ViT models and finetune them on smaller datasets, this work focuses on pretraining optimization and we therefore train and test on the same distribution. The models’ hyperparameters are detailed in App. B.3 , we use a batch size of 4096 4096 for all our experiments. We used training hyperparameters from Dosovitskiy et al. (2021b) 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 } \alpha\in\{1,5,10,15,20\} and β 3 ∈ { 0.999 , 0.9999 } \beta_{3}\in\{0.999,0.9999\} . We train models for 320 320 and 37 37 epochs for resp. the ImageNet-1k and ImageNet-21k datasets, corresponding in both cases to 100 ​ k 100k iterations. Data-augmentation techniques have been shown to be central to the efficient training of ViTs Touvron et al. (2021) ; Touvron et al. (2022) . We use simple data-augmentations, which includes mixup ( Zhang et al., 2018 ) . We train on 8 8 A100 NVIDIA GPUs using Pytorch Fully Sharded Data-Parallelism ( Zhao et al., 2023 , FSDP) .

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

## 5 Conclusion

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

## 6 Acknowledgements

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

## Appendix A Implementation details

### A.1 Deriving the β 3 \beta_{3} scheduler

Let’s consider S ⁡ ( t ) S(t) , the sum of the weights given to the last t t gradients by an EMA parameterized by β ∈ [ 0 , 1 [ \beta\in[0,1[ : S ⁡ ( t ) = ( 1 − β ) ​ ∑ i = 0 t β i S(t)=(1-\beta)\sum_{i=0}^{t}\beta^{i} We want to know which timestep t t would correspond to a cumulative weight of 0.5 0.5 : ( 1 − β ) ​ ∑ i = 0 t β i = 0.5 ⇔ β t + 1 = 0.5 ⇔ t = ln ⁡ ( 0.5 ) ln ⁡ ( β ) − 1 (1-\beta)\sum_{i=0}^{t}\beta^{i}=0.5\Leftrightarrow\beta^{t+1}=0.5\Leftrightarrow t=\frac{\ln(0.5)}{\ln(\beta)}-1 Let f ⁡ ( β ) = ln ⁡ ( 0.5 ) ln ⁡ ( β ) − 1 f(\beta)=\frac{\ln(0.5)}{\ln(\beta)}-1 . This function provides how many past consecutive gradients receive a cumulative weight of 0.5 0.5 .

Its inverse is: f − 1 ​ ( t ) = 0.5 1 t + 1 f^{-1}(t)=0.5^{\frac{1}{t+1}}

We want a scheduler which increases β \beta from β s ​ t ​ a ​ r ​ t \beta_{start} to β e ​ n ​ d \beta_{end} such that f ⁡ ( β ) f(\beta) increases linearly. Given an interpolating parameter μ ∈ [ 0 , 1 ] \mu\in[0,1] , this scheduler can be written as: β ⁡ ( μ ) = f − 1 ​ ( ( 1 − μ ) ​ f ​ ( β s ​ t ​ a ​ r ​ t ) + μ ​ f ​ ( β e ​ n ​ d ) ) \beta(\mu)=f^{-1}((1-\mu)f(\beta_{start})+\mu f(\beta_{end})) By replacing f f and f − 1 f^{-1} by their respective formula, one can arrive to: β ⁡ ( μ ) = exp ⁡ ( ln ⁡ ( β s ​ t ​ a ​ r ​ t ) ​ ln ⁡ ( β e ​ n ​ d ) ( 1 − μ ) ​ ln ⁡ ( β e ​ n ​ d ) + μ ​ ln ⁡ ( β s ​ t ​ a ​ r ​ t ) ) \beta(\mu)=\exp\big(\frac{\ln(\beta_{start})\ln(\beta_{end})}{(1-\mu)\ln(\beta_{end})+\mu\ln(\beta_{start})}\big) By setting β e ​ n ​ d = β 3 \beta_{end}=\beta_{3} and μ = t T β 3 \mu=\frac{t}{T_{\beta_{3}}} , we arrive to the β 3 \beta_{3} -scheduler introduced in § 3 . We show the shape of our scheduler and compare it to a linear scheduler in Fig. 7 .

### A.2 Pytorch implementation

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

### A.3 Optax implementation

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

## Appendix B Experimental details

### B.1 Transformer LLM experiments

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

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

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

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

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

### B.2 Mamba LM experiments

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

Hyperparameter details. We used parameters mostly taken from Gu & Dao (2023) : • Learning rate: η = 0.0006 \eta=0.0006 ,

• Warmup steps: 3 ​ k 3k ,

• η \eta -scheduler type: cosine decay to 10 − 5 10^{-5} ,

• Weight-decay: λ = 0.1 \lambda=0.1 ,

• Gradient clipping value: 1 1 ,

• Total training steps T ∈ { 64 ​ k , 128 ​ k } T\in\{64k,128k\} .

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

### B.3 ViT experiments

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

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

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

## Appendix C Additional results

### C.1 LLM experiments

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

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

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

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

#### C.1.2 More results on forgetting

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