A heterogeneous federated learning training method based on multi-branch neural network model

Abstract

The present invention relates to a heterogeneous federated learning training method based on a multi-branch neural network model. A multi-branch neural network model is introduced as a shared global model, which can match suitable sub-branch models according to the computing power of different devices, can adapt well to the scenario of heterogeneous computing resources, and make full use of the computing resources of different devices, thereby effectively improving the performance and efficiency of the entire heterogeneous federated learning training system; in view of the characteristics of the multi-branch model, the present invention proposes a multi-branch model aggregation method based on shared layer parameters, which can aggregate different sub-branch models into a global multi-branch model, so that model parameters can be effectively shared between different device models; on the basis of parameter aggregation, the present invention also introduces distillation learning to solve the problem of performance fluctuations after model aggregation, thereby accelerating the convergence speed of the global model, reducing the number of required training rounds, and saving communication consumption.

Description

Technical Field

[0001] The present invention relates to the technical field of federated learning, and more specifically, to a heterogeneous federated learning training method based on a multi-branch neural network model. Background Art

[0002] Existing federated learning algorithms such as FedAvg and its improved algorithms can effectively perform distributed machine learning while satisfying data privacy. However, due to the limitation of model gradient aggregation, the models deployed on all devices participating in federated learning are required to be isomorphic. However, in reality, the existence of device heterogeneity is inevitable, which leads to the size and capacity of the training model depending on the device with the weakest computing power in the system, resulting in system resource waste and performance bottlenecks. At present, many scholars have also tried to solve the federated learning training of heterogeneous models by introducing knowledge distillation. For example, FedMD uses labeled public data sets and average aggregation of the output logit of device models to transfer knowledge, thereby personalizing the training of each device-side model, but it cannot aggregate to form a public model on the central server like classic federated learning. In addition, in the work of FedDF, by averaging and aggregating all device models and their output logits, an integrated distillation for model aggregation is proposed, which can aggregate multiple heterogeneous shared models in the cloud. However, if the device heterogeneity changes, the types of device-side models will inevitably increase, resulting in a sharp increase in the cost of maintaining and training heterogeneous models on the server side. Summary of the invention

[0003] In order to overcome the defects in the above-mentioned prior art, the present invention provides a heterogeneous federated learning training method based on a multi-branch neural network model, which effectively improves the performance and efficiency of the entire heterogeneous federated learning training system.

[0004] In order to solve the above technical problems, the technical solution adopted by the present invention is: a heterogeneous federated learning training method based on a multi-branch neural network model, comprising the following steps:

[0005] S1. Initialization training of multi-branch models on the cloud: There are global branch models with multiple branches and a public data set D0 in the cloud server. Different global branch models retain their own output layers on the basis of sharing some public hidden layers. First, multiple global branch models are pre-trained based on the public data set D on the cloud to initialize them.

[0006] S2. Matching of sub-branch neural network models: Before federated learning begins, all devices that request to participate in federated learning report the available computing and storage information of the devices to the cloud server. Then, the cloud server calculates and determines the best branch model suitable for the device based on the collected device information, and distributes the weighted single-branch model to the corresponding device.

[0007] S3. Local training of device models: After receiving the single-branch model sent by the cloud, each participating device replaces it with the local model of the device, and then trains the current device model based on the local private dataset;

[0008] S4. Aggregation of multi-branch lifting network models in the cloud: After completing local training, each participating device uploads the parameters of its own model to the cloud server, and performs weighted averaging of the parameters of the same parts of the device model and multiple global branch models, thereby aggregating and updating the global branch model in the cloud;

[0009] S5. Distillation training of multi-branch neural network model on the cloud: Based on the public data set on the cloud and the model gradients uploaded by each device, knowledge distillation learning is performed on the aggregated multi-branch network model; after the distillation training is completed, the model parameters of each branch in the multi-branch model are sent to the corresponding device for the next round of local training;

[0010] S6. Application of the multi-branch neural network model in the cloud: Repeat steps S3 to S5 until a predetermined number of training rounds is reached, and finally a global branch neural network model is obtained.

[0011] In one embodiment, the overall parameter of the global branch model is W s , where the parameters of each sub-branch model are recorded as On the device side, there are N participating devices, and each device i has a private local dataset D i and local models In step S1, the cross entropy loss value between the output of each branch model and the ground truth is calculated, and then the weighted average of the loss values ​​of each sub-branch is taken as the overall loss value:

[0012]

[0013] Finally, backpropagation is performed according to the overall loss value to update the parameters until the model converges:

[0014]

[0015] In one embodiment, in step S2, when each device i reports its maximum satisfiable parameter quantity P iThen, match a i ≥P k The largest sub-branch model k of P is then sent to the selected sub-branch model k to device i, where P k is the parameter quantity of sub-branch model k.

[0016] In one of the embodiments, in step S3, constraints are imposed on the difference between the parameters of the current local model and the original cloud model, so that under the constraints of the cloud model, the shared hidden layer of each device model can obtain parameter distribution that is as similar as possible during training.

[0017] In one embodiment, the constraint on the difference between the parameters of the current local model and the original cloud model specifically includes: adding an L2 regularization loss function to the original cross entropy loss function to measure the difference in parameter distribution between the local model and the cloud model; when the device i receives the sub-branch model k, then initialize before local training Assume that the number of shared network layers between the sub-branch model k and the cloud main branch model is H, then the L2 regularized loss function of the local training of device i is expressed as:

[0018]

[0019] Combined with the loss function based on the local data set, the total loss function of device i is:

[0020]

[0021] Among them, η is a hyperparameter used to measure the proportion of L2 regular loss function to the overall loss function;

[0022] Finally, the local model parameters of device i are optimized and updated according to the following formula:

[0023]

[0024] In one embodiment, in step S4, a multi-branch joint average aggregation method is used, specifically including: when the cloud server receives the model parameter set uploaded by all participating devices Then start multi-branch aggregation; first obtain the dictionary set W of all parameter layers in the global multi-branch model S , the number of layers of the model is recorded as K, and the index of the kth layer is recorded as l k , and the corresponding parameters of the kth layer are expressed as W S [l k ]; then traverse W S All parameter layers in: For each layer k, the layer parameters W S [l k]First initialize it to zero value and count the index l of this layer k Upload model collection The total number of occurrences in is denoted as Count[l k ], and for every existence l k Device model of the index layer Accumulate the parameters of its corresponding layer into the global model, that is, After traversing and accumulating each layer of the global multi-branch model, the new global model parameter W is obtained and output S .

[0025] In one embodiment, the loss function in the distillation training includes: a cross entropy loss function with the true label, as shown in Formula 1; and a KL divergence loss function with the soft label output by the device model, as shown in Formula 6 below:

[0026]

[0027] The total loss function of distillation training is shown in Formula 7, where the hyperparameter α is used to set the weight ratio between the cross entropy function and the KL divergence loss;

[0028]

[0029] The most optimized global branch model in distillation training is calculated by formula 8:

[0030]

[0031] In one of the embodiments, in the step S6, the output global multi-branch neural network model can match the appropriate branch model according to the resource limitations and accuracy requirements of different devices to meet the needs of different types of device applications; when the computing resources on the device side are tight or insufficient, the device side can only process the calculation of the shared parameter part of the local model, and then send the intermediate results to the cloud global network for the remaining calculations, and finally return the cloud results to the edge device.

[0032] The present invention also provides an electronic device, comprising: a memory, a processor, and a computer program stored in the memory and executable on the processor, wherein the processor executes the computer program to implement the above-mentioned heterogeneous federated learning training method based on a multi-branch neural network model.

[0033] The present invention also provides a computer-readable storage medium having a computer program stored thereon. When the computer program is executed by a processor, the above-mentioned heterogeneous federated learning training method based on a multi-branch neural network model is implemented.

[0034] Compared with the prior art, the beneficial effects are:

[0035] 1. In the present invention, a multi-branch neural network model is introduced as a shared global model, and suitable sub-branch models can be matched according to the computing power of different devices. Compared with the traditional federated learning method, the present invention can well adapt to the scenario of heterogeneous computing resources and make full use of the computing resources of different devices, thereby effectively improving the performance and efficiency of the entire heterogeneous federated learning training system;

[0036] 2. In view of the characteristics of multi-branch models, the present invention proposes a multi-branch model aggregation method based on shared layer parameters, which can aggregate different sub-branch models into a global multi-branch model, so that model parameters can be effectively shared between different device models;

[0037] 3. On the basis of parameter aggregation, the present invention also introduces distillation learning to solve the problem of performance fluctuation after model aggregation, thereby accelerating the convergence speed of the global model, reducing the number of required training rounds, and saving communication consumption;

[0038] 4. In the present invention, the multi-branch neural network model obtained by cloud-based training can flexibly match the application requirements of various types of devices and support end-edge collaborative reasoning. BRIEF DESCRIPTION OF THE DRAWINGS

[0039] Figure 1 It is a schematic flow chart of the method of the present invention.

[0040] Figure 2 It is a schematic diagram of the workflow of the multi-branch model training system of the present invention.

[0041] Figure 3 It is a schematic diagram of the loss function calculation in the local training of the present invention. DETAILED DESCRIPTION

[0042] 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. The present invention is described in one of the embodiments in combination with the specific implementation methods. Among them, the drawings are only used for exemplary descriptions, and only schematic diagrams are shown, not physical drawings, and cannot be understood as limitations on this patent; in order to better illustrate the embodiments of the present invention, some parts of the drawings will be omitted, enlarged or reduced, and do not represent the size of the actual product; for those skilled in the art, it is understandable that some well-known structures and their descriptions in the drawings may be omitted.

[0043] In the description of the present invention, it should be understood that if the terms "upper", "lower", "left", "right" and the like indicate an orientation or positional relationship based on the orientation or positional relationship shown in the drawings, it is only for the convenience of describing the present invention and simplifying the description, rather than indicating or implying that the device or element referred to must have a specific orientation, be constructed and operated in a specific orientation. Therefore, the terms describing the positional relationship in the drawings are only used for exemplary explanations and cannot be understood as a limitation of this patent. For ordinary technicians in this field, the specific meanings of the above terms can be understood according to specific circumstances. In addition, if there are descriptions involving "first", "second", etc. in the embodiments of the present invention, the descriptions of "first", "second", etc. are only used for descriptive purposes, and cannot be understood as indicating or implying their relative importance or implicitly indicating the number of technical features indicated. Therefore, the features defined as "first" and "second" can explicitly or implicitly include at least one of the features. In addition, the meaning of "and / or" appearing in the full text is to include three parallel schemes. Taking "A and / or B" as an example, it includes scheme A, or scheme B, or schemes that satisfy both A and B.

[0044] Embodiment 1:

[0045] like Figure 1 As shown, a heterogeneous federated learning training method based on a multi-branch neural network model includes the following steps:

[0046] Step 1. Initialization training of the cloud multi-branch model. In the present invention, most of the data is stored in the private data sets of each parameter party device, and only a small amount of public data sets are stored in our cloud server. There is a global multi-branch neural network model in the cloud server such as Figure 2 As shown in the figure, different branch models retain their own output layers on the basis of sharing some common hidden layers. Therefore, the model sizes of each branch are different, which can meet the needs of different computing resource devices. Then, firstly, the global multi-branch model is pre-trained based on the public data set in the cloud, so that it has certain initial performance, which helps to improve the convergence speed and training accuracy of subsequent federated learning.

[0047] Specifically, in the system of the present invention, it can be divided into the cloud side and the device side. In the cloud server, there is a global multi-branch model with K branches, and the overall parameter of the model is recorded as W s , where the parameters of each sub-branch model are recorded as And a public data set, denoted as D0. On the device side, there are N participating devices, and each device i has a private local data set D i and local models

[0048] Since the target tasks of each branch in the multi-branch model are the same, joint optimization training will be performed based on the public dataset D0. First, the cross entropy loss value between the output of each sub-branch model and the ground truth is calculated, and then the weighted average of the loss values ​​of each sub-branch is used as the overall loss value as shown in Formula 1:

[0049]

[0050] Finally, back propagation is performed according to the overall loss value to update the parameters until the model converges, as shown in Formula 2:

[0051]

[0052] Step 2. Matching of sub-branch neural network models: Before the start of federated learning, all devices requesting to participate in federated learning should report the available computing resource information of the device to the cloud server, and in the system of the present invention, the model parameter quantity is used as a measure of computing power. Therefore, when each device i reports its maximum satisfiable parameter quantity P i Then, match a i ≥P k The maximum sub-branch model k, P k It is the parameter quantity of sub-branch model k, and then the selected sub-branch model k is sent to device i.

[0053] Step 3. Local training of device models: After receiving the single-branch model sent by the cloud, each participating device replaces it with the local model of the device, and then trains the current device model based on the local private dataset; however, in each round of local training, Figure 3 As shown, constraints are imposed on the difference between the parameters of the current local model and the original cloud model. Therefore, under the constraints of the cloud model, the shared hidden layers of each device model can obtain parameter distributions as similar as possible during training, which can effectively reduce noise parameters in the subsequent parameter aggregation step.

[0054] In the traditional federated learning algorithm, the device uses the local model to infer the local data to obtain the predicted value, and then calculates the loss value between the predicted value and the true value according to the traditional loss function (such as the cross entropy function), and finally updates the model parameters according to the obtained loss value. However, since the local models of the devices in the system of the present invention are heterogeneous, the distribution of the model parameters obtained by their local training may be inconsistent. In order to make the models of different devices have parameter distributions as similar as possible after local training, it is necessary to impose constraints on the update of model parameters during the training process. The specific approach is to add an L2 regularized loss function to the original cross entropy loss function to measure the difference in parameter distribution between the local model and the cloud model. When device i receives the sub-branch model k, it is initialized before local training. Assuming that the number of shared network layers between the sub-branch model k and the cloud main branch model is H, the L2 regularized loss function of the local training of device i can be expressed as:

[0055]

[0056] Combined with the loss function based on the local data set, the total loss function of device i is shown in Formula 4, where η is a hyperparameter used to measure the proportion of the L2 regularized loss function to the overall loss function.

[0057]

[0058] Finally, the local model parameters of device i are optimized and updated according to formula 5:

[0059]

[0060] Step 4. Aggregation of multi-branch lifting network models in the cloud: After completing local training, each participating device uploads the parameters of its own model to the cloud server, and performs weighted averaging on the parameters of the same parts of the device model and multiple global branch models, thereby aggregating and updating the global branch model in the cloud.

[0061] Since the models uploaded by each device belong to different sub-branches, and the sub-branch models are heterogeneous, it is impossible to directly use the traditional federated learning method to aggregate the models. Considering that the device models are heterogeneous, but at the same time they are part of the global multi-branch model, a multi-branch joint average aggregation method is proposed. The core approach is to average and aggregate the shared layer parameters of each device model. The specific process is as follows:

[0062] When the cloud server receives the model parameter set uploaded by all participating devices Then start multi-branch aggregation; first obtain the dictionary set W of all parameter layers in the global multi-branch model S , the number of layers of the model is recorded as K, and the index of the kth layer is recorded as l k , and the corresponding parameters of the kth layer are expressed as W S [l k ]; then traverse W S All parameter layers in: For each layer k, the layer parameters W S [l k ]First initialize it to zero value and count the index l of this layer k Upload model collection The total number of occurrences in is denoted as Count[l k ], and for every existence l k Device model of the index layer Accumulate the parameters of its corresponding layer into the global model, that is, After traversing and accumulating each layer of the global multi-branch model, the new global model parameter W is obtained and output S .

[0063] Step 5. Distillation training of multi-branch neural network model in the cloud: Due to the inconsistent structures of different branch models, the parameter distribution of each device model will be inconsistent. After the parameters are aggregated, it will inevitably bring some parameter noise to the global model, causing the performance fluctuation of the global model. Therefore, based on the public data set in the cloud and the model gradients uploaded by each device, the aggregated multi-branch network model will be subjected to knowledge distillation learning to restore and improve the prediction accuracy of the multi-branch model. After the distillation training is completed, the model parameters of each branch in the multi-branch model will be sent to the corresponding device for the next round of local training.

[0064] Specifically, after the new global multi-branch model is generated by aggregation in the cloud, distillation training is then performed based on the public data set D0 in the cloud and the model parameters of each device. Similarly, the loss function in distillation training is also divided into two parts: one is the cross entropy loss function with the true label, as shown in Formula 1, and the other is the KL divergence loss function between the soft label output by the device model, as shown in Formula 6.

[0065]

[0066] Therefore, the total loss function of distillation training is shown in Formula 7, where the hyperparameter α is used to set the weight ratio between the cross entropy function and the KL divergence loss.

[0067]

[0068] Then the most optimized global branch model in distillation training can be calculated by formula 8:

[0069]

[0070] After the distillation training is completed, the model parameters of each branch in the global multi-branch model are sent to the corresponding device for the next round of local training.

[0071] Step 6. Application of multi-branch neural network model in the cloud: Repeat steps S3 to S5 until the predetermined number of training rounds is reached, and finally a global branch neural network model is obtained. When it is necessary to deploy or add new device applications, sub-branch models with different reasoning accuracy can be matched according to the resource limitations and requirements of different devices, so as to achieve flexibility and diversity in model deployment. In addition, the multi-branch neural network model can also support end-edge collaborative reasoning. For example, when the computing resources on the device side are insufficient or low latency is required, the device side can only process the calculation of the shared parameter part of the local model, and then send the intermediate results to the cloud global network for the remaining calculations, and finally return the cloud results to the edge device, thereby realizing end-edge collaborative accelerated reasoning.

[0072] Example 2

[0073] This embodiment provides an electronic device, including: a memory, a processor, and a computer program stored in the memory and executable on the processor, wherein the processor executes the computer program to implement the heterogeneous federated learning training method based on a multi-branch neural network model described in Embodiment 1.

[0074] Example 3

[0075] This embodiment provides a computer-readable storage medium on which a computer program is stored. When the computer program is executed by a processor, the heterogeneous federated learning training method based on a multi-branch neural network model described in Example 1 is implemented.

[0076] Obviously, the above embodiments of the utility model are only examples for clearly explaining the utility model, and are not intended to limit the implementation methods of the utility model. For ordinary technicians in the relevant field, other different forms of changes or modifications can be made on the basis of the above description. It is not necessary and impossible to list all the implementation methods here. Any modifications, equivalent substitutions and improvements made within the spirit and principles of the utility model should be included in the protection scope of the claims of the utility model.

Claims

1. A heterogeneous federated learning training method based on a multi-branch neural network model, characterized in that: The following steps are involved: S1. Initialization training of multi-branch models on the cloud: There are global branch models with multiple branches and a public data set D0 in the cloud server. Different global branch models retain their own output layers on the basis of sharing some public hidden layers. First, multiple global branch models are pre-trained based on the public data set D on the cloud to initialize them. S2. Matching of sub-branch neural network models: Before federated learning begins, all devices that request to participate in federated learning report the available computing and storage information of the devices to the cloud server. Then, the cloud server calculates and determines the best branch model suitable for the device based on the collected device information, and distributes the weighted single-branch model to the corresponding device. S3. Local training of device models: After receiving the single-branch model sent by the cloud, each participating device replaces it with the local model of the device, and then trains the current device model based on the local private dataset; S4. Aggregation of multi-branch lifting network models in the cloud: After completing local training, each participating device uploads the parameters of its own model to the cloud server, and performs weighted averaging of the parameters of the same parts of the device model and multiple global branch models, thereby aggregating and updating the global branch model in the cloud; S5. Distillation training of multi-branch neural network model on the cloud: Based on the public data set on the cloud and the model gradients uploaded by each device, knowledge distillation learning is performed on the aggregated multi-branch network model; After the distillation training is completed, the model parameters of each branch in the multi-branch model are sent to the corresponding device for the next round of local training; S6. Application of the multi-branch neural network model in the cloud: Repeat steps S3 to S5 until a predetermined number of training rounds is reached, and finally a global branch neural network model is obtained.

2. The heterogeneous federated learning training method based on a multi-branch neural network model according to claim 1 is characterized in that: The overall parameter of the global branching model is W s , where the parameters of each sub-branch model are recorded as W s k , k = 1, ..., K; On the device side, there are N participating devices, and each device i has a private local data set D i and local models i=1,…,N; in the step S1, the cross entropy loss value between the output of each branch model and the ground truth is calculated, and then the weighted average of the loss values ​​of each sub-branch is taken as the overall loss value: Finally, backpropagation is performed according to the overall loss value to update the parameters until the model converges:

3. The heterogeneous federated learning training method based on a multi-branch neural network model according to claim 2 is characterized in that: In step S2, when each device i reports its maximum satisfiable parameter quantity P i Then, match a i ≥P k The largest sub-branch model k of P is then sent to the selected sub-branch model k to device i, where P k is the parameter quantity of sub-branch model k.

4. The heterogeneous federated learning training method based on a multi-branch neural network model according to claim 2 is characterized in that: In step S3, constraints are imposed on the difference between the parameters of the current local model and the original cloud model, so that under the constraints of the cloud model, the shared hidden layers of each device model can obtain parameter distributions that are as similar as possible during training.

5. The heterogeneous federated learning training method based on a multi-branch neural network model according to claim 4 is characterized in that: The constraints imposed on the difference between the parameters of the current local model and the original cloud model specifically include: adding an L2 regularization loss function to the original cross entropy loss function to measure the difference in parameter distribution between the local model and the cloud model; when device i receives the sub-branch model k, then initialize before local training Assume that the number of shared network layers between the sub-branch model k and the cloud main branch model is H, then the L2 regularized loss function of the local training of device i is expressed as: Combined with the loss function based on the local data set, the total loss function of device i is: Among them, η is a hyperparameter used to measure the proportion of L2 regular loss function to the overall loss function; Finally, the local model parameters of device i are optimized and updated according to the following formula:

6. The heterogeneous federated learning training method based on a multi-branch neural network model according to claim 5 is characterized in that: In step S4, a multi-branch joint average aggregation method is used, specifically including: when the cloud server receives the model parameter set uploaded by all participating devices After i∈{1,…,N}, multi-branch aggregation begins; first obtain the dictionary set W of all parameter layers in the global multi-branch model. S , the number of layers of the model is recorded as K, and the index of the kth layer is recorded as l k , and the corresponding parameters of the kth layer are expressed as W S [l k ]; then traverse W S All parameter layers in: For each layer k, the layer parameters W s [l k ]First initialize it to zero value and count the index l of this layer k Upload model collection The total number of occurrences in is denoted as Count[l k ], and for every existence l k Device model of the index layer Accumulate the parameters of its corresponding layer into the global model, that is, After traversing and accumulating each layer of the global branch model, the new global model parameter W is obtained and output s .

7. The heterogeneous federated learning training method based on a multi-branch neural network model according to claim 6 is characterized in that: The loss functions in distillation training include: the cross entropy loss function with the true label, as shown in Formula 1; and the KL divergence loss function between the soft label output by the device model, as shown in Formula 6 below: The total loss function of distillation training is shown in Formula 7, where the hyperparameter α is used to set the weight ratio between the cross entropy function and the KL divergence loss; The most optimized global branch model in distillation training is calculated by formula 8:

8. The heterogeneous federated learning training method based on a multi-branch neural network model according to claim 7 is characterized in that: In the step S6, the output global multi-branch neural network model can match the appropriate branch model according to the resource limitations and accuracy requirements of different devices to meet the needs of different types of device applications; when the computing resources on the device side are tight or insufficient, the device side can only process the calculation of the shared parameter part of the local model, and then send the intermediate results to the cloud global network for the remaining calculations, and finally return the cloud results to the edge device.

9. An electronic device, comprising: A memory, a processor, and a computer program stored in the memory and executable on the processor, wherein the processor executes the computer program to implement the heterogeneous federated learning training method based on a multi-branch neural network model as described in any one of claims 1 to 8.

10. A computer-readable storage medium having a computer program stored thereon, characterized in that: When the computer program is executed by a processor, the heterogeneous federated learning training method based on a multi-branch neural network model described in any one of claims 1 to 8 is implemented.

Images