The AdEMAMix Optimizer: Better, Faster, Older

arxiv:2409.03137

Original source · Paper JSON

Data, training, optimization · 2024-09-05

Cites (7) · Cited by (1)

Complete cleaned Markdown · Source Markdown including bibliography

Read in parts

1. Part 1
2. Part 2
3. Part 3
4. Part 4
5. Part 5

Contents

• The AdEMAMix Optimizer: Better, Faster, Older
• Abstract
• 1 Introduction
• 2 Related Work
• 3 Our method: AdEMAMix
• 4 Results
• 4.1 Transformer LLM Training
• 4.2 Mamba LM Training
• 4.3 ViT Training
• 5 Conclusion
• 6 Acknowledgements
• Appendix A Implementation details
• A.1 Deriving the β 3 \beta_{3} scheduler
• A.2 Pytorch implementation
• A.3 Optax implementation
• Appendix B Experimental details
• B.1 Transformer LLM experiments
• B.2 Mamba LM experiments
• B.3 ViT experiments
• Appendix C Additional results
• C.1 LLM experiments
• C.1.1 In-Context Learning (ICL) results
• C.1.2 More results on forgetting
• C.1.3 Removing the second EMA mid-training
• C.1.4 Hyperparameter sensitivity
• C.1.5 Impact of training for fewer iterations
• C.1.6 Limitations of a single EMA
• C.1.7 Removing m 1 m_{1} by setting β 1 = 0 \beta_{1}=0
• 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}
• C.1.9 Miscellaneous results
• C.2 ViT experiments
• C.3 Comparison with other methods
• C.3.1 Comparison with AdMeta and DEMA
• C.3.2 Comparison with Lion
• C.3.3 Adding a third momentum term ( ∼ \sim AggMo)
• Instructions for reporting errors