Projected Stochastic Momentum Methods for Nonlinear Equality-Constrained Optimization for Machine Learning
Qi Wang, Christian Piermarini, Yunlang Zhu, Frank E. Curtis
TL;DR
This work tackles nonlinear equality-constrained stochastic optimization in ML by extending stochastic momentum methods to constrained settings using projected gradient estimates within a stochastic Newton-SQP framework. It introduces two algorithms—Projected Stochastic Heavy-Ball SQP and Projected Stochastic Adam SQP—with convergence guarantees that match unconstrained counterparts, achieved by updating momentum with projected gradient estimates $P_k g_k$ rather than raw gradients. Theoretical results include sublinear convergence bounds (e.g., an $O(1/K)$ rate) on the average merit decrease, and the heavy-ball variant provides a merit-function-based guarantee under standard assumptions. Empirical studies on informed supervised learning problems show the constrained approach can outperform regularized unconstrained Adam and projection-less variants, highlighting the practical benefits of hard constraints in learning.
Abstract
Two algorithms are proposed, analyzed, and tested for solving continuous optimization problems with nonlinear equality constraints. Each is an extension of a stochastic momentum-based method from the unconstrained setting to the setting of a stochastic Newton-SQP-type algorithm for solving equality-constrained problems. One is an extension of the heavy-ball method and the other is an extension of the Adam optimization method. Convergence guarantees for the algorithms for the constrained setting are provided that are on par with state-of-the-art guarantees for their unconstrained counterparts. A critical feature of each extension is that the momentum terms are implemented with projected gradient estimates, rather than with the gradient estimates themselves. The significant practical effect of this choice is seen in an extensive set of numerical experiments on solving informed supervised machine learning problems. These experiments also show benefits of employing a constrained approach to supervised machine learning rather than a typical regularization-based approach.