An image classification method based on personalized federated learning

Abstract

The present application relates to an image classification method based on personalized federated learning, belonging to the technical field of computer image processing. The method applies model decoupling and model contrast learning to the image classification task in personalized federated learning, enabling the federated learning to train a personalized global model that can adapt to all local data of the clients. Since the global model trained on the overall data set can learn better representation than the local model trained on the skewed subset, the present application adds a regularization loss function of model contrast to the local model update of the client, so that the feature extractor of the local model constantly approaches the feature extractor of the global model that learns better representation, which can reduce the performance difference between the local models and thus produce a better personalized global model. Compared with the prior art, the present application can reach convergence faster in terms of communication efficiency and has better scalability.

Description

Technical Field

[0001] This invention relates to an image classification method based on personalized federated learning, belonging to the field of computer image processing technology. Background Technology

[0002] With the enactment of data privacy laws such as the General Data Protection Regulation (GDPR), people have begun to pay attention to the privacy protection of user data in the field of artificial intelligence. To protect user data privacy, Google pioneered federated learning technology. This process allows participating parties to train their local data using a global model sent by the server in each round, and then upload the trained parameters to the server for aggregation, without sharing the original data. These participating parties can be different institutions, organizations, companies, or individuals, each holding a different dataset, which may contain sensitive information such as personal privacy data or trade secrets. This approach can protect user data privacy while ensuring the accuracy and robustness of the model, thus showing broad application prospects in tasks such as image classification.

[0003] However, federated learning methods still face several challenges. One key challenge is the data heterogeneity problem. Each client's local data may have different distributions, data types, sizes, and qualities, all of which can negatively impact model performance. In cases where cross-device data is non-independent and individually distributed (non-IID), a single globally shared model struggles to adapt to differing global and local distributions. This affects the global model's training performance and its performance on the participants' local data, potentially leading to participants with poor training performance ceasing further participation in federated learning model training.

[0004] To address the data heterogeneity problem in federated learning, several federated learning methods have proposed corresponding solutions. For example, Federated Averaging (FedAvg), proposed by McMahan et al., is currently the most commonly used federated learning algorithm, which aggregates model parameters through weighted averaging. However, when FedAvg learns on non-IID data on the client side, the global model after training may deviate from the global optimum after averaging, resulting in a significant reduction in post-training accuracy. FedProx introduces a local regularization term to penalize the distribution differences of the model among different participants. SCAFFOLD alleviates client-drift by adding a correction term to the update of the local model. MOON incorporates a model local loss regularization term that compares the local model and the global model during local updates, ensuring that each client's local model learns representations close to the global model during training. FedRep trains the entire model sequentially during the local update phase. Each participant updates the classifier using the aggregated representation. Within an epoch, it updates the feature extractor using only its own classifier. In the aggregation phase, only the feature extractor is aggregated, and both the global model's classifier and feature extractor are updated to reflect the contributions of all participants. FedBABU trains only the model with randomly initialized classifiers during the local update phase and aggregates only the feature extractor during the aggregation phase. FedPer trains the entire model during the local update phase and aggregates only the bottom-level feature extractors during the aggregation phase, excluding the top-level classifiers. It allows for personalized settings of the classifier in the local model, while the feature extractor is shared across all participants.

[0005] Existing federated learning methods have the following drawbacks:

[0006] 1. Existing federated learning methods use large-scale gradient information in each iteration, which leads to high computational and communication overhead. For example, the MOON method involves computation for the entire model's contrastive learning. Especially when there are many participants and a large amount of data, processing and transmitting large-scale gradient information can become complex and time-consuming.

[0007] 2. Regarding model scalability, the reliance on the merging and aggregation of local models can limit its scalability in some cases. Especially when facing large-scale participants and complex models, effectively merging and aggregating model parameters can be challenging, making it difficult to achieve high performance on image datasets with deep learning models.

[0008] 3. Existing learning methods face the challenge of choosing the right hyperparameters. Some methods, such as SCAFFOLD and FedProx, involve selecting hyperparameters, such as the weight of the regularization term and the learning rate. Choosing appropriate hyperparameters is crucial for the performance and safety of the algorithm, but it requires adjustment based on the specific problem and dataset, which may require a lot of trial and error and optimization. Summary of the Invention

[0009] The purpose of this invention is to address the problems and shortcomings of existing technologies. To reduce the impact of non-IID data issues on the performance of the global model in federated learning and improve the accuracy and efficiency of federated learning in image classification tasks, this invention creatively proposes an image classification method based on personalized federated learning. This method applies model decoupling and model contrastive learning to the image classification task in personalized federated learning, enabling federated learning to train a personalized global model that can adapt to local data from all clients.

[0010] Research on federated learning models reveals that, in cases of imbalanced data distribution, the neural network classifier's learning based on linear decision boundaries is unfair in classifying imbalanced datasets. Therefore, this invention proposes decoupling the model into a classifier and a feature extractor, fixing the classifier during local model training. Since a global model trained on the entire dataset learns better representations than a local model trained on a skewed subset, this invention incorporates a regularized loss function for model comparison during local model updates on the client side. This allows the feature extractor of the local model to continuously approach the feature extractor of the global model, which has better representation learning, thereby reducing performance differences between local models and resulting in a better personalized global model.

[0011] like Figure 1 As shown, the method of this invention trains a powerful personalized global model. Specifically, it uses the feature extractor w of the global model ResNet. ext and classifier w cls Training is done separately, and only the feature extractor w of the local model of the participating party is updated. ext Do not update its classifier w cls When updating the feature extractor of the local model of the participating party, the representation distance between the local model and the global model is reduced, and the representation distance between the local model and the local model in the previous training round is increased.

[0012] When the client updates locally, it adjusts the feature extractors of the local models respectively. and classifier The learning rate is used to update both.

[0013] Among them, the local model in the client's t-th round During training, the classifier of the local model in round t will be... The learning rate is set to 0, and the classifier of the local model in round t is not evaluated. Update; change the feature extractor of the local model in round t. The learning rate is set to η, and the feature extractor of the local model in the t-th round is updated only. Feature extractor of the local model in round t During the update, the loss values ​​of the two parts, cross-entropy loss l1 and model contrast loss l2, are added for supervised learning.

[0014] During model training, for each sample's input feature x, the feature extractor of the global model in round t is used. Extracting the representation of x Feature extractor from the local model in round t-1 Extracting the representation of x Feature extractor from the local model in round t Extracting the representation of x

[0015] To facilitate the representation of the calculation of the contrast loss l2 of the model, let Since the global model can learn a better representation than the local model, the goal of this method is to reduce p and p glob The distance between them, and increase p and p prev The distance between them.

[0016] Feature extractor of the local model in the t-th round on the client side During the update, client C i The local loss l includes the cross-entropy loss l1 and the model contrastive loss l2 from supervised learning, that is: Where μ represents a hyperparameter used to control the weights of the model's contrastive loss l2; (x, y) represents the sample data, and x represents the input feature of the sample. The model's contrastive loss l2 is defined as:

[0017]

[0018] Where τ is a temperature parameter used to control the smoothness of the model output distribution.

[0019] The client's goal is to minimize the local loss l during local model training, defined as:

[0020]

[0021] Among them, D i This represents the local data of the i-th client.

[0022] During the model aggregation phase, the server is only responsible for aggregating and updating the feature extractors of the local model in round t. The client uses local data Update the feature extractor of the local model in round t. Classifier without updating local model initialization The goal of globalization is to solve this problem:

[0023]

[0024] Where w represents the global model; This represents the average loss value generated by the server aggregating the training of N client models. This represents the feature extractor for the local model in round t. Represents the classifier used for local model initialization; This represents the local data of N clients participating in the training. sum; This represents the loss value calculated by each client for each sample data (x, y) on the local data. Represents local data for each client The loss value calculated for all sample data (x, y) in L. i It is a feature extractor using the t-th round local model. Classifier initialized with local model Calculate l i The obtained expected loss reflects the model How well it fits the local dataset.

[0025] Beneficial effects

[0026] Compared with existing technologies, this invention has the following advantages:

[0027] 1. The FedBOON method proposed in this invention is a simple and effective federated learning framework. Its core idea is to improve the training process of federated learning at the model level. By decoupling the model and using the similarity between the main representations of the model to correct the local training of the client, it can alleviate the data heterogeneity problem in the federated learning architecture, reduce the weight difference between the local model and the global model, and improve the generalization ability and personalization ability of the global model.

[0028] 2. This invention achieves faster convergence in terms of communication efficiency. During local client updates, only the feature extractor of the local model is updated. After the update, the local model's representation continuously learns from the global model's representation, improving the performance and adaptability of the local model. During server parameter aggregation, only the model's feature extractor is aggregated, allowing for rapid updates to the global model and accelerating the server's parameter aggregation and model iteration process.

[0029] 3. This invention has certain advantages in terms of scalability. Because the FedBOON method fixes the model classifier during local updates, it reduces the negative impact of linear decision boundary learning in image classification tasks, mitigating overfitting issues when clients have limited data or uneven data distribution. This provides scalability in terms of data volume and distribution. Regarding model complexity, this invention reduces model complexity and improves scalability by utilizing model decoupling and contrastive learning. In terms of server-side model aggregation, only the model's feature extractor is aggregated. Therefore, when more clients participate in model training, it reduces complex parameter aggregation and communication overhead, providing scalability in terms of the number of clients. Attached Figure Description

[0030] Figure 1 This is a schematic diagram of the architecture of this method;

[0031] Figure 2 This is a schematic diagram of the model training process;

[0032] Figure 3 This diagram illustrates a comparison between our method and existing federated learning methods.

[0033] Figure 4 This is a schematic diagram of the image classification process. Detailed Implementation

[0034] The specific implementation process of the present invention will be further described in detail below with reference to the accompanying drawings.

[0035] Before training the federated learning model, input the model training parameters, including the number of communication rounds T, the number of clients N, the local batch size B, the local epochs E, the learning rate η, the temperature parameter τ, and the hyperparameter μ.

[0036] like Figure 2 As shown, an image classification method based on personalized federated learning includes the following steps:

[0037] Step 1: The server randomly initializes the parameters of the global model and sends them to each client participating in the training.

[0038] Specifically, in federated learning, N clients are set up, denoted as C1, ..., C2.N Each client C i There is a local data Where (x, y) is a sample data, x is the input feature of the sample, and y is the label corresponding to the sample; K (i) It is local data The number of samples in the data represents the local data. Size.

[0039] Before model training, the server first randomly initializes the global model in round t, obtaining... This represents the classifier used for global model initialization. This represents the feature extractor of the global model in round t. During model training, the server will use the global model from round t. Send to the clients participating in the training.

[0040] Step 2: After receiving the global model, the client participating in the training updates the model on the local dataset using stochastic gradient descent (SGD).

[0041] Specifically, the client obtains the local model initialized in round t. i∈[1,N], where N represents the number of clients participating in model training, and t∈[1,T], where T represents the total number of communications during model training. Let D represent the global model in round t-1. Each client uses local data D. i Perform an update to epochsE. Client C i Receive server sends to C i Global model in round t-1 This represents the classifier used for global model initialization. This represents the feature extractor for the global model in round t-1. The client updates the feature extractor for the local model in round t. Initialize the classifier of the local model and set it as the classifier initialized for the global model. get This represents the classifier used for local model initialization. Client C i Update the feature extractor of the local model in round t using local data. Right now The client's local loss l includes the cross-entropy loss l1 and the model contrast loss l2 from supervised learning.

[0042] Step 3: After updating the model, the client uploads the trained local model parameters to the server.

[0043] Specifically, after each client independently trains a local model using local data, the clients participating in model training will use the feature extractor of the local model in round t. All data is sent to the server for aggregation.

[0044] Step 4: The server performs average aggregation on the model parameters uploaded by the client and updates the global model for the next round of training.

[0045] Specifically, the server aggregates the feature extractors of the local models from the N clients participating in the training in the t-th round. The feature extractor of the global model in round t. Updated to:

[0046]

[0047] in, ∪ represents a set It equals the set of all i in N. The union of .

[0048] Repeat steps 1-4 until the global model converges.

[0049] After T rounds of model training, we obtain This represents the classifier used for global model initialization. This represents the feature extractor of the global model in round T. This represents the global model for round T.

[0050] Step 5: Use the trained global model to perform image preprocessing and feature extraction to complete image classification.

[0051] The accuracy of this method on image datasets was evaluated through experiments. Experimental results are as follows: Figure 3 As shown, this method significantly outperforms other federated learning optimization algorithms. On the CIFAR-10 dataset, the accuracy is improved by approximately 1.8% compared to the compared algorithms, and on the CIFAR-100 and Tiny-ImageNet datasets, the accuracy is improved by at least 4.2%. Notably, on the CIFAR100 dataset with 100 clients, this method achieves a peak accuracy of 73.02%. Figure 4 The process of an image classification task is shown.

Claims

1. An image classification method based on personalized federated learning, characterized in that, The method comprises the following steps: Step 1: the server randomly initializes the parameters of the global model and sends them to each client participating in training; In federated learning, N clients are set up, denoted as C1, ..., C2. N Each client C i There is a local data Where (x, y) is a sample data, x is the input feature of the sample, and y is the label corresponding to the sample; K (i) It is local data The number of samples in the data represents the local data. Size; Before model training, the server first randomly initializes the global model in round t, obtaining... This represents the classifier used for global model initialization. Let represent the feature extractor of the global model in round t; during model training, the server will use the global model from round t. Send to the clients participating in the training; Step 2: after the client participating in training receives the global model, the model is updated using stochastic gradient descent on the local data set; The client obtains the local model initialized in the tth round N represents the number of clients participating in model training, t ∈ [1, T], T represents the total number of communications for model training, Gt-1 represents the global model of the t-1th round; each client uses local data D i Update for epochs E Client C i Receive the server sends to C i Global model of the t-1th round Classifier representing the initialization of the global model, Feature extractor representing the global model of the t-1th round; Feature extractor of the tth round local model updated by the client Initialize the classifier of the local model and set it to the classifier initialized by the global model Get Classifier representing the initialization of the local model; Client C i Update the feature extractor of the tth round local model using local data That is Where the local loss of the client Contains the cross-entropy loss in supervised learning And the model comparison loss Step 3: after the client updates the model, the trained local model parameters are uploaded to the server; After each client independently trains a local model using local data, the clients participating in the model training send the feature extractor of the t-th round local model to the server for aggregation to the server for aggregation; Step 4: the server aggregates the model parameters uploaded by the client and updates the global model for the next round of training; The server aggregates the feature extractor of the t-th round local model of N clients participating in training on average The feature extractor of the t-th round global model is updated as: The feature extractor of the t-th round global model is updated as: wherein, ∪ denotes the set union is equal to the set union of all i in N is the union of the sets Repeat steps 1-4 until the global model reaches convergence; after T rounds of model training, obtain a classifier representing the initialization of the global model, a feature extractor representing the Tth round of the global model, a global model representing the Tth round; Step 5: using the trained global model, image preprocessing and feature extraction are performed to complete image classification.

2. The image classification method based on personalized federated learning according to claim 1, wherein, Feature extractor w of the global model ext and classifier w cls Separately trained, only updating the feature extractor w of the local model of the participant ext , not updating its classifier w cls ; when the feature extractor of the local model of the participant is updated, let the local model reduce the representation distance with the global model and increase the representation distance with the local model in the last round of training; The client adjusts learning rates of the feature extractor and the classifier of the local model respectively when updating locally and realizes the update of both ​ wherein during the training of the client's t-th round local model the learning rate of the classifier of the t-th round local model is set to 0, and the classifier of the t-th round local model is not updated; the learning rate of the feature extractor of the t-th round local model is set to η, and only the feature extractor of the t-th round local model is updated; during the update of the feature extractor of the t-th round local model, the cross-entropy loss of supervised learning and the model contrast loss are added to the loss value of the two parts; During the model training process, for the input feature x of each sample, the representation of x is extracted from the feature extractor of the global model in the t-th round the representation of x is extracted from the feature extractor of the local model in the t-1-th round the representation of x is extracted from the feature extractor of the local model in the t-th round Let ​​​ Feature extractor of the t-th round local model at the client At the update time, the client C i The local loss Includes cross-entropy loss in supervised learning And model comparison loss Where μ represents a hyperparameter used to control the weight of the model comparison loss ; (x, y) represents sample data, x represents the input features of the sample; The model comparison loss Is defined as: Where τ is a temperature parameter used to control the smoothness of the model output distribution; The goal of the client is to minimize the local loss of the local model training is defined as: wherein, represents the local data of the i-th client; In the aggregation phase of the model, the server is only responsible for aggregating the feature extractor of the t-th round of local model update The client utilizes local data The feature extractor of the t-th round of local model update The classifier of the local model initialization is not updated The global target solves this problem: where w denotes the global model; represents the server average aggregation of the loss values trained by N client models; represents the feature extractor of the t-th round local model; represents the classifier initialized by the local model; represents the local data of N participating training clients summed up; represents the loss value calculated by each client on each sample data (x, y) on the local data, represents the loss value calculated by each client on all sample data (x, y) in the local data L i is calculated using the feature extractor of the t-th round local model and the classifier initialized by the local model The expected loss obtained by reflects the degree of fitting of the model to the local data set.

Images