SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
\begin{array}{l} \mathbb {E} \left[ \frac {1}{n _ {1}} \sum_ {i \in S _ {1}} \frac {1}{n _ {2}} \sum_ {j \in S _ {2}} \| v _ {i, j, t + 1} - h _ {i, j} (\mathbf {w} _ {t + 1}) \| \right] \\ \leq \left(1 - \frac {B _ {1} B _ {2} \tau_ {1}}{2 n _ {1} n _ {2}}\right) ^ {t + 1} \frac {1}{n _ {1}} \sum_ {i \in S _ {1}} \frac {1}{n _ {2}} \sum_ {j \in S _ {2}} \| v _ {i, j, 0} - h _ {i, j} (\mathbf {w} _ {0}) \| + \frac {2 \tau_ {1} ^ {1 / 2} \sigma}{B _ {3} ^ {1 / 2}} + \frac {4 n _ {1} n _ {2} C _ {h} M \eta}{B _ {1} B _ {2} \tau_ {1} ^ {1 / 2}} \\ \mathbb {E} \left[ \frac {1}{n _ {1}} \sum_ {i \in S _ {1}} \frac {1}{n _ {2}} \sum_ {j \in S _ {2}} \| v _ {i, j, t + 1} - h _ {i, j} (\mathbf {w} _ {t + 1}) \| ^ {2} \right] \\ \leq \left(1 - \frac {B _ {1} B _ {2} \tau_ {1}}{2 n _ {1} n _ {2}}\right) ^ {2 (t + 1)} \frac {1}{n _ {1}} \sum_ {i \in S _ {1}} \frac {1}{n _ {2}} \sum_ {j \in S _ {2}} \| v _ {i, j, 0} - h _ {i, j} (\mathbf {w} _ {0}) \| ^ {2} + \frac {4 \tau_ {1} \sigma^ {2}}{B _ {3}} + \frac {1 6 n _ {1} ^ {2} n _ {2} ^ {2} C _ {h} ^ {2} M ^ {2} \eta^ {2}}{B _ {1} ^ {2} B _ {2} ^ {2} \tau_ {1}} \\ \end{array}
$$
Lemma A.4. Consider MSVR update for  $u$ . Assume  $g_{i}(\cdot)$  is  $C_g$ -Lipshitz for all  $i \in S_1$ . With  $\gamma_{2} = \frac{n_{+} - B_{1}}{B_{1}(1 - \tau_{2})} + (1 - \tau_{2})$  and  $\tau_{2} \leq \frac{1}{2}$ , we have
$$
\begin{array}{l} \mathbb {E} \left[ \frac {1}{n _ {1}} \sum_ {i \in S _ {1}} \| u _ {i, t + 1} - \frac {1}{n _ {2}} \sum_ {j \in S _ {2}} g _ {i} (v _ {i, j, t + 1}) \| \right] \\ \leq \left(1 - \frac {B _ {1} \tau_ {2}}{2 n _ {1}}\right) ^ {t + 1} \frac {1}{n _ {1}} \sum_ {i \in S _ {1}} \| u _ {i, 0} - \frac {1}{n _ {2}} \sum_ {j \in S _ {2}} g _ {i} (v _ {i, j, 0}) \| + \frac {2 \tau_ {2} ^ {1 / 2} \sigma}{B _ {2} ^ {1 / 2}} + \frac {C _ {2} n _ {1} ^ {1 / 2} B _ {2} ^ {1 / 2} \tau_ {1}}{B _ {1} ^ {1 / 2} n _ {2} ^ {1 / 2} \tau_ {2} ^ {1 / 2}} + \frac {C _ {2} n _ {1} ^ {3 / 2} n _ {2} ^ {1 / 2} \eta}{B _ {1} ^ {3 / 2} B _ {2} ^ {1 / 2} \tau_ {2} ^ {1 / 2}} \\ \end{array}
$$
where  $C_2$  is a constant defined in the proof.
Proof of Theorem A.2. 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 {\bar {\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 {18} \\ \leq F (\hat {\mathbf {w}} _ {t}) + \frac {\bar {\rho}}{2} \left(\| \hat {\mathbf {w}} _ {t} - \mathbf {w} _ {t} \| ^ {2}\right) + \bar {\rho} \mathbb {E} _ {t} \left[ \eta \langle \hat {\mathbf {w}} _ {t} - \mathbf {w} _ {t}, G _ {t} \rangle \right] + \frac {\eta^ {2} \bar {\rho} M ^ {2}}{2} \\ = F _ {1 / \bar {\rho}} (\mathbf {w} _ {t}) + \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} \\ \end{array}
$$
$$
\mathbb {E} _ {t} [ G _ {t} ] = \frac {1}{n _ {1}} \sum_ {i = 1} ^ {n _ {1}} \left[ \frac {1}{n _ {2}} \sum_ {j = 1} ^ {n _ {2}} \nabla h _ {i, j} (\mathbf {w} _ {t}) \partial g _ {i} (v _ {i, j, t}) \right] \partial f _ {i} (u _ {i, t}),
$$
Define  $\hat{\mathbf{w}}_t\coloneqq \mathrm{prox}_{F / \bar{\rho}}(\mathbf{w}_t)$ . For a given  $i\in \{1,\ldots ,m\}$ , we have
where (a) follows from the convexity of  $f_{i}$ , (b) uses the assumption that  $f_{i}(\cdot)$  is non-decreasing and  $g_{i}$  is weak convex, (c) is due to  $0 \leq \partial f_{i}(u_{i,t}) \leq C_{f}$ .
The  $L_{h}$ -smoothness assumption of  $h_{i,j}(\mathbf{w})$  (or weakly-convexity of  $h_{i,j}(\mathbf{w})$ , then only the second inequality holds) for all  $i, \mathbf{w}$  implies
$$
h _ {i, j} (\hat {\mathbf {w}} _ {t}) \leq h _ {i, j} \left(\mathbf {w} _ {t}\right) + \nabla h _ {i, j} \left(\mathbf {w} _ {t}\right) ^ {\top} \left(\hat {\mathbf {w}} _ {t} - \mathbf {w} _ {t}\right) + \frac {L _ {h}}{2} \| \hat {\mathbf {w}} _ {t} - \mathbf {w} _ {t} \| ^ {2}, \tag {20}
$$
$$
h _ {i, j} (\hat {\mathbf {w}} _ {t}) \geq h _ {i, j} (\mathbf {w} _ {t}) + \nabla h _ {i, j} (\mathbf {w} _ {t}) ^ {\top} (\hat {\mathbf {w}} _ {t} - \mathbf {w} _ {t}) - \frac {L _ {h}}{2} \| \hat {\mathbf {w}} _ {t} - \mathbf {w} _ {t} \| ^ {2}.
$$
We first assume that  $g_{i}(\cdot)$  is non-increasing. Since  $\partial f_{i}(u_{i,t}) \geq 0$  and  $\partial g_{i}(v_{i,j,t}) \leq 0$ , we bound  $A_{1}$  as following
$$
\begin{array}{l} A _ {1} = \partial f _ {i} \left(u _ {i, t}\right) \partial^ {\top} g _ {i} \left(v _ {i, j, t}\right) ^ {\top} \left(h _ {i, j} \left(\hat {\mathbf {w}} _ {t}\right) - v _ {i, j, t}\right) \\ \stackrel {(a)} {\geq} \left\langle \partial f _ {i} \left(u _ {i, t}\right) ^ {\top} \partial g _ {i} \left(v _ {i, j, t}\right) ^ {\top} \left(h _ {i, j} \left(\mathbf {w} _ {t}\right) - v _ {i, j, t}\right) + \partial f _ {i} \left(u _ {i, t}\right) ^ {\top} \partial g _ {i} \left(v _ {i, j, t}\right) ^ {\top} \nabla h _ {i, j} \left(\mathbf {w} _ {t}\right) ^ {\top} \left(\hat {\mathbf {w}} _ {t} - \mathbf {w} _ {t}\right) \right. \\ + \partial f _ {i} (u _ {i, t}) ^ {\top} \partial g _ {i} (v _ {i, j, t}) ^ {\top} \frac {L _ {h}}{2} \| \hat {\mathbf {w}} _ {t} - \mathbf {w} _ {t} \| ^ {2} \rangle \\ \stackrel {(b)} {\geq} - C _ {f} C _ {g} \left\| h _ {i, j} (\mathbf {w} _ {t}) - v _ {i, j, t} \right\| + \partial f _ {i} \left(u _ {i, t}\right) ^ {\top} \partial g _ {i} \left(v _ {i, j, t}\right) ^ {\top} \nabla h _ {i, j} \left(\mathbf {w} _ {t}\right) ^ {\top} \left(\hat {\mathbf {w}} _ {t} - \mathbf {w} _ {t}\right) \\ - \frac {C _ {f} C _ {g} L _ {h}}{2} \| \hat {\mathbf {w}} _ {t} - \mathbf {w} _ {t} \| ^ {2} \tag {21} \\ \end{array}
$$
where inequality (a) follows from the first inequality in (20), (b) follows from the Lipschitz continuity and monotone assumptions on  $f_{i}, g_{i}, h_{i,j}$ . On the other hand, if we assume  $g_{i}(\cdot)$  is non-decreasing, we may use the second inequality in (20) and obtain the same result as (21). Now plugging the new formulation of  $A_{1}$  back to inequality 19 yields
$$
\begin{array}{l} \frac {1}{n _ {1}} \sum_ {i \in S _ {1}} f _ {i} \left(\frac {1}{n _ {2}} \sum_ {j \in S _ {2}} g _ {i} \left(h _ {i, j} \left(\hat {\mathbf {w}} _ {t}\right)\right)\right) - \frac {1}{n _ {1}} \sum_ {i \in S _ {1}} f _ {i} \left(u _ {i, t}\right) \\ \geq \frac {1}{n _ {1}} \sum_ {i \in S _ {1}} \partial f _ {i} (u _ {i, t}) ^ {\top} \left[ \frac {1}{n _ {2}} \sum_ {j \in S _ {2}} g _ {i} (v _ {i, j, t}) - u _ {i, t} \right] + \frac {1}{n _ {1}} \sum_ {i \in S _ {1}} \frac {1}{n _ {2}} \sum_ {j \in S _ {2}} - C _ {f} C _ {g} \| h _ {i, j} (\mathbf {w} _ {t}) - v _ {i, j, t} \| \\ + \frac {1}{n _ {1}} \sum_ {i \in S _ {1}} \frac {1}{n _ {2}} \sum_ {j \in S _ {2}} \partial f _ {i} (u _ {i, t}) ^ {\top} \partial g _ {i} (v _ {i, j, t}) ^ {\top} \nabla h _ {i, j} (\mathbf {w} _ {t}) ^ {\top} (\hat {\mathbf {w}} _ {t} - \mathbf {w} _ {t}) - \frac {C _ {f} C _ {g} L _ {h}}{2} \| \hat {\mathbf {w}} _ {t} - \mathbf {w} _ {t} \| ^ {2} \\ - \frac {1}{n _ {1}} \sum_ {i \in S _ {1}} \frac {1}{n _ {2}} \sum_ {j \in S _ {2}} \left(2 \rho_ {f} C _ {g} ^ {2} + \rho_ {g} C _ {f}\right) \| h _ {i, j} (\mathbf {w} _ {t}) - v _ {i, j, t} \| ^ {2} \\ - \left(2 \rho_ {f} C _ {g} ^ {2} + \rho_ {g} C _ {f} C _ {h} ^ {2}\right) \| \hat {\mathbf {w}} _ {t} - \mathbf {w} _ {t} \| ^ {2} - \frac {1}{n _ {1}} \sum_ {i \in S _ {1}} \rho_ {f} \left\| \frac {1}{n _ {2}} \sum_ {j \in S _ {2}} g _ {i} (v _ {i, j, t}) - u _ {i, t} \right\| ^ {2} \\ \geq \frac {1}{n _ {1}} \sum_ {i \in S _ {1}} - C _ {f} \left\| \frac {1}{n _ {2}} \sum_ {j \in S _ {2}} g _ {i} \left(v _ {i, j, t}\right) - u _ {i, t} \right\| + \frac {1}{n _ {1}} \sum_ {i \in S _ {1}} \frac {1}{n _ {2}} \sum_ {j \in S _ {2}} - C _ {f} C _ {g} \| h _ {i, j} (\mathbf {w} _ {t}) - v _ {i, j, t} \| \\ + \left\langle \mathbb {E} _ {t} \left[ G _ {t} \right], \hat {\mathbf {w}} _ {t} - \mathbf {w} _ {t} \right\rangle - \frac {1}{n _ {1}} \sum_ {i \in S _ {1}} \frac {1}{n _ {2}} \sum_ {j \in S _ {2}} \left(2 \rho_ {f} C _ {g} ^ {2} + \rho_ {g} C _ {f}\right) \| h _ {i, j} (\mathbf {w} _ {t}) - v _ {i, j, t} \| ^ {2} \\ - \left(2 \rho_ {f} C _ {g} ^ {2} + \rho_ {g} C _ {f} C _ {h} ^ {2} + \frac {C _ {f} C _ {g} L _ {h}}{2}\right) \| \hat {\mathbf {w}} _ {t} - \mathbf {w} _ {t} \| ^ {2} - \frac {1}{n _ {1}} \sum_ {i \in S _ {1}} \rho_ {f} \left\| \frac {1}{n _ {2}} \sum_ {j \in S _ {2}} g _ {i} (v _ {i, j, t}) - u _ {i, t} \right\| ^ {2} \\ \end{array}
$$
It follows
$$
\begin{array}{l} \langle \mathbb {E} _ {t} [ G _ {t} ], \hat {\mathbf {w}} _ {t} - \mathbf {w} _ {t} \rangle \\ \leq \frac {1}{n _ {1}} \sum_ {i \in S _ {1}} f _ {i} \left(\frac {1}{n _ {2}} \sum_ {j \in S _ {2}} g _ {i} \left(h _ {i, j} \left(\hat {\mathbf {w}} _ {t}\right)\right)\right) - \frac {1}{n _ {1}} \sum_ {i \in S _ {1}} f _ {i} \left(u _ {i, t}\right) + \frac {1}{n _ {1}} \sum_ {i \in S _ {1}} C _ {f} \left\| \frac {1}{n _ {2}} \sum_ {j \in S _ {2}} g _ {i} \left(v _ {i, j, t}\right) - u _ {i, t} \right\| \\ + \frac {1}{n _ {1}} \sum_ {i \in S _ {1}} \frac {1}{n _ {2}} \sum_ {j \in S _ {2}} C _ {f} C _ {g} \| h _ {i, j} (\mathbf {w} _ {t}) - v _ {i, j, t} \| + \frac {1}{n _ {1}} \sum_ {i \in S _ {1}} \frac {1}{n _ {2}} \sum_ {j \in S _ {2}} (2 \rho_ {f} C _ {g} ^ {2} + \rho_ {g} C _ {f}) \| h _ {i, j} (\mathbf {w} _ {t}) - v _ {i, j, t} \| ^ {2} \\ + \left(2 \rho_ {f} C _ {g} ^ {2} + \rho_ {g} C _ {f} C _ {h} ^ {2} + \frac {C _ {f} C _ {g} L _ {h}}{2}\right) \| \hat {\mathbf {w}} _ {t} - \mathbf {w} _ {t} \| ^ {2} + \frac {1}{n _ {1}} \sum_ {i \in S _ {1}} \rho_ {f} \left\| \frac {1}{n _ {2}} \sum_ {j \in S _ {2}} g _ {i} \left(v _ {i, j, t}\right) - u _ {i, t} \right\| ^ {2} \tag {22} \\ \end{array}
$$
Combining inequality 22 and 18 yields
where (a) follows from the Lipschitz continuity of  $f_{i}, g_{i}, h_{i,j}$ .
$$
\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 {23} \\ \end{array}
$$
Plugging inequality 23 back into A.2, we obtain
$$
\mathbb {E} _ {t} \left[ F _ {1 / \bar {\rho}} \left(\mathbf {w} _ {t + 1}\right) \right]
$$
$$
\begin{array}{l} \leq F _ {1 / \bar {\rho}} (\mathbf {w} _ {t}) + \frac {\eta^ {2} \bar {\rho} M ^ {2}}{2} + \bar {\rho} \eta \left\{\frac {1}{n _ {1}} \sum_ {i \in S _ {1}} \left[ \left(\frac {\rho_ {F}}{2} - \bar {\rho}\right) \| \mathbf {w} _ {t} - \hat {\mathbf {w}} _ {t} \| ^ {2} + 2 C _ {f} \right. \right\| \frac {1}{n _ {2}} \sum_ {j \in S _ {2}} g _ {i} (v _ {i, j, t}) - u _ {i, t} \| \\ + \frac {1}{n _ {2}} \sum_ {j \in S _ {2}} 2 C _ {f} C _ {g} \| h _ {i, j} (\mathbf {w} _ {t}) - v _ {i, j, t} \| + \frac {1}{n _ {2}} \sum_ {j \in S _ {2}} (2 \rho_ {f} C _ {g} ^ {2} + \rho_ {g} C _ {f}) \| h _ {i, j} (\mathbf {w} _ {t}) - v _ {i, j, t} \| ^ {2} \\ \left. \left. + (2 \rho_ {f} C _ {g} ^ {2} + \rho_ {g} C _ {f} C _ {h} ^ {2} + \frac {C _ {f} C _ {g} L _ {h}}{2}) \| \hat {\mathbf {w}} _ {t} - \mathbf {w} _ {t} \| ^ {2} + \rho_ {f} \left\| \frac {1}{n _ {2}} \sum_ {j \in S _ {2}} g _ {i} (v _ {i, j, t}) - u _ {i, t} \right\| ^ {2} \right] \right\} \\ \end{array}
$$
$$
\begin{array}{l} \stackrel {(a)} {\leq} F _ {1 / \bar {\rho}} (\mathbf {w} _ {t}) + \frac {\eta^ {2} \bar {\rho} M ^ {2}}{2} + \bar {\rho} \eta \left\{\frac {1}{n _ {1}} \sum_ {i \in S _ {1}} \left[ - \frac {\bar {\rho}}{2} \| \mathbf {w} _ {t} - \hat {\mathbf {w}} _ {t} \| ^ {2} + C _ {1} \right. \left\| \frac {1}{n _ {2}} \sum_ {j \in S _ {2}} g _ {i} (v _ {i, j, t}) - u _ {i, t} \right. \right\| \\ + \frac {1}{n _ {2}} \sum_ {j \in S _ {2}} C _ {1} \| h _ {i, j} (\mathbf {w} _ {t}) - v _ {i, j, t} \| + \frac {1}{n _ {2}} \sum_ {j \in S _ {2}} C _ {1} \| h _ {i, j} (\mathbf {w} _ {t}) - v _ {i, j, t} \| ^ {2} \\ \left. + C _ {1} \left\| \frac {1}{n _ {2}} \sum_ {j \in S _ {2}} g _ {i} \left(v _ {i, j, t}\right) - u _ {i, t} \right\| ^ {2} \right] \Bigg \} \\ \end{array}
$$
$$
\begin{array}{l} \stackrel {(b)} {=} 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} + C _ {1} \bar {\rho} \eta \frac {1}{n _ {1}} \sum_ {i \in S _ {1}} \left\| \frac {1}{n _ {2}} \sum_ {j \in S _ {2}} g _ {i} (v _ {i, j, t}) - u _ {i, t} \right\| \\ + C _ {1} \bar {\rho} \eta \frac {1}{n _ {1}} \sum_ {i \in S _ {1}} \frac {1}{n _ {2}} \sum_ {j \in S _ {2}} \| h _ {i, j} (\mathbf {w} _ {t}) - v _ {i, j, t} \| + C _ {1} \bar {\rho} \eta \frac {1}{n _ {1}} \sum_ {i \in S _ {1}} \frac {1}{n _ {2}} \sum_ {j \in S _ {2}} \| h _ {i, j} (\mathbf {w} _ {t}) - v _ {i, j, t} \| ^ {2} \\ + C _ {1} \bar {\rho} \eta \frac {1}{n _ {1}} \sum_ {i \in S _ {1}} \left\| \frac {1}{n _ {2}} \sum_ {j \in S _ {2}} g _ {i} \left(v _ {i, j, t}\right) - u _ {i, t} \right\| ^ {2} \\ \end{array}
$$
where in inequality (a) we use  $\bar{\rho} = \rho_F + 4\rho_fC_g^2 + 2\rho_gC_fC_h^2 + C_fCgL_h$  and  $C_1 = \max \{2C_fC_g, 2C_f, (2\rho_fC_g^2 + \rho_gC_f), \rho_f\}$ , and equality (b) uses Lemma 3.2.
With general error bounds
$$
\frac {1}{n _ {1}} \sum_ {i \in S _ {1}} \frac {1}{n _ {2}} \sum_ {j \in S _ {2}} \mathbb {E} [ \| h _ {i, j} (\mathbf {w} _ {t}) - v _ {i, j, t} \| ] \leq (1 - \mu_ {1}) ^ {t} \frac {1}{n _ {1}} \sum_ {i \in S _ {1}} \frac {1}{n _ {2}} \sum_ {j \in S _ {2}} \| h _ {i, j} (\mathbf {w} _ {0}) - v _ {i, j, 0} \| + R _ {1},
$$
$$
\frac {1}{n _ {1}} \sum_ {i \in S _ {1}} \frac {1}{n _ {2}} \sum_ {j \in S _ {2}} \mathbb {E} [ \| h _ {i, j} (\mathbf {w} _ {t}) - v _ {i, j, t} \| ^ {2} ] \leq (1 - \mu_ {1}) ^ {t} \frac {1}{n _ {1}} \sum_ {i \in S _ {1}} \frac {1}{n _ {2}} \sum_ {j \in S _ {2}} \| h _ {i, j} (\mathbf {w} _ {0}) - v _ {i, j, 0} \| ^ {2} + R _ {2},
$$
$$
\frac {1}{n _ {1}} \sum_ {i \in S _ {1}} \mathbb {E} \left[ \left\| \frac {1}{n _ {2}} \sum_ {j \in S _ {2}} g _ {i} \left(v _ {i, j, t}\right) - u _ {i, t} \right\| \right] \leq (1 - \mu_ {2}) ^ {t} \frac {1}{n _ {+}} \sum_ {i \in S _ {+}} \left\| \frac {1}{n _ {2}} \sum_ {j \in S _ {2}} g _ {i} \left(v _ {i, j, 0}\right) - u _ {i, 0} \right\| + R _ {3},
$$
$$
\frac {1}{n _ {1}} \sum_ {i \in S _ {1}} \mathbb {E} \left[ \left\| \frac {1}{n _ {2}} \sum_ {j \in S _ {2}} g _ {i} (v _ {i, j, t}) - u _ {i, t} \right\| ^ {2} \right] \leq (1 - \mu_ {2}) ^ {t} \frac {1}{n _ {+}} \sum_ {i \in S _ {+}} \left\| \frac {1}{n _ {2}} \sum_ {j \in S _ {2}} g _ {i} (v _ {i, j, 0}) - u _ {i, 0} \right\| ^ {2} + R _ {4},
$$
$$
\begin{array}{l} \mathbb {E} \left[ F _ {1 / \bar {\rho}} (\mathbf {w} _ {t + 1}) \right] \\ \leq \mathbb {E} \left[ F _ {1 / \bar {\rho}} (\mathbf {w} _ {t}) \right] + \frac {\eta^ {2} \bar {\rho} M ^ {2}}{2} - \frac {\eta}{2} \mathbb {E} \left[ \| \nabla F _ {1 / \bar {\rho}} (\mathbf {w} _ {t}) \| ^ {2} \right] + C _ {1} \bar {\rho} \eta (1 - \mu_ {\min }) ^ {t} \left[ \frac {1}{n _ {1}} \sum_ {i \in S _ {1}} \left\| \frac {1}{n _ {2}} \sum_ {j \in S _ {2}} g _ {i} (v _ {i, j, 0}) - u _ {i, 0} \right\| \right. \\ + \frac {1}{n _ {1}} \sum_ {i \in S _ {1}} \left\| \frac {1}{n _ {2}} \sum_ {j \in S _ {2}} g _ {i} \left(v _ {i, j, 0}\right) - u _ {i, 0} \right\| ^ {2} + \frac {1}{n _ {1}} \sum_ {i \in S _ {1}} \frac {1}{n _ {2}} \sum_ {j \in S _ {2}} \left\| h _ {i, j} \left(\mathbf {w} _ {0}\right) - v _ {i, j, 0} \right\| \\ \left. + \frac {1}{n _ {1}} \sum_ {i \in S _ {1}} \frac {1}{n _ {2}} \sum_ {j \in S _ {2}} \| h _ {i, j} (\mathbf {w} _ {0}) - v _ {i, j, 0} \| ^ {2} \right] + C _ {1} \bar {\rho} \eta (R _ {1} + R _ {2} + R _ {3} + R _ {4}), \\ \end{array}
$$
where  $\mu_{min} = \min \{\mu_1,\mu_2\}$
$$
\begin{array}{l} \mathbb {E} \left[ F _ {1 / \hat {\rho}} (\mathbf {w} _ {T}) \right] \\ \leq \mathbb {E} \left[ F _ {1 / \bar {\rho}} (\mathbf {w} _ {0}) \right] + \frac {\eta^ {2} \bar {\rho} M ^ {2} T}{2} - \frac {\eta}{2} \sum_ {t = 0} ^ {T - 1} \mathbb {E} \left[ \| \nabla F _ {1 / \bar {\rho}} (\mathbf {w} _ {t}) \| ^ {2} \right] + C _ {1} \bar {\rho} \eta \sum_ {t = 0} ^ {T - 1} \left(1 - \mu_ {\min}\right) ^ {t} \Delta_ {0} \\ + T C _ {1} \bar {\rho} \eta \left(R _ {1} + R _ {2} + R _ {3} + R _ {4}\right) \\ \leq \mathbb {E} \left[ F _ {1 / \bar {\rho}} (\mathbf {w} _ {0}) \right] + \frac {\eta^ {2} \bar {\rho} M ^ {2} T}{2} - \frac {\eta}{2} \sum_ {t = 0} ^ {T - 1} \mathbb {E} \left[ \| \nabla F _ {1 / \bar {\rho}} (\mathbf {w} _ {t}) \| ^ {2} \right] + \frac {C _ {1} \bar {\rho} \eta \Delta_ {0}}{\mu_ {m i n}} + T C _ {1} \bar {\rho} \eta (R _ {1} + R _ {2} + R _ {3} + R _ {4}) \\ \end{array}
$$
where we use  $\sum_{t=0}^{T-1}(1 - \mu_{min})^t \leq \frac{1}{\mu_{min}}$  and define constant  $\Delta_0$  such that
$$
\begin{array}{l} \left[ \frac {1}{n _ {1}} \sum_ {i \in S _ {1}} \left\| \frac {1}{n _ {2}} \sum_ {j \in S _ {2}} g _ {i} \left(v _ {i, j, 0}\right) - u _ {i, 0} \right\| + \frac {1}{n _ {1}} \sum_ {i \in S _ {1}} \left\| \frac {1}{n _ {2}} \sum_ {j \in S _ {2}} g _ {i} \left(v _ {i, j, 0}\right) - u _ {i, 0} \right\| ^ {2} \right. \\ \left. + \frac {1}{n _ {1}} \sum_ {i \in S _ {1}} \frac {1}{n _ {2}} \sum_ {j \in S _ {2}} \| h _ {i, j} (\mathbf {w} _ {0}) - v _ {i, j, 0} \| + \frac {1}{n _ {1}} \sum_ {i \in S _ {1}} \frac {1}{n _ {2}} \sum_ {j \in S _ {2}} \| h _ {i, j} (\mathbf {w} _ {0}) - v _ {i, j, 0} \| ^ {2} \right] \leq \Delta_ {0}. \\ \end{array}
$$
Then it follows
$$
\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}) - \mathbb {E} \left[ F _ {1 / \bar {\rho}} (\mathbf {w} _ {T}) \right] + \frac {\eta^ {2} \bar {\rho} M ^ {2} T}{2} + \frac {C _ {1} \bar {\rho} \eta \Delta_ {0}}{\mu_ {\min }} + T C _ {1} \bar {\rho} \eta \left(R _ {1} + R _ {2} + R _ {3} + R _ {4}\right) \right] \\ \leq \frac {2 \Delta}{\eta T} + \eta \bar {\rho} M ^ {2} + \frac {2 C _ {1} \bar {\rho} \Delta_ {0}}{\mu_ {\text {m i n}} T} + 2 C _ {1} \bar {\rho} (R _ {1} + R _ {2} + R _ {3}) \\ = \mathcal {O} \left(\frac {1}{T} \left(\frac {1}{\eta} + \frac {1}{\mu_ {\text {m i n}}}\right) + \eta + R _ {1} + R _ {2} + R _ {3} + R _ {4}\right) \\ \end{array}
$$
where we define constant  $\Delta$  such that  $F_{1 / \bar{\rho}}(\mathbf{w}_0,\mathbf{s}_0,s_0^{\prime}) - \mathbb{E}[F_{1 / \bar{\rho}}(\mathbf{w}_T,\mathbf{s}_T,s_T^{\prime})]\leq \Delta$
With MSVR updates for  $v_{i,j,t}$  and  $u_{i,t}$ , following from Lemma A.3 and Lemma A.4, we have
$$
\mu_ {1} = \frac {B _ {1} B _ {2} \tau_ {1}}{2 n _ {1} n _ {2}}, \quad \mu_ {2} = \frac {B _ {1} \tau_ {2}}{2 n _ {1}}, \quad R _ {1} = \frac {2 \tau_ {1} ^ {1 / 2} \sigma}{B _ {3} ^ {1 / 2}} + \frac {4 n _ {1} n _ {2} \sqrt {C _ {h}} M \eta}{B _ {1} B _ {2} \tau_ {1} ^ {1 / 2}}
$$
$$
R _ {2} = \frac {4 \tau_ {1} \sigma^ {2}}{B _ {3}} + \frac {1 6 n _ {1} ^ {2} n _ {2} ^ {2} C _ {h} M ^ {2} \eta^ {2}}{B _ {1} ^ {2} B _ {2} ^ {2} \tau_ {1}}, R _ {3} = \frac {2 \tau_ {2} ^ {1 / 2} \sigma}{B _ {2} ^ {1 / 2}} + \frac {C _ {2} n _ {1} ^ {1 / 2} B _ {2} ^ {1 / 2} \tau_ {1}}{B _ {1} ^ {1 / 2} n _ {2} ^ {1 / 2} \tau_ {2} ^ {1 / 2}} + \frac {C _ {2} n _ {1} ^ {3 / 2} n _ {2} ^ {1 / 2} \eta}{B _ {1} ^ {3 / 2} B _ {2} ^ {1 / 2} \tau_ {2} ^ {1 / 2}},
$$
$$
R _ {4} = \frac {4 \tau_ {2} \sigma^ {2}}{B _ {2}} + \frac {C _ {2} ^ {2} n _ {1} B _ {2} \tau_ {1} ^ {2}}{B _ {1} n _ {2} \tau_ {2}} + \frac {C _ {2} ^ {2} n _ {1} ^ {3} n _ {2} \eta^ {2}}{B _ {1} ^ {3} B _ {2} \tau_ {2}}.
$$
$$
\begin{array}{l} \frac {1}{T} \sum_ {t = 0} ^ {T - 1} \mathbb {E} [ \| \nabla F _ {1 / \bar {\rho}} (\mathbf {w} _ {t}) \| ^ {2} ] \\ \leq \mathcal {O} \left(\frac {1}{T} \left(\frac {1}{\eta} + \frac {1}{\mu_ {\text {m i n}}}\right) + \eta + \frac {\tau_ {1} ^ {1 / 2}}{B _ {3} ^ {1 / 2}} + \frac {\tau_ {1}}{B _ {3}} + \frac {\tau_ {2} ^ {1 / 2}}{B _ {2} ^ {1 / 2}} + \frac {\tau_ {2}}{B _ {2}} \right. \\ + \frac {n _ {1} n _ {2} \eta}{B _ {1} B _ {2} \tau_ {1} ^ {1 / 2}} + \frac {n _ {1} ^ {2} n _ {2} ^ {2} \eta^ {2}}{B _ {1} ^ {2} B _ {2} ^ {2} \tau_ {1}} + \frac {n _ {1} ^ {1 / 2} B _ {2} ^ {1 / 2} \tau_ {1}}{B _ {1} ^ {1 / 2} n _ {2} ^ {1 / 2} \tau_ {2} ^ {1 / 2}} + \frac {n _ {1} ^ {3 / 2} n _ {2} ^ {1 / 2} \eta}{B _ {1} ^ {3 / 2} B _ {2} ^ {1 / 2} \tau_ {2} ^ {1 / 2}} + \frac {n _ {1} B _ {2} \tau_ {1} ^ {2}}{B _ {1} n _ {2} \tau_ {2}} + \frac {n _ {1} ^ {3} n _ {2} \eta^ {2}}{B _ {1} ^ {3} B _ {2} \tau_ {2}}) \\ \leq \mathcal {O} \bigg (\frac {1}{T} (\frac {1}{\eta} + \frac {1}{\mu_ {m i n}}) + \frac {\tau_ {1} ^ {1 / 2}}{B _ {3} ^ {1 / 2}} + \frac {\tau_ {2} ^ {1 / 2}}{B _ {2} ^ {1 / 2}} + \frac {n _ {1} n _ {2} \eta}{B _ {1} B _ {2} \tau_ {1} ^ {1 / 2}} + \frac {n _ {1} ^ {1 / 2} B _ {2} ^ {1 / 2} \tau_ {1}}{B _ {1} ^ {1 / 2} n _ {2} ^ {1 / 2} \tau_ {2} ^ {1 / 2}} + \frac {n _ {1} ^ {3 / 2} n _ {2} ^ {1 / 2} \eta}{B _ {1} ^ {3 / 2} B _ {2} ^ {1 / 2} \tau_ {2} ^ {1 / 2}} \bigg). \\ \end{array}
$$
Algorithm 3 Stochastic Optimization algorithm for Non-smooth FCCO with coordinate moving average
1: Initialization:  $\mathbf{w}_0, \{u_{i,0} : i \in \mathcal{S}\}$ .  
2: for  $t = 0, \dots, T - 1$  do  
3: Draw sample batches  $\mathcal{B}_1^t \sim \mathcal{S}$ , and  $\mathcal{B}_{2,i}^t \sim \mathcal{D}_i$  for each  $i \in \mathcal{B}_1^t$ .  
4:  $u_{i,t + 1} = \begin{cases} (1 - \tau) u_{i,t} + \tau g_i(\mathbf{w}_t; \mathcal{B}_{2,i}^t), & i \in \mathcal{B}_1^t \\ u_{i,t}, & i \notin \mathcal{B}_1^t \end{cases}$   
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  
8: return  $\mathbf{w}_{\bar{t}}$  with uniformly sampled  $\bar{t} \in \{0, T - 1\}$ .
$$
\begin{array}{l} \tau_ {1} = \mathcal {O} \left(\min \{B _ {3}, \frac {B _ {1} ^ {1 / 2} n _ {2} ^ {1 / 2}}{n _ {1} ^ {1 / 2}} \} \epsilon^ {4}\right), \quad \tau_ {2} = \mathcal {O} (B _ {2} \epsilon^ {4}), \\ \eta = \mathcal {O} \left(\min  \left\{\frac {B _ {1} B _ {2}}{n _ {1} n _ {2}} \tau_ {1} ^ {1 / 2} \epsilon^ {2}, \frac {B _ {1} ^ {3 / 2} B _ {2} ^ {1 / 2}}{n _ {1} ^ {3 / 2} n _ {2} ^ {1 / 2}} \tau_ {2} ^ {1 / 2} \right\}\right) \\ = \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), \\ \end{array}
$$
then with
$$
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 _ {1} n _ {2}}{B _ {1} B _ {2}} \epsilon^ {- 6}\right),
$$
$$
\frac {1}{T} \sum_ {t = 0} ^ {T - 1} \mathbb {E} [ \| \nabla F _ {1 / \bar {\rho}} (\mathbf {w} _ {t}) \| ^ {2} ] \leq \epsilon^ {2}
$$
B Solving Non-smooth FCCO and TCCO with Coordinate Moving Average
In this section we consider solving non-smooth weakly-convex FCCO and TCCO without variance reduction method. To be specific, we use coordinate moving average updates for function values estimations instead of MSVR. This allows us to weaken the assumption on the Lipschitz continuity, i.e. the Lipschitz continuity of the stochastic function value estimation is not required, and can be replaced by the Lipschitz continuity of the function value. Moreover, compared with MSVR, coordinate moving average update does not need the stochastic evaluation from the previous iteration, and thus has a simpler implementation. However, as a result of not using variance reduction technique, the algorithms suffer from worse convergence rates in terms of  $\epsilon$ .