A medical field-oriented federated learning domain generalization method

Abstract

The application discloses a medical field-oriented federated learning domain generalization method, which firstly receives global model parameter initialization local model from a central server, adds random disturbance to samples x in a local data set to obtain a disturbed version of x, and secondly inputs x and x into a feature extractor to obtain the n-layer output z of the feature extractor n and constructs a loss function to train. Then, the local client calculates feature prototypes for each category, uploads the prototypes and the trained local model parameters to the central server. Finally, the central server aggregates the local model parameters, randomly samples a batch of new feature samples for each category using the feature prototypes, inputs the new feature samples into the parameter aggregated classifier to obtain a prediction result, and calculates cross-entropy loss to optimize the classifier. The application makes the feature extractor have sufficient robustness to low-level feature deviation and improves the generalization ability.

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 next generation of artificial intelligence (AI) presents new requirements and challenges for data privacy and the use of public data. Promoting the openness and utilization of public data resources requires expanding the scope of access to public data that does not carry personal information and does not impact public safety, based on its intended purpose. Furthermore, public data involving personal information must be desensitized and anonymized to avoid infringing on personal information rights. Medical data, as a crucial component of public data, plays a key role in the development of the next generation of AI. It is crucial for integrating medical resources, developing treatment plans, and promoting the standardization of medical data.

[0003] Federated learning, an emerging machine learning technology, allows multiple participants to collaboratively train a shared global model in a way that makes data "available but invisible" and "controllable and measurable." This can be used for applications 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 practice, data quality and characteristics vary between institutions, leading to domain drift. This can result in poor generalization of medical models trained using the federated learning framework to hospitals that did not participate in the training.

[0004] Therefore, the problem of generalizing federated domains 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 to address 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 phase (Step 1 to Step 10) and the server-side optimization phase (Step 11 to Step 15), as follows:

[0007] Step 1: Each client receives global model parameters from a 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 trained locally using a local dataset. Feature extractor F maps the input medical image x to a d-dimensional feature vector, and classifier C maps this feature vector to a predicted label distribution. In this paper, the feature extractor is implemented using a convolutional neural network (CNN), while the classifier is a two-layer fully connected network that maps the extracted features to the classification results.

[0008] Step 2: Obtain medical image samples x and the corresponding true labels y in the local dataset and apply normalization 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 of

[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 cross entropy loss and consistency loss. The local model is then trained by backpropagation using the total loss L.

[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 into the classifier to get the prediction result where z new is the feature sample 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: Observations show that despite varying quality, medical images sharing the same label retain consistent semantic information. We hypothesize that the degradation in classification performance is due to the feature extractor's lack of robustness to subtle variations in low-level features present in medical images. Our approach incorporates consistency learning into the training process for each domain, forcing the extractor to produce more reliable representations and making the feature extractor robust to shifts in low-level features.

[0024] Stronger generalization: is a convex combination of all source domain prototypes. Each prototype vector represents the center 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 and can therefore be considered a cross-domain class feature vector, preserving the key characteristics of the class while taking into account the contribution of each source domain. And because 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 characteristics. This process does not obscure 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, as it helps the model capture common features across domains and reduces the impact of domain-specific differences, or it can be considered that these new samples may come from potential target domains.

[0025] Class imbalance mitigation: Class imbalance refers to the uneven distribution of samples across different classes in the training dataset. Models trained on imbalanced data often bias towards the majority class, leading to poor generalization when encountering minority class samples during testing or real-world deployment. In the present invention, by generating new feature samples for each class, the server can ensure that even classes with fewer original samples receive adequate feature representation during training. This balanced sampling strategy allows the classifier to learn more representative decision boundaries for minority classes, mitigating bias towards the majority class. BRIEF DESCRIPTION OF DRAWINGS

[0026] Figure 1 Flowchart of the method of the present invention;

[0027] Figure 2 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] As shown in Figure 1 and Figure 2 a federated learning domain generalization method for the medical field, comprising the following steps:

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

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

[0032] Step 3: Add random perturbations to the samples x in Step 2, here we use Gaussian noise as perturbation, then get the perturbed version of x

[0033] Step 4: Obtain the shallow feature representation z of the model for x (1) = F (1) (x) and the predicted result where F (1) (x) represents the output of the first convolutional layer of the feature extractor F.

[0034] Step 5: Obtain the shallow feature representation of the model for

[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 the perturbed feature The channel-wise discrepancy between them is calculated as follows: Take the mean of it as the consistency loss L c The structural similarity (SSIM) is used as the method to measure the channel-wise discrepancy of the feature map Discrepancy, and its calculation formula is as follows:

[0038]

[0039] Where a and b are the convolution feature maps, μ a and μ b are the mean of a and b, and are 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 Where α is a hyperparameter that balances the weight between cross-entropy loss and consistency loss, and is set to 0.8 through experiments. The total loss L is used for backpropagation to 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: Each local client calculates the feature prototype for each class, and the prototype calculation formula is as follows:

[0043]

[0044] Where k is the current client number, is the feature prototype of class θ, is a subset of the local data set S k , containing all samples belonging to class θ, and F is the 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: The 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 denote 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 from 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 batch of new feature samples of size 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 into the classifier to get the prediction result where z new For the explored feature samples. Calculate the true label y and the predicted label The cross entropy loss between The cross entropy loss is used to backpropagate and optimize the classifier.

[0053] Step 15: If the number of global communication rounds has not reached the predetermined value, that is, the training has not ended, the new global model parameters are distributed to each local client and the process returns 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 were collected to form two classification tasks. During the training process, the present invention adopted a leave-one-out strategy: one dataset was selected as the target domain and the remaining datasets were used as the source domains.

[0057] Diabetic Retinopathy Image Classification Task: The first task is to classify diabetic retinopathy 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 retinopathy classification dataset, namely APTOS-2019, IDRiD, and Messidor-2, which include 3662, 516, and 1744 samples respectively.

[0058] Glaucoma Image Classification Task: The second task is a binary classification task for glaucoma, labeling images as either glaucoma or non-glaucoma. Four glaucoma classification datasets were collected: G1020, BEH, ORIGA-light, and PAPILA, containing 1020, 634, 650, and 488 samples, respectively. Sixty-eight samples in the PAPILA dataset were labeled "suspect" and were removed to maintain consistency with the other datasets.

[0059] 2. Baseline 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 applying 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 retinal image classification task: Table 1 shows the comparison results of the diabetic retinal classification task. Compared with other methods, the proposed method 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 comparative results for the glaucoma classification task. The results show that all other methods achieve varying degrees of accuracy improvement compared to FedAvg. It can be observed that, except for the case where the target domain is G1020, where our method is 0.5% lower than FedGMA, our method 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 by: The following steps are involved: S1: 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 dataset. 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, and the classifier is a two-layer fully connected network that maps the extracted features to the classification results. S2: Obtain medical image samples x and corresponding true labels y in the local dataset and apply normalization to x; S3: Add random perturbations to the standardized sample x to obtain a perturbed version of x S4: 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; S5: Get local model pair Shallow feature representation of S6: Calculate the true label y and the predicted label The cross entropy loss between S7: Structural similarity SSIM is used as a method to measure the channel-by-channel difference of feature maps and calculate the original feature z n and disturbance characteristics Channel-by-channel differences The mean is taken as the consistency loss L c ; S8: Calculate total loss Where α is a hyperparameter that balances the weight between cross entropy loss and consistency loss, and then the local model is trained by backpropagation using the total loss L; S9: Repeat S2 to S7 until the number of training rounds reaches the set value; S10: The local client calculates feature prototypes for each category and then uploads these prototypes along with the trained local model parameters to the central server; S11: Aggregate model parameters on the central server; S12: 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, as shown in the following formula: where w i ≥0 and w i is a randomly generated weight, H represents the total number of clients; is a new feature sample belonging to category θ; S13: Send the new sample into the classifier to get the prediction result where z new is the feature sample explored; S14: Calculate the true label y and the predicted label The cross entropy loss between S15: Use cross-entropy loss back propagation to optimize the classifier; return to S1 and repeat the above steps until the number of communication rounds reaches the set value.

2. A method for generalizing federated learning domains in the medical field according to claim 1, characterized in that: The channel-by-channel difference is calculated as follows: 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.

3. A method for generalizing federated learning domains in the medical field according to claim 2, characterized in that: The calculation formula of the feature prototype is as follows: Where k is the current client number, is the feature prototype of category θ, is the current local dataset S k The subset of contains all samples belonging to category θ, and F is the feature extractor.

4. A method for generalizing federated learning domains in the medical field according to claim 3, characterized in that: The model parameter aggregation is as follows: Among them, F g and C g is the new global model of aggregation, F i and C i denote the parameters of the i-th feature extractor and classifier respectively.

Images