SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
\begin{array}{l} \mathbb {E} \left\| w _ {p, m} ^ {(r + 1, K)} - w _ {p, m} ^ {(r, K)} \right\| ^ {2} = \mathbb {E} \left\| w _ {p, m} ^ {(r + 1, K)} - w _ {p, m} ^ {(r + 1, 0)} + w _ {p, m} ^ {(r + 1, 0)} - w _ {p, m} ^ {(r, K)} \right\| ^ {2} (153) \\ \leq 2 \mathbb {E} \| w _ {p, m} ^ {(r + 1, K)} - w _ {p, m} ^ {(r + 1, 0)} \| ^ {2} + 2 \mathbb {E} \| w _ {p, m} ^ {(r + 1, 0)} - w _ {p, m} ^ {(r, K)} \| ^ {2} (154) \\ \end{array}
$$
Using the Lemmas E.18 and E.14, we proceed as
$$
\begin{array}{l} \leq 7 2 \left(1 + \eta_ {\ell} ^ {2}\right) K ^ {3} \eta_ {\ell} ^ {2} \left(5 K \left(G ^ {2} + B ^ {2} \mathbb {E} | | \nabla F \left(w _ {g} ^ {(r)}\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 \left(w _ {g} ^ {(r)}\right) | | ^ {2}\right) \tag {155} \\ \end{array}
$$
Theorem E.20 (Convergence of the Personalized Model for Non-convex Cases). If each client's objective function  $f_{m}$  satisfies Assumptions E.2, E.3, E.4 using the learning rate  $\eta_{\ell} \leq \frac{1}{K\sqrt{12\beta}}$  in Algorithm 2, then the following convergence holds:
$$
\begin{array}{l} \frac {1}{R} \sum_ {r = 1} ^ {R} \mathbb {E} | | \nabla f _ {m} \left(w _ {p, m} ^ {(r, K)}\right) | | ^ {2} \leq \frac {2}{R} \left(\mathbb {E} \left[ f _ {m} \left(w _ {p, m} ^ {(1, K)}\right) \right] - \mathbb {E} \left[ f _ {m} \left(w _ {p, m} ^ {(R, K)}\right) \right]\right) \\ + 6 \left(1 + \eta_ {\ell} ^ {2}\right) K \left(5 K \left(G ^ {2} + B ^ {2} \frac {1}{R} \sum_ {r = 1} ^ {R} \mathbb {E} | | \nabla F \left(w _ {g} ^ {(r)}\right) | | ^ {2}\right) + 1 2 \sigma_ {\ell} ^ {2}\right) \\ + 3 K \eta_ {\ell} ^ {2} \left(\sigma_ {\ell} ^ {2} + \left(\delta_ {m} ^ {\psi} + \frac {\delta^ {\psi}}{M}\right) K\right) \left(\sigma_ {\ell} ^ {2} + \left(\delta_ {m} ^ {w _ {g}} + \frac {\delta^ {w _ {g}}}{M}\right) K\right) \\ + \left(1 + \eta_ {\ell} ^ {2} K ^ {2} \beta^ {4}\right) \eta_ {\ell} ^ {2} K \left(\sigma_ {\ell} ^ {2} + \left(\delta_ {m} ^ {\psi} + \frac {\delta^ {\psi}}{M}\right) K\right) \left(G ^ {2} + B ^ {2} \frac {1}{R} \sum_ {r = 1} ^ {R} \mathbb {E} | | \nabla F (w _ {g} ^ {(r)}) | | ^ {2}\right) \\ \end{array}
$$
Proof. According to the update rule of Equation 10 and  $\beta$ -smoothness of  $f_{m}$ , we have,
$$
\left. f _ {m} \left(w _ {p, m} ^ {(r + 1, K)}\right) \leq f _ {m} \left(w _ {p, m} ^ {(r, K)}\right) + \left\langle \nabla f _ {m} \left(w _ {p, m} ^ {(r, K)}\right), w _ {p, m} ^ {(r + 1, K)} - w _ {p, m} ^ {(r, K)} \right\rangle + \frac {\beta}{2} \left\| w _ {p, m} ^ {(r + 1, K)} - w _ {p, m} ^ {(r, K)} \right\| ^ {2} \right. \tag {156}
$$
$$
\begin{array}{l} \mathbb {E} \left[ f _ {m} \left(w _ {p, m} ^ {(r + 1, K)}\right) \right] \leq \mathbb {E} \left[ f _ {m} \left(w _ {p, m} ^ {(r, K)}\right) \right] + \frac {\beta}{2} \mathbb {E} \| w _ {p, m} ^ {(r + 1, K)} - w _ {p, m} ^ {(r, K)} \| ^ {2} \\ + \mathbb {E} \left\langle \nabla f _ {m} \left(w _ {p, m} ^ {(r, K)}\right), w _ {p, m} ^ {(r + 1, K)} - w _ {p, m} ^ {(r, K)} \right\rangle \tag {157} \\ \end{array}
$$
Using  $\langle a,b\rangle = \frac{1}{2} ||a||^2 +\frac{1}{2} ||b||^2 -\frac{1}{2} ||a - b||^2$
$$
\begin{array}{l} \mathbb {E} \left[ f _ {m} \left(w _ {p, m} ^ {(r + 1, K)}\right) \right] \leq \mathbb {E} \left[ f _ {m} \left(w _ {p, m} ^ {(r, K)}\right) \right] + \frac {1}{2} \mathbb {E} | | \nabla f _ {m} \left(w _ {p, m} ^ {(r, K)}\right) | | ^ {2} + \left(\frac {\beta + 1}{2}\right) \mathbb {E} | | w _ {p, m} ^ {(r + 1, K)} - w _ {p, m} ^ {(r, K)} | | ^ {2} \\ - \frac {1}{2} \mathbb {E} | | \nabla f _ {m} \left(w _ {p, m} ^ {(r, K)}\right) - \left(w _ {p, m} ^ {(r + 1, K)} - w _ {p, m} ^ {(r, K)}\right) | | ^ {2} (158) \\ \leq \mathbb {E} \left[ f _ {m} \left(w _ {p, m} ^ {(r, K)}\right) \right] + \frac {1}{2} \mathbb {E} | | \nabla f _ {m} \left(w _ {p, m} ^ {(r, K)}\right) | | ^ {2} + \left(\frac {\beta + 1}{2}\right) \mathbb {E} | | w _ {p, m} ^ {(r + 1, K)} - w _ {p, m} ^ {(r, K)} | | ^ {2} \\ - \mathbb {E} \| \nabla f _ {m} \left(w _ {p, m} ^ {(r, K)}\right) \| ^ {2} - \mathbb {E} \| \left(w _ {p, m} ^ {(r + 1, K)} - w _ {p, m} ^ {(r, K)}\right) \| ^ {2} (159) \\ \leq \mathbb {E} \left[ f _ {m} \left(w _ {p, m} ^ {(r, K)}\right) \right] - \frac {1}{2} \mathbb {E} \left| \left| \nabla f _ {m} \left(w _ {p, m} ^ {(r, K)}\right) \right| \right| ^ {2} + \left(\frac {\beta - 1}{2}\right) \mathbb {E} \left| \left| w _ {p, m} ^ {(r + 1, K)} - w _ {p, m} ^ {(r, K)} \right| \right| ^ {2} (160) \\ \end{array}
$$
Rearranging the terms to put  $\frac{1}{2}\mathbb{E}||\nabla f_m(w_{p,m}^{(r,K)})||^2$  at LHS,
$$
\frac {1}{2} \mathbb {E} \| \nabla f _ {m} \left(w _ {p, m} ^ {(r, K)}\right) \| ^ {2} \leq \mathbb {E} \left[ f _ {m} \left(w _ {p, m} ^ {(r, K)}\right) \right] - \mathbb {E} \left[ f _ {m} \left(w _ {p, m} ^ {(r + 1, K)}\right) \right] + \left(\frac {\beta - 1}{2}\right) \underbrace {\mathbb {E} \| w _ {p , m} ^ {(r + 1 , K)} - w _ {p , m} ^ {(r , K)} \| ^ {2}} _ {\text {L e m m a E . 1 9}} \tag {161}
$$
$$
\begin{array}{l} \left. \mathbb {E} \left| \left| \nabla f _ {m} \left(w _ {p, m} ^ {(r, K)}\right) \right| \right| ^ {2} \leq 2 \left(\mathbb {E} \left[ f _ {m} \left(w _ {p, m} ^ {(r, K)}\right) \right] - \mathbb {E} \left[ f _ {m} \left(w _ {p, m} ^ {(r + 1, K)}\right) \right]\right) \right. \\ + 7 2 \beta \left(1 + \eta_ {\ell} ^ {2}\right) K ^ {3} \eta_ {\ell} ^ {2} \left(5 K \left(G ^ {2} + B ^ {2} \mathbb {E} | | \nabla F \left(w _ {g} ^ {(r)}\right) | | ^ {2}\right) + 1 2 \sigma_ {\ell} ^ {2}\right) \\ + 3 6 \beta \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 \beta \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) \\ \cdot \left(G ^ {2} + B ^ {2} \mathbb {E} \| \nabla F \left(w _ {g} ^ {(r)}\right) \| ^ {2}\right) \tag {162} \\ \end{array}
$$
Taking an average over all the rounds  $r \in [R]$ ,
$$
\begin{array}{l} \frac {1}{R} \sum_ {r = 1} ^ {R} \mathbb {E} | | \nabla f _ {m} \left(w _ {p, m} ^ {(r, K)}\right) | | ^ {2} \leq \frac {2}{R} \left(\mathbb {E} \left[ f _ {m} \left(w _ {p, m} ^ {(1, K)}\right) \right] - \mathbb {E} \left[ f _ {m} \left(w _ {p, m} ^ {(R, K)}\right) \right]\right) \\ + 7 2 \beta \left(1 + \eta_ {\ell} ^ {2}\right) K ^ {3} \eta_ {\ell} ^ {2} \left(5 K \left(G ^ {2} + B ^ {2} \frac {1}{R} \sum_ {r = 1} ^ {R} \mathbb {E} | | \nabla F \left(w _ {g} ^ {(r)}\right) | | ^ {2}\right) + 1 2 \sigma_ {\ell} ^ {2}\right) \\ + 3 6 \beta K ^ {2} \eta_ {\ell} ^ {4} \left(\sigma_ {\ell} ^ {2} + \left(\delta_ {m} ^ {\psi} + \frac {\delta^ {\psi}}{M}\right) K\right) \left(\sigma_ {\ell} ^ {2} + \left(\delta_ {m} ^ {w _ {g}} + \frac {\delta^ {w _ {g}}}{M}\right) K\right) \\ + 1 2 \beta \left(1 + \eta_ {\ell} ^ {2} K ^ {2} \beta^ {4}\right) \eta_ {\ell} ^ {4} K ^ {3} \left(\sigma_ {\ell} ^ {2} + \left(\delta_ {m} ^ {\psi} + \frac {\delta^ {\psi}}{M}\right) K\right) \left(G ^ {2} + B ^ {2} \frac {1}{R} \sum_ {r = 1} ^ {R} \mathbb {E} | | \nabla F \left(w _ {g} ^ {(r)}\right) | | ^ {2}\right) \tag {163} \\ \end{array}
$$
Assuming  $12K^{2}\eta_{\ell}^{2}\beta \leq 1\leq 1\Rightarrow \eta_{\ell}\leq \frac{1}{K\sqrt{12\beta}}$
$$
\begin{array}{l} \frac {1}{R} \sum_ {r = 1} ^ {R} \mathbb {E} | | \nabla f _ {m} \left(w _ {p, m} ^ {(r, K)}\right) | | ^ {2} \leq \frac {2}{R} \left(\mathbb {E} \left[ f _ {m} \left(w _ {p, m} ^ {(1, K)}\right) \right] - \mathbb {E} \left[ f _ {m} \left(w _ {p, m} ^ {(R, K)}\right) \right]\right) \\ + 6 \left(1 + \eta_ {\ell} ^ {2}\right) K \left(5 K \left(G ^ {2} + B ^ {2} \frac {1}{R} \sum_ {r = 1} ^ {R} \mathbb {E} | | \nabla F \left(w _ {g} ^ {(r)}\right) | | ^ {2}\right) + 1 2 \sigma_ {\ell} ^ {2}\right) \\ + 3 K \eta_ {\ell} ^ {2} \left(\sigma_ {\ell} ^ {2} + \left(\delta_ {m} ^ {\psi} + \frac {\delta^ {\psi}}{M}\right) K\right) \left(\sigma_ {\ell} ^ {2} + \left(\delta_ {m} ^ {w _ {g}} + \frac {\delta^ {w _ {g}}}{M}\right) K\right) \\ + \left(1 + \eta_ {\ell} ^ {2} K ^ {2} \beta^ {4}\right) \eta_ {\ell} ^ {2} K \left(\sigma_ {\ell} ^ {2} + \left(\delta_ {m} ^ {\psi} + \frac {\delta^ {\psi}}{M}\right) K\right) \left(G ^ {2} + B ^ {2} \frac {1}{R} \sum_ {r = 1} ^ {R} \mathbb {E} | | \nabla F \left(w _ {g} ^ {(r)}\right) | | ^ {2}\right) \tag {164} \\ \end{array}
$$
Plugging in Theorem E.12 to get bounds on  $\sum_{r=1}^{R} \mathbb{E}||\nabla F(w_g^{(r)})|^2$  would get us bounds on  $\frac{1}{R}\sum_{r=1}^{R}\mathbb{E}||\nabla f_m(w_{p,m}^{(r,K)})||^2$ .