PaperlayerView exact revision on arXiv

Source-linked AI summary

High-probability Bounds for Non-Convex Stochastic Optimization with Heavy Tails

Ashok Cutkosky, Harsh Mehta

arXiv:2106.14343v2cs.LGmath.OCstat.ML

TL;DR

Heavy-tailed stochastic gradients complicate high-probability convergence analysis for non-convex optimization. The paper analyzes clipped normalized SGD with momentum, including under second-order smoothness and non-Hilbert norms. It reports high-probability critical-point and last-iterate guarantees, and finds that an initial zero-learning-rate period helps training relative to no special treatment.

  • Problem

    For non-convex optimization with heavy-tailed gradients, prior clipped-SGD analysis provides an in-expectation result, while high-probability guarantees are needed to characterize individual runs.

    Source: §1 · p. 2

  • Method

    The paper combines gradient clipping, momentum, and normalized updates to analyze critical-point finding under heavy-tailed gradients, including second-order smooth losses.

    Source: §2 · p. 3, §3 · p. 7

  • Results

    The algorithms provide high-probability critical-point bounds and last-iterate guarantees, and keeping the learning rate zero for the first iterations empirically helps training compared with no special treatment.

    Source: §5 · p. 10

  • Takeaways & Limitations

    After burn-in, the analysis predicts monotonic objective decrease until an approximate critical point, offering an optimization-based perspective on learning-rate warm-up.

    Source: §1 · p. 2

    The norm setup restricts the analysis to cases where the dual norm is Frechet differentiable with p = 2 and C ≥ 1.

    Source: §1.1 · p. 3

Abstract

from arXiv · show

We consider non-convex stochastic optimization using first-order algorithms for which the gradient estimates may have heavy tails. We show that a combination of gradient clipping, momentum, and normalized gradient descent yields convergence to critical points in high-probability with best-known rates for smooth losses when the gradients only have bounded $\mathfrak{p}$th moments for some $\mathfrak{p}\in(1,2]$. We then consider the case of second-order smooth losses, which to our knowledge have not been studied in this setting, and again obtain high-probability bounds for any $\mathfrak{p}$. Moreover, our results hold for arbitrary smooth norms, in contrast to the typical SGD analysis which requires a Hilbert space norm. Further, we show that after a suitable "burn-in" period, the objective value will monotonically decrease for every iteration until a critical point is identified, which provides intuition behind the popular practice of learning rate "warm-up" and also yields a last-iterate guarantee.

1 SGD and Heavy Tails

SGD’s standard convergence analysis relies on bounded variance, while practical gradients may have heavy tails with only bounded pth moments. The section motivates high-probability methods using clipping and momentum, and broadens the analysis beyond Hilbert-space norms.

1 SGD and Heavy Tails

  • Heavy-tailed gradients may have extremely large or effectively infinite variance while retaining a bounded pth moment for some p ∈ (1, 2].Empirical evidence cited in the paper reports heavy-tailed gradients in large NLP models based on attention and transformers.Source: §1 · p. 2
  • Clipping truncates heavy-tailed values, and prior analysis showed clipped SGD can find critical points in expectation.The paper contrasts that expectation-only result with its own high-probability guarantee.Source: §1 · p. 2
  • The proposed clipped normalized SGD with momentum finds an ϵ-critical point with high probability under heavy-tailed gradients.The method improves on prior heavy-tail analysis that provided only an in-expectation guarantee.Source: §1 · p. 2
  • Momentum-based updates extend critical-point analysis beyond Hilbert spaces to arbitrary smooth norms, despite the bias that nonlinear conversion of dual gradient estimates may add.The method avoids relying on the usual unbiasedness argument for SGD.Source: §1 · p. 2
  • For both smooth and second-order smooth cases, the paper gives a high-probability last-iterate guarantee: the final objective value is below that at some approximate critical point.The introduction also connects normalized SGD with momentum to monotonic objective decrease after a burn-in period and to learning-rate warm-up.Source: §1 · p. 2

1.1 Setup and assumptions

  • The setup assumes a Frechet-differentiable objective on a Banach space, an L-smooth objective, and bounded pth moments of stochastic-gradient noise for p ∈ (1, 2].The paper also considers second-order smoothness in its later analysis.Source: §1.1 · p. 3

2 Normalized SGD with Momentum

The section presents clipped, momentum-based normalized SGD for high-probability critical-point guarantees under heavy-tailed gradients. After a burn-in period, the momentum estimate is likely close to the true gradient, supporting monotonic objective decrease and a last-iterate guarantee.

2 Normalized SGD with Momentum

  • Clipping limits the effect of outliers, momentum averages clipped gradients, and normalized updates control bias introduced by the moving average.The analysis is formulated for smooth norms in Banach spaces, where converting a dual gradient estimate to a primal vector can introduce bias.Source: §2 · p. 3
  • Theorem 2 gives a high-probability analogue of the standard in-expectation non-convex result, matching Zhang et al. up to logarithmic factors.The following theorem establishes monotonic decrease after burn-in.Source: §2 · p. 4
  • After burn-in, the objective decreases monotonically with high probability until an approximate critical point is reached.The momentum estimate is then likely to be a high-quality estimate of the true gradient.Source: §2 · p. 4
  • After burn-in, returning the iterate with the smallest momentum-norm identifies a critical point; the objective can increase by only O(T^(1−p/(3p−2))) during burn-in.The paper attributes the post-burn-in progress to the momentum estimate being close to the true gradient.Source: §2 · p. 4
  • Corollary 4 gives a non-asymptotic guarantee for the final iterate, rather than only for a randomly selected iterate.The paper relates this guarantee to practical settings where the final iterate is the one deployed.Source: §2 · p. 7

3 Second-Order Smooth Losses

For second-order smooth losses, the section extends the clipped, normalized momentum approach using implicit gradient transport. It gives results analogous to those for smooth losses, including a last-iterate guarantee.

3 Second-Order Smooth Losses

  • The method uses normalized SGD with momentum and implicit gradient transport, implemented through the NIGT algorithm’s extrapolation procedure.Clipping is added to this procedure; second-order smoothness lets the momentum estimate be less biased than with standard momentum.Source: §3 · p. 7
  • Theorems 5 and 6 provide high-probability bounds under bounded pth moments, L-smoothness, and second-order smoothness.The clipping threshold is set using the moment bound and momentum parameter.Source: §3 · p. 7
  • The section also states a last-iterate guarantee for Algorithm 2, analogous to Corollary 4.Corollary 7 gives the guarantee with probability at least 1 − δ.Source: §3 · p. 8

4 Experiments

The experiments compare burn-in with learning-rate warm-up on BERT pretraining, using Algorithm 1 and LAMB. Burn-in nearly matches warm-up for Algorithm 1, while LAMB performs better with warm-up than with burn-in.

4 Experiments

  • The experiment tests whether burn-in can account for part of warm-up’s success, motivated by a phase when the momentum estimate may be far from the true gradient.The BERT pretraining study uses BERT-Large and also includes LAMB; BERT is cited as having empirical heavy tails.Source: §4 · p. 9
  • Final accuracy was 62.95% with burn-in versus 63.45% with warm-up, supporting the intuition that initial optimization steps may be unhelpful or hurtful.Performance was very poor without or with smaller burn-in values, then improved sharply with enough burn-in, with smaller gains as burn-in increased further.Source: §4 · pp. 9–10
  • For LAMB, warm-up performs considerably better than burn-in, though the gap between warm-up and neither strategy is smaller than for Algorithm 1.The authors note a remaining gap between linear learning-rate increase and the coarser burn-in policy.Source: §4 · p. 10

5 Conclusion

The paper presents clipped, momentum-based normalized SGD algorithms with high-probability rates for finding critical points under heavy-tailed gradients. It also identifies an early burn-in phase and reports that initially setting the learning rate to zero helps training empirically.

5 Conclusion

  • The proposed normalized SGD algorithms combine momentum and gradient clipping to provide high-probability rates for finding critical points with heavy-tailed gradients.As a by-product, the last iterate has a lower function value than a critical point.Source: §5 · p. 10
  • During an early burn-in phase, the optimizer may make poor decisions, and empirically, setting the learning rate to zero initially helps training versus no special treatment.The benefits of the increasing learning-rate schedule used during warm-up remain an open question.Source: §5 · p. 10

A Proofs for Section 3

The appendix supplies proofs for Theorems 5 and 6 under bounded pth-moment gradients and smoothness assumptions. The proofs use a recursion with Hessian-dependent terms that cancel, and the stated bound constants depend on problem parameters but not T or δ.

A Proofs for Section 3

  • Theorem 5 assumes bounded pth moments of stochastic gradients and an objective that is L-smooth and ρ-second-order smooth.Its proof defines a recursion in which the Hessian-dependent terms cancel.Source: A Proofs for Section 3 · p. 14
  • Theorem 5 sets the clipping threshold to G α^(1/p) under α ≤ 1 and gives a bound with probability at least 1 − δ.The hidden constants depend on L, ρ, G, b, s, C, and the objective decrease, but not T or δ.Source: A Proofs for Section 3 · p. 14
  • The proof tracks gradient-estimation errors and uses the recursion for all t ≥ 1, extending it to t = 0 by setting w_0 = w_1.It then unwinds the recursion.Source: A Proofs for Section 3 · p. 14
  • Theorem 6 selects the iteration with the smallest momentum norm among iterations t ≥ T, and its proof follows the case-work argument of Theorem 3.The proof reuses the Theorem 2 argument to obtain a probability-at-least-1 − δ/T bound.Source: A Proofs for Section 3 · pp. 14–16

B A Concentration Bound for Truncated Random Vectors

This section develops concentration bounds for truncated random vectors, first in Hilbert spaces and then in Banach spaces. The proof controls truncation bias and variance before applying a dimension-free Freedman-style concentration bound.

B A Concentration Bound for Truncated Random Vectors

C Dimension-Free Martingale concentration from 1-d concentration

This section constructs a dimension-free Freedman-style martingale concentration bound in a Banach space by reducing to the one-dimensional Freedman inequality. The construction uses a recursively defined scalar sequence whose key property supports the bound.

C Dimension-Free Martingale concentration from 1-d concentration

  • The section presents a Banach-space Freedman-style concentration bound by reduction to the one-dimensional Freedman inequality, while noting it does not improve on existing bounds.The proof is offered for completeness and as a potentially more user-friendly alternative to advanced bounds.Source: C Dimension-Free Martingale… · p. 18
  • The Banach-space norm is required to be Frechet differentiable and satisfy the stated smoothness inequality; Hilbert spaces allow p = 2 and C = 1.These are the conditions used for the recursive construction and its concentration result.Source: C Dimension-Free Martingale… · p. 18
  • The recursively defined scalar sequence has the key property that |s_t| is bounded by ||X_t||.The lemma also gives a bound for the sequence's accumulated sum.Source: C Dimension-Free Martingale… · p. 18

C.1 Warm-up: Hilbert space

This section builds a Hilbert-space concentration bound by applying a scalar Freedman inequality to martingale quantities. The resulting bound holds simultaneously for all iterations with high probability.

C.1 Warm-up: Hilbert space

  • Lemma 11 states the scalar Freedman inequality for bounded martingale differences with conditional variance bounds.The proof invokes the inequality after setting its right-hand side to δ and bounding ε.Source: §C.1 · p. 20
  • With probability at least 1 − 3δ, Lemma 12 bounds the martingale sequence for all k when increments are bounded by R and conditional second moments by σ_t^2.The argument extends the scalar Freedman inequality to Hilbert spaces.Source: §C.1 · p. 20
  • The proof applies scalar Freedman to a martingale difference sequence s_t bounded by ∥X_t∥, yielding a bound with probability at least 1 − 2δ.The proof also controls a second martingale formed from ∥X_t∥^2 and its conditional expectation.Source: §C.1 · p. 20

C.2 Extension to Banach space

The section extends the concentration argument from Hilbert spaces to Banach spaces under a smoothness condition. It gives a high-probability bound for martingale differences with bounded increments and conditional pth moments.

C.2 Extension to Banach space

  • With probability at least 1 − 3δ, Lemma 14 bounds the Banach-space martingale sequence for all k under its bounded-increment and conditional-moment assumptions.The stated range is p ∈ (1, 2], and the space satisfies condition (7).Source: §C.2 · p. 22
  • The proof controls a martingale formed from ∥X_t∥^p and its conditional expectation, then applies Young’s inequality.The passage also states that the bound is obtained when C ≥ 1.Source: §C.2 · pp. 22–23
  • Figure 2 compares warm-up and burn-in strategies using evaluation accuracy, with separate plots for Algorithm 1 and LAMB.The supplied caption identifies the task as BERT-Base masked-language-model pre-training.Source: Figure 2 · p. 23

D Experiments

The BERT-Base experiments find that burn-in is less crucial than in BERT-Large and that the model converges across the tested burn-in durations. The section also reports no convergence issues without warm-up.

D Experiments

  • BERT-Base converges regardless of burn-in duration, and the experiments report no convergence issues without warm-up.Burn-in is less crucial for BERT-Base than for BERT-Large in these experiments.Source: D Experiments · p. 24
  • The experiments use Algorithm 1 with base learning rate 0.5 and batch size 512, selecting the learning rate through a grid search.Burn-in durations tested were 0, 250, 500, 1000, 2000, and 3125 steps; the warm-up baseline uses 3125 linear warm-up steps before polynomial decay.Source: D Experiments · p. 24

E Smoothness Lemmas in Banach Spaces

This section sets out Banach-space definitions and smoothness properties used in the paper’s analysis. It also presents a descent lemma for an update along a unit vector aligned with a dual vector.

E Smoothness Lemmas in Banach Spaces

  • A Banach space’s dual consists of bounded linear functions, and its dual norm is defined by the supremum over unit-norm vectors.The paper writes ⟨v, w⟩ for applying a dual vector v to a vector w.Source: E Smoothness Lemmas in Banach Spaces · p. 24
  • The section defines the gradient as a bounded linear operator in the dual space and the Hessian as a map from the space to its dual.It also states that the paper considers spaces whose dual norm satisfies a smoothness condition.Source: E Smoothness Lemmas in Banach Spaces · p. 24
  • The paper assumes a Frechet-differentiable function whose gradient is L-Lipschitz in the dual norm relative to the primal norm.Proposition 15 states the resulting smoothness property for all x and y in the Banach space.Source: E Smoothness Lemmas in Banach Spaces · p. 24
  • The second-order smoothness condition bounds Hessian variation, and Propositions 16 and 17 state consequences for all x and y.The condition bounds ∥(∇²F(x) − ∇²F(y))v∥⋆ by ρ∥v∥∥x − y∥.Source: E Smoothness Lemmas in Banach Spaces · p. 25
  • Lemma 1 generalizes a smoothness lemma to more general norms and improves its constants.It considers an L-smooth function, an update w′ = w − ηg, and an error ϵ = g⋆ − ∇F(w).Source: E Smoothness Lemmas in Banach Spaces · p. 26

Related papers

Loading 2106.14343v2…