Adam with model exponential moving average is effective for nonconvex optimization

TL;DR

This paper provides a theoretical analysis showing that a clipped Adam optimizer with model exponential moving average achieves optimal convergence in nonconvex settings, highlighting its advantages in complex model training.

Contribution

It offers the first theoretical proof of Adam with model EMA's effectiveness and optimal convergence rates in nonconvex optimization.

Findings

Clipped Adam with model EMA achieves optimal convergence rates.

Coordinate-wise adaptivity of Adam is provably beneficial.

Analysis emphasizes the importance of momentum and discounting in Adam.

Abstract

In this work, we offer a theoretical analysis of two modern optimization techniques for training large and complex models: (i) adaptive optimization algorithms, such as Adam, and (ii) the model exponential moving average (EMA). Specifically, we demonstrate that a clipped version of Adam with model EMA achieves the optimal convergence rates in various nonconvex optimization settings, both smooth and nonsmooth. Moreover, when the scale varies significantly across different coordinates, we demonstrate that the coordinate-wise adaptivity of Adam is provably advantageous. Notably, unlike previous analyses of Adam, our analysis crucially relies on its core elements -- momentum and discounting factors -- as well as model EMA, motivating their wide applications in practice.