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

arXiv: https://arxiv.org/abs/2506.02504 · HTML: https://mlsanity.com/paper/2506.02504/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 38 of its 56 references that are in the map.
- First posted on arXiv 3 Jun 2025.

## Builds on (most cited first, 30 of 38)

- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [Learning Transferable Visual Models From Natural Language Supervision](https://mlsanity.com/paper/2103.00020/) (2021)
- [Densely Connected Convolutional Networks](https://mlsanity.com/paper/1608.06993/) (2016)
- [DistilBERT, a distilled version of BERT: smaller, faster, cheaper and lighter](https://mlsanity.com/paper/1910.01108/) (2019)
- [Deep Learning Face Attributes in the Wild](https://mlsanity.com/paper/1411.7766/) (2014)
- [LAION-400M: Open Dataset of CLIP-Filtered 400 Million Image-Text Pairs](https://mlsanity.com/paper/2111.02114/) (2021)
- [WILDS: A Benchmark of in-the-Wild Distribution Shifts](https://mlsanity.com/paper/2012.07421/) (2020)
- [BDD100K: A Diverse Driving Dataset for Heterogeneous Multitask Learning](https://mlsanity.com/paper/1805.04687/) (2018)
- [Stochastic model-based minimization of weakly convex functions](https://mlsanity.com/paper/1803.06523/) (2018)
- [PAGE: A Simple and Optimal Probabilistic Gradient Estimator for Nonconvex Optimization](https://mlsanity.com/paper/2008.10898/) (2020)
- [Stochastic Compositional Gradient Descent: Algorithms for Minimizing Compositions of Expected-Value Functions](https://mlsanity.com/paper/1411.3803/) (2014)
- [Weakly-Convex Concave Min-Max Optimization: Provable Algorithms and Applications in Machine Learning](https://mlsanity.com/paper/1810.02060/) (2018)
- [Biased Stochastic First-Order Methods for Conditional Stochastic Optimization and Applications in Meta Learning](https://mlsanity.com/paper/2002.10790/) (2020)
- [Accelerating Stochastic Composition Optimization](https://mlsanity.com/paper/1607.07329/) (2016)
- [Provable Stochastic Optimization for Global Contrastive Learning: Small Batch Does Not Harm Performance](https://mlsanity.com/paper/2202.12387/) (2022)
- [Oracle Complexity in Nonsmooth Nonconvex Optimization](https://mlsanity.com/paper/2104.06763/) (2021)
- [Stochastic First-order Methods for Convex and Nonconvex Functional Constrained Optimization](https://mlsanity.com/paper/1908.02734/) (2019)
- [Stochastic Optimization of Areas Under Precision-Recall Curves with Provable Convergence](https://mlsanity.com/paper/2104.08736/) (2021)
- [Proximally Guided Stochastic Subgradient Method for Nonsmooth, Nonconvex Problems](https://mlsanity.com/paper/1707.03505/) (2017)
- [When AUC meets DRO: Optimizing Partial AUC for Deep Learning with Non-Convex Convergence Guarantee](https://mlsanity.com/paper/2203.00176/) (2022)
- [Multi-Level Composite Stochastic Optimization via Nested Variance Reduction](https://mlsanity.com/paper/1908.11468/) (2019)
- [Finite-Sum Coupled Compositional Stochastic Optimization: Theory and Applications](https://mlsanity.com/paper/2202.12396/) (2022)
- [Quadratically Regularized Subgradient Methods for Weakly Convex Optimization with Weakly Convex Constraints](https://mlsanity.com/paper/1908.01871/) (2019)
- [Not All Semantics are Created Equal: Contrastive Self-supervised Learning with Automatic Temperature Individualization](https://mlsanity.com/paper/2305.11965/) (2023)
- [Rate-improved Inexact Augmented Lagrangian Method for Constrained Nonconvex Optimization](https://mlsanity.com/paper/2007.01284/) (2020)
- [Stochastic Inexact Augmented Lagrangian Method for Nonconvex Expectation Constrained Optimization](https://mlsanity.com/paper/2212.09513/) (2022)
- [First-Order Methods for Nonsmooth Nonconvex Functional Constrained Optimization with or without Slater Points](https://mlsanity.com/paper/2212.00927/) (2022)
- [Multi-block-Single-probe Variance Reduced Estimator for Coupled Compositional Optimization](https://mlsanity.com/paper/2207.08540/) (2022)
- [Complexity of Single Loop Algorithms for Nonlinear Programming with Stochastic Objective and Constraints](https://mlsanity.com/paper/2311.00678/) (2023)
- [A Primal-Dual Smoothing Framework for Max-Structured Non-Convex Optimization](https://mlsanity.com/paper/2003.04375/) (2020)

## Abstract (arXiv)

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/ε6) 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/ε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.