A classifier-assisted federated contrastive learning method for solving the Non-IID problem

Abstract

The present invention provides a federated contrastive learning method for solving the Non-IID problem with the assistance of a classifier. The present invention learns feature representation by shortening the distance between similar data and widening the distance between different data in the feature space. Similar data in the present invention refers to sample features belonging to the same category obtained by a global model. Because contrastive learning can make better use of the mutual information between different sample features. On a data set with imbalanced data, contrastive learning can balance the proportion of different categories in the loss through a class averaging strategy, while reducing the importance of multiple categories, making the model pay more attention to minority categories, so it has broad application prospects in processing imbalanced data. However, due to the lack of a unified feature representation in federated learning, contrastive learning often widens the differences between different local samples. Therefore, a unified feature representation is the key to the application of contrastive learning in federated learning.

Description

Technical Field

[0001] The present invention relates to the technical field of artificial intelligence, and in particular to a federated contrastive learning method for solving the Non-IID problem using a classifier as an aid. Background Art

[0002] Federated Learning [1] It is a distributed machine learning framework designed to improve the accuracy of local models while preserving the privacy of local data by leveraging model parameters from other clients. However, in real-world scenarios, data from different participants often exhibit significant variability, and differences in data distribution between clients can lead to deviations in the direction of client-side model optimization. This imbalance in local model gradient updates can adversely affect the convergence of the global model and overall communication overhead.

[0003] The key to solving the problem of non-identically distributed data in federated learning is to reduce the deviation between the optimization directions when updating the local model gradients across clients. Transfer learning, knowledge extraction, and adaptive aggregation of similar clients are effective methods to solve the data heterogeneity in federated learning. [2] An adaptive local aggregation module is designed to adaptively integrate the downloaded global model with the local model on each client to initialize the local model before each training iteration. This is an effective aggregation strategy to improve the performance of local models. However, this approach focuses on training personalized local models rather than generating a well-performing global model.

[0004] There are many algorithms that use contrastive learning to solve the problem of non-independent and independently distributed (Non-IID) data in federated learning. [3] FedPCL is the first algorithm to apply contrastive learning to the federated framework. This algorithm is designed to solve the Non-IID data problem based on model contrastive learning. Unlike standard contrastive algorithms, it maximizes the feature similarity between the current local model and the global model by comparing the feature representation of the current local model with the feature representation of the local model and the global model in the previous round, thereby preventing the forgetting of global knowledge of locally missing data. However, due to the need to calculate multiple contrast models, this algorithm has high time complexity and memory requirements. [4] By sharing the prototype feature vectors of the class instead of the model parameters, a pre-trained model is used to build a federated model customized for a specific task. However, this is just a fine-tuning of the pre-trained model and has significant limitations. [5]The class prototype feature vector is used as the global feature representation to correct local training. This algorithm calculates the similarity between the class prototype and the client's local data features, ensuring that the local feature representation of each client model does not deviate from the global prototype. This will alleviate the drift of the optimization direction of the local model to a certain extent.

[0005] In addition, contrastive learning has also been applied to solve other problems in federated learning. [6] It proposes to solve the client model drift problem caused by multimodal differences through intra-modal and inter-modal contrastive learning. This is a federated learning framework suitable for heterogeneous models and multimodal data, demonstrating the learning ability of contrastive learning in handling joint multimodal scenarios. This algorithm also provides important inspiration for the application of contrastive learning in solving data heterogeneity problems in federated learning. [7] It aims to learn a unified and balanced feature representation for decentralized data in a privacy-preserving manner, where local feature representations are aggregated into a unified representation for the target task, and adversarial learning is applied to eliminate the bias of sensitive features. However, this algorithm is not suitable for horizontal federated learning and is an algorithm applied to vertical federated learning.

[0006] References:

[0007] [1]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.

[0008] [2] Zhang J, Hua Y, Wang H, et al. Fedala: Adaptive local aggregation for personalized federated learning [C] / / Proceedings of the AAAI Conference on Artificial Intelligence. 2023, 37(9): 11237-11244.

[0009] [3]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.

[0010] [4]Tan Y, Long G, Ma J, et al. Federated learning from pre-trained models: A contrastive learning approach [J]. Advances in neural information processing systems, 2022, 35: 19332-19344.

[0011] [5]Mu

[0012] [6]Yu Q, Liu Y, Wang Y, et al.Multimodal federated learning viacontrastive representation ensemble[J].arxiv preprint arxiv:2302.08888,2023.

[0013] [7]Qi T, Wu F, Wu C, et al. Fairvfl: A fair vertical federated learning framework with contrastive adversarial learning [J]. Advances in neural information processing systems, 2022, 35: 7852-7865. Summary of the Invention

[0014] In view of the technical problems mentioned in the above background technology, a federated contrastive learning method for solving the Non-IID problem using a classifier as an aid is provided.

[0015] The technical means adopted in the present invention are as follows:

[0016] A federated contrastive learning method for solving the Non-IID problem using a classifier as an aid, characterized by comprising the following steps:

[0017] 1. A federated contrastive learning method for solving the Non-IID problem using a classifier as an aid, characterized by comprising the following steps:

[0018] Step 1: Design and use a two-branch network framework for federated contrastive learning and initialize the global model, which is then sent to the local client.

[0019] Step 2: Input sample x undergoes two different data augmentations to obtain two different augmented samples x a ,x b ; Put two samples x a ,x b At the same time, input the feature extractor h to obtain different features v a ,v b ; Then use batch normalization to process the feature vector to obtain the normalized sample z a ,z b ;

[0020] Step 3: Set the sample feature z a Input the classifier f to obtain the output probability vector Q of the sample for different categories, and calculate the cross entropy loss with bias between the output probability vector Q and its label y;

[0021] Step 4: For sample feature z b After passing the feature extractor, use the projection head P to obtain its feature representation p x ; For sample feature z a Get p′ through the projection head of the global model x , and at the same time, the weights of the classifier of the global model [ω1,ω2,...,ω l ] After splitting into class vectors, each vector is normalized to a modulus of 1, and the feature vector p is obtained through the projection head of the global model w ;

[0022] Step 5: Calculate the contrast loss and cross entropy loss of the current local model;

[0023] Step 6: Update the gradient based on the loss and return to step 2, repeating multiple rounds;

[0024] Step 7: Upload the local models of multiple clients to the server;

[0025] Step 8: The server performs a weighted summation of the model parameters from multiple clients based on the amount of data from different clients to obtain an averaged global model;

[0026] Step 9: Send the global model to the local machine and repeat step 1.

[0027] Furthermore, the shared feature extraction module is an arbitrary neural network for extracting features from samples, and the parameters of the shared feature extraction module are shared for the projection layer and the classifier head.

[0028] Furthermore, the normalization process is:

[0029]

[0030] Among them, μ represents the mean of the batch samples, σ represents the standard deviation of the batch samples, ∈ represents a small amount to prevent the denominator from being 0, and γ and β represent the learnable parameters in the network;

[0031] The projection head consists of a multi-layer linear layer network, where the weight parameters of the multi-layer linear layer are set to fixed values ​​that cannot be updated by gradients. The initialization formula is:

[0032]

[0033] Among them, E represents a The size of the matrix, j represents the input dimension, k represents the output dimension, and I represents a The identity matrix, and E and I satisfy the equation E T E=I, Represents an all-1

[0034] The matrix of .

[0035] Furthermore, the cross entropy loss formula with bias term between the output probability vector Q and its label y is:

[0036]

[0037] Among them, n i represents the number of classes i; τ represents the temperature coefficient, which is used to control the smoothness of the output probability, y i Indicates the predicted probability of whether the sample belongs to class i.

[0038] Furthermore, the contrast loss between the current local model and the global model is:

[0039]

[0040] Among them, z irepresents the feature p from the local model x , z p Represents the features p′ of all data in the current batch size from the global model x The collection and class vector p w , τ′ represents a temperature coefficient, which is used to control the degree of attention paid to negative labels by contrastive learning.

[0041] Furthermore, the loss in step 6 is:

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

[0043] This invention addresses the issue of non-IID data in horizontal federated classification. This algorithm addresses the slow convergence of federated learning when cross-client data is non-IID. Applying contrastive learning to federated learning significantly improves its convergence speed. Furthermore, under these conditions, the accuracy of the test set reaches a state-of-the-art level. BRIEF DESCRIPTION OF THE DRAWINGS

[0044] In order to more clearly illustrate the embodiments of the present invention or the technical solutions in the prior art, the following briefly introduces the drawings required for use in the embodiments or the description of the prior art. Obviously, the drawings described below are some embodiments of the present invention. For ordinary technicians in this field, other drawings can be obtained based on these drawings without paying any creative labor.

[0045] Figure 1 This is the local gradient update process of the present invention.

[0046] Figure 2 This is the overall process of the federated learning algorithm of the present invention.

[0047] In the figure: ① is the global model downloaded to the local server; ② is the local model updated using local data; ③ is the local model uploaded to the global server. DETAILED DESCRIPTION

[0048] In order to enable those skilled in the art to better understand the solutions of the present invention, the technical solutions in the embodiments of the present invention will be clearly and completely described below in conjunction with the drawings in the embodiments of the present invention. Obviously, the embodiments described are only part of the embodiments of the present invention, not all of the embodiments. Based on the embodiments of the present invention, all other embodiments obtained by ordinary technicians in this field without making creative efforts should fall within the scope of protection of the present invention.

[0049] It should be noted that the terms "first", "second", etc. in the description and claims of the present invention and the above-mentioned drawings are used to distinguish similar objects and are not necessarily used to describe a specific order or sequence. It should be understood that the numbers used in this way can be interchanged where appropriate, so that the embodiments of the present invention described herein can be implemented in an order other than those illustrated or described herein. In addition, the terms "including" and "having" and any variations thereof are intended to cover non-exclusive inclusions. For example, a process, method, system, product or device that includes a series of steps or units is not necessarily limited to those steps or units clearly listed, but may include other steps or units that are not clearly listed or inherent to these processes, methods, products or devices.

[0050] Assume there are n samples Data owned by each participant The purpose of federated learning is to m ) to conduct joint training of the model. Since the client data distribution is not independent and identically distributed, the total number of samples n for each client and the number of samples for each category are are different.

[0051] The contrastive learning loss function is:

[0052]

[0053] Where q is the characteristic of the sample, k + It is the feature of samples of the same type as E. Generally, it is a different enhanced representation of the same sample in unsupervised learning. τ represents the temperature coefficient, which can be used to adjust the degree of attention of the loss to negative samples. i It is the feature representation of all samples involved in the comparison of the entire batch size.

[0054] Contrastive learning multiplies sample features with those of other samples, bringing samples with the same characteristics closer together and separating features of different categories further apart in feature space. This leverages the mutual information between samples to improve the model's feature extraction capabilities. Federated learning leverages this characteristic of contrastive learning to improve the model's ability to handle non-IID data.

[0055] In this algorithm, the overall framework of federated learning remains the same as that of the federated averaging algorithm. The global server first generates a random global model, and then the local client downloads the global model and performs gradient updates.

[0056] After that, the global server performs weighted averaging to obtain the global model. The overall process is as follows Figure 2 shown.

[0057]

[0058] A federated contrastive learning method for solving the Non-IID problem using a classifier as an aid, characterized by comprising the following steps:

[0059] Step 1: Design a two-branch network framework for federated contrastive learning. The two-branch network framework for federated contrastive learning includes a shared feature extraction module, a projection module for contrastive learning, and a classification head for classifier training. The shared feature extraction module is an arbitrary neural network that extracts features from samples. The parameters of the shared feature extraction module are shared between the projection layer and the classification head.

[0060] Step 2: Input sample x undergoes two different data augmentations to obtain two different augmented samples x a ,x b ; Put two samples x a ,x b At the same time, input the feature extractor h to obtain different features v a ,v b ; Then use batch normalization to process the feature vector to obtain the normalized sample z a ,z b ; The normalization process is:

[0061]

[0062] Among them, μ represents the mean of the batch samples, τ represents the standard deviation of the batch samples, ∈ represents a small amount to prevent the denominator from being 0, and γ and β represent the learnable parameters in the network;

[0063] The projection head consists of a multi-layer linear layer network, where the weight parameters of the multi-layer linear layer are set to fixed values ​​that cannot be updated by gradients. The initialization formula is:

[0064]

[0065] Among them, E represents a The size of the matrix, j represents the input dimension, k represents the output dimension, and I represents a The identity matrix, and E and I satisfy the equation E T E=I, Represents an all-1

[0066] The matrix of .

[0067] Step 3: Set the sample feature z aInput the classifier f to obtain the output probability vector Q of the sample for different categories, and calculate the cross entropy loss with bias between the output probability vector Q and its label y. The formula for calculating the cross entropy loss with bias between the output probability vector Q and its label y is:

[0068]

[0069] Among them, n i represents the number of classes i; τ represents the temperature coefficient, which is used to control the smoothness of the output probability, y i Indicates the predicted probability of whether the sample belongs to class i.

[0070] Step 4: For sample feature z b After passing the feature extractor, use the projection head P to obtain its feature representation p x ; For sample feature z a Get p′ through the projection head of the global model x , and at the same time, the weights of the classifier of the global model [ω1,ω2,...,ω l ] After splitting into class vectors, each vector is normalized to a modulus of 1, and the feature vector p is obtained through the projection head of the global model w ;

[0071] Step 5: Calculate the contrast loss between the current local model and the global model; the contrast loss between the current local model and the global model is:

[0072]

[0073] Among them, z i represents the feature p from the local model x , z p Represents the features p′ of all data in the current batch size from the global model x The collection and class vector p w , τ′ represents a temperature coefficient, which is used to control the degree of attention paid to negative labels by contrastive learning.

[0074] Step 6: Calculate the loss, perform gradient update based on the loss, and return to step 2, repeating multiple rounds. The loss in step 6 is:

[0075] Step 7: Upload the local models of multiple clients to the server;

[0076] Step 8: The server performs a weighted summation of the model parameters from multiple clients based on the amount of data from different clients to obtain an averaged global model; the weighted formula is:

[0077]

[0078] where n k is the total number of samples of the k-th client;

[0079] Step 9: Send the global model to the local machine and repeat step 1.

[0080] In federated learning scenarios, the local data distribution of different clients often exhibits significant differences due to various factors such as the device resources, geographic location, and time differences of different participants. These differences in data distribution can cause local models to tend toward the local optimal solution during updates and deviate from the global optimal model. This paper considers a federated learning scenario where the data is not independent and identically distributed: Consider a large medical group with multiple hospitals seeking to build a cross-hospital early cancer diagnosis model. Due to differences in location, resources, and expertise, each hospital has a distinct patient population and data distribution.

[0081] like Figure 2 As shown in the figure, different participants have their own private data and cannot share data with each other. When their local data differs significantly, the locally updated model will deviate significantly from the global optimal solution, and the weighted results will often have poor test accuracy. Therefore, we consider using federated contrastive learning, using the classifier weights as auxiliary feature vectors for contrastive learning. This allows different clients to converge on the same class center during local updates, resulting in a unified class feature representation.

[0082] This algorithm significantly improves model convergence performance for solving the federated Non-IID problem. It first establishes a two-branch network model and then designs a cross-entropy loss for imbalanced data. Furthermore, it utilizes the classifier weights of the global model to obtain a unified class vector, ensuring that the local model does not deviate from the feature representation of the global model during the update process.

[0083] The serial numbers of the above embodiments of the present invention are for description only and do not represent the advantages or disadvantages of the embodiments.

[0084] In the above embodiments of the present invention, the description of each embodiment has its own focus. For parts that are not described in detail in a certain embodiment, reference can be made to the relevant descriptions of other embodiments.

[0085] In the several embodiments provided in this application, it should be understood that the disclosed technical content can be implemented in other ways. Among them, the device embodiments described above are only exemplary. For example, the division of the units can be a logical function division. In actual implementation, there may be other division methods, such as multiple units or components can be combined or integrated into another system, or some features can be ignored or not executed. Another point is that the mutual coupling or direct coupling or communication connection shown or discussed can be through some interfaces, indirect coupling or communication connection of units or modules, which can be electrical or other forms.

[0086] The units described as separate components may or may not be physically separate, and the components shown as units may or may not be physical units, that is, they may be located in one place or distributed across multiple units. Some or all of the units may be selected according to actual needs to achieve the purpose of the present embodiment.

[0087] In addition, the functional units in the various embodiments of the present invention may be integrated into a single processing unit, each unit may exist physically separately, or two or more units may be integrated into a single unit. The aforementioned integrated units may be implemented in the form of hardware or software functional units.

[0088] If the integrated unit is implemented in the form of a software functional unit and sold or used as an independent product, it can be stored in a computer-readable storage medium. Based on this understanding, the technical solution of the present invention, or the part that contributes to the prior art, or all or part of the technical solution can be embodied in the form of a software product. The computer software product is stored in a storage medium and includes several instructions for enabling a computer device (which can be a personal computer, server or network device, etc.) to perform all or part of the steps of the method described in each embodiment of the present invention. The aforementioned storage medium includes: U disk, read-only memory (ROM, Read-Only Memory), random access memory (RAM, Random Access Memory), mobile hard disk, magnetic disk or optical disk, etc. Various media that can store program codes.

[0089] Finally, it should be noted that the above embodiments are only used to illustrate the technical solutions of the present invention, rather than to limit it. Although the present invention has been described in detail with reference to the above embodiments, those skilled in the art should understand that they can still modify the technical solutions described in the above embodiments, or replace some or all of the technical features therein with equivalents. However, these modifications or replacements 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 federated contrastive learning method for solving the Non-IID problem using classifier assistance, characterized in that: The following steps are involved: Step 1: Design and use a two-branch network framework for federated contrastive learning and initialize the global model, which is then sent to the local client. Step 2: Input sample x undergoes two different data augmentations to obtain two different augmented samples x a ,x b ; Put two samples x a ,x b At the same time, input the feature extractor h to obtain different features v a ,v b ; Then use batch normalization to process the feature vector to obtain the normalized sample z a ,z b ; Step 3: Set the sample feature z a Input the classifier f to obtain the output probability vector Q of the sample for different categories, and calculate the cross entropy loss with bias between the output probability vector Q and its label y; Step 4: For sample feature z b After passing the feature extractor, use the projection head P to obtain its feature representation p x ; For sample feature z a Get p′ through the projection head of the global model x , and at the same time, the weights of the classifier of the global model [ω1,ω2,...,ω l ] After splitting into class vectors, each vector is normalized to a modulus of 1, and the feature vector p is obtained through the projection head of the global model w ; Step 5: Calculate the contrast loss and cross entropy loss of the current local model; Step 6: Update the gradient based on the loss and return to step 2, repeating multiple rounds; Step 7: Upload the local models of multiple clients to the server; Step 8: The server performs a weighted summation of the model parameters from multiple clients based on the amount of data from different clients to obtain an averaged global model; Step 9: Send the global model to the local machine and repeat step 1.

2. The federated contrastive learning method for solving the Non-IID problem using classifier assistance according to claim 1, characterized in that: The dual-branch network framework for federated contrastive learning includes: a shared feature extraction module, a projection head for contrastive learning, and a classification head for classifier training.

3. The federated contrastive learning method for solving the Non-IID problem using classifier assistance according to claim 2, characterized in that: The shared feature extraction module is an arbitrary neural network for extracting features from samples, and the parameters of the shared feature extraction module are shared by the projection layer and the classifier head.

4. The method for solving the Non-IID problem using classifier-assisted federated contrastive learning according to claim 1, characterized in that: The normalization process is: Among them, μ represents the mean of the batch samples, τ represents the standard deviation of the batch samples, ∈ represents a small amount to prevent the denominator from being 0, and γ and β represent the learnable parameters in the network; The projection head consists of a multi-layer linear layer network, where the weight parameters of the multi-layer linear layer are set to fixed values ​​that cannot be updated by gradients. The initialization formula is: Among them, E represents a The size of the matrix, j represents the input dimension, k represents the output dimension, and I represents a The identity matrix, and E and I satisfy the equation E T E=I, Represents an all-1 The matrix of .

5. The federated contrastive learning method for solving the Non-IID problem using classifier assistance according to claim 1, characterized in that: The cross entropy loss formula with bias between the output probability vector Q and its label y is: Among them, n i represents the number of classes i; τ represents the temperature coefficient, which is used to control the smoothness of the output probability, y i Indicates the predicted probability of whether the sample belongs to class i.

6. The federated contrastive learning method for solving the Non-IID problem using classifier assistance according to claim 1, characterized in that: The contrast loss between the current local model and the global model is: Among them, z i represents the feature p from the local model x , z p Represents the features p′ of all data in the current batch size from the global model x The collection and class vector p w The union of τ ′ Represents a temperature coefficient that controls the degree to which contrastive learning focuses on negative labels.

7. The federated contrastive learning method for solving the Non-IID problem using classifier assistance according to claim 1, characterized in that: The loss in step 6 is:

Images