SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Quanqi Hu
quanqi-hu@tamu.edu
Dixian Zhu
Department of Genetics
dixian-zhu@stanford.edu
tianbao-yang@tamu.edu
This paper investigates new families of compositional optimization problems, called non-smooth weakly-convex finite-sum coupled compositional optimization (NSWC FCCO). There has been a growing interest in FCCO due to its wide-ranging applications in machine learning and AI, as well as its ability to address the shortcomings of stochastic algorithms based on empirical risk minimization. However, current research on FCCO presumes that both the inner and outer functions are smooth, limiting their potential to tackle a more diverse set of problems. Our research expands on this area by examining non-smooth weakly-convex FCCO, where the outer function is weakly convex and non-decreasing, and the inner function is weakly-convex. We analyze a single-loop algorithm and establish its complexity for finding an  $\epsilon$ -stationary point of the Moreau envelop of the objective function. Additionally, we also extend the algorithm to solving novel non-smooth weakly-convex tri-level finite-sum coupled compositional optimization problems, which feature a nested arrangement of three functions. Lastly, we explore the applications of our algorithms in deep learning for two-way partial AUC maximization and multi-instance two-way partial AUC maximization, using empirical studies to showcase the effectiveness of the proposed algorithms.
In this paper, we consider two classes of non-convex compositional optimization problems. The first class is formulated as following:
$$
\min  _ {\mathbf {w} \in \mathbb {R} ^ {d}} F (\mathbf {w}) := \frac {1}{n} \sum_ {i \in \mathcal {S}} f _ {i} \left(\mathbb {E} _ {\xi \sim \mathcal {D} _ {i}} [ g _ {i} (\mathbf {w}; \xi) ]\right), \tag {1}
$$
where  $S$  denotes a finite set of  $n$  items and  $\mathcal{D}_i$  denotes a distribution that could depend on  $i$ . The second class is given by:
$$
\min  _ {\mathbf {w} \in \mathbb {R} ^ {d}} F (\mathbf {w}) := \frac {1}{n _ {1}} \sum_ {i \in \mathcal {S} _ {1}} f _ {i} \left(\frac {1}{n _ {2}} \sum_ {j \in \mathcal {S} _ {2}} g _ {i} \left(\mathbb {E} _ {\xi \sim \mathcal {D} _ {i, j}} [ h _ {i, j} (\mathbf {w}; \xi) ]\right)\right), \tag {2}
$$
where  $S_{1}$  denotes a finite set of  $n_{1}$  items and  $S_{2}$  denotes a finite set of  $n_{2}$  items and  $\mathcal{D}_{ij}$  denotes a distribution that could depend on  $(i,j)$ . For simplicity of discussion, we denote by  $g_{i}(\mathbf{w}) =$
$\mathbb{E}_{\xi \sim \mathcal{D}_i}[g_i(\mathbf{w};\xi)]:\mathbb{R}^d\to \mathbb{R}^{d_1}$  and by  $h_{i,j}(\mathbf{w}) = \mathbb{E}_{\xi \sim \mathcal{D}_{i,j}}[h_{i,j}(\mathbf{w};\xi)]:\mathbb{R}^d\to \mathbb{R}^{d_2}$ . For both classes of problems, we focus our attention on non-convex  $F$  with non-smooth non-convex functions  $f_{i}$  and  $g_{i}$ , which, to the best of our knowledge, has not been studied in any prior works.
The first problem (1) with smooth functions  $f_{i}$  and  $g_{i}$  has been explored in previous works [26, 15, 21, 33], which is known as finite-sum coupled compositional optimization (FCCO). It is subtly different from standard stochastic compositional optimization (SCO) [27] and conditional stochastic optimization (CSO) [13]. FCCO has been successfully applied to optimizing a wide range of X-risks [33] with convergence guarantee, including smooth surrogate losses of areas under the curves [20] and ranking measures [21], listwise losses [21], and contrastive losses [36]. The second problem (2) is a novel class and is referred to as tri-level finite-sum coupled compositional optimization (TCCO). Both problems differ from traditional two-level or multi-level compositional optimization due to the coupling of variables  $i, \xi$  in (1) or the coupling of variables  $i, j, \xi$  in (2) at the inner most level.
One limitation of prior works about non-convex FCCO is that their convergence analysis heavily relies on the smoothness conditions of  $f_{i}$  and  $g_{i}$  [26, 15]. This raises a concern about whether existing techniques can be leveraged for solving non-smooth non-convex FCCO problems with non-asymptotic convergence guarantee. Non-smooth non-convex FCCO and TCCO problems have important applications in ML and AI, e.g., group distributionally robust optimization [4] and two-way partial AUC maximization for deep learning [44]. We defer discussions and formulations of these problems to Section 5. The difficulty for solving smooth FCCO lies at high costs of computing a stochastic gradient  $\nabla g_{i}(\mathbf{w})\nabla f_{i}(g_{i}(\mathbf{w}))$  for a randomly sampled  $i$  and the overall gradient  $\nabla F(\mathbf{w})$ . To approximate the stochastic gradient, a variance-reduced estimator of  $g_{i}(\mathbf{w}_{t})$  denoted by  $u_{i,t}$  is usually maintained and updated for sampled data in the mini-batch  $i\in \mathcal{B}_t$ . As a result, the stochastic gradient can be approximated by  $\nabla g_{i}(\mathbf{w}_{t};\xi_{t})\nabla f_{i}(u_{i,t})$ , where  $\xi_t\sim \mathcal{D}_i$  is a random sample. The overall gradient can be estimated by averaging the stochastic gradient estimator over the mini-batch or using variance-reduction techniques. A key insight of the convergence analysis for smooth FCCO is to bound the following error using the  $L$ -smoothness of  $f_{i}$ , which reduces to bounding the error of  $u_{i,t}$  for estimating  $g_{i}(\mathbf{w}_{t})$ :
$$
\left\| \nabla g _ {i} \left(\mathbf {w} _ {t}; \xi_ {t}\right) \nabla f _ {i} \left(u _ {i, t}\right) - \nabla g _ {i} \left(\mathbf {w} _ {t}; \xi_ {t}\right) \nabla f _ {i} \left(g _ {i} \left(\mathbf {w} _ {t}\right)\right) \right\| ^ {2} \leq \left\| \nabla g _ {i} \left(\mathbf {w} _ {t}; \xi_ {t}\right) \right\| ^ {2} L \| u _ {i, t} - g _ {i} \left(\mathbf {w} _ {t}\right) \| ^ {2}.
$$
A central question to be addressed in this paper is "Can these gradient estimators be used in stochastic optimization for solving non-smooth non-convex FCCO with provable convergence guarantee"? To address this question we focus our attention on a specific class of FCCO/TCCO called non-smooth weakly-convex (NSWC) FCCO/TCCO. This approach aligns with many established works on NSWC optimization [6-9]. Nevertheless, NSWC FCCO/TCCO is more complex than a standard weakly-convex optimization problem because an unbiased stochastic subgradient is not readily accessible. In addition, the convergence measure in terms of the gradient norm of smooth non-convex objectives is not applicable to weakly convex optimization, which will complicate the analysis involving the biased stochastic gradient estimator  $\partial g_{i}(\mathbf{w}_{t};\xi_{t})\partial f_{i}(u_{i}^{t})$  1.
Contributions. A major contribution of this paper is to present novel convergence analysis of single-loop stochastic algorithms for solving NSWC FCCO/TCCO problems, respectively. In particular,
- For non-smooth FCCO, we analyze the following single-loop updates:
$$
\mathbf {w} _ {t + 1} = \mathbf {w} _ {t} - \eta \frac {1}{B} \sum_ {i \in \mathcal {B} _ {t}} \partial g _ {i} \left(\mathbf {w} _ {t}; \xi_ {t}\right) \partial f _ {i} \left(u _ {i, t}\right), \tag {3}
$$
where  $\mathcal{B}_t$  is a random mini-batch of  $B$  items, and  $u_{i,t}$  is an appropriate variance-reduced estimator of  $g_i(\mathbf{w}_t)$  that is updated only for  $i\in \mathcal{B}_t$  at the  $t$ -th iteration. To overcome the non-smoothness, we adopt the tool of Moreau envelop of the objective as in previous works [6, 7]. The key difference of our convergence analysis from previous ones for smooth FCCO is that we bound the inner product  $\langle \mathbb{E}_i\partial g_i(\mathbf{w})\partial f_i(u_{i,t}),\widehat{\mathbf{w}}_t - \mathbf{w}_t\rangle$ , where  $\widehat{\mathbf{w}}_t$  is the solution of the proximal mapping of the objective at  $\mathbf{w}_t$ . To this end, specific conditions of  $f_{i},g_{i}$  are imposed, i.e.,  $f_{i}$  is weakly convex and non-decreasing and  $g_{i}(\mathbf{w})$  is weakly convex, under which we establish an iteration complexity of  $T = \mathcal{O}(\epsilon^{-6})$  for finding an  $\epsilon$ -stationary point of the Moreau envelope of  $F(\cdot)$ .
- For non-smooth TCCO, we analyze the following single-loop updates:
$$
\mathbf {w} _ {t + 1} = \mathbf {w} _ {t} - \eta \frac {1}{B _ {1}} \sum_ {i \in \mathcal {B} _ {1} ^ {t}} \left[ \frac {1}{B _ {2}} \sum_ {j \in \mathcal {B} _ {2} ^ {t}} \partial h _ {i, j} (\mathbf {w} _ {t}; \xi_ {t}) \partial g _ {i} (v _ {i, j, t}) \right] \partial f _ {i} (u _ {i, t}), \tag {4}
$$
<sup>1</sup>We use  $\nabla$  to denote gradient of a differentiable function and  $\partial$  to denote a subgradient of a non-smooth function.
Table 1: Comparison with prior works for solving (1) and (2). In the monotonicity column, notation  $\uparrow$  means the given function is required to be non-decreasing. If not specified, the given function is only required to be monotone.
Method	Objective	Smoothness	Weak Convexity	Monotonicity	Complexity
SOX [26]	(1)	fi, gi	none	none	O(ε-4)
MSVR [15]	(1)	fi, gi	none	none	O(ε-3)
SONX (Ours)	(1)	none	fi, gi	fi↑	O(ε-6)
SONT (Ours)	(2)	none	fi, gi, hi,j	fi↑, gi↑	O(ε-6)
SONT (Ours)	(2)	hi,j	fi, gi	fi↑, gi	O(ε-6)
where  $\mathcal{B}_t^1$  and  $\mathcal{B}_t^2$  are random mini-batches of  $B_{1}$  and  $B_{2}$  items, respectively, and  $u_{i,t}$  is an appropriate variance-reduced estimator of  $\frac{1}{n_2}\sum_{j\in S_2}g_i(h_{ij}(\mathbf{w}_t))$  that is updated only for  $i\in \mathcal{B}_t^1$ , and  $v_{i,j,t}$  is an appropriate variance-reduced estimator of  $h_{i,j}(\mathbf{w}_t)$  that is updated only for  $i\in \mathcal{B}_t^1$ ,  $j\in \mathcal{B}_t^2$ . To prove the convergence, we impose conditions of  $f_{i}, g_{i}, h_{i,j}$ , i.e.,  $f_{i}$  is weakly convex and non-decreasing and  $g_{i}(\cdot)$  is weakly convex and non-decreasing (or monotonic),  $h_{ij}$  is weakly convex (or smooth), and establish an iteration complexity of  $T = \mathcal{O}\left(\epsilon^{-6}\right)$  for finding an  $\epsilon$ -stationary point of the Moreau envelope of  $F(\cdot)$ .
- We extend the above algorithms to solving (multi-instance) two-way partial AUC maximization for deep learning, and conduct extensive experiments to verify the effectiveness of the both algorithms.
Smooth SCO. There are many studies about two-level smooth SCO [27, 38, 10, 19, 3, 28] and multi-level smooth SCO [32, 32, 1, 39]. The complexities of finding an  $\epsilon$ -stationary point for two-level smooth SCO have been improved from  $O(\epsilon^{-5})$  [27] to  $O(\epsilon^{-3})$  [19], and that for multi-level smooth SCO have been improved from a level-dependent complexity of  $O(\epsilon^{-(7 + K) / 2})$  [32] to a level-independent complexity of  $O(\epsilon^{-3})$  [32], where  $K$  is the number of levels. The improvements mostly come from using advanced variance reduction techniques for estimating each level function or its Jacobian and for estimating the overall gradient. Two stochastic algorithms have been developed in [13] for CSO but suffer a limitation of requiring large batch sizes.
Smooth FCCO. FCCO was first introduced in [20] for optimizing average precision. Its algorithm and convergence analysis was improved in [26] and [15]. The former work [26] proposed an algorithm named SOX by using moving average (MA) to estimate the inner function values and the overall gradient. In the smooth non-convex setting, SOX is proved to achieve an iteration complexity of  $\mathcal{O}(\epsilon^{-4})$ . The latter work [15] proposed a novel multi-block-single-probe variance reduced (MSVR) estimator for estimating the inner function values, which helps achieve a lower iteration complexity  $\mathcal{O}(\epsilon^{-3})$ . Recently, [11] proposed an extrapolation based estimator for the inner function, which yields a method with a complexity that matches MSVR when  $n \leq \epsilon^{2/3}$ . These techniques have been employed for optimizing various X-risks, including contrastive losses [36], ranking measures and listwise losses [21], and other objectives [26, 15]. However, all of these prior works assume the smoothness of  $f_{i}$  and  $g_{i}$ . Hence, their analysis is not applicable to NSWC FCCO problems. Our novel analysis of a simple algorithm for NSWC FCCO problems yields an iteration complexity of  $\mathcal{O}(\epsilon^{-6})$  for using the MSVR estimators of the inner functions. The comparison with [26, 15] is shown in Table 1.
Non-smooth Weakly Convex Optimization. Analysis of weakly convex optimization with unbiased stochastic subgradients was pioneered by [6, 7]. Optimization of compositional functions that are weakly convex have been tackled in earlier works [8, 9], where the inner function is deterministic or does not involve coupling between two random variables. A closely related work to our NSWC FCCO is weakly-convex concave minimax optimization [22]. Assuming  $f_{i}$  is convex, (1) can be written as:  $\min_{\mathbf{w}} \max_{\pi \in \mathbb{R}^n} \frac{1}{n} \sum_{i \in S} \langle \pi_i, g_i(\mathbf{w}) \rangle - f_i^*(\pi_i)$ , where  $f_i^*(\cdot)$  is the convex conjugate of  $f_{i}$ . It can be solved using existing methods [22, 31, 41, 43, 17] but with several limitations: (i) the algorithms in [22, 31, 41, 43] have a comparable complexity of  $O(1/\epsilon^6)$  but have unnecessary double loops which require setting the number of iterations for the inner loop; (ii) the algorithm in [17] is single loop but has a worse complexity of  $O(1/\epsilon^8)$ ; (iii) these existing algorithms and analysis does not account for complexity of updating all coordinates of  $\pi$ , which could be prohibitive in
many applications; iv) these approaches are not applicable to NSWC FCCO/TCCO with weakly convex  $f_{i}$ . In fact, the double loop algorithm has been leveraged and extended to solving the two-way partial AUC maximization problem, a special case of NSWC FCCO [44], by sampling and updating a batch of coordinates of  $\pi$  at each iteration. However, it is less practical thus not implemented and its analysis did not explicitly show the convergence rate dependency on  $n_{+}, n_{-}$  and the block batch size.
A special case of NSWC SCO problem was considered in [46], which is given by
$$
\min  _ {x \in \mathcal {X}} f (x, g (x)), \text {w i t h} f (x, u) = \mathbb {E} _ {\zeta} [ u + \varkappa \max  (0, g (x; \zeta) - u) ], \quad g (x) = \mathbb {E} _ {\xi} [ g (x; \xi) ].
$$
They proposed two methods, SCS for smooth  $g(x)$  and SCS with SPIDER for non-smooth  $g(x)$ . For both proposed methods, they proved a sample complexity of  $\mathcal{O}(1 / \epsilon^6)$  for achieving an  $\epsilon$ -stationary point of the objective's Moreau envelope. We would like to remark that the above problem with a non-smooth  $g(x)$  is a special case of NSWC FCCO with only a convex outer function, one block and no coupled structure. Nevertheless, their algorithm for non-smooth  $g(\cdot)$  suffers a limitation of requiring a large batch size in the order of  $O(1 / \epsilon^2)$  for achieving the same convergence.
Finally, we would like to mention that non-smooth convex or strongly convex SCO problems have been considered in [27, 42, 26], which, however, are out of scope of the present work.
Let  $\| \cdot \|$  be the Euclidean norm of a vector and spectral norm of a matrix. We use  $\Pi_C[\cdot]$  to denote the Euclidean projection onto  $\{v \in \mathbb{R}^m : \| v \| \leq C\}$ . For vectors, inequality notations including  $\leq, \geq, >, <$  are used to denote element-wise inequality. For an expectation function  $f(\cdot) = \mathbb{E}_{\xi}[f(\cdot; \xi)]$ , let  $f(\cdot; \mathcal{B}) = \frac{1}{|\mathcal{B}|} \sum_{\xi \in \mathcal{B}} f(\cdot; \xi)$  be its stochastic unbiased estimator evaluated on a sample batch  $\mathcal{B}$ . A stochastic unbiased estimator is said to have bounded variance  $\sigma^2$  if  $\mathbb{E}_{\xi}[ \| f(\cdot) - f(\cdot; \xi) \|^2 ] \leq \sigma^2$ . The Jacobian matrix of function  $f: \mathbb{R}^{m_1} \to \mathbb{R}^{m_2}$  is in dimension  $\mathbb{R}^{m_1 \times m_2}$ . We recall the definition of general subgradient and subdifferential following [6, 24].
Definition 3.1 (subgradient and subdifferential). Consider a function  $f: \mathbb{R}^n \to \mathbb{R} \cup \{\infty\}$  and a point with  $f(x)$  finite. A vector  $v \in \mathbb{R}^n$  is a general subgradient of  $f$  at  $x$ , if
$$
f (y) \geq f (x) + \langle v, y - x \rangle + o (\| y - x \|), \quad \text {a s} y \rightarrow x.
$$
The subdifferential  $\partial f(x)$  is the set of subgradients of  $f$  at point  $x$ .
For simplicity, we abuse the notation and also use  $\partial f(x)$  to denote one subgradient from the corresponding subgradient set when no confusion could be caused. We use  $\partial f(x;\mathcal{B})$  to represent a stochastic unbiased estimator of the subgradient  $\partial f(x)$  that is evaluated on a sample batch  $\mathcal{B}$ . A function is called  $C^1$ -smooth if it is continuously differentiable. A function  $f = (f_{1},\dots ,f_{m_{2}}):\mathbb{R}^{m_{1}}\to \mathbb{R}^{m_{2}}$  is called monotone if  $\forall i\in \{1,\ldots ,m_2\}$ ,  $f_{i}:\mathbb{R}^{m_{1}}\to \mathbb{R}$  is monotone with respect to each element of the input. Note that if a Lipschitz continuous function  $f:O\to \mathbb{R}^{m_2}$  is assumed to be non-increasing (resp. non-decreasing), where the domain  $O\subset \mathbb{R}^{m_1}$  is open, then all subgradients of  $f$  are element-wise non-positive (resp. non-negative). We refer the details to Appendix D.1.
A function  $f$  is  $C$ -Lipschitz continuous if  $\| f(x) - f(y) \| \leq C \| x - y \|$ . A differentiable function  $f$  is  $L$ -smooth if  $\| \nabla f(x) - \nabla f(y) \| \leq L \| x - y \|$ . A function  $f: \mathbb{R}^d \to \mathbb{R} \cup \{\infty\}$  is  $\rho$ -weakly-convex if the function  $f(\cdot) + \frac{\rho}{2} \| \cdot \|^2$  is convex. A vector-valued function  $f: \mathbb{R}^d \to \{\mathbb{R} \cup \{\infty\}\}^m$  is called  $\rho$ -weakly-convex if it is  $\rho$ -weakly-convex for each output. It is difficult sometimes impossible to find an  $\epsilon$ -stationary point of a non-smooth weakly-convex function  $F$ , i.e.,  $\mathrm{dist}(0, \partial F(\mathbf{w})) \leq \epsilon$ . For example, an  $\epsilon$ -stationary point of function  $f(x) = |x|$  does not exist for  $0 \leq \epsilon < 1$  unless it is the optimal solution. To tackle this issue, [6] proposed to use the stationarity of the problem's Moreau envelope as the convergence metric, which has become a standard metric for solving weakly-convex problems [7, 22, 31, 41, 43, 17]. Given a weakly-convex function  $\varphi: \mathbb{R}^m \to \mathbb{R}$ , its Moreau envelope and proximal map with  $\lambda > 0$  are constructed as
$$
\varphi_ {\lambda} (x) := \min  _ {y} \{\varphi (y) + \frac {1}{2 \lambda} \| y - x \| ^ {2} \}, \quad \operatorname {p r o x} _ {\lambda \varphi} (x) := \arg \min  _ {y} \{\varphi (y) + \frac {1}{2 \lambda} \| y - x \| ^ {2} \}.
$$
The Moreau envelope is an implicit smoothing of the original problem. Thus it attains a continuous differentiation. As a formal statement, the following lemma follows from standard results [6, 18].
Lemma 3.2. Given a  $\rho$ -weakly-convex function  $\varphi$  and  $\lambda < \rho^{-1}$ , the envelope  $\varphi_{\lambda}$  is  $C^1$ -smooth with gradient given by  $\nabla \varphi_{\lambda}(x) = \lambda^{-1}(x - \operatorname{prox}_{\lambda \varphi}(x))$ .
It is notable that we use a slightly different definition of  $\epsilon$ -stationary point with  $\| \nabla F_{\rho}(\mathbf{w})\|^{2}\leq \epsilon^{2}$ .
Algorithm 1 Stochastic Optimization algorithm for Non-smooth FCCO (SONX)
1: Initialization:  $\mathbf{w}_0$ ,  $\{u_{i,0}: i \in S\}$ .
2: for  $t = 0, \dots, T - 1$  do
3: Draw sample batches  $\mathcal{B}_1^t\sim S$  , and  $\mathcal{B}_{2,i}^{t}\sim \mathcal{D}_{i}$  for each  $i\in \mathcal{B}_1^t$
$$
4: \quad u _ {i, t + 1} = \left\{ \begin{array}{l} (1 - \tau) u _ {i, t} + \tau g _ {i} (\mathbf {w} _ {t}; \mathcal {B} _ {2, i} ^ {t}) + \gamma (g _ {i} (\mathbf {w} _ {t}; \mathcal {B} _ {2, i} ^ {t}) - g _ {i} (\mathbf {w} _ {t - 1}; \mathcal {B} _ {2, i} ^ {t})), \quad i \in \mathcal {B} _ {1} ^ {t} \\ u _ {i, t}, \quad i \notin \mathcal {B} _ {1} ^ {t} \end{array} \right.
$$
5: Compute  $G_{t} = \frac{1}{B_{1}}\sum_{i\in \mathcal{B}_{1}^{t}}\partial g_{i}(\mathbf{w}_{t};\mathcal{B}_{2,i}^{t})\partial f_{i}(u_{i,t})$
6: Update  $\mathbf{w}_{t + 1} = \mathbf{w}_t - \eta G_t$
7: end for
Moreover, for any point  $x \in \mathbb{R}^m$ , the proximal point  $\hat{x} \coloneqq \operatorname{prox}_{\lambda \varphi}(x)$  satisfies [6]
$$
\| \hat {x} - x \| = \lambda \| \nabla \varphi_ {\lambda} (x) \|, \quad \varphi (\hat {x}) \leq \varphi (x), \quad \operatorname {d i s t} (0, \partial \varphi (\hat {x})) \leq \| \nabla \varphi_ {\lambda} (x) \|.
$$
Thus if  $\| \nabla \varphi_{\lambda}(x)\| \leq \epsilon$ , we can say  $x$  is close to a point  $\hat{x}$  that is  $\epsilon$ -stationary, which is called nearly  $\epsilon$ -stationary solution of  $\varphi(x)$ .
4 Algorithms and Convergence
4.1 Non-Smooth Weakly-Convex FCCO
In this section, we assume the following conditions hold for the FCCO problem (1).
Assumption 4.1. For all  $i\in S$  , we assume that
-  $f_{i}$  is  $\rho_{f}$ -weakly-convex,  $C_f$ -Lipschitz continuous and non-decreasing;
-  $g_{i}(\cdot)$  is  $\rho_{g}$ -weakly-convex and  $g_{i}(\cdot; \xi)$  is  $C_{g}$ -Lipschitz continuous;
- Stochastic gradient estimators  $g_{i}(\mathbf{w};\xi)$  and  $\partial g_{i}(\mathbf{w};\xi)$  have bounded variance  $\sigma^2$
Proposition 4.2. Under Assumption 4.1,  $F(\mathbf{w})$  in (1) is  $\rho_{F}$  weakly convex with  $\rho_{F} = \sqrt{d_{1}}\rho_{g}C_{f} + \rho_{f}C_{g}^{2}$ .
One challenge in solving FCCO is the lack of access to unbiased estimation of the subgradients  $\frac{1}{n}\sum_{i\in S}\partial g_i(\mathbf{w})\partial f_i(g_i(\mathbf{w}))$  due to the expectation form of  $g_{i}(\mathbf{w})$  inside a non-linear function  $f_{i}$ . A common solution in existing works for solving smooth FCCO is to maintain function value estimators  $\{u_{i}:i\in S\}$  for  $\{g_i(\mathbf{w}):i\in S\}$ , and approximate the true gradient by a stochastic version  $\frac{1}{B_1}\sum_{i\in B_1}\partial g_i(\mathbf{w};B_2)\partial f_i(u_i)$  [26, 15], where  $B_{1},B_{2}$  are sampled mini-batches. Simply using a mini-batch estimator of  $g_{i}$  inside  $f_{i}$  does not ensure convergence if mini-batch size is small.
Inspired by existing algorithms of smooth FCCO, a simple method for solving non-smooth FCCO is presented in Algorithm 1 referred to as SONX. A key step is the step 4, which uses the multi-block-single-probe variance reduced (MSVR) estimator proposed in [15] to update  $\{u_i:i\in S\}$  in a block-wise manner. It is an advanced variance reduced update strategy for multi-block variable inspired by STORM [5]. In the update of MSVR estimator, for each sampled  $i\in \mathcal{B}_1^t$ $u_{i,t}$  is updated following a STORM-like rule with a specialized parameter  $\gamma = \frac{n - B_1}{B_1(1 - \tau)} +(1 - \tau)$  for the error correction term. For the unsampled  $i\notin \mathcal{B}_1^t$  , no update for  $u_{i,t}$  is needed. When  $\gamma = 0$  , the estimator becomes the moving average estimator analyzed in [26] for smooth FCCO, which is also analyzed in the Appendix. With the function values of  $\{g_i(\mathbf{w}_t):i\in \mathcal{S}\}$  well-estimated, the gradient can be approximated by  $G_{t}$  in step 5. Next, we directly update  $\mathbf{w}_t$  by subgradient descent using the stochastic gradient estimator  $G_{t}$  . Note that unlike existing works on smooth FCCO that often maintain a moving average estimator [26] or a STORM estimator [15] for the overall gradient to attain better rates, this is not possible in the non-smooth case as those variance reduction techniques for the overall gradient critically rely on the Lipschitz continuity of  $\nabla F$  , i.e., the smoothness of  $F$
4.2 Non-Smooth Weakly-Convex TCCO
In this section, we consider non-smooth TCCO problem and aim to extend Algorithm 1 to solve it. First of all, for convergence analysis and to ensure the weak convexity of  $F(\mathbf{w})$  in (2), we make the following assumptions.