Stochastic Momentum Methods for Non-smooth Non-Convex Finite-Sum Coupled Compositional Optimization
Authors: Xingyu Chen
Bokun Wang
Ming Yang
Qihang Lin
Tianbao Yang
Venue: NeurIPS 2025 poster ()
Keywords: finite-sum coupled compositional optimization
non-convex constrained optimization
Primary area: 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.

Stochastic Momentum Methods for Non-smooth Non-Convex Finite-Sum Coupled Compositional Optimization
Xingyu Chen" tamucxy2025@tamu edu
Bokun Wang" bokun- waann-aang@tam... edu
Ming Yang*
myangStamu, edu
Qihang Lin qihang-lin@uicn wwwwuiowwaww....m.. edu
Tianbao Yang*
tianbao- Baanbao-yanggtame edu.
Abstract
Finite-sum Coupled Compositional Optimization (FCCO), characterized by its coupled compositional objective structure, emerges as an important optimization paradigm for addressing wide range of machine learning problems. In this paper, we focus on challenging class of non-convex non-smooth FCCO, where the outer functions are non-smooth weakly convex convex and the inner functions are smooth or weakly convex. Existing state-of-the-arr result face two key limitations: (I) a high iteration complexity of O(1/E") 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/c".. Moreover, we apply our algorithms multiple inequality constrained non-convex optimization problems involving smooth or weakly convex functional inequality constraints. By optimizing smoothed hinge penalty based formulation, we achieve new state-of-the-art complexity of ((1/e²) for finding an (nearly) e-level KKT solu- tion. Experiments on three tasks demonstrate the effectiveness of the proposed algorithms.
Introduction
In this paper, we consider the finite-sum coupled compositional optimization (FCCO) problem
min F(w) = - f((a(w)),
weRd
where ft d + R is continuous, gs Rd + R1 is continuous and satisfies (w()) = Eg:( (w, E.).
and the expectation is taken over the random variable &, for ,n.
FCCO has been effectively applied to optimizing wide range of risk functions known as X- and group distributionally robust optimization (GDRO) 5 6 Recently, it was
risks
also applied to solving non-convex inequality constrained optimization problems [7 8]. Several optimization algorithms have been developed for non-convex FCCO ① under different conditions.
Department of CSE, Texas A&M University, College Station, USA Tippie College of Business, The University of Iowa, Iowa City, USA. Coorespondence to: tianbao- saabaa--yanggtauu.com edu.
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

Table 1: Comparison between our algorithms and prior works for the non-smooth non-convex FCCO problem. The complexity of SONX is for finding a nearly c-stationary solution (Definition 3.1). and that of SONEX and ALEXR2 for smooth inner functions are for finding approximate c-stationary solution (Definition 4.Ip which implies nearly c-stationary solution under a verifiable condition of the inner functions. The complexity of ALEXR2 for weakly convex inner functions is for finding a nearly e-stationary solution to the outer smoothed objective. this table, "WC" means weakly convex, "C" means convex," 70h means monotonically non-decreasing. "SM" means smooth, "PMS" means that the proximal mapping can be easily computed, "MLCO" means that the function is mean Lipschitz continuous of zero-order (Assumption
Algorithm
fi
9i WC, MLCO
Complexity
Loop
Update type
SONX [6]
WC, A
Single
SGD
SONEX (Ours)
WC, PMS
SM, MLCO
Single
Momentum or Adam
ALEXR2 (Ours)
C, PMS
SM
O O(c-5)
Double
Momentum Adam
ALEXR2 (Ours)
C. PMS. /
WC
Double
Momentum Adam
Table 2: Comparison between our algorithms and prior works for solving non-convex inequality constrained optimization problem. 9o is the objective and gi are constraint functions.
Algorithm
9o
Constraint Sampling
Complexity
Loop
Update type
OSS 9
WC
WC
No
O
Double
SGD
ICPPAC 10
WC
WC
No
O(c-5)
Double
SGD
SSG |1]
WC
C
No
O(e-8)
Single
SGD
Li et al. 7
SM
SM
Yes
O(e-7)
Single
Momentum or Adam
Yang al. 8
WC
WC, MLCO
Yes
O(c-6)
Single
SGD
Liu et al. 12
WC
WC, MI.CO
No
O(e-O)
Single
SGD
SONEX (Ours)
SM
SM. MLCO
Yes
(
Single
Momentum or Adam
ALEXR2 (Ours)
WC
WC
Yes
0
Double
Momentum or Adam
Most of them require the smoothness of f. and gi. Hu et al. 66] has initiated the study of non-smooth non-convex FCCO where both the inner and outer functions could be non-smooth. However, their results face two key limitations: (1) a high iteration complexity of O(1/E") 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.
Novelty. This paper addresses these limitations by proposing stochastic momentum methods for non-smooth FCCO, where Si is non-smooth and gi could be smooth or weakly convex, and improving the convergence rate to O(1/e).. To the best of our knowledge, this is the first work to propose stochastic momentum methods for solving non-smooth FCCO problems. Unlike [6] that directly solves the original problem, the key to our methods is to smooth the outer non-smooth functions based on the Moreau envelope smoothing or equivalently the Nesterov smoothing when they are convex, which is referred as outer smoothing. When the inner functions are non-smooth weakly convex, we further smooth the transformed objective using another layer of Moreau envelope smoothing, which referred to as nested smoothing. Then we propose stochastic momentum algorithms to optimize these smoothed objectives.
Contributions and Significance. We establish two main results regarding non-smooth FCCO. In the first result, we consider non-smooth FCCO with non-smooth outer functions and smooth inner functions. With outer smoothing, we propose a single-loop stochastic momentum method named SONEX to solve the resulting smoothed objective. Our main contribution here lies at a theory that guarantees a convergence rate ((1/e) under meaningful convergence measure. Specially, we show that when the outer smoothing parameter is small enough, the proposed algorithm is guaranteed to find novel notion of approximate e-stationary solution to the original problem, which implies the standard nearly c-stationary solution to the original problem under verifiable condition of inner functions. In the second result, we consider non-smooth FCCO with non- mon-smoohh convex outer functions and non-smooth weakly inner functions. With nested smoothing, we propose novel double-loop stochastic momentum method named ALEXR2 for solving the resulting smoothed objective and establish convergence rate of O(1/e5) for finding an e-stationary solution of the smoothed objective. Table TaaleToomme compares our methods with those of the existing works for non-smooth FCCO problems from different aspects.

---

Then we consider novel applications of these two algorithms to non-convex inequality constrained optimization. By optimizing a smoothed hinge penalty function with the proposed stothastic mo- mentum methods, we derive new state-of-the-arr rate of O(1 (11/e)) for finding an e-K solution when objective and constraint functions are smooth, and for finding a nearly K- solution when objective and the constraint functions are weakly convex. Table summarizes Tamliiissm the complexity of our methods and existing works for solving non-convex inequality constrained optimization problems.
2 Related work
FCCO. FCCO is a special case of stochastic composite optimization (SCO) 13 and conditional stochastic optimization (CSO) [15] when the outer expectation is a finite sum. FCCO was first intro- duced in 4 as a model for optimizing the average precision. Later, (2 proposed the SOX algorithm and improved the convergence rate by adding gradient momentum and further being improved in [[6] by replacing exponential moving average (EMA) with MSVR and STORM estimators for inner function and overall gradient, respectively. All these techniques have gain success in optimizing Deep X-Risks [ such as smooth surrogate losses of AUC and contrastive loss [17 18] For convex FCCO problem, [2] reformulate FCCO saddle point problem and use restarting technique to boost convergence rate. For non-smooth weakly-convex FCCO (NSWC FCCO), 6 leverages MSVR technique and requires ( 6)) iterations to achieve nearly e-stationary point.
Smoothing techniques: An effective approach for solving non-smooth optimization is to approximate the original problem by smoothed one using different smoothing techniques, including Nesterov's s smoothing [19], randomized smoothing [20), and Moreau envelope [2T [2]. Our method is related to the one based on Moreau envelope, which is a prevalent technique for weakly convex non-smooth problems. This method approximates the objective function by its smooth Moreau envelope and then computes an c-stationary point of the Moreau envelope, which is also a nearly c-stationary solution of the original problem [22 223. The similar technique has been extended to min-max problems 124 25
and constrained problems [9 10 26 Our methods are different from these works because we consider the Moreau envelope of each outer function and design stochastic momentum methods. Similar outer smoothing techniques have been developed in 27 for compositional objectives. However, different from our work (1) their algorithm cannot handle multiple outer functions of a finite sum structure as ( (2) their algorithm is not momentum method; (3) their convergence theory requires mean Lipchitz continuity of the gradient of the inner functions for finding a primal-dual c-stationary (KKT) solution of the min-max formulation of the original problem.
Non-convex constrained optimization: Most existing works for stochastic non-conved constrained optimization focus on the smooth problems. For example, (28) proposes double-loop inexact
augmented Lagrangian method for stochastic smooth nonconvex equality constrained optimization
under a regularity assumption on the gradients of the constraints. A quadratic penalty method developed by (29) under similar assumptions but uses only a single loop. Both [29] and [28]
consider only smooth problems while we also consider non-smooth problems. For non-smooth but weakly convex problems, the existing methods can compute a nearly e- KKK point based on the Moreau envelopes of the original problem with iteration complexity O(E 5)
9
10
[T] 26 g.
The penalty-base -base method by [1] reduces the number of evaluations of the subgradients to Occ
but still needs to evaluate O 0) function values. In contrast, our method only needs complexity Occ=5) to compute a nearly k-K point. For smooth constrained optimization with non-convex equality constraints, variance-reduced stochastic gradient methods have been developed based on quadratic penalty [29 [30. One can convert inequality constraints into equality constraints by adding slack variables x2 (requiring 8 > 0) to each inequality constraint 30 However, adding s² will introduce spurious stationary points [31 32), while adding 8 2 O will make the domain or the
constraint functions unbounded which violate the assumptions in [29 30 More importantly, we
consider non-smooth constraint functions while [29 30) only consider smooth cases.
3 Preliminaries
Let ||- || be the Euclidean norm. For g(-) Rd + Rd, let Vg(-) E IRdxd, be its Ja- cobian and let VVVl maxpeed. -1 [-]u be the operator norm. Let gi(w, B)
[8 EeeB (w, i = 1,2,-.. n be an estimator of g, based on samples from a minibatch B. A function f(-) is monotonically non-decreasing if for any X1, E 1d satisfying that
3

---

X1, S 82,k, & E [d1] where X denote the k-th element of x, it holds that f(x)) I f(x2). A function f()) is p-weakly convex for p 2 if f()) + 2 I1-11 is convex. A function g(.) is L-smooth for L if g is differentiable and V-(-) is L-Lipschitz continuous. A vector-valued mapping g(-) is L-smooth for L 2 0 g is differentiable each component and V-(() is L-Lipschitz continuous w.r.t to the operator norm. Given a p-weakly convex function f(.), we denote its regular subdifferential as ()()). For p-weakly convex function ʃ(-) and constant (0, (0,p="), the proximal mapping and the Moreau envelope of f()) with parameter À are defined as follows, respectively,
2A ||u vll and fa == min f(v) + 1 |
p (u) == arg min f(v) +
v
2X
It is known that, for x (0, fa(+) by == max +. TSxp) smoohhh 33 and (u) = pr (u) E 0 (proxx (u))
(2)
For non-smooth weakly convex objective, we follow existing works 21 6 and aim to find nearly e-stationary solution defined below.
Definition 3.1. A solution is e-stationary solution of I if dist(O (FF(( I A solution w is nearly c-stationary solution of (T) if there exists w' such that |w w'll I c and dist(0, (FF00,FF("")) c.
We make the following assumption on (T throughout the paper. Assumption 3.2. We assume that is Cy-Lipschitz continuous, 9: is Cy-Lipschitz continuous
for i = 1,..., n, xr (w) and its subgradient are easily computable, and one of the following conditions hold:
A1. Sa py-weakly convex, and 9.(w) is Lg -mooth for all i.
A2. f. is convex, and 9, (w) is Lg -mmoohh for all i.
A3. f. is convex and monotonically non-decreasing, and g.(w) is Py Py-weakly convex for all i.
The Lipschitz continuity of f and in Assumption 3.2 are fairly standard for FCCO problems. Since we target at non-smooth / and g, Lipchitz continuity is a minimal assumption. In the considered applications of GDRO and learning with fairness constraints in section 6 f is hinge and hence Lipschitz continuous.
Remark:
Smoothing of Non-smooth FCCO
We first describe the main idea of our methods and then present detailed algorithms and their convergence. Under the Assumption 3.2 we can show F is also weakly convex. To improve the convergence rates and accommodate deep learning applications, we develop stochastic momentum methods, which can be easily modified to incorporate with adaptive step sizes (cf. Appendix D The challenges of developing provable stochastic momentum methods for solving (I) lie at the non-smoothness of Si and potentially gr.
Let us first consider smooth 9i and defer the discussion for weakly convex 9; to subsection The key idea is to use a smoothed version of ft. denoted by fi,,, and to approximate ① by
min FA(w) : --=""" " fex(e((W))-
(3)
WERd
n2
We can easily show that when 9i is Lggsmooth, then Fx() is Lp-smooth with LF O² m -ー + Cy L,g (ef. Lemma A.I As a result, the problem becomes a smooth FCCO where both inner functions and outer functions are smooth, which makes it possible to employ existing techniques to develop stochastic momentum methods find an e-stationary point of (3). which is point w satisfying V[[| (x()]l] < However, what this implies for solving the original problem and what constitutes good choice of À remain unclear. Below, we present a theory to address these questions. We introduce the following notion of stationarity of the original problem.
Definition 4.1. A solution w is an approximate c-stationary solution of [I] if there exist t bm and y; E 0ft(t;) for i I,..., such that ||t: g:(w)] < i ==1, 1, n, and 1 E Vg.() S.
This definition implies that when 6 - O. the solution converges to a stationary solution of F(-). Given this definition, our first theorem is stated below whose proof is given in Section A.3

---

Algorithm SONEX for solving
Input: T, A, Wo W-1, Vo. 7/ 7. 7', Vt.
2: Sample a batch of data 2, from the distribution of &, for = 1,...,
3: Set u1,0 = 9: (Wo. B,,)) for I....,
4: for = 0, 1, 1.... T
do
5: Sample Bi C ............... n and batch of data Bi from the distribution of & for E Bi.
6: for 1,.... do
7: Update
Wi,+1 = {{
(1 - yu,,
(g,,,, 9.(w.-I....)
if E Bi if $ Bi
Wi,t
8:
end for Compute G+ [81 Eiegy
9:
Update Vi+1
8) + 3G,
11: Update Wi+1 = Wt 7/ or Adam-type update
12:
end for Output: W, with randomly sampled from (1,2, .T)
Theorem 4.2. If is e-stationary solution Fx(+) with X c/C, such that (VV s. then w is an approximate e-stationary solution to the original objective F(-).
remaining question is whether an approximate e-stationary point is also a nearly e-stationary solution of the original problem ( We present a theorem below under the following assumption.
Assumption 4.3 (EVLB assumption). There exist c > 0 and 8 > such that, w is an e- stationary point of Fx with À e/C, and I c. it holds that Amin (Vg(w)Vg(w))) 8, where Vg(w) = [Vi (w) ..... (g)(w)
Remark: In Assumption 4.3 once is smaller than c, the lower bound ó does not depend on The empirical justification for this assumption is provided in subsection F.5
Theorem 4.4. Suppose Assumptions
3.2
(AI) and 4.3 hold. If is approximate e-stationary
solution of (1 with e S. min {C, 16nC24,
w is also nearly e-stationary solution of (
is easily verifiable (cf. Appendix [F55 for empirical evidence). In addition, when we consider the applications in non-convex constrained optimization, this condition is commonly used in existing analyses (7] 8
Remark: The proof is given in Appendix A.4 It is notable that the Assumption
4.1 Single-loop Methods for Smooth Inner Functions
Next, we present single-loop algorithm for finding e-stationary point of 3}. The key ingredients of the algorithm include two parts: (1) maintaining and updating n sequences ú₁ for tracking each ge (:(w), 1, ...... which updated in coordinate-wiss manner; (2) stochastic momentum update. We present the detailed updates Algorithm|] I which is referred to as SONEX. We note that the algorithm is inspired by existing stochastic momentum methods for smooth FCCO. First, Step 7 for updating W1,+1 is from the MSVR algorithm [66. Different from MSVR, we directly utilize the stochastic momentum update in Step 10, 11, which are similar to SOX 2 In contrast, MSVR also leverages variance reduction technique (STROM) to compute an estimate of the stochastic gradient, which requires stronger assumption that V (w; {) is Lipschitz continuous in expectation. Our analysis shows that this is not helpful for improving the convergence, as the complexity will be dominated by term related the smoothness of frx, which is in the order of O(1/c).
We assume the following conditions of the stochastic estimators of 9: and their gradients. Assumption 4.5. There exist constants du and σı > 0 such that the following statements hold for yi (w) and gi(w, (Ww,.W)) Ea): for i 1, ... and any E Rd
(w &.) g (w) 2| < of E [1 (w,, Ex) 0g. (w)l2 < oo
Assumption 4.6. There exist constant Cp such that Eg (w E) 9.(w', . S C,llw 1
5

---

It is notable that these assumptions have been made in [6] which is important for analyzing variance reduction technique such as MSVR. We refer Assumption 4.6 as mean Lipchitz continuity of zero-order (MICO). Let B1 and B2 denote outer and inner batch sizes.
Theorem 4.7. Under Assumption 3.2 (AI).
4.5 and 4.6 by setting X = + =
(minn B2)e), O(B2e*), O
SONEX with 7 = -1- B1(1-2)
Bi and
n
7 7 converges to an approximate e-stationary solution of 1 within T O B.BE itera- tions.
Combining the above theorem with Theorem 4.4
we obtain the following guarantee:
Corollary 4.8. Under Assumption AI), 4.3 4.3 and 4.6 with the same setting as in Theorem 4.7
SONEX converges to a nearly e-stationary solution of within T 0 BIVB2
12 - - iterations.
We compare our proposed SONEX with SONX for solving T. Under assumption 4.6 the rate of SONX is ((1/c)) which worse than our result of O(1/e)) in Corollary 4.8
4.2 A Double-loop Algorithm for Smooth or Weakly-Conver Inner Functions
In this subsection, we further improve our results to achieve the same convergence rate of O(1/c)) by (1) removing the MCL0 assumption instead assuming that fa are convex; (2) designing stochastic momentum method with a convergence guarantee for weakly convex 9:- It is worth mentioning that the convexity of f. holds for broad range of real applications such as group DRO and the application in non-convex constrained optimization as discussed in next section. Before introducing our new algorithm, need first reformulate the problem.
(A3). For simplicity of notation, we denote fi, by ft. When fs is convex, cast the problem into minimax formulation:
Let us first consider the scenario that satisfies Assumption 3.2
min Fx (w)
Σ T
max y.ERd,
(y. g(w) - (v:)),
WERd
where ft is the conjugate function of fi. In Appendix A.2 we show that this is also equivalent the classical Nesterov's 's smoothing [1]]. Under Assumption 3.2 (A3), Jt is also Cy-Lipschitz continuous, convex, non-decreasing because of ( As consequence, dom(f?)
{y, E R R4llly S Cy)- This implies that, for any y, E dom( f ). function y, 9, (w) is Plyylla/g-weake convex in w and also
ydgCyp--weakl convex because ||y |1 S Vdilyy S Cd This further implies that Fx(-) is
PF == Pr,
convex.
Different from last subsection, there is another challenge we need to deal with in order to develop a stochastic momentum method, which comes from that 9i are non-smooth. To tackle this challenge, we use another Moreau envelope smoothing of (-)-) and solve the following problem:
||z - w|l w(}}}
min weRd
(w)(w) = min ZER
(z) + 2v
(5)
min weRd
min max
y. (z)) (r(y)) + 2v 1|z w/1
(6)
2e3d yeRndi
1=1
where E (0, p7)) and y = [yT .y........
The benefit of doing this is that (i) the resulting objective (x,(w) (w) smooth (cf. Lemma C.2 which allows us employ stochastic momentum method; (ii) the inner min, maxy becomes strongly-conve and strongly-concave problem due to that F. is smooth and its conjugate is strongly convex, and E (0, (0,P7) (cf. Lemma C.3
1 (ww- prox,r, (w)). Given approximate solution of prox, (w). denoted by of the inner minimization problem in (5), we can use L(w 2) as an inexact gradient Fx.(wW) to update w with momentum method in order to solve (5). To obtain an estimate 2, at the t-th iteration, we employ a recent algorithm ALEXR 55) for solving the inner minimax problem. With an estimate 2,, we will update Wi+1 by momentum update. We present the detailed steps of our double-loop algorithm named ALEXR2 in Algorithm 2 where Uy: (yi, y') the Bregman (y.). ft(y() - (0ft (yl), y()- For
Recall that (w)
divergence induced by ft, namely, Uz. yY) the sake of convergence analyses we need Uy: to be bounded, as stated in the assumption below.

---

Algorithm 2 ALEXR2 for solving 66)
1: Input: T. Wo E R Vo = 0, K2 a, B, v, 0,.00.7 0, 77000 0
2: for = 0,1.....
1 do
3: Set 21,0 = 4, = Wt and initialize Yı,o = 1v ....:
(n) JT E Rndd.
4: for k = 0, ..... K,
do
5: Sample | c {1. ...................... n) and two independent batches of data B1 and B from the
distribution of E for i B.
6: Compute gi.k = (z
(22) + ((g. 11,2) -
for B,
7: for I....... n do
y Di.k (y 1 (y( ,
if Bi if E Bit
(6) arg maxx (1
8:
Vik+1
Yikk
9:
end for
10:
Compute GG | Eiegy. (2, B (1:))yy...
y(() -
Update Zt,k+1 = arg min {(GERs 2) + 2 l|z well 2ell2= + 2nj
Ze 112)
11:
12:
end for
13:
Let G2 = (w, kt,
14:
Update V2+1 = (1 - B3v, + BG+
15:
Update W++1 = Wt 0V++1 or use Adam-type update
16:
end for
17:
Output: w, with randomly sampled from (1, ...2....)) T}
Besides, we would like point out that the sequences of y; can be updated similar u, sequences in SONEX similar to [. i.e., vy (u(,))))) where ui,t1 = (1 --) x + (z, B
(( (z (, (z B ), where 4 = 7/(1 7).
Assumption 4.9. We assume that fi is a function s.t. Uz. (21,32) is bounded for any Vıs y2 dom( C Rdd
Remark: We point out that this condition is not strong under our setting where J, is lipschitz continuous and dom(ff") is consequently bounded and it is satisfied by many practical convex lipschitz-comtnuouus functions such as hinge function, smoothed hinge function(i.e. fi, or f. in this paper), etc.
The following theorem states the convergence of ALEXR2 with proof given in Section C
Theorem 4.10. Suppose assumption 3.2 (A3), 4.5 and 4.9 hold and X = e/Cr in 4 For any > 0, 1-8 L, and 7
there exists 0 E (0, 1) with 1 A ( such that, by setting 1 9
(1-8)n
B < 1 B,
2'
B, B + ) ALEXR2 returns w, as nearly e-stationary solution of in expectation within T (e) and total iteration complexity of
8 and K, = Ô 2
n
B, +
Remark: We can easily extend the above result to Assumption 3.2 (A2). When 9i are smooth, we do not need the monotonicity of fi as the minimax problem minl maxy is still guaranteed to be strongly-conver and strongly-concave when E (0. where PFx = Cd is the weak convexity parameter of Fx.
Finally, we show that when are smooth, we can also obtain a nearly e-stationary solution to the original problem [T]. Different from the result Theorem 4.7 the above theorem only guarantees a nearly e-stationary solution to Fx. To address this gap, we can recover a nearly e-stationary solution to (T) from w, by running ALEXR starting with w, with K Ö(c-5) iterations. This result is stated as the following corollary with the proof given in Section C
82 hold. For any 6 S min {c 16nC31, Lo S
Corollary 4.11. Suppose assumption 3.2 (A2),
4.3
4.5
lei W ALEXR(W w,, K) with K Ô
By Byes
+
n
= 1 L, and ? (1-01mm
0
B, for 0 E (0, 1)
7

---

with 1 - - O(e)). w, is nearly e-stationary solution of 1 and it is found with a total iteration complexity of O B, Bgc
n + ).
5
Smoothed Hinge Penalty Method for Constrained Optimization
In this section, we consider a constrained optimization problem with m 1 inequality constraints:
min (o s.t. (:(w) < 0, = 1, -, ,
where gi Rd IR is continuous and satisfies gi (w) (w:(w, Ex) and the expectation is taken the random variable & for = 0, 1,.... m. Following 7 by solving the following unconstrained problem:
we consider an exact penalty method for
min ((w) : WEBd
(o()) +
P
9a [w]]
(8)
where p is sufficiently large number. Let f(-) := p|-|+ so the penalty term E (w)]+ has the same structure as F(w) in ① Yang al. [8] have employed SONX for solving the above
problem and established a complexity of (c) for finding nearly kKKT solution. Below, we apply the smoothing idea and ALEXR2 or SONEX for improving the convergence rate. Our key idea
is to optimize an outer smoothed problem:
m
A 9:9. (w) 2
min 2 (w) = go ()) weRd
(x()) go (w) +
max [0,p))
m
m
where fx(2) == min f(=") + 2x (= z²)2. Except for the term go. the objective function of (9)
has the same structure as 3 and with J (w:) v/ 1 lu (v:). (v.))
Assumption 5.1. There exists a constant
such that Jmin (J(w) 2 8 for any matrix J J(w) = [h, (w), (w)] E dR with (,(w) 0g. (w) and any w satisfying max,=1,.... 9r >
Under this assumption, we have the following proposition whose proof is given in Section D.I
Proposition 5.2. Suppose Assumption 5.1 hold. If p mlCg+1) and X 5. a nearly e-stationary
solution w to (9) is also a nearly ((e)-KTT solution to the original problem
7 in the sense that there
exist w and v 0 for 1, I,.... m such that ||w w S e, dist(0, dgo (w) Em 0g (w)())))) S Oce), max;=1...... g: S O(e) and g (()) < O(e), vi 1,2, 1.2..... m
To find nearly e-stationary solution to p we can adapt either SONEX if gi are smooth or ALEXR2 if gi are weakly convex with a minor change by including the stochastic gradient of go in the calculation of G, The complexity of this approach is presented below with proofs given in Section D.2
Theorem 5.3. Suppose assumption 4.5 and 5.7 hold, and the stochastic (sub)gradient of go has bounded variance. Let p > m(C, + 1 and c/p.
if f9 are weakly convex, then with δ
iterations ALEXR2 find nearly c-KKT
B. Bg845
solution for 7 satisfying that there exist w and v 0 1,..., such that E ||w w|| S e and (0, dgo wi >m Oge (we W):))) < O(e) and it holds with probability x maxi=1......
O(e) that
9e( (w) S (()) and Ig: W W))k < O(e), Vi 1, m
if Ym are smooth, then with δ
iterations SONEX can find an c-KKT solution B. 55
for similar as above except for w w.
Notably, our method not only improve the rate but also enjoy an improved dependence on 6, compared to O 68-6)
6) in prior work [8].
'If Ig. (w)l < +00 for any and this high probability result can be replaced by [8] g. [w)] - O(c) 1=1...m
and O Vi 1,2, m.
8

---

Camelyon
Amazon
CelebA
SOREX
SONEXX
SOMEX SONX
SCRI
SONA Perriiuu
Dugl
Eerations
Teeaiooss
terstions
Cametyon
CelebA
Arrzzon
Figure 1: Training loss curves (left three) and testing accuracy (right onc) of different methods for Group DRO with CVaR ratio = 0.15 on different datasets.
Adult
COMPAS
COMPAS
6.70
S00r 20000
3.01
Spack
Eanch
Figure 2: Training curves of AUC values (fig 1,3) and constraint violation (fig 2,4) of different methods. The format of legend is "Aloorithm(penale function, p)", and SH, H, smH mean square hinge, hinge and smoothed hinge, respectively.
toggy
Feggy
Target Ovarcest
20051 pa
20901
Figure 3: Training curves of Target AACC values (fig 1,3) and constraint violation (fig 2,4) of different methods. The format of label "Algoritmmppenaty function, p)", and SH, smH mean square hinge and smoothed hinge, respectively.
Numerical Experiments
We conduct experiments to verify the effectiveness of the proposed algorithms. For non-smooth non-convex FCCO, we consider GDRO with CVaR divergence [5]. For non-convex constrained optimization, we follow [8 and consider two tasks, namely AUC maximization with ROC faimesss constraints and continual learning with forgettggg constraints.
GDRO with CVaR divergence. We consider the following GDRO with CVaR divergence, which minimizes top-k worst groups' losses [5]
min min +
E (x,)) {(0; (x, $]+
(10)
where n is the number of groups, r = k/n, and D, denotes the data of the g-th group.
We use 3 datasets: Camelyon17. 17, Amazon 34 and CelebA [35 The first two datasets are from WILDS Benchmark for evaluating methods tackling the distributional shift (36), where Camelyon17 has 30 groups and Amazon has 1252 groups. CelebA is a large-scale facial attribute dataset containing over 200,000 celebrity images annotated with 40 binary attributes. We select 4 binary attributes 'Attractive', 'Mouth_S0 'Mouth_Slightly _Openn,__p'' 'Male' and 'Blonde_ _air" and construct 16 groups, where 'Blonde_Hai also serves as the label for classification. We compare SONEX with SONX and Primal- We use pretrained Densenet121 38 Distilbert (39] and ResNet50 (40] for Camelyon11 17, Amazon and CelebA, respectively. We perform Adam-type update for SONEX Amazon dataset and momentum-type update Camelyon17 and CelebA datasets. We perform Adam-type update for Duil on Amazon dataset and SGD-type update Camelyon17 and CelebA. We each experiment for random seeds and report their average performance. The hyperparameter tuning is discussed in Appendix E.
Dual, where PrimalDual is stochastic primal-dual algorithm F]
'Naeed 'OOA' in [5) it's derived by extending the algorithm in 37 GDRO with CVaR setting

---

Results. We report our result under setting of = 0.15 in Figure ! where the first 3 figures are the loss curves on the training dataset while the last one is the test accuracy. Our experiment results shows that SONEX performs better than SONX and PrimalDual on GDRO tasks.
AUC maximization with ROC fairness constraints. The formulation is given in Appendix F.2 where the objective is AUC loss and constraints specify the tolerance of the gap between false positive rates (true positive rates) of two sensitive groups different classification thresholds. For this experiment, we follow almost the same setting as in g Two datasets are used, namely Adult [T] and
COMPAS 42,, which contain male/female, Caucasian/non- Caucasian groups, respectively, We set thresholds at I = {-3, -2, -1, ... 33 and the tolerance K = 0.005, which gives us 14 constraints. We learn a simple neural network with 2 hidden layers. We compare ALEXR2 with the method in 7 that optimizes a squared- -higee penalty function with the SOX algorithm [2, the method that optimizes a hinge penalty function with the SONX algorithm [], the double-loop method ICPPAC [I0]. We perform Adam-type update for ALEXR2. We run each method for totally 60 epoches with batch size of 128, repeat five times with different seeds and then report average of the AUC scores and constraint values. Hyperparameter tuning is presented in Appendix F2
Results. We compare the training curves of objective AUC values and the constraint violation measured by the worst constraint function value at each epoch for different methods in Figure 2 on different datasets. These results demonstrate that ALEXR2 optimizing smoothed hinge penalty function can better maximize AUC value on both datasets than the baseline methods while still having similar constraint satisfaction.
Continual Learning with non-forgetting constraints. We follow a similar experimental setup to ? for fine-tuning a CLIP model for autonomous driving the BDD100K dataset (3 large-scale, multi-task driving image dataset. The formulation is given in Appendix F3 where the objective is the contrastive loss on a target task (e.g., classifying foggy and overcast condition). and the constraints specify the logistic loss of the new model should not be larger than that of the old model for other different classes (e.g., clear, snowy, rainy, partly cloudy). Since the objective function is another FCCO, we modify SONEX such that the gradient estimator of the objective is computed similar as the smoothed square penalty function. We compare with SOX that optimizes the squared hinge loss use the adam-type update for both methods. We do not compare with SONX that optimizes the hinge penalty function with SGD-ttpe update as it fails for learning the Transformer network used in this experiment. Hyperparameter tuning is presented in Appendix F3
Results. We present training curves of accuracy improvement and constraint violation on different target tasks in Figure 3 It shows that our method can achieve higher accuracy on target tasks than the baseline method of using squared hinge function while retaining similar constraint satisfaction.
7 Conclusions and Discussion
In this paper, have considered non-smooth non-convex finite-sum coupled compositional opti- mization, where the out functions are non non-smooth and inner functions are smooth or weakly convex. We proposed stochastic momentum methods to improve the convergence rate and accelerate the training for deep neural networks. We also considered the applications of the proposed stochastic momentum methods solving non-convex inequality constrained optimization problems and derived ofate-of-the-anr art results for finding an c-K solution. Our experiments demonstrate the effectiveness of our methods. One limitation of our results is that our single-loop method SONEX still has worse complexity than the double-loop method ALEXR2 in terms of dependence on the inner batch size.
Acknowledgments and Disclosure of Funding
We thank Quanqi Hu for the assistance in developing the codehase for group DRO. X. Chen, B. Wang, M. Yang and T. Yang were partially supported by the National Science Foundation Career Award 2246753, the National Science Foundation Award 2246757 and 2306572. Q. Lin was partially supported by National Science Foundation Award 2147253.
10

---

References
[1] Tianbao Yang. Algorithmic foundations of empirical X-risk minimization. arXiv:2206.0044 00439, 2022.
arXiv preprint
[2] Bokun Wang and Tianbao Yang. Finite-sum coupled compositional stochastic optimization: Theory and applications. In International Conference on Machine Learning, pages 23292-23317. PMIR, 2022.
[3] Dixian Zhu, Gang Li, Bokun Wang, Xiaodong Wu, and Tianbao Yang. When auc meets dro: Optimizing partial auc for deep learning with non-convex convergence guarantee. In International Conference on Machine Learning, pages 27548 27573. PMLR, 2022.
[4] Qi Qi. Youzhi Luo, Zhao Xu, Shuiwang Ji, and Tianbao Yang. Stochastic optimization of areas under precision-recall curves with provable convergence. Advances in neural information processing systems. 34:1752-1765, 2021.
[5] Bokun Wang and Tianbao Yang. Alexr: An optimal single-loop algorithm for convex finite-sum coupled compositional stochastic optimization. arXiv preprint arXiv:2312.02277 2023.
[6] Quanqi Hu, Dixian Zhu, and Tianbao Yang. Non-smooth weakly-convex finite-sum coupled compositional optimization. Advances in Neural Information Processing Systems, 36:5348-54033, 2023.
[7] Gang Li, Wendi Yu, Yao Yao, Wei Tong, Yingbin Liang, Qihang Lin, and Tianbao Yang. Model devel- opmental safety: A safety-centric method and applications in vision-lanuuage models. arXiv preprint arXiv:2410.03955, 2024.
[8] Ming Yang. Gang Li, Quanqi Hu, Qihang Lin, and Tianbao Yang. Single-loop algorithms for stochastic non-convex optimization with weakly-convex constraints. arXiv preprint arXiv:2504. 15243, 2025.
Runchao Ma, Qihang Lin, and Tianbao Yang. Quadratically regularized subgradicnt methods for weakly convex optimization with weakly convex constraints. In International Conference on Machine Learning. pages 6554-6564. PMLR, 2020.
[10] Digvijay Boob, Qi Deng, and Guanghui Lan. Stochastic first-order methods for convex and nonconvex functional constrained optimization. Mathematical Programming, 197(1):255-279, 2023.
[11] Yankun Huang and Qihang Lin. Oracle complexity of single-loop switching subgradient methods for non- smooth weakly convex functional constrained optimization. Advances Neural Information Processing Systems, 36:61 323-61347, 2023.
[12] Wei Liu and Yangyang Xu. A single-loop spider-yyp stochastic subgradient method for expectation- constrained nonconvex nonsmooth optimization. arXiv preprint arXiv:2501. 192144 2025.
[13] Mengdi Wang, Fthan X Fang, and Han Liu. Stochastic compositional gradient descent: algorithms for minimizing compositions of expected-value functions. Mathematical Programming, 161:419-449, 2017.
[14] Mengdi Wang. Ji Liu, and Ethan X Fang. Accelerating stochastic composition optimization. Journal of Machine Learning Research, 18(105):1 22, 2017.
[15] Yifan Hu, Siqi Zhang, Xin Chen, and Niao He. Biased stochastic first-order methods for conditional stochastic optimization and applications in meta learning. Advances in Neural Information Processing Systems, 33:2759-22777 2770, 2020.
[16] Wei Jiang, Gang Li, Yibo Wang. Lijun Zhang, and Tianbao Yang. Multi-block-single-pee variance reduced estimator for coupled compositional optimization. Advances in Neural Information Processing Systems, 35:32499-32511, 2022.
[17] Zhuoning Yuan, Yuexin Wu, Zi-Hao Qiu, Xianzhi Du, Lijun Zhang, Denny Zhou, and Tianbao Yang. Provable stochastic optimization for global contrastive learning: Small batch does not harm performance. In International Conference Machine Learning, pages 15760-25788 PMLR, 2022.
[18] Zi-Hao Qiu, Quanqi Hu, Zhuoning Yuan, Denny Zhou, Lijun Zhang, and Tianbao Yang. Not all semanties are created equal: contrastive supersuud learning with automatic temperature individualization. In Proceedings of the 40th International Conference on Machine Learning, pages 28389-28421, 2023.
[19] Yu Nesterov. Smooth minimization of mon-shooth functions. Math. Program.. 103(1):127-1g) 152, May 2005.

---

[20] Guy Kornowski and Ohad Shamir. Oracle complexity in nonsmooth nonconvex optimization. Journal of Machine Learning Research, 23(314):1-44, 2022.
[21] Damek Davis and Dmitriy Drusvyatskiy. Stochastic model-based minimization of weakly convex functions. SIAM Journal on Optimization, 29(1):207-239, 2019.
[22] Damek Davis and Benjamin Grimmer. Proximally guided stochastic subgradient method for nonsmooth, nonconvex problems. SIAM Journal on Optimization, 1908)19908-193,, 2019.
[23] Ahmet Alacaoglu, Yura Malitsky, and Volkan Cevher. Convergence of adaptive algorithms for constrained weakly convex optimization. Advonces in Neural Information Processing Systems, 34:14214 14225, 2021.
[24] Hassan Rafique, Mingrui Liu, Qihang Lin, and Tianbao Yang. Weakly-convex-conee conccve min-mas opti- mization: provable algorithms and applications in machine learning. Optimization Methods and Software 37(3):188-11221, 2022.
[25] Renbo Zhao. A primal-dual smoothing framework for max max-structured non-convex optimization. Mathe- matics of operations research, 19(3):1535-156) 2024.
[26] Zhichao Jia and Benjamin Grimmer. First-order methods for nonsmooth nonconvex functional constrained optimization with without slater points. arXiv preprint arXiv:22/2.0092) 2022.
[27] Quoe Tran Dinh, Deyi Liu, and Lam Nguyen. Hybrid variance-reduced sgd algorithms for minimax problems with nonconvex -nconvex-lincar function. Advances in Neural Information Processing Systems, 33:1 1009- 11107, 2020.
[28] Zichong Li, Pin- Yu Chen, Sijia Liu, Songtao Lu, and Yangyang Xu. Stochastic inexact augmented Lagrangian methed for nonconvex expectation constrained optimization. Computational Oprimization and Applications, 8771)1117-147, 2024.
[29] Ahmet Alacaoglu and Stephen J Wright. Complexity single loop algorithms for nonlinear programming with stochastic objective and constraints. In International Conference on Artificial Intelligence and Statistics, pages 4627-4635. PMLR, 2024.
[30] Zichong Li. Pin-Ye Yu Chen, Sijia Liu, Songtao Lu, and Yangyang Xu. Rate-improved inexact augmented lagrangian method for constrained nonconvex optimization. In International Conference Artificial Intelligence and Statistics, pages 21 2178.117. PMIR, 2021.
[31] Ellen Fukuda and Masao Fukushima. A note on the squared slack variables technique for nonlinear optimization. Journal of the Operations Research Society of Japan, 60:262-217 07 2017.
[32] Lijun Ding and Stephen J Wright. On squared-variable formulations. arXiv preprint arXiv:2310.0178. 2023.
[33] Marien Renaud, Valentin De Bortoli, Arthur Leclaire, and Nicolas Papadakis. From stability of langevin dif- fusion to convergence of proximal meme for non-log-concave sampling. arXiv preprint arXiv:2505.1417 2025.
[34] Peter Bandi, Oscar Geessink, Quirine Manson, Marcory Van Dijk, Maschenka Balkenhol, Meyke Hermsen, Babak Ehteshami Bejnordi, Byungjae Lee, Kyunghyun Paeng, Aoxiao Zhong, et al. From detection of individual metastases to classification of lymph node status the patient level: the camelyon17 challenge. IEEE Transactions on Medical Imaging, 2018
[35] Ziwei Liu, Ping Luo, Xiaogang Wang, and Xiaoou Tang. Deep learning face attributes in the wild. In Proceedings of the IEEE international conference on computer vision, pages 3730-3738, 2015.
[36] Pang Wei Koh, Shiori Sagawa, Henrik Marklund, Sang Michael Xie, Marvin Zhang, Akshay Balsubramani, Weihua Hu, Michihiro Yasunaga, Richard Lanas Phillips, Irena Gao, et al. Wilds: A henchmark of in-the-wild distribution shifts. In International conference machine learning, pages 5637-5664. PMIR, 2021.
[37] Shiori Sagawa, Pang Wei Koh, Tatsunori Hashimoto, and Percy Liang. Distributionally robust neural networks. In International Conference on Learning Representations, 2019.
[38] Gao Huang, Zhuang Liu, Laurens Van Der Maaten, and Kilian Q Weinberger. Densely connected convolutional networks. In Proceedings of the IEEE conference on computer vision and pattern recognition. pages 4708, 2017.
12

---

[39] Victor Sanh, Lysandre Debut, Julien Chaumond. and Thomas Wolf. Distilbert, distilled version of bert: smaller, faster, cheaper and lighter. arXiv preprini urXiv:1910.01006 2019.
[40] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE conference computer vision and pattern recognition, pages 770-775 2016.
[41] Barry Becker and Ronny Kohavi.
Adult.
UCI Machine Learning Repository, 1996.
DOI:
https//dbo.org/11.2444//////
[42] Julia Angwin, Jeff Larson, Surya Mattu, and Lauren Kirchner. Machine bias. ProPublica, 2016.
[43] Fisher Yu, Haofeng Chen, Xin Wang, Wenqi Xian, Yingying Chen, Fangchen Liu, Vashisht Madhavan, and Trevor Darrell. Bdd100k: A diverse driving dataset for heterogeneous multitask learning. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pages 2636-2645, 2020.
[44] Junyu Zhang and Lin Xiao. Multilevel composite stochastic optimization via nested variance reduction. SIAM Journal Optimization, 31(2)1113-1117, 2021.
[45] Maurice Sion. On general minimax theorems. 1958.
[46] Zhize Li, Hongyan Bao, Xiangliang Zhang, and Peter Richtarik. Page: A simple and optimal probabilistic gradient estimator for nonconvex optimization. In Proceedings of the 38th International Conference on Machine Learning, pages 6286-6295, 2021.
[47] Heinz H Bauschke, Patrick I. Combettes, Heinz H Bauschke, and Patrick I Combettes. Correction to: convex analysis and monotone operator theory in Hilbert spaces. Springer, 2017.
[48] Zhishuai Guo, Yi Xu, Wotao Yin, Rong Jin, and Tianbao Yang. Unified convergence analysis for adaptive optimization with moving average estimator. Machine Learning, 2021.
[49] R Tyrrell Rockafellar, Stanislav Uryasev, et al. Optimization of conditional valuc-at-risk. Journal of risk, 2:21 42, 2000.
[50] Jianmo Ni, Jiacheng Li. and Julian McAuley. Justifying recommendations using Eestantly-abblle reviews and fine-grained aspects. In Proceedings the 2019 Conference on Empirical Methods in Natural Language Processing and the International Joint Conference on Natural Language Processing (EMNLP. IJNLP)), 2019.
[51] Robin Vogel, Aurelien Bellet, and Stephan Clemencon. Learning fair scoring functions: Bipartite ranking under roc-hased fairness constraints. In International conference on artificial intelligence and statistics, pages 784-792. PMLR, 2021.
[52] Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agarwal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, al. Learning transferable visual models from natural language supervision. In International conference on machine learning, pages 8748-8763. PMLR, 2021.
[53] Christoph Schuhmann, Richard Vencu, Romain Beaumont, Robert Kaczmarzzyk, Clayton Mullis, Aarush Katta, Theo Coombes, Jenia Jitsev, Aran Komatsuzaki. Laion-400m: Open dataset of clip-filtered 400 million image-text pairs. arXiv preprint arXiv:2111.02114, 2021.
13

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: We have detailed explanation all items described in the abstract presented in the paper.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
is fine to include aspirational goals as motivation as long as is clear that these goals are not attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification:
We discussed the limitation of this in the conclusion section.
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g.. independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged judgment and recognize that individual actions in favor of transparency play impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
the paper. The authors should use their best
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer:[Yes
14

---

Justification: Assumptions are presented in the main paper. The detailed theorem statements and proofs are presented in the appendix. Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Question:
Answer: [Yes]
Justification: We provide all the information needed to reproduce the experimental results in the "Numerical Experiments" section and related sections in the Appendix.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution a specific model and empirical evaluation, it may be necessary either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of large language model), releasing of a model checkpoint, or other means that appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how reproduce that algorithm. (b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome deseribe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
15

---

Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, as described in supplemental material?
Question:
Answer: [Yes]
The data we used are public datasets. Code is included in Supplementary Material.
Justification:
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips... cc/
for more details.
While encourage the release of code and data, we understand that this might not be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https
for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing much information possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code permitted.
6. Experimental setting/details
Does the paper specify all the training and test details (c.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Question:
Answer: [Yes]
Justification: We provided all the implementation details in the "Numerical Experiments" section and related sections in Appendix.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment statistical significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [Yes]
Justification: We ran multiple trails for each experiment setting and present the error bars in the plot or table. The error bar in this paper is 1-sigma standard deviation over the random seeds.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by emor bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
16

---

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
It should be clear whether the error bar is the standard deviation the standard error of the mean.
It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not show in tables figures symmetric error bars that would yield results that are out range (e.g. negative error rates).
error bars are reported in tables or plots, The authors should explain the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Question:
Answer: [Yes]
Justification: Compute resources information is provided in the Appendix section.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount compute required for each of the individual experimental runs as well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make into the paper).
9. Code ethics
Question: Does the research conducted the paper conform, in every respect, with the NeurlPS Code of Ethics https: httpas://eeuuiiiee hww.rpsu...../ cwtturps..ccpuulleee ctc/ppublc/thiiccdddeell....
Answer: [Yes]
Justification: The research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics.
Guidelines:
The answer NA means that the authors have reviewed the NeurIPS Code of Ethics.
the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (c.g.. there special consid- eration due to laws or regulations in their jurisdiction..
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification: There is societal impact of the work performed.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
17

---

If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification:
This paper poses no such risks.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets (e.g.. code, data, models), used in the paper, properly credited and the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification:
For all assets we used in this paper, we cited their original source.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset used and, possible, include URL
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
18

---

If assets are released, the license, copyright information, and terms of use in the thcode has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
package should be provided. For popular datasets, paperswi
latasets.paperssi
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information not available online, the authors are encouraged to reach out to the asset's s creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification: This paper does not release new assets.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/moode as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsoureing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification: This paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based the requirements of your country institution) were obtained?
Answer: [NA]
Justification: This paper does involve crowdsoureing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing research with human subjects.
Depending the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
19

---

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors adhere the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
16. Declaration of LLM usage
Does the paper describe the usage of LLMs if it is an important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, formatting purposes and does not impact the core methodology, scientific rigorousness, or originality of the research, declaration is not required.
Question:
Answer: [NA]
Justification: The core method development in this research does involve LLMs as any important, original, or non-standard components.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, non-standard components.
Please refer to our LLM for what should or should not be described.
20

---

A Proofs of Technical Lemmas
A.1 Proof of Lemma A.1
Lemma A.I. Under Assumption A.),, F() in ① is pr-weekly convex with pr Cg
DgCg. If is by smooth, then Fx (x(w) is Lp-smooth with I.p C3 max 3, 1-x0; Cylig.
Proof. Consider any w and w' in RR, any i E {I...... n and any V @ft(t;) at t; = g((). By
Assumption 3.2 A1), v. (i( is livell I -mooohh in w and si is P Py-weekly convex, we have
f(lg((')) 2 [i()((w)) + v, (e:(w)) 9:(w)) Pf
(ls((w)) 9:wl]]
2 f((a:(w)) + vv (9:(w) 9:(w)) CJC l| - w|
2
2 [[) V. (g)(ww)( - -w) livelleTig 1| - -w1
PrC
||' - w|?
2
2 f((a(w)) Vg.(w)) (w -w) VTlCGLo + C lw'-wll
2
where the second and the last inequalities hold because lipschitzness of fi and 9s from Assump- tion 3.2. the third because v. (i(w) is |vs Lggsmoothh -smooh in w.
By the week convexity of f. and lemma 13 in 333, we know that fr. is Lis = max Ps }- X T-
smooth. Since f.,. (w prox (w))/A ( (w)). we have [w))]] S
Cj by Assumption 3.2 Therefore, by Assumption 3.2 again, we can easily show that
(f.x( is LL + C.I,g))-ipsshhii continuous (see Lemma 4.2 in 44 so
(Fx(w) = E (l)a((())VV
A.2 Proof of Proposition A.2
Proposition A.2. When fo is convex, the smoothed problem smoothing, i.e.,
is equivalent to the Nesterov's ss
1 Σ v g:(w) ((1) 11
A
min max yel n2
(11)
where Ω dom(sf)) 17 (16) + A = Si (w:). In addition, fi,1(0) S f((9) S f(,(g) + A02
-
Proof.
We prove by deriving 11 from
1
F2(w)=
(w(())
min yel n
Σ M(y)) 2X
lly: - (/(
Σ yenl" 2EZ" n d=1
min max-
y, 2, (z.) + Ily. (w))
1
2x
1
max min 2E2" yEl
Yy 2: (2,)) + 2X
lly. g:(()l
where the dual domian 2 = (2. E domsi) C (z: ||2|| < Cy)is bounded; the second equality from definition of JT,,, the third equality holds from convexity of Si and the fourth equality holds from Sion's minimax theorem [45] and the convexity of fr
21

---

Note that the the inner-level minimization problem can be solved exactly as y = -:(()) Az since it's a quadratic function, we plug y? into the RHS, then we have:
R
X
Fx(w max 2=2
Σ z (w() (z.) 2 l|z:1?
(12)
72
Besides, note that ft fx(xxxx u) = fi(u) + 2x l|u x|2 is jointly convex, fLA(x) ming fi(u)
2X u x|| is also convex. Then we have
Σ (ix(s(()))
= Σ 2 g:(w fix(zi)
i=1
Compare this with 12 we have
[x(Zz)== = fi(z;) + A 1|2:1?
2
Now already have that
X |1:1| and fi(g) = -max 2 (1,)
f()(g) = max n g f1 (2;) z,E2
2
2,E2
then (1,1()) S f((g) naturally holds. Besides,
A ||2:1
||2.1| S s(,( + A
[i(9) max 2 9 (Z:) S max 2, 9 fi 2,62
-
max 2,
2,
2
A.3 Proof of Theorem 4.2
Proof of Theorem
4.2
Suppose is an e-stationary solution to Fa(+) with À e/Cy. It holds that
[[FF(w)]]=
(ee(W) proxAs, (g: Vg))) Sc. Vi = nd
n
By (2) and Cy-Lipschitz continuity of fú, we have ([g((w) (oa & AC S c. Let t, = prox >,. (w)) and y; = (g:(w) xro (g:( (w:(w)))/> E @pr (0:(w)) = (t.) is approximate
holds that ||t: gi (w)l S E, i= 1,..,, n, and Σ (w)(w S e, e-stationary solution to the original objective F(-) by definition.
A.4 Proof of Theorem 4.4
We need the following lemma.
Lemma A.3. Suppose Assumption 3.2 (AI) or (A2) holds. If Amin (Vg(()VVg(()]] 2 6, it holds for any w' satisfying l|w' w|l I0,4,
that Amin (Va(w')V
Proof. Consider any w' satisfying ||w'- - -wll- S 40,7," Note that
Amin (Vg: uER*1
min VV(
Let
Uw E
min u [VVg:W'VVg u. ue |||||=1
It holds that
- Ami (Vg. (V Vg( )
Amin V W)
Su Vgg ()) VW) uw Th V ( Vg::'''
Vg) -Vo(W)Va( uw
Vg((W)
<2C, I. wwwwll)
22

---

This implies
8
Amin (Vo((")Vo(w))) 2 Ami (Va(w)Vo((w))) - 20, 1|w - w/l wwwll|l-- 8
2 2
Proof of Theorem 4.4 4.4 Suppose w is an approximate e-stationary solution of i with S
min c 16nC2 f. There exist ti,..., ty and ys E 0f((t.) for ., n such that ||t, g:will S e, i= I,...,, n, and || #/ Ei_ Vai(w)y;ll se.
8²
Consider the following optimization problem
min q(v) = 2n
1
Σ |t: [(v))
(13)
We want to show that the optimal objective value of (13) is zero and there exists an optimal solution w' such that || |w - w| < nC.ne
By the condition on c, holds that 4C,ne S 40, holds that
Hence, by Lemma A.3 if ||v w|l S 4C,ne
6
(l - tt|l²= 5 g(v).
[2(v)] n²
V -gi(v)-ti]
-til
2
2n²
Moreover, always holds that
n (l.(v) - -til2 -+ll1==2000 =2Cgg(v).
[Vg[v)]
40,ne Let V, for = 0, l.... be generated by the gradient descent method using Vo = w and a step size of - i.e. V = V, v/v).. Suppose ||v, - w/l < 4C,me for 8 = 0, 0, (which is true at least when 0). By the standard convergence analysis of the gradient descent method for smooth objective function, we have
Let Ly be the smoothness parameter of q(v) the compact set {v Rd ||v w||
T,
g(v(+1) Sg(ve) + (Vg(ve), V6+1 - ve) Lg 2
| vel1
1 s((v)) 8
Sg(vi) 2T.
2I. v((v)).
Applying this inequality recursively gives
1
2La ) q(vo).
Vg(v.)) S g(v))) =
2C? 9
Recall that ||t, - g: (w)ll < such that q(vo)
1
Σm ||, gi(w)]l s S.. By the triangular
2n
inequality, we have
6 1 e
4n1Cglyn
4Cgne
|| w/
nl Vg(vv)l < Σ n, 1
2La
8
By induction, we have [|v,
-wll < 4G for t 0, l.....
Let w' be any limiting point of {v.]-2o- holds that || - w/l 5 1Gyne and
g(w)) E lim (1- 2L, g(vo) = O.
t
23

---

meaning that w' is the optimal solution to 13 with the objective value equal to zero, which implies t, go (') for i 1, ,.
Since w is an approximate e-stationary point of ( and yı E @fi(t;), we have
(lilo(()))
0fi(t,)
1 Σ (V(()) Vel(w))
1 E Vg
+
1/
Σ 0f((t)) (Vgi(w))
(Va(() -Vo Vol(w))
Sc+ bg ||w - w/l (1+ 4C
where the last inequality by Assumptions 3.2 and thus w is nearly Sttatinary solution of
This means w' is an O(e) stationat solution of
B Convergence Analysis of SONEX
In this section, we present the proofs for Theorems 4.7 Let E, be the conditional expectation conditioning B and 2 for 1,..., and following quantities are defined based on the notations in Algorithm 1 for = o,...,
0, 1.....f. 1. To facilitate the proofs, the
V |[ (w:)l))
(14)
U1+1 =
||ut+1 g(
(15)
u, =- Un.el
(16)
g(w.) = -[g1 (w::) ... Ir (w:))
(17)
B.1 Proof of Theorem 4.7
Lemma B.1 (Lemma 2 in 46
Suppose S 2L 1 in Algorithm with LF defined in lemma A.1 Ir
holds that
(W & Fx(w:) + " V1+1 " VV:((::) ||
(18)
where V441 is defined in (1
A recursion for V4+1 is provided below.
Lemma B.2 (Lemma 9 in (2)). Let Ve+1 be defined as in Suppose 3 S 3 in Algorithm 1 It holds that
217
|e+111
3L30?
E [V+2] S (1 BBE
E [|u - 11+
3
n
28 18 of)
+ 55833CBB1112)
(19)
min Bi B2}
where C7 C2 + of B and U4+ u, and g(w:) are defined in 15). 16 and 17
respectively.
The following lemma provides the recursion for the estimation error U4a1 of the MSVR estimator of g(w).
24

---

Lemma B.3 (Lemma 1 in [6].. Suppose 7 is used in Algorithm 7 It holds that
n-Br
+ (1 - 7) and S and the MSVR update
Bi(1-7)
E [1242 (1 BB
8n B1
E W wrll2+ 2B1
DEIU++1]
(20)
nB2
The following lemma handles the term || uell?
Lemma B.4 (Lemma 6 [16].. Under the same assumptions as Lemma B.3 it holds that
4B
9n²
2Brr208
E Blllu+22 - u -
ElU++1]
+
2 B W141
w.ll +
(21)
B1
B2
Now we are ready to proof Theorem 4.7
Proof of Theorem 4.7 Note that |We+1 well /
11.. Let P be non-negative constant
to determined later. Taking the weighted summation of both sides specified in the following formula
20 and (1 as
3
x
+
nB
we have
+ (Fx(1(1)
EV++2 + (P- (P-51e -5I707) EU1+2
B
-1-1)
1
(P(1-12)
7B1
121,,79021
(wi) +
EVH++++
+
BU++1
3
12
Bn²
1
1
22
27
E
||
4
32
B1
B13
CC +
2Br
6
+
(w:))
E
min B1, B2)
nB2
3nB²
2
We choose 77 such that
2T
SnO2p12
27m13,09C31 0.
(22)
4
82
B1
B13
Moreover, we set P to be the solution of the following linear equation
((1
BB1
1
P
P P -.
Bn²
which means
12L30792B.
p= BB
51307+
Bn²
Combining the results above gives
1 3 BV + (P-5L3C EU
EFx(w.+1)
11
EFx(w.) EVi-1+ (P- (P-51}c)) EU++1
"
CB + of)
2B 6L B1
nB, +
E ([
min B1, B2}
BnBB
2
Summing both sides of the inequality above overt 0, 1, 0,1...., T and organizing term lead
2
L (F( RA(WT))
1
(V.)(w))) S
EV + (P (P- 51 RU,
T
3
4309 (C2 +ff)
4BiT2OiP
1213079708BB
(23)
min B2}
nB2
BnB²
25

---

Note that Lp O L, = = and P= (
Brye + n We require that all the terms in
the right-hand side to be in the order of ((c2). To ensure 13C7 +of)
min B, .B} = Oc22), we must have
((, To ensure 4B1>0o] 0 Bgr²
nBg
7
+
n B2
B1g
O(²).. we must have 7 =
12
O (B With these orders of B and 7, also have P = 0 Bi and
(
3nB2
Therefore the last three terms in (3 are (
To satisfy
22 we will need to set
B1 B nP
= 0
min
LF
L
n
Bi min Bi, B23 (2) 1)
Bi BL
B2 2
min
min B1, B₂ c2
n
n
B2
0
n
We also set the sizes of Br,2 to be O ) so that EU1 = O Also, it is easy to see that Fx (wo)
Fx (1) and EV1 O1)) Then the first term in (3) becomes
-
o
- 1 + P(.
P
1 -
n2
n
+
BI B2
min (n, B2)
B1 Bge
which will become O when 7' 0
Big
C Convergence Analysis of ALEXR2
To establish the complexity of Algorithm 2 we first present the convergence result of ALEXR 5 which is the inner loop of ALEXR2, and it shows the complexity of Algorithm 2 to produce 2; such that /2: prox,F, (w,) || < O(e)..
Theorem C.1 (A Variant of Theorem 1
Pg- Py-weakly convex and ydgC Suppose Assumptions 3.2 (A3) and 4.5 hold. For any > 0, there ex- B
5 when g. non-smooth
ists 0 E (0, 1) with 1 - = =((c') such that, by setting (40 1y and 7 = n(1 1)-
the inner loop of Algorithm 2 (i.e., ALEXR) guarantees
EE 22 prox (wr)l1 < é for any t after
K, δ n Bi iterations.
Coy B. B.B.C
nyy nLyor
Cyoî Baddipge
C,C² + C
Ō not B.Baee
Bivdipge
Jaip
In this section, we provide the convergence analysis for ALEXR2 in Algorithm 2 Lemma C.2. (w) is LF.., -smoohh with LFx., max -oppy)) and pF. VdrC
Lemma C.3. The inner min-max problem in 6 G Prs -Prt)--rraaa >-strongly-coomee and >-strongly-concave problem when v E (0, (F....
Proof. As discussed in Section 4.2 the term y. gi (w) is pr, -weally convex with pF, Vdl
Therefore, when E (0, ), the objective in (6) becomes strongly convex in z for any fixed y md
To establish strong concavity in y, we note that the term Jt is the Fenchel conjugate of a -moooth
function ft. It is standard result in convex analysis (see, e.g., Theorem 18.15 in 47) that the conjugate of an L-smooth function is -strongly concave. Thus, f? A-strongly convex.
The remaining terms in the objective are either independent of y or linear in y, and hence do not affect strong concavity. Therefore, the objective is strongly convex in z and strongly concave in y when / E (0, PF)
D
26

---

C.1 Proof of Technical Lemma
Note that our proof for theorem C.1 is almost the same as that of theorem 1 in S we only show the difference this and the following subsections.
Lemma C.4 (Generalization of Lemma in ALEXR (5) to weakly-convex gi). Suppose that Assumptions 3.2 A2). 4.5 hold. Then, the following holds for Algorithm 2
n
E
gi(W*++)> gi(W.), Met1) E (Gk, Wk+1 -w.)
71
1=1
c
c703 LgCg
+
vdiC,Lg
B1
wall2 + 2
w w'12
(24)
1
2
7
Proof. For this lemma we only highlight the difference to the proof of lemma 7 in ALEXR 5 We
define Ao = E. Eient (wi: []] Vi+1 Vi=i-Vee [W)() 0
(8)
(Gh; Wk+1 - w.) ge(w (e(w++1) - (w.), w )
(w((W<+1) - Dett
(g((Ww) g.(w.), (
7
n2
+
[W(( Det1 + D.. W. Wx+1
n
smoot
9i (W - (wx), mh-1) +
[g:(w:)] T
Tot11 We - w.
11
Lg
+
Ww - 11
Uotti 1)
+
Σ
(w Toht + Ak. w. - Wh+1
2
1=1
Detti (w g((ww)) +
V (wx)]
1etr. We - wa+1
VTiC,Ly
+ (Ak. wi+1)
| w*1l
(25)
Note that have extra term in blue in 25 comparing with lemma 7 in (5) due to the weakly convexity of gr- And the remaining part is just the same so we omit it.
Lemma C.5 (Generalization of Lemma 8 in ALEXR weakly-convex 9:). Suppose that assump- hold. Then, the following holds for Algorithm 2
tions] 3.2 A3), 4.5
gi (w<++1)) 9.ww),, Vlot1) E (Gx. Wk+1 - W.)
n
1=1
c7c² + 4C702
1 ++
| - wxll3 + VTiC
+
W w w=1
(26)
+ 1
4
2
VTCC | w 121 in the upper bound for
Proof. Similar to lemma C.4 have extra term weakly convex gi-
"his inequality also holds with assumption A.3 so lemma C.5 can use the same technique to handle weakly convexity of
27

---

C.2 Proof of Inner Loop
Proof of theorem C.7 we have
Similar to theorem in [5] but use lemma C.5 instead of lemma 8 in ALEXR,
E Wk 11 y+) L(.., [k++)]
1 + (1
B.)
1 +
E (W (10.)
3/4k]
E (yh, yk+1)l
B1
B1
1
1
1
E |w. -well
E
||
w.
W 11
2n
2
2n
2v
A2 + AaO
IF U (1;(ykk+) (Jk++1, mnm
g
+ v
E |
wxll2
OC
E ||
||
w
ym
An
20
2X2
2X3
C90
2(1 + 20
+
+ 407009
+ [T[T+T k+1 kT
B2
B
(27)
BX(1 + +)
1 +
Define TV = E |w. 2
11 and x = BU (v-, w.). Note that L(W1-1,, --..-..-........ Mh++1 2
O. Multiply both sides of (27) by 0-k and do telescoping sum from k = to K, Add 170 Kr K.W. to both sides.
0 K. TW
K,-1
Σ 0-h ((++ VTCCINNTT + ( (1-
11=(1-1 B1 )) r% - 0 -
k=0
- + 1 ry
++ + {
+ 1) k+1 ETx+1] k++1
7
r-K:Tw
2(1 + 2220)a0
c7o7 B2 +
+ 4070g
K,
Ke
B.
Σ g
+
7
AB(1 +
1+1
v
K.-1
1
(A2 + X36)
EUUFF (Uk-1.
y
An
K
1 +
C2
C²
C2
0-k
EE [|K+++ - wxll2 - 0 g-K++1
E |Ww. -x.-1l22
21
4
2
2
2\g
(28)
and 1/0> 4VTTC such that > 1+V 1+
Let 2 0/v Vd
1-0
0 1
and 1 +1 S m(188
2
7
1 1+1
such that
>
Then,
$ +1
K.-1
( 1 B1 )) ry 8ET,) - ((2 + 1
1 r + _ + +1)r4 - -ET+22)
Σ x-
+
7
k=
< { + VTTCipo)TT + (++ (1- +++11-1
B1
n )) ry - OET, - 0 ki+1 (1 + + + + - -ETx,
By setting w-1 = Wo. we have To = 0. C, [.. wk, lly. yk,ll |l
Besides, we have l'k
EL [g::(w.)) - (<(w,, | y
(i)
S
ca
|WK w Z f- (y+, JK.).
Note that the first term in the RHS here can-
2X
28

---

celled out with the last term of RHS in 28 Thus.
8-Ki KK
1 I - +
B1
))
/d/C/p,)
To +
- +
1
r
7
11
?
+
1
A3.5
g- 1xyy
17
+
( +
2Am²
K=
n70 K.
2(1 + 00
Cjoi
B2
+
B,
+ +0000g
K,-1
Σ k=0
0
K,-1
Σ
0
( 1
(12 + +20)
BEU;; (1h+1, 36c]
AB(1+ 5)
+ v
7m
An
k=0
K,-1
0 0-k
1
1
+
9
C2
E [[x++11 wx
(29)
21
4
2
2
To make the V terms in
29 be non-negative, we choose X2 x x3x x
C,VBIA
while ensuring that
VavdiCypp
V Vnvāaccppd
An
A
VBAA
1
Y S 0
A
0
75O
A
AB
A2
)30
CaVBI
Cov CV PoCs
(30)
By selecting 1) = 0/v- Viu-vd;Cypp'
1-0
1
B,
and 1 = AVdiCspp we have that 1 + (1
B1)
7
m[1-0]
8,0 m[--8)
( 7 + (1-
B1) M(1 - 0)
M. K. S p(1 + AV Mo
g
8/v Cipo
2(1 20)00
0707 Br
C.C.
+ 40709
(0-K: - -1)0k
(1 0)μ
B.
AB(1+ +
1 1+1 1
0-1-1
8/v - SdiC
c7o7
+ C302
B1 0p gk. rg- 2/10
2(1 + 20)08
+ 10702
(11 /u d. (10pp)
B2
B.
5/v
AB2(1 +
1 ++
2(1 + 20)0 +11 \B2( +1)
(Coof Bg
S.G.2
40703)
S s(1 + (TCP +20 +2
B1
We select
B 1 n
BIAVTiCIPo AB,Bie on
BIVTTPPEE Cjo7
Bivdangee C₁G2
VTapge
1-0 0
Cg
n
C₁O2
to make
30 hold and
2(1 + 20)a3 AB2 +1 +n
(ffff Bg
C7C%% B1
+ Ac700) <
2(1 + 20)
0)azn
(1-
Goo
B2
B,
+ 40303)
SC.
XB,B1
VT
Besides, we show that TV can be bounded by constant..boundeeee of Ty is already guaranteed by assumption 4.9 Note that w = prox,F, (We,)) we have:
:. (DFX(W(W) (w" Wo)
wo w =VOT((W")
|o w 11 - { ( 2/ =
1. TW =
2
29

---

Since Lg = O the number of inner loop iterations needed by Algorithm
to make
k. K. S is
K, Ō
n
nL, + oL BLVdiPgGG B2Bie
Cjoî Baydiage
C Bı Bivdinge
C₁O. Vdipge
+ CaN
+
BI
where Ó(-) hides the polylog(C,Copg/e factor and the green term is the dominant term so K; nog Ō
BrBiee
C.3 Proof of Outer loop
Proof of Theorem 4.10 Let w-1 Wo. By lemma
C.2
(wi( is LFx.- -mmooth. Since 8 1
we have
Vi+1 (wil(
=|u [v... (1
(2 (W(: 2.) -Phholm))
VFn. B
S(1 + 2 (1 B)²
(u+ 2 llv. (W-
+1+ 3 [VF3,, -VB.--)) (.B..(
(1 / )0²
2
- (we 2c) - -VFh.(ww))
1
8
3L3
1 (w. 2.) VFx (wi)
<(1 2
//v. VB. (w--111
|| - - 33
3a 3a=Li
1 (we 2c) (we)
5(1 2) 1/v. VF (wr-1))l
lIv + 33
(31)
8
Since S 21 by Lemma 2 in 46
we can obtain result similar to
(8) in Lemma B.1 that is,
(wo())+1) S (x,(
a
V - Fw.((w))l a
a VF (w.) Il
Vi+1l (32)
2
2
4
Multiplying both sides of (31
by a and adding them to both sides of (32). we have
2 (w.)ll
{Fxo(we) - (W +
)
a
(Ivv - V FF
|| - (W:.)()))l)))
3
- 11
a.
|eel1
+ 3a
1
(w. 22.) -VFxo(W)))
B2
1
<Fx,(We)) Px.((we
-
/
[1vv
VF2 || VFA.
3
a
4 (Ivell2-lvvealll
30 y
122 - prox (w:)l1
where the second inequality is because V FA. = 1 prox (wi) and a S 2LFx
89
so
1 A..
By Theorem C.1 for é' uu?.² there exists 0 (0,1) with - 0 =((2) such that, by setting 17
1-0
and 1100 inner loop of Algorithm L.. ALEXR) guarantees E ||2, prox, (w:)
2,2
for any after K, = mLy
+ 4 iterations.
18
B, Byc
30

---

Recall that Vo = 0 and w-1 Wo- Summing the inequality above over t = -., T and dividing
both sides by af 2 we have:
E E, FF (w+)))
E
V
T 1=0
2F (wo) FA.((WT))
(3 -
V FA (wo)l?
T-1
EE
prox, (we)ll
af
T
D
t=
2B( 2B(Fxx(wo) - (WT)) aT
( | FV (wo)l
--1
6²
T
3
By setting 7 O(E-2), we'll have E[dist(0, x vF, (w-))) v e² meaning that w, is nearly c-stationary point of (4) in expectation. Recall that L = 0( so K; = O 3)- Hence, the
total complexity for finding such solution is Ei Kr = 5)
Proof of Corollary 4.7 17 When is L smooth, Fx(.) is I. ---soooth where LF = = is
defined in lemma
A.I By Theorem 4.10
E v Fx (prox (w-) (w+))ll Bllprox (W+) w- 0(().
According Theorem 1-0
C.1 for c' = Oee*), there exists 0 (0. with 1-0 = O(e*) such that, by set-
ting 1) 9 L, and 7
(1-0)m in inner loop of Algorithm 2 i.e. ALEXR), w, =ALEXR(w,, K)
B:
satisfies "E || Fr (w>)] < e and thus
E W prox DFA (w+) 50(2)
after K Ō -0( B.
nhe
4 iterations. Then have
V S V (prox,, (w+))ll + Lp prox (w,) - w- w
S O(e),
which means w, is an e-stationary solution of (3). By Proposition 4.4 w, is nearly e-stationary solution of (T). By Theorem 4.10 w, is found within complexity O and w, =ALEXR(w,, K) has complexity K δ
aly
+ 2 The total complexity for computing w, is still
B1 Bgck
δ B. + }
D Convergence Analysis of SONEX and ALEXR2 applying to Constraint
Optimization Problems
In this section, we present the complexity analysis of Algorithm 2 when it is applied to to solve the constrained optimization problem 7
D.1 Proof of Proposition 5.2 and its variant
Proof of Proposition 5.2
By the definition of fx(+), we have
f-x(-) = min min
Suppose w is a nearly c-stationary point of (9). which means there exists w such that |/w w/ Sc and dist(0, 0Dx (w)) This means there exist (w()) dg. (w) and
1 Vfl(g(W)) m
1 = Am
1i
min (w)]+,
for ..... such that
[h(w) + +JWW Sc sc
(33)
31

---

where J(w [h1 (w),.. hm (w)l E R and v = (V1, vm) E R.
Suppose max Fi( (W)
Then there exists such that [ox (w]]+ e. Recall that C p
We have
1 min[[g((W]]++ App Am m
1 P
Uk =
Am
By Assumption 5.1 we have
||||| mmin (J(W)) bo (wrr)) óp
[w)) + J |l 2 [J(W)]]| [[ho(w)]] 2 mmin
Co L,
which contradicts with 33 Therefore, we must have
max i
gi(w) S E.
(34)
Finally, when max;=1..... (w) S 6, we have for Vi = ,12,...,1.
ge(()es Σ [o((w)]+ min (W))
99 ( (W(W)s
[9((w)] 3 O(c). (35)
Am
With 33
and 35 w is an k-K point of so is nearly e-KKT point of
Proposition 5.2 is given when w is a deterministic nearly e-stationary solution to 9 However, the solution w found by our algorithms is only nearly c-stationary solution to (9) in expectation. This means there exists w such that E w w| S and Edist(O,, x)))(w)))))) I e. In this case, we cannot prove the four inequalities only show that the first two inequality hold in expectation while the last two hold in high probabilities. We present this variant of Proposition 5.2
the conclusion of Proposition 5.2 hold deterministically, Instead, we can
below with its proof.
Proposition D.I. Suppose Assumptions 3.2 A3) and 5.1 hold. pf m(C,+1) and X f. nearly
e-stationary solution w to (9) in expectation is also a nearly ((e)-KKT solution to the original problem 7 in the sense that there exist w and v;
1 ....., such that [[w w|| S and Edist(O, 0go (w) + Em (g.(w) (w)) & Oc)) and holds with probability (c) that max (:( S (c) and 9: (w)vl O(c), 1,2,- m. 1
Proof. By the definition of fx(+), we have
V = min min(())) Ap)
Suppose is nearly c-stationary point of (9) expecation, which means there exists such that E | w|| < c and Edist(O, ( S This means there exist h, h,ww) Ög: (w) and
1 = vsa(a((W) = Am
1
min min((g((W))
for 0, 1,..,, such that
E hho + J(W) Sc
where J(w) [hy (W), hm (w) Rd and "m,"",,,mm) E m.
Suppose max;=l,.., (:) e. Then there exists such that [gx (w]]+ > e. Recall that E
We have
1
1 P.
min min((g Ap} Am m
= Am
By Assumption 5.1 we have
- [V)]( Omin (J(W)) ||z|| [bb(w)]] 2 8p
[o (W) J(W)
C,
m
32

---

Therefore,
2 ho(w) + J(w)
=E ho w) + JWW)
max ge(w) > o) Prob max 9:(W) 2() 6)
i=I....m
mll....
+ E | + J(w)
max g:(w) S
Prob
max gi(w)) S c)
:=....m
Öp Cg
Prrbb max gi(w) > e) il....mm
As a result,
Prob max (w) E) Se 0p C. =(e)).
(36)
ml....m
Finally, when max;=l,.... ge (w( S e, have for Vi 1, 2, , m
min { w +. Ap}
(w:(W) S
0e (w)jss = Σ g (w)]+
[w((w)] + s O(c).
Am
i=1
It then follows from 36) that for Vi 1.2,..., m
Prob (19: (w):s 2 ((())s((()))
D.2 Sketch Proof of Theorem 5.3
Since the proof is almost the same as that of Theorem 4.10 we only highlight the difference in the proof.
We first consider the case where {9r He are weakly convex. We will slightly modify Algorithm 2 solve the following problem
to
min wERd
(x)(w)=
min #CRd
T xz(Z)
+
1
||z
w/l w((}]
(37)
2v
min
min
max
90 (z)
+
1 m
Vigi(z)
A yi
+ 2v
1 ||
w/2
(38)
welt
zeled yo (0
m
where 2x. defined in (). According to Proposition D.2 in [6], the second (compositional) term in &x (w) is (pp,))-eakkl convex. Hence, @x(w) is Pa. -weakly convex with pv,, Pg PPg
If Pa (x,(w)) is ( pis,) )-strongly convex. As done in Section (4.2) in the tth iteration of
Algorithm we apply ALEXR to solve the inner min-max problem in (38) with w = We, namely,
1
A yig((w)
min max yE m
Lzzzy) = (z(z)
+ m
of
+
we
2
2v
we only need to replace Gt, in Algorithm
to
GE G Gl Ögo Zr B +
Σ
(0g. (z,,
B1,21 (0)
Mektr
(39)
1
*EB
where Bi is a batch of data sampled from the distribution of Eo.
We then consider the case where {9i Am are smooth. this case, we can directly apply Algorithm
to g which is different from 33) only in the extra term go. To handle this difference, we only need
to replace G2 Algorithm [ to
G G. + G. = Vo Wr, (b,)) +
1
Σ
(40)
|B4
iEB
where Bi a batch of data sampled from the distribution of Eo- With these changes, we can prove Theorem as follows.
33

---

Sketch Proof of Theorem 5.3 When {g:) are weakly convex:
For simplicity of notation, we drop the dependence on /. in the notation of all variables and parameters in the proof below. First, we need to modify Lemma in [S). We have
Le(Z++1,8) (z Vk+1)
(41)
U 1: (x) 7 +
UJ (x41) _
--
yy (FR+1,
m
m
m
m
g:((k+11) - 9k Yi
(()
(1)
||z Zu 11
(42)
171
2n
1
1)
1
+
1|z 2x+1ll
|xx+1 Z5 12
21
4v
2nj
(2((x+1)= -9i(z), 5k+1 (Gk, Zx++1 -z) z)
(2)
(43)
m
+ (Z (0(2)) -G2 Zk+1 -2),
where the blue terms above are included due to the modifications to Algorithm
2
Next, we modify
Lemma 8 in 5 to handle the additional terms above. Denote that Ag == 96 (zk)
G2- Note that
GG.. Z++1 2) = (zx), Zk+1 - z) + ( Zk+1 z)
where (Z+) E 0g(z))) is a subgradient at zg. The term ( ( Z++1 z) can be handled in the same
way as in (24) (5) Due to the weak convexity of Do. we have
(o(z) go (zx) 2 (g6(Zz), z 2
Pg 2
Zz 113.
Then, the remaining blue terms in 41) can be upper bounded as
(0(+++1)= - 9o (z) - (g6 (z)). Zx+1 2)
(1)+)) - 90 [zz) + (o(z)) - -u(z)) - (Zx) 12++1 - zz)
Z z + (zb( 2x 2) (96 (Z+). 26-11 - 22) + Pg ||z 2x112
<Ggl M6++1
2
4C2
() + 1 ||++11 - 2x113
+ (g6( Zz =x+1) + Pg ||z 2
16
2
8C²
() + v)
|z - 2x112
Pg 1 2x.12.
T
8
Let yz) be the saddle point of Le(z,)) v). Then, we can get
E T, (x) - [2,2,,, Fb+1))
+ (1
B:)
_2- 7- +
ElUr. (y.,
E(U f:
yk+1)l
B1
B1
+
1
+
Pg
E
1|.z. 2x11/
1
1
E 1|z.-
26++ 11/
2n
2
2
21
4v
1
d
+ Ax0
1
C2
3
( + 1
1
E
FF Elzz41 -2u112
my
Am
2n
2)²
8 11
C2c
OC2
2(1 + 20)08
+
2
+ +80003 +
+80
+
|
-
112
ETTT+
k+
07 +
B2
B,
2
B₂ + 1)
(44)
Following the similar proof as in Theorem C.I the modified inner loop of ALEXR2 described in
Section can guarantee _ l| prox,
2 << é' for any t after K, = Ō AB.
Because
34

---

A= and > as specified Proposition D.1 we have K₁ = Ō( directly leverage the result from Appendix C
B2 B2
C.m
For outer loop, we
L.e.
Cgpo C2 lc-
T ( (Pg + PPg. pCg pC 0 +
8 8
By setting e' PPt
Pst Sage the total iteration will be
mp B3 B 8
When {gi) are smooth:
The proof follows almost the same procedure the proof of Theorem 4.7 In particular, since v, and Ge will contain the additional stochastic gradient from go. we need to replace (14) by
Ve+1 = l[ve+1 VV,(ww)1
(45)
while U4+1, un and g(w:) are still defined as in
15
16
and
117
Like Lemma B.1
it holds that
T (W S (w:) + ++
"
| x (wi)l
" ||
(46)
2
we can still show (19) for V in the new definition (45). Moreover, (20) and (2T) still hold they are related to 9o- Finally, we can still take the weighted summation of both sides (46). 19). 20 and (1) as in the proof of Theorem 4.7 show that
1 2F By proof similar to Lemma 3 in 2
Σ Iv -
m
+ Bi B2c
B BivRo B
min [1, B2)
This means Algorithm iterations. Then the conclusion follows from Proposition D.1 and the fact that Cy (
I finds nearly e-stationary point of Φ (w) in expectation in T 0
B BIVB.e
E SONEX with Adaptive Learning Rates
In this section we will show that our algorithms can also be easily extended to adaptive learning rate while still retaining the same complexity under an additional assumption introduced later. We consider adaptive step size update such as adam:
7) v + St11 (1 BB)s) G'G. o Gr
W++1 W. 1j v
where G2 is the overall gradient estimator mentioned in algorithm and o denotes Hladmmard(element wise) product. The following assumption has been justified for the adaptive step size by (48). Assumption E.1. We assume that the adaptive learning rates î) (element-wisee boseddddddddddd....... i.c.
nct S ñ 1/c
for any element fi of ij.
Below provide lemmas which are straightforward extensions of lemma
B.I B.4 to adaptive
learning setting. The proof is similar except that || w.ll < 1 o v.41
1 1" v
v 11.
Lemma E.2. (Lemma 3 in (48)) Under assumption E. for W = W. f VEL with na S ij S mca, and ?I. E 20g. we have:
F(W(1)) < F(w.)
+
7cuu
V4+
1701
[F((2))
170
|
(47)
2
2
4
where V4+1 ||1411 VF(w:)]1
35

---

Lemma E.3. 1 S 7. the gradient variance Ve+1 = V VF(w,)ll can be bounded as
2 V142|1
3130P
E [[++2] ( (1-
Σ
N - uc+1, 1
(48)
2 2807 C2)
min( B3} + 53L 0 [U+2])
where U2+1 = A |1++1 - g(w:))l² u1 = [u,1,..., ut, T g(w)) = [91 (w.))....,k. ........99 (w:))]
Lemma E.4. Suppose S' n-Br 11(1-) (1 7) and y S 2/2 and the MSVR update is used in
Algorithm 7 holds that
Blu S BB
2B1
DE[U++1]+ +
E M 12
(49)
B
nB3
Lemma E.5. Under the same assumptions as Lemma B.3 holds that
Blllu++2 Uc+ uu++111]] [1]] 4B11T²
On
2
2B17200
B[U++I]+
Ell
E
112
(50)
B1
B2
We then give the following theorem with similar proof technique.
Theorem E.6. Under Assumption 3.2
(AI),
4.5
and 4.6 by setting À O(E),
O (c min By, Ba)c ),7 O( (() ). O
B2
e²). SONEX with 7 1 7 B1(1-7)
n-B
Cn
converges to an approximate e-stationary solution of within T = O
2 n
5) iterations.
of BiBB.BB
Combining the above theorem with Theorem 4.4
we obtain the following guarantee:
Corollary E.7. Under Assumption (nttppttion..((111. A1),
4.3 and 4.6 with the same setting as in Theorem E.6
aa 2.5
SONEX converges to nearly e-stationary solution of within T O
iterations.
o2 BB
Proof of theorem
E.6 and corollary E.7 Since the proof is almost the same as the theorem 4.7 and
corollary 4.8 we only highlight the difference.
Let P
positive constant to decided later. Taking the weighted summation of both sides
49 and
500 as specified in the following formula
3
48
P
x
49
50
T/Cu
3
n23
we have
EF (WF+1) + BVi42 (P-5L307) BU142
17Cu
3
1
EV++ (p (p(1-78)) 11-1B1
12170992B,
(F((
EU--1
n/uu
8
2
n
3n2
222
8nC2Pipel
27
a
Ellve+1l1
E
4cn
g2
Bi
BI3
23C + +ff)
2Bi
612.077208B, +
Ci
+
min min B2}
nB2
BnB2
2c
We choose 17 such that
22
81C2 B
271330903122 0.
(51)
4cg
B13
36

---

We follow proof of 4.7 to set
12
P n
5L309 +
BBi
Bn²
which satisfies the following linear equation
P p 2
121,
P P-51301
Bn2
Combining the results above gives
1
EFx(W++1) + VV -P-55339) EU142
7/uu
(w,) + EV + (P-513C EU,+1
7cuu
2BC + ff) +
2B17
6L 077
Ct E (V:)))
Bi, B2}
nB2
BnB2
2c,
- Summing both sides of the inequality above over t = 0, , and organizing term lead to
[w:(()))1
20n,
(Fx(wo) - --((w))
- EV. (PV -57,399) BUS
1
T
17Cu
(48
g + of)
4B170j
121707
(52)
min B1, B2]
nB2
3nB2
Note that LF = 0(2), Ly 0 and P == P=0 Biye m We require that all the terms
4c, +o = O( we must
in the right-hand side to be in the order of ( To ensure min B2
have B O St m/inm Ba)e")
To ensure B'off su 0
B1 T2 = O we
Cu
cmnB2
must have 7 ( Sr
With these orders of 3 and 7. we also have P = 0 B, Bze
Can
and
12o, o7 O(e)) ). Therefore the last three terms 52 are (c2))
a,5nB3
To satisfy
51 will need set
3
B
1
B1
17
min
LF
np
Lg
n
=0
min
Ci
min B1, e2
B1 BIVB2 3
B min (Bin B2)
e })
Cu
n
Cu
BIVBBB
B2
=0
2
We also set the sizes Bi,2 to be O so that EU; = ( Also, it is easy to see that FX(wo) (Wo)
Fx(w)) O(1) and EV, (1). Then the first term in 52) becomes
1 0 ( 17
pe3
1
25p n
Cu a mi
+
Cu
2.5 B
B
B2e3
which will become O(E?) when T' O
2,n
o7 B.e
Then the remaining part is just the same as the proof of corollary
4.8
F More Experiment Details
F1 GDRO with CVaR divergence
Motivation of GDRO Modern machine learning models are typically trained under the empirical risk minimization (ERM) framework, which treats all samples in the dataset equally. Although ERM
37

---

often achieves strong average performance on test sets drawn from distributions similar to the training data, can perform poorly on rare or underrepresented subpopulations, i.e., it lacks robustness to distributional shifts. The motivation of GDRO is to address the spurious correlation between features and labels, and distributional shift. GDRO partitions the dataset into groups representing different distributions and applies a robust optimization scheme that assigns more weight on the worst groups. Specifically, GDRO considers the following objective:
12
min max pell 71
Σ Pg (xxx)) (x. y)) g=1
where ? S A is the set of distribution under consideration and A is the simplex. We consider here a popular choice of {p A P: I vi E [n]), which, by solving the inner-level maximization problem, leads to an equivalent reformulation IT corresppnn to the so-called CVaR divergence [ 49.. Intuitively GDRO with CVaR divergence aims to minimize averaged loss of the top-k worst groups.
CAMELYONI7-WILD (34] is part of the WILDS benchmark suite and consists of histopathology whole-slide images from five medical centers, with the goal of detecting metastatic tissue in lymph node biopsies. Following the WILDS setup, we frame this as binary classification task on image patches, where the primary challenge lies in distribution shift across hospitals (domains). We construct group with attributes 'hospital' and 'slide" which generates 30 groups
Amazon- WILDS [50] is a text classification dataset derived from Amazon product reviews, where the goal is predict binary sentiment (positive negative) based on TF-IDF features of review text. The data spans multiple product categories. We construct group with the attribute 'user' which generates 1252 groups.
CelebA (35] is large-scale facial attribute dataset containing over 200,000 celebrity images annotated with 40 binary attributes. We select 4 attributes 'Attractive', 'Mouth Moott__Slghtly___ SlightlyOOeen', 'Male' and 'Blonde Hair' and construct 16 groups, where 'Blonde_Han also serves as the target attribute for us to do classification.
Hyperparameter tuning. We tune the same hyperparameters of different methods from the same candidates as follows for fair comparison. For the three tasks we train the models for 10, 4, 15 epochs with batch size and the number of groups within a mini batch of 256(8), 32(8), 64(4), respectively. We set = 0.15 for all the three dataset. We tune learning rate in (le-5, 2e-5, 5e-5, le-4, 2e-4, 5e-4, 1e-3, 2e-3, 5e-3), A in (1, 0.1, 0.01), 7 and ,3 in 10.1, 0.2, 0.5, 0.8) and y (0.01, 0.02, 0.05, 0.1, 0.2). We set weight decay to be 0.01, 0.01, 0.02 for the three tasks, respectively. We use step decay (decay by 0.3x for every epochs), linear decay with 1st epoch warmup, step decay (decay by 0.2x for every 3 epochs) for learning rate for the three tasks, respectively.
F22 AUC Maximization with ROC Fairness Constraints
In this part we perform experiment on learning a model with ROC faimess constraint [1] following the same experiment setting as [8] Suppose the data are divided into two demographic groups D, = [(af, , H/ - and D, = {aa ; by )e where denotes the input data and b E {1, -13 denotes the class label. ROC fairness is to ensure the ROC curves for classification of the two groups are the same.
Since the ROC curve is constructed with all possible thresholds, we follow (51) by using a set of thresholds I' = ....... 3 to define the ROC faimesss. For each threshold т, we impose constraint that the false positive rate (FPR) and true positive rate (TPR) of the two groups are close, formulated as follow:
np
h+():= =
Σ npp
Mb. =
-T)
Σ I6b = (sw - 7) - k 0
na
up
Ru
1r (w) :=
- -l00sssa(aa2)) T)
110 = 110 - -7) - -K S 5.0
nip 1=1
1=1
38

---

SONISH, pi 1002
p-
SCNISH, 0 4000
SONSA O- 9000
Oher
Party
SOREEX/emEss yom 201
Ovaraad
Training curves of 5 constraint values in zero-one loss of different methods for continual learning with non-forgetting constraints when targeting the foggy class. Top: squared-hings penalty method with different Bottom: smoothed hinge penalty method with different p.
Figure 4:
where Sw( (.) denotes parameterized model, (2) is the sigmoid function, and is >> O is a tolerance parameter. We use the pairwise AUC loss as the objective function:
F(w)= = D
Σ X;ED XjED
Σ
(ss(w,,,) - x,),,
where D_ and D_ denote the set of positive and negative examples regardless their groups, respec- tively. We follow g to recast original constraint optimization as the following hinge- pengee-enalized objective:
min ((w) +
Σ TET
3l1 (w)]+ + B(h, (w)]+
(53)
2T
is notable that the penalty terms in the above problem can be formulated as [h+ (w]]+ ((g(w; (o(w::))) +)).
where [g) [lg: gal x]+. 91 (w; = n Σn I(b 1sw T), 92(w; T)
Tu 1 = 1 ( (sw T). As a result, f is convex, g(w) is smooth. As result, Algo- rithm 2 is applicable.
Hyperparameter tuning. For ALEXR2, we tune K; in {5, 10), tune a (i.e. learning rate for outer loop) and 17 (i.e. learning rate for inner loop) in (le-3, le-2, 0.1, 1}. and tune À (i.e. smoothing coefficient for outer function) and v (i.e. smoothing coefficient for overall objective) in (2e-2, 2e-3, 2e-4] and (0.1, 0,01), respectively. We fix the hyperparameters in MSVR update for both ALEXR2 and SONX, i.e. 3 = 0.8 and 9' = = = 0.1. For the two baseline algorithms, we tune the initial learning rate in (0.1, 1c-2, 1c-3, lc-4). We decay learning rate(outer Ir for ALEXR2) at 50% and 75% epochs by a factor of 10. We tuned p = (4, 6, 8, 10, 20, 40) for ALEXR2 and SONX 40, and p {10, 40, 80, 100, 200, 400, 800, 1000) for SOX. We also compare a double-loop method (ICPPAC) [T0] Algorithm 4], where we tune their 7) in (0.1, 0.01), T in (1, 10, 100), y in (1e-2, 1e-3, 1e-4), and fix 0, 0.1.
F.3 Continual learning with non-forgetting constraints
We follow 7 and consider training a CLIP model [52 with global contrastive loss (GCL) 17 as the objective and a so-called model developmental safety (MDS) constraints on protected tasks, given by:
min F(w, D) = 1 n
Σ
LOCL (w: (x,, (T)
(Tx II, ))
w
s.t. hg = Lw Da.) (w.w.l Da) S 0, k- = 1,
As discuss in section
4.2 update of y in inner loop is equivalent MSVR update, i.c. Vi.1
v Vix(uu,)), wist1 = (1 wi. (z, 2 (0(g: (2 B%
39

---

SONISHL p- 10)
SONISH, p 500
CONUSA 1000
Eairy
p-
SONEXISmiI
SONEXtemil P.
Co
Counsr
Fairy
A.
Opechh
Figure 5: Training curves of 4 constraint values in zero-one loss of different methods for continual learning with non non-forgetee forgettinn constraints when targeting the overcast class. Top: squared-hing penalty method with different p; Bottom: smoothed hinge penalty method with different p.
where L-CL (w; (xi, t.), (T. 1) is the two-way GCL for each image-text pair (x;, tz),
Loai(w;; T IT/ = T log (sp(E)(w,))) xx) Ey (4)/T)
E.JET, exp(F1 (F1 (x:) (T(())/T)
exp(ET((4) Br
T log
Ex,er, exp(Er(t,) (x())//))
where (x) and (t) denote the normalized feature representation of an image x and text t, generated by visual encoder and text encoder of CLIP model, respectively. T denotes the set of all texts to be contrasted with respect to (w.r.)) (including itself) and I denotes the set of all images to be contrasted w.r.t & (including itself). Here, the data D is a target dataset. ay Exx.y.)>D, L (W, xi, Us) is the loss of the model w on the k-th protected task, where la is logistic loss for the k-th classification task.
(w( Dx) 1
Hyperparameter tuning. The training data for the target tasks sampled from the training set of BDDI00K and LAION400M [53], while the data used to construct the MDS constraints are sampled exclusively from BDD 100K. For all the method we use the same learning rate of lc-6 and the same weight decay 0.1. For each experiment we run for 60 epochs 400 iterations per epoch and the batch size 256. We tune p in (0.01, 0. 1, 5, 10, 20, 40, 80) for SONEX+smooth Hinge and in 20, 50, 100, 200, 400, 800) for SOXasuarred hinge; we set where 71.7 the hyperparameter MSVR update for objective term (i.e. LOCL [ while 72. 75 is the hyperparameter in MSVR update for the penalty terms. We set À 0.1 which is the smoothing coefficient for smooth Hinge function in SONEX.
72 = 0.8, 7 % 0.1, = -0.05
F.4 Ablation Study
We conduct series of ablation studies to examine the effects of key hyperparameters and design choices.
Results for varying 3 We investigate the impact of different 3 on model performance. We evaluate 3 E {0.1. 0.3, 0.5, 1.0} and conduct experiments for SONEX on Amazon dataset for GDRO task and ALEXR2 on Adult dataset for AUC maximization with ROC fairness constraints. The results are shown in table 3 which indicates the importance of setting 3 < 1.0.
Results for varying 0 We analyze the influence of different 0 of ALEXR2 on model performance,
testing 0 E 16: 8 2: 1} and conduct experiments for ALEXR2 on Adult dataset. As shown in Table 4 the final AUC of ALEXR2 remains relatively stable across this range, suggesting that
[1' 1 11
not highly sensitive to 6.
note that GCL is also type of FCCO problem, can also apply MSVR update to track the inner functions
40

---

Finally we compare Adam-typ update with momentum- type update. We conduct experiments for SONEX on Amazon dataset and ALEXR2 on Adult dataset and the results are summarized in table [
Benefit of Adam-type update
They clearly demonstrate the advantage of using Adam-type updates, which consistently yield superior performance.
Table 3: Final loss and AUC of SONEX and ALEXR2 on Amazon dataset and Adult dataset respectively, with varying /
SONEX w/ varying 3
0.1
0.3
0.5
1.0
Final loss
0.5657
0.5563
0.5768
0.578
ALEXR2 w/ varying 3
0.1
0.3
0.5
1.0
Final AUC
0.8975
0.8976
0.8973
0.8969
Table 4: Final AUC of ALEXR2 with varying on Adult dataset
varying 0
0.1
0.3
0.5
1.0
Final AUC
0.8975
0,8976
0.8973
0.8969
Table 5: Comparison between Adam-type update and Momentum-type update for SONEX and ALEXR2 on Amazon dataset and Adult dataset, respectively
Algorithm SONEX(Final Loss)
Adam-type
Momentum-type
0.5657
0.9666
ALEXR2(Final AUC)
0.8975
0.861
F,5 Verification of Assumption 4.3
We verify the assumption 4.3 by computing minimum eigenvalue Amin of Vg(w (g()) T in group DRO experiment. We compute Amin for models trained on Camelyon17 7 and CelebA from the last epoch and report it in 6 Our experiment results demonstrate that the minimum eigenvalue of Vg(w) Vg(w) remains positive after training process finishes. Besides, assumption 5.1 has been verified empirically in Appendix A. 1. of
F.6 Other Details of Experiments
Computing Resource and Running time: The experiments of AUC Maximization with ROC Faimnsss Constraints and the experiments of group DRO of Camelyon17 dataset and CelebA dataset in our paper run on an A30 24G GPU, among which the first experiment takes less than 10 minutes for each run while for the second one, Camelyon 17 takes about 4 hours and CelebA takes about hours, for each run. The Amazon dataset of group DRO experiment is run on one A100 40GB GPUs and takes about 12 hours each run. The experiment of continual learning with non-forgetting constraints is run on two A100 40GB GPUs and takes about hours each run.
Data Split: We perform data split of CelebA dataset ourselves: within each group, samples are divided into training, validation, and test sets in an 8:1:1 ratio. For all the other datasets mentioned in paper
default split.
Other Experiment results We show the training curves of individual constraint values for our two experiments about constraint optimization, as shown in Figure (continual learning with forgeetinnn constraints) and Figure 6 7 (AUC maximization with ROC fairness constraints). Note that other weather conditions except foggy construct nor-ferget constraints since there no foggy data in BDD100k for defining such a constraint.
41

---

Table 6: Minimum eigen values Amin of (gg())V((W)) different solution.
Camelyon17
seed
seed 2
seed
Amin
0.0206
0.0349
0.0261
ColebA
seed 10
seed 20
Amin
0.9000
0.2713
0.000
SOXISH, 15001
0.000
0.000
0.000
0.035
0.005
-0.005
143122
8.000
23
20
30
0.100
10
20
23
-0.000
20
40
SONDOH,
SONKIH,
0.000
6015
0.015
0.000
Granue Waa
0.000
Condrea Vikae
0.000
0.000
0.000
-2.000
-0.000
-0.005
20
-0.000
20
20
30
-0.030
Eaach
Opecn
Epech
ALEXI2I0MM, 2 10)
ALEX1215MH, 20))
ALEXR2(smit,
3.33
0.03
04
3.00
6010
0.000
0.005
0200
-2.005
-3.000
bzach
Opoch
Figure 6: Training curves of 14 constraint values of different methods adult dataset for AUC maximization with ROC fairness constraints. Top row: SOX with squared- -inge penalty method; Middle: SONX with Hinge penalty method; Bottom: ALEXR2 with smoothed hinge penalty method.
S0W0SIL, p.
SOXISIL, 0=
6015
0.055
0.032
Castrort VALL
2,300
0000
0.000
000
-0.000
-3.005
tsach
SpocT
Epodh
SONNOM, p
SONDCH, p 5)
p
6015
0.05
0.012
Oestarns 1014
0.000
0.000
0.00
0.000
0:000
-2.000
-0.000
.0.00
Ezech
Spech
Epech
3.320
ALL ALEXR2IsmiL, pm 2)
6.000
ALEXRZIsmIL, p 4)
ALFMROGmmH
3015
6.15
0.15
3.85
8405
3.00
0000
0,000
Eoch
Figure 7: Training curves of 14 constraint values of different methods on COMPAS dataset for AUC maximization with ROC fairness constraints. Top row: SOX with squared-hinge penalty method; Middle: SONX with Hinge penalty method; Bottom: ALEXR2 with smoothed hinge penalty method.
42


Reviews:
[{"review_id": "UigXnGgpyO", "reviewer": "Reviewer_3ZyT", "summary": "This 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.", "questions": "1. **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?", "limitations": "yes", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 4, "significance": 4, "originality": 3, "strengths_and_weaknesses": "**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.", "was_revised": true, "final_justification": "I will keep my score after discussion with authors and other reviewers.", "rebuttal": "We sincerely thank the reviewer for your positive rating and valuable reviews.\n\n**Q1**: the novelty primarily lies in combining it with momentum methods.\n\n**A**: We'd like to point out that our second algorithm leverages **two levels** of smoothing, including the smoothing of the outer function $f$ and the Moreau envelope smoothing of the overall objective. It is this **nested smoothing** allows us to establish the state-of-the art complexity of $O(1/\\epsilon^5)$ for non-smooth weakly-convex FCCO and non-smooth weakly-convex constrained optimization. If we just apply the Moreau envelope smoothing of the overall objective, the best rate is $O(1/\\epsilon^6)$ [26]. \n\n**Q2**:**Generality of Assumption 4.3**: Is the condition $\\lambda_{min}(\\nabla g_i\\nabla g_i^\\top)\\geq \\delta$ realistic for deep networks (e.g., ResNet)? If $\\delta$ depends on $\\epsilon$, how would that affect the overall complexity? Empirical validation on ImageNet-scale models would substantially strengthen the claim.\n\n**A:** Yes to the first question! We kindly refer the reviewer to Table 3 in Appendix F.4 about verifying this condition on ResNet50 (25.5 million paras) for GDRO on CelebA and DenseNet121 (8 million paras) for GDRO on Camelyon17. Both models are deep neural networks widely used for ImageNet-scale experiments. Also, Table 2 in [r1] has verified an equivalent condition for hinge-penalized objective of constraint optimization problem, which is also a special case of FCCO, on the CLIP model where the encoders are Transformers.\n\nFor the second question, Theorem 4.4 (for finding a nearly $\\epsilon$-stationary solution) can tolerate $\\delta$ as small as $O(\\sqrt{\\epsilon})$ without sacrificing any loss in complexity. Moreover, even the condition $\\lambda_{min}(\\nabla g_i\\nabla g_i^\\top)\\geq \\delta$  does not hold, Theorem 4.2 still guarantees that our algorithm can find a meaningful **approximate $\\epsilon$-stationary solution**. \n\n\n**Q3**:**Efficiency of the Double-Loop Design**: 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\n**A:**   **On the wall-clock time**. The per-epoch running time of ALEXR2 is longer than SONX due to the momentum update, e.g., 1.87s of ALEXR2 vs 1.8s of SONX on COMPAS dataset.  However, since ALEXR2 converges faster than SONX (cf. Fig. 2), the overall clock time of ALEXR2 is still better for achieving the same level of training performance in our experiments. We give an example below for runtime (secs) reaching a set of given AUC level on the COMPAS dataset for AUC maximization with non-smooth ROC fairness constraints where \"N/A\" means that the algorithm cannot reach the given AUC. \n\n\n| Algorithm\\training AUC | 0.69 | 0.70 | 0.71 | 0.72 |\n| -------- | -------- | -------- |  -------- |  -------- |\n| SONX (runtime:s)     |  10.8    |  21.6    | 37.8 | N/A |\n| ALEXR2 (runtime:s)   |  9.35    | 13.09    |22.44 | 52.36 |\n\n\n\n**On the Single-Loop Design.** \nIt is indeed possible to design a single-loop algorithm for the weakly-convex strongly-concave min-max problem under consideration, similar to the approach in [r2]. However, such algorithms typically exhibit much worse dependence on the strong concavity parameter $\\mu$ compared to double-loop methods. For instance, the worst-case iteration complexity in [r2] scales as $O(1/\\mu^5)$. Even in the smooth setting, single-loop algorithms may suffer from an $O(1/\\mu^3)$ dependence in iteration complexity [r3, r4]. \n\nIn our setting, $\\mu$ is on the order of $\\epsilon$, since the smoothed outer function has a smoothness of $O(1/\\epsilon)$. Consequently, any dependence worse than $O(1/\\mu)$ would lead to an iteration complexity exceeding the $O(1/\\epsilon^5)$ achieved by our method.\n \n\n\n**Q4**:**Adaptive Learning Rate Assumptions**: Is Assumption E.1 satisfied in practical scenarios involving Adam? Could the theory be extended to relax this assumption?\n\n**A:** We acknowledge that this assumption may be restrictive for the extension, although our main contribution lies in the momentum updates. Our analysis in the adaptive learning rate setting including Assumption E.1 follows [r4], which empirically verified this assumption in their experiments (see Appendix K of their paper).\n\nFrom a theoretical standpoint, if the stochastic gradient is upper bounded and $\\varepsilon$ is a constant, then the adaptive learning rate term $\\frac{1}{\\sqrt{\\mathbf{s}_t} + \\varepsilon}$ would also be both lower and upper bounded by constants.\n\nIn practice, however, we typically set $\\varepsilon$ to a very small value (e.g., $1\\text{e}{-8}$). As a result, for coordinates with near-zero second-order moments, $\\frac{1}{\\sqrt{\\mathbf{s}_t} + \\varepsilon}$ can become extremely large. \n\nIn such cases, the upper bound assumption may no longer hold. One practical solution to mitigate this issue is to apply **clipping** to the adaptive learning rate. Specifically, for coordinates with very small second-order moments, the corresponding first-order moment $\\mathbf{v}_{t+1}$ is usually also very small. \n\nTherefore, using a clipped update rule of the form  $\\mathbf w_{t+1} = \\mathbf w_t - \\eta \\min(\\frac{1}{\\sqrt{\\mathbf s_t}+\\varepsilon}, c_u)\\mathbf v_{t+1}$ for some constant $c_u$ may not impact performance significantly, while ensuring the assumption is satisfied.\n\nWe also note that general convergence analysis of Adam is challenging. Many existing works have to assume some kind of assumptions or suffer other limitations, e.g., [r5] assumes strong growth condition of the gradient and [r6] has a dependence on $O(1/\\varepsilon^2)$ in their complexity, which is non-vacuous only when $\\varepsilon$ is a constant, translating to the upper bound $c_u$ in Assumption E.1.  \n\n**Q5**: A clearer intuition or motivation for nested smoothing of ALEXR2 would improve accessibility.\n\n**A**: The motivation is as follows. The Moreau envelope smoothing is to ensure the overall objective $F_{\\lambda, \\nu}$ is smooth so that its stationary point can be found. This is a typical approach for handling non-smooth weakly convex objectives. The smoothing on the outer function $f$ is to ensure the inner min-max problem of (6) strongly-convex and **strongly-concave** so that an existing algorithm ALEXR can be used to solve the inner min-max problem, which naturally leads to a double loop algorithm. \n\n\n[r1] Yang et al. Single-loop Algorithms for Stochastic Non-convex Optimization with Weakly-Convex Constraints. arxiv 2025.\n\n[r2] Hu et al. Single-Loop Stochastic Algorithms for Difference of Max-Structured Weakly Convex Functions. NeurIPS 2024. \n\n[r3] Lin et al. On Gradient Descent Ascent for Nonconvex-Concave\nMinimax Problems.  ICML 2020. \n\n[r4] Guo et al. Unified convergence analysis for adaptive optimization with moving average estimator. ML 2025. \n\n[r5] Zhang et al. Adam Can Converge Without Any Modification On\nUpdate Rules.  NeurIPS 2022. \n\n[r6] Li et al. Convergence of Adam Under Relaxed Assumptions. NeurIPS 2023."}, {"review_id": "ATEO3znj1L", "reviewer": "Reviewer_hTnd", "summary": "In 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.", "questions": "1. 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 “]”.", "limitations": "N/A", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths:\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.", "was_revised": true, "final_justification": "Thank you for the detailed response. I will maintain my score.", "rebuttal": "We sincerely thank the reviewer for your valuable reviews.\n\n**Q1**: What do you see as the biggest difficulty in generalizing the finite sum form to the stochastic expectation form in the problem model?\n\n**A:** The core of the algorithms for finite-sum setting is to track the inner function values $g_1(\\mathbf w_t), \\ldots, g_n(\\mathbf w_t)$ at all iterations such that the estimation error is decreasing over iterations. The biggest difficulty in generalizing the finite-sum form to the stochastic expectation form $\\mathbb E_\\xi f_{\\xi}(\\mathbb E_{\\zeta}[g_{\\xi}(\\mathbf w, \\zeta)])$ is how to compute an estimator of the inner function $g_{\\xi_t}(\\mathbf w_t) = \\mathbf E_{\\zeta}g_{\\xi_t}(\\mathbf w_t, \\zeta)$ for a random variable $\\xi_t$ with decreasing estimation error. Since $\\xi_t$ is a fresh random variable that may not be encountered before, hence the moving average or MSVR technique for the finite-sum case is not applicable.  Existing works like BSGD [r7] has to use a very large batch size to control the estimation error, which can be prohibitive in practice.\n\n[r7] Hu et al. Biased stochastic first-order methods for conditional  stochastic optimization and applications in meta learning. NeurIPS, 2020. \n\n\n**Q2**: Elaborate the design of the novel steps, and what motivated the selection of the steps borrowed from prior work?\n\n**A:** We thank the reviewer for the suggestion.  We highlight the novel design below. \n\n**Algorithm 1:** The key novel step is to compute a gradient estimator using the smoothed outer function $f_{i,\\lambda}$ at Step 9. It is this step that allows us to borrow the momentum update (Step 10, 11) of an existing algorithm SOX for smooth FCCO problems [2] to tackle non-smooth FCCO problems.\n\n**Algorithm 2:** The key lies at the gradient estimator in Step 13 for the nested smooth objective $F_{\\lambda, \\nu}$, which allows us to build a recursion of its estimation error for the Moreau envelope defined in (5). The inner loop is borrowed from an existing algorithm ALEXR [5]. \n\n\n**Q3**: In numerical experiments, what method is used to solve the subproblem in Algorithm 2?\n\n**A:** As mentioned in section 4.2, we use ALEXR [5] to solve each subproblem. Note that the $\\mathbf y$-step update (line 7-9 in algorithm 2) and $\\mathbf z$-step update (line 11 in algorithm 2) have closed form updates, i.e.\n$y^i_{t,k} = \\nabla f_{i,\\lambda}(u^i_{t,k})$, where $u^i_{t,k+1} =(1-\\hat\\gamma) u^i_{t,k} + \\hat\\gamma   g_i(\\mathbf z_{t,k};B_{i,2}^{t,k}) + \\hat\\gamma \\theta(g_i(\\mathbf z_{t,k};B_{i,2}^{t,k}) - g_i(\\mathbf z_{t,k-1};B_{i,2}^{t,k}))$ and $\\hat\\gamma=  \\gamma/(1+\\gamma)$, \nand \n$\\mathbf z_{t,k+1} = \\frac{\\eta \\mathbf w_t +\\nu \\mathbf z_{t,k}}{\\eta+\\nu} - \\frac{\\eta\\nu}{\\nu+\\eta}G_{t,k}$\n\n**Q4**: Some typos.\n\n**A:** We thank reviewer for pointing out all these typos and we will correct them in the revision.\n\n**Q5**: Elaborate in detail on at least one representative scenario of FCCO. \n\n**A:** Thank you for the suggestion!  We have studied multiple applications of FCCO in Section 6, e.g., GDRO with CVaR divergence. The motivation of GDRO is to address the spurious correlation between features and labels, and distributional shift. We will add more details in the revision about the formulation and motivation."}, {"review_id": "M3EMjH2NNE", "reviewer": "Reviewer_soxr", "summary": "This 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.", "questions": "(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.", "limitations": "Yes", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths: 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.", "was_revised": true, "final_justification": "I will keep my score after discussion with authors and other reviewers.", "rebuttal": "We sincerely thank the reviewer for your constructive comments.\n\n**Q1**:It's better to provide an ablation study that isolates the effect of momentum ($\\beta$), smoothing parameter ($\\theta$), and Adam-type updates. How critical are they to the observed improvements?\n\n**A:** Thanks for the suggestion. This is an easy fix. We have done some ablation studies.  \n\n- Using $\\beta<1$ is important for SONEX. The ablation study on the Amazon dataset for GDRO shows that SONEX with $\\beta=1$  converges slower, as shown in the table below. \n\n| Setting for $\\beta$ |0.1 | 0.3 | 0.5 | 1.0 |\n| -------- | -------- | -------- | -------- |-------- |\n| Final Loss   | 0.5657 | **0.5563** | 0.5768 | 0.578|\n\n-  The ablation study shows that ALEXR2 is not sensitive to $\\theta\\in\\{0, 1/16, 1/8, 1/4, 1/2, 1\\}$ and $\\beta\\in\\{0.1, 0.3, 0.5, 0.7, 1\\}$.  The ablation results of ALEXR2 on the Adult dataset for fair AUC maximization are summarized in the following table.  \n\n| Setting for $\\beta$ | 0.1 | 0.3 | 0.5 | 0.7 | 1.0 |\n| -------- | -------- | -------- | -------- |-------- |-------- |\n| Final AUC   | 0.8975 | 0.8976 | 0.8973 | 0.8968 | 0.8969 |\n\n| Setting for $\\theta$ | 0 | 1/16 | 1/8 | 1/4 | 1/2 | 1 |\n| -------- | -------- | -------- | -------- |-------- |-------- |-------- |\n| Final AUC   | 0.8976 | 0.8981 | 0.8975 | 0.8973 | 0.8972 | 0.8971 |\n\n- Using the Adam-type update is beneficial for both SONEX and ALEXR2 in the ablation studies.  \n\n| Setting | Adam-type |momentum-type |\n| -------- | -------- | -------- |\n| ALEXR2 (Final AUC on Adult data)     | **0.8975**     | 0.861     |\n| SONEX (Final Loss on Amazon data)     | **0.5657**| 0.9666     |\n\n\n\n**Q2**: 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\n**A:** This reason is that SONEX does not leverage the convexity of the outer function, while ALEXR2 assumes the convexity of the outer function. Without the convexity of outer function, SONEX directly estimates the inner function values using MSVR, which inherits its dependence on the inner batch size i.e.., $O(1/\\sqrt{B_2})$. Using the convexity of the outer function, ALEXR2 leverage the min-max formulation leading to a standard dependence of $O(1/B_2)$ on the inner batch size. \n\nThese differences in dependence between the non-convex and convex outer function settings are consistent with prior works [18, 5]."}, {"review_id": "0w3342tALL", "reviewer": "Reviewer_gzXU", "summary": "This 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.", "questions": "- 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 ?", "limitations": "Yes", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 1, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strength :\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", "was_revised": true, "final_justification": "This paper is technically solid and provide new accelerated algorithms both in theory and in practice. I read and try to understand all the proofs/experiments. However, it is a paper for expert and I found it hard to read and follow as a non-expert of FCCO. In particular, I did not catch the motivation and the intuitions to choose these particular algorithms designs (Algorithm 1-2) for having accelerated rates.\n\nTherefore, I choose to keep my score. I think that this work could interest the FCCO community but it is not written to be easily understandable for non-expert of the field.", "rebuttal": "We sincerely thank the reviewer for your valuable comments.\n\n**Q1**: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\n**A:** This choice offers practical benefits by enabling adaptive learning rates for each parameter, which is especially advantageous in training deep neural networks—particularly Transformer architectures. In our continual learning experiments with CLIP models, which are based on Transformers, this adaptive behavior proved highly beneficial. Please also see the response to Q1 of reviewer soxr. \n\n\n**Q2**: Can you generalize your analysis without the restrictive Lipschitz condition on $f$ and $g$?\n\n**A**: Unfortunately, no if we target non-smooth problems. The Lipschitz continuity of $f$ and $g$ is essential and widely assumed in all existing works on compositional optimization. Even in the standard non-smooth stochastic optimization $\\mathbb E_\\xi[f(\\mathbf w; \\xi)]$, the Lipschitz continuity is necessary for the convergence analysis [r8, r9].  In fact, our assumptions are minimal compared to prior studies of FCCO, many of which additionally require smoothness of $g$ and/or $f$.\n\n\n\n\n**Q3**: Why can you apply directly results of Lemma B.1-B.4 that are derive for different algorithms to SONEX algorithm ?\n\n**A**: Indeed, while SONEX is novel in addressing non-smooth FCCO through the use of a smoothed outer function $f_{i,\\lambda}$ for gradient estimation (as in Step 9), its other components are inspired by existing algorithms such as SOX and MSVR for smooth FCCO. Thanks to its modular design and analysis, we can draw on prior results for analyzing individual components. Specifically, the $u$-update in Step 7 is identical to that of MSVR, so the recursion for its estimator error (Lemmas B.3 and B.4) is directly borrowed from MSVR [18]. The momentum update in Steps 10–11 resembles that of SOX, allowing us to use Lemma B.2 from [2] to analyze the recursion of the momentum estimator's error. Lemma B.1 is general and applies to any smooth function, regardless of the update mechanism. This modularity enables us to leverage well-established analyses for each component.\n\n**Q4**: About the assumptions.\n\n**A**: Regarding the assumptions: \n- We have verified Assumption 4.3 in our experiment in Appendix F.4.\n- Assumption 4.6 is only used for analyzing SONEX. Our second algorithm ALEXR2 does not need this assumption but still enjoys the same rate. It is necessary for SONEX to analyze the MSVR estimator $u$ following [6] as it does not leverage the convexity of $f$ as ALEXR2. Without this assumption, we can only use the moving average estimator for $u$ and establish a worse complexity of $O(1/\\epsilon^7)$ for SONEX and [6] will suffer from an even worse complexity of $O(1/\\epsilon^8)$. \n\n- Assumption 4.5 is standard for analyzing stochastic algorithms. Assumption 5.1 is a standard assumption for analyzing penalty-based methods of non-convex constrained optimization (see [r10, r11]). It has been empirically verified in [r10] (Table 2). \n- The Lipschitz continuity of $f$ and $g$ in Assumption 3.1 are fairly standard for FCCO problems. Since we target at **non-smooth $f$ and $g$**, Lipchitz continuity is a minimal assumption. Quadratic functions are smooth and convex, then the problem becomes easier and it is possible to extend existing algorithms for smooth min-max problems without the Lipschitz continuity, e.g., [r12]. That is out of scope of this paper.  \n- In the  considered applications of GDRO and  learning with fairness constraints, $f$ is hinge and hence Lipschitz continuous. \n\n\n**Q5**: There exists a tighter lipschitz constant for $f_\\lambda$ in line 113.\n\n**A**: We thank reviewer for introducing a new result improving our lemma C.2 and the subsequent theorems, and we will update them in the revision.\n\n\n**Q6**: Some typos.\n\n**A**: We thank reviewer for pointing out all these typos and we will correct them in the revision.\n\n\n[r8] Davis & Drusvyatskiy. Stochastic model-based minimization of weakly convex functions. SIAM OPT, 2019. \n\n[r9] Cutkosky et al. Optimal Stochastic Non-smooth Non-convex Optimization through Online-to-Non-convex Conversion. ICML 2023. \n\n[r10] Yang et al. Single-loop algorithms for stochastic non-convex optimization with weakly-convex constraints. arXiv 2025.\n\n[r11] Liu and  Xu. A single-loop spider-type stochastic subgradient method for expectation constrained nonconvex nonsmooth optimization. arXiv 2025.\n\n[r12] Guo et al. Unified convergence analysis for adaptive optimization with moving average estimator. ML, 2025."}]

Author rebuttal:

Meta-review: This paper studied non-smooth non-convex finite-sum coupled compositional optimization. The proposed stochastic momentum method achieves the best know iteration complexity of $O(1/\epsilon^5)$. It can also be applied to solve the inequality constrained non-convex problem and achieve the new state-of-the-art complexity. 

All reviewers are positive to this paper and the questions have been address in rebuttal and discussion. Therefore, I recommend acceptance.
Decision: Accept (poster)