Research map / Optimization / Optimizers
The AdEMAMix Optimizer: Better, Faster, Older
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes AdEMAMix, an Adam variant that adds a second slow EMA of gradients (beta3 ~ 0.9999) alongside the fast EMA, letting old gradients remain useful for tens of thousands of steps and yielding faster convergence and lower loss for LLM and ViT pretraining.
Problem
A single exponential moving average of gradients cannot simultaneously weight recent gradients heavily and retain non-negligible weight on very old gradients, so momentum-based optimizers like Adam discard useful long-range gradient information.
Approach
Augment Adam with a second, very slow EMA of gradients (beta3 close to 1) added to the fast momentum in the update numerator, with warmup schedulers for alpha and beta3 to avoid early instabilities, so the optimizer uses both recent and very old gradients.
Key results
- A 1.3B AdEMAMix LLM trained on 101B tokens performs comparably to an AdamW model trained on 197B tokens (+95% tokens).
- For 1.3B models, AdEMAMix at 770k steps matches the AdamW baseline at 1.5M iterations, with 51% fewer tokens.
- AdEMAMix's slow-EMA gradients remain useful over tens of thousands of steps and AdEMAMix forgets training batches more slowly than AdamW.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Replace AdamW with AdEMAMix | 1.3B LLM · In-context learning tasks | score | Mixed |
| Replace AdamW with AdEMAMix | Transformer LLM training · RedPajama; matched training iterations | validation loss | Helped |
| Replace AdamW with AdEMAMix | 168M Mamba LM · FineWeb | loss | Helped |
| Replace AdamW with AdEMAMix | 24M ViT · ImageNet-21k, 11M images | test loss | Helped |
| Replace AdamW with AdEMAMix | 86M ViT · ImageNet-21k, 11M images | test loss | Helped |
| Replace AdamW with AdEMAMix | 86M ViT · ImageNet-1k, 1.3M images | test loss | No clear effect |
| Increase AdamW beta1 above 0.9, including with a scheduler (ablation) | AdamW · 110M LLM training | performance | No clear effect |
| Add a second slow EMA and mix it with the fast EMA | AdamW · Language modeling | validation loss | Helped |
| Switch from AdamW to AdEMAMix during training (ablation) | AdamW-trained LLM · RedPajama | final loss | Helped |
| Switch from AdEMAMix to AdamW during training (ablation) | AdEMAMix-trained LLM · RedPajama | final loss | Hurt |
| Remove the second EMA mid-training (ablation) | AdEMAMix · RedPajama | final loss | Mixed |
| Use AdEMAMix instead of AdamW | LLM optimizer · Held-out training batch after injection | batch loss retention | Helped |
Surprising findings
- Simply increasing Adam's beta1 (including with schedulers) never outperformed the beta1=0.9 baseline, while adding a separate slow EMA did.
- On ImageNet-1k (86M params, 320 epochs), outperforming AdamW on test loss was difficult despite lower training loss, indicating overfitting limits gains.
Limitations
- Retaining very old gradients may hurt domains requiring fast adaptation to sudden distribution shifts.
- Best suited to long training runs; smaller beta3 (e.g. 0.999) is better for low iteration counts.
- Adds up to four hyperparameters (alpha, beta3, T_alpha, T_beta3) and extra optimizer memory.
Experimental setup
Datasets: RedPajama v2 FineWeb ImageNet-1k ImageNet-21k
Benchmarks: HellaSwag Winogrande ARC BoolQ LogiQA MathQA MMLU OpenbookQA PIQA PubmedQA RewardBench Sciq
Metrics: validation loss perplexity test loss accuracy training loss
Baselines: AdamW Lion Sophia Grokfast AdMeta
Builds on: Adam AdamW Lion DiLoCo Mamba
Closest papers by meaning
- Adam with model exponential moving average is effective for nonconvex optimization (NeurIPS 2024 · similarity 0.86)
- MADA: Meta-Adaptive Optimizers Through Hyper-Gradient Descent (ICML 2024 · similarity 0.85)
- MGUP: A Momentum-Gradient Alignment Update Policy for Stochastic Optimization (NeurIPS 2025 · similarity 0.85)
- SMMF: Square-Matricized Momentum Factorization for Memory-Efficient Optimization (AAAI 2025 · similarity 0.85)
- Implicit Bias of AdamW: $\ell_\infty$-Norm Constrained Optimization (ICML 2024 · similarity 0.85)
- Delve into the Applicability of Advanced Optimizers for Multi-Task Learning (NeurIPS 2026 · similarity 0.85)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.