A self-supervised pre-training method based on federated learning
By combining federated learning and self-supervised pre-training methods with masked autoencoders and ViT encoders, the problems of insufficient data annotation and privacy protection in multi-center medical image segmentation are solved, improving the accuracy of liver cancer image segmentation and the model's generalization ability, while ensuring patient privacy and security.
Patent Information
- Application Number
- CN202411423926.1
- Authority / Receiving Office
- CN · China
- Patent Type
- Patents(China)
- Current Assignee / Owner
- Filing Date
- 2024-10-12
- Publication Date
- 2026-01-30
- Estimated Expiration
- 2044-10-12
AI Technical Summary
Existing technologies face challenges in multi-center medical image segmentation, including a lack of data annotation and the need to protect patient privacy. Traditional methods struggle to efficiently segment liver cancer images in multi-center settings, especially given the growing emphasis on data privacy and ethical issues.
A self-supervised pre-training method based on federated learning is adopted, using mask autoencoders and ViT encoders for medical image analysis. The model is trained in a multi-center environment through self-supervised contrastive learning and reconstruction loss function. The client uploads the processed model and embeds it into the central server for parameter updates and aggregation.
It significantly improves the accuracy of lesion segmentation models, enhances the model's generalization ability and data diversity, ensures patient privacy and security, and provides more reliable support for early detection and diagnosis of liver cancer.
Smart Images

Figure CN119358628B_ABST
Abstract
Description
Technical Field
[0001] This invention belongs to the field of artificial intelligence, specifically a self-supervised pre-training method based on federated learning, focusing on medical image segmentation, especially for the detection and diagnosis of liver cancer. Background Technology
[0002] Medical image segmentation is a key task in computer vision, aiming to accurately extract regions of interest such as tumors and organs from medical images to aid in clinical diagnosis and treatment decisions. For liver cancer image segmentation, early detection and treatment planning are crucial and can significantly improve patient survival rates. However, traditional medical image segmentation methods typically rely on supervised learning, requiring large amounts of labeled data from multiple medical centers, which is often difficult to implement in medical practice, especially given the increasing emphasis on data privacy and ethical issues.
[0003] To address this challenge, self-supervised learning and federated learning techniques have gained increasing attention. Self-supervised learning utilizes large amounts of unlabeled data for feature self-learning, enhancing the model's generalization ability and adapting to different data distributions. Federated learning, on the other hand, allows for model training on distributed data from multiple medical centers, protecting patient privacy while improving model accuracy on diverse data. Combining these two advanced techniques enables efficient self-supervised pre-training of liver cancer images in multi-center environments, significantly improving model performance in downstream tasks, thereby advancing medical image analysis and promoting the development of precision medicine.
[0004] The advantage of federated learning frameworks for medical image segmentation lies in their ability to jointly train on data from multiple centers, addressing both the shortage of labeled data and protecting patient privacy. Commonly used federated methods include FedAvg, FedProx, FedBN, and MOON. FedAvg combines data privacy protection with model training by decomposing the task into local model training and global model aggregation. Each node generates weights locally and updates the global model through average aggregation. Its advantages include efficient use of distributed data for training, reduced data storage requirements, and improved model generalization ability. FedProx introduces a regularization term on top of FedAvg to alleviate heterogeneity issues among clients, improving model consistency and system robustness, especially when data distributions differ significantly, enabling more effective convergence to a better global model. FedBN addresses statistical inconsistencies caused by data heterogeneity by adding batch normalization (BN) layers to local models, mitigating training instability. Finally, MOON uses the similarity between model representations for comparative learning to correct local training biases, improve the feature extraction capability of the model in non-IID data scenarios, control drift, and bridge the gap between local and global models.
[0005] While federated learning can protect data privacy and enable joint training across multiple centers, its training process still relies on a large number of fully labeled images. However, in practical applications, especially in multi-center scenarios, there is often a situation where data is abundant but labels are scarce. To address this, self-supervised methods can be combined to perform self-supervised training on medical data from different centers, extensively learning semantic features. This not only increases the number of client participants and improves the model's generalization ability but also reduces the burden on hospitals in processing data. Currently, in self-supervised tasks on medical data, SelfMedMAE is a self-supervised pre-training framework based on reconstruction learning. It divides a large amount of unlabeled data into equal-volume blocks and treats them as sequences, randomly occluding parts of the blocks as reconstruction targets. During the process of restoring the complete image, the ViT-based encoder can learn rich semantic knowledge without labels, thereby achieving more accurate segmentation in subsequent tasks. Summary of the Invention
[0006] In view of the shortcomings of existing technologies, the purpose of this invention is to provide a self-supervised pre-training method based on federated learning to address the challenges of data labeling scarcity and patient privacy protection in multi-center environments. This method aims to increase the generalization ability of the model and improve the accuracy of lesion segmentation in medical images, thereby providing more reliable support for the early detection and diagnosis of serious diseases such as liver cancer. Simultaneously, by utilizing the federated learning framework, participating hospital centers can conduct joint training without sharing original patient data, effectively ensuring the privacy and security of patient information. This innovative method not only promotes the efficient use of data but also lays a solid foundation for achieving precision medicine, providing clinicians with more accurate image analysis tools and helping to improve the overall quality of medical services.
[0007] To address the problems existing in the prior art, this invention provides a self-supervised pre-training method based on federated learning. This method uses a masked autoencoder for medical image analysis and includes the following steps:
[0008] 1) First, 75% of the input medical image for each client is randomly masked. The pre-trained model uses the ViT encoder to aggregate the contextual information of the unmasked area to infer the content of the masked part. The MSE function is used as the reconstruction loss to reconstruct the complete image of the masked part.
[0009] 2) The client uses the ViT encoder embedding vector obtained in step 1) to convert it into a higher-dimensional feature representation through the MLP multilayer perceptron layer, uploads the embedding vector to the central server for comparison and learning to update the client parameters, and then aggregates to obtain the global model and distributes the updated parameters to each client.
[0010] 3) The client uses these parameters for local training and optimization, repeatedly uploading updated parameters to further improve the model until the model reaches the expected performance level.
[0011] As a preferred embodiment of the present invention, the client-side training model framework includes a ViT encoder for learning medical image features, a Transformer decoder for image reconstruction, and an MLP for transforming the embedding vectors into higher-dimensional, more complex feature representations.
[0012] As a preferred embodiment of the present invention, the ViT encoder transforms medical image data into a high-dimensional feature representation. Specifically, it segments the randomly masked image into small blocks and linearly embeds them into a fixed-dimensional vector space, combining a self-attention mechanism to capture the relationships between the small blocks and extract global features of the image.
[0013] As a preferred embodiment of the present invention, in step 2), the client transmits the high-dimensional features obtained through MLP to the central server, and updates the client using the contrastive loss obtained through self-supervised contrastive learning.
[0014] As a preferred embodiment of the present invention, the reconstruction loss, self-supervised contrastive loss, and cross-entropy loss are added together according to a certain parameter ratio to obtain a hybrid loss, and the hybrid loss is used to update the local client model parameters.
[0015] As a preferred embodiment of the present invention, for the current t-th round of communication training, the following steps are included:
[0016] Step 1: The central server sends the global model of the current round t to each client, and the client processes the current global model to train the local model of the current round.
[0017] Step 2: The client stores the local model that needs to be trained now, the global model returned from the central server in the previous round, and the local model obtained from the previous training round. The client performs a reconstruction task on the local masked image, and the reconstruction loss function is shown in formula (1):
[0018]
[0019] Where n represents the number of client-side image data participants, y i It is the actual target value of the i-th image data. It is the model's prediction of the i-th image data;
[0020] Meanwhile, the high-dimensional features from the ViT encoder after passing through the MLP are used for training using a hybrid loss function of self-supervised contrastive loss and cross-entropy loss. The self-supervised contrastive loss function is shown in Equation (2):
[0021]
[0022] Where sim(·) represents the cosine similarity function, and τ is a temperature parameter used to control the loss value and ensure the convergence of the loss. This represents the feature representation learned by the local model in round t. The feature representation represents the local model learning in the (t-1)th round. The feature representation learned by the global model in the (t-1)th round;
[0023] The training hybrid loss function is obtained by adding the three loss functions according to a certain parameter ratio, as shown in formula (3):
[0024]
[0025] Where s is the local dataset sample currently being trained. For the local model to be trained in the current t-th round, For the local model in round t-1, Let L be the global model for round t-1, μ be a hyperparameter used to control the weights of the self-supervised contrastive loss, and L be the weights of the model. CE Cross-entropy loss;
[0026] Step 3: After each client finishes training locally, it sends the trained local model to the central server for model aggregation.
[0027] Step 4: After collecting all the local models from all clients, the central server aggregates the models using a weighted average to obtain the global model for round t. At this point, round t of communication training is complete; the process then returns to step 1 for the next round of communication training. The overall optimization objective is shown in formula (4):
[0028]
[0029] Where N represents the total amount of data from all clients, n k This represents the number of datasets from the k-th client. For the mixed loss of the k-th client, w glob Since this is the global model to be trained, the overall optimization objective of the model is to use a mixed loss for each client. After training, a weighted average is then used to aggregate the results, ultimately yielding the optimal global model w. glob .
[0030] Based on the ideas provided by the present invention, see the specific implementation technical solutions below. Figure 1It mainly consists of two parts: first, self-supervised pre-training based on image reconstruction; second, using the self-supervised pre-training method under the federated learning framework and comparing the feature embeddings after ViT encoding to update the parameters, thereby reducing the encoding gap between different clients.
[0031] The self-supervised pre-training method uses a masked autoencoder (MAE) for image analysis. A portion of the input image's volumetric blocks is masked and fed as a sequence. The model is then trained to reconstruct the complete image from the partially masked image. During reconstruction, a ViT-based encoder and decoder are employed, which effectively aggregates contextual information from the unmasked portions from a sequence perspective, thereby inferring the content of the masked parts. This approach not only allows the model to better understand the overall structure of the image but also improves prediction accuracy in cases of missing information, resulting in a significant performance improvement for image analysis tasks.
[0032] Meanwhile, the system's self-supervised pre-training is based on a federated learning framework. The embedding vectors encoded by each client in the previous step are uploaded to the central server for comparative learning. The parameters of each client are updated using the combined contrastive and reconstruction losses. The central server then aggregates these parameters to obtain a global model, which is then distributed to each client. Each client then uses these updated parameters for local training and optimization, repeatedly uploading updated embeddings to further improve the model. This process is repeated until the model reaches the expected performance level, thus completing the self-supervised pre-training in a multi-center environment.
[0033] As described above, the self-supervised pre-training method based on federated learning of the present invention has the following beneficial effects: through joint self-supervised pre-training on multi-center data, the accuracy of the lesion segmentation model in downstream tasks is significantly improved; the self-supervised framework allows hospital centers that do not have labeled medical images to participate in training, enhancing the generalization ability of the model and the diversity of data; under the federated learning framework, each hospital center does not need to share the original data, but only needs to upload the processed model embedding, ensuring the privacy and security of patient data. Attached Figure Description
[0034] Figure 1 The flowchart of the self-supervised pre-training method based on federated learning in this invention is shown.
[0035] Figure 2 The diagram shows the network structure of the self-supervised pre-training method based on federated learning in this invention.
[0036] Figure 3 A schematic diagram illustrating the update of privacy parameters in federated learning according to the present invention.
[0037] Figure 4The relationship between Epoch and Loss Score in the self-supervised pre-training method based on federated learning in this invention is shown in the graph. Detailed Implementation
[0038] Federal Learning
[0039] Federated learning is a distributed machine learning technique that aims to enable collaborative training among partner institutions (such as hospitals) by keeping data locally. In this process, institutions only exchange model parameters without sharing local data, thus achieving cross-institutional model training while ensuring data privacy.
[0040] The training process for federated learning mainly consists of the following steps:
[0041] (1) Initialize the global model: Initialize a global model containing initial parameters on the central server;
[0042] (2) Model transmission: such as Figure 3 As shown, the central server transmits the current global model to each client; the current global model The model is distributed to each client as a local model for local training. Let the local model of the k-th (k = 1, 2, ..., K) local client be .
[0043] (3) Perform local training: Client S k The model is trained using local data; after receiving the model from the central server, the client trains it using the cross-entropy loss function and stochastic gradient descent, let L... CE Indicates client S k The loss function and the model update process can be expressed as formula (1), where η represents the learning rate during training:
[0044]
[0045] (4) Upload local model: After training is completed, each client uploads its updated local model to the central server;
[0046] (5) Aggregating local models: The central server is responsible for aggregating the local models uploaded by each client, usually by weighted average. The approach, where weights are typically based on the amount of data from each client, is a global model. The aggregation can be expressed as a formula:
[0047]
[0048] Where |D| represents the total number of samples in each client dataset, |D k| Indicates that it comes from client S k The number of samples in the local data.
[0049] (6) Obtaining the final global model: The main goal of federated learning is to learn a global model w from each client through a central server without exchanging data. Therefore, the overall optimization objective of the model is as shown in formula (3):
[0050]
[0051] Among them, L k (w) represents the client S k The cross-entropy loss during model training is given by the formula (4):
[0052]
[0053] In formula (4), D k Indicates client S k The corresponding dataset, E, represents the expectation for dataset D. k The expected value of all samples (x,y) in the dataset, (x,y) ~ D k Representing data D k In the data, there is a data sample (x, y), where x represents the sample in the data and y represents the label in the data.
[0054] Self-supervised pre-training based on contrastive learning
[0055] This invention primarily studies a self-supervised pre-training framework based on contrastive learning, a derivative technique of self-supervised learning. It enhances the understanding of semantic context features by different client models from the perspective of image reconstruction and the shared constraints of contrastive learning. Its main purpose is to enable the encoder to learn feature embeddings from a large number of unlabeled samples, allowing it to more fully understand the semantic context features of samples when encoding in downstream segmentation tasks. The introduction of contrastive learning ensures that the feature embeddings encoded by encoders from different clients exhibit similar representations for the same feature embedding within a certain feature space, while different feature embeddings exhibit disparate representations. This method, combining self-supervised pre-training and contrastive learning, is particularly suitable for multi-center feature learning in the absence of labeled data.
[0056] During the encoding phase, 75% of the input medical image for each client is randomly masked. The pre-trained model uses the ViT Encoder to aggregate contextual information of the unmasked regions to infer the content of the masked parts. The MSE function is used as the reconstruction loss to reconstruct the complete image of the masked parts. For the input hepatocellular carcinoma CT image data x of size h×w×d... i Cut it into n pieces of size The data is divided into sub-volume blocks, with 75% of the region masked as masked blocks and the remaining 25% visible as visible blocks. Positional information is added to these visible blocks, and the resulting sequence is used as input to the ViT encoder. The ViT encoder used in this invention stacks 12 Transformer blocks with 12 heads and an embedding dimension of 384, using sincos positional embedding. After encoding, the liver cancer CT image data is mapped to deeper feature embeddings as a deeper representation of semantic context information.
[0057] set up Representing the image x i For image x i Generate feature embedding vector z i Its positive and negative pair feature embedding vectors are and The definition of its self-supervised contrastive loss function is shown in formula (5):
[0058]
[0059] in Indicate z i The positive feature embedding vector, Indicate z i The negative pair feature embedding vector, where τ represents the temperature parameter.
[0060] The goal of the decoding stage is to embed the aforementioned features into the vector z. iAfter the Transformer blocks are restored, the original complete input sequence is reconstructed, with the masked blocks as the learning target. Then, the obtained ViT encoder embedding vectors are transformed into higher-dimensional feature representations using MLP (Multilayer Perceptron) layers. These embedding vectors are uploaded to a central server for comparative learning to update client parameters. Finally, the models are aggregated to obtain a global model, and the updated parameters are distributed to each client. Since the decoder needs to reconstruct the masked pixels, the input contains two parts: feature vectors generated by the encoder for the undisturbed pixel blocks and information about the masked pixel blocks. All masked pixel block information is represented by the same learnable feature vector. The decoder here is essentially an architecture containing Transformer blocks; therefore, the input vector includes positional encoding. Like the encoder, the Transformer cannot learn positional information, so positional information is needed to indicate the location of a pixel block in the original image. This decoder is only used in the pre-training stage. When transferring to downstream tasks, the decoder can be replaced with any architecture according to the user's actual needs. The last layer of the decoder is a linear layer. Assuming the block size is 16×16, the output dimension of the linear layer is 256. After reshaping to 16×16, the reconstructed pixel block is obtained. Finally, the MSE function is used as the loss function to measure the similarity between the reconstructed image and the real image, as shown in formula (6).
[0061]
[0062] Where n represents the number of client-side image data participants, y i It is the actual target value of the i-th image data. It is the model's prediction of the i-th image data;
[0063] The final loss in the self-supervised pre-training process can be defined as:
[0064] L mix =L MSE +μL SC (7)
[0065] Example:
[0066] like Figure 2As shown, a self-supervised pre-training method based on federated learning is presented. The client-side training model framework includes a ViT encoder for learning medical image features, a Transformer decoder for image reconstruction, and an MLP for transforming embedded vectors into higher-dimensional, more complex feature representations. The ViT encoder transforms medical image data into high-dimensional feature representations by segmenting randomly masked images into small patches and linearly embedding them into a fixed-dimensional vector space, combined with a self-attention mechanism to capture the relationships between the patches and extract global features of the image. This invention improves upon the self-supervised mask pre-training method, aiming to learn features of unlabeled images. By using the ViT encoder and ViT decoder for masked image reconstruction, it enhances the learning ability of the self-supervised model on unlabeled images. This can solve the problem of insufficient medical image data and improve the accuracy of lesion segmentation.
[0067] The communication training for the current t-th round includes the following steps:
[0068] Step 1: The central server sends the global model of the current round t to each client, and the client processes the current global model to train the local model of the current round.
[0069] Step 2: The client stores the local model that needs to be trained now, the global model returned from the central server in the previous round, and the local model obtained from the previous training round. The client performs a reconstruction task on the local masked image, and the reconstruction loss function is shown in Equation (8):
[0070]
[0071] Where n represents the number of client-side image data participants, y i It is the actual target value of the i-th image data. It is the model's prediction of the i-th image data;
[0072] Meanwhile, the high-dimensional features from the ViT Encoder after passing through the MLP are trained using a hybrid loss function of self-supervised contrastive loss and cross-entropy loss. The self-supervised contrastive loss function is shown in Equation (9):
[0073]
[0074] Where sim(·) represents the cosine similarity function, and τ is a temperature parameter used to control the loss value and ensure the convergence of the loss. This represents the feature representation learned by the local model in round t. The feature representation represents the local model learning in the (t-1)th round. The feature representation learned by the global model in the (t-1)th round;
[0075] The training mixed loss function is obtained by adding the three loss functions according to a certain parameter ratio, as shown in formula (10):
[0076]
[0077] Where s is the local dataset sample currently being trained. For the local model to be trained in the current t-th round, For the local model in round t-1, Let L be the global model for round t-1, μ be a hyperparameter used to control the weights of the self-supervised contrastive loss, and L be the weights of the model. CE Cross-entropy loss;
[0078] Step 3: After each client finishes training locally, it sends the trained local model to the central server for model aggregation.
[0079] Step 4: After collecting all the local models from all clients, the central server aggregates the models using a weighted average to obtain the global model for round t. At this point, round t of communication training is complete; the process then returns to step 1 for the next round of communication training. The overall optimization objective is shown in formula (11):
[0080]
[0081] Where n represents the total amount of data from all clients, n k This represents the number of datasets from the k-th client. For the mixed loss of the k-th client, w glob Since this is the global model to be trained, the overall optimization objective of the model is to use a mixed loss for each client. After training, a weighted average is then used to aggregate the results, ultimately yielding the optimal global model w. glob .
[0082] Summarizing the above steps, the pseudocode of the algorithm structure of the self-supervised pre-training method based on federated learning proposed in this invention is as follows:
[0083]
[0084]
[0085] As described above, the self-supervised pre-training method based on federated learning of the present invention has the following beneficial effects: it significantly improves the accuracy of the lesion segmentation model through joint training with self-supervised learning and multi-center data; it allows hospital centers that do not have labeled medical images to participate in training, enhancing the model's generalization ability and data diversity; under the federated learning framework, each hospital center does not need to share the original data, but only needs to upload the processed model embedding, ensuring the privacy and security of patient data.
[0086] Comparative experiment:
[0087] 1. Experimental Preparation
[0088] To demonstrate the superiority of this invention for lesion segmentation in multicenter medical images, experiments were conducted on a private liver cancer dataset. This private dataset, sourced from Guangdong Provincial People's Hospital, contains 247 computed tomography (CT) images of liver cancer. These images were annotated by two radiologists with over ten years of experience. The loss of the reconstructed images was calculated using the annotated images to learn their features.
[0089] To simulate the experimental setup of federated learning, the dataset needs to be partitioned to simulate various local clients. Before training, this invention divides the dataset images and their label information for each client into a training set and a test set at a ratio of approximately 4:1, with their quantities being 197:50 respectively. Then, the images in the training set are divided into a dataset conforming to the federated learning framework under three centers at a ratio of 65:65:67.
[0090] 2. Comparative Experiment
[0091] This invention not only compares experimental results with Central learning but also with several methods in the federated learning field. These include the FedAvg algorithm, a commonly used federated method that updates the global model by training locally on each participant and averaging the model weights. Furthermore, it is compared with the MOON algorithm, which addresses the heterogeneity of client data, to comprehensively evaluate the performance advantages and applicability of this invention. Comparisons with other researchers' methods demonstrate that this invention can effectively segment lesions in liver cancer CT medical images, achieving higher segmentation accuracy than other methods.
[0092] Table 1. Comparative analysis results of this invention with other network models on a private dataset (Test Dice)
[0093]
[0094] The data analysis in Table 1 shows that the FedAvg and MOON methods perform relatively well on private datasets, achieving good results. However, despite demonstrating their potential within the federated learning framework, their performance in most cases still fails to surpass that of traditional centralized methods. This indicates that federated learning still faces certain performance bottlenecks in practical applications, especially in complex tasks involving the integration of multi-center data.
[0095] However, it is particularly noteworthy that, compared to the FedAvg and MOON methods, the federated contrastive self-supervised pre-training method of this invention demonstrates significant advantages in medical image segmentation tasks. Specifically, when dealing with lesion segmentation, this method not only improves the accuracy of the segmentation model but also exhibits strong model generalization ability. This performance improvement may be attributed to the enhanced feature representation capabilities of the self-supervised pre-training mechanism under limited labeled data conditions, as well as the better adaptability of the federated learning framework to different center data distributions. This result demonstrates that the method of this invention can effectively improve the accuracy of lesion segmentation in a federated learning environment and provides a powerful solution for multi-center medical image analysis tasks.
[0096] from Figure 4 As can be seen, the loss function value exhibits a significant and rapid decreasing trend in the early epochs during training. This phenomenon indicates that the model parameters effectively optimize the training error in the initial stage, demonstrating strong learning ability and rapid convergence. As training continues, the rate of decrease in the loss function gradually slows down and eventually stabilizes. This stable decreasing curve generally indicates that the model's learning rate is properly configured, and the training process is efficient and stable. Furthermore, the low final stable value of the loss function further illustrates that the model has reached a good local optimum on the training data, effectively capturing patterns in the data. These observations demonstrate that the model method proposed in this invention performs excellently during training, rapidly adapting to data and stably optimizing its predictive performance.
[0097] In summary, this invention provides a self-supervised pre-training method based on federated learning, aiming to improve the effectiveness of medical image analysis. By designing a self-supervised mask pre-training process, this method improves data utilization, enabling hospitals and centers without labeled medical images to participate in training. This not only enhances the model's generalization ability but also enriches data diversity, promoting a better understanding of the characteristics of different patient groups. This method combines privacy-preserving federated learning with self-supervised learning techniques, aiming to improve the accuracy and robustness of liver cancer image segmentation. By jointly training on medical data from multiple medical centers, the image segmentation model is optimized to achieve more efficient and accurate liver cancer image analysis, thereby better supporting clinical decision-making and treatment planning. Leveraging the advantages of federated learning, hospitals and centers do not need to share raw data; they only need to upload processed model embeddings, effectively protecting the privacy and security of patient data. This design is particularly important in today's increasingly data-protection and ethical medical environment. Furthermore, based on this, contrastive learning is used to correct the client's local model through the similarity between model representations, further improving the accuracy and stability of the segmentation model in practical applications. This method provides a new approach to advance medical image segmentation technology and helps to achieve more efficient decision support in clinical diagnosis.
[0098] Therefore, this invention effectively overcomes the various shortcomings of the prior art and has high industrial application value.
[0099] The above embodiments are merely illustrative of the principles and effects of the present invention and are not intended to limit the invention. Any person skilled in the art can modify or alter the above embodiments without departing from the spirit and scope of the present invention. Therefore, all equivalent modifications or alterations made by those skilled in the art without departing from the spirit and technical concept disclosed in the present invention should still be covered by the claims of the present invention.
Claims
1. A federated learning based self-supervised pre-training method, characterized in that, The method comprises the following steps for model training of the current t-th round: 1) The central server sends the global model of the current t-th round to each client, and the client processes the current global model to train the local model of the current round; 2) After receiving the global model, the client uses the ViT encoder to aggregate the context information of the unmasked area of the local randomly masked image, infers the masked part of the image for the reconstruction task, and calculates the reconstruction loss using the mean square error (MSE) function, as shown in formula (1): wherein n represents the number of clients participating in image data, y i is an actual target value of the i-th image data, is a predicted value of the i-th image data by the model; At this time, the client saves the local model to be trained, the global model returned by the central server of the previous round, and the local model obtained by training of the previous round, transforms the embedding vector output by the ViT encoder into a high-dimensional feature representation through a multi-layer perceptron (MLP), and constrains the high-dimensional feature based on a self-supervised contrastive learning loss, as shown in formula (2): wherein sim(·) represents a cosine similarity function, τ is a temperature parameter used to control the loss value to ensure the convergence of the loss, represents the feature representation learned by the local model in the t-th round, represents the feature representation learned by the local model in the t-1-th round, represents the feature representation learned by the global model in the t-1-th round. The reconstruction loss function, the cross-entropy loss function and the self-supervised contrastive loss function are added according to a certain parameter ratio to obtain a training mixed loss function formula (3): wherein s is the current training local data set sample, is the current t-th round of local model to be trained, is the t-1-th round of local model, is the t-1-th round of global model, and μ is a hyperparameter for controlling the weight of self-supervised contrastive loss. 3) Each client updates the local model parameters according to the mixed loss of step 2), and uploads the trained local model to the central server for aggregation; 4) After collecting the local models of all clients, the central server aggregates the models using a weighted average method to obtain the global model of the t-th round, and the communication training of the t-th round is completed; 5) Repeat steps 1 to 4, and the overall optimization objective is shown in formula (4): where N represents the total number of client data, n k represents the number of data sets from the kth client, is the mixed loss of the kth client, w glob is the current global model to be trained, therefore, the model overall optimization target is to use the mixed loss of each client After training, then in the form of weighted average aggregation, the final optimal global model w glob .
2. The federated learning based self-supervised pre-training method of claim 1, wherein: The model framework trained by the client includes a ViT encoder for learning medical image features, a Transformer decoder for image reconstruction, and an MLP for converting embedding vectors into higher-dimensional and more complex feature representations.
3. The federated learning based self-supervised pre-training method according to claim 1 or 2, characterized in that: The ViT encoder converts medical image data into high-dimensional feature representations, specifically by dividing the randomly masked image into small blocks and linearly embedding them into a fixed-dimensional vector space, combining a self-attention mechanism to capture the relationship between the small blocks, and extracting global features of the image.
4. The federated learning based self-supervised pre-training method of claim 3, wherein: In step 2), the client transmits the high-dimensional feature after the MLP to the central server, and updates the client using the contrastive loss obtained by self-supervised contrastive learning.
5. The federated learning based self-supervised pre-training method according to claim 4, wherein: The reconstruction loss, the self-supervised contrastive loss and the cross-entropy loss are added according to a certain parameter ratio to obtain a mixed loss, and the local client model parameters are updated using the mixed loss.
Citation Information
Patent Citations
Method and device for constructing medical image large model based on federated learning
CN115408377A
Personalized federal learning method based on comparative learning and conditional calculation
CN118396082A