A Personalized Federated Learning Algorithm Based on Dynamic Weight Allocation

Abstract

Federated learning is a privacy-preserving distributed machine learning paradigm. However, due to client data heterogeneity, the global model trained by a traditional federated averaging algorithm often exhibits poor generalization ability. To mitigate the impact of data heterogeneity, some existing research has proposed clustered federated learning, where clients with similar data distributions are grouped together to reduce interference from dissimilar clients. However, since the data distribution of clients is unknown, determining the optimal number of clusters is difficult, leading to reduced model convergence efficiency. To address this issue, this paper proposes a personalized federated learning algorithm based on dynamic weight allocation. First, each client is allowed to obtain a global model tailored to fit its local data distribution. During the client model aggregation process, the server first computes the similarity of model updates between clients and dynamically allocates aggregation weights to client models based on these similarities. Secondly, clients use the received exclusive global model to train their local models via the personalized federated learning algorithm. Extensive experimental results demonstrate that, compared to other personalized federated learning algorithms, the proposed method effectively improves model accuracy and convergence speed.

1. Introduction

With the rapid development of artificial intelligence, it has become common for neural network models to use large-scale data for performance optimization. However, the data in the device usually contain the sensitive information of many users. Centrally collecting data to a central server for model training will cause serious data privacy issues. To this end, Google proposed federated learning (FL) in 2017 [1,2]. As a new distributed machine learning paradigm, FL collaborates with multiple clients to train a global model while ensuring that the original data on the device are not exposed to third parties [3].

In real-world scenarios, data between different clients are often non-independent and identically distributed (Non-IID) [4], which poses a significant challenge to federated learning. Global models trained by traditional federated averaging (FedAvg) [1] often struggle to generalize effectively when the data distribution varies greatly between clients. Such differences in data distribution cause global models to perform poorly in capturing the unique features of each client’s local data, resulting in less than optimal model performance [5].

To address the problem of the statistical heterogeneity of data among clients, the existing research proposes personalized federated learning (PFL) [6] to mitigate the Non-IID problem by training a personalized model for each client that fits its own data distribution. There are various schemes for PFL, such as multi-task learning, knowledge distillation, and clustered federated learning (CFL). Among them, clustered federated learning [7] groups clients with similar data distributions and trains global models independently within each cluster to adequately fit the data distributions of similar clients within each cluster. However, the performance of some CFL algorithms [8,9] depends on the number of clusters they choose, and this value is usually obtained through empirical settings or hyperparameter optimization, leading to their limited usefulness in dealing with complex data distributions in real scenarios. In addition, most CFL algorithms mainly focus on intra-cluster client collaboration while neglecting inter-cluster collaboration, failing to fully integrate the data distribution characteristics of the clusters, and thus limiting the generalization ability of the model. Finally, existing CFL algorithms still use traditional weighted average aggregation; however, this method makes it difficult to fully capture the data distribution differences among clients, and it also makes it difficult to accurately assess the contribution of each client to the aggregation model, resulting in the performance of the global model favoring some clients. To address the above problems, we consider assigning aggregation weights based on the similarity between client models and aggregating an exclusive global model for each client that adapts to its own data distribution. This dynamic weight assignment method avoids the limitation of a pre-defined number of clusters in traditional clustering methods and also solves the difficulty of inter-group collaboration. By flexibly assigning aggregation weights to each client, it can more accurately enhance collaboration between similar clients and attenuate interference between dissimilar clients.

To this end, this paper proposes a personalized federated learning algorithm (FedDWA) based on dynamic weight assignment. The algorithm adaptively assigns aggregation weights to clients, fully considering the potential connection between client data distribution and model parameters, so that each client obtains an exclusive global model adapted to its local data distribution. Compared with the global model obtained by aggregation within each cluster in CFL, the aggregation of each client-exclusive global model takes into account the contributions among all clients, while avoiding introducing the problem of setting the number of clusters. Specifically, during the aggregation of client-exclusive global models, the server calculates the cosine similarity of model updates among clients to assign aggregation weights to the clients. The higher the inter-client similarity, the larger the corresponding aggregation weight. Subsequently, the server aggregates exclusive global models for clients based on the calculated aggregation weights. This aggregation method not only overcomes the limitation of a single global model in handling complex data but also effectively copes with the difference in data distribution among clients, thus achieving the matching of client-specific models and client-local data. In addition, considering that the exclusive global model is still affected by slight data heterogeneity, the client uses the exclusive global model as a guide model to train an additional personalized model locally, which further attenuates the effect of data heterogeneity. In summary, the main contributions of this paper are as follows:

• We propose a model aggregation algorithm based on dynamic weight assignment, which assigns aggregation weights based on the cosine similarity of model updates between clients, ensuring collaboration between similar clients and attenuating interference between dissimilar clients. On this basis, additional personalized models are trained for each client, which attenuates the influence of data heterogeneity among clients.
• We propose a personalized federated learning algorithm. In the local update phase, additional personalization models are trained for each client to further attenuate the impact of data heterogeneity.
• A performance evaluation on five image datasets, CIFAR-10, CIFAR-100, Tiny-Imag-eNet, MNIST, and Fashion-MNIST, shows that the proposed algorithm has higher accuracy compared with seven comparison baseline algorithms.

The paper structure is organized as follows: Section 1 provides the background and a description of the problem, Section 2 discusses the related work, Section 3 describes the problem formulation in detail, Section 4 presents our solution to the problem, Section 5 details the experimental setup and results, and Section 6 concludes the whole paper.

2. Related Work

2.1. Federated Learning for Data Heterogeneous Scenarios

In order to reduce the impact of data heterogeneity, the existing research mainly focuses on optimizing the two aspects of enhancing the generalization ability of the global model and developing personalized federated learning algorithms. For example, in order to enhance the generalization ability of the global model, FedProx, proposed by Li et al. [10], uses an optimization objective with a regularization term to reduce the negative impact of system heterogeneity and data heterogeneity on the performance of the global model. The scaffold proposed by Karimireddy et al. [11] corrects the drift phenomenon when the client model is updated by controlling variables. FedAvgM, proposed by Hsu et al. [12], uses a momentum-based optimization method to update the global model to reduce the adverse impact of Non-IID data on the model training process. FedRoD, proposed by Chen et al. [13], eliminates the impact of label distribution offset by introducing a balanced Softmax loss function and improves the generalization ability of the global model. For personalized FL research, FedAvg-FT, proposed by Cheng et al. [14], regards the local model updated by the client as a personalized model and evaluates its performance. Ditto, proposed by Li et al. [15], introduces the idea of federated multi-task learning. While updating the global model, it trains an additional personalized model for each client to reduce the impact of data heterogeneity. Flow, proposed by Panchal et al. [16], improves prediction accuracy by creating a dynamic personalized model that dynamically chooses between local and global parameters based on each client’s data distribution and specific instance.

2.2. Clustered Federated Learning

Clustered federated learning not only promotes cooperation between similar clients but also ensures that the aggregated model can better adapt to the data distribution of the clients [17]. Specifically, in CFL, the server first groups the clients according to the similarity of their data distribution, and then similar clients collaborate to train the global model within the cluster. CFL can be divided into one-time clustering and iterative clustering.

One-time clustering methods cluster clients only once during the entire federated learning training process. These methods require setting a specific number of clusters before application. For example, Ghosh et al. [18] proposed a clustered federated learning algorithm based on K-Means. The algorithm first obtains the model parameters of the client through preliminary training, then uses the K-Means algorithm to cluster the clients once, and finally performs federated average aggregation within each cluster. Briggs et al. [19] proposed a federated learning algorithm based on hierarchical clustering, which uses a hierarchical clustering algorithm to cluster clients. Liu et al. [20] divided clients into preset K clusters based on the similarity of sparse vectors between clients. Yang et al. [21] proposed the G-FML algorithm, which combines personalized federated learning with clustering learning by first clustering the customers through the K-means method and then improving the accuracy of the model through PFL on the basis of clustering. The limitation of the above algorithm is that it only ensures the collaboration between clients within a cluster but does not consider the knowledge sharing of clients between clusters. It is difficult to use the global characteristics of data between clients to improve the performance of the overall model.

To address the above problems, Ghosh et al. [9] proposed an improved K-Means clustered federated learning algorithm, IFCA. In each iteration, the server dynamically adjusts the client’s cluster according to the client’s loss function value to adapt to the client’s real-time data distribution. However, this method requires the distribution of the global model corresponding to K clusters to each client to determine the client’s cluster, and the resulting communication overhead is K times that of the FedAvg algorithm. In order to further improve the scalability of CFL, Sattler et al. [17] proposed a more advanced adaptive iterative clustering algorithm, FMTL, which does not require the number of clusters to be determined in advance. In each iteration, the algorithm performs binary clustering based on the similarity of the clients and automatically adjusts the number of clusters during the iteration. Although the algorithm improves the performance of the model to a certain extent, it still introduces a lot of computational overhead. At the same time, in real scenarios, due to the server’s unknown distribution of client data, it is difficult for the clustered federated learning algorithm to infer the optimal number of clusters.

Different from the above methods, this paper proposes a personalized federated learning algorithm FedDWA, which can adaptively assign aggregation weights to each client so that each client can obtain an exclusive global model suitable for its local data distribution, thereby effectively integrating global knowledge and avoiding limiting knowledge to specific clusters. In addition, this method trains an additional personalized model for each client so that it can better fit the client’s data distribution.

3. Problem Formulation

3.1. Federated Learning Objective

In the standard federated learning framework, a central server coordinates $N$ clients to jointly train a global model $w$. It is assumed that each client $i\in N$ has a local dataset $D_i$. Federated learning aims to solve the following problems:

$$\underset{w}{\min }G(F_1(w),\dots ,F_i(w\left)\right)$$

(1)

where $F_i(w)={\displaystyle \frac{1}{\left|D_i\right|}}{\displaystyle \sum _{x\in D_i}f(w,x)}$, $F_i(w)$ represents the local objective of client $i$, and $f(w,x)$ represents the loss function of client $i$. In FedAvg, $G(\cdot )$ represents the aggregation function, which is used to weigh and aggregate the local objectives of each client. In FedAvg, $G(\cdot )$ represents the aggregation function used to weigh and aggregate the local objectives of each client, where $G={\displaystyle {\sum }_{i=1}^N{p}_i{F}_i(w)}$, $p_i$ is a predefined weight, $p_i>0$, and the sum of all $p_i$ is equal to 1.

Due to the diversity of client device environments, usage scenarios, and user behaviors, data distribution among clients is statistically heterogeneous. The weighted average aggregation used by the traditional federated averaging algorithm allocates aggregation weights based on the number of samples owned by the client but fails to consider the contributions of clients, resulting in the global model obtained by aggregation favoring some clients, meaning that a single global model cannot fit the data distribution of all clients.

3.2. FedDWA Objective

In real-world scenarios, compared with traditional FL to train a single global model, training additional personalized models for each client can better meet its specific data distribution and requirements. Therefore, this paper proposes a personalized federated learning algorithm for Non-IID scenarios. The algorithm determines the optimal weighted aggregation weights for each client by calculating the similarity between client models for customized aggregation. In this work, two main tasks are considered: (a) optimizing the global objective $G_g(\cdot )$ for updating the exclusive global model $w_{g,i}$ of client $i$; and (b) optimizing the local objective $F_i(v_i)$ for updating the personalized model $v_i$ of client $i$. To relate these tasks, a regularization coefficient λ is introduced to control the distance between the personalized model $v_i$ of client $i$ and its exclusive global model $w_{g,i}$. The optimization problem for each client $i$ is formulated by the following equations:

$$\underset{v_i}{\min } u_i(v_i;w_{g,i})=F_i(v_i)+{\displaystyle \frac{\lambda }{2}}\parallel v_i-w_{g,i}{\parallel }^2$$

(2)

$$s.t. w_{g,i}\in \underset{w}{\arg \min } G_g(F_1(w),\dots ,F_N(w\left)\right)$$

(3)

where, when $\lambda$ is equal to 0, the FedDWA objective is transformed to focus on the individual client’s local objective $F_i(v_i)$, and the model training process is equivalent to training only locally. On the contrary, if $\lambda$ is infinite, the FedDWA objective is close to the global objective $G_g(\cdot )$, at which point the model training process is equivalent to FedAvg.

4. Overview and Implementation

4.1. Overview

The overall process of FedDWA is shown in Figure 1. First, the server sends the exclusive global model to the corresponding client, which updates the local personalized model and the exclusive global model in turn. The training process of the personalized model takes the exclusive global model as a reference and dynamically adjusts the extent to which the personalized model refers to the exclusive global model through hyperparameters $\lambda$. Next, the client uploads the model updates corresponding to the updated exclusive global model to the server. The server calculates the cosine similarity between the model updates uploaded by the clients and computes the corresponding aggregation weights. Finally, the server aggregates the exclusive global model for each client.

4.2. Implementation and Algorithm Description

To address the above FedDWA objective, this work alternately updates the exclusive global model ${\left\{w_{g,i}\right\}}_{i\in [N]}$ and the personalized model ${\left\{v_i\right\}}_{i\in [N]}$ of each client.

For the exclusive global model $w_{g,i}$ of client $i$, when the server receives the model updates from each client, it calculates the aggregation weight based on the similarity between the model updates. Clients with higher similarity have more aggregation weight in the aggregation process. In machine learning, there are many similarity measures such as Euclidean distance, cosine distance, and Jaccard distance. Among them, Euclidean distance and cosine similarity are widely used. Euclidean distance measures the absolute distance between points in space, which is directly related to the position of the points, while cosine distance is more reflective of the difference in direction by measuring the angle between vectors. Since model parameters usually reflect the convergence direction of the model, and cosine similarity is more robust than Euclidean distance in evaluating the similarity of high-dimensional model parameters [22], cosine similarity is more appropriate than Euclidean distance in evaluating the similarity of model update vectors [23,24]. Therefore, in this paper, cosine similarity is used to measure the similarity between models. Taking the exclusive global model aggregation of client $i$ as an example, the server adopts the cosine similarity to measure the similarity between the model update vectors $\Delta w^i$ and $\Delta w^j$, which can be expressed by the following formula:

$$simi(i,j)={\displaystyle \frac{(\Delta w^i,\Delta w^j)}{\left|\right|\Delta w^i\left|\right|\cdot \left|\right|\Delta w^j\left|\right|}}(i,j\in C)$$

(4)

where $\Delta w^i$ and $\Delta w^j$ denote the model updates of client $i$ and client $j$ in a particular training round, respectively. The value of similarity $simi(i,j)$ is in the range of [−1,1]. The closer the value is to 1, the more similar the direction of model updates between clients; the closer the value is to −1, the larger the difference in direction between model updates. $C$ is the set of participating users in the current phase.

For the choice of normalization method for similarity values, the characteristics of direct scaling and Softmax are taken into account; Softmax, due to its exponential function property, can amplify the differences between larger values, thus highlighting more significantly the contribution of high similarity clients while ignoring the influence of low similarity clients. This feature is particularly useful in situations where the task needs to differentiate the importance of different clients. In contrast, direct scaling, while also capable of normalization, is insensitive to differences in similarity values, resulting in a more even distribution of weights, which may ignore the contributions of high similarity clients. Therefore, we chose Softmax as the normalization method to ensure that the differences in similarity between clients are more accurately reflected. The cosine similarity between client $i$ and the model updates of all other clients is normalized by applying the softmax function [25] to obtain the aggregation weight ${\beta }_{ij}$:

$${\beta }_{ij}={\displaystyle \frac{e^{simi(i,j)}}{{\displaystyle \sum _{k\in C}{e}^{simi(i,k)}}}}$$

(5)

where ${\beta }_{ij}$ denotes the aggregation weight of client $j$ over client $i$ at the time of aggregation.

Meanwhile, considering the weight of its own model update, the aggregation of the exclusive global model of client $i$ in round $t+1$ is shown below:

$$w_{g,i}^{t+1}=\alpha \cdot w_i^t+{\displaystyle \sum _{j\in C,j\ne i}{\beta }_{ij}\cdot w_j^t}$$

(6)

where $w_{g,i}^{t+1}$ denotes the exclusive global model of client $i$ in round $t+1$. $w_i^t$ is the model trained locally by the exclusive global model $w_{g,i}^t$, and $w_j^t$ is the model trained by the exclusive global model of the other client $j(j\ne i)$. $\alpha$ denotes the aggregated weights of client $i$ itself, and ${\beta }_{ij}$ is the weights from the other clients $j$. Here, the sum of $\alpha$ and all ${\beta }_{ij}$ is 1:

$$\alpha +{\displaystyle \sum _{j\in C,j\ne i}{\beta }_{ij}=1}$$

(7)

After aggregation, the server sends the exclusive global model to the client. Subsequently, the client performs personalization training locally. Personalization training is divided into two steps: personalized model update and exclusive global model update. In addition, to improve the efficiency of model training, the data are split into small batches and fed into the network model for training in an iterative manner. The training method used is the stochastic gradient descent (SGD) algorithm, which has high training efficiency on small batches of data. SGD first computes the error $F_i(w_{g,i})$ between the model-predicted labels and the true labels by forward propagation, and then computes the gradient $\nabla F_i(w_{g,i})$ of the model parameters by backpropagation in order to update the model parameters to reduce the error. The exclusive global model update process for client $i$ is shown below. The updates to the exclusive global model for client $i$ are as follows:

$$w_i^t=w_{g,i}^t-\eta \nabla F_i(w_{g,i}^t)$$

(8)

where $\eta$ is the learning rate and $\nabla F_i(w_{g,i}^t)$ denotes the gradient of the loss function for client $i$.

In order to better fit the client’s data distribution, the server trains additional personalized models for each client. The training process of the personalized model uses the exclusive global model as a reference, and the distance between the personalized model and the exclusive global model is dynamically adjusted through hyperparameters $\lambda$ to ensure that the model finds a balance between global and personalized learning. Meanwhile, a small batch gradient descent algorithm is used in the training process to further improve the training efficiency. In addition, a regularization term is added to the cross-entropy loss function of the personalization model for updating. The local gradient update part $\nabla F_i(v_i^t)$ enables the model to better fit the client’s own data distribution, while the regularization term $\lambda (v_i^t-w_{g,i}^t)$, on the other hand, achieves a balance between personalization needs and global optimization by controlling the distance between the personalization model and the global model, thus avoiding the model from overfitting the local data or deviating from the direction of global optimization. The updates to the personalized model $v_i^t$ for client $i$ are as follows:

$$v_i^{t+1}=v_i^t-\eta (\nabla F_i(v_i^t)+\lambda (v_i^t-w_{g,i}^t))$$

(9)

where $v_i^t$ denotes the personalized model for training client $i$ in round $t$, $\lambda$ is the regularization coefficient, and $w_{g,i}^t$ denotes the exclusive global model received by client $i$ in round $t$.

The specific process of FedDWA is shown in Algorithms 1 and 2. The flowchart of the FedDWA algorithm is shown in Figure 2.

Algorithm 1. FedDWA Client Algorithm.

Input: Exclusive global model $w_{g,i}^t$ Output: $w_i^t$//exclusive global model after training

1: Client $i$ receives its corresponding exclusive global model

2: for each local epoch do

$3:\hspace{1em}\hspace{1em}\mathrm{for}\mathrm{each}\mathrm{batch}b=\{x,y\}$$\mathrm{of}{D}_{\mathrm{i}}$ do

$4:\hspace{1em}\hspace{1em}\hspace{1em}\hspace{1em}{w}_i^t=w_{g,i}^t-\eta \nabla F_i(w_{g,i}^t)$

$5:\hspace{1em}\hspace{1em}\hspace{1em}\hspace{1em}{v}_i^{t+1}=v_i^t-\eta (\nabla F_i(v_i^t)+\lambda (v_i^t-w_{g,i}^t))$

6:   end for

7: end for

8: return $w_i^t$

Algorithm 2. FedDWA Server Algorithm.

Input: client set $C$, where $\left|C\right|=N$, total communication rounds T, learning rate $\eta$, local epochs $E$, Client participation rate $F$. Output: personalized models ${\left\{v_i\right\}}_{i\in [N]}$ Initialize: exclusive global model ${\mathrm{w}}_g^0$, personalized models ${\left\{v_i\right\}}_{i\in [N]}$

$1:C_F\leftarrow$$\mathrm{randomly}\mathrm{select}N\times F$ clients

$2:\mathrm{for}t=0,\dots ,T-1$ do

3:  for each client i in parallel do

$4:\hspace{1em}\hspace{1em}\hspace{1em}\hspace{1em}\mathrm{Send}\mathrm{the}\mathrm{exclusive}\mathrm{global}\mathrm{model}{w}_{g,i}^t$ to client $i$.

$5:\hspace{1em}\hspace{1em}\hspace{1em}\hspace{1em}{w}_i^t$$\leftarrow$$\mathrm{Algorithm}1(w_{g,i}^t$)

$6:\hspace{1em}\hspace{1em}\hspace{1em}\hspace{1em}\mathrm{for}\mathrm{client}i\in C_p$ do

$7:\hspace{1em}\hspace{1em}\hspace{1em}\hspace{1em}\hspace{1em}\hspace{1em}simi(i,j)\leftarrow similarity(w_i^t,w_j^t)$$//\mathrm{Calculate}\mathrm{the}\mathrm{similarity}\mathrm{between}{w}_i^t$$\mathrm{and}{w}_j^t$.

$8:\hspace{1em}\hspace{1em}\hspace{1em}\hspace{1em}\hspace{1em}\hspace{1em}{\beta }_{ij}={\displaystyle \frac{e^{simi(i,j)}}{{\displaystyle \sum _{k\in C_p}{e}^{simi(i,k)}}}}$//Normalize to get the aggregation weight.

9:     end for

$10:\hspace{1em}\hspace{1em}\hspace{1em}\hspace{1em}{w}_{g,i}^{t+1}=\alpha \cdot w_i^t+{\displaystyle \sum _{j\in C,j\ne i}{\beta }_{ij}\cdot w_j^t}$

11:  end for

12: end for

4.3. Computational Complexity Analysis

In the following section, we will analyze the computational complexity of the FedDWA algorithm. The FedDWA algorithm consists of two main phases: model training and exclusive global model aggregation.

Time complexity:

1.

Model Training:

In FedDWA, each client needs to use its local dataset $D_i$ to train the personalized model $v_i$ and the exclusive global model $w_{g,i}$. Specifically, if the number of local iterations for each client is $E$, the number of samples in the client dataset $D_i$ is $n_i$, and the complexity of computing the model gradient is $O(d)$ (where $d$ denotes the dimensionality of the model parameter); the time complexity of training personalized models and exclusive global models for each client is $O(2\cdot E\cdot n_i\cdot d)$.

2.

Exclusive global model aggregation:

The server aggregates exclusive global models from $K$ clients. The aggregation process involves a weighted summation of model parameters across all clients, with a time complexity of $O(K\cdot d)$. Therefore, the overall time complexity of FedDWA is as follows:

$$\begin{array}{l}O(2\cdot {\displaystyle \sum _{i=1}^{K}E}\cdot n_i\cdot d+K\cdot d)\\ =O(2\cdot E\cdot d\cdot {\displaystyle \sum _{i=1}^K{n}_i}+K\cdot d)\\ \approx O(2\cdot E\cdot d\cdot K\cdot \overline{n}+K\cdot d)\\ =O(K\cdot d\cdot (2\cdot E\cdot \overline{n}+1))\\ \approx O(K\cdot d\cdot 2\cdot E\cdot \overline{n}))\\ \approx O(K\cdot E\cdot \overline{n}\cdot d))\end{array}$$

(10)

where in Equation (10), $\overline{n}$ denotes the average number of samples across all clients.

Space complexity:

1.

Clients:

Each client needs to store the personalized model parameters $v_i$ and exclusive global model parameters $w_{g,i}$. The storage space complexity of both parts is $O(d)$. In addition, during the local training process, the client may also occupy some space for storing gradient information and intermediate computation results, but these space requirements are usually small compared to the space for storing $v_i$ and $w_{g,i}$.

Therefore, the total space complexity of each client is $O(d)$.

2.

Server:

The server needs to store the exclusive global model parameters $w_{g,i}$ and the aggregated weights ${\beta }_{ij}$ for each client. The space complexity of storing the exclusive global model parameters is $O(d)$ and the space complexity of storing the aggregated weights is $O(K)$. Therefore, the total space complexity of the server is $O(d+K)$.

To summarize, the overall space complexity of the FedDWA algorithm is as follows:

$$O(K\cdot d+K)\approx O(K\cdot d)$$

(11)

Therefore, the time complexity of the FedDWA algorithm is $O(K\cdot E\cdot \overline{n}\cdot d)$, and the space complexity is $O(K\cdot d)$. Compared to such algorithms like FedAvg, Ditto, and IFCA, FedDWA has the same time complexity and gives preference to other algorithms in terms of performance and convergence speed.

5. Experiments and Analysis

5.1. Datasets

We used five real-world datasets to evaluate the performance of the algorithms: CIFAR-10, CIFAR-100, Tiny-ImageNet, MNIST, and Fashion-MNIST(FMNIST). The CIFAR-10 dataset consists of 60,000 RGB color images in 10 categories with a 32 × 32 pixel resolution and contains 6000 images per category, including 5000 training images and 1000 test images. CIFAR-100 is an extended version of the CIFAR-10 dataset, which is widely used in image classification tasks. This dataset contains 60,000 32 × 32 pixel color images divided into 100 categories, each category contains 600 images, and the dataset is divided into 50,000 training images and 10,000 test images. The Tiny-ImageNet dataset is a subset of the ImageNet dataset, and the categories are more diversified compared with the previous two datasets, with a total of 200 categories of image data and image resolution. The Tiny-ImageNet dataset is a subset of the ImageNet dataset with more diverse categories than the previous two, containing a total of 200 categories of image data with a resolution of 64 × 64 pixels, and each category contains 500 training samples and 50 validation samples.

The MNIST dataset consists of 70,000 grayscale images of handwritten digits in 10 classes, each with a resolution of 28 × 28 pixels. The dataset includes 60,000 training images and 10,000 test images. The FMNIST dataset consists of 70,000 grayscale images of fashion products in 10 classes, each with a resolution of 28 × 28 pixels. Similar to the MNIST dataset, FMNIST includes 60,000 training images and 10,000 test images.

In this paper, we considered two methods of data partitioning that are widely used in the field of federated learning to simulate practical heterogeneous data distribution and pathological heterogeneous data distribution, respectively. In the practical heterogeneous distribution, we controlled the distribution of data on each client through Dirichlet distribution. Among them, for the three datasets of CIFAR-10, CIFAR-100, and Tiny-ImageNet, we chose the same Dirichlet coefficient (denoted as dir) for the experiment and set the value of dir to 0.3 and 0.5, where smaller dir values indicate higher heterogeneity. Pathological heterogeneous data distribution is the earliest heterogeneous data scenario considered and studied in federated learning. In pathological heterogeneous data distribution, we only assigned 2 classes to each client from the 10 classes of CIFAR-10, MNIST, and FMNIST, and ensured that the data samples owned by each client were disjoint [1]. Due to space limitations, the complete experimental setup and experimental results under pathological heterogeneous data distribution will be provided in Appendix A and Appendix B.

5.2. Model Settings

For all datasets, a convolutional neural network model was used in this paper for pattern recognition. The model consists of a feature extractor and a classifier. The feature extractor consists of two convolutional layers, each followed by a ReLU activation function and a 2 × 2 maximum pooling layer. The first convolutional layer uses 32 5 × 5 convolutional kernels to extract basic features from the input features, and the second convolutional layer uses 64 5 × 5 convolutional kernels for feature extraction. After feature extraction, the model maps the extracted features onto a 512-dimensional vector through a fully connected layer. This vector is passed through the ReLU activation function and finally converted into the output of the model by another fully connected layer.

5.3. Baselines

In this paper, seven FL algorithms were selected for comparison, including a classical FL algorithm: FedAvg, and eight PFL algorithms, FedAvg-FT, deploy the trained global model to the client and performs a step of fine-tuning on the local data; Ditto achieves fairness and robustness by introducing a regularization term to induce the client model to be aligned with the global model; IFCA minimizes the loss function by estimating the client’s cluster identity; FedGH [26] achieves personalization by introducing a uniform global header to reconcile model differences across clients and personalization tuning on top of that; FedPAC [27] achieves personalized federated learning through feature alignment and collaborative classifiers; FedSoft [7] allows each client to belong to multiple clusters at the same time, thus enabling a more flexible collaboration structure; FedCollab [28] improves model performance by clustering clients based on the data distribution distance and the data volume to optimize the client collaboration structure; and FedALA [29] adapts to aggregate global and local models through the adaptive local aggregation module that optimizes the local objective of each client, thus solving the problem of statistical heterogeneity in federated learning.

5.4. Parameter Settings

In all experiments, this paper used the SGD optimizer to train the models, and the loss functions are all cross-entropy loss functions. The number of local training rounds for all algorithms was uniformly set to 5, the learning rate was uniformly set to 0.01, the batch size was set to 16, the total number of clients was set to 20, and the participation rate $F$ was set to 1. In addition, we set the number of communication rounds to 100, 200, and 300 in the CIFAR-10, CIFAR-100, and Tiny-ImageNet datasets, respectively. In Ditto and FedDWA, the parameter $\lambda$ was set to 1. When FedDWA performs model aggregation, the model aggregation weight $\alpha$ was set to 0.2. In addition, the number of clusters was set to 4 in the experiments of the IFCA algorithm.

5.5. Results and Discussion

In this section, we conducted extensive experiments to determine the advantages of FedDWA in terms of model convergence speed and performance. At the same time, we demonstrate the superiority of FedDWA over other algorithms under various datasets and data heterogeneity conditions.

Table 1. A comparison of the average test accuracy (%) of the FedDWA algorithm with baseline algorithms on the CIFAR-10, CIFAR-100, and Tiny-ImageNet datasets.

Algorithm	CIFAR-10		CIFAR-100		Tiny-ImageNet
	dir = 0.3	dir = 0.5	dir = 0.3	dir = 0.5	dir = 0.3	dir = 0.5
FedAvg	53.44	58.32	22.46	26.13	16.37	18.06
FedAvg-FT	82.87	76.14	42.95	29.09	39.79	31.28
Ditto	85.41	77.53	39.66	28.36	39.71	33.12
IFCA	85.84	79.7	42.9	30.05	36.76	33.28
FedGH	79.07	73.56	41.31	36.69	39.71	31.87
FedPAC	84.67	78.65	42.28	34.77	39.54	33.86
FedCollab	82.66	76.34	43.84	36.64	36.21	31.17
FedSoft	80.4	77.24	40.58	33.21	36.36	33.37
FedALA	83.05	77.97	42.18	36.3	34.79	33.32
FedDWA (Ours)	86.97	81.74	45.63	38.13	41.83	35.23

shows the average test accuracies achieved by FedDWA and other algorithms on CIFAR-10, CIFAR-100, and Tiny-ImageNet datasets under different data heterogeneity scenarios. Among them, the bold in the Table 1 indicates the best accuracy.The experimental results show that the personalized federated learning algorithm generally outperforms the traditional federated learning algorithm FedAvg under different datasets and data heterogeneity conditions, which further proves the effectiveness of the personalized federated learning algorithm in dealing with data heterogeneity problems. Additionally, FedDWA achieves the best average test accuracy in most cases. Even in Non-IID conditions with Dirichlet coefficients of 0.3 and 0.5, FedDWA achieves the highest accuracies of 41.83% and 35.23%, respectively, on the large and complex dataset Tiny-ImageNet.

Figure 3, Figure 4 and Figure 5 show the average test accuracy of different algorithms on the CIFAR-10, CIFAR-100, and Tiny-ImageNet datasets. It is obvious from these figures that FedDWA significantly outperforms other algorithms in terms of accuracy. This is because the FedDWA algorithm flexibly assigns aggregation weights to each client based on the cosine similarity between client model updates, thereby aggregating an exclusive global model for each client. Compared with other algorithms, this personalized aggregation method can more accurately enhance collaboration between similar clients and reduce the interference between dissimilar clients, thereby effectively solving the Non-IID data problem. In addition, in the local update phase, FedDWA introduces the hyperparameter $\lambda$ to adjust the distance between the local model and the exclusive global model to avoid over-reliance on the exclusive global model or the ignoring of the characteristics of local data, so as to optimize the local training process. At the same time, Figure 3, Figure 4 and Figure 5 also show that FedDWA outperforms other algorithms in terms of convergence speed. This is because its adaptive aggregation weights and hyperparameter adjustment in the local training phase significantly accelerate the convergence of the model. In addition, for relatively simple datasets like CIFAR-10, the vast majority of the algorithms are able to converge quickly within the first 20 rounds and reach an accuracy close to 75% in a short period of time. However, when facing the more complex Tiny-ImageNet dataset, the convergence speed of most algorithms slows down significantly and the convergence curves show large fluctuations. This suggests that the complexity and heterogeneity of the dataset significantly increase the difficulty of convergence of the federated learning models, leading to the difficulty of stable convergence of the models within a limited number of communication rounds. By observing Figure 4 and Figure 5, it can be seen that FedDWA demonstrates a strong advantage in dealing with complex datasets, and its final accuracy is usually better than other algorithms. This suggests that FedDWA is able to adapt more effectively to the differences in client data distribution and promote effective collaboration between similar clients through the dynamic weight allocation mechanism, thus significantly improving the accuracy of the model.

Table 2. Communication rounds to reach target test accuracy (C_acc(x)) under Non-IID settings.

Algorithm	Non-IID (dir = 0.3)	Non-IID (dir = 0.5)
	Caac (36%)	Caac (30%)
	Round	Round
FedAvg-FT	82	149
Ditto	106	62
IFCA	125	72
FedGH	131	192
FedPAC	110	91
FedCollab	153	198
FedSoft	183	85
FedALA	-	84
FedDWA (Ours)	51	48

shows the number of communication rounds required for FedDWA and the comparison algorithm to achieve the target test accuracy on the Tiny-ImageNet dataset. Where the best results are marked in bold. The symbol ‘-’ indicates that the target value is not reached within a limited number of communication rounds. As can be seen from Table 2, FedDWA achieves the target accuracy using the minimum number of communication rounds. This result shows that FedDWA exhibits significant advantages in terms of communication efficiency, which not only takes full advantage of the personalized model but also achieves competitive performance on Non-IID data while reducing the communication overhead.

5.6. Hyperparametric Analysis

In the FedDWA algorithm, the hyperparameter $\lambda$ is used to control the degree of dependence of the personalized model on the exclusive global model to better fit the client’s local data distribution. In the experiments, the initial setting is $\lambda =1$. To investigate the effect of $\lambda$ on model performance, experiments were conducted on three datasets—CIFAR-10, CIFAR-100, and Tiny-ImageNet—using different values of $\lambda$: {0.3, 0.5, 1, 1.5, 2}.

In Figure 6, we show how the model performance varies for different values of the hyperparameter $\lambda$. Overall, $\lambda =1$ produces better results. When the value of $\lambda$ is too low or too high, the performance of FedDWA decreases. When $\lambda$ is too low, the personalized model finds it difficult to effectively utilize the information from the exclusive global model; when $\lambda$ is too high, the model is overly reliant on the exclusive global model and ignores the properties of the local data, leading to a decrease in model adaptability. The experimental results show that, therefore, choosing a reasonable value of $\lambda$ ensures a balance between the independent global model and the personalized model, thus improving the overall accuracy of the model.

In the FedDWA algorithm, hyperparameter $\alpha$ denotes the aggregation weight of the client’s own model when the exclusive global model is aggregated. In Figure 7, we show how the model performance varies with different values of the hyperparameter $\alpha$. From Figure 7, we can see that either too large or too small a value of hyperparameter $\alpha$ will reduce the model testing accuracy of FedDWA. Under the Non-IID setting, FedDWA performs best when $\alpha$ is taken to be 0.2 on the CIFAR-10, CIFAR-100, and Tiny-ImageNet datasets, which is probably due to the fact that there is an implicit grouping of the clients at this point in time situation; when $\alpha$ takes a smaller value, the collaboration between similar clients gains more on the model performance compared to the clients themselves. Therefore, in the experimental part of this paper, we set $\alpha$ to 0.2.

5.7. Impact of Participating Rates

Our experiment focused on the CIFAR-10 dataset with parameter dir = 0.3. In this experiment, the total number of clients was 60. We set the participation rate of clients to {0.05, 0.1, 0.2, 0.5}. Figure 8 and Figure 9 represent the accuracy curves of the personalized model and the global model at different participation rates, respectively. As shown in Figure 8 and Figure 9, both the personalized model and the global model of our algorithm outperform the other algorithms in terms of accuracy under different participation rates and exhibit faster convergence. It is worth noting that many algorithms show obvious convergence instability when the participation rate decreases, while our algorithm is still able to maintain a relatively stable convergence trend in this case, which fully proves its robustness in dealing with different participation rates.

5.8. Discussion on Model Similarity Measures

In this section, we compare the FedDWA algorithm using cosine similarity and Euclidean distance as model similarity metrics through experiments on the CIFAR-10 dataset and analyze their performance in average test accuracy to evaluate the effectiveness of these two metrics. The experimental results are shown in Figure 10. The experimental results show that the FedDWA algorithm using cosine similarity as the similarity metric performs better in accuracy, verifying the effectiveness of cosine similarity in measuring model similarity. This is because cosine similarity can more accurately reflect the similarity between client model updates by considering the directional differences between client model updates rather than relying solely on the absolute size of the distance. In contrast, the FedDWA algorithm using Euclidean distance is greatly affected by data heterogeneity, resulting in lower model accuracy. Therefore, this paper used cosine similarity to measure the similarity between model updates.

6. Conclusions

In this paper, we propose a personalized federated learning algorithm based on dynamic weight allocation, FedDWA. In the server aggregation phase, the algorithm enhances collaboration among similar clients and encourages them to learn more similar features by calculating the cosine similarity of model updates among clients and dynamically allocating aggregation weights based on the similarity. In the local update phase, each client trains a personalized model by adjusting the distance between the local model and the global model after receiving the exclusive global model. This not only improves the generalization ability of the model but also enables the model to better adapt to the local data distribution. The experimental results show that FedDWA shows significant improvements in convergence speed and model accuracy compared to all benchmark algorithms on the CIFAR-10, CIFAR-100, and Tiny-ImageNet datasets. Considering the diversity of real-world data, future research will focus on using more complex and diverse non-IID datasets to further analyze the adaptability and robustness of FedDWA. In addition, considering more realistic application scenarios, we will expand the scale of experiments in the future to verify the scalability of FedDWA in large-scale client environments, especially in the case of hundreds or thousands of clients, evaluate its performance in large-scale distributed scenarios, and explore how to further reduce the algorithm’s computational and communication overhead when the number of clients surges. Finally, we also plan to introduce appropriate privacy-preserving techniques into FedDWA to further enhance the security of the joint learning system and ensure its effectiveness in real-world applications.