# Fast and Accurate Deep Leakage from Gradients Based on Wasserstein Distance

Authors: Xing He; Changgen Peng; Weijie Tan
Journal: International Journal of Intelligent Systems
 (2023)

---

## Abstract

Shared gradients are widely used to protect the private information of training data in distributed machine learning systems. However, Deep Leakage from Gradients (DLG) research has found that private training data can be recovered from shared gradients. The DLG method still has some issues such as the “Exploding Gradient,” low attack success rate, and low fidelity of recovered data. In this study, a Wasserstein DLG method, named WDLG, is proposed; the theoretical analysis shows that under the premise that the output layer of the model has a “bias” term, predicting the “label” of the data by whether the “bias” is “negative” or not is independent of the approximation of the shared gradient, and thus, the label of the data can be recovered with 100% accuracy. In the proposed method, the Wasserstein distance is used to calculate the error loss between the shared gradient and the virtual gradient, which improves model training stability, solves the “Exploding Gradient” phenomenon, and improves the fidelity of the recovered data. Moreover, a large learning rate strategy is designed to improve model training convergence speed in-depth. Finally, the WDLG method is validated on datasets from MNIST, Fashion MNIST, SVHN, CIFAR-100, and LFW. Experiments results show that the proposed WDLG method provides more stable updates for virtual data, a higher attack success rate, faster model convergence, higher image fidelity during recovery, and support for designing large learning rate strategies.


---

## Body

## 1. Introduction

Shared gradients are common in computationally and time-intensive task systems. Machine learning (ML) model training is frequently used in distributed training and collaborative learning due to its similar task characteristics. The gradients are shared and aggregated across multiple microprocessors (nodes) for machine learning model training; each worker node operates in parallel to accelerate ML model training. Therefore, distributed training methods are critical for speeding up model training on large-scale datasets.In recent years, research using shared gradient methods has yielded excellent results. McMahan et al. [1] were the first to propose the Federated Learning (FL) architecture. FL builds iterable aggregated models by training distributed models across multiple data sources with local data, only exchanging model parameters or intermediate results models, and thus learning a shared target model. Improved approaches based on FL methods [2–9] have carried out a lot of research work in achieving a balance between data privacy protection and data sharing computation. Currently, more researchers are using cryptographic privacy-preserving methods and differential privacy-preserving methods to achieve privacy-preserving security for local gradients in the federal learning security problem [10, 11]. Meanwhile, for large-scale ML tasks, other distributed training [12–14] and collaborative learning [15, 16] methods are widely used.In distributed ML systems, it is widely assumed that “shared gradients” do not leak the training dataset. The output of machine learning model decisions frequently contains a large amount of inference information from training samples; recent studies have shown that these outputs are vulnerable to membership inference attacks [17–21] and gradient reversal attacks [16, 22–25], which result in leaking training datasets. To the best of our knowledge, the most dominant Deep Leakage from Gradients (DLG) [24] and Improved DLG (iDLG) [25] are currently among the most effective gradient-based inverse attacks. Nevertheless, the DLG method still generates random noisy images for training data recovery, and “Exploding Gradients” is one of the main reasons for the low success rate of DLG method attacks. In addition, the DLG method has problems such as slow convergence speed of model training, low fidelity of recovered data (the original training data can be recovered; however, the recovered training data still has random noise in some positions of the image), and an undesirable design strategy of a large learning rate.For starters, data label prediction is more accurate and faster. We theoretically prove that the prediction of data labels is independent of the recovery of training data to predict data labels more accurately and improve attack accuracy, assuming that the activation function is ReLU and using cross-entropy to calculate the loss value. Our method’s time complexity is much lower than that of iDLG’s method, especially when the last hidden layer contains more neurons. Second, model training is more stable, and convergence occurs more quickly. The WDLG method can stabilize model training and play the role of fast convergence, and the main reason why the analysis can solve the phenomenon of “Exploding Gradient” is that the original gradient calculated by random initialization of model parameters has certain randomness, and the gradient calculated during training still has certain randomness. The superiority of calculating the distance between two random variables according to the Wasserstein distance can better stabilize the model training to better calculate the distance between two random variables, which is the first attempt to improve the original DLG method as far as we know. Third, large learning rate model training is more stable; we attempt to provide a large learning rate design strategy to improve model training speed and achieve better results; however, the balance between model training stability and learning rate design must be fully considered.Our proposed WDLG method has been validated on MNIST [26], Fashion MNIST [27], SVHN [28], CIFAR-100 [29], and LFW [30] datasets. The experiments demonstrate the benefits of the WDLG method for more stable virtual data updates for model training, higher attack success rates, faster model convergence, and support for large learning rate design strategies. This paper’s main contributions are as follows:(i)
We assume that the output layer employs a biased ReLU activation function and that the model employs cross-entropy loss. Based on a differentiable model, the WDLG method differentiates the bias term and theoretically demonstrates that the prediction of label data is only related to the bias of the output layer, independent of the recovery of the training data. As a result, the true labels are extracted with 100% accuracy from the shared gradients, allowing for more accurate data extraction.(ii)
Wasserstein distance is used to calculate the gradient loss, and a new loss function is selected to solve the stability of the model training. The model training converges faster and solves the “Exploding Gradient” phenomenon.(iii)
Through experiments, we analyze the relationship between the stability of model training and the fidelity of recovered data with different learning rates and conclude that the WDLG method can support large learning rate design strategies.The rest of this paper is organized as follows: Section2 describes the related work in detail. Section 3 presents our method, which includes the extremely important prediction proofs for training data labels, the choice of the loss function, and the approximation of the Wasserstein distance calculation. Section 4 gives three valid comparison experiments, which are the comparison of prediction accuracy of training data labels, the comparison of training model stability and convergence speed, and the comparison of large learning rate training. Finally, Section 6 concludes the paper.


## 2. Related Works

Many studies have been conducted to improve the scalability of distributed training of ML models during distributed training of ML models. Sharing data gradients is a common practice in existing distributed ML systems, and multiple parties’ data are jointly trained using the shared gradient information. Because of stochastic gradient descent’s stable scalability, the algorithm [31–34] and framework [35–37] have essentially adopted synchronized stochastic gradient descent as the dominant method for parameter updating. However, both centralized [38–40] and decentralized [34, 41] training methods have the potential technical security risk of data and model parameters being compromised during training and application. One of the most common attacks for model leakage is the model reversal attack, in which the attacker extracts information about the training data from the model prediction results. Model reversal attacks are classified into two types based on the different model output information possessed by attackers: model output-based data leakage and gradient update-based data leakage.


### 2.1. Model Output-Based Data Leakage

The output of machine learning model decisions frequently contains many inferences about the training data, and recent studies have shown that these outputs can be used to steal information about the original training data. Fredrikson et al. [17] proposed a confidence-based model reversal attack method and achieved better attack results on the model [42–44]. It demonstrated that confidence could be used as a measure to recover training data and learn sensitive private information about individuals for adversarial access to ML models. The membership inference attack was first proposed by Shokri et al. [18]; the basic principle of the attack is to statistically analyze the discrepancy between the prediction results of ML models on training and nontraining data; the attacker can implement the membership inference attack if he can capture this discrepancy. On this basis, the equivalent attack effect of Shokri’s proposed method is achieved by fitting the membership inference attack [45] and membership inference attack in the black box scenario [46]. The literature [6, 47] investigated the membership inference attack in white-box scenarios and successfully built more robust attacks by exploiting the difference in model gradients between training and nontraining data, resulting in better attack results. Song et al. [48] discovered that adding adversarial training during the training of ML models increases the risk of the model leaking members’ private information through an in-depth study of membership inference attacks. By using generative adversarial networks [49] to learn the statistics of real data distributions, Hayes et al. [19] completed the construction of membership inference attacks. Furthermore, Salem et al. [20] improved the membership inference attack by extending it to the domain of online learning.


### 2.2. Gradient Update-Based Data Leakage

The gradients continuously generated by machine learning models during training still imply rich privacy information about the training dataset. Melis et al. [16] used user-updated model parameters as features for training attack model inputs for inferring relevant attributes of other user datasets. The literature [7, 50, 51] employs generative adversarial networks to generate methods for recovering training data from other users, and Mahendran et al. [22] investigate gradient inversion information maximization to synthesize real data from training networks, but both rely on a priori information from auxiliary datasets. Mordvintsev et al. [23] use only the gradients in the input to enable the separation of noise and image, making it difficult to obtain higher-fidelity information on large datasets. Geiping et al. [8] were the first to push decision boundaries to ImageNet-level gradient inversion, but there is an issue with an unpredictable gradient averaging emergence. To the best of our knowledge, the most popular Deep Leakage from Gradients (DLG) [24] and Improved Deep Leakage from Gradients (iDLG) [25] methods are currently among the most effective methods for attacking gradient-based data leakage methods. The method generates virtual data with the same size as the training samples and labels at random and then uses the virtual data as the model input and error direction propagation to calculate the virtual gradient. The virtual data are learned using an optimization algorithm in such a way that the gradient obtained by backpropagation on the common model is similar to the real gradient, and the training data and labels are obtained after several rounds of iterative optimization. At the moment, this is one of the hottest topics in the study of variants of DLG-based methods [9, 52–54].Although the DLG method has produced good attack results, its attack capability is limited, and the attack succeeds only when the predicted label is correct, and the model converges slowly; if the predicted label fails, the attack never converges, resulting in the DLG method’s low accuracy. The improved iDLG, on the other hand, theoretically proves that data label prediction is independent of training data recovery under the condition that the activation function is a ReLU activation function or a sigmoid function and the loss value is calculated by cross-entropy, allowing the iDLG method to always attack successfully and thus recover the training data with high fidelity. Although the iDLG method produces good experimental results, the “Exploding Gradient” phenomenon occurs when training the model, resulting in slow model convergence and an undesirable large learning rate design strategy.We analyze the causes of the more unstable training process of the DLG method and the improved iDLG method; the least squares method to calculate loss values typically suffers from the uncertainty of the gradient calculated by randomly initializing the weights. There may be multiple identical possible outcomes for the same pseudodata input during iterative training of the model, and in these cases, the least-squares loss function aims for fuzzy prediction. When viewed through the lens of stochasticity, the original gradients computed by randomly initializing the model parameters suffer from a certain degree of randomness, as do the gradients of the parameters generated during the training process. This results in inefficient DLG method attacks, a lack of stability in model training, the “Exploding Gradient” in training data recovery, and undesirable design strategies for high learning rates.


## 2.1. Model Output-Based Data Leakage

The output of machine learning model decisions frequently contains many inferences about the training data, and recent studies have shown that these outputs can be used to steal information about the original training data. Fredrikson et al. [17] proposed a confidence-based model reversal attack method and achieved better attack results on the model [42–44]. It demonstrated that confidence could be used as a measure to recover training data and learn sensitive private information about individuals for adversarial access to ML models. The membership inference attack was first proposed by Shokri et al. [18]; the basic principle of the attack is to statistically analyze the discrepancy between the prediction results of ML models on training and nontraining data; the attacker can implement the membership inference attack if he can capture this discrepancy. On this basis, the equivalent attack effect of Shokri’s proposed method is achieved by fitting the membership inference attack [45] and membership inference attack in the black box scenario [46]. The literature [6, 47] investigated the membership inference attack in white-box scenarios and successfully built more robust attacks by exploiting the difference in model gradients between training and nontraining data, resulting in better attack results. Song et al. [48] discovered that adding adversarial training during the training of ML models increases the risk of the model leaking members’ private information through an in-depth study of membership inference attacks. By using generative adversarial networks [49] to learn the statistics of real data distributions, Hayes et al. [19] completed the construction of membership inference attacks. Furthermore, Salem et al. [20] improved the membership inference attack by extending it to the domain of online learning.


## 2.2. Gradient Update-Based Data Leakage

The gradients continuously generated by machine learning models during training still imply rich privacy information about the training dataset. Melis et al. [16] used user-updated model parameters as features for training attack model inputs for inferring relevant attributes of other user datasets. The literature [7, 50, 51] employs generative adversarial networks to generate methods for recovering training data from other users, and Mahendran et al. [22] investigate gradient inversion information maximization to synthesize real data from training networks, but both rely on a priori information from auxiliary datasets. Mordvintsev et al. [23] use only the gradients in the input to enable the separation of noise and image, making it difficult to obtain higher-fidelity information on large datasets. Geiping et al. [8] were the first to push decision boundaries to ImageNet-level gradient inversion, but there is an issue with an unpredictable gradient averaging emergence. To the best of our knowledge, the most popular Deep Leakage from Gradients (DLG) [24] and Improved Deep Leakage from Gradients (iDLG) [25] methods are currently among the most effective methods for attacking gradient-based data leakage methods. The method generates virtual data with the same size as the training samples and labels at random and then uses the virtual data as the model input and error direction propagation to calculate the virtual gradient. The virtual data are learned using an optimization algorithm in such a way that the gradient obtained by backpropagation on the common model is similar to the real gradient, and the training data and labels are obtained after several rounds of iterative optimization. At the moment, this is one of the hottest topics in the study of variants of DLG-based methods [9, 52–54].Although the DLG method has produced good attack results, its attack capability is limited, and the attack succeeds only when the predicted label is correct, and the model converges slowly; if the predicted label fails, the attack never converges, resulting in the DLG method’s low accuracy. The improved iDLG, on the other hand, theoretically proves that data label prediction is independent of training data recovery under the condition that the activation function is a ReLU activation function or a sigmoid function and the loss value is calculated by cross-entropy, allowing the iDLG method to always attack successfully and thus recover the training data with high fidelity. Although the iDLG method produces good experimental results, the “Exploding Gradient” phenomenon occurs when training the model, resulting in slow model convergence and an undesirable large learning rate design strategy.We analyze the causes of the more unstable training process of the DLG method and the improved iDLG method; the least squares method to calculate loss values typically suffers from the uncertainty of the gradient calculated by randomly initializing the weights. There may be multiple identical possible outcomes for the same pseudodata input during iterative training of the model, and in these cases, the least-squares loss function aims for fuzzy prediction. When viewed through the lens of stochasticity, the original gradients computed by randomly initializing the model parameters suffer from a certain degree of randomness, as do the gradients of the parameters generated during the training process. This results in inefficient DLG method attacks, a lack of stability in model training, the “Exploding Gradient” in training data recovery, and undesirable design strategies for high learning rates.


## 3. WDLG Method

We give the structure of the WDLG method, as shown in Figure1. Among the key problems, we need to solve is the prediction of the data labels and the construction based on the Wasserstein distance loss value calculation, which leads to the recovery of the training data. Therefore, we first describe the proof process of data label prediction, followed by the selection of the Wasserstein distance-based loss function, and finally the approximate solution of the Wasserstein distance.Figure 1
Overview of our WDLG algorithm. Normal participant inputs training samplex to obtain the original gradient ∇W, and malicious attacker randomly initializes the virtual data x′ to obtain the virtual gradient ∇W′ and minimizes the Wasserstein distance between the original gradient ∇W and the virtual gradient ∇W′. When the iterative optimization is completed, the virtual data x′ randomly initialized by a malicious attacker converges to x with a preset threshold, thereby obtaining the training sample x from the normal participant.


### 3.1. Data Label Prediction

There is some randomness in the DLG method for the label recovery. In the iDLG method [25], it is found that shared gradients can leak the true labels, and for the labels encoded by unique heat, it is proposed that any differentiable model trained based on cross-entropy loss can extract the true labels effectively. Based on the iDLG approach, we give a more concise method for true label recovery. For the classification scenario, where the neural network model is usually trained using the cross-entropy loss of the unique heat-encoded true labels, it is defined as follows:(1)Lossx,c=−logeyc∑kCeyk=−yc+log∑kCeyk,where x denotes the training sample, C denotes the total number of classification categories, and yc denotes the confidence that the training sample x is predicted to be the c-th class.(2)yc=f∑j=1tyj+bc,where t denotes the number of neurons in the penultimate hidden layer and ∑j=1tyj denotes the input of the c-th neuron in the last layer, whose corresponding bias is bc. For the ReLU activation function, fx=max0,x and f′x=1. The loss gradient for bias bc is(3)gi=∂Lossx,c∂bi=−∂yc∂bi+∂log∑kCeyk∂bi=−∂f∑j=1tyj+bc∂bi+∂log∑kCefΣj=1tyj+bk∂bi=−∂f∑j=1tyj+bc∂bi+ef∑j=1tyj+bc∑kCef∑j=1tyj+bk=−1+ef∑j=1tyj+bc∑kCef∑j=1tyj+bk,ifi=c,ef∑j=1tyj+bc∑kCef∑j=1tyj+bk,otherwise.Sinceef∑j=1tyj+bc/∑kCef∑j=1tyj+bk∈0,1,gi∈−1,0 when i=c, otherwise gi∈0,1.Therefore, we can get the same conclusion of true label prediction as in the iDLG method, but we are more efficient in terms of computational complexity. We only need to obtain the negative bias term of the last layer of neurons to obtain the corresponding true labels of the input training samplex.


### 3.2. WDLG Loss Function

The DLG method makes the virtual gradient approximate the original gradient by optimization, and in theory, the virtual data will also be closer to the training data. Given the gradient of the original training data, the training data are obtained by minimizing the following objective:(4)x′∗,y′∗=argminx′,y′∇W′−∇W2=argminx′,y′∂LossFx′,W,y′∂W−∇W2,where Fx′,W is a neural network parameterized by W with x′ as input. The ∇W′−∇W2 distance is differentiable concerning the training data x′ and the virtual label y′ and thus can be optimized using a standard gradient-based approach.Based on the superior performance of the Wasserstein distance, we propose the WDLG method and recompute the distance between∂LossFx′,W,y′/∂W−∇W2. Our WDLG method obtains the training data by minimizing the following objectives:(5)x′∗=argminx′Wd∇W′,∇W=argminx′Wd∂LossFx′,W,y′∂W,∇W,where Wd denotes the Wasserstein distance.We refer to the DLG [24] method and proposed the WDLG method in Figure 1, where the color “red” is not the same. The specific process is given as follows:Step 1: normal participant inputs the training samplex into the differentiable deep network model F to obtain the predicted value Pred and at the same time calculates the label corresponding to the input sample according to equation (3), noting that the prediction of the label is independent of the training process of the whole network, and our method can recover 100% of the correct label. We record the gradient ∇W (the original gradient) of the real sample corresponding to each parameter.Step 2: malicious attacker randomly initializes the virtual datax′, which is sampled from the standard positive-terrestrial distribution. x′ is used as the input of the deep network model F to obtain the virtual gradient ∇W′ corresponding to the virtual data.Step 3: we calculate the Wasserstein distance between∇W′ and ∇W to calculate the loss value according to equation (5) and iteratively update the input virtual data x′ cyclically until the fidelity index MSE between x′ and x reaches the preset threshold.In Algorithm1, Fx,W denotes the differentiable deep learning model, and the model is parametrically represented by W. bc′c′=1,2,⋯,C denotes the bias of the c′ neuron in the output layer, and under the assumption of cross-entropy output, the subscript corresponding to the negative bias term, i.e., is the class of the sample prediction.Algorithm 1: WDLG Algorithm.
Input: ∇W: real gradient generated by real training sample x. Epochs: maximum number of iterations. η: learning rate.(1)
c′←minb1,b2,⋯,bC≤0//Get the subscript c′ of the last layer of negative bias(2)
cpred← One_Hot c′//Convert c′ to One_Hot code(3)
x′←N0,1//Initialize virtual data with the same dimensions as x(4)
fori←1 to Ndo(5)
∇W′←∂LossFx′,W,cpred/∂W//Calculating virtual gradients(6)
LossWd←WDCA∇W′,∇W//WDCA is Algorithm 2(7)
x′←x′−η∇x′LossWd(8)
end forOutput: x′


### 3.3. Wasserstein Distance Calculation

The Wasserstein distance measure is widely used in econometrics, machine learning, and other fields. The proposal of the WGAN [55] model, in particular, has reignited a surge in its research and applications.


#### 3.3.1. Wasserstein Distance Definition

LetΠP,Q be the set of all possible joint probability distributions of the combination of two probability distributions Px and Qy, and for any joint probability distribution γx,y, the distance dx,y of the sample x,y∼γ distribution is defined as follows:(6)WP,Q=infγ∼ΠP,Q∬γx,ydx,y,=infγ∼ΠP,QEx,y∼γdx,y,where dx,y is a cost function. For each possible joint distribution γ, samples x and y can be obtained by sampling x,y∼γ from it and calculating the cost dx,y between the pair, so the expectation Ex,y∼γdx,y of the sample pair cost under this joint distribution γ can be calculated. The lower bound that can be taken on this expectation in all possible joint distributions is the Wasserstein distance. However, it is intuitive to interpret Ex,y∼γdx,y as the consumption required to move the P distribution to the Q distribution under this path planning of γ. Therefore, the Wasserstein distance is the minimum consumption under optimal path planning. Therefore, the optimal transmission Sinkhorn algorithm [56] is used to approximate the Wasserstein distance. Regularization encourages transmission using mostly small traffic paths while penalizing sparse paths. Therefore, introducing entropy regularization can limit the complexity of the solution of the optimal transmission problem, which leads to an approximate solution of the optimal transmission problem. To obtain the Sinkhorn algorithm, the solution of the regularized Kantorovich problem can be written in the following form:(7)Pi,j=uiKi,jvj,∀i,j∈n×m,where Ci,j denotes the cost of transferring a unit mass from ai to bj and i,j denotes the matrix subscript in the cost matrix C. P=diaguKdiagv, u,v∈R+n×R+m, and Ki,j=e−Ci,j/ε. ε is the regularization factor whose magnitude determines the strength of the regularization effect. The vectors u and v are the variables to be required by Sinkhorn’s algorithm to join the mass conservation condition for optimal transmission.(8)a=u⊙Kvandb=v⊙KTv,where ⊙ is the Hadamard product of the vector. Therefore, we can solve for u and v iteratively. u is updated at each step, then v is updated, and finally, the iterations converge and both sides of the equation are satisfied simultaneously, and we get the following iterative equation:(9)ut+1=aKvt,vt+1=bKTut+1.In algorithm2, we give an approximate solution algorithm WDCA for Wd.Algorithm 2: Wasserstein distance calculation algorithm (WDCA).
Input: input original gradient original_grad, virtual gradient dummy_grad.(1)
Initializationu0,v0(2)
C←cost_matrixoriginal_grad,dummy_grad//Calculate the cost matrix C(3)
K←e−C/ε(4)
a←u0⊙Kv0(5)
b←v0⊙KTu0(6)
fori←0 to max_iterdo(7)
ut+1←a/Kvt+ut(8)
vt+1←b/KTut+1+vt(9)
Ifsumabsut+1−ut≤err−thresh, then(10)
u∗←ut+1(11)
v∗←vt+1(12)
break(13)
end if(14)
end for(15)
Wd←DC,u∗,v∗//D obtains an approximation of the Wasserstein distance WdOutput: Wd


## 3.1. Data Label Prediction

There is some randomness in the DLG method for the label recovery. In the iDLG method [25], it is found that shared gradients can leak the true labels, and for the labels encoded by unique heat, it is proposed that any differentiable model trained based on cross-entropy loss can extract the true labels effectively. Based on the iDLG approach, we give a more concise method for true label recovery. For the classification scenario, where the neural network model is usually trained using the cross-entropy loss of the unique heat-encoded true labels, it is defined as follows:(1)Lossx,c=−logeyc∑kCeyk=−yc+log∑kCeyk,where x denotes the training sample, C denotes the total number of classification categories, and yc denotes the confidence that the training sample x is predicted to be the c-th class.(2)yc=f∑j=1tyj+bc,where t denotes the number of neurons in the penultimate hidden layer and ∑j=1tyj denotes the input of the c-th neuron in the last layer, whose corresponding bias is bc. For the ReLU activation function, fx=max0,x and f′x=1. The loss gradient for bias bc is(3)gi=∂Lossx,c∂bi=−∂yc∂bi+∂log∑kCeyk∂bi=−∂f∑j=1tyj+bc∂bi+∂log∑kCefΣj=1tyj+bk∂bi=−∂f∑j=1tyj+bc∂bi+ef∑j=1tyj+bc∑kCef∑j=1tyj+bk=−1+ef∑j=1tyj+bc∑kCef∑j=1tyj+bk,ifi=c,ef∑j=1tyj+bc∑kCef∑j=1tyj+bk,otherwise.Sinceef∑j=1tyj+bc/∑kCef∑j=1tyj+bk∈0,1,gi∈−1,0 when i=c, otherwise gi∈0,1.Therefore, we can get the same conclusion of true label prediction as in the iDLG method, but we are more efficient in terms of computational complexity. We only need to obtain the negative bias term of the last layer of neurons to obtain the corresponding true labels of the input training samplex.


## 3.2. WDLG Loss Function

The DLG method makes the virtual gradient approximate the original gradient by optimization, and in theory, the virtual data will also be closer to the training data. Given the gradient of the original training data, the training data are obtained by minimizing the following objective:(4)x′∗,y′∗=argminx′,y′∇W′−∇W2=argminx′,y′∂LossFx′,W,y′∂W−∇W2,where Fx′,W is a neural network parameterized by W with x′ as input. The ∇W′−∇W2 distance is differentiable concerning the training data x′ and the virtual label y′ and thus can be optimized using a standard gradient-based approach.Based on the superior performance of the Wasserstein distance, we propose the WDLG method and recompute the distance between∂LossFx′,W,y′/∂W−∇W2. Our WDLG method obtains the training data by minimizing the following objectives:(5)x′∗=argminx′Wd∇W′,∇W=argminx′Wd∂LossFx′,W,y′∂W,∇W,where Wd denotes the Wasserstein distance.We refer to the DLG [24] method and proposed the WDLG method in Figure 1, where the color “red” is not the same. The specific process is given as follows:Step 1: normal participant inputs the training samplex into the differentiable deep network model F to obtain the predicted value Pred and at the same time calculates the label corresponding to the input sample according to equation (3), noting that the prediction of the label is independent of the training process of the whole network, and our method can recover 100% of the correct label. We record the gradient ∇W (the original gradient) of the real sample corresponding to each parameter.Step 2: malicious attacker randomly initializes the virtual datax′, which is sampled from the standard positive-terrestrial distribution. x′ is used as the input of the deep network model F to obtain the virtual gradient ∇W′ corresponding to the virtual data.Step 3: we calculate the Wasserstein distance between∇W′ and ∇W to calculate the loss value according to equation (5) and iteratively update the input virtual data x′ cyclically until the fidelity index MSE between x′ and x reaches the preset threshold.In Algorithm1, Fx,W denotes the differentiable deep learning model, and the model is parametrically represented by W. bc′c′=1,2,⋯,C denotes the bias of the c′ neuron in the output layer, and under the assumption of cross-entropy output, the subscript corresponding to the negative bias term, i.e., is the class of the sample prediction.Algorithm 1: WDLG Algorithm.
Input: ∇W: real gradient generated by real training sample x. Epochs: maximum number of iterations. η: learning rate.(1)
c′←minb1,b2,⋯,bC≤0//Get the subscript c′ of the last layer of negative bias(2)
cpred← One_Hot c′//Convert c′ to One_Hot code(3)
x′←N0,1//Initialize virtual data with the same dimensions as x(4)
fori←1 to Ndo(5)
∇W′←∂LossFx′,W,cpred/∂W//Calculating virtual gradients(6)
LossWd←WDCA∇W′,∇W//WDCA is Algorithm 2(7)
x′←x′−η∇x′LossWd(8)
end forOutput: x′


## 3.3. Wasserstein Distance Calculation

The Wasserstein distance measure is widely used in econometrics, machine learning, and other fields. The proposal of the WGAN [55] model, in particular, has reignited a surge in its research and applications.


### 3.3.1. Wasserstein Distance Definition

LetΠP,Q be the set of all possible joint probability distributions of the combination of two probability distributions Px and Qy, and for any joint probability distribution γx,y, the distance dx,y of the sample x,y∼γ distribution is defined as follows:(6)WP,Q=infγ∼ΠP,Q∬γx,ydx,y,=infγ∼ΠP,QEx,y∼γdx,y,where dx,y is a cost function. For each possible joint distribution γ, samples x and y can be obtained by sampling x,y∼γ from it and calculating the cost dx,y between the pair, so the expectation Ex,y∼γdx,y of the sample pair cost under this joint distribution γ can be calculated. The lower bound that can be taken on this expectation in all possible joint distributions is the Wasserstein distance. However, it is intuitive to interpret Ex,y∼γdx,y as the consumption required to move the P distribution to the Q distribution under this path planning of γ. Therefore, the Wasserstein distance is the minimum consumption under optimal path planning. Therefore, the optimal transmission Sinkhorn algorithm [56] is used to approximate the Wasserstein distance. Regularization encourages transmission using mostly small traffic paths while penalizing sparse paths. Therefore, introducing entropy regularization can limit the complexity of the solution of the optimal transmission problem, which leads to an approximate solution of the optimal transmission problem. To obtain the Sinkhorn algorithm, the solution of the regularized Kantorovich problem can be written in the following form:(7)Pi,j=uiKi,jvj,∀i,j∈n×m,where Ci,j denotes the cost of transferring a unit mass from ai to bj and i,j denotes the matrix subscript in the cost matrix C. P=diaguKdiagv, u,v∈R+n×R+m, and Ki,j=e−Ci,j/ε. ε is the regularization factor whose magnitude determines the strength of the regularization effect. The vectors u and v are the variables to be required by Sinkhorn’s algorithm to join the mass conservation condition for optimal transmission.(8)a=u⊙Kvandb=v⊙KTv,where ⊙ is the Hadamard product of the vector. Therefore, we can solve for u and v iteratively. u is updated at each step, then v is updated, and finally, the iterations converge and both sides of the equation are satisfied simultaneously, and we get the following iterative equation:(9)ut+1=aKvt,vt+1=bKTut+1.In algorithm2, we give an approximate solution algorithm WDCA for Wd.Algorithm 2: Wasserstein distance calculation algorithm (WDCA).
Input: input original gradient original_grad, virtual gradient dummy_grad.(1)
Initializationu0,v0(2)
C←cost_matrixoriginal_grad,dummy_grad//Calculate the cost matrix C(3)
K←e−C/ε(4)
a←u0⊙Kv0(5)
b←v0⊙KTu0(6)
fori←0 to max_iterdo(7)
ut+1←a/Kvt+ut(8)
vt+1←b/KTut+1+vt(9)
Ifsumabsut+1−ut≤err−thresh, then(10)
u∗←ut+1(11)
v∗←vt+1(12)
break(13)
end if(14)
end for(15)
Wd←DC,u∗,v∗//D obtains an approximation of the Wasserstein distance WdOutput: Wd


## 3.3.1. Wasserstein Distance Definition

LetΠP,Q be the set of all possible joint probability distributions of the combination of two probability distributions Px and Qy, and for any joint probability distribution γx,y, the distance dx,y of the sample x,y∼γ distribution is defined as follows:(6)WP,Q=infγ∼ΠP,Q∬γx,ydx,y,=infγ∼ΠP,QEx,y∼γdx,y,where dx,y is a cost function. For each possible joint distribution γ, samples x and y can be obtained by sampling x,y∼γ from it and calculating the cost dx,y between the pair, so the expectation Ex,y∼γdx,y of the sample pair cost under this joint distribution γ can be calculated. The lower bound that can be taken on this expectation in all possible joint distributions is the Wasserstein distance. However, it is intuitive to interpret Ex,y∼γdx,y as the consumption required to move the P distribution to the Q distribution under this path planning of γ. Therefore, the Wasserstein distance is the minimum consumption under optimal path planning. Therefore, the optimal transmission Sinkhorn algorithm [56] is used to approximate the Wasserstein distance. Regularization encourages transmission using mostly small traffic paths while penalizing sparse paths. Therefore, introducing entropy regularization can limit the complexity of the solution of the optimal transmission problem, which leads to an approximate solution of the optimal transmission problem. To obtain the Sinkhorn algorithm, the solution of the regularized Kantorovich problem can be written in the following form:(7)Pi,j=uiKi,jvj,∀i,j∈n×m,where Ci,j denotes the cost of transferring a unit mass from ai to bj and i,j denotes the matrix subscript in the cost matrix C. P=diaguKdiagv, u,v∈R+n×R+m, and Ki,j=e−Ci,j/ε. ε is the regularization factor whose magnitude determines the strength of the regularization effect. The vectors u and v are the variables to be required by Sinkhorn’s algorithm to join the mass conservation condition for optimal transmission.(8)a=u⊙Kvandb=v⊙KTv,where ⊙ is the Hadamard product of the vector. Therefore, we can solve for u and v iteratively. u is updated at each step, then v is updated, and finally, the iterations converge and both sides of the equation are satisfied simultaneously, and we get the following iterative equation:(9)ut+1=aKvt,vt+1=bKTut+1.In algorithm2, we give an approximate solution algorithm WDCA for Wd.Algorithm 2: Wasserstein distance calculation algorithm (WDCA).
Input: input original gradient original_grad, virtual gradient dummy_grad.(1)
Initializationu0,v0(2)
C←cost_matrixoriginal_grad,dummy_grad//Calculate the cost matrix C(3)
K←e−C/ε(4)
a←u0⊙Kv0(5)
b←v0⊙KTu0(6)
fori←0 to max_iterdo(7)
ut+1←a/Kvt+ut(8)
vt+1←b/KTut+1+vt(9)
Ifsumabsut+1−ut≤err−thresh, then(10)
u∗←ut+1(11)
v∗←vt+1(12)
break(13)
end if(14)
end for(15)
Wd←DC,u∗,v∗//D obtains an approximation of the Wasserstein distance WdOutput: Wd


## 4. Experiments

In purpose to give comparative experiments with fairness, we choose PyTorch [57] as the experimental platform and give the same network model and hyperparameter settings uniformly except for different algorithms. In Algorithm 1, batch_size=1, regularization factor ε=0.01, maximum number of iterations max_iter=100, and error threshold errthresh=0.1 are given. We perform comparative validation and analysis of DLG, iDLG, and WDLG methods on MNIST, Fashion MNIST, CIFAR-100, SVHN, and LFW public datasets for accuracy of data label prediction, stability and convergence speed of virtual data training to recover data fidelity (MSE indicator), and large learning rate design strategies.


### 4.1. Comparison of the Accuracy of Predicted Labels

The DLG method that uses the least squares method to calculate the gradient loss values for data label recovery is initialized with a randomization method for iterative learning, and there is certain randomness in its ability to recover data labels, which makes the attack efficiency (accuracy of label prediction) not high. In contrast, the iDLG method theoretically demonstrates that under the condition of using the ReLU activation function as well as cross-entropy to calculate the loss value, the label prediction is independent of the network iterative training and the training labels can be fully recovered.In our WDLG method, it is assumed that the last hidden layer of the deep network model hasN neurons and the training dataset has C categories. According to the basic premise assumption of the iDLG method, it is further assumed that the i-th neuron of the output layer has bias bi. We can still prove theoretically that the recovery of data labels is independent of the iterative training process of the network and refer to the proof of (equation (3)) for details. This allows the WDLG method to always attack successfully and thus recover the training data with high fidelity. In Table 1, we give a comparison of the accuracy of DLG, iDLG, and WDLG in predicting the true labels. Both WDLG and iDLG methods have a 100% success rate in predicting the data labels, while DLG often makes errors in predicting the true labels. However, for the computational time complexity of prediction of data labels, the computation step iDLG_Num of the iDLG method is much higher than that of our WDLG method, especially when the number N of neurons in the last hidden layer is very large; the advantage of our method is more significant.Table 1
Comparison of DLG, iDLG, and WDLG prediction labels and prediction complexity.
DatasetsDLGiDLGiDLG_NumWDLG (%)WDLG_NumMNIST89.9%100.0%N×C100.0CFashion MNIST××N×C100.0CSVHN××N×C100.0CCIFAR-10083.3%100.0%N×C100.0CLFW79.1%100.0%N×C100.0CThe accuracy of DLG [24], iDLG [25], and WDLG in predicting data labels. Note that iDLG and WDLG always predict the correct labels, while DLG often makes errors in predicting the true labels.In general,C≪N. The number of computational steps for the iDLG method to recover the labels is N×C, while our method requires only C computational steps. Therefore, it can be theoretically demonstrated that the computational time complexity of the WDLG method on the task of recovering labels is only 1/N that of the iDLG method, which is significant from the computational analysis, as shown in Table 1 for detailed comparison. Our proposed data label recovery method assumes a stronger conditional generalization capability.


### 4.2. Comparison of the Stability and Convergence Speed of Model Training

The training process of the DLG method is less stable, and the loss value is often very large in the initial iteration step, which makes the whole iterative computation process difficult to converge. The improved iDLG method uses the cross-computed loss function and also uses the least squares method to calculate the gradient loss value, which still suffers from the difficulty of convergence of the loss value. Analyzing the main reason, the least squares method to calculate the loss values usually suffers from the uncertainty of the gradient calculated by random initialization weights. When iteratively training the model, there may be various equally possible results for the same random data input. In all these cases, the least squares loss function aims to accommodate the uncertainty in the prediction by fuzzy prediction, or simply, by taking the average of the possible outputs. This is because the average of all possible outcomes will lead to overall minimization of the parameter space of the inputs during training.For the comparison of model training stability and model convergence speed experiments, we experimentally validate the DLG, iDLG, and WDLG methods on simple datasets with a single feature, MNIST, Fashion MNIST, and datasets with complex features, SVHN, CIFAR-100, and LFW. In Figure2, the visualization of depth leakage on MNIST and Fashion MNIST (from top to bottom) images is shown separately; the first image is the original training image, and from left to right are the recovered views of training data on the same number of iterations for DLG, iDLG, and WDLG methods, respectively. On the images with clear training images and single features, there is not much difference in the number of iterations between DLG and iDLG methods to achieve the same image fidelity quality, while our method WDLG only requires fewer iterations, and the model converges faster. With the same number of iterations, our method recovers more realistic training data with higher fidelity.Figure 2
Comparison of the stability of model training on a dataset with a single feature.To compare the comparison on feature-complex datasets, in Figure3, the visualization of depth leakage on SVHN, CIFAR-100, and LFW (top-to-bottom) images is shown separately, with the first image being the original training image, and from left to right, the recovered data views of the DLG, iDLG, and WDLG methods on the same number of iterations, respectively, are given. On the training image with insufficient sharpness and feature-rich images, the number of iterations required by the DLG, iDLG, and WDLG methods increases substantially to achieve the same image fidelity quality. However, our method requires fewer iterations, and the model training still maintains a significant advantage in terms of convergence speed.Figure 3
Comparison of stability of model training on datasets with complex features.In particular, we note that the loss values and fidelity during training are carried out as follows to compare the experimental results better; the output loss values and fidelity are taken logarithmically and then shifted up, and such operation does not change the nature of the function. The purpose is to make a better comparison in detail on the same graph. Finally, schematic diagrams of the relationship between loss value and the number of iterations, fidelity and number of iterations, and fidelity and loss value on complex feature datasets are given, respectively. In Figure4, the relationship between loss value and fidelity versus the number of iterations for DLG, iDLG, and WDLG methods on three datasets, SVHN, CIFAR-100, and LFW, is given. In the (a), (b), and (c) plots, our WDLG method has the smallest loss value and better convergence than the other two methods on the same iteration steps. In the (d), (e), and (f) plots, our WDLG method has the smallest MSE and higher fidelity of recovered data on the same iteration steps. In the (g), (h), and (i) plots, to achieve the same image fidelity and loss values, the DLG method requires the most iteration steps, the iDLG method is the second, and our WDLG method requires the least iteration steps. The proposed WDLG method consistently outperforms the DLG and iDLG methods in recovering data and has significant advantages in all three tasks. In the plot of (f), the advantage of WDLG is significant on the difficult task of LFW, with MSE close to 0 and higher fidelity of recovered data.Figure 4
The relationship between loss values and fidelity versus the number of iterations for DLG, iDLG, and WDLG methods on three datasets, SVHN, CIFAR-100, and LFW.
(a)(b)(c)(d)(e)(f)(g)(h)(i)Through the abovementioned analysis, to solve the problems existing in the least squares method to calculate the loss value, the proposed Wasserstein distance to calculate the loss value plays a relatively good and stable training effect in the training process, and the main reason for its success, analyzed from the perspective of randomness, the original gradient∇W calculated by randomly initializing the model parameters also has a certain degree of randomness, and the parameter gradient generated in the training process ∇W′ also has some randomness. The aim is to better calculate the distance between two random variables. This is because the KL-divergence does not satisfy the symmetry of the distance definition, and the JS-divergence always keeps the distance as log2 when the two distributions do not overlap at all, leading to difficulties in providing an effective gradient calculation during the model training process. The key to solving the problem lies in using the Wasserstein distance to measure the distance between two distributions, and the WGAN [55] theoretically demonstrates that the Wasserstein distance can reflect the distance between two distributions and provide an effective gradient calculation even if they do not have any overlap. Therefore, our proposed WDLG method can better assign singular values to other random components when singular values appear in the iterative gradient calculation ∇W′ and thus will be more stable for model training.


### 4.3. Comparison of Large Learning Rate Training

To analyze our WDLG method supporting large learning rate strategy design in detail, we test and validate the loss value vs. the number of iterations and fidelity vs. the number of iteration curves at different learning rates only on the complex dataset LFW.In (a) of Figure5, the loss value decreases with increasing the learning rate, achieving a surprising effect. However, we should pay more attention to the existence of an inflection point in the loss function curve for lr=0.01 and lr=0.1 at iterations<20, which causes an abrupt change in the slope of the curve and is one of the reasons for the difficulty in the convergence of the model training, corresponding to the difficulty in the convergence of the recovered data fidelity (MSE) at a finite number of iterations in (b). However, for the cases of lr=0.5 and lr=1, especially at lr=1, the rate of decline of the loss function curve is quite stable throughout the training process, and there is no case of singularity, and the corresponding MSE curve decreases steadily, and the fidelity continues to increase. Therefore, we have reasons to believe that our WDLG method can support a large learning strategy design.Figure 5
Variation of loss value and fidelity curves by changing the learning rate lr strategy on the LFW dataset.
(a)(b)In the course of our experiments, we also tried to provide larger learning rates to the DLG and iDLG methods, but they failed frequently to recover the fidelity of the data. To analyze the reason, the optimization path direction is dominated by the gradient in the direction of the singular value when the gradient∇W′ of the iterative computation appears singular, resulting in the optimization path deviating from the feasible solution path. Therefore, the learning rate of DLG and iDLG methods is set too large to cause the network to fail to converge or to hover around the feasible solution, thus neglecting the location of the feasible solution.In our WDLG method, the iterative generation of the singular gradient∇W′ assigned to other random components of the distance calculation is still the main reason why we can set a larger learning rate successfully. We still try to provide a larger learning rate during the training of the network, and Figure 5 shows our experimental results. In (a), it is clear that the loss value decreases faster with increasing the learning rate on the same iteration step, while MSE, which portrays the fidelity metric, also shows a significant decrease in (b), so we have reasons to believe that our WDLG method for learning rate generalization ability is stronger. Although the larger the learning rate, the faster the model training reaches the optimal value, the learning rate increases to a certain threshold, and the whole network no longer converges. Therefore, in our approach, the design of a large learning rate strategy is encouraged and the stability of model training should be considered more.


## 4.1. Comparison of the Accuracy of Predicted Labels

The DLG method that uses the least squares method to calculate the gradient loss values for data label recovery is initialized with a randomization method for iterative learning, and there is certain randomness in its ability to recover data labels, which makes the attack efficiency (accuracy of label prediction) not high. In contrast, the iDLG method theoretically demonstrates that under the condition of using the ReLU activation function as well as cross-entropy to calculate the loss value, the label prediction is independent of the network iterative training and the training labels can be fully recovered.In our WDLG method, it is assumed that the last hidden layer of the deep network model hasN neurons and the training dataset has C categories. According to the basic premise assumption of the iDLG method, it is further assumed that the i-th neuron of the output layer has bias bi. We can still prove theoretically that the recovery of data labels is independent of the iterative training process of the network and refer to the proof of (equation (3)) for details. This allows the WDLG method to always attack successfully and thus recover the training data with high fidelity. In Table 1, we give a comparison of the accuracy of DLG, iDLG, and WDLG in predicting the true labels. Both WDLG and iDLG methods have a 100% success rate in predicting the data labels, while DLG often makes errors in predicting the true labels. However, for the computational time complexity of prediction of data labels, the computation step iDLG_Num of the iDLG method is much higher than that of our WDLG method, especially when the number N of neurons in the last hidden layer is very large; the advantage of our method is more significant.Table 1
Comparison of DLG, iDLG, and WDLG prediction labels and prediction complexity.
DatasetsDLGiDLGiDLG_NumWDLG (%)WDLG_NumMNIST89.9%100.0%N×C100.0CFashion MNIST××N×C100.0CSVHN××N×C100.0CCIFAR-10083.3%100.0%N×C100.0CLFW79.1%100.0%N×C100.0CThe accuracy of DLG [24], iDLG [25], and WDLG in predicting data labels. Note that iDLG and WDLG always predict the correct labels, while DLG often makes errors in predicting the true labels.In general,C≪N. The number of computational steps for the iDLG method to recover the labels is N×C, while our method requires only C computational steps. Therefore, it can be theoretically demonstrated that the computational time complexity of the WDLG method on the task of recovering labels is only 1/N that of the iDLG method, which is significant from the computational analysis, as shown in Table 1 for detailed comparison. Our proposed data label recovery method assumes a stronger conditional generalization capability.


## 4.2. Comparison of the Stability and Convergence Speed of Model Training

The training process of the DLG method is less stable, and the loss value is often very large in the initial iteration step, which makes the whole iterative computation process difficult to converge. The improved iDLG method uses the cross-computed loss function and also uses the least squares method to calculate the gradient loss value, which still suffers from the difficulty of convergence of the loss value. Analyzing the main reason, the least squares method to calculate the loss values usually suffers from the uncertainty of the gradient calculated by random initialization weights. When iteratively training the model, there may be various equally possible results for the same random data input. In all these cases, the least squares loss function aims to accommodate the uncertainty in the prediction by fuzzy prediction, or simply, by taking the average of the possible outputs. This is because the average of all possible outcomes will lead to overall minimization of the parameter space of the inputs during training.For the comparison of model training stability and model convergence speed experiments, we experimentally validate the DLG, iDLG, and WDLG methods on simple datasets with a single feature, MNIST, Fashion MNIST, and datasets with complex features, SVHN, CIFAR-100, and LFW. In Figure2, the visualization of depth leakage on MNIST and Fashion MNIST (from top to bottom) images is shown separately; the first image is the original training image, and from left to right are the recovered views of training data on the same number of iterations for DLG, iDLG, and WDLG methods, respectively. On the images with clear training images and single features, there is not much difference in the number of iterations between DLG and iDLG methods to achieve the same image fidelity quality, while our method WDLG only requires fewer iterations, and the model converges faster. With the same number of iterations, our method recovers more realistic training data with higher fidelity.Figure 2
Comparison of the stability of model training on a dataset with a single feature.To compare the comparison on feature-complex datasets, in Figure3, the visualization of depth leakage on SVHN, CIFAR-100, and LFW (top-to-bottom) images is shown separately, with the first image being the original training image, and from left to right, the recovered data views of the DLG, iDLG, and WDLG methods on the same number of iterations, respectively, are given. On the training image with insufficient sharpness and feature-rich images, the number of iterations required by the DLG, iDLG, and WDLG methods increases substantially to achieve the same image fidelity quality. However, our method requires fewer iterations, and the model training still maintains a significant advantage in terms of convergence speed.Figure 3
Comparison of stability of model training on datasets with complex features.In particular, we note that the loss values and fidelity during training are carried out as follows to compare the experimental results better; the output loss values and fidelity are taken logarithmically and then shifted up, and such operation does not change the nature of the function. The purpose is to make a better comparison in detail on the same graph. Finally, schematic diagrams of the relationship between loss value and the number of iterations, fidelity and number of iterations, and fidelity and loss value on complex feature datasets are given, respectively. In Figure4, the relationship between loss value and fidelity versus the number of iterations for DLG, iDLG, and WDLG methods on three datasets, SVHN, CIFAR-100, and LFW, is given. In the (a), (b), and (c) plots, our WDLG method has the smallest loss value and better convergence than the other two methods on the same iteration steps. In the (d), (e), and (f) plots, our WDLG method has the smallest MSE and higher fidelity of recovered data on the same iteration steps. In the (g), (h), and (i) plots, to achieve the same image fidelity and loss values, the DLG method requires the most iteration steps, the iDLG method is the second, and our WDLG method requires the least iteration steps. The proposed WDLG method consistently outperforms the DLG and iDLG methods in recovering data and has significant advantages in all three tasks. In the plot of (f), the advantage of WDLG is significant on the difficult task of LFW, with MSE close to 0 and higher fidelity of recovered data.Figure 4
The relationship between loss values and fidelity versus the number of iterations for DLG, iDLG, and WDLG methods on three datasets, SVHN, CIFAR-100, and LFW.
(a)(b)(c)(d)(e)(f)(g)(h)(i)Through the abovementioned analysis, to solve the problems existing in the least squares method to calculate the loss value, the proposed Wasserstein distance to calculate the loss value plays a relatively good and stable training effect in the training process, and the main reason for its success, analyzed from the perspective of randomness, the original gradient∇W calculated by randomly initializing the model parameters also has a certain degree of randomness, and the parameter gradient generated in the training process ∇W′ also has some randomness. The aim is to better calculate the distance between two random variables. This is because the KL-divergence does not satisfy the symmetry of the distance definition, and the JS-divergence always keeps the distance as log2 when the two distributions do not overlap at all, leading to difficulties in providing an effective gradient calculation during the model training process. The key to solving the problem lies in using the Wasserstein distance to measure the distance between two distributions, and the WGAN [55] theoretically demonstrates that the Wasserstein distance can reflect the distance between two distributions and provide an effective gradient calculation even if they do not have any overlap. Therefore, our proposed WDLG method can better assign singular values to other random components when singular values appear in the iterative gradient calculation ∇W′ and thus will be more stable for model training.


## 4.3. Comparison of Large Learning Rate Training

To analyze our WDLG method supporting large learning rate strategy design in detail, we test and validate the loss value vs. the number of iterations and fidelity vs. the number of iteration curves at different learning rates only on the complex dataset LFW.In (a) of Figure5, the loss value decreases with increasing the learning rate, achieving a surprising effect. However, we should pay more attention to the existence of an inflection point in the loss function curve for lr=0.01 and lr=0.1 at iterations<20, which causes an abrupt change in the slope of the curve and is one of the reasons for the difficulty in the convergence of the model training, corresponding to the difficulty in the convergence of the recovered data fidelity (MSE) at a finite number of iterations in (b). However, for the cases of lr=0.5 and lr=1, especially at lr=1, the rate of decline of the loss function curve is quite stable throughout the training process, and there is no case of singularity, and the corresponding MSE curve decreases steadily, and the fidelity continues to increase. Therefore, we have reasons to believe that our WDLG method can support a large learning strategy design.Figure 5
Variation of loss value and fidelity curves by changing the learning rate lr strategy on the LFW dataset.
(a)(b)In the course of our experiments, we also tried to provide larger learning rates to the DLG and iDLG methods, but they failed frequently to recover the fidelity of the data. To analyze the reason, the optimization path direction is dominated by the gradient in the direction of the singular value when the gradient∇W′ of the iterative computation appears singular, resulting in the optimization path deviating from the feasible solution path. Therefore, the learning rate of DLG and iDLG methods is set too large to cause the network to fail to converge or to hover around the feasible solution, thus neglecting the location of the feasible solution.In our WDLG method, the iterative generation of the singular gradient∇W′ assigned to other random components of the distance calculation is still the main reason why we can set a larger learning rate successfully. We still try to provide a larger learning rate during the training of the network, and Figure 5 shows our experimental results. In (a), it is clear that the loss value decreases faster with increasing the learning rate on the same iteration step, while MSE, which portrays the fidelity metric, also shows a significant decrease in (b), so we have reasons to believe that our WDLG method for learning rate generalization ability is stronger. Although the larger the learning rate, the faster the model training reaches the optimal value, the learning rate increases to a certain threshold, and the whole network no longer converges. Therefore, in our approach, the design of a large learning rate strategy is encouraged and the stability of model training should be considered more.


## 5. Conclusion

In this paper, a WDLG method is proposed, and the theoretical analysis shows that the prediction of data labels is independent of the learning process of the virtual data used, and the labels of the data can be recovered accurately. In the proposed method, the Wasserstein distance is used to calculate the error loss between the shared gradient and the virtual gradient, which solves the stable model training and “Exploding Gradient” phenomena and improves the fidelity of the recovered data. Meanwhile, it is experimentally demonstrated that our method fully supports large learning rate strategies and therefore improves the convergence speed of model training.Based on this research, we will consider extending the WDLG method to deeper network models and more complex training datasets in the future. Meanwhile, a more secure machine learning algorithm based on differential privacy is proposed based on the attack capability of the WDLG method. The purpose is to allow researchers to reacquaint themselves with the hidden potential security issues in distributed ML, to provide thoughts on security research for the next generation of AI model training, and to provide a research idea to solve the conflict between AI applications and security.



---

*Source: 5510329-2023-03-21.xml*