[Paper Review] A Diffusion Approximation Theory of Momentum SGD in Nonconvex Optimization

Tianyi Liu, Zhehui Chen+2 more|arXiv (Cornell University)|Feb 14, 2018

Stochastic Gradient Optimization Techniques38 references4 citations

TL;DR

This paper develops a diffusion approximation theory for Momentum Stochastic Gradient Descent (MSGD) in nonconvex optimization, showing that momentum accelerates escape from strict saddle points by enhancing exploration but hinders convergence near local optima without step size or momentum annealing. The analysis reveals that MSGD with small step size behaves like VSGD with a larger effective step size, explaining its empirical success in deep learning.

ABSTRACT

Momentum Stochastic Gradient Descent (MSGD) algorithm has been widely applied to many nonconvex optimization problems in machine learning, e.g., training deep neural networks, variational Bayesian inference, and etc. Despite its empirical success, there is still a lack of theoretical understanding of convergence properties of MSGD. To fill this gap, we propose to analyze the algorithmic behavior of MSGD by diffusion approximations for nonconvex optimization problems with strict saddle points and isolated local optima. Our study shows that the momentum helps escape from saddle points, but hurts the convergence within the neighborhood of optima (if without the step size annealing or momentum annealing). Our theoretical discovery partially corroborates the empirical success of MSGD in training deep neural networks.

Motivation & Objective

• To address the lack of theoretical understanding of Momentum SGD (MSGD) in nonconvex optimization, especially its empirical success in training deep neural networks.
• To analyze the role of momentum in escaping strict saddle points and converging to local optima in nonconvex landscapes with isolated optima and strict saddle points.
• To establish a theoretical link between MSGD's behavior and its performance in practice, particularly in avoiding sharp local minima.
• To extend existing convergence results for VSGD and MSGD by incorporating diffusion approximations and weak convergence to SDEs for nonconvex problems.
• To explain why MSGD can use larger effective step sizes than VSGD while maintaining generalization performance, despite theoretical limitations in convergence near optima.

Proposed method

• Use continuous-time diffusion approximation via weak convergence to model the discrete MSGD trajectory, transforming the algorithm into a stochastic differential equation (SDE) in the limit as step size → 0.
• Apply interpolation techniques to map the discrete MSGD iterates into a continuous-time process, enabling weak convergence analysis to an ODE and SDE under nonconvexity.
• Analyze global dynamics via ODE approximation in non-stationary regions where gradient magnitude dominates stochastic variance, showing momentum improves descent speed by a factor of $\frac{1}{1-\mu}$.
• Study local behavior near stationary points using normalized error sequences, showing that the continuous interpolation converges weakly to a solution of an SDE with non-degenerate diffusion coefficient.
• Leverage the strict saddle property and isolated local optima assumption to characterize the algorithm’s behavior in both escaping and converging regimes.
• Extend the analysis to connected local optima by using rotationally invariant metrics such as principal angles in low-rank matrix problems, enabling bounded error characterization.

Experimental results

Research questions

• RQ1 How does momentum in MSGD affect the escape from strict saddle points in nonconvex optimization?
• RQ2 Why does MSGD outperform VSGD in practice despite theoretical convergence guarantees being similar?
• RQ3 What is the impact of momentum on convergence near local optima, particularly when step size is not annealed?
• RQ4 How does the effective step size of MSGD compare to VSGD under the same hyperparameters, and what explains the observed performance gap?
• RQ5 Can the diffusion approximation framework explain why MSGD tends to converge to flatter, better-generalizing minima in deep learning?

Key findings

• Momentum in MSGD accelerates escape from strict saddle points by enhancing descent direction exploration, effectively increasing the algorithm’s progress per step by a factor of $\frac{1}{1-\mu}$ in non-stationary regions.
• Near local optima, where stochastic gradient variance dominates, momentum degrades convergence unless step size or momentum is annealed, indicating a trade-off between acceleration and stability.
• MSGD with step size $\eta$ and momentum $\mu$ behaves asymptotically like VSGD with effective step size $\frac{\eta}{1-\mu}$, explaining its ability to use larger effective learning rates.
• The normalized error sequence of MSGD converges weakly to a solution of an SDE around stationary points, enabling characterization of the algorithm’s local behavior and noise-induced exploration.
• MSGD tends to avoid sharp local optima due to higher variance induced by momentum, which encourages exploration out of small basins of attraction—aligning with empirical observations on generalization.
• The theory explains why MSGD can achieve better generalization than VSGD with the same effective step size, as it explores flatter minima more effectively due to increased noise from momentum.