Generic Attention-model Explainability for Interpreting Bi-Modal and Encoder-Decoder Transformers

In plain words

Transformers now read pictures together with text, and some of them output things such as object boxes. When one of these models gives an answer, it is hard to say which parts of the input drove that answer. The usual tools look at the attention inside one input only. They miss the links between two inputs and the links from an output back to an input, which these models rely on.

This paper proposes a way to follow influence through every attention layer, including those links. It scores each attention map by how much the output depends on it, keeps the positive part, averages over the attention heads, and adds the result up layer by layer. The output is a heat map over image regions or words that shows what the model used for a given answer or detected object. The authors test it on a visual question answering model with two inputs, on an object detector that writes its outputs by querying the image, and on a single-image Transformer. They report that it beats the earlier methods they compare against.

TL;DR

• A method for explaining any attention-based Transformer, including models that mix two modalities (co-attention) and encoder-decoder models such as DETR §1 §3.
• Each attention map is weighted by the gradient of the target score, its positive part is kept, and heads are averaged eq. 5. Relevancy matrices are then updated layer by layer with the rules in eq. 6 to eq. 10.
• The main method uses raw attention, not LRP. The LRP variant is optional and changes little app. E tab. 5.
• Headline numbers are author-reported: DETR weakly supervised segmentation AP 13.1 against 7.2 for the best baseline tab. 1; ViT perturbation results close to Transformer attribution tab. 2.
• Some details in the released code differ from the paper or are not stated in it. See Paper vs code and Pitfalls.

Problem

Attention maps alone are a poor explanation for a Transformer. They ignore the other parts of each layer and the way the tokens are mixed across layers §2. Methods that follow self-attention, such as rollout and Transformer attribution, do not give a rule for attention between two modalities or from decoder queries to encoder tokens §1 §2. Co-attention models (for example LXMERT and ViLBERT) and encoder-decoder models (for example DETR) are common in vision, and the paper says no earlier method handled them completely §1 §2.

The method needs the gradient of the target score with respect to each attention map eq. 5. Classification relevance is read from the row of the classification token §3.1. The method is applied to an already trained model. The paper does not train anything.

Method

Relevancy matrices §3. For two modalities with t text tokens and i image tokens, the method keeps one relevancy matrix per interaction: \mathbf{R}^{tt} and \mathbf{R}^{ii} for self-attention within a modality, and \mathbf{R}^{ti} and \mathbf{R}^{it} for attention across modalities. Self-attention relevancies start as the identity, because a token only relates to itself before any attention. Cross-modal relevancies start at zero, because the modalities are separate before the first cross-attention layer §3 eq. 1 eq. 2.

\mathbf{R}^{ii} = \mathbb{I}^{i\times i}, \quad \mathbf{R}^{tt} = \mathbb{I}^{t\times t} \quad [\text{eq. 1}]

\mathbf{R}^{it} = \mathbf{0}^{i\times t}, \quad \mathbf{R}^{ti} = \mathbf{0}^{t\times i} \quad [\text{eq. 2}]

The attention being explained is the standard one, \mathbf{A} = \mathrm{softmax}(\mathbf{Q}\mathbf{K}^\top / \sqrt{d_h}) and \mathbf{O} = \mathbf{A}\mathbf{V}, with h heads eq. 3 eq. 4 §3. In co-attention, the queries come from one modality and the keys and values from the other.

Head averaging §3 eq. 5. The paper follows Chefer et al. (2020) and uses gradients to combine heads. For each layer, the gradient of the target output y_t with respect to the attention map \mathbf{A} multiplies \mathbf{A} element by element. The negative entries are removed, and the result is averaged over heads:

\bar{\mathbf{A}} = \mathbb{E}_h\big((\nabla \mathbf{A} \odot \mathbf{A})^{+}\big), \quad \nabla \mathbf{A} := \frac{\partial y_t}{\partial \mathbf{A}} \quad [\text{eq. 5}]

Here \odot is the element-wise product, (\cdot)^{+} keeps the positive part, and \mathbb{E}_h is the mean over the head dimension §3.

Self-attention update §3 eq. 6 eq. 7. For a same-modality layer with s the attending side, the accumulated relevancies are updated as follows. Eq. 6 mixes the tokens already contextualised by earlier layers. Eq. 7 mixes the cross-modal context carried in \mathbf{R}^{sq}.

\mathbf{R}^{ss} = \mathbf{R}^{ss} + \bar{\mathbf{A}} \cdot \mathbf{R}^{ss} \quad [\text{eq. 6}]

\mathbf{R}^{sq} = \mathbf{R}^{sq} + \bar{\mathbf{A}} \cdot \mathbf{R}^{sq} \quad [\text{eq. 7}]

Normalisation for cross-attention §3 eq. 8 eq. 9. The identity from initialisation is removed from each accumulated self-attention matrix. Each row of the remainder is divided by its sum, and the identity is added back. The paper notes that these layer contributions are typically much smaller than the identity after gradient-based head averaging, and wants a token's influence on itself and its contextualisation by attention to count equally §3.

\hat{\mathbf{S}}^{xx}_{m,n} = \sum_{k} \hat{\mathbf{R}}^{xx}_{m,k}, \quad \hat{\mathbf{R}}^{xx} = \mathbf{R}^{xx} - \mathbb{I}^{x\times x} \quad [\text{eq. 8}]

\bar{\mathbf{R}}^{xx} = \hat{\mathbf{R}}^{xx} / \hat{\mathbf{S}}^{xx} + \mathbb{I}^{x\times x} \quad [\text{eq. 9}]

Cross-attention update §3 eq. 10 eq. 11. For a cross-attention layer between s and q, the cross term uses the normalised self-attention of both sides. The paper also gives rule 11, which adds context from \mathbf{R}^{qs} into \mathbf{R}^{ss}.

\mathbf{R}^{sq} = \mathbf{R}^{sq} + (\bar{\mathbf{R}}^{ss})^{\top} \cdot \bar{\mathbf{A}} \cdot \bar{\mathbf{R}}^{qq} \quad [\text{eq. 10}]

\mathbf{R}^{ss} = \mathbf{R}^{ss} + \bar{\mathbf{A}} \cdot \mathbf{R}^{qs} \quad [\text{eq. 11}]

The paper applies eq. 8 to 10 symmetrically, to both modalities §3. Rule 11 is needed for co-attention models (for example LXMERT, where both streams update each other). For the DETR decoder, \mathbf{R}^{qs} is never computed, so rule 11 is not used there §3.2.

Reading the explanation §3.1. For classification, the row of the classification token gives the relevance of each token. In LXMERT the first text token is the classification token §3.2. The text-to-image relevancy \mathbf{R}^{ti} then gives the image relevance per region §3.1 §3.2.

Architecture types §3.2 fig. 2.

• (a) Pure self-attention (VisualBERT). Text and image tokens are concatenated, so there is one matrix \mathbf{R}^{(t+i,t+i)}, updated with eq. 6 only. The classification row gives text and image relevance §3.2 fig. 2a.
• (b) Self-attention plus co-attention (LXMERT). Both modalities use all the rules: eq. 6, 7 for self-attention and eq. 10, 11 for co-attention. The first text token is the classification token §3.2 fig. 2b.
• (c) Encoder-decoder (DETR). The encoder uses eq. 6 on the image tokens, so \mathbf{R}^{ee} is updated. The decoder uses eq. 6 and 7 for query self-attention, then eq. 10 for query-to-image attention. \mathbf{R}^{ed} is not computed. The relevance of detected object j is row j of \mathbf{R}^{de}, which gives one score per image region §3.2 fig. 2c.

DETR segmentation procedure app. C alg. 1. The paper turns DETR's boxes into masks. Queries whose class probability is above 50% are kept alg. 1 L.3. For each kept query, row j of \mathbf{R}^{qi} is used as a heat map over the image features alg. 1 L.6. Otsu thresholding separates foreground from background alg. 1 L.7. The mask is bilinearly upsampled to the target size, passed through a sigmoid, and kept where the value is above 0.5 alg. 1 L.8-9. Finally it is upsampled to the original image size alg. 1 L.10.

Pseudocode. The paper gives no code for the general method. The blocks below are written from the equations above, with the DETR pass as the worked example. The DETR pass follows the released code's order ExplanationGenerator.py.

# Head-averaged map for one attention layer  [§3] [eq. 5]
# Inputs: attention and grad, each (h, s, q); grad of the target score
# Output: (s, q)

def head_averaged_map(attention, grad):
    # Keep only the positive part of the gradient-weighted attention  [eq. 5]
    weighted = (grad * attention).clamp(min=0)
    # Mean over heads, after the positive part  [eq. 5]
    return weighted.mean(dim=0)

# Row-normalise an accumulated self-attention matrix  [eq. 8] [eq. 9]
# Input: R, square (n, n), starting from the identity
# Output: R_bar, (n, n)

def normalise_self(R):
    # Remove the identity from initialisation, keep layer contributions  [eq. 8]
    R_hat = R - eye(R.shape[0])
    # Divide each row by its sum  [eq. 8] [eq. 9]
    row_sums = R_hat.sum(dim=-1, keepdim=True)
    R_hat = R_hat / row_sums
    # Add the identity back  [eq. 9]
    return R_hat + eye(R.shape[0])

# Update for a same-modality attention layer  [eq. 6] [eq. 7]
# Inputs: R_ss (s, s), R_sq (s, q), A_bar (s, s)
# Outputs: updated R_ss and R_sq

def update_self_attention(R_ss, R_sq, A_bar):
    # Tokens already mixed by earlier layers: multiply by the map  [eq. 6]
    R_ss = R_ss + A_bar @ R_ss
    # Cross-modal context in R_sq is mixed in the same way  [eq. 7]
    R_sq = R_sq + A_bar @ R_sq
    return R_ss, R_sq

# Update for a cross-attention layer  [eq. 10]
# Inputs: R_sq (s, q), R_ss (s, s), R_qq (q, q), A_bar (s, q)
# Output: updated R_sq

def update_cross_attention(R_sq, R_ss, R_qq, A_bar):
    # Normalised self-attention on each side  [eq. 8] [eq. 9]
    R_ss_bar = normalise_self(R_ss)
    R_qq_bar = normalise_self(R_qq)
    # Add the cross-modal relevance this layer creates  [eq. 10]
    R_sq = R_sq + R_ss_bar.T @ A_bar @ R_qq_bar
    return R_sq

# Relevance of image regions for one detected query  [§3.2] [eq. 1] [eq. 2]
# Inputs: per-layer maps and grads, target query index
# Output: (num_regions,)

def detr_relevancy(enc_layers, dec_layers, target_query):
    # Identity within each stream, zeros across streams  [eq. 1] [eq. 2]
    R_enc = eye(num_regions)
    R_dec = eye(num_queries)
    R_dec_enc = zeros(num_queries, num_regions)
    # Encoder: self-attention over image regions  [eq. 6]
    for layer in enc_layers:
        A_bar = head_averaged_map(layer.attn, layer.grad)
        R_enc = R_enc + A_bar @ R_enc
    # Decoder: query self-attention, then query-to-region attention  [§3.2]
    for layer in dec_layers:
        A_bar = head_averaged_map(layer.self_attn, layer.self_grad)
        # Rules 6 and 7 for the query stream  [eq. 6] [eq. 7]
        R_dec, R_dec_enc = update_self_attention(R_dec, R_dec_enc, A_bar)
        A_bar = head_averaged_map(layer.cross_attn, layer.cross_grad)
        # Rule 10 for query-to-region attention  [eq. 10]
        R_dec_enc = update_cross_attention(R_dec_enc, R_dec, R_enc, A_bar)
    # Row of the detected query gives one score per region  [§3.2]
    return R_dec_enc[target_query]

Hyperparameters

The paper trains nothing, so there are no learning rates, epochs or batch sizes. The values below are the ones the text and the appendix state.

Results

All numbers are author-reported. Setup: a pre-trained model is explained, then each method's maps are tested. Baselines are rollout, raw attention, Grad-CAM, partial LRP and Transformer attribution §4. Table 1 uses the COCO validation set §5 tab. 1.

DETR weakly supervised segmentation (COCO validation) tab. 1. Higher is better. The gray column is DETR's own box-detection score, and the figure in parentheses is the paper's printed gain over the best baseline.

ViT-Base perturbation AUC (ImageNet validation) tab. 2. For negative perturbation a larger AUC is better; for positive perturbation a smaller AUC is better. The table column is labelled "LRP"; the text does not define it further §4. The table reports the predicted and the target class.

The paper states that the only difference from Transformer attribution here is the LRP head averaging, and that the gap is mild §5.

VisualBERT and LXMERT perturbation (VQA) fig. 3 fig. 5. The paper says that for LXMERT the image perturbation shows the largest margin over the other methods. Negative perturbation keeps accuracy after removing more than 80% of image tokens, and positive perturbation falls steeply §5 fig. 3. For VisualBERT, the paper reports results close to Transformer attribution, and it says the remaining difference comes from the choice of LRP in head averaging. On the text perturbation it reports a sizeable margin over Transformer attribution §5 fig. 5. The LXMERT text perturbation is in fig. 3 (c) and (d). The numeric AUC values appear only in these figures and are not transcribed here.

LXMERT ablations (VQA AUC) tab. 4. Higher is better for negative perturbation, lower for positive. Removing the normalisation or the self-attention term moves each value by about one point or less; removing aggregation changes the values by tens of points.

DETR ablations (segmentation) tab. 3. Higher is better.

LRP variant on LXMERT (VQA AUC) tab. 5. Our method with LRP app. E.

The paper's conclusion is that LRP is not needed, since the variant with LRP is close to the version without it app. E.

Paper vs code

The official repository is hila-chefer/Transformer-MM-Explainability at the commit above. Its LICENSE is MIT (copyright 2021 Hila Chefer). The repository also contains vendored subprojects: lxmert/lxmert/LICENSE (MIT) and VisualBERT/LICENSE (BSD). The DETR folder has no LICENSE file in this checkout. No code is quoted in this card.

• Main DETR method matches the paper. The DETR command-line default is ours_no_lrp main.py, which uses raw attention as in eq. 5. The LRP variant is the separate method ours_with_lrp mask_generator.py.
• Default flag differs. generate_ours defaults to use_lrp=True ExplanationGenerator.py. The LXMERT version does the same ExplanationGenerator.py. The LXMERT perturbation script constructs its model with use_lrp=True perturbation.py. Method ours_no_lrp is selected separately perturbation.py.
• Rules match. Rules 6, 7 and 10 match eq. 6, 7 and 10 in the DETR code ExplanationGenerator.py ExplanationGenerator.py ExplanationGenerator.py. Eq. 8 and 9 match ExplanationGenerator.py. Rule 11 is implemented in the LXMERT code only ExplanationGenerator.py, consistent with §3.2.
• Head averaging matches. Positive part first, then mean over heads ExplanationGenerator.py.
• Target score is the logit. The code back-propagates the pre-softmax score pred_logits ExplanationGenerator.py ExplanationGenerator.py. The paper's "model's output for class t" does not say which.
• Target class is chosen by the code. The class is the argmax over the non-background classes for the given query ExplanationGenerator.py.
• LXMERT zeroes the classification self-relevance. The code sets R_t_t[0,0] = 0 ExplanationGenerator.py. The paper does not mention this. It does not change the ranking used in the text perturbation, because index 0 is excluded there perturbation.py and min-max normalisation is monotone perturbation.py.
• Perturbation grids differ from each other. LXMERT uses [0, 0.25, 0.5, 0.75, 0.8, 0.85, 0.9, 0.95, 1] perturbation.py. VisualBERT uses a different image grid and the same text grid evaluation_loop.py. The paper gives no grid.
• Accuracy is the VQA score. LXMERT adds the VQA label score of the top answer perturbation.py. The paper says "top-1 accuracy" §5.
• Minimum IoU 0.2 is not in the code. The paper says the minimum IoU was lowered to 0.2 §5. The only IoU setting in the DETR evaluation is a commented-out line engine.py engine.py. The released code therefore does not show the 0.2 setting.
• Masks are Otsu on quantised values. The relevance is rescaled to 0-255 and cast to 8-bit before Otsu mask_generator.py. The paper does not mention the 8-bit step.
• Mask post-processing matches Alg. 1 L.8-9. The DETR segmentation post-processing bilinearly upsamples and thresholds the sigmoid at 0.5 segmentation.py segmentation.py. The final nearest-neighbour step alg. 1 L.10 was not located in the files read.
• Compute. Each kept query triggers a full forward and backward pass mask_generator.py ExplanationGenerator.py.
• Hard-coded CUDA. The code calls .cuda() in the DETR path ExplanationGenerator.py ExplanationGenerator.py and .to("cuda") in the LXMERT perturbation perturbation.py.
• Not checked. The ViT notebook and the CLIP notebook are in the repository. The paper's ViT results were not checked against the notebook, and the CLIP notebook is outside this paper.

Pitfalls

• Use the raw-attention path for the main method. generate_ours defaults to LRP, but eq. 5 and the paper's main results use raw attention §3 app. E. Set use_lrp=False to match the paper and the DETR command-line default.
• Clamp before the mean. Eq. 5 removes the negative part before averaging over heads §3. Averaging first gives a different map. The code clamps first ExplanationGenerator.py.
• Logit or probability. The paper does not say whether y_t is a logit or a probability. The code uses the logit ExplanationGenerator.py ExplanationGenerator.py. A softmax probability would change the gradients, so pick the logit to match the code.
• Division by zero. If a row of the accumulated matrix sums to zero, eq. 9 divides by zero. The code then zeroes only the NaN entries after the matrix product ExplanationGenerator.py. Entries that depend on a NaN are then set to zero, which silently drops those contributions. Add an explicit guard. The guard is not in the paper.
• The accumulation must be cumulative. Dropping the accumulation drops DETR AP from 13.1 to 0.1 tab. 3. Each update in eq. 6 and eq. 7 must add to the previous matrix.
• Normalisation matters for DETR. Without it, DETR AP falls from 13.1 to 11.7 tab. 3. The effect on LXMERT is smaller tab. 4.
• Rule 11 is not used for DETR. The decoder never computes \mathbf{R}^{qs}, so eq. 11 does not apply there §3.2. It applies to co-attention models such as LXMERT ExplanationGenerator.py.
• The 0.2 IoU is unverified. The reported AP depends on the minimum IoU §5. The released evaluation code does not show the 0.2 setting engine.py engine.py. Do not assume the released code reproduces Table 1 without that change.
• Masks grow after upsampling. The mask is 0 or 255 at feature resolution mask_generator.py. Bilinear upsampling makes any pixel with a non-zero weight on a foreground cell positive, and sigmoid greater than 0.5 is the same as a positive value segmentation.py. The masks therefore extend up to one feature-map cell beyond the Otsu boundary. This follows from reading the code, not from a measurement.
• Unfilled queries count as background. Queries below the probability threshold keep the value -1 mask_generator.py. Its sigmoid is about 0.27, below 0.5, so they are background.
• Perturbation grids are not in the paper. The step grid changes the curve and therefore the AUC. Use the code's grid for each modality perturbation.py evaluation_loop.py.
• Off-by-one in the VisualBERT average. The loop adds the sample to the total before it checks the limit evaluation_loop.py evaluation_loop.py. It then divides by the configured count evaluation_loop.py. With 10,000 samples the bias is one part in 10,000, but it means the reported mean is not exactly the average over the stated sample count.
• AUC is not computed in the scripts. The LXMERT and VisualBERT scripts print accuracy per step. The integration rule for the AUC is not stated in the paper §5 or in the code read here.
• Text perturbation keeps the special tokens. The LXMERT text perturbation always keeps the first and last tokens perturbation.py perturbation.py. The paper gives the reason accuracy never drops to zero, but not this detail §5.
• Results are from figures for two of the four experiments. The VisualBERT and LXMERT perturbation results are shown only in figures fig. 3 fig. 5. Do not cite numbers for them.
• Cost scales with the number of queries. The DETR segmentation runs one forward and backward pass for each kept query mask_generator.py. Total cost therefore grows with the number of kept queries.

Lineage

• Builds on: Chefer et al. 2020, "Transformer Interpretability Beyond Attention Visualization". The paper reuses its self-attention rule and its head averaging, except that eq. 5 uses raw attention where that method uses LRP relevance eq. 13. It adds the cross-attention rules §1 §2 §3.
• Builds on: layer-wise relevance propagation (Bach et al. 2015), used only in the optional LRP variant §2 eq. 13.
• Compared against: rollout (Abnar and Zuidema, arXiv:2005.00928), raw attention, Grad-CAM, partial LRP (Voita et al. 2019) and Transformer attribution §4 tab. 1 tab. 2.
• Superseded by: not stated in the paper, and no later card in this wiki covers it.