PaperlayerView exact revision on arXiv

Source-linked AI summary

Clipping Improves Adam-Norm and AdaGrad-Norm when the Noise Is Heavy-Tailed

Savelii Chezhegov, Yaroslav Klyukin, Andrei Semenov, Aleksandr Beznosikov, Alexander Gasnikov, Samuel Horváth, Martin Takáč, Eduard Gorbunov

arXiv:2406.04443v3cs.LGmath.OC

TL;DR

High-probability convergence guarantees for AdaGrad- and Adam-type methods under heavy-tailed noise were limited, leaving unclear whether clipping was necessary. The paper proves that unclipped methods can have poor confidence dependence and derives polylogarithmic-confidence bounds for clipped variants. Its results also extend to delayed and coordinate-wise methods.

  • Problem

    Whether AdaGrad and Adam achieve high-probability convergence with polylogarithmic confidence dependence under heavy-tailed noise, and whether clipping is needed, remained unresolved.

    Source: §1 · p. 2

  • Method

    The paper analyzes negative examples for unclipped methods and proves upper bounds for clipped norm methods, clipping gradients in both the momentum update and scaling-factor computation.

    Source: §1.1 · p. 2, §3 · p. 5

  • Results

    Unclipped Adam and AdaGrad variants can have inverse-power confidence dependence, while clipped variants attain polylogarithmic-confidence bounds for the stated smooth convex and non-convex settings.

    Source: §1.1 · p. 2

  • Takeaways & Limitations

    The analysis extends to delayed methods and to coordinate-wise delayed methods, whose bounds explicitly depend on problem dimension.

    Source: §1.1 · p. 2

    The negative result for Adam(D) is established only for β2 = 1−1/T, and the examples require initialization sufficiently far from the optimum.

    Source: §2 · p. 5

Abstract

from arXiv · show

Methods with adaptive stepsizes, such as AdaGrad and Adam, are essential for training modern Deep Learning models, especially Large Language Models. Typically, the noise in the stochastic gradients is heavy-tailed for the later ones. Gradient clipping provably helps to achieve good high-probability convergence for such noises. However, despite the similarity between AdaGrad/Adam and Clip-SGD, the current understanding of the high-probability convergence of AdaGrad/Adam-type methods is limited in this case. In this work, we prove that AdaGrad/Adam (and their delayed version) can have provably bad high-probability convergence if the noise is heavy-tailed. We also show that gradient clipping fixes this issue, i.e., we derive new high-probability convergence bounds with polylogarithmic dependence on the confidence level for AdaGrad-Norm and Adam-Norm with clipping and with/without delay for smooth convex/non-convex stochastic optimization with heavy-tailed noise. We extend our results to the case of AdaGrad/Adam with delayed stepsizes. Our empirical evaluations highlight the superiority of clipped versions of AdaGrad/Adam in handling the heavy-tailed noise.

1. Introduction

AdaGrad and Adam are widely used because standard SGD is sensitive to stepsize choice, yet whether adaptive methods need clipping or achieve strong high-probability convergence under heavy-tailed noise remains unclear. The paper gives negative results for unclipped methods and proves high-probability guarantees for clipped variants.

1. Introduction

  • The literature lacked polylogarithmic-confidence high-probability guarantees for Adam under heavy-tailed noise, and it remained unclear whether clipping was needed for AdaGrad and Adam.High-probability bounds matter because converting expectation bounds can introduce inverse-power dependence on the confidence level.Source: §1 · p. 2, §1.2 · p. 3

1.1. Our Contributions

  • Unclipped Adam, AdaGrad, and delayed variants can have inverse-power dependence on confidence, while clipping yields polylogarithmic-confidence bounds for norm methods under heavy-tailed noise.The clipped results cover smooth convex problems for delayed methods and non-convex problems with and without delay; non-convex complexities are optimal up to logarithmic factors.Source: §1.1 · p. 2
  • Experiments report that clipped Adam and AdaGrad outperform their unclipped versions on a heavy-tailed quadratic problem, and clipped Adam significantly outperforms Adam on ALBERT fine-tuning.On RTE, Clip-Adam performs similarly to Adam; the paper also reports similar results for RoBERTa Large fine-tuning.Source: §1.1 · p. 2

1.2. Preliminaries

  • The setup accesses differentiable f through a stochastic first-order oracle that returns ∇fξ(x), an approximation to ∇f(x), where ξ may depend on x and time.High-probability guarantees bound the iterations or oracle calls needed to find x with P{P(x) ≤ ε} ≥ 1−δ, where δ is the confidence level or failure probability.Source: §1.2 · pp. 2–3

1.3. Related Work

  • Prior AdaGrad high-probability analyses relied on restrictive noise assumptions or incurred inverse-power dependence on δ, whereas the paper's clipped-method results have polylogarithmic dependence on δ.The cited prior approaches include sub-Gaussian noise and bounded empirical risk with a non-central moment condition.Source: §1.3 · p. 4

2. Failure of Adam/AdamD and AdaGrad/AdaGradD with Momentum

The section establishes that Adam, AdaGrad with momentum, and their delayed versions can fail to attain desirable high-probability convergence rates. Its examples yield inverse-power confidence dependence, motivating the paper's clipped methods.

2. Failure of Adam/AdamD and AdaGrad/AdaGradD with Momentum

  • Theorem 2.1 gives examples where Adam, AdaGrad with momentum, and delayed variants have high-probability complexity with inverse-power dependence on confidence.The theorem's setting includes α = 2, so the result does not require infinite-variance noise.Source: §2 · p. 5
  • The examples use different time-dependent noise sequences: an initial perturbation inflates the scaling buffer for non-delayed methods, while a final-step perturbation spoils delayed methods.The objective is a one-dimensional Huber loss; the paper notes the examples apply to coordinate-wise and norm versions.Source: §2 · p. 5
  • The negative result for Adam(D) is shown for β2 = 1−1/T and sufficiently distant initialization; results for arbitrary initialization or broader parameter choices are not established.The authors conjecture similar results for more general γ and worse dependence, or non-convergence, when α < 2.Source: §2 · p. 5
  • Replacing the example's objective with a modified function satisfying Assumption 1.4 preserves the proofs for the methods without delay.The modification chooses D greater than |x0|.Source: §2 · p. 5

3. New Upper Bounds

Clipping in both the momentum update and the scaling-factor computation yields high-probability convergence bounds for AdaGrad-Norm and Adam-Norm variants under heavy-tailed noise. The results cover convex and non-convex settings, with and without delayed stepsizes, while identifying an additional assumption for one non-delayed analysis.

Methods

  • Clipping both the momentum update and scaling-factor computation controls large steps and prevents the scaling factor from growing too quickly.The two clipping roles address different risks: heavy-tailed gradients can cause overly large steps, while rapid scaling-factor growth can worsen high-probability guarantees.Source: §3 · p. 5

Convergence results

  • The paper derives new high-probability convergence bounds for the generalized method formalized as Algorithm 2 in both convex and non-convex cases, with a stated main theorem for Clip-AdamD/Clip-AdaGradD-Norm in the convex case.Source: §3 · p. 6

Discussion of the results

  • The non-convex leading terms match the best-known Clip-SGD complexities up to logarithmic factors and are optimal up to logarithmic factors.The first terms of the non-convex bounds can be improved; the first convex term is not optimal, and optimality of its second term remains open.Source: §3 · p. 7
  • The delayed-method results are among the first high-probability bounds with polylogarithmic confidence dependence for heavy-tailed noise without extra Assumption 1.4.The authors state that Theorems 3.1 and 3.2 provide this guarantee.Source: §3 · p. 7
  • The non-delayed analysis requires relatively restrictive Assumption 1.4 to handle the dependence between stochastic gradients and stepsizes.The authors use the assumption to estimate a weighted sum of function values arising in their analysis.Source: §3 · p. 8

Analysis of coordinate-wise methods

  • The coordinate-wise delayed variants also receive non-convex bounds, with explicit dependence on problem dimension.The analysis and bounds are described as similar to those for the non-coordinate-wise delayed methods.Source: §3 · p. 8

4. Numerical Experiments

The experiments compare clipped and unclipped adaptive methods on a heavy-tailed quadratic problem and ALBERT fine-tuning tasks. Clipping improves observed convergence on the quadratic problem and CoLa, while Adam and Clip-Adam perform similarly on RTE.

Quadratic problem

  • On the heavy-tailed quadratic problem, clipped AdaGrad variants with and without delay achieve better convergence with higher probability than their unclipped counterparts.The plotted trajectories show squared distance to the solution; bands span the 10th–90th percentiles, and the vertical axis is logarithmic.Source: §4 · p. 8, Figure 1 · p. 9

ALBERT Base v2 fine-tuning

  • For ALBERT Base v2 fine-tuning, the study uses Adam, Clip-Adam, and Clip-AdamD, with layer-wise clipping shown in the main text.Clipped methods use Adam's batch size and stepsize, with clipping levels tuned for two clipping types.Source: §4 · pp. 8–9
  • The fine-tuning experiments measure gradient-noise heavy-tailedness along the trajectory and find heavier tails for CoLa than for RTE.Source: §4 · p. 9
  • On CoLa, Adam shows poor high-probability convergence, while Clip-Adam performs much better; on RTE, they perform similarly, where the noise is less heavy-tailed.The comparison uses validation-loss medians and 5th–95th percentile bands across 100 runs.Source: §4 · p. 9, Figure 2 · p. 9

Additional RoBERTa Large experiments

  • For the 355M-parameter RoBERTa Large model, clipped variants consistently outperform their unclipped counterparts on QNLI and CoLa.Source: §4 · p. 9

A. Technical Details and Auxiliary Results

This appendix section gathers auxiliary tools used in the proofs, including inequalities, bounds for clipped stochastic gradients, and a martingale concentration inequality.

A. Technical Details and Auxiliary Results

B. Missing Proofs from Section 2

The appendix gives lower-bound proofs showing that adaptive methods including M-AdaGrad, M-AdaGradD, Adam, and AdamD can have inverse-power dependence on the confidence level under heavy-tailed noise, with constructions spanning coordinate-wise and norm variants and a convex smooth problem with an  = 2 noise assumption.

B. Missing Proofs from Section 2

  • The examples use a one-dimensional convex, 1-smooth Huber loss, and their conclusions extend to coordinate-wise and norm versions of the considered methods.The appendix treats the noise construction and proof separately for Adam, M-AdaGrad, AdamD, and M-AdaGradD.Source: B. Missing Proofs from Section 2 · p. 15

B.1. Failure of M-AdaGrad

  • M-AdaGrad can have high-probability complexity with inverse-power dependence on δ under an oracle satisfying Assumption 1.1 with α = 2.Theorem B.4 establishes this for a constructed convex L-smooth problem, subject to its stated parameter conditions.Source: §B.1 · p. 16

B.2. Failure of M-AdaGradD

  • M-AdaGradD can likewise have high-probability complexity with inverse-power dependence on δ under Assumption 1.1 with α = 2.Theorem B.8 gives this result for a constructed convex L-smooth problem.Source: §B.2 · p. 19

Adam

  • Adam can have high-probability complexity with inverse-power dependence on δ under Assumption 1.1 with α = 2, using β2 = 1 − 1/K.Theorem B.12 gives the result for a constructed convex L-smooth problem, with its stated stepsize and starting-point conditions.Source: §B.3 · p. 22

B. Missing Proofs from Section 2

  • The proof uses deterministic AdamD’s endpoint behavior: if |xT| ≤ ν, it invokes an iteration-count lower-bound implication.The surrounding argument first studies iterates that remain positive and outside the ν-neighborhood, then applies the endpoint implication.Source: §B.4 · pp. 24–25
  • AdamD’s high-probability complexity has inverse-power dependence on δ for a constructed convex L-smooth problem with an α = 2 stochastic gradient oracle.Theorem B.16 gives conditions on ε, δ, σ, γ, x0, b0, and the number of steps K; its proof constructs the stochasticity adversarially.Source: §B.4 · p. 25

C. Missing Proofs from Section 3

The appendix provides proof-supporting details for the clipped methods, including an equivalent reweighted algorithm and lemmas on its scaling and momentum sequences.

C. Missing Proofs from Section 3

  • Algorithm 3 reweights the clipped Adam and M-AdaGrad updates with parameter η, yet is equivalent to Algorithm 2 after rescaling b_t and γ by √η.The two algorithms produce the same iterates under the same initialization and stochasticity because γ/b_t is unchanged.Source: C. Missing Proofs from Section 3 · p. 27

C.1. Technical Lemmas

  • The technical lemmas bound relationships within the scaling sequence b_t and momentum sequence m_t; the scaling constant c_m is 1 for clipped M-AdaGrad variants and 1/2 for clipped Adam variants.The appendix also bounds the momentum norm using its update rule and Jensen’s inequality.Source: §C.1 · pp. 27–28

C.2. Non-Convex Case: Methods with Delay

For the non-convex delayed setting, the section establishes high-probability results for clipped AdaGrad- and Adam-Norm variants. Its theorem gives a bound after K + 1 iterations with probability at least 1 − δ.

C.2. Non-Convex Case: Methods with Delay

  • The analysis derives the result using a descent lemma and bounds on terms involving the gradient and stochastic gradient.The proof passage identifies the scalar product ⟨∇f(xr), gr⟩ and squared norm ∥gr∥2 as key terms to estimate.Source: §C.2 · pp. 29–30
  • The theorem gives a high-probability bound after K + 1 iterations for Clip-M-AdaGradD-Norm and Clip-AdamD-Norm.The stated event holds with probability at least 1 − δ.Source: §C.2 · pp. 30–31
  • A further bound is stated for the Clip-M-AdaGradD-Norm and Clip-AdamD-Norm methods with probability at least 1 − δ.The supplied passage identifies distinct bounds for the two methods but does not preserve their complete formulas.Source: §C.2 · p. 35

C.3. Convex Case: Methods with Delay

For the convex delayed setting, the section establishes high-probability bounds for clipped AdaGrad- and Adam-Norm variants. The theorem’s result holds with probability at least 1 − δ.

C.3. Convex Case: Methods with Delay

  • The descent lemma is stated for the delayed clipped methods under the section’s assumptions and an initial-scale condition on b0.The displayed condition includes b0 ≥ 8γL/((1 − β1)^2 c_m^2).Source: §C.3 · p. 36
  • The convex-case theorem gives a bound for Clip-M-AdaGradD-Norm and Clip-AdamD-Norm with probability at least 1 − δ.The theorem assumes the iterates lie in the specified ball around x∗.Source: §C.3 · pp. 38, 39, 44
  • The proof builds probability events inductively to establish the high-probability result over the iterations.It aims to show P{Ek} ≥ 1 − kδ/(K + 1) for successive k.Source: §C.3 · p. 39

C.4. Non-Convex Case: Methods without Delay

For the non-convex setting without delay, the section establishes high-probability results for clipped AdaGrad- and Adam-Norm variants. The theorem states its bound holds with probability at least 1 − δ.

C.4. Non-Convex Case: Methods without Delay

  • The section derives descent lemmas for the two clipped methods before obtaining the final bound.It treats Clip-M-AdaGrad-Norm and Clip-Adam-Norm separately in constructing those lemmas.Source: §C.4 · pp. 44, 46
  • The theorem gives a high-probability result for Clip-M-AdaGrad-Norm and Clip-Adam-Norm after K iterations.The result holds with probability at least 1 − δ.Source: §C.4 · p. 47
  • The proof uses a bounded-gradient assumption, so it does not need to introduce probabilistic events for the gradient bound.The proof passage explicitly attributes this to Assumption 1.4.Source: §C.4 · p. 47

C.5. Non-Convex Case: Coordinate-wise Methods with Delay

For non-convex coordinate-wise methods with delay, the section gives high-probability results for clipped AdaGrad and Adam. It also introduces an assumption intended to capture tighter dependencies on coordinate variances.

C.5. Non-Convex Case: Coordinate-wise Methods with Delay

  • The coordinate-wise proof derives its descent argument by bounding terms separately and summing across coordinates.It applies Cauchy–Schwarz and smoothness before combining the resulting inequalities.Source: §C.5 · pp. 52–53
  • Assumption C.11 is introduced to capture tighter dependencies on the variances of different stochastic-gradient coordinates.The supplied assumption passage also states that the oracle is unbiased.Source: §C.5 · p. 54
  • The theorem gives a high-probability bound after K + 1 iterations for Clip-M-AdaGradD and Clip-AdamD.The stated result holds with probability at least 1 − δ.Source: §C.5 · pp. 54–55

D.1. Quadratic Problem

On the quadratic problem, clipped AdaGrad and Adam variants outperform their non-clipped counterparts in the reported high-probability comparisons. This pattern also holds for AdaGrad with a smaller stepsize and when the noise and clipping level are scaled down.

AdaGrad comparisons

  • Clipped AdaGrad methods consistently outperform non-clipped methods at γ = 1/128.The comparison is reported for the quadratic problem.Source: §D.1 · p. 60

Scaled-noise comparison

  • With noise and clipping levels scaled down 100-fold, clipped methods still achieve better optimization error with high probability.Stepsizes were tuned separately for AdaGrad, Clip-AdaGrad, Adam, and Clip-Adam.Source: §D.1 · p. 60

D.2. ALBERT Base v2 Fine-tuning

In ALBERT Base v2 fine-tuning, the reported comparisons generally favor clipping, though the delayed Adam variant performs worse than Clip-Adam, especially on CoLa. The noise is substantially heavier-tailed on CoLa than on RTE, and the AdaGrad comparison shows little overall difference between clipped and unclipped methods in one setting.

Noise histograms

  • The noise distribution has much heavier tails for CoLa than for RTE.For CoLa, tail metrics remain large; for RTE, ρeR decreases to zero while ρmR remains between 5 and 10.Source: §D.2 · p. 61, Figure 6 · p. 62

Adam clipping results

  • Clipped Adam improves high-probability convergence in the ALBERT fine-tuning comparisons, and coordinate-wise clipping also improves it.The coordinate-wise results are described as similar to the layer-wise results.Source: §D.2 · pp. 61–62

Delayed Adam

  • Clip-AdamD performs worse than Clip-Adam, especially on CoLa.Both used the same clipping level; the authors note that bias correction and further ϵ tuning were not tried for Clip-AdamD and could improve its performance.Source: §D.2 · p. 62, Figure 8 · p. 63

AdaGrad comparisons

  • In the reported AdaGrad setting, clipping makes little overall difference, although Clip-AdaGrad has a much smaller error band on CoLa.The comparison concerns AdaGrad-based methods with and without clipping or delay.Source: §D.2 · pp. 62–63

D.3. Scaling Up: Fine-Tuning of 355M Model

In fine-tuning the 355M-parameter RoBERTa Large model on QNLI and CoLa, clipped Adam and AdaGrad consistently outperform their unclipped counterparts. Noise histograms also indicate substantial heavy-tailed behavior throughout training.

Setup

  • The fine-tuning setup uses batch size 16 and layer-wise clipping with λ = 1 for both datasets.The model is fine-tuned on QNLI and CoLa with identical hyperparameters, including the learning-rate scheduler, warmup ratio, and optimizer parameters.Source: §D.3 · p. 63

Noise histograms

  • Both noise metrics remain consistently large throughout training, and the histogram profiles closely resemble a Lévy α-stable distribution.The noise is sampled using 1000 mini-batched gradient estimators for CoLa and 100 for QNLI.Source: §D.3 · p. 63, Figure 10 · p. 64

Comparison of methods with and without clipping

  • The comparison uses five random seeds rather than 100 because scaling to the 355M-parameter model is computationally constrained.At each step, the authors report median validation loss and the 5th and 95th percentiles.Source: §D.3 · p. 63
  • Clipped Adam and AdaGrad consistently outperform their unclipped counterparts on both QNLI and CoLa.The comparisons use validation loss for RoBERTa Large fine-tuning.Source: §D.3 · p. 63, Figure 11 · p. 65, Figure 12 · p. 65
Loading 2406.04443v3…