# Amos: An Adam-style Optimizer with Adaptive Weight Decay towards Model-Oriented Scale

arXiv: https://arxiv.org/abs/2210.11693 · HTML: https://mlsanity.com/paper/2210.11693/

## Key facts

- Cited by 7 papers in the map (data through 14 Oct 2026).
- Builds on 25 of its 41 references that are in the map.
- First posted on arXiv 21 Oct 2022.

## Builds on (most cited first, 25 of 25)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://mlsanity.com/paper/1810.04805/) (2018)
- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer](https://mlsanity.com/paper/1910.10683/) (2019)
- [RoBERTa: A Robustly Optimized BERT Pretraining Approach](https://mlsanity.com/paper/1907.11692/) (2019)
- [Batch Normalization: Accelerating Deep Network Training by Reducing Internal Covariate Shift](https://mlsanity.com/paper/1502.03167/) (2015)
- [Scaling Laws for Neural Language Models](https://mlsanity.com/paper/2001.08361/) (2020)
- [Layer Normalization](https://mlsanity.com/paper/1607.06450/) (2016)
- [PaLM: Scaling Language Modeling with Pathways](https://mlsanity.com/paper/2204.02311/) (2022)
- [RoFormer: Enhanced Transformer with Rotary Position Embedding](https://mlsanity.com/paper/2104.09864/) (2021)
- [On the Opportunities and Risks of Foundation Models](https://mlsanity.com/paper/2108.07258/) (2021)
- [A Broad-Coverage Challenge Corpus for Sentence Understanding through Inference](https://mlsanity.com/paper/1704.05426/) (2017)
- [Self-Attention with Relative Position Representations](https://mlsanity.com/paper/1803.02155/) (2018)
- [Aligning Books and Movies: Towards Story-like Visual Explanations by Watching Movies and Reading Books](https://mlsanity.com/paper/1506.06724/) (2015)
- [Adafactor: Adaptive Learning Rates with Sublinear Memory Cost](https://mlsanity.com/paper/1804.04235/) (2018)
- [On the Convergence of Adam and Beyond](https://mlsanity.com/paper/1904.09237/) (2019)
- [Introduction to Online Convex Optimization](https://mlsanity.com/paper/1909.05207/) (2019)
- [Fine-Tuning can Distort Pretrained Features and Underperform Out-of-Distribution](https://mlsanity.com/paper/2202.10054/) (2022)
- [Stochastic First- and Zeroth-order Methods for Nonconvex Stochastic Programming](https://mlsanity.com/paper/1309.5549/) (2013)
- [Large Batch Optimization for Deep Learning: Training BERT in 76 minutes](https://mlsanity.com/paper/1904.00962/) (2019)
- [Mind the Gap: Understanding the Modality Gap in Multi-modal Contrastive Representation Learning](https://mlsanity.com/paper/2203.02053/) (2022)
- [Shampoo: Preconditioned Stochastic Tensor Optimization](https://mlsanity.com/paper/1802.09568/) (2018)
- [The Power of Interpolation: Understanding the Effectiveness of SGD in Modern Over-parametrized Learning](https://mlsanity.com/paper/1712.06559/) (2017)
- [Shatter: An Efficient Transformer Encoder with Single-Headed Self-Attention and Relative Sequence Partitioning](https://mlsanity.com/paper/2108.13032/) (2021)

## Built on by (most cited first, 7 of 7)

- [The Falcon Series of Open Language Models](https://mlsanity.com/paper/2311.16867/) (2023)
- [Accelerating Neural Network Training: An Analysis of the AlgoPerf Competition](https://mlsanity.com/paper/2502.15015/) (2025)
- [RMNP: Row-Momentum Normalized Preconditioning for Scalable Matrix-Based Optimization](https://mlsanity.com/paper/2603.20527/) (2026)
- [Improving Deep Learning Optimization through Constrained Parameter Regularization](https://mlsanity.com/paper/2311.09058/) (2023)
- [Evolution of Optimization Methods: Algorithms, Scenarios, and Evaluations](https://mlsanity.com/paper/2604.12968/) (2026)
- [A Theoretical and Empirical Study on the Convergence of Adam with an "Exact" Constant Step Size in Non-Convex Settings](https://mlsanity.com/paper/2309.08339/) (2023)
- [Poisson Process for Bayesian Optimization](https://mlsanity.com/paper/2402.02687/) (2024)

## Abstract (arXiv)

We present Amos, a stochastic gradient-based optimizer designed for training deep neural networks. It can be viewed as an Adam optimizer with theoretically supported, adaptive learning-rate decay and weight decay. A key insight behind Amos is that it leverages model-specific information to determine the initial learning-rate and decaying schedules. When used for pre-training BERT variants and T5, Amos consistently converges faster than the state-of-the-art settings of AdamW, achieving better validation loss within <=70% training steps and time, while requiring <=51% memory for slot variables. Our code is open-sourced at: https://github.com/google-research/jestimator