Federal learning aggregation method based on Dirichlet ideal distribution and dynamic attenuation

Abstract

The invention discloses a federated learning aggregation method based on Dirichlet ideal distribution and dynamic attenuation. The central server broadcasts a concentration parameter alpha which is used for defining symmetric Dirichlet ideal distribution, and the mathematical expectation of the symmetric Dirichlet ideal distribution is uniform distribution. Each client side locally calculates the KL divergence between the category distribution and the uniform distribution of the client side and generates a distribution quality weight according to the KL divergence, the weight is exponentially attenuated along with the increase of the KL divergence, and the attenuation rate is controlled by alpha; and the client uploads the scalar weight to a server. And the server constructs and normalizes an initial aggregation weight in combination with the data volume weight and the time-varying attenuation coefficient, and is used for weighting and averaging the local model of each client to update the global model. Wherein the larger the alpha is, the lower the tolerance of the system to the data distribution deflection is. According to the method, an aggregation strategy can be flexibly regulated and controlled only through alpha without additional super parameters, and a client only uploads a scalar, so that privacy is effectively protected, and the model convergence and generalization ability in a non-independent identically distributed scene are remarkably improved.

Description

Technical Field

[0001] This invention relates to the fields of artificial intelligence and privacy computing technology, and in particular to a federated learning aggregation method for non-independent identically distributed (Non-IID) scenarios. Specifically, it is a client-side weight adjustment method based on Dirichlet ideal distribution modeling and dynamic decay mechanism, which is suitable for distributed machine learning systems where multiple participants collaboratively train deep learning models and the original data does not leave the local machine. Background Technology

[0002] Federated Learning (FL) is a distributed machine learning paradigm that enables collaborative modeling among multiple parties while protecting data privacy. Its typical process includes: a central server initializes a global model and distributes it to participating clients; each client trains its model using local data and uploads the updated model parameters; the server performs weighted aggregation of the client models and updates the global model. This process iterates until the model converges. Because the original data remains locally, federated learning is widely used in privacy-sensitive scenarios such as healthcare, finance, and smart terminals.

[0003] However, in practical applications, the data distribution of each client often exhibits non-independent and identically distributed (Non-IID) characteristics, such as class skew (some clients contain only a few classes) and imbalanced data distribution. This leads to significant differences in the optimization direction of local models. Traditional aggregation strategies based on data volume ratios (such as FedAvg) are easily affected by clients with extremely skewed distributions, causing model oscillations, slow convergence, and even performance degradation.

[0004] To alleviate this problem, researchers have proposed several improved aggregation methods. One type of method dynamically adjusts weights by evaluating the similarity or gradient consistency of client model updates; another type introduces a data distribution balance metric, such as using KL divergence to measure the difference between the local distribution and the global or ideal distribution, and designs a penalty mechanism accordingly. However, existing schemes generally have the following shortcomings: (1) They lack a probabilistic modeling basis for the "ideal distribution," and the uniform distribution assumption lacks theoretical support; (2) The balance weights are decoupled from the aggregation strategy, requiring the introduction of additional hyperparameters, which are complex to tune; (3) They cannot adaptively adjust the tolerance for distribution skew according to the training process, making it difficult to balance early convergence speed and later generalization performance.

[0005] Therefore, there is an urgent need for a federated learning aggregation method that can both theoretically model the ideal data distribution and dynamically adjust the aggregation strategy through a single controllable parameter, so as to effectively improve the robustness and generalization ability of the model in Non-IID scenarios while ensuring privacy. Summary of the Invention

[0006] To overcome the problems of slow model convergence, weak generalization ability, complex hyperparameters, and lack of theoretical modeling in existing federated learning aggregation methods in non-independent and identically distributed (Non-IID) scenarios, this invention provides a federated learning aggregation method based on the Dirichlet ideal distribution and dynamic decay. This method introduces an ideal distribution model with clear probabilistic meaning and utilizes a single concentration parameter to achieve flexible and adaptive control of the aggregation strategy, significantly improving the global model performance while ensuring data privacy.

[0007] The objective of this invention is achieved through the following technical solution:

[0008] Each round of federal training includes the following steps:

[0009] S1. The central server broadcasts a concentration parameter α, where α > 0, to multiple clients participating in federated learning, to define a parameterized Dirichlet ideal distribution model P. ideal ~Dir(α,α,…,α);

[0010] S2. After receiving the α, each client performs the following operations locally:

[0011] (a) The category distribution P is obtained based on local data statistics. i =(p i1 p i2 , ..., p iK ), where K is the total number of categories;

[0012] (b) Calculate the category distribution P i The mathematical expectation E[P] of the Dirichlet ideal distribution model ideal KL divergence between (1 / K, 1 / K, ..., 1 / K):

[0013]

[0014] (c) Calculate the distribution quality weights based on the aforementioned α:

[0015]

[0016] (d) The distribution quality weights Uploaded to the central server;

[0017] S3, the central server is based on the data sample size n of each client. i Calculate data volume weights

[0018]

[0019] Where N is the number of clients in the current round;

[0020] S4. The central server obtains the time-varying decay coefficient γ(t) corresponding to the current training round t, and calculates the initial aggregation weight for each client:

[0021]

[0022] S5. Normalize the initial aggregate weights of all clients to obtain the final aggregate weights:

[0023]

[0024] S6. Based on the final aggregation weight, perform a weighted average of the local model parameters uploaded by each client in the current round, and update the global model; wherein, the concentration parameter α controls the concentration of the Dirichlet ideal distribution model: the larger α is, the more concentrated the distribution is to a uniform distribution, and the lower the system's tolerance for data distribution skewness; the smaller α is, the larger the distribution variance is, and the system allows for higher distribution diversity.

[0025] Furthermore, the concentration parameter α is determined by the central server according to any of the following methods:

[0026] As a preset fixed value, it is configured according to the data distribution characteristics of the federated learning task;

[0027] As a function of training round t, α(t) is dynamically adjusted during the training process;

[0028] The model is adaptively updated based on the rate of change of its performance on the validation set.

[0029] Furthermore, when a dynamic adjustment method is adopted, α(t) increases monotonically with the training round t, so as to tolerate the skewness of client data distribution in the early stage of training and strengthen the requirement for data balance in the later stage of training.

[0030] Furthermore, when using the adaptive update method, if the rate of change of the global model's accuracy on the validation set is lower than a preset threshold, or if the loss function oscillates, then the value of α is increased.

[0031] Furthermore, the time-varying decay coefficient γ(t) decreases monotonically with the training round t, and the decay mode includes exponential decay, linear decay, or step decay.

[0032] Furthermore, when calculating the KL divergence, each client does not sample a specific distribution from the Dirichlet ideal distribution model, but directly uses its mathematical expectation as a reference target to ensure the determinism and reproducibility of the calculation results.

[0033] The present invention also provides an aggregation system for federated learning, comprising multiple clients and a central server;

[0034] The central server is used to broadcast the concentration parameter α in each round of training to define the Dirichlet ideal distribution model, receive the distribution quality weights uploaded by the client, calculate the final aggregate weights, and update the global model.

[0035] The client is used to calculate the KL divergence locally and calculate the distribution quality weights based on α. And upload;

[0036] In this context, α indirectly adjusts the requirements of federated aggregation on the balance of client data by controlling the concentration of the Dirichlet ideal distribution.

[0037] The present invention also provides a computer-readable storage medium having stored thereon computer program instructions, characterized in that, when the instructions are executed by a processor, they implement the method as described in any one of claims 1 to 7.

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

[0039] (1) Solid theoretical foundation: The Dirichlet distribution is used as an ideal model to provide a probabilistic interpretation of "uniform expectation" and avoid empirical assumptions;

[0040] (2) Extremely simple and efficient parameters: The ideal distribution shape and penalty intensity are controlled simultaneously by α issued by the server, without the need to introduce additional hyperparameters;

[0041] (3) Strong privacy protection: The client only uploads scalar weights, and the original distribution and data do not leave the local machine;

[0042] (4) Strong adaptive capability: Supports dynamic adjustment of α, taking into account both convergence speed and final performance;

[0043] (5) Convenient engineering deployment: the calculation is deterministic, there is no random sampling, and it is easy to integrate into the existing federated learning framework. Attached Figure Description

[0044] Figure 1 A flowchart of the federated learning aggregation process of the present invention. Detailed Implementation

[0045] 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.

[0046] Example

[0047] like Figure 1 As shown, this embodiment takes image classification as an example. The federated learning system includes one central server and N clients (N≥2). Each client holds a local non-independent identically distributed (Non-IID) image dataset, with a total number of categories K=10 (e.g., CIFAR-10). The total number of training rounds T=100.

[0048] Step S1: Server Initialization and Parameter Broadcast

[0049] The central server initializes the global model (e.g., ResNet-18) and sets the concentration parameter α = 2.0 (a fixed value) to define the Dirichlet ideal distribution Dir(2.0, ..., 2.0). At the start of the t=1th training round, the server broadcasts the current global model parameters and α to all clients.

[0050] Step S2: Local computation on the client side

[0051] After each client i receives α, it executes the following locally:

[0052] (a) Calculate the percentage of samples from each of the 10 categories in the training set to obtain the category distribution P. i =(p i1 p i2 , ..., p i10 );

[0053] (b) Calculate P i The KL divergence with uniform distribution u ~ (0.1,..., 0.1):

[0054]

[0055] (c) Calculate the distribution quality weights:

[0056]

[0057] (d) Train the model locally using the data to obtain local model parameters. And and the amount of data n i Uploaded to the server.

[0058] Steps S3-S6: Server Aggregation

[0059] Server execution:

[0060] (a) Calculate data volume weights

[0061] (b) Obtain the current round decay coefficient γ(t) = 0.8 t-1 (Exponential decay);

[0062] (c) Calculate the initial aggregation weights:

[0063] (d) Normalization yields the final weights

[0064] (e) Update the global model:

[0065] Dynamic α Example (Comparative)

[0066] In another embodiment, α(t) = 0.5 + 0.05t, i.e., linearly increasing from 0.5 to 5.5. This strategy aims to tolerate greater data distribution skew in the early rounds for rapid convergence, and then gradually emphasize the importance of data balance to improve generalization ability.

[0067] 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 federated learning aggregation method based on Dirichlet ideal distribution and dynamic decay, applied to the current training round of federated learning, characterized in that, Includes the following steps: S1. The central server broadcasts a concentration parameter α, where α > 0, to multiple clients participating in federated learning, to define a parameterized Dirichlet ideal distribution model P. ideal ~Dir(α, α, ..., α); S2. After receiving the α, each client performs the following operations locally: (a) The category distribution P is obtained based on local data statistics. i =(p i1 p i2 , ..., p iK ), where K is the total number of categories; (b) Calculate the category distribution P i The mathematical expectation E[P] of the Dirichlet ideal distribution model ideal KL divergence between (1 / K, 1 / K, ..., 1 / K): (c) Calculate the distribution quality weights based on the aforementioned α: (d) The distribution quality weights Uploaded to the central server; S3, the central server is based on the data sample size n of each client. i Calculate data volume weights Where N is the number of clients in the current round; S4. The central server obtains the time-varying decay coefficient γ(t) corresponding to the current training round t, and calculates the initial aggregation weight for each client: S5. Normalize the initial aggregate weights of all clients to obtain the final aggregate weights: S6. Based on the final aggregation weight, perform a weighted average of the local model parameters uploaded by each client in the current round, and update the global model; wherein, the concentration parameter α controls the concentration of the Dirichlet ideal distribution model: the larger α is, the more concentrated the distribution is to a uniform distribution, and the lower the system's tolerance for data distribution skewness; the smaller α is, the larger the distribution variance is, and the system allows for higher distribution diversity.

2. The method according to claim 1, characterized in that, The concentration parameter α is determined by the central server according to any of the following methods: As a preset fixed value, it is configured according to the data distribution characteristics of the federated learning task; As a function of training round t, α(t) is dynamically adjusted during the training process; The model is adaptively updated based on the rate of change of its performance on the validation set.

3. The method according to claim 2, characterized in that, When a dynamic adjustment method is adopted, α(t) increases monotonically with the training round t to tolerate the skewness of client data distribution in the early stage of training and to strengthen the requirement for data balance in the later stage of training.

4. The method according to claim 2, characterized in that, When using the adaptive update method, if the rate of change of the global model's accuracy on the validation set is lower than the preset threshold, or if the loss function oscillates, then the value of α is increased.

5. The method according to claim 1, characterized in that, The time-varying decay coefficient γ(t) decreases monotonically with the training round t, and the decay method includes exponential decay, linear decay or step decay.

6. The method according to claim 1, characterized in that, When calculating the KL divergence, each client does not sample a specific distribution from the Dirichlet ideal distribution model, but directly uses its mathematical expectation as a reference target to ensure the determinism and reproducibility of the calculation results.

7. A federated learning system, characterized in that, It includes multiple clients and a central server; The central server is used to broadcast the concentration parameter α in each round of training to define the Dirichlet ideal distribution model, receive the distribution quality weights uploaded by the client, calculate the final aggregate weights, and update the global model. The client is used to calculate the KL divergence locally and calculate the distribution quality weights based on α. And upload; In this context, α indirectly adjusts the requirements of federated aggregation on the balance of client data by controlling the concentration of the Dirichlet ideal distribution.

8. A computer-readable storage medium having computer program instructions stored thereon, characterized in that, When the instructions are executed by the processor, they implement the method as described in any one of claims 1 to 7.

Images