AttExplainer: Explain Transformer via Attention by Reinforcement Learning

1 Introduction

With the development of deep learning methods in natural language processing (NLP), model explanation has been increasingly investigated to interpret the black-box model in an intelligible manner. As a recent representative neural network architecture, Transformer [Vaswani and others, 2017] is leading the current NLP studies. For instance, the Transformerbased BERT [Devlin and others, 2019] is believed to start a groundbreaking era and has been proven effective in many downstream tasks. Basically, it comprises multiple stacked layers of self-attention and nonlinear feed-forward modules. As the key to the explanation of the attention mechanism, the attention matrix (also known as attention weight) is particularly investigated in previous research. It has been used to reveal the latent semantic relationships between input tokens [Hewitt and Manning, 2019], or find some regular patterns in attention matrix images [Kovaleva et al., 2019].

∗ Corresponding authors

Typically, the attention matrix is regarded as an inherent explanation for Transformer. Since [Jain and Wallace, 2019] threw a question that the relationship between attention weights and model outputs is unclear, researchers have begun to investigate how to explain attention matrices and harvest the potential information in them. Though widely discussed in previous works, none of them address the open problem as presented in Figure 1. In particular, there are two major issues: (1) Most existing perturbation-based interpretation methods use predict probabilities as the basis for interpretation. It is still an open question whether some other inherent information of the model, such as attention matrix, can be employed as a basis to explain the model's predictions; (2) Not just at the level of analysis, it is still under-explored how to extract potential information from attention matrix for benefiting the downstream tasks, such as model explanation and model adversarial attack.

On the other hand, Perturbed Masking is an unsupervised method to measure the dependency degrees between input tokens [Wu et al., 2020]. By successively masking token pairs in all pairwise combinations and sorting the degree of changes before and after the masking operations, the probing method can draw an unlabeled dependency tree from the input sentence. In light of this, we argue that the attention matrix could be explained by observing the changes introduced by the masking operation. Specially, we try to extend the possibility from pair to an unlimited number of token masking combinations. However, this will lead to an explosion of combination spaces. Some previous methods used greedy search, such as beam search or genetic algorithms to select crucial positions [Ren et al., 2019; Li et al., 2019]. But the time cost of searching is still high. Nonetheless, there is still a lack of this kind of method and an effective model is needed.

To address the aforementioned issues, we introduce a reinforcement learning (RL)-based method for perturbation sampling and build a game environment for two downstream tasks: model explanation and text adversarial attack. Perturbation sampling is required in both tasks. Compared to other combination search algorithms, the RL is more efficient to model the Markov Decision Process when searching in huge action space. Meanwhile, the attention matrix can be considered as a kind of input observation for the RL agent for further efficiencies. Specifically, in each game step, a deep Q-network (DQN) based agent observes the attention matrix, then selects a position to change the token's visibility to the Transformer model by masking or unmasking the token. The main contributions are summarized as follows:

• We propose an RL-based model interpretation method named ATTEXPLAINER to explain the Transformer via attention matrix quantitatively by finding the key/important tokens in a sentence. Especially, attention matrices are employed to provide heuristic information to speed up the process.
• We prove that the attention matrix does contain information about token importance. Also, we demonstrate the strong transferability of our method within different well-trained Transformer models and datasets.
• We conduct extensive experiments on three public datasets, demonstrating the interpretability of the attention matrix on two popular downstream tasks, namely, model explanation and model adversarial attack. In addition, we perform case studies to visualize the step-by-step masking process for better understanding.

2 RelatedWork

2.1 Model Explanation

Typically, model explanation aims to accurately deliver the true reasons behind the model's decision/prediction by identifying which parts of the input are important. The two dominant factions are gradient-based and perturbation-based. In this article, we focus on perturbation-based explanation methods. As introduced in [Ribeiro et al., 2016], the goal of LIME is to use a smaller and simpler model to explain the original big model within a local scope of an input sample. [Lundberg and Lee, 2017] proposed the SHAP framework, which assigns each feature an important value for one particular prediction. The Shapley value of a feature is the average marginal contribution of the feature in all features. Different methods are proposed for the model explanation, but how to uniformly define model interpretation is still an open question [Feng and others, 2018].

2.2 Attention Explanation

Attention matrices, which are considered important inherent information from Transformer models, have been extensively studied to interpret model predictions in NLP. For example, some research works tried to study whether the attention matrix is interpretable by designing perturbation experiments on model parameters [Jain and Wallace, 2019; Serrano and Smith, 2019; Wiegreffe and Pinter, 2019]. [Kovaleva et al., 2019] summarized several BERT's self-attention patterns through manual observation. The model proposed in [Htut et al., 2019] compared the consistency of attention heads and human-labeled syntactic dependency trees. In [Abnar and Zuidema, 2020], the authors proposed a method to track the flow of information through self-attention layers to better explain the attention matrix. However, there is still a lot of controversy about whether attention can indeed deliver important information about the input tokens. [Chefer et al., 2021] proposes a gradient-based method for interpreting the attention matrix, which considers both the gradient of the attention matrix and the relevancy score between tokens reflected by the attention matrix.

2.3 Model Adversarial Attack

Adversarial attack of text model is a fast-developing research direction in the field of NLP. Normally, the attack process has two steps: finding the key positions to conduct the perturb operations, then generating the replacement text to these positions. The adversarial attacker usually utilizes different types of information as its observation, such as gradient, probability, prediction decision, or value of loss function [Papernot et al., 2016; Li et al., 2019]. Some automated model adversarial attack toolkits integrate existing methods [Wallace and others, 2019; Zeng and others, 2021]. As far as we know, there is still no previous work that uses the attention matrix as the observation.

2.4 Reinforcement Learning (RL) in NLP

There have already been various NLP tasks that are solved by the Reinforcement Learning (RL) techniques, such as question answering [Xiong et al., 2018], dialogue generation [Li et al., 2016b], and so on. In general, the RL agent model does not carry any language knowledge on itself. It is generally used for auxiliary tasks, such as ensuring the output is properly formatted [Zhong et al., 2017]. Some recent works focused on using RL methods to improve the performance of NLP models. [Li et al., 2016a] proposed an RL-based method by removing the minimum number of words to change the model's prediction. In [Xu et al., 2019], the authors improved the robustness of the text classification model by using RL to generate semantically similar samples. Also, RL can be used to generate adversarial examples to improve the robustness of machine translations [Zou et al., 2020].

3 Framework

In this section, we introduce the technical details of the proposed ATTEXPLAINER framework. The overview is shown in Figure 2.

3.1 Transformer Model

In our task setting, a Transformer classifier is used as a model being explained or a victim model for the adversarial attack.

Given an input token sequence X = (x1, ..., xn), where n is the length of the sequence. We use a well-trained BERT model as an instance of the Transformer, which is formalized as a function T(·). ˆ y is the model's predicted label with ˆ y = T(X). p(ˆ y|X) is the output probability for label ˆ y. The cross-entropy loss for classification is denoted as L(ˆ y, y), where y is the golden label.

Since BERT adopts dot-product style attention, in each layer, the score of the i-th token xi in the attention matrix A ∈ Rn×n is computed by a similarity function ϕ(Q,K) on the input query Q ∈ Rn×dim and key K ∈ Rn×dim followed by an activation layer of softmax:

where the similarity function is ϕ(Q, K) = Q×KT /√dim, and dim is the hidden dimension of Transformer. The attention matrix that averaged over all layers is Att ∈ R(n,n) .

3.2 RLGame Environment

The environment treats one sample's analysis process as a game round. There are two visibility for a single token to the model: masked or unmasked. For a complete sequence, all tokens' status can be combined into one boolean mask sequence et = (e(t,1), ..., e(t,n)), e(t,i) ∈ {True, False}, in which e(t,i) = True means token i is visible to the Transformer model at step t. Applying mask sequence e on token sequence X is formalized as X ⋆ e. So the input sequence to the model at step t is Xt = X ⋆ et. In the initial step, e.g., t = 0, all e(0,i) is set to 1.

Environment. The environmental state s is composed of the current token masking state et and attention matrix Attt. A flag bit d indicates whether the game is over or not.

The model interpretation aims to find an optimal combination of token masks that maximizes the reduction of the model's output probability on a given label. We denote this reduction in probability as ∆p = p(ˆ y|X) - p(ˆ y|X′ ). The game ends when ∆p is reached to a certain threshold ε.

The adversarial attack aims to find an adversarial sample X′ as quickly as possible to flip the output of the model. The game ends when T(X) ̸= T(X′ ), which means the attack action is successful.

Action. The agent action is denoted as an integer at, which indicates the next position index to inverse. The action of step t will affect the input of the next step, e(t+1,at) is the reverse of e(t,at). In this environment, the agent is encouraged to select actions that can obtain the largest overall rewards. In order not to affect the text classification task-specific sequence input format, it is not allowed to choose positions of special tokens, such as [CLS], [SEP], and [PAD].

Reward. The environment reward is made up of three aspects of the game: the model classification loss or label's probability, token mask rate mt = 1n ∑n et, and game status dt.

In the model interpretability task, we design the reward function as Eq.(2). In the adversarial attack task, the reward function is Eq.(3).

where α is the final reward for game completion and β is time punishment. The reward functions encourage a lower token mask rate and lower time cost.

3.3 DQNAgent

We adopt a deep Q-learning network (DQN), denoted as Q(s, a), to predict the state-action value of expected return at each step and the goal is to maximize J = E [∑∞t=0 γ t rt], where γ is the discount factor.

As presented in Figure 3, we extract numerical distributions and statistical features by performing histogram statistics along rows and columns from the attention matrix Att. Specifically, We fix the bin size of the histogram to b. For an input sentence with length n, we compute histograms along the two dimensions separately, and the value is divided by n to normalize. Stitch them together with the mean and variance of the rows and columns, resulting in a statistical feature matrix F in the shape of n × (2b + 4). OutputLayer and HiddenLayer are two linear layers. Therefore, the structure of the value network is:

The ϵ-greedy policy is adopted to balance the explorationexploitation, in which the action at = argmax Q(st) is either determined with probability ϵ, or a random position is selected as action with probability 1 - ϵ. The agent chooses one position to change its status in each game step.

We use prioritized experience replay (PER) [Schaul et al., 2016] to store transitions and sample DQN training batch. The PER importance-sampling weights for one sampling example is indicated by w. Therefore, the overall TD loss is:

The detailed of DQN training progress is presented in Algorithm 1.

Algorithm 1 DQN training progress

Input: Transformer model T , input token sequence X. Variables: DQN eval net Q and target net Q′ , agent replay buffer, maximum game steps M, value net parameters replace frequency C.

• 1: Initialize token sequence X0 =X, mask sequence e0.
• 2: Compute ˆ y0 = T(X0), get Att0,L0, d0, then s0 = [Att0, e0, d0].
• 3: Set current step t = 1.
• 4: while t <M and dt-1 ̸= True do
• 5: The DQN agent gives an action at by observing st-1.
• 6: Update the environment stage et with at, then Xt = X ⋆ et.
• 7: Compute ˆ yt = T(Xt), get Attt,Lt. Set dt = True if the end condition has been reached. Compose new environment stage st = [Attt, et, dt]. Compute the action reward rt.
• 8: Store (st-1, at-1,rt,st,at) into replay buffer.
• 9: Sample a mini-batch of transitions using PER from replay buffer. Perform one step training on Q.
• 10: Reset target net with the parameters of A value net Q′ ← Q every C steps.
• 11: end while

Output: The DQN value net Q.

4 Experiment

In this section, we conduct extensive experiments to verify the effectiveness of the proposed method in terms of two different tasks, namely, model explanation and adversarial attack. More precisely, we focus on investigating the following research questions (RQs):

• RQ1: How does the proposed model perform in terms of model explanation and text adversarial attack tasks?
• RQ2: Does the attention matrix necessarily deliver the key information of input tokens?
• RQ3: Are parameters of trained agents transferable within different Transformer models and datasets?

4.1 Experimental Settings

Datasets. We used two types of text classification settings: single sentence classification and Natural Language Inference (NLI). Specifically, single sentence classification datasets include the Emotion dataset [Saravia et al., 2018] and the Stanford Sentiment Treebank (SST2) dataset [Wang and others, 2019]. The NLI dataset we used is the SNLI corpus [Bowman and others, 2015]. The details of these datasets are presented in Table 1. We utilize the fine-tuned model pretrained by Huggingface [Wolf et al., 2019] 1 .

Baselines. For model interpretation scenarios, we selected five classical approaches for comparison: FeatureAblation, Occlusion [Zeiler and Fergus, 2014], KernelShap [Lundberg and Lee, 2017], Shapley [Castro et al., 2009] and LIME [Ribeiro et al., 2016]. These methods are implemented by Captum toolkit [Kokhlikyan et al., 2020].

For textual adversarial attack scenarios, we selected the following methods for comparison: PWWS [Ren et al., 2019], Genetic [Alzantot and others, 2018], SCPN [Iyyer et al., 2018], HotFlip [Ebrahimi et al., 2018], GAN [Zhao et al., 2018], DeepWordBug [Gao et al., 2018], TextBugger[Li et al., 2019], PSO [Zang et al., 2020], BERTAttacker [Li et al., 2020]. These attack methods are implemented by OpenAttack [Zeng and others, 2021] toolkit.

Evaluation Metrics. We use 4 indicators to evaluate the model explanation methods: Fidelity, Model Query Times, ∆p and Token Masked Rate. We use Fidelity to measure the effectiveness of the masking operation:

1 We use BertForSequenceClassification as the base model architecture. The names of the well-trained model are EM-1: AdapterHub/bert-base-uncased-pf-emotion, EM-2: nateraw/bertbase-uncased-emotion, SSTM: textattack/bert-base-uncased-SST-2 and SNLIM: textattack/bert-base-uncased-snl. These well-trained models have reached a high accuracy on the corresponding datasets, the average evaluation accuracy is shown in Table 1.

We adopt several commonly used metrics in model adversarial attack scenarios: Attack Success Rate, Victim Model Query Times, and Word Modification Rate.

Implementation Details. We trained our model on multiple Titan RTX graphics cards. The max game step is limited to 100. Max size of replacing buffer for PER is 100,000. The Adam optimizer [Kingma and Ba, 2015] was used for training the DQN model at a fixed learning rate 10-4 , training batch size is 256. We set α = 10, β = 0.2, ϵ = 0.7, γ = 0.9. The number of feature bins b is fixed at 32. The parameters of the DQN target net are replaced by the eval net every 100 learning steps. We use the special token [MASK] to replace the original token in masking operation.

4.2 Experimental Results

Model Explanation Results (RQ1)

The performance of ATTEXPLAINER serving as a model explainer is presented in Table 2. It can be observed that our method achieves surprisingly competitive performance in terms of Fidelity and ∆p with a much lower Token Masked Rate. We argue that, given two model interpretation methods, if their performances in ∆p are comparatively close, while one of them can find fewer key positions (i.e., with lower Token Masked Rate) that determine the classification result, it means the corresponding method is more effective than the other one. For example, according to the experimental results on the Emotion dataset, the LIME tends to explain that 66.9% of tokens in a sentence are supporting the classification decision. In contrast, the results of the ATTEXPLAINER show that we only need 28.7% of the tokens to be masked while keeping a comparable ∆p of 83.6%, indicating the superiority of the proposed method in finding more precise key positions. Meanwhile, we can also find that our proposed method spends less query time to achieve a better balance between ∆p and Token Masked Rate, surpassing all the baselines. Note that our approach does not perform optimally on the SST2 dataset, which is also challenging for all baselines. Though the Shapley method achieves a higher ∆p metric, it is exponential in time cost, but the marginal benefits will diminish when increasing query times. Overall, compared to existing methods on all three datasets, the ATTEXPLAINER generally makes the searching both faster and more accurate, and we attribute this to the design of introducing the attention matrix as heuristic information in our method.

Adversarial Attack Results (RQ1)

The results of the model adversarial attack are shown in Table 3. Overall, the proposed ATTEXPLAINER achieved the highest attack success rate, relatively lower victim model query times and word modification rate in three datasets compared to baseline methods. More precisely, for the single sentence setting, our method surpasses the second runner-ups with only 1.5% and 0.2% absolute increment in terms of attack success rate on both Emotion and SST2 detests, individually. In contrast, there is a comparatively significant performance boost of 13.1% between the existing state-of-the-art method (e.g. PWWS) and our method in the NLI setting, indicating the superiority of our proposed model in more complicated scenarios.

Interpretability of Attention Matrix (RQ2)

We conduct extensive experiments to validate the effectiveness of the proposed method in providing interpretable information based on the attention matrix. Specifically, we generate a random matrix as environment observation to replace the real attention matrix in each game step, while other settings are kept unchanged.

As can be observed from Figure 4, compared to the normal training setting (Emotion@1M), the agents represented by the remaining two curves fall into a poor solution that masks most of the tokens to let the decision flip. On one hand, this reveals that there is potential token importance information in the attention matrix guiding the agent to choose a better gaming strategy. On the other hand, it also highlights the necessity of introducing the RL method due to the difficulty in finding a reasonable masking combination in the huge searching space without any textual information, other model-related observation, or external knowledge.

Transferability Study (RQ3)

Since the decision-making process of the agent does not refer to the text information, we studied whether the agent's parameters can be transferred to other well-trained models. In the scenario of the adversarial attack, the transferability result is shown in Table 4. Agents trained in three training settings will be tested in four different evaluation settings. The table shows that, in a simpler dataset (Emotion), all settings can achieve well results. But the agent which is trained on a more difficult dataset (SST2) has a better transferability to other datasets. There is also a significant effect when migrating to different well-trained models, from EM-1 to EM-2.

4.3 Case Study

In this section, we give some practical examples for analysis. For the same input, we use two marking colors to indicate the different positions found by the explainer and attacker. The examples are shown in Table 5. In most cases, the key positions found by the explainer and attacker overlap with each other. The explainer found more tokens than the attacker. This coincides with the training aims of both. We find that the explainer prefers to choose content words (e.g., nouns, verbs), while the attacker tends to choose function words (e.g., prepositions, articles) to finish faster. Such results inspire us to consider whether these pre-trained language models are indeed robust.

Two visualization examples are shown in Figure 5. At each time step, we plotted a set of charts, each chart containing four subplots: attention matrix, label prediction probability, heat-map of extracted features, and Q value distribution. The agent chooses the next mask position based on the maximum Q value.

Step Mask Action

Figure 5: Visualization of the game process in adversarial attack scenario. The original image of BERT attention matrix in subplot is marked by 1⃝, output probability of label classification is marked by 2⃝, the numerical statistical features is marked by 3⃝, and Q value as the measure for next action selection is marked by 4⃝.

5 Conclusion

In this paper, we present a reinforcement learning-based method, named ATTEXPLAINER, for the Transformer explanation. Especially, a deep Q-network (DQN) based agent is designed to select the crucial position in a sentence stepby-step, with the benefit of observing attention matrices. By designing different reward functions, we have tested our approach in two tasks: perturbation-based model explanation and text adversarial attack. In addition, we found that the parameters of the DQN agent can be migrated within the same Transformer architecture. Through comparative experiments, we believe that the attention matrix can provide heuristic information for the agent to speed up the task. We also perform in-depth case studies to visualize the game processes intuitively. In line with human's intuition, we discover that the attention matrix indeed contains potential token importance information, which can be used as a basis for various applications, such as model explanation and model attack. These can be the direction of our future work.

Acknowledgments

This work has been funded by the National Natural Science Foundation of China (No.62072212), the Development Project of Jilin Province of China (20200401083GX, 20200403172SF).