A multi-task federated learning system and method for enabling multi-client collaboration

Abstract

The application provides a multi-task federated learning system for realizing multi-client cooperation, comprising: a segmentation client S and a classification client C; the number of the segmentation client S or the classification client C is one or more, each segmentation client S or each classification client C holds a local model; and the segmentation client S and the classification client C are both U-net network structures. The method combines federated learning and multi-task learning, so that the clients for realizing different tasks can be jointly trained, the existing data is maximally utilized, and the client model effect is improved. Compared with the traditional single-task federated learning method, the present task more widely utilizes the existing data, provides a new method for applying deep learning to process medical images in a real scene, and provides a new idea for further development of federated learning.

Description

Technical Field

[0001] This invention relates to the field of medical image processing technology, and more particularly to a multi-task federated learning method for realizing multi-client collaboration. Background Technology

[0002] Multi-task learning [1] Multi-task learning is a machine learning paradigm that learns multiple tasks simultaneously by sharing a model. By learning shared representations across related tasks, it can alleviate overfitting to some extent, improve generalization ability, and address the cold start problem. Multi-task learning can be categorized from the model framework into soft parameter sharing. [2] Sharing with hard parameters [1] Hard parameter sharing, which involves sharing hidden layers across all tasks, is the most commonly used multi-task learning mechanism in neural networks. Multi-task learning allows the discovery of shared information across similar tasks, improving model performance.

[0003] Federated learning is a groundbreaking, privacy-preserving distributed machine learning approach that shares learned knowledge with each client while storing each client's data locally. The concept of federated learning was first proposed by Google in 2016, followed by the Federated Averaging (FedAvg) method proposed by B. McMahan et al. [3] Federated learning is a widely used basic algorithm, upon which various higher-performance federated learning algorithms have been derived. However, general federated learning methods face challenges such as poor convergence with highly heterogeneous data and a lack of personalized solutions. To address these challenges, personalized federated learning methods have been proposed to obtain high-quality, personalized models for individual clients holding non-independent and identically distributed local data.

[0004] Currently, classic personalized federated learning methods include:

[0005] 1. FedProx [4] Introducing a proximal term in local customer training limits local model bias caused by non-independent and identically distributed data, thereby increasing the stability and convergence of model training.

[0006] 2. HarmoFL [5] The image in the client is converted to the frequency domain by fast Fourier transform, and then amplitude normalization is performed by combining amplitude perturbation, so as to build a personalized local model that is more adapted to local data for each client.

[0007] 3. MOON [6] This approach combines federated learning with contrastive learning and adds a contrastive loss to reduce local weight differences, addressing the situation where the local data held by each client is not independent and identically distributed.

[0008] In addition to the methods mentioned above, many works combine federated learning with knowledge distillation, meta-learning, model interpolation, clustering, and other methods to develop better personalized federated learning models. Unlike existing methods, this invention uses a novel information-sharing model to collaboratively train clients performing different tasks to generate personalized federated learning models, making full use of existing data.

[0009] Reference documents:

[0010] [1]Caruana R.Multitask learning[J].Machine learning,1997,28:41-75.

[0011] [2]Ruder S.An overview of multi-task learning in deep neural networks[J].arXiv preprint arXiv:1706.05098,2017.

[0012] [3]McMahan B, Moore E, Ramage D, et al. Communication-efficient learning of deep networks from decentralized data [C] / / Artificial intelligence and statistics. PMLR, 2017: 1273-1282.

[0013] [4]Li T,Sahu AK,Zaheer M,et al.Federated optimization inheterogeneous networks[J].Proceedings of Machine learning and systems,2020,2:429-450.

[0014] [5]Jiang M, Wang Z, Dou Q. Harmofl: Harmonizing local and global drifts in federated learning on heterogeneous medical images[C] / / Proceedings of the AAAI Conference on Artificial Intelligence.2022,36(1):1087-1095.

[0015] [6]Li Q,He B,Song D.Model-contrastive federated learning[C] / / Proceedings of the IEEE / CVF conference on computer vision and patternrecognition.2021:10713-10722. Summary of the Invention

[0016] In response to the technical problems mentioned in the background section, this invention provides a multi-task federated learning method for multi-client collaboration. By leveraging the similarities between image classification and image segmentation tasks through multi-client joint learning, this invention aims to improve the performance of both classification and segmentation client models while protecting patient privacy.

[0017] The technical means employed in this invention are as follows:

[0018] A multi-task federated learning system for multi-client collaboration includes: a segmentation client S and a classification client C; the number of segmentation clients S or classification clients C is one or more, and each segmentation client S or each classification client C holds a local model; both segmentation clients S and classification clients C are U-net network structures.

[0019] The local model of the segmentation client S includes: a main task image segmentation unit and an auxiliary task image reconstruction unit I; the main task image segmentation unit extracts the image representation of the patient's lung CT image through an encoder, then inputs the image representation into a decoder, and supplements the edge information missing during the downsampling process during the decoding process;

[0020] The local model of the classification client C includes: a main task image classification unit and an auxiliary task image reconstruction unit II.

[0021] Furthermore, the forward propagation process of the segmentation client is defined as follows: For the encoder part of the network:

[0022] and

[0023] Image segmentation loss function L seg By Dice loss function L Dice and MSE loss function L seg-MSE The composition is shown in formulas (1) and (2); simultaneously, through image reconstruction to assist in image segmentation training, the loss function L reAs shown in formula (3), where the subscript j represents the j-th pixel in the image; the total loss function L of the segmentation client. seg As shown in formula (4):

[0024]

[0025] Among them, y seg^ The segmented image is the predicted output of the segmentation model; y seg Represents the ground truth image of the image segmentation data; ε represents the Laplacian smoothing term;

[0026]

[0027] in, This represents the j-th pixel in the predicted segmented image output by the segmentation model; This represents the j-th pixel in the ground truth image of the image segmentation data;

[0028]

[0029] in, Represents the j-th pixel in the predicted reconstructed image from the output of the reconstruction network; The j-th pixel in the ground truth image of the image reconstruction data represents N; N represents the total number of pixels in an image.

[0030] L seg =L Dice +L seg-MSE +L re (4).

[0031] Furthermore, the forward propagation process of the segmentation client is defined as follows:

[0032] and

[0033] Using the cross-entropy loss function L CE As the loss function for image classification, it is shown in Equation (5); consistent with the image reconstruction network in the segmentation client, hard parameter sharing is also used in the classification client. The reconstruction and classification networks share a single encoder to assist in training the image classification network. The loss function L of the image reconstruction network is... re As shown in formula (3); the total loss function L of the classification client cls As shown in (6);

[0034] L CE =-[y cls log(y cls^ )+(1-y cls )log (1-y cls^ (5);

[0035] Among them, y cls^ This represents the predicted classification result output by the classification model; y cls Labels representing image classification data;

[0036] L cls =L CE +L re (6).

[0037] This invention also includes a multi-task federated learning method for implementing multi-client collaboration, comprising the following steps:

[0038] Step 1: Each client updates its local model using local data;

[0039] Step 2: Each client sends its local model parameters to the server via the network. The server performs a weighted average aggregation of the shared encoder portion of all clients: by adding the parameters of the encoder portion of the local model of all clients and then dividing by the total number of clients; the resulting model parameters are used as the new global encoder.

[0040] Step 3: The server sends all the model parameters of the newly generated global encoder to each client; each client uses the received model parameters of the global encoder to replace the encoder part of the local model, and uses the local data to train the local model and update the local model parameters.

[0041] Step 4: Aggregate all model parameters of local models of the same type using a weighted average to generate a global model for a single-task classification task and a global model for a single-task segmentation task. Then, send the parameters of the global model for the single-task classification task to all classification clients and the parameters of the global model for the single-task segmentation task to all segmentation clients via the network.

[0042] Step 5: Guide the training of the local model using the single-task global model sent to the local client in the previous step: Input local data into both the local model and the single-task global model, and freeze the parameters of the single-task global model (i.e., do not update the parameters of the single-task global model), and perform a third local update.

[0043] Step 6: Repeat the above steps until each local model converges.

[0044] Furthermore, the KL divergence loss L was added to the loss function. KL As shown in formula (7), where the superscript "main" refers to the output result of the main task of each client, i.e., the segmentation result or the classification result, y main^ This is the prediction output of the main task in the single-task global model. The prediction output for the main task of the i-th client model:

[0045]

[0046] Therefore, the loss function for segmenting clients in the third local update is shown in Equation (8), and the loss function for classifying clients is shown in Equation (9):

[0047] L seg-KL =L seg +L KL (8);

[0048] L cls-KL =L cls +L KL (9).

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

[0050] This method fully utilizes existing data while ensuring data privacy and security, improving the segmentation performance of medical images. Image reconstruction is incorporated as an auxiliary task into each client network model, effectively assisting the training and convergence of the main task and improving its performance. This method combines federated learning with multi-task learning, enabling clients implementing different tasks to be jointly trained, maximizing the use of existing data and improving client model performance. Compared to traditional single-task federated learning methods, this method makes broader use of existing data, providing a new approach for applying deep learning to medical image processing in real-world scenarios and offering new ideas for the further development of federated learning. Attached Figure Description

[0051] To more clearly illustrate the technical solutions in the embodiments of the present invention or the prior art, the drawings used in the description of the embodiments or the prior art will be briefly introduced below. Obviously, the drawings described below are some embodiments of the present invention. For those skilled in the art, other drawings can be obtained based on these drawings without creative effort.

[0052] Figure 1 This invention segments the client network structure.

[0053] Figure 2 This invention categorizes client network structures.

[0054] Figure 3 This is the overall network structure of the present invention.

[0055] Figure 4 This forms the basic framework for horizontal federated learning in this invention. Detailed Implementation

[0056] To enable those skilled in the art to better understand the present invention, the technical solutions of the present invention will be clearly and completely described below with reference to the accompanying drawings of the embodiments of the present invention. Obviously, the described embodiments are only some embodiments of the present invention, and not all embodiments. Based on the embodiments of the present invention, all other embodiments obtained by those skilled in the art without creative effort should fall within the scope of protection of the present invention.

[0057] It should be noted that the terms "first," "second," etc., in the specification, claims, and accompanying drawings of this invention are used to distinguish similar objects and are not necessarily used to describe a specific order or sequence. It should be understood that such data can be interchanged where appropriate so that the embodiments of the invention described herein can be implemented in orders other than those illustrated or described herein. Furthermore, the terms "comprising" and "having," and any variations thereof, are intended to cover a non-exclusive inclusion; for example, a process, method, system, product, or apparatus that comprises a series of steps or units is not necessarily limited to those steps or units explicitly listed, but may include other steps or units not explicitly listed or inherent to such processes, methods, products, or apparatus.

[0058] like Figure 1-4 As shown, this invention provides a single-task federated learning method where existing federated learning methods involve joint learning among clients performing the same task. This invention, however, includes two clients performing different tasks: a segmentation client S and a classification client C. Both network structures are further designs based on the U-net network structure.

[0059] Client segmentation: Network structure as follows Figure 1 As shown. The segmentation client consists of two parts: the main task of image segmentation and the auxiliary task of image reconstruction. The image segmentation network uses an encoder to extract image representations of the patient's lung CT images, then inputs these image representations into a decoder, supplementing the edge information missing during downsampling during the decoding process. In this invention, the forward propagation process of the segmentation client is defined as: and Image segmentation loss function L seg By Dice loss function L Dice and MSE loss function L seg-MSE The composition is shown in formulas (1) and (2). Simultaneously, we use image reconstruction to assist in the training of image segmentation, with a loss function L... re As shown in formula (3), where the subscript j represents the j-th pixel in the image. The total loss function L for the segmentation client... seg As shown in formula (4).

[0060]

[0061]

[0062]

[0063] L seg =L Dice +L seg-MSE +L re (4);

[0064] Classified clients: Network structure such as Figure 2 As shown. To achieve joint learning of the classification model and the segmentation model, a classification client model is designed based on the network structure of the segmentation client. The classification client consists of two parts: the main task of image classification and the auxiliary task of image reconstruction. In this invention, the forward propagation process of the segmentation client is defined as: and Using the cross-entropy loss function L CE The loss function for image classification is shown in Equation (5). Consistent with the image reconstruction network in the segmentation client, the classification client also uses hard parameter sharing, with the reconstruction and classification networks sharing a single encoder to assist in training the image classification network. The loss function L... re As shown in formula (3). The total loss function L for the classification client. cls As shown in (6).

[0065] L CE =-[y cls log(y cls^ )+(1-y cls log(1-y) cls^ (5);

[0066] L cls =L CE +L re (6);

[0067] To achieve collaboration between clients, this invention proposes a novel federated learning model aggregation and update method. Each round of communication includes three local updates and two model aggregations, using different loss functions and model aggregation methods. The training process is as follows: Figure 3As shown. The single-task global model is generated by aggregating all local models from the same client. First local update: Each client updates its local model using local data so that the local model can learn from the local data. First model aggregation: A weighted average aggregation of the shared encoder parts from all clients is performed to generate an encoder with global features to promote cooperation among different clients. Second local update: Clients fine-tune the encoder with global features using local data to improve the generalization ability of the local model on local data. Second model aggregation: The corresponding local models are aggregated into a classification single-task global model and a segmentation single-task global model, which are used to guide local models implementing the same task, ensuring that the optimization objective of the local model does not deviate from the corresponding global objective and task. Third local update: Local data is input into the local model and the corresponding single-task global model, while the parameters of the single-task global model are frozen. KL divergence loss is added to the total loss function to guide the local model utilizing the single-task global model.

[0068] In the first two local updates, the loss functions for the segmentation client and the classification client are shown in equations (4) and (5), respectively. The loss function for the third local update incorporates the KL divergence loss L. KL As shown in formula (7), the superscript main refers to the output result of the main task of each client, that is, the segmentation result or the classification result, y main^ This is the prediction output of the main task in the single-task global model. This is the prediction output of the main task of the i-th client model.

[0069]

[0070] Therefore, the loss function for segmenting the client in the third local update is shown in Equation (8), and the loss function for classifying the client is shown in Equation (9). The addition of KL divergence loss is mainly to limit the distribution of parameters of the updated local model so that it does not deviate too much from the corresponding single-task global model, thereby ensuring the convergence of the local model.

[0071] L seg-KL =L seg +L KL (8);

[0072] L cls-KL =L cls +L KL (9);

[0073] Example 1

[0074] Federated machine learning, also known as federated learning, is a distributed machine learning framework with privacy protection and secure encryption technologies. It aims to enable decentralized participants to collaboratively train machine learning models without disclosing data to other participants. Using federated learning to process medical data avoids the data storage and enormous communication overhead associated with uploading large amounts of medical data to a single server. It also leverages more data to generate better models while ensuring data security and patient privacy, achieving modeling results comparable to traditional deep learning. Federated learning effectively addresses current privacy and security concerns in data sharing and analysis within the medical field, as well as the insufficient availability of local data for certain diseases.

[0075] Medical image processing is an important medical diagnostic task, involving the processing of various types of images such as CT, MRI, and X-rays. Disease diagnosis and lesion severity assessment are two crucial components of medical image processing. Disease diagnosis involves classifying organs based on the identification of lesion areas in images. Lesion severity assessment involves segmenting the images of lesion areas to identify the affected regions, observing the size and progression of the lesions, and analyzing the severity of the lesions. This process greatly improves the diagnostic efficiency of physicians.

[0076] The invention will be further illustrated below with examples. Multi-task federated learning, which involves joint training of classification and segmentation clients, greatly helps improve the diagnostic efficiency of hospitals and ensure the effective conduct of medical diagnostic work. Specific application scenarios are as follows:

[0077] In real-world scenarios, medical institutions vary significantly in size. Larger institutions possess far more data, both in quantity and quality, than smaller ones. Due to limitations in equipment and patient flow, smaller institutions typically suffer from poor data quality and limited data volume. However, training deep learning models requires both sufficient data volume and quality; otherwise, the trained model will perform poorly. Furthermore, because labeling high-quality segmentation data is costly, even in large-scale medical institutions, the amount of segmentation data is far less than the amount of classification data.

[0078] Therefore, the multi-task federated learning algorithm proposed in this invention enables mutual assistance among medical institutions of different sizes in classifying and segmenting data within the same organization. Each hospital can set up a client for classification and a client for segmentation, and conduct collaborative training with other hospitals. In use, an application interface can be designed as the model interface; doctors only need to input the obtained patient medical images into the system to obtain diagnostic results regarding the disease. Furthermore, since the symptoms of the same disease are similar but different, federated learning can enhance the generalization of disease diagnosis among various medical institutions, better diagnosing unseen lesions and facilitating the diagnosis process for doctors and the patient's treatment.

[0079] The sequence numbers of the above embodiments of the present invention are for descriptive purposes only and do not represent the superiority or inferiority of the embodiments. In the above embodiments of the present invention, the descriptions of each embodiment have their own emphasis; parts not described in detail in a certain embodiment can be referred to in the relevant descriptions of other embodiments. It should be understood that the disclosed technical content in the several embodiments provided in this application can be implemented in other ways.

[0080] Finally, it should be noted that the above embodiments are only used to illustrate the technical solutions of the present invention, and not to limit them; although the present invention has been described in detail with reference to the foregoing embodiments, those skilled in the art should understand that modifications can still be made to the technical solutions described in the foregoing embodiments, or equivalent substitutions can be made to some or all of the technical features; and these modifications or substitutions do not cause the essence of the corresponding technical solutions to deviate from the scope of the technical solutions of the embodiments of the present invention.

Claims

1. A multi-task federated learning system for realizing multi-client collaboration, characterized in that, include: Split Client And Classification Client The segmentation client Or categorized clients The number is one or more, each segment client Or each category client Each holds a local model; the segmentation client And Classification Client All are U-net network structures; The segmentation client The local model includes: a main task image segmentation unit and an auxiliary task image reconstruction unit I; the main task image segmentation unit extracts the image representation of the patient's lung CT image through an encoder, then inputs the image representation into a decoder, and supplements the edge information missing during the downsampling process during the decoding process; The classification client The local model includes: a primary task image classification unit and an auxiliary task image reconstruction unit II; the forward propagation process of the segmentation client is defined as follows: For the encoder part of the network: and ; Loss function for image segmentation By Dice loss function and MSE loss function The composition is shown in formulas (1) and (2); simultaneously, image segmentation is trained by image reconstruction, and the loss function is... As shown in formula (3), where the subscript j Represents the first in the image j 1 pixel; total loss function for segmenting the client As shown in formula (4): ； in, This represents the predicted segmented image output by the segmentation model. A ground truth image representing image segmentation data; ε Represents the Laplace smoothing term; ； in, The segmentation model outputs the predicted segmented image of the th j 1 pixel; The ground truth image of image segmentation data j 1 pixel; ； in, The first image in the reconstructed image is represented by the predicted output of the reconstruction network. j 1 pixel; The first true image in the image reconstruction data j 1 pixel; N Represents the total number of pixels in an image. ； The forward propagation process of the classification client is defined as follows: and ; Through cross-entropy loss function As the loss function for image classification, it is shown in Equation (5); consistent with the image reconstruction network in the segmentation client, the classification client also uses hard parameter sharing, and the reconstruction and classification networks share a single encoder to assist in training the image classification network. The loss function of the image reconstruction network is... As shown in formula (3); the total loss function of the classification client As shown in (6); ； in, This represents the predicted classification result output by the classification model. Labels representing image classification data; 。 2. A multi-task federated learning method for implementing multi-client collaboration using the system described in claim 1, characterized in that, Includes the following steps: Step 1: Each client updates its local model using local data; Step 2: Each client sends its local model parameters to the server via the network. The server then performs a weighted average aggregation of the shared encoder portion of all clients: by adding up the parameters of the encoder portion of the local models of all clients and then dividing by the total number of clients. The model parameters obtained from this will be used as the new global encoder; Step 3: The server sends all model parameters of the newly generated global encoder to each client; Each client replaces the encoder part of its local model with the model parameters received from the global encoder, and trains its local model using local data to update the local model parameters. Step 4: Generate a global model for a single classification task and a global model for a single segmentation task by aggregating all model parameters of local models of the same type using a weighted average; and send the parameters of the global model for a single classification task to all classification clients and the parameters of the global model for a single segmentation task to all segmentation clients via the network. Step 5: Guide the training of the local model using the single-task global model sent to the local client in the previous step: Input local data into both the local model and the single-task global model, and freeze the parameters of the single-task global model (i.e., do not update the parameters of the single-task global model), and perform a third local update. Step 6: Repeat the above steps until each local model converges.

3. The multi-task federated learning method for realizing multi-client collaboration according to claim 2, characterized in that, The loss function was added KL Divergence loss As shown in formula (7), where the superscript main This refers to the output of each client's main task, i.e., the segmentation or classification results. This is the prediction output of the main task in the single-task global model. The prediction output for the main task of the i-th client model: ； Therefore, the loss function for segmenting clients in the third local update is shown in formula (8), and the loss function for classifying clients is shown in formula (9): ； 。

Images