SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
\begin{array}{l} \mathbb {E} \left\| w _ {p, m} ^ {(r, K)} - w _ {p, m} ^ {(r, k)} \right\| ^ {2} = \mathbb {E} \left\| \sum_ {i = k} ^ {K} w _ {p, m} ^ {(r, k + 1)} - w _ {p, m} ^ {(r, k)} \right\| ^ {2} (123) \\ \leq 3 \eta_ {\ell} ^ {2} \sum_ {i = k} ^ {K} \mathbb {E} | | \nabla f _ {m} (w _ {p, m} ^ {(r, i)}) | | ^ {2} + 3 \eta_ {\ell} ^ {2} K \sum_ {i = k} ^ {K} \mathbb {E} | | \psi_ {g, m} ^ {(r, i)} | | ^ {2} \\ + 6 \eta_ {\ell} ^ {2} \sum_ {i = k} ^ {K} \left(6 K ^ {2} \eta_ {\ell} ^ {2} \mathbb {E} | | \nabla f _ {m} \left(w _ {g} ^ {(r)}\right) | | ^ {2} + 3 K \eta_ {\ell} ^ {2} \sigma_ {\ell} ^ {2}\right) \\ + 6 \eta_ {\ell} ^ {2} \sum_ {i = k} ^ {K} \left(8 K ^ {3} \eta_ {\ell} ^ {2} \mathbb {E} | | \psi_ {g, m} ^ {(r, i)} | | ^ {2} \mathbb {E} | | \nabla f _ {m} \left(w _ {g} ^ {(r)}\right) | | ^ {2} + 4 K \eta_ {\ell} ^ {2} \sigma_ {\ell} ^ {2}\right) (124) \\ \end{array}
$$
$$
\begin{array}{l} \leq 6 \beta \eta_ {\ell} ^ {2} \left(\mathbb {E} \left[ f _ {m} \left(w _ {p, m} ^ {(r, K)}\right) \right] - f \left(w _ {p, m} ^ {*}\right)\right) + 3 \eta_ {\ell} ^ {2} K \sum_ {i = k} ^ {K} \mathbb {E} | | \psi_ {g, m} ^ {(r, i)} | | ^ {2} \\ + 3 6 K ^ {3} \eta_ {\ell} ^ {4} \mathbb {E} | | \nabla f _ {m} (w _ {g} ^ {(r)}) | | ^ {2} + 1 8 K ^ {2} \eta_ {\ell} ^ {4} \sigma_ {\ell} ^ {2} \\ + 4 8 K ^ {4} \eta_ {\ell} ^ {4} \mathbb {E} | | \nabla f _ {m} \left(w _ {g} ^ {(r)}\right) | | ^ {2} \sum_ {i = k} ^ {K} \mathbb {E} | | \psi_ {g, m} ^ {(r, i)} | | ^ {2} + 2 4 K ^ {2} \eta_ {\ell} ^ {4} \sigma_ {\ell} ^ {2} \tag {125} \\ \end{array}
$$
$$
\begin{array}{l} \therefore \mathbb {E} | | w _ {p, m} ^ {(r, K)} - w _ {p, m} ^ {(r, k)} | | ^ {2} \leq 6 \beta \eta_ {\ell} ^ {2} \left(\mathbb {E} [ f _ {m} (w _ {p, m} ^ {(r, K)}) ] - f (w _ {p, m} ^ {*})\right) + 3 \eta_ {\ell} ^ {2} K \sum_ {i = k} ^ {K} \mathbb {E} | | \psi_ {g, m} ^ {(r, i)} | | ^ {2} \\ + 3 6 K ^ {3} \eta_ {\ell} ^ {4} \mathbb {E} | | \nabla f _ {m} \left(w _ {g} ^ {(r)}\right) | | ^ {2} + 4 0 K ^ {2} \eta_ {\ell} ^ {4} \sigma_ {\ell} ^ {2} \\ + 4 8 K ^ {4} \eta_ {\ell} ^ {4} \mathbb {E} | | \nabla f _ {m} \left(w _ {g} ^ {(r)}\right) | | ^ {2} \sum_ {i = k} ^ {K} \mathbb {E} | | \psi_ {g, m} ^ {(r, i)} | | ^ {2} \tag {126} \\ \end{array}
$$
Theorem E.16 (Convergence of the Personalized Model for Convex (Strong and General) Cases). If each client's objective function  $f_{m}$  satisfies Assumptions E.2, E.3, E.4, using the learning rate  $\frac{1}{\mu R} \leq \eta_{\ell} \leq \frac{1}{K\beta^{2}}$  in Algorithm 2, then the following convergence holds: (Strong Convex Case)
$$
\begin{array}{l} \mathbb {E} \left[ f _ {m} (w _ {p, m} ^ {(R, 0)}) \right] - f _ {m} (w _ {p, m} ^ {*}) \leq \frac {3 6 \mu^ {2}}{R K ^ {3}} \mathbb {E} | | w _ {p, m} ^ {(1, K)} - w _ {p, m} ^ {*} | | ^ {2} \exp \left(\frac {1}{K - 1} - \eta_ {\ell} \mu K R\right) \\ + 1 2 K ^ {2} \eta_ {\ell} ^ {2} \delta_ {m} ^ {w _ {g}} + 1 2 K ^ {2} \eta_ {\ell} ^ {2} \frac {1}{R} \sum_ {r = 1} ^ {R} \mathbb {E} | | \nabla F (w _ {g} ^ {(r + 1)}) | | ^ {2} + 4 K \eta_ {\ell} ^ {2} \sigma_ {\ell} ^ {2} + \frac {\mathbf {q} _ {0} ^ {2}}{2} + 1 6 K ^ {3} \eta_ {\ell} ^ {2} \mathbf {q} _ {0} ^ {2} \delta_ {m} ^ {w _ {g}} \\ + 1 6 K ^ {3} \eta_ {\ell} ^ {2} \mathbf {q} _ {0} ^ {2} \frac {1}{R} \sum_ {r = 1} ^ {R} \mathbb {E} | | \nabla F (w _ {g} ^ {(r + 1)}) | | ^ {2} + \frac {K ^ {2} \eta_ {\ell} ^ {2}}{2} \left(\frac {\sigma_ {\ell} ^ {2}}{K} + \left(\delta_ {m} ^ {\psi} + \frac {\delta^ {\psi}}{M}\right)\right) \left(\frac {\sigma_ {\ell} ^ {2}}{K} + \left(\delta_ {m} ^ {w _ {g}} + \frac {\delta^ {w _ {g}}}{M}\right)\right) \\ + \frac {1 + \eta_ {\ell} ^ {2} K ^ {2} \beta^ {4}}{6} \left(K \sigma_ {\ell} ^ {2} \eta_ {\ell} ^ {2} + \left(\delta_ {m} ^ {\psi} + \frac {\delta^ {\psi}}{M}\right) K ^ {2} \eta_ {\ell} ^ {2}\right) \left(\frac {2}{R} \sum_ {r = 1} ^ {R} \mathbb {E} | | \nabla F (w _ {g} ^ {(r)}) | | ^ {2} + 2 \delta_ {m} ^ {w _ {g}}\right) \\ \end{array}
$$
$$
\begin{array}{l} \mathbb {E} \left[ f _ {m} \left(w _ {p, m} ^ {(R, 0)}\right) \right] - f _ {m} \left(w _ {p, m} ^ {*}\right) \leq \frac {1}{3 6 \eta_ {\ell} ^ {2} K ^ {2} R} \left(1 + \frac {1}{K - 1}\right) \mathbb {E} | | w _ {p, m} ^ {(1, K)} - w _ {p, m} ^ {*} | | ^ {2} \\ + \eta_ {\ell} ^ {2} \left(1 2 K ^ {2} \frac {1}{R} \sum_ {r = 1} ^ {R} \mathbb {E} | | \nabla F (w _ {g} ^ {(r)}) | | ^ {2}\right) ^ {1 / 3} + \eta_ {\ell} ^ {2} \left(1 2 K ^ {2} \delta_ {m} ^ {w _ {g}}\right) ^ {1 / 3} + \eta_ {\ell} ^ {2} \left(4 K \sigma_ {\ell} ^ {2}\right) ^ {1 / 3} \\ + \eta_ {\ell} ^ {2} (1 6 K ^ {3} \mathbf {q} _ {0} ^ {2} \frac {1}{R} \sum_ {r = 1} ^ {R} \mathbb {E} | | \nabla F (w _ {g} ^ {(r + 1)}) | | ^ {2}) ^ {1 / 3} + \frac {\mathbf {q} _ {0} ^ {2}}{2} + \eta_ {\ell} ^ {2} (1 6 K ^ {3} \mathbf {q} _ {0} ^ {2} \delta_ {m} ^ {w _ {g}}) ^ {1 / 3} \\ \left. \right. + \eta_ {\ell} ^ {2} \left(\frac {K ^ {2}}{2} \left(\frac {\sigma_ {\ell} ^ {2}}{K} + \left(\delta_ {m} ^ {\psi} + \frac {\delta^ {\psi}}{M}\right)\right)\left(\frac {\sigma_ {\ell} ^ {2}}{K} + \left(\delta_ {m} ^ {w _ {g}} + \frac {\delta^ {w _ {g}}}{M}\right)\right)\right) ^ {1 / 3} \\ + \eta_ {\ell} ^ {2} \left(\frac {K ^ {2}}{3} \left(\frac {\sigma_ {\ell} ^ {2}}{K} + \left(\delta_ {m} ^ {\psi} + \frac {\delta^ {\psi}}{M}\right)\right) \left(\frac {2}{R} \sum_ {r = 1} ^ {R} \mathbb {E} | | \nabla F \left(w _ {g} ^ {(r)}\right) | | ^ {2} + 2 \delta_ {m} ^ {w _ {g}}\right)\right) ^ {1 / 3} \\ \end{array}
$$
where  $\frac{1}{R}\sum_{r = 1}^{R}\mathbb{E}||\nabla F(w_g^{(r)})||^2$  is bounded as shown in Theorem E.12.
Proof. We restate the update rules of the personalized model in Algorithm 2,
1. For all samples  $x_{m}$ , define  $\tilde{w}_{p,m}^{(r,k)}(x_m) \gets \psi_{g,m}^{(r,k)}(x_m) w_{g,m}^{(r,k)}(x_m) + (1 - \psi_{g,m}^{(r,k)}(x_m)) w_{\ell,m}^{(r,K)}(x_m)$ .
2. Train policy parameters  $\psi_{g,m}^{(r,k + 1)}\gets \psi_{g,m}^{(r,k)} - \eta_{\ell}\nabla_{\psi_{g,m}^{(r,k)}}f_m(\tilde{w}_{p,m}^{(r,k)}(x_m),y_m)$
3. For all samples  $x_{m}$ , define
$$
w _ {p, m} ^ {(r, k)} \left(x _ {m}\right) \leftarrow \psi_ {g, m} ^ {(r, k + 1)} \left(x _ {m}\right) w _ {g, m} ^ {(r, k)} \left(x _ {m}\right) + \left(1 - \psi_ {g, m} ^ {(r, k + 1)} \left(x _ {m}\right)\right) w _ {\ell , m} ^ {(r, K)} \left(x _ {m}\right)
$$
4. Train global parameters  $w_{g,m}^{(r,k)} \gets w_{g,m}^{(r,k-1)} - \eta_{\ell} \nabla_{w_{g,m}^{(r,k-1)}} f_m(w_{p,m}^{(r,k)}(x_m), y_m)$
$$
\begin{array}{l} \mathbb {E} \left\| w _ {p, m} ^ {(r + 1, K)} - w _ {p, m} ^ {*} \right\| ^ {2} = \mathbb {E} \left\| w _ {p, m} ^ {(r + 1, K)} - \tilde {w} _ {p, m} ^ {(r + 1, 0)} + \tilde {w} _ {p, m} ^ {(r + 1, 0)} - w _ {p, m} ^ {(r, K)} + w _ {p, m} ^ {(r, K)} - w _ {p, m} ^ {*} \right\| ^ {2} (127) \\ \leq 2 K \underbrace {\mathbb {E} | | w _ {p , m} ^ {(r + 1 , K)} - \tilde {w} _ {p , m} ^ {(r + 1 , 0)} | | ^ {2}} _ {\text {L e m m a E . 1 3}} + 2 K \underbrace {\mathbb {E} | | \tilde {w} _ {p , m} ^ {(r + 1 , 0)} - w _ {p , m} ^ {(r , K)} | | ^ {2}} _ {\text {L e m m a E . 1 4}} \\ + \left(1 + \frac {1}{K - 1}\right) \mathbb {E} \left\| w _ {p, m} ^ {(r, K)} - w _ {p, m} ^ {*} \right\| ^ {2} (128) \\ \end{array}
$$
And using Assumption E.4,
$$
\begin{array}{l} \leq 3 6 K ^ {2} \eta_ {\ell} ^ {2} \left[ f _ {m} \left(w _ {p, m} ^ {*}\right) - \mathbb {E} \left[ f _ {m} \left(w _ {p, m} ^ {(r + 1, 0)}\right) \right] \right] + 2 1 6 K ^ {4} \eta_ {\ell} ^ {4} \mathbb {E} | | \nabla f _ {m} \left(w _ {g} ^ {(r + 1)}\right) | | ^ {2} + 1 2 6 K ^ {3} \eta_ {\ell} ^ {4} \sigma_ {\ell} ^ {2} \\ + 1 8 K ^ {2} \eta_ {\ell} ^ {2} \mathbb {E} \left\| \psi_ {g, m} ^ {(r + 1, K)} \right\| ^ {2} + 2 8 8 K ^ {5} \eta_ {\ell} ^ {4} \mathbb {E} \| \psi_ {g, m} ^ {(r + 1, K)} \| ^ {2} \mathbb {E} \| \nabla f _ {m} (w _ {g} ^ {(r + 1)}) \| ^ {2} \\ + 1 8 \left(K \sigma_ {\ell} ^ {2} \eta_ {\ell} ^ {2} + \left(\delta_ {m} ^ {\psi} + \frac {\delta^ {\psi}}{M}\right) K ^ {2} \eta_ {\ell} ^ {2}\right) \left(K \sigma_ {\ell} ^ {2} \eta_ {\ell} ^ {2} + \left(\delta_ {m} ^ {w _ {g}} + \frac {\delta^ {w _ {g}}}{M}\right) K ^ {2} \eta_ {\ell} ^ {2}\right) \\ + 6 \left(1 + \eta_ {\ell} ^ {2} K ^ {2} \beta^ {4}\right) \eta_ {\ell} ^ {2} K ^ {2} \left(K \sigma_ {\ell} ^ {2} \eta_ {\ell} ^ {2} + \left(\delta_ {m} ^ {\psi} + \frac {\delta^ {\psi}}{M}\right) K ^ {2} \eta_ {\ell} ^ {2}\right) \mathbb {E} | | \nabla f _ {m} \left(w _ {g} ^ {(r)}\right) | | ^ {2} \\ + \left(1 + \frac {1}{K - 1} - \mu \eta_ {\ell}\right) \mathbb {E} | | w _ {p, m} ^ {(r, K)} - w _ {p, m} ^ {*} | | ^ {2} \tag {129} \\ \end{array}
$$
Rearranging the terms,
$$
\begin{array}{l} 3 6 K ^ {2} \eta_ {\ell} ^ {2} \left[ \mathbb {E} \left[ f _ {m} \left(w _ {p, m} ^ {(r + 1, 0)}\right) \right] - f _ {m} \left(w _ {p, m} ^ {*}\right) \right] \leq \left(1 + \frac {1}{K - 1} - \mu \eta_ {\ell}\right) \mathbb {E} | | w _ {p, m} ^ {(r, K)} \\ + 2 1 6 K ^ {4} \eta_ {\ell} ^ {4} \mathbb {E} | | \nabla f _ {m} \left(w _ {g} ^ {(r + 1)}\right) | | ^ {2} + 1 2 6 K ^ {3} \eta_ {\ell} ^ {4} \sigma_ {\ell} ^ {2} + 1 8 K ^ {2} \eta_ {\ell} ^ {2} \mathbb {E} \left| | \psi_ {g, m} ^ {(r + 1, K)} \right| ^ {2} \\ - w _ {p, m} ^ {*} \| ^ {2} - \mathbb {E} \| w _ {p, m} ^ {(r, K + 1)} - w _ {p, m} ^ {*} \| ^ {2} + 2 8 8 K ^ {5} \eta_ {\ell} ^ {4} \mathbb {E} \| \psi_ {g, m} ^ {(r + 1, K)} \| ^ {2} \mathbb {E} \| \nabla f _ {m} (w _ {g} ^ {(r + 1)}) \| ^ {2} \\ + 1 8 \left(K \sigma_ {\ell} ^ {2} \eta_ {\ell} ^ {2} + \left(\delta_ {m} ^ {\psi} + \frac {\delta^ {\psi}}{M}\right) K ^ {2} \eta_ {\ell} ^ {2}\right) \left(K \sigma_ {\ell} ^ {2} \eta_ {\ell} ^ {2} + \left(\delta_ {m} ^ {w _ {g}} + \frac {\delta^ {w _ {g}}}{M}\right) K ^ {2} \eta_ {\ell} ^ {2}\right) \\ + 6 \left(1 + \eta_ {\ell} ^ {2} K ^ {2} \beta^ {4}\right) \eta_ {\ell} ^ {2} K ^ {2} \left(K \sigma_ {\ell} ^ {2} \eta_ {\ell} ^ {2} + \left(\delta_ {m} ^ {\psi} + \frac {\delta^ {\psi}}{M}\right) K ^ {2} \eta_ {\ell} ^ {2}\right) \mathbb {E} | | \nabla f _ {m} \left(w _ {g} ^ {(r)}\right) | | ^ {2} \tag {130} \\ \end{array}
$$
$$
\begin{array}{l} \therefore \mathbb {E} \left[ f _ {m} \left(w _ {p, m} ^ {(r + 1, 0)}\right) \right] - f _ {m} \left(w _ {p, m} ^ {*}\right) \leq \frac {1}{3 6 \eta_ {\ell} ^ {2} K ^ {2}} \left(1 + \frac {1}{K - 1} - \mu \eta_ {\ell}\right) \mathbb {E} | | w _ {p, m} ^ {(r, K)} - w _ {p, m} ^ {*} | | ^ {2} \\ - \frac {1}{3 6 \eta_ {\ell} ^ {2} K ^ {2}} \mathbb {E} \| w _ {p, m} ^ {(r, K + 1)} - w _ {p, m} ^ {*} \| ^ {2} + 6 K ^ {2} \eta_ {\ell} ^ {2} \mathbb {E} \| \nabla f _ {m} (w _ {g} ^ {(r + 1)}) \| ^ {2} + 4 K \eta_ {\ell} ^ {2} \sigma_ {\ell} ^ {2} + \frac {1}{2} \mathbb {E} \| \psi_ {g, m} ^ {(r + 1, K)} \| ^ {2} \\ + 8 K ^ {3} \eta_ {\ell} ^ {2} \mathbb {E} | | \psi_ {g, m} ^ {(r + 1, K)} | | ^ {2} \mathbb {E} | | \nabla f _ {m} \left(w _ {g} ^ {(r + 1)}\right) | | ^ {2} + \frac {K ^ {2} \eta_ {\ell} ^ {2}}{2} \left(\frac {\sigma_ {\ell} ^ {2}}{K} + \left(\delta_ {m} ^ {\psi} + \frac {\delta^ {\psi}}{M}\right)\right) \left(\frac {\sigma_ {\ell} ^ {2}}{K} + \left(\delta_ {m} ^ {w _ {g}} + \frac {\delta^ {w _ {g}}}{M}\right)\right) \\ + \frac {1 + \eta_ {\ell} ^ {2} K ^ {2} \beta^ {4}}{6} \left(K \sigma_ {\ell} ^ {2} \eta_ {\ell} ^ {2} + \left(\delta_ {m} ^ {\psi} + \frac {\delta^ {\psi}}{M}\right) K ^ {2} \eta_ {\ell} ^ {2}\right) \mathbb {E} | | \nabla f _ {m} \left(w _ {g} ^ {(r)}\right) | | ^ {2} \tag {131} \\ \end{array}
$$
For strong convex  $(\mu > 0)$  case, using the linear convergence rate lemma from [18] (Lemma 1) and Definition E.5,
$$
\begin{array}{l} \mathbb {E} \left[ f _ {m} \left(w _ {p, m} ^ {(R, 0)}\right) \right] - f _ {m} \left(w _ {p, m} ^ {*}\right) \leq \frac {3 6 \mu^ {2}}{R K ^ {3}} \mathbb {E} \left\| w _ {p, m} ^ {(1, K)} - w _ {p, m} ^ {*} \right\| ^ {2} \exp \left(\frac {1}{K - 1} - \eta_ {\ell} \mu K R\right) \\ + 1 2 K ^ {2} \eta_ {\ell} ^ {2} \delta_ {m} ^ {w _ {g}} + 1 2 K ^ {2} \eta_ {\ell} ^ {2} \frac {1}{R} \sum_ {r = 1} ^ {R} \mathbb {E} | | \nabla F (w _ {g} ^ {(r + 1)}) | | ^ {2} + 4 K \eta_ {\ell} ^ {2} \sigma_ {\ell} ^ {2} + \frac {\mathbf {q} _ {0} ^ {2}}{2} + 1 6 K ^ {3} \eta_ {\ell} ^ {2} \mathbf {q} _ {0} ^ {2} \delta_ {m} ^ {w _ {g}} \\ + 1 6 K ^ {3} \eta_ {\ell} ^ {2} \mathbf {q} _ {0} ^ {2} \frac {1}{R} \sum_ {r = 1} ^ {R} \mathbb {E} | | \nabla F (w _ {g} ^ {(r + 1)}) | | ^ {2} + \frac {K ^ {2} \eta_ {\ell} ^ {2}}{2} \left(\frac {\sigma_ {\ell} ^ {2}}{K} + \left(\delta_ {m} ^ {\psi} + \frac {\delta^ {\psi}}{M}\right)\right) \left(\frac {\sigma_ {\ell} ^ {2}}{K} + \left(\delta_ {m} ^ {w _ {g}} + \frac {\delta^ {w _ {g}}}{M}\right)\right) \\ + \frac {1 + \eta_ {\ell} ^ {2} K ^ {2} \beta^ {4}}{6} \left(K \sigma_ {\ell} ^ {2} \eta_ {\ell} ^ {2} + \left(\delta_ {m} ^ {\psi} + \frac {\delta^ {\psi}}{M}\right) K ^ {2} \eta_ {\ell} ^ {2}\right) \left(\frac {2}{R} \sum_ {r = 1} ^ {R} \mathbb {E} | | \nabla F \left(w _ {g} ^ {(r)}\right) | | ^ {2} + 2 \delta_ {m} ^ {w _ {g}}\right) \tag {132} \\ \end{array}
$$
For general convex  $(\mu = 0)$  case, using the sublinear convergence rate lemma from [18] (Lemma 2), and conditioning on  $\eta_{\ell}^{2}K^{2}\beta^{4}\leq 1\Rightarrow \eta_{\ell}\leq \frac{1}{K\beta^{2}}$
$$
\begin{array}{l} \mathbb {E} \left[ f _ {m} \left(w _ {p, m} ^ {(R, 0)}\right) \right] - f _ {m} \left(w _ {p, m} ^ {*}\right) \leq \frac {1}{3 6 \eta_ {\ell} ^ {2} K ^ {2} R} \left(1 + \frac {1}{K - 1}\right) \mathbb {E} | | w _ {p, m} ^ {(1, K)} - w _ {p, m} ^ {*} | | ^ {2} \\ + \eta_ {\ell} ^ {2} \left(1 2 K ^ {2} \frac {1}{R} \sum_ {r = 1} ^ {R} \mathbb {E} | | \nabla F \left(w _ {g} ^ {(r)}\right) | | ^ {2}\right) ^ {1 / 3} + \eta_ {\ell} ^ {2} \left(1 2 K ^ {2} \delta_ {m} ^ {w _ {g}}\right) ^ {1 / 3} + \eta_ {\ell} ^ {2} \left(4 K \sigma_ {\ell} ^ {2}\right) ^ {1 / 3} + \frac {\mathbf {q} _ {0} ^ {2}}{2} \\ + \eta_ {\ell} ^ {2} \left(1 6 K ^ {3} \mathbf {q} _ {0} ^ {2} \delta_ {m} ^ {w _ {g}}\right) ^ {1 / 3} + \eta_ {\ell} ^ {2} \left(1 6 K ^ {3} \mathbf {q} _ {0} ^ {2} \frac {1}{R} \sum_ {r = 1} ^ {R} \mathbb {E} | | \nabla F (w _ {g} ^ {(r + 1)}) | | ^ {2}\right) ^ {1 / 3} \\ + \eta_ {\ell} ^ {2} \left(\frac {K ^ {2}}{2} \left(\frac {\sigma_ {\ell} ^ {2}}{K} + \left(\delta_ {m} ^ {\psi} + \frac {\delta^ {\psi}}{M}\right)\right) \left(\frac {\sigma_ {\ell} ^ {2}}{K} + \left(\delta_ {m} ^ {w _ {g}} + \frac {\delta^ {w _ {g}}}{M}\right)\right)\right) ^ {1 / 3} t z o \\ + \eta_ {\ell} ^ {2} \left(\frac {K ^ {2}}{3} \left(\frac {\sigma_ {\ell} ^ {2}}{K} + \left(\delta_ {m} ^ {\psi} + \frac {\delta^ {\psi}}{M}\right)\right) \left(\frac {2}{R} \sum_ {r = 1} ^ {R} \mathbb {E} | | \nabla F \left(w _ {g} ^ {(r)}\right) | | ^ {2} + 2 \delta_ {m} ^ {w _ {g}}\right)\right) ^ {1 / 3} \tag {133} \\ \end{array}
$$
E.6 Convergence Proof for the Personalized Model: Non-convex Case
Lemma E.17 (One round progress of the local model). If  $m^{th}$  client's objective function  $f_{m}$  satisfies Assumptions E.3, E.4, in Algorithm 2, the following is satisfied:
$$
\mathbb {E} \left\| w _ {\ell , m} ^ {(r + 1, K)} - w _ {\ell , m} ^ {(r, K)} \right\| ^ {2} \leq \left(1 - 2 \eta_ {\ell} K \beta^ {2} + \eta_ {\ell} ^ {2} K ^ {2} \beta^ {4}\right) \eta_ {\ell} ^ {2} K ^ {2} \left(G ^ {2} + B ^ {2} \mathbb {E} \| \nabla F (w _ {g} ^ {(r)}) \| ^ {2}\right)
$$
$$
\begin{array}{l} \mathbb {E} \| w _ {\ell , m} ^ {(r + 1, K)} - w _ {\ell , m} ^ {(r, K)} \| ^ {2} = \mathbb {E} \| w _ {g} ^ {(r + 1)} - \eta_ {\ell} \sum_ {k = 1} ^ {K} \nabla f _ {m} \left(w _ {g} ^ {(r + 1)}\right) - w _ {g} ^ {(r)} + \eta_ {\ell} \sum_ {k = 1} ^ {K} \nabla f _ {m} \left(w _ {g} ^ {(r)}\right) \| ^ {2} (134) \\ = \mathbb {E} | | w _ {g} ^ {(r + 1)} - w _ {g} ^ {(r)} - \eta_ {\ell} \sum_ {k = 1} ^ {K} \left[ \nabla f _ {m} \left(w _ {g} ^ {(r + 1)}\right) - \nabla f _ {m} \left(w _ {g} ^ {(r)}\right) \right] | | ^ {2} (135) \\ \leq \mathbb {E} \left| \left| w _ {g} ^ {(r + 1)} - w _ {g} ^ {(r)} - \eta_ {\ell} K \beta^ {2} \left(w _ {g} ^ {(r + 1)} - w _ {g} ^ {(r)}\right) \right| \right| ^ {2} (136) \\ = \mathbb {E} | | (1 - \eta_ {\ell} K \beta^ {2}) \left(w _ {g} ^ {(r + 1)} - w _ {g} ^ {(r)}\right) | | ^ {2} (137) \\ \leq \left(1 - \eta_ {\ell} K \beta^ {2}\right) ^ {2} \mathbb {E} \left\| \frac {1}{M} \sum_ {c \in [ M ]} w _ {g, m} ^ {(r, K)} - w _ {g} ^ {(r)} \right\| ^ {2} (138) \\ = \left(1 - \eta_ {\ell} K \beta^ {2}\right) ^ {2} \mathbb {E} \left\| \frac {1}{M} \sum_ {c \in [ M ]} \left(w _ {g} ^ {(r)} - \eta_ {\ell} \sum_ {k = 1} ^ {K} \nabla f _ {m} \left(w _ {g} ^ {(r)}\right)\right) - w _ {g} ^ {(r)} \right\| ^ {2} (139) \\ = \left(1 - \eta_ {\ell} K \beta^ {2}\right) ^ {2} \mathbb {E} \bigg \| - \frac {\eta_ {\ell}}{M} \sum_ {c \in [ M ]} \sum_ {k = 1} ^ {K} \nabla f _ {m} \left(w _ {g} ^ {(r)}\right) \bigg \| ^ {2} (140) \\ \leq \left(1 - \eta_ {\ell} K \beta^ {2}\right) ^ {2} \eta_ {\ell} ^ {2} \mathbb {E} \left\| \frac {K}{M} \sum_ {c \in [ M ]} \nabla f _ {m} \left(w _ {g} ^ {(r)}\right) \right\| ^ {2} (141) \\ \leq \left(1 - \eta_ {\ell} K \beta^ {2}\right) ^ {2} \eta_ {\ell} ^ {2} K ^ {2} \left(G ^ {2} + B ^ {2} \mathbb {E} | | \nabla F \left(w _ {g} ^ {(r)}\right) | | ^ {2}\right) (142) \\ = \left(1 - 2 \eta_ {\ell} K \beta^ {2} + \eta_ {\ell} ^ {2} K ^ {2} \beta^ {4}\right) \eta_ {\ell} ^ {2} K ^ {2} \left(G ^ {2} + B ^ {2} \mathbb {E} | | \nabla F \left(w _ {g} ^ {(r)}\right) | | ^ {2}\right) (143) \\ \end{array}
$$
The last inequality follows from Assumption E.4.
We proceed with a lemma which binds the deviation of the personalized model  $w_{p}$  of an arbitrary client  $m$  over one round, i.e.,  $w_{p,m}^{(r + 1)}$  and  $w_{p,m}^{(r)}$ , for non-convex case.
Lemma E.18 (Local progress of personalized model). If  $m^{th}$  client's objective function  $f_{m}$  satisfies Assumptions E.3, E.4, and  $\eta_{\ell} \leq \frac{1}{K\sqrt{12\beta(K - 1)}}$ , in Algorithm 2, the following is satisfied:
$$
\begin{array}{l} \mathbb {E} \| w _ {p, m} ^ {(r, k + 1)} - w _ {p, m} ^ {(r, 0)} \| ^ {2} \leq 1 8 K ^ {5} \eta_ {\ell} ^ {2} \mathbb {E} \| \nabla f _ {m} (w _ {g} ^ {(r)}) \| ^ {2} + 9 K ^ {4} \eta_ {\ell} ^ {2} \sigma_ {\ell} ^ {2} + 3 6 K ^ {3} \eta_ {\ell} ^ {2} \sigma_ {\ell} ^ {2} \mathbb {E} \| 1 - \psi_ {g, m} ^ {(r, k)} \| ^ {2} \\ + 2 4 K ^ {4} \eta_ {\ell} ^ {2} \mathbb {E} \| \nabla f _ {m} (w _ {g} ^ {(r)}) \| ^ {2} \mathbb {E} \| \psi_ {g, m} ^ {(r, k)} \| ^ {2} \\ \end{array}
$$
Proof. We start with using the update rule stated for the personalized model at the beginning of Theorem E.16,
$$
\mathbb {E} \| w _ {p, m} ^ {(r, k + 1)} - w _ {p, m} ^ {(r, 0)} \| ^ {2} = \mathbb {E} \| \psi_ {g, m} ^ {(r, k + 1)} w _ {g, m} ^ {(r, k + 1)} + (1 - \psi_ {g, m} ^ {(r, k + 1)}) w _ {\ell , m} ^ {(r, K)} - w _ {p, m} ^ {(r, 0)} \| ^ {2} \tag {144}
$$
Expanding by one iterate,
$$
\begin{array}{l} = \mathbb {E} | | \left(\psi_ {g, m} ^ {(r, k)} - \eta_ {\ell} \nabla_ {\psi_ {g, m} ^ {(r, k)}} f _ {m} (w _ {p, m} ^ {(r, k)})\right) \left(w _ {g, m} ^ {(r, k)} - \eta_ {\ell} \nabla_ {w _ {g, m} ^ {(r, k)}} f _ {m} (w _ {p, m} ^ {(r, k)})\right) \\ + \left(1 - \psi_ {g, m} ^ {(r, k)} + \eta_ {\ell} \nabla_ {\psi_ {g, m} ^ {(r, k)}} f _ {m} \left(w _ {p, m} ^ {(r, k)}\right)\right) w _ {\ell , m} ^ {(r, K)} - w _ {p, m} ^ {(r, 0)} \| ^ {2} \tag {145} \\ \end{array}
$$
$$
\begin{array}{l} = \mathbb {E} | | \psi_ {g, m} ^ {(r, k)} w _ {g, m} ^ {(r, k)} - w _ {g, m} ^ {(r, k)} \eta_ {\ell} \nabla_ {\psi_ {g, m} ^ {(r, k)}} f _ {m} (w _ {p, m} ^ {(r, k)}) - \psi_ {g, m} ^ {(r, k)} \eta_ {\ell} \nabla_ {w _ {g, m} ^ {(r, k)}} f _ {m} (w _ {p, m} ^ {(r, k)}) \\ + \eta_ {\ell} ^ {2} \nabla_ {\psi_ {g, m} ^ {(r, k)}} f _ {m} (w _ {p, m} ^ {(r, k)}) \nabla_ {w _ {g, m} ^ {(r, k)}} f _ {m} (w _ {p, m} ^ {(r, k)}) + (1 - \psi_ {g, m} ^ {(r, k)}) w _ {\ell , m} ^ {(r, K)} \\ + w _ {\ell , m} ^ {(r, K)} \eta_ {\ell} \nabla_ {\psi_ {g, m} ^ {(r, k)}} f _ {m} \left(w _ {p, m} ^ {(r, k)}\right) - w _ {p, m} ^ {(r, 0)} \| ^ {2} \tag {146} \\ \end{array}
$$
$$
\begin{array}{l} = \mathbb {E} | | w _ {p, m} ^ {(r, k)} - w _ {p, m} ^ {(r, 0)} + (w _ {\ell , m} ^ {(r, K)} - w _ {g, m} ^ {(r, k)}) \eta_ {\ell} \nabla_ {\psi_ {g, m} ^ {(r, k)}} f _ {m} (w _ {p, m} ^ {(r, k)}) \\ \left(- \psi_ {g, m} ^ {(r, k)} + \eta_ {\ell} \nabla_ {\psi_ {g, m} ^ {(r, k)}} f _ {m} \left(w _ {p, m} ^ {(r, k)}\right)\right) \eta_ {\ell} \nabla_ {w _ {g, m} ^ {(r, k)}} f _ {m} \left(w _ {p, m} ^ {(r, k)}\right) \| ^ {2} \tag {147} \\ \end{array}
$$
$$
\leq 3 \mathbb {E} \| w _ {p, m} ^ {(r, k)} - w _ {p, m} ^ {(r, 0)} \| ^ {2} + 3 \mathbb {E} \| w _ {\ell , m} ^ {(r, K)} - w _ {g, m} ^ {(r, k)} \| ^ {2} + 3 \eta_ {\ell} ^ {2} \mathbb {E} \| \nabla_ {w _ {g, m} ^ {(r, k)}} f _ {m} (w _ {p, m} ^ {(r, k)}) \| ^ {2} \tag {148}
$$
The inequality was derived from the fact that  $\mathbb{E}|| - \eta_{\ell}\nabla_{\psi_{g,m}^{(r,k)}}f_m(w_{p,m}^{(r,k)})||^2 = \mathbb{E}||\psi_{g,m}^{(r,k + 1)} - \psi_{g,m}^{(r,k)}||^2\leq 1$  Unrolling the recursion across  $r\in [R]$ , then using Lemmas E.6 and E.10 and Assumption E.4,
$$
\begin{array}{l} \mathbb {E} \left\| w _ {p, m} ^ {(r, K)} - w _ {p, m} ^ {(r, 0)} \right\| ^ {2} \leq \sum_ {k = 1} ^ {K} \left(\left(1 + \frac {1}{K - 1}\right) \mathbb {E} \left\| w _ {\ell , m} ^ {(r, K)} - w _ {g, m} ^ {(r, k)} \right\| ^ {2} + K \eta_ {\ell} ^ {2} \mathbb {E} \left\| \nabla_ {w _ {g, m} ^ {(r, k)}} f _ {m} \left(w _ {p, m} ^ {(r, k)}\right) \right\| ^ {2}\right) (149) \\ \leq \left(6 K ^ {4} \eta_ {\ell} ^ {2} \mathbb {E} | | \nabla f _ {m} \left(w _ {g} ^ {(r)}\right) | | ^ {2} + 3 K ^ {3} \eta_ {\ell} ^ {2} \sigma_ {\ell} ^ {2} + 1 2 K ^ {2} \eta_ {\ell} ^ {2} \sigma_ {\ell} ^ {2} \mathbb {E} | | 1 - \psi_ {g, m} | | ^ {2} \right. \\ + 4 \left(1 + \frac {1}{K - 1}\right) K ^ {3} \eta_ {\ell} ^ {2} \mathbb {E} | | \nabla f _ {m} \left(w _ {g} ^ {(r)}\right) | | ^ {2} \mathbb {E} | | \psi_ {g, m} | | ^ {2}) \\ \cdot \sum_ {k = 1} ^ {K} \left(1 + \frac {1}{K - 1} + 1 2 K ^ {2} \eta_ {\ell} ^ {2} \beta\right) ^ {k} (150) \\ \end{array}
$$
Assuming  $\frac{1}{K - 1} \geq 12K^2\eta_\ell^2\beta \Rightarrow \eta_\ell \leq \frac{1}{K\sqrt{12(K - 1)\beta}}$
$$
\begin{array}{l} \mathbb {E} \| w _ {p, m} ^ {(r, K)} - w _ {p, m} ^ {(r, 0)} \| ^ {2} \leq \left(6 K ^ {4} \eta_ {\ell} ^ {2} \mathbb {E} \| \nabla f _ {m} \left(w _ {g} ^ {(r)}\right) \| ^ {2} + 3 K ^ {3} \eta_ {\ell} ^ {2} \sigma_ {\ell} ^ {2} + 1 2 K ^ {2} \eta_ {\ell} ^ {2} \mathbb {E} \| 1 - \psi_ {g, m} ^ {(r, k)} \| ^ {2} \right. \\ + 4 \left(1 + \frac {1}{K - 1}\right) K ^ {3} \eta_ {\ell} ^ {2} \mathbb {E} | | \nabla f _ {m} \left(w _ {g} ^ {(r)}\right) | | ^ {2} \mathbb {E} | | \psi_ {g, m} ^ {(r, k)} | | ^ {2} \sigma_ {\ell} ^ {2}\left. \right) 3 K (151) \\ = 1 8 K ^ {5} \eta_ {\ell} ^ {2} \mathbb {E} | | \nabla f _ {m} (w _ {g} ^ {(r)}) | | ^ {2} + 9 K ^ {4} \eta_ {\ell} ^ {2} \sigma_ {\ell} ^ {2} + 3 6 K ^ {3} \eta_ {\ell} ^ {2} \sigma_ {\ell} ^ {2} \mathbb {E} | | 1 - \psi_ {g, m} ^ {(r, k)} | | ^ {2} \\ + 2 4 K ^ {4} \eta_ {\ell} ^ {2} \mathbb {E} | | \nabla f _ {m} \left(w _ {g} ^ {(r)}\right) | | ^ {2} \mathbb {E} | | \psi_ {g, m} ^ {(r, k)} | | ^ {2} (152) \\ \end{array}
$$
Lemma E.19 (One round progress of personalized model). If  $m^{th}$  client's objective function  $f_{m}$  satisfies Assumptions E.3, E.4, in Algorithm 2, the following is satisfied:
$$
\begin{array}{l} \mathbb {E} \left\| w _ {p, m} ^ {(r + 1, K)} - w _ {p, m} ^ {(r, K)} \right\| ^ {2} \leq 7 2 \left(1 + \eta_ {\ell} ^ {2}\right) K ^ {3} \eta_ {\ell} ^ {2} \left(5 K \left(G ^ {2} + B ^ {2} \mathbb {E} \left\| \nabla F \left(w _ {g} ^ {(r)}\right) \right\| ^ {2}\right) + 1 2 \sigma_ {\ell} ^ {2}\right) \\ + 3 6 \left(K \sigma_ {\ell} ^ {2} \eta_ {\ell} ^ {2} + \left(\delta_ {m} ^ {\psi} + \frac {\delta^ {\psi}}{M}\right) K ^ {2} \eta_ {\ell} ^ {2}\right) \left(K \sigma_ {\ell} ^ {2} \eta_ {\ell} ^ {2} + \left(\delta_ {m} ^ {w _ {g}} + \frac {\delta^ {w _ {g}}}{M}\right) K ^ {2} \eta_ {\ell} ^ {2}\right) \\ + 1 2 \left(1 + \eta_ {\ell} ^ {2} K ^ {2} \beta^ {4}\right) \eta_ {\ell} ^ {2} K ^ {2} \left(K \sigma_ {\ell} ^ {2} \eta_ {\ell} ^ {2} + \left(\delta_ {m} ^ {\psi} + \frac {\delta^ {\psi}}{M}\right) K ^ {2} \eta_ {\ell} ^ {2}\right) \\ \left(G ^ {2} + B ^ {2} \mathbb {E} | | \nabla F (w _ {g} ^ {(r)}) | | ^ {2}\right) \\ \end{array}
$$