Federal learning acceleration method based on multi-agent reinforcement learning and knowledge distillation

Abstract

The invention provides an FL (Federated Learning) acceleration method based on multi-agent reinforcement learning (MARL) and knowledge distillation (KD), and the method comprises the steps: firstly, because local data of edge nodes are generally non-independent identically distributed, the model training efficiency of the FL is reduced, and the local data of the edge nodes are not independent identically distributed, so that the local data of the edge nodes are not independent identically distributed, and the local data of the edge nodes are not independent identically distributed, so that the local data of the edge nodes are not independent identically distributed, and the local data of the edge nodes are not independent identically distributed, and the local data of the edge nodes are not independent identically distributed, and the local data of the edge nodes are not independent identically distributed; according to the number of nodes participating in training in each round, nodes with close data distribution are divided into the same cluster by using a clustering algorithm on the basis of the model weight, and the nodes in different clusters are mutually combined to finally obtain approximate global data distribution; then, each cluster of a clustering result is regarded as an independent agent, an MARL algorithm is adopted, appropriate nodes are selected from each cluster to participate in training, an MARL reward value is maximized, and excessive communication rounds are punished; and finally, due to different computing storage capacities of edge node equipment, the model is globally trained through a bidirectional KD method, the local model is further personalized, the transmission content is compressed, the communication pressure is reduced, and the purpose of federal learning acceleration is achieved.

Description

Technical Field

[0001] The present invention relates to the fields of edge computing and machine learning, and in particular to a federated learning (FL) acceleration method based on multi-agent reinforcement learning and knowledge distillation. Background Art

[0002] Federated learning acceleration method based on multi-agent reinforcement learning and knowledge distillation, based on federated learning, multi-agent reinforcement learning and knowledge distillation technology. First, since the local data of edge nodes are usually non-independent and identically distributed, the model training efficiency of FL is reduced. According to the number of nodes participating in each round of training, a clustering algorithm is used to divide the nodes with similar data distribution into the same cluster based on the model weight. The nodes in different clusters are combined to finally obtain an approximate global data distribution; then, each cluster of the clustering results is regarded as an agent, and a multi-agent reinforcement learning algorithm is used to select appropriate nodes from each cluster to participate in the training, maximize the reinforcement learning reward value, and punish too many communication rounds; finally, due to the different computing and storage capabilities of edge node devices, the model is globally trained through the method of two-way knowledge distillation, further personalizing the local model, compressing the transmission content, and reducing the communication pressure, so as to achieve the purpose of federated learning acceleration. The technologies closest to the present invention are:

[0003] (1) DQN-based federated learning node selection algorithm: In 2020, Wang Hao et al. pointed out that for non-independent and identically distributed data, a DQN-based node selection strategy is used to offset the impact of non-independent and identically distributed data on model training. This method uses model weights as states and selected participants as actions, which can reduce communication rounds and accelerate federated learning to a certain extent. However, due to the limitations of the DQN algorithm, when the dimension of the action space is large, selecting the optimal action becomes more complicated and may not be feasible within a reasonable time frame.

[0004] (2) Federated learning node selection algorithm based on DDPG: In 2021, Zhang Peiying et al. proposed a federated learning node selection algorithm based on DDPG, which uses a policy gradient-based deep reinforcement learning algorithm and needs to map continuous action values ​​to discrete action sets. This mapping process may be inaccurate or incomplete, resulting in the inability to accurately select the optimal discrete action.

[0005] (3) Federated learning acceleration algorithm based on model compression: In 2020, Felix Sattler et al. proposed a federated learning acceleration algorithm based on model compression, which compresses the transmitted gradient by sparse ternary compression. However, this gradient compression method may be highly sensitive to the model, and the gradient characteristics of different models may be different. Therefore, the compression method suitable for one model may not be suitable for another model. It is necessary to select an appropriate compression method according to the specific situation.

[0006] Among them, (1) the federated learning node selection algorithm based on DQN and (2) the federated learning node selection algorithm based on DDPG are relatively sensitive to the action space dimension. In edge network scenarios, as the total number of participants increases, traditional federated learning algorithms are not applicable; (3) the federated learning acceleration algorithm based on model compression considers the universality of the algorithm. At the same time, the existing technologies for federated learning node selection methods based on multi-agent reinforcement learning have not considered and designed this aspect. Summary of the invention

[0007] In order to solve the shortcomings and deficiencies in the prior art, the present invention proposes a federated learning acceleration method based on multi-agent reinforcement learning and knowledge distillation. First, since the local data of edge nodes are usually non-independent and identically distributed, the model training efficiency of FL is reduced. According to the number of nodes participating in each round of training, a clustering algorithm is used to divide the nodes with similar data distribution into the same cluster based on the model weight. The nodes in different clusters are combined to finally obtain an approximate global data distribution; then, each cluster of the clustering results is regarded as an independent agent, and a multi-agent reinforcement learning algorithm is used to select appropriate nodes from each cluster to participate in the training, maximize the reinforcement learning reward value, and punish too many communication rounds; finally, due to the different computing and storage capabilities of edge node devices, the model is globally trained through the method of two-way knowledge distillation, and the local model is further personalized, the transmission content is compressed, and the communication pressure is reduced, so as to achieve the purpose of federated learning acceleration.

[0008] The technical solution of the present invention is:

[0009] A federated learning acceleration method based on multi-agent reinforcement learning and knowledge distillation, including a node partitioning module, a participant selection module, a local training module and a distillation aggregation module, including the following steps:

[0010] Step (1) Node division module: The server sends the same initial model to N participants. All participants upload the model weights after a round of local training. The server uses a clustering algorithm to perform clustering based on the uploaded model weights and divides the N participants into K clusters.

[0011] Step (2) Participant selection module: each cluster in the clustering result of the node partitioning module is regarded as an independent agent, the weight uploaded to the model is regarded as the state in reinforcement learning, the selected participant in each cluster is the action, and the reward value is positively correlated with the accuracy of the final model. K participants are selected using the multi-agent reinforcement learning algorithm.

[0012] Step (3) Local training module, starting from the second round of training, different from uploading once after each training, the selected participants upload the parameters after multiple local training iterations. The number of iterations is related to the reward value of the previous round of reinforcement learning, and the upper limit is set to 10.

[0013] Step (4) Distillation Aggregation Module, unlike traditional FL, participants will adaptively set the size and structure of the local model according to their needs. At the beginning of FL training, the server will send the global model, and the participants will train the local model based on the local data, and then the local model and the global model will be distilled with each other. In this way, the local model can obtain local information based on local data and obtain global information through mutual distillation, while achieving model personalization, compressing transmission content, and improving communication efficiency.

[0014] Beneficial effects of the present invention:

[0015] (1) Clustering the participants through clustering algorithms, obtaining global data distribution in non-IID scenarios, and ensuring global model stability. The federated learning paradigm is adopted to protect data privacy;

[0016] (2) At the same time, in order to further improve the convergence speed of the model, the algorithm uses a multi-agent reinforcement learning algorithm to select participants in each round of training based on the clustering results. Specifically, each cluster in the clustering results is regarded as an agent, and the results selected by each agent are combined into the participants of this round of training. This method can significantly reduce the computational complexity and accelerate training;

[0017] (3) When participants conduct local training, they use knowledge distillation to personalize local models, compress transmission content, reduce federated learning communication volume, and improve communication efficiency. BRIEF DESCRIPTION OF THE DRAWINGS

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

[0019] Figure 1This is a structural diagram of the federated learning acceleration algorithm based on multi-agent reinforcement learning and knowledge distillation of the present invention;

[0020] Figure 2 This is a flow chart of the federated learning acceleration algorithm based on multi-agent reinforcement learning and knowledge distillation of the present invention. DETAILED DESCRIPTION

[0021] The following will be combined with the drawings in the embodiments of the present invention to clearly and completely describe the technical solutions in the embodiments of the present invention. Obviously, the described embodiments are only part of the embodiments of the present invention, not all of the embodiments. Based on the embodiments of the present invention, all other embodiments obtained by ordinary technicians in this field without creative work are within the scope of protection of the present invention.

[0022] like Figure 1 As shown, the federated learning acceleration method based on multi-agent reinforcement learning and knowledge distillation of the present invention includes four modules: a node partitioning module, a participant selection module, a local training module and a distillation aggregation module.

[0023] Combine the following Figure 1 and Figure 2 , the specific process of the federated learning acceleration algorithm based on multi-agent reinforcement learning and knowledge distillation is described in detail:

[0024] A federated learning acceleration method based on multi-agent reinforcement learning and knowledge distillation, including a node partitioning module, a participant selection module, a local training module and a distillation aggregation module, including the following steps:

[0025] Step (1) Node division module: The server sends the same initial model to N participants. All participants perform a round of local training and then upload the model weights. The server uses a clustering algorithm to perform clustering based on the uploaded model weights and divides the N participants into K clusters.

[0026] Step (2) Participant selection module: each cluster in the clustering result of the node partitioning module is regarded as an independent agent, the weight uploaded to the model is regarded as the state in reinforcement learning, the selected participant in each cluster is the action, and the reward value is positively correlated with the accuracy of the final model. K participants are selected using the multi-agent reinforcement learning algorithm.

[0027] Step (3) Local training module, starting from the second round of training, different from uploading once after each training, the selected participants upload the parameters after multiple local training iterations. The number of iterations is related to the reward value of the previous round of reinforcement learning, and the upper limit is set to 10.

[0028] Step (4) Distillation Aggregation Module, unlike traditional FL, participants will adaptively set the size and structure of the local model according to their needs. At the beginning of FL training, the server will send the global model, and the participants will train the local model based on the local data, and then the local model and the global model will be distilled with each other. In this way, the local model can obtain local information based on local data and obtain global information through mutual distillation, while achieving model personalization, compressing transmission content, and improving communication efficiency.

[0029] The present invention clusters participants through a clustering algorithm, obtains global data distribution in a non-independent and identically distributed scenario, and ensures the stability of the global model. The federated learning paradigm is adopted to protect data privacy; at the same time, in order to further improve the convergence speed of the model, the algorithm selects participants for each round of training according to the clustering results through a multi-agent reinforcement learning algorithm. Specifically, each cluster in the clustering results is regarded as an agent, and the results selected by each agent are combined into the participants of this round of training. In this way, the computational complexity can be significantly reduced and the training can be accelerated; when participants conduct local training, knowledge distillation is used to personalize the local model, compress the transmission content, reduce the amount of federated learning communication, and improve communication efficiency.

[0030] The above description is only a preferred embodiment of the present invention and is not intended to limit the present invention. Any modifications, equivalent substitutions, improvements, etc. made within the spirit and principle of the present invention should be included in the protection scope of the present invention.

Claims

1. A federated learning acceleration method based on multi-agent reinforcement learning and knowledge distillation, including a node partitioning module, a participant selection module, a local training module and a distillation aggregation module, including the following steps: Step (1) Node division module: The server sends the same initial model to N participants. All participants upload the model weights after a round of local training. The server uses a clustering algorithm to perform clustering based on the uploaded model weights and divides the N participants into K clusters. Step (2) Participant selection module: each cluster in the clustering result of the node partitioning module is regarded as an agent, the weight uploaded to the model is regarded as the state in reinforcement learning, the selected participant in each cluster is the action, and the reward value is positively correlated with the accuracy of the final model. K participants are selected using the multi-agent reinforcement learning algorithm. Step (3) Local training module, starting from the second round of training, different from uploading once after each training, the selected participants upload the parameters after multiple local training iterations. The number of iterations is related to the reward value of the previous round of reinforcement learning, and the upper limit is set to 10. Step (4) Distillation Aggregation Module, unlike traditional FL, participants will adaptively set the size and structure of the local model according to their needs. At the beginning of FL training, the server will send the global model, and the participants will train the local model based on the local data, and then the local model and the global model will be distilled with each other. In this way, the local model can obtain local information based on local data and obtain global information through mutual distillation, while achieving model personalization, compressing transmission content, and improving communication efficiency.

Images