Adversarial Self-Attention for Language Understanding

Adversarial training (AT) (Goodfellow, Shlens, and Szegedy 2015) encourages model robustness by pushing the perturbed model prediction to $y$:

$$\mathop{\min}_{\theta}\mathcal{D}\left[y,p(y|\mathbf{x}+\delta^{*},\theta)\right]$$

(1)

where $\mathcal{D}[\cdot]$ refers to KL divergence and $p(y|\mathbf{x}+\delta^{*},\theta)$ refers to the perturbed model prediction under an adversarial perturbation $\delta^{*}$. In this paper, we apply the idea of virtual adversarial training (VAT) (Miyato et al. 2019), where $y$ is smoothed by the model prediction $p(y|\mathbf{x},\theta)$. However, the smoothness can hold only if the amount of training samples is sufficiently large so that $p(y|\mathbf{x},\theta)$ can be so close to $y$. This assumption is tenable for PrLMs. Even when fine-tuning on limited samples, the model can not predict so badly in the support of large-scale pre-trained weights.

The adversarial perturbation $\delta^{*}$ is defined to maximize the empirical risk of training:

$$\delta^{*}=\mathop{\arg\max}_{\delta;\parallel\delta\parallel\leq\epsilon}\mathcal{D}\left[p(y|\mathbf{x},\theta),p(y|\mathbf{x}+\delta,\theta)\right]$$

(2)

where $\parallel\delta\parallel\leq\epsilon$ refers to the decision boundary restricting the adversary $\delta$. We expect that $\delta^{*}$ is so minor but greatly fools the model. Eq. 1 and Eq. 2 make an adversarial “game”, in which the adversary seeks to find out the vulnerability while the model is trained to overcome the malicious attack.

Standard self-attention (SA, or vanilla SA) (Vaswani et al. 2017) can be formulated as:

$${\rm SA}(Q,K,V)={\rm Softmax}\left(\frac{Q\cdot K^{T}}{\sqrt{d}}\right)\cdot V$$

(3)

where $Q$, $K$, and $V$ refer to the query, key, and value components respectively, and $\sqrt{d}$ is a scaling factor. In this paper, we define the pair-wise matrix ${\rm Softmax}\left(\frac{Q\cdot K^{T}}{\sqrt{d}}\right)$ as the attention topology $\mathcal{T}(Q,K)$. In such a topology, each unit refers to the attention score between a token pair, and every single token is allowed to attend to all the others (including itself). The model learns such a topology so that it can focus on particular pieces of text.

However, empirical results show that manually biasing this process (i.e. determining how each token can attend to the others) can lead to better SA convergence and generalization, e.g. enforcing sparsity (Shi et al. 2021), strengthening local correlations (You, Sun, and Iyyer 2020), incorporating structural clues (Wu, Zhao, and Zhang 2021a). Basically, they smooth the output distribution around the attention structure with a certain priority $\mu$. We call $\mu$ the structure bias. Therefore, it leads us to a more general form of self-attention:

$${\rm SA}(Q,K,V,\mu)=\mathcal{T}(Q,K,\mu)\cdot V$$

(4)

where $\mathcal{T}(Q,K,\mu)$ is the biased attention topology ${\rm Softmax}\left(\frac{Q\cdot K^{T}}{\sqrt{d}}+\mu\right)$. In standard SA, $\mu$ equals an all-equivalent matrix (all elements are equivalent). That means the attentions on all token pairs are unbiased.

The general form in Eq. 4 indicates that we are able to manipulate the way the attentions unfold between tokens via overlapping a specific structure bias $\mu$. The corresponding sketches are in Figure 2. We focus on the masking case, which is commonly used to mask out the padding positions, where $\mu$ refers to a binary matrix with elements in $\{0,-\infty\}$. When an element equals $-\infty$, the attention score of that unit is off ($=0$). Note that the mask here is different from that in dropout (Srivastava et al. 2014) since the masked units will not be discarded but will be redistributed to other units.

The idea of ASA is to mask out those attention units to which the model is most vulnerable. Specifically, ASA can be regarded as an instance of general self-attention with an adversarial structure bias $\mu^{*}$. However, those vulnerable units vary from inputs. Thus, we let $\mu^{*}$ be a function of $\mathbf{x}$, denoted as $\mu_{\mathbf{x}}^{*}$. ASA can be eventually formulated as:

$${\rm ASA}(Q,K,V,\mu_{\mathbf{x}}^{*})=\mathcal{T}(Q,K,\mu_{\mathbf{x}}^{*})\cdot V$$

(5)

where $\mu_{\mathbf{x}}$ is parameterized by $\eta$ namely $\mu_{\mathbf{x}}^{*}=\mu(\mathbf{x},\eta^{*})$. We also call $\mu_{\mathbf{x}}$ the “adversary” in the following. Eq. 5 indicates that $\mu_{\mathbf{x}}^{*}$ acts as “meta-knowledge” learned from the input data itself, which sets ASA apart from other variants where the bias is predefined based on a certain priority.

Similar to adversarial training, the model is trained to minimize the following divergence:

$$\mathop{\min}_{\theta}\mathcal{D}\left[p(y|\mathbf{x},\theta),p(y|\mathbf{x},\mu(\mathbf{x},\eta^{*}),\theta)\right]$$

(6)

where $p(y|\mathbf{x},\mu(\mathbf{x},\eta^{*}),\theta)$ refers to the model prediction under the adversarial structure bias. We evaluate $\eta^{*}$ by maximizing the empirical risk:

$$\eta^{*}=\mathop{\arg\max}_{\eta;\|\mu(\mathbf{x},\eta)\|\leq\epsilon}\mathcal{D}\left[p(y|\mathbf{x},\theta),p(y|\mathbf{x},\mu(\mathbf{x},\eta),\theta)\right]$$

(7)

where $\|\mu(\mathbf{x},\eta)\|\leq\epsilon$ refers to the new decision boundary for $\eta$. The design of this constraint is necessary for keeping ASA from hurting model training.

Generally, researchers use $L_{2}$ or $L_{\infty}$ norm to make the constraint in adversarial training. Considering that $\mu(\mathbf{x},\eta)$ is in form of a binary mask, it is more reasonable to constrain it by limiting the proportion of the masked units, which comes to $L_{0}$ or $L_{1}$ norm (since $\mu(\mathbf{x},\eta)$ is binary, they are the same), namely $\|\mu(\mathbf{x},\eta)\|_{1}\leq\epsilon$. The question is that it is cumbersome to heuristically develop the value of $\epsilon$. As an alternative, we transform the problem with a hard constraint into an unconstrained one with a penalty:

	$\displaystyle\eta^{*}=$	$\displaystyle\mathop{\arg\max}_{\eta}\>\ \mathcal{D}\left[p(y|\mathbf{x},\theta),p(y|\mathbf{x},\mu(\mathbf{x},\eta),\theta)\right]$		(8)
		$\displaystyle+\tau\|\mu(\mathbf{x},\eta)\|_{1}$

where we use a temperature coefficient $\tau$ to control the intensity of the adversary. Eq. 8 indicates that the adversary needs to maximize the training risk and at the same time mask the least number of units as possible. Our experiments show that it is much easier to adjust $\tau$ than to adjust $\epsilon$ as in adversarial training. We find good performances when $\tau=0.1\sim 0.3$.

Eventually, we generalize Eq. 8 to a model with $n$ self-attention layers (e.g. BERT). Let $\mu(\mathbf{x},\eta)=\{\mu(\mathbf{x},\eta)^{1},\cdots,\mu(\mathbf{x},\eta)^{n}\}$, where $\mu(\mathbf{x},\eta)^{i}$ refers to the adversary for the $i^{\rm th}$ layer. The penalty term thus becomes the summation $\|\mu(\mathbf{x},\eta)\|_{1}=\sum_{i=1}^{n}\|\mu(\mathbf{x},\eta)^{i}\|_{1}$.

Adversarial training is naturally expensive. To remedy this, we propose a fast and simple implementation of ASA. There are two major points below.

We eventually present the training objective for training an ASA-empowered model. Given a task with labels, we let $\mathcal{L}_{e}(\theta)$ be the task-specific loss (e.g. classification, regression). The model needs to make the right predictions against the ASA adversary so that we obtain the ASA loss $\mathcal{L}_{asa}(\theta,\eta)=\mathcal{D}\left[p(y|\mathbf{x},\theta),p(y|\mathbf{x},\mu(\mathbf{x},\eta),\theta)\right]$. At the same time, the adversary is subject to the penalty term $\mathcal{L}_{c}(\eta)=\|\mu(\mathbf{x},\eta)\|_{1}$. The final training objective thus consists of three components:

$$\mathcal{L}_{e}(\theta)+\alpha\mathcal{L}_{asa}(\theta,\eta)+\tau\mathcal{L}_{c}(\eta)$$

(9)

where we find ASA performs just well when simply fixing the balancing coefficient $\alpha$ to 1 so that $\tau$ is the only hyperparameter of learning ASA.

Eq. 9 explains how ASA works. One point is that the penalty term is paralleled to the model since it is not associated with $\theta$. We seek to find the optimal model parameters $\theta$ to minimize the first two terms. On the other hand, because of GRL, we seek to find the optimal adversary parameters $\eta$ to maximize the last two terms.

Though Eq. 9 covers all cases when fine-tuning the ASA-empowered model on downstream tasks. We will discuss more on pre-training. ASA is consistent with the current trend of self-supervised language modeling like MLM (Devlin et al. 2019) and RTD (Clark et al. 2020), where we construct the negative samples based on the super large corpus itself without additional labeling.

To be more concrete, we rely on MLM pre-training setting (Devlin et al. 2019) in what follows and the other situations can be easily generalized. MLM intends to recover a number of masked pieces within the sequence and the loss of it is obtained from the cross-entropy on those selected positions, denoted as $\mathcal{L}_{mlm}$. Thus, we can compute the divergence between the two model predictions before and after biased by ASA on those positions and obtain the token-level ASA loss $\mathcal{L}^{t}_{asa}(\theta,\eta)$.

Aside from the masked positions, the beginning position is also crucial to PrLMs (e.g [CLS] in BERT), which is always used as an indicator for the relationship within the sequence (e.g. sentence order, sentiment). Thus, we pick this position out from the final hidden states and calculate the divergence on it as another part of the ASA loss $\mathcal{L}^{s}_{asa}(\theta,\eta)$. Finally, pre-training with ASA can be formulated as:

$$\mathcal{L}_{mlm}(\theta)+\mathcal{L}^{t}_{asa}(\theta,\eta)+\mathcal{L}^{s}_{asa}(\theta,\eta)+\tau\mathcal{L}_{c}(\eta)$$

(10)

where $\mathcal{L}^{t}_{asa}$ and $\mathcal{L}^{s}_{asa}$ refer to the token-level and sentence-level ASA loss respectively.

Based on Eq. 10, we may touch on the idea of ASA pre-training from two perspectives: (a) Structural loss: ASA acts as a regularizer on the empirical loss of language modeling (the same in Eq. 9); (b) Multiple objectives: ASA can be viewed as two independent self-supervised pre-training objectives in addition to MLM.

We experiment on five NLP tasks down to 10 datasets:

$\bullet$ Sentiment Analysis: Stanford Sentiment Treebank (SST-2) (Socher et al. 2013), which is a single-sentence binary classification task; $\bullet$ Natural Language Inference (NLI): Multi-Genre Natural Language Inference (MNLI) (Williams, Nangia, and Bowman 2018) and Question Natural Language Inference (QNLI) (Wang et al. 2019), where we need to predict the relations between two sentences; $\bullet$ Semantic Similarity: Semantic Textual Similarity Benchmark (STS-B) (Cer et al. 2017) and Quora Question Pairs (QQP) (Wang et al. 2019), where we need to predict how similar two sentences are; $\bullet$ Named Entity Recognition (NER): WNUT-2017 (Aguilar et al. 2017), which contains a large number of rare entities; $\bullet$ Machine Reading Comprehension (MRC): Dialogue-based Reading Comprehension (DREAM) (Sun et al. 2019), where we need to choose the best answer from the three candidates given a question and a piece of dialogue; $\bullet$ Robustness learning: Adversarial NLI (ANLI) (Nie et al. 2020) for NLI, PAWS-QQP (Zhang, Baldridge, and He 2019) for semantic similarity, and HellaSWAG (Zellers et al. 2019) for MRC.

We verify the gain of ASA on the top of two different self-attention (SA) designs: vanilla SA in BERT-base (Devlin et al. 2019) and its stronger variant RoBERTa-base (Liu et al. 2019), and disentangled SA in DeBERTa-large (He et al. 2021). In addition, we do the experiments on both pre-training ($\tau=0.1$, Eq. 10) and fine-tuning ($\tau=0.3$, Eq. 9) stages (the training details can be found in Appendix). For pre-training, we continue to pre-train BERT based on MLM with ASA on the English Wikipedia corpus. Besides, for fair enough comparison, we train another BERT with vanilla SA (BERT$\dagger$ in Table 1). We set the batch size to 128 and train both models for 20K steps with FP16. Note that we directly fine-tune them without ASA.

We compare ASA with three naive masking strategies. Bernoulli distribution is a widely-used priority in network dropout (Srivastava et al. 2014). We report the best results with the masking probability selected in {0.05, 0.1}. Besides, we introduce another two potentially stronger strategies. For the first one, we dynamically schedule the masking probability for each step following the learned pattern by ASA. Different from ASA, the masked units here are Bernoulli chosen. For the second one, we always choose to mask those units with the most significant attention scores (a magnitude-based strategy). Similar to ASA, we apply the masking matrices to all self-attention layers, and the training objective corresponds to the first two terms of Eq. 9.

From Table 3, we can see that pure Bernoulli works the best among the three naive strategies, slightly better than scheduled Bernoulli. However, ASA outweighs them even by a large margin, suggesting that the worst-case masking can better facilitate model training. Besides, the magnitude-based masking turns out to be harmful. Since ASA acts as a gradient-based adversarial strategy, it may not always mask those globally most significant units in the attention matrix (this pattern can be seen in previous Figure 1).

On the other hand, it has been found in previous work that adversarial training on word embeddings can appear mediocre compared to random perturbations (Aghajanyan et al. 2021). We are positive that adversarial training benefits model training, but hypothesize that the current spatial optimization of embedding perturbations suffers from shortcomings so that it sometimes falls behind random perturbations. However, the optimization of ASA is carefully designed in our paper.

ASA is close to conventional adversarial training, but there are two main differences. In the text domain, adversarial training works on the input space, imposing perturbations on word embeddings, while ASA works on model structures. Besides, adversarial training normally leverages projected gradient descent (PGD) (Madry et al. 2018) to learn the adversary, while ASA is optimized through an unconstrained manner. We compare the performances on different tasks between ASA and FreeLB (Zhu et al. 2020), one of the state-of-the-art adversarial training approaches in the text domain.

From Table 4, we can see that ASA and FreeLB are competitive on MNLI and QNLI, while ASA outperforms by 1.4 and 1.2 points on PAWS-QQP and HellaSWAG. It may leave a new line for future research, where ASA can be superior to conventional adversarial training on certain tasks. Another advantage of ASA is that it only introduces one hyperparameter $\tau$, while for FreeLB, we need to sweep through different adversarial step sizes and boundaries.

On the other hand, we find that both FreeLB and SMART can hardly induce a significant variation in the attention maps, even if the input embeddings are perturbed. The model remains focused on those pieces that are supposed to be focused on before being perturbed. This can be detrimental when one tries to explain the adversary’s behavior.

Figure 4 (a) summarizes the speed performances of ASA and other two state-of-the-art adversarial training algorithms in the text domain. FreeLB (Zhu et al. 2020) is currently the fastest algorithm, which requires at least two inner steps (FreeLB-2) to complete its learning process. SMART (Jiang et al. 2020) leverages the idea of virtual adversarial training, which thus requires at least one more forward passes. We turn off FP16, fix the batch size to 16, and do the experiments under different sequence lengths. We can see that ASA is slightly faster than FreeLB-2, taking about twice the period of naive training when the sequence length is up to 512. However, training with SMART and FreeLB-3 is much more expensive, taking about three and four times that period of naive training (SMART-1 is very close to FreeLB-3, so we omit it from the figure).

The only hyperparameter for ASA is the temperature coefficient $\tau$, which controls the intensity of the adversary, a lower $\tau$ corresponding to a stronger attack (higher masking proportion). In practice, $\tau$ balances the model generalization and robustness. We conduct experiments on a benign task (DREAM) and an adversarial task (HellaSWAG) respectively with $\tau$ selected in {0.3, 0.6, 1.0}. As in Figure 4 (b), the trends of two curves are opposite (we offset them vertically to make them close). A stronger adversary might lead to a decrease in generalization but benefit robustness. For example, we see the peak DREAM result when $\tau=1.0$, while the peak HellaSWAG result when $\tau=0.3$.

A higher masking proportion in ASA implies that the layer is more vulnerable. As in Figure 5 (a), we observe that the masking proportion almost decreases layer by layer from the bottom to the top. We attribute this to information diffusion (Goyal et al. 2020), which states that the input vectors are progressively assimilating through continuously making self-attention. Consequently, the attention scores in the bottom layers are more important so that more vulnerable, while those in the top layers become less important after their assimilation. The feed-forward layers are more contributing this time (You, Sun, and Iyyer 2020).

In Figure 5 (b), we calculate the average masking proportion of all layers and see that the situations can also be different between tasks even with the same temperature. Take sentiment analysis as an instance, the adversary merely needs to focus on specific keywords, which is enough to lead to misclassification. For NER and MRC, however, there are more sensitive words that are scattered across the sentence, and therefore a greater attack of the adversary is needed.