Towards Understanding Convergence and Generalization of AdamW
Pan Zhou,Xingyu Xie,Zhouchen Lin,Shuicheng Yan
DOI: https://doi.org/10.1109/tpami.2024.3382294
IF: 23.6
2024-01-01
IEEE Transactions on Pattern Analysis and Machine Intelligence
Abstract:AdamW modifies Adam by adding a decoupled weight decay to decay network weights per training iteration. For adaptive algorithms, this decoupled weight decay does not affect specific optimization steps, and differs from the widely used l<sub>2</sub>-regularizer which changes optimization steps via changing the first- and second-order gradient moments. Despite its great practical success, for AdamW, its convergence behavior and generalization improvement over Adam and l<sub>2</sub>-regularized Adam ( l<sub>2</sub>-Adam) remain absent yet. To solve this issue, we prove the convergence of AdamW and justify its generalization advantages over Adam and l<sub>2</sub>-Adam. Specifically, AdamW provably converges but minimizes a dynamically regularized loss that combines vanilla loss and a dynamical regularization induced by decoupled weight decay, thus yielding different behaviors with Adam and l<sub>2</sub>-Adam. Moreover, on both general nonconvex problems and PŁ-conditioned problems, we establish stochastic gradient complexity of AdamW to find a stationary point. Such complexity is also applicable to Adam and l<sub>2</sub>-Adam, and improves their previously known complexity, especially for over-parametrized networks. Besides, we prove that AdamW enjoys smaller generalization errors than Adam and l<sub>2</sub>-Adam from the Bayesian posterior aspect. This result, for the first time, explicitly reveals the benefits of decoupled weight decay in AdamW. Experimental results validate our theory.
computer science, artificial intelligence,engineering, electrical & electronic