A federated learning optimization method and system for imbalanced datasets

Abstract

This invention provides a federated learning optimization method and system for imbalanced datasets. In each round of communication, the server sends a global model, data distribution divergence, and a global set of difficult categories to the client. The client adjusts its learning rate accordingly and adds weights for difficult categories to the loss function. After completing local training, the client uploads its local model and local state vector. The server aggregates the state vectors of each client, calculates the new set of global difficult categories and the data distribution divergence of each client, and adjusts the model aggregation weights based on factors such as whether the client's divergence is abnormal and the number of samples in the difficult categories. When aggregating classification layer parameters, the average loss of each client in each category is further combined to achieve refined model optimization for class imbalance and data distribution differences.

Description

Technical Field

[0001] This application belongs to the field of artificial intelligence, and in particular relates to a federated learning optimization method and system for imbalanced datasets. Background Technology

[0002] Federated learning, as an emerging distributed machine learning paradigm, collaborates with a central server to train a global model, thereby protecting data privacy. Classic federated learning algorithms, such as federated averaging, work by having the server distribute the global model to each client. Each client trains its model using local data, uploads updates, and the server then performs a weighted average based on the amount of data from each client. In real-world applications, the data distribution across clients is often heterogeneous, i.e., non-independent and identically distributed, with class imbalance or skewed label distribution being typical manifestations. This data heterogeneity can lead to conflicts in the optimization objectives of the local models on different clients, causing client drift and severely impacting the convergence speed and performance of the global model. Corrections can be made during the local training phase, for example, by using regularization terms in the local loss function to limit the difference between the local and global models, thus mitigating client drift. The impact of heterogeneous data can also be mitigated by adjusting the aggregation weights of each client, such as reducing the weights of clients with higher local losses or larger model update magnitudes. However, most of these methods lack in-depth diagnostics of data heterogeneity. These methods typically fail to identify which data categories pose common challenges to model learning, and struggle to represent and identify which clients are statistical outliers due to the uniqueness of their data distribution. Furthermore, they fail to perform more targeted, fine-grained aggregation of classification layer parameters directly related to the categories. Therefore, a federated learning method is urgently needed that can comprehensively perceive the global data distribution and learning status, and collaboratively optimize local training and server aggregation. Summary of the Invention

[0003] To address the problem that existing federated learning methods struggle to fully perceive global data distribution and learning status, and to collaboratively optimize local training and server aggregation.

[0004] In the first aspect, the present invention proposes a federated learning optimization method for imbalanced datasets, comprising the following steps:

[0005] In any communication round of federated learning, the client receives the global model from the server, as well as the data distribution divergence and global difficulty category set calculated by the server in the previous communication round. The client adjusts the learning rate of local training based on the data distribution divergence and increases the weight of the corresponding category in the loss function based on the global difficulty category set for local training. After local training is completed, the client uploads the local model and the local state vector containing the number of samples for each data category and the average loss to the server.

[0006] The server receives and aggregates the local state vectors from each client to obtain a global state vector that includes the global data distribution and the global average loss. Based on the global state vector, the server determines the global set of difficulty categories and the data distribution divergence of each client for the next communication round, and calculates the divergence of each client's data distribution relative to the global data distribution.

[0007] The server determines the aggregation weight of each client, wherein the aggregation weight of a client is the ratio of the total number of samples owned by the client to the total number of samples of all clients participating in the aggregation; the server then performs weighted aggregation on the local models uploaded by each client based on the aggregation weight to generate a new round of global models.

[0008] In a second aspect, this invention proposes a federated learning optimization system for imbalanced datasets, comprising the following modules:

[0009] The adjustment module is used in any communication round of federated learning. The client receives the global model from the server, as well as the data distribution divergence and global difficulty category set calculated by the server in the previous communication round. The client adjusts the learning rate of local training based on the data distribution divergence and increases the weight of the corresponding category in the loss function based on the global difficulty category set for local training. After local training is completed, the client uploads the local model and the local state vector containing the number of samples for each data category and the average loss to the server.

[0010] The calculation module is used by the server to receive and aggregate the local state vectors of each client to obtain a global state vector containing the global data distribution and the global average loss. Based on the global state vector, the server determines the global difficulty category set and the data distribution divergence of each client for the next communication round, and calculates the divergence of the data distribution of each client relative to the global data distribution.

[0011] The generation module is used by the server to determine the aggregation weight of each client, wherein the aggregation weight of the client is the ratio of the total number of samples owned by the client to the total number of samples of all clients participating in the aggregation; the server then performs weighted aggregation on the local models uploaded by each client based on the aggregation weight to generate a new round of global models.

[0012] This invention improves the reliability and focus of local training on the client side by utilizing a local training adjustment mechanism that combines data distribution divergence and global difficulty categories. On the server side, clients with large differences in data distribution are suppressed to avoid the negative impact of local outliers on the global model, while clients containing important samples of the difficulty category are incentivized, accelerating the learning of the difficulty category. At the classification layer of the aggregation model, the loss performance of each client on specific categories is further incorporated, allowing the global model to fully leverage the expertise of each client in its advantageous categories. This improves the convergence stability and classification accuracy of the global model in imbalanced data scenarios, enhances the model's ability to identify minority or difficulty categories, and strengthens the overall performance and generalization ability of the model. Attached Figure Description

[0013] Figure 1 A flowchart of the first embodiment;

[0014] Figure 2 A visual diagram illustrating the divergence of the data distribution;

[0015] Figure 3 This is a schematic diagram of the server-side aggregation weight adjustment process;

[0016] Figure 4 This is a schematic diagram of category-aware aggregation of classification layer parameters. Detailed Implementation

[0017] The technical solutions of the embodiments of this application will be clearly and completely described below with reference to the accompanying drawings. Obviously, the described embodiments are only some embodiments of this application, and not all embodiments. Based on the embodiments of this application, all other embodiments obtained by those skilled in the art without creative effort are within the scope of protection of this application.

[0018] Example 1

[0019] A federated learning optimization method for imbalanced datasets, such as Figure 1 This includes the following steps:

[0020] S1, In any communication round of federated learning, the client receives the global model from the server, as well as the data distribution divergence and global difficulty category set calculated by the server in the previous communication round; the client adjusts the learning rate of local training based on the data distribution divergence, and increases the weight of the corresponding category in the loss function based on the global difficulty category set for local training; after local training is completed, the client uploads the local model and the local state vector containing the number of samples and the average loss for each data category to the server;

[0021] Specifically, the client receives serialized global model parameters, a floating-point value representing the data distribution divergence, and a list containing difficulty category IDs from the server through a network interface such as HTTP or gRPC; the client then loads the received model parameters into its local neural network structure.

[0022] The client sets the local learning rate to a base learning rate multiplied by one minus the normalized data distribution divergence. For example, if the divergence is 0.8 and the base learning rate is 0.01, the adjusted learning rate is 0.002. When calculating the cross-entropy loss, for a training sample, if the true label belongs to one of the global hard-class categories, the loss value is multiplied by a preset weight coefficient greater than 1, such as 1.5; otherwise, the weight coefficient is 1. The learning rate adjustment is used to stabilize the local training process, and the aggregated weight penalty is used to reduce the impact of statistical bias on the global model.

[0023] The client serializes the locally trained and updated model parameters and creates a data structure such as a dictionary. The key of the dictionary is the data category ID, and the value is a tuple containing the total number of samples in that category and the average loss value calculated on those samples. The client sends the serialized model parameters and the state vector dictionary to the server through the network interface.

[0024] To adjust the learning intensity based on the heterogeneity of client data, in one optional embodiment, the client adjusts the learning rate for local training based on the data distribution divergence, including:

[0025] The client's local learning rate η k Based on the data distribution divergence D k The formula for adjustment is: ;

[0026] Where η0 is the initial learning rate and β is the preset decay coefficient.

[0027] At the beginning of each round or a predetermined number of rounds of federated learning, the data distribution divergence D is calculated for each client k. k The divergence value represents the difference between the local data distribution on the client and the global data distribution. Simultaneously, a global initial learning rate η0 and a decay coefficient β are pre-set, for example, η0 is set to 0.1 and β to 1.0.

[0028] The obtained divergence value D k Substitute the values ​​into the adjustment formula to calculate the local learning rate η for each client. k For example, client A's data distribution is very close to the global distribution, with a divergence value D. A If the value is 0.05, then the learning rate is adjusted to η. A≈0.095. However, client B's data is highly skewed, with a divergence value of D. B If the value is 0.9, then the learning rate is adjusted to η. B 0.041. Client A, whose data is representative, will update with larger step sizes, while client B, whose data may be biased, will update cautiously with smaller step sizes. This protects the model from extreme data distributions while ensuring training stability and efficiency. This process does not change the structure of the model used by the client, such as a convolutional neural network; it only adjusts the training parameters at the optimizer level.

[0029] The convolutional neural network (CNN) model typically includes an input layer, several alternating stacked convolutional and pooling layers, several fully connected layers, and an output layer. The convolutional layers extract local features of the image, such as edges or textures, using learnable filters. The pooling layers downsample the feature maps to reduce dimensionality and enhance the translation invariance of the features. The fully connected layers are responsible for integrating and classifying the extracted features. In a federated learning scenario, the network's training set is distributed across various clients. Each client holds a portion of local data, such as image samples and corresponding class labels; this data is not uploaded to a central server. The training process is iterative. In each round of communication, the client downloads the current global model, uses its local data, and performs multiple rounds of training using backpropagation and a stochastic gradient descent optimizer to minimize the cross-entropy loss function. The updated model parameters are then uploaded to the server for aggregation, generating a new global model. The network's input is an image tensor, such as a 32×32×3 color image. The output is a probability distribution vector, where the dimension of the vector equals the total number of classes, and each element represents the predicted probability that the input image belongs to the corresponding class.

[0030] To enable the model to focus more on categories that the global model generally struggles to classify correctly during local training, in an optional embodiment, the step of adding weights for the corresponding categories in the loss function based on the global set of difficult categories for local training includes:

[0031] For any category belonging to the global difficulty category set, the client will multiply the weight in the loss function by a preset gain factor.

[0032] Suppose that in an image classification task, cats and birds are identified as difficult categories. A gain factor greater than 1, such as 1.5, is set. This gain factor amplifies the contribution of the difficult category to the loss calculation, forcing the model to allocate more resources to learning the features of that category.

[0033] During local training on each client, the client's model calculates the loss for each sample in the batch of training data. Before calculating the total loss, the true label of the sample is checked. If the sample label does not belong to the hard category set, such as car or ship, the loss is calculated normally. If the sample label belongs to the hard category set, such as cat, the calculated loss value for that sample is multiplied by a preset gain factor of 1.5. The total loss for the entire batch is the sum of the adjusted losses for all samples. This method modifies the calculation process of the loss function but does not change the architecture of the model itself; for example, the network layers and connections of a ResNet classification model remain unchanged.

[0034] S2, the server receives and aggregates the local state vectors of each client to obtain a global state vector containing the global data distribution and the global average loss; based on the global state vector, the server determines the global difficulty category set and the data distribution divergence of each client for the next communication round, and calculates the divergence of each client's data distribution relative to the global data distribution.

[0035] The server initializes a global state vector; it iterates through the local state vectors uploaded by all clients. For each data category, the server sums the number of samples from all clients in that category to obtain the total number of global samples for that category. At the same time, the server multiplies the average loss of each client in that category by the corresponding number of samples, sums them up, and divides them by the total number of global samples for that category to obtain the global average loss for that category. The resulting global state vector records the total number of global samples and the global average loss for each category.

[0036] The server extracts the global average loss value for all categories from the global state vector and calculates the 75th percentile of the loss value. All category IDs with a global average loss value higher than this percentile form a new global set of difficult categories. For each client, the server constructs a local data distribution vector P, where each element is the number of samples in a certain category for that client divided by the total number of samples. The server then constructs a global data distribution vector Q, where each element is the number of samples in a certain category globally divided by the total number of samples globally. The server calculates the Jensen-Shannon divergence between P and Q as the data distribution divergence for that client; optionally, the data distribution divergence is a normalized divergence metric. The data distribution vector P for each client k is determined. k and global data distribution vector P G In a classification task with C classes, the vector is C-dimensional, where each element represents the proportion of samples in the corresponding class out of the total number of samples. For example, for client k, the distribution vector P... k =[n k,1 / N k ,n k,2 / N k,…,n k,C / N k Where n k,c N is the number of samples in category c. k That is the total number of samples.

[0037] Calculate D using the Jensen-Shannon divergence formula k This formula calculates a mixture distribution M=(P k +P G ) / 2. Calculate P respectively. k The Körbek-Leibler divergence D relative to M KL (P k ||M) and P G The Körbek-Leibler divergence D relative to M KL (P G ||M). The Jensen-Shannon divergence is the arithmetic mean of the two values. The Jensen-Shannon divergence is symmetric and has a finite range; a larger value indicates a greater difference between the data distribution of client k and the global distribution, and vice versa. Figure 2 .

[0038] On the server side, the average loss of all clients on each class is collected and aggregated. Taking a 10-class classification task as an example, for class c, the server will obtain a list L containing the average loss of all K clients on that class. c ={L 1,c ,L 2,c ,…,L K,c}

[0039] For list L c Sort all loss values ​​in ascending order, and calculate the 75th quantile of the sorted list, i.e., find the value located at the 75th percentile of the list. For example, if there are 100 clients, the loss value of the 75th client after sorting is the loss threshold T for that category. L The threshold T L This is then used to determine whether category c is a globally difficult category. Setting it to the 75th percentile means that if a category shows a high loss for more than a quarter of the clients, i.e., poor learning performance, then it is very likely to be identified as a globally difficult category. This setting takes into account the performance of most clients while giving sufficient attention to clients with poor performance.

[0040] S3, the server determines the aggregation weight of each client, wherein the aggregation weight of a client is the ratio of the total number of samples owned by the client to the total number of samples of all clients participating in the aggregation; the server then performs weighted aggregation on the local models uploaded by each client based on the aggregation weight to generate a new round of global models. When the data distribution divergence of a client exceeds an outlier threshold set based on the median of the data distribution divergence of all clients, a penalty is applied to the weight; when the number of global hard category samples held by a client exceeds a preset threshold, a gain is applied to the weight; and when aggregating the classification layer parameters of the model, for each data category, the adjusted aggregation weight is multiplied by the inverse of the average loss of the client on that category to obtain the weight used to aggregate the parameters corresponding to that category.

[0041] Specifically, the server calculates an initial aggregation weight for each client, equal to the total number of local samples divided by the total number of samples from all clients. The server calculates the median of the divergence of the data distribution across all clients and multiplies this median by a constant, such as 1.5, to obtain the outlier threshold. If a client's divergence exceeds this threshold, the initial weight is multiplied by a penalty factor less than 1, such as 0.7. The server counts the total number of samples belonging to the global difficulty category set held by that client. If this total exceeds a preset value, such as 50, the client's weight is multiplied by a gain factor greater than 1, such as 1.2. After the penalty and gain steps, the adjusted aggregation weight is obtained, as shown below. Figure 3 .

[0042] For the fully connected classification layer of the model, parameters are aggregated according to category; for example, when aggregating parameters corresponding to category C, for client k, the server takes the aggregated weight W adjusted by the client. k And obtain the client's average loss L on category C from the uploaded local state vector. k_c The server calculates the weight as equal to W. k Multiply by L k_c The reciprocal of the order; the server uses weights to calculate a weighted average of the parameters corresponding to category C in all client-uploaded classification layers, generating the parameters for category C in the new global model; this process is repeated for all categories, such as... Figure 4 Preferably, the reciprocal of the client's average loss in said category is numerically stable calculated using the following formula: 1 / (L k,c +ε), where L k,c Let L be the average loss for client k on class c, and ε be a preset positive constant to prevent the denominator from being zero. In federated learning frameworks, low loss usually indicates high specialization. The average loss L is optimal when the model performs perfectly on class c for client k. k,cThe value could be 0, and calculating the reciprocal could lead to a division by zero error, causing the program to crash or produce an invalid infinity value. To solve this problem, a very small preset positive constant ε is used, such as 10. -8 When calculating the professionalism score, first consider the average loss L. k,c Add this to ε and then take the reciprocal. For example, if client A has an average loss of 0.01 in category c, their expertise score is approximately 100. If client B performs perfectly in category c, with a loss of 0, their expertise score is a very large but finite value of 10. 8 This operation accurately reflects the high level of professionalism of client B while avoiding calculation errors.

[0043] In an optional embodiment, the step of penalizing the weights when the data distribution divergence of a client exceeds an outlier threshold set based on the median of the data distribution divergence of all clients includes:

[0044] Outlier threshold T D Based on the data distribution divergence of all clients, D={D1,D2,...,D...} K The median (median(D)) and median absolute deviation (MAD(D)) of the data are determined using the following formula: T D =median(D) + c × MAD(D); where c is a preset constant; for the data distribution divergence D k >T D The client multiplies the client's aggregate weight by a preset penalty factor.

[0045] This method is used to identify and mitigate the impact of clients with extremely anomalous data distributions during model aggregation. On the server side, the data distribution divergence values ​​of all clients participating in the current training round are collected, forming a set D. The median (median(D)) of this set is calculated, representing the central tendency of the divergence values. The median absolute deviation (MAD(D)) is calculated, which is the median of the set of absolute differences between each divergence value and the median; it is a stable measure of dispersion.

[0046] Using a preset constant c, for example, c is set to 3, the outlier threshold T is calculated using a formula. D The threshold can be adjusted based on the divergence distribution of all clients in the current round. Before performing global model aggregation, the server checks the divergence value D of each client k. k If D k Greater than T D If a client is considered an outlier, its default aggregation weight, such as a weight based on data volume, will be multiplied by a preset penalty factor less than 1, such as 0.5. This reduces the potential negative impact of clients with extremely uneven data distribution on the global model.

[0047] To enhance the contributions of clients that have potentially high value in learning global difficulty categories, in an optional embodiment, applying a gain to the weights when the number of global difficulty category samples held by a client exceeds a preset threshold includes:

[0048] When the total number of samples belonging to the global difficulty category set held by the client exceeds a preset sample number threshold, the client's aggregation weight is multiplied by a preset gain factor.

[0049] It is necessary to determine the global set of hard-point categories, for example, identifying cats and dogs as hard-point categories. Simultaneously, a sample size threshold, such as 100, and a gain factor greater than 1, such as 1.2, are set. Before each round of model aggregation, the server counts the total number of samples belonging to the hard-point category set possessed by each client k. For example, client A has 70 cat samples and 50 dog samples, with a total of 120 hard-point samples. Since 120 is greater than the preset threshold of 100, client A is considered to play a significant role in improving the model's performance on hard-point categories. Therefore, when calculating the global model update, the aggregation weight of client A will be multiplied by the gain factor 1.2. For client B, whose total number of hard-point samples is less than 100, the aggregation weight remains unchanged. In this way, the federated learning system can rely more heavily on clients with more hard-point data, thereby accelerating the overcoming of the global model's performance bottleneck.

[0050] Example 2

[0051] A federated learning optimization system for imbalanced datasets includes the following modules:

[0052] The adjustment module is used in any communication round of federated learning. The client receives the global model from the server, as well as the data distribution divergence and global difficulty category set calculated by the server in the previous communication round. The client adjusts the learning rate of local training based on the data distribution divergence and increases the weight of the corresponding category in the loss function based on the global difficulty category set for local training. After local training is completed, the client uploads the local model and the local state vector containing the number of samples for each data category and the average loss to the server.

[0053] The calculation module is used by the server to receive and aggregate the local state vectors of each client to obtain a global state vector containing the global data distribution and the global average loss. Based on the global state vector, the server determines the global difficulty category set and the data distribution divergence of each client for the next communication round, and calculates the divergence of the data distribution of each client relative to the global data distribution.

[0054] The generation module is used by the server to determine the aggregation weight of each client, wherein the aggregation weight of the client is the ratio of the total number of samples owned by the client to the total number of samples of all clients participating in the aggregation; the server then performs weighted aggregation on the local models uploaded by each client based on the aggregation weight to generate a new round of global models.

[0055] Example 3

[0056] A federated learning optimization method for imbalanced datasets in image classification, characterized in that the method includes:

[0057] In any communication round of federated learning, the client receives a global model for the image classification task from the server, as well as the data distribution divergence and global set of difficult image categories calculated by the server based on the image data of each client in the previous communication round; wherein, the global model is an image classification model, and the image categories correspond to the classification labels in the image classification task;

[0058] The client dynamically adjusts the learning rate of local training based on the data distribution divergence, and adds the weight of the corresponding image category to the image classification loss function based on the global set of difficult image categories, and trains the model on the local image data; after local training is completed, the client uploads the updated local image classification model parameters and the local state vector containing the number of samples of each image category and the average loss of each image category to the server.

[0059] The server receives and aggregates the local state vectors from each client to obtain a global state vector containing global image data distribution information and global average classification loss. Based on the global state vector, the server determines a new set of global difficult image categories for the next communication round and calculates the distribution divergence of each client's image data distribution relative to the global image data distribution.

[0060] The server determines the aggregation weight of each client, wherein the aggregation weight of a client is the ratio of the total number of samples owned by the client to the total number of samples of all clients participating in the aggregation; the server then performs weighted aggregation on the local models uploaded by each client based on the aggregation weight to generate a new round of global image classification models.

[0061] Taking cat / dog binary classification as an example, assume the system has multiple clients, each with a different number of cat and dog images stored locally. After the (t-1)th round of communication, each client uploads a local state vector to the server, which includes at least: the number of cat images, the number of dog images, and the average classification loss for cats and dogs during this round of local training. The server sums the sample counts of all clients to obtain the global cat sample ratio and the global dog sample ratio. Simultaneously, the server calculates the distribution divergence between the class distribution of each client and the global distribution, for example, using KL divergence or L1 / L2 distance, to obtain the data distribution divergence value for each client. Then, based on the global average loss, categories with higher average losses, such as cats whose overall loss is significantly higher than that of dogs, are identified as globally difficult image categories and are sent to the clients in round t. In round t, each client first receives the global cat / dog classification model, its own data distribution divergence value, and the global difficult category set sent by the server. If a client's cat / dog ratio is heavily biased towards cats, for example, 90% are cats, its distribution dispersion will be large. The client will adjust its base learning rate accordingly, for example, η. k =η0 / (1+αD k ), where D k Let α be the divergence of the client's k distribution, and α be the scaling factor. If the server labels cats as the global difficulty category, the client will increase the weight of the cat class when constructing the loss function, for example, by setting a weight w in the cross-entropy loss. cat >w dog , where w cat Calculated using the global average loss ratio. The client uses the adjusted learning rate and weighted loss function to complete several rounds of training on local cat and dog images, obtaining new local state vectors which are then uploaded to the server.

[0062] After receiving the model parameters and local state vectors from all clients in round t, the server first aggregates the sample count and loss, and updates the global cat / dog distribution and global average loss. Then, it recalculates the divergence of each client relative to the latest global distribution, and sets an outlier threshold based on the median of all divergences: if a client's divergence is significantly higher than this threshold, its model aggregation base weights are multiplied by a penalty coefficient (e.g., a penalty coefficient less than 1); if a client's sample count in a globally difficult category, such as cat, exceeds a preset threshold, it is multiplied by a gain coefficient greater than 1. When aggregating classification layer parameters, the server calculates the final weights for cat and dog classes separately, ensuring that clients with lower losses and more reasonable distributions in the cat category contribute more to the cat classification parameters.

[0063] The above description of the disclosed embodiments enables those skilled in the art to make or use this application. Various modifications to these embodiments will be readily apparent to those skilled in the art, and the general principles defined herein may be implemented in other embodiments without departing from the spirit or scope of this application. Therefore, this application is not to be limited to the embodiments shown herein, but is to be accorded the widest scope consistent with the principles and novel features disclosed herein.

Claims

1. A federated learning optimization method for imbalanced datasets, characterized in that, Includes the following steps: In any communication round of federated learning, the client receives the global model from the server, as well as the data distribution divergence and global difficulty category set calculated by the server in the previous communication round. The client adjusts the learning rate of local training based on the data distribution divergence and increases the weight of the corresponding category in the loss function based on the global difficulty category set for local training. After local training is completed, the client uploads the local model and the local state vector containing the number of samples for each data category and the average loss to the server. The server receives and aggregates the local state vectors from each client to obtain a global state vector that includes the global data distribution and the global average loss. Based on the global state vector, the server determines the global set of difficulty categories and the data distribution divergence of each client for the next communication round, and calculates the divergence of each client's data distribution relative to the global data distribution. The server determines the aggregation weight of each client, wherein the aggregation weight of a client is the ratio of the total number of samples owned by the client to the total number of samples of all clients participating in the aggregation; the server performs weighted aggregation on the local models uploaded by each client based on the aggregation weight to generate a new round of global models.

2. The method according to claim 1, characterized in that, The method further includes: When the data distribution divergence of the client exceeds the outlier threshold set based on the median of the data distribution divergence of all clients, a penalty is applied to the aggregation weight; when the number of global hard category samples held by the client exceeds a preset threshold, a gain is applied to the aggregation weight; and when aggregating the classification layer parameters of the model, for each data category, the adjusted aggregation weight is multiplied by the inverse of the average loss of the client on that category to obtain the weight used to aggregate the parameters corresponding to that category.

3. The method according to claim 1, characterized in that, The client adjusts the learning rate for local training based on the data distribution divergence, including: The client's local learning rate η k Based on the data distribution divergence D k The formula for adjustment is: ; Where η0 is the initial learning rate and β is the preset decay coefficient.

4. The method according to claim 1, characterized in that, The step of adding weights for the corresponding categories in the loss function based on the global set of difficult categories for local training includes: For any category belonging to the global difficulty category set, the client will multiply the weight in the loss function by a preset gain factor.

5. The method according to claim 2, characterized in that, When the data distribution divergence of a client exceeds an outlier threshold set based on the median of the data distribution divergence of all clients, a penalty is imposed on the weights, including: Outlier threshold T D Based on the data distribution divergence of all clients, D={D1,D2,...,D...} K The median (median(D)) and median absolute deviation (MAD(D)) of the data are determined using the following formula: T D =median(D) + c × MAD(D); where c is a preset constant; for the data distribution divergence D k >T D The client multiplies the client's aggregate weight by a preset penalty factor.

6. The method according to claim 2, characterized in that, When the number of global difficulty category samples held by the client exceeds a preset threshold, a gain is applied to the weights, including: When the total number of samples belonging to the global difficulty category set held by the client exceeds a preset sample number threshold, the client's aggregation weight is multiplied by a preset gain factor.

7. A federated learning optimization system for imbalanced datasets, characterized in that, Includes the following modules: The adjustment module is used in any communication round of federated learning. The client receives the global model from the server, as well as the data distribution divergence and global difficulty category set calculated by the server in the previous communication round. The client adjusts the learning rate of local training based on the data distribution divergence and increases the weight of the corresponding category in the loss function based on the global difficulty category set for local training. After local training is completed, the client uploads the local model and the local state vector containing the number of samples for each data category and the average loss to the server. The calculation module is used by the server to receive and aggregate the local state vectors of each client to obtain a global state vector containing the global data distribution and the global average loss. Based on the global state vector, the server determines the global difficulty category set and the data distribution divergence of each client for the next communication round, and calculates the divergence of the data distribution of each client relative to the global data distribution. The generation module is used by the server to determine the aggregation weight of each client, wherein the aggregation weight of the client is the ratio of the total number of samples owned by the client to the total number of samples of all clients participating in the aggregation; the server performs weighted aggregation on the local models uploaded by each client based on the aggregation weight to generate a new round of global models.

Images