SlowGuess/NeurIPS_2023_OCR huggingface dataset:
The global distribution is defined as  $\mathcal{D} = \sum_{m\in [M]}q_m\mathcal{D}_m$  where  $q_{m}$  is the weight associated with  $m^{th}$  client and  $\sum_{m\in [M]}q_m = 1$
Note that  $w_{p,m}$  is a combination of outputs of  $w_{g,m}$  (Global parameters) and  $w_{\ell,m}$  (Local parameters) on each layer. For tractability of analysis, we will assume that the combination is only after the last layer. Hence,
$$
w _ {p, m} (x _ {m}) \leftarrow \psi_ {g, m} (x _ {m}) w _ {g, m} (x _ {m}) + (1 - \psi_ {g, m} (x _ {m})) w _ {\ell , m} (x _ {m}).
$$
The local model update rule is,
$$
w _ {\ell , m} ^ {(r, k)} \leftarrow w _ {\ell , m} ^ {(r, k - 1)} - \eta_ {\ell} \nabla f _ {m} \left(w _ {\ell , m} ^ {(r, k - 1)} \left(x _ {m}\right), y _ {m}\right)
$$
where  $w_{\ell,m}^{(r,0)} = w_{g,m}^{(r,0)} = w_g^{(r)}$ . Indices  $r \in [R]$  and  $k \in [K]$  are the global round and the local epoch indices.
The policy update rule is,
$$
\psi_ {g, m} ^ {(r, k)} \leftarrow \psi_ {g, m} ^ {(r, k - 1)} - \eta_ {\ell} \nabla_ {\psi_ {g}} f _ {m} \left(w _ {p, m} ^ {(r, k - 1)} \left(x _ {m}\right), y _ {m}\right).
$$
The global model update rule is,
$$
w _ {g, m} ^ {(r, k)} \leftarrow w _ {g, m} ^ {(r, k - 1)} - \eta_ {\ell} \nabla_ {w _ {g}} f _ {m} \big (w _ {p, m} ^ {(r, k - 1)} (x _ {m}), y _ {m} \big).
$$
We list out all the optimization problems relevant to Flow:
- Local true risk of the personalized model
$$
F _ {m} \left(w _ {p, m}\right) := \mathbb {E} _ {\left(x _ {m}, y _ {m} \sim \mathcal {D} _ {m}\right)} \left[ f _ {m} \left(w _ {p, m} \left(x _ {m}\right), y _ {m}\right) \right]
$$
where  $f_{m}$  is a loss function associated with the  $m^{th}$  client.
- Local empirical risk of the personalized model
$$
\hat {F} _ {m} (w _ {p, m}) := \frac {1}{n _ {m}} \sum_ {i \in [ n _ {m} ]} f _ {m} (w _ {p, m} (x _ {m} ^ {(i)}), y _ {m} ^ {(i)})
$$
- Local true risk of the global model
$$
F _ {m} \left(w _ {g, m}\right) := \mathbb {E} _ {\left(x _ {m}, y _ {m} \sim \mathcal {D} _ {m}\right)} \left[ f _ {m} \left(w _ {g, m} \left(x _ {m}\right), y _ {m}\right) \right]
$$
- Local empirical risk of the global model
$$
\hat {F} _ {m} (w _ {g, m}) := \frac {1}{n _ {m}} \sum_ {i \in [ n _ {m} ]} f _ {m} (w _ {g, m} (x _ {m} ^ {(i)}), y _ {m} ^ {(i)})
$$
- Local minimizer of local empirical risk of the personalized model
$$
w _ {p, m} ^ {*} \in \mathcal {H} \text {s u c h t h a t} \hat {F} _ {m} \left(w _ {p, m}\right) \geq \hat {F} _ {m} \left(w _ {p, m} ^ {*}\right); \forall w _ {p, m} \in \mathcal {H}, \exists \epsilon > 0, \| w _ {p, m} - w _ {p, m} ^ {*} \| <   \epsilon
$$
- Global true risk of the global model
$$
F (w _ {g}) = \frac {1}{n M} \sum_ {m \in [ M ]} n _ {m} \mathbb {E} _ {(x _ {m}, y _ {m}) \sim \mathcal {S} _ {m}} \left[ f _ {m} \left(w _ {g, m} \left(x _ {m}\right), y _ {m}\right) \right] \text {w h e r e} n = | \mathcal {S} | = \left| \bigcup_ {m \in [ M ]} \mathcal {S} _ {m} \right|
$$
- Global empirical risk of the global model
$$
\hat {F} (w _ {g}) = \frac {1}{n M} \sum_ {m \in [ M ]} n _ {m} \hat {F} _ {m} (w _ {g, m} (x _ {m}), y _ {m}) = \frac {1}{n M} \sum_ {m \in [ M ]} \sum_ {i \in [ n _ {m} ]} f _ {m} (w _ {g, m} (x _ {m} ^ {(i)}), y _ {m} ^ {(i)})
$$
- Local minimizer of global empirical risk
$$
w _ {g} ^ {*} \in \mathcal {H} \text {s u c h t h a t} \hat {F} \left(w _ {g}\right) \geq \hat {F} \left(w _ {g} ^ {*}\right); \forall w _ {g} \in \mathcal {H}, \exists \epsilon > 0, | | w _ {g} - w _ {g} ^ {*} | | <   \epsilon
$$
We also use the following assumptions similar to [11, 18, 6]:
Assumption E.1 (Strong Convexity).  $f_{m}$  is  $\mu$ -convex for  $\mu \geq 0$ . Hence,
$$
\langle \nabla f _ {m} (w), v - w \rangle \leq f _ {m} (v) - f _ {m} (w) - \frac {\mu}{2} | | w - v | | ^ {2}, \forall m \in [ M ] \text {a n d} w, v \in \mathcal {H}.
$$
We also generalize our convergence analysis for  $\mu = 0$ , general convex cases.
Assumption E.2 (Smoothness). The gradient of  $f_{m}$  is  $\beta$ -Lipschitz,
$$
\left| \left| \nabla f _ {m} (w) - \nabla f _ {m} (v) \right| \right| \leq \beta | | w - v | |, \forall m \in [ M ] \text {a n d} w, v \in \mathcal {H}.
$$
Assumption E.3 (Bounded Local Variance).  $h_m(w) \coloneqq \nabla f_m(w(x_m), y_m)$  is an unbiased stochastic gradient of  $f_m$  with variance bounded by  $\sigma_\ell^2$ .
$$
\mathbb {E} _ {(x _ {m}, y _ {m} \sim \mathcal {D} _ {m})} | | h _ {m} (w) - \nabla f _ {m} (w (x _ {m}), y _ {m}) | | ^ {2} \leq \sigma_ {\ell} ^ {2}, \forall m \in [ M ] \text {a n d} w \in \mathcal {H}.
$$
Assumption E.4  $((G,B)$ -Bounded Gradient Dissimilarity). There exists constants  $G \geq 0$  and  $B \geq 1$  such that
$$
\frac {1}{M} \sum_ {m \in [ M ]} | | \nabla f _ {m} (w) | | ^ {2} \leq G ^ {2} + 2 \beta B ^ {2} (F (w) - F (w ^ {*}))
$$
for a convex  $f_{m}$ . And for a non-convex  $f_{m}$ ,
$$
\frac {1}{M} \sum_ {m \in [ M ]} | | \nabla f _ {m} (w) | | ^ {2} \leq G ^ {2} + B ^ {2} | | \nabla F (w) | | ^ {2}.
$$
The derivation is given in Section D.1 of Scaffold [18].
We also use a definition to quantify the diversity of a client's gradient with respect to the global gradient as defined in [29]:
Definition E.5 (Gradient Diversity). The difference between gradients of the  $m^{th}$  client's true risk and the global true risk based on the global model  $w$  is,
$$
\delta_ {m} = \sup _ {w \in \mathcal {H}} | | \nabla f _ {m} (w) - \nabla F (w) | | ^ {2}
$$
E.3 Convergence Proof for the Global Model: Convex (Strong and General) Cases
A client's local update for one local epoch on the global model, starting with  $w_{g,m}^{(r,0)} \gets w_g^{(r)}$ , is
$$
w _ {g, m} ^ {(r, k + 1)} = w _ {g, m} ^ {(r, k)} - \eta_ {\ell} h _ {m} \left(w _ {p, m} ^ {(r, k)}\right). \tag {8}
$$
And a client's local update for  $K$  epochs on the global model, would be
$$
\begin{array}{l} w _ {g, m} ^ {(r, K)} = w _ {g, m} ^ {(r, 0)} - \eta_ {\ell} \sum_ {k = 1} ^ {K} h _ {m} \left(w _ {p, m} ^ {(r, k - 1)}\right) (9) \\ = w _ {g, m} ^ {(r, 0)} - \eta_ {\ell} \sum_ {k = 1} ^ {K} h _ {m} \left(\psi_ {g, m} ^ {(r, k)} \left(x _ {m}\right) w _ {g, m} ^ {(r, k - 1)} \left(x _ {m}\right) + \left(1 - \psi_ {g, m} ^ {(r, k)} \left(x _ {m}\right)\right) w _ {\ell , m} ^ {(r, K)} \left(x _ {m}\right), y _ {m}\right). (10) \\ \end{array}
$$
In both the above cases, the gradient is with respect to  $w_{g}$  parameters.
The global model update is,
$$
w _ {g} ^ {(r + 1)} = \frac {1}{n M} \sum_ {m \in [ M ]} n _ {m} w _ {g, m} ^ {(r, K)} \tag {11}
$$
We first start with a lemma which binds the deviation between the local model  $w_{\ell,m}^{(r,K)}$  and the global model starting point  $w_g^{(r)}$  for it at round  $r$ .
Lemma E.6 (Local model progress). If  $m^{th}$  client's objective function  $f_{m}$  satisfies Assumptions E.2, E.3, and condition  $\eta_{\ell} \leq \frac{1}{\beta \sqrt{2K(K - 1)}}$  in Algorithm 2, the following is satisfied:
$$
\mathbb {E} | | w _ {\ell , m} ^ {(r, K)} - w _ {\ell , m} ^ {(r, 0)} | | ^ {2} \leq 6 K ^ {2} \eta_ {\ell} ^ {2} \mathbb {E} | | \nabla f _ {m} (w _ {g} ^ {(r)}) | | ^ {2} + 3 K \eta_ {\ell} ^ {2} \sigma_ {\ell} ^ {2}
$$
$$
\begin{array}{l} \mathbb {E} \left| \left| w _ {\ell , m} ^ {(r, K)} - w _ {\ell , m} ^ {(r, 0)} \right| \right| ^ {2} = \mathbb {E} \left| \left| w _ {\ell , m} ^ {(r, K - 1)} - \eta_ {\ell} \nabla f _ {m} \left(w _ {\ell , m} ^ {(r, K - 1)}\right) - w _ {\ell , m} ^ {(r, 0)} \right| \right| ^ {2} (12) \\ \leq \left(1 + \frac {1}{K - 1}\right) \mathbb {E} \| w _ {\ell , m} ^ {(r, K - 1)} - w _ {\ell , m} ^ {(r, 0)} \| ^ {2} + K \eta_ {\ell} ^ {2} \mathbb {E} \| \nabla f _ {m} \left(w _ {\ell , m} ^ {(r, K - 1)}\right) \| ^ {2} + \eta_ {\ell} ^ {2} \sigma_ {\ell} ^ {2} (13) \\ \end{array}
$$
Here we have used triangle inequality and variance separation.
$$
\begin{array}{l} \leq \left(1 + \frac {1}{K - 1}\right) \mathbb {E} | | w _ {\ell , m} ^ {(r, K - 1)} - w _ {\ell , m} ^ {(r, 0)} | | ^ {2} + K \eta_ {\ell} ^ {2} \mathbb {E} | | \nabla f _ {m} \left(w _ {\ell , m} ^ {(r, K - 1)}\right) | | ^ {2} + \eta_ {\ell} ^ {2} \sigma_ {\ell} ^ {2} (14) \\ \leq \left(1 + \frac {1}{K - 1}\right) \mathbb {E} | | w _ {\ell , m} ^ {(r, K - 1)} - w _ {\ell , m} ^ {(r, 0)} | | ^ {2} + \eta_ {\ell} ^ {2} \sigma_ {\ell} ^ {2} \\ + K \eta_ {e} ^ {2} \mathbb {E} \| \nabla f _ {m} \left(w _ {\ell , m} ^ {(r, K - 1)}\right) - \nabla f _ {m} \left(w _ {\ell , m} ^ {(r, 0)}\right) + \nabla f _ {m} \left(w _ {\ell , m} ^ {(r, 0)}\right) \| ^ {2} (15) \\ \end{array}
$$
$$
\begin{array}{l} \leq \left(1 + \frac {1}{K - 1}\right) \mathbb {E} | | w _ {\ell , m} ^ {(r, K - 1)} - w _ {\ell , m} ^ {(r, 0)} | | ^ {2} + 2 K \eta_ {\ell} ^ {2} \mathbb {E} | | \nabla f _ {m} (w _ {\ell , m} ^ {(r, K - 1)}) - \nabla f _ {m} (w _ {\ell , m} ^ {(r, 0)}) | | ^ {2} \\ + 2 K \eta_ {\ell} ^ {2} \mathbb {E} \| \nabla f _ {m} \left(w _ {\ell , m} ^ {(r, 0)}\right) \| ^ {2} + \eta_ {\ell} ^ {2} \sigma_ {\ell} ^ {2} \tag {16} \\ \end{array}
$$
$$
\begin{array}{l} \leq \left(1 + \frac {1}{K - 1}\right) \mathbb {E} | | w _ {\ell , m} ^ {(r, K - 1)} - w _ {\ell , m} ^ {(r, 0)} | | ^ {2} + 2 K \beta^ {2} \eta_ {\ell} ^ {2} \mathbb {E} | | w _ {\ell , m} ^ {(r, K - 1)} - w _ {\ell , m} ^ {(r, 0)} | | ^ {2} \\ + 2 K \eta_ {\ell} ^ {2} \mathbb {E} \| \nabla f _ {m} \left(w _ {\ell , m} ^ {(r, 0)}\right) \| ^ {2} + \eta_ {\ell} ^ {2} \sigma_ {\ell} ^ {2} \tag {17} \\ \end{array}
$$
Assuming  $\eta_{\ell} \leq \frac{1}{\beta \sqrt{2K(K - 1)}}$ , we get
$$
\begin{array}{l} \mathbb {E} | | w _ {\ell , m} ^ {(r, K)} - w _ {\ell , m} ^ {(r, 0)} | | ^ {2} \leq \left(1 + \frac {2}{K - 1}\right) \mathbb {E} | | w _ {\ell , m} ^ {(r, K - 1)} - w _ {\ell , m} ^ {(r, 0)} | | ^ {2} \\ + 2 K \eta_ {\ell} ^ {2} \mathbb {E} \| \nabla f _ {m} \left(w _ {\ell , m} ^ {(r, 0)}\right) \| ^ {2} + \eta_ {\ell} ^ {2} \sigma_ {\ell} ^ {2} \tag {18} \\ \end{array}
$$
Unrolling the above recursion,
$$
\begin{array}{l} \mathbb {E} \left| \left| w _ {\ell , m} ^ {(r, K)} - w _ {\ell , m} ^ {(r, 0)} \right| \right| ^ {2} \leq \sum_ {i = 1} ^ {K} \left(2 K \eta_ {\ell} ^ {2} \mathbb {E} \left| \left| \nabla f _ {m} \left(w _ {\ell , m} ^ {(r, 0)}\right) \right| \right| ^ {2} + \eta_ {\ell} ^ {2} \sigma_ {\ell} ^ {2}\right) \left(1 + \frac {2}{K - 1}\right) ^ {i} (19) \\ \leq 3 K \left(2 K \eta_ {\ell} ^ {2} \mathbb {E} \| \nabla f _ {m} \left(w _ {\ell , m} ^ {(r, 0)}\right) \| ^ {2} + \eta_ {\ell} ^ {2} \sigma_ {\ell} ^ {2}\right) (20) \\ = 6 K ^ {2} \eta_ {\ell} ^ {2} \mathbb {E} \| \nabla f _ {m} \left(w _ {\ell , m} ^ {(r, 0)}\right) \| ^ {2} + 3 K \eta_ {\ell} ^ {2} \sigma_ {\ell} ^ {2} (21) \\ \end{array}
$$
Now we move forward to a lemma which binds the deviation between the local version of the global model  $w_{g,m}^{(r,k)}$  and the global model starting point  $w_{g}^{(r)}$  for it round  $r$ .
Lemma E.7 (Local version of the global model progress). If  $m^{th}$  client's objective function  $f_{m}$  satisfies Assumptions E.1, E.2, E.3, and condition  $\eta_{\ell} \leq \frac{1}{\beta \sqrt{2k}}$  in Algorithm 2, the following is satisfied:
$$
\mathbb {E} | | w _ {g, m} ^ {(r, k)} - w _ {g, m} ^ {(r, 0)} | | ^ {2} \leq 8 k ^ {3} \eta_ {\ell} ^ {2} \mathbb {E} | | \psi_ {g, m} ^ {(r, k)} | | ^ {2} \mathbb {E} | | \nabla f _ {m} (w _ {g} ^ {(r)}) | | ^ {2} + 4 k \eta_ {\ell} ^ {2} \sigma_ {\ell} ^ {2}
$$
Proof. We start by expanding  $w_{g,m}^{(r,k)}$  in terms of its previous epoch iterate.
$$
\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 {22}
$$
Using triangle inequality and separation of variance, we get,
$$
\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} (w _ {p, m} ^ {(r, k - 1)}) | | ^ {2} + \eta_ {\ell} ^ {2} \sigma_ {\ell} ^ {2} \tag {23}
$$
Using the convex property of  $f_{m}$ , we get
$$
\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} | | \psi_ {g, m} ^ {(r, k)} \nabla f _ {m} (w _ {g, m} ^ {(r, k - 1)}) | | ^ {2} + \eta_ {\ell} ^ {2} \sigma_ {\ell} ^ {2} (24) \\ \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} | | \psi_ {g, m} ^ {(r, k)} (\nabla f _ {m} (w _ {g, m} ^ {(r, k - 1)}) - \nabla f _ {m} (w _ {g, m} ^ {(r, 0)}) + \nabla f _ {m} (w _ {g, m} ^ {(r, 0)})) | | ^ {2} (25) \\ \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} | | \psi_ {g, m} ^ {(r, k)} | | ^ {2} \mathbb {E} | | \nabla f _ {m} (w _ {g, m} ^ {(r, 0)})) | | ^ {2} \\ + 2 k \eta_ {\ell} ^ {2} \mathbb {E} \left\| \psi_ {g, m} ^ {(r, k)} \right\| ^ {2} \mathbb {E} \left\| \left(\nabla f _ {m} \left(w _ {g, m} ^ {(r, k - 1)}\right) - \nabla f _ {m} \left(w _ {g, m} ^ {(r, 0)}\right)\right) \right\| ^ {2} (26) \\ \leq \left(1 + \frac {1}{k - 1} + 2 k \eta_ {\ell} ^ {2} \beta^ {2} \mathbb {E} | | \psi_ {g, m} ^ {(r, k)} | | ^ {2}\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} \left\| \psi_ {g, m} ^ {(r, k)} \right\| ^ {2} \mathbb {E} \left\| \nabla f _ {m} \left(w _ {g, m} ^ {(r, 0)}\right) \right\| ^ {2} (27) \\ \end{array}
$$
Unrolling the recursion,
$$
\begin{array}{l} \mathbb {E} | | w _ {g, m} ^ {(r, k)} - w _ {g, m} ^ {(r, 0)} | | ^ {2} \leq \sum_ {i = 1} ^ {k} \left( \right.\left(2 k \eta_ {\ell} ^ {2} \mathbb {E} | | \psi_ {g, m} ^ {(r, k)} | | ^ {2} \mathbb {E} | | \nabla f _ {m} (w _ {g, m} ^ {(r, 0)}) | | ^ {2} + \eta_ {\ell} ^ {2} \sigma_ {\ell} ^ {2}\right) \\ \left. \cdot \left(1 + \frac {1}{k - 1} + 2 k \eta_ {\ell} ^ {2} \beta^ {2} \mathbb {E} | | \psi_ {g, m} ^ {(r, k)} | | ^ {2}\right) ^ {i}\right) \tag {28} \\ \end{array}
$$
Assuming that  $\eta_{\ell} \leq \frac{1}{\beta \sqrt{2k}}$  we get  $k \eta_{\ell}^{2} \beta^{2} \leq 1$ ,
$$
\begin{array}{l} \mathbb {E} \left| \left| w _ {g, m} ^ {(r, k)} - w _ {g, m} ^ {(r, 0)} \right| \right| ^ {2} \leq \left(2 k \eta_ {\ell} ^ {2} \sum_ {i = 1} ^ {k} \mathbb {E} \left| \left| \psi_ {g, m} ^ {(r, i)} \right| \right| ^ {2} \mathbb {E} \left| \left| \nabla f _ {m} \left(w _ {g, m} ^ {(r, 0)}\right) \right| \right| ^ {2} + \eta_ {\ell} ^ {2} \sigma_ {\ell} ^ {2}\right) \sum_ {i = 1} ^ {k} \left(1 + \frac {1}{k - 1} + 2\right) ^ {i} (29) \\ \leq 4 k \left(2 k ^ {2} \eta_ {\ell} ^ {2} \mathbb {E} | | \psi_ {g, m} ^ {(r, k)} | | ^ {2} \mathbb {E} | | \nabla f _ {m} \left(w _ {g} ^ {(r)}\right) | | ^ {2} + \eta_ {\ell} ^ {2} \sigma_ {\ell} ^ {2}\right) (30) \\ \leq 8 k ^ {3} \eta_ {\ell} ^ {2} \mathbb {E} \| \psi_ {g, m} ^ {(r, k)} \| ^ {2} \mathbb {E} \| \nabla f _ {m} \left(w _ {g} ^ {(r)}\right) \| ^ {2} + 4 k \eta_ {\ell} ^ {2} \sigma_ {\ell} ^ {2} (31) \\ \end{array}
$$
Lemma E.8 (Deviation of the personalized model from the global model). If  $m^{th}$  client's objective function  $f_{m}$  satisfies Assumptions E.1, E.2, E.3, and condition  $\eta_{\ell} \leq \min \left( \frac{1}{\beta \sqrt{2K(K - 1)}}, \frac{1}{\beta \sqrt{2K}} \right)$  in Algorithm 2, the following is satisfied:
$$
\begin{array}{l} \mathbb {E} \| w _ {p, m} ^ {(r, k)} - w _ {g, m} ^ {(r, 0)} \| ^ {2} \leq 1 6 k ^ {3} \eta_ {\ell} ^ {2} \mathbb {E} \| 1 - \psi_ {g, m} ^ {(r, k)} \| ^ {2} \mathbb {E} \| \nabla f _ {m} (w _ {g} ^ {(r)}) \| ^ {2} + 8 k \eta_ {\ell} ^ {2} \sigma_ {\ell} ^ {2} \mathbb {E} \| \psi_ {g, m} ^ {(r, k)} \| ^ {2} \\ + 1 2 K ^ {2} \eta_ {\ell} ^ {2} \mathbb {E} | | 1 - \psi_ {g, m} ^ {(r, k)} | | ^ {2} \mathbb {E} | | \nabla f _ {m} (w _ {g} ^ {(r)}) | | ^ {2} + 6 K \eta_ {\ell} ^ {2} \sigma_ {\ell} ^ {2} \mathbb {E} | | \psi_ {g, m} ^ {(r, k)} | | ^ {2} \\ \end{array}
$$
$$
\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} (32) \\ = \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} (33) \\ = \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} (34) \\ \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} (35) \\ \end{array}
$$
Using lemmas E.6 and E.7,
$$
\begin{array}{l} \mathbb {E} | | w _ {p, m} ^ {(r, k)} - w _ {g, m} ^ {(r, 0)} | | ^ {2} \leq 2 \mathbb {E} | | \psi_ {g, m} ^ {(r, k)} | | ^ {2} \left(8 k ^ {3} \eta_ {\ell} ^ {2} \mathbb {E} | | \psi_ {g, m} ^ {(r, k)} | | ^ {2} \mathbb {E} | | \nabla f _ {m} (w _ {g} ^ {(r)}) | | ^ {2} + 6 K ^ {2} \eta_ {\ell} ^ {2} \mathbb {E} | | \nabla f _ {m} (w _ {g} ^ {(r)}) | | ^ {2}\right) \\ + 2 \mathbb {E} \| 1 - \psi_ {g, m} ^ {(r, k)} \| ^ {2} \left(4 k \eta_ {\ell} ^ {2} \sigma_ {\ell} ^ {2} + 3 K \eta_ {\ell} ^ {2} \sigma_ {\ell} ^ {2}\right) (36) \\ \leq 1 6 k ^ {3} \eta_ {\ell} ^ {2} \mathbb {E} \| 1 - \psi_ {g, m} ^ {(r, k)} \| ^ {2} \mathbb {E} \| \nabla f _ {m} (w _ {g} ^ {(r)}) \| ^ {2} + 8 k \eta_ {\ell} ^ {2} \eta_ {\ell} ^ {2} \mathbb {E} \| \psi_ {g, m} ^ {(r, k)} \| ^ {2} \\ + 1 2 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} + 6 K \eta_ {\ell} ^ {2} \sigma_ {\ell} ^ {2} \mathbb {E} \| \psi_ {g, m} ^ {(r, k)} \| ^ {2} (37) \\ \end{array}
$$
Theorem E.9 (Convergence of the Global Model for Convex Cases). If each client's objective function  $f_{m}$  satisfies Assumptions E.1, E.2, E.3, E.4 using the learning rate  $\frac{1}{\mu R} \leq \eta_{\ell} \leq \min \left( \frac{1}{4\sqrt{10}\beta BK^2}, \frac{1}{8B^2\beta} \right)$  in Algorithm 2, then the following convergence holds:
(Strong Convex Case)
$$
\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} \\ \end{array}
$$
(General Convex Case)
$$
\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 / 3} + \eta_ {\ell} ^ {3} \left(4 0 K ^ {2} \beta^ {3} \frac {\mathbf {q} _ {1} ^ {2}}{\mathbf {q} _ {0} ^ {2}} G ^ {2}\right) ^ {1 / 4} + \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} \\ \end{array}
$$
where  $\mathbf{q}_{0/1}^2$  are the probabilities of picking global/local routes averaged over all the instances sampled from the global distribution.
Proof. From the update rules stated in Equations 10 and 11:
$$
\begin{array}{l} w _ {g} ^ {(r + 1)} - w _ {g} ^ {*} = \frac {1}{n M} \sum_ {m \in [ M ]} n _ {m} \left[ w _ {g, m} ^ {(r)} - \eta_ {\ell} \sum_ {k = 1} ^ {K} h _ {m} \left(w _ {p, m} ^ {(r, k - 1)}\right) \right] - w _ {g} ^ {*} (38) \\ = \frac {1}{n M} \sum_ {m \in [ M ]} n _ {m} w _ {g, m} ^ {(r)} - \frac {\eta_ {\ell}}{n M} \sum_ {m \in [ M ]} n _ {m} \sum_ {k = 1} ^ {K} h _ {m} \left(w _ {p, m} ^ {(r, k - 1)}\right) - w _ {g} ^ {*} (39) \\ = w _ {g} ^ {(r)} - w _ {g} ^ {*} - \frac {\eta_ {\ell}}{n M} \sum_ {m \in [ M ]} n _ {m} \sum_ {k = 1} ^ {K} h _ {m} \left(w _ {p, m} ^ {(r, k - 1)}\right) (40) \\ \end{array}
$$
Taking squared norm and expectation on both sides with respect to the choice of  $h_m$ ,
$$
\begin{array}{l} \mathbb {E} \left[ \left| \left| w _ {g} ^ {(r + 1)} - w _ {g} ^ {*} \right| \right| ^ {2} \right] \leq \mathbb {E} \left[ \left| \left| w _ {g} ^ {(r)} - w _ {g} ^ {*} \right| \right| ^ {2} \right] - 2 \eta_ {\ell} \left\langle \frac {1}{n M} \sum_ {m \in [ M ]} n _ {m} \sum_ {k = 1} ^ {K} \mathbb {E} \left[ h _ {m} \left(w _ {p, m} ^ {(r, k - 1)}\right) \right], w _ {g} ^ {(r)} - w _ {g} ^ {*} \right\rangle \\ + \eta_ {\ell} ^ {2} \mathbb {E} \left[ \left\| \frac {1}{n M} \sum_ {m \in [ M ]} n _ {m} \sum_ {k = 1} ^ {K} h _ {m} \left(w _ {p, m} ^ {(r, k - 1)}\right) \right\| ^ {2} \right] \tag {41} \\ \end{array}
$$
Separating mean and variance according to Lemma 4 of Scaffold [18],
$$
\begin{array}{l} \leq \mathbb{E}\left[||w_{g}^{(r)} - w_{g}^{*}||^{2}\right]\underbrace{-2\eta_{\ell}\left\langle\frac{1}{nM}\sum_{m\in[M]}n_{m}\sum_{k = 1}^{K}\mathbb{E}[\nabla_{w_{g,m}^{(r,k - 1)}}f_{m}(w_{p,m}^{(r,k - 1)})],w_{g}^{(r)} - w_{g}^{*}\right\rangle}_{T_{1}} \\ \underbrace {+ \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]} _ {T _ {2}} + \frac {\eta_ {\ell} ^ {2} \sigma_ {\ell} ^ {2} K}{M} \tag {42} \\ \end{array}
$$
Bounding  $T_{1}$
$$
\begin{array}{l} T _ {1} = - 2 \eta_ {\ell} \left\langle \frac {1}{n M} \sum_ {m \in [ M ]} n _ {m} \sum_ {k = 1} ^ {K} \mathbb {E} \left[ \nabla_ {w _ {g, m} ^ {(r, k - 1)}} f _ {m} \left(w _ {p, m} ^ {(r, k - 1)}\right) \right], w _ {g} ^ {(r)} - w _ {g} ^ {*} \right\rangle (43) \\ = 2 \eta_ {\ell} \left\langle \frac {1}{n M} \sum_ {m \in [ M ]} n _ {m} \sum_ {k = 1} ^ {K} \mathbb {E} \left[ \nabla_ {w _ {g, m} ^ {(r, k - 1)}} f _ {m} \left(w _ {p, m} ^ {(r, k - 1)}\right) \right], w _ {g} ^ {*} - w _ {g} ^ {(r)} \right\rangle (44) \\ \end{array}
$$
Using perturbed strong convexity lemma (Lemma 5) from [18], we get,
$$
T _ {1} \leq \frac {2 \eta_ {\ell}}{n M} \sum_ {m \in [ M ]} n _ {m} \sum_ {k = 1} ^ {K} \left(\mathbb {E} \left[ \nabla f _ {m} \left(w _ {g} ^ {*}\right) \right] - \nabla f _ {m} \left(w _ {g} ^ {(r)}\right) - \frac {\mu}{4} \mathbb {E} \| w _ {g} ^ {(r)} - w _ {g} ^ {*} \| ^ {2} + \beta \underbrace {\mathbb {E} \| w _ {p , m} ^ {(r , k - 1)} - w _ {g} ^ {(r)} \| ^ {2}} _ {\text {L e m m a E . 8}}\right) \tag {45}
$$
$$
\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} \left| \left| w _ {g} ^ {(r)} - w _ {g} ^ {*} \right| \right| ^ {2} \\ + \frac {2 \eta_ {\ell} \beta}{n M} \sum_ {m \in [ M ]} n _ {m} \sum_ {k = 1} ^ {K} \left(1 6 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} + 8 k \eta_ {\ell} ^ {2} \sigma_ {\ell} ^ {2} \mathbb {E} \| \psi_ {g, m} ^ {(r, k)} \| ^ {2} \right. \\ \left. + 1 2 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} + 6 K \eta_ {\ell} ^ {2} \sigma_ {\ell} ^ {2} \mathbb {E} | | \psi_ {g, m} ^ {(r, k)} | | ^ {2}\right) \tag {46} \\ \end{array}
$$