Provable Adaptivity of Adam under Non-uniform Smoothness

TL;DR

The paper addresses the theoretical gap for Adam-type optimizers by proving trajectory-wise convergence of randomly reshuffled Adam under the relaxed $(L_0,L_1)$-smoothness condition, which permits the gradient’s Hessian to grow with the gradient norm. It introduces and analyzes an RR Adam scheme with diminishing learning rates, avoiding the traditional bounded-smoothness assumption and gradient bounds. The authors show that Adam can converge to a neighborhood of stationary points at a faster rate than SGD in regimes where $D_0=0$ and can outperform SGD in certain $D_0>0$ scenarios, highlighting adaptivity as a practical theoretical advantage. They also provide a detailed comparison to SGD with a refined lower bound for SGD under the same setting and discuss practical implications, including hyperparameter choices and limitations. The work lays groundwork for understanding the benefits of Adam’s local adaptivity and motivates future exploration of momentum’s role within adaptive methods.

Abstract

Adam is widely adopted in practical applications due to its fast convergence. However, its theoretical analysis is still far from satisfactory. Existing convergence analyses for Adam rely on the bounded smoothness assumption, referred to as the \emph{L-smooth condition}. Unfortunately, this assumption does not hold for many deep learning tasks. Moreover, we believe that this assumption obscures the true benefit of Adam, as the algorithm can adapt its update magnitude according to local smoothness. This important feature of Adam becomes irrelevant when assuming globally bounded smoothness. This paper studies the convergence of randomly reshuffled Adam (RR Adam) with diminishing learning rate, which is the major version of Adam adopted in deep learning tasks. We present the first convergence analysis of RR Adam without the bounded smoothness assumption. We demonstrate that RR Adam can maintain its convergence properties when smoothness is linearly bounded by the gradient norm, referred to as the \emph{$(L_0, L_1)$-smooth condition. We further compare Adam to SGD when both methods use diminishing learning rate. We refine the existing lower bound of SGD and show that SGD can be slower than Adam. To our knowledge, this is the first time that Adam and SGD are rigorously compared in the same setting and the advantage of Adam is revealed.

Figures (3)

• Figure 1: Experiments on the WMT 2014 dataset trained with the transformer. (a): The training loss of SGD and Adam. (b): The gradient norm vs. the local smoothness on the training trajectory. The blue line in (b) stands for $\log(\text{local smoothness})= \log(\text{gradient norm})+1.4$. It can be observed that $(e^{1.4},0)$-smooth condition holds in this task. Similar results can be seen in zhang2019gradient.
• Figure 2: Reconduct of experimental results from Zhang2022Adam. The objective function is defined in Eq. (\ref{['eq: counter_example']}). One can observe that while letting $\beta_2$ closer to $1$ can make the limiting gradient norm smaller, the limiting gradient norm always stabilizes beyond 0.
• Figure 3: Performance of Adam with different shuffling orders. We respectively plot the training loss and the training accuracy of Adam together with their variances over 10 runs with different random shuffling order. The result indicate the performance of Adam is robust w.r.t. the shuffling order.

Theorems & Definitions (36)

• Theorem 1
• Lemma 1: Informal
• Remark 1
• Remark 2: Difficulty compared to the analysis under $L$-smooth condition
• Remark 3: On Why State-of-the-Art Results Do Not Achieve Trajectory-Wise Convergence as Ours
• Remark 4: Similar potential functions in the existing literature.
• Theorem 2
• proof
• Theorem 3: Theorem \ref{['thm: sgd_diminishing']}, restated
• Lemma 2