Stochastic Momentum Methods for Non-smooth Non-Convex Finite-Sum Coupled Compositional Optimization

Xingyu Chen, Bokun Wang, Ming Yang, Qihang Lin, Tianbao Yang

finite-sum coupled compositional optimization non-convex constrained optimization

Abstract

Finite-sum Coupled Compositional Optimization (FCCO), characterized by its coupled compositional objective structure, emerges as an important optimization paradigm for addressing a wide range of machine learning problems. In this paper, we focus on a challenging class of non-convex non-smooth FCCO, where the outer functions are non-smooth weakly convex or convex and the inner functions are smooth or weakly convex.

Existing state-of-the-art result face two key limitations:

a high iteration complexity of $O (1/ ϵ^{6})$ under the assumption that the stochastic inner functions are Lipschitz continuous in expectation;
reliance on vanilla SGD-type updates, which are not suitable for deep learning applications.

Our main contributions are two fold:

We propose stochastic momentum methods tailored for non-smooth FCCO that come with provable convergence guarantees;
We establish a new state-of-the-art iteration complexity of $O (1/ ϵ^{5})$ .

Moreover, we apply our algorithms to multiple inequality constrained non-convex optimization problems involving smooth or weakly convex functional inequality constraints. By optimizing a smoothed hinge penalty based formulation, we achieve a new state-of-the-art complexity of

O (1/ ϵ^{5})

for finding an (nearly)

ϵ

-level KKT solution. Experiments on three tasks demonstrate the effectiveness of the proposed algorithms.