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

2025/06/03 by Xingyu Chen, Bokun Wang, Chen, Xingyu +7
Computer Science · Mathematics · #Advanced Optimization Algorithms Research #FOS: Computer and information sciences #Geochemistry and Geologic Mapping #Machine Learning (cs.LG) #Stochastic Gradient Optimization Techniques

paper · pdf · doi:10.48550/arxiv.2506.02504

openalex publication_date 2025/06/03 · openalex created_date 2025/10/14 · openalex updated_date 2026/07/28

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: (1) a high iteration complexity of O(1/ε⁶) under the assumption that the stochastic inner functions are Lipschitz continuous in expectation; (2) reliance on vanilla SGD-type updates, which are not suitable for deep learning applications. Our main contributions are two fold: (i) We propose stochastic momentum methods tailored for non-smooth FCCO that come with provable convergence guarantees; (ii) We establish a new state-of-the-art iteration complexity of O(1/ε⁵). 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/ε⁵) for finding an (nearly) ε-level KKT solution. Experiments on three tasks demonstrate the effectiveness of the proposed algorithms.

Citations

• From stability of Langevin diffusion to convergence of proximal MCMC for non-log-concave sampling
• Single-loop Algorithms for Stochastic Non-convex Optimization with Weakly-Convex Constraints
• A single-loop SPIDER-type stochastic subgradient method for expectation-constrained nonconvex nonsmooth optimization
• Complexity of Single Loop Algorithms for Nonlinear Programming with Stochastic Objective and Constraints
• Non-Smooth Weakly-Convex Finite-sum Coupled Compositional Optimization
• On Squared-Variable Formulations
• Not All Semantics are Created Equal: Contrastive Self-supervised Learning with Automatic Temperature Individualization
• Oracle Complexity of Single-Loop Switching Subgradient Methods for Non-Smooth Weakly Convex Functional Constrained Optimization
• Stochastic Inexact Augmented Lagrangian Method for Nonconvex Expectation Constrained Optimization
• First-Order Methods for Nonsmooth Nonconvex Functional Constrained Optimization with or without Slater Points
• Multi-block-Single-probe Variance Reduced Estimator for Coupled Compositional Optimization
• Algorithmic Foundations of Empirical X-risk Minimization
• Finite-Sum Coupled Compositional Stochastic Optimization: Theory and Applications
• Provable Stochastic Optimization for Global Contrastive Learning: Small Batch Does Not Harm Performance
• LAION-400M: Open Dataset of CLIP-Filtered 400 Million Image-Text Pairs
• Unified Convergence Analysis for Adaptive Optimization with Moving Average Estimator
• Stochastic Optimization of Areas Under Precision-Recall Curves with Provable Convergence
• Oracle Complexity in Nonsmooth Nonconvex Optimization
• Learning Transferable Visual Models From Natural Language Supervision
• WILDS: A Benchmark of in-the-Wild Distribution Shifts
• PAGE: A Simple and Optimal Probabilistic Gradient Estimator for Nonconvex Optimization
• Rate-improved Inexact Augmented Lagrangian Method for Constrained Nonconvex Optimization
• A Primal-Dual Smoothing Framework for Max-Structured Non-Convex Optimization
• DistilBERT, a distilled version of BERT: smaller, faster, cheaper and lighter
• Multi-Level Composite Stochastic Optimization via Nested Variance Reduction
• Stochastic First-order Methods for Convex and Nonconvex Functional Constrained Optimization
• Quadratically Regularized Subgradient Methods for Weakly Convex Optimization with Weakly Convex Constraints
• From Detection of Individual Metastases to Classification of Lymph Node Status at the Patient Level: The CAMELYON17 Challenge
• BDD100K: A Diverse Driving Dataset for Heterogeneous Multitask Learning
• BDD100K: A Diverse Driving Dataset for Heterogeneous Multitask Learning
• Stochastic model-based minimization of weakly convex functions
• Proximally Guided Stochastic Subgradient Method for Nonsmooth, Nonconvex Problems
• Densely Connected Convolutional Networks
• Accelerating Stochastic Composition Optimization
• Deep Residual Learning for Image Recognition
• Deep Learning Face Attributes in the Wild
• Stochastic Compositional Gradient Descent: Algorithms for Minimizing Compositions of Expected-Value Functions
• Smooth minimization of non-smooth functions

Related

• Finite-Sum Coupled Compositional Stochastic Optimization: Theory and Applications
• Non-Smooth Weakly-Convex Finite-sum Coupled Compositional Optimization
• A Near-Optimal Single-Loop Stochastic Algorithm for Convex Finite-Sum Coupled Compositional Optimization
• Stochastic Recursive Momentum Method for Non-Convex Compositional Optimization
• Katyusha Acceleration for Convex Finite-Sum Compositional Optimization