English

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

Machine Learning 2022-11-22 v2

Abstract

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

Keywords

Cite

@article{arxiv.2210.11693,
  title  = {Amos: An Adam-style Optimizer with Adaptive Weight Decay towards Model-Oriented Scale},
  author = {Ran Tian and Ankur P. Parikh},
  journal= {arXiv preprint arXiv:2210.11693},
  year   = {2022}
}
R2 v1 2026-06-28T04:08:40.805Z