DropAttack: A Masked Weight Adversarial Training Method to Improve Generalization of Neural Networks

Deep neural networks (DNNs) (LeCun et al., 2015) have achieved state-of-the-art performance in many artificial intelligence applications, such as natural language processing and computer vision. Regularization methods such as L1 Tibshirani (1996) and L2 (Tikhonov, 1943) regularization, early stopping (Morgan & Bourlard, 1989) and Dropout (Srivastava et al., 2014), play an important role in the impressive performance of deep networks by controlling the model complexity, and thus preventing overfitting and improving generalization. Adversarial training (Goodfellow et al., 2015) was originally proposed as a method to improve the security of machine learning systems in order to train a neural network that is robust to attack samples. And adversarial training is the process of training a model, which minimizes the maximal risk for label-preserving input perturbations. It improves not only robustness to adversarial examples, but also generalization performance for original examples. Goodfellow et al. (2015) have demonstrated that adversarial training can result in regularization; even further regularization than dropout.

In this work, we mainly focus on improving the generalization performance of the model and preventing the model from overfitting, rather than enhancing the robustness of model to attack samples. Miyato et al. (2017) applies adversarial training to text classification tasks, and finds that adversarial training can effectively improve the generalization of text or RNN models on the test set. Most of the recent adversarial training is aimed at the input of the model to attack. Moreover, these existing adversarial training methods add perturbation to every element in the input tensor during the attack. (It should be noted that the input in the NLP field is the embeddings of the text, and in the CV field it is the value of each pixel of the picture. In this paper, it is uniformly called the input for the convenience of expression.) In order to increase the breadth of the attack, we expanded the attack target from the input to the weight parameters of other layers, that is, in the process of adversarial training, while attacking the input, it also attacks the weight parameters of other layers. In each iteration of the attack, we randomly mask the attack on a certain proportion elements instead of attacking all the elements in the input or weight tensor. In this way, exponentially different attack combinations can be obtained, and the internal adversarial loss of the model can be maximized.

In this paper, we show the impact of DropAttack and various other well-known regularization methods on the generalization performance of the model. We experiment with different neural network models on five different public datasets to prove the effectiveness of the proposed method. We visually analyze the training and verification accuracy of some models of different architectures as the training progresses. And we also analyzed the impact of hyperparameters and optimization of multi-forward-backward propagation on DropAttack. Finally, we also conducted a theoretical analysis of the proposed method from another perspective to prove the effectiveness of DropAttack.

Adversarial training can be traced back to (Goodfellow et al., 2015), in which the model improves the robustness and generalization of the model by generating adversarial examples and injecting them into training data. The effectiveness of adversarial training largely depends on the direction of the attack, so it is necessary to find the perturbation value that maximizes the adversarial loss. Due to its linear characteristics, neural networks are easily attacked by linear perturbation. Therefore, Goodfellow et al. (2015) proposed the Fast Gradient Sign Method (FGSM) to calculate the perturbation of the input sample. They linearized the cost function around the current value of parameters, obtaining an optimal max-norm constrained pertubation of:

Where $\bm{\theta}$ is the model parameter, $\bm{x}$ is the input of the model, y is the label corresponding to the input, $L$ is the cost used to train the neural network. sgn is the symbolic function, and $\epsilon$ is the perturbation coefficient. In order to find a better perturbation, Miyato et al. (2017) proposed the Fast Gradient Method (FGM), which made a simple modification to the calculation of perturbation in FGSM. And FGM is the first time that adversarial training has been applied to text classification tasks. The formula is as follows:

Athalye et al. (2018) proposed a confrontation training method called Projected Gradient Descent (PGD), which obtains the final perturbation value through multiple forward and backward propagation iterations, and the perturbation obtained in each iteration is limited to a set range, if it exceeds this range will be mapped to the “sphere” of the range. To put it simply, “walk in small steps, take a few more steps.” The formula is as follows:

where $S=r\in\mathbb{R}^{d};\|\bm{r}\|_{2}\leq\epsilon$ is the constraint space of the perturbation, and a is the step length of the “small step”.

After that, although many defenses were broken by Athalye et al. (2018), PGD-based adversarial training is one of the few that can withstand powerful attacks. Athalye et al. (2018) proved that PGD can largely avoid the problem of gradient confusion, but it will still cause high convolution and non-linear loss surface. When K is very small, it will be lightly broken under powerful adversaries. To get more efficient PGD-based adversarial training, it must iteratively calculate the gradient many times, but this will consume a lot of computing resources. Shafahi et al. (2019) proposed a “free” adversarial training algorithm that eliminates the overhead cost of calculating adversarial perturbations by recycling the gradient information computed when updating model parameters. Zhang et al. (2019) effectively reduce the total number of full forward and backward propagations by restricting most of the forward and back propagation within the first layer of the network during adversary updates. Miyato et al. (2019) proposed virtual adversarial training as a regularization method for semi-supervised learning in the text field. Zhu et al. (2020) propose FreeLB for improving the generalization of language models, which performs multiple PGD iterations to attack the embeddings, and simultaneously accumulates the “free” parameter gradients in each iteration.

In this work, we are the first to propose an adversarial training method that simultaneously attacks the input of the model and the weight parameters of other layers to improve the generalization of the model. And our method, DropAttack, is the first to use random masking of some elements to increase the diversity of adversarial attack combinations.

The proposed new adversarial training method is inspired by Dropout, so we named it DropAttack. Standard adversarial training is to explore the optimal parameters to minimize the maximum risk of adversarial attacks. The Min-Max formula is as follows:

where $D$ is the data distribution, $y$ is the label, and L is loss function. $r_{adv}$ is the perturbation under maximizing internal risk. $S$ is the perturbation constraint space. Here we propose a new adversarial training method, DropAttack, which simultaneously attacks the input x of the model and the weight parameters of other layers, and randomly masks some attacks. The overall procedure is shown in Algorithm 1. The Min-Max formula of DropAttack can be expressed as:

where $\bm{r_{x}}$ and $\bm{r_{\theta}}$ are the perturbation of the input x and the parameter $\bm{\theta}$ under maximizing the internal risk. We respectively approximate these values by linearizing $\nabla_{\bm{x}}L(\bm{\theta},\bm{x_{t}},y)$ and $\nabla_{\bm{{\theta}}}L(\bm{\theta},\bm{x_{t}},y)$ around $\bm{x}$ and $\bm{\theta}$. Using the linear approximation in equation (6) and the L2 norm constraint, the resulting adversarial perturbation is

These perturbation can be easily calculated using backpropagation in a neural network.

$\bm{M_{x}}$ and $\bm{M_{\theta}}$ are the Random attack masks of $r_{x}$ and $r_{\theta}$ respectively. For any attack mask, $\bm{M_{ij}}$ is a matrix of independent Bernoulli random variables with the same dimension as the perturbation value, and the probability of each Bernoulli random variable is 1. Multiplying the perturbation matrix and the attack mask matrix will randomly mask a part of the element values in the perturbation matrix.

The attack on the input in the text model is essentially an attack on the weight parameters of the embedding layer. We can see that Figure 1 (a) is a standard neural network with 12 weight parameters. Figure 1 (b) is an new neural network produced by applying standard adversarial Training on the left, which attacks all inputs. And Figure 1 (c) is an new neural network after DropAttack is applied, and some of the weights are added with the perturbation calculated by the corresponding gradient. Compared with standard adversarial training, DropAttack maximizes internal risk through a wider range of attacks (not limited to the input layer). And randomly mask some dimensional perturbation, which will generate a more robust and diversified embedding space.

In this section, we test and analyze the effect of DropAttack on three NLP data sets and two CV data sets. In addition, we also analyze the ability of DropAttack to prevent overfitting under different sizes of training data. Additional experimental details and results are provided in the Appendix A.

We provide another theoretical perspectives to explain why the adversarial training method DropAttack can be used as regularization to improve the generalization of the model and prevent overfitting. According to Section 3, the task of DropAttack can be to minimize the maximum internal adversarial risk, that is, to approximately optimize the following goals:

For formula 8, we Taylor expand functions $f(x)=L(\bm{\theta},\bm{x}+\bm{M_{x}}\cdot\bm{r_{x}},y)$ and $f(\theta)=L(\bm{\theta}+\bm{M_{\theta}}\cdot\bm{r_{\theta}},\bm{x},y)$ at points $(\bm{x}+\bm{M_{x}}\cdot\bm{r_{x}})$ and $(\bm{\theta}+\bm{M_{\theta}}\cdot\bm{r_{\theta}})$ respectively:

Then, after substituting the values $\bm{r_{x}}=\epsilon_{x}\cdot\nabla_{\bm{x}}L(\bm{\theta},\bm{x},y)/\|\nabla_{\bm{x}}L(\bm{\theta},\bm{x},y)\|_{2},~{}\bm{r_{\theta}}=\epsilon_{\theta}\cdot\nabla_{\bm{{\theta}}}L(\bm{\theta},\bm{x},y)/\|\nabla_{\bm{{\theta}}}L(\bm{\theta},\bm{x},y)\|_{2}$ of the perturbation that maximize the antagonistic Loss, the following formula 10 is obtained:

We can see the final optimization function formula 13, which actually adds the implicit gradient regularization $\epsilon_{x}\cdot\|\bm{M_{x}}\cdot\bm{g_{x}}\|_{2}$ and $\epsilon_{\theta}\cdot\|\bm{M_{\theta}}\cdot\bm{g_{\theta}}\|_{2}$ to a certain proportion of input $\bm{x}$ and parameters $\bm{\theta}$ after loss every time the parameter $\bm{\theta}$ is updated. Gradient penalty pushes the gradient of some parameters and inputs to approach zero, so that the model is likely to be optimized to a flatter minimum.

In this work, we propose a masked weight adversarial training method, DropAttack, to improve the generalization ability of neural network models and prevent overfitting. The proposed algorithm uses the gradient method to attack the input and weight parameters according to a certain probability, and enhances the generalization of the model by minimizing the resultant adversarial risk. Experimental results prove that DropAttack can effectively improve the generalization of models and prevent overfitting, especially in the field of NLP. In addition, we theoretically proved that our algorithm can regularize the gradient of model parameters. Therefore, DropAttack can improve the robustness and generalization of the model. The current adversarial training still consumes more computing resources and time than the standard stochastic gradient descent, so it is a valuable research direction to accelerate the adversarial training while improving generalization in the future.