Attention as Explanation - What Transformers Are (and Aren't) Looking At

Reading time: 55 min | Interview relevance: Very High - transformers dominate industry; attention interpretability is a frequent interview topic | Target roles: ML Engineer, AI Engineer, Research Engineer, NLP Engineer

---

The Clinical NLP Team That Trusted the Wrong Evidence

It is 2020. A clinical NLP team deploys a BERT model for medical note classification - specifically, identifying notes that indicate a high risk of hospital readmission. The model achieves 87% AUC on the held-out test set, which represents a meaningful improvement over the previous rule-based system. Before deployment, the team does due diligence: they visualize the attention weights for several positive predictions using BERTViz.

The visualizations look exactly right. For a note that the model flags as high readmission risk, the attention heads strongly attend to tokens like "pneumonia," "respiratory failure," "discharged against medical advice," and "O2 saturation 84%." The team presents these visualizations to the hospital's clinical committee. "See? The model is reasoning about the clinically relevant terms." The committee approves deployment.

Three months later, a follow-up analysis reveals something uncomfortable. When the team applies the findings from Jain & Wallace (2019) - a paper showing that attention weights can be dramatically altered without changing predictions - they discover that they can permute the attention distribution almost arbitrarily for many notes and the classification barely changes. The model's output is driven by something other than what the attention maps suggest. The "pneumonia" attention was a coincidence, not a cause.

The team digs further. They apply probing classifiers to the internal representations of the model, and the picture becomes clearer: clinical risk information is encoded diffusely across multiple layers and heads, not concentrated in the high-attention tokens. The model learned to extract the right signal, but the attention mechanism is not the pathway it used. The visualization showed a correlation; the team mistook it for causation.

This story plays out across the industry throughout 2019–2022. Teams using transformers for high-stakes decisions - clinical NLP, legal document review, financial analysis, content moderation - visualize attention and declare understanding. The visualizations are compelling. They confirm what the domain experts already believe. But the underlying research suggests the relationship between attention and explanation is far more complicated than it first appears.

This lesson is the definitive guide to that complication: what attention actually computes, when attention weights are informative, when they mislead, what the academic debate resolved, and what methods actually work for explaining transformer decisions. By the end, you will be able to deploy attention-based explanations responsibly - knowing which situations warrant raw attention, which warrant rollout, and which demand gradient-based attribution.

---

Why This Matters

Before transformers, explaining neural networks for text meant explaining RNNs. That was hard but structured: you could trace hidden state activations, apply LIME, or use gradient-based methods directly on the input sequence. Transformers introduced a new structure - attention - that seemed to offer an obvious explanation channel. Every transformer layer explicitly computes which tokens it is "attending to" and by how much. Those weights are right there in the computation graph.

The natural intuition: if the model attends strongly to token A when making its decision, then token A must be important for that decision. This intuition is powerful, visually compelling, and widely shared. It is also - as Jain & Wallace showed - not necessarily correct.

Understanding why requires understanding what attention actually computes, what the gradient-attention interaction reveals, and what properties a faithful explanation must have. The stakes are not academic. When companies use attention visualizations to claim their AI systems are explainable - in medical, financial, or legal contexts - they need to know whether those claims are defensible.

---

What Attention Actually Computes

The scaled dot-product attention mechanism is:

$\text{Attn}(Q, K, V) = \text{softmax}\!\left(\frac{QK^\top}{\sqrt{d_k}}\right)V$

where $Q \in \mathbb{R}^{n \times d_k}$ is the query matrix, $K \in \mathbb{R}^{n \times d_k}$ is the key matrix, $V \in \mathbb{R}^{n \times d_v}$ is the value matrix, $n$ is the sequence length, and $d_k$ is the key/query dimension.

The weight matrix we visualize is:

$\alpha = \text{softmax}\!\left(\frac{QK^\top}{\sqrt{d_k}}\right) \in \mathbb{R}^{n \times n}$

where $\alpha_{ij}$ is the weight that token $i$ assigns to token $j$. The output for token $i$ is then:

$\text{output}_i = \sum_j \alpha_{ij} V_j$

This means: the output representation for token $i$ is a weighted sum of the value vectors, weighted by the attention scores. The attention weight $\alpha_{ij}$ controls how much of value vector $V_j$ flows into the representation of position $i$.

The critical observation: what the model "uses" is the output representation, not the attention weights themselves. The attention weights determine which value vectors are combined, but the downstream computation uses only the combined output. This creates the core problem with attention-as-explanation: even if the attention weights change dramatically, if the value vectors are similar enough (which happens when similar tokens appear), the output barely changes.

Multi-Head Attention Adds Another Layer of Complexity

In practice, BERT uses multi-head attention with $h = 12$ heads:

$\text{MultiHead}(Q, K, V) = \text{Concat}(\text{head}_1, \ldots, \text{head}_h)W^O$

$\text{head}_i = \text{Attn}(QW_i^Q, KW_i^K, VW_i^V)$

Each head has its own $Q$, $K$, $V$ projections and produces its own attention distribution. The outputs are concatenated and linearly projected. This means there are $h \times L$ attention weight matrices (heads × layers) for a model with $L$ layers. BERT-base has $12 \times 12 = 144$ such matrices per input. Which one do you visualize? How do you aggregate across heads?

This is not a minor implementation detail - it goes to the heart of the explanation problem. Researchers have found that different heads specialize in different linguistic phenomena: some heads track syntactic dependencies, others track coreference, others track positional proximity. Collapsing these into a single attention map requires choices that are not theoretically grounded.

---

The Faithfulness vs Plausibility Distinction

Before the Jain & Wallace paper, the field lacked precise language for what a "good" explanation means. Jacovi & Goldberg (2020) formalized the distinction that now organizes the entire debate:

Plausibility: the explanation agrees with human intuition about what the model should be looking at. High plausibility means the highlighted tokens are the ones a domain expert would consider relevant.

Faithfulness: the explanation accurately reflects the actual computational process that led to the prediction. A faithful explanation would change proportionally if the model's reasoning changed.

These are independent properties. An explanation can be:

• Plausible and faithful: ideal
• Plausible but unfaithful: looks right but is wrong - the clinical NLP case
• Faithful but implausible: reveals the model is doing something unexpected (a bug-finding tool)
• Neither: useless

Raw attention visualization is often plausible - the highlighted tokens tend to be semantically relevant. But Jain & Wallace showed it is often not faithful. The field shifted to chasing faithfulness rather than plausibility, which drove the development of gradient-based methods.

---

The Jain & Wallace (2019) Result

The paper "Attention is not Explanation" (Jain & Wallace, NAACL 2019) makes two core arguments, each backed by experiments on text classification and question answering.

Argument 1: Attention weights do not correlate with gradient-based feature importance. Gradient-based methods (like integrated gradients) provide a principled measure of how much each input token affects the output. Jain & Wallace measure the correlation between attention weights and gradient attributions across multiple tasks and models. They use the Wilcoxon rank-sum test to compare the distribution of correlation values to zero. The test statistic shows that gradient-attention correlation is not significantly different from zero for most tasks - the tokens receiving high attention are frequently not the tokens with the highest gradient attribution.

Argument 2: Adversarial attention distributions exist. More dramatically, they show that for a trained model you can often find an alternative attention distribution $\tilde{\alpha}$ that is dramatically different from $\alpha$ (measured by total variation distance) but produces the same prediction:

$\|\tilde{\alpha} - \alpha\|_{TV} > \delta \quad \text{yet} \quad |f_{\tilde{\alpha}}(x) - f_\alpha(x)| < \epsilon$

This is the adversarial attention result. If you can change the attention distribution to something completely different while barely affecting the prediction, then the original attention distribution cannot be a faithful explanation of why the model made that prediction. A faithful explanation, by definition, would change the prediction if you altered it.

The mechanism: in models with high-dimensional value vectors, different attention distributions can produce similar output representations. The model's prediction depends on the output representation, not on the path through the attention weights. When value vectors for different tokens cluster together (semantically similar tokens have similar value vectors), any attention distribution over those tokens produces nearly the same weighted sum.

The Wiegreffe & Pinter (2019) Response

Almost simultaneously, Wiegreffe & Pinter (2019) published "Attention is not not Explanation" (yes, the double negative is intentional). Their argument is more nuanced: the question is not whether attention weights are a complete explanation, but whether they are informative at all.

They point out that Jain & Wallace's adversarial construction is not the only test of faithfulness. They propose a different framework: can you train a model to match a trained model's predictions while using a different attention distribution, starting from scratch? They call this a "decision tree" test - if attention is truly uninformative, you should be able to learn the same function with random attention distributions.

Their finding: it is often difficult or impossible to train a model that matches the original's predictions while using a uniform or random attention distribution, especially for complex tasks. This suggests attention does encode something important - just not necessarily in the way we naively visualize it.

They also propose that gradient × attention is a more faithful measure than raw attention alone. By weighting attention values by the gradient of the loss with respect to each attention weight, you filter out attention weights that the model is not sensitive to - producing attributions that are both more faithful and often more aligned with gradient-only methods.

The Current Resolution

By 2021, the field had converged on a nuanced position: attention is a useful diagnostic tool, not a faithful causal explanation.

The distinction matters:

• Diagnostic: attention patterns reveal what the model has learned to correlate. High attention to certain tokens tells you the model has learned relationships between those tokens and the output. This is informative and useful for debugging.
• Faithful causal: changing the attention weights would change the prediction proportionally. This is generally not true.

As a diagnostic tool, attention visualization remains valuable - but it must be interpreted carefully, as correlation rather than causation.

---

Gradient-Weighted Attention: A More Faithful Measure

The most principled attention-based approach combines attention weights with gradient information. The gradient-weighted attention score for attention weight $\alpha_i$ is:

$\tilde{\alpha}_i = \alpha_i \cdot \left\|\nabla_{\alpha_i}\mathcal{L}\right\|$

The gradient $\nabla_{\alpha_i}\mathcal{L}$ measures how much the loss would change if attention weight $\alpha_i$ changed. This directly addresses the adversarial attention problem: a high attention weight with zero gradient means the model's output is insensitive to this weight - so the gradient kills the attribution. A token only receives high attribution when both:

1. The attention weight is high (the token is attended to), and
2. The gradient is high (the model is actually sensitive to changes in this attention weight)

This is directly analogous to GradCAM for convolutional networks. For multi-head models, aggregate across heads by either taking the head with highest average gradient attribution, or summing:

$\tilde{\alpha}_{\text{agg}} = \sum_{h=1}^{H} \alpha_i^{(h)} \cdot \left\|\nabla_{\alpha_i^{(h)}}\mathcal{L}\right\|$

---

Attention Rollout: Tracing Information Flow Across Layers

Abnar & Zuidema (2020) proposed attention rollout - a method to track how information flows from input tokens to any position across all layers simultaneously.

The intuition: in a transformer with $L$ layers, information can flow through both attention weights and residual connections. The residual connection in each layer adds an identity term, which means each token's representation is partly its original representation (identity path) and partly information gathered from other tokens (attention path).

Attention rollout models this as:

$\hat{A}^{(l)} = 0.5 \cdot A^{(l)} + 0.5 \cdot I$

where $I$ is the identity matrix (modeling the residual connection) and $A^{(l)}$ is the average attention matrix at layer $l$ (averaged across heads).

The rollout from layer $l$ to the input is:

$R^{(l)} = \hat{A}^{(l)} \cdot \hat{A}^{(l-1)} \cdots \hat{A}^{(1)}$

The final rollout matrix $R^{(L)}$ gives the "information flow" from each input token to each output position, accounting for the multi-layer aggregation.

Limitations of rollout: it assumes all attention heads contribute equally and that the residual connection contributes exactly 50% - both simplifications that are not theoretically justified. It also does not account for the value matrix projections within each head. Think of rollout as a better heuristic than single-layer attention, not as a provably faithful attribution method.

---

Circuit Analysis in Transformers

Elhage et al. (2021) introduced mechanistic interpretability - the effort to reverse-engineer transformers into interpretable computational circuits. Their key framework: view attention heads as key-value lookup operations.

Each attention head implements:

1. A query computation: what is this position looking for?
2. A key computation: what does each position offer as a key?
3. A value computation: what information does each position contribute if selected?

From this perspective, an attention head is a differentiable associative memory: given a query, find the key that matches best, and copy the corresponding value. This is not a metaphor - the mathematics is exact.

Induction heads: one of the most important discovered circuits. An induction head is a two-layer attention pattern where:

• Head A in layer 1 is a "previous token" head - it attends to the token immediately before each position
• Head B in layer 2 is an induction head - it attends to tokens that appeared after the same token earlier in the context

Together, these two heads implement in-context learning: if the sequence contains [A][B] ... [A], the induction circuit will predict [B] after the second [A]. This is the mechanism behind few-shot prompting - the model learns the pattern from examples in the prompt and applies it.

The existence of induction heads has profound implications for attention visualization: the "important" attention weights for in-context learning predictions are in head B's attention to positions after previous occurrences of the current token. If you visualize the wrong head, you see nothing. If you know which circuit is active, the attention visualization is precisely interpretable.

---

Probing Classifiers: What Is Encoded Where?

Probing classifiers (also called diagnostic classifiers) are a systematic method for understanding what information is encoded in a transformer's internal representations at each layer.

The procedure:

1. Take a pre-trained transformer and freeze its weights
2. Extract internal representations (hidden states) at each layer for a dataset with known labels
3. Train a small linear classifier on these representations to predict the labels
4. If the linear classifier achieves high accuracy, the representation at that layer linearly encodes that information

For BERT, probing has revealed:

• Lower layers (1–4): syntactic information - part of speech, dependency relations
• Middle layers (5–9): semantic information - named entities, coreference
• Upper layers (10–12): task-specific information - for fine-tuned models, the target label is encoded here

import torch
import torch.nn as nn
from transformers import AutoTokenizer, AutoModel
from sklearn.linear_model import LogisticRegression
from sklearn.preprocessing import LabelEncoder
import numpy as np

# ─── PROBING CLASSIFIER FOR POS TAGS ─────────────────────────────────────────

model_name = "bert-base-uncased"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModel.from_pretrained(model_name, output_hidden_states=True)
model.eval()

def extract_layer_representations(
    texts: list,
    layer_idx: int,
    model,
    tokenizer,
) -> np.ndarray:
    """
    Extract token-level representations from a specific layer.
    Returns: (num_tokens, hidden_dim) array
    """
    all_reps = []
    for text in texts:
        inputs = tokenizer(text, return_tensors="pt", truncation=True, max_length=128)
        with torch.no_grad():
            outputs = model(**inputs)
        # hidden_states: tuple of (batch, seq, hidden) per layer
        # Index 0 = embedding layer, 1 = layer 1, ..., 12 = layer 12
        layer_rep = outputs.hidden_states[layer_idx]  # (1, seq, 768)
        # Take mean across tokens (sentence-level probing)
        all_reps.append(layer_rep[0].mean(dim=0).numpy())
    return np.array(all_reps)

# Compare probing accuracy across layers for sentiment
# Layer 1 (surface): low accuracy on sentiment
# Layer 12 (task-specific): high accuracy on sentiment
def probe_all_layers(texts, labels, model, tokenizer, num_layers=13):
    """
    Probe each transformer layer for a classification task.
    Returns: list of (layer_idx, accuracy) pairs.
    """
    le = LabelEncoder()
    y = le.fit_transform(labels)
    results = []

    for layer_idx in range(num_layers):
        X = extract_layer_representations(texts, layer_idx, model, tokenizer)
        # Simple 80/20 train/test split
        split = int(0.8 * len(X))
        X_train, X_test = X[:split], X[split:]
        y_train, y_test = y[:split], y[split:]

        clf = LogisticRegression(max_iter=500, random_state=42)
        clf.fit(X_train, y_train)
        acc = clf.score(X_test, y_test)
        results.append((layer_idx, acc))
        print(f"  Layer {layer_idx:2d}: accuracy = {acc:.4f}")

    return results

---

BERTViz: Multi-Head Attention Visualization

BERTViz (Vig, 2019) provides interactive visualizations for multi-head attention patterns in BERT and other transformers. It offers three views:

1. Head view: one attention head at a time - colored arcs from each token to the tokens it attends to
2. Model view: all heads in a grid - lets you visually identify head specialization patterns
3. Neuron view: shows individual query and key vector dimensions

import torch
import numpy as np
from transformers import AutoTokenizer, AutoModel, AutoModelForSequenceClassification
from typing import List, Tuple, Optional

model_name = "bert-base-uncased"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModel.from_pretrained(model_name, output_attentions=True)
model.eval()

text = "The patient was admitted with severe pneumonia and respiratory failure."
inputs = tokenizer(text, return_tensors="pt")
tokens = tokenizer.convert_ids_to_tokens(inputs["input_ids"][0])

# ─── 1. EXTRACT RAW ATTENTION WEIGHTS ────────────────────────────────────────

with torch.no_grad():
    outputs = model(**inputs)

# outputs.attentions: tuple of (batch, heads, seq, seq) per layer
attentions = outputs.attentions  # 12 layers for BERT-base
print(f"Number of layers: {len(attentions)}")
print(f"Attention shape per layer: {attentions[0].shape}")
# Shape: (1, 12, seq_len, seq_len) = (batch, heads, from, to)

# Average across heads for single-layer visualization
last_layer_attn = attentions[-1][0]  # (12, seq, seq)
avg_last_layer = last_layer_attn.mean(dim=0).numpy()  # (seq, seq)

# CLS token attention to all other tokens (last layer, avg heads)
cls_attention = avg_last_layer[0]  # [CLS] row
print("\n[CLS] attention to tokens (last layer, avg heads):")
for token, weight in zip(tokens, cls_attention):
    bar = "█" * int(weight * 50)
    print(f"  {token:20s} {weight:.4f} {bar}")

# ─── 2. HEAD SPECIALIZATION ANALYSIS ─────────────────────────────────────────

def analyze_head_patterns(attentions, tokens):
    """
    Identify which heads show interesting attention patterns:
    - Diagonal (attend to adjacent tokens)
    - CLS-focused (token attends to [CLS])
    - Long-range (attend across many positions)
    """
    n_layers = len(attentions)
    n_heads = attentions[0].shape[1]
    seq_len = attentions[0].shape[2]

    for layer_idx in range(n_layers):
        for head_idx in range(n_heads):
            attn = attentions[layer_idx][0, head_idx].numpy()

            # Check for diagonal pattern (adjacent token attention)
            diag_score = np.diag(attn, k=1).mean() if seq_len > 1 else 0

            # Check for CLS attention (how much does each token attend to [CLS]?)
            cls_score = attn[:, 0].mean()  # col 0 = [CLS] position

            if diag_score > 0.3:
                print(f"  Layer {layer_idx}, Head {head_idx}: ADJACENT pattern (score={diag_score:.3f})")
            elif cls_score > 0.3:
                print(f"  Layer {layer_idx}, Head {head_idx}: CLS-focused (score={cls_score:.3f})")

# analyze_head_patterns(attentions, tokens)

# ─── 3. ATTENTION ROLLOUT ─────────────────────────────────────────────────────

def attention_rollout(attentions: tuple) -> np.ndarray:
    """
    Compute attention rollout following Abnar & Zuidema (2020).
    Accounts for residual connections by mixing 50% identity.
    Returns: rollout matrix R (seq_len, seq_len)
    """
    seq_len = attentions[0].shape[-1]
    rollout = np.eye(seq_len)  # start with identity

    for layer_attn in attentions:
        # Average across heads: (seq, seq)
        avg_attn = layer_attn[0].mean(dim=0).numpy()

        # Add residual: 0.5 * attention + 0.5 * identity
        avg_attn_with_residual = 0.5 * avg_attn + 0.5 * np.eye(seq_len)

        # Normalize rows to sum to 1
        row_sums = avg_attn_with_residual.sum(axis=-1, keepdims=True)
        avg_attn_with_residual = avg_attn_with_residual / row_sums

        # Chain: R = A_l * R_{l-1}
        rollout = avg_attn_with_residual @ rollout

    return rollout

rollout = attention_rollout(attentions)
cls_rollout = rollout[0]  # [CLS] token's information sources

print("\nAttention Rollout - [CLS] token (multi-layer aggregated):")
for token, weight in zip(tokens, cls_rollout):
    bar = "█" * int(weight * 200)
    print(f"  {token:20s} {weight:.4f} {bar}")

# ─── 4. GRADIENT × ATTENTION ──────────────────────────────────────────────────

def grad_attention_attribution(
    model,
    inputs: dict,
    class_idx: int = 1,
) -> np.ndarray:
    """
    Compute Gradient × Attention attribution: α · ||∇α L||
    Returns token importance vector for the [CLS] token.

    This is more faithful than raw attention because it weights
    each attention coefficient by the model's sensitivity to it.
    Tokens with high attention but zero gradient → low attribution.
    """
    model.zero_grad()
    outputs = model(**inputs, output_attentions=True)
    logit = outputs.logits[0, class_idx]

    # Enable gradient computation on attention weights
    for layer_attn in outputs.attentions:
        layer_attn.retain_grad()

    logit.backward(retain_graph=True)

    # Compute α · |∂y/∂α| for last 4 layers, average across heads
    grad_attns = []
    for layer_attn in outputs.attentions[-4:]:
        if layer_attn.grad is not None:
            # (batch, heads, seq, seq)
            grad_attn = layer_attn * layer_attn.grad.abs()
            grad_attn = grad_attn[0].mean(dim=0)  # avg across heads: (seq, seq)
            grad_attns.append(grad_attn.detach().numpy())

    if not grad_attns:
        return np.zeros(inputs["input_ids"].shape[-1])

    combined = np.stack(grad_attns).sum(axis=0)
    return combined[0]  # [CLS] row attribution

# ─── 5. BERTVIZ INTEGRATION ───────────────────────────────────────────────────

# Install BERTViz for interactive visualization:
# pip install bertviz
#
# from bertviz import head_view, model_view, neuron_view
#
# head_view(attentions, tokens)    # one head at a time
# model_view(attentions, tokens)   # all heads in a grid
# neuron_view(model, tokenizer, text, layer=10, head=0)  # Q·K dot product

# ─── 6. LAYER-WISE ATTENTION ENTROPY ─────────────────────────────────────────

def attention_entropy(attentions):
    """
    Compute entropy of attention distributions per layer per head.
    High entropy = diffuse attention (not looking at specific tokens)
    Low entropy = focused attention (strong preference for few tokens)
    """
    for layer_idx, layer_attn in enumerate(attentions):
        # (batch, heads, seq, seq)
        attn_np = layer_attn[0].detach().numpy()  # (heads, seq, seq)
        # Entropy of each row: H = -sum(p log p)
        eps = 1e-10
        entropy = -(attn_np * np.log(attn_np + eps)).sum(axis=-1)  # (heads, seq)
        mean_entropy = entropy.mean()
        print(f"  Layer {layer_idx+1:2d}: mean entropy = {mean_entropy:.3f}")

print("\nAttention entropy by layer (higher = more diffuse):")
attention_entropy(attentions)

# ─── 7. COMPARING METHODS ─────────────────────────────────────────────────────

print("\n" + "="*60)
print("COMPARISON: Raw Attention vs Rollout")
print("="*60)
print(f"{'Token':<20} {'Raw Attn':>12} {'Rollout':>12}")
print("-"*44)
for token, raw, roll in zip(tokens, cls_attention, cls_rollout):
    print(f"{token:<20} {raw:12.4f} {roll:12.4f}")

---

Captum: Production-Grade Attention Attribution

Facebook's Captum library provides validated, production-ready implementations of attention-based and gradient-based attribution methods. For transformers, LayerIntegratedGradients applied to attention outputs is the recommended approach for faithful attribution.

# pip install captum transformers

from captum.attr import LayerIntegratedGradients, LayerAttribution
from transformers import AutoTokenizer, AutoModelForSequenceClassification
import torch

model_name = "textattack/bert-base-uncased-SST-2"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSequenceClassification.from_pretrained(model_name)
model.eval()

def predict_proba(input_ids, attention_mask, token_type_ids):
    """Wrapper function for Captum."""
    logits = model(
        input_ids=input_ids,
        attention_mask=attention_mask,
        token_type_ids=token_type_ids,
    ).logits
    return torch.softmax(logits, dim=-1)

text = "The patient had an excellent recovery and was discharged."
inputs = tokenizer(text, return_tensors="pt")
input_ids = inputs["input_ids"]
attention_mask = inputs["attention_mask"]
token_type_ids = inputs.get("token_type_ids", torch.zeros_like(input_ids))
tokens = tokenizer.convert_ids_to_tokens(input_ids[0])

# ─── LAYER INTEGRATED GRADIENTS ON ATTENTION OUTPUT ──────────────────────────

# Target: last attention layer's output
target_layer = model.bert.encoder.layer[-1].attention.self

lig = LayerIntegratedGradients(predict_proba, target_layer)

# Baseline: [PAD] token everywhere
baseline_ids = torch.zeros_like(input_ids)

attributions, delta = lig.attribute(
    inputs=(input_ids, attention_mask, token_type_ids),
    baselines=(baseline_ids, attention_mask, token_type_ids),
    target=1,          # class 1 = positive sentiment
    n_steps=50,
    return_convergence_delta=True,
)

print(f"Convergence delta (completeness error): {delta.item():.6f}")
print(f"(Should be near 0 for faithful attribution)")

# Aggregate attributions across attention head dimensions
# attributions shape: (1, seq_len, heads * head_dim)
attr_summed = attributions.sum(dim=-1).squeeze(0)  # (seq_len,)
attr_norm = attr_summed / attr_summed.abs().max()   # normalize to [-1, 1]

print("\nLayer Integrated Gradients - attention output attributions:")
for token, attr in zip(tokens, attr_norm.detach().numpy()):
    direction = "+" if attr > 0 else "-"
    bar = "█" * int(abs(attr) * 30)
    print(f"  {token:20s} {direction}{abs(attr):.4f} {bar}")

# ─── GUIDED INTEGRATED GRADIENTS ON TOKEN EMBEDDINGS ─────────────────────────
# For the most faithful token-level attribution, apply IG to embeddings

from captum.attr import IntegratedGradients

def forward_with_embeddings(embeddings, attention_mask, token_type_ids):
    """Pass pre-computed embeddings through the model (skipping embedding layer)."""
    # This requires patching - simplified version shown
    outputs = model(
        inputs_embeds=embeddings,
        attention_mask=attention_mask,
        token_type_ids=token_type_ids,
    )
    return torch.softmax(outputs.logits, dim=-1)

embedding_layer = model.bert.embeddings

# Get actual embeddings
with torch.no_grad():
    embeddings = embedding_layer(input_ids, token_type_ids=token_type_ids)

baseline_embeddings = torch.zeros_like(embeddings)

ig = IntegratedGradients(forward_with_embeddings)
ig_attrs, ig_delta = ig.attribute(
    inputs=embeddings,
    baselines=baseline_embeddings,
    additional_forward_args=(attention_mask, token_type_ids),
    target=1,
    n_steps=100,
    return_convergence_delta=True,
)

# Sum across embedding dimensions → per-token attribution
ig_per_token = ig_attrs.sum(dim=-1).squeeze(0).detach().numpy()
ig_norm = ig_per_token / (abs(ig_per_token).max() + 1e-10)

print("\nIntegrated Gradients on embeddings (most faithful):")
for token, attr in zip(tokens, ig_norm):
    direction = "+" if attr > 0 else "-"
    bar = "█" * int(abs(attr) * 30)
    print(f"  {token:20s} {direction}{abs(attr):.4f} {bar}")

print(f"\nIG Convergence delta: {ig_delta.item():.6f}")
print(f"Sum of attributions: {ig_per_token.sum():.4f}")
print(f"(Should match f(x) - f(x') for completeness)")

---

TransformerLens: Mechanistic Interpretability Toolkit

TransformerLens (Nanda, 2022) is the primary toolkit for mechanistic interpretability research. It allows you to:

• Hook any intermediate computation (attention patterns, activations, residual stream)
• Perform activation patching (ablate specific heads to measure their causal effect)
• Identify circuits by systematically zeroing out heads and measuring performance change

# pip install transformer-lens

# from transformer_lens import HookedTransformer
# import transformer_lens.utils as utils

# model = HookedTransformer.from_pretrained("gpt2")

# ─── ACTIVATION PATCHING ─────────────────────────────────────────────────────
# Find which heads are causally responsible for a prediction by patching
# activations from a corrupted input into a clean forward pass

# tokens = model.to_tokens("The Eiffel Tower is located in")
# corrupted_tokens = model.to_tokens("The Colosseum is located in")

# # Run corrupted forward pass, cache all activations
# _, corrupted_cache = model.run_with_cache(corrupted_tokens)

# # For each attention head, patch its output from corrupted into clean run
# # and measure how much the prediction changes (logit difference)
# def patch_head(head_layer, head_idx, corrupted_cache):
#     def hook_fn(value, hook):
#         value[:, :, head_idx, :] = corrupted_cache[hook.name][:, :, head_idx, :]
#         return value
#     return hook_fn

# # Systematic patching reveals which heads are critical
# # Heads with high logit difference drop = causally important

The mechanistic interpretability workflow is:

1. Identify a behavior: find inputs where the model does something interesting (completes a rhyme, solves an analogy, applies a grammatical rule)
2. Ablate heads: systematically zero out or patch attention heads and measure the effect on the behavior
3. Name the circuit: the minimal set of heads whose ablation disrupts the behavior is the circuit for that behavior
4. Verify: confirm the circuit generalizes across many examples of the same behavior

---

Sparse Attention and Longformer

Standard BERT attention is $O(n^2)$ in sequence length - infeasible for documents longer than ~512 tokens. Longformer (Beltagy et al., 2020) uses sparse attention: each token attends to a local window of $w$ tokens around it, plus selected "global" tokens (like the [CLS] token) that attend to everything.

The Longformer attention pattern has two components:

• Local attention: $O(n \cdot w)$ - each position attends to $w/2$ tokens on each side
• Global attention: $O(n \cdot g)$ - $g$ global tokens (typically [CLS] + task-specific tokens) attend to all positions

For a document of 4096 tokens with $w=512$ and $g=1$:

• Full attention: $4096^2 = 16.7M$ attention weights
• Longformer: $4096 \times 512 + 4096 = 2.1M$ - an 8× reduction

The explanation implications: local attention is more interpretable (each token only interacts with nearby tokens), while global attention concentrates all long-range information in the [CLS] token. Visualizing [CLS] attention in Longformer is even more meaningful than in BERT because [CLS] is explicitly designed as the global information aggregator.

---

ViT Attention Maps: Transformers for Vision

Vision Transformers (ViT, Dosovitskiy et al., 2020) apply the same multi-head attention mechanism to image patches. A 224×224 image divided into 16×16 patches gives 196 patches, each treated as a "token." The [CLS] token aggregates information for classification.

ViT attention visualization is both more and less interpretable than BERT attention:

• More: image patches are spatial - high attention to a patch clearly maps back to a region of the image
• Less: ViT attention still has the multi-head, multi-layer complexity; the same adversarial attention concerns apply

Rollout attention for ViTs (Abnar & Zuidema 2020) is particularly useful because the attention propagation across layers directly traces which image patches flow their information into the [CLS] token at the end:

# ViT Attention Rollout
# For ViT, the [CLS] token is at position 0
# Attention from [CLS] in the last layer tells us which patches are used

# from transformers import ViTForImageClassification, ViTFeatureExtractor
# model = ViTForImageClassification.from_pretrained("google/vit-base-patch16-224")
# ...
# Then apply the same attention_rollout() function defined earlier
# The resulting cls_rollout[1:] (excluding CLS row) = (196,) spatial patch attributions
# Reshape to (14, 14) and upsample to (224, 224) for visualization

def vit_attention_rollout(attentions):
    """
    ViT-specific rollout: returns (H, W) spatial attribution map.
    For patch size 16 and image 224x224: 14×14 patches.
    """
    seq_len = attentions[0].shape[-1]  # 197 = 1 [CLS] + 196 patches
    rollout = np.eye(seq_len)

    for layer_attn in attentions:
        avg_attn = layer_attn[0].mean(dim=0).numpy()
        avg_attn_with_residual = 0.5 * avg_attn + 0.5 * np.eye(seq_len)
        row_sums = avg_attn_with_residual.sum(axis=-1, keepdims=True)
        avg_attn_with_residual /= row_sums
        rollout = avg_attn_with_residual @ rollout

    # [CLS] row: how much each patch contributed to the [CLS] representation
    cls_row = rollout[0, 1:]  # (196,) - skip [CLS] position itself
    return cls_row.reshape(14, 14)  # spatial layout

# Compare with GradCAM for ViT:
# GradCAM uses the attention output from the last block as the "feature map"
# GradCAM tends to produce tighter, more class-discriminative localization
# Rollout tends to capture broader, more complete regions

---

Production Engineering Notes

Use attention visualization for debugging, not auditing. Attention patterns help you quickly identify whether a model has learned plausible associations. They are too unreliable for compliance documentation or explaining individual decisions to users.

For high-stakes explanations, use Integrated Gradients on token embeddings. This is slower (50+ forward passes) but satisfies completeness and sensitivity axioms. The captum library provides a production-ready IG implementation for HuggingFace models.

Specify which layer and which head. "Attention visualization" without specifying layer and head is not reproducible. Different heads encode different things; different layers represent different levels of abstraction. Always document: layer index (0-indexed or 1-indexed - clarify), head index, and aggregation method (mean, max, or specific head).

Attention rollout is a better default than raw attention. It is almost no additional compute and accounts for multi-layer information mixing. Use rollout instead of last-layer attention as your default visualization baseline.

For mechanistic interpretability research, use TransformerLens. The activation patching API makes it straightforward to identify which heads are causally responsible for a behavior, going beyond correlation-based attention visualization.

warning

Do not use attention weights in regulatory compliance documents or user-facing explanations for high-stakes decisions (lending, hiring, medical) without a qualified ML engineer signing off that faithfulness has been validated. Attention visualization that looks correct is not the same as attention visualization that is correct.

---

YouTube Resources

Resource	Creator	Focus
Attention in Transformers Visually Explained	3Blue1Brown	Intuition-first attention mechanics
BERTViz Demo and Walkthrough	Jesse Vig	BERTViz tool walkthrough, head specialization
Attention is not Explanation - Paper Review	Yannic Kilcher	Jain & Wallace paper deep dive
Mechanistic Interpretability in Transformers	Neel Nanda	Circuits, induction heads, TransformerLens
Probing Classifiers and What BERT Learns	Stanford NLP Group	Layer-wise probing for linguistic structure

---

Common Mistakes

:::danger Common Mistake 1: Reporting a single attention head as "the" model's attention Multi-head transformers have $H \times L$ attention matrices. Selecting one that looks interpretable and presenting it as representative is p-hacking for attention maps. Always report the aggregation method and check that your conclusions hold across multiple heads and layers. :::

:::danger Common Mistake 2: Treating attention weight = feature importance Attention weight $\alpha_{ij}$ measures how much of the value vector $V_j$ flows into position $i$'s representation. It is not a direct measure of how much token $j$ affects the output. The output depends on the value matrix, subsequent layers, and the residual connection. High attention does not equal high importance. The adversarial attention result proves this: you can find very different attention patterns that produce the same output. :::

:::warning Common Mistake 3: Visualizing attention for encoder-only models as if it were cross-attention In encoder-only models (BERT for classification), self-attention serves multiple purposes: some heads track position, some track syntax, some track semantics. Cross-attention in encoder-decoder models is more directly interpretable because it is the only information channel from encoder to decoder. :::

:::warning Common Mistake 4: Not normalizing for attention rollout Attention rollout without row normalization at each layer can produce matrices with row sums not equal to 1, leading to attributions that do not sum to a consistent total. Always renormalize rows after adding the residual identity matrix. :::

:::danger Common Mistake 5: Conflating plausibility and faithfulness A plausible explanation looks correct to a domain expert. A faithful explanation accurately reflects the model's computation. These are not the same thing. The clinical NLP case is the canonical example: highly plausible attention visualization, unfaithful to the actual computational pathway. Always validate with gradient-based methods before claiming an attention visualization is a faithful explanation. :::

---

Interview Q&A

Q1: Explain the Jain & Wallace (2019) result. What exactly did they show and why does it matter?

Jain & Wallace showed two things. First, that attention weights often do not correlate with gradient-based feature importance measures - they used the Wilcoxon rank-sum test and found that the correlation between attention weights and gradient attributions was not significantly different from zero for most tasks. Second, more dramatically, that you can construct an adversarial attention distribution that is very different from the learned one (measured by total variation distance) but produces nearly identical predictions. If $\|\tilde{\alpha} - \alpha\|_{TV} > \delta$ yet $|f_{\tilde{\alpha}}(x) - f_\alpha(x)| < \epsilon$, the original attention cannot be causally responsible for the prediction - a causal factor would change the output when you changed it. The implication is that attention visualization can be misleading. Many teams were using attention maps for model auditing, and those audits may have been certifying the wrong thing. The mechanism: high-dimensional value vectors for semantically similar tokens cluster together, meaning different attention distributions over those tokens produce nearly the same weighted sum.

Q2: What is attention rollout and how does it improve on raw attention visualization?

Attention rollout (Abnar & Zuidema, 2020) computes how information flows from input tokens to output positions across all layers of the transformer, not just one layer. The core insight is that transformers have residual connections, which means information can bypass any given attention layer. Rollout models this by mixing 50% of the identity matrix into each layer's attention matrix before chaining the matrices: $R^{(L)} = \hat{A}^{(L)} \cdot \hat{A}^{(L-1)} \cdots \hat{A}^{(1)}$ where $\hat{A}^{(l)} = 0.5 A^{(l)} + 0.5 I$. This is more principled than looking at a single layer's attention because it accounts for multi-layer information aggregation. The limitation is that it still does not account for value matrix projections and uses a fixed 50/50 split for the residual that is not theoretically justified. Use rollout as a better heuristic than single-layer attention, not as a provably faithful method.

Q3: What is gradient-weighted attention and why is it more faithful than raw attention?

Gradient-weighted attention computes $\tilde{\alpha}_i = \alpha_i \cdot \|\nabla_{\alpha_i}\mathcal{L}\|$ - the product of the attention weight and the magnitude of the gradient of the loss with respect to that attention weight. The gradient measures how sensitive the model's output is to changes in that attention weight. This directly addresses the adversarial attention problem: an attention weight that is high but has zero gradient means the model's output is insensitive to this weight - so the gradient zeros out the attribution. A token can only receive high attribution if both the attention weight is high (the token is attended to) and the model is sensitive to that attention weight (changing it would change the prediction). Wiegreffe & Pinter showed empirically that gradient × attention correlates more strongly with other feature attribution methods than raw attention alone. This is analogous to GradCAM: using the gradient to weight the feature map channels produces more faithful localization than raw class activation maps.

Q4: Explain the faithfulness vs plausibility distinction. Why does it matter for model auditing?

Faithfulness (Jacovi & Goldberg, 2020): the explanation accurately reflects the actual computational process that produced the prediction. A faithful explanation would change proportionally if the model's reasoning changed - you could modify the highlighted features and the prediction would shift accordingly. Plausibility: the explanation agrees with human intuition about what should be important. Plausible explanations look right to domain experts even if they do not reflect the model's actual computation. For model auditing, plausibility is insufficient - you need faithfulness. A plausible but unfaithful explanation can provide false assurance: the model appears to be using the right features, but is actually using something entirely different (a watermark, a spurious correlation, a demographic proxy). Audits based on plausible-but-unfaithful explanations certify the wrong thing and miss actual failure modes. Gradient-based methods (IG, GradCAM) are more likely to be faithful because they directly measure the computational sensitivity; attention visualization is more likely to be merely plausible.

Q5: What are induction heads and what do they reveal about transformer circuits?

Induction heads (Elhage et al., 2021) are a two-head circuit that implements in-context pattern matching. Head A in an early layer is a "previous token" head - for each position, it attends to the immediately preceding position. Head B in a later layer is the induction head - it attends to positions that are one step after wherever the current token appeared earlier in the context. Together, if the sequence contains [A][B]...[A], after [A] appears for the second time, head A puts [A]'s information at [B]'s position in the first occurrence, and head B attends to that position and predicts [B]. This is the mechanistic basis for few-shot learning from examples in the prompt. The implication for attention visualization: the "important" attention weights for in-context predictions are in head B's pattern, not in the final-layer CLS attention. Understanding which circuit is active for a given behavior determines which attention heads to look at. Circuits are identified through activation patching in TransformerLens - ablate a head and measure how much the target behavior degrades.

Q6: How would you explain a BERT classification decision to a regulator, and what would you not use?

For regulatory compliance, I would use Integrated Gradients on the input token embeddings, not attention visualization. IG satisfies three formal axioms: completeness ($\sum_i \text{IG}_i(x) = f(x) - f(x')$ - attributions sum to the prediction difference from baseline), sensitivity (tokens that affect the prediction get nonzero attribution), and implementation invariance (equivalent models give the same attributions). The procedure: choose a baseline (the [PAD] token embedding), compute 50–300 interpolated inputs between baseline and actual input, average the gradients at each interpolated point, and multiply by the input-baseline difference. The resulting per-token attribution scores can be presented as "how much did each token contribute to this classification decision" with a formal guarantee. I would document: (1) baseline choice and why, (2) number of integration steps, (3) numerical completeness error (must be below 5%), (4) model randomization sanity check to confirm maps reflect learned behavior. I would explicitly not use raw attention visualization or guided backpropagation as primary evidence - attention fails faithfulness criteria, and guided backpropagation fails model sanity checks.

---

Key Takeaways

Method	Faithfulness	Speed	Best Use Case
Raw attention (single layer)	Low	Fastest	Quick debugging, hypothesis generation
Raw attention ([CLS] last layer)	Low-Medium	Fastest	BERT classification exploration
Attention rollout	Medium	Fast	Better default for diagnostics
Gradient × Attention	Medium-High	Medium	Production debugging
Layer Integrated Gradients (Captum)	High	Medium-Slow	Production attribution
IG on token embeddings	Highest	Slow (50–300 passes)	Compliance, auditing
Probing classifiers	N/A (diagnostic)	Medium	Understanding what is encoded where
Circuit analysis (TransformerLens)	Causal	Slow	Research, mechanistic understanding

The rule of thumb: use attention visualization to generate hypotheses about model behavior, then validate those hypotheses with gradient-based methods before acting on them. Attention is a window into the model - but like all windows, what you see depends on the angle, and sometimes you are looking at a reflection.

---

Historical Context: How Attention Became the Dominant Explanation Tool

Attention mechanisms were introduced in neural machine translation by Bahdanau et al. (2015) - not primarily for interpretability, but to solve the encoder bottleneck problem. In seq2seq models without attention, the encoder must compress an entire source sentence into a single fixed-size vector. Attention allowed the decoder to selectively focus on different parts of the source sentence at each decoding step. A convenient side effect: the attention weights formed a natural alignment between source and target tokens, visually resembling word-by-word translation.

This alignment property made attention immediately compelling for interpretation. By 2018, with BERT's release, attention visualization was the default explanation method for transformer-based NLP. The BERTViz paper (Vig, 2019) formalized the tooling, and dozens of papers analyzed BERT's attention patterns in terms of linguistic phenomena.

The Jain & Wallace paper (2019) was received as controversial precisely because the field had invested heavily in attention-as-explanation. It challenged a widespread practice with evidence. The Wiegreffe & Pinter response (2019) was similarly influential because it clarified the more nuanced question: attention may not be a complete causal explanation, but it is not uninformative either.

The mechanistic interpretability movement (2021–present), led by Anthropic's interpretability team and researchers like Neel Nanda, has pushed beyond attention visualization to ask more fundamental questions: what computations are individual neurons performing? Which circuits implement which behaviors? This is a more ambitious and more rigorous program, but it requires significantly more effort than visualizing attention weights.

---

Integrated Gradients on Token Embeddings: The Full Implementation

For completeness, here is the full Integrated Gradients implementation on token embeddings - the method with the strongest formal guarantees for BERT-based text classification explanations.

import torch
import numpy as np
from transformers import AutoTokenizer, AutoModelForSequenceClassification

def integrated_gradients_text(
    model,
    tokenizer,
    text: str,
    target_class: int,
    n_steps: int = 100,
    baseline_type: str = "pad",
) -> dict:
    """
    Integrated Gradients on input token embeddings for text classification.

    IG_i(x) = (x_i - x'_i) * (1/m) * Σ_k ∂f(x' + k/m*(x-x')) / ∂x_i

    Satisfies:
    - Completeness: Σ IG_i = f(x) - f(x')
    - Sensitivity: important tokens get nonzero attribution

    Args:
        baseline_type: 'pad' (embed [PAD] tokens) or 'zero' (zero vectors)

    Returns:
        dict with 'tokens', 'attributions', 'completeness_error'
    """
    inputs = tokenizer(text, return_tensors="pt", truncation=True, max_length=128)
    tokens = tokenizer.convert_ids_to_tokens(inputs["input_ids"][0])

    # Get embedding layer
    embedding_layer = model.bert.embeddings.word_embeddings

    # Actual input embeddings
    with torch.no_grad():
        actual_embeddings = embedding_layer(inputs["input_ids"])  # (1, seq, 768)

    # Baseline embeddings
    if baseline_type == "pad":
        pad_id = tokenizer.pad_token_id
        baseline_ids = torch.full_like(inputs["input_ids"], pad_id)
        with torch.no_grad():
            baseline_embeddings = embedding_layer(baseline_ids)
    else:  # zero baseline
        baseline_embeddings = torch.zeros_like(actual_embeddings)

    # Compute f(x) and f(x')
    def forward_from_embeddings(embeddings):
        outputs = model(
            inputs_embeds=embeddings,
            attention_mask=inputs["attention_mask"],
        )
        return torch.softmax(outputs.logits, dim=-1)[0, target_class]

    with torch.no_grad():
        f_x = forward_from_embeddings(actual_embeddings).item()
        f_baseline = forward_from_embeddings(baseline_embeddings).item()

    # Accumulate gradients along the interpolation path
    accumulated_grads = torch.zeros_like(actual_embeddings)

    for step in range(1, n_steps + 1):
        alpha = step / n_steps
        interp = (baseline_embeddings + alpha * (actual_embeddings - baseline_embeddings))
        interp.requires_grad_(True)

        score = forward_from_embeddings(interp)
        score.backward()
        accumulated_grads += interp.grad.detach()

    # IG = (x - x') * avg_grad, summed across embedding dimension
    avg_grads = accumulated_grads / n_steps
    ig = (actual_embeddings - baseline_embeddings) * avg_grads  # (1, seq, 768)
    ig_per_token = ig[0].sum(dim=-1).detach().numpy()  # (seq,)

    # Verify completeness
    total_attr = ig_per_token.sum()
    completeness_error = abs(total_attr - (f_x - f_baseline))
    relative_error = completeness_error / (abs(f_x - f_baseline) + 1e-8)

    print(f"f(x) = {f_x:.4f}, f(x') = {f_baseline:.4f}")
    print(f"Sum of attributions = {total_attr:.4f}")
    print(f"Completeness error = {relative_error:.4%}")
    if relative_error < 0.05:
        print("PASSED: IG satisfies completeness axiom")
    else:
        print(f"WARNING: {relative_error:.1%} error - increase n_steps")

    return {
        "tokens": tokens,
        "attributions": ig_per_token,
        "f_x": f_x,
        "f_baseline": f_baseline,
        "completeness_error": relative_error,
    }

# Usage example
# model = AutoModelForSequenceClassification.from_pretrained("textattack/bert-base-uncased-SST-2")
# tokenizer = AutoTokenizer.from_pretrained("textattack/bert-base-uncased-SST-2")
# result = integrated_gradients_text(model, tokenizer, "The movie was fantastic", target_class=1)
# for token, attr in zip(result["tokens"], result["attributions"]):
#     direction = "+" if attr > 0 else "-"
#     bar = "█" * int(abs(attr) * 40)
#     print(f"  {token:20s} {direction}{abs(attr):.4f} {bar}")

The completeness check is not optional - it is the only way to verify that your IG implementation is correct. A common bug: not multiplying by $(x_i - x_i')$ at the end, which violates completeness. Another common bug: using the wrong baseline, which changes what the attribution is "relative to." Document baseline choice explicitly in any compliance report.

:::tip 🎮 Interactive Playground

Visualize this concept: Try the Transformer Attention demo on the EngineersOfAI Playground - no code required.

:::