SlowGuess/NeurIPS_2023_OCR huggingface dataset:
This paper introduces a new data&system-aware paradigm for learning from multiple heterogeneous data sources to achieve optimal statistical accuracy across all data distributions without imposing stringent constraints on computational resources shared by participating devices. The proposed PERM schema, though simple, provides an efficient solution to enable each client to learn a personalized model by learning who to learn with via personalizing the aggregation of data sources through an efficient empirical statistical discrepancy estimation module. To efficiently solve all aggregated personalized losses, we propose a model shuffling idea to optimize all losses in parallel. PERM can also be employed in other learning settings with multiple sources of data such as domain adaptation and multi-task learning to entail optimal statistical accuracy.
We would like to embark on the scalability of PERM. The compute burden on clients and servers is roughly the same as existing methods thanks to shuffling (except for extra overhead due to estimating mixing parameters which is the same as running FedAvg in a two-stage approach and an extra communication in an interleaved approach). The only hurdle would be the required memory at server to maintain mixing parameters, which scales proportionally to the square of the number of clients, which can be alleviated by clustering devices which we leave as a future work.
This work was partially supported by NSF CAREER Award #2239374 NSF CNS Award #1956276.
[56] Yae Jee Cho, Pranay Sharma, Gauri Joshi, Zheng Xu, Satyen Kale, and Tong Zhang. On the convergence of federated averaging with cyclic client participation. arXiv preprint arXiv:2302.03109, 2023.
[57] Mehryar Mohri, Afshin Rostamizadeh, and Ameet Talwalkar. Foundations of machine learning. MIT press, 2018.
[58] Tian Li, Anit Kumar Sahu, Manzil Zaheer, Maziar Sanjabi, Ameet Talwalkar, and Virginia Smith. Federated optimization in heterogeneous networks. arXiv preprint arXiv:1812.06127, 2018.
[59] Sebastian Caldas, Peter Wu, Tian Li, Jakub Konečny, H Brendan McMahan, Virginia Smith, and Ameet Talwalkar. Leaf: A benchmark for federated settings. arXiv preprint arXiv:1812.01097, 2018.
[60] Tianyi Lin, Chi Jin, and Michael I Jordan. On gradient descent ascent for nonconvex-concave minimax problems. arXiv preprint arXiv:1906.00331, 2019.
[61] Markus Schneider. Probability inequalities for kernel embeddings in sampling without replacement. In Artificial Intelligence and Statistics, pages 66-74. PMLR, 2016.
In addition to experiments on synthetic and CIFAR10 datasets reported before, we have also conducted experiments on the EMNIST dataset, highlighting PERM's capability to derive superior personalized models by tapping into inter-client data similarities. Additionally, further insights emerged from our tests on the MNIST dataset, revealing how PERM's learned mixture weights adeptly respond to both homogeneous and highly heterogeneous data scenarios.
Experiment on EMNIST dataset In addition to the synthetic and CIFAR10 datasets discussed in the main body, we run experiments on the EMNIST dataset [59], which is naturally distributed in a federated setting. In this case, we chose 50 clients and use a 2-layer MLP model, each with 200 neurons. We compare the PERM algorithm with the localized model in FedAvg and perFedAvg [9]. As it can be seen in Figure 4, PERM can learn a better personalized model by attending to each client's data according to the similarity of the data distribution between clients. The learned values of  $\alpha$ , in Figure 5, show that the clients are learning from each others' data, and not focused on their own data only. This signifies that the distribution of data among clients in this dataset is not highly heterogeneous. Note that, since we are using a subset of clients in the EMNIST dataset for the training (only 50 clients for 100 rounds of communication), the results would be sub-optimal. Nonetheless, the experiments are designed to show the effectiveness of different algorithms. As it can be concluded, in terms of performance, PERM consistently excels beyond its peers, demonstrating exemplary results on various benchmark datasets.
Figure 4: Comparative Analysis of Personalization methods, including our single-loop PERM algorithm, localized FedAvg, and perFedAg, with EMNIST dataset. The disparity in personalized accuracy and loss highlights PERM's capability in leveraging relevant client correlations.
Figure 5: The heat map of the learned  $\alpha$  values for the PERM algorithms on the EMNIST dataset with a 2-layer MLP model. The weights signify that clients mutually benefiting from one another's data, which also highlight that the distribution of data is not significantly heterogeneous in this dataset.
The effectiveness of learned mixture weights To show the effectiveness of the two-stage PERM algorithm, as well as the effects of heterogeneity on the distribution of data among clients on the learned weights  $\alpha$  in the algorithm, we run this algorithm on MNIST dataset. We use 50 clients, and the model is an MLP, similar to the EMNIST experiment. In this case, we consider two cases: distributing the data randomly across clients (homogeneous) and only allocating 1 class per client (highly heterogeneous). As it can be seen from Figure 6, when the data distribution is homogeneous the learned values of  $\alpha$  as diffused across clients. However, when the data is highly heterogeneous, the learned  $\alpha$  values will be highly sparse, indicating that each client is mostly learning from its own data and some other clients with partial distribution similarity. Notably, the matrix predominantly exhibits sparsity, indicating that each client selectively leverages information solely from a subset of other clients. This discernible pattern reinforces the inherent confidence that each client is effectively learning from a limited but strategically chosen group of clients.
(a) Homogeneous distribution
(b) Highly heterogeneous distribution
Figure 6: Comparing the performance of two-stage PERM algorithm in learning  $\alpha$  values on heterogeneous and homogeneous data distributions. We use MNIST dataset across 50 clients with homogeneous and heterogeneous distributions.
B Proof of Two Stages Algorithm
In this section we provide the proof of convergence of two-stage implementation of PERM (computing mixing parameters followed by learning personalized models via model shuffling using permutation-based variant of distributed SGD with periodic communication).
B.1 Technical Lemmas
Lemma 1. Define  $\pmb{v}^{*}(\alpha) \coloneqq \arg \min_{\pmb{v} \in \mathcal{W}} \Phi(\pmb{\alpha}, \pmb{v})$ , and assume  $\Phi(\pmb{\alpha}, \cdot)$  is  $\mu$ -strongly convex and  $\nabla_{\pmb{v}} \Phi(\pmb{\alpha}, \pmb{v})$  is  $L$  Lipschitz in  $\pmb{\alpha}$ . Then,  $\pmb{v}^{*}(\cdot)$  is  $\kappa$ -Lipschitz where  $\kappa = L / \mu$ .
Proof. The proof is similar to Lin et al's result on minimax objective [60]. First, according to optimality conditions we have:
$$
\langle \boldsymbol {v} - \boldsymbol {v} ^ {*} (\boldsymbol {\alpha}), \nabla_ {2} \Phi (\boldsymbol {\alpha}, \boldsymbol {v} ^ {*} (\boldsymbol {\alpha})) \rangle \geq 0,
$$
$$
\langle \boldsymbol {v} - \boldsymbol {v} ^ {*} (\boldsymbol {\alpha} ^ {\prime}), \nabla_ {2} \Phi (\boldsymbol {\alpha} ^ {\prime}, \boldsymbol {v} ^ {*} (\boldsymbol {\alpha} ^ {\prime})) \rangle \geq 0
$$
Substituting  $\pmb{v}$  with  $v^{*}(\alpha^{\prime})$  and  $v^{*}(\alpha)$  in the above first and second inequalities respectively yields:
$$
\langle \boldsymbol {v} ^ {*} (\boldsymbol {\alpha} ^ {\prime}) - \boldsymbol {v} ^ {*} (\boldsymbol {\alpha}), \nabla_ {2} \Phi (\boldsymbol {\alpha}, \boldsymbol {v} ^ {*} (\boldsymbol {\alpha})) \rangle \geq 0,
$$
$$
\langle \boldsymbol {v} ^ {*} (\boldsymbol {\alpha}) - \boldsymbol {v} ^ {*} (\boldsymbol {\alpha} ^ {\prime}), \nabla_ {2} \Phi (\boldsymbol {\alpha} ^ {\prime}, \boldsymbol {v} ^ {*} (\boldsymbol {\alpha} ^ {\prime})) \rangle \geq 0
$$
Adding up the above two inequalities yields:
$$
\langle \boldsymbol {v} ^ {*} (\boldsymbol {\alpha} ^ {\prime}) - \boldsymbol {v} ^ {*} (\boldsymbol {\alpha}), \nabla_ {2} \Phi (\boldsymbol {\alpha}, \boldsymbol {v} ^ {*} (\boldsymbol {\alpha})) - \nabla_ {2} \Phi (\boldsymbol {\alpha} ^ {\prime}, \boldsymbol {v} ^ {*} (\boldsymbol {\alpha} ^ {\prime})) \rangle \geq 0, \tag {4}
$$
Since  $\Phi (\alpha ,\cdot)$  is  $\mu$  strongly convex, we have:
$$
\left\langle \boldsymbol {v} ^ {*} (\boldsymbol {\alpha} ^ {\prime}) - \boldsymbol {v} ^ {*} (\boldsymbol {\alpha}), \nabla_ {2} \Phi (\boldsymbol {\alpha}, \boldsymbol {v} ^ {*} (\boldsymbol {\alpha} ^ {\prime})) - \nabla_ {2} \Phi (\boldsymbol {\alpha}, \boldsymbol {v} ^ {*} (\boldsymbol {\alpha})) \geq \mu \| \boldsymbol {v} ^ {*} (\boldsymbol {\alpha} ^ {\prime}) - \boldsymbol {v} ^ {*} (\boldsymbol {\alpha}) \| ^ {2}. \right. \tag {5}
$$
Adding up (4) and (5) yields:
$$
\left\langle \boldsymbol {v} ^ {*} (\boldsymbol {\alpha} ^ {\prime}) - \boldsymbol {v} ^ {*} (\boldsymbol {\alpha}), \nabla_ {2} \Phi (\boldsymbol {\alpha}, \boldsymbol {v} ^ {*} (\boldsymbol {\alpha} ^ {\prime})) - \nabla_ {2} \Phi (\boldsymbol {\alpha} ^ {\prime}, \boldsymbol {v} ^ {*} (\boldsymbol {\alpha} ^ {\prime})) \geq \mu \| \boldsymbol {v} ^ {*} (\boldsymbol {\alpha} ^ {\prime}) - \boldsymbol {v} ^ {*} (\boldsymbol {\alpha}) \| ^ {2} \right.
$$
Finally, using  $L$  smoothness of  $\Phi$  will conclude the proof:
$$
\begin{array}{l} L \left\| \boldsymbol {v} ^ {*} (\boldsymbol {\alpha} ^ {\prime}) - \boldsymbol {v} ^ {*} (\boldsymbol {\alpha}) \right\| \left\| \boldsymbol {\alpha} - \boldsymbol {\alpha} ^ {\prime} \right\| \geq \mu \left\| \boldsymbol {v} ^ {*} (\boldsymbol {\alpha} ^ {\prime}) - \boldsymbol {v} ^ {*} (\boldsymbol {\alpha}) \right\| ^ {2} \\ \Longleftrightarrow \kappa \| \boldsymbol {\alpha} - \boldsymbol {\alpha} ^ {\prime} \| \geq \| \boldsymbol {v} ^ {*} (\boldsymbol {\alpha} ^ {\prime}) - \boldsymbol {v} ^ {*} (\boldsymbol {\alpha}) \| \\ \end{array}
$$
Lemma 2 (Optimality Gap). Let  $\Phi (\alpha ,\pmb {v})$  be defined in (3). Let  $\hat{\pmb{v}} = \mathcal{P}_{\mathcal{W}}(\tilde{\pmb{v}} -\frac{1}{L}\nabla_{\pmb{v}}\Phi (\hat{\pmb{\alpha}},\tilde{\pmb{v}}))$ . If we assume each  $f_{i}$  is  $L$ -smooth,  $\mu$ -strongly convex and with gradient bounded by  $G$ , then the following statement holds true:
$$
\Phi (\boldsymbol {\alpha} ^ {*}, \hat {\boldsymbol {v}}) - \Phi (\boldsymbol {\alpha} ^ {*}, \boldsymbol {v} ^ {*}) \leq 2 L \| \tilde {\boldsymbol {v}} - \boldsymbol {v} ^ {*} (\hat {\boldsymbol {\alpha}}) \| ^ {2} + \left(2 \kappa_ {\Phi} ^ {2} L + \frac {4 N G ^ {2}}{L}\right) \| \hat {\boldsymbol {\alpha}} - \boldsymbol {\alpha} ^ {*} \| ^ {2}.
$$
where  $\kappa_{\Phi} = \frac{\sqrt{N}G}{\mu}$ ,  $\pmb{v}^{*} = \arg \min_{\pmb{v}\in \mathcal{W}}\Phi (\pmb{\alpha}^{*},\pmb{v})$
Proof. First we show that  $\nabla_{\pmb{v}}\Phi (\pmb {\alpha},\pmb {v})$  is  $\sqrt{NG}$  Lipschitz in  $\pmb{\alpha}$ . To see this:
$$
\begin{array}{l} \| \nabla_ {\boldsymbol {v}} \Phi (\boldsymbol {\alpha}, \boldsymbol {v}) - \nabla_ {\boldsymbol {v}} \Phi (\boldsymbol {\alpha} ^ {\prime}, \boldsymbol {v}) \| = \left\| \sum_ {j = 1} ^ {N} \alpha_ {i} (j) \nabla f _ {j} (\boldsymbol {v}) - \sum_ {j = 1} ^ {N} \alpha_ {i} ^ {\prime} (j) \nabla f _ {j} (\boldsymbol {v}) \right\| \\ \leq \sqrt {N} G \| \boldsymbol {\alpha} _ {i} - \boldsymbol {\alpha} _ {i} ^ {\prime} \|. \\ \end{array}
$$
Hence due to Lemma 1, we know  $\pmb{v}^{*}(\pmb{\alpha})$  is  $\kappa_{\Phi} := \frac{\sqrt{N}G}{\mu}$  Lipschitz. According to property of projection, we have:
$$
\begin{array}{l} 0 \leq \langle \boldsymbol {v} - \hat {\boldsymbol {v}}, L (\hat {\boldsymbol {v}} - \tilde {\boldsymbol {v}}) + \nabla_ {\boldsymbol {v}} \Phi (\hat {\boldsymbol {\alpha}}, \tilde {\boldsymbol {v}}) \rangle \\ = \underbrace {\langle \boldsymbol {v} - \hat {\boldsymbol {v}} , L (\hat {\boldsymbol {v}} - \tilde {\boldsymbol {v}}) + \nabla_ {\boldsymbol {v}} \Phi (\boldsymbol {\alpha} ^ {*} , \tilde {\boldsymbol {v}}) \rangle} _ {T _ {1}} + \underbrace {\langle \boldsymbol {v} - \hat {\boldsymbol {v}} , \nabla_ {\boldsymbol {v}} \Phi (\hat {\boldsymbol {\alpha}} , \tilde {\boldsymbol {v}}) - \nabla_ {\boldsymbol {v}} \Phi (\boldsymbol {\alpha} ^ {*} , \tilde {\boldsymbol {v}}) \rangle} _ {T _ {2}}. \\ \end{array}
$$
For  $T_{1}$ , we notice:
$$
\begin{array}{l} \langle \boldsymbol {v} - \hat {\boldsymbol {v}}, L (\hat {\boldsymbol {v}} - \tilde {\boldsymbol {v}}) + \nabla_ {\boldsymbol {v}} \Phi (\boldsymbol {\alpha} ^ {*}, \tilde {\boldsymbol {v}}) \rangle = L \langle \boldsymbol {v} - \tilde {\boldsymbol {v}}, \hat {\boldsymbol {v}} - \tilde {\boldsymbol {v}} \rangle + \frac {1}{\eta} \langle \tilde {\boldsymbol {v}} - \hat {\boldsymbol {v}}, \hat {\boldsymbol {v}} - \tilde {\boldsymbol {v}} \rangle + \langle \boldsymbol {v} - \hat {\boldsymbol {v}}, \nabla_ {\boldsymbol {v}} \Phi (\boldsymbol {\alpha} ^ {*}, \tilde {\boldsymbol {v}}) \rangle \\ = L \left\langle \pmb {v} - \tilde {\pmb {v}}, \hat {\pmb {v}} - \tilde {\pmb {v}} \right\rangle - L \left\| \tilde {\pmb {v}} - \hat {\pmb {v}} \right\| ^ {2} + \left\langle \pmb {v} - \hat {\pmb {v}} _ {i}, \nabla_ {\pmb {v}} \Phi (\pmb {\alpha} ^ {*}, \tilde {\pmb {v}}) \right\rangle \\ \leq L \left(\| \boldsymbol {v} - \tilde {\boldsymbol {v}} \| ^ {2} + \frac {1}{4} \| \hat {\boldsymbol {v}} - \tilde {\boldsymbol {v}} \| ^ {2}\right) - L \| \tilde {\boldsymbol {v}} - \hat {\boldsymbol {v}} \| ^ {2} + \underbrace {\langle \boldsymbol {v} - \hat {\boldsymbol {v}} , \nabla_ {\boldsymbol {v}} \Phi (\boldsymbol {\alpha} ^ {*} , \tilde {\boldsymbol {v}}) \rangle} _ {\triangleleft} \\ \end{array}
$$
where at last step we used Young's inequality. To bound  $\spadesuit$ , we apply the  $L$  smoothness and  $\mu$  strongly convexity of  $\Phi(\alpha, \cdot)$ :
$$
\begin{array}{l} \langle \boldsymbol {v} - \hat {\boldsymbol {v}} _ {i}, \nabla_ {\boldsymbol {v}} \Phi (\boldsymbol {\alpha} _ {i} ^ {*}, \tilde {\boldsymbol {v}} _ {i}) \rangle = \langle \boldsymbol {v} - \tilde {\boldsymbol {v}} _ {i}, \nabla_ {\boldsymbol {v}} \Phi (\boldsymbol {\alpha} ^ {*}, \tilde {\boldsymbol {v}}) \rangle + \langle \tilde {\boldsymbol {v}} - \hat {\boldsymbol {v}}, \nabla_ {\boldsymbol {v}} \Phi (\boldsymbol {\alpha} ^ {*}, \tilde {\boldsymbol {v}}) \rangle \\ \leq \Phi (\boldsymbol {\alpha} ^ {*}, \boldsymbol {v}) - \Phi (\boldsymbol {\alpha} ^ {*}, \tilde {\boldsymbol {v}}) - \frac {\mu}{2} \| \tilde {\boldsymbol {v}} - \boldsymbol {v} \| ^ {2} + \Phi (\boldsymbol {\alpha} ^ {*}, \tilde {\boldsymbol {v}}) - \Phi (\boldsymbol {\alpha} ^ {*}, \hat {\boldsymbol {v}}) + \frac {L}{2} \| \tilde {\boldsymbol {v}} - \hat {\boldsymbol {v}} \| ^ {2} \\ \leq \Phi (\boldsymbol {\alpha} ^ {*}, \boldsymbol {v}) - \Phi (\boldsymbol {\alpha} ^ {*}, \hat {\boldsymbol {v}}) - \frac {\mu}{2} \| \tilde {\boldsymbol {v}} - \boldsymbol {v} \| ^ {2} + \frac {L}{2} \| \tilde {\boldsymbol {v}} - \hat {\boldsymbol {v}} \| ^ {2} \\ \end{array}
$$
Putting above bound back yields:
$$
\left\langle \boldsymbol {v} - \hat {\boldsymbol {v}}, \frac {1}{\eta} (\hat {\boldsymbol {v}} - \tilde {\boldsymbol {v}}) + \nabla_ {\boldsymbol {v}} \Phi (\boldsymbol {\alpha} ^ {*}, \tilde {\boldsymbol {v}}) \right\rangle \leq \Phi (\boldsymbol {\alpha} ^ {*}, \boldsymbol {v}) - \Phi (\boldsymbol {\alpha} ^ {*}, \hat {\boldsymbol {v}}) + \frac {1}{2 \eta} \| \tilde {\boldsymbol {v}} - \boldsymbol {v} \| ^ {2} - \left(\frac {3 L}{4} - \frac {L}{2}\right) \| \tilde {\boldsymbol {v}} - \hat {\boldsymbol {v}} \| ^ {2}
$$
Now we switch to bounding  $T_{2}$ . Applying Cauchy-Schwartz yields:
$$
\begin{array}{l} \langle \boldsymbol {v} - \hat {\boldsymbol {v}} _ {i}, \nabla_ {\boldsymbol {v}} \Phi (\hat {\boldsymbol {\alpha}} _ {i}, \tilde {\boldsymbol {v}} _ {i}) - \nabla_ {\boldsymbol {v}} \Phi (\boldsymbol {\alpha} _ {i} ^ {*}, \tilde {\boldsymbol {v}} _ {i}) \rangle \leq \frac {L}{4} \| \boldsymbol {v} - \tilde {\boldsymbol {v}} _ {i} \| ^ {2} + \frac {L}{4} \| \tilde {\boldsymbol {v}} - \hat {\boldsymbol {v}} _ {i} \| ^ {2} + \frac {4}{L} \| \nabla_ {\boldsymbol {v}} \Phi (\hat {\boldsymbol {\alpha}} _ {i}, \tilde {\boldsymbol {v}} _ {i}) - \nabla_ {\boldsymbol {v}} \Phi (\boldsymbol {\alpha} _ {i} ^ {*}, \tilde {\boldsymbol {v}} _ {i}) \| ^ {2} \\ \leq \frac {L}{4} \left\| \boldsymbol {v} - \tilde {\boldsymbol {v}} _ {i} \right\| ^ {2} + \frac {L}{4} \left\| \tilde {\boldsymbol {v}} - \hat {\boldsymbol {v}} _ {i} \right\| ^ {2} + \frac {4 N G ^ {2}}{L} \left\| \hat {\boldsymbol {\alpha}} _ {i} - \boldsymbol {\alpha} _ {i} ^ {*} \right\| ^ {2} \\ \end{array}
$$
where at last step we apply  $\sqrt{NG}$  smoothness of  $\Phi (\cdot ,\pmb {v})$  . Putting pieces together yields:
$$
0 \leq \Phi (\boldsymbol {\alpha} _ {i} ^ {*}, \boldsymbol {v}) - \Phi (\boldsymbol {\alpha} _ {i} ^ {*}, \hat {\boldsymbol {v}} _ {i}) + \frac {L}{2} \| \tilde {\boldsymbol {v}} _ {i} - \boldsymbol {v} \| ^ {2} + \frac {L}{2} \| \boldsymbol {v} - \tilde {\boldsymbol {v}} _ {i} \| ^ {2} + \frac {4 N G ^ {2}}{L} \| \hat {\boldsymbol {\alpha}} _ {i} - \boldsymbol {\alpha} _ {i} ^ {*} \| ^ {2}
$$
Re-arranging terms and setting  $\pmb{v} = \pmb{v}^{*}(\pmb{\alpha}^{*}) = \arg \min_{\pmb{v}\in \mathcal{W}}\Phi (\pmb{\alpha}^{*},\pmb{v})$  yields:
$$
\Phi (\boldsymbol {\alpha} ^ {*}, \boldsymbol {\hat {v}}) - \Phi (\boldsymbol {\alpha} ^ {*}, \boldsymbol {v} ^ {*}) \leq L \| \tilde {\boldsymbol {v}} - \boldsymbol {v} ^ {*} \| ^ {2} + \frac {4 N G ^ {2}}{L} \| \hat {\boldsymbol {\alpha}} - \boldsymbol {\alpha} ^ {*} \| ^ {2}.
$$
At last, due to the  $\kappa_{\Phi}$ -Lipschitzness property of  $v^{*}(\cdot)$  as shown in Lemma 1, it follows that:
$$
\begin{array}{l} L \| \tilde {\boldsymbol {v}} - \boldsymbol {v} ^ {*} (\boldsymbol {\alpha} ^ {*}) \| ^ {2} \leq L \| \tilde {\boldsymbol {v}} - \boldsymbol {v} ^ {*} (\hat {\boldsymbol {\alpha}}) \| ^ {2} + L \| \boldsymbol {v} ^ {*} (\hat {\boldsymbol {\alpha}}) - \boldsymbol {v} ^ {*} (\boldsymbol {\alpha} ^ {*}) \| ^ {2} \\ \leq 2 L \| \tilde {\boldsymbol {v}} - \boldsymbol {v} ^ {*} (\hat {\boldsymbol {\alpha}}) \| ^ {2} + 2 \kappa_ {\Phi} ^ {2} L \| \hat {\boldsymbol {\alpha}} - \boldsymbol {\alpha} ^ {*} \| ^ {2}, \\ \end{array}
$$
Algorithm A1: Discrepancy Estimation at Optimum
Input: Number of clients  $N$ , number of local steps  $K$ , number of communications rounds  $R$  for  $r = 0, \dots, R - 1$  do  
parallel for client  $i = 0, \dots, N - 1$  do  
Client  $i$  initializes model  $\boldsymbol{w}_i^{r,0} = \boldsymbol{w}_i^r$ .  
for  $t = 0, \dots, K - 1$  do  
 $\boldsymbol{w}_i^{r,t + 1} = \boldsymbol{w}_i^{r,t} - \gamma \nabla f_i(\boldsymbol{w}_i^{r,t}; \xi_i^{r,t})$  where  $\xi_i^{r,t}$  is a mini-batch sampled from  $\mathcal{S}_i$ .  
Client  $i$  sends  $\boldsymbol{w}_i^{r,K}$  to Server.  
Server computes  $\boldsymbol{w}^{r + 1} = \mathcal{P}_{\mathcal{W}}\left(\frac{1}{N} \sum_{i = 1}^{N} \boldsymbol{w}_i^{r,K}\right)$   
Server broadcasts  $\boldsymbol{w}^{r + 1}$  to all clients.  
Server computes  $\hat{\alpha}_i, i = 1, 2, \dots, N$  by running  $T_\alpha$  steps of GD on  $g_i(\boldsymbol{w}^R, \alpha)$ .  
Output:  $\hat{\alpha}_1, \dots, \hat{\alpha}_N$ .
B.2 Proof of Convergence of Theorem 1
In this section we are going to prove the result in Theorem 1. To this end, we need to show that mixing parameters we compute by first learning the global model and then solving the optimization problem in objective (2) (as depicted in Algorithm A1) converges to optimal values. Notice that in Algorithm A1 we do not solve  $g_{i}(\boldsymbol{w}^{*}, \boldsymbol{\alpha})$  directly, but optimize  $g_{i}(\boldsymbol{w}^{R}, \boldsymbol{\alpha})$  on  $\boldsymbol{\alpha}$  for  $T_{\boldsymbol{\alpha}}$  iterations of GD. Hence, firstly we need to show that optimizing the surrogate function will also guarantee the convergence of output of algorithm  $\widehat{\alpha}$  to  $\boldsymbol{\alpha}^{*}$  by deriving a property of the objective in (2). Formally the property is captured by the following lemma.
Lemma 3. Let  $g(\pmb{w}, \pmb{\alpha}) \coloneqq \sum_{j=1}^{N} \alpha_j \left\| \nabla f_i(\pmb{w}) - \nabla f_j(\pmb{w}) \right\|^2 + \lambda \sum_{j=1}^{N} \alpha_j^2 / n_j$  and  $\pmb{\alpha}_g^*(\pmb{w}) = \arg \min_{\pmb{\alpha} \in \Delta_N} g(\pmb{w}, \pmb{\alpha})$ . Let  $\pmb{w}^R$  be the output of Algorithm A1. Then the following statement holds:
$$
\left\| \boldsymbol {\alpha} _ {g} ^ {*} (\boldsymbol {w} ^ {R}) - \boldsymbol {\alpha} _ {g} ^ {*} (\boldsymbol {w} ^ {*}) \right\| \leq \kappa_ {g} ^ {2} \sum_ {j = 1} ^ {N} \left(2 \| \nabla f _ {i} (\boldsymbol {w} ^ {*}) - \nabla f _ {j} (\boldsymbol {w} ^ {*}) \| ^ {2} + 4 L ^ {2} \| \boldsymbol {w} ^ {R} - \boldsymbol {w} ^ {*} \| ^ {2}\right) 4 L \| \boldsymbol {w} ^ {R} - \boldsymbol {w} ^ {*} \| ^ {2}
$$
where  $\kappa_{g}:= \frac{n_{\mathrm{max}}}{2\lambda}$
Proof. Define function
$$
W (\boldsymbol {z}, \boldsymbol {\alpha}) = \sum_ {j = 1} ^ {N} \alpha_ {j} z _ {j} + \lambda \sum_ {j = 1} ^ {N} \alpha_ {j} ^ {2} / n _ {j} \tag {6}
$$
Apparently,  $W(z,\alpha)$  is linear in  $z$  and  $2\frac{\lambda}{n_{\mathrm{max}}}$  strongly convex in  $z$ . Next we show that  $\nabla_{\alpha}W(z,\alpha)$  is Lipschitz in  $w$ . To see this,
$$
\begin{array}{l} \| \nabla_ {\boldsymbol {\alpha}} W (\boldsymbol {z}, \boldsymbol {\alpha}) - \nabla_ {\boldsymbol {\alpha}} W (\boldsymbol {z} ^ {\prime}, \boldsymbol {\alpha}) \| = \| [ z _ {1}, \dots , z _ {N} ] - [ z _ {1} ^ {\prime}, \dots , z _ {N} ^ {\prime} ] \| \\ \leq \left\| z - z ^ {\prime} \right\|. \\ \end{array}
$$
Then, according to Proposition 1,  $\pmb{\alpha}_{W}^{*}(\pmb{z})\coloneqq \arg \min_{\pmb {\alpha}\in \Delta_{N}}W(\pmb {z},\pmb {\alpha})$  is  $\kappa_{g}$  lipschitz in  $\pmb{z}$  where  $\kappa_{g} = \frac{n_{\mathrm{max}}}{2\lambda}$ , i.e.,  $\| \pmb{\alpha}_{W}^{*}(\pmb {z}) - \pmb{\alpha}_{W}^{*}(\pmb{z}^{\prime})\| \leq \kappa_{g}\| \pmb {z} - \pmb{z}^{\prime}\|$ . Now, let us consider the objective (2):
$$
g (\boldsymbol {w}, \boldsymbol {\alpha}) := \sum_ {j = 1} ^ {N} \alpha_ {j} \| \nabla f _ {i} (\boldsymbol {w}) - \nabla f _ {j} (\boldsymbol {w}) \| ^ {2} + \lambda \sum_ {j = 1} ^ {N} \alpha_ {j} ^ {2} / n _ {j}
$$
We define  $\alpha_{g}^{*}(\boldsymbol {w}) = \arg \min_{\boldsymbol {\alpha}\in \Delta_{N}}g(\boldsymbol {w},\boldsymbol {\alpha})$
We set
$$
\begin{array}{l} \boldsymbol {z} ^ {R} = \left[ \left\| \nabla f _ {i} (\boldsymbol {w} ^ {R}) - \nabla f _ {1} (\boldsymbol {w} ^ {R}) \right\| ^ {2}, \dots , \left\| \nabla f _ {i} (\boldsymbol {w} ^ {R}) - \nabla f _ {N} (\boldsymbol {w} ^ {R}) \right\| ^ {2} \right], \\ \boldsymbol {z} ^ {*} = \left[ \| \nabla f _ {i} (\boldsymbol {w} ^ {*}) - \nabla f _ {1} (\boldsymbol {w} ^ {*}) \| ^ {2}, \dots , \| \nabla f _ {i} (\boldsymbol {w} ^ {*}) - \nabla f _ {N} (\boldsymbol {w} ^ {*}) \| ^ {2} \right]. \\ \end{array}
$$