Method for improving federal learning robustness in extreme data heterogeneous scene

Abstract

The invention discloses a method for improving federal learning robustness in an extreme data heterogeneous scene. The server maintains an independent global auxiliary data set and an independent verification set, and receives data distribution reported by the client; for each uploaded local model, identifying a missing or sparse target category, and calculating missing severity; screening corresponding samples from the auxiliary set, and adaptively setting a fine tuning learning rate and a round number; during fine adjustment, only the target category output dimension participates in gradient return, and other gradients are zeroed; the model is evaluated on the global verification set, and a compensation credibility score is generated; and finally, by taking the product of the local sample number and the credibility score as a weight, performing weighted aggregation to update the global model. According to the method, model deviation and knowledge interference are effectively relieved through class-aware gradient mask fine tuning and credibility driven aggregation, and convergence and generalization ability are remarkably improved.

Description

TECHNICAL FIELD

[0001] The present application relates to the technical field of artificial intelligence and distributed machine learning, in particular to a method for improving the robustness of federated learning in an extreme data heterogeneous scenario, which is suitable for a federated learning environment where the local data distribution of the clients is highly uneven, and some data categories are severely or completely missing, and can be widely applied to medical health, intelligent finance, Internet of Vehicles, edge computing and other scenarios with high requirements for data privacy protection and extremely heterogeneous data distribution. BACKGROUND

[0002] As a privacy-preserving distributed machine learning paradigm, federated learning allows multiple clients to collaboratively train a global model without sharing raw data, and has been widely applied in medical, financial, and Internet of Things fields which are highly sensitive to data privacy. In a typical federated learning framework, the server coordinates the local model training of each client and iteratively optimizes the global model by aggregating local model updates, such as the classic FedAvg algorithm.

[0003] However, the data of clients in real-world scenarios often presents high heterogeneity, especially when there are significant differences in the types and category distributions of data collected by different clients. In an extreme data heterogeneous scenario (e.g., some clients completely lack samples of a certain category, or there are very few samples of a certain category), the local model is prone to severe category bias during training, causing the model parameter update direction to deviate from the global optimal solution. This bias accumulates after multiple rounds of aggregation, easily causing the global model to oscillate, converge slowly, or even diverge, severely weakening the robustness and practicality of federated learning.

[0004] To alleviate the challenges brought by data heterogeneity, existing technologies have proposed some improvement schemes. For example, FedProx constrains the local update magnitude by introducing a proximal term; some methods also attempt to introduce auxiliary datasets on the server side to fine-tune the global model. However, these methods mostly assume that the server is only used for aggregation or global fine-tuning, and fail to compensate for the specific missing conditions of each client. In particular, when the server has some global data, if the same data or strategy is used to retrain all client models, it not only is inefficient, but also may introduce new bias, and cannot effectively solve the model instability problem caused by extreme heterogeneity.

[0005] Therefore, there is an urgent need for a federated learning mechanism that can dynamically and accurately compensate for the missing features of each client's local data using server-side auxiliary data for optimization, to significantly improve the convergence stability and robustness of the model in an extreme heterogeneous environment while ensuring privacy. SUMMARY

[0006] The application aims to solve the technical problems of model divergence, convergence difficulty and poor robustness of existing federated learning in an extreme data heterogeneous scene, and provides a federated learning method and system capable of dynamically sensing missing features of client data and precisely compensating with server-side auxiliary data.

[0007] To achieve the above-mentioned purpose, the application provides a method for improving the robustness of federated learning in an extreme data heterogeneous scene, comprising the following steps:

[0008] S1, the server maintains a global auxiliary data set with limited scale, complete categories and independent of all client private data, and an independent global validation set; wherein the global auxiliary data set is derived from public standard data sets, generated model synthetic data or user authorized samples, and is used for compensatory fine-tuning of the client model, and the global validation set is used to evaluate the generalization performance of the fine-tuned model; the server simultaneously receives the local data distribution information reported by each client;

[0009] S2, the server distributes the global model parameters of the current round to all clients;

[0010] S3, each client performs local training on the global model parameters based on its local highly heterogeneous data set, obtains a local updated model, and uploads the local model to the server;

[0011] S4, after the server receives the local models uploaded by all clients, for each client's local model, the following operations are performed:

[0012] (1) according to the local data distribution information reported by the client, determine the target data category whose data is missing or sample size is lower than the preset threshold, and calculate the missing severity index of the client;

[0013] (2) from the global auxiliary data set, filter out a sample subset belonging to the target data category;

[0014] (3) based on the missing severity index, adaptively determine the fine-tuning learning rate and the fine-tuning round number;

[0015] (4) using the sample subset, compensatory fine-tuning training is performed on the local model of the client, wherein in the fine-tuning process, only the model output dimension corresponding to the target data category is allowed to participate in gradient back propagation, and the gradients of the remaining categories are set to zero to avoid interfering with the knowledge learned by the client on non-missing categories;

[0016] (5) the client model after compensatory fine-tuning is used for inference on the global validation set, the average loss is calculated, and a compensation credibility score is generated based on the loss;

[0017] S5. The server performs weighted aggregation on all compensated and optimized client models based on the product of the number of local samples of each client and its corresponding compensation confidence score, and generates a new global model.

[0018] S6. Use the new global model as the global model for the next round, return to step S2 to continue iterating until the preset convergence condition is met.

[0019] Furthermore, in step S1, the local data distribution information reported by the client includes:

[0020] The sample size statistics for each data category, the histogram of sample proportion distribution, or Boolean markers indicating the existence of a specific data category.

[0021] Furthermore, in step S4(1), the missing severity index is calculated as follows:

[0022]

[0023] Where C is the total number of categories; N missing The number of data categories in client i with zero sample size (satisfying 0 ≤ N) missing ≤C-1); Let be the number of samples of category c in client i; α∈[0,1] is a preset smoothing weight coefficient used to balance the contributions of the two heterogeneous factors of "complete category missing" and "centralized data distribution"; the missing severity index S i The value range is [0, 1]. The larger the value, the more severe the client data is missing or the distribution is skewed.

[0024] Furthermore, in step S4(3), the learning rate η is fine-tuned. i for:

[0025] η i =η base ·S i (2)

[0026] Where η base The preset baseline learning rate;

[0027] Fine-tuning the number of wheels:

[0028]

[0029] Among them, T base The base number of rounds; T max This represents the maximum number of allowable fine-tuning rounds. This indicates rounding down to the nearest integer.

[0030] Furthermore, in step S4(4), the specific implementation of setting the gradient to zero is as follows:

[0031] Define a category mask vector m∈{0,1} C ,in:

[0032]

[0033] During the backpropagation process of compensation fine-tuning, the loss function is calculated for the final unnormalized output layer (i.e., the logits layer) of the classification model. The gradient vector of the output y of this layer Then, the gradient vector is multiplied element-wise with the class mask vector m to generate the masked gradient:

[0034]

[0035] Where ⊙ represents element-wise multiplication;

[0036] After using the mask, the gradient continues to backpropagate in order to update only the model parameters related to the target class.

[0037] Furthermore, in step S4(5), the compensation credibility score is calculated as follows:

[0038] The compensated and finely tuned client model is used for inference on the global validation set, and its average cross-entropy loss is calculated. And generate the compensation credibility score according to the following formula:

[0039]

[0040] in, is the average loss of the fine-tuned model on the global validation set; exp(·) is the natural exponential function; the range of the compensation credibility score is (0, 1], and the lower the loss, the stronger the model's generalization ability and the higher the credibility.

[0041] Furthermore, the weighted aggregation method in step S5 is as follows:

[0042] The aggregate weight w of the i-th client i for:

[0043]

[0044] The global model parameters are updated as follows:

[0045]

[0046] Among them, |D i | represents the total number of local samples for client i; θ′ i The model parameters are optimized after compensation; N is the total number of clients participating in this round of aggregation.

[0047] Accordingly, the present invention also provides a system for implementing the above method, comprising:

[0048] (1) Client cluster, used to perform local model training and upload local model and data distribution information;

[0049] (2) A server, which is communicatively connected to the client cluster, the server comprising:

[0050] The data management module is used to maintain the global auxiliary dataset and the global validation set, and to receive data distribution information from the client.

[0051] The model distribution module is used to distribute global models to clients;

[0052] The compensation training module is used to: a) determine the target missing category and missing severity based on the client data distribution; b) select samples of the corresponding category from the global auxiliary dataset; c) adaptively set fine-tuning hyperparameters based on missing severity; d) perform compensation fine-tuning under gradient mask constraints; and e) evaluate the model on the global validation set and generate a compensation confidence score.

[0053] The aggregation module is used to perform weighted aggregation of the compensated model based on the product of the number of local samples and the confidence score;

[0054] The compensation training module is the core component that distinguishes this system from the prior art, and its workflow implements the method described in any one of claims 1 to 7.

[0055] Furthermore, the present invention also covers electronic devices that perform the method and computer-readable storage media storing corresponding programs.

[0056] The present invention has the following advantages and effects compared with the prior art:

[0057] (1) Significantly improve robustness: By specifically compensating for missing category information, model divergence is effectively alleviated and training stability in extreme heterogeneous environments is improved.

[0058] (2) Accelerate convergence speed: reduce invalid updates caused by client bias, and enable the global model to approach the optimal solution faster;

[0059] (3) Maintain local characteristics: Adaptive learning rate and fine-tuning rounds are used to avoid server data dominating model updates and to retain the characteristics of local client data;

[0060] (4) Strong compatibility and flexible deployment: No need to modify the client architecture, only the server side needs to add a compensation training module, which is easy to integrate into the existing federated learning framework;

[0061] (5) Privacy compliance: Server auxiliary data can be public or synthetic data, does not rely on original user data, and complies with privacy protection requirements. Attached Figure Description

[0062] Figure 1 This is a schematic diagram of a system architecture for improving the robustness of federated learning in extreme data heterogeneous scenarios, provided by an embodiment of the present invention.

[0063] Figure 2 A flowchart of the federated learning method provided in an embodiment of the present invention. Detailed Implementation

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

[0065] Example

[0066] The specific embodiments of the present invention will now be described in detail with reference to the accompanying drawings.

[0067] I. System Architecture

[0068] like Figure 1 As shown, the system of the present invention includes a client cluster (101) and a server (102).

[0069] The client cluster (101) contains multiple clients (e.g., client 1, client 2, ..., client N), each client holding a highly heterogeneous local dataset (e.g., client 1 is missing "category A", client 2 is missing "category C"). Each client is configured with a local training module to perform model training and report local data distribution information (e.g., the proportion of samples in each category or whether a label exists) to the server.

[0070] Server (102) includes:

[0071] Data Management Module (103): Maintains a global auxiliary dataset (104) containing samples of the complete categories, which may be derived from public datasets, de-identified historical data, synthetic data from generative models, or authorized centralized data subsets;

[0072] Model distribution module (105): Distributes the global model parameters for the current round to all clients;

[0073] Compensation training module (106): It is used to identify the missing or scarce target data categories based on the data distribution information reported by the client, and extract the sample subset of the corresponding category from the global auxiliary dataset. It performs one or more rounds of fine-tuning on the local model uploaded by the client with a learning rate lower than that of local training to generate a compensation optimization model.

[0074] Aggregation module (107): performs weighted aggregation on all compensation optimization models to generate a new global model.

[0075] II. Method and Flow

[0076] This example uses an image classification task, with data categories of {cat, dog, bird, fish, horse}, a total of 5 categories. Assumptions:

[0077] Client 1's local data does not contain "bird" or "fish";

[0078] The number of "horse" samples in the local data of client 2 is extremely small;

[0079] The global auxiliary dataset maintained by the server contains balanced samples for each class.

[0080] like Figure 2 As shown, the specific steps are as follows:

[0081] Step S201: The server initializes the global model (such as ResNet-18 or a lightweight CNN) and distributes the model parameters to all clients.

[0082] Step S202: Each client performs several rounds of local training using local data (e.g., using a stochastic gradient descent optimizer with a learning rate of 0.01). After training, the data is uploaded to the server.

[0083] 1) Locally updated model parameters;

[0084] 2) Local data distribution information (e.g., [cat: 500, dog: 400, bird: 0, fish: 0, horse: 300]).

[0085] Step S203: After receiving all client models, the server performs compensatory fine-tuning for each client:

[0086] For client 1:

[0087] 1) Identify "bird" and "fish" as missing categories (sample count is 0 or below the preset threshold);

[0088] 2) Calculate the missing severity index S1 using formula (1):

[0089]

[0090] Where C = 5 is the total number of categories; N missing =2 represents the number of data categories with zero samples in client 1; α represents the number of samples of category c in client 1; α∈[0,1] is the preset smoothing weight coefficient;

[0091] 3) Calculate the fine-tuning learning rate η1 using formula (2):

[0092] η1=η base ·S1 (2)

[0093] Where η base The preset baseline learning rate;

[0094] 4) Calculate the number of fine-tuning wheels T1 using formula (3):

[0095]

[0096] Among them, T base The base number of rounds; T max This represents the maximum number of allowable fine-tuning rounds. Indicates rounding down;

[0097] 5) Select all “bird” and “fish” samples from the global auxiliary dataset to form a sample subset for fine-tuning; perform T1 rounds of fine-tuning on the client model with a learning rate η1;

[0098] During fine-tuning, when backpropagating the loss function, the gradient of the output dimension corresponding to non-target categories is set to zero, and only the parameters related to the target category are updated;

[0099] For client 2:

[0100] 1) Identify "horse" as a rare category and extract corresponding samples accordingly;

[0101] 2) Calculate S2, η2 and T2 using formulas (1), (2) and (3) respectively, and fine-tune T2 in rounds with learning rate η2; and during the fine-tuning process, when backpropagating the loss function, set the gradient of the output dimension corresponding to the non-target category to zero, and only update the parameters related to the target category.

[0102] Step S204: The server performs a weighted aggregation of all compensated and optimized client models. The weights are set based on the total number of local samples for each client and the compensation confidence score, generating a new round of global models;

[0103] The calculation method for the compensation credibility score is as follows:

[0104] For client 1:

[0105] The compensation credibility score Trust1 is:

[0106]

[0107] in exp(·) represents the average loss of the client 1 model on the global validation set after fine-tuning; exp(·) is the natural exponential function.

[0108] The aggregate weight of client 1 is:

[0109]

[0110] Where |D1| is the total number of local samples for client 1; |D j | and Trust j These represent the total number of local samples and the compensation credibility score for client j, respectively.

[0111] For client 2:

[0112] The compensation credibility score Trust2 and the aggregate weight w2 of client 2 are calculated by formulas (6) and (7), respectively.

[0113] Then, update the global model parameters as follows:

[0114]

[0115] Where θ′ i The model parameters are optimized and compensated for client i (1 or 2).

[0116] Step S205: Determine whether the preset convergence conditions are met, such as reaching the maximum number of iterations or the change in loss of the global model on the server validation set being less than a threshold. If not met, return to step S201 to continue the next round of training; otherwise, terminate the training process.

[0117] III. Technical Effects Description

[0118] In this implementation, because the server can dynamically identify the data gaps of each client and only use corresponding minority class samples from the global auxiliary dataset for compensation and fine-tuning, it can effectively alleviate the deviation in model update direction caused by local data bias. This mechanism avoids the aggregation oscillation or divergence problems caused by excessive deviation of client models in extreme heterogeneous scenarios in traditional federated learning, and is expected to significantly improve the stability of the training process and the robustness of the global model.

[0119] Furthermore, by introducing a gradient masking mechanism, interference with non-missing category knowledge is further prevented; and by using dynamic weighted aggregation driven by credibility scores, the negative impact of low-quality compensation models on the global model is suppressed, thereby achieving more stable and robust collaborative training under extreme heterogeneity.

[0120] IV. Variations and Extensions

[0121] Global auxiliary datasets can be synthesized using generative adversarial networks or diffusion models to further reduce the reliance on real centralized data.

[0122] The data distribution information reported by the client can be subjected to differential privacy perturbation before uploading, thereby enhancing privacy protection;

[0123] The compensation training module can be combined with knowledge distillation technology to use a server-side teacher model to guide the fine-tuning process and improve compensation efficiency.

[0124] The above embodiments are preferred embodiments of the present invention, but the embodiments of the present invention are not limited to the above embodiments. Any changes, modifications, substitutions, combinations, or simplifications made without departing from the spirit and principle of the present invention shall be considered equivalent substitutions and shall be included within the protection scope of the present invention.

Claims

1. A method for improving the robustness of federated learning in extreme data heterogeneous scenarios, characterized in that, Includes the following steps: S1. The server maintains a global auxiliary dataset that is limited in size, complete in categories, and independent of all client private data, as well as an independent global validation set. The global auxiliary dataset is derived from public standard datasets, generated model synthetic data, or user-authorized samples, and is used to perform compensatory fine-tuning on client models. The global validation set is used to evaluate the generalization performance of the fine-tuned model. The server also receives local data distribution information reported by each client. S2. The server distributes the global model parameters for the current round to all clients; S3. Each client trains the global model parameters locally based on its highly heterogeneous local dataset to obtain the locally updated model, and then uploads the local model to the server. S4. After receiving all the local models uploaded by the clients, the server performs the following operations for each client's local model: (1) Based on the local data distribution information reported by the client, determine the target data category with missing data or sample size below the preset threshold, and calculate the client's missing severity index. (2) Select a subset of samples belonging to the target data category from the global auxiliary dataset; (3) Based on the missing severity index, adaptively determine the fine-tuning learning rate and the number of fine-tuning rounds; (4) Using the sample subset, the local model of the client is trained for compensatory fine-tuning. During the fine-tuning process, only the model output dimension corresponding to the target data category is allowed to participate in gradient backpropagation, and the gradients of other categories are set to zero to avoid interfering with the knowledge learned by the client on non-missing categories. (5) Perform inference on the global validation set using the compensated and fine-tuned client model, calculate its average loss, and generate a compensated credibility score based on the loss. S5. The server performs weighted aggregation on all compensated and optimized client models based on the product of the number of local samples of each client and its corresponding compensation confidence score, and generates a new global model. S6. Use the new global model as the global model for the next round, return to step S2 to continue iterating until the preset convergence condition is met.

2. The method for improving the robustness of federated learning in extreme data heterogeneous scenarios according to claim 1, characterized in that, In step S1, the local data distribution information reported by the client includes: The sample size statistics for each data category, the histogram of sample proportion distribution, or Boolean markers indicating the existence of a specific data category.

3. The method for improving the robustness of federated learning in extreme data heterogeneous scenarios according to claim 1, characterized in that, In step S4(1), the missing severity index is calculated as follows: Where C is the total number of categories; N missing The number of data categories in client i with zero sample size (satisfying 0 ≤ N) missing ≤C-1); Let be the number of samples of category c in client i; α∈[0,1] is a preset smoothing weight coefficient used to balance the contributions of the two heterogeneous factors of "complete category missing" and "centralized data distribution"; the missing severity index S i The value range is [0, 1]. The larger the value, the more severe the client data is missing or the distribution is skewed.

4. The method for improving the robustness of federated learning in extreme data heterogeneous scenarios according to claim 1, characterized in that, In step S4(3), the learning rate η is fine-tuned. i for: or i =the base ·S i (2) Where η base The preset baseline learning rate; Fine-tuning the number of wheels: T i =max(1,min(T max ,[T base ·S i ])) (3) Among them, T base The base number of rounds; T max This represents the maximum number of allowable fine-tuning rounds. This indicates rounding down to the nearest integer.

5. The method for improving the robustness of federated learning in extreme data heterogeneous scenarios according to claim 1, characterized in that, In step S4(4), the specific implementation of setting the gradient to zero is as follows: Define a category mask vector m∈{0,1} C ,in: During the backpropagation process of compensation fine-tuning, the loss function is calculated for the final unnormalized output layer (i.e., the logits layer) of the classification model. The gradient vector of the output y of this layer Then, the gradient vector is multiplied element-wise with the class mask vector m to generate the masked gradient: Where ⊙ represents element-wise multiplication; After using the mask, the gradient continues to backpropagate in order to update only the model parameters related to the target class.

6. The method for improving the robustness of federated learning in extreme data heterogeneous scenarios according to claim 1, characterized in that, In step S4(5), the compensation credibility score is calculated as follows: The compensated and finely tuned client model is used for inference on the global validation set, and its average cross-entropy loss is calculated. And generate the compensation credibility score according to the following formula: in, is the average loss of the fine-tuned model on the global validation set; exp(·) is the natural exponential function; the range of the compensation credibility score is (0, 1], and the lower the loss, the stronger the model's generalization ability and the higher the credibility.

7. A method for improving the robustness of federated learning in extreme data heterogeneous scenarios according to claim 1, characterized in that, The weighted aggregation method in step S5 is as follows: The aggregate weight w of the i-th client i for: The global model parameters are updated as follows: Among them, |D i | represents the total number of local samples for client i; θ′ i The model parameters are optimized after compensation; N is the total number of clients participating in this round of aggregation.

8. A system for improving the robustness of federated learning in extreme data heterogeneous scenarios, characterized in that, include: (1) Client cluster, used to perform local model training and upload local model and data distribution information; (2) A server, which is communicatively connected to the client cluster, the server comprising: The data management module is used to maintain the global auxiliary dataset and the global validation set, and to receive data distribution information from the client. The model distribution module is used to distribute global models to clients; The compensation training module is used to: a) determine the target missing category and missing severity based on the client data distribution; b) select samples of the corresponding category from the global auxiliary dataset; c) adaptively set fine-tuning hyperparameters based on missing severity; d) perform compensation fine-tuning under gradient mask constraints; and e) evaluate the model on the global validation set and generate a compensation confidence score. The aggregation module is used to perform weighted aggregation of the compensated model based on the product of the number of local samples and the confidence score; The compensation training module is the core component that distinguishes this system from the prior art, and its workflow implements the method described in any one of claims 1 to 7.

9. An electronic device comprising a memory, a processor, and a computer program stored in the memory and executable on the processor, characterized in that, When the processor executes the program, it implements the steps of the method as described in any one of claims 1 to 7.

10. A computer-readable storage medium having a computer program stored thereon, characterized in that, When the program is executed by the processor, it implements the steps of the method as described in any one of claims 1 to 7.

Images