A federated learning training method and system based on a prototype representation

Abstract

The application discloses a kind of federated learning training methods and systems based on prototype representation, based on neural network training model, central server initializes global model and feature adapter network model, and global model parameter is synchronized to all participating training edge device;Local prototype is formed in each round iteration by edge device executing local training, local model is updated, and local prototype and local model parameter are uploaded to server, server is based on local model parameter, local prototype training adapter network updates unbiased prototype and global prototype, global model, and the global model parameter and prototype after updating are sent to edge device, and edge device utilizes unbiased prototype and global prototype to optimize local training.The method and the corresponding training system of the application solve the problem that the model trained by federated learning has weak generalization ability in the feature offset scene, and can improve the accuracy of the model.

Description

Technical Field

[0001] This invention relates to the fields of distributed computing and computer vision, and in particular to a federated learning training method and system based on prototype representation. Background Technology

[0002] In some big data statistical analyses, it is often necessary to aggregate data from edge devices scattered across various locations onto a server for comprehensive analysis. However, due to the privacy concerns surrounding edge device data, traditional centralized deep learning paradigms are no longer sufficient. To leverage geographically dispersed private data, Federated Learning (FL) has emerged. Federated Learning is a privacy-preserving machine learning paradigm that allows multiple edge devices to collaboratively train a global model without disclosing or uploading their local private data to a central server, effectively addressing data privacy and security issues. In recent years, Federated Learning has demonstrated its potential for practical applications in many fields, including medical image analysis, recommender systems, natural language processing, and computer vision.

[0003] For example, in the field of medical image analysis, hospitals and medical institutions in different locations have stored a large number of images related to diseases. Medical institutions hope to use this image data to train a neural network model for disease diagnosis. However, because this data involves sensitive patient information, it is not convenient for different hospitals to share the data. Therefore, federated learning can be used to train the neural network model.

[0004] However, data models obtained through federated learning suffer from data heterogeneity. Data heterogeneity (also known as non-independent and identically distributed, non-IID data) means that the private data held by the various edge devices participating in training have significant differences in statistical characteristics. This heterogeneity is mainly reflected in two aspects: label shift and feature shift. Label shift refers to the inconsistent label distribution of data across different edge devices. Some labels appear more frequently on some edge devices, while they are relatively scarce on others. This uneven distribution may cause the model to have strong recognition ability for some labels but weak recognition ability for others, thus affecting the overall performance of the model. Feature shift refers to the difference in the distribution of data in the feature space across different edge devices. This difference may stem from various factors, such as differences in user behavior in different regions, different types of devices or sensors, etc. Feature shift is particularly critical because it directly affects the quality and relevance of the model input. When the model faces data with large differences in feature distribution, the knowledge it learns may not be effectively transferred to data with different feature distributions, resulting in the model performing well on some edge devices but poorly on others. Feature shift has a more profound impact on federated learning systems than label shift because features are the foundation for model learning and decision-making.

[0005] Taking medical image analysis as an example, due to differences in the model and quality of imaging equipment among different hospitals, the acquired image features vary greatly. For instance, for the same type of tumor, some hospitals may use color images, while others may use black and white images. This leads to neural network models trained based on existing federated learning methods failing to classify accurately, resulting in performance degradation, weak model generalization ability, and insufficient adaptability and robustness to some new scenarios. Summary of the Invention

[0006] Purpose of the invention: To address the weak generalization ability of models trained by existing federated learning methods, this invention provides a federated learning training method and corresponding training system based on prototype representation. By minimizing the loss function, it extracts representative feature representations that are unaffected by the data distribution bias of edge devices, thereby improving the generalization ability of the model.

[0007] Technical solution: To solve the above problems, the federated learning training method based on prototype representation in this invention includes the following steps:

[0008] (1) Initialization steps: Based on the neural network training model, the central server initializes the global model, feature adapter network model, global prototype and unbiased prototype, and broadcasts the global model parameters, global prototype and unbiased prototype to all edge devices participating in the training.

[0009] (2) In each iteration, the following steps are performed:

[0010] (21) Each edge device updates its local model using the received global model parameters;

[0011] (22) Each edge device optimizes the local model using local data, global prototype and unbiased prototype, then uses local data and the optimized local model to obtain the local prototype, and uploads the local prototype and local model parameters to the central server.

[0012] (23) The central server updates the global model based on the local model parameters uploaded by each edge device, and updates the global prototype based on the local prototype uploaded by each edge device; the feature adapter network model is trained based on the local prototype of each edge device, and the feature adapter network model is used to process the local prototype to update the unbiased prototype.

[0013] (24) The central server broadcasts the updated global model parameters, global prototype and unbiased prototype to each edge device and enters the next round of iteration.

[0014] This invention also discloses a federated learning training system based on prototype representation, comprising a central server and several edge devices. The central server trains a global model, a feature adapter network model, a global prototype, and an unbiased prototype. The edge devices train local models and local prototypes. In each iteration, the central server updates the global model based on the local model parameters uploaded by each edge device, updates the global prototype based on the local prototypes uploaded by each edge device, trains the feature adapter network model based on the local prototypes of each edge device, uses the feature adapter network model to process the local prototypes to update the unbiased prototype, and broadcasts the updated global model parameters, global prototype, and unbiased prototype to each edge device.

[0015] The edge device optimizes the local model using local data, global prototype, and unbiased prototype, then uses the local data and the optimized local model to obtain a local prototype, and uploads the local prototype and local model parameters to the central server.

[0016] Compared with the prior art, the present invention has the following advantages:

[0017] This invention proposes a federated learning training method and system based on prototype representation. In each iteration, edge devices perform local training and compute local prototypes, then upload model parameters and prototypes to a central server for global model updates and computation of unbiased prototypes. The central server not only aggregates prototypes to obtain a global prototype, but also utilizes a feature adapter network to learn and extract representative unbiased prototypes from local prototypes, unaffected by the data distribution bias of specific edge devices. During local training, edge devices are guided by the global prototype and the unbiased prototype. By learning the unbiased prototype, the influence of feature shifts is mitigated, which not only improves the model's generalization ability but also achieves more accurate global model training and updates while protecting user privacy. Attached Figure Description

[0018] Figure 1 This is a schematic diagram of the federated learning training system based on prototype representation in this invention;

[0019] Figure 2 This is a schematic diagram of local training on edge devices;

[0020] Figure 3 This is a schematic diagram of the central server training adapter network model. Detailed Implementation

[0021] The present invention will be further described below with reference to embodiments.

[0022] This invention relates to a standard federated learning scenario, which generally includes N edge devices and one central server, where each edge device i has an independent private dataset D.i Each dataset exhibits a different distribution. This invention focuses on addressing the problem of non-independent, identically distributed (Non-iid) data caused by feature shift in federated learning scenarios, where the data heterogeneity caused by feature shift is defined as:

[0023]

[0024] Here, P(x|y) represents the feature conditional distribution of data x given label y, i represents edge device i, and k represents the k-th edge device. In other words, different edge devices have different feature distributions for data with the same label. This scenario is called a multi-domain scenario, where the private data of each edge device has feature skew and belongs to different domains. For example, images of the same type of tumor may correspond to both color and black-and-white domains.

[0025] In this invention, each edge device employs a neural network model with identical structure and parameterized as θ. This model consists of two parts: a feature extractor parameterized as W and a decision classifier parameterized as H, denoted as θ = (W, H). The feature extractor W transforms the input data x into an embedding h in a d-dimensional feature space, i.e., h = f(W; x) ∈ R. d The decision classifier H maps the feature embedding h to the logits vector z, i.e., z = f(H; h) ∈ R. I , where I represents the number of categories.

[0026] like Figure 1 As shown, the federated learning training system based on prototype representation of this invention mainly includes a central server and N edge devices. The central server is responsible for coordinating the entire federated learning process and trains a global model, a feature adapter network model, a global prototype, and an unbiased prototype. The edge devices train local models, performing local training based on their respective data and generating local prototypes. In each iteration, the central server and edge devices communicate to synchronize their trained model parameters and prototypes. The central server updates the global model based on the local model parameters uploaded by each edge device and updates the global prototype based on the local prototypes uploaded by each edge device. It trains a feature adapter network model based on the local prototypes of each edge device, uses the feature adapter network model to process the local prototypes to update the unbiased prototype, and broadcasts the updated global model parameters, global prototype, and unbiased prototype to each edge device. The edge devices optimize their local models using local data, the global prototype, and the unbiased prototype, then use the local data and the optimized local model to obtain a local prototype, and upload the local prototype and local model parameters to the central server.

[0027] Taking medical imaging applications as an example, the aforementioned edge devices can be deployed in various hospitals, holding local image data. The central server serves as a comprehensive processing center for the image data, used to form target training models and continuously update them.

[0028] The prototype-based federated learning training method of this invention includes the following steps:

[0029] (1) Initialization steps: Based on the neural network training model, the central server initializes the global model, feature adapter network model, global prototype and unbiased prototype, and broadcasts the global model parameters, global prototype and unbiased prototype to all edge devices participating in the training.

[0030] (2) In each iteration, the following steps are performed:

[0031] (21) Each edge device updates its local model using the received global model parameters;

[0032] (22) Each edge device optimizes the local model using local data, global prototype and unbiased prototype, then uses local data and the optimized local model to obtain the local prototype, and uploads the parameters of the local prototype and the optimized local model to the central server.

[0033] (23) The central server updates the global model based on the local model parameters uploaded by each edge device, updates the global prototype based on the local prototype uploaded by each edge device, trains the feature adapter network model based on the local prototype uploaded by each edge device, and uses the feature adapter network model to process the local prototype to update the unbiased prototype.

[0034] (24) The central server broadcasts the updated global model parameters, global prototype and unbiased prototype to each edge device, and enters the next round of iteration.

[0035] The above process will be further explained below with reference to the accompanying drawings and specific parameters.

[0036] (1) Initialization steps: The central server initializes the global model parameters. Initialize the feature adapter network model parameters F and the global prototype. and unbiased prototypes The feature adapter network model aims to learn how to extract unbiased prototypes from local prototypes through training.

[0037] (2) Taking the current round t as an example, the iteration steps are introduced as follows:

[0038] (21) All edge devices participating in the training receive the global model parameters θ sent by the central server. t Taking edge device i as an example, the global model parameters θ of the current round are used.t Update the local model for this round, i.e.

[0039] (22) Local model optimization and local prototype calculation process

[0040] Taking edge device i as an example, local model training is performed based on a neural network model, and the neural network model θ i Including feature extractor model W i and classifier model H i In other words, the local model is a neural network model.

[0041] Edge device i inputs private data x into neural network model θ i To obtain the classification probability of the data, the cross-entropy loss function is calculated based on the true label of the data, and the neural network model parameters θ are optimized. i This yields optimized local model parameters. The optimization process for these local model parameters will be described in detail below.

[0042] Feature extractor model W based on optimized local model i The features of the data samples are extracted to form a local prototype. The specific calculation method is as follows:

[0043]

[0044] Among them, C i j D represents the local prototype corresponding to the j-th type of data of edge device i. i j f(W) represents the data of the j-th class of edge device i. i ;x) indicates that the data x is input into the feature extractor model W. i The calculations are performed in the middle. Finally, each edge device uploads the acquired local prototype and optimized local model parameters to the server.

[0045] like Figure 1 As shown, edge device 1 holds local data D1, and after training and optimization, uploads the local model parameters (W1, H1) and local prototype C1 to the central server; and edge device N holds local data D N After training and optimization, the local model parameters (W) N H N ) and the local prototype C N Uploaded to the central server.

[0046] (23) Optimization and updating of the central server's global model, global prototype, and unbiased prototype.

[0047] (a) Optimization and updating of the global model

[0048] The central server collects local model parameters uploaded by each edge device and updates the global model θ for the next round by aggregating the local model parameters uploaded by each edge device. t+1 The specific calculations are as follows:

[0049]

[0050] Where |D| is the total size of all data, |D i | represents the data size of edge device i, and N represents the number of edge devices participating in the training.

[0051] (b) Optimization and updating of the global prototype

[0052] Meanwhile, the central server aggregates local prototypes to obtain the global prototype. The specific calculation method is as follows:

[0053]

[0054]

[0055] Among them, C i j Let I represent the local prototype corresponding to the j-th data type of edge device i, where I is the total number of data categories and N is the number of edge devices participating in the training.

[0056] (c) Optimization and updating of unbiased prototypes

[0057] The central processing unit (CPU) trains the feature adapter network model F using the classifier model H and the local prototypes uploaded by each edge device. During this process, the classifier parameters are frozen, and only the parameters of the feature adapter network model F are optimized. The optimization objectives are as follows:

[0058] L loss =l cls +λl cos

[0059] Among them, l cls For the true classification loss corresponding to the prototype, l cos The consistency loss between classifiers is represented by λ, which is a balancing coefficient used to adjust the relative importance between classification loss and consistency loss.

[0060] The central server uses an optimized feature adapter network model F to process the local prototype and generate an unbiased prototype. The unbiased prototype is calculated as follows:

[0061]

[0062]

[0063] Among them, C i j This represents the local prototype corresponding to the j-th type of data from edge device i. This indicates that the prototype is input into the feature adapter network model for computation. This represents the unbiased prototype corresponding to the j-th type of data.

[0064] (24) The central server broadcasts the updated global model parameters, global prototype, and unbiased prototype to each edge device, and enters the next round of iteration, repeating step (2). This process is repeated, and each iteration aims to optimize the global prototype and unbiased prototype simultaneously to better guide the local training of the edge devices.

[0065] The method of this invention optimizes the neural network model by minimizing the loss function. The following section will focus on how edge devices can optimize the local model.

[0066] like Figure 2 As shown, edge device i receives global model parameters and global prototype from the central server. and unbiased prototypes Update the local model (W) with global model parameters i H i Edge device i uses local data D i Perform local training, with data (x,y)∈D i Input the local model and obtain the feature embedding h = f(W) through the feature extractor. i x), and then obtain the corresponding logits vector z through a classifier. i =f(H i The feature embedding h is used to compute the loss about the prototype, and the logits vector z i Used to calculate the classification loss (cross-entropy loss) with the true label y.

[0067] The training process not only relies on local data, but is also guided by global and unbiased prototypes. This ensures that the model learned by the edge device not only adapts to the local data distribution, but also generalizes to a wider range of data distributions.

[0068] Optimization of edge devices during local training is performed using the following loss function:

[0069] L loss =l ce +l unbias +l global

[0070] Among them, l ce For cross-entropy loss, l unbias For the alignment loss of the unbiased prototype, lglobal This represents the contrast loss of the global prototype.

[0071] l ce The specific calculation method is as follows:

[0072]

[0073] Here, M is the total number of samples, I is the total number of categories, and y mj It is an indicator function; it is 1 if sample m belongs to class j, and 0 otherwise. m Let η(·) represent the model's logits output for sample m, and let η(·) represent the softmax function.

[0074] Alignment loss of unbiased prototype l unbias Comparison loss with global prototype l global The specific calculation method is as follows:

[0075]

[0076]

[0077]

[0078] Where h is the feature embedding output of the feature extractor, h j For the feature embedding of category j, v is the corresponding dimension index, d is the total number of dimensions, and j is the category of the current data. This represents the unbiased prototype corresponding to category j. This represents the global prototype corresponding to category j. represents the global prototype of other classes that do not belong to the current j category, sim is the cosine similarity calculation function with temperature parameter, and T is the representation strength of hyperparameter temperature control learning.

[0079] In short, edge devices combine global and unbiased prototypes to adjust and optimize the updated local model (which has been synchronized with the global model on the central server) to reduce the impact of feature shifts and improve the model's accuracy and generalization ability. The training optimization on edge devices not only considers the characteristics of local data but also integrates information provided by the global and unbiased prototypes, improving the performance and generalization ability of the local model.

[0080] In this invention, the central server optimizes the adapter network model by minimizing classification loss and consistency loss, thereby generating an unbiased prototype. The following section will focus on how the central server optimizes the parameters of the feature adapter network model F.

[0081] like Figure 3As shown, after the central server receives the local prototypes C uploaded by each edge device, it inputs these prototypes as input data into the feature adapter network model F. The goal of the adapter network is to transform these local prototypes into unbiased prototypes. By minimizing the loss function, it optimizes the parameters of the feature adapter network model so that the prototypes after passing through the feature adapter network model can be classified by all classifiers H1,…,H… i ,…H N It correctly classifies data and generates consistent output. This approach leverages prior knowledge from all classifiers to guide training.

[0082] This process is achieved by minimizing the loss function, which is calculated using the following formula:

[0083] L loss =l cls +λl cos

[0084] Among them, l cls For the true classification loss corresponding to the prototype, l cos The consistency loss between classifiers is represented by λ, which is a balancing coefficient used to adjust the relative importance between classification loss and consistency loss.

[0085] The classification loss ensures that the feature adapter network can extract unbiased prototypes while retaining key semantic information about the category in the prototypes, so that the unbiased prototypes can represent the common features of their respective categories, thereby improving the accuracy of the model's classification based on these prototypes.

[0086] Classification loss cls The specific calculation method is as follows:

[0087]

[0088] Where N is the total number of edge devices, L n The cross-entropy loss of the nth classifier is calculated as follows:

[0089]

[0090] M represents the total number of uploaded prototypes, I represents the total number of data categories, and y represents the total number of data categories. mj It is an indicator function; its value is 1 if the actual category of the prototype m is j, and 0 otherwise. The logits output of the prototype m corresponding to the i-th classifier is used to calculate the classification loss (cross-entropy loss) with the prototype's true label y, and η(·) represents the softmax function.

[0091] Consistency loss reduces the differences between local prototypes uploaded from different edge devices, ensuring that the unbiased prototypes generated by the feature adapter network remain consistent across different edge devices. By minimizing consistency loss, the consistency of the unbiased prototypes is ensured globally, and universally applicable feature representations are extracted, mitigating model bias caused by heterogeneous data distribution.

[0092] Consistency loss cos The specific calculation method is as follows:

[0093]

[0094]

[0095] Where i is the edge device index, N is the number of edge devices participating in training, and z i The classifier H represents the edge device i. i The output before softmax classification is the logistic regression. For example, in the MNIST dataset, j has 10 classes. The output before softmax classification is the output of the last linear layer (output dimension is 10). Each value represents the predicted value for that class. After ordinary softmax classification, it becomes the normalized probability of each class. Let D represent the average output of all classifiers before softmax classification, T be the hyperparameter temperature control consistency strength, and η(·) represent the softmax function. KL (·) Use KL divergence to calculate the similarity between the current classifier output and the global average output.

[0096] In short, the key to this step is to leverage the feature adapter network to learn and extract representative global feature representations from local prototypes that are unaffected by the biases in the data distribution of specific edge devices. For example, in medical image classification scenarios, for images of the same type of tumor, it is necessary to extract the essential features that represent this type of tumor from images with color and black and white or other differences. A crucial step in training the feature adapter network is to fix the parameters of the classifiers in the local models uploaded from each edge device. The purpose of this is to allow the adapter network to focus on learning and extracting unbiased prototypes from the uploaded prototypes, rather than attempting to adjust the classifier to fit a specific data distribution. By fixing the classifier parameters, it is ensured that the training process of the adapter network depends solely on the quality and representation of the prototype itself, rather than the decision boundaries of the classifier, thus helping to extract more representative and generalizable unbiased prototypes.

[0097] The preferred embodiments of the present invention have been described in detail above. However, the present invention is not limited to the specific details of the above embodiments. Within the scope of the technical concept of the present invention, various equivalent transformations can be made to the technical solutions of the present invention, and these equivalent transformations all fall within the protection scope of the present invention.

Claims

1. A federated learning training method based on prototype representation for medical image analysis, characterized in that, The method includes the following steps: (1) Initialization steps: Based on the neural network training model, the central server initializes the global model, feature adapter network model, global prototype and unbiased prototype, and broadcasts the global model parameters, global prototype and unbiased prototype to all edge devices participating in the training. The edge devices are deployed in various hospitals and hold local image data. The local image data has feature skew. (2) In each iteration, the following steps are performed: (21) Each edge device updates its local model using the received global model parameters; (22) Each edge device optimizes its local model using local image data, a global prototype, and an unbiased prototype. The loss function for optimizing the local model is: , For cross-entropy loss, The alignment loss is for the unbiased prototype. The contrast loss is for the global prototype; Alignment loss of the unbiased prototype and The calculation formula is: Where h is the feature embedding output of the feature extractor. For the feature embedding of category j, v is the corresponding dimension index, d is the total number of dimensions, and j is the category of the current data. This represents the unbiased prototype corresponding to category j. This represents the global prototype corresponding to category j. This represents the global prototype of other classes that do not belong to the current class j. sim is the cosine similarity calculation function with temperature parameter, and T is the hyperparameter temperature. Then, a local prototype is obtained using local image data and an optimized local model, and the local prototype and local model parameters are uploaded to the central server. (23) The central server updates the global model based on the local model parameters uploaded by each edge device, and updates the global prototype based on the local prototype uploaded by each edge device; trains the feature adapter network model based on the local prototype of each edge device, and uses the feature adapter network model to process the local prototype to update the unbiased prototype; when training the feature adapter network model, the classifier parameters uploaded by all edge devices are frozen and not updated. Training the feature adapter network model includes: the central server inputting the local prototypes of each edge device as input data into the feature adapter network model to obtain the true classification loss. Consistency loss between classifier and classifier By minimizing the loss function Optimize the feature adapter network model, where λ is the balance coefficient; Consistency loss among classifiers The calculation formula is: ） Where N is the number of edge devices participating in the training. This represents the output of the classifier corresponding to edge device i before classification. This represents the average output of all classifiers before classification, where T is the hyperparameter temperature. (·) represents the softmax function. (·) represents the KL divergence; (24) The central server broadcasts the updated global model parameters, global prototype and unbiased prototype to each edge device and enters the next iteration process.

2. The federated learning training method based on prototype representation according to claim 1, characterized in that, The central server aggregates the local models uploaded by each edge device to obtain the global model, and the calculation formula is as follows: in, For all data sizes, Let N be the data size of edge device i, and N be the number of edge devices participating in the training. This represents the local model parameters of edge device i in round t. This represents the global model parameters for round t+1.

3. The federated learning training method based on prototype representation according to claim 1, characterized in that, The neural network model includes a feature extractor model and a classifier model. The calculation formula for the edge device to obtain the local prototype is as follows: in, Let j represent the j-th local prototype of edge device i. This represents the j-th type of data at edge device i. This means inputting data x into the feature extractor model. Calculated in the middle.

4. The federated learning training method based on prototype representation according to claim 3, characterized in that, The central server aggregates the local prototypes uploaded by each edge device to obtain the global prototype, and the calculation formula for the global prototype is as follows: in, This is the global prototype for the j-th type of data. This is the global prototype, where N is the number of edge devices participating in the training. The total number of data categories.

5. A federated learning training system based on prototype representation for medical image analysis, comprising a central server and several edge devices deployed in various hospitals, wherein the central server is trained with a global model, a feature adapter network model, a global prototype, and an unbiased prototype, and the edge devices are trained with local models and local prototypes, wherein the edge devices hold local image data, and the local image data exhibits feature skew; characterized in that, This training system executes the federated learning training method based on prototype representation as described in any one of claims 1-4; in each iteration, the central server updates the global model based on the local model parameters uploaded by each edge device, updates the global prototype based on the local prototype uploaded by each edge device, trains the feature adapter network model based on the local prototype of each edge device, uses the feature adapter network model to process the local prototype to update the unbiased prototype, freezes all classifier parameters uploaded by edge devices and does not update them during the training of the feature adapter network model, and broadcasts the updated global model parameters, global prototype and unbiased prototype to each edge device.

6. The federated learning training system based on prototype representation according to claim 5, characterized in that, The edge device optimizes the local model using local image data, a global prototype, and an unbiased prototype, then uses the local data and the optimized local model to obtain a local prototype, and uploads the local prototype and local model parameters to the central server.

Images