SlowGuess/NeurIPS_2023_OCR huggingface dataset:
For MIL TPAUC maximization, we train a simple Feed Forward Neural Network (FFNN) with one hidden layer (the number of neurons equals to data dimension) for the two tabular datasets and ResNet20 for the two medical image datasets. Sigmoid transformation is adopted for the output layer to generate prediction score. The training epoch number is fixed as 100 epochs for all methods; the bag batch size is fixed as 16 (resp. 8) and the number of sampled instances per bag is fixed as 4 (resp. 128) for tabular (resp. medical image) datasets; the learning rate is tuned in  $\{1\mathrm{e} - 2,1\mathrm{e} - 3,1\mathrm{e} - 4\}$  and decreased by 10 folds at the end of 50-th and 75-th epoch for all baselines. For SONT (att), we set moving average parameter  $\tau_{1} = \tau_{2}$  as 0.9; tune the parameter  $\gamma_{1} = \gamma_{2} = \gamma$  in  $\{0,1\mathrm{e} - 1,1\mathrm{e} - 2,1\mathrm{e} - 3\}$  and fix the margin parameter of the surrogate loss  $\ell$  as 0.5, and the parameter  $\alpha ,\beta$  in  $\{0.1,0.5,0.9\}$ . Similar parameters in baselines are set the same or tuned similarly. For all experiments, we utilize 5-fold-cross-validation to evaluate the testing performance based on the best validation performance with possible early stopping choice.
Results. The testing results for the regular and MIL TPUAC maximization with different TPAUC measures are summarized in the Table 2. From Table 2, we observe that our method SONX achieves the best performance for regular TPAUC maximization. It is better than the state-of-the-art method SOTAs for TPAUC maximization. We attribute the better performance of SONX to the fact that the objective of SONX is an exact estimator of TPAUC while the smoothed objective of SOTAs is an inexact estimator of TPAUC. We also observe that SONT (att) achieves the best performance in all cases, which is not surprising since it is the only one that directly optimizes the TPAUC surrogate. In contrast, other baselines either optimizes a different objective (MIDAM) or does not ensure convergence due to the use of mini-batch pooling (AUC-M, SOTAs).
Ablation Study. We conduct ablation studies to demonstrate the effect of the error correction term on the training convergence by varying the  $\gamma$  value for SONX and SONT, where  $\gamma_{1} = \gamma_{2} = \gamma$  is set as the same value in SONT. The training convergence results are presented in Figure 1. We can see that an appropriate value of  $\gamma > 0$  can yield a faster convergence than  $\gamma = 0$ , which verifies the faster convergence of using MSVR estimators than using moving average estimators. However, we do observe a gap between theory and practice, as setting a large value of  $\gamma > 1$  as in the theory might not yield convergence. This phenomenon is also observed in [12]. We conjecture that the gap could be fixed by considering convex objectives [40], which is left as future work.
In this paper, we have considered non-smooth weakly-convex two-level and tri-level finite-sum coupled compositional optimization problems. We presented novel convergence analysis of two stochastic algorithms and established their complexity. Applications in deep learning for two-way partial AUC maximization was considered and great performance of proposed algorithms were demonstrated through experiments on multiple datasets. A future work is to prove the convergence of both algorithms for convex objectives.
We thank anonymous reviewers for constructive comments. Q. Hu, D. Zhu and T. Yang were partially supported by NSF Career Award 2246753, NSF Grant 2246757, 2246756 and 2306572.
[33] Tianbao Yang. Algorithmic foundation of deep x-risk optimization. CoRR, abs/2206.00439, 2022.
[34] Tianbao Yang and Yiming Ying. AUC maximization in the era of big data and AI: A survey. ACM Comput. Surv., 55(8):172:1-172:37, 2023.
[35] Zhiyong Yang, Qianqian Xu, Shilong Bao, Yuan He, Xiaochun Cao, and Qingming Huang. When all we need is a piece of the pie: A generic framework for optimizing two-way partial auc. In Marina Meila and Tong Zhang, editors, Proceedings of the 38th International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pages 11820-11829. PMLR, 18-24 Jul 2021.
[36] 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 Kamalika Chaudhuri, Stefanie Jegelka, Le Song, Csaba Szepesvari, Gang Niu, and Sivan Sabato, editors, Proceedings of the 39th International Conference on Machine Learning, volume 162 of Proceedings of Machine Learning Research, pages 25760-25782. PMLR, 17-23 Jul 2022.
[37] Zhuoning Yuan, Yan Yan, Milan Sonka, and Tianbao Yang. Large-scale robust deep auc maximization: A new surrogate loss and empirical studies on medical image classification, 2021.
[38] Junyu Zhang and Lin Xiao. A stochastic composite gradient method with incremental variance reduction. In Advances in Neural Information Processing Systems, pages 9075-9085, 2019.
[39] Junyu Zhang and Lin Xiao. Multilevel composite stochastic optimization via nested variance reduction. SIAM J. Optim., 31(2):1131-1157, 2021.
[40] Xuan Zhang, Necdet Serhat Aybat, and Mert Gurbuzbalaban. Robust accelerated primal-dual methods for computing saddle points. 2021.
[41] Xuan Zhang, Necdet Serhat Aybat, and Mert Gurbuzbalaban. Sapd+: An accelerated stochastic method for nonconvex-concave minimax problems, 2023.
[42] Zhe Zhang and Guanghui Lan. Optimal algorithms for convex nested stochastic composite optimization, 2022.
[43] Renbo Zhao. A primal-dual smoothing framework for max-structured non-convex optimization, 2022.
[44] 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 Kamalika Chaudhuri, Stefanie Jegelka, Le Song, Csaba Szepesvari, Gang Niu, and Sivan Sabato, editors, Proceedings of the 39th International Conference on Machine Learning, volume 162 of Proceedings of Machine Learning Research, pages 27548–27573. PMLR, 17–23 Jul 2022.
[45] Dixian Zhu, Bokun Wang, Zhi Chen, Yaxing Wang, Milan Sonka, Xiaodong Wu, and Tianbao Yang. Provable multi-instance deep auc maximization with stochastic pooling. In International Conference on Machine Learning. PMLR, 2023.
[46] Landi Zhu, Mert Gurbuzbalaban, and Andrzej Ruszczyński. Distributionally robust learning with weakly convex losses: Convergence rates and finite-sample guarantees, 2023.
A Proofs of Theorem 4.6 and Theorem 4.7
In this section, we provide the detailed proofs for Theorem 4.6 and Theorem 4.7. We first give a basic property for weakly-convex functions.
Proposition A.1 (Proposition 2.1 in [7]). Suppose function  $g: \mathbb{R}^d \to \mathbb{R} \cup \{\infty\}$  is lower-semicontinuous. Then  $g$  is  $\rho$ -weakly-convex if and only if
$$
g (y) \geq g (x) + \langle v, y - x \rangle - \frac {\rho}{2} \| y - x \| ^ {2} \tag {10}
$$
holds for all vectors  $v\in \partial g(x)$  and  $x,y\in \mathbb{R}^d$
A.1 Proof of Theorem 4.6
Note that the proof of Lemma 4.5 also implies the following squared-norm error bound,
$$
\mathbb {E} \bigg [ \frac {1}{n} \sum_ {i \in \mathcal {S}} \| u _ {i, t + 1} - g _ {i} (\mathbf {w} _ {t + 1}) \| ^ {2} \bigg ] \leq (1 - \frac {B _ {1} \tau}{2 n}) ^ {t + 1} \frac {1}{n} \sum_ {i \in \mathcal {S}} \| u _ {i, 0} - g _ {i} (\mathbf {w} _ {0}) \| ^ {2} + \frac {4 \tau \sigma^ {2}}{B _ {2}} + \frac {1 6 n ^ {2} C _ {g} ^ {2} M ^ {2} \eta^ {2}}{B _ {1} ^ {2} \tau}.
$$
Proof of Theorem 4.6. Define  $\hat{\mathbf{w}}_t\coloneqq \mathrm{prox}_{F / \bar{\rho}}(\mathbf{w}_t)$ . For a given  $i\in S$ , we have
$$
f _ {i} \left(g _ {i} (\hat {\mathbf {w}} _ {t})\right) - f _ {i} \left(u _ {i, t}\right)
$$
$$
\stackrel {(a)} {\geq} \partial f _ {i} (u _ {i, t}) ^ {\top} \left(g _ {i} (\hat {\mathbf {w}} _ {t}) - u _ {i, t}\right) - \frac {\rho_ {f}}{2} \left\| g _ {i} (\hat {\mathbf {w}} _ {t}) - u _ {i, t} \right\| ^ {2}
$$
$$
\geq \partial f _ {i} (u _ {i, t}) ^ {\top} (g _ {i} (\hat {\mathbf {w}} _ {t}) - u _ {i, t}) - \rho_ {f} \| g _ {i} (\hat {\mathbf {w}} _ {t}) - g _ {i} (\mathbf {w} _ {t}) \| ^ {2} - \rho_ {f} \| g _ {i} (\mathbf {w} _ {t}) - u _ {i, t} \| ^ {2}
$$
$$
\geq \partial f _ {i} (u _ {i, t}) ^ {\top} \left(g _ {i} (\hat {\mathbf {w}} _ {t}) - u _ {i, t}\right) - \rho_ {f} C _ {g} ^ {2} \| \hat {\mathbf {w}} _ {t} - \mathbf {w} _ {t} \| ^ {2} - \rho_ {f} \| g _ {i} (\mathbf {w} _ {t}) - u _ {i, t} \| ^ {2}
$$
$$
\begin{array}{l} \stackrel {(b)} {\geq} \partial f _ {i} (u _ {i, t}) ^ {\top} \left[ g _ {i} (\mathbf {w} _ {t}) - u _ {i, t} + \partial g _ {i} (\mathbf {w} _ {t}) ^ {\top} (\hat {\mathbf {w}} _ {t} - \mathbf {w} _ {t}) - \frac {\rho_ {g}}{2} \| \hat {\mathbf {w}} _ {t} - \mathbf {w} _ {t} \| ^ {2} \right] \\ - \rho_ {f} C _ {g} ^ {2} \| \hat {\mathbf {w}} _ {t} - \mathbf {w} _ {t} \| ^ {2} - \rho_ {f} \| g _ {i} (\mathbf {w} _ {t}) - u _ {i, t} \| ^ {2} \\ \end{array}
$$
$$
\begin{array}{l} \stackrel {(c)} {\geq} \partial f _ {i} (u _ {i, t}) ^ {\top} (g _ {i} (\mathbf {w} _ {t}) - u _ {i, t}) + \partial f _ {i} (u _ {i, t}) ^ {\top} \partial g _ {i} (\mathbf {w} _ {t}) ^ {\top} (\hat {\mathbf {w}} _ {t} - \mathbf {w} _ {t}) - (\frac {\rho_ {g} C _ {f}}{2} + \rho_ {f} C _ {g} ^ {2}) \| \hat {\mathbf {w}} _ {t} - \mathbf {w} _ {t} \| ^ {2} \\ - \rho_ {f} \| g _ {i} (\mathbf {w} _ {t}) - u _ {i, t} \| ^ {2} \\ \end{array}
$$
where (a) follows from the  $\rho_f$ -weak-convexity of  $f_i$ , (b) follows from that  $f_i(\cdot)$  is non-decreasing and the weak convexity of  $g_i$ , (c) is due to  $0 \leq \partial f_i(u_{i,t}) \leq C_f$ . Then it follows
$$
\begin{array}{l} \frac {1}{n} \sum_ {i \in \mathcal {S}} \partial f _ {i} (u _ {i, t}) ^ {\top} \partial g _ {i} (\mathbf {w} _ {t}) ^ {\top} (\hat {\mathbf {w}} _ {t} - \mathbf {w} _ {t}) \\ \leq \frac {1}{n} \sum_ {i \in \mathcal {S}} \left[ f _ {i} \left(g _ {i} (\hat {\mathbf {w}} _ {t})\right) - f _ {i} (u _ {i, t}) - \partial f _ {i} (u _ {i, t}) ^ {\top} \left(g _ {i} (\mathbf {w} _ {t}) - u _ {i, t}\right) + \left(\frac {\rho_ {g} C _ {f}}{2} + \rho_ {f} C _ {g} ^ {2}\right) \| \hat {\mathbf {w}} _ {t} - \mathbf {w} _ {t} \| ^ {2} \right. \\ \left. + \rho_ {f} \left\| g _ {i} \left(\mathbf {w} _ {t}\right) - u _ {i, t} \right\| ^ {2} \right] \tag {11} \\ \end{array}
$$
Now we consider the change in the Moreau envelope:
$$
\begin{array}{l} \mathbb {E} _ {t} \left[ F _ {1 / \bar {\rho}} (\mathbf {w} _ {t + 1}) \right] = \mathbb {E} _ {t} \left[ \min  _ {\tilde {\mathbf {w}}} F (\tilde {\mathbf {w}}) + \frac {\rho}{2} \| \tilde {\mathbf {w}} - \mathbf {w} _ {t + 1} \| ^ {2} \right] \\ \leq \mathbb {E} _ {t} \left[ F (\hat {\mathbf {w}} _ {t}) + \frac {\bar {\rho}}{2} \| \hat {\mathbf {w}} _ {t} - \mathbf {w} _ {t + 1} \| ^ {2} \right] \\ = F \left(\hat {\mathbf {w}} _ {t}\right) + \mathbb {E} _ {t} \left[ \frac {\bar {\rho}}{2} \left\| \hat {\mathbf {w}} _ {t} - \left(\mathbf {w} _ {t} - \eta G _ {t}\right) \right\| ^ {2} \right] \tag {12} \\ \leq F (\hat {\mathbf {w}} _ {t}) + \frac {\bar {\rho}}{2} \| \hat {\mathbf {w}} _ {t} - \mathbf {w} _ {t} \| ^ {2} + \bar {\rho} \mathbb {E} _ {t} [ \eta \langle \hat {\mathbf {w}} _ {t} - \mathbf {w} _ {t}, G _ {t} \rangle ] + \frac {\eta^ {2} \bar {\rho} M ^ {2}}{2} \\ = F _ {1 / \bar {\rho}} (\mathbf {w} _ {t}) + \bar {\rho} \eta \langle \hat {\mathbf {w}} _ {t} - \mathbf {w} _ {t}, \mathbb {E} _ {t} [ G _ {t} ] \rangle + \frac {\eta^ {2} \bar {\rho} M ^ {2}}{2} \\ \end{array}
$$
$$
\mathbb {E} _ {t} [ G _ {t} ] = \frac {1}{n} \sum_ {i \in \mathcal {S}} \partial g _ {i} (\mathbf {w} _ {t}) \partial f _ {i} (u _ {i, t}),
$$
and the second inequality uses the bound of  $\mathbb{E}[\| G_t\|^2]$ , which follows from the Lipschitz continuity and bounded variance assumptions and is denoted by  $M$ .
Combining inequality 29 and 30 yields
$$
\begin{array}{l} \mathbb {E} _ {t} \left[ F _ {1 / \bar {\rho}} \left(\mathbf {w} _ {t + 1}\right) \right] \\ \leq F _ {1 / \bar {\rho}} (\mathbf {w} _ {t}) + \frac {\eta^ {2} \bar {\rho} M ^ {2}}{2} + \frac {\bar {\rho} \eta}{n} \sum_ {i \in \mathcal {S}} \left[ f _ {i} \left(g _ {i} \left(\hat {\mathbf {w}} _ {t}\right)\right) - f _ {i} \left(u _ {i, t}\right) \right. \\ \left. - \partial f _ {i} \left(u _ {i, t}\right) ^ {\top} \left(g _ {i} \left(\mathbf {w} _ {t}\right) - u _ {i, t}\right) + \left(\frac {\rho_ {g} C _ {f}}{2} + \rho_ {f} C _ {g} ^ {2}\right) \left\| \hat {\mathbf {w}} _ {t} - \mathbf {w} _ {t} \right\| ^ {2} + \rho_ {f} \left\| g _ {i} \left(\mathbf {w} _ {t}\right) - u _ {i, t} \right\| ^ {2} \right] \tag {13} \\ = F _ {1 / \bar {\rho}} (\mathbf {w} _ {t}) + \frac {\eta^ {2} \bar {\rho} M ^ {2}}{2} + \frac {\bar {\rho} \eta}{n} \sum_ {i \in \mathcal {S}} \left[ F _ {i} (\hat {\mathbf {w}} _ {t}) - F _ {i} (\mathbf {w} _ {t}) + f _ {i} (g _ {i} (\mathbf {w} _ {t})) - f _ {i} (u _ {i, t}) \right. \\ \left. - \partial f _ {i} \left(u _ {i, t}\right) ^ {\top} \left(g _ {i} \left(\mathbf {w} _ {t}\right) - u _ {i, t}\right) + \left(\frac {\rho_ {g} C _ {f}}{2} + \rho_ {f} C _ {g} ^ {2}\right) \| \hat {\mathbf {w}} _ {t} - \mathbf {w} _ {t} \| ^ {2} + \rho_ {f} \| g _ {i} \left(\mathbf {w} _ {t}\right) - u _ {i, t} \| ^ {2} \right] \\ \end{array}
$$
Due to the  $\rho_F$ -weak convexity of  $F_{i}(\mathbf{w})$ , we have  $(\bar{\rho} -\rho_F)$ -strong convexity of  $\mathbf{w}\mapsto F_i(\mathbf{w}) + \frac{\bar{\rho}}{2}\| \mathbf{w}_t - \mathbf{w}\|^2$ . Then it follows
$$
\begin{array}{l} F _ {i} (\hat {\mathbf {w}} _ {t}) - F _ {i} (\mathbf {w} _ {t}) = \left[ F _ {i} (\hat {\mathbf {w}} _ {t}) + \frac {\bar {\rho}}{2} \| \mathbf {w} _ {t} - \hat {\mathbf {w}} _ {t} \| ^ {2} \right] - \left[ F _ {i} (\mathbf {w} _ {t}) + \frac {\bar {\rho}}{2} \| \mathbf {w} _ {t} - \mathbf {w} _ {t} \| ^ {2} \right] - \frac {\bar {\rho}}{2} \| \mathbf {w} _ {t} - \hat {\mathbf {w}} _ {t} \| ^ {2} \\ \leq \left(\frac {\rho_ {F}}{2} - \bar {\rho}\right) \| \mathbf {w} _ {t} - \hat {\mathbf {w}} _ {t} \| ^ {2} \tag {14} \\ \end{array}
$$
Plugging inequality 32 into inequality 31 yields
$$
\begin{array}{l} \mathbb {E} _ {t} \left[ F _ {1 / \bar {\rho}} \left(\mathbf {w} _ {t + 1}\right) \right] \leq \mathbb {E} \left[ F _ {1 / \bar {\rho}} \left(\mathbf {w} _ {t}\right) \right] + \frac {\eta^ {2} \bar {\rho} M ^ {2}}{2} + \frac {\bar {\rho} \eta}{n} \sum_ {i \in S} \left[ \left(\frac {\rho_ {F}}{2} - \bar {\rho}\right) \| \mathbf {w} _ {t} - \hat {\mathbf {w}} _ {t} \| ^ {2} \right. \\ + f _ {i} \left(g _ {i} \left(\mathbf {w} _ {t}\right)\right) - f _ {i} \left(u _ {i, t}\right) - \partial f _ {i} \left(u _ {i, t}\right) ^ {\top} \left(g _ {i} \left(\mathbf {w} _ {t}\right) - u _ {i, t}\right) \tag {15} \\ \left. + \left(\frac {\rho_ {g} C _ {f}}{2} + \rho_ {f} C _ {g} ^ {2}\right) \| \hat {\mathbf {w}} _ {t} - \mathbf {w} _ {t} \| ^ {2} + \rho_ {f} \| g _ {i} (\mathbf {w} _ {t}) - u _ {i, t} \| ^ {2} \right] \\ \end{array}
$$
Set  $\bar{\rho} = \rho_{F} + \rho_{g}C_{f} + 2\rho_{f}C_{g}^{2}$ . We have
$$
\begin{array}{l} \mathbb {E} _ {t} \big [ F _ {1 / \bar {\rho}} (\mathbf {w} _ {t + 1}) \big ] \leq F _ {1 / \bar {\rho}} (\mathbf {w} _ {t}) + \frac {\eta^ {2} \bar {\rho} M ^ {2}}{2} + \frac {\bar {\rho} \eta}{n _ {+}} \sum_ {i \in \mathcal {S}} \left[ - \frac {\bar {\rho}}{2} \| \mathbf {w} _ {t} - \hat {\mathbf {w}} _ {t} \| ^ {2} + f _ {i} (g _ {i} (\mathbf {w} _ {t})) - f _ {i} (u _ {i, t}) \right] \\ \left. - \partial f _ {i} \left(u _ {i, t}\right) ^ {\top} \left(g _ {i} \left(\mathbf {w} _ {t}\right) - u _ {i, t}\right) + \rho_ {f} \| g _ {i} \left(\mathbf {w} _ {t}\right) - u _ {i, t} \| ^ {2} \right] \\ \stackrel {(a)} {\leq} F _ {1 / \bar {\rho}} (\mathbf {w} _ {t}) + \frac {\eta^ {2} \bar {\rho} M ^ {2}}{2} - \frac {\eta}{2} \| \nabla F _ {1 / \bar {\rho}} (\mathbf {w} _ {t}) \| ^ {2} + \frac {\bar {\rho} \eta}{n} \sum_ {i \in \mathcal {S}} \left[ f _ {i} (g _ {i} (\mathbf {w} _ {t})) - f _ {i} (u _ {i, t}) \right. \\ \left. - \partial f _ {i} \left(u _ {i, t}\right) ^ {\top} \left(g _ {i} \left(\mathbf {w} _ {t}\right) - u _ {i, t}\right) + \rho_ {f} \| g _ {i} \left(\mathbf {w} _ {t}\right) - u _ {i, t} \| ^ {2} \right] \\ \end{array}
$$
where inequality (a) follows from Lemma 3.2.
Using the Lipschitz continuity of  $f_{i}$ , we have
$$
\begin{array}{l} \mathbb {E} _ {t} \left[ F _ {1 / \bar {\rho}} (\mathbf {w} _ {t + 1}) \right] \leq F _ {1 / \bar {\rho}} (\mathbf {w} _ {t}) + \frac {\eta^ {2} \bar {\rho} M ^ {2}}{2} - \frac {\eta}{2} \| \nabla F _ {1 / \bar {\rho}} (\mathbf {w} _ {t}) \| ^ {2} + \frac {\bar {\rho} \eta}{n} \sum_ {i \in \mathcal {S}} 2 C _ {f} \| g _ {i} (\mathbf {w} _ {t}) - u _ {i, t} \| \\ + \frac {\bar {\rho} \eta}{n} \sum_ {i \in \mathcal {S}} \rho_ {f} \| g _ {i} (\mathbf {w} _ {t}) - u _ {i, t} \| ^ {2} \\ \end{array}
$$
By Lemma 4.5, the error bound of the MSVR update gives
$$
\begin{array}{l} \mathbb {E} \left[ \frac {1}{n} \sum_ {i \in \mathcal {S}} \| u _ {i, t} - g _ {i} (\mathbf {w} _ {t}) \| \right] \leq (1 - \mu) ^ {t} \frac {1}{n} \sum_ {i \in \mathcal {S}} \| u _ {i, 0} - g _ {i} (\mathbf {w} _ {0}) \| + R _ {1}, \\ \mathbb {E} \left[ \frac {1}{n} \sum_ {i \in \mathcal {S}} \| u _ {i, t} - g _ {i} (\mathbf {w} _ {t}) \| ^ {2} \right] \leq (1 - \mu) ^ {t} \frac {1}{n} \sum_ {i \in \mathcal {S}} \| u _ {i, 0} - g _ {i} (\mathbf {w} _ {0}) \| ^ {2} + R _ {2}, \\ \end{array}
$$
$$
\mu = \frac {B _ {1} \tau}{2 n}, R _ {1} = \frac {2 \tau^ {1 / 2} \sigma}{B _ {2} ^ {1 / 2}} + \frac {4 n C _ {g} M \eta}{B _ {1} \tau^ {1 / 2}}, R _ {2} = \frac {4 \tau \sigma^ {2}}{B _ {2}} + \frac {1 6 n ^ {2} C _ {g} ^ {2} M ^ {2} \eta^ {2}}{B _ {1} ^ {2} \tau}
$$
$$
\begin{array}{l} \mathbb {E} \left[ F _ {1 / \bar {\rho}} \left(\mathbf {w} _ {t + 1}\right) \right] \leq F _ {1 / \bar {\rho}} \left(\mathbf {w} _ {t}\right) + \frac {\eta^ {2} \bar {\rho} M ^ {2}}{2} - \frac {\eta}{2} \mathbb {E} \left[ \| \nabla F _ {1 / \bar {\rho}} \left(\mathbf {w} _ {t}\right) \| ^ {2} \right] \\ + 2 C _ {f} \bar {\rho} \eta \left((1 - \mu) ^ {t} \frac {1}{n} \sum_ {i \in \mathcal {S}} \| g _ {i} (\mathbf {w} _ {0}) - u _ {i, 0} \| + R _ {1}\right) \tag {16} \\ + C \rho_ {f} \bar {\rho} \eta \left((1 - \mu) ^ {t} \frac {1}{n} \sum_ {i \in S} \| g _ {i} (\mathbf {w} _ {0}) - u _ {i, 0} \| ^ {2} + R _ {2}\right) \\ \end{array}
$$
Taking summation from  $t = 0$  to  $T - 1$  yields
$$
\mathbb {E} \left[ F _ {1 / \bar {\rho}} (\mathbf {w} _ {T}) \right]
$$
$$
\begin{array}{l} \leq F _ {1 / \bar {\rho}} (\mathbf {w} _ {0}) + \frac {\eta^ {2} \bar {\rho} M ^ {2} T}{2} - \frac {\eta}{2} \sum_ {t = 0} ^ {T - 1} \mathbb {E} [ \| \nabla F _ {1 / \bar {\rho}} (\mathbf {w} _ {t}) \| ^ {2} ] \\ + 2 C _ {f} \bar {\rho} \eta \left(\sum_ {t = 0} ^ {T - 1} (1 - \mu) ^ {t} \frac {1}{n} \sum_ {i \in S} \| g _ {i} (\mathbf {w} _ {0}) - u _ {i, 0} \| + R _ {1} T\right) \\ + C \rho_ {f} \bar {\rho} \eta \left((1 - \mu) ^ {t} \frac {1}{n} \sum_ {i \in \mathcal {S}} \| g _ {i} (\mathbf {w} _ {0}) - u _ {i, 0} \| ^ {2} + R _ {2} T\right) \\ \end{array}
$$
$$
\begin{array}{l} \stackrel {(a)} {\leq} F _ {1 / \bar {\rho}} (\mathbf {w} _ {0}) + \frac {\eta^ {2} \bar {\rho} M ^ {2} T}{2} - \frac {\eta}{2} \sum_ {t = 0} ^ {T - 1} \mathbb {E} [ \| \nabla F _ {1 / \bar {\rho}} (\mathbf {w} _ {t}) \| ^ {2} ] \\ + \frac {2 C _ {f} \bar {\rho} \eta}{n \mu} \sum_ {i \in S} \| g _ {i} (\mathbf {w} _ {0}) - u _ {i, 0} \| + 2 C _ {f} \bar {\rho} \eta R _ {1} T + \frac {\rho_ {f} \bar {\rho} \eta}{n \mu} \sum_ {i \in S} \| g _ {i} (\mathbf {w} _ {0}) - u _ {i, 0} \| ^ {2} + 2 \rho_ {f} \bar {\rho} \eta R _ {2} T, \tag {17} \\ \end{array}
$$
where (a) uses  $\sum_{t=0}^{T-1}(1-\mu)^{t} \leq \frac{1}{\mu}$ .
Lower bounding the left-hand-side by  $\min_{\mathbf{w}} F(\mathbf{w})$ , we obtain
$$
\begin{array}{l} \frac {1}{T} \sum_ {t = 0} ^ {T - 1} \mathbb {E} [ \| \nabla F _ {1 / \bar {\rho}} (\mathbf {w} _ {t}) \| ^ {2} ] \\ \leq \frac {2}{\eta T} \left[ F _ {1 / \bar {\rho}} (\mathbf {w} _ {0}) - \min  _ {\mathbf {w}} F (\mathbf {w}) + \frac {\eta^ {2} \bar {\rho} M ^ {2} T}{2} + \frac {2 C _ {f} \bar {\rho} \eta}{n} \sum_ {i \in \mathcal {S}} \| g _ {i} (\mathbf {w} _ {0}) - u _ {i, 0} \| + 2 C _ {f} \bar {\rho} \eta R _ {1} T \right. \\ \left. + \frac {\rho_ {f} \bar {\rho} \eta}{n} \sum_ {i \in S} \| g _ {i} \left(\mathbf {w} _ {0}\right) - u _ {i, 0} \| ^ {2} + \rho_ {f} \bar {\rho} \eta R _ {2} T \right] \\ \leq \frac {2 \Delta}{\eta T} + \eta \bar {\rho} M ^ {2} + \frac {4 C _ {f} \bar {\rho}}{\mu T n} \sum_ {i \in \mathcal {S}} \| g _ {i} (\mathbf {w} _ {0}) - u _ {i, 0} \| + 4 C _ {f} \bar {\rho} R _ {1} + \frac {2 \rho_ {f} \bar {\rho}}{\mu T n} \sum_ {i \in \mathcal {S}} \| g _ {i} (\mathbf {w} _ {0}) - u _ {i, 0} \| ^ {2} + 2 \rho_ {f} \bar {\rho} R _ {2} \\ \leq \frac {C}{T} \left(\frac {1}{\eta} + \frac {1}{\mu}\right) + C \left(\eta + R _ {1} + R _ {2}\right) \\ \end{array}
$$
where we assume  $F_{1 / \bar{\rho}}(\mathbf{w}_0,\mathbf{s}_0,s_0^{\prime}) - \min_{\mathbf{w},\mathbf{s},s^{\prime}}F(\mathbf{w},\mathbf{s},s^{\prime})\leq \Delta$  and
$$
C = \max \{8 \Delta , 1 2 \bar {\rho} M ^ {2}, \frac {1 6 C _ {f} \bar {\rho}}{n} \sum_ {i \in S} \| g _ {i} (\mathbf {w} _ {0}) - u _ {i, 0} \|, \frac {8 \rho_ {f} \bar {\rho}}{n} \sum_ {i \in S} \| g _ {i} (\mathbf {w} _ {0}) - u _ {i, 0} \| ^ {2}, 1 6 C _ {f} \bar {\rho}, 8 \rho_ {f} \bar {\rho} \}.
$$
$$
\begin{array}{l} \frac {1}{T} \sum_ {t = 0} ^ {T - 1} \mathbb {E} \left[ \| \nabla F _ {1 / \bar {\rho}} (\mathbf {w} _ {t}) \| ^ {2} \right] \\ \leq \frac {C}{T} (\frac {1}{\eta} + \frac {2 n}{B _ {1} \tau}) + C (\eta + \frac {2 \tau^ {1 / 2} \sigma}{B _ {2} ^ {1 / 2}} + \frac {4 n C _ {g} M \eta}{B _ {1} \tau^ {1 / 2}} + \frac {4 \tau \sigma^ {2}}{B _ {2}} + \frac {1 6 n ^ {2} C _ {g} ^ {2} M ^ {2} \eta^ {2}}{B _ {1} ^ {2} \tau}) \\ = \mathcal {O} \left(\frac {1}{T} \left(\frac {1}{\eta} + \frac {n}{B _ {1} \tau}\right) + \left(\eta + \frac {\tau^ {1 / 2} \sigma}{B _ {2} ^ {1 / 2}} + \frac {n \eta}{B _ {1} \tau^ {1 / 2}} + \frac {\tau \sigma^ {2}}{B _ {2}} + \frac {n ^ {2} \eta^ {2}}{B _ {1} ^ {2} \tau}\right)\right) \\ \end{array}
$$
$$
\tau = \mathcal {O} (B _ {2} \epsilon^ {4}), \quad \eta = \mathcal {O} \left(\frac {B _ {1} B _ {2} ^ {1 / 2} \epsilon^ {4}}{n}\right)
$$
To reach an  $\epsilon$ -stationary point, we need
$$
T = \mathcal {O} \left(\frac {n}{B _ {1} B _ {2} ^ {1 / 2} \epsilon^ {6}}\right)
$$
A.2 Proof of Theorem 4.7
A formal statement in given below.
Theorem A.2. Under Assumption 4.3, with  $\gamma_{1} = \frac{n_{1}n_{2} - B_{1}B_{2}}{B_{1}B_{2}(1 - \tau_{1})} + (1 - \tau_{1}), \gamma_{2} = \frac{n_{1} - B_{1}}{B_{1}(1 - \tau_{2})} + (1 - \tau_{2}), \tau_{1} = \mathcal{O}\left(\min\{B_{3}, \frac{B_{1}^{1/2}n_{2}^{1/2}}{n_{1}^{1/2}}\} \epsilon^{4}\right) \leq \frac{1}{2}, \tau_{2} = \mathcal{O}(B_{2}\epsilon^{4}) \leq \frac{1}{2}, \eta = \mathcal{O}\left(\min\left\{B_{3}^{1/2}, \frac{B_{1}^{1/4}n_{2}^{1/4}}{n_{1}^{1/4}}, \frac{B_{1}^{1/2}n_{2}^{1/2}}{n_{1}^{1/2}}\right\} \frac{B_{1}B_{2}}{n_{1}n_{2}} \epsilon^{4}\right)$ , and  $\bar{\rho} = \rho_F + 4\rho_fC_g^2 + 2\rho_gC_fC_h^2 + C_fCgL_h$ . Algorithm 2 converges to an  $\epsilon$ -stationary point of the Moreau envelope  $F_{1/\bar{\rho}}$  in  $T = \mathcal{O}\left(\max\left\{\frac{1}{B_3^{1/2}}, \frac{n_1^{1/4}}{B_1^{1/4}n_2^{1/4}}, \frac{n_1^{1/2}}{B_1^{1/2}n_2^{1/2}}\right\} \frac{n_1n_2}{B_1B_2}\epsilon^{-6}\right)$  iterations.
We first define constant  $M^2 \geq \max \left\{ \frac{3C_f^2 C_g^2 \sigma^2}{B_3} + \frac{3C_f^2 C_g^2 C_h^2}{B_2} + \frac{3C_f^2 C_g^2 C_h^2}{B_1}, \tilde{C}_h^2 + \sigma^2 \right\}$  so that  $\mathbb{E}_t[\| G_t\|^2] \leq M^2$  and  $\| v_{i,j,t}\|^2 \leq M^2$  for all  $i \in S_1, j \in S_2$  and  $t$ . Then to prove Theorem A.2, we need the following Lemmas.
Lemma A.3. Consider MSVR update for  $v$ . Assume  $h_{i,j}(\mathbf{w};\xi)$  is  $C_h$ -Lipshitz for all  $(i,j) \in S_1 \times S_2$ , and  $\mathbb{E}[\| G_t\|^2] \leq M^2$ . With  $\gamma_1 = \frac{n_1n_2 - B_1B_2}{B_1B_2(1 - \tau_1)} + (1 - \tau_1)$ , and  $\tau_1 \leq \frac{1}{2}$ , we have