SlowGuess/NeurIPS_2023_OCR huggingface dataset:
The  $L_{h} = 2\tilde{L}_{h}$ -weakly-convexity of  $h_{i,j}(\mathbf{w})$  and  $C_h = 2\tilde{C}_h$ -Lipschitz continuity of  $h_{i,j}(\mathbf{w};\xi ,\zeta)$  directly follows from the  $\tilde{L}_h$ -smoothness of  $h_i(\mathbf{w})$  and  $\tilde{C}_h$ -Lipschitz continuity of  $h_i(\mathbf{w};\xi)$ . Finally,
we show the weakly-convexity of  $f_{i}(g_{i}(h_{i,j}(\mathbf{w}),s_{i}),s^{\prime})$ :
$$
f _ {i} \left(g _ {i} \left(h _ {i, j} (\tilde {\mathbf {w}}), \tilde {s} _ {i}\right) \tilde {s} ^ {\prime}\right)
$$
$$
\begin{array}{l} \stackrel {(a)} {\geq} f _ {i} \left(g _ {i} \left(h _ {i, j} (\mathbf {w}), s _ {i}\right), s ^ {\prime}\right) + \left\langle \partial_ {s ^ {\prime}} f _ {i} \left(g _ {i} \left(h _ {i, j} (\mathbf {w}), s _ {i}\right), s ^ {\prime}\right), \tilde {s} ^ {\prime} - s ^ {\prime} \right\rangle \\ + \left\langle \partial_ {u} f _ {i} \left(g _ {i} \left(h _ {i, j} (\mathbf {w}), s _ {i}\right), s ^ {\prime}\right), g _ {i} \left(h _ {i, j} (\tilde {\mathbf {w}}), \tilde {s} _ {i}\right) - g _ {i} \left(h _ {i, j} (\mathbf {w}), s _ {i}\right) \right\rangle \\ \end{array}
$$
$$
\begin{array}{l} \stackrel {(b)} {\geq} f _ {i} \left(g _ {i} \left(h _ {i, j} (\mathbf {w}), s _ {i}\right), s ^ {\prime}\right) + \left\langle \partial_ {s ^ {\prime}} f _ {i} \left(g _ {i} \left(h _ {i, j} (\mathbf {w}), s _ {i}\right), s ^ {\prime}\right), \tilde {s} ^ {\prime} - s ^ {\prime} \right\rangle \\ + \left\langle \partial_ {u} f _ {i} \left(g _ {i} \left(h _ {i, j} (\mathbf {w}), s _ {i}\right), s ^ {\prime}\right), \left\langle \partial_ {\ell} g _ {i} \left(h _ {i, j} (\mathbf {w}), s _ {i}\right), \ell \left(h _ {i, j} (\tilde {\mathbf {w}})\right) - \ell \left(h _ {i, j} (\mathbf {w})\right) \right\rangle \right\rangle \\ + \left\langle \partial_ {u} f _ {i} \left(g _ {i} \left(h _ {i, j} (\mathbf {w}), s _ {i}\right) s ^ {\prime}\right), \left\langle \partial_ {s _ {i}} g _ {i} \left(h _ {i, j} (\mathbf {w}), s _ {i}\right), \tilde {s} _ {i} - s _ {i} \right\rangle \right\rangle \\ \end{array}
$$
$$
\begin{array}{l} \stackrel {(c)} {\geq} f _ {i} \left(g _ {i} \left(h _ {i, j} (\mathbf {w}), s _ {i}\right), s ^ {\prime}\right) + \left\langle \partial_ {s ^ {\prime}} f _ {i} \left(g _ {i} \left(h _ {i, j} (\mathbf {w}), s _ {i}\right), s ^ {\prime}\right), \tilde {s} ^ {\prime} - s ^ {\prime} \right\rangle \\ + \left\langle \partial_ {u} f _ {i} \left(g _ {i} \left(h _ {i, j} (\mathbf {w}), s _ {i}\right), s ^ {\prime}\right) \partial_ {\ell} g _ {i} \left(h _ {i, j} (\mathbf {w}), s _ {i}\right) \partial \ell \left(h _ {i, j} (\mathbf {w})\right), h _ {i, j} (\tilde {\mathbf {w}}) - h _ {i, j} (\mathbf {w}) \right\rangle \\ - \frac {C _ {f} C _ {g} L _ {\ell}}{2} \| h _ {i, j} (\tilde {\mathbf {w}}) - h _ {i, j} (\mathbf {w}) \| ^ {2} + \langle \partial_ {u} f _ {i} (s ^ {\prime}, g _ {i} (h _ {i, j} (\mathbf {w}), s _ {i})) \partial_ {s _ {i}} g _ {i} (h _ {i, j} (\mathbf {w}), s _ {i}), \tilde {s} _ {i} - s _ {i} \rangle \\ \end{array}
$$
$$
\begin{array}{l} \stackrel {(d)} {\geq} f _ {i} \left(g _ {i} \left(h _ {i, j} (\mathbf {w}), s _ {i}\right), s ^ {\prime}\right) + \left\langle \partial_ {s ^ {\prime}} f _ {i} \left(g _ {i} \left(h _ {i, j} (\mathbf {w}), s _ {i}\right), s ^ {\prime}\right), \tilde {s} ^ {\prime} - s ^ {\prime} \right\rangle \\ + \left\langle \partial_ {u} f _ {i} \left(g _ {i} \left(h _ {i, j} (\mathbf {w}), s _ {i}\right), s ^ {\prime}\right) \partial_ {\ell} g _ {i} \left(h _ {i, j} (\mathbf {w}), s _ {i}\right) \partial \ell \left(h _ {i, j} (\mathbf {w})\right) \nabla h _ {i, j} (\mathbf {w}), \tilde {\mathbf {w}} - \mathbf {w} \right\rangle \\ + \left\langle \partial_ {u} f _ {i} \left(g _ {i} \left(h _ {i, j} (\mathbf {w}), s _ {i}\right), s ^ {\prime}\right) \partial_ {s _ {i}} g _ {i} \left(h _ {i, j} (\mathbf {w}), s _ {i}\right), \tilde {s} _ {i} - s _ {i} \right\rangle - \left(\frac {C _ {f} C _ {g} C _ {h} ^ {2} L _ {\ell}}{2} + \frac {C _ {f} C _ {g} L _ {h}}{2}\right) \| \tilde {\mathbf {w}} - \mathbf {w} \| ^ {2} \\ \end{array}
$$
where (a) uses the convexity of  $f_{i}$ , (b) uses the monotonicity of  $f_{i}$  w.r.t.  $u$  and convexity of  $g_{i}(\ell, s_{i})$  w.r.t.  $(\ell, s_{i})$ , (c) uses monotonicity of  $f_{i}$  w.r.t.  $u$ , monotonicity of  $g_{i}$  w.r.t.  $\ell$  and  $L_{\ell}$ -weak-convexity of  $\ell$ , (d) uses the smoothness of  $h_{i,j}$ . Thus  $f_{i}(g_{i}(h_{i,j}(\mathbf{w}), s_{i}), s')$  is  $\rho_{F} = (C_{f} C_{g} C_{h}^{2} L_{\ell} + C_{f} C_{g} L_{h})$ -weakly-convex w.r.t.  $(\mathbf{w}, s_{i}, s')$ . Therefore,  $\frac{1}{n_{+}} \sum_{i \in S^{+}} f_{i}(g_{i}(h_{i,j}(\mathbf{w}), s_{i}), s')$  is  $\rho_{F}$ -weakly-convex w.r.t.  $(\mathbf{w}, \mathbf{s}, s')$ .
C.2 Algorithms for TPAUC and Multi-instance TPAUC Maximization
Algorithm 5 SONX for TPAUC
1: Initialization:  $\mathbf{w}_0, \{u_{i,0} : i \in S_+\}, \{s_{i,0} : i \in S_+\}, s_0'$
2: for  $t = 0, \dots, T - 1$  do
3: Sample batches  $\mathcal{B}_1^t\subset S_+$  and  $\mathcal{B}_2^t\subset S_{-}$
$$
\begin{array}{l} 4: \quad u _ {i, t + 1} = \left\{ \begin{array}{l} (1 - \tau) u _ {i, t} + \tau \psi_ {i} (\mathbf {w} _ {t}, s _ {i, t}; \mathcal {B} _ {2} ^ {t}) + \gamma (\psi_ {i} (\mathbf {w} _ {t}, s _ {i, t}; \mathcal {B} _ {2} ^ {t}) - \psi_ {i} (\mathbf {w} _ {t - 1}, s _ {i, t - 1}; \mathcal {B} _ {2} ^ {t})), i \in \mathcal {B} _ {1} ^ {t} \\ u _ {i, t}, \quad i \notin \mathcal {B} _ {1} ^ {t} \end{array} \right. \\ 5: \hskip 1 4. 2 2 6 3 7 8 p t s _ {i, t + 1} = \left\{ \begin{array}{l l} s _ {i, t} - \eta \frac {1}{B _ {1}} \partial_ {s} \psi_ {i} (\mathbf {w} _ {t}, s _ {i, t}; \mathcal {B} _ {2} ^ {t}) \partial_ {u} f (u _ {i, t}, s _ {t} ^ {\prime}), & i \in \mathcal {B} _ {1} ^ {t} \\ s _ {i, t}, & i \not \in \mathcal {B} _ {1} ^ {t} \end{array} \right. \\ 6: \quad s _ {t + 1} ^ {\prime} = s _ {t} ^ {\prime} - \eta \frac {1}{B _ {1}} \sum_ {i \in B _ {1} ^ {t}} \partial_ {s ^ {\prime}} f \left(u _ {i, t}, s _ {t} ^ {\prime}\right) \\ 7: \quad \text {C o m p u t} G _ {t} = \frac {1}{B _ {1}} \sum_ {i \in \mathcal {B} _ {1} ^ {t}} \partial_ {w} \psi_ {i} (\mathbf {w} _ {t}, s _ {i, t}; \mathcal {B} _ {2} ^ {t}) \partial_ {u} f (u _ {i, t}, s _ {t} ^ {\prime}) \\ 8: \quad \text {U p d a t e} \mathbf {w} _ {t + 1} = \mathbf {w} _ {t} - \eta G _ {t} \\ 9: \text {e n d f o r} \\ 0: \operatorname {r e t u r n} \mathbf {w} _ {\bar {t}} \text {w i t h} \bar {t} \text {u n i f o r m l y s a m p l e d f r o m} \{0, \dots , T - 1 \}. \\ \end{array}
$$
Algorithm 6 SONT for Multi-instance TPAUC
1: Initialization:  $\mathbf{w}_0, \{u_{i,0} : i \in S_+\}, \{s_{i,0} : i \in S_+\}, s_0', \{v_{i,j,0} : (i,j) \in S_+ \times S_-\}$
2: for  $t = 0, \dots, T - 1$  do
3: Sample batches  $\mathcal{B}_1^t\subset S_+$ $\mathcal{B}_2^t\subset S_{-}$  , and  $\mathcal{B}_{3,i}^{t}\subset X_{i}$  for  $i\in \mathcal{B}_1^t\cup \mathcal{B}_2^t$
$$
\begin{array}{l} 4: \quad v _ {i, t + 1} = \left\{ \begin{array}{l l} \Pi_ {\tilde {C} _ {h}} [ (1 - \tau_ {1}) v _ {i, t} + \tau_ {1} h _ {i} (\mathbf {w} _ {t}; \mathcal {B} _ {3, i} ^ {t}) + \gamma_ {1} (h _ {i} (\mathbf {w} _ {t}; \mathcal {B} _ {3, i} ^ {t}) - h _ {i} (\mathbf {w} _ {t - 1}; \mathcal {B} _ {3, i} ^ {t})) ], & i \in \mathcal {B} _ {1} ^ {t} \\ \Pi_ {\tilde {C} _ {h}} [ (1 - \tau_ {1}) v _ {i, t} + \tau_ {1} h _ {i} (\mathbf {w} _ {t}; \mathcal {B} _ {3, i} ^ {t}) + \gamma_ {2} (h _ {i} (\mathbf {w} _ {t}; \mathcal {B} _ {3, i} ^ {t}) - h _ {i} (\mathbf {w} _ {t - 1}; \mathcal {B} _ {3, i} ^ {t})) ], & i \in \mathcal {B} _ {2} ^ {t} \\ v _ {i, t}, \quad i \notin \mathcal {B} _ {1} ^ {t} \text {a n d} i \notin \mathcal {B} _ {2} ^ {t} \end{array} \right. \\ 5: \quad u _ {i, t + 1} = \left\{ \begin{array}{l l} (1 - \tau_ {2}) u _ {i, t} + \frac {1}{B _ {2}} \sum_ {j \in \mathcal {B} _ {2} ^ {t}} [ \tau_ {2} g (v _ {j, t} - v _ {i, t}, s _ {i, t}) \\ \qquad + \gamma_ {3} (g (v _ {j, t} - v _ {i, t}, s _ {i, t}) - g (v _ {j, t - 1} - v _ {i, t - 1}, s _ {i, t - 1})) ], & i \in \mathcal {B} _ {1} ^ {t} \\ u _ {i, t}, & i \notin \mathcal {B} _ {1} ^ {t} \end{array} \right. \\ 6: \hskip 1 4. 2 2 6 3 7 8 p t s _ {i, t + 1} = \left\{ \begin{array}{l l} s _ {i, t} - \eta_ {1} \frac {1}{B _ {1}} \left[ \frac {1}{B _ {2}} \sum_ {j \in \mathcal {B} _ {2} ^ {t}} \partial_ {s _ {i}} g (v _ {j, t} - v _ {i, t}, s _ {i, t}) \right] \partial_ {u} f (s _ {t} ^ {\prime}, u _ {i, t}), & i \in \mathcal {B} _ {1} ^ {t} \\ s _ {i, t}, & i \notin \mathcal {B} _ {1} ^ {t} \end{array} \right. \\ 7: \quad s _ {t + 1} ^ {\prime} = s _ {t} ^ {\prime} - \eta_ {2} \frac {1}{B _ {1}} \sum_ {i \in \mathcal {B} _ {1} ^ {t}} \partial_ {s ^ {\prime}} f (u _ {i, t}, s _ {t} ^ {\prime}) \\ 8: \quad G _ {t} = \frac {1}{B _ {1}} \sum_ {i \in \mathcal {B} _ {1} ^ {t}} \partial_ {u} f \left(u _ {i, t}, s _ {t} ^ {\prime}\right) \\ 9: \quad \left[ \frac {1}{B _ {2}} \sum_ {j \in \mathcal {B} _ {2} ^ {t}} \left(\nabla h _ {j} \left(\mathbf {w} _ {t}; \mathcal {B} _ {3, j} ^ {t}\right) - \nabla h _ {i} \left(\mathbf {w} _ {t}; \mathcal {B} _ {3, i} ^ {t}\right)\right) \partial_ {v} g \left(v _ {j, t} - v _ {i, t}, s _ {i, t}\right) \right] \\ 1 0: \quad \text {U p d a t e} \mathbf {w} _ {t + 1} = \mathbf {w} _ {t} - \eta G _ {t} \\ 1 1: \text {e n d f o r} \\ 1 2: \mathbf {r e t u r n} \mathbf {w} _ {\bar {t}} \text {w i t h} \bar {t} \text {u n i f o r m l y s a m p l e d f r o m} \{0, \dots , T - 1 \}. \\ \end{array}
$$
C.3 TPAUC in MIL with smoothed-max pooling and attention-based pooling
We can extend our results to smoothed-max pooling and attention-based pooling.
Smoothed-max Pooling. The smoothed-max pooling can be written as [45]:
$$
h _ {\mathbf {w}} (X) = \tau \log \left(\frac {1}{| X |} \sum_ {\mathbf {x} \in X} \exp \left(\phi (\mathbf {w}; \mathbf {x}) / \tau\right)\right), \tag {24}
$$
where  $\tau > 0$  is a hyperparameter and  $\phi(\mathbf{w};\mathbf{x}) = e(\mathbf{w}_e,\mathbf{x})^\top \mathbf{w}_c$  is the prediction score for instance  $\mathbf{x}$ . We can see that  $h_{\mathbf{w}}(X)$  itself is a compositional function. To map the problem into TCCO, we define  $h_i(\mathbf{w}) = \frac{1}{|X_i|} \sum_{\mathbf{x} \in X_i} \exp(\phi(\mathbf{w};\mathbf{x}) / \tau) + C$ , where  $C > 0$  is a constant. Then the objective function becomes
$$
\min  _ {\mathbf {w}, \mathbf {s} ^ {\prime}, \mathbf {s}} \frac {1}{n _ {+}} \sum_ {X _ {i} \in \mathcal {S} _ {+}} f _ {i} (\psi_ {i} (\mathbf {w}, s _ {i}), s ^ {\prime}),
$$
$$
\text {w h e r e} f _ {i} \left(g, s ^ {\prime}\right) = s ^ {\prime} + \frac {\left(g - s ^ {\prime}\right) _ {+}}{\alpha}, \tag {25}
$$
$$
\psi_ {i} (\mathbf {w}, s _ {i}) = \frac {1}{n _ {-}} \sum_ {X _ {j} \in \mathcal {S} _ {-}} s _ {i} + \frac {\left(\ell \left(\tau \log h _ {j} (\mathbf {w}) - \tau \log h _ {i} (\mathbf {w})\right) - s _ {i}\right) _ {+}}{\beta},
$$
In this case we define  $g_{i}(\ell (\mathbf{v}),s_{i}) = s_{i} + \frac{(\ell(\tau\log v_{1} - \tau\log v_{2}) - s_{i})_{+}}{\beta}$  and  $h_{i,j}(\mathbf{w}) = [h_i(\mathbf{w}),h_j(\mathbf{w})]$ . We can still prove that  $g_{i}(\ell (\mathbf{v}),s_{i})$  is monotone w.r.t to each component of  $\mathbf{v}$ . It is not difficult to prove that  $\ell (\tau \log v_1 - \tau \log v_2)$  is weakly convex w.r.t  $\mathbf{v}$  because  $\tau \log v_{1} - \tau \log v_{2}$  is a smooth mapping of  $\mathbf{v}$  due to  $\mathbf{v}\geq C$  and  $\ell$  is a convex function [8]. As a result, since  $g_{i}(\ell ,s_{i})$  is non-decreasing and convex w.r.t to  $\ell$ , it is easy to prove that  $g_{i}(\ell (\mathbf{v}),s_{i})$  is weakly convex w.r.t  $\mathbf{v}$  and is monotone (either non-decreasing or non-increasing) w.r.t to each component of  $\mathbf{v}$ . Hence, assuming  $h_i(\mathbf{w})$  is a smooth and Lipschitz continuous function, we can prove that  $g_{i}(h_{i,j}(\mathbf{w}),s_{i})$  is weakly convex w.r.t. to  $\mathbf{w}$ .
Attention-based Pooling. Attention-based pooling was recently introduced for deep MIL [14], which aggregates the feature representations using attention, i.e.,
$$
E (\mathbf {w}; X) = \sum_ {\mathbf {x} \in X} \frac {\exp (g (\mathbf {w} ; \mathbf {x}))}{\sum_ {\mathbf {x} ^ {\prime} \in X} \exp (g (\mathbf {w} ; \mathbf {x} ^ {\prime}))} e \left(\mathbf {w} _ {e}; \mathbf {x}\right) \tag {26}
$$
where  $g(\mathbf{w};\mathbf{x})$  is a parametric function, e.g.,  $g(\mathbf{w};\mathbf{x}) = \mathbf{w}_a^\top \tanh(Ve(\mathbf{w}_e;\mathbf{x})) + C$ , where  $V \in \mathbb{R}^{m \times d_o}$  and  $\mathbf{w}_a \in \mathbb{R}^m$ . Based on the aggregated feature representation, the bag level prediction can be computed by
$$
\begin{array}{l} h _ {\mathbf {w}} (\mathbf {w}, X) = \left(\mathbf {w} _ {c} ^ {\top} E (\mathbf {w}; X)\right) \tag {27} \\ = \left(\sum_ {\mathbf {x} \in X} \frac {\exp (g (\mathbf {w} ; \mathbf {x})) \delta (\mathbf {w} ; \mathbf {x})}{\sum_ {\mathbf {x} ^ {\prime} \in X} \exp (g (\mathbf {w} ; \mathbf {x} ^ {\prime}))}\right), \\ \end{array}
$$
where  $\delta (\mathbf{w};\mathbf{x}) = \mathbf{w}_c^\top e(\mathbf{w}_e;\mathbf{x})$
We can see that  $h_{\mathbf{w}}(X)$  itself is a compositional function. To map the problem into TCCO, we define  $h_i^1(\mathbf{w}) = \frac{1}{|X_i|} \sum_{\mathbf{x} \in X_i} \exp(g(\mathbf{w}; \mathbf{x})) \delta(\mathbf{w}; \mathbf{x})$ , and  $h_i^2(\mathbf{w}) = \frac{1}{|X_i|} \sum_{\mathbf{x}' \in X_i} \exp(g(\mathbf{w}; \mathbf{x}')')$ . Assume  $|\mathbf{w}_a^\top \tanh(Ve(\mathbf{w}_e; \mathbf{x}))| \leq C_b$  then  $h_i^2(\mathbf{w}) \geq \exp(C - C_b)$ . Then the objective function becomes
$$
\min  _ {\mathbf {w}, \mathbf {s} ^ {\prime}, \mathbf {s}} \frac {1}{n _ {+}} \sum_ {X _ {i} \in \mathcal {S} _ {+}} f _ {i} \left(\psi_ {i} (\mathbf {w}, s _ {i}), s ^ {\prime}\right),
$$
$$
\text {w h e r e} f _ {i} \left(g, s ^ {\prime}\right) = s ^ {\prime} + \frac {\left(g - s ^ {\prime}\right) _ {+}}{\alpha}, \quad \psi_ {i} (\mathbf {w}, s _ {i}) = \frac {1}{n _ {-}} \sum_ {X _ {j} \in \mathcal {S} _ {-}} s _ {i} + \frac {\left(\ell \left(\frac {h _ {j} ^ {1} (\mathbf {w})}{h _ {j} ^ {2} (\mathbf {w})} - \frac {h _ {i} ^ {1} (\mathbf {w})}{h _ {i} ^ {2} (\mathbf {w})}\right) - s _ {i}\right) _ {+}}{\beta}, \tag {28}
$$
In this case we define  $g_{i}(\ell(\mathbf{v}), s_{i}) = s_{i} + \frac{(\ell(\frac{v_{3}}{v_{4}} - \frac{v_{1}}{v_{2}}) - s_{i})_{+}}{\beta}$  and  $h_{i,j}(\mathbf{w}) = [h_{i}^{1}(\mathbf{w}), h_{i}^{2}(\mathbf{w}), h_{j}^{1}(\mathbf{w}), h_{j}^{2}(\mathbf{w})]$ . We can still prove that  $g_{i}(\ell(\mathbf{v}), s_{i})$  is monotone w.r.t to each component of  $\mathbf{v}$ . It is not difficult to prove that  $\ell\left(\frac{v_{3}}{v_{4}} - \frac{v_{1}}{v_{2}}\right)$  is weakly convex w.r.t  $\mathbf{v}$  because  $\frac{v_{3}}{v_{4}} - \frac{v_{1}}{v_{2}}$  is a smooth mapping of  $\mathbf{v}$  when  $v_{2}, v_{4}$  are lower bounded and  $\ell$  is a convex function [8]. As a result, since  $g_{i}(\ell, s_{i})$  is non-decreasing and convex w.r.t to  $\ell$ , it is easy to prove that  $g_{i}(\ell(\mathbf{v}), s_{i})$  is weakly convex w.r.t  $\mathbf{v}$  and is monotone (either non-decreasing or non-increasing) w.r.t to each component of  $\mathbf{v}$ . Hence, assuming  $h_{i}^{1}(\mathbf{w}), h_{i}^{2}(\mathbf{w})$  are smooth and Lipschitz continuous, we can prove that  $g_{i}(h_{i,j}(\mathbf{w}), s_{i})$  is weakly convex w.r.t. to  $\mathbf{w}$ .
C.4 Convergence Analysis of TPAUC Maximization
C.4.1 Convergence analysis for Algorithm 5
We first consider TPAUC maximization in the regular learning setting. Define  $F(\mathbf{w}, \mathbf{s}, s') \coloneqq \frac{1}{n_+} \sum_{X_i \in S_+} f_i(\psi_i(\mathbf{w}, s_i), s')$ . Due to the weak-convexity of  $F(\mathbf{w}, \mathbf{s}, s')$  w.r.t.  $(\mathbf{w}, \mathbf{s}, s')$ , we consider the following Moreau envelope and proximal map defined as
$$
F _ {\lambda} (\mathbf {w}, \mathbf {s}, s ^ {\prime}) = \min _ {\tilde {\mathbf {w}}, \tilde {\mathbf {s}}, \tilde {s} ^ {\prime}} F (\tilde {\mathbf {w}}, \tilde {\mathbf {s}}, \tilde {s} ^ {\prime}) + \frac {1}{2 \lambda} \left(\| \tilde {\mathbf {w}} - \mathbf {w} \| ^ {2} + \| \tilde {\mathbf {s}} - \mathbf {s} \| ^ {2} + \| \tilde {s} ^ {\prime} - s ^ {\prime} \| ^ {2}\right),
$$
$$
\operatorname {p r o x} _ {\lambda F} (\mathbf {w}, \mathbf {s}, s ^ {\prime}) = \underset {\tilde {\mathbf {w}}, \tilde {\mathbf {s}}, \tilde {s} ^ {\prime}} {\arg \min } F (\tilde {\mathbf {w}}, \tilde {\mathbf {s}}, \tilde {s} ^ {\prime}) + \frac {1}{2 \lambda} \left(\| \tilde {\mathbf {w}} - \mathbf {w} \| ^ {2} + \| \tilde {\mathbf {s}} - \mathbf {s} \| ^ {2} + \| \tilde {s} ^ {\prime} - s ^ {\prime} \| ^ {2}\right).
$$
Following the same proof of Lemma 4.5, we have the following error bound
Lemma C.3. Consider the update for  $\{u_{i,t} : X_i \in \mathcal{S}_+\}$  in Algorithm 5. Assume  $\psi_i(\mathbf{w}, s_i)$  is  $C_\psi$ -Lipschitz continuous for all  $X_i \in \mathcal{S}_+$ . Assume  $\mathbb{E}_t[\| G_t\|^2] \leq M^2$  and  $\mathbb{E}_t[\| \frac{1}{B_1} \sum_{X_i \in \mathcal{B}_1^t} \partial_s \psi_i(\mathbf{w}_t, s_{i,t}; \mathcal{B}_2^t) \partial_u f(u_{i,t}, s_t') e_i\|^2] \leq M^2$ , where  $e_i$  is the  $n_+$ -dimensional vector with 1 at the  $i$ -th entry and 0 everywhere else. With  $\gamma = \frac{n_+ - B_1}{B_1(1 - \tau)} + (1 - \tau)$  and  $\tau \leq \frac{1}{2}$ , we have
$$
\begin{array}{l} \mathbb {E} \left[ \frac {1}{n _ {+}} \sum_ {X _ {i} \in \mathcal {S} _ {+}} \| u _ {i, t + 1} - \psi_ {i} (\mathbf {w} _ {t + 1}, s _ {i, t + 1}) \| \right] \\ \leq (1 - \frac {B _ {1} \tau}{2 n _ {1}}) ^ {t + 1} \frac {1}{n} \sum_ {X _ {i} \in \mathcal {S} _ {+}} \| u _ {i, 0} - \psi_ {i} (\mathbf {w} _ {0}, s _ {i, 0}) \| + \frac {2 \tau^ {1 / 2} \sigma}{B _ {2} ^ {1 / 2}} + \frac {8 n _ {+} C _ {\psi} M \eta}{B _ {1} \tau^ {1 / 2}}. \\ \end{array}
$$
Then we have following convergence guarantee.
Theorem C.4. Under the assumptions given in Lemma C.1, with  $\gamma = \frac{n_{+} - B_{1}}{B_{1}(1 - \tau)} + (1 - \tau)$ ,  $\tau = \mathcal{O}(B_2\epsilon^4) \leq \frac{1}{2}$ ,  $\eta = \mathcal{O}\left(\frac{B_1B_2^{1/2}\epsilon^4}{n_+}\right)$ , and  $\bar{\rho} = \rho_F + \rho_\psi C_f$ , Algorithm 5 converges to an  $\epsilon$ -stationary point of the Moreau envelope  $F_{1/\bar{\rho}}$  in  $T = \mathcal{O}\left(\frac{n_{+}}{B_1B_2^{1/2}}\epsilon^{-6}\right)$  iterations.
Proof of Theorem C.4. Define  $(\hat{\mathbf{w}}_t,\hat{\mathbf{s}}_t,\hat{s}_t^{\prime}):= \mathrm{prox}_{F / \bar{\rho}}(\mathbf{w}_t,\mathbf{s}_t,s_t^{\prime})$  . For a given  $X_{i}\in S_{+}$  , we have  $f_{i}(\psi_{i}(\hat{\mathbf{w}}_{t},\hat{s}_{i,t}),\hat{s}_{t}^{\prime}) - f_{i}(u_{i,t},s_{t}^{\prime})$
$$
\begin{array}{l} \stackrel {(a)} {\geq} \partial_ {s ^ {\prime}} f _ {i} (u _ {i, t}, s _ {t} ^ {\prime}) (\hat {s} _ {t} ^ {\prime} - s _ {t} ^ {\prime}) + \partial_ {u} f _ {i} (u _ {i, t}, s _ {t} ^ {\prime}) (\psi_ {i} (\hat {\mathbf {w}} _ {t}, \hat {s} _ {i, t}) - u _ {i, t}) \\ \stackrel {(b)} {\geq} \partial_ {s ^ {\prime}} f _ {i} (u _ {i, t}, s _ {t} ^ {\prime}) (\hat {s} _ {t} ^ {\prime} - s _ {t} ^ {\prime}) + \partial_ {u} f _ {i} (u _ {i, t}, s _ {t} ^ {\prime}) \bigg [ \psi_ {i} (\mathbf {w} _ {t}, s _ {i, t}) - u _ {i, t} + \langle \partial_ {w} \psi_ {i} (\mathbf {w} _ {t}, s _ {i, t}), \hat {\mathbf {w}} _ {t} - \mathbf {w} _ {t} \rangle \\ - \frac {\rho_ {\psi}}{2} \left\| \hat {\mathbf {w}} _ {t} - \mathbf {w} _ {t} \right\| ^ {2} + \left\langle \partial_ {s _ {i}} \psi_ {i} (\mathbf {w} _ {t}, s _ {i, t}), \hat {s} _ {i, t} - s _ {i, t} \right\rangle - \frac {\rho_ {\psi}}{2} \left\| \hat {s} _ {i, t} - s _ {i, t} \right\| ^ {2} \Biggr ] \\ \stackrel {(c)} {\geq} \partial_ {s ^ {\prime}} f _ {i} (u _ {i, t}, s _ {t} ^ {\prime}) (\hat {s} _ {t} ^ {\prime} - s _ {t} ^ {\prime}) + \partial_ {u} f _ {i} (u _ {i, t}, s _ {t} ^ {\prime}) [ \psi_ {i} (\mathbf {w} _ {t}, s _ {i, t}) - u _ {i, t} ] + \langle \partial_ {u} f _ {i} (u _ {i, t}, s _ {t} ^ {\prime}) \partial_ {w} \psi_ {i} (\mathbf {w} _ {t}, s _ {i, t}), \hat {\mathbf {w}} _ {t} - \mathbf {w} _ {t} \rangle \\ + \langle \partial_ {u} f _ {i} (u _ {i, t}, s _ {t} ^ {\prime}) \partial_ {s _ {i}} \psi_ {i} (\mathbf {w} _ {t}, s _ {i, t}), \hat {s} _ {i, t} - s _ {i, t} \rangle - \frac {\rho_ {\psi} C _ {f}}{2} \left(\| \hat {\mathbf {w}} _ {t} - \mathbf {w} _ {t} \| ^ {2} + \| \hat {s} _ {i, t} - s _ {i, t} \| ^ {2}\right) \\ \end{array}
$$
where (a) follows from the convexity of  $f_{i}$ , (b) follows from the monotonicity of  $f_{i}(\cdot, s')$  and weak convexity of  $\psi_{i}$ , (c) is due to  $0 \leq \partial_{u}f_{i}(u_{i,t}, s_{t}') \leq C_{f}$ . Then it follows
$$
\begin{array}{l} \frac {1}{n _ {+}} \sum_ {X _ {i} \in S _ {+}} \left[ \partial_ {s ^ {\prime}} f _ {i} (u _ {i, t}, s _ {t} ^ {\prime}) (\hat {s} _ {t} ^ {\prime} - s _ {t} ^ {\prime}) + \langle \partial_ {u} f _ {i} (u _ {i, t}, s _ {t} ^ {\prime}) \partial_ {w} \psi_ {i} (\mathbf {w} _ {t}, s _ {i, t}), \hat {\mathbf {w}} _ {t} - \mathbf {w} _ {t} \rangle \right. \\ \left. + \left\langle \partial_ {u} f _ {i} \left(u _ {i, t}, s _ {t} ^ {\prime}\right) \partial_ {s _ {i}} \psi_ {i} \left(\mathbf {w} _ {t}, s _ {i, t}\right), \hat {s} _ {i, t} - s _ {i, t} \right\rangle \right] \tag {29} \\ \leq \frac {1}{n _ {+}} \sum_ {X _ {i} \in S _ {+}} \left[ f _ {i} \big (\psi_ {i} \big (\hat {\mathbf {w}} _ {t}, \hat {s} _ {i, t} \big), \hat {s} _ {t} ^ {\prime} \big) - f _ {i} (u _ {i, t}, s _ {t} ^ {\prime}) - \partial_ {u} f _ {i} (u _ {i, t}, s _ {t} ^ {\prime}) \big [ \psi_ {i} (\mathbf {w} _ {t}, s _ {i, t}) - u _ {i, t} \big ] \right. \\ \left. + \frac {\rho_ {\psi} C _ {f}}{2} \left(\| \hat {\mathbf {w}} _ {t} - \mathbf {w} _ {t} \| ^ {2} + \| \hat {s} _ {i, t} - s _ {i, t} \| ^ {2}\right) \right] \\ \end{array}
$$
$$
\begin{array}{l} \mathbb {E} _ {t} \left[ F _ {1 / \bar {\rho}} \left(\mathbf {w} _ {t + 1}, \mathbf {s} _ {t + 1}, s _ {t + 1} ^ {\prime}\right) \right] \\ = \mathbb {E} _ {t} \left[ \min  _ {\tilde {\mathbf {w}}, \tilde {\mathbf {s}}, \tilde {s} ^ {\prime}} F (\tilde {\mathbf {w}}, \tilde {\mathbf {s}} _ {t}, \tilde {s} _ {t} ^ {\prime}) + \frac {\bar {\rho}}{2} \left(\| \tilde {\mathbf {w}} - \mathbf {w} _ {t + 1} \| ^ {2} + \| \tilde {\mathbf {s}} - \mathbf {s} _ {t + 1} \| ^ {2} + \| \tilde {s} ^ {\prime} - \mathbf {s} _ {t + 1} ^ {\prime} \| ^ {2}\right) \right] \\ \leq \mathbb {E} _ {t} \left[ F (\hat {\mathbf {w}} _ {t}, \hat {\mathbf {s}} _ {t}, \hat {s} _ {t} ^ {\prime}) + \frac {\bar {\rho}}{2} \left(\| \hat {\mathbf {w}} _ {t} - \mathbf {w} _ {t + 1} \| ^ {2} + \| \hat {\mathbf {s}} _ {t} - \mathbf {s} _ {t + 1} \| ^ {2} + \| \hat {s} _ {t} ^ {\prime} - s _ {t + 1} ^ {\prime} \| ^ {2}\right) \right] \\ = F \left(\hat {\mathbf {w}} _ {t}, \hat {\mathbf {s}} _ {t}, \hat {s} _ {t} ^ {\prime}\right) + \mathbb {E} _ {t} \left[ \frac {\bar {\rho}}{2} \left(\| \hat {\mathbf {w}} _ {t} - \left(\mathbf {w} _ {t} - \eta G _ {t}\right) \| ^ {2} + \| \hat {\mathbf {s}} _ {t} - \left(\mathbf {s} _ {t} - \eta G _ {t} ^ {1}\right) \| ^ {2} \right. \right. \\ \left. + \left\| \hat {s} _ {t} ^ {\prime} - \left(s _ {t} ^ {\prime} - \eta G _ {t} ^ {2}\right) \right\| ^ {2}\right) \Bigg ] \tag {30} \\ \leq F (\hat {\mathbf {w}} _ {t}, \hat {\mathbf {s}} _ {\mathbf {t}}, \hat {s} _ {t} ^ {\prime}) + \frac {\bar {\rho}}{2} \left(\| \hat {\mathbf {w}} _ {t} - \mathbf {w} _ {t} \| ^ {2} + \| \hat {\mathbf {s}} _ {t} - \mathbf {s} _ {t} \| ^ {2} + \| \hat {s} _ {t} ^ {\prime} - s _ {t} ^ {\prime} \| ^ {2}\right) \\ + \bar {\rho} \mathbb {E} _ {t} [ \eta \langle \hat {\mathbf {w}} _ {t} - \mathbf {w} _ {t}, G _ {t} \rangle + \eta \langle \hat {\mathbf {s}} _ {t} - \mathbf {s} _ {t}, G _ {t} ^ {1} \rangle + \eta \langle \hat {s} _ {t} ^ {\prime} - s _ {t} ^ {\prime}, G _ {t} ^ {2} \rangle ] + \frac {3 \eta^ {2} \bar {\rho} M ^ {2}}{2} \\ = F _ {1 / \bar {\rho}} (\mathbf {w} _ {t}, \mathbf {s} _ {t}, s _ {t} ^ {\prime}) + \bar {\rho} \mathbb {E} _ {t} [ \eta \langle \hat {\mathbf {w}} _ {t} - \mathbf {w} _ {t}, G _ {t} \rangle + \eta \langle \hat {\mathbf {s}} _ {t} - \mathbf {s} _ {t}, G _ {t} ^ {1} \rangle + \eta \langle \hat {s} _ {t} ^ {\prime} - s _ {t} ^ {\prime}, G _ {t} ^ {2} \rangle ] \\ + \frac {3 \eta^ {2} \bar {\rho} M ^ {2}}{2} \\ \end{array}
$$