SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
\begin{array}{l} \leq - 2 \eta_ {\ell} K \left(\mathbb {E} \left[ F \left(w _ {g} ^ {(r)}\right) \right] - F \left(w _ {g} ^ {*}\right)\right) - \frac {\eta_ {\ell} \mu K}{2 M} \mathbb {E} | | w _ {g} ^ {(r)} - w _ {g} ^ {*} | | ^ {2} \\ + \frac {2 \eta_ {\ell} \beta}{n M} \sum_ {m \in [ M ]} n _ {m} \left(1 6 K ^ {4} \eta_ {\ell} ^ {2} \mathbb {E} | | \nabla f _ {m} \left(w _ {g} ^ {(r)}\right) | | ^ {2} \mathbb {E} | | 1 - \psi_ {g, m} ^ {(r, K)} | | ^ {2} + 8 K ^ {2} \eta_ {\ell} ^ {2} \sigma_ {\ell} ^ {2} \mathbb {E} | | \psi_ {g, m} ^ {(r, K)} | | ^ {2} \right. \\ \left. + 1 2 K ^ {3} \eta_ {\ell} ^ {2} \mathbb {E} | | \nabla f _ {m} \left(w _ {g} ^ {(r)}\right) | | ^ {2} \mathbb {E} | | 1 - \psi_ {g, m} ^ {(r, K)} | | ^ {2} + 6 K ^ {2} \eta_ {\ell} ^ {2} \sigma_ {\ell} ^ {2} \mathbb {E} | | \psi_ {g, m} ^ {(r, K)} | | ^ {2}\right), \tag {47} \\ \end{array}
$$
Next, using Assumption E.4,
$$
\begin{array}{l} T _ {1} \leq - 2 \eta_ {\ell} K \left(\mathbb {E} \left[ F \left(w _ {g} ^ {(r)}\right) \right] - F \left(w _ {g} ^ {*}\right)\right) - \frac {\eta_ {\ell} \mu K}{2 M} \mathbb {E} | | w _ {g} ^ {(r)} - w _ {g} ^ {*} | | ^ {2} \\ + 3 2 \eta_ {\ell} ^ {3} K ^ {4} \beta \mathbb {E} | | 1 - \psi_ {g} ^ {(r)} | | ^ {2} \left(G ^ {2} + 2 \beta B ^ {2} \left(\mathbb {E} \left[ F \left(w _ {g} ^ {(r)}\right) \right] - F \left(w _ {g} ^ {*}\right)\right)\right) + 1 6 \eta_ {\ell} ^ {3} K ^ {2} \beta \sigma_ {\ell} ^ {2} \mathbb {E} | | \psi_ {g} ^ {(r)} | | ^ {2} \\ + 2 4 \eta_ {\ell} ^ {3} K ^ {3} \beta \mathbb {E} | | 1 - \psi_ {g} ^ {(r)} | | ^ {2} \left(G ^ {2} + 2 \beta B ^ {2} \left(\mathbb {E} \left[ F \left(w _ {g} ^ {(r)}\right) \right] - F \left(w _ {g} ^ {*}\right)\right)\right) + 1 2 \eta_ {\ell} ^ {3} K ^ {2} \beta \sigma_ {\ell} ^ {2} \mathbb {E} | | \psi_ {g} ^ {(r)} | | ^ {2} (48) \\ \leq - 2 \eta_ {\ell} K \left(\mathbb {E} \left[ F \left(w _ {g} ^ {(r)}\right) \right] - F \left(w _ {g} ^ {*}\right)\right) - \frac {\eta_ {\ell} \mu K}{2 M} \mathbb {E} | | w _ {g} ^ {(r)} - w _ {g} ^ {*} | | ^ {2} \\ + 1 6 \eta_ {\ell} ^ {3} K ^ {3} \beta^ {2} B ^ {2} (4 K + 3) \mathbb {E} \| 1 - \psi_ {g} ^ {(r)} \| ^ {2} \left(\mathbb {E} \left[ F \left(w _ {g} ^ {(r)}\right) \right] - F \left(w _ {g} ^ {*}\right)\right) \\ + 8 \eta_ {\ell} ^ {3} K ^ {3} \beta (4 K + 3) \mathbb {E} | | 1 - \psi_ {g} ^ {(r)} | | ^ {2} G ^ {2} + 2 8 \eta_ {\ell} ^ {3} K ^ {2} \beta \sigma_ {\ell} ^ {2} \mathbb {E} | | \psi_ {g} ^ {(r)} | | ^ {2} (49) \\ \end{array}
$$
Bounding  $T_{2}$
$$
\begin{array}{l} T _ {2} = \eta_ {\ell} ^ {2} \mathbb {E} \left[ \left\| \frac {1}{n M} \sum_ {m \in [ M ]} n _ {m} \sum_ {k = 1} ^ {K} \nabla_ {w _ {g, m} ^ {(r, k - 1)}} f _ {m} \left(w _ {p, m} ^ {(r, k - 1)}\right) \right\| ^ {2} \right] (50) \\ = \eta_ {\ell} ^ {2} \mathbb {E} \left[ \left| \left| \frac {1}{n M} \sum_ {m \in [ M ]} n _ {m} \sum_ {k = 1} ^ {K} \left(\nabla_ {w _ {g, m} ^ {(r, k - 1)}} f _ {m} \left(w _ {p, m} ^ {(r, k - 1)}\right) - \nabla f _ {m} \left(w _ {g} ^ {(r)}\right) + \nabla f _ {m} \left(w _ {g} ^ {(r)}\right)\right) \right| \right| ^ {2} \right] (51) \\ \leq 2 \eta_ {\ell} ^ {2} \mathbb {E} \left[ \left| \left| \frac {1}{n M} \sum_ {m \in [ M ]} n _ {m} \sum_ {k = 1} ^ {K} \left(\nabla_ {w _ {g, m} ^ {(r, k - 1)}} f _ {m} \left(w _ {p, m} ^ {(r, k - 1)}\right) - \nabla f _ {m} \left(w _ {g} ^ {(r)}\right)\right) \right| \right| ^ {2} \right] \\ + 2 \eta_ {\ell} ^ {2} \mathbb {E} \left[ \left\| \frac {1}{n M} \sum_ {m \in [ M ]} n _ {m} \sum_ {k = 1} ^ {K} \nabla f _ {m} \left(w _ {g} ^ {(r)}\right)\right) \right\| ^ {2} (52) \\ \leq 2 \eta_ {\ell} ^ {2} \beta^ {2} K \cdot \frac {1}{n M} \sum_ {m \in [ M ]} n _ {m} \sum_ {k = 1} ^ {K} \underbrace {\mathbb {E} \left[ \left| \left| w _ {p , m} ^ {(r , k - 1)} - w _ {g} ^ {(r)} \right| \right| ^ {2} \right]} _ {\text {L e m m a E . 8}} \\ + 2 \eta_ {\ell} ^ {2} K \cdot \frac {1}{n M} \sum_ {m \in [ M ]} n _ {m} \sum_ {k = 1} ^ {K} \mathbb {E} \left[ \left| \left| \nabla f _ {m} \left(w _ {g} ^ {(r)}\right)\right) \right| ^ {2} \right] (53) \\ \leq 1 6 \eta_ {\ell} ^ {4} K ^ {4} \beta^ {3} B ^ {2} (4 K + 3) \mathbb {E} \| 1 - \psi_ {g} ^ {(r)} \| ^ {2} \left(\mathbb {E} [ F (w _ {g} ^ {(r)}) ] - F (w _ {g} ^ {*})\right) \\ + 8 \eta_ {\ell} ^ {4} K ^ {4} \beta^ {3} (4 K + 3) \mathbb {E} | | 1 - \psi_ {g} ^ {(r)} | | ^ {2} G ^ {2} \\ + 5 6 \eta_ {\ell} ^ {5} K ^ {3} \beta^ {3} \sigma_ {\ell} ^ {2} \mathbb {E} | | \psi_ {g} ^ {(r)} | | ^ {2} + 2 \eta_ {\ell} ^ {2} K \left(G ^ {2} + 2 \beta B ^ {2} \mathbb {E} [ F \left(w _ {g} ^ {(r)}\right) ] - F \left(w _ {g} ^ {*}\right)\right) (54) \\ \end{array}
$$
Plugging in  $T_{1}$  and  $T_{2}$  bounds,
$$
\begin{array}{l} \mathbb {E} \left[ | | w _ {g} ^ {(r + 1)} - w _ {g} ^ {*} | | ^ {2} \right] \leq \mathbb {E} \left[ | | w _ {g} ^ {(r)} - w _ {g} ^ {*} | | ^ {2} \right] - 2 \eta_ {\ell} K \left(\mathbb {E} [ F (w _ {g} ^ {(r)}) ] - F (w _ {g} ^ {*})\right) - \frac {\eta_ {\ell} \mu K}{2 M} \mathbb {E} | | w _ {g} ^ {(r)} - w _ {g} ^ {*} | | ^ {2} \\ + 1 6 \eta_ {\ell} ^ {3} K ^ {3} \beta^ {2} B ^ {2} (4 K + 3) (\eta_ {\ell} \beta + 1) \mathbb {E} | | 1 - \psi_ {g} ^ {(r)} | | ^ {2} \left(\mathbb {E} [ F (w _ {g} ^ {(r)}) ] - F (w _ {g} ^ {*})\right) \\ + 8 \eta_ {\ell} ^ {3} K ^ {3} \beta (4 K + 3) (\eta_ {\ell} \beta^ {2} + 1) \mathbb {E} | | 1 - \psi_ {g} ^ {(r)} | | ^ {2} G ^ {2} + 2 8 \eta_ {\ell} ^ {3} K ^ {2} \beta \sigma_ {\ell} ^ {2} \mathbb {E} | | \psi_ {g} ^ {(r)} | | ^ {2} \\ + 5 6 \eta_ {\ell} ^ {5} K ^ {3} \beta^ {3} \sigma_ {\ell} ^ {2} \mathbb {E} | | \psi_ {g} ^ {(r)} | | ^ {2} + 2 \eta_ {\ell} ^ {2} K \left(G ^ {2} + 2 \beta B ^ {2} \mathbb {E} [ F (w _ {g} ^ {(r)}) ] - F (w _ {g} ^ {*})\right) \tag {55} \\ \end{array}
$$
Rearranging the terms, and replacing  $\mathbb{E}||\psi_g^{(r)}||^2$  and  $\mathbb{E}||1 - \psi_g^{(r)}||^2$  with  $\mathbf{q}_0^2$  (probability of picking global route averaged over the instances sampled from the global distribution) and  $\mathbf{q}_1^2$  respectively,
$$
\begin{array}{l} \mathbb {E} \left[ | | w _ {g} ^ {(r + 1)} - w _ {g} ^ {*} | | ^ {2} \right] \leq \left(1 - \frac {\eta_ {\ell} \mu K}{2 M}\right) \mathbb {E} \left[ | | w _ {g} ^ {(r)} - w _ {g} ^ {*} | | ^ {2} \right] \\ - \left(2 \eta_ {\ell} K - 8 0 \eta_ {\ell} ^ {3} K ^ {4} \beta^ {2} B ^ {2} (\eta_ {\ell} \beta + 1) \mathbf {q} _ {1} ^ {2} - 4 \eta_ {\ell} ^ {2} K \beta B ^ {2}\right) \left(\mathbb {E} \left[ F \left(w _ {g} ^ {(r)}\right) \right] - F \left(w _ {g} ^ {*}\right)\right) \\ + 4 0 \eta_ {\ell} ^ {3} K ^ {3} \beta \left(\eta_ {\ell} \beta^ {2} + 1\right) \mathbf {q} _ {1} ^ {2} G ^ {2} + 2 \eta_ {\ell} ^ {2} K G ^ {2} + 2 8 \eta_ {\ell} ^ {3} K ^ {2} \beta \left(2 \eta_ {\ell} ^ {2} \beta^ {2} K + 1\right) \mathbf {q} _ {0} ^ {2} \sigma_ {\ell} ^ {2} \tag {56} \\ \end{array}
$$
Assuming  $\frac{\eta_{\ell}K}{2}\geq 80\eta_{\ell}^{3}K^{4}\beta^{2}B^{2}(\eta_{\ell}\beta +1)\Rightarrow \eta_{\ell}\leq \frac{1}{4\sqrt{10}\beta BK^{2}}$  and  $\frac{\eta_{\ell}K}{2}\geq 4\eta_{\ell}^{2}K\beta B^{2}\Rightarrow \eta_{\ell}\leq \frac{1}{8B^{2}\beta},$  we get
$$
\begin{array}{l} \mathbb {E} \left[ | | w _ {g} ^ {(r + 1)} - w _ {g} ^ {*} | | ^ {2} \right] \leq \left(1 - \frac {\eta_ {\ell} \mu K}{2 M}\right) \mathbb {E} \left[ | | w _ {g} ^ {(r)} - w _ {g} ^ {*} | | ^ {2} \right] - \eta_ {\ell} K (1 - \mathbf {q} _ {1}) ^ {2} \left(\mathbb {E} \left[ F (w _ {g} ^ {(r)}) \right] - F (w _ {g} ^ {*})\right) \\ + 4 0 \eta_ {\ell} ^ {3} K ^ {3} \beta \left(\eta_ {\ell} \beta^ {2} + 1\right) \mathbf {q} _ {1} ^ {2} G ^ {2} + 2 \eta_ {\ell} ^ {2} K G ^ {2} + 2 8 \eta_ {\ell} ^ {3} K ^ {2} \beta \left(2 \eta_ {\ell} ^ {2} \beta^ {2} K + 1\right) \mathbf {q} _ {0} ^ {2} \sigma_ {\ell} ^ {2} \tag {57} \\ \end{array}
$$
Moving  $\mathbb{E}\left[F(w_g^{(r)})\right] - F(w_g^*)$  to the left-hand side, and rest of the terms on right-hand side,
$$
\begin{array}{l} \eta_ {\ell} K \mathbf {q} _ {0} ^ {2} \left(\mathbb {E} \left[ F (w _ {g} ^ {(r)}) \right] - F (w _ {g} ^ {*})\right) \leq \left(1 - \frac {\eta_ {\ell} \mu K}{2 M}\right) \mathbb {E} \left[ | | w _ {g} ^ {(r)} - w _ {g} ^ {*} | | ^ {2} \right] - \mathbb {E} \left[ | | w _ {g} ^ {(r + 1)} - w _ {g} ^ {*} | | ^ {2} \right] \\ + 4 0 \eta_ {\ell} ^ {3} K ^ {3} \beta \left(\eta_ {\ell} \beta^ {2} + 1\right) \mathbf {q} _ {1} ^ {2} G ^ {2} + 2 \eta_ {\ell} ^ {2} K G ^ {2} + 2 8 \eta_ {\ell} ^ {3} K ^ {2} \beta \left(2 \eta_ {\ell} ^ {2} \beta^ {2} K + 1\right) \mathbf {q} _ {0} ^ {2} \sigma_ {\ell} ^ {2} \tag {58} \\ \end{array}
$$
$$
\begin{array}{l} \therefore \mathbb {E} \left[ F \left(w _ {g} ^ {(r)}\right) \right] - F \left(w _ {g} ^ {*}\right) \leq \frac {1}{\eta_ {\ell} K \mathbf {q} _ {0} ^ {2}} \left(1 - \frac {\eta_ {\ell} \mu K}{2 M}\right) \mathbb {E} \left[ \left| \left| w _ {g} ^ {(r)} - w _ {g} ^ {*} \right| \right| ^ {2} \right] - \frac {1}{\eta_ {\ell} K \mathbf {q} _ {0} ^ {2}} \mathbb {E} \left[ \left| \left| w _ {g} ^ {(r + 1)} - w _ {g} ^ {*} \right| \right| ^ {2} \right] \\ + 4 0 \eta_ {\ell} ^ {2} K ^ {2} \beta \left(\eta_ {\ell} \beta^ {2} + 1\right) \frac {\mathbf {q} _ {1} ^ {2}}{\mathbf {q} _ {0} ^ {2}} G ^ {2} + \frac {2 \eta_ {\ell} G ^ {2}}{\mathbf {q} _ {0} ^ {2}} + 2 8 \eta_ {\ell} ^ {2} K \beta \left(2 \eta_ {\ell} ^ {2} \beta^ {2} K + 1\right) \sigma_ {\ell} ^ {2} \tag {59} \\ \end{array}
$$
Unrolling the recursion over  $R$  rounds and then using the linear convergence lemma (Lemma 1) for strong convex case from Scaffold [18],
$$
\begin{array}{l} \mathbb {E} \left[ F (w _ {g} ^ {(R)}) \right] - F (w _ {g} ^ {*}) \leq \frac {\mu}{\mathbf {q} _ {0} ^ {2} K} \mathbb {E} | | w _ {g} ^ {(0)} - w _ {g} ^ {*} | | ^ {2} \exp \left(- \frac {\eta_ {\ell} \mu K R}{2 M}\right) + \frac {2 G ^ {2}}{\mathbf {q} _ {0} ^ {2} \mu R} \\ + \frac {4 0 K ^ {2} \beta}{\mu^ {2} R ^ {2}} \left(\frac {\beta^ {2}}{\mu R} + 1\right) \frac {\mathbf {q} _ {1} ^ {2}}{\mathbf {q} _ {0} ^ {2}} G ^ {2} + \frac {2 8 K \beta}{\mu^ {2} R ^ {2}} \left(\frac {2 \beta^ {2} K}{\mu^ {2} R ^ {2}} + 1\right) \sigma_ {\ell} ^ {2} \tag {60} \\ \end{array}
$$
Unrolling the recursion over  $R$  rounds and then using the sublinear convergence lemma (Lemma 2) for general convex case from Scaffold [18],
$$
\begin{array}{l} \mathbb {E} \left[ F (w _ {g} ^ {(R)}) \right] - F (w _ {g} ^ {*}) \leq \frac {1}{\eta_ {\ell} K \mathbf {q} _ {0} ^ {2} (R + 1)} \mathbb {E} | | w _ {g} ^ {(0)} - w _ {g} ^ {*} | | ^ {2} + \eta_ {\ell} \left(\frac {2 G ^ {2}}{\mathbf {q} _ {0} ^ {2}}\right) ^ {1 / 2} + \\ + \eta_ {\ell} ^ {2} \left(4 0 K ^ {2} \beta \frac {\mathbf {q} _ {1} ^ {2}}{\mathbf {q} _ {0} ^ {2}} G ^ {2}\right) ^ {1 / 2} + \eta_ {\ell} ^ {3} \left(4 0 K ^ {2} \beta^ {3} \frac {\mathbf {q} _ {1} ^ {2}}{\mathbf {q} _ {0} ^ {2}} G ^ {2}\right) ^ {1 / 3} + \eta_ {\ell} ^ {2} \left(2 8 K \beta \sigma_ {\ell} ^ {2}\right) ^ {1 / 3} + \eta_ {\ell} ^ {4} \left(5 6 K \beta^ {3} \sigma_ {\ell} ^ {2}\right) ^ {1 / 5} \tag {61} \\ \end{array}
$$
E.4 Convergence Proof for the Global Model: Non-convex Case
We start with a non-convex version of Lemmas E.7 and E.8,
Lemma E.10 (Local version of the global model progress). If  $m^{th}$  client's objective function  $f_{m}$  satisfies Assumptions E.2, E.3, in Algorithm 2, the following is satisfied:
$$
\mathbb {E} | | w _ {g, m} ^ {(r, k)} - w _ {g, m} ^ {(r, 0)} | | ^ {2} \leq 4 k ^ {2} \eta_ {\ell} ^ {2} \mathbb {E} | | \nabla f _ {m} (w _ {g} ^ {(r)}) | | ^ {2} + 2 k \eta_ {\ell} ^ {2} \sigma_ {\ell} ^ {2} + 4 k ^ {2} \eta_ {\ell} ^ {2} \beta^ {2} \sum_ {i = 1} ^ {k} \mathbb {E} | | w _ {p, m} ^ {(r, i - 1)} - w _ {g} ^ {(r)} | | ^ {2}
$$
$$
\mathbb {E} \| w _ {g, m} ^ {(r, k)} - w _ {g, m} ^ {(r, 0)} \| ^ {2} = \mathbb {E} \| w _ {g, m} ^ {(r, k - 1)} - \eta_ {\ell} \nabla_ {w _ {g, m} ^ {(r, k - 1)}} f _ {m} \left(w _ {p, m} ^ {(r, k - 1)}\right) - w _ {g, m} ^ {(r, 0)} \| ^ {2} \tag {62}
$$
$$
\begin{array}{l} \leq \left(1 + \frac {1}{k - 1}\right) \mathbb {E} | | w _ {g, m} ^ {(r, k - 1)} - w _ {g, m} ^ {(r, 0)} | | ^ {2} + k \eta_ {\ell} ^ {2} \mathbb {E} | | \nabla_ {w _ {g, m} ^ {(r, k - 1)}} f _ {m} \left(w _ {p, m} ^ {(r, k - 1)}\right) | | ^ {2} + \eta_ {\ell} ^ {2} \sigma_ {\ell} ^ {2} (63) \\ \leq \left(1 + \frac {1}{k - 1}\right) \mathbb {E} | | w _ {g, m} ^ {(r, k - 1)} - w _ {g, m} ^ {(r, 0)} | | ^ {2} + \eta_ {\ell} ^ {2} \sigma_ {\ell} ^ {2} \\ + k \eta_ {\ell} ^ {2} \mathbb {E} \| \nabla f _ {m} \left(w _ {p, m} ^ {(r, k - 1)}\right) - \nabla f _ {m} \left(w _ {g, m} ^ {(r, 0)}\right) + \nabla f _ {m} \left(w _ {g, m} ^ {(r, 0)}\right) \| ^ {2} (64) \\ \end{array}
$$
$$
\begin{array}{l} \leq \left(1 + \frac {1}{k - 1}\right) \mathbb {E} | | w _ {g, m} ^ {(r, k - 1)} - w _ {g, m} ^ {(r, 0)} | | ^ {2} + \eta_ {\ell} ^ {2} \sigma_ {\ell} ^ {2} + 2 k \eta_ {\ell} ^ {2} \mathbb {E} | | \nabla f _ {m} (w _ {g} ^ {(r)})) | | ^ {2} \\ + 2 k \eta_ {\ell} ^ {2} \mathbb {E} | | \nabla f _ {m} \left(w _ {p, m} ^ {(r, k - 1)}\right) - \nabla f _ {m} \left(w _ {g} ^ {(r)}\right) | | ^ {2} \tag {66} \\ \end{array}
$$
$$
\begin{array}{l} \leq \left(1 + \frac {1}{k - 1}\right) \mathbb {E} | | w _ {g, m} ^ {(r, k - 1)} - w _ {g, m} ^ {(r, 0)} | | ^ {2} + \eta_ {\ell} ^ {2} \sigma_ {\ell} ^ {2} + 2 k \eta_ {\ell} ^ {2} \mathbb {E} | | \nabla f _ {m} (w _ {g} ^ {(r)})) | | ^ {2} \\ + 2 k \eta_ {\ell} ^ {2} \beta^ {2} \mathbb {E} \| w _ {p, m} ^ {(r, k - 1)} - w _ {g} ^ {(r)} \| ^ {2} \tag {67} \\ \end{array}
$$
$$
\begin{array}{l} \mathbb {E} \| w _ {g, m} ^ {(r, k)} - w _ {g, m} ^ {(r, 0)} \| ^ {2} \leq \sum_ {i = 1} ^ {k} \left(2 k \eta_ {\ell} ^ {2} \mathbb {E} \| \nabla f _ {m} \left(w _ {g} ^ {(r)}\right) \| ^ {2} + \eta_ {\ell} ^ {2} \sigma_ {\ell} ^ {2} + 2 k \eta_ {\ell} ^ {2} \beta^ {2} \mathbb {E} \| w _ {p, m} ^ {(r, k - 1)} - w _ {g} ^ {(r)} \| ^ {2}\right) \\ \cdot \left(1 + \frac {1}{k - 1}\right) ^ {i} \tag {69} \\ \end{array}
$$
$$
\begin{array}{l} \mathbb {E} \left\| w _ {g, m} ^ {(r, k)} - w _ {g, m} ^ {(r, 0)} \right\| ^ {2} \leq 2 k \left(2 k \eta_ {\ell} ^ {2} \mathbb {E} \left\| \nabla f _ {m} \left(w _ {g} ^ {(r)}\right) \right\| ^ {2} + \eta_ {\ell} ^ {2} \sigma_ {\ell} ^ {2} + 2 k \eta_ {\ell} ^ {2} \beta^ {2} \sum_ {i = 1} ^ {k} \mathbb {E} \left\| w _ {p, m} ^ {(r, i - 1)} - w _ {g} ^ {(r)} \right\| ^ {2}\right) (70) \\ = 4 k ^ {2} \eta_ {\ell} ^ {2} \mathbb {E} | | \nabla f _ {m} \left(w _ {g} ^ {(r)}\right) | | ^ {2} + 2 k \eta_ {\ell} ^ {2} \sigma_ {\ell} ^ {2} + 4 k ^ {2} \eta_ {\ell} ^ {2} \beta^ {2} \sum_ {i = 1} ^ {k} \mathbb {E} | | w _ {p, m} ^ {(r, i - 1)} - w _ {g} ^ {(r)} | | ^ {2} (71) \\ \end{array}
$$
Lemma E.11 (Deviation of the personalized model from the global model). If  $m^{th}$  client's objective function  $f_{m}$  satisfies Assumptions E.2, E.3, and condition  $\eta_{\ell} \leq \frac{1}{2\sqrt{2}\beta K}$  in Algorithm 2, the following is satisfied:
$$
\mathbb {E} | | w _ {p, m} ^ {(r, k)} - w _ {g, m} ^ {(r, 0)} | | ^ {2} \leq 2 0 K ^ {3} \eta_ {\ell} ^ {2} \mathbb {E} | | 1 - \psi_ {g, m} ^ {(r, k)} | | ^ {2} \mathbb {E} | | \nabla f _ {m} (w _ {g} ^ {(r)}) | | ^ {2} + 1 0 K ^ {2} \eta_ {\ell} ^ {2} \sigma_ {\ell} ^ {2} \mathbb {E} | | \psi_ {g, m} ^ {(r, k)} | | ^ {2}.
$$
$$
\begin{array}{l} \mathbb {E} \| w _ {p, m} ^ {(r, k)} - w _ {g, m} ^ {(r, 0)} \| ^ {2} = \mathbb {E} \| \psi_ {g, m} ^ {(r, k)} w _ {g, m} ^ {(r, k)} + (1 - \psi_ {g, m} ^ {(r, k)}) w _ {\ell , m} ^ {(r, K)} - w _ {g, m} ^ {(r, 0)} \| ^ {2} (72) \\ = \mathbb {E} \| \psi_ {g, m} ^ {(r, k)} \left(w _ {g, m} ^ {(r, k)} - w _ {\ell , m} ^ {(r, k)}\right) + \left(w _ {\ell , m} ^ {(r, K)} - w _ {g, m} ^ {(r, 0)}\right) \| ^ {2} (73) \\ = \mathbb {E} | | \psi_ {g, m} ^ {(r, k)} \left(w _ {g, m} ^ {(r, k)} - w _ {g, m} ^ {(r, 0)} + w _ {g, m} ^ {(r, 0)} - w _ {\ell , m} ^ {(r, k)}\right) + \left(w _ {\ell , m} ^ {(r, K)} - w _ {g, m} ^ {(r, 0)}\right) | | ^ {2} (74) \\ \leq 2 \mathbb {E} \| \psi_ {g, m} ^ {(r, k)} \left(w _ {g, m} ^ {(r, k)} - w _ {g, m} ^ {(r, 0)}\right) \| ^ {2} + 2 \mathbb {E} \| \left(1 - \psi_ {g, m} ^ {(r, k)}\right) \left(w _ {\ell , m} ^ {(r, K)} - w _ {\ell , m} ^ {(r, 0)}\right) \| ^ {2} (75) \\ \end{array}
$$
Using lemmas E.6 and E.10,
$$
\begin{array}{l} \mathbb {E} \| w _ {p, m} ^ {(r, k)} - w _ {g, m} ^ {(r, 0)} \| ^ {2} \leq 2 \mathbb {E} \| 1 - \psi_ {g, m} ^ {(r, k + 1)} \| ^ {2} \left(4 K ^ {2} \eta_ {\ell} ^ {2} \mathbb {E} \| \nabla f _ {m} \left(w _ {g} ^ {(r)}\right) \| ^ {2} + 6 K ^ {2} \eta_ {\ell} ^ {2} \mathbb {E} \| \nabla f _ {m} \left(w _ {g} ^ {(r)}\right) \| ^ {2}\right) \\ \left. + 2 \mathbb {E} \left| \left| \psi_ {g, m} ^ {(r, k + 1)} \right| \right| ^ {2} \left(2 K \eta_ {\ell} ^ {2} \sigma_ {\ell} ^ {2} + 3 K \eta_ {\ell} ^ {2} \sigma_ {\ell} ^ {2} + 4 K ^ {2} \eta_ {\ell} ^ {2} \beta^ {2} \sum_ {i = 1} ^ {k} \mathbb {E} \left| \left| w _ {p, m} ^ {(r, i - 1)} - w _ {g} ^ {(r)} \right| \right| ^ {2}\right) \right. \tag {76} \\ \end{array}
$$
Assuming  $8K^2\eta_\ell^2\beta^2 \leq 1 \Rightarrow \eta \leq \frac{1}{2\sqrt{2}\beta K}$  and unrolling the recursion over  $w_{p,m}^{(r,i-1)} - w_g^{(r)}$ ,
$$
\begin{array}{l} \leq \sum_ {i = 1} ^ {k} \left(2 0 K ^ {2} \eta_ {\ell} ^ {2} \mathbb {E} | | 1 - \psi_ {g, m} ^ {(r, k)} | | ^ {2} \mathbb {E} | | \nabla f _ {m} \left(w _ {g} ^ {(r)}\right) | | ^ {2} + 1 0 K \eta_ {\ell} ^ {2} \sigma_ {\ell} ^ {2} \mathbb {E} | | \psi_ {g, m} ^ {(r, k)} | | ^ {2}\right) (77) \\ \leq 2 0 K ^ {3} \eta_ {\ell} ^ {2} \mathbb {E} \| 1 - \psi_ {g, m} ^ {(r, k)} \| ^ {2} \mathbb {E} \| \nabla f _ {m} \left(w _ {g} ^ {(r)}\right) \| ^ {2} + 1 0 K ^ {2} \eta_ {\ell} ^ {2} \sigma_ {\ell} ^ {2} \mathbb {E} \| \psi_ {g, m} ^ {(r, k)} \| ^ {2} (78) \\ \end{array}
$$
Theorem E.12 (Convergence of the Global Model for Non-convex Case). If each client's objective function  $f_{m}$  satisfies Assumptions E.2, E.3, E.4, using the learning rate  $\frac{1}{2\beta} \leq \eta_{\ell} \leq \min \left( \frac{1}{2\sqrt{5}\beta B K^2}, \frac{1}{\sqrt[3]{40K^4\beta^3B^2}} \right)$  in Algorithm 2, then the following convergence holds:
$$
\begin{array}{l} \frac {1}{R} \sum_ {r = 1} ^ {R} \mathbb {E} \left\| \nabla F \left(w _ {g} ^ {(r)}\right) \right\| ^ {2} \leq \frac {2}{\eta_ {\ell} \mathbf {q} _ {0} ^ {2} R} \left[ \mathbb {E} \left[ F \left(w _ {g} ^ {(1)}\right) \right] - \mathbb {E} \left[ F \left(w _ {g} ^ {(R + 1)}\right) \right] \right] + \frac {\eta_ {\ell} \beta \sigma_ {\ell} ^ {2} K}{M \mathbf {q} _ {0} ^ {2}} \\ + 4 0 \frac {\mathbf {q} _ {1} ^ {2}}{\mathbf {q} _ {0} ^ {2}} K ^ {4} \beta^ {2} \eta_ {\ell} G ^ {2} \left(\frac {2 \beta \eta_ {\ell} ^ {2} - \eta_ {\ell}}{2}\right) + 2 0 K ^ {3} \beta^ {2} \eta_ {\ell} \sigma_ {\ell} ^ {2} \left(\beta \eta_ {\ell} ^ {2} - \frac {\eta_ {\ell}}{2}\right). \\ \end{array}
$$
Proof. From the update rule stated in Equation 11, and  $\beta$ -smoothness of  $f_{m}$ , we have
$$
F \left(w _ {g} ^ {(r + 1)}\right) \leq F \left(w _ {g} ^ {(r)}\right) + \left\langle \nabla F \left(w _ {g} ^ {(r)}\right), w _ {g} ^ {(r + 1)} - w _ {g} ^ {(r)} \right\rangle + \frac {\beta}{2} \left| \left| w _ {g} ^ {(r + 1)} - w _ {g} ^ {(r)} \right| \right| ^ {2} \tag {79}
$$
Taking expectation on both sides,
$$
\mathbb {E} \left[ F \left(w _ {g} ^ {(r + 1)}\right) \right] \leq \mathbb {E} \left[ F \left(w _ {g} ^ {(r)}\right) \right] + \mathbb {E} \left[ \left\langle \nabla F \left(w _ {g} ^ {(r)}\right), w _ {g} ^ {(r + 1)} - w _ {g} ^ {(r)} \right\rangle \right] + \frac {\beta}{2} \| w _ {g} ^ {(r + 1)} - w _ {g} ^ {(r)} \| ^ {2} \tag {80}
$$
Using Equation 10 for second and third terms, and using the fact that the expectation is with respect to the choice of  $h_m$ ,
$$
\begin{array}{l} \leq \mathbb {E} \left[ F \left(w _ {g} ^ {(r)}\right) \right] - \eta_ {\ell} \left\langle \nabla F \left(w _ {g} ^ {(r)}\right), \frac {1}{M} \sum_ {m \in [ M ]} \alpha_ {m} \sum_ {k = 1} ^ {K} \mathbb {E} \left[ h _ {m} \left(w _ {p, m} ^ {(r, k - 1)}\right) \right] \right\rangle \\ + \frac {\beta \eta_ {\ell} ^ {2}}{2} \mathbb {E} \left\| \frac {1}{M} \sum_ {m \in [ M ]} \alpha_ {m} \sum_ {k = 1} ^ {K} h _ {m} \left(w _ {p, m} ^ {(r, k - 1)}\right) \right\| ^ {2}, \tag {81} \\ \end{array}
$$
where  $\alpha_{m} = \frac{n_{m}}{n}$ , which are the weights for weighted aggregation according to the sample count, as shown in Equation 11.
Separating mean and variance according to Assumption E.3,
$$
\begin{array}{l} \mathbb {E} \left[ F (w _ {g} ^ {(r + 1)}) \right] \leq \mathbb {E} \left[ F (w _ {g} ^ {(r)}) \right] - \eta_ {\ell} \left\langle \nabla F (w _ {g} ^ {(r)}), \frac {1}{M} \sum_ {m \in [ M ]} \alpha_ {m} \sum_ {k = 1} ^ {K} \mathbb {E} \left[ \nabla_ {w _ {g, m} ^ {(r, k - 1)}} f _ {m} (w _ {p, m} ^ {(r, k - 1)}) \right] \right\rangle \\ + \frac {\beta \eta_ {\ell} ^ {2}}{2} \mathbb {E} \left[ \left\| \frac {1}{M} \sum_ {m \in [ M ]} \alpha_ {m} \sum_ {k = 1} ^ {K} \nabla_ {w _ {g, m} ^ {(r, k - 1)}} f _ {m} \left(w _ {p, m} ^ {(r, k - 1)}\right) \right\| ^ {2} \right] + \frac {\eta_ {\ell} ^ {2} \beta \sigma_ {\ell} ^ {2} K}{2 M} \tag {82} \\ \end{array}
$$
Using  $\langle a,b\rangle = -\frac{1}{2} ||a - b||^2 +\frac{1}{2} ||a||^2 +\frac{1}{2} ||b||^2$
$$
\begin{array}{l} \mathbb {E} \left[ F \left(w _ {g} ^ {(r + 1)}\right) \right] \leq \mathbb {E} \left[ F \left(w _ {g} ^ {(r)}\right) \right] - \eta_ {\ell} \left[ - \frac {1}{2} \mathbb {E} \left\| \nabla F \left(w _ {g} ^ {(r)}\right) - \frac {1}{M} \sum_ {m \in [ M ]} \alpha_ {m} \sum_ {k = 1} ^ {K} \nabla_ {w _ {g, m} ^ {(r, k - 1)}} f _ {m} \left(w _ {p, m} ^ {(r, k - 1)}\right) \right\| ^ {2} \right] \\ - \eta_ {\ell} \left[ \frac {1}{2} \mathbb {E} \left\| \nabla F \left(w _ {g} ^ {(r)}\right) \right\| ^ {2} + \frac {1}{2} \mathbb {E} \left\| \frac {1}{M} \sum_ {m \in [ M ]} \alpha_ {m} \sum_ {k = 1} ^ {K} \nabla_ {w _ {g, m} ^ {(r, k - 1)}} f _ {m} \left(w _ {p, m} ^ {(r, k - 1)}\right) \right\| ^ {2} \right] \\ + \frac {\beta \eta_ {\ell} ^ {2}}{2} \mathbb {E} \left[ \left\| \frac {1}{M} \sum_ {m \in [ M ]} \alpha_ {m} \sum_ {k = 1} ^ {K} \nabla_ {w _ {g, m} ^ {(r, k - 1)}} f _ {m} \left(w _ {p, m} ^ {(r, k - 1)}\right) \right\| ^ {2} \right] + \frac {\eta_ {\ell} ^ {2} \beta \sigma_ {\ell} ^ {2} K}{2 M} (83) \\ \leq \mathbb {E} \left[ F \left(w _ {g} ^ {(r)}\right) \right] - \frac {\eta_ {\ell}}{2} \mathbb {E} \left\| \nabla F \left(w _ {g} ^ {(r)}\right) \right\| ^ {2} \\ - \left(\frac {\eta_ {\ell}}{2} - \frac {\beta \eta_ {\ell} ^ {2}}{2}\right) \mathbb {E} \left\| \frac {1}{M} \sum_ {m \in [ M ]} \alpha_ {m} \sum_ {k = 1} ^ {K} \nabla_ {w _ {g, m} ^ {(r, k - 1)}} f _ {m} (w _ {p, m} ^ {(r, k - 1)}) \right\| ^ {2} \\ + \frac {\eta_ {\ell}}{2} \mathbb {E} \left\| \nabla F \left(w _ {g} ^ {(r)}\right) - \frac {1}{M} \sum_ {m \in [ M ]} \alpha_ {m} \sum_ {k = 1} ^ {K} \nabla_ {w _ {g, m} ^ {(r, k - 1)}} f _ {m} \left(w _ {p, m} ^ {(r, k - 1)}\right) \right\| ^ {2} + \frac {\eta_ {\ell} ^ {2} \beta \sigma_ {\ell} ^ {2} K}{2 M} (84) \\ \end{array}
$$