Attention is not not Explanation

In plain words

Neural text models often use an attention step that gives each word a score for how much it matters to the answer. A prior paper argued that these scores cannot be trusted as explanations, because you can find different scores that lead to almost the same prediction. This paper says that the answer depends on what you mean by "explanation". It proposes a few checks that tell you whether attention is doing real work in a given model and dataset.

The checks are simple. Replace attention with evenly spread weights and see whether accuracy drops. Train the same model with different random seeds and see how much the attention changes anyway. Freeze the attention scores and train a much simpler model on top of them. Finally, train a second model to produce very different attention while keeping its predictions close to the original.

The authors find that such alternative attention can be built for some classification tasks. But the alternatives are poor guides when a simple model is trained with them, so the earlier claim is not overturned. The paper is a guide to testing attention in your own model, not a verdict on attention in general.

TL;DR

Jain and Wallace argued that attention weights are not explanations because different attention distributions can give similar predictions. This paper says that claim depends on the definition of explanation. It adds four tests: a uniform-attention baseline, a seed-variance baseline, a frozen-weight MLP diagnostic and a consistent adversary. Adversarial attention can be found for some binary LSTM tasks, but it performs poorly as a guide, so the paper does not disprove the usefulness of attention. Results are author-reported and limited to binary tasks with a single-layer bidirectional LSTM.

Problem

Attention scores are often read as explanations of which inputs drove a prediction. Jain and Wallace's claim is that this reading fails when alternative attention distributions give near-identical predictions §2. This paper argues that their experiments leave too much freedom in the setup: the adversary is trained per instance, the base attention is detached from the model, and there is no baseline for how much attention varies between ordinary training runs §2.1. It also proposes checks for whether attention is needed at all §3.2 and whether it carries token importance that a non-contextual model can use §3.4.

Method

The paper studies a single-layer bidirectional LSTM with tanh activation and additive attention, which matches the setup of Jain and Wallace §3.1. It uses the binary subset of their tasks and analyses on the test set §3.1. The authors state that the attention module can be replaced at four points, each covered by one test fig. 1.

Metrics. Predictions are compared with total variation distance (TVD) and attention with Jensen-Shannon divergence (JSD). The equations are unnumbered in the source, so they are cited by section §3.1:

\operatorname{TVD}(\hat y_1, \hat y_2) = \frac{1}{2} \sum_{i=1}^{|\mathcal{Y}|} \left|\hat y_{1i} - \hat y_{2i}\right| §3.1

\operatorname{JSD}(\alpha_1, \alpha_2) = \frac{1}{2} \mathrm{KL}[\alpha_1 \parallel \bar{\alpha}] + \frac{1}{2}\mathrm{KL}[\alpha_2 \parallel \bar{\alpha}], \quad \bar{\alpha} = \tfrac{\alpha_1 + \alpha_2}{2} §3.1

Test 1, uniform attention. The attention is frozen to uniform weights over the hidden states §3.2. For AG News and 20 Newsgroups the uniform model matches the base model, so the paper drops those two tasks; SST is kept as a borderline case §3.2, tab. 2.

Test 2, seed variance. Eight models are trained from the main setup with different initialisation seeds §3.3. The spread of attention between these seeds is the baseline against which the adversary's divergence is judged. SST attention is robust to seed changes, while the negative class of Diabetes is already arbitrary across seeds §3.3, fig. 3.

Test 3, frozen-weight MLP guide. A non-contextual MLP is trained with attention weights imposed per instance. Each token gets its own affine layer with tanh and no access to neighbouring tokens §3.4, fig. 4. Four weight sources are compared: uniform, a trained MLP attention, the base LSTM attention and the adversary's attention §3.4, tab. 3.

Test 4, consistent adversary. A second model is trained with the loss below, where the prediction distance to the base model is traded against the KL divergence between the two attention distributions §4. The adversary is trained on the whole training set rather than per instance §4.

\mathcal{L}(\mathcal{M}_a, \mathcal{M}_b)^{(i)} = \operatorname{TVD}(\hat{y}^{(i)}_a,\hat{y}^{(i)}_b) - \lambda~ \mathrm{KL}(\boldsymbol{\alpha}^{(i)}_a \parallel \boldsymbol{\alpha}^{(i)}_b) §4

The paper sweeps \lambda per dataset tab. 5. For each dataset it reports the best F1 among models whose average JSD exceeds 0.4 tab. 4. The TVD/JSD trade-off curves are in fig. 5, and an IMDb example comparing base, Jain-and-Wallace and consistent-adversary attention is in fig. 2.

Definitions. The paper separates plausible explanations from faithful ones. Under the plausibility reading, multiple valid attention distributions do not invalidate attention as an explanation §5. The four tests are described as functionally grounded, meaning they need no human judgements §5.

Pseudocode for the adversary's loss, written from the paper's equation and the code's implementation. The paper has no algorithm block §4.

# Adversary loss for one minibatch  [§4]
# Input:  tokens (B, L); base prediction p_base (B,); base attention a_base
#         (B, L); labels (B,); lambda_ (scalar); pos_weight (scalar)
# Output: scalar loss to minimise

def adversary_loss(tokens, p_base, a_base, labels, lambda_, pos_weight):
    # Own encoder and tanh attention  [§3.1] [code: Attention.py#L20-L25]
    states = encoder(tokens)                      # (B, L, 2H), H = 128
    a_adv = tanh_attention(states)                # (B, L), zero on padding
    context = sum_over_positions(a_adv * states)  # (B, 2H)
    p_adv = sigmoid(output_layer(context))        # (B,), positive score
    # Distance on the positive score only  [§4] [code: metrics.py#L9-L10]
    tvd = sum(0.5 * abs(p_adv - p_base))          # scalar
    # KL(a_adv || a_base), summed over batch and positions  [§4]
    # Treat 0 log 0 as 0 on padding positions
    kl = sum(a_adv * (log(a_adv) - log(a_base)))  # scalar
    # Keep predictions close, push attention apart  [§4]
    loss_orig = tvd - lambda_ * kl                # scalar
    # Positives weighted by pos_weight  [code: Binary_Classification.py#L132]
    weights = labels * pos_weight + (1 - labels)  # (B,)
    loss = loss_orig * sum(weights)               # scalar
    return loss

Hyperparameters

The paper defers most training settings to Jain and Wallace §3.1. The rows marked "code" were read from the authors' repository at code_commit, not from the paper.

Results

Author-reported. Setup: binary tasks, a single-layer bidirectional LSTM, all analysis on the test set, F1 on the positive class §3.1. Hardware: one local Tesla K40 app. A. The figures are not reproduced here.

Dataset statistics tab. 1:

Base attention versus uniform attention, F1 on the positive class tab. 2:

"Reported" is Jain and Wallace's number; "reproduced" is this paper's run tab. 2.

Attention weights imposed on the MLP guide, F1 on the positive class tab. 3:

Best adversary per dataset, among models with average JSD above 0.4 tab. 4:

Full adversary sweep, every λ tested tab. 5. The caption says no post-selection was made, and that SST was trained for 80 epochs, the other datasets for 40 app. B:

Headline reading: adversarial attention with JSD above 0.4 exists on each of the four datasets kept for this test tab. 4, but when used as the MLP's weights it scores well below base attention on Diabetes, SST and IMDb. Anemia is the exception tab. 3. The paper attributes Anemia to its skewed label distribution and to positive instances that contain a few indicative tokens §4, footnote.

Paper vs code

The official repository is https://github.com/sarahwie/attention, owned by the first author. It is GPL-3.0, so only the citations and pseudocode above come from it; no code is quoted.

• Adversary epochs. The paper says 40 epochs, with 80 for SST app. B. run_adversarial.sh passes --n_iters 80 for every dataset run_adversarial.sh.
• Adversary model selection. The appendix caption says the SST epoch is chosen by the loss objective on the test set app. B. The code saves the best adversary by training loss and stops early after 10 non-improving iterations TrainerBC.py. The code does not select on the test set here.
• TVD definition. The paper's TVD sums over all classes §3.1. The code computes 0.5·|Δp| on the positive-class score only metrics.py Binary_Classification.py. For a binary task the paper's formula equals |Δp|, so the code's TVD is half of it. Under the paper's formula, the code's λ = 0.001 corresponds to λ = 0.002. This is derived from the formulas, not confirmed by running them.
• Metric aggregation. The paper reports an "instance-average" JSD tab. 4. The code's TVD and JSD are sums over instances Binary_Classification.py metrics.py. I found no division by N in the code. Check this before comparing against the tables.
• Uniform attention. The paper says uniform weights over the hidden states §3.2. The code gives 1/(L-2) over the interior positions and masks the first and last position modelUtils.py modelUtils.py modelUtils.py. The paper does not mention this masking.
• Base-model epochs. The paper gives none. The code default is 8 train_and_run_experiments_bc.py.
• Hyperparameters. The paper says it uses Jain and Wallace's hyperparameters §3.1. The values in the Hyperparameters table are the code's, not the paper's.
• Adversary initialisation. The paper does not say how the adversary is initialised. The code builds a fresh model for it TrainerBC.py.

Pitfalls

• The paper leaves the training recipe to another paper. Optimiser, epochs, batch size and length filters are all absent from the text. An implementer who follows only the paper will not reproduce the numbers. The code's values are in the Hyperparameters table, with citations.
• Base-model selection uses the test set. The base model is chosen by test-set ROC-AUC TrainerBC.py TrainerBC.py. The analysis is also on the test set §3.1. Reported numbers are therefore optimistic.
• Table 4 is picked on test metrics. It keeps the highest-F1 model among those with JSD above 0.4 tab. 4. This is a selection, not a fixed setting.
• Two tasks were dropped after seeing the uniform results. AG News and 20 Newsgroups were removed because uniform attention matched base attention §3.2. This is a post-hoc choice.
• The adversary's objective is scaled by the batch's class weights. loss_orig is one scalar per batch, and the code multiplies it by a per-instance class weight and then sums Binary_Classification.py Binary_Classification.py. The same factor applies to the TVD and KL terms, so λ's trade-off is unchanged, but the overall scale depends on the label mix in each batch.
• Zero attention on padding. Masked positions are zero in both attention distributions. The KL term needs 0·log 0 = 0 there. The code's kld helper instead adds 1e-10 and clamps to [0, 1] modelUtils.py. Using a different convention may change the JSD values.
• Scale of JSD. JSD is bounded by ln 2 ≈ 0.693 fig. 5. Several adversarial runs sit at that ceiling tab. 5.
• Preprocessing is not in the repository. The repo loads pickled datasets and embeddings from a preprocess/ directory that is not included [code: Trainers/DatasetBC.py#L42]. Data splits and embeddings cannot be checked from the repository alone.
• The Anemia exception is explained by data, not tested. The paper's reason for the Anemia outlier is the label skew and the sparse indicative tokens §4, footnote. It is a plausible reading, not a controlled test.
• Scope is narrow. Only binary tasks and a single-layer bidirectional LSTM were tested. Sequence-to-sequence models, QA and other languages are named as future work §6.

Lineage

• Builds on: Jain and Wallace's "Attention is not Explanation", which the paper reproduces and critiques §2 (named and linked on arXiv: 1902.10186; no card). The paper thanks them for code that made reproduction fast [§6, acknowledgments].
• Uses: additive attention of Bahdanau et al. §3.1 (no card).
• Compared against: Jain and Wallace's per-instance adversary §4, fig. 5.
• Superseded by: not stated.