A federated learning method based on diffusion model and weight adaptive knowledge distillation

Abstract

The present application relates to the technical field of multi-source heterogeneous data fusion of federated learning, and discloses a federated learning algorithm based on a diffusion model and weight adaptive knowledge distillation, comprising: a client locally trains a diffusion model and generates data conforming to an original image distribution at a server end, and expands the local data of the client by the generated image to achieve data enhancement; and by dynamically assigning weights to each participant when knowledge distillation is performed, knowledge migration is completed at the server end.The present application trains a diffusion model at each client and directly generates data conforming to a global image distribution at the server end, without the need to transmit original data to a central server, so that user privacy can be better protected; and by using generated data for weight adaptive knowledge distillation, the present application can alleviate knowledge forgetting caused when a model is aggregated, so that the generalization capability of a global model is improved.

Description

TECHNICAL FIELD

[0001] The present application relates to the technical field of multi-source heterogeneous data fusion of federated learning, and particularly relates to a federated learning method based on a diffusion model and weight adaptive knowledge distillation. BACKGROUND

[0002] With the rapid development of digitalization of human society, machine learning plays a very important role in the current society. Through machine learning technology, the treasures hidden in data can be automatically mined. Machine learning models trained by a large amount of data have been applied in various scenarios and are deeply changing our world. However, multiple data source parties cannot directly exchange data, forming a "data island" phenomenon, which restricts the further improvement of the ability of artificial intelligence models. With the increasingly strict regulation of privacy data by laws and regulations in recent years, combined with the inherent data fragmentation problem between various departments of companies, all industries are facing the dilemma of data approximation and data islands. How to find an effective means to protect the data privacy of users and improve the generalization ability of the model by combining the data of each participant has become a hot research problem in the field of artificial intelligence. A feasible method is that each institution with a data source trains a model independently using its own data, and then the models of each institution interact with each other, and finally a global model is obtained through model aggregation. In order to ensure user privacy and data security, the process of exchanging model information between institutions will be carefully designed so that no institution can guess the private data content of any other institution. At the same time, when building a global model, its effect is almost the same as that of centralized training when the data sources are integrated together. This is the motivation and core idea of federated learning (FL). Federated learning is a distributed machine learning paradigm that uses data sets distributed among various participants to integrate multi-party data information through privacy protection technology and collaboratively build a global model. The core concept emphasized by federated learning is that data does not move and models move. It can effectively solve the data island problem, allowing participants to jointly model without sharing data, breaking down data islands technically, and realizing multi-party collaboration.

[0003] To solve the problem of insufficient data in machine learning, Connor Shorten et al. summarized some traditional data augmentation methods. Including augmentation methods based on geometric transformation, horizontal flip, cropping, translation, color space transformation, etc. Hongyi Zhang et al. proposed a data augmentation method called Mixup, and later research applied it to federated learning, but because it needs to send mixed data to the server, this method is very expensive in terms of communication; In addition, if a small number of samples are selected for mixing, it may infringe privacy. Another method of data augmentation is to generate data using a generative model, such as Generative Adversial Networks (GAN), Variational Auto-Encoders (VAE), and Flow-based Generative Model. To solve the problem of data privacy in federated learning, some related research has combined Generative Adversial Networks, Variational Auto-Encoders with federated learning to train the generative model and generate pictures under the premise of protecting privacy. Another research on generative models is Diffusion model. The research on Diffusion model was first proposed by Jascha Sohl-Dickstein in 2015; In 2020, Jonathan Ho et al. published a research called Denoising Diffusion Probabilistic Models, which made it truly useful. Diffusion model has recently become a research hotspot, but this research has not yet been applied to federated learning. This patent applies Diffusion model to federated learning, aiming to train Diffusion model to generate data and use it for knowledge distillation under the premise of ensuring data privacy.

[0004] Knowledge distillation is a model compression method proposed by Hinton et al. in 2015. The roles of classical knowledge distillation can be divided into Teacher and Student. Through the "distillation" of the feature representation learned by the complex and powerful network, it is passed to the network with small parameters and weak learning ability. Thus, a fast and powerful network is obtained. In the research of combining federated learning and knowledge distillation, T Shen et al. proposed a FedDML algorithm based on knowledge distillation. Each client has two local models and distills each other. But the calculation cost is high, the communication overhead is large, and the convergence speed of the global network is affected. Zhuangdi Zhu et al. proposed a personalized federated learning algorithm: FedGEN, which aims to make each client get a model suitable for its own data distribution. In FedGEN, the server no longer needs the mean of all local model parameters. The server is only a medium for clients to share knowledge. Shen, Tao, et al. proposed a personalized knowledge distillation method. Each client has a personalized model in addition to the local model. During local training, the local model and the client's personalized model distill each other's knowledge, and upload the distilled local model to the server after training; finally, the server aggregates to generate a global model and sends it to the client. However, these methods require a long training time, increasing the communication overhead. Subsequently, some scholars proposed a Data-Free federated knowledge distillation method. Eunjeong Jeong et al. applied Codistillation (Rohan Anil et al.) to federated learning. Each client uploads the average logit vector of each category in its local data to the server; for client mi, the server averages the logit vectors of all other clients and sends them to mi as distilled knowledge to assist its training. From the perspective of model aggregation, Lin Zhang et al. believe that in the traditional federated learning framework, the global model is directly aggregated by the local model, which will cause knowledge forgetting, resulting in the performance of the global model being worse than that of the local model. Then they proposed the FedFTC framework: explore the input space of the local model through the generator, and use it to transfer knowledge from the local model to the global model. But this requires strict super parameters, and is easy to fall into overfitting. Tao Lin et al. proposed FedDF: by minimizing the KL divergence of the data output by the local model and the global model, the distance between the local model and the global model is narrowed. But in FedDF, due to the too small setting of the distillation temperature, the dark knowledge in the local model cannot be effectively transferred to the global model.

[0005] Privacy protection is the premise of federated learning and the primary challenge faced in federated learning. In addition, in federated learning applications, the accuracy and generalization ability of the model depend on the learning of a large amount of data, but multiple data sources cannot directly exchange data, ultimately leading to the emergence of data heterogeneity. Finally, due to the different network structures of each participant, federated learning may not be able to effectively aggregate the model. Therefore, model heterogeneity has also become one of the main challenges faced by federated learning. SUMMARY

[0006] To solve the above problems, the present patent proposes a federated learning method based on diffusion model and weight adaptive knowledge distillation. The main purposes are as follows:

[0007] 1. Protect privacy: The present patent trains a diffusion model on each client and generates data that conforms to the global image distribution on the server side, without the need to transmit the original data to the center server, which can better protect user privacy;

[0008] 2. Enhance data diversity: By generating images that conform to the original data distribution, not only can the local data be expanded, but also the diversity of the data can be increased, thereby improving the performance and generalization ability of the model;

[0009] 3. Improve the robustness of the model: By generating images that conform to the original data distribution, the model can learn more data distribution and features, thereby improving the robustness of the model;

[0010] 4. Improve the generalization ability of the model: By using generated data for weight adaptive knowledge distillation, the knowledge forgetting caused by model aggregation can be alleviated, thereby improving the generalization ability of the global model.

[0011] To achieve the above purposes, the present application adopts the following technical solutions:

[0012] A federated learning method based on diffusion model and weight adaptive knowledge distillation, comprising:

[0013] The client trains a diffusion model locally and generates data that conforms to the original image distribution on the server side, expands the local data of the client with the generated images, and realizes data enhancement;

[0014] By dynamically assigning weights to each participant during knowledge distillation, and completing the migration of knowledge on the server side.

[0015] Further, the client trains a diffusion model locally and generates data that conforms to the original image distribution on the server side, expands the local data of the client with the generated images, and realizes data enhancement, comprising:

[0016] The training of the diffusion model is completed by minimizing the following loss:

[0017]

[0018] where N represents the total number of clients; w is a local model parameter; is the local data set of the i-th client; D is the data set on all clients; is the cross-entropy loss function; is the target loss of each client:

[0019]

[0020] where T is the total number of training steps, and α t is linear interpolation; z t is the noise diffused at the t-th step; x0 is the real picture of the input; is the expectation of z t and x0; z θ is the noise predicted by the model;

[0021] After training, the server side iterates according to the following formula, that is, the generation of the picture can be completed:

[0022]

[0023] where x t-1 is the picture at the t-1-th step, x t is the picture at the t-th step, σ t is a constant in the derivation process, and z is the variance at the t-th step;

[0024] Through the above process, the shared data set is finally generated.

[0025] Further, the dynamic allocation of weights to each participant during knowledge distillation and the completion of knowledge migration at the server side include:

[0026] The server randomly selects some clients, and the client set is recorded as S t , and sends the global model ω t to the client.

[0027] The client receives the global model and performs local training as a local model;

[0028] The client uploads the trained local model to the server;

[0029] The server aggregates the local models of each client to generate a new round of global model according to the following formula:

[0030]

[0031] wherein is the local dataset of the nth client; is the updated local model of the nth client in the tth round;

[0032] The server dynamically assigns distillation weights to each local model using the shared dataset, and uses the weights to perform knowledge distillation on the global model.

[0033] Further, the optimization objective of the client receiving the global model and performing local training as a local model is:

[0034]

[0035] wherein L i (w) is the empirical loss of the client P i , denoted as:

[0036]

[0037] where w is the model parameter of the client P i , (x, y) is the sample on the local data of the client P i , l i (w; (x, y)) represents the loss of the client P i for the sample (x, y).

[0038] Further, the server dynamically assigns distillation weights to each local model using the shared dataset, and uses the weights to perform knowledge distillation on the global model includes:

[0039] First, the predicted value of the local model on the shared dataset is calculated by formula (7) and the cross-entropy loss between the predicted value and the actual label

[0040]

[0041] wherein y c represents the real label of the data class c; w m represents the local model parameter of the client P m ; is the logits output of the client P m on the shared dataset for class c data; τ is the distillation temperature;

[0042] and the distillation weight γ m is assigned to each participant by formula (8):

[0043]

[0044] where M is the number of clients selected for each round of training;

[0045] The goal of knowledge distillation is to minimize the following loss:

[0046]

[0047] where, is the overall loss for knowledge distillation; γ m is the distillation weight for client P m ; is the logits output of the global model for data of class c on the dataset ;

[0048] The cross-entropy loss of the global model on the shared dataset is calculated by formula (10)

[0049]

[0050] where, is the logits output of the global model for data of class c;

[0051] The final optimization goal is:

[0052]

[0053] The distilled global model is obtained by minimizing formula (11).

[0054] Compared with the prior art, the present application has the beneficial effects of:

[0055] 1. Protecting privacy: The present patent trains the diffusion model at each client and directly generates data conforming to the global image distribution at the server side, without the need to transmit the original data to the center server, which can better protect user privacy;

[0056] 2. Enhancing data diversity: By generating images conforming to the original data distribution, not only can the local data be expanded, but also the diversity of the data can be increased, thereby improving the performance and generalization ability of the model;

[0057] 3. Improving the robustness of the model: By generating images conforming to the original data distribution, the model can learn more data distribution and features, thereby improving the robustness of the model;

[0058] 4. Improving the generalization ability of the model: By using the generated data for weight adaptive knowledge distillation, the knowledge forgetting caused by model aggregation can be alleviated, thereby improving the generalization ability of the global model. BRIEF DESCRIPTION OF DRAWINGS

[0059] Figure 1A flowchart of a federated learning method based on a diffusion model and weight adaptive knowledge distillation according to an embodiment of the present application is shown in FIG. 1.

[0060] Figure 2 A data example for generating a diffusion model according to an embodiment of the present application is shown in FIG. 2. DETAILED DESCRIPTION

[0061] The present application will be further explained in conjunction with the accompanying drawings and specific embodiments:

[0062] As shown in FIG. 3, a federated learning method based on a diffusion model and weight adaptive knowledge distillation (referred to as FedAdKD) includes the following steps: Figure 1

[0063] 1. A data augmentation module is designed. The client locally trains the diffusion model and directly generates data conforming to the original image distribution on the server. The local data is augmented by generating images, thereby achieving data augmentation.

[0064] 2. A federated learning optimization method based on weight adaptive knowledge distillation is proposed. By dynamically assigning weights to each participant during knowledge distillation, and completing knowledge migration on the server side, the knowledge forgetting caused by federated averaging is reduced, and the accuracy of the global model is improved.

[0065] Specifically, the present application first pre-trains the diffusion model. Specifically, the training of the diffusion model is completed by minimizing the following loss:

[0066]

[0067] where N represents the total number of clients, w is the local model parameter, is the local data set of the i-th client, and D is the data set on all clients. is the cross-entropy loss function, is the target loss of each client:

[0068]

[0069] where (T is the total number of training steps), and t is a linear interpolation; z t is the noise diffused at the t-th step, x0 is the input real picture, is the expected value of z t and x0, and z θ is the noise predicted by the model. After training, the server side is iterated according to the following formula, and the generation of the picture can be completed:

[0070] ​

[0071] where x t-1 is the picture at the t-1th step, x t is the picture at the tth step, σ t is a constant in the derivation process, and z is the variance at the tth step.

[0072] Pictures generated by the diffusion model are shown in (a), (b), and (c) in FIG. 1. Figure 2

[0073] Suppose the shared dataset generated by the diffusion model is denoted as V = (x, y) K , where K represents the number of samples. After the above preparation is completed, the present application includes the following five steps:

[0074] Step 1: The server randomly selects some clients, and denotes the set of clients as S t , and sends the global model ω t to them.

[0075] Step 2: The clients receive the global model and perform local training as their local model.

[0076] Specifically, suppose there are N clients, denoted as P1, …, P N , where P1 represents the first client, and so on, and where the client P i has a local dataset w t represents the global model parameters at the tth round, is the local model parameters of the i-th client at the tth round. is the cross-entropy loss function. The optimization objective of the client for local training is:

[0077]

[0078] where L i (w) is the empirical loss of the client P i , and is expressed as:

[0079]

[0080] where w is the model parameters of the client P i , (x, y) is the sample on the local data of the client P i , and l i (w; (x, y)) represents the loss of the client P i for the sample (x, y).

[0081] Step 3: The client uploads the trained local model to the server.

[0082] ​Step 4: The server aggregates the local models of each client to generate a new round of global model by the following formula.

[0083]

[0084] where, is the updated local model of the nth client in the tth round.

[0085] Step 5: The server dynamically assigns the distillation weight to each local model using the shared dataset, and uses this weight to perform knowledge distillation on the global model.

[0086] Specifically, assume that the shared dataset generated by the diffusion model is denoted as where K represents the number of samples. The distillation temperature is denoted as τ. The logits output of the model is denoted as Z = [z 1 ,...,z C ], where C represents the sample category, z c is the value of the cth category in the logits output of the model, c ∈ [1, C]. The selected client in each round of training is denoted as P1,..P m ..,P M , where M is the number of selected clients in each round of training. The present application first calculates the prediction value of each local model (w m represents the local model parameters of client P m ) on between the actual label and the prediction value by formula (7):

[0087]

[0088] and assigns a distillation weight γ m to each participant by formula (8):

[0089]

[0090] The goal of knowledge distillation is to minimize the following loss:

[0091]

[0092] where, is the overall loss of knowledge distillation; τ is the distillation temperature; γ m is the distillation weight of client P m ; is the logits output of client P m on the dataset ; is the logits output of the global model on the dataset The logits output of the global model for data of class c.

[0093] The cross-entropy loss of the global model on the data set

[0094]

[0095] Wherein, y c represents the true label of data class c, is the logits output of the global model for data of class c.

[0096] The final optimization target is:

[0097]

[0098] Table 1. Pseudocode of FedAdKD algorithm

[0099]

[0100] Table 2. Pseudocode of local update algorithm

[0101]

[0102]

[0103] Table 3. Pseudocode of weighted adaptive distillation algorithm

[0104]

[0105] Table 4. Highest accuracy when the total number of training rounds is 200 rounds and the number of local update rounds is 20 rounds

[0106]

[0107]

[0108] As can be seen from the above, the federated learning method (FedAdKD) based on diffusion model and weight adaptive knowledge distillation disclosed in the present application has the best effect, and the accuracy can reach 95.92±0.10 on the Mnist data set. And for Mnist and Cifar-10, the smaller the Dirichlet distribution control data heterogeneity α is, the higher the data heterogeneity is.

[0109] The above only shows the preferred embodiments of the present application, and it should be noted that for ordinary skilled in the art, without departing from the principles of the present application, a number of improvements and refinements can be made, and these improvements and refinements should be considered as the protection scope of the present application.​

Claims

1. A federated learning method based on a diffusion model and weighted adaptive knowledge distillation, characterized in that, include: The client trains a diffusion model locally and generates data on the server that conforms to the distribution of the original image. The generated images are then used to augment the client's local data to achieve data augmentation. By dynamically assigning weights to each participant during knowledge distillation and completing the knowledge transfer on the server side, the process includes: The server-side dynamically assigns distillation weights to each local model using a shared dataset, and then uses these weights to perform knowledge distillation on the global model, including: First, calculate the local model using formula (7). Cross-entropy loss between predicted values ​​and actual labels on a shared dataset Among them, y c The true label representing data category c; w m Indicates client P m Local model parameters; For client P m The logits output for data of category c on the shared dataset; τ is the distillation temperature; And the distillation weight γ is assigned to each participant using formula (8). m : Where M is the number of clients selected in each round of training; The goal of knowledge distillation is to minimize the following loss: in, The overall loss of knowledge distillation; γ m For client P m Distillation weight; For the global model on the dataset The logits output for data of category c; The cross-entropy loss of the global model on the shared dataset is calculated using formula (10). in, This is the logits output of the global model for data of category c; The final optimization goal is: The global model after distillation is obtained by minimizing formula (11).

2. The federated learning method based on diffusion model and weight adaptive knowledge distillation according to claim 1, characterized in that, The client trains a diffusion model locally and generates data on the server that conforms to the original image distribution. The generated images are then used to augment the client's local data, achieving data augmentation through: The diffusion model is trained by minimizing the following loss: Where N represents the total number of clients; w is the local model parameter; Let be the local dataset of the i-th client; D is the dataset on all clients. The cross-entropy loss function; Target loss for each client: in T is the total number of training steps, α t For linear interpolation; z t x represents the noise diffused at step t; x0 represents the input real image. For z t The expected value of x0; z θ The noise in the model prediction; After training is complete, the image can be generated by iteratively performing the following formula on the server side: Where x t-1 For the image at step t-1, x t For the image at step t, σ t Let z be a constant in the derivation process, and z be the variance at step t. The above process ultimately generates a shared dataset.

3. The federated learning method based on diffusion model and weighted adaptive knowledge distillation according to claim 2, characterized in that, The process of dynamically assigning weights to each participant during knowledge distillation and completing the knowledge transfer on the server side includes: The server randomly selects a set of clients, denoted as S. t And send the global model ω to the client. t ; The client receives the global model and trains it locally as a local model. The client uploads the trained local model to the server. The server aggregates the local models from each client using the following formula to generate a new global model: in This is the local dataset for the nth client; This is the local model updated in round t for the nth client; The server dynamically assigns distillation weights to each local model using a shared dataset, and then uses these weights to perform knowledge distillation on the global model.

4. The federated learning method based on diffusion model and weight adaptive knowledge distillation according to claim 3, characterized in that, The optimization objective of the client receiving the global model and training it locally is: Among them, L i (w) is the client P i The loss of experience is expressed as: Where w is the client P i The model parameters, (x,y) are the client P i Samples on local data, l i (w;(x,y)) represents the client P i Loss for sample (x,y).

Images