[Paper Review] Unified Convergence Analysis of Stochastic Momentum Methods for Convex and Non-convex Optimization

Tianbao Yang, Qihang Lin+1 more|arXiv (Cornell University)|Apr 12, 2016

Stochastic Gradient Optimization Techniques7 references90 citations

TL;DR

This paper presents a unified convergence analysis framework for stochastic momentum methods—stochastic heavy-ball (HB) and stochastic Nesterov’s accelerated gradient (S-NAG)—in both convex and non-convex optimization. It establishes that both methods achieve an expected convergence rate of $O(1/ frac{1}{\sqrt{t}})$ for objective value (convex) and gradient norm (non-convex), with the framework revealing a continuous transition from gradient descent to NAG to HB via a free parameter, explaining observed training and testing error dynamics in deep learning.

ABSTRACT

Recently, {\it stochastic momentum} methods have been widely adopted in training deep neural networks. However, their convergence analysis is still underexplored at the moment, in particular for non-convex optimization. This paper fills the gap between practice and theory by developing a basic convergence analysis of two stochastic momentum methods, namely stochastic heavy-ball method and the stochastic variant of Nesterov's accelerated gradient method. We hope that the basic convergence results developed in this paper can serve the reference to the convergence of stochastic momentum methods and also serve the baselines for comparison in future development of stochastic momentum methods. The novelty of convergence analysis presented in this paper is a unified framework, revealing more insights about the similarities and differences between different stochastic momentum methods and stochastic gradient method. The unified framework exhibits a continuous change from the gradient method to Nesterov's accelerated gradient method and finally the heavy-ball method incurred by a free parameter, which can help explain a similar change observed in the testing error convergence behavior for deep learning. Furthermore, our empirical results for optimizing deep neural networks demonstrate that the stochastic variant of Nesterov's accelerated gradient method achieves a good tradeoff (between speed of convergence in training error and robustness of convergence in testing error) among the three stochastic methods.

Motivation & Objective

• To bridge the gap between the widespread use of stochastic momentum methods in deep learning and their underdeveloped theoretical convergence analysis.
• To provide a unified theoretical framework that reveals the structural similarities and differences between stochastic heavy-ball, stochastic Nesterov’s method, and stochastic gradient descent.
• To establish convergence guarantees for stochastic momentum methods in both convex and non-convex settings with Lipschitz continuous gradients.
• To explain the observed convergence behavior in deep learning—particularly the trade-off between training speed and testing error robustness—using a continuous parameterized family of momentum methods.
• To serve as a foundational reference for future development and comparison of stochastic momentum algorithms.

Proposed method

• Proposes a unified framework that parameterizes momentum methods via a free parameter, enabling a continuous transition from gradient descent to Nesterov’s accelerated method to heavy-ball method.
• Analyzes the stochastic heavy-ball (HB) and stochastic Nesterov’s accelerated gradient (S-NAG) methods under the same theoretical setup, using stochastic gradients with bounded variance.
• Derives convergence rates for the expected objective value in convex optimization and expected gradient norm in non-convex optimization using Lyapunov function analysis and recursive inequalities.
• Establishes that for convex, Lipschitz continuous functions, the objective value converges at rate $O(1/\sqrt{t})$ in expectation.
• For non-convex functions with Lipschitz continuous gradients, the expected norm of the gradient converges at rate $O(1/\sqrt{t})$.
• Uses a parameterized update rule that generalizes both HB and NAG, showing that the momentum term evolves continuously with the parameter, mirroring empirical behavior in deep learning.

Experimental results

Research questions

• RQ1 How do stochastic momentum methods like HB and S-NAG converge in convex and non-convex settings, and what are their theoretical convergence rates?
• RQ2 What is the structural relationship between stochastic heavy-ball, stochastic Nesterov’s method, and stochastic gradient descent in terms of algorithmic design and convergence behavior?
• RQ3 Can a unified theoretical framework explain the observed transition in training and testing error dynamics in deep neural network optimization?
• RQ4 Why does the stochastic variant of Nesterov’s method often achieve a better trade-off between training speed and generalization robustness than HB and SGD?
• RQ5 How does the momentum parameter influence the convergence trajectory and stability of stochastic momentum methods?

Key findings

• The stochastic heavy-ball method and the stochastic variant of Nesterov’s accelerated gradient method both achieve an expected convergence rate of $O(1/\sqrt{t})$ for the objective value in convex optimization with Lipschitz continuous functions.
• For non-convex optimization with Lipschitz continuous gradients, both methods achieve an expected convergence rate of $O(1/\sqrt{t})$ for the norm of the gradient.
• The proposed unified framework reveals a continuous transformation from gradient descent to Nesterov’s method to heavy-ball method as the momentum parameter varies, explaining similar transitions in testing error dynamics observed in deep learning.
• Empirical results show that the stochastic Nesterov’s method achieves a favorable trade-off between fast training error reduction and robust testing error convergence, outperforming both stochastic HB and stochastic gradient descent in generalization.
• The theoretical analysis is general and does not rely on strong convexity or bounded gradients, making it applicable to a wide range of non-convex deep learning problems.
• The framework provides a theoretical basis for understanding the behavior of momentum methods in deep learning and serves as a benchmark for future algorithmic developments.