SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
\left\| \boldsymbol {\alpha} _ {g _ {i}} ^ {*} (\boldsymbol {w} ^ {R}) - \boldsymbol {\alpha} _ {g _ {i}} ^ {*} (\boldsymbol {w} ^ {*}) \right\| ^ {2} \leq 4 L ^ {2} \kappa_ {g} ^ {2} \sum_ {j = 1} ^ {N} \left(2 \left\| \nabla f _ {i} (\boldsymbol {w} ^ {*}) - \nabla f _ {j} (\boldsymbol {w} ^ {*}) \right\| ^ {2} + 4 L ^ {2} \left\| \boldsymbol {w} ^ {R} - \boldsymbol {w} ^ {*} \right\| ^ {2}\right) \left\| \boldsymbol {w} ^ {R} - \boldsymbol {w} ^ {*} \right\| ^ {2}
$$
where  $\alpha_{g_i}^*(\boldsymbol{w}) \coloneqq \arg \min_{\boldsymbol{\alpha} \in \Delta_N} g_i(\boldsymbol{w}, \boldsymbol{\alpha})$  and  $\kappa_g \coloneqq n_{\max} / (2\lambda)$  is the condition number of  $g_i(\boldsymbol{w}, \cdot)$  where  $n_{\max} = \max_{i \in [N]} n_i$ . The Lipschitz constant mainly depends on gradient dissimilarity at optimum. As  $\boldsymbol{w}^R$  tends to  $\boldsymbol{w}^*$ , the  $\alpha_{g_i}^*(\cdot)$  becomes more Lipschitz continuous, i.e., the coefficient in front of  $\left\| \boldsymbol{w}^R - \boldsymbol{w}^* \right\|^2$  getting smaller, thus leading to more accurate mixing parameters.
To establish the convergence, we make the following standard assumptions.
Assumption 1 (Smoothness and strong convexity). We assume  $f_{i}(\pmb {x})$ 's are  $L$ -smooth and  $\mu$ -strongly convex, i.e.,
$$
\forall \boldsymbol {x}, \boldsymbol {y}: \| \nabla f _ {i} (\boldsymbol {x}) - \nabla f _ {i} (\boldsymbol {y}) \| \leq L \| \boldsymbol {x} - \boldsymbol {y} \|.
$$
$$
\forall \boldsymbol {x}, \boldsymbol {y}: f _ {i} (\boldsymbol {y}) \geq f _ {i} (\boldsymbol {x}) + \left\langle \nabla f _ {i} (\boldsymbol {x}), \boldsymbol {y} - \boldsymbol {x} \right\rangle + \frac {1}{2} \mu \| \boldsymbol {y} - \boldsymbol {x} \| ^ {2}
$$
We denote the condition number by  $\kappa = L / \mu$
Assumption 2 (Bounded variance). The variance of stochastic gradients computed at each local function is bounded, i.e.,  $\forall i\in [N],\forall \pmb {w}\in \mathcal{W},\mathbb{E}[\| \nabla f_i(\pmb {w};\xi) - \nabla f_i(\pmb {w})\| ^2 ]\leq \delta^2$
Assumption 3 (Bounded domain). The domain  $\mathcal{W} \subset \mathbb{R}^d$  is a bounded convex set, with diameter  $D$  under  $\ell_2$  metric, i.e.,  $\forall \boldsymbol{w}, \boldsymbol{w}' \in \mathcal{W}, \| \boldsymbol{w} - \boldsymbol{w}' \| \leq D$ .
Definition 2 (Gradient dissimilarity). We define the following quantities to measure the gradient dissimilarity among local functions:
$$
\zeta_ {i, j} (\boldsymbol {w}) := \left\| \nabla f _ {i} (\boldsymbol {w}) - \nabla f _ {j} (\boldsymbol {w}) \right\| ^ {2}, \quad \bar {\zeta} _ {i} (\boldsymbol {w}) := \frac {1}{N} \sum_ {j = 1} ^ {N} \zeta_ {i, j} (\boldsymbol {w}),
$$
$$
\zeta := \sup _ {\boldsymbol {w} \in \mathcal {W}} \max _ {i \in [ N ]} \| \nabla f _ {i} (\boldsymbol {w}) - (1 / N) \sum_ {j = 1} ^ {N} \nabla f _ {j} (\boldsymbol {w}) \| ^ {2}.
$$
The following theorem gives the convergence rate of estimated discrepancies to optimal counterparts.
Theorem 1. Under Assumptions 1-3, if we run Algorithm A1 on  $F(\boldsymbol{w}) \coloneqq \frac{1}{N} \sum_{j=1}^{N} f_j(\boldsymbol{w})$  with  $\gamma = \Theta\left(\frac{\log(RK)}{\mu RK}\right)$  for  $R$  rounds with synchronization gap  $K$ , for  $\kappa_g = 1 / (\lambda n_{\min})$ , it holds that
$$
\mathbb {E} \| \pmb {\alpha} _ {i} ^ {R} - \pmb {\alpha} _ {i} ^ {*} \| ^ {2} \leq \tilde {O} \left(\exp \left(- \frac {T _ {\pmb {\alpha}}}{\kappa_ {g}}\right) + \kappa_ {g} ^ {2} \bar {\zeta} _ {i} (\pmb {w} ^ {*}) L ^ {2} \left(\frac {D ^ {2}}{R K} + \frac {\kappa \zeta^ {2}}{\mu^ {2} R ^ {2}} + \frac {\delta^ {2}}{\mu^ {2} N R K}\right)\right) \quad \forall i \in [ N ].
$$
An immediate implication of Theorem 1 is that even we solve (2) at  $\boldsymbol{w}^R$ , the algorithm will eventually converge to optimal solution of (2) at  $\boldsymbol{w}^*$ . The core technique in the proof, as we mentioned, is to show that for a parameter within a small region centered at  $\boldsymbol{w}^*$ , the function  $\alpha_{g_i}^*(\boldsymbol{w})$  becomes 'more Lipschitz'. The rigorous characterization of this property is captured by Lemma 3 in appendix.
Stage 2: Scalable personalized optimization with model shuffling. After obtaining the per client mixing parameters, in the second stage we aim at solving  $N$  different personalized variants of PERM denoted by  $\Phi (\hat{\alpha}_1,v),\Phi (\hat{\alpha}_2,v),\ldots \Phi (\hat{\alpha}_N,v)$  to learn local models where
$$
\min  _ {\boldsymbol {v} \in \mathcal {W}} \Phi (\hat {\boldsymbol {\alpha}} _ {i}, \boldsymbol {v}) := \frac {1}{N} \sum_ {j = 1} ^ {N} \hat {\alpha} _ {i} (j) f _ {j} (\boldsymbol {v}). \tag {3}
$$
Here we devise an iterative algorithm based on distributed SGD with periodic averaging (a.k.a. Local SGD [48]) to solve these  $N$  optimization problems in parallel with no extra overhead. The idea is to replace the model averaging in vanilla distributed (Local) SGD with model shuffling. Specifically, as shown in Algorithm 1 the algorithm proceeds for  $R$  epochs where each epoch runs for  $N$  communication rounds. At the beginning of each epoch  $r$  the server generates a random permutation  $\sigma_r$  over  $N$  clients. At each communication round  $j$  within the epoch, the server sends the model of client  $i$  to client  $i_j = (i + j) \mod N$  in the permutation  $\sigma_r$  along with  $\alpha_i(i_j)$ . After receiving a model from the server, the client updates the received model for  $K$  local steps and returns it back to the server. As it can be seen, the updates of each loss  $\Phi(\hat{\alpha}_i, \mathbf{v}), i = 1, 2, \dots, N$  during an epoch is equivalent to sequentially processing individual losses in (3) which can be considered as permutation-based SGD but with the different that each component now is updated for  $K$  steps. By interleaving the permutations, we are able to simultaneously optimize all  $N$  objectives. We note that the computation and communication complexity of the proposed algorithm is the same as Local SGD with two differences: the model averaging is replaced with model shuffling, and the algorithms run over a permutation of devices. The convergence rate of Local SGD is well-established in literature [50, 51, 52, 53, 54], but here we establish the convergence of permutation-based variant which is interesting by its own.
Assumption 4 (Bounded Gradient). The variance of stochastic gradients computed at each local function is bounded, i.e.,  $\forall i\in [N],\sup_{\boldsymbol {v}\in \mathcal{W}}\| \nabla f_i(\boldsymbol {v})\| \leq G$
We note that the Assumption 4 can be realized since we work with a bounded domain  $\mathcal{W}$ .
Theorem 2. Let Assumptions 1-4 hold. Assume  $\alpha_{i}^{*}$  is the solution of (2). Then if we run Algorithm 1 on the  $\hat{\alpha}_{i}$  obtained from Algorithm A1, then Algorithm 1 with  $\eta = \Theta\left(\frac{\log(NKR^3)}{\mu R}\right)$  will output the solution  $\hat{v}_i$ ,  $\forall i\in [N]$ , such that with probability at least  $1 - p$ , the following statement holds:
$$
\begin{array}{l} \mathbb {E} \left[ \Phi \left(\boldsymbol {\alpha} _ {i} ^ {*}, \hat {\boldsymbol {v}} _ {i}\right) - \Phi \left(\boldsymbol {\alpha} _ {i} ^ {*}, \boldsymbol {v} ^ {*} \left(\boldsymbol {\alpha} _ {i} ^ {*}\right)\right) \right] \leq \tilde {O} \left(\frac {D ^ {2} L}{N K R ^ {2}}\right) + \frac {L \delta^ {2}}{\mu^ {2} R} + \left(\frac {L ^ {4} + N}{\mu^ {4} R ^ {2}}\right) L G ^ {2} N \log (1 / p) \\ + \kappa_ {\Phi} ^ {2} L \tilde {O} \left(\exp \left(- \frac {T _ {\boldsymbol {\alpha}}}{\kappa_ {g}}\right) + \kappa_ {g} ^ {2} \bar {\zeta} _ {i} (\boldsymbol {w} ^ {*}) L ^ {2} \left(\frac {\kappa \zeta^ {2}}{\mu^ {2} R ^ {2}} + \frac {\delta^ {2}}{\mu^ {2} N R K}\right)\right), \\ \end{array}
$$
Algorithm 1: Shuffling Local SGD
Input: Clients 1,...,N, Number of Local Steps  $K$  , Number of Epoch  $\mathcal{R}$  mixing parameter  $\hat{\alpha}_{1},\dots,\hat{\alpha}_{N}$    
Epoch for  $r = 0,\ldots ,R - 1$  do Server generates permutation  $\sigma_r:[N]\mapsto [N]$  parallel for  $i = 1,\dots,N$  do Client i sets initial model  $\pmb {v}_i^{r,0} = \pmb {v}_i^r$  . for  $j = 1,\dots,N$  do Set indices  $i_j = \sigma_r((i + j)\mod N)$  Server sends  $\pmb {v}_i^{r,j}$  to Client  $i_j$ $\pmb {v}_i^{r,j + 1} = \mathrm{SGD - Update}(\pmb {v}_i^{r,j},\eta ,i_j,K,\hat{\alpha}_i)$  Client i does projection:  $\pmb {v}_i^{r + 1} = \mathcal{P}_{\mathcal{W}}(\pmb {v}_i^{r,N})$    
Output:  $\hat{v}_i = \mathcal{P}_{\mathcal{W}}(\pmb {v}_i^R -(1 / L)\nabla_\pmb {v}\Phi (\hat{\alpha}_i,\pmb {v}_i^R)),\forall i\in [N].$    
SGD-Update(v,  $\eta ,j,K,\alpha$  ) Initialize  $\pmb{v}^{0} = \pmb{v}$  for  $t = 0,\dots,K - 1$  do  $\pmb{v}^{t} = \pmb{v}^{t - 1} - \eta \alpha (j)\pmb {N}\nabla f_{j}(\pmb{v}^{t - 1};\xi_{j}^{t - 1})$  Output:  $\pmb{v}^{K}$
where  $\kappa = \frac{L}{\mu}, \kappa_{g} = \frac{n_{\max}}{2\lambda}$  and  $\kappa_{\Phi} = \frac{\sqrt{NG}}{\mu}$ , and the expectation is taken over randomness of Algorithm A1. That is, to guarantee  $\mathbb{E}[\Phi(\pmb{\alpha}_{i}^{*}, \hat{\pmb{v}}_{i}) - \Phi(\pmb{\alpha}_{i}^{*}, \pmb{v}_{i}^{*})] \leq \epsilon$ , we choose  $R = O\left(\max\left\{\frac{L\delta^{2}}{\mu\epsilon}, \frac{\kappa_{\Phi}^{2}\kappa_{g}^{2}L^{3}\bar{\zeta}_{i}(\pmb{w}^{*})D^{2}}{\epsilon}\right\}\right)$  and  $T_{\alpha} = O\left(\kappa_{g}\log\left(\frac{L\kappa_{\Phi}^{2}}{\epsilon}\right)\right)$ .
The above theorem shows that even though we run the optimization on  $\Phi (\hat{\alpha}_i,v)$ , our obtained model  $\hat{v}_i$  will still converge to the optimal solution of  $\Phi (\pmb {\alpha}_i^*,\pmb {v})$ . The convergence rate is contributed from two parts: convergence of  $\hat{\alpha}_i$  (Algorithm A1) and convergence of personalized model  $\hat{v}_i$  (Algorithm 1). Notice that, for the convergence rate of  $\hat{v}_i$ , we roughly recover the optimal rate of shuffling SGD [55], which is  $O(1 / R^2)$ . However, we suffer from a  $O(\delta^2 /R)$  term since each client runs vanilla SGD on their local data (the SGD-Update procedure in Algorithm 1). One medication for this variance term could be deploying variance reduction or shuffling data locally at each client before applying SGD. We notice that there is a recent work [56] also considering the client-level shuffling idea, but our work differs from it in two aspects: 1) they work with local SGD type algorithm and the shuffling idea is employed for model averaging within a subset of clients, while in our algorithm, during each local update period, each client runs shuffling SGD directly on other's model 2) from a theoretical perspective, we are mostly interested in investigating whether the algorithm can converge to the true optimal solution of  $\Phi (\pmb {\alpha}_i^*,\pmb {v})$  if we only optimize on a surrogate function  $\Phi (\hat{\alpha}_i,\pmb {v})$ .
One drawback of Algorithm 1 is that we have to wait for Algorithm A1 to finish and output  $\hat{\alpha}_i$ , so that we can proceed with Algorithm 1. However, if we are not satisfied with the precision of  $\hat{\alpha}_i$ , we may not have a chance to go back to refine it. Hence in the next subsection, we propose to interleave Algorithm 1 and Algorithm A1, and introduce a single-loop variant of PERM which will jointly optimize mixture weights and learn personalized models in an interleaving fashion.
3.2 A unified single loop algorithm
We now turn to introducing a single-stage algorithm that jointly optimizes  $\alpha_{i}$ s and  $v_{i}$ s as depicted in Algorithm 2 by intertwining the two stages in Algorithm A1 and Algorithm 1 in a single unified method. The idea is to learn the global model, which is used to estimate mixing parameters, concurrent to personalized models. At each communication round, the clients compute gradients on the global model, on their data, after the server collects these gradients does a step mini-batch SGD update on the global model, and then updates the mixing parameters. Then we proceed to update the personalized models similar to Algorithm 1. We note that, unlike the two-stage method where the mixing parameters are computed at the final global model, here the mixing parameters are updated adaptively based on intermediate global models.
Theorem 3. Let Assumptions 1 to 4 to be satisfied. Assume  $\alpha_{i}^{*}$  is the solution of (2). Then if we run Algorithm 2 with  $\eta = \Theta\left(\frac{\log(NKR^3)}{\mu R}\right)$  and  $\gamma = \Theta\left(\frac{\log(NKR^3)}{\mu R}\right)$ , it will output the solution  $\hat{v}_i$
Algorithm 2: Single Loop PERM
Input: Clients 1,...,  $N$ , Number of Local Steps  $K$ , Number of Epoch  $R$ , Initial mixing parameter  $\alpha_{1}^{0} = \dots, \alpha_{N}^{0} = \bar{\alpha} = [1 / N, \dots, 1 / N]$ .
Epoch for  $r = 0,\dots,R - 1$  do
Server generates permutation  $\sigma_r:[N]\mapsto [N]$
parallel for Client  $i = 1,\dots,N$  do
Client  $i$  sets initial model  $\pmb{v}_i^{r,0} = \pmb{v}_i^r$ .
for  $j = 1,\dots,N$  do
Set indices  $i_j = \sigma_r((i + j) \mod N)$ .
Server sends  $\pmb{v}_i^{r,j}$  to client  $i_j$ .
$\pmb{v}_i^{r,j+1} = \text{SGD-Update}(\pmb{v}_i^{r,j}, \eta, i_j, K, \pmb{\alpha}_i^r)$ . // Personalized model update
Client  $i$  does projection:  $\pmb{v}_i^{r + 1} = \mathcal{P}_{\mathcal{W}}(\pmb{v}_i^{r,N})$
$\pmb{w}^{r+1} = \mathcal{P}_{\mathcal{W}}(\pmb{w}^r - \gamma \frac{1}{N} \sum_{i=1}^{N} \frac{1}{M} \sum_{j=1}^{M} \nabla f_i(\pmb{w}^r, \xi_{i,j}^r))$  // Global model update
Compute  $\alpha_{i}^{r + 1}$  by running  $T_{\alpha}$  steps GD on  $g_{i}(\pmb{w}^{r + 1},\alpha)$  //  $\alpha$  update
Output:  $\hat{\pmb{v}}_i = \mathcal{P}_{\mathcal{W}}(\pmb {v}_i^R -(1 / L)\nabla_{\pmb{v}}\Phi (\pmb {\alpha}_i^R,\pmb {v}_i^R)),\hat{\alpha}_i = \pmb {\alpha}_i^R,\forall i\in [N].$
SGD-Update  $(\pmb {v},\eta ,j,K,\pmb {\alpha})$
Initialize  $\pmb{v}_j^0 = \pmb{v}$
for  $t = 0,\dots,K - 1$  do
$\pmb{v}^{t} = \pmb{v}^{t - 1} - \eta \alpha (j)N\nabla f_{j}(\pmb{v}_{j}^{t - 1};\xi_{j}^{t - 1})$
Output:  $v^{K}$
$\forall i \in [N]$ , such that with probability at least  $1 - p$ , the following statement holds:
$$
\begin{array}{l} \mathbb {E} \left[ \Phi \left(\boldsymbol {\alpha} _ {i} ^ {*}, \hat {\boldsymbol {v}} _ {i}\right) - \Phi \left(\boldsymbol {\alpha} _ {i} ^ {*}, \boldsymbol {v} _ {i} ^ {*}\right) \right] \leq O \left(\frac {L D ^ {2}}{N K R ^ {3}}\right) + \tilde {O} \left(\left(\frac {\kappa^ {4} L}{R ^ {2}} + \frac {N L}{\mu^ {2} R ^ {2}}\right) G ^ {2} N \log (1 / p) + \frac {L \delta^ {2}}{\mu R}\right) \\ + \kappa_ {\Phi} ^ {2} L \tilde {O} \left(\frac {\kappa^ {2} \kappa_ {g} ^ {2} L ^ {2} \bar {\zeta} _ {i} (\boldsymbol {w} ^ {*}) D G}{R} + R ^ {2} \exp \left(- \frac {T _ {\boldsymbol {\alpha}}}{\kappa_ {g}}\right) + \frac {L \kappa^ {2} \kappa_ {g} ^ {2} \bar {\zeta} _ {i} (\boldsymbol {w} ^ {*}) \delta^ {2}}{\mu^ {2} M}\right), \\ \end{array}
$$
where  $\kappa = \frac{L}{\mu}, \kappa_g = \frac{n_{\mathrm{max}}}{2\lambda}, \kappa_\Phi = \frac{\sqrt{N}G}{\mu}$  and the expectation is taken over the randomness of stochastic samples in Algorithm 2. That is, to guarantee  $\mathbb{E}[\Phi(\pmb{\alpha}_i^*, \hat{\pmb{v}}_i) - \Phi(\pmb{\alpha}_i^*, \pmb{v}_i^*)] \leq \epsilon$ , we choose  $M = O\left(\frac{L^2\kappa^2\kappa_g^2\kappa_\Phi^2\bar{\zeta}_i(\pmb{w}^*)\delta^2}{\mu^2\epsilon}\right)$ ,  $R = O\left(\max\left\{\frac{L\delta^2}{\mu\epsilon}, \frac{\kappa_\Phi^2\kappa^2\kappa_g^2L^3\bar{\zeta}_i(\pmb{w}^*)DG}{\epsilon}\right\}\right)$  and  $T_\alpha = O\left(\kappa_g\log\left(\frac{LR^2}{\epsilon}\right)\right)$ .
Compared to Theorem 2, we achieve a slightly worse rate, since we need a large mini-batch when we update global model  $\boldsymbol{w}$ . However, the advantages of the single-loop algorithm are two-fold. First, as we mentioned in the previous subsection, we have the freedom to optimize  $\hat{\alpha}_i$  to arbitrary accuracy, while in double loop algorithm (Algorithm A1 + Algorithm 1), once we get  $\hat{\alpha}_i$ , we do not have the chance to further refine it. Second, in practice, a single-loop algorithm is often easier to implement and can make better use of caches by operating on data sequentially, leading to improved performance, especially on modern processors with complex memory hierarchies.
3.3 Extension to heterogeneous model setting
In the homogeneous model setting, we assumed a shared model space  $\mathcal{W}$  for clients and the server. However, in real-world FL applications, devices have diverse resources and can only train models that match their capacities. We demonstrate that the PERM paradigm can be extended to support learning in model-heterogeneous settings, where different models with varying capacities are used by the server and clients. Focusing on learning the global optimal model to estimate pairwise statistical discrepancies, we note that by utilizing partial training methods [28], where at each communication round a sub-model with a size proportional to resources of each client is sampled from the server's global model (extracted either random, static, or rolling) and is transmitted to be updated locally. Upon receiving updated sub-models, the server can simply aggregate (average) heterogeneous sub-model updates sent from the clients to update the global model. We can consider the complexity of
(a) Personalized Accuracy
(b) Personalized Loss
Figure 1: Comparative analysis of personalization methods, including our single-loop PERM algorithm, localized FedAvg, perFedAvg, and pFedME, with synthetic data. The disparity in personalized accuracy and loss highlights PERM's capability to leverage relevant client correlations.
models used by clients when estimating mixing parameters by solving a modified version of (2) as:
$$
\sum_ {j = 1} ^ {N} \alpha (j) \sqrt {\mathsf {V C} (\mathcal {H} _ {j}) / n _ {j}} + \sum_ {j = 1} ^ {N} \alpha (j) \| \nabla f _ {i} (\boldsymbol {m} _ {i} \odot \boldsymbol {w} ^ {*}) - \nabla f _ {j} (\boldsymbol {m} _ {j} \odot \boldsymbol {w} ^ {*}) \| ^ {2} + \lambda \sum_ {j = 1} ^ {N} \alpha (j) ^ {2} / n _ {j},
$$
where we simply upper bounded the Rademacher complexity w.r.t. each data source in (GEN) with VC dimension [57]. Here  $m_{i}$  is the masking operator to extract a sub-model of the global model to compute local gradients at client  $i$  based on its available resources. By doing so, we can adjust mixing parameters based on the complexity of underlying models, as different sub-models of the global model (i.e.,  $m_{i} \odot w^{*}$  versus  $m_{j} \odot w^{*}$ ) are used to compute drift between pair of gradients at the optimal solution. With regards to training personalized models with heterogeneous local models, as we solve a distinct aggregated empirical loss for each client by interleaving permutations and shuffling models, we can utilize different model spaces  $\mathcal{W}_{i}, i = 1,\dots ,N$  for different clients that meet their available resources with aforementioned partial training strategies.
In this section we benchmark the effectiveness of PERM on synthetic data with 50 clients, where it notably outshone other renowned methods as evident in Figure 1. Our experiments concluded with the CIFAR10 dataset, employing a 2-layer convolutional neural network, where PERM, despite a warm-up phase, demonstrated unmatched convergence performance (Figure 2). Additional experiments are reported in the appendix. Across all datasets, the PERM algorithm consistently showcased its robustness and unmatched efficiency in the realm of personalized federated learning.
Experiment on synthetic data. To demonstrate the superior effectiveness of our proposed single-loop PERM algorithm compared to other existing personalization methods, we conducted an experiment using synthetic data generated according to the following specifications. We consider a scenario with a total of  $N$  clients, where we draw samples from the distribution  $\mathcal{N}(\pmb{\mu}_1, \pmb{\Sigma}_i)$  for half of the clients, denoted by  $i \in [1, \frac{N}{2}]$ , and from  $\mathcal{N}(\pmb{\mu}_2, \pmb{\Sigma}_i)$  for the remaining clients, denoted by  $i \in (\frac{N}{2}, N]$ . Following the approach outlined in [58], we adopt a uniform variance for all samples, with  $\Sigma_{k,k} = k^{-1.2}$ . Subsequently, we generate a labeling model using the distribution  $\mathcal{N}(\pmb{\mu}_w, \pmb{\Sigma}_w)$ .
Given a data sample  $\pmb{x} \in \mathbb{R}^d$ , the labels are generated as follows: clients 1, ...,  $\frac{N}{2}$  assign labels based on  $y = \mathrm{sign}(\pmb{w}^\top \pmb{x})$ , while clients  $\frac{N}{2} + 1, \dots, N$  assign labels based on  $y = \mathrm{sign}(-\pmb{w}^\top \pmb{x})$ . For this specific experiment, we set  $\mu_1 = 0.2$ ,  $\mu_2 = -0.2$ , and  $\mu_w = 0.1$ . The data dimension is  $d = 60$ , and there are 2 classes in the output. We have a total of 50 clients, each generating 500 samples following the aforementioned guidelines. We train a logistic regression model on each client's data.
To demonstrate the superiority of our PERM algorithm, we conducted a performance comparison against other prominent personalized approaches, including the fined-tuned model of FedAvg [14] (referred to as localized FedAvg), perFedAg [9], and pFedME [7]. The results in Figure 1 highlight PERM's efficient learning of personalized models for individual clients. In contrast, competing methods relying on globally trained models struggle to match PERM's effectiveness in highly heterogeneous scenarios, as seen in personalized accuracy and loss. This showcases PERM's exceptional ability to leverage relevant client learning.
(b) Personalized Loss
Experiment on CIFAR10 dataset. We extend our experimentation to the CIFAR10 dataset using a 2-layer convolutional neural network. During this test, 50 clients participate, each limited to data from just 2 classes, resulting in a pronounced heterogeneous data distribution. We benchmark our algorithm against PerFedAvg, PFedMe, and the localized FedAvg. As illustrated in Figure 2, PERM demonstrates superior convergence performance compared to other personalized strategies. It's noteworthy that PERM's initial personalized validation is significantly lower than that of approaches like PerFedAvg and PFedMe. This discrepancy stems from our choice to implement 10 communica
Figure 2: Comparative analysis of our single-loop PERM algorithm, localized FedAvg, PFedMe, and perFedAg, on CIFAR10 dataset and a 2-layer CNN model. Each client has access to only 2 classes of data. PERM rapidly catches up after 10 rounds of warmup without personalization involved.
Figure 3: Runtime of different algorithms in a limited environment. We compare PERM (single loop), PerFedAvg, FedAvg, and pFedMe. PERM has a minimal overhead over FedAvg and is comparable to other personalization methods.
tion rounds as a warm-up phase before initiating personalization, whereas other models embark on personalization right from the outset.
Computational overhead. In demonstrating the computational efficiency of the proposed PERM algorithm, we present a comparison of wall-clock time of completing one round of communication of PERM and other methods. Each method undertakes 20 local steps along with their distinct computations for personalization. As depicted in Figure 3, the PERM (single loop) algorithm's runtime is compared against personalization methods such as PerFedAvg, FedAvg, and pFedMe. Remarkably, PERM maintains a notably minimal computational overhead. The run-time is slightly worse due to overhead of estimating mixing parameters.
5 Discussion & Conclusion