Title: Stochastic Momentum Methods for Non-smooth Non-Convex Finite-Sum Coupled Compositional Optimization
Venue: NeurIPS 2025
Accepted: yes
Keywords: 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: (1) a high iteration complexity of  $O(1/\epsilon^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/\epsilon^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/\epsilon^5)$ for finding an (nearly) $\epsilon$-level KKT solution. Experiments on three tasks demonstrate the effectiveness of the proposed algorithms.


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper focuses on the finite-sum coupled compositional optimization problem in the case of non-convex and non-smooth functions. There approach combines stochastic estimation of the gradient and momentum. The new propose algorithms SONEX or ALEXR2 have an accelerate convergence rate compared to existing algorithms. Finally, the experimental validation validates the empirical benefit of the method.\n\nSTRENGTHS AND WEAKNESSES:\nStrength :\nThe paper is complete and detailed. The propose method has a provable accelerate rate and perform well in practice. \n\nWeaknesses :\nMajors:\n- The paper is technical and hard to read for non-expert. A few intuitions are given.\n- The assumptions are not commented. In which case are verifying this assumption ? What are the limitations of this assumption ? Can you give examples where these assumptions are verified. In particular, the sense and consequences of Assumption 4.3 and 4.6 are not clear for me. Moreover in Assumption 3.2, the functions are supposed to be Lipschitz. So quadratic functions are not include is this analysis.\n- In line 113, $f_{\\lambda}$ is said to be $\\frac{2 - \\lambda\\rho}{\\lambda - \\lambda^2\\rho}$ which is true. However, a tighter Lipschitz constant can be derive in $\\max\\left(\\frac{1}{\\lambda}, \\frac{\\rho}{1-\\lambda\\rho}\\right)$ as demonstrated for instance in Lemma 13 of [1]. Thus, Lemma C.2 can also be improved.\n\n\nMinors:\n\t- Line 158: Can you put references for the \"existing analyses\" ?\n\t- In Theorem 4.7/ SONEX algorithm, I think that you need the inequalities $0<\\gamma\\le \\frac{1}{2}$ (to apply Lemma B.3), could you precise this constraint ?\n\nTypos/minors suggestions :\n- In line 473, I think there is no \"-\" before the subgradient.\n- Between line 481 and line 482, the indice of the max is $z_i \\in \\mathcal{Z}$ however the $z_i$ is outside the sum on $i$, can you make the notation more precise ?\n- In line 487, I suggest to precise that here $f_i$ is convex. In fact, if $f_i$ is only weakly convex, the function $f^+$ is not necessarly jointly convex and $f_{i, \\lambda}$ not necessarly convex.\n- In equation between line 491/492, I suggest to add $z_i \\in \\mathcal{Z}$ for clarity.\n- Line 113, I might be $f_{\\lambda}$ instead of $f_{\\lambda}(g)$.\n- Line 7 of SONEX algorithm : Can you precise the notation of $g_i(w_{t-1}, B_{i, 2}^t)$, the application of $g_i$ on a batch of random samples is not introduce in the paper.\n- In Line 124, could you precise the definition of monotonically non-decreasing, it is not a trivial definition for a function $f_i : \\mathbb{R}^{d_1} \\to \\mathbb{R}$.\n\n\n[1] From stability of Langevin diffusion to convergence of proximal MCMC for non-log-concave sampling, Marien Renaud, Valentin De Bortoli, Arthur Leclaire, Nicolas Papadakis, arXiv 2025\n\nQUESTIONS:\n- In line 11 of SONEX algorithm, if you apply a ADAM update, you make a second momentum update after the line 10 computation. What is the interest to do so ?\n- Can you generalize your analysis without the restrictive Lipschitz condition on $f$ and $g$ ?\n- Why can you apply directly results of Lemma B.1-B.4 that are derive for different algorithms to SONEX algorithm ?\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "CLARITY": 1, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-06-10"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper addresses Finite-sum Coupled Compositional Optimization (FCCO), which is non-smooth non-convex, where the out functions are non-smooth and inner functions are smooth or weakly convex. The currently existing state-of-the-art result face two key limitations: (1) iteration complexity O(1/epsilon^6) which is high; (2) reliance on vanilla SGD-type updates, which are not suitable for deep learning applications. \n\nThe authors propose two new stochastic momentum methods: SONEX and ALEXR2. They provide theoretical analysis that proves a faster iteration complexity of O(1/epsilon^5), improving upon previous O(1/epsilon^6) results (for finding a nearly epsilon-level KKT solution). This is achieved by applying these algorithms to multiple inequality constrained non-convex optimization problems involving smooth or weakly convex functional inequality constraints, and then optimizing a smoothed hinge penalty based formulation. \n\nEmpirical validation is conducted across GDRO, ROC-fair AUC maximization, and continual learning tasks, demonstrating improved performance over several baselines.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths: Provide a rigorous theoretical proof for proposed method that the convergence rate improves from O(1/epsilon^6) to O(1/epsilon^5).\nProvide a lot experiments to show that the proposed methods outperform existing approaches on datasets like Camelyon17, CelebA, and BDD100K.\n\nWeaknesses: The empirical analysis lacks detailed ablation studies isolating the effects of momentum, smoothing, and adaptive learning.\n\nQUESTIONS:\n(1) It's better to provide an ablation study that isolates the effect of momentum (β), smoothing parameter (θ), and Adam-type updates. How critical are they to the observed improvements?\n\n(2) Maybe it is better to have more explanation regarding why single-loop method SONEX still has a worse complexity than the double-loop method ALEXR2 in terms of dependence on the inner batch size.\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 3, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-06-28"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nIn this paper, the authors consider a 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. First, they present stochastic momentum methods for non-smooth FCCO that come with provable convergence guarantees; Second, they establish a new state-of-the-art iteration complexity of O(1/epsilon^5). Moreover, they apply the algorithms to multiple inequality constrained non-convex optimization problems involving smooth or weakly convex functional inequality constraints, and achieve a new state-of-the-art complexity of O(1/epsilon^5) for finding an (nearly) epsilon-level KKT solution. They conclude the paper by presenting results on the empirical performance of their algorithms.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths:\n\n1.The problem they address is interesting and significant for the finite-sum coupled compositional optimization. \n\n2.This is the first work to propose stochastic momentum methods for solving non-smooth FCCO problems. \n\n3.The assumptions are fairly mild. \n\nWeaknesses: \nIn the introduction, the authors list application scenarios of FCCO, but only briefly mention them without strong motivation. It is recommended to elaborate in detail on at least one representative scenario.\n\nQUESTIONS:\n1. What do you see as the biggest difficulty in generalizing the finite sum form to the stochastic expectation form in the problem model?\n\n2. Prior to formally introducing the algorithm, the authors should elaborate in detail on the design logic of each step, specifying which steps are their own innovations and which are inspired by or borrowed from existing literature. What considerations led to the design of the novel steps, and what motivated the selection of the steps borrowed from prior work? \n\n3. Some of the equations are labeled redundant and are not cited below, so please remove them.\n\n4. In numerical experiments, what method is used to solve the subproblem in Algorithm 2? \n\n5. The format of some references needs to be modified, e.g. [1]. x-risks -- X-risks. [30]. lagrangian -- Lagrangian. In addition, the format of proceedings is sometimes inconsistent. \n\n6. In Assumption 4.6, there is an extra “]”.\n\nLIMITATIONS:\nN/A", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 3, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-07-01"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper tackles the challenging problem of non-smooth, non-convex finite-sum coupled compositional optimization (FCCO), where the objective function is defined as $F(\\mathbf{w}) = \\frac{1}{n}\\sum_{i=1}^n f_i(g_i(\\mathbf{w}))$. Here, the outer functions $f_i$ are non-smooth (weakly convex or convex), and the inner functions $g_i$ are either smooth or weakly convex. The authors identify two major limitations in prior work (e.g., SONX): (1) a high iteration complexity of $O(\\epsilon^{-6})$ under restrictive assumptions, and (2) the use of vanilla SGD updates, which are not well-suited for deep learning applications. \n\nTo address these issues, the paper proposes two novel stochastic momentum-based methods:  \n1. **SONEX** (single-loop): For smooth $g_i$, leverages outer smoothing (via the Moreau envelope) combined with momentum updates.  \n2. **ALEXR2** (double-loop): For weakly convex $g_i$, applies nested smoothing and a minimax reformulation.\n\nTheoretically, both methods achieve a **new state-of-the-art complexity of $O(\\epsilon^{-5})**, representing a notable improvement over the existing $O(\\epsilon^{-6})$ rate. The authors further extend these algorithms to constrained non-convex optimization using a smoothed hinge penalty, maintaining the same complexity for obtaining approximate KKT points. Experimental results on group DRO, AUC maximization under fairness constraints, and continual learning show significant performance improvements over established baselines.\n\nSTRENGTHS AND WEAKNESSES:\n**Strengths:**   \n- 1. First to incorporate momentum-based methods (Adam-compatible) into non-smooth FCCO. The complexity reduction is practically meaningful for scalable deep learning.  \n- 2. The proposed nested smoothing strategy for weakly convex inner functions and the formalization of \"approximate $\\epsilon$-stationary points\" represent important technical innovations.  \n- 3. Rigorous theoretical analysis (17-page appendix); the convergence guarantees under weak convexity are sound. Experiments span 6 datasets across 3 diverse tasks.  \n- 4. Clear organization; Figures 1–3 effectively illustrate key comparisons. Detailed pseudocode enhances reproducibility. \n\n**Weaknesses:**   \n- 1. The complexity results rely on the inner functions being smooth or weakly convex; fully non-smooth inner functions remain an open challenge.  \n- 2. While effective, the use of Moreau envelope for outer smoothing is a known approach; the novelty primarily lies in combining it with momentum methods.\n- 3. Assumption 4.3, which places a lower bound on $\\lambda_{\\min}(\\nabla g_i \\nabla g_i^\\top)$, lacks empirical validation beyond small-scale experiments (Table 3). Its applicability to large-scale models (e.g., Transformers) is uncertain.  \n- 4. The double-loop structure of ALEXR2 (Algorithm 2) is relatively complex; a clearer intuition or motivation for nested smoothing would improve accessibility.\n\nQUESTIONS:\n1. **Generality of Assumption 4.3**:  \n   Theorem 4.4 depends on the condition $\\lambda_{\\min}(\\nabla g_i \\nabla g_i^\\top) \\geq \\delta$. Is this assumption realistic for deep networks (e.g., ResNet), where Hessians can be nearly singular? If $\\delta$ depends on $\\epsilon$, how would that affect the overall complexity? *Empirical validation on ImageNet-scale models would substantially strengthen the claim.*  \n\n2. **Efficiency of the Double-Loop Design**:  \n   ALEXR2 requires inner-loop iterations with $K_t = \\tilde{O}(\\epsilon^{-3})$. Despite the theoretical gains, is the wall-clock time actually better than SONX in practice? Could a modified single-loop variant handle weakly convex $g_i$ effectively?  \n\n3. **Adaptive Learning Rate Assumptions**:  \n   Appendix E introduces Assumption E.1, which bounds the adaptive learning rates. Is this assumption satisfied in practical scenarios involving Adam? Could the theory be extended to relax this assumption?\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "CLARITY": 4, "IMPACT": 4, "ORIGINALITY": 3, "DATE": "2025-07-03"}