Federal learning data balance training method and device based on CVAE and multi-head attention mechanism

Abstract

The invention discloses a federal learning data balance training method and device based on CVAE and a multi-head attention mechanism, and aims to solve the problem of unbalanced data distribution of local equipment and the problems of low precision and slow convergence of a federal learning model caused by Non-IID data problems in the fields of data security technology and mobile edge computing. The invention provides a federal learning data balance training method based on CVAE and a multi-head attention mechanism. Compared with an existing federal learning algorithm, the method has the advantages that local missing feature sample data is generated according to global feature distribution, so that the unbalance degree of data distribution between local devices is reduced. Considering that the local model deviates to feature extraction under Non-IID data, so that model training is trapped in a local optimal solution and even becomes worse, the teacher model is used to help the local model learn the features of the data, so that the student model is ensured to learn richer feature representation, and the generalization ability of the model can be improved.

Description

Technical Field

[0001] The present invention belongs to the field of data security technology and mobile edge computing, and specifically to a federated learning data balancing training method and device based on CVAE and multi-head attention mechanism. Background Art

[0002] As the amount of data continues to increase, more and more local data needs to be uploaded to the cloud for processing and calculation. The centralized data processing method can easily cause the leakage of user privacy and threaten user security. Edge computing, as an emerging computing model, complements federated learning to solve this problem well. The process is that the cloud server sends the global model to the terminal device, and the terminal device uses the local data set to train the model. It only needs to upload the trained model parameters to the edge server for global aggregation, and then continuously iterate the process to finally obtain a high-quality global model. The process does not involve the sharing of the data itself, so that the user's privacy is protected to a certain extent.

[0003] Federated learning has been widely used in many industries, such as industrial Internet and smart healthcare. Through federated learning, devices can collect and analyze data locally, realize the prediction and preventive maintenance of equipment failures, and reduce downtime and maintenance costs. At the same time, various production equipment can collaboratively optimize the production process without sharing commercially sensitive data, improve production efficiency and product quality. The currently popular federated learning algorithm is the Federated Average Algorithm (FedAvg), which can deal well with independent and identically distributed (IID) data. However, in the actual scenarios of most distributed mobile devices, the data distribution is non-independent and identically distributed (Non-IID), that is, the data distribution of local devices is different, so that when the model updates the weights, different devices will update in different directions. FedAvg cannot obtain ideal results by simply performing average aggregation processing. In cross-user Non-IID data, the usual approach is to impose additional rounds on federated learning to meet the accuracy required by the global model. However, this consumes a lot of communication. The present invention hopes to reduce the imbalance of data between local devices by increasing the corresponding category data. At the same time, the present invention found in the research that since the global model is an aggregation of all local training models, it has the ability to more accurately represent data features. However, the distribution of Non-IID local device data is unbalanced, and the feature representation of model training has a certain bias, which makes the model training results far away from the actual optimal results, and ultimately leads to worse global model performance.

[0004] In order to solve this statistical challenge of federated learning, the present invention quantifies the weight difference between centralized training and distributed training with the same weight divergence. Then, the present invention finds that the distributed weight divergence comes from the difference between the distribution of each device and the global distribution and the deviation of the learning of data features from the optimal solution caused by the imbalance of local data. Taking the above factors into consideration, the present invention proposes a data balanced training method for federated learning based on CVAE and multi-head attention mechanism, which can deal with the weight divergence problem caused by Non-IID data, thereby improving the accuracy and convergence speed of the global model. Summary of the invention

[0005] In response to the problem of federated learning weight divergence in the field of data security technology and mobile edge computing, the present invention proposes a federated learning data balanced training method based on CVAE and multi-head attention mechanism. It is a method that optimizes data imbalance through data synthesis and reduces the gradient divergence between different clients by adjusting the gradient control variables at the global and local levels. It can deal with the problems of slow convergence and low accuracy of the global model caused by Non-IID data.

[0006] The solution of the present invention realizes the synthesis of local data, making the data distribution between local devices more balanced, thereby solving the problem of model divergence.

[0007] The local device data distribution is approximated to the global data distribution, and the multi-head attention mechanism and knowledge distillation are used to help the local model learn the data features.

[0008] In order to achieve the above objectives, the present invention provides a federated learning data balancing training method based on CVAE and multi-head attention mechanism, which is used to suppress weight divergence, including:

[0009] Step 1: Train a global Conditional Variational Autoencoder (CVAE) model on the server side and distribute it to local devices.

[0010] Step 2: The local device uses the CVAE model to obtain the latent variable z for each category and uploads it to the server.

[0011] Step 3: The server uses the locally uploaded latent variables z to calculate the global latent features for each category And the initial parameter w 0 and global latent features Distributed to each local device. Local devices use global latent features The decoder of the CVAE model synthesizes data for unbalanced categories and realizes CVAE data reconstruction.

[0012] Step 4: The local device combines the synthetic data with the local data, uses the multi-head attention mechanism to help the Convolutional Neural Network (CNN) obtain the output of the student model (student_out) in the federated learning model, and adds the loss of the student model output and the teacher model output (teacher_out) to the optimization target, and updates the local model parameter w i ; After the local model training is completed, the model parameters w i Upload to the server.

[0013] Step 5: The server aggregates parameters to obtain the new global model parameter w t+1 .

[0014] Repeat steps 2-5 until the global model converges.

[0015] Furthermore, the specific process of data balancing in steps 1-3 includes the following:

[0016] 1) Use the CVAE model to sample potential feature variables:

[0017] The CVAE model includes an encoder and a decoder. The encoder is used to infer the posterior probability q(z|x,y), where y represents the "condition", which contains label information, y is encoded as a one-hot vector, x represents the sample feature, and z represents the latent variable.

[0018] To obtain the sample feature x, introducing the Multi-Head Attention (MHA) mechanism in the Convolutional Neural Network (CNN) can significantly improve the model feature extraction capability, especially for important data features. The process of feature weighting by the Multi-Head Attention mechanism is as follows:

[0019] First, the local features of category c in all devices are stacked into a feature matrix:

[0020]

[0021] Among them, H c ∈R N*d , represents the local feature representation of device i for category c, the dimension is d, and N represents the number of devices.

[0022] Secondly, perform linear transformation to obtain the query Q with different heads i , key K i Sum value V i matrix:

[0023] Q i =H c W iQ ,K i =H c W i K ,V i =H c W i V for i=1,2,...,h

[0024] Where W i Q , W i K and W i V is the weight matrix of the i-th head, and h is the number of heads;

[0025] Next, calculate the attention output of each head:

[0026]

[0027] Among them, d k is the dimension of the key vector;

[0028] Next, connection and linear transformation:

[0029] Concatenate the outputs of all heads and integrate them through a linear transformation:

[0030] x=MultiHead(Q,K,V)=Concat(head1,head2,...,head h )W 0

[0031] Among them, head i =Attention(Q i ,K i ,V i ), W 0 is the output weight matrix.

[0032] The local participants in federated learning need to initially upload the latent variables z of each category of the local device to the server. The latent variable z is inferred from the data x and label y through the conditional distribution q(z|x,y). Assume that this conditional distribution is a Gaussian distribution, whose mean and log variance are calculated by a neural network. Assume that μ(x,y) and logσ 2 (x, y) are the mean and logarithmic variance calculated from x and y respectively, then the calculation formula for the latent variable z is as follows:

[0033] z~q(z|x,y)=Ν(z;μ(x,y),σ 2 (x,y)

[0034] By reparameterization technique, it can be expressed as:

[0035]

[0036] Among them, ξ~N(0,I) is the sampling noise, Represents element-wise multiplication.

[0037] 2) Local participant initialization: The server calculates the mean of each class latent variable z uploaded locally, which is defined as the global latent feature And the global latent features of the latent variables and the initial parameters w of the global model 0 Distributed to the local participants’ devices. The global latent features The calculation formula is as follows:

[0038]

[0039] in is the potential characteristic of class c of the ith device.

[0040] 3) CVAE data reconstruction: local participants use their own potential features and the global potential features distributed by the server to reconstruct the data. Generate missing features of its own category and finally achieve data reconstruction.

[0041] During the reconstruction process, the local device uses the decoder to Model the conditional likelihood and output labeled samples. One-hot encode the missing labels, combine them with the latent feature variables, and use CVAE to maximize the variational lower bound of the marginal likelihood p(y|x,z) to generate high-quality specified samples.

[0042] The calculation formula for reconstructing data is:

[0043]

[0044] Among them, z is the latent variable, x is the sample feature, It is the reconstructed data, that is, the generated data is obtained by decoding x and z.

[0045] 4) The design process of the optimization objective function for data balance is as follows:

[0046] Design reconstruction error (BCE):

[0047]

[0048] Here x i is the ith element of the original input, is the i-th element of the reconstructed output, and N is the total number of samples.

[0049] Design KL Divergence (KLD):

[0050]

[0051] where μ j and is the mean and variance of the jth latent variable, and J is the total number of latent variables.

[0052] In summary, the objective function can be expressed as:

[0053] l(θ,φ;x)=min(BCE+KLD)

[0054] where θ is the decoder parameter and φ is the encoder parameter, i.e. they represent all the network parameters involved, e.g., μ j , N, J, min() means taking the minimum value.

[0055] Furthermore, the feature training method of the multi-head attention mechanism in step 4-5 is different from the initialized multi-head attention mechanism in that in addition to extracting the local original data features, it also extracts the generated data to improve the learning ability of local features. Specific process:

[0056] 1) Build a CNN network with a multi-head attention mechanism: Here, the multi-head attention mechanism will be applied after the convolution layer and before the fully connected layer of the CNN network. The present invention uses CNN to obtain the primary features of the local device, and then uses the multi-head attention mechanism to generate enhanced feature representations. The CNN model includes three outputs, namely the output features of the feature extraction network, the intermediate feature representation, and the final output. The multi-head attention mechanism performs multiple independent attention operations (called "heads") in parallel, then connects their outputs and integrates them through linear transformations, so that the features of different subspaces can be captured.

[0057] Assuming that the input has three matrices: query matrix Q, key matrix K and value matrix K, the operation steps of the multi-head attention mechanism are as follows:

[0058] 1. Linear transformation:

[0059] Perform linear transformation on Q, K and V to obtain query, key and value matrices of different heads:

[0060] Q i =QW i Q ,K i =KW i K ,V i=VW i V for i=1,2,...,h

[0061] Among them, W i Q , W i K and W i V is the weight matrix of the ith head, and h is the number of heads.

[0062] 2. Calculate the attention output of each head:

[0063]

[0064] Among them, d k is the dimension of the key vector, and the subscript i represents the i-th head.

[0065] 3. Connection and linear transformation:

[0066] Concatenate the outputs of all heads and integrate them through a linear transformation:

[0067] MultiHead(Q,K,V)=Concat(head1,head2,...,head h )W 0

[0068] Among them, head i =Attention(Q i ,K i ,V i ), W 0 is the output weight matrix.

[0069] 2) The student model learns the features extracted from the teacher model:

[0070] In federated learning, the teacher model has a better representation of feature learning, while the local representation of features is biased due to data imbalance. By learning from the teacher model, the classification performance and generalization ability of the student model can be improved.

[0071] 3) Establishing the optimization objective function: The present invention combines the cross entropy loss function output by the local model with the knowledge distillation loss, that is, its overall objective function It consists of two parts, which are defined as follows:

[0072] Cross entropy loss function:

[0073]

[0074] Among them, y i is a one-hot vector of the true label, is the probability distribution predicted by the model.

[0075] Knowledge distillation loss function:

[0076]

[0077] Among them, h s,i is the feature output of the i-th sample of the student model, h t,i is the feature output of the i-th sample of the teacher model.

[0078] In summary, the overall objective function is:

[0079]

[0080] In the formula, α is a hyperparameter used to smooth the two parts of the loss value.

[0081] 4) Update local model parameters: The local device calculates the local stochastic gradient (SGD) through training, updates the local model parameters, and sends the updated model parameters to the server to participate in the aggregation of global model parameters.

[0082]

[0083] Among them, here represents the model parameters of the participant device i in the t+1 round of communication, and η is the learning rate. 5) Aggregate the local model to update the global model parameters:

[0084]

[0085] Among them, n represents the total amount of data of all participants, n k represents the data volume of the kth participant. t+1 It is the global model parameter of the t+1th round.

[0086] In the federated learning, the present invention needs to perform multiple rounds of communication to ensure that the global model obtains excellent performance. The federated learning repeats the above steps 2-5 until the global model converges.

[0087] Based on the above method, the present invention also proposes a device for federated learning data balance training, including a server device and several local devices; the server device can execute the contents of steps 1, 3, and 5 above; the local device can execute the contents of steps 2 and 4 above.

[0088] Beneficial effects of the present invention:

[0089] (1) The present invention takes into account the problem of unbalanced data distribution in local devices, and provides a federated learning data balance training method based on CVAE and multi-head attention mechanism to solve the problem of low accuracy and slow convergence of federated learning models caused by Non-IID data problems in the fields of data security technology and mobile edge computing. Compared with the existing federated learning algorithms, the present invention generates locally missing feature sample data based on the global feature distribution, thereby reducing the imbalance of data distribution between local devices.

[0090] (2) Considering that the local model is biased in feature extraction under Non-IID data, the model training falls into the local optimal solution or even becomes worse. Therefore, the present invention uses a teacher model to help the local model learn the features of the data, thereby ensuring that the student model learns richer feature representations, which can improve the generalization ability of the model. BRIEF DESCRIPTION OF THE DRAWINGS

[0091] Figure 1 Workflow diagram for implementing the data balance training method for federated learning based on CVAE and multi-head attention mechanism;

[0092] Figure 2 Data balancing flowchart for training and generating data using the CVAE model;

[0093] Figure 3 Flowchart of the feature training method for the multi-head attention mechanism; DETAILED DESCRIPTION

[0094] In order to further illustrate the technical solution of the present invention, the technical solution of the present invention is analyzed in detail below in conjunction with relevant drawings, and detailed implementation steps are given. The examples in the present invention are only a part of the relevant examples, so the protection scope of the present invention is not limited to this implementation example.

[0095] like Figure 1 As shown in the figure, the present invention shows the three stages of federated learning: local model initialization, local training and model aggregation. The present invention optimizes the initialization and local training of the local model of federated learning by implementing optimization schemes for the two parts. The process includes data enhancement to balance local data and multi-head attention mechanism feature extraction, and the specific steps are as follows:

[0096] Step S1: Initialization: The local device uses CVAE to sample local data, obtains the latent variable z of its own data, and uploads it to the server. The server aggregates the uploaded latent features to obtain the feature distribution of the global category. Next, the server initializes the local model parameters and initializes the global model parameters w 0 and the global average latent feature Distributed to each local device. The calculation formula for CVAE to obtain the potential feature z of local data is as follows:

[0097] z~q(z|x,y)=Ν(z;μ(x,y),σ 2 (x,y)

[0098] By reparameterization technique, it can be expressed as:

[0099]

[0100] Among them, ξ~N(0,I) is the sampling noise, Represents element-wise multiplication.

[0101] Step S2: Data balancing: The process is as follows Figure 2 As shown in the figure, due to the imbalance of local data, the data features are missing. The local device uses the CVAE model to achieve the balance of local device data. Using CVAE to maximize the marginal likelihood The variational lower bound of , generates high-quality specified samples.

[0102] The calculation formula for reconstructing data is:

[0103]

[0104] Reconstruction error (BCE):

[0105]

[0106] KL Divergence (KLD):

[0107]

[0108] To sum up, the optimization objective function of the local model can be expressed as:

[0109] l(θ,φ;x)=min(BCE+KLD)

[0110] Among them, min() means taking the minimum value.

[0111] Step S3: Local training: This method adds a multi-head attention mechanism to the existing CNN network model to help the local CNN model pay attention to more data features and reduce the deviation of the local model feature representation. After training the balanced data, the teacher model is used to train the feature representation of the student model. The process is as follows: Figure 3 As shown, its objective function It consists of two parts:

[0112] Cross entropy loss function:

[0113]

[0114] Knowledge distillation loss function:

[0115]

[0116] In summary, the objective function is:

[0117]

[0118] Here α is a hyperparameter used to smooth the two loss values.

[0119] The local device calculates the local stochastic gradient (SGD) through training, updates the local model parameters, and sends the updated model parameters to the server. The local model parameters are updated as follows:

[0120]

[0121] Step S4: Server-side model aggregation: The server aggregates the model parameters uploaded by the local device using the FedAvg method. The aggregated model parameters are used to update the global model parameters w of the previous round. t . Repeat the above training until the global model converges.

[0122] The global model is updated as follows:

[0123]

[0124] The series of detailed descriptions listed above are only specific descriptions of feasible implementation methods of the present invention. They are not intended to limit the scope of protection of the present invention. All equivalent methods or changes that do not deviate from the technical creation of the present invention should be included in the scope of protection of the present invention.

Claims

1. A federated learning data balance training method based on CVAE and multi-head attention mechanism, characterized by: These include: Step 1: Train a global CVAE model on the server and distribute it to local devices; Step 2: The local device uses the CVAE model to obtain the latent variable z for each category and uploads it to the server; Step 3: The server uses the latent variable z uploaded by the local device to calculate the average feature of each category And the initial parameters of the global model and the global potential features Distributed to each local device, the local device uses The decoder of CVAE synthesizes data for imbalanced categories and realizes CVAE data reconstruction; Step 4: The local device combines the synthetic data with the local data, uses the multi-head attention mechanism to help the CNN network obtain the student model output, and adds the loss of the student model output and the teacher model output to the optimization target, and updates the local model parameters; after the local model training is completed, the model parameters are uploaded to the server; Step 5: The server aggregates the global model parameters and finally obtains new global model parameters; Step 6: Repeat steps 2-5 until the global model converges.

2. The method for balanced training of federated learning data based on CVAE and multi-head attention mechanism according to claim 1, characterized in that: The CVAE model in steps 1-3 includes an encoder and a decoder; the encoder is used to infer the posterior probability q(z|x,y), where y represents "condition", which contains label information and is encoded as a one-hot vector, x represents sample features, and z represents latent variables; the decoder is used to locally model the conditional likelihood through p(y|x,z) and output labeled samples, one-hot encode the missing labels, combine them with the latent variable z, and use CVAE to maximize the variational lower bound of the marginal likelihood p(y|x,z) to generate high-quality samples.

3. The method for balanced training of federated learning data based on CVAE and multi-head attention mechanism according to claim 1, characterized in that: To obtain the feature x in step 2, a multi-head attention mechanism (MHA) is introduced in the convolutional neural network (CNN). The process of weighting the feature by the multi-head attention mechanism is as follows: First, the local features of category c in all devices are stacked into a matrix: Among them, H c ∈R N*d , represents the local feature representation of device i for category c, with dimension d; Perform linear transformation to obtain the query, key, and value matrices of different headers: in and is the weight matrix of the i-th head, and h is the number of heads; Next, calculate the attention output of each head: Among them, d k is the dimension of the key vector; Next, concatenation and linear transformation: concatenate the outputs of all heads and integrate them through linear transformation: x=Concat(head1,head2,...,head h )W 0 Among them, head i =Attention(Q i ,K i ,V i ), W 0 is the output weight matrix. In step 2, the latent variable z is inferred from the feature x and the label y through the conditional distribution q(z|x,y), assuming that this conditional distribution is a Gaussian distribution with μ(x,y) and logσ 2 (x, y) are the mean and logarithmic variance calculated from x and y respectively, then the calculation formula for the latent variable z is as follows: z~q(z|x,y)=Ν(z:μ(x,y),σ 2 (x,y)) By reparameterization technique, it can be expressed as: Among them, ξ~N(0,I) is the sampling noise, Represents element-wise multiplication.

4. The method for balanced training of federated learning data based on CVAE and multi-head attention mechanism according to claim 1, characterized in that: In step 3, the global latent feature The calculation formula is as follows: in is the potential feature of class c of the ith device, and N is the total number of local devices.

5. The method for balanced training of federated learning data based on CVAE and multi-head attention mechanism according to claim 1, characterized in that: In step 3, the calculation formula for CVAE data reconstruction is as follows: Where z is the latent variable, x is the conditional variable, It is to reconstruct the data.

6. The method for balanced training of federated learning data based on CVAE and multi-head attention mechanism according to claim 1, characterized in that: In step 3, the design process of optimizing the objective function during data reconstruction is as follows: Construct reconstruction error BCE: Among them, x i is the ith element of the original input, is the i-th element of the reconstructed output, and N is the total number of samples; Construct KL divergence KLD: Among them, μ j and is the mean and variance of the jth latent variable; In summary, the objective function can be expressed as: l(θ,φ;x)=min(BCE+KLD) Here, θ is the parameter of the decoder and φ is the parameter of the encoder.

7. The method for balanced training of federated learning data based on CVAE and multi-head attention mechanism according to claim 1, characterized in that: The specific implementation of step 4 includes: Step 4.1 builds a CNN network with a multi-head attention mechanism: the multi-head attention mechanism is applied after the convolution layer and before the fully connected layer; CNN is used to obtain the primary features of the local device, and then the obtained primary features are used to generate enhanced feature representations using the multi-head attention mechanism; the CNN network model includes three outputs, namely the output features of the feature extraction network, the intermediate feature representation, and the final output; the multi-head attention mechanism performs multiple independent attention operations (called "heads") in parallel, then connects their outputs and integrates them through linear transformation to capture features of different subspaces; Step 4.2 Use the student model to learn the feature extraction of the teacher model: In federated learning, the teacher model has a better representation of feature learning, while the local device has a deviation in the representation of features due to data imbalance. By learning from the teacher model, the classification performance and generalization ability of the student model are improved; Step 4.3 Establish the optimization objective function: Combine the cross entropy loss output by the local student model and the teacher model with the knowledge distillation loss to obtain the optimization objective function The details are as follows: Design cross entropy loss function: Among them, y i is a one-hot vector of the true label, is the probability distribution predicted by the model; Design knowledge distillation loss function: Among them, h s,i is the feature output of the i-th sample of the student model, h t,i is the feature output of the i-th sample of the teacher model; The final local optimization objective function is:

8. The method for balanced training of federated learning data based on CVAE and multi-head attention mechanism according to claim 7, characterized in that: The multi-head attention mechanism in step 4.1 is specifically designed as follows: Assuming that the input has three matrices: query matrix Q, key matrix K and value matrix K, the operation process of the multi-head attention mechanism is as follows: Step 4.1.1 Linear transformation: Perform linear transformation on Q, K and V to obtain query, key and value matrices of different heads: Q i =QW i Q ,K i =KW i K ,V i =VW i V for i=1,2,...,h Among them, W i Q , W i K and W i V is the weight matrix of the i-th head, and h is the number of heads; Step 4.1.2 calculates the attention output of each head: Among them, d k is the dimension of the key vector; Step 4.1.3 Connection and linear transformation: Concatenate the outputs of all heads and integrate them through a linear transformation: MultiHead(Q,K,V)=Concat(head1,head2,...,head h )W 0 Among them, head i =Attention(Q i ,K i ,V i ), W 0 is the output weight matrix.

9. The method for balanced training of federated learning data based on CVAE and multi-head attention mechanism according to claim 1, characterized in that: In step 5, the aggregation method of the global model parameters is as follows: in, represents the model parameters of the participant device i in the t+1 round of communication, and η is the learning rate; The new global model parameters are obtained as follows: Among them, n represents the total amount of data of all participants, n k represents the amount of data of the kth participant, w t+1 is the global model parameter in the t+1th round.

10. A device for federated learning data balance training, characterized in that: It comprises a server device and several local devices; the server device can execute the contents of steps 1, 3, and 5 described in any one of claims 1-9; the local device can execute the contents of steps 2 and 4 described in any one of claims 1-9.

Images