Adam Converges Without Any Modification On Update Rules

Abstract

Adam is the default algorithm for training neural networks, including large language models (LLMs). However, \citet{reddi2019convergence} provided an example that Adam diverges, raising concerns for its deployment in AI model training. We identify a key mismatch between the divergence example and practice: \citet{reddi2019convergence} pick the problem after picking the hyperparameters of Adam, i.e., $(\beta_1,\beta_2)$; while practical applications often fix the problem first and then tune $(\beta_1,\beta_2)$. In this work, we prove that Adam converges with proper problem-dependent hyperparameters. First, we prove that Adam converges when $\beta_2$ is large and $\beta_1 < \sqrt{\beta_2}$. Second, when $\beta_2$ is small, we point out a region of $(\beta_1,\beta_2)$ combinations where Adam can diverge to infinity. Our results indicate a phase transition for Adam from divergence to convergence when changing the $(\beta_1, \beta_2)$ combination. To our knowledge, this is the first phase transition in $(\beta_1,\beta_2)$ 2D-plane reported in the literature, providing rigorous theoretical guarantees for Adam optimizer. We further point out that the critical boundary $(\beta_1^*, \beta_2^*)$ is problem-dependent, and particularly, dependent on batch size. This provides suggestions on how to tune $\beta_1$ and $\beta_2$: when Adam does not work well, we suggest tuning up $\beta_2$ inversely with batch size to surpass the threshold $\beta_2^*$, and then trying $\beta_1< \sqrt{\beta_2}$. Our suggestions are supported by reports from several empirical studies, which observe improved LLM training performance when applying them.