A federated image classification method for class imbalance
By introducing a Gaussian distribution prototype and virtual feature set into federated learning, the class imbalance problem is solved, improving the classification accuracy of the minority class, model robustness, adaptability, and privacy protection capabilities, making it suitable for scenarios such as medical image classification.
Patent Information
- Application Number
- CN202411223253.5
- Authority / Receiving Office
- CN · China
- Patent Type
- Patents(China)
- Current Assignee / Owner
- Filing Date
- 2024-09-03
- Publication Date
- 2025-11-21
- Estimated Expiration
- 2044-09-03
AI Technical Summary
Existing federated learning methods suffer from poor class imbalance, especially in medical diagnostic tasks, where models perform poorly in class classification. Furthermore, traditional methods struggle to effectively improve minority class classification performance while protecting data privacy.
By introducing a Gaussian distribution prototype and a virtual feature set, the client generates intra-class variance information and trains the model. The server aggregates and updates the Gaussian distribution prototype to generate a new global model. The client then retrains the classifier to ensure a balanced class distribution and reduce classification bias.
It effectively reduces classification bias caused by class imbalance, improves the classification accuracy of minority classes, enhances the robustness and generalization ability of the model, and reduces the data transmission burden, thus meeting privacy protection requirements.
Smart Images

Figure CN119339125B_ABST
Abstract
Description
TECHNICAL FIELD
[0001] The present application relates to a federated image classification method, belonging to the technical field of image classification. BACKGROUND
[0002] Image classification is an important task in the field of computer vision, widely applied in scenarios such as autonomous driving, intelligent security, and medical image analysis. In recent years, deep learning models have shown excellent performance in centralized data environments. However, in practical applications, data is often distributed among different devices or institutions, making it difficult to train centrally. To solve this problem, federated learning has become an important technical means.
[0003] Federated learning is a distributed learning paradigm that allows multiple participants to jointly train a global model without sharing data. It not only protects data privacy but also effectively trains with scattered data, thus being widely applied in fields such as finance and medicine where privacy protection is highly required. Through federated learning, each participant only needs to share local model updates without transmitting actual data, thus effectively reducing the risk of data leakage.
[0004] Although federated learning improves model training efficiency while protecting data privacy, it still faces many challenges. One of the most significant problems is class imbalance. Class imbalance refers to the phenomenon where the number of samples in certain classes in the data set is significantly less than that in other classes. This phenomenon is particularly common in the field of medical image classification, for example, in medical diagnosis tasks, cancer samples are usually much less than healthy samples, leading to poor classification performance on minority classes, thus affecting the accuracy and reliability of diagnosis results.
[0005] Existing federated learning methods mainly focus on data heterogeneity, and the study of class imbalance is relatively less. Compared with centralized learning, the class imbalance problem in federated learning is more complex, which may exist in the local data of clients or may differ between clients. Due to privacy protection requirements, traditional class imbalance handling methods such as oversampling and undersampling cannot be directly applied. Although some methods have attempted to address class imbalance, the effect is limited, especially in improving the classification performance of minority classes. For example, some studies suggest uploading local class distribution information for client selection, which violates the privacy protection principle. Another method attempts to retrain using auxiliary balanced data sets on the server side, but such data sets are often difficult to obtain in practical applications.
[0006] Therefore, there is an urgent need to design an image classification method that can effectively address the class imbalance problem in a federated learning environment while protecting data privacy. This method should be able to improve the classification performance of minority classes without relying on external data, thus addressing the shortcomings of existing technologies and better meeting the needs of practical applications. SUMMARY
[0007] The present application is to solve the problem that the existing federated learning method mainly focuses on data heterogeneity and cannot reduce the influence of class imbalance on model performance, and further proposes a federated image classification method for class imbalance.
[0008] The technical solution adopted by the present application to solve the above problem is that the steps of the present application include:
[0009] Step 1, the server initializes a global Gaussian distribution prototype, and the client initializes a local model parameter;
[0010] Step 2, the global Gaussian distribution prototype is distributed to the client;
[0011] Step 3, the client generates intra-class variance information based on the local data set using the Gaussian prototype, performs local model training through the configured loss function, and generates a local prototype;
[0012] Step 4, the client generates a balanced virtual feature set by random sampling, and re-trains the classifier combined with the global and local prototypes;
[0013] Step 5, the client uploads the updated local prototype to the server;
[0014] Step 6, the server aggregates the local prototypes uploaded by each client to generate a new global Gaussian distribution prototype, which is used for the next round of training;
[0015] Step 7, determine whether the set training round or model convergence is reached, if yes, the training is ended, if not, repeat steps 2 to 6.
[0016] Further, the local Gaussian prototype The mean and covariance of each class of samples are calculated to capture the intra-class variance information and enhance the robustness of the classification model:
[0017] (1)
[0018] (2)
[0019] In formulas (1) and (2), denotes the feature vector of the i-th input sample of the j-th class of the k-th client. i k j
[0020] Further, the class-aware Gaussian prototype learning is introduced at the client side, and the loss function is as follows:
[0021] (3).
[0022] Further, the virtual feature set is generated by fusing global and local Gaussian prototypes through random sampling, ensuring that the fused features have both global information and retain local features;
[0023] (4),
[0024] In formula (4), represents the local virtual feature, represents the global virtual feature, is a fusion coefficient, used to balance global and local features.
[0025] Further, the client classifier is retrained using the fused features, ensuring balanced class distribution to reduce the bias of the classifier.
[0026] Further, when the server aggregates the uploaded local prototypes, it generates new global Gaussian prototypes using the global mean and covariance of each class of samples.
[0027] (5)
[0028] (6)
[0029] In formulas (5) and (6), represents the total number of samples of all clients in the i-th class. k represents the total number of samples of all clients in the i-th class.
[0030] The beneficial effects of the present application are: 1) reducing the impact of class imbalance: by introducing a prototype based on Gaussian distribution and a classifier retraining mechanism, the present application effectively reduces the classification bias caused by class imbalance, especially improving the classification accuracy of minority classes, solving the problem of deviation and misclassification in traditional methods when dealing with class imbalance data. 2) Improve the adaptability of federated learning: through the generation and balancing of virtual feature sets, the present application improves the adaptability of the model to heterogeneous data while maintaining data privacy, enhancing the robustness and generalization ability of the federated learning system. 3) Enhance the classification performance of minority classes: using Gaussian prototypes and intra-class variance information, the present application improves the discriminant ability of the classifier in the case of fewer minority class samples. 4) Reduce data transmission burden: the present application only uploads the updated local prototype, reducing the data transmission bandwidth requirement and reducing the risk of data leakage, meeting the privacy protection requirements of federated learning. 5) Wide application prospect: the present application can adapt to various class imbalance problems, such as medical image classification, financial risk analysis, etc., and has wide application value and market potential. These beneficial effects not only improve the performance of the federated learning model in the class imbalance scenario, but also exhibit significant advantages in protecting data privacy and improving computing efficiency BRIEF DESCRIPTION OF DRAWINGS
[0031] Fig. 1 is a flowchart of the present application;
[0032] Fig. 2 is a schematic diagram of the global model training framework of the present application;
[0033] Fig. 3 is a comparison diagram of whether the classifier of the method of the present application is retrained. DETAILED DESCRIPTION
[0034] Specific implementation one: as shown in the figure, a federated image classification method for class imbalance, the specific steps include: Figs. 1-3
[0035] Step 1, the server initializes the global Gaussian distribution prototype, and the client initializes the local model parameter;
[0036] Step 2, the global Gaussian distribution prototype is distributed to the client;
[0037] Step 3, the client generates the intra-class variance information based on the local data set, using the Gaussian prototype, and trains the local model through the configured loss function, to generate the local prototype;
[0038] Step 4, the client generates a balanced virtual feature set by random sampling, and re-trains the classifier in combination with the global and local prototypes;
[0039] Step 5, the client uploads the updated local prototype to the server;
[0040] Step 6, the server aggregates the local prototypes uploaded by each client to generate a new global Gaussian distribution prototype, which is used for the next round of training;
[0041] Step 7, determine whether the set training round or model convergence is reached, if yes, the training is ended, if not, repeat steps 2 to 6.
[0042] Specific implementation two: as shown in the figure, the local Gaussian prototype Figs. 1-3 is obtained by calculating the mean and covariance of each class of samples, which is used to capture the intra-class variance information and enhance the robustness of the classification model:
[0043] (1)
[0044] (2)
[0045] In formulas (1) and (2), indicates the i th client k th class j input samples of the feature vector.
[0046] Embodiment three: As shown in the figure, a class-aware Gaussian prototype learning is introduced at the client side, and its loss function is as follows: Figs. 1-3
[0047] (3).
[0048] Embodiment four: As shown in the figure, the virtual feature set is generated by fusing global and local Gaussian prototypes through random sampling, ensuring that the fused features have both global information and local features; Figs. 1-3
[0049] (4),
[0050] In formula (4), represents the local virtual feature, represents the global virtual feature, is a fusion coefficient for balancing global and local features.
[0051] Embodiment five: As shown in the figure, the client classifier is retrained using the fused features to ensure balanced class distribution to reduce the bias of the classifier. Figs. 1-3
[0052] Embodiment six: As shown in the figure, when the server aggregates the uploaded local prototypes, it generates new global Gaussian prototypes using the global mean and covariance of each class of samples. Figs. 1-3
[0053] (5)
[0054] (6)
[0055] In formulas (5) and (6), represents the total number of samples of the i-th class of all clients. k
[0056] Embodiment
[0057] A class imbalance-oriented federated image classification method, comprising the following steps:
[0058] Step one: the server initializes the global Gaussian prototype , and the client initializes the local model parameter , wherein the local models of various clients can be heterogeneous;
[0059] Step two: the server distributes the global Gaussian prototype to each client;
[0060] Step three: The client trains locally based on the global Gaussian prototype, especially for the processing of minority class samples.
[0061] In the first t communication, the client uses the local private data and extracts the feature vector of each sample through the feature extraction network (such as ResNet18). Specifically, for the i input sample j of the client, its feature vector is , where is the feature extractor of the i client. Since the features generated by deep neural networks can be approximated by a mixture of Gaussian distributions, in order to better model the intra-class features, we assume that the features of each class k are subject to Gaussian distribution , where, for class k , is the mean of the local Gaussian prototype, is the covariance matrix. The calculation process is as follows:
[0062] (7)
[0063] (8)
[0064] where denotes the number of samples of class i in the k client.
[0065] In order to learn discriminative representations, not only do we need to ensure that the prototype boundary is clear, but we also need to align between clients to preserve semantic information. Therefore, class-aware Gaussian prototype learning is introduced in the client, and its loss function is as follows:
[0066] (9)
[0067] where denotes the loss function of the class-aware Gaussian prototype learning of the i client. For of the i-th client, and . is a similarity measure function used to measure the distance k between the input sample and the class mean of the global prototype
[0068] (10)
[0069] where denotes the hyper-parameter, which is set to 0.1 in this experiment.
[0070] Thus, the local objective function of the ith client is
[0071] (11)
[0072] where, is the cross-entropy loss function, is a hyper-parameter used to balance the influence of class-aware Gaussian prototype learning.
[0073] Step four: The client uploads the updated local Gaussian prototype to the server.
[0074] Step five: The server performs Gaussian fusion on the uploaded local Gaussian prototype, generates a new Gaussian distribution for each class, and updates the global Gaussian prototype. The calculation process is as follows:
[0075] (12)
[0076] where, denotes the total number of samples of the jth class of all clients. The global covariance of the jth class is calculated as follows: k k
[0077] (13)
[0078] Step six: Feature fusion and training of the trainer.
[0079] Determine whether the set training round or global model convergence is reached; if not, repeat steps one to five; if so, use the local Gaussian prototype and the global Gaussian prototype to generate a virtual feature set and M represents the number of virtual features sampled for each class. This virtual feature set is used to balance the number of samples of each class, especially to enhance the number of samples of the minority class. To make the virtual samples have global information and retain local features, the client performs weighted fusion
[0080] (14)
[0081] where, denotes the local virtual feature, denotes the global virtual feature, is the fusion coefficient used to balance the global and local features. Since the classifier is the main cause of model bias, in order to further alleviate the class imbalance problem, the client fixes the feature extractor and re-trains the classifier using the generated fusion features.
[0082] To demonstrate the effectiveness of the present application, the following experiments are conducted. First, introduce the practical application of medical image data set used in the experiment: dermoscopy data set Fed-ISIC2019. The data set contains a total of 23247 images, divided into six clients, including eight types of skin lesions. The image size is 200x200x3. We use FedAvg, FedProx, FedBN, FedProto, PRR-Imbalance, FedProc, FPL and FedIIC as benchmark methods. We use BACC (balanced accuracy) as the classification index;
[0083] As Fig. 3 As shown in Table 1, the federated image classification method for class imbalance proposed by the present application can obtain a higher classification accuracy, which is 15.49%, 16.90%, 8.72%, 7.00%, 12.83%, 19.84%, 9.25% and 12.89% higher than FedAvg, FedProx, FedBN, FedProto, PRR-Imbalance, FedProc, FPL and FedIIC respectively.
[0084] Table 1 Comparison of BACC (%) values of the proposed method and eight benchmark methods
[0085]
[0086] The above is only the preferred embodiment of the present application, and does not limit the present application in any form. Although the present application has been disclosed as above with the preferred embodiment, it is not intended to limit the present application. Any person skilled in the art can make some changes or modifications to the above disclosed technical content without departing from the scope of the technical solution of the present application, and any simple modification, equivalent replacement and improvement of the above embodiments within the scope of the technical solution of the present application, according to the technical essence of the present application, within the spirit and principles of the present application, are all within the protection scope of the present application.
Claims
1. A federated image classification method for class imbalance, characterized in that, The specific steps include: Step 1, the server initializes a global Gaussian distribution prototype, and the client initializes local model parameters; Step 2, the global Gaussian distribution prototype is distributed to the client; Step 3, the client generates intra-class variance information based on a local image data set using the Gaussian prototype, performs local model training through a configured loss function, and generates a local prototype; where the local Gaussian prototypes The mean and covariance of each class are computed to capture the intra-class variance information and enhance the robustness of the classification model. (1) (2) In formulas (1) and (2), denotes the number of classes of the i first client, k denotes the number of input samples of the j first client, denotes the feature vector of the first input sample of the first client, denotes the mean of the local Gaussian prototype, denotes the covariance matrix, i denotes the number of classes of the k first client, The client introduces a class-aware Gaussian prototype learning, and the loss function is as follows: (3) In formula (3), represents the loss function of the Gaussian prototype learning of the category perception of the i-th client, i , and , is a similarity measurement function used to measure the distance between the input sample and the i-th global prototype; and k is a similarity measurement function used to measure the distance between the input sample and the i-th global prototype. Step 4, the client generates a balanced virtual feature set through random sampling, and re-trains the classifier in combination with the global and local prototypes; The virtual feature set is generated by fusing the global and local Gaussian prototypes through random sampling, ensuring that the fused features have both global information and local features; (4) In formula (4), denotes a local virtual feature, denotes a global virtual feature, is a fusion coefficient for balancing global and local features; Step 5, the client uploads the updated local prototype to the server; Step 6, the server aggregates the local prototypes uploaded by each client to generate a new global Gaussian distribution prototype, which is used for the next round of training; Step 7, determine whether the set training round or model convergence is reached, if yes, the training is ended, if not, repeat steps 2 to 6.
2. The class imbalance oriented federated image classification method according to claim 1, characterized in that, The client classifier re-trains using the fused features to ensure balanced class distribution to reduce the bias of the classifier.
3. The class imbalance oriented federated image classification method according to claim 1, characterized in that, When the server aggregates the uploaded local prototypes, the global mean and covariance of each class of samples are used to generate a new global Gaussian prototype (5) (6) In formulas (5) and (6), represents the total number of samples of all clients k classes.
Citation Information
Patent Citations
Federal learning-based clustering model training system and method
CN115994226A
Federal learning image classification method and system oriented to class imbalance
CN116994054A