A Fast Single-Loop Primal-Dual Algorithm for Non-Convex Functional Constrained Optimization

Abstract: Non-convex functional constrained optimization problems have gained substantial attention in machine learning and signal processing. This paper develops a new primal-dual algorithm for solving this class of problems. The algorithm is based on a novel form of the Lagrangian function, termed {\em Proximal-Perturbed Augmented Lagrangian}, which enables us to develop an efficient and simple first-order algorithm that converges to a stationary solution under mild conditions. Our method has several key features of differentiation over existing augmented Lagrangian-based methods: (i) it is a single-loop algorithm that does not require the continuous adjustment of the penalty parameter to infinity; (ii) it can achieves an improved iteration complexity of $\widetilde{\mathcal{O}}(1/\epsilon^2)$ or at least ${\mathcal{O}}(1/\epsilon^{2/q})$ with $q \in (2/3,1)$ for computing an $\epsilon$-approximate stationary solution, compared to the best-known complexity of $\mathcal{O}(1/\epsilon^3)$; and (iii) it effectively handles functional constraints for feasibility guarantees with fixed parameters, without imposing boundedness assumptions on the dual iterates and the penalty parameters. We validate the effectiveness of our method through numerical experiments on popular non-convex problems.