PaperlayerView exact revision on arXiv

Source-linked AI summary

Momentum Improves Normalized SGD

Ashok Cutkosky, Harsh Mehta

arXiv:2002.03305v2cs.LGmath.OCstat.ML

TL;DR

The paper addresses limitations of non-convex SGD and normalized updates, including worst-case rates, variance sensitivity, and dimension dependence. It analyzes normalized SGD with momentum, introduces an implicit-gradient-transport modification and an adaptive variant, and reports dimension-free convergence results alongside strong performance on BERT and ResNet-50. The method matches the performance of task-specific optimizers while using less memory than Adam or LAMB.

  • Problem

    Non-convex optimization for large-scale deep learning lacks normalized-SGD guarantees that avoid small-variance or large-batch requirements and dimension-dependent rates.

  • Method

    The paper combines normalized SGD with momentum, modifies momentum using implicit gradient transport for second-order smooth objectives, and proposes an adaptive variant.

  • Results

    The method attains dimension-free O(1/T^2/7) convergence under second-order smoothness, automatically improves with low gradient variance, and matches task-specific methods on BERT pretraining and ResNet-50.

  • Takeaways & Limitations

    A single momentum-based normalized method performs well on both BERT pretraining and ResNet-50, where Adam and SGD are typically used respectively.

  • Takeaways & Limitations

    The adaptive analysis imposes an upper bound on F and uses an extra gradient sample per iteration; the authors identify removing these requirements as open problems.

Abstract

from arXiv · show

We provide an improved analysis of normalized SGD showing that adding momentum provably removes the need for large batch sizes on non-convex objectives. Then, we consider the case of objectives with bounded second derivative and show that in this case a small tweak to the momentum formula allows normalized SGD with momentum to find an $ε$-critical point in $O(1/ε^{3.5})$ iterations, matching the best-known rates without accruing any logarithmic factors or dependence on dimension. We also provide an adaptive method that automatically improves convergence rates when the variance in the gradients is small. Finally, we show that our method is effective when employed on popular large scale tasks such as ResNet-50 and BERT pretraining, matching the performance of the disparate methods used to get state-of-the-art results on both tasks.

1 Non-Convex Stochastic Optimization

Non-convex stochastic optimization targets critical points in high-dimensional, large-scale settings where SGD is effective but has worst-case convergence limits. The paper addresses variance sensitivity and dimension dependence by combining normalized updates with momentum, implicit gradient transport, and adaptation.

  • Motivation: Deep-learning objectives are high-dimensional, large-scale, and non-convex, motivating dimension-free methods for finding critical points.Modern neural networks can have billions of parameters and are trained on datasets containing millions of examples.
  • Limitations of SGD: O(1/T^1/4) is optimal for SGD in the worst case, so faster rates require additional assumptions or interpolation across problem regimes.Adaptive methods can improve when stochastic-gradient variance is small.
  • Momentum: Momentum maintains a running average of past gradients, yet prior theoretical bounds did not establish a benefit over base SGD.Its empirical success motivated analysis of whether it improves stability or conditioning in normalized methods.
  • Normalized Updates: Normalized updates use gradient direction while reducing reliance on gradient magnitude, but existing analyses require very small variance or extremely large batches.Normalization can amplify small estimation errors, making the normalized direction unreliable.
  • Adaptivity and Experiments: The adaptive algorithm improves its convergence guarantee automatically when stochastic-gradient variance is small, and the method matches strong baselines on BERT pretraining and ResNet-50.The method uses substantially less memory than Adam because it keeps only momentum as a significant slot variable.

2 Normalized SGD with Momentum

Normalized SGD with momentum matches ordinary SGD’s optimal convergence rate on non-convex objectives without requiring large batch sizes. The analysis explains how momentum controls normalization error and establishes the resulting gradient-norm guarantee.

  • Theorem 1: Momentum lets normalized SGD match ordinary SGD’s optimal convergence rate without requiring large batch sizes.Theorem 1 analyzes the momentum update under assumptions (A1)–(A3).
  • Motivation: Large batches are otherwise needed because normalization can amplify stochastic-gradient error when that error exceeds the true gradient.A one-dimensional example shows the unnormalized gradient estimate can be biased away from a critical point.
  • Theorem 1: Theorem 1 uses constant η and βt = 1 − α, with α selected as a function of the horizon and problem parameters.The theorem initializes momentum with one stochastic gradient and applies the normalized recurrences throughout T iterations.
  • Proof strategy: The proof bounds progress using the smoothness of F and the accumulated momentum estimation error.The analysis separates cases according to whether the true gradient dominates the estimation error, then sums a telescoping inequality.

3 Faster Convergence with Second-Order Smoothness

Under bounded second derivative, a modified momentum scheme uses implicit gradient transport to reduce the gradient-transport remainder from linear to quadratic in the step size. This yields an average expected gradient-norm dependence of O(1/T^2/7).

  • Second-order smoothness: Quadratic remainder control, ∥Z(a, b)∥ ≤ ρ∥a − b∥^2, replaces the linear smoothness bound and enables faster convergence.Taylor’s theorem supplies the quadratic bound under assumption (A4).
  • Algorithm 1: NIGT introduces an auxiliary point xt, transports gradients implicitly, and normalizes the momentum update to produce the next iterate.The auxiliary point can be recomputed from consecutive iterates, while mt combines the previous momentum with a stochastic gradient at xt.
  • Algorithm 1: The momentum setting β = 1 − O(T^-4/7) extends implicit gradient transport from convex quadratics with constant Hessian to non-convex objectives.The paper contrasts this setting with βt = t/(t+1), which exactly recovers the earlier implicit-gradient-transport update.
  • Theorem 3: O(1/T^2/7) is the resulting dependence of the average expected gradient norm on the iteration horizon.The rate is obtained for Algorithm 1 under the stated smoothness, variance, and second-order assumptions.
  • Proof strategy: The proof cancels the Hessian terms through the auxiliary-point construction and tracks the remaining non-constant-Hessian error in Z.The resulting recursion is bounded using the quadratic remainder estimate before applying the earlier lemma.

4 Adaptive Algorithm

The adaptive algorithm varies its learning rate and momentum using observed gradients, automatically improving its convergence guarantee when gradient variance is small without knowing that variance in advance. It uses two independent gradient samples per point and assumes bounded function values above.

  • Adaptive schedule: The adaptive schedule adjusts η and β on the fly from observed gradients and improves convergence automatically when σ is small.The schedule does not require prior knowledge of σ.
  • Adaptive schedule: Two independent gradient samples are taken at each auxiliary point xt to estimate gradient variability within the stochastic-oracle model.The second sample is denoted ∇f(xt, ξ′t).
  • Theorem 4: The adaptive analysis additionally assumes bounded stochastic-gradient norms and an upper bound M on F(w).These assumptions accompany (A1)–(A4) in Theorem 4.
  • Algorithm 2: Algorithm 2 updates Gt using the squared difference between the two gradient samples and a time-dependent g^2 correction.It then sets βt = 1 − αt, updates the transported momentum, and normalizes the step.
  • Theorem 4: The theorem’s ˜O notation hides constants depending on R, G, and M and includes a factor of log(T).This qualification describes the displayed adaptive convergence guarantee.

5 Experiments

The experiments evaluate NIGT on BERT pretraining and ResNet-50, using practical learning-rate heuristics and comparing against the stronger task-specific baselines. NIGT slightly outperforms Adam on BERT validation accuracy and SGD on ResNet-50 top-1 accuracy.

  • Experimental setup: NIGT is evaluated on BERT pretraining and ResNet-50 using per-layer gradient normalization and momentum β=0.9.The implementation also uses task-specific learning-rate schedules and weight-norm scaling.
  • Baselines: NIGT is compared with Adam for BERT and SGD for ResNet-50 because these are the stronger commonly used baselines for the respective tasks.The paper notes that Adam is typically used for BERT, while SGD is used for ResNet-50 because Adam performs poorly there.
  • BERT pretraining: 70.91 vs 70.76 masked language modeling validation accuracy is achieved by NIGT versus Adam on BERT.BERT-Base is trained for 500k steps with batch size 256; the reported comparison uses the validation accuracy shown in Figure 1.
  • ResNet-50: 76.37% vs 76.2% top-1 validation accuracy is achieved by NIGT versus SGD on ResNet-50 over ImageNet after 90 epochs.The ResNet-50 experiment uses batch size 1024 and compares against SGD with momentum and a standard warmup-and-decay schedule.

6 Conclusion and Future Work

The paper concludes that momentum improves normalized SGD theoretically and practically, including dimension-free second-order-smooth rates and performance across BERT and ResNet-50. It also identifies unresolved assumptions, algorithmic questions, and the empirical role of learning-rate heuristics.

  • Theoretical conclusions: O(1/ε^3.5) iterations are sufficient for normalized SGD with momentum to find an ε-critical point under second-order smoothness, matching the best-known dimension-free rate.The conclusion presents this as the paper’s main second-order-smooth theoretical result.
  • Empirical conclusions: NIGT matches the performance of the methods used for state-of-the-art training of ResNet-50 on ImageNet and BERT pretraining.These tasks commonly use different optimizers—SGD and Adam, respectively—yet NIGT performs well on both.
  • Adaptive method: The adaptive algorithm automatically improves its convergence guarantee when the underlying gradient distribution has low variance.This extends the theoretical analysis beyond fixed-rate behavior.
  • Interpretation: The results suggest that momentum can accelerate convergence when second-order smoothness is assumed, offering theoretical intuition for its practical success.Prior analyses had not demonstrated a significant theoretical benefit for momentum.
  • Open problems: The adaptive analysis may unnecessarily require an upper bound on F and an extra gradient sample per iteration.These are identified as open problems rather than established failures of the method.
  • Practical limitations: The implementation relies on linear warm-up, polynomial decay, and weight-norm learning-rate scaling, which performed better than the theoretically motivated adaptive algorithm.The authors leave the principles underlying these heuristics for future investigation.

A Proof of Theorem 4

The proof of Theorem 4 establishes a convergence guarantee for an adaptive normalized SGD method under bounded-gradient and objective assumptions. It combines recursive momentum-error analysis, potential-function arguments, and summation of per-iteration bounds, with logarithmic factors hidden in the final rate notation.

  • Theorem assumptions: Theorem 4 assumes conditions (A1)–(A4), uniformly bounded stochastic gradients, and an objective bounded above by M.The theorem states that Algorithm 2 then guarantees the displayed convergence bound.
  • Rate qualification: The final ˜O rate hides constants depending on R, G, and M as well as a factor of log(T).This is an explicit qualification on the theorem’s asymptotic notation.
  • Momentum-error analysis: The proof analyzes the recursive momentum-noise update using a potential-function approach because direct unrolling is not tractable.The recurrence includes transported stochastic-gradient terms and accumulated error terms.
  • Error control: The analysis controls cross terms with a quadratic inequality and uses the stochastic error’s lack of correlation with quantities independent of the current sample.These steps produce bounds on the recursive error contributions.
  • Adaptive-rate bounds: The proof bounds the adaptive quantities through monotonicity, convexity, Jensen’s inequality, and the definitions of G_t and η_t.It also uses that η_t is non-increasing and that α_t≤1.
  • Conclusion of the proof: Theorem 4 follows after assembling the intermediate inequalities, summing over t, and taking expectations.The displayed final guarantee is obtained from the accumulated per-iteration bounds.
Loading 2002.03305v2…