Medical federal learning method for decentration of medical scene

Abstract

The invention discloses a medical federal learning method for decentralization of a medical scene. The method comprises the following steps: 1) deploying a local model, an expert model and an expert mixing module on each client of a federal learning system; the local model comprises a Body module and a Head module, and the expert model comprises a feature space conversion module and lightweight Head modules corresponding to the Head modules on other clients; 2) the Body module extracts features of the data and sends the features to the feature space conversion module and the Head module; (3) the Head module executes downstream tasks according to the features to obtain a prediction weight and sends the prediction weight to the expert mixing module; the feature space conversion module converts the features into features which can be processed by the lightweight Head modules and executes downstream tasks to obtain a prediction weight; and 4) the expert mixing module outputs a final prediction result according to each prediction weight and sends the final prediction result to the Head module of each client for knowledge sharing.

Description

TECHNICAL FIELD

[0001] The present application relates to the field of medical federated learning, in particular to a decentralized medical federated learning method for medical scenarios. BACKGROUND

[0002] Federated learning has a wide range of applications in the medical field. A well-designed federated system can share high-level knowledge among different clients while protecting data privacy. This enables each client's network to gain additional support, resulting in better performance and generalization capabilities. In the medical scenario, patient data is difficult to collect and has strong privacy protection requirements. The federated learning system can effectively solve the limitations of data collection for medical institutions, improve the performance and generalization capabilities of the model while ensuring privacy.

[0003] Existing federated learning systems, such as FedAvg, FedRep, MH-pFLID, and FedMD, are designed in a centralized manner. In each training round, each client needs to upload specific client knowledge (e.g., model parameters) to the central server for aggregation, and then distribute it back to each client. Regarding the aggregation method, these systems require a uniform model structure (such as FedAvg, FedRep), a centralized messenger model (such as MH-pFLID), or a uniform public dataset (such as FedMD). Although this centralized design can achieve good results, it may also cause performance bottlenecks. The centralized federated learning framework extracts knowledge from each client's local data and sends it to the centralized server for aggregation, and then distributes the aggregated knowledge back to each client. However, the aggregation process usually combines the information of all client models into a single aggregated model, usually using weighted sum or other sample merging methods. Considering the differences in domain and data distribution between clients, using the same aggregation process for all clients may result in knowledge loss, or even damage before the aggregated knowledge returns to each client. The centralized aggregation method, especially the weighted sum method, may not be able to well preserve the knowledge of each client, and may even harm the performance of the federated learning framework.

[0004] In addition, the centralized federated framework also heavily relies on the stability of the central server and its connection. If the central server fails or the connection is unstable, it will significantly affect the training stability of each client. SUMMARY

[0005] In view of the problems in the prior art, the present application aims to provide a decentralized federated learning method for medical scenarios, aiming to deal with the loss of knowledge in centralized aggregation and reduce the dependence on the central server. The main purpose is to minimize the damage of knowledge in the aggregation process, eliminating the centralized model aggregation operation. At the same time, in the knowledge exchange process, the knowledge that each client originally sends to the server is now directly transmitted to other clients. In this way, each client can receive the complete knowledge sent by other clients without loss.

[0006] The technical solution of the present application is:

[0007] A medical federated learning method for medical scenarios, comprising the following steps:

[0008] 1) Deploy a local model, an expert model and an expert mixing module on each selected client of the federated learning system; the local model comprises a Body module and a Head module, a corresponding lightweight Head module is created for the Head module in the local model of each selected client, and the expert model on the client A comprises a feature space conversion module and a plurality of lightweight Head modules corresponding to the Head modules on other clients;

[0009] 2) The Body module on the client A converts input data into high-level semantic features and sends them to the feature space conversion module and the Head module, respectively;

[0010] 3) The Head module on the client A performs a set downstream task according to the high-level semantic features to obtain a prediction weight and sends it to the expert mixing module; the feature space conversion module on the client A converts the high-level semantic features into features that can be processed by each lightweight Head module, respectively; each lightweight Head module performs a set downstream task according to the features to obtain a prediction weight and sends it to the expert mixing module;

[0011] 4) The expert mixing module on the client A outputs a final prediction result according to each prediction weight and sends it to the Head module of each selected client for knowledge sharing.

[0012] Further, the expert mixing module splices each prediction weight, sends the spliced result and the high-level semantic features to the cross-attention layer to output a final prediction result and sends it to the Head module of each selected client.

[0013] Further, the expert mixing module splices each prediction weight as a key value, and sends the high-level semantic features as a query to the cross-attention layer to output a final prediction result.

[0014] Furthermore, the feature space conversion module is a multi-layer MLP.

[0015] Furthermore, the downstream task is a classification task, and the Head module is a classifier.

[0016] The advantages of the present invention are as follows:

[0017] The method of the present invention directly transfers the knowledge of each client to other clients and makes local decisions on each client, effectively avoiding the knowledge loss caused by centralized server aggregation and eliminating the dependence on the central server.

[0018] The present invention designs a lightweight Mixture of Experts (MoE) module for each client. This local MoE module can adaptively make client-specific decisions by using lightweight experts from local and other clients, which helps to better adapt to the knowledge of other clients to improve performance and generalization ability without significantly increasing communication costs. BRIEF DESCRIPTION OF THE DRAWINGS

[0019] Figure 1 This is a framework flow chart of the method of the present invention.

[0020] Figure 2 Schematic diagram of the designed feature space transformation. DETAILED DESCRIPTION

[0021] The present invention will be described in further detail below with reference to the accompanying drawings. The examples given are only used to explain the present invention and are not used to limit the scope of the present invention.

[0022] The process of the decentralized federated learning method for medical scenarios of the present invention is as follows: Figure 1 As shown, each client includes a local model, a mixture of experts (MoE) module and an invited expert model.

[0023] The local model includes a Body module and a Head module; the expert model includes a feature space transformation module (Featurespace transform) and multiple lightweight Head modules. The Body module is used to extract features. The feature space transformation module converts local features into feature spaces corresponding to respective experts (heads). The Head module uses features to generate network outputs, and the head module of each client is also shared among all clients, and by inviting heads from other clients, an expert mixture (MoE) module of each client is formed. The present invention regards each Head module as an expert and uses the expert mixture (MoE) method to obtain the final output.

[0024] Body module: also the feature extraction module, its main purpose is to convert the original data into high-level semantic features, to facilitate the further analysis of the Head.

[0025] Head module: use the features generated by the Body module to perform downstream tasks, for example, in the classification task, Head represents the classifier, in the segmentation task, Head represents the Decoder of the segmentation model, etc.

[0026] Feature space conversion module: composed of multiple layers of MLP, their main function is to help the head module of client A to extract features and other clients corresponding to the expert model (lightweight head module) to complete the domain adaptation, and convert the extracted features into features that can be processed by each lightweight Head module.

[0027] MoE module: the present application adds a cross-attention mechanism to learn the weight of the prediction generated by each client expert. It should be noted that the local model not only captures the key information in the respective local data set, but also may inherit some bias, leading to overfitting. Through MoE, experts from all clients are used as candidates, and local features are used as queries to extract relevant information from all experts.

[0028] The training process of the present application includes three steps: 1) local model training, 2) local knowledge exchange, 3) mixed expert decision. Next, the present application will explain each step in detail. In the training process of the present application, the local client is trained using local data, wherein the medical image data (pathological sections, OCT data, etc.) is used in the medical image classification and segmentation task.

[0029] 1) Local model training, at this stage, the goal of the present application is to obtain a local model containing local knowledge through local data. Therefore, the present application only trains the head (Head module) and body (Body module) of the local model, and freezes the parameters of the feature space conversion module and the MoE module.

[0030] 2) Local knowledge exchange. In the communication phase of existing decentralized federated learning, each of the N participating clients needs to share its local model with the other N-1 clients. Therefore, there are a total of N(N-1) communications, which is much higher than centralized federated learning, which only requires 2N communications (uploads and downloads). In order to reduce the communication cost of decentralized federated learning, in the stage of sharing local headers between clients, the present invention only shares the header of the local model instead of the entire local model. The parameters of the header are several orders of magnitude smaller than the parameters of the local model, which greatly reduces the computational cost. Compared with centralized federated learning, this method does not introduce a significant communication burden. Experiments of the present invention show that compared with sharing the entire local model to each client, the communication overhead of the present invention is only 0.02%, and the performance of the present invention is comparable to it.

[0031] 3) Mixture of Experts Decision, this stage aims to learn the combined weights of all experts based on local data. During the training process of this stage, the present invention uses local data to fine-tune the parameters of feature space transformation and MoE again, while freezing other parameters. Feature space transformation in MoE. Before the final Mixture of Experts decision, the present invention designs a feature space transformation module to adapt local features to different expert models. It transforms local features into the feature space of the corresponding expert through a multi-layer MLP ( Figure 2 As shown). After the feature space transformation, the features generate predictions through their respective expert models in the corresponding space. The present invention uses the Mixture of Experts (MoE) framework to effectively aggregate the collected expert predictions. In order to enhance the MoE's attention to key experts, the present invention combines the cross-attention mechanism to learn the weights of the predictions generated by each client expert. It should be noted that the local model not only captures the key information in the respective local datasets, but may also inherit some biases, resulting in overfitting. By taking experts from all clients as candidates and using local features as queries to extract relevant information from all experts, the present invention ensures that the selected information reflects the common knowledge shared across clients. The present invention believes that this common information from different datasets has stronger generalization ability, and local biases are effectively suppressed in the selection process. Therefore, the present invention proposes to use a cross-attention design to filter local biases and enhance the overall generalization ability of the model.

[0032] The present invention has multiple advantages: 1) a decentralized federated learning method is designed; 2) a lightweight mixture of experts (MoE) module; 3) wide applicability.

[0033] While specific embodiments of the application have been disclosed in order to illustrate the application and to assist those skilled in the art in practicing the application, it is to be understood that various substitutions, modifications and changes can be made by those skilled in the art without departing from the spirit of the application and the scope of the appended claims. Accordingly, it is intended that the application not be limited, except by the scope of the claims.

Claims

1. A decentralized medical federated learning method for medical scenarios, comprising the following steps: 1) Deploy a local model, an expert model, and an expert hybrid module on each selected client of the federated learning system; The local model includes a Body module and a Head module. A corresponding lightweight Head module is created for the Head module in the local model on each selected client. The expert model on client A includes a feature space conversion module and lightweight Head modules corresponding to the Head modules on multiple other clients. 2) The Body module on client A converts the input data into high-level semantic features and sends them to the feature space conversion module and the Head module respectively; 3) The Head module on client A performs the set downstream task based on the high-level semantic features to obtain a prediction weight and sends it to the expert mixture module; the feature space conversion module on client A converts the high-level semantic features into features that can be processed by each lightweight Head module; each lightweight Head module performs the set downstream task based on the features to obtain a prediction weight and sends it to the expert mixture module; 4) The expert hybrid module on client A outputs the final prediction result according to each prediction weight and sends it to the Head module of each selected client for knowledge sharing.

2. The method according to claim 1, characterized in that The expert mixture module splices the prediction weights, sends the splicing result and the high-level semantic features to the cross-attention layer to output the final prediction result and sends it to the Head module of each selected client respectively.

3. The method according to claim 2, characterized in that The expert mixture module concatenates the prediction weights as key values, takes the high-level semantic features as queries, and sends them to the cross-attention layer to output the final prediction results.

4. The method according to claim 1, 2 or 3, characterized in that: The feature space conversion module is a multi-layer MLP.

5. The method according to claim 1, 2 or 3, characterized in that: The downstream task is a classification task, and the Head module is a classifier.

Images