Lesson 87: Transformer Interpretability

Discussion prompt

A has shape (T, T) per head; A[i, j] is the weight query i places on key j. Each row sums to 1 (Lesson 17 softmax). This matrix is what we visualize as a heatmap.

That is stated as though it always holds. Do one of two things: produce a case where it fails, or say precisely what rules such a case out. "It just does" is not on the menu.

Hint: Hunt at the extremes first — zero, one, negative, empty, equal. If every extreme survives, the reason they survive is the proof.

Predict first

The table runs: The | 0.7695 | 0.1041 | 0.0632 | 0.0632 · cat | 0.0410 | 0.8228 | 0.1114 | 0.0248 · sat | 0.0209 | 0.2541 | 0.6907 | 0.0344

In Extracting attention weights in PyTorch, given the rows so far: what is the next one — the row where query \ key is mat?

Correct: mat | 0.0273 | 0.0551 | 0.2468 | 0.6708

Why: The relationship between the columns, not the individual numbers, is what generates the next row. Scale by 1/√d_k to prevent softmax saturation in high-dimensional key spaces (Lesson 82).

Comparison matrix

From Extracting attention weights in PyTorch: refill the The column from what you know. The rest of the table is as it appeared.

Comparison matrix

From Head specialization patterns: every row here is a choice with a cost. Fill the what it attends to column, then say which row you would actually pick and what you give up for it.

Fill in the blanks

From Attention rollout: 2-layer toy example — one line has had its right-hand side removed. Put it back.

import torch, torch.nn.functional as F
torch.manual_seed(7)

def random_attn(T):
return F.softmax(torch.rand(T, T), dim=-1)

def rollout(A_list):
result = torch.eye(A_list[0].shape[0])
for A in A_list:
A_res = (A + torch.eye(A.shape[0])) / 2 # residual mix
A_res = A_res / A_res.sum(-1, keepdim=True) # re-normalize
result = A_res @ result
return result

A1, A2 = random_attn(4), random_attn(4)
R = rollout([A1, A2])
print('Token-0 attribution to inputs:', R[0].numpy().round(4))

Why: result is what everything below it consumes, so the wrong expression here fails later and somewhere else. Adding 0.5·I and re-normalizing models the fact that residual connections carry raw token information alongside the attended information.

Annotate

Walk the callouts on Attention rollout: 2-layer toy example. Each one is a place this is easy to get subtly wrong.

• Adding 0.5·I and re-normalizing models the fact that residual connections carry raw token information alongside the attended information.
• Row sums = 1.0 (verified). Off-diagonal values reveal which original input positions influenced each output token after 2 layers.

Predict first

A student writes this, and it looks reasonable:

Wrong approach: read layer-12 attention weights and say "token A attends to token B, therefore B caused the prediction."

It is wrong. Say what breaks — and say it before you turn the page.

Correct: Sounds intuitive — high weight = high importance.

Correct approach: use attention rollout (or gradient attribution) to trace all the way back to input positions.

Why: Sounds intuitive — high weight = high importance.

Discussion prompt

The rule this trap just fixed:

Each intermediate representation is a mix of attended tokens AND raw pass-through (residual). Rollout models that correctly.

Now break it on purpose. Build a case that violates it and follow the consequences until something visibly fails. Where does the failure first show up — and would you have noticed it if you had not been looking?

Hint: The dangerous rules are the ones whose violation still produces an answer. If yours fails loudly, try to find one that fails quietly.

Discussion prompt

Explain The probing paradigm by analogy to something with no Machine Learning in it at all — a queue, a recipe, a map, a bank balance, whatever fits. Then say where your analogy breaks.

Hint: An analogy that never breaks is not an analogy, it is the same idea wearing a hat. Find the seam — that is the part that is actually new.

Fill in the blanks

Train a linear POS probe, with the scaffolding fading: two lines are gone now — fill both.

import numpy as np
from sklearn.linear_model import LogisticRegression
from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score

np.random.seed(42)
reps = np.random.randn(200, 32)
reps[:100, 0] += 2.0 # noun signal in dim 0
reps[100:, 0] -= 2.0 # verb signal
labels = np.array([0]100 + [1]100) # 0=noun, 1=verb

X_tr, X_te, y_tr, y_te = train_test_split(reps, labels,
test_size=0.25, random_state=0)
probe = LogisticRegression(max_iter=200).fit(X_tr, y_tr)
acc = accuracy_score(y_te, probe.predict(X_te))
print(f'Probe accuracy: ___') # 0.9200
print(f'Coeff dim-0: ___') # dominant

Why: Reproducing these unaided, rather than reading them, is what tells you the method has transferred. We simulate 200 tokens (100 nouns, 100 verbs) with d=32 representations; noun reps have a +2 shift in dimension 0 modeling the syntactic signal BERT encodes.

Draw it

Draw what Train a linear POS probe just did — the shape of it, not the line-by-line working. One picture, labels only where you need them. Then check it against the steps: anything you could not draw is a step you followed rather than understood.

Discussion prompt

In Gradient × input: one backward pass this move is made:

Token 3 is the most influential: highest raw attribution 0.7319, normalized share 24.2%

Why is that legal? Name the rule or definition it rests on before you read on.

Hint: If you can only say "because that is what you do", the rule is the thing to go and find.

Comparison matrix

From Gradient × input: one backward pass: refill the most influential? column from what you know. The rest of the table is as it appeared.

Fill in the blanks

From Integrated gradients: 50-step Riemann approximation — one line has had its right-hand side removed. Put it back.

import torch
torch.manual_seed(5)
T, d = 3, 4
inp = torch.randn(T, d)
baseline = torch.zeros(T, d)
w = torch.randn(d) # simple linear model

n_steps = 50
alphas = torch.linspace(0, 1, n_steps)
grads = []
for alpha in alphas:
x_a = (baseline + alpha*(inp - baseline)).requires_grad_(True)
out = **(x_a.mean(0) * w).sum()**
out.backward()
grads.append(x_a.grad.detach().clone())

avg_g = torch.stack(grads).mean(0) # (T, d)
ig = (inp - baseline) * avg_g # (T, d)
token_ig = ig.abs().sum(-1) # (T,)
print('IG per token (raw):', token_ig.numpy().round(4))
print('IG normalized:', (token_ig/token_ig.sum()).numpy().round(4))

Why: out is what everything below it consumes, so the wrong expression here fails later and somewhere else. At α=0 the input is all zeros (neutral baseline); at α=1 it is the actual embedding.

Annotate

Walk the callouts on Integrated gradients: 50-step Riemann approximation. Each one is a place this is easy to get subtly wrong.

• At α=0 the input is all zeros (neutral baseline); at α=1 it is the actual embedding. Averaging gradients along this path captures the contribution even when the endpoint gradient is near zero.
• IG integrates over the entire path, so even tokens that hit a flat gradient region at α=1 still get credited for the work done at smaller α values.

Discussion prompt

Probing classifiers at every layer reveal a consistent pattern in BERT-family models: lower layers encode surface and syntactic features; upper layers encode semantics and task-relevant features.

That is stated as though it always holds. Do one of two things: produce a case where it fails, or say precisely what rules such a case out. "It just does" is not on the menu.

Hint: Hunt at the extremes first — zero, one, negative, empty, equal. If every extreme survives, the reason they survive is the proof.

Discussion prompt

Run Transformer interpretability: the four-tool recipe with this step confiscated:

Attention rollout: fold in residual connections: à = 0.5·A + 0.5·I, re-normalize, then R = Ã_L · … · Ã_1; row i of R is token i's attribution to each input token

Is it still possible? If it is, say what takes its place and what it costs you. If it is not, say exactly what that step was providing that nothing else does.

Hint: A step you can drop for free was never load-bearing. If you cannot drop it, name the thing that goes wrong the moment it is gone.

Discussion prompt

Transformer interpretability: the four-tool recipe works on the cases you have just seen. Push it to the edge: what is the most degenerate input it still handles — empty, zero, one item, everything equal — and what is the first case where it stops being true? Name the case, not just "it breaks".

Hint: Try the smallest legal input, then the largest, then the one where two things collide. Methods are specified at their edges; the middle takes care of itself.

Eliminate the wrong options

Why does attention rollout add 0.5·I to each layer's attention matrix before multiplying?

3 of these 4 are wrong. Strike them one at a time, and say what rules each one out before you strike the next. The survivor is the answer.

• A. To model the residual (skip) connections that pass token information unchanged through each layer
• B. To prevent numerical underflow when multiplying many small attention probabilities together
• C. To normalize the rows to sum to 1 after the matrix product
• D. To zero out attention weights below a confidence threshold

Survives elimination: A

Why: Residual connections in a transformer add the token's own representation directly to the attention output at each layer. The 0.5·A + 0.5·I formula models this: 50% of the signal passes through attention, 50% bypasses it via the skip path. Without the identity term, rollout would attribute everything to heavily attended tokens even when residuals dominate the signal.

Check your understanding

Why does attention rollout add 0.5·I to each layer's attention matrix before multiplying?

• A. To model the residual (skip) connections that pass token information unchanged through each layer (correct)
• B. To prevent numerical underflow when multiplying many small attention probabilities together
• C. To normalize the rows to sum to 1 after the matrix product
• D. To zero out attention weights below a confidence threshold

Answer: A

Why: Residual connections in a transformer add the token's own representation directly to the attention output at each layer. The 0.5·A + 0.5·I formula models this: 50% of the signal passes through attention, 50% bypasses it via the skip path. Without the identity term, rollout would attribute everything to heavily attended tokens even when residuals dominate the signal.

Why B tempts people

Underflow is a real concern when multiplying many matrices, but the identity addition serves no numerical stabilization purpose — a log-space formulation would address underflow instead. The correct reason is residual modeling.

Why C tempts people

Re-normalization is a separate step applied after adding the identity; the addition itself does not enforce that rows sum to 1 (they need an explicit row-normalization pass).

Why D tempts people

Thresholding would require a comparison, not an additive identity. The 0.5·I adds equal probability mass to every self-position, the opposite of zeroing out low weights.

Predict first

A linear probe trained on layer-6 hidden states achieves 89% accuracy at predicting dependency-role labels (subject vs object). Which conclusion is best supported?

Answer it in your own words, now, with nothing to choose from. The options are on the next slide — and picking the right one off a list is an easier skill than producing it.

Correct: Dependency-role information is linearly decodable from layer-6 representations

Why: Probing classifiers measure the geometry of representations — specifically whether a given property is linearly separable in that space. 89% linear probe accuracy means the property is encoded in a linearly accessible way. It does NOT prove the model uses this information causally (that requires intervention experiments like activation patching), nor that training was supervised on it (BERT learns it without explicit labels).

Check your understanding

A linear probe trained on layer-6 hidden states achieves 89% accuracy at predicting dependency-role labels (subject vs object). Which conclusion is best supported?

• A. The model uses syntactic dependency roles to compute its predictions
• B. Dependency-role information is linearly decodable from layer-6 representations (correct)
• C. The model was explicitly trained with dependency-role supervision
• D. A deeper probe (MLP) trained on layer-6 would achieve lower accuracy

Answer: B

Why: Probing classifiers measure the geometry of representations — specifically whether a given property is linearly separable in that space. 89% linear probe accuracy means the property is encoded in a linearly accessible way. It does NOT prove the model uses this information causally (that requires intervention experiments like activation patching), nor that training was supervised on it (BERT learns it without explicit labels).

Why A tempts people

Correlation in representations ≠ causal use. The model could encode dependency roles without using them — the probe reveals encoding, not computation. Causal claims require patching or ablation experiments.

Why C tempts people

BERT and GPT learn syntactic structure from raw text via the masked-LM or next-token objectives alone; no dependency-role supervision is needed. The probe measures emergent structure.

Why D tempts people

A nonlinear (MLP) probe typically achieves higher or equal accuracy on the same representations, not lower. If anything, a deep MLP probe would overfit noise and not test linear separability — which is why linear probes are preferred.

Eliminate the wrong options

Integrated gradients satisfies the completeness axiom. What does completeness guarantee?

3 of these 4 are wrong. Strike them one at a time, and say what rules each one out before you strike the next. The survivor is the answer.

• A. Attributions are always non-negative, so each token's contribution is interpretable as a probability
• B. The sum of all token attributions equals f(input) − f(baseline)
• C. Each attribution equals the gradient of f at the input point, scaled by the embedding norm
• D. The attribution of any token not in the baseline vocabulary is exactly zero

Survives elimination: B

Why: Completeness (also called 'efficiency') states that Σᵢ IG(i) = f(x) − f(x_baseline). This is provable from the fundamental theorem of calculus applied to the path integral. It means attributions account for the full output difference with no unexplained residual. Neither gradient × input nor raw attention satisfies this. Attributions can be negative (token pushed output away from the prediction) so A is false.

Check your understanding

Integrated gradients satisfies the completeness axiom. What does completeness guarantee?

• A. Attributions are always non-negative, so each token's contribution is interpretable as a probability
• B. The sum of all token attributions equals f(input) − f(baseline) (correct)
• C. Each attribution equals the gradient of f at the input point, scaled by the embedding norm
• D. The attribution of any token not in the baseline vocabulary is exactly zero

Answer: B

Why: Completeness (also called 'efficiency') states that Σᵢ IG(i) = f(x) − f(x_baseline). This is provable from the fundamental theorem of calculus applied to the path integral. It means attributions account for the full output difference with no unexplained residual. Neither gradient × input nor raw attention satisfies this. Attributions can be negative (token pushed output away from the prediction) so A is false.

Why A tempts people

IG attributions can be negative — a token may actively suppress a class logit. Requiring non-negativity would violate the completeness sum (since f(x) − f(baseline) can be positive while some individual terms are negative).

Why C tempts people

That describes gradient × input saliency, not integrated gradients. IG averages gradients along the interpolation path from baseline to input, not just the gradient at the endpoint.

Why D tempts people

IG makes no statement about vocabulary membership. The baseline is a zero embedding (or pad token), not a vocabulary restriction. Tokens present in both input and baseline can still receive non-zero attribution.

Match the pairs

From Project brief — match each one to what it actually does. The descriptions have been shuffled.

• c1. Milestone 1
• c2. Milestone 2
• c3. Milestone 3
• c4. Show it off

• b1. Toy transformer forward pass — extract per-head attention tensors
• b2. Attention rollout across both layers
• b3. Gradient × input token attribution with a single backward pass
• b4. Compare heatmap vs rollout vs gradient attribution for the same input

Why: Milestone 1, Milestone 2, Milestone 3, Show it off are easy to tell apart while they are sitting next to their descriptions and much harder afterwards, which is what this checks.

Sort into buckets

Sort these by shape, from memory, without looking back at Milestone 1: toy transformer + attention…. Telling them apart on the spot is the skill; the table is only where the answer happens to be written down.

(1, 6, 16)

x input; out

(1, 2, 6, 6)

attn_maps[0]; attn_maps[1]

g1

shape is "(1, 6, 16)" for x input, out — that is what the table on "Milestone 1: toy transformer +…" records, and it is the single property separating this group from the rest.

g2

shape is "(1, 2, 6, 6)" for attn_maps[0], attn_maps[1] — that is what the table on "Milestone 1: toy transformer +…" records, and it is the single property separating this group from the rest.

Fill in the blanks

Milestone 2: rollout across 2 layers, with the scaffolding fading: two lines are gone now — fill both.

import torch, torch.nn.functional as F

def rollout(attn_maps):
# attn_maps: list of (B, n_heads, T, T)
result = None
for w in attn_maps:
A = w[0].mean(0) # (T, T): average over heads
A_res = (A + torch.eye(A.shape[0])) / 2
A_res = A_res / A_res.sum(-1, keepdim=True)
if result is None:
result = A_res
else:
result = A_res @ result
return result # (T, T)

R = rollout(model.attn_maps) # model from Milestone 1
print('Rollout shape:', R.shape) # (6, 6)
print('Row sums:', R.sum(-1).detach().numpy().round(4))
print('Row 0 (token-0 attribution):', R[0].detach().numpy().round(3))

Why: Reproducing these unaided, rather than reading them, is what tells you the method has transferred. Multiple heads are averaged to get a single (T×T) map per layer before rollout — a standard simplification that treats all heads as equally important.

Predict first

The table runs: 0 | computed | ~1/6 · 1 | computed | ~1/6 · 2 | computed | ~1/6 · 3 | computed | ~1/6 · 4 | computed | ~1/6

In Milestone 3: gradient × input attribution, given the rows so far: what is the next one — the row where token idx is 5?

Correct: 5 | computed | ~1/6

Why: The relationship between the columns, not the individual numbers, is what generates the next row. We treat the mean of the output as a proxy scalar.

Comparison matrix

From Milestone 3: gradient × input attribution: every row here is a choice with a cost. Fill the norm attr column, then say which row you would actually pick and what you give up for it.

Fill in the blanks

From Show it off: compare all three tools — one line has had its right-hand side removed. Put it back.

# Run after Milestones 1-3 above
import torch
torch.manual_seed(42)
model3 = TinyTransformer(d=16, n_heads=2, T=6)
x_full = torch.randn(1, 6, 16, requires_grad=True)
out_full = model3(x_full)
out_full.mean().backward()

# 1. Head-0 of layer-1 heatmap (row 2)
head0_row2 = model3.attn_maps[0][0, 0, 2].detach().numpy().round(3)

# 2. Rollout row 2
R2 = rollout(model3.attn_maps)
rollout_row2 = R2[2].detach().numpy().round(3)

# 3. Gradient x input token 2
g = x_full.grad
attr_t2 = (g * x_full.detach()).abs().sum(-1)[0].detach().numpy().round(3)

print('Head-0 L1 attn (row 2):', head0_row2)
print('Rollout attn (row 2):', rollout_row2)
print('Grad x input (all T):', attr_t2)

Why: rollout_row2 is what everything below it consumes, so the wrong expression here fails later and somewhere else. Heatmap shows layer-local patterns; rollout shows propagated input-space attribution; gradient × input shows output-sensitivity.

Comparison matrix

From Show it off: compare all three tools: refill the question answered column from what you know. The rest of the table is as it appeared.

Draw it

One page, no notation unless you need it: draw how these connect — Attention heatmaps · Attention rollout · Linear probing classifiers · Gradient-based attribution · Your turn: attention inspector. Put an arrow wherever one of them is what makes another possible, and label the arrow with why.