A federated learning method for device heterogeneity based on knowledge distillation

Abstract

The present invention belongs to the technical field of data information privacy protection, and specifically relates to a device-heterogeneous federated learning method based on knowledge distillation. The system to which the present invention relates includes K clients and one server; each client has a classification model; to efficiently perform knowledge distillation between the client model representation layers, each round of communication is divided into a server training phase and a local training phase; in the server training phase, a generative model is first established on the server with the goal of inferring the posterior distribution of the low-dimensional representation of the sample, and then the trained generative model is transmitted to the client; in the local training phase, the client uses private samples to calculate the task loss, and uses the mean sample output by the generative model to calculate the loss of the tuning representation layer. In this way, after multiple rounds of iteration, each client can obtain a model with higher accuracy than traditional training methods.

Description

Technical Field

[0001] The present invention belongs to the technical field of data information privacy protection, and specifically relates to a federated learning method for device heterogeneity. Background Art

[0002] With the rapid increase in data volumes and the need for privacy protection, federated learning has emerged as a promising area of ​​research. Federated learning typically consists of many clients with access only to private data and a central server that coordinates the learning process but does not have access to any of the original data. Its goal is to train a global model on a central server using distributed data without explicitly sharing private data. One challenge facing this approach is data heterogeneity: model performance degrades when different clients have different data distributions. Existing approaches to address data heterogeneity are mostly based on gradient fusion, which must be performed under the condition of local model homogeneity.

[0003] This invention focuses on the problem of federated learning with heterogeneous devices. In this context, the varying storage, computing, and communication capabilities of individual clients result in varying local model structures, presenting significant challenges to existing approaches. In some practical federated learning scenarios, training is required on devices with significantly different hardware. When the model structure is complex, devices with fewer resources cannot participate in training; when the model structure is simple, devices with more abundant resources are underutilized.

[0004] To this end, this paper proposes a federated learning algorithm based on knowledge distillation. This algorithm allows each client to build a personalized model, addressing both model and data heterogeneity. The algorithm divides each round of communication into two phases. During the server training phase, a generative model is first built on the server, aiming to infer the posterior distribution of the sample's low-dimensional representation. This trained generative model is then transferred to the client. During the local training phase, the client calculates the task loss using private samples and the loss of the fine-tuned representation layer using the mean sample output from the generative model. This approach, after multiple rounds of iteration, allows each client to obtain a model with higher accuracy than traditional training methods. Summary of the Invention

[0005] The purpose of this invention is to propose a federated learning method for device heterogeneity so as to perform federated learning in scenarios where client resources vary greatly, thereby providing a strong guarantee for mining data information.

[0006] The device heterogeneity-oriented federated learning method proposed in the present invention is based on knowledge distillation technology; the system involved includes K clients and 1 server; wherein;

[0007] Each client has a classification model configured based on its hardware and software resources. The client's classification model is divided into a representation layer and a decision layer. The representation layer is used to map samples into low-dimensional representations, and the decision layer is used to map low-dimensional representations into probability vectors. The objective function for knowledge distillation between clients is defined as follows:

[0008]

[0009] Where K is the number of clients; X k is the private dataset of the kth client, x and y are samples and labels; L k is the loss function of the classification task; φ(·) is the decision layer function of the client model, f k (·) is the presentation layer function of the client model, θ k is the representation layer parameter; γ is a hyperparameter; the representation layer structure and parameters of the classification model on each client are different, while the structure and initialization parameters of the decision layer are the same.

[0010] A generative model is set up on the server, consisting of a mean function and a variance function of shared inputs. The goal of the system is to efficiently solve equation (1). To this end, the method of the present invention divides each round of communication into two stages: a server training stage and a local training stage. In the server training stage, a generative model is first established on the server with the goal of inferring the posterior distribution of the low-dimensional representation of the sample, and then the trained generative model is transmitted to the client. In the local training stage, the client calculates the task loss using private samples and calculates the loss of the tuning representation layer using the mean sample output by the generative model. In this way, after multiple rounds of iteration, each client can obtain a model with higher accuracy than traditional training methods.

[0011] Specifically:

[0012] During the server training phase, the server first collects the representation layers of all client models; collects the low-dimensional representations collected by the client to form a set Z. Then, the variational inference method is used to solve the posterior distribution Specifically:

[0013] Assumptions It is a Gaussian distribution. The generative model takes Z as input, and the mean function and variance function output respectively. The mean and variance of , that is:

[0014]

[0015] Where Z is the set of low-dimensional representations, is the sample set estimated by the generative model; U(·) and θ U are the mean function and parameters of the generative model; V(·) and θ Vare the variance function and parameters of the generative model respectively; is a Gaussian distribution with mean u and variance v;

[0016] from The samples sampled in are input into all the collected representation layers respectively to obtain K low-dimensional representations; if the label of the sample is j, then Find the weighted mean of these low-dimensional representations for the weights, and calculate the Euclidean distance loss between the weighted mean and Z (where N j is the total number of class j samples of all clients, N k，j is the total number of samples of the jth category on the kth client).

[0017] On the other hand, calculation The KL divergence loss with the standard Gaussian distribution. The above process uses formula (3) as the loss function to train the generative model using stochastic gradient descent:

[0018]

[0019] Among them, f k (·) is the presentation layer function of the client model, θ k is the representation layer parameter; λ is a hyperparameter; is the standard Gaussian distribution; KL(·) is the KL divergence (Kullback-Leibler divergence); w k is the client weight; the meanings of other symbols are the same as in (2);

[0020] Finally, the server transmits the mean function of the trained generative model to all clients.

[0021] In the local training phase, the client first receives the mean function from the server. Then, on the one hand, it uses the private dataset X k Calculate the classification task loss; on the other hand, collect the low-dimensional representations generated in the process of calculating the classification task loss, input the low-dimensional representation into the mean function to obtain the mean sample, and calculate the mean sample and X k The Euclidean distance loss is used. The above process uses formula (4) as the loss function to train the local classification model using stochastic gradient descent method:

[0022]

[0023] Among them, X k is the private dataset of the kth client, x and y are samples and labels; L k is the loss function of the classification task; φ(·) is the decision layer function of the client model, f k (·) is the presentation layer function of the client model, θ kis the representation layer parameter; λ′ is a hyperparameter; the meanings of other symbols are the same as in (2);

[0024] Next, the client collects low-dimensional representations. Specifically, the client continues to execute the stochastic gradient descent process. After every q rounds, the low-dimensional representations obtained during this period are averaged by label. The collection stops after obtaining at least c low-dimensional representation means.

[0025] Finally, the client uploads the classification model representation layer and the collected low-dimensional representation mean to the server.

[0026] Repeat the server training phase and the local training phase. After multiple rounds of iterations, each client can obtain a model with higher accuracy than the traditional training method.

[0027] The main features and advantages of the present invention are:

[0028] First, it allows different clients to have differently distributed data and differently structured model representation layers, addressing both model and data heterogeneity challenges and expanding application scenarios.

[0029] Second, it allows each client to build a personalized model. In scenarios with heterogeneous data, it can achieve higher average accuracy than other methods that build a global model.

[0030] Third, it optimizes the representation layer based on knowledge distillation technology, obtaining information from relevant clients while reducing interference from irrelevant clients, making the local model more accurate than the method of optimizing the decision layer or other traditional methods; Fourth, during the communication process, the client does not upload the original value of the low-dimensional representation but the mean, and does not upload the entire client model but only the representation layer, which avoids the server using the client model parameters or low-dimensional representation to infer client data. BRIEF DESCRIPTION OF THE DRAWINGS

[0031] Figure 1 It is a network structure diagram of the present invention.

[0032] Figure 2 This is an example of image data.

[0033] Figure 3 is the client-side sampled low-dimensional representation.

[0034] Figure 4 It is the client that generates the global representation.

[0035] Figure 5 is the client-side CNN model.

[0036] Figure 6 is the client-side MLP model.

[0037] Figure 7It is a generative model. DETAILED DESCRIPTION

[0038] The specific embodiments of the present invention will be described below with reference to the accompanying drawings.

[0039] The network structure diagram of the present invention is as follows: Figure 1 shown.

[0040] First, deploy an appropriate classification model based on the client's hardware and software conditions. Generally speaking, you can set up a complex model on a resource-rich client and a simple model on a resource-limited client.

[0041] The client then trains a local model using local image data and samples low-dimensional representations during training. After a round of training, the client uploads the obtained low-dimensional representations and representation layer parameters to the server.

[0042] The server collects the low-dimensional representations and representation layer parameters uploaded by the client to train a generative model and sends the trained generative model to each client. After the client updates the generative model, it uses local image data to train the local model. Furthermore, the low-dimensional representations obtained during training are fed into the generative model to generate mean samples, which are then used to tune the model's representation layer.

[0043] The following describes the classification task of the MNIST dataset. Figure 2 The images shown are examples from the MNIST dataset. The MNIST dataset consists of 10 digit labels, 60,000 training examples, and 10,000 test examples. The training examples are randomly distributed to 100 clients with varying resources. Each client has 300 images from two categories. The test examples are also distributed to the same 100 clients, maintaining the same distribution as the training examples. A classification model is trained for each client using federated learning.

[0044] First, because the resources of the 100 clients are different, it is necessary to deploy a suitable classification model based on the hardware and software conditions. For example, for clients with relatively abundant resources, a CNN model can be set up with a model structure such as Figure 5 As shown, Representor identifies the representation layer and Predictor identifies the decision layer. Similarly, for clients with relatively insufficient resources, an MLP model is set, and the model structure is as follows: Figure 6 As shown in Figure 2, the CNN model and the MLP model have the same decision layer, so the output vectors of their representation layers have the same length.

[0045] Then, before training begins, the parameters of the decision layers of different clients need to be synchronized. During training, the client randomly reads the local image data input model, reading 50 images each time, calculates the loss and gradient, and updates the model parameters. After training 8*75 / 50=12 times, the client continues to randomly read the local image data input model representation layer, reading 50 images each time, calculates the low-dimensional representation and classifies it by label, and calculates the representation mean of each label every 4 readings. After obtaining 10 representation means, upload them together with the representation layer parameters to the server. Figure 3 shown.

[0046] Assume that in this round of communication, the server connects to 100 clients and collects 100*10=1000 low-dimensional representations uploaded by the clients. Use these 1000 representation vectors and representation layer parameters to train the generative model. The structure of the generative model is as follows: Figure 7 As shown in Figure 1, the labels (f1, f21) constitute the mean function and the labels (f1, f22) constitute the variance function, which share the input layer f1. Finally, the mean function part of the trained generative model is sent to the connected clients.

[0047] After the client updates the generative model, it uses local image data to train the local model. On the other hand, it inputs the low-dimensional representation obtained during the training process into the generative model to obtain the mean sample, such as Figure 4 As shown in Figure 2, the mean sample looks like noise and does not leak information. The Euclidean distance between the mean sample and the local sample is calculated as the regularization term of the loss function to tune the representation layer of the model.

[0048] Repeating the above training steps, after 10 rounds of iterations, the average accuracy of the client model can reach 98.65% in the test.

Claims

1. A device-heterogeneous federated learning method based on knowledge distillation, wherein the system includes K clients and one server; the characteristics are: Each client has a classification model configured based on its hardware and software resources. The client's classification model is divided into a representation layer and a decision layer. The representation layer is used to map samples into low-dimensional representations, and the decision layer is used to map low-dimensional representations into probability vectors. The objective function for knowledge distillation between the client model representation layers is defined as follows: Where K is the number of clients; X k is the private dataset of the kth client, x and y are samples and labels; L k is the loss function of the classification task; φ(·) is the decision layer function of the client model, f k (·) is the presentation layer function of the client model, θ k is the representation layer parameter; γ is a hyperparameter; the representation layer structure and parameters of the classification model on each client are different, while the structure and initialization parameters of the decision layer are the same; A generative model is set on the server, which consists of a mean function and a variance function of the shared input. The goal of the system is to efficiently solve equation (1). To this end, each round of communication is divided into two phases: a server training phase and a local training phase. During the server training phase, a generative model is first built on the server with the goal of inferring the posterior distribution of the low-dimensional representation of the sample, and then the trained generative model is passed to the client. During the local training phase, the client uses private samples to calculate the task loss and uses the mean sample output by the generative model to calculate the loss of the tuning representation layer. In this way, after multiple rounds of iterations, each client can obtain a model with higher accuracy than traditional training methods.

2. The device heterogeneity-oriented federated learning method based on knowledge distillation according to claim 1 is characterized in that: During the server training phase, the server first collects the representation layers of all client models; collects the low-dimensional representations collected by the client to form a set Z; then, uses the variational inference method to solve the posterior distribution Specifically: Assumptions It is a Gaussian distribution. The generative model takes Z as input, and the mean function and variance function output respectively. The mean and variance of , that is: Where Z is the set of low-dimensional representations, is the sample set estimated by the generative model; U(·) and θ U are the mean function and parameters of the generative model; V(·) and θ V are the variance function and parameters of the generative model respectively; is a Gaussian distribution with mean u and variance v; from The samples sampled in are input into all the collected representation layers respectively to obtain K low-dimensional representations; if the label of the sample is j, then Find the weighted mean of these low-dimensional representations for the weights, and calculate the Euclidean distance loss between the weighted mean and Z, where N j is the total number of class j samples of all clients, N k，j is the total number of samples of the jth class on the kth client; On the other hand, calculation The KL divergence loss with the standard Gaussian distribution is used; the generative model is trained using the stochastic gradient descent method with formula (3) as the loss function: Among them, f k (·) is the presentation layer function of the client model, θ k is the representation layer parameter; λ is a hyperparameter; is the standard Gaussian distribution; KL(·) is the KL divergence; w k is the client weight; the meanings of other symbols are the same as in (2); Finally, the server transmits the mean function of the trained generative model to all clients.

3. The device heterogeneity-oriented federated learning method based on knowledge distillation according to claim 2, characterized in that: In the local training phase, the client first receives the mean function from the server; then, on the one hand, it uses the private dataset X k Calculate the classification task loss; on the other hand, collect the low-dimensional representations generated in the process of calculating the classification task loss, input the low-dimensional representation into the mean function to obtain the mean sample, and calculate the mean sample and X k The Euclidean distance loss is used to train the local classification model using stochastic gradient descent method with formula (4) as the loss function: Among them, X k is the private dataset of the kth client, x and y are samples and labels; L k is the loss function of the classification task; φ(·) is the decision layer function of the client model, f k (·) is the presentation layer function of the client model, θ k is the representation layer parameter; λ′ is a hyperparameter; the meanings of other symbols are the same as in (2); Next, the client collects low-dimensional representations. Specifically, the client continues to perform the stochastic gradient descent process, and after every q rounds, the low-dimensional representations obtained during this period are averaged by label. The collection stops after obtaining at least c low-dimensional representation means. Finally, the client uploads the classification model representation layer and the collected low-dimensional representation mean to the server; Repeat the server training phase and the local training phase. After multiple rounds of iterations, each client will obtain a model with higher accuracy than the traditional training method.

Images