A Momentum-based Stochastic Algorithm for Linearly Constrained Nonconvex Optimization

TL;DR

This paper introduces a momentum-based stochastic augmented Lagrangian algorithm for linearly constrained nonconvex optimization, achieving efficient convergence with only one gradient evaluation per iteration.

Contribution

It presents a novel momentum-based method with Polyak-type gradient estimator that guarantees convergence and improves efficiency over existing methods.

Findings

Achieves an $oldsymbol{ ext{O}( ext{epsilon}^{-4})}$ complexity for $oldsymbol{ ext{epsilon}}$-stationary solutions.

Requires only one stochastic gradient evaluation per iteration.

Demonstrates competitive iteration complexity and better wall-clock time in experiments.

Abstract

This paper studies a stochastic algorithm for linearly constrained nonconvex optimization, where the objective function is smooth but only unbiased stochastic gradients with bounded variance are available. We propose a momentum-based augmented Lagrangian method that employs a Polyak-type gradient estimator and requires only one stochastic gradient evaluation per iteration. Under the standard stochastic oracle model and the smoothness condition of the expected objective, we establish a convergence guarantee in terms of the first-order KKT residual of the original constrained problem. In particular, the proposed method computes an $\epsilon$-stationary solution in expectation within $O(\epsilon^{-4})$ stochastic gradient evaluations. Numerical experiments further show that the proposed method achieves competitive iteration complexity and improved wall-clock efficiency compared with representative recursive-momentum baselines.