Federal learning domain generalization method for medical field

Abstract

The invention discloses a medical field-oriented federal learning domain generalization method, which comprises the following steps of: firstly, receiving a global model parameter from a central server by a client to initialize a local model, adding random disturbance to a sample x in a local data set to obtain a disturbance version # imgabs0 # of the x, and then respectively inputting the x and # imgabs1 # into a feature extractor, and obtaining the nth layer output zn and # imgabs2 of the feature extractor, and constructing a loss function for training. Then the local client calculates a feature prototype for each category, and uploads the prototype and the trained local model parameters to the central server; and finally, the central server performs local model parameter aggregation, randomly samples a batch of new feature samples for each category by using a feature prototype, sends the new feature samples to the classifier after parameter aggregation, obtains a prediction result, and calculates a cross entropy loss back propagation optimization classifier. According to the method, the feature extractor has enough robustness for low-level feature offset, and the generalization ability is improved.

Description

Technical Field

[0001] The present invention belongs to the medical field, and in particular relates to a federated learning domain generalization method for the medical field. Background Art

[0002] The development of the new generation of artificial intelligence has put forward new requirements and challenges for data privacy and the use of public data. To promote the opening and use of public data resources, it is necessary to increase the scope of supply and use of public data that does not carry personal information and does not affect public safety according to the purpose. At the same time, public data involving personal information should be desensitized and anonymized to avoid infringement of personal information rights. As an important part of public data elements, medical data plays a key role in the development of the new generation of artificial intelligence. It is related to important aspects such as the integration of medical resources, the formulation of treatment plans, and the standardization of medical data.

[0003] As an emerging machine learning technology, federated learning allows multiple participants to collaboratively train a shared global model in a way that the data is "available but invisible" and "controllable and measurable", such as electronic medical record similarity search and patient representation learning. Therefore, federated learning can promote collaboration and development among medical institutions while protecting patient privacy. However, in actual applications, there are differences in data quality and characteristics between different medical institutions, and there is a domain drift phenomenon, which will lead to poor generalization ability of the medical model trained under the federated learning framework in hospitals that did not participate in the training.

[0004] Therefore, the problem of generalizing the federated domain for medical scenarios needs to be solved urgently. Summary of the invention

[0005] In response to the above problems, the present invention proposes a federated learning domain generalization technology for the medical field, which combines prototype-based data exploration and local consistency learning, aiming to solve the shortcomings of traditional federated learning in generalization performance in unknown domains. This method can not only effectively improve the adaptability of the model in different medical data environments, but also provide strong technical support for the standardization and sharing of public data. First, the present invention proposes a prototype-based server-side sample exploration strategy, which can directly explore more potential samples in the feature space without exposing privacy, and use these samples to further optimize the global model to make it have stronger generalization performance. In addition, the present invention introduces a consistency learning method locally, so that the feature extractor has sufficient robustness to low-level feature offsets.

[0006] The proposed overall method framework includes the local learning stage (Step 1 to Step 10) and the server-side optimization stage (Step 11 to Step 15), as follows:

[0007] Step 1: Each client receives global model parameters from the central server. The global model includes a feature extractor F and a classifier C. The local model is initialized with the global model parameters and then local training is started using the local data set. The feature extractor F maps the input medical image x to a d-dimensional feature vector, and the classifier C maps this feature vector to the predicted label distribution. In the present invention, the feature extractor is implemented using a convolutional neural network (CNN), and the classifier is a two-layer fully connected network that maps the extracted features to the classification results.

[0008] Step 2: Obtain the medical image samples x and the corresponding true labels y in the local dataset, and apply standardization to x.

[0009] Step 3: Add random perturbations to the standardized sample x to obtain the perturbed version of x

[0010] Step 4: Get the shallow feature representation z of x n =F n (x) and predicted results Among them, F n (x) represents the output of the nth convolutional layer of the feature extractor F.

[0011] Step 5: Get the local model pair Shallow feature representation

[0012] Step 6: Calculate the true label y and the predicted label The cross entropy loss between

[0013] Step 7: Calculate the original feature z n and disturbance characteristics Channel-by-channel differences The mean is taken as the consistency loss L c .

[0014] Step 8: Calculate the total loss Where α is a hyperparameter that balances the weight between the cross entropy loss and the consistency loss. The local model is then trained using the total loss L through backpropagation.

[0015] Step 9: Repeat Step 2 to Step 7 until the number of training rounds reaches the set value.

[0016] Step 10: The local client computes feature prototypes for each category and then uploads these prototypes along with the trained local model parameters to the central server.

[0017] Step 11: Aggregate model parameters on the central server.

[0018] Step 12: The central server uses the feature prototypes from each client to randomly sample a batch of new feature samples for each category in the feature space. A new feature sample is obtained by a random convex combination of the prototypes of each client.

[0019] Step 13: Send the new sample to the classifier to get the prediction result where z new is a sample of features explored.

[0020] Step 13: Calculate the true label y and the predicted label The cross entropy loss between

[0021] Step 15: Use cross entropy loss L ce Back propagation optimizes the classifier. Return to Step 1 and repeat the above steps until the number of communication rounds reaches the set value.

[0022] The present invention has the following beneficial effects:

[0023] Stronger robustness: It is observed that despite different qualities, when medical images share the same label, they retain consistent semantic information. It is hypothesized that the degradation in classification performance is due to the feature extractor being insufficiently robust to subtle variations in low-level features present in medical images. This approach incorporates consistency learning during the training process of each domain, which forces the extractor to produce more reliable representations and makes the feature extractor sufficiently robust to low-level feature shifts.

[0024] Stronger generalization: is a convex combination of all source domain prototypes. Each prototype vector represents the central position in the feature space of a specific class in a specific source domain. By applying convex combination to these prototype vectors, a weighted average fused prototype is essentially created. This new vector still lies within the convex hull of these prototypes, so it can be regarded as a cross-domain class feature vector, which retains the main characteristics of the class while considering the contribution of each source domain. And since each prototype represents the average feature vector of a specific class, combining these prototypes through convex combination still results in a weighted average of the class features. This process does not blur class-specific information. Convex combination allows the new vector to integrate knowledge from different source domains, thereby improving generalization ability. This is particularly useful for domain generalization because it helps the model capture common features across domains, reducing the impact of domain-specific differences, or it can be considered that these new samples may come from potential target domains.

[0025] Mitigating class imbalance: Class imbalance refers to the situation where samples in the training dataset are unevenly distributed between different classes. Models trained on imbalanced data are usually biased towards the majority class, resulting in poor generalization ability when encountering minority class samples during testing or actual deployment. In the present invention, by generating new feature samples for each class, the server can ensure that even classes with fewer original samples can obtain sufficient feature representation during training. This balanced sampling strategy allows the classifier to learn a more representative decision boundary for the minority class, thereby mitigating the bias towards the majority class. BRIEF DESCRIPTION OF THE DRAWINGS

[0026] Figure 1 is a flow chart of the method of the present invention;

[0027] Figure 2 It is a schematic diagram of the framework of the present invention. DETAILED DESCRIPTION

[0028] The present invention will be further described below in conjunction with the accompanying drawings and specific implementation steps.

[0029] like Figure 1 and Figure 2 A method for generalizing a federated learning domain in the medical field is shown, comprising the following steps:

[0030] Step 1: Each client receives the global model parameters from the central server, uses them to initialize the local model and then starts local training using the local dataset.

[0031] Step 2: Get local training data (x, y). First, apply min-max normalization to the medical image sample x in the local dataset to ensure that the pixel values ​​are scaled between 0 and 1, which helps improve training stability.

[0032] Step 3: Add random perturbations to the sample x in Step 2. Here, Gaussian noise is used. As a perturbation, we then get the perturbation version of x

[0033] Step 4: Get the model's shallow feature representation z for x (1) =F (1) (x) and predicted results Among them, F (1) (x) represents the output of the first convolutional layer of the feature extractor F.

[0034] Step 5: Get model pair Shallow feature representation

[0035] Step 6: Calculate the true label y and the predicted label The cross entropy loss between them is calculated as follows:

[0036]

[0037] Step 7: Calculate the original feature z (1) and disturbance characteristics Channel-by-channel differences The mean is taken as the consistency loss L c The present invention uses structural similarity (SSIM) as a method to measure the channel-by-channel discrepancy of feature maps, and its calculation formula is as follows:

[0038]

[0039] Among them, a and b are convolution feature maps, μ a and μ b is the mean of a and b, and is the variance of a and b, σ ab is the covariance of a and b, and C1 and C2 are constants to maintain numerical stability.

[0040] Step 8: Calculate the total loss Among them, α is a hyperparameter that balances the weight between cross entropy loss and consistency loss. After experimental selection, this method sets it to 0.8. The total loss L is used to back-propagate and optimize the model.

[0041] Step 9: Repeat Step 2 to Step 8 until the number of training rounds reaches the set value.

[0042] Step 10: The local client calculates the feature prototype for each category. The prototype calculation formula is as follows:

[0043]

[0044] Where k is the current client number, is the feature prototype of category θ, is the local dataset S k contains all samples belonging to category θ, and F is a feature extractor.

[0045] Step 11: Each local client uploads all the calculated prototypes along with the model parameters to the central server.

[0046] Step 12: Model parameters are aggregated on the server side. The aggregation method is as follows:

[0047]

[0048] Among them, H represents the total number of clients, F g and C g is the new global model of aggregation, F i and C i Represent the parameters of the i-th feature extractor and classifier respectively.

[0049] Step 13: The server uses the feature prototypes from each client to randomly sample a batch of new feature samples for each category in the feature space. A new feature sample is obtained by a random convex combination of the prototypes of each client:

[0050]

[0051] where w i ≥0 and w i are randomly generated weights. is a new feature sample belonging to category θ. In this step, the server samples a new feature sample with a batch size of B for each category, where B is a predefined hyperparameter. In this method, it is set to 32.

[0052] Step 14: Send the new sample to the classifier to get the prediction result where z new is the feature sample explored. Calculate the true label y and the predicted label The cross entropy loss between The cross entropy loss is used to back-propagate and optimize the classifier.

[0053] Step 15: If the number of global communication rounds does not reach the preset value, that is, the training is not completed, the new global model parameters are distributed to each local client and return to Step 1. Otherwise, the above process ends and the global model trained under the federated learning framework is obtained.

[0054] Experimental verification

[0055] 1. Dataset:

[0056] Different medical datasets are collected to form two classification tasks. During the training process, the present invention adopts a leave-one-domain strategy: one dataset is selected as the target domain and the remaining datasets are used as the source domain.

[0057] Diabetic retina image classification task: The first task is to classify diabetic retina into five levels: no DR, mild DR, moderate DR, severe DR, and proliferative DR. Three datasets were collected from different sources to construct the diabetic retina classification dataset, namely APTOS-2019, IDRiD, and Messidor-2, including 3662, 516, and 1744 samples respectively.

[0058] Glaucoma image classification task: The second task is a binary classification task of glaucoma, where images are labeled as glaucoma or non-glaucoma. Four glaucoma classification datasets were collected, namely G1020, BEH, ORIGA-light, and PAPILA, containing 1020, 634, 650, and 488 samples, respectively. 68 samples in the PAPILA dataset were marked as "suspected" and removed to maintain consistency with other datasets.

[0059] 2. Benchmark method for comparison:

[0060] FedAvg: A classic federated learning algorithm.

[0061] ELCFS: A pioneering approach to exchanging image magnitude spectra across domains to locally synthesize new samples.

[0062] FedSR: A method that aims to learn simple representations by adding an L2-norm regularizer and a conditional mutual information regularizer to the representation.

[0063] FedAD: A method that proposes to locally train a feature extractor and a discriminator in an adversarial manner to help the extractor gradually learn domain-invariant features.

[0064] FedGM: A method that proposes to apply gradient masking to emphasize consistent updates between clients, thereby reducing the impact of gradient conflicts.

[0065] 3. Implementation details:

[0066] Backbone network: ResNet-18 pre-trained on ImageNet is used as the feature extractor, and the output dimension of its last layer is modified to 512. The classifier is a two-layer fully connected network with a first layer of 512×512, followed by a final output layer that maps the intermediate output to a prediction result of Θ dimension.

[0067] Hyperparameters: The number of communication rounds is set to 50. The number of local training rounds is set to 5, and the number of global rounds for optimizing the classifier is set to 3. The consistency weight α is set to 0.8. The batch size for local updates is set to 128, while the sampling size B of our method is 32. The learning rate of our method is set to 0.0005, while the learning rate and all other hyperparameters specific to each baseline are configured according to their original papers.

[0068] 4. Compare the results:

[0069] Diabetic retina image classification task: Table 1 shows the comparison results of the diabetic retina classification task. Compared with other methods, the present invention achieves the best performance on this task.

[0070] Table 1 Accuracy of diabetic retinal image classification task (%)

[0071]

[0072] Glaucoma image classification task: Table 2 lists the comparison results of the glaucoma classification task. The results show that all other methods have different degrees of improvement in accuracy compared with FedAvg. It can be observed that except when the target domain is G1020, the present invention is 0.5% lower than FedGMA, and the present invention achieves the highest accuracy in all other cases.

[0073] Table 2 Accuracy of glaucoma image classification task (%)

[0074]

[0075] Judging from the results, the method of the present invention significantly improves the generalization ability of the federated learning model in the face of unknown domains in medical scenarios, provides favorable assistance for the deployment of the model in actual medical scenarios, and has a very broad application prospect.

Claims

1. A method for generalizing federated learning domains in the medical field, characterized in that: The following steps are involved: S1: Each client receives the global model parameters including feature extractor F and classifier C from the central server and initializes the local model with the global model parameters; S2: Get the medical image sample x and the corresponding true label y in the local dataset, apply standardization to x, and add random perturbations to the standardized sample x to obtain the perturbed version of x S3: Medical image samples x and Input the feature extractor F respectively, and get the n-th layer output z of the feature extractor F n and And construct a loss function to train the local model; S4: The local client calculates the feature prototype for each category and uploads the prototype and the trained local model parameters to the central server; S5: The central server aggregates local model parameters and uses the feature prototypes of each client to randomly sample a batch of new feature samples for each category in the feature space; S6: Send the new sample to the classifier after parameter aggregation to obtain the prediction result, calculate the cross entropy loss between the real label and the real label, back propagate and optimize the classifier, repeat the above operations to obtain the global model trained under the federated learning framework.

2. According to claim 1, a method for generalizing a federated learning domain in the medical field is characterized in that: In step S1, the feature extractor F maps the input medical image x to a d-dimensional feature vector, and the classifier C maps this feature vector to the predicted label distribution. The feature extractor is implemented using a convolutional neural network CNN, and the classifier is a two-layer fully connected network that maps the extracted features to the classification results.

3. The method for generalizing a federated learning domain in the medical field according to claim 2, characterized in that: The specific implementation process of step S3 is as follows: S3.1: Get the shallow feature representation z of x n =F n (x) and predicted results Among them, F n (x) represents the output of the nth convolutional layer of the feature extractor F; S3.2: Get local model pair Shallow feature representation S3.3: Calculate the true label y and the predicted label The cross entropy loss between S3.4: Calculate z n and The channel-by-channel difference is taken as the consistency loss L c ; S3.5: Calculate total loss Where α is a hyperparameter that balances the weight between the cross entropy loss and the consistency loss. Finally, the total loss L is used to back-propagate and train the local model.

4. The method for generalizing a federated learning domain in the medical field according to claim 3, characterized in that: The channel-by-channel difference is calculated as follows: Structural similarity SSIM is used as a method to measure the channel-by-channel differences in feature maps. The calculation formula is as follows: Among them, a and b are convolutional feature maps, μ a and μ b is the mean of a and b, and is the variance of a and b, σ ab is the covariance of a and b, and C1 and C2 are constants to maintain numerical stability.

5. The method for generalizing a federated learning domain in the medical field according to claim 4, characterized in that: The calculation formula of the feature prototype described in S4 is as follows: Where k is the current client number, is the feature prototype of category θ, is the current local dataset S k contains all samples belonging to category θ, and F is a feature extractor.

6. The method for generalizing a federated learning domain in the medical field according to claim 5, characterized in that: The local model parameter aggregation in step S5 is specifically formulated as follows: Among them, H represents the total number of clients, F g and C g is the new global model of aggregation, F i and C i Represent the parameters of the i-th feature extractor and classifier respectively.

7. The method for generalizing a federated learning domain in the medical field according to claim 6, characterized in that: In step S5, a new feature sample is obtained by a random convex combination of each client prototype: where w i ≥0 and w i is a randomly generated weight; is a new feature sample belonging to category θ.

Images