Attention Analysis & Probing

"The question is not whether neural networks are black boxes. The question is whether we have the right flashlights."

Probe, Flashlight Wielding AI Agent

1. Attention Visualization

In Chapter 04, we learned that every transformer layer computes attention weights: a matrix that specifies how much each token attends to every other token (recall the query-key-value mechanism from Section 04.2). Visualizing these weights provides an immediate, intuitive window into the model's computation. However, interpreting attention requires care: attention weights do not directly indicate which tokens are "important" for the prediction (as we discuss in Section 18.4).

The following snippet demonstrates how to extract attention weights from a GPT-2 model and visualize a specific attention head as a heatmap.

# Extracting and visualizing attention patterns
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
import matplotlib.pyplot as plt
import numpy as np

model_name = "gpt2"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(
 model_name, output_attentions=True
)

text = "The cat sat on the mat because it was tired"
inputs = tokenizer(text, return_tensors="pt")

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

# outputs.attentions is a tuple of (num_layers,) tensors
# Each tensor has shape (batch, num_heads, seq_len, seq_len)
attentions = outputs.attentions
print(f"Number of layers: {len(attentions)}")
print(f"Attention shape per layer: {attentions[0].shape}")

tokens = tokenizer.convert_ids_to_tokens(inputs["input_ids"][0])

def plot_attention_head(attention_matrix, tokens, layer, head):
 """Plot a single attention head as a heatmap."""
 fig, ax = plt.subplots(figsize=(8, 6))
 attn = attention_matrix[0, head].numpy() # (seq_len, seq_len)
 im = ax.imshow(attn, cmap="Blues", vmin=0, vmax=1)
 ax.set_xticks(range(len(tokens)))
 ax.set_yticks(range(len(tokens)))
 ax.set_xticklabels(tokens, rotation=45, ha="right", fontsize=8)
 ax.set_yticklabels(tokens, fontsize=8)
 ax.set_xlabel("Key (attending to)")
 ax.set_ylabel("Query (attending from)")
 ax.set_title(f"Layer {layer}, Head {head}")
 plt.colorbar(im)
 plt.tight_layout()
 return fig

# Visualize a specific head
fig = plot_attention_head(attentions[5], tokens, layer=5, head=1)
plt.savefig("attention_head_L5_H1.png", dpi=150)

from bertviz import head_view
from transformers import AutoModel, AutoTokenizer

model = AutoModel.from_pretrained("gpt2", output_attentions=True)
tokenizer = AutoTokenizer.from_pretrained("gpt2")
inputs = tokenizer("The cat sat on the mat", return_tensors="pt")
outputs = model(**inputs)

head_view(outputs.attentions, tokenizer.convert_ids_to_tokens(inputs["input_ids"][0]))
# Opens an interactive HTML widget in Jupyter

# Production equivalent using scikit-learn
from sklearn.linear_model import LogisticRegression
from sklearn.model_selection import cross_val_score

# X = extracted hidden states (n_tokens, hidden_dim), y = labels
probe = LogisticRegression(max_iter=1000)
scores = cross_val_score(probe, X, y, cv=5)
print(f"Probe accuracy: {scores.mean():.3f} +/- {scores.std():.3f}")

1.1 Common Attention Patterns

Research has identified several recurring attention patterns across transformer models. These patterns appear consistently regardless of model size, training data, or architecture variant, suggesting they represent fundamental computational primitives. Figure 18.1.1 catalogs the most common types.

Figure 18.1.2: Common attention head types observed across transformer models. Each pattern represents a distinct computational role.

The following table summarizes these common attention head types, noting where each typically appears and what computational role it serves.

1.2 The Attention Debate: 2024-2025 Resolution

For years, the interpretability community was stuck in a binary argument: "attention is explanation" versus "attention is not explanation." Jain and Wallace (2019) argued the latter; Wiegreffe and Pinter (2019) pushed back with evidence that attention can be meaningful. By 2024-2025, a more nuanced consensus has emerged that transcends this binary framing. The short answer: attention weights alone are not explanations, but attention patterns combined with complementary methods (probing, ablation, circuit analysis) are genuinely informative.

Why does this resolution matter? If you dismiss attention entirely, you lose the cheapest and fastest window into model behavior. If you trust attention uncritically, you build explanations on shaky ground. The modern consensus gives practitioners a principled middle path: use attention as a hypothesis generator, then validate with causal methods.

In practice, this consensus translates to a three-step workflow for using attention in interpretability. First, visualize attention to generate hypotheses ("Head L5H3 seems to track coreference"). Second, validate the hypothesis using ablation or activation patching ("zeroing out L5H3 breaks coreference resolution but not other tasks"). Third, connect the finding to the broader circuit by examining which other heads and MLP layers interact with L5H3 using tools like TransformerLens (see Section 18.2 for the tooling).

# Attention + ablation validation workflow
import transformer_lens
from transformer_lens import HookedTransformer
import torch

model = HookedTransformer.from_pretrained("gpt2-small")

prompt = "The nurse told the doctor that she was ready"
clean_logits, clean_cache = model.run_with_cache(prompt)

# Step 1: Identify candidate head via attention pattern
# (e.g., head L5H3 attends from "she" to "nurse")

# Step 2: Validate causally by ablating the head
def ablate_head(activation, hook):
 """Zero out a specific attention head's output."""
 activation[:, :, 3, :] = 0 # head index 3
 return activation

ablated_logits = model.run_with_hooks(
 prompt,
 fwd_hooks=[("blocks.5.attn.hook_z", ablate_head)]
)

# Step 3: Compare predictions with and without the head
she_token = model.to_single_token(" she")
nurse_token = model.to_single_token(" nurse")

clean_prob = torch.softmax(clean_logits[0, -1], dim=-1)
ablated_prob = torch.softmax(ablated_logits[0, -1], dim=-1)

print(f"With L5H3: P(nurse-related) = {clean_prob[nurse_token]:.4f}")
print(f"Without L5H3: P(nurse-related) = {ablated_prob[nurse_token]:.4f}")
# A significant drop confirms the head is causally important

2. Probing Classifiers

Probing classifiers provide a more rigorous way to test what information is encoded in a model's hidden representations. The idea is simple: extract hidden states from a specific layer, freeze them, and train a lightweight classifier on top to predict some property of interest (part of speech, syntactic dependency, semantic role, entity type). If the classifier succeeds, the property is encoded in the hidden states.

Below, we implement both a linear and nonlinear probing classifier for testing what linguistic information is encoded in each transformer layer. Code Fragment 18.1.3 shows this approach in practice.

# Probing classifier for linguistic properties
import torch
import torch.nn as nn
from torch.utils.data import DataLoader, TensorDataset
from transformers import AutoModel, AutoTokenizer
from sklearn.metrics import accuracy_score
import numpy as np

class LinearProbe(nn.Module):
 """A simple linear probe for testing representation content."""
 def __init__(self, hidden_dim, num_classes):
 super().__init__()
 self.classifier = nn.Linear(hidden_dim, num_classes)

 def forward(self, hidden_states):
 return self.classifier(hidden_states)

class MLPProbe(nn.Module):
 """A nonlinear probe with one hidden layer."""
 def __init__(self, hidden_dim, num_classes, probe_dim=256):
 super().__init__()
 self.net = nn.Sequential(
 nn.Linear(hidden_dim, probe_dim),
 nn.ReLU(),
 nn.Dropout(0.1),
 nn.Linear(probe_dim, num_classes),
 )

 def forward(self, hidden_states):
 return self.net(hidden_states)

def extract_hidden_states(model, tokenizer, texts, layer_idx):
 """Extract hidden states from a specific layer."""
 model.eval()
 all_hidden = []

 for text in texts:
 inputs = tokenizer(text, return_tensors="pt", truncation=True)
 with torch.no_grad():
 outputs = model(**inputs, output_hidden_states=True)

 # Get hidden states from the specified layer
 # Shape: (1, seq_len, hidden_dim)
 hidden = outputs.hidden_states[layer_idx]
 all_hidden.append(hidden.squeeze(0))

 return all_hidden

def train_probe(
 hidden_states, # list of (seq_len, hidden_dim) tensors
 labels, # list of (seq_len,) label tensors
 num_classes,
 probe_type="linear",
 epochs=10,
 lr=1e-3,
):
 """Train a probing classifier on frozen hidden states."""
 # Flatten all tokens into a single dataset
 X = torch.cat(hidden_states, dim=0)
 y = torch.cat(labels, dim=0)

 hidden_dim = X.shape[1]
 if probe_type == "linear":
 probe = LinearProbe(hidden_dim, num_classes)
 else:
 probe = MLPProbe(hidden_dim, num_classes)

 optimizer = torch.optim.Adam(probe.parameters(), lr=lr)
 criterion = nn.CrossEntropyLoss()

 dataset = TensorDataset(X, y)
 loader = DataLoader(dataset, batch_size=256, shuffle=True)

 for epoch in range(epochs):
 total_loss = 0
 for batch_x, batch_y in loader:
 logits = probe(batch_x)
 loss = criterion(logits, batch_y)
 optimizer.zero_grad()
 loss.backward()
 optimizer.step()
 total_loss += loss.item()

 if (epoch + 1) % 5 == 0:
 print(f"Epoch {epoch+1}: loss={total_loss/len(loader):.4f}")

 return probe

# Example: Probe for part-of-speech tags across layers
model = AutoModel.from_pretrained("gpt2")
tokenizer = AutoTokenizer.from_pretrained("gpt2")

# For each layer, train a probe and measure accuracy
layer_accuracies = {}
for layer in range(model.config.num_hidden_layers + 1):
 hidden = extract_hidden_states(model, tokenizer, train_texts, layer)
 probe = train_probe(hidden, pos_labels, num_pos_tags, "linear")
 # evaluate_probe: run probe on validation hidden states,
 # compare predictions to val_labels, return accuracy
 acc = evaluate_probe(probe, val_hidden, val_labels)
 layer_accuracies[layer] = acc
 print(f"Layer {layer}: POS accuracy = {acc:.3f}")

# Control task for validating probe results
import random

def run_probing_experiment_with_control(
 model, tokenizer, texts, labels, num_classes, layer_idx
):
 """Run probing with control task to measure selectivity."""

 hidden = extract_hidden_states(model, tokenizer, texts, layer_idx)

 # Real task probe
 # val_hidden and val_labels are computed from a held-out set
 # using extract_hidden_states() and the same label pipeline
 real_probe = train_probe(hidden, labels, num_classes, "linear")
 real_acc = evaluate_probe(real_probe, val_hidden, val_labels)

 # Control task: shuffle labels to create random assignment
 control_labels = [
 torch.randint(0, num_classes, label.shape)
 for label in labels
 ]
 control_probe = train_probe(hidden, control_labels, num_classes, "linear")
 # val_control_labels: same random mapping applied to validation set
 control_acc = evaluate_probe(control_probe, val_hidden, val_control_labels)

 selectivity = real_acc - control_acc

 return {
 "real_accuracy": real_acc,
 "control_accuracy": control_acc,
 "selectivity": selectivity,
 "meaningful": selectivity > 0.1, # threshold
 }

2.1 Control Tasks

A common criticism of probing is that a powerful enough probe might learn the task itself rather than reflecting what the model has learned. Control tasks (Hewitt and Liang, 2019) address this by training the same probe on a random labeling of the data. If the probe achieves high accuracy on both the real task and the control task, the probe is too powerful, and the result is not meaningful.

# Using the tuned lens package
# pip install tuned-lens
from tuned_lens import TunedLens
from transformers import AutoModelForCausalLM, AutoTokenizer

model = AutoModelForCausalLM.from_pretrained("gpt2")
tokenizer = AutoTokenizer.from_pretrained("gpt2")

# Load pre-trained tuned lens for GPT-2
tuned = TunedLens.from_model_and_pretrained(model)

text = "The capital of France is"
inputs = tokenizer(text, return_tensors="pt")

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

hidden_states = outputs.hidden_states

# Apply tuned lens at each layer
for layer_idx in range(len(hidden_states) - 1):
 # Tuned lens applies a learned affine transform
 logits = tuned(hidden_states[layer_idx], layer_idx)
 probs = F.softmax(logits[0, -1], dim=-1)
 top_token = tokenizer.decode(probs.argmax())
 top_prob = probs.max().item()
 print(f"Layer {layer_idx:2d}: '{top_token}' (p={top_prob:.3f})")

3. The Logit Lens

The logit lens (nostalgebraist, 2020) is a technique for inspecting what a transformer "thinks" at each intermediate layer. The idea is to take the hidden state at any layer and project it through the model's final unembedding matrix (the same matrix used to produce the output logits). This reveals the model's current "best guess" for the next token at each point in the computation. Figure 18.1.3 illustrates how predictions sharpen from noisy early layers to confident later layers. Code Fragment 18.1.5 shows this approach in practice.

Figure 18.1.3: The logit lens projects hidden states from each layer through the unembedding matrix. Earlier layers produce noisy predictions that sharpen as computation progresses.

The following implementation projects each layer's hidden state through the unembedding matrix, revealing the model's intermediate "best guess" at each stage of computation.

Code Fragment 18.1.5 applies the logit lens.

# Logit Lens Implementation
import torch
import torch.nn.functional as F
from transformers import AutoModelForCausalLM, AutoTokenizer

def logit_lens(model, tokenizer, text, top_k=5):
 """
 Apply the logit lens to see what the model predicts
 at each intermediate layer.
 """
 inputs = tokenizer(text, return_tensors="pt")
 with torch.no_grad():
 outputs = model(**inputs, output_hidden_states=True)

 hidden_states = outputs.hidden_states # (num_layers+1,) tuple
 # Get the unembedding matrix
 unembed = model.lm_head.weight # (vocab_size, hidden_dim)

 tokens = tokenizer.convert_ids_to_tokens(inputs["input_ids"][0])
 last_pos = len(tokens) - 1 # predict next token after last

 print(f"Input: {text}")
 print(f"Predicting token after: '{tokens[last_pos]}'")
 print("-" * 50)

 for layer_idx, hidden in enumerate(hidden_states):
 # Project hidden state through unembedding
 # hidden shape: (1, seq_len, hidden_dim)
 logits = hidden[0, last_pos] @ unembed.T # (vocab_size,)
 probs = F.softmax(logits, dim=-1)

 top_probs, top_ids = probs.topk(top_k)
 top_tokens = [tokenizer.decode(tid) for tid in top_ids]

 top_str = ", ".join(
 f"'{t}' ({p:.3f})" for t, p in zip(top_tokens, top_probs)
 )
 print(f"Layer {layer_idx:2d}: {top_str}")

# Run logit lens on GPT-2
model = AutoModelForCausalLM.from_pretrained("gpt2")
tokenizer = AutoTokenizer.from_pretrained("gpt2")
logit_lens(model, tokenizer, "The Eiffel Tower is in")

3.1 The Tuned Lens

The tuned lens (Belrose et al., 2023) improves on the logit lens by training a learned affine transformation for each layer. The raw logit lens assumes that intermediate representations are approximately in the same space as the final layer, which is only roughly true. The tuned lens trains a small per-layer probe to account for the differences in representation spaces across layers.

pip install transformers torch matplotlib seaborn numpy scikit-learn

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

model_name = "gpt2"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(
 model_name, output_attentions=True)
model.eval()

text = "The cat sat on the mat because it was comfortable"
inputs = tokenizer(text, return_tensors="pt")
tokens = tokenizer.convert_ids_to_tokens(inputs['input_ids'][0])

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

# outputs.attentions: tuple of (num_layers,) tensors
attentions = outputs.attentions
print(f"Layers: {len(attentions)}")
print(f"Shape per layer: {attentions[0].shape}")
print(f"Tokens: {tokens}")

# TODO: Stack into (layers, heads, seq, seq) numpy array
all_attn = torch.stack(attentions).squeeze(1).numpy()
print(f"All attention shape: {all_attn.shape}")

import matplotlib.pyplot as plt
import seaborn as sns

def plot_attention_head(attn_matrix, tokens, layer, head, ax):
 sns.heatmap(attn_matrix, xticklabels=tokens, yticklabels=tokens,
 cmap='Blues', ax=ax, vmin=0, vmax=1)
 ax.set_title(f'L{layer}H{head}', fontsize=9)
 ax.tick_params(labelsize=6)

# TODO: Create a 3x4 grid for all 12 heads in layer 0
fig, axes = plt.subplots(3, 4, figsize=(20, 15))
for h in range(12):
 row, col = h // 4, h % 4
 plot_attention_head(all_attn[0, h], tokens, 0, h, axes[row, col])

plt.tight_layout()
plt.savefig('attention_heads_layer0.png', dpi=100)
plt.show()

# Find token positions
it_pos = tokens.index("Ġit") if "Ġit" in tokens else tokens.index("it")
cat_pos = tokens.index("Ġcat") if "Ġcat" in tokens else tokens.index("cat")
print(f"'it' at position {it_pos}, 'cat' at position {cat_pos}")

# Find heads that show strong it-to-cat attention
print("\nAttention from 'it' to 'cat' by layer:")
for layer in range(len(attentions)):
 for head in range(12):
 attn_score = all_attn[layer, head, it_pos, cat_pos]
 if attn_score > 0.1:
 print(f" Layer {layer}, Head {head}: {attn_score:.3f}")

from sklearn.linear_model import LogisticRegression
from sklearn.model_selection import cross_val_score

probe_data = [
 ("The dog chased the cat", ["DET","NOUN","VERB","DET","NOUN"]),
 ("A bird flew over trees", ["DET","NOUN","VERB","ADP","NOUN"]),
 ("She quickly ran home today", ["PRON","ADV","VERB","NOUN","ADV"]),
 ("Big waves crashed on shore", ["ADJ","NOUN","VERB","ADP","NOUN"]),
 ("They slowly walked to school", ["PRON","ADV","VERB","ADP","NOUN"]),
 ("My friend reads many books", ["DET","NOUN","VERB","DET","NOUN"]),
 ("The tall man opened doors", ["DET","ADJ","NOUN","VERB","NOUN"]),
 ("We often eat fresh food", ["PRON","ADV","VERB","ADJ","NOUN"]),
]

def extract_hidden(model, tokenizer, text, layer):
 inputs = tokenizer(text, return_tensors="pt")
 with torch.no_grad():
 out = model(**inputs, output_hidden_states=True)
 return out.hidden_states[layer][0].numpy()

# TODO: For each layer, collect hidden states and POS labels,
# then train a LogisticRegression probe and report accuracy.
for layer in [0, 3, 6, 9, 11]:
 X, y = [], []
 for text, pos_tags in probe_data:
 hidden = extract_hidden(model, tokenizer, text, layer)
 toks = tokenizer.convert_ids_to_tokens(
 tokenizer(text)['input_ids'])
 for i, tag in enumerate(pos_tags):
 if i < hidden.shape[0]:
 X.append(hidden[i])
 y.append(tag)

 X = np.array(X)
 clf = LogisticRegression(max_iter=1000, random_state=42)
 scores = cross_val_score(clf, X, y, cv=3, scoring='accuracy')
 print(f"Layer {layer}: POS probe accuracy = "
 f"{scores.mean():.3f} (+/- {scores.std():.3f})")

import torch, numpy as np, matplotlib.pyplot as plt, seaborn as sns
from transformers import AutoModelForCausalLM, AutoTokenizer
from sklearn.linear_model import LogisticRegression
from sklearn.model_selection import cross_val_score

tokenizer = AutoTokenizer.from_pretrained("gpt2")
model = AutoModelForCausalLM.from_pretrained("gpt2", output_attentions=True)
model.eval()

text = "The cat sat on the mat because it was comfortable"
inputs = tokenizer(text, return_tensors="pt")
tokens = tokenizer.convert_ids_to_tokens(inputs['input_ids'][0])
with torch.no_grad(): outputs = model(**inputs)
all_attn = torch.stack(outputs.attentions).squeeze(1).numpy()

# Heatmaps
fig, axes = plt.subplots(3, 4, figsize=(20, 15))
for h in range(12):
 sns.heatmap(all_attn[0,h], xticklabels=tokens, yticklabels=tokens,
 cmap='Blues', ax=axes[h//4, h%4], vmin=0, vmax=1)
 axes[h//4, h%4].set_title(f'L0H{h}', fontsize=9)
 axes[h//4, h%4].tick_params(labelsize=6)
plt.tight_layout(); plt.savefig('attention_heads.png'); plt.show()

# Coreference
it_pos = tokens.index("Ġit") if "Ġit" in tokens else tokens.index("it")
cat_pos = tokens.index("Ġcat") if "Ġcat" in tokens else tokens.index("cat")
for l in range(12):
 for h in range(12):
 a = all_attn[l, h, it_pos, cat_pos]
 if a > 0.1: print(f"L{l}H{h}: it->cat = {a:.3f}")

# Probing
probe_data = [
 ("The dog chased the cat", ["DET","NOUN","VERB","DET","NOUN"]),
 ("A bird flew over trees", ["DET","NOUN","VERB","ADP","NOUN"]),
 ("She quickly ran home today", ["PRON","ADV","VERB","NOUN","ADV"]),
 ("Big waves crashed on shore", ["ADJ","NOUN","VERB","ADP","NOUN"]),
 ("They slowly walked to school", ["PRON","ADV","VERB","ADP","NOUN"]),
 ("My friend reads many books", ["DET","NOUN","VERB","DET","NOUN"]),
 ("The tall man opened doors", ["DET","ADJ","NOUN","VERB","NOUN"]),
 ("We often eat fresh food", ["PRON","ADV","VERB","ADJ","NOUN"]),
]

for layer in [0, 3, 6, 9, 11]:
 X, y = [], []
 for txt, tags in probe_data:
 inp = tokenizer(txt, return_tensors="pt")
 with torch.no_grad():
 out = model(**inp, output_hidden_states=True)
 h = out.hidden_states[layer][0].numpy()
 for i, tag in enumerate(tags):
 if i < h.shape[0]: X.append(h[i]); y.append(tag)
 scores = cross_val_score(LogisticRegression(max_iter=1000), np.array(X), y, cv=3)
 print(f"Layer {layer}: {scores.mean():.3f} +/- {scores.std():.3f}")

Exercises