On the $O(\frac{\sqrt{d}}{K^{1/4}})$ Convergence Rate of AdamW Measured by $\ell_1$ Norm

TL;DR

The paper analyzes the convergence of AdamW (and its variant NAdamW) in nonconvex stochastic optimization under standard smoothness and noise assumptions, establishing a concrete rate in the $\ell_1$-norm of the gradient: $\frac{1}{K}\sum_{k=1}^K \mathbb{E}[\|\nabla f(x^k)\|_1] \leq O\left( \frac{\sqrt{d}}{K^{1/4}} \right)$ times problem-dependent constants. The rate is shown to be near-optimal with respect to the iteration budget $K$, gradient noise norm $\sigma_s$, and Lipschitz constant $L$, and it separates cleanly into regimes governed by the noise variance, yielding faster rates when gradient noise is small. The analysis introduces an infinity-norm constraint $\|x^k\|_{\infty} < 1/\lambda$ via weight decay and derives concrete parameter guidelines for $\lambda$, $\varepsilon$, and initialization; it additionally extends the results to NAdamW with the same rate. Empirically, the authors demonstrate $\,\|\nabla f(x)\|_1 = \Theta(\sqrt{d})\,\|\nabla f(x)\|_2$ on real DL tasks, supporting the interpretation that the $\ell_1$ rate mirrors SGD’s optimal $\ell_2$ rate in idealized settings, and they provide detailed experimental validation on ResNet and GPT-2 tasks. Overall, the work fills a theoretical gap for AdamW and shows that, under practical conditions, its convergence behavior aligns with the best-known rates for first-order stochastic methods.

Abstract

As the default optimizer for training large language models, AdamW has achieved remarkable success in deep learning. However, its convergence behavior is not theoretically well-understood. This paper establishes the convergence rate $\frac{1}{K}\sum_{k=1}^KE\left[||\nabla f(x^k)||_1\right]\leq O(\frac{\sqrt{d}C}{K^{1/4}})$ for AdamW measured by $\ell_1$ norm, where $K$ represents the iteration number, $d$ denotes the model dimension, and $C$ matches the constant in the optimal convergence rate of SGD. Theoretically, we have $||\nabla f(x)||_2\ll ||\nabla f(x)||_1\leq \sqrt{d}||\nabla f(x)||_2$ for any high-dimensional vector $x$ and $E\left[||\nabla f(x)||_1\right]\geq\sqrt{\frac{2d}π}E\left[||\nabla f(x)||_2\right]$ when each element of $\nabla f(x)$ is generated from Gaussian distribution $\mathcal N(0,1)$. Empirically, our experimental results on real-world deep learning tasks reveal $||\nabla f(x)||_1=\varTheta(\sqrt{d})||\nabla f(x)||_2$. Both support that our convergence rate can be considered to be analogous to the optimal $\frac{1}{K}\sum_{k=1}^KE\left[||\nabla f(x^k)||_2\right]\leq O(\frac{C}{K^{1/4}})$ convergence rate of SGD in the ideal case. We also extend our result to NAdamW, an AdamW variant that employs a double-momentum mechanism, and demonstrate that it maintains the same convergence rate.

Figures (5)

• Figure 1: Illustration of average training loss $f(\mathbf{x}^k)$ for AdamW over epochs/steps, and at the initialization, $f(\mathbf{x}^1)\leq 8$.
• Figure 2: Illustration of $\|\nabla f(\mathbf{x}^k)\|_1=\varTheta(\sqrt{d})\|\nabla f(\mathbf{x}^k)\|_2$ for AdamW over epochs/steps. The gradient norm ratio shows $\frac{\|\nabla f(\mathbf{x}^k)\|_1}{\|\nabla f(\mathbf{x}^k)\|_2}$, and $\sqrt{d}= 4868$, $5060$, and $11136$, respectively.
• Figure 3: Illustration of $\|\mathbf{x}^k\|_{\infty}<\frac{1}{\lambda}$ for AdamW over epochs/steps. The model $\ell_{\infty}$ norm shows $\|\mathbf{x}^k\|_{\infty}$, and $\lambda=0.01$, $0.1$, and $0.05$, respectively.
• Figure 4: Illustration of small $\frac{\sigma_s^2}{d}$ over epochs/steps. The magnitude $\sigma_s^2$ is approximated by $\|\mathbf{g}^k-\nabla f(\mathbf{x}^k)\|^2$ for AdamW without taking expectation, and $d=2.37\times 10^7$, $2.56\times 10^7$, and $1.24\times 10^8$, respectively.
• Figure 5: Illustrations of $\frac{1}{k}\sum_{t=1}^k|\nabla f(x^t)|$ (left) and $x^k$ (right) over steps on the toy example.

Theorems & Definitions (17)

• Theorem 1
• Corollary 1
• Lemma 1
• Theorem 2
• Corollary 2
• Proof 1
• Remark 1
• Remark 2
• Lemma 2
• Lemma 3