PaperlayerView exact revision on arXiv

Source-linked AI summary

Convergence of Adam Under Relaxed Assumptions

Haochuan Li, Alexander Rakhlin, Ali Jadbabaie

arXiv:2304.13972v3math.OCcs.LGstat.ML

TL;DR

Adam is widely used, but its convergence theory has relied on strong assumptions such as bounded gradients. This paper proves convergence under generalized smoothness and introduces variance-reduced Adam, obtaining O(ε^-4) and O(ε^-3) gradient complexities, respectively.

  • Problem

    Despite Adam’s widespread use, its convergence analyses have limited guarantees and may require bounded gradients or objective Lipschitzness.

    Source: §1 · p. 1

  • Method

    The analysis proves a high-probability constant bound on gradients along Adam’s trajectory under generalized smoothness, and the paper modifies its momentum update for variance reduction.

    Source: §1 · p. 2

  • Results

    Adam has O(ε^-4) gradient complexity under relaxed assumptions, while variance-reduced Adam achieves O(ε^-3) under an additional component-function smoothness assumption.

    Source: §1 · p. 2

  • Takeaways & Limitations

    The paper concludes that Adam and its variance-reduced version converge under less restrictive assumptions than those used in existing analyses.

    Source: §7 · p. 13

    The stochastic analysis assumes bounded or sub-Gaussian noise, stronger than bounded variance; extending it to bounded variance remains challenging.

    Source: §3.2.2 · p. 6

Abstract

from arXiv · show

In this paper, we provide a rigorous proof of convergence of the Adaptive Moment Estimate (Adam) algorithm for a wide class of optimization objectives. Despite the popularity and efficiency of the Adam algorithm in training deep neural networks, its theoretical properties are not yet fully understood, and existing convergence proofs require unrealistically strong assumptions, such as globally bounded gradients, to show the convergence to stationary points. In this paper, we show that Adam provably converges to $\epsilon$-stationary points with ${O}(\epsilon^{-4})$ gradient complexity under far more realistic conditions. The key to our analysis is a new proof of boundedness of gradients along the optimization trajectory of Adam, under a generalized smoothness assumption according to which the local smoothness (i.e., Hessian norm when it exists) is bounded by a sub-quadratic function of the gradient norm. Moreover, we propose a variance-reduced version of Adam with an accelerated gradient complexity of ${O}(\epsilon^{-3})$.

1 Introduction

The paper addresses gaps in Adam’s convergence theory by relaxing bounded-gradient and global-smoothness assumptions. It proves convergence under generalized smoothness and introduces a variance-reduced Adam method with improved gradient complexity.

1. Introduction

  • The generalized smoothness condition bounds local smoothness by a sub-quadratic function of the gradient norm, rather than requiring global gradient Lipschitzness.The paper describes this condition as a relaxation of standard global smoothness.Source: §1 · pp. 1–2, §1.1 · p. 2
  • A high-probability constant bound on gradients along Adam’s trajectory is proved without assuming objective Lipschitzness.This trajectory-bounding technique is key to the convergence analysis.Source: §1 · p. 2

1.1 Contributions

  • O(ε^-3) gradient complexity is obtained for variance-reduced Adam under an additional generalized smoothness assumption on each component function.The method modifies Adam’s momentum update, inspired by STORM.Source: §1 · p. 2, §1.1 · p. 2
  • O(ε^-4) dimension-free gradient complexity is established for Adam with centered bounded gradient noise under relaxed assumptions.The analysis does not assume bounded gradients or objective Lipschitzness.Source: §1 · pp. 1–2, §1.1 · p. 2
  • The analysis generalizes to gradient noise that is centered and has sub-Gaussian norm, showing that Adam converges with gradient complexity O(ε^-4 log^3.25(1/ε)).Source: §1.1 · p. 2

2 Related work

Prior Adam analyses include convergence gaps, counterexamples, and results that rely on bounded gradients or provide limited guarantees. Related work also covers convergent Adam variants, generalized smoothness analyses, and variance reduction methods.

Convergence of Adam

  • Reddi et al. (2018) identified a gap in Adam’s original convergence proof and constructed non-convergence examples for selected hyper-parameters, without ruling out problem-dependent choices.Shi et al. (2021) and Zhang et al. (2022) also noted the possibility of guarantees with problem-dependent hyper-parameters.Source: §2 · pp. 2–3
  • Several non-convex Adam analyses assume bounded gradients, while weaker-assumption results may converge only to a constant-radius neighborhood or have slower rates.The cited discussion also notes that bounded adaptive stepsizes may require bounded gradients or AdaBound.Source: §2 · p. 3

Variants of Adam

  • Convergent Adam variants change the second-moment update, bound the stepsize, or correct its bias; Wang and Klabjan (2022) study variance-reduced Adam but assume bounded gradients and show asymptotic convergence.Examples include AMSGrad, AdaFom, AdaBound, and AdaShift.Source: §2 · p. 3

Generalized smoothness condition

  • The paper studies smoothness conditions that bound local smoothness by a function of gradient size, beyond the affine (L0, L1) condition used in earlier analyses.The related-work discussion also notes prior variance-reduction methods reaching O(ε^-3) complexity for non-convex optimization.Source: §2 · p. 3

Variance reduction methods

  • The proposed variance-reduced Adam method is motivated by STORM, which adds a momentum-update term to correct bias and reduce variance.Source: §2 · p. 4

3 Preliminaries

The preliminaries specify Adam’s adaptive update and the function and stochastic-gradient assumptions used in its analysis. They replace global smoothness with a gradient-dependent local condition and allow centered bounded or sub-Gaussian gradient noise.

3 Preliminaries

  • The notation defines vector norms and coordinate-wise operations used to express the algorithm and assumptions.The paper uses Euclidean vector norms, spectral matrix norms, coordinate-wise products and quotients, and coordinate-wise inequalities.Source: §3 · p. 4

3.1 Description of the Adam algorithm

  • Adam updates first- and second-moment estimates, corrects their initialization bias, and uses an adaptive stepsize vector to update each iterate.The update combines the re-scaled first moment with the stepsize η/(√v̂_t + λ).Source: §3.1 · p. 4

3.2.1 Function class

  • The objective is differentiable, closed on its open domain, and bounded below; for continuous functions, closedness constrains behavior at the domain boundary.The paper notes this boundary condition helps ensure Adam iterates remain in the domain with high probability when η is small enough.Source: §3.2.1 · p. 4
  • The objective’s local smoothness may grow as a sub-quadratic function of the gradient norm, with 0 ≤ ρ < 2.This condition includes standard smoothness when ρ = 0 or Lρ = 0, and reduces to (L0, L1) smoothness when ρ = 1.Source: §3.2.1 · p. 5

3.2.2 Stochastic gradient

  • The analysis allows either centered, almost surely bounded gradient noise or centered gradient noise with sub-Gaussian norm.The sub-Gaussian assumption is weaker than almost-sure boundedness, but yields a slightly worse convergence rate up to poly-log factors.Source: §3.2.2 · pp. 5–6

4 Results

Under Assumptions 1–3, Adam achieves O(ǫ^-4) leading-term gradient complexity, matching the state of the art among existing Adam analyses. Under the sub-Gaussian noise assumption, the bound includes a polylogarithmic factor; the results also identify λ-dependent complexity and report preliminary CIFAR-10 comparisons across λ values.

4 Results

  • O(ǫ^-4) leading-term gradient complexity is obtained under Assumptions 1–3, matching the state of the art among existing Adam analyses.The theorem also bounds gradients along the trajectory with probability at least 1 − δ.Source: §4 · p. 6
  • The rate is dimension-free but has O(λ^-2) dependence, which may be large when practical choices of λ are small.The paper gives λ = 10^-8 as the default choice in the PyTorch implementation.Source: §4 · p. 6
  • On CIFAR-10, increasing λ from 10^-8 to 0.01 makes only minor differences in test error across the three evaluated architectures.Figure 1 compares a 6-layer CNN, 20-layer ResNet-Small, and 110-layer ResNet110.Source: Figure 1 · p. 7, §4 · p. 7
  • O(ǫ^-4 log^p(1/ǫ)) gradient complexity is bounded under the sub-Gaussian noise assumption.The theorem gives a high-probability trajectory gradient bound; the passage does not provide the value of p.Source: §4 · pp. 6–7

5 Analysis

The analysis proves bounded gradients along Adam’s trajectory without assuming globally bounded gradients, then uses that bound to apply local smoothness and establish convergence. It extends the argument from a deterministic warm-up to stochastic noise settings.

5.1 Bounding the gradients along the optimization trajectory

  • Bounding gradients also controls Adam’s second-order momentum and step sizes; under the generalized smoothness assumption, it implies bounded Hessians at iterate points.The resulting smoothness is local near iterate points, not the stronger global smoothness condition.Source: §5.1 · p. 7
  • A contradiction argument breaks the circularity between assuming bounded gradients to prove convergence and using convergence to infer bounded gradients.It bounds the suboptimality gap and gradient before a stopping time, then contradicts the gap’s crossing a chosen threshold.Source: §5.1 · p. 8

5.2 Local smoothness

  • With bounded gradients, the generalized smoothness condition yields local smoothness, and sufficiently small updates keep consecutive iterates within its required neighborhood.The local inequalities can then be applied like standard smoothness inequalities along the trajectory.Source: §5.1 · p. 7, §5.2 · pp. 8–9

5.3 Warm-up: analysis in the deterministic setting

  • O(ε^-2) gradient complexity follows in the deterministic warm-up, where a Lyapunov function bounds the suboptimality gap and the stopping-time contradiction ensures the bound holds throughout.The warm-up replaces stochastic gradients with exact gradients and assumes Assumptions 1–3.Source: §5.1 · p. 8, §5.3 · pp. 9–10

5.4 Extension to the stochastic setting

  • In the stochastic setting, noise prevents the Lyapunov function from being non-increasing, so the proof instead uses stopping times and concentration to obtain high-probability bounds.The argument covers almost-surely bounded noise and sub-Gaussian noise, with additional analysis for the latter case.Source: §5.4 · pp. 10–11

6 Variance-reduced Adam

The paper introduces VRAdam, which modifies Adam’s momentum update with a same-sample gradient difference. Under stronger smoothness and sampling assumptions, it proves an O(ε^-3) gradient-complexity guarantee, while discussing analysis and initialization trade-offs.

6 Variance-reduced Adam

  • VRAdam adds a same-sample gradient difference to Adam’s momentum update, correcting momentum bias and reducing variance.The added term is inspired by STORM; the method uses a mega-batch or full-batch gradient evaluation at initialization.Source: §6 · p. 11
  • Fixed hyperparameters make tuning more familiar, but require a more accurate initialization gradient estimate from a mega-batch or full batch.The authors leave analysis with time-decreasing hyperparameters as future work.Source: §6 · p. 11

6.2 Convergence guarantees for VRAdam

  • The guarantee assumes generalized smoothness for both the objective and each component function, plus i.i.d. samples; its failure-probability dependence is polynomial.The paper notes that this δ dependence is worse than Adam’s poly-log dependence.Source: §6 · p. 11, §6.2 · p. 12

6.1 Analysis

  • VRAdam’s analysis handles difficult update bounds by defining a stopping time that also bounds the difference between the gradient and momentum estimate.This stopping-time choice creates additional challenges in lower-bounding the objective gap at the stopping time.Source: §6.1 · p. 12

6.2 Convergence guarantees for VRAdam

  • Under Assumptions 1, 5, and 6, VRAdam attains leading-term gradient complexity O(ε^-3), using η=O(ε) versus Adam’s O(ε²); with probability at least 1−δ, every trajectory gradient norm is at most G.Source: §6.2 · p. 12

7 Conclusion and future works

The paper concludes that Adam and VRAdam converge under less restrictive assumptions using generalized non-uniform smoothness and a contradiction argument to bound trajectory gradients. It identifies unresolved noise assumptions and the question of why Adam can outperform SGD, while pointing to possible extensions of the technique.

7 Conclusion and future works

  • The paper’s convergence analysis uses generalized non-uniform smoothness and a contradiction argument to prove bounded gradients along Adam’s trajectory without assuming Lipschitzness of the objective.The stated smoothness bounds the Hessian norm by a sub-quadratic function of the gradient norm almost everywhere.Source: §7 · p. 13
  • The analysis assumes bounded or sub-Gaussian noise, leaving unclear whether Adam’s O(ε^-4) complexity is tight under bounded noise.The authors identify relaxing the noise assumption or proving a matching lower bound as open directions.Source: §7 · p. 13
  • The gradient-bounding techniques may extend to other optimization problems and algorithms when the function class is well behaved and the objective gap at a stopping time can be bounded.The appendix also states the Azuma–Hoeffding inequality and optional stopping theorem as probabilistic tools.Source: §7 · p. 13, A Probabilistic lemmas · p. 18
  • These results do not explain why Adam is better than SGD for training transformers, since non-adaptive SGD has the same O(ε^-4) complexity under weaker conditions.The authors call for a setting where adaptive methods provably outperform SGD.Source: §7 · p. 13

B Proofs related to (ρ, L0, Lρ) smoothness

The appendix provides examples of functions covered by generalized (ρ, L0, Lρ) smoothness, including classes that need not satisfy the paper’s (L0, L1) smoothness condition. It outlines the proof approach for these examples and related smoothness results.

B Proofs related to (ρ, L0, Lρ) smoothness

  • The appendix collects proofs of the smoothness examples and of Lemma 3.4, Lemma 5.1, and Corollary 5.2.These results support the generalized smoothness arguments used in the main analysis.Source: B Proofs related to (ρ, L0, Lρ)… · p. 18

B.1 Examples

  • Univariate rational functions and double exponential functions are (ρ, L0, Lρ) smooth for some 1 < ρ < 2, but need not be (L0, L1) smooth.The appendix gives 1/x and e^(e^x) as examples that fail the latter condition.Source: §B.1 · pp. 18–19
  • The rational-function proof uses partial-fraction decomposition and treats finitely many singular points separately from regions where the second derivative is bounded.For ρ between 3/2 and 2, the local derivative bounds near singular points establish the generalized smoothness condition.Source: §B.1 · pp. 18–19

B.2 Proof of Lemma 3.4

The proof uses a generalized Grönwall comparison to bound gradient norms near a point, then applies that bound along a segment to establish the neighborhood result.

B.2 Proof of Lemma 3.4

  • Lemma B.2 compares u, with u′(t) ≤ ℓ(u(t)), against the solution of v′(t) = ℓ(v(t)) to bound u.The comparison follows by applying the increasing function φ to u and v.Source: §B.2 · p. 20
  • Lemma B.3 applies this comparison to the gradient norm along the segment from x to y.The argument assumes second-order differentiability and handles zeros of the gradient norm by restricting the interval.Source: §B.2 · p. 20
  • For points within distance a, Lemma B.3 bounds the endpoint gradient by ∥∇f(x)∥ + a.The proof applies this bound along the segment to rule out a boundary exit and establish y ∈ dom(f).Source: §B.2 · pp. 20–21

B.3 Proof of Lemma 5.1

These proof sections establish a local gradient bound, set up Adam’s convergence analysis, and give an equivalent update rule alongside technical bounds used in that analysis.

B.3 Proof of Lemma 5.1

  • The proof of Lemma 5.1 defines G = ∥∇f(x)∥ and L = 3L0 + 4LρG^ρ, then uses Corollary 5.2 to complete the argument.Its displayed choice of y and final inequality are not fully stated in these passages.Source: §B.3 · p. 21

B.4 Proof of Corollary 5.2

  • Corollary 5.2 chooses a = max{G, (L0/Lρ)^(1/ρ)} when ∥∇f(x)∥ ≤ G and applies Lemma 3.4 along the segment to y.The proof then uses the resulting bound along z(t) = (1 − t)x + ty.Source: §B.4 · p. 22

C Convergence analysis of Adam

  • Before the stopping time τ, the analysis bounds the suboptimality gap by F and the gradient norm by G.τ is the first time the gap exceeds F, truncated at T + 1.Source: C Convergence analysis of Adam · p. 22

C.1 Equivalent update rule of Adam

  • Adam’s bias-corrected updates have an equivalent, simpler form using re-scaled momentum and second-moment variables.The re-scaled variables are initialized from the first stochastic gradient and its elementwise square.Source: §C.1 · p. 23

C.2 Useful lemmas for Adam

  • For t ≤ τ, the error terms satisfy ∥ε_t∥ ≤ 2σ and ∥γ_{t−1}∥ ≤ 2σ under the stated stepsize condition.Under the parameter choices in either main theorem, I1 ≤ I2 and I1/T ≤ ε².Source: §C.2 · p. 24

C.3 Proof of Theorem 4.1

The proof establishes a high-probability convergence event by combining a descent argument with bounds on the update and accumulated error before a stopping time.

C.3 Proof of Theorem 4.1

  • The descent analysis combines a bound on each update with a bound on the sum of squared error terms before τ.The resulting lemma simultaneously bounds the suboptimality gap and accumulated error.Source: §C.3 · pp. 24–25
  • With probability at least 1 − δ, the proof rules out early stopping and obtains τ = T + 1 on the event E.The contradiction uses I1 ≤ I2; the accompanying bound is then supplied by Lemma C.8.Source: §C.3 · pp. 25–26

C.4 Proof of Lemma C.7

  • Before τ, the proof bounds the accumulated error using martingale differences and Azuma–Hoeffding under the bounded-noise assumption.The stopping-time indicator and γ_{t−1} depend only on prior samples, yielding the martingale-difference sequence.Source: §C.3 · p. 25, §C.4 · pp. 26–27

C.5 Omitted proofs for Adam

  • The omitted proofs bound the stochastic gradient, re-scaled moments, and stepsize before τ, and verify inequalities required by the parameter choices.The moment bounds follow by induction from the gradient bound before τ.Source: §C.5 · pp. 28–29

C.6 Proof of Theorem 4.2

The proof extends the stopping-time analysis to sub-Gaussian noise, while the accompanying VRAdam analysis bounds local smoothness and per-step updates before its stopping time.

C.6 Proof of Theorem 4.2

  • The sub-Gaussian-noise proof bounds the probability of early stopping so that the event of reaching T + 1 has probability at least 1 − δ.It splits early stopping into two cases and combines their probability bounds.Source: §C.6 · pp. 30–31

D Convergence analysis of VRAdam

  • The VRAdam analysis defines its momentum deviation and stopping times, then bounds the objective gap, gradient norm, and deviation before τ.Before τ, the update is bounded by ηD with D = 2G/λ.Source: D Convergence analysis of VRAdam · p. 31

D.1 Useful lemmas

  • For VRAdam, Lemma D.1 gives a local smoothness bound when the gradient is bounded, and Lemma D.3 bounds each update by ηD, where D = 2G/λ.The local smoothness result applies within distance r/2 of x.Source: §D.1 · p. 32

D.2 Proof of Theorem 6.2

The proof combines lemmas for VRAdam to bound the probability of stopping early or failing to converge. It also adjusts component-function smoothness bounds using the gradient bound in Assumption 6.

D.2 Proof of Theorem 6.2

  • The proof uses that the descent lemma still holds for VRAdam and verifies the theorem's parameter choices satisfy the required lemmas.The parameter check is attributed to Lemma D.5.Source: §D.2 · pp. 32, 34
  • P(F ∪ {τ ≤ T}) ≤ δ, so the proof bounds the probability of either nonconvergence or stopping by time T by δ.Here F is defined as the event of not converging to stationary points.Source: §D.2 · p. 34

D.3 Proofs of lemmas in Appendix D.1

  • Assumption 6 gives ∥∇f(x, ξ)∥ ≤ G + σ ≤ 2G, leading to locality size r/2 and smoothness constant 4L for component functions.These bounds follow by applying Corollary 5.2 with 2G in place of G.Source: §D.3 · p. 35
Loading 2304.13972v3…