[Paper Review] Convergence of a Stochastic Gradient Method with Momentum for Non-Smooth Non-Convex Optimization

Vien V. Mai, Mikael Johansson|arXiv (Cornell University)|Feb 13, 2020

Stochastic Gradient Optimization Techniques4 citations

TL;DR

This paper establishes the first convergence rate for a stochastic subgradient method with Polyak momentum on non-smooth, non-convex, and constrained optimization problems. By constructing a novel Lyapunov function, it achieves an O(1/√K) convergence rate to a stationary point without tuning the momentum parameter, extending known complexity bounds to the constrained and non-smooth case.

ABSTRACT

Stochastic gradient methods with momentum are widely used in applications and at the core of optimization subroutines in many popular machine learning libraries. However, their sample complexities have not been obtained for problems beyond those that are convex or smooth. This paper establishes the convergence rate of a stochastic subgradient method with a momentum term of Polyak type for a broad class of non-smooth, non-convex, and constrained optimization problems. Our key innovation is the construction of a special Lyapunov function for which the proven complexity can be achieved without any tuning of the momentum parameter. For smooth problems, we extend the known complexity bound to the constrained case and demonstrate how the unconstrained case can be analyzed under weaker assumptions than the state-of-the-art. Numerical results confirm our theoretical developments.

Motivation & Objective

• To close the gap in theoretical understanding of stochastic momentum methods for non-smooth, non-convex problems.
• To establish sample complexity bounds for stochastic subgradient methods with Polyak momentum beyond convex or smooth settings.
• To develop a convergence analysis framework that does not require tuning the momentum parameter.
• To extend known complexity results to constrained, non-smooth, and non-convex optimization problems.

Proposed method

• The method uses a stochastic subgradient with Polyak-type momentum, defined by recursive updates of the iterate and momentum vector.
• A special Lyapunov function is constructed to analyze the convergence, combining the objective function and momentum-related terms.
• The analysis leverages the concept of ρ-weakly convex functions, where f(x) + ρ‖x‖² is convex, to handle non-smoothness.
• Key inequalities are derived using norm expansions and convexity properties of the regularized objective Fλ(x) = f(x) + λ⁻¹‖x‖².
• The convergence rate is derived by summing inequalities over iterations and bounding the expected norm of the gradient of Fλ.
• A uniform random index k* is used to analyze the expected gradient norm, leading to a high-probability convergence bound.

Experimental results

Research questions

• RQ1 Can stochastic gradient methods with momentum achieve convergence rates in non-smooth, non-convex problems?
• RQ2 Is it possible to establish sample complexity bounds for momentum-based methods without assuming smoothness or convexity?
• RQ3 Can a Lyapunov function be constructed to ensure convergence without tuning the momentum parameter?
• RQ4 How does the momentum method perform in the constrained, non-smooth, non-convex setting compared to unconstrained cases?
• RQ5 Can the analysis be extended to weakly convex functions using regularized objectives?

Key findings

• The method achieves an O(1/√K) convergence rate to a stationary point, measured by the expected squared norm of the gradient of the regularized objective Fλ.
• The convergence bound holds without tuning the momentum parameter, relying on the constructed Lyapunov function.
• For smooth problems, the analysis extends known complexity bounds to the constrained case under weaker assumptions than prior work.
• The result is tight in the sense that the convergence rate matches known bounds for smooth, convex problems.
• Numerical results confirm the theoretical convergence behavior and validate the effectiveness of the momentum method in non-smooth, non-convex settings.