FedDWA: Personalized Federated Learning with Dynamic Weight Adjustment

1 Introduction

Federated Learning (FL), as an emerging distributed machine learning paradigm, allows decentralized clients to collaboratively train a global machine learning model without exposing their private data [McMahan et al., 2017]. However, one of the most challenging problems confronted by FL is the performance degradation caused by the heterogeneity of data distribution on decentralized clients [Li et al., 2020b]. Specifically, data distribution on clients is non-independent and identically distributed (non-IID) such that a single global model cannot meet personalized needs of all clients. It has been reported in [Li et al., 2020b; Yu et al., 2020] that data heterogeneity can result in slow convergence and poor model accuracy. For example, the global next-word prediction model trained by FedAvg [McMahan et al., 2017] is not always effective for all clients because of their personalized habits. Such a single global model may significantly deviate from personalized optimal models [Yu et al., 2020].

∗An extended version of this paper (with the Appendix included) can be found in http://arxiv.org/abs/2305.06124.

†Corresponding author.

To address the above problem, Personalized Federated Learning (PFL) has been proposed and studied in [Smith et al., 2017; T Dinh et al., 2020; Fallah et al., 2020; Li et al., 2021b; Collins et al., 2021]. PFL aims to handle non-IID data distribution by training personalized models for each client, so as to improve model accuracy. In essence, PFL can either incorporate personalized components into the global model or train multiple models to obtain personalized models. For example, the works [T Dinh et al., 2020; Li et al., 2021b] added regularization terms to the global model in order to train personalized models. A distance metric to shrink the search space of personalized models around the global model is applied. However, such an approach fails to optimize personalized models because distance metrics usually cannot exactly capture the heterogeneity of data distribution among clients. Later on, more radical approaches (e.g., [Zhang et al., 2021b; Li et al., 2022]) were proposed, which train multiple models to meet personalized requirements by distributing other clients' models to each individual client. Based on its local dataset, each client can decide how to aggregate models from other clients to obtain a personalized model. Despite that the performance of PFL is improved, this approach will make communication traffic explode and users' privacy may be compromised.

In this paper, we propose a novel PFL method called FedDWA (Federated Learning with Dynamic Weight Adjustment), which can improve PFL performance by encouraging collaborations among clients with similar data distributions. In existing works [Zhang et al., 2021b; Li et al., 2022], each client needs to collect models from all other clients and evaluate similarities between clients with extra local validation sets. Different from these works, our framework characterizes client similarity in an analytical way instead of empirical searching via the validation dataset. In addition, there is no need to share local models among clients in FedDWA avoiding excessive communication traffic and potential privacy leakage [Hu et al., 2021]. In FedDWA, individual clients can obtain personalized models computed by the PS based on collected model parameters and guidance models from clients. Guidance models are obtained using the onestep ahead adaptation method by individual clients. Based on guidance models, the PS can tune aggregation weights to minimize the distance between each model with its guidance model. Thus, FedDWA can improve PFL performance without incurring heavy overhead by avoiding exchanging information between clients.

In summary, our main contributions in this paper can be summarized as follows:

• We propose a new personalized federated learning framework called FedDWA. FedDWA can effectively exploit clients owning data with a similar distribution to improve personalized model accuracy.
• We theoretically analyze the properties of the FedDWA algorithm, and show how the weights are dynamically adjusted to achieve personalization.
• By conducting experiments using five real datasets, we demonstrate that FedDWA outperforms other methods under three heterogeneous FL settings.

2 RelatedWork

In this section, we discuss related works from two perspectives: data-based PFL and model-based PFL. Data-based PFL focuses on reducing data heterogeneity among clients while model-based PFL focuses on designing a personalized model for each client. Typically, data-based PFL shares a global dataset that is balanced across all clients [Zhao et al., 2018] (or private statistical information [Shin et al., 2020; Yoon et al., 2021] among clients) to realize personalized learning. However, sharing a global dataset may potentially violate privacy policies since it is at the risk of privacy leakage. To address the above problem, model-based PFL was proposed, which can be divided into two types: single-model PFL and multi-model PFL.

Most single-model PFL methods are extensions of conventional FL algorithms (e.g., FedAvg [McMahan et al., 2017]). For example, FedProx [Sahu et al., 2020] employed a proximal term to formulate clients' optimization objectives so as to mitigate the adverse influence of systematic and statistical heterogeneity on FL. FedAvg FT and FedProx FT [Wang et al., 2019] obtained personalized models by fine tuning the global model generated by FedAvg and FedProx, respectively. FedAvgM proposed by [Hsu et al., 2019] adopted a momentum method to update the global model, so as to alleviate the adverse influence of non-IID data distribution on FL. An alternative single-model approach for PFL is based on meta-learning. Recent works [Khodak et al., 2019; Yue et al., 2021; Acar et al., 2021] extended Model Agnostic Meta-learning (MAML) for FL under non-IID data distribution. However, the personalized learning ability of singlemodel methods is limited because it is hard to fit all heterogeneous data distributions with a single model very well.

Multi-model methods outperform single-model methods by training multiple models to better adapt to the personalized requirements of clients. Cluster FL [Sattler et al., 2021; Ghosh et al., 2020; Mansour et al., 2020] assumed that clients can be partitioned into multiple clusters, and clients are grouped based on loss values or gradients. A customized model can be trained for each cluster. However, clusterbased client grouping may not be able to effectively improve PFL performance. FedEM [Marfoq et al., 2021] refined the cluster-based client group method by proposing a soft client clustering algorithm. However, it requires each client to download multiple models, which can considerably increase the communication overhead.

Other than clustering clients, more advanced multi-model PFL methods were developed including additive model mixture between local and global models (such as L2GD [Hanzely and Richt´arik, 2020] and APFL [Deng et al., 2020]), multi-task learning methods with model similarity penalization (such as MOCHA [Smith et al., 2017], pFedMe [T Dinh et al., 2020] and Ditto [Li et al., 2021b]). More PFL methods were developed by leveraging Gaussian processes [Achituve et al., 2021] and knowledge transfer [Zhang et al., 2021a]. However, these methods inevitably need public shared data or inducing points set. It is also possible to achieve PFL by decomposing FL models into a global part and multiple personalized parts. Inspired by representation learning, the works [Arivazhagan et al., 2019; Collins et al., 2021; Tan et al., 2022; Chen and Chao, 2022; Oh et al., 2021; Mills et al., 2022] decomposed the FL model into a shared feature extractor part and a personalized part to realize PFL. Nevertheless, how to decompose FL models is only heuristically designed and discussed by existing works.

PFL can be achieved by customizing weights of model aggregation for each client as well, e.g., FedAMP [Huang et al., 2021], FedFomo [Zhang et al., 2021b] and L2C [Li et al., 2022]. These customizing weights represent potential similarities between clients. [Chen et al., 2022] proposed that graph neural networks can also be used to learn similarities among clients to realize personalization. FedAMP proposed an attentive message passing mechanism to compute personalized models, which is not flexible enough, because all clients need to participate in training in every round. FedFomo and L2C computed personalized aggregation weights via minimizing the validation loss on each client based on the model information collected from other clients, resulting in heavy communication traffic and concerns on privacy leakage. Although the effectiveness of customizing aggregation weights has been validated in existing works, their design is empirical based, not communication-efficient for large-scale real-world FL systems. Our work aligns with the line of work to customize aggregation weights in an analytical way without incurring heavy communication overhead.

3 ProblemFormulation

In federated learning, each client i owns a local private dataset denoted by Di drawn from a distinct distribution Pi. The objective of FL is to train a single global model w for all clients by solving the following problem:

where the function fi : Rd → R represents the expected loss over the data distribution of client i, i.e.,

In the above equation, ξi is a random sample generated according to the local distribution Pi and ˜ fi (w; ξi) represents the loss function corresponding to sample ξi and w. Since clients' data possibly come from different environments, they likely have non-IID data distributions, i.e., for i ̸= j, Pi ̸= Pj .

In conventional FL, the target is to train a global model through conducting multiple global iterations. In the (t - 1)- th global iteration, the PS distributes the latest global model parameters wt-1 to all participating clients. The clients train locally and send the trained model ˆwt i to PS for aggregation. The whole process can be shown as

Here, we suppose that there are N participating clients and pi is a pre-defined non-negative weight typically proportional to |Di| with ∑Ni=1 pi = 1. wt represents the global model for the t-th round and η is the learning rate.

From another perspective, the single global model trained by conventional FL is to minimize the L2 distance between a global model and all local models, which can be expressed as

However, if we consider the optimization of personalized models for individual clients, the optimization problem should be revised as

Here w⋆ i represents the optimal target model for client i. If the data distribution is IID, it implies that w⋆ i ≈ w⋆ i′ for any two clients, which means that the optimal model applicable for each individual client can be derived by Eq. (5). However, if the data distribution is non-IID, it has been investigated in [Li et al., 2019; Li et al., 2020a] that a global model cannot satisfy all clients very well, resulting in poor model accuracy.

4 Methodology

In this section, we elaborate the design of FedDWA and prove its effectiveness through analysis.

Algorithm 1 FedDWA algorithm

Input: Communication Round T, learning rate η, local epochs E, number of clients N, init model parameter w0 . Output: Personalized model parameters w1, w2, . . . , wN .

Server

• 1: for t = 1, . . . , T do
• 2: Server randomly selects a subset of clients St and sends wt1 , wt2 , ..., wtm to them.
• 3: for each client i ∈ St in parallel do
• 4: ˆwt i , ˆw⋆ i ←Client(i,wt i )
• 5: end for
• 6: Compute pi,j according to Eq. (15) for each client i.
• 7: for each client i ∈ St do
• 8: Select top-K clients {Ki}.
• 9: Aggregate newmodel according to Eq. (8).
• 10: end for
• 11: end for

Client

• 1: B ←(split D into batches of size B)
• 2: for each local epoch i from 1 to E do
• 3: for batch b ∈ B do
• 4: w =w- η∇˜ f (w, b)
• 5: end for
• 6: end for
• 7: Train one more local iteration (one epoch).
• 8: ˆw =w- η∇˜ f (w, D)
• 9: return w and ˆw

4.1 Optimization Objective

In the previous section, it has been pointed out that the aggregation rule defined by Eq. (4) cannot meet personalized requirement with non-IID data distribution. Rather than training a single global model, we propose to train a model for each individual client by customizing aggregation weights so as to deduce each individual model.

Specifically, we define pi,j as the weight to aggregate the model for client i using the local model from client j . Here, ∑Nj=1 pi,j = 1. Then, the PS can generate the personalized model for client i in the (t-1)-th global iteration as follows:

Intuitively speaking, Eq. (8) is very flexible. When aggregating the model for client i, we can set a larger value for pti,j if the data distribution of client j is closer to that of client i such that the PS can explore optimal personalized model for client i1 . To specify how to set the value of pi,j , we formulate the optimization problem as below:

1 pi,j and ηi can be time-dependent, but when context allows, we write pti,j as pi,j and ηt i as ηi for simplicity.

Here ˆw⋆ i is the guidance model for client i in the (t - 1)-th global iteration, and it tells client i which clients to cooperate with. The main challenge for solving Eq. (9) lies in that ˆw⋆ i is unknown in advance. We can solve Eq. (9) with two steps. Firstly, we need to find a high quality guidance model ˆw⋆ i and fix ˆw⋆ i to derive how to optimally set pi,j. Secondly, after deriving the optimal weight pi,j, client i can get its own personalized model at the t-th round, and then its guidance model ˆw⋆ i can be further updated. In the following, we elaborate how to solve Eq. (9).

Tuning Aggregation Weights

Since ∑Nj=1 pi,j = 1, we can rewrite Eq. (9) as:

Let the vector pi = [pi,1, pi,2, ..., pi,N]T denote aggregation weights for obtaining client i's personalized model. Let Wi denote the cross distance between the guidance model ˆw⋆ i and local models contributed by clients. The (j, k)-th entry of Wi can be written as:

Then the optimization problem for client i defined in Eq. (9) can be expressed as follows:

Note that the PS can optimize personalized models for all clients via Eq. (12). Suppose that Wi is invertible, then the solution is given by:

It is very difficult to directly calculate Eq. (13) due to the following three challenges. First, it involves the inner product of model parameters among clients which can be seen from Eq. (11). For advanced neural networks, there may exist millions of parameters making the computation cost unaffordable. Second, training models in federated learning is an iterative process with multiple rounds of communications. It implies that computing the inversion of Wi is cumbersome, especially when the dimension of Wi is very large. Third, since Wi is just a symmetric matrix, W-1 i may not exist at all. As a consequence, the solution of Eq. (12) is not unique. Thus, trying to solve Eq. (12) directly cannot guarantee that a high-quality solution will be yielded. A toy example is discussed in Appendix A.1.

To make the problem tractable, we simplify the objective in Eq. (12) by only reserving the diagonal elements of Wi. The effectiveness of such simplification has been verified in previous works [Chen et al., 2015; Zhao and Sayed, 2012]. The simplified problem is presented as follows:

subject to 1TN pi = 1, pi,j ≥ 0.

It is easy to find that there is a unique solution to the simplified problem. It is worth noting that Eq. (14) is very alike to the aggregation rule in conventional FL defined in Eq. (5). The difference of our method can be explained from two perspectives. First, aggregation weights are tunable parameters in our method, which however are fixed in traditional FL. Second, our method can search the optimal aggregation weights for individual clients to derive personalized models. The solution of Eq. (14) is:

where the detailed derivation can be found in Appendix A.2.

One-step Ahead Adaptation

Executing the combination rule in Eq. (15) by an individual client requires the knowledge of the guidance model ˆw⋆ i , which is generally not available beforehand or not. Intuitively speaking, whether we can realize personalization depends on similar clients identified by our algorithm. In other words, the weighted combination of models for client i should align with client i's personal data distribution. It implies that the guidance model ˆw⋆ i should capture the local data distribution of client i. We have tried several options for ˆw⋆ i including using the last iteration model ˆw⋆ i = ˆwt-1 i , current model ˆw⋆ i = wt i , and one-step ahead adaptation. More discussion can be found in Appendix C.6. In this work, we employ an instantaneous adaptation argument, a.k.a. local one-step ahead adaptation, to accommodate this issue as follows:

The validity of this adaptation can be found in previous work [Jin et al., 2020; Chen et al., 2015]. It is important to note that this is fundamentally different from traditional FedAvg training and then local fine-tuning, because our method will only select other clients that are beneficial to client i for aggregation while the fine-tuning approach treats all clients equally during the aggregation phase. In fact, after a step of adaptation, Eq. (16) in advance, ˆw⋆ i can characterize its local data distribution well and therefore it can screen out other clients with similar data distribution and give them a higher weight for cooperation. We can also use two steps or even more steps, the specific experimental results will be shown in the following sections. By substituting Eq. (16) into Eq. (15), we can finally derive the aggregation weights by deriving personalized models as

Similar to the previous works [Zhang et al., 2021b; Li et al., 2022], the top-K technique can be used. We rank pi,j in descending order, and only the top K aggregation weights are selected and they will be normalized such that ∑Nj=1 pi,j = 1. By wrapping up our analysis, we present the detailed FedDWA algorithm in Algorithm. 1.

4.2 Analysis of FedDWA

Communication Overhead

The FedDWA algorithm will incur 2Σ traffic in the uplink communication and the traffic in the downlink communication is the same as that of the original FedAvg, where Σ denotes the model size. It is worth noting that FedDWA can incur significantly less communication traffic than other similar baselines, and more results can be found in Appendix C.4.

Computational Cost

Suppose that there are N clients participating in training and the number of model parameters is d, the computation complexity of FedDWA is O(N 2d) in the server. We also test the total FLOPs required by FedDWA and the experimental results show that the computational amount required to calculate the similarity (Eq.(15)) is negligible compared with model training. More results can be found in Appendix C.5.

Personalized Learning

In this part, we illustrate how FedDWA can make clients with similar data distribution collaborate to train personalized models. Considering the inverse of the numerator of Eq. (17), it can be expanded as

The first term ∥∥ ˆwt i - ˆwtj ∥∥2 refers to the distance between current models of client i and client j. This term will lower the aggregation weight pi,k if the distance ∥∥ ˆwt i - ˆwtj ∥∥2 is large, and thereby prohibit their collaborations. Using the first-order Taylor series to expand fi(w) at ˆwt i , we have:

For the second term 2ηt-1 i ( ˆwtj - ˆwt i )T ∇fi( ˆwt i ), we can find that it is proportional to fi( ˆwtj )-fi( ˆwt i ) which also decreases the aggregation weight pi,k if fi( ˆwtj ) is far away from fi( ˆwt i ). The last term (ηt-1 i )2 ∥∇fi( ˆwt i )∥2 can be perceived as a constant when optimizing aggregation weights. In summary, Eq. (17) provides the aggregation weights for deriving a personalized model for client i based on the similarity distance between client models ˆwj's and the guidance model ˆ wi ⋆ .

5 Experiments

5.1 Experiment Setups

Datasets and Models

We evaluate our algorithm on five benchmark datasets, namely, EMNIST [Cohen et al., 2017], CIFAR10, CIFAR100 [Krizhevsky et al., 2009], CINIC10 [Darlow et al., 2018] and Tiny-ImageNet (TINY) [Chrabaszcz et al., 2017]. For EMNIST, we use the same model as that used in [Sattler et al., 2021]. For CIFAR10, CIFAR100 and CINIC10, we use the CNN model which is the same as that in [Mills et al., 2022]. To evaluate the effectiveness of FedDWA on a highdimensional model, we use ResNet-8 for the Tiny-ImageNet, and the model architecture is the same as that in [He et al., 2020]. More details can be found in Appendix B.1.

Data Partitioning

We simulate the heterogeneous settings with three widely used scenarios, including a pathological setting and two practical settings.

• Pathological Heterogeneous Setting. Each client is randomly assigned with a small number of classes with the same amount of data on each class [McMahan et al., 2017; Shamsian et al., 2021]. We sample 4, 2, 2, 6 and 10 classes for EMNIST, CIFAR10, CINIC10, CIFAR100, Tiny-ImageNet from a total of 62, 10, 10, 100, 200 classes for each client, respectively. There is no group-wise similarity between clients in this setting.
• Practical Heterogeneous Setting 1. All clients have the same data size but different distributions. For each client, s% of data (80% by default) are selected from a

set of dominant classes, and the remaining (100 - s)% are uniformly sampled from all classes [Karimireddy et al., 2020; Huang et al., 2021]. All clients are divided into multiple groups. Clients in each group share the same dominant classes implying that there is an underlying clustering structure between clients.

• Practical Heterogeneous Setting 2. Each client contains most of the classes but the data in each class is not uniformly distributed [Hsu et al., 2019; Li et al., 2021a; Chen and Chao, 2022]. We create the federated version by randomly partitioning datasets among N clients using a symmetric Dirichlet distribution Dir(α) (α = 0.07 by default). For example, for each class c, we sample a vector pc from Dir(α) and allocate to client m a fraction pc,m of all training instances of class c.

Baselines

We compare the performance of our algorithm with that of FedAvg [McMahan et al., 2017], FedAvgM [Hsu et al., 2019], FedProx [Li et al., 2020a] and a few latest personalization approaches including a personalized model trained only on each client's local dataset (Local Training), FedAvg with local tuning (FedAvg FT) [Wang et al., 2019], pFedMe [T Dinh et al., 2020], per-FedAvg [Fallah et al., 2020], ClusterFL [Sattler et al., 2021], FedAMP [Huang et al., 2021], FedFomo [Zhang et al., 2021b], SFL[Chen et al., 2022], L2C [Li et al., 2022] and FedRoD [Chen and Chao, 2022]. The settings of hyper-parameters for each method can be found in Appendix B.3.

Evaluation Metrics

We use the same evaluation metrics as that widely used by previous works which report the test accuracy of the best single global model for the single-model PFL methods and the average test accuracy of the best personalized models for other PFL methods.

Training Settings

Similar to [Zhang et al., 2021b], we evaluate the performance in two settings, i.e., (1) N = 20 clients, 100% participation and (2) N = 100 clients, 20% participation for all datasets. The number of local training epochs is set to E = 1 and the number of global communication rounds is set to 100. We employ the mini-batch SGD as a local optimizer for all approaches. The batch size for each client is set as 20 and the learning rate η is set as 0.01. We test all methods over three runs and average the results.

5.2 Performance Evaluation and Analysis

Pathological Heterogeneous Setting

Table 1 shows the average test accuracy for all methods under the pathological heterogeneous setting. Over all datasets and client setups, our FedDWA algorithm outperforms other baseline methods. However, the performance of methods (such as FedAvg and FedProx) that only train a single global model degrade significantly on CIFAR10, CIFAR100, CINIC10 and Tiny-ImageNet, since a single global model cannot well accommodate statistical heterogeneity of clients. Other personalized methods achieve comparable accuracy and the local training method achieves rather high performance due to the small number of classes on each client.

Practical Heterogeneous Setting

Table 1 shows the average test accuracy for all methods under the practical heterogeneous setting 1 in which each client has a primary data class with a small number of samples from all other classes. In this setting, we find that our method is significantly superior to all other baseline methods. For example, when using Tiny-ImageNet, our method outperforms FedFomo and L2C by up to 9.76%, 7.88% in test accuracy respectively, which means that the aggregation weights obtained by our method are better than those with the aggregation weights obtained by empirical searching through the validation set. To test the applicability of our algorithm, we also evaluate the performance of our approach for a large-scale FL scenario under the practical heterogeneous setting 2. We set N = 100 clients with 20% participation in each round. The final results are shown in Table 2. It is worth noting that our approach still achieves competitive performance though there is no clear similarity between clients in this scenario.

Personalized Model Weighting

We use the practical heterogeneous setting 1 in which clients are divided into multiple groups to explain why FedDWA outperforms existing works by showing how FedDWA helps clients quickly identify other similar clients for conducting personalized model aggregation. To ease our visualization, we group clients first before clients are indexed such that clients in the same group will have consecutive indices. For example, clients of the same data distribution are grouped in the first group who are indexed by 0-4. In Figure 1, we show the K most similar clients selected by different methods according to personalized weights in each training round with 20 clients divided into 4 groups. Since there are five clients in each group, we set K = 5. The result manifests that FedDWA can identify similar clients faster than FedFomo and L2C. In addition, under this setting, L2C cannot accurately identify the similarity between users, indicating that weights found with validation datasets are not very effective. More experiment results such as the heat maps showing the aggregation weights can be found in Appendix C.1.

Selection of K

Figure 4 in Appendix C shows that when there is a potential cluster structure among clients participating in FL, the weight calculated by FedDWA at the first round can reflect the similarity between clients well. At this time, server can easily determine the value of K according to the weight matrix. However, as reported in [Ghosh et al., 2020; Zhang et al., 2021b; Li et al., 2022; Marfoq et al., 2022], it is not trivial to set K if there is no obvious data similarity between clients. Fortunately, the performance of FedDWA is not very sensitive with respect to K. We use the practical heterogeneous setting 2 to simulate different degrees of data heterogeneity with 100 clients, and present the test performance in Figure 2. The results indicate that FedDWA is not very sensitive to the value of K when data distributions are more heterogeneous (with a smaller α). When α is large, the discrepancy between users' data distributions is smaller, which is not a typical PFL scenario. At this time it becomes more important to increase the value of K so as to cooperate with more clients.

Effect of Guidance Model

We explore the influence of guidance models obtained by Eq. (16) on the final personalized model accuracy by tuning the number of local epochs. According to Eq. (16), ˆw⋆ i will be different if we set a larger number of local epochs. We evaluate this effect with 20 clients on four datasets, respectively. As shown in Figure 3, the number of local iterations to compute guidance models has little influence on the final model accuracy of FedDWA indicating that one-step adaptation is sufficient for FedDWA. Moreover, we compare the performance of guidance model ˆw⋆ i and personalized model wt i in Appendix C.3. The results further show that the performance of wt i is almost the same as that of ˆw⋆ i under the practical setting 2 and better than that of ˆw⋆ i under the practical setting 1.

6 Conclusions

In this paper, we propose a novel FedDWA algorithm which can identify similarities between clients with much less communication overhead than other relevant works since no information will be exchanged between clients. Meanwhile, personalized models are generated based on uploaded model information and its effectiveness is guaranteed with theoretical analysis. Comprehensive experiments on five datasets demonstrate the superb performance of FedDWA which can achieve the highest model accuracy compared to the state-ofthe-art baselines under three heterogeneous FL settings.

Acknowledgements

This work was supported by the National Natural Science Foundation of China under Grants U1911201, U2001209, 62072486, and the Natural Science Foundation of Guangdong Province under Grant 2021A1515011369.