Year: 2021

Abstract:
Federated learning (FL) is vulnerable to heterogeneously distributed data, since a common global model in FL may not adapt to the heterogeneous data distribution of each user. To counter this issue, personalized FL (PFL) was proposed to produce dedicated local models for each individual user. However, PFL is far from its maturity, because existing PFL solutions either demonstrate unsatisfactory generalization towards different model architectures or cost enormous extra computation and memory. In this work, we propose federated learning with personalized sparse mask (FedSpa), a novel PFL scheme that employs personalized sparse masks to customize sparse local models on the edge. Instead of training an intact (or dense) PFL model, FedSpa only maintains a fixed number of active parameters throughout training (aka sparse-to-sparse training), which enables users’ models to achieve personalization with cheap communication, computation, and memory cost. We theoretically show that with the rise of data heterogeneity, setting a higher sparsity of FedSpa may potentially result in a better convergence on its personalized models, which also coincides with our empirical observations. Comprehensive experiments demonstrate that FedSpa significantly saves communication and computation costs, while simultaneously achieves higher model accuracy and faster convergence speed against several state-of-the-art PFL methods.

text: 1 INTRODUCTION
Data privacy raises increasingly intensive concerns, and governments have enacted legislation to regulate the privacy intrusion behavior of mobile users, e.g., the General Data Protection Regulation (Voigt & Von dem Bussche, 2017). Traditional distributed learning approaches, requiring massive users’ data to be collected and transmitted to a central server for training model, soon may no longer be realistic under the increasingly stringent regulations on users’ private data. On this ground, federated learning (FL), a distributed training paradigm emerges as a successful solution to cope with privacy concerns, which allows multiple clients to perform model training within the local device without the necessity to exchange the data to other entities. In this way, the data privacy leakage problem could be potentially relieved.
Despite the promising prospect, several notorious issues are afflicting practical performance of FL: • The global model produced by weight average (or FedAvg and its non-personalized variants)
exhibits unsatisfactory performance in a Non-IID data distribution setting. To alleviate this problem, the most popular idea is to integrate personalized features into the global model, and produce dedicated model for each local distribution. However, how to make this integration is an open problem that remains unresolved. Prior works on personalized FL (PFL) zero in this issue, but the existing methods either demonstrate weak generalization towards different model architectures (Arivazhagan et al., 2019), or require extra computation and storage (Li et al., 2021). • The communication and training overhead is prohibitively high for both the FL and PFL. Clients in FL/PFL responsible for model training are mostly edge-devices with limited computation capacity and low bandwidth, and may not be powerful enough to fulfill a modern machine learning
task with large deep neural networks. Existing studies (Li et al., 2020; Vahidian et al., 2021) integrate model compression into FL/PFL to save communication and computation overhead. However, both methods embrace the technique of dense-to-sparse training, which still requires a large amount of communication at the beginning of training. In addition, how to effectively aggregate the dynamic sparse models is another challenging problem that remains unresolved.
In this work, we propose FedSpa (see Figure 2), which has two key features to counter the above two challenges: (i) FedSpa does not deploy a single global model, but allows each client to own its unique sparse model masked by a personalized mask, which successfully alleviates the Non-IID challenge. (ii) FedSpa allows each client to train over an evolutionary sparse model with constant sparsity1 throughout the whole federated training process, which consistently alleviates the computation overhead of clients. Besides, all the local models in FedSpa are sparse models, which requires a smaller amount of communication cost in each communication round. Theoretically, we conclude that with the rise of Non-IID extent, setting a higher sparsity may result in a better convergence on the personalized models of FedSpa. Empirically, in the Non-IID setting, we demonstrate that FedSpa accelerates the convergence (respectively 76.2% and 38.1% less communication rounds to reach the best accuracy of FedAvg (McMahan et al., 2016) and Ditto (Li et al., 2021)), in-
creases the final accuracy (up to 21.9% and 4.4% higher accuracy than FedAvg and Ditto, respectively), reduces the communication overhead (50% less parameters communicated than the dense solutions), and lowers the computation (15.3% lower floating-point operations (FLOPs) than algorithms trained with fully dense model). To the end, we summarize our contribution as:
• We present a novel formulation of the sparse personalized FL (SPFL) problem, which can be applied to various network architectures by enforcing personalized sparse masks to a global model.
• We propose a solution dubbed as FedSpa to solve the SPFL problem. By our novel design, FedSpa reduces the communication and computation overhead of the general FL solution.
• Two sparse-to-sparse mask searching techniques are integrated as plugins of our solution. To adapt our PFL training context, we modify the DST-based mask searching technique to enable a warm-start of the searching process, which achieves superior performance.
• We theoretically present the convergence analysis of the personalized models. Experimental results demonstrate the superiority of FedSpa and also coincides with the theoretical conclusion – with the rise of data heterogeneity, setting a higher sparsity of FedSpa may potentially result in a better convergence on its personalized models.

heading: 2 RELATED WORKS
text: Federated learning (FL) (McMahan et al., 2016) is seriously afflicted by the issue of heterogeneously distributed (or Non-IID) data. Personalized FL (PFL), initiated by recent literature (Li et al., 2021; Arivazhagan et al., 2019), is shown to be effective to counter this issue of FL. In this work, we propose an alternative yet effective way to enhance PFL with personalized sparse models.

heading: 2.1 PERSONALIZED FEDERATED LEARNING
text: We categorize PFL into five genres. Firstly, PFL via layer partition, e.g., FedPer (Arivazhagan et al., 2019), LG-FedAvg (Liang et al., 2020), FedRep (Collins et al., 2021), is to divide the global model layers into shared layers and personalized layers. For the shared layers, weights average as in FedAvg is adopted, while for personalized layers, models are trained only locally and will not
1Sparsity specifies the ratio of parameters that are set to 0 (or inactive) in a model.
be exchanged with others. Secondly, PFL via regularization, e.g., Ditto (Li et al., 2021), L2GD (Hanzely & Richtárik, 2020) is to add a proximal term on the local model to force the local model and global model closely in the local model fine-tuning stage. Similarly, Sarcheshmehpour et al. (2021); SarcheshmehPour et al. (2021) propose a total variation (TV) regularization to form the network lasso (nLasso) problem, and primal-dual methods adapted from (Jung, 2020) are proposed to solve the nLasso problems. Thirdly, PFL via model interpolation, e.g., MAPPER (Mansour et al., 2020), APFL (Deng et al., 2020) achieves personalization by linearly interpolating the weights of the cluster (global) model and local model as the personalized model. Fourthly, PFL via transfer learning, e.g., FedMD (Li & Wang, 2019), FedSteg (Yang et al., 2020), and Fedhealth (Chen et al., 2020), is to either use model and domain-specific local fine-tuning or knowledge distillation to adapt the global model into the personalized model. Finally, PFL via model compression, e.g., LotteryFL (Li et al., 2020) and Sub-FedAvg (Vahidian et al., 2021), achieves personalization via employing principle model compression techniques, such as weight pruning and channel pruning, over the shared global model.

heading: 2.2 SPARSE DEEP NEURAL NETWORKS
text: Methods to Sparsify neural networks can be classified into two genres: dense-to-sparse methods and sparse-to-sparse methods. Dense-to-sparse methods train from a dense model, and compress the model along the training process. Iterative pruning, first proposed by Frankle & Carbin (2018), shows promising performance in dynamically searching for a sparse yet accurate network. Recently, sparse-to-sparse methods have been proposed to pursue training efficiency. Among them, dynamic sparse training (DST) (Bellec et al., 2018; Evci et al., 2020; Liu et al., 2021) is the most successful technique that allows sparse networks, trained from scratch, to match the performance of their dense equivalents. Stemming from the first work – sparse evolutionary training (Mocanu et al., 2018; Liu et al., 2020), DST has evolved as a class of sparse training methods absorbing many advanced techniques, e.g., weight redistribution (Mostafa & Wang, 2019; Dettmers & Zettlemoyer, 2019), gradient-based regrowth (Dettmers & Zettlemoyer, 2019; Evci et al., 2020), and extra weight exploration (Jayakumar et al., 2020; Liu et al., 2021).
Our work also achieves personalization via model compression. We emphasize that three main progresses are made towards SOTA compression-based PFL: (i) We rigorously formulate the sparse personalized FL problem, filling the gap left by the prior works. (ii) While prior works either vaguely describe their model aggregation as ”aggregating the Lottery Ticket Network via FedAvg” (Li et al., 2020), or ”taking the average on the intersection of unpruned parameters in the network” (Vahidian et al., 2021), we explicitly formulate the aggregation as averaging the sparse update from clients. (iii) Both the two prominent prior works use the idea of iterative pruning to prune the network from dense to sparse. We instead provide two sparse-to-sparse training alternatives to plug in our solution, which largely reduces the costs of communication at the beginning of the training process, and exhibits remarkable performance.

heading: 3 PROBLEM FORMULATION
text: We assume a total number of K clients within our FL system, and we consistently use k to index a specific client. First, we give a preliminary introduction on the general FL and PFL problem.
General FL problem. Letw ∈ Rd be the global weight. General FL takes the formulation as below
(P1) min w
f̃(w) = 1
K K∑ k=1 { F̃k(w) := E[L(x,y)∼Dk(w; (x, y))] } ,
whereD = D1∪· · ·∪DK is the joint distribution of k local heterogeneous distributions, data (x, y) is uniformly sampled according to distribution Dk wrt. client k, and L(·; ·) is the loss function. Ultimate PFL Problem. Let {wk} be the personalized models. The ultimate PFL is defined as:
(P2) min {w1,··· ,wK}
f̂(w1, · · · ,wK) = 1
K K∑ k=1 { F̂k (wk) := E[L(x,y)∼Dk(wk; (x, y))] } ,
according to (Zhang et al., 2020; Hanzely et al., 2021). The problem could be separately solved by individual client with no communication. However, if the local data is insufficient, poor performance could be observed by this direct solution, since the local models cannot be boosted by other clients.
Regularized PFL problem. Regularized PFL can ease the heterogeneous challenges encountered by general FL, while escaping the curse of insufficient samples encountered by the ultimate PFL problem. Inspired by (Chen & Chao, 2021), we formulate the Regularized PFL problem as follows:
(P3) min {w1,··· ,wK}
f̄(w1, · · · ,wK) = 1
K K∑ k=1 { F̄k (wk) := E[L(x,y)∼Dk(wk; (x, y))] } +R (·) ,
where wk denote the personalized models and R(·) is the regularizer that enables information exchange between clients, making the problem tractable even the local samples are insufficient. However, it remains controversial about how to define the regularizer. Also, the gap between regularized PFL problem (P3) and the ultimate PFL problem (P2) still remains unspecified in existing PFL study.
In this work, our ultimate goal is to solve problem (P2), which requires information exchange between clients to ensure effective solution. Below, instead of utilizing regularizer as in (P3), we alternatively propose a novel Sparse PFL (SPFL) problem to reach the same goal.
Sparse PFL problem. By introducing personalized masks into FL, we derive the SPFL problem as
(P4) min w
f(w) = 1
K K∑ k=1 { Fk(m ∗ k w) := E[L(x,y)∼Dk(m ∗ k w; (x, y))] } ,
where m∗k ∈ {0, 1}d is a personalized sparse binary mask for k-th client. denotes the Hadamard product for two given vectors. Our ultimate goal is to find a global model w, such that the personalized model for k-th client can be extracted from the global model by personalized mask m∗k, i.e., m∗k w. The element of m∗k being 1 means that the weight in the global model is active for k-th personalized model, otherwise, remains dormant. Thus, the information exchange between all personalized models is enforced by a shared global model w.
Compared with existing PFL algorithms, solving our SPFL problem (P4) does not sacrifice additional computation and storage overhead of clients, since we do not maintain both personalized local models and global model in clients as in (Li et al., 2021; Mansour et al., 2020). On contrary, the solution to our problem could potentially lower the communication and computation overhead. Moreover, SPFL problem (P4) can be applied to most of the model architectures without modelspecific hyper-parameter tuning, since we do not make model-specific separation of the public and personalized layer in (Arivazhagan et al., 2019; Liang et al., 2020; Collins et al., 2021), or domainspecific fine-tuning in (Chen et al., 2020; Yang et al., 2020).

heading: 4 FEDSPA: SOLUTION FOR SPFL
text: In this section, we first introduce our proposed FedSpa in Algorithm 1. Then, we specify the update rule of global model, and two sparse-to-sparse mask searching methods that can be plug in the update process. At last, we give a theoretical analysis on evaluating the quality of the iterates of FedSpa with respect to the ultimate PFL problem (P2).

heading: 4.1 GLOBAL MODEL UPDATE RULE FOR FEDSPA
text: Data Parallel-based Update. We first propose the following iterative update to solve problem (P4)
wt+1 = wt − η
K K∑ k=1 m∗k ∇w̃k,tL(w̃k,t; ξk,t), (1)
where ξk,t is a batch of data that is uniformly sampled from the k-th client’s local distribution Dk, η is the learning rate, and w̃k,t = m∗k wt is the sparse weights sparsified by maskm∗k. However,
the optimal personalized masks {m∗k} are generally not accessible to us in the solution process. Let mk,t be an intermediate surrogate personalized mask ofm∗k, we rewrite Eq. (1) as:
wt+1 = wt − η
K K∑ k=1 mk,t ∇w̃k,tL(w̃k,t; ξk,t). (2)
For our proposed update rule, it is worth mentioned that: (i) Some coordinates of the model weights have been made zero before doing the forward process, i.e., not all the parameters have to be involved when calculating L(w̃k,t; ξk,t). This means that the computation overhead in the forward process could be potentially saved. (ii) In the backward process, the stochastic gradient ∇w̃k,tL(w̃k,t; ξk,t) is masked again bymk,t, which means that we do not need to backward the gradient for those sparse coordinates. Thus, the computation cost can be largely saved.
FL-adapted Update. To save the communication overhead, we integrate the idea from local SGD (Stich, 2018) and partial participation to our solution. Let w̃k,t,τ denote the weights before doing (τ+1)-th step of local SGD and set w̃k,t,0 = mk,t wt (i.e., the local weights will be synchronized every N steps with the global weights). Then, each client k ∈ St updates its model as below:
w̃k,t,τ+1 = w̃k,t,τ − ηmk,t ∇w̃k,t,τL(w̃k,t,τ ; ξk,t,τ ), τ = 0, 1, . . . , N − 1, (3)
where ξk,t,τ is a batch of sampled data in τ -th step at round t. After the local training is finished, the models of participated clients are updated and aggregated to the global model in server as follows:
wt+1 = wt − 1 |St| ∑ k∈St (w̃k,t,0 − w̃k,t,N ), (4)
where St is the set of clients selected to be participant in round t. According to Eq. (3), the update synchronized to the server (i.e., w̃k,t,0 − w̃k,t,N ), and the model distributed to clients (i.e., w̃k,t,0) are all sparse with a constant sparsity. Therefore, the communication overhead over synchronization could be largely saved. At last, we summarize our proposed FedSpa in Algorithm 1.
Algorithm 1 FedSpa Input Training iteration T ; Learning rate η; Local Steps N ; Random seed seed;
1: procedure SERVER’S MAIN LOOP 2: Randomly initialize global model w0 3: mk,0 = Initialization(seed) for k=0, 1, . . . ,K . Initialize masks via a mask-searching
solution 4: for t = 0, 1, . . . , T − 1 do 5: Uniformly sample a fraction of client into St 6: for each client k /∈ St do 7: mk,t+1 = mk,t . Inherit masks for round t+ 1 if not chosen 8: for k ∈ St do 9: Send w̃k,t,0 = mk,t wt to client k
10: Call Client k’s main loop and receive Uk,t andmk,t+1 11: wt+1 = wt − 1S ∑ k∈St Uk,t . Average and apply the update 12: procedure CLIENT’S MAIN LOOP 13: for τ = 0, 1, . . . , N − 1 do 14: Sample a batch of data ξk,t,τ from local dataset 15: gk,t,τ (w̃k,t,τ ) = ∇w̃k,t,τL(w̃k,t,τ ; ξk,t,τ ) 16: w̃k,t,τ+1 = w̃k,t,τ − ηmk,t gk,t,τ (w̃k,t,τ ) . Local update with fixed maskmk,t 17: Uk,t = w̃k,t,0 − w̃k,t,N 18: mk,t+1 = Next Masks(·) . Plug in a mask-searching solution to produce next masks 19: Send back Uk,t andmk,t+1

heading: 4.2 SPARSE-TO-SPARSE MASK SEARCHING TECHNIQUE
text: The framework of FedSpa is extensible. Specifically, we use the mask surrogate mk,t in Eq. (3) to perform the update, which allows us to plug in arbitrary mask searching techniques to determine
the iterating process of mk,t. In this work, we nominate two kinds of sparse-to-sparse training techniques: modified dynamic sparse training (DST) (Liu et al., 2021) and Random Static Masks (RSM), into FedSpa to reduce the computation and communication cost.
Modified DST for FedSpa. Our modified DST solution (see Algorithm 2) for FL follows these procedures. Firstly, randomly initialize the same mask for each client based on Erdós-Rényi Kernel (ERK) (Evci et al., 2020). Secondly, after local training, each client prunes out a number of unpruned weights with the smallest magnitude, and the number of weights being pruned is determined by a decayed pruned rate. Thirdly, recover the same amount of weights pruned in the last step. We follow the recovery process as in (Evci et al., 2020) by utilizing the gradient information to do the recovery. By our DST method, the number of sparse weights (aka. sparse volume) remains a constant (i.e., β) throughout the whole training process.
Algorithm 2 Modified DST for FedSpa Input Initial pruning rate α0; Set of Model layers J ;
1: procedure INITIALIZATION(seed) 2: Randomly initializemk,0 using the same random seed seed. 3: procedure NEXT MASKS(w̃k,t,N ) 4: Decay αt using cosine annealing with initial pruning rate α0 5: Sample a batch of data and backward the dense gradient g(w̃k,t,N ) 6: for layer j ∈ J do 7: Update maskm(j)
k,t+ 12 by zeroing out αt-proportion of weights with magnitude pruning
8: Update maskm(j)k,t+1 via recovering weights with gradient information g(w̃k,t,N )
9: Returnmk,t+1
Remark. We highlight our main modification over traditional DST techniques like Rigl (Evci et al., 2020) and Set (Mocanu et al. (2018)) to an FL context exist in two main aspects: (i) The pruning is performed individually by each client based on their local models, and the gradient used for weights recovery is derived using the client’s local training data. (ii) Once the next masks are generated, existing DST solutions immediately apply them to the local model weight. Indicated by Liu et al. (2021), by doing so, the recovered coordinate may need extra training steps to grow from 0 to a dense value. Our solution relieves this problem by applying the new mask on the global weights (which are dense), such that the recovered coordinates could have a dense initial value to warm-start.
Algorithm 3 RSM for FedSpa 1: procedure INITIALIZATION(seed) 2: Randomly initialize mk,0 us-
ing the same random seed seed 3: procedure NEXT MASKS 4: mk,t+1 = mk,t 5: Returnmk,t+1 RSM for FedSpa. RSM (shown in Algorithm 3) is basically fixingmk,t for all k ∈ [K] to the same randomly initialized mask, which remains unchanged during the whole training session. This solution also ensures the same sparse volume for all the clients throughout the training process, and could also reduce the computation and communication overhead as DST. Interestingly, within the setting of the homogeneous data distribution, we empirically show that RSM is more effective than DST in FedSpa.

heading: 4.3 THEORETICAL ANALYSIS
text: In this section, we shall introduce the convergence property of personalized models produced by FedSpa. We first give the following assumptions. Assumption 1 (Coordinate-wise bounded gradient dissimilarity). For any w̃ ∈ Rd, there exists a constant G ≥ 0 bounding the coordinate-wise gradient dissimilarity over all clients, i.e.,∥∥∥∇Fk(w̃)− 1K ∑Kk′=1∇Fk′(w̃)∥∥∥∞ ≤ G. Assumption 2 (Bounded variance). Assume that gk,t,τ (w̃) :=∇L(w̃; ξk,t,τ ) is an unbiased estimator of∇Fk(w̃) with bounded variance, i.e., E [ ‖gk,t,τ (w̃)−∇Fk(w̃)‖2 ] ≤σ2,∀k, t, τ, w̃ ∈ Rd.
Assumption 3 (L-smoothness). We assume L-smoothness over the client’s loss function, i.e., ‖∇Fk(w̃1)−∇Fk(w̃2)‖ ≤ L‖w̃1 − w̃2‖ holds for arbitrary w̃1, w̃2 ∈ Rd.
Assumption 4 (Coordinate-wise bounded gradient). Suppose the coordinate-wise gradient of each client is upper-bounded, i.e., ‖∇w̃Fk(w̃)‖∞ ≤ B.
Assumptions 2 and 3 are commonly used for characterizing the convergence of FL algorithms. We modify the other two assumptions slightly from their counterparts in existing FL literature (Xu et al., 2021, see Assumptions 3 and 5) by replacing the L2 norm with L∞ norm. These modifications are made to reveal the coordinate-wise gradient information in our analysis. Theorem 1 (Convergence of personalized models). Given the above assumptions, suppose the learning rate satisfies η ≤ 116LN , the optimal personalized masks m ∗ k ∈ {0, 1}d and the evolutionary masks mk,t ∈ {0, 1}d maintain the same sparse volume constraint, i.e., ||1 −m∗k||0 = ||1 −mk,t||0 = β, the personalized sparse models {w̃k,t} generated via FedSpa exhibits the following convergence property towards the optimal solution of ultimate PFL problem (P2):
1
TK T−1∑ t=0 K∑ k=1 E [ ‖∇Fk (w̃k,t) ‖2 ] ≤ 3(f (w0)− f (w ∗)) TηNκ + 3 T T−1∑ t=0 ρt + Υ, (5)
where κ = 12 − 150N 3η3L3 − 15N2η2L2 − 5NηL and ρt = (25N3η4L3 + 5N2η3L2
2 )(σ 2 + 18NΦt) + 4N2η2L+Nη 2K ∑ k dist(mk,t,m ∗ k)B 2 + 9N2η2LΦt + Nη2Lσ2
S , Φt = 1 K ∑ k((d− β)G
2︸ ︷︷ ︸ ≈0 if β→d
+ 1K ∑ k′ B
2(dist(mk,t,mk′,t︸ ︷︷ ︸ ≈0 if mk,t→mk′,t ) + dist(mk′,t,m ∗ k′)︸ ︷︷ ︸ ≈0 if mk′,t→m∗k′ )), Υ =
3(d− β)G2︸ ︷︷ ︸ ≈0 if β→d + 3βB2︸ ︷︷ ︸ ≈0 if β→0
+ 3K2 ∑ k ∑ k′ dist(m ∗ k,m ∗ k′)B
2︸ ︷︷ ︸ ≈0 if m∗k→m ∗ k′ , dist(m1,m2) is the hamming distance2
between masksm1 andm2.
Below, we give several comments on the above theorem. We focus our analysis on the on the second and third residuals, which do not vanish as communication round T →∞ • Impact of sparse volume. The result in Eq. (5) is highly related to setting of sparse volume,
i.e., β. The term βB2 would vanish as β → 0, while the term (d − β)G2 would disappear as β → d. When the data is highly heterogeneous, i.e., gradient dissimilarity is drastic, G could be prohibitively large to dominate the error. Then setting the sparsity level to a relatively large value, which results in a lower (d−β)G2, would potentially give a smaller non-vanished error. On contrary, when the data heterogeneity is mild, the error could be dominated by βB2. Then setting an excessively high sparsity might even hurt the performance. This conclusion is also evidenced by our experimental results (see Figure 1 in Section 1, where the optimal sparsity setting of FedSpa that achieves the highest accuracy is 0.5).
• Error by evolutionary mask searching. Recall that FedSpa uses evolutionary masks as a surrogate of optimal masks to perform update on the global model wt (see Eq. (3)). This replacement may bring additional errors to the convergence bound, as the hamming distance between the evolutionary masks and optimal masks (i.e., the term dist(mk,t,m∗k)) exists in our bound. • Error by dissimilar evolutionary masks. The theoretical result on the term dist(mk,t,mk′,t) in Φt indicates that the dissimilarity among the evolutionary masks may also play a role in optimizing the ultimate PFL problem (P2). This term could be minimized if the evolutionary masks remain the same throughout the training process, and would dominate the error if the masks are too distinct. This observation exactly motivates us to present the naive mask searching technique RSM in Algorithm 3, which forces all the clients to share the same random masks.

heading: 5 EXPERIMENTS
text: In this section, we conduct extensive experiments on verifying the efficacy of the proposed FedSpa. Our implementation of FedSpa is based on an open-source FL simulator FedML (He et al., 2020).

heading: 5.1 EXPERIMENTAL SETUP
text: Dataset. We evaluate the efficacy of FedSpa on EMNIST-Letter (EMNIST-L henceforth), CIFAR10, and CIFAR100 datasets. We simulate the client’s data distribution on Non-IID and IID setting. We
2Hamming distance measures the number of positions at which the two masks are different.
simulate two groups of Non-IID settings via γ-Dirichlet distribution, named setting A and setting B. Setting A and setting B respectively specify γ = 0.2, 0.1 for both EMNIST-L and CIFAR100, while specify γ = 0.5, 0.3 for CIFAR10. Details of our simulation setting are available in Appendix B.1.
Baselines. We compare our proposed FedSpa with four baselines, including FedAvg (McMahan et al. (2016)), Sub-FedAvg (Vahidian et al. (2021)), Ditto (Li et al. (2021)) and Local. We tune the hyper-parameters of the baselines to their best states. Specifically, the regularization factor of Ditto is set to 0.5. The prune rate each round, distance threshold, and accuracy threshold of Fed-Subavg are fixed to 0.05, 0.0001, 0.5, respectively. We ran 3 random seeds in our comparison.
Models and hyper-parameters. We use LeNet5 for EMNIST-L, VGG11 for CIFAR10, and ResNet18 for CIFAR100 in our experiment. We use a SGD optimizer with weight decayed parameter 0.0005. The learning rate is initialized with 0.1 and decayed with 0.998 after each communication round. We simulate 100 clients in total, and in each round 10 of them are picked to perform local training (the setting follows McMahan et al. (2016)). For all the methods except Ditto, local epochs are fixed to 5. For Ditto, in order to ensure a fair comparison, each client uses 3 epochs for training of the local model, and 2 epochs for global model training. The batch size of all the experiments is fixed to 128. For FedSpa, the pruning rate (i.e., αt) is decayed using cosine annealing with an initial pruned rate 0.5. The initial sparsity of layers is initialized by ERK with scale parameter 1.

heading: 5.2 MAIN PERFORMANCE EVALUATION
text: We fix the dense ratio of FedSpa (DST), FedSpa (RSM), and the final dense ratio of Fed-SubAvg both to 0.5 (i.e., 50% of parameters are pruned) in our main evaluation. Other hyper-parameters are fixed as default. Figure 3 and Table 1 illustrate the training performance of different algorithms on three datasets. We evaluate the performance based on the following metrics: Final Accuracy. In the Non-IID setting, we show that FedSpa (DST) achieves remarkable performance. Specifically, in Non-IID setting B of CIFAR100, FedSpa (DST) achieves respectively 4.4%, 11.9% and 21.9% higher final model accuracy, compared with Ditto, Sub-FedAvg and FedAvg. FedSpa (DST) seems to achieve better performance as the FL tasks becoming difficult (since better performance is observed in a higher Non-IID extent, and in datasets that are intrinsically more difficult). Interestingly, in the IID setting, we show that all the personalized solutions exhibit some extents of performance degradation, which become more significant as the dataset becomes challenging. The compression-based methods seem to be especially vulnerable in this setting. Our interpretation for this phenomenon is that: since the information exchange between clients would be limited by employing different sub-networks for training, the clients could not efficiently make an effective fusion on their models through parameter averaging. This hypothesis is substantiated by our experiment on FedSpa (RSM), an alternative implementation of FedSpa, which forces all the masks to maintain the same sub-network. FedSpa (RSM) achieves commensurate performance with FedAvg in the IID setting, outperforming the personalized solutions. The reason FedSpa (DST)
performing better in Non-IID setting can also be explained by Theorem 1, which concludes that in a Non-IID setting, setting proper sparsity for personalized models may promise a better convergence.
Convergence. As shown in Table 2, FedSpa achieves significantly faster convergence, which potentially saves the communication rounds to train a model from scratch to a specific accuracy.
Training FLOPs and Communication. From Table 1, FedSpa (DST) achieves 15.4%∼42.9% lower FLOPs than the dense solutions (e.g., Ditto, FedAvg), 13.0%∼28.2% lower communication overhead than another model compression solution Sub-FedAvg, and 50% lower communication than the dense solution. The edge of FedSpa (DST) stems from its training pattern – it is trained from a sparse model, with constant sparsity throughout the training process. However, it is interesting to see that the training FLOPS of Fed-SubAvg is considerably lower than FedSpa, even under the same sparsity setting. This phenomenon stems from our ERK initialization, which is essential for the high performance of our solution, for which we will have a further discussion in Appendix B.4.

heading: 6 CONCLUSIONS
text: In this paper, we propose FedSpa, a personalized FL solution that enables sparse-to-sparse training and efficient sub-model aggregation. As demonstrated by our experiments, FedSpa exhibits outstanding performance in the Non-IID setting, outperforming other existing solutions in terms of accuracy, convergence speed as well as communication overhead. Additionally, we present theoretical analysis to evaluate the error bound of FedSpa towards the ultimate PFL problem. Future direction includes designing new model aggregation solutions for the sparse sub-network, and new mask searching techniques specifically targeting on federated learning process.

heading: 7 REPRODUCIBILITY STATEMENT
text: For sake of reproducibility of our solution, we make the following efforts: (i) In Appendix B.1, we clearly state the data splitting method for IID and Non-IID data distribution. (ii) In Appendix B.2, the network models we used in our experiment are clearly described. (iii) In Appendix B.4, the proposed FedSpa with various hyper-parameters during the implementation are also systemically tested to demonstrate its stability and superiority. (iv) In Appendix C, we give the self-contained proof of Theorem 1. (v) At last, the source-code of FedSpa is enclosed in the supplementary material.

heading: A CRITICAL COMPONENTS OF FEDSPA (DST)
text: ERK initialization. In Algorithm 2, we use Erdós-Rényi Kernel (ERK) originally proposed by Evci et al. (2020) to initialize the sparsity of each layer. Specifically, the active parameters of the convolutional layer initialized by ERK is proportional to 1n 1−1+nl+wl+hl
nl−1∗nl∗wl∗hl , where n l−1, nl wl and
hl respectively specify number of input channels, output channels and kernel’s width and height in the l-th layer. For the linear layer, the number of active parameters scale with 1n l−1+nl
nl−1÷nl where n l−1
and nl are the number of neurons in the l − 1-th and l-th layer. This initialization basically allows the layer with less parameters have more proportion of active parameters.
Cosine annealing. Recall that we set the initial pruning rate as α0 and gradually decay it to 0 with cosine annealing Liu et al. (2021). The update of pruning rate with cosine annealing can be formalized as: αt = 0.5× α0 × ( 1 + cos ( t T−1π ))
. We perform this decay in order to ensure that the network (specifically, its active coordinates) would not experience drastic change on the later stage of training while ensuring that the mask searching is effective on the early stage of training.

heading: B DETAILED SETTING AND ADDITIONAL EXPERIMENTS

heading: B.1 DATA SPLITTING SETTING
text: In our implementation, we first split the training data (60k pieces of data for CIFAR10 and CIFAR100, and 145.6k for EMNIST-L, respectively) to clients for IID setting and Non-IID setting. For the IID setting, data are uniformly sampled for each client. For the Non-IID setting, we use γDirichlet distribution on the label ratios to ensure uneven label distributions among devices as (Hsu et al., 2019). The lower the distribution parameter γ is, the more uneven the label distribution will be, and would be more challenging for FL. After the initial splitting of training data, we sample 100 pieces of testing data from the testing set to each client. To simulate the personalized setting, each client’s testing data has the same proportion of labels as its training data. Testing of the personalized model is performed by each client based on their personalized data, and the overall testing accuracy is calculated as the average of all the client’s testing accuracy. In our experiment, we simulate different Non-IID settings. For CIFAR10, Non-IID setting A and B respectively specify γ = 0.5 and γ = 0.3. For EMNIST-L and CIFAR100, since the number of the total labels are bigger3, we use smaller γ, wherein setting A and B respectively specify γ = 0.2 and γ = 0.1.

heading: B.2 NETWORK ARCHITECTURES
text: We follow the Caffe’s implementation of LeNet5 4 (LeCun et al., 2015), VGG11 (Simonyan & Zisserman, 2014) and ResNet18 (He et al., 2016) to do the evaluation. Suggested by Hsieh et al. (2020), DNNs with batch normalization layers (Ioffe & Szegedy, 2015) are particularly vulnerable to the Non-IID setting, suffering significant model quality loss in the FL process. Following the recommendation from Hsieh et al. (2020), we use group normalization (Wu & He, 2018) to substitute the original batch normalization layer in both ResNet18 and VGG11.

heading: B.3 BASELINE DESCRIPTION
text: Below, we give a brief introduction of the baselines compared in our evaluations:
• FedAvg (McMahan et al., 2016) is the vanilla solution of FL. It utilizes weights average to enable all the clients to collaboratively train a global model, which efficiently absorbs knowledge from personal data resided in clients.
• Ditto (Li et al., 2021) is a personalized FL solution aiming to smooth the tension brought by the data heterogeneity problem of FL. Ditto achieves personalization via maintaining both the local models and global model. Specifically, within each round of iteration, each client first trains the global model based on its local empirical loss (which shares the same procedure as FedAvg).
326 and 100 labels respectively in EMNIST-L and CIFAR100, while only 10 labels in CIFAR10. 4Available in https://github.com/mi-lad/snip/blob/master/train.py
After the global model is updated, each client additionally trains its local model based on a loss function involving its local empirical loss and the proximal term towards the global model. This local training phase is used to extract the global knowledge into each client’s local model. Since each client has to maintain and train both local model and global model, Ditto might need extra computation and storage overhead to achieve its personalization.
• Local is the direct solution to the ultimate PFL problem (P2). Each client performs SGD based on its local data, and there is no communication between clients. To mimic the FL setting, we sample 10 out of 100 clients to do the local update on its local model after every 5 epochs of training (same with the number of local epochs in a communication round that is performed by other solutions). For sake of consistency, we still use 1 communication round to represent 5 local epochs of Local in our evaluation.
• Sub-FedAvg (Vahidian et al., 2021) is a prominent model compression-based PFL. Sub-FedAvg maintains personalized sub-networks for each client. Training of Sub-FedAvg starts from a fully dense model, and this solution iteratively prunes out the parameters and channels as the training progresses. Finally, the commonly shared parameters of each layer are removed, and only the personalized parameters that can represent the features of local data are kept.
• Subsampling (Konečnỳ et al., 2016) is a gradient-compression solution aiming to reduce the communication overhead of FL. The local training procedure is the same with FedAvg. The difference is that Subsampling does not communicate the intact model for aggregation, but only communicates the sparse gradient update to the server for aggregation. Explicitly, in each round, the sparse gradient update is produced through element-wisely multiplying a random mask. Different from FedSpa, the randomized mask is independently generated in each round, and would only be used to compress the gradient when uploading the gradient update (which means in the model distribution phase, the model distributed would not be sparsified, and therefore would not save the downlink communication cost).

heading: B.4 ABLATION STUDY
text: In this sub-section, we focus on the ablation study of FedSpa. Specifically, we study the impacts of dense ratio, different mask initialization methods, and the gradient-involved weight recovery procedure. Additionally, we present an interesting observation on the performance of the global model trained by our personalized solution. Our ablation study is done with ResNet-18 on CIFAR100.
Impact of sparsity (aka. sparse ratio). Fixing other components and hyper-parameters to the default value in our setup, we change the sparsity of FedSpa to 0.2, 0.4, 0.5, 0.6 and 0.8, to show its impact on the algorithm performance. Experimental results are available in Figure 4 and Table 3. By our report, we observe that sparsity may impact learning performance under different data distribution settings. For the IID setting, a higher sparsity seems to seriously degrade the training performance, while for the Non-IID setting, properly sparsifying the model may even enhance the final accuracy, and with a higher Non-IID extent, the benefit of sparsification reinforces. This observation surprisingly coincides with our theoretical conclusion given in Theorem 1, that setting the sparse volume to a proper value when gradient dissimilarity is large (i.e., under a high extent of Non-IID) could foster a better performance of personalized models. But too much sparsity, even in the highly non-iid setting (e.g. γ = 0.1) leads to performance degradation. On the contrary, while it is iid, the convergence could be dominated by the errors brought by sparsification, and setting the mask to a higher sparsity could possibly enlarge these existing errors.
Another more intuitive interpretation for the impact of sparse ratio is from the perspective of information exchange. Too much sparsification may limit the information exchange between the local sparse models. If all the clients maintain an extremely high sparsity, the intersected coordinates between clients’ local sparse models (or identically, their masks) would be small. Then the local update averaging process (see Eq. (4)), the only way to extract global knowledge into the local models, would not be effective. On contrary, while the sparsity is set to an extremely low value, the personalized features of local models could be eliminated, since only limited coordinates in their models are different.
ERK vs. Uniform sparsity initialization. Recall our mask initialization procedure in Algorithm 3 that the layer-wise sparsity is initialized by ERK. This in essence ensures that the layer-wise sparsity of a model is scaled with the number of parameters in a layer. Liu et al. (2021) confirms the out-
standing effect of ERK initialization in improving overall training performance over the centralized training primitive, but it remains unexplored how it performs in our proposed distributed training framework. Below, we show in Figure 5 how accuracy evolves with communication rounds under ERK and Uniform 5 initialization. As shown, a drastic drop of accuracy is observed by replacing ERK with Uniform, by which we conclude that ERK is an essential component for FedSpa (DST).
However, though a significant accuracy enhancement is observed, we note that integrating ERK may sacrifice potentially more FLOPS reduction. This observation can be found in Table 6, wherein our results show that initialization with Uniform can save 34.3% FLOPs of that with ERK.
Different or same mask initialization. Recall that based on the layer-wise sparsity calculated by ERK, FedSpa uses the same random seed to initialize the mask, so as to make the mask exploration
5Uniform enforces the same sparsity for all the layers in a model.
of all clients started from the same mask. In the following, we give another implementation that allows each client to share different masks in the beginning.
As shown in Figure 6, we surprisingly find that for FedSpa (DST), maintaining different masks in initialization may slightly enhance its training performance in IID and Non-IID (γ = 0.2) setting. We hypothesize that by different mask initialization, each client could more efficiently search for their optimal masks that better represents the features and labels of the personal data.
For FedSpa (RSM), compared with initialization using the same mask, different mask initialization may result in a drastic performance loss in the IID setting, and a significant improvement in the NonIID setting. With the same mask initialization of RSM, each client consistently trains based on the same sub-network, which completely eliminates personalization. So this setting shares a similar performance with FedAvg – with satisfactory performance in IID setting and rather weak performance in Non-IID setting. On contrary, by initializing different masks in the beginning, FedSpa (RSM) reserves some degrees of personalization, since only the intersected coordinates in their local models are shared and updated by the information exchange (i.e., average) process. Consequently. FedSpa (RSM) with different mask initializations has a similar performance pattern with FedSpa (DST).
Another interesting observation is that FedSpa (RSM) with different mask initialization cannot outperform FedSpa (DST) in both the two groups of Non-IID settings. This indicates that the DST mask searching process is effective to achieve a superior performance of FedSpa in Non-IID setting.
Weight recovery w/ or w/o gradient information. Recall that in FedSpa (DST), we proposed to use gradient information to recover the pruned weights, which is empirically proven in (Evci et al., 2020) to outperform its random recovery counterpart in Set (Mocanu et al., 2018). Specifically, for gradient information-based recovery, the weight coordinates with the top-αt magnitude of the gradient would be recovered, while for random recovery, the coordinates are recovered randomly. To demonstrate the impact of the weight recovery method over FedSpa (DST), in Figure 7, we compare the gradient information-based recovery with random recovery. Our experimental result demonstrates that recovery with gradient information could slightly accelerate the convergence and enhance the final accuracy in our FedSpa framework.
Global model vs. Personalized model. In our main experimental result, all the testings are conducted by clients based on their own personalized models. But it is interesting to evaluate whether
the global model trained by FedSpa itself could converge, or even could achieve commensurate performance with the global model trained by general FL solution (e.g. FedAvg). As demonstrated by Figure 8, we empirically find that in the IID setting, the global model trained by FedSpa cannot recover the performance of that trained by FedAvg, and a considerable performance drop is also observed in the Non-IID setting. Another observation is that the global model of FedSpa surprisingly maintains roughly the same performance as its personalized models in the IID setting, but conceivably suffers significant performance loss in the Non-IID setting. This corroborates our conclusion that sub-networks extracted from a global model may potentially outperform the full model, under the condition that the data distributions of clients are skewed (or heterogeneous).
Wall time. Recall that we do an additional mask searching procedure in FedSpa (DST), which might possibly induce extra wall time on the local devices. We show in Table 5 the wall time used for training and mask searching in one single local round. The sparsity used for this experiment is fixed to 0.5, while other parameters remain the default setting (see Section 5.1). We use one single 1080Ti to perform training for the GPU-based experiment, while the CPU-based experiment is conducted on an Intel(R) Xeon(R) CPU E5-2620 v4 @ 2.10GHz with 8 cores. Our experimental results confirm that the mask searching process only accounts for a small portion of wall time (approximately 5%− 20%) for the entire computation time on the local devices.

heading: C PROOF OF THEOREM 1
text: In this section, for sake of readability, we first clarify the notations and reiterate several facts that we use in our proof. Then we present several lemmas that are commonly used in the FL literature (see (Karimireddy et al., 2020; Xu et al., 2021)). Later, several key lemmas are listed with exhaustive proof, and finally, the proof of our main theorem is given by exploiting the listed lemmas and facts.

heading: C.1 NOTATIONS AND FACTS
text: Throughout the proof, we assume ∑ τ equivalent to ∑N−1 τ=0 , ∑ k equivalent to ∑K k=1, and ∑ k,τ
equivalent to ∑K k=1 ∑N−1 τ=0 unless otherwise specified. In our proof, we reuse most of the notations
from our problem formulation part. We use gk,t,τ (w̃k,t,τ ) = ∇w̃k,t,τL(w̃k,t,τ ; ξk,t,τ ) to denote the stochastic gradient of client k in round t and at step τ .
Then as per our formulation in Section 4.1, we reiterate the following facts, which would be heavily used in our proof. Fact 1 (Local step). As per Eq. (3), one local step of client’s update can be formalized as follows:
w̃k,t,τ+1 = w̃k,t,τ − ηmk,t gk,t,τ (w̃k,t,τ ), (6)
where gk,t,τ (w̃k,t,τ ) is the stochastic gradient over the sparse model weights w̃k,t,τ , and w̃k,t,0 = mk,t wt is the synchronized local weights at the beginning of a communication round. Fact 2 (Local update from client k). The local update of clients can be formalized as follows:
Uk,t = w̃k,t,0 − w̃k,t,N = η ∑ τ mk,t gk,t,τ (w̃k,t,τ ). (7)
Fact 3 (Server’s update). The server aggregates the sparse update by averaging, which can be formalized as follows:
wt+1 = wt − 1
S ∑ k∈St Uk,t = wt − η S ∑ k∈St,τ mk,t gk,t,τ (w̃k,t,τ ). (8)
Fact 4 (Global and local loss). The local loss of a client is denoted by Fk(w̃k), and is formulated as: Fk(w̃k) = E[L(x,y)∼Dk(w̃k; (x, y))] (9) where w̃k = m∗k w, and the global loss can be formalized as follows:
f(w) = 1
K K∑ k=1 Fk(w̃k) = 1 K K∑ k=1 Fk(m ∗ k w) (10)
and {m∗k} in the function are viewed as optimal personalized masks which satisfies the sparse volume constraint ‖1−m∗k‖0 = β for all k.

heading: C.2 AUXILIARY LEMMAS
text: In the following. we shall present several common lemmas that are heavily used in the FL literature. Lemma 1 (Cauchy-Schwarz). Assume arbitrary vector sequences {ak}k=1,...,K and {bk}k=1,...,K , Cauchy-Schwarz inequality implies:∥∥∥∥∥ K∑ k=1 akbk ∥∥∥∥∥ 2 ≤ ( K∑ k=1 ‖ak‖2 )( K∑ k=1 ‖bk‖2 ) , (11)
by taking bk = 1, we also have: ∥∥∥∥∥ K∑ k=1 ak ∥∥∥∥∥ 2 ≤ K ( K∑ k=1 ‖ak‖2 ) , (12)
Lemma 2 (Separating mean and variance, Lemma B.3 (Xu et al., 2021)). Let {a1, . . . ,aτ} be τ random vectors in Rd . Suppose that {ai − ξi} form a martingale difference sequence, i.e. E [ai − ξi | a1, . . . ,ai−1] = 0, and suppose that their variance is bounded by E [ ‖ai − ξi‖2 ] ≤
σ2. Then, the following inequality holds:
E ∥∥∥∥∥ τ∑ i=1 ai ∥∥∥∥∥ 2  ≤ 2 ∥∥∥∥∥ τ∑ i=1 ξi ∥∥∥∥∥ 2 + 2τσ2.
Lemma 3 (Relaxed triangle inequality, Lemma 3 (Karimireddy et al., 2020)). Let vi and vj be vectors in Rd. Then the following inequality holds true for any a > 0:
‖vi + vj‖2 ≤ (1 + a) ‖vi‖2 + ( 1 + 1
a
) ‖vj‖2 . (13)
Lemma 4. For random vector v1 satisfying E[v1] =  0·· 0 , and assume another random vector v2 is independent with v1, we have:
E[‖v1 + v2‖2] = E[‖v1‖2] + E‖v2‖2 (14)
Proof.
E[‖v1 + v2‖2] = E〈v1 + v2,v1 + v2〉 = E‖v1‖2 + E‖v2‖2 + 2E〈v1,v2〉 = E‖v1‖2 + E‖v2‖2 + 2〈Ev1,Ev2〉 = E‖v1‖2 + E‖v2‖2
(15)
This completes the proof.

heading: C.3 KEY LEMMAS
text: In this section, we present several important lemmas that would be used in our formal proof. All the presented claims are rigorously proved.
Lemma 5 (Smoothness of f(w)). Assume w̃k = m∗k w for any m∗k ∈ {0, 1}d, we have Lsmoothness for f(w) = 1K ∑ k Fk(w̃k), i.e., for any w1,w2 ∈ Rd, we have:
‖∇f(w1)−∇f(w2)‖ ≤ L‖w1 −w2‖. (16)
Proof.
‖∇f(w1)−∇f(w2)‖ = ∥∥∥∥∥ 1K∑ k (m∗k ∇Fk(m∗k w1)−m∗k ∇Fk(m∗k w2)) ∥∥∥∥∥ ≤ 1 K ∑ k ‖m∗k (∇Fk(m∗k w1)−∇Fk(m∗k w2))‖
≤ 1 K ∑ k ‖∇Fk(m∗k w1)−∇Fk(m∗k w2)‖
(a) ≤ L 1 K ∑ k ‖m∗k (w1 −w2)‖
≤ L‖w1 −w2‖
(17)
where the first equality holds by the Fact 4 and inequality (a) is due to Assumption 3. This completes the proof.
Lemma 6 (Separating mean and variance of stochastic gradient). For mk,t ∈ {0, 1}d, We have the following bounding for the expected average gradient:
Et  ∥∥∥∥∥∥ 1S ∑ k∈St,τ mk,t gk,t,τ (w̃k,t,τ ) ∥∥∥∥∥∥ 2  ≤ 2Et  ∥∥∥∥∥∥ 1S ∑ k∈St,τ mk,t ∇Fk(w̃k,t,τ ) ∥∥∥∥∥∥ 2 + 2Nσ2 S .
(18) where Et[·] denotes the expectation over all the randomness of round t.
Proof. We view 1Smk,t gk,t,τ (w̃k,t,τ ) for all k and τ as stochastic vectors. By the unbiasedness of stochastic gradient, we know their variance satisfies:
|| 1 S mk,t gk,t,τ (w̃k,t,τ )− E[ 1 S mk,t gk,t,τ (w̃k,t,τ )||2]
=|| 1 S mk,t (gk,t,τ (w̃k,t,τ )−∇Fk(w̃k,t,τ ))||2 ≤ 1 S2 ||gk,t,τ (w̃k,t,τ )−∇Fk(w̃k,t,τ )||2 ≤σ 2
S2
(19)
where the first inequality holds sincemk,t ∈ {0, 1}d and the last inequality is due to Assumption 2. As per the variance given above, and directly apply Lemma 2, the claim immediately shows.
Lemma 7. We have the following error bound by introducing personalized masks:
‖mk,t ∇Fk(w̃k,t,τ )−m∗k∇Fk(w̃k,t)‖2 ≤ ‖∇Fk(w̃k,t,τ )−∇Fk(w̃k,t)‖2 + dist(mk,t,m∗k)B2 (20)
and,
‖mk,t ∇Fk(w̃k,t)−∇f(wt)‖2 ≤ 3(d−β)G2+ 3
K ∑ k′ B2(dist(mk,t,mk′,t)+dist(mk′,t,m ∗ k′))
(21) where dist(m1,m2) is the hamming distance between two masks.
Proof. Let ∇(j)Fk(w̃k,t,τ ) and ∇(j)Fk(w̃k,t) be the derivative over the j-th weight coordinates. Define the support of a vector v ∈ Rd as supp(v) = {j : v(j) 6= 0}. The left hand side of the first claim can be rewritten and bounded as follows:
‖mk,t ∇Fk(w̃k,t,τ )−m∗k ∇Fk(w̃k,t)‖2 = ∑
j∈supp(mk,t)∩supp(m∗k)
‖∇(j)Fk(w̃k,t,τ )−∇(j)Fk(w̃k,t)‖2
+ ∑
j∈supp(mk,t)andj /∈supp(m∗k)
‖∇(j)Fk(w̃k,t,τ )‖2 + ∑
j /∈supp(mk,t)andj∈supp(m∗k)
‖∇(j)Fk(w̃k,t)‖2
(a) ≤ ∑
j∈supp(mk,t)∩supp(m∗k)
‖∇(j)Fk(w̃k,t,τ )−∇(j)Fk(w̃k,t)‖2 + ∑
j∈supp(mk,t⊕m∗k)
B2
(b) = ∑ j∈supp(mk,t)∩supp(m∗k) ‖∇(j)Fk(w̃k,t,τ )−∇(j)Fk(w̃k,t)‖2 + dist(mk,t,m∗k)B2
≤ ∑ j ‖∇(j)Fk(w̃k,t,τ )−∇(j)Fk(w̃k,t)‖2 + dist(mk,t,m∗k)B2
≤‖∇Fk(w̃k,t,τ )−∇Fk(w̃k,t)‖2 + dist(mk,t,m∗k)B2 (22)
where inequality (a) is true since we have ‖∇(j)Fk(w̃k,t,τ )‖2 ≤ B2 and ‖∇(j)Fk(w̃k,t)‖2 ≤ B2 by coordinate-wise bounded gradient Assumption 4. Inequality (b) is due to the definition of hamming distance, dist(v1,v2) = ∑ j∈supp(v1⊕v2). This shows our first claim.
Following the same technique, we expand the left hand side of the second claim using Cauchy Schwarz inequality in Lemma 1, as follows:
‖mk,t ∇Fk(w̃k,t)−∇f(wt)‖2
=‖mk,t ∇Fk(w̃k,t)− 1
K ∑ k′ m∗k′ ∇Fk′(w̃k′,t)‖2
≤3‖mk,t ∇Fk(w̃k,t)− 1
K ∑ k′ mk,t ∇Fk′(w̃k′,t)‖2
+ 3‖ 1 K ∑ k′ (mk,t ∇Fk′(w̃k′,t)−mk′,t ∇Fk′(w̃k′,t))‖2
+ 3‖ 1 K ∑ k′ (mk′,t ∇Fk′(w̃k′,t)−m∗k′ ∇Fk′(w̃k′,t))‖2
≤3‖mk,t ∇Fk(w̃k,t)− 1
K ∑ k′ mk,t ∇Fk′(w̃k′,t)‖2
+ 3
K ∑ k′ ‖mk,t ∇Fk′(w̃k′,t)−mk′,t ∇Fk′(w̃k′,t)‖2
+ 3
K ∑ k′ ‖(mk′,t ∇Fk′(w̃k′,t)−m∗k′ ∇Fk′(w̃k′,t))‖2
≤3(d− β)G2 + 3 K ∑ k′ dist(mk,t,mk′,t)B 2 + 3 K ∑ k′ dist(mk′,t,m ∗ k′)B 2
=3(d− β)G2 + 3 K ∑ k′ B2(dist(mk,t,mk′,t) + dist(mk′,t,m ∗ k′))
(23)
For the last inequality, we use the same technique as in the first claim to bound the second term and the third term, and the first term is bounded by Lemma 8.
Lemma 8. Let ‖1 −m‖0 = β be the sparse volume of a mask m, and suppose m ∈ {0, 1}d, by assumption 4, we have the following relation:
‖m (∇Fk(w̃k,t)− 1
K ∑ k′ ∇Fk′(w̃k′,t))‖2 ≤ (d− β)G2 (24)
and,
‖∇Fk(w̃k,t)− 1
K ∑ k′ m ∇Fk′(w̃k′,t)‖2 ≤ (d− β)G2 + βB2 (25)
where d is the dimension of the models (or identically, the dimension of the maskm).
Proof. The proof of this lemma follows a similar technique in Lemma 7. Let ∇(j)Fk(w̃k,t) be the derivative over the j-th weight coordinate. Define the support of a vector v ∈ Rd as supp(v) = {j : v(j) 6= 0}. The left hand side (L.H.S) of our first claim can be bounded by:
‖m (∇Fk(w̃k,t)− 1
K ∑ k′ ∇Fk′(w̃k′,t))‖2
= ∑
j∈supp(m)
‖∇(j)Fk(w̃k,t)− 1
K ∑ k′ ∇(j)Fk′(w̃k′,t)‖2
≤(d− β)G2
(26)
where the last inequality is obtained from the sparsity constraint and the coordinate-wise gradient dissimilarity Assumption 1. This shows our first claim.
For our second claim, we expand its L.H.S as follows:
‖∇Fk(w̃k,t)− 1
K ∑ k′ m ∇Fk′(w̃k′,t)‖2
= ∑
j∈supp(m)
‖∇(j)Fk(w̃k,t)− 1
K ∑ k′ ∇(j)Fk′(w̃k′,t)‖2 + ∑ j /∈supp(m) ‖∇(j)Fk(w̃k,t)‖2
≤(d− β)G2 + βB2,
(27)
where the last inequality is due to Assumption 1 and Assumption 4. It completes the proof.
Lemma 9 (Drift towards Synchronized Point). For any t ∈ {1, . . . , T}, τ ∈ {0, . . . , N}, and learning rate satisfies η ≤ 116LN , we have the following claim:
1
K ∑ k Et [ ‖w̃k,t,τ − w̃k,t‖2 ] ≤5Nη2(σ2 + 18NΦt) + 30N2η2‖∇f(wt)‖2 (28)
where Φt = 1K ∑ k((d − β)G2 + 1 K ∑ k′ B 2(dist(mk,t,mk′,t) + dist(mk′,t,m ∗ k′))) and Et[·] denotes the expectation over all the randomness of round t.
Proof. We follow the basic techniques from (Reddi et al., 2020) to prove this lemma.
We first assume that:
• O = mk,t gk,t,τ−1(w̃k,t,τ−1)−mk,t ∇Fk(w̃k,t,τ−1)
• P = mk,t ∇Fk(w̃k,t,τ−1)−mk,t ∇Fk(w̃k,t)
• Q = mk,t ∇Fk(w̃k,t)−∇f(wt)
We can expand w̃k,t,τ as follows: 1
K ∑ k Et [ ‖w̃k,t,τ − w̃t‖2 ] Fact 1 = 1
K ∑ k Et [ ‖w̃k,t,τ−1 − w̃t − ηmk,t gk,t,τ−1(w̃k,t,τ−1)‖2 ] = 1
K ∑ k Et [ ‖w̃k,t,τ−1 − w̃t − η(O + P +Q+∇f(wt))||2 ] ≤
1 + 12N−1 K ∑ k Et‖w̃k,t,τ−1 − w̃t‖2 + η2 K ∑ k Et‖O‖2 + 6Nη2 K ∑ k Et‖P‖2
+ 6Nη2
K ∑ k Et‖Q‖2 + 6Nη2 K ∑ k ‖∇f(wt)‖2
(29)
where the last inequality follows from Lemma 3 and Lemma 4. Explicitly, we use Lemma 4 to treat the stochastic term with Et‖O‖2, and then we use Lemma 3 with a = 2N − 1 to separate the other four terms.
Then we proceed by separately bounding the components in the above inequality.
Bounding the second term: η2
K ∑ k Et‖O‖2 = η2 K ∑ k Et‖mk,t gk,t,τ−1(w̃k,t,τ−1)−mk,t ∇Fk(w̃k,t,τ−1)‖2
= η2
K ∑ k Et‖mk,t (gk,t,τ−1(w̃k,t,τ−1)−∇Fk(w̃k,t,τ−1))‖2
≤ η 2
K ∑ k Et‖gk,t,τ−1(w̃k,t,τ−1)−∇Fk(w̃k,t,τ−1)‖2
≤ η2σ2
(30)
where the last inequality holds by Assumption 2.
Bounding the third term:
6Nη2
K ∑ k Et‖P‖2 = 6Nη2 K ∑ k Et‖mk,t ∇Fk(w̃k,t,τ−1)−mk,t ∇Fk(w̃k,t)‖2
≤6Nη 2
K ∑ k Et‖∇Fk(w̃k,t,τ−1)−∇Fk(w̃k,t)‖2
≤(6Nη2L2) 1 K ∑ k Et‖w̃k,t,τ−1 − w̃k,t‖2
(31)
Bounding the fourth term:
6Nη2
K ∑ k Et‖Q‖2
= 6Nη2
K ∑ k Et‖mk,t ∇Fk(w̃k,t)−∇f(wt)‖2
≤18Nη 2
K ∑ k ((d− β)G2 + 1 K ∑ k′ B2(dist(mk,t,mk′,t) + dist(mk′,t,m ∗ k′))),
(32)
where we use the second claim in Lemma 7 in the last inequality. Let Φt = 1K ∑ k((d−β)G2 + 1 K ∑ k′ B 2(dist(mk,t,mk′,t) +dist(mk′,t,m ∗ k′))). The bound can be simplified as follows:
6Nη2
K ∑ k Et‖Q‖2 ≤ 18Nη2Φt. (33)
Putting together: Plugging all the components into Eq.(29) , the following result immediately follows:
1
K ∑ k Et [ ‖w̃k,t,τ − w̃k,t‖2 ] ≤(1 + 1
2N − 1 + 6Nη2L2)
1
K ∑ k Et‖w̃k,t,τ−1 − w̃k,t‖2 + η2σ2 + 18Nη2Φt + 6Nη2‖∇f(wt)‖2
≤(1 + 1 N − 1 ) 1 K ∑ k Et‖w̃k,t,τ−1 − w̃k,t‖2 + η2(σ2 + 18NΦt) + 6Nη2‖∇f(wt)‖2,
(34)
where the last inequlity holds by our assumption η ≤ 116LN . Unrolling the recursion, we obtain the following results:
1
K ∑ k Et [ ‖w̃k,t,τ − w̃k,t‖2 ] ≤ N−1∑ τ=0 (1 + 1 N − 1 )τ [ η2(σ2 + 18NΦt) + 6Nη 2‖∇f(wt)‖2 ]
≤(N − 1)× ( (1 + 1
N − 1 )N − 1
)[ η2(σ2 + 18NΦt) + 6Nη 2‖∇f(wt)‖2 ]
≤5Nη2(σ2 + 18NΦt) + 30N2η2‖∇f(wt)‖2
(35)
The last inequality holds since (
(1 + 1N−1 ) N − 1
) ≤ 5 for N ≥ 1. This completes the proof.

heading: C.4 FORMAL PROOF
text: We start our proof by expanding f(wt+1) under its smoothness condition (see Lemma 5), which indicates that:
Et [f (wt+1) | wt]
≤f(wt)− 〈∇f(wt),Et[wt+1 −wt]〉+ L
2 Et||wt+1 −wt||2
Fact 3 = f (wt)− ηEt [〈 ∇f (wt) , 1
S ∑ k∈St Uk,t
〉] + L
2 Et ∥∥∥∥∥ 1S ∑ k∈St Uk,t ∥∥∥∥∥ 2
Fact 2 = f (wt)−
η
N
〈 N∇f (wt) ,Et  1 K ∑ k,τ mk,t ∇Fk(w̃k,t,τ ) 〉+ L 2 Et ∥∥∥∥∥ 1S ∑ k∈St Uk,t ∥∥∥∥∥ 2
≤f (wt)− ηN
2 ||∇f (wt) ||2 +
η
2N Et||
1
K ∑ k,τ
mk,t ∇Fk(w̃k,t,τ )−N∇f (wt) ||2︸ ︷︷ ︸ T1
+ L
2 Et ∥∥∥∥∥ 1S ∑ k∈St Uk,t ∥∥∥∥∥ 2
︸ ︷︷ ︸ T2
(36)
where the last inequality holds since −ab ≤ 12 ((b − a) 2 − a2), and Et[·] is the expectation over all the randomness in round t.
In the following, we shall separately bound T1 and T2.
Bounding T1:
T1 = η
2N Et||
1
K ∑ k,τ mk,t ∇Fk(w̃k,t,τ )−N∇f (wt) ||2
Fact 4 =
η
2N Et||
1
K ∑ k,τ mk,t ∇Fk(w̃k,t,τ )−N 1 K ∑ k m∗k ∇Fk(w̃k,t)||2
= η
2N Et||
1
K ∑ k,τ (mk,t ∇Fk(w̃k,t,τ )−m∗k ∇Fk(w̃k,t))||2
(a) ≤ η 2K ∑ k,τ Et||mk,t ∇Fk(w̃k,t,τ )−m∗k ∇Fk(w̃k,t)||2
(b) ≤ η 2K ∑ k,τ (‖∇Fk(w̃k,t,τ )−∇Fk(w̃k,t)‖2 + dist(mk,t,m∗k)B2)
≤ηL 2
2K ∑ k,τ Et||w̃k,t,τ − w̃k,t||2 + ηN 2K ∑ k dist(mk,t,m ∗ k)B 2
(37)
where inequality (a) is due to Cauchy-Schwarz inequality (i.e., Lemma 1), (b) follows from the first claim in Lemma 7, and the last inequality holds by Assumption 3.
Plugging the results of Lemma 9, we obtain that:
T1 ≤ ηL2N
2 (5Nη2(σ2 + 18NΦt) + 30N
2η2‖∇f(wt)‖2) + ηN
2K ∑ k dist(mk,t,m ∗ k)B 2
≤ 5N 2η3L2
2 (σ2 + 18NΦt) + 15N
3η3L2‖∇f(wt)‖2 + ηN
2K ∑ k dist(mk,t,m ∗ k)B 2
(38)
Bounding T2:
T2 = L
2 Et ∥∥∥∥∥ 1S ∑ k∈St Uk,t ∥∥∥∥∥ 2
Fact 2 =
Lη2
2 Et  ∥∥∥∥∥∥ 1S ∑ k∈St,τ mk,t gk,t,τ (w̃k,t,τ ) ∥∥∥∥∥∥ 2 
(a) ≤Lη2Et  ∥∥∥∥∥∥ 1S ∑ k∈St,τ mk,t ∇Fk(w̃k,t,τ ) ∥∥∥∥∥∥ 2 + NLη2σ2 S
≤Lη2N ∑ τ Et ∥∥∥∥∥ 1S ∑ k∈St mk,t ∇Fk(w̃k,t,τ ) ∥∥∥∥∥ 2 
︸ ︷︷ ︸ T3
+ NLη2σ2
S
(39)
where (a) is obtained as per Lemma 6.
Bounding T3:
T3
= 1
S2 Et 〈∑ i∈[K] I{i∈St}mi,t ∇Fi(w̃i,t,τ ), ∑ j∈[K] I{j∈St}mj,t ∇Fj(w̃j,t,τ ) 〉
= 1
S2 Et  ∑ i,j∈[K],j 6=i,τ ESt [I{i∈St∩j∈St}] 〈mi,t ∇Fi(w̃i,t,τ ),mj,t ∇Fj(w̃j,t,τ )〉
+ ∑ i ESt [I{i∈St}]‖mi,t ∇Fi(w̃i,t,τ )‖ 2
]
= 1
S2 Et  ∑ i,j∈[K],j 6=i S(S − 1) K(K − 1) 〈mi,t ∇Fi(w̃i,t,τ ),mj,t ∇Fj(w̃j,t,τ )〉+ ∑ i S K ‖mi,t ∇Fi(w̃i,t,τ )‖2  = 1
S2 Et  ∑ i,j∈[K] S(S − 1) K(K − 1) 〈mi,t ∇Fi(w̃i,t,τ ),mj,t ∇Fj(w̃j,t,τ )〉
+ ∑ i S(K − S) K(K − 1)
‖mi,t ∇Fi(w̃i,t,τ )‖2 ]
≤Et  1K2 ‖ ∑ k
mk,t ∇Fk(w̃k,t,τ )‖2︸ ︷︷ ︸ T4
+ (K − S) SK(K − 1) ∑ k
‖mk,t ∇Fk(w̃k,t,τ )‖2︸ ︷︷ ︸ T5  (40)
where the last inequality holds since S−1SK(K−1) = S−1 SK2−SK ≤ S SK2 = 1 K2 .
Bounding T4:
T4 = ∥∥∥∥∥( 1K∑ k mk,t ∇Fk(w̃k,t,τ )−∇f(wt)) +∇f(wt)) ∥∥∥∥∥ 2
Lemma 1 ≤ 2 ∥∥∥∥∥ 1K∑ k mk,t ∇Fk(w̃k,t,τ )−∇f(wt) ∥∥∥∥∥ 2 + 2 ‖∇f(wt)‖2
Fact 2 = 2 ∥∥∥∥∥ 1K∑ k (mk,t ∇Fk(w̃k,t,τ )−m∗k ∇Fk(w̃k,t)) ∥∥∥∥∥ 2 + 2 ‖∇f(wt)‖2
Lemma 1 ≤ 2
K ∑ k ‖mk,t ∇Fk(w̃k,t,τ )−m∗k∇Fk(w̃k,t)‖ 2 + 2 ‖∇f(wt)‖2
Lemma 7 ≤ 2
K ∑ k (‖∇Fk(w̃k,t,τ )−∇Fk(w̃k,t)‖2 + dist(mk,t,m∗k)B2) + 2 ‖∇f(wt)‖ 2
≤2L 2
K ∑ k ‖w̃k,t,τ − w̃k,t‖2 + 2 K ∑ k dist(mk,t,m ∗ k)B 2 + 2 ‖∇f(wt)‖2 ,
(41)
where the last inequality is obtained by the L-smoothness Assumption 3.
Bounding T5:
T5 = (K − S) SK(K − 1) ∑ k ‖mk,t ∇Fk(w̃k,t,τ )‖2
≤ 3(K − S) SK(K − 1) ∑ k (‖mk,t ∇Fk(w̃k,t,τ )−mk,t ∇Fk(w̃k,t)‖2
+ ‖mk,t ∇Fk(w̃k,t)−∇f(wt)‖2 + ‖∇f(wt)‖2) (a)
≤ 3L 2(K − S) SK(K − 1) ∑ k ‖w̃k,t,τ − w̃k,t‖2 + 3(K − S) S(K − 1)K ∑ k ( 3(d− β)G2
+ 3
K ∑ k′ B2(dist(mk,t,mk′,t) + dist(mk′,t,m ∗ k′)) ) + 3(K − S) S(K − 1) ‖∇f(wt)‖2
≤3L 2
K ∑ k ‖w̃k,t,τ − w̃k,t‖2 + 9Φt + 3‖∇f(wt)‖2,
(42)
where the last inequality holds since K−SS(K−1) ≤ 1 under the condition S ≥ 1. Inequality (a) is obtained by Assumption 3 and the second claim in Lemma 7 .
Summing T4 and T5, we have the following bounding for T3:
T3 ≤ 5L2
K ∑ k ‖w̃k,t,τ − w̃k,t‖+ 2 K ∑ k dist(mk,t,m ∗ k)B 2 + 9Φt + 5‖∇f(wt)‖2 (43)
Plugging Lemma 9 into the above inequality, we have: T3
≤5L2 ( 5Nη2(σ2 + 18NΦt) + 30N 2η2‖∇f(wt)‖2 ) + 2
K ∑ k dist(mk,t,m ∗ k)B 2 + 9Φt + 5‖∇f(wt)‖2
=25η2L2N(σ2 + 18NΦt) + (150N 2η2L2 + 5)‖∇f(wt)‖2 +
2
K ∑ k dist(mk,t,m ∗ k)B 2 + 9Φt
(44) Plugging T3 into Inequality (39), we bound T2 as follows:
T2 ≤(150N4η4L3 + 5Lη2N2)‖∇f(wt)‖2 + 25η4L3N3(σ2 + 18NΦt)
+ 2η2N2L
K
∑ k dist(mk,t,m ∗ k)B 2 + 9η2N2LΦt + NLη2σ2 S .
(45)
Plugging T2 and T1 into R.H.S of Inequality (36), we obtain that:
E [f (wt+1 | wt)] ≤ f (wt)− ηN( 1
2 − 150N3η3L3 − 5LηN − 15N2η2L2)||∇f (wt) ||2
+(25η4L3N3 + 5N2η3L2
2 )(σ2 + 18NΦt) +
4η2N2L+ ηN
2K
∑ k dist(mk,t,m ∗ k)B 2
+ 9η2N2LΦt + NLη2σ2
S .
(46)
Taking expectation over the randomness before round t towards both sides of the inequality, it yields:
E [f (wt+1)] ≤ E[f (wt)]− ηNκE[||∇f (wt) ||2] + ρt. (47)
where κ = 12 − 150N 3η3L3 − 15N2η2L2 − 5NηL and ρt = (25N3η4L3 + 5N
2η3L2
2 )(σ 2 +
18NΦt) + 4N2η2L+Nη
2K
∑ k dist(mk,t,m ∗ k)B 2 + 9N2η2LΦt + Nη2Lσ2 S .
By rearanging, and summing from t = 0, . . . , T − 1, the following result immediately shows:
1
T T−1∑ t=0 E[||∇f (wt) ||2 ≤ E [f (w0)]− E [f (wT )] TηNκ + 1 T T−1∑ t=0 ρt
≤ f (w0)− f (w ∗)
TηNκ +
1
T T−1∑ t=0 ρt
(48)
Now we bound the local gradient over the personalized sparse model:
1
TK T−1∑ t=0 K∑ k=1 E[||∇Fk (w̃k,t) ||2
= 1
TK T−1∑ t=0 K∑ k=1 E[||∇Fk (w̃k,t)−∇f(wt) +∇f(wt)||2]
Lemma 1 ≤ 3
TK T−1∑ t=0 K∑ k=1 E
[ ‖∇Fk(w̃k,t)− 1
K ∑ k′ m∗k ∇Fk′(w̃k′,t)‖2
+‖ 1 K ∑ k′ (m∗k ∇Fk′(w̃k′,t)−m∗k′ ∇Fk′(w̃k′,t))‖2 + ‖∇f(wt)‖2 ]
(a) ≤ 3 TK T−1∑ t=0 K∑ k=1 ((d− β)G2 + βB2 + 1 K ∑ k′ dist(m∗k,m ∗ k′)B 2) + 3 TK T−1∑ t=0 K∑ k=1 E‖∇f(wt)‖2
≤3(d− β)G2 + 3βB2 + 3 K2 ∑ k ∑ k′ dist(m∗k,m ∗ k′)B 2 + 3 TK T−1∑ t=0 K∑ k=1 E‖∇f(wt)‖2
(49)
In inequality (a), we use the second claim in Lemma 8 to treat the first term, and the second term is bounded using a similar technique as in the first claim in Lemma 8.
Let Υ = 3(d− β)G2 + 3βB2 + 3K2 ∑ k ∑ k′ dist(m ∗ k,m ∗ k′)B
2 and plugging inequality (48) into the above inequality, it further implies that:
1
TK T−1∑ t=0 K∑ k=1 E[||∇Fk (w̃k,t) ||2 ≤ 3(f (w0)− f (w∗)) TηNκ + 3 T T−1∑ t=0 ρt + Υ (50)
This shows the claim.

heading: D DISCUSSION ON VARIOUS BOUNDS IN EXISTING STUDY
text: Under the non-convex setting, Theorem 1 bounds the squared gradient norm of the personalized iterates over the local loss. Existing studies on PFL zero in various different kinds of convergence bound, and under different assumptions. Below are some concrete examples:
• Li et al. (2021) study the upper bound of E [ ‖w̃k,T −w∗k‖ 2 ]
for any device k (see their Corollary 1), wherewk,T is the personalized weights at iteration T andw∗k is the optimal weights under the client k’s local loss. Under the assumption of strongly convex, smoothness, bounded gradient, and gradient dissimilarity, as well as an additional assumption on the distance between personalized models and the optimal global model, it is theoretically proved that the proposed solution, named Ditto, achieves O(1/T ) convergence rate. Their derived convergence rate is on the same scale as the convergence rate of the global model produced by FedAvg. However, it is important to note that an extra assumption on the distance between personalized models and the optimal global model is required to achieve this rate.
• Hanzely et al. (2021) study the convergence of the proposed regularization problem (see their Theorem 4.5). Specifically, they propose to bound E [∑K k=1 ‖w̃k,T −w∗k(λ)‖ 2 ]
where w̃k,T is the personalized model, and w∗k(λ) is the optimal model for their defined regularization problem under the regularized hyper-parameter λ. Under the strongly convex and smoothness assumption, the authors derive the bound as E [∑K k=1 ‖w̃k,T −w∗k(λ)‖ 2 ] ≤(
1− αµn )T ∑K k=1 ‖w̃k,T −w∗k(λ)‖ 2 + 2nασ 2
µ . Further, by applying the variance reduction technique on their proposed L2GD solution, they eliminate the constant term (i.e., the second term) in the above bound. However, their convergence analysis are made towards a special case of the regularized problem. It remains unspecified about how and whether the personalized iteration produced by L2GD can achieve convergence towards the optimum for each client’s local loss (i.e., the optimum of the ultimate PFL problem (P2)). In other words, the gap between w̃k,T and w∗k = minwk F̂k(wk) remains unknown.
• Under the non-convex setting, Deng et al. (2020) study the upper bound of squared gradient norm, or formally, 1T ∑T t=1 E [ ‖∇Fk (w̃k,t)‖2 ] (see their Theorem 6) , where w̃k,t is the personalized
model for client k and∇Fk(·) is the gradient of a model over the local loss. Under the smoothness and bounded variance conditions, they derive a bound with rate O( 1√
T ) to convergence. Similar
to our results, a non-vanished residual presents in their bound. This non-vanished residual is related to: i) the gradient dissimilarity between clients’ local loss function; and (ii) the gradient discrepancy, The form of this non-vanished term is highly similar to our derived residual term (i.e., Υ in our Theorem 5), since both the gradient dissimilarity factor (i.e., G in Assumption 1) and gradient bound (i.e.,B in Assumption 4) would present in the non-vanished residual. Additionally, when the personalized factor is completely eliminated i.e., αk → 0 in their formulation or β → d in our formulation, the non-vanished residual (i.e., Γ in their formulation and Υ in ours) would both become significant to dominate the bound.

Review summaries:
This paper proposes an interesting approach to FL with heterogeneous client data. The main idea is  to learn a common weight vector for all clients but allow them to individually mask this global weight vector.   **** I appreciate the effort that authors put into adressing my concerns which have been partially covered. Thus, i have slighltly raised the grade of my recommendation. However, i would ask the area chair to verify if the authors have addressed all my concerns satisfactorily. In particular, im still not sure about the relation between (P4) and regularized PFL (P3) which is quite generic and also allows for reguarlizers that are indicator functions of sets.   

This paper studies personalized sparse training for federated learning. The proposed method FedSpa is a novel personalized federated learning scheme that employs personalized sparse masks to customize sparse local models on the edge.  The authors provide the theoretical result with regard to the error bound and empirical results with several personalized federated learning methods.

This work studies personalized federated learning through sparse local masks. For each client, the local model is subset of the global model. The paper proposes and analysis a new method, named FedSpa, which trains sparse local models for all clients reducing communication size, the amount of computation and memory cost.