Looking inside a transformer with TransformerLens

Interpretability work is mostly reading numbers out of a running model and checking whether they mean what you think. This section sets up the tooling for that: loading a small model with every internal value exposed, capturing them on a forward pass, and using three standard views, the residual stream, the logit lens and attention patterns, to watch one prediction form layer by layer. Everything later in this topic assumes you can get these values out and know which name refers to which point in the model.

Every activation, by name#

TransformerLens loads models such as GPT-2 small with a hook at every interesting point, so you can read or change any internal value. It is the standard tool for learning this field on small models.

The naming scheme is the part worth learning first, because every later technique is a matter of naming a location. A cache key is a component path and a hook point: blocks.5.hook_resid_post is the residual stream after block 5, blocks.5.attn.hook_z is the per-head output of block 5's attention before it is projected back, and the shorthand cache["pattern", 5] gives the attention weights. Shapes follow a convention too. Residual stream tensors are batch by position by model width; attention patterns are batch by head by query position by key position. Getting those two shapes into your fingers removes most of the confusion in the first week.

from transformer_lens import HookedTransformer

model = HookedTransformer.from_pretrained("gpt2-small")     # 12 layers, 12 heads
logits, cache = model.run_with_cache("The Eiffel Tower is in the city of")

print(model.to_string(logits[0, -1].argmax()))   # " Paris"
print(cache["blocks.5.hook_resid_post"].shape)   # [batch, position, d_model]
print(cache["pattern", 5].shape)                 # [batch, head, query_pos, key_pos]

The residual stream#

Think of a transformer as a stream of vectors, one per position, flowing from the embedding to the output. Each attention layer and each MLP reads from the stream and adds its result back. Nothing overwrites anything; everything is added. This view is what makes the rest of the field workable: you can ask what each component added.

Two consequences are worth stating explicitly. Because the final vector is a sum of contributions, the final logit for a token is a sum of each component's contribution to that logit, and you can attribute an output to components by decomposing that sum rather than by guessing. And because every component reads the whole stream, the stream works as a shared bus: a head in layer 2 can write something that only a head in layer 9 ever reads, with nothing in between touching it. That is why circuits span layers and why reading a single layer in isolation is rarely enough.

The logit lens#

Take the stream partway through the model and push it straight through the final layer norm and unembedding, as if the model stopped there. You get the model's best guess at each depth, and can watch the answer appear.

for layer in range(model.cfg.n_layers):
    resid = cache[f"blocks.{layer}.hook_resid_post"][:, -1:, :]
    guess = model.unembed(model.ln_final(resid))[0, -1].argmax()
    print(layer, repr(model.to_string(guess)))

The lens is useful and it rests on an assumption that is only roughly true: that an intermediate residual state lives in the same space as the final one, so the final unembedding is a fair way to read it. Early layers break that assumption badly, and a lens reading of layer 1 is close to meaningless. Treat the layer at which an answer first appears as a rough marker, not a measurement, and check anything you plan to claim from it by patching instead.

Attention patterns and induction heads#

An induction head completes repeated patterns: having seen 'Jack Kowalczyk' once, on seeing 'Jack' again it attends to the token after the earlier 'Jack' and predicts 'Kowalczyk'. It is one of the best understood circuits, and a good first target because you can find it yourself.

import torch

seq = torch.randint(1000, 10000, (1, 50))
tokens = torch.cat([seq, seq], dim=1)            # random tokens, repeated once
_, cache = model.run_with_cache(tokens)

for layer in range(model.cfg.n_layers):
    pattern = cache["pattern", layer][0]          # [head, q, k]
    # an induction head at position i attends to position i - 49 in the second half
    score = pattern.diagonal(offset=-49, dim1=-2, dim2=-1).mean(-1)
    for head in (score > 0.4).nonzero().flatten().tolist():
        print(f"layer {layer} head {head}: {score[head]:.2f}")

What an attention pattern does not tell you#

An attention pattern is the most seductive picture in the field. It is a heatmap, it is easy to produce, and it appears to show the model paying attention to the right words. What it actually shows is which positions a head reads from. It says nothing about what the head does with what it read, and nothing about whether removing the head would change the answer. Those are separate measurements, and the honest way to learn the difference is to make both on a model small enough to train in ten seconds.

The task below is three tokens: two items, then a cue saying which of the two to answer about. The answer is the colour of the cued item. Attention is written out by hand rather than taken from a library, because the point of the exercise is to switch one head off, and a head has to be visible before it can be removed.

import torch, torch.nn as nn

torch.manual_seed(0)
N_ITEMS, D, N_HEADS = 48, 32, 4
D_HEAD = D // N_HEADS
CUE_FIRST, CUE_SECOND = N_ITEMS, N_ITEMS + 1

def colour(item): return item % 4

class Block(nn.Module):
    """Attention written out so individual heads can be read and switched off."""
    def __init__(self):
        super().__init__()
        self.q, self.k, self.v, self.o = (nn.Linear(D, D) for _ in range(4))
        self.mlp = nn.Sequential(nn.Linear(D, 4 * D), nn.ReLU(), nn.Linear(4 * D, D))

    def forward(self, x, ablate=()):
        B, T, _ = x.shape
        q, k, v = (f(x).view(B, T, N_HEADS, D_HEAD).transpose(1, 2)
                   for f in (self.q, self.k, self.v))
        pattern = (q @ k.transpose(-1, -2) / D_HEAD ** 0.5).softmax(-1)   # [B,head,q,k]
        z = pattern @ v
        for h in ablate:
            z = z.clone(); z[:, h] = 0                    # switch one head off
        x = x + self.o(z.transpose(1, 2).reshape(B, T, D))
        return x + self.mlp(x), pattern

class Toy(nn.Module):
    def __init__(self):
        super().__init__()
        self.embed = nn.Embedding(N_ITEMS + 2, D)
        self.pos = nn.Parameter(torch.randn(3, D) * 0.02)
        self.blocks = nn.ModuleList([Block(), Block()])
        self.out = nn.Linear(D, 4)

    def forward(self, tokens, cache=None, ablate=()):
        x = self.embed(tokens) + self.pos
        for i, blk in enumerate(self.blocks):
            x, pattern = blk(x, ablate=[h for (layer, h) in ablate if layer == i])
            if cache is not None:
                cache[f"pattern_{i}"] = pattern.detach()
        return self.out(x[:, -1])

def make_batch(n):
    """[item 1, item 2, cue] -> the colour of whichever item the cue names."""
    a = torch.randint(0, N_ITEMS, (n,)); b = torch.randint(0, N_ITEMS, (n,))
    cue = torch.randint(0, 2, (n,))
    toks = torch.stack([a, b, torch.where(cue == 0, CUE_FIRST, CUE_SECOND)], 1)
    return toks, torch.where(cue == 0, colour(a), colour(b))

model = Toy()
opt = torch.optim.Adam(model.parameters(), lr=3e-3)
for step in range(800):
    toks, labels = make_batch(256)
    loss = nn.functional.cross_entropy(model(toks), labels)
    opt.zero_grad(); loss.backward(); opt.step()

Now measure both things for every head. How much of its attention, from the final position, goes to the item the cue selected; and what it costs to switch the head off. The cost needs two metrics, because one of them turns out to be useless.

toks, labels = make_batch(4000)
cued_first = toks[:, 2] == CUE_FIRST
cache = {}
with torch.no_grad(): model(toks, cache=cache)

def score(ablate=()):
    """Accuracy, and the margin between the right logit and the best wrong one."""
    with torch.no_grad(): lg = model(toks, ablate=ablate)
    right = lg.gather(1, labels[:, None]).squeeze(1)
    rest = lg.clone(); rest.scatter_(1, labels[:, None], -1e9)
    return (lg.argmax(-1) == labels).float().mean().item(), (right - rest.max(-1).values).mean().item()

base_acc, base_margin = score()
print(f"baseline: accuracy {base_acc:.3f}  margin {base_margin:.2f}\n")
print("head   attends to the cued item   accuracy   margin   margin lost")
attn, lost = [], []
for layer in range(2):
    pattern = cache[f"pattern_{layer}"][:, :, 2]           # queries from the last position
    for h in range(N_HEADS):
        w = torch.where(cued_first, pattern[:, h, 0], pattern[:, h, 1]).mean().item()
        a, m = score([(layer, h)])
        attn.append(w); lost.append(base_margin - m)
        print(f"L{layer}H{h}          {w:6.2f}             {a:.3f}   {m:6.2f}        {base_margin - m:5.2f}")
for name, group in (("all four heads in layer 0", [(0, h) for h in range(4)]),
                    ("all four heads in layer 1", [(1, h) for h in range(4)])):
    a, m = score(group)
    print(f"{name:26s}             {a:.3f}   {m:6.2f}        {base_margin - m:5.2f}")

baseline: accuracy 1.000  margin 18.41

head   attends to the cued item   accuracy   margin   margin lost
L0H0            0.67             1.000    16.05         2.36
L0H1            0.70             1.000    17.21         1.20
L0H2            0.56             1.000    17.39         1.02
L0H3            0.82             0.998    16.33         2.08
L1H0            0.93             0.993    14.10         4.31
L1H1            0.97             1.000    13.23         5.18
L1H2            0.59             1.000    16.64         1.77
L1H3            0.92             1.000    14.26         4.15
all four heads in layer 0              0.987    11.22         7.19
all four heads in layer 1              0.654     2.20        16.21

The accuracy column is the first lesson and it is a negative one. Seven of the eight heads can be deleted with accuracy still reading 1.000, and the eighth costs 0.007. If you had chosen accuracy as your metric you would have concluded that no head in this model matters, which is plainly false, because removing layer 1 entirely takes accuracy to 0.654. The margin column resolves every one of those ties. This is why interpretability work reports logit differences rather than accuracy, and it is worth having met the problem once.

The attention column is the second lesson. L1H1 attends 0.97 of its weight to the cued item and is also the most expensive head to remove, at 5.18, so here the picture and the intervention agree. But L1H2 attends only 0.59, less than L0H3's 0.82, and costs more than any head in layer 0. Across the eight heads the rank correlation between attention weight and margin lost is 0.88, which is high and is also the whole point: you do not know it is 0.88 rather than 0.1 until you run the ablation, and in a model with a genuine division of labour it is often far lower. Attention tells you where a head reads. Ablation tells you whether it matters. Publishing the first while describing the second is the commonest error in this area.

from scipy.stats import spearmanr
print(f"\nrank correlation, attention weight against margin lost: "
      f"{spearmanr(attn, lost).statistic:.2f}")
print(f"accuracy range across all eight single-head ablations: "
      f"{min(score([(l, h)])[0] for l in range(2) for h in range(4)):.3f} to 1.000")
# rank correlation, attention weight against margin lost: 0.88
# accuracy range across all eight single-head ablations: 0.993 to 1.000

One more thing the last two rows are telling you. The four layer-0 heads cost 1.02, 1.20, 2.08 and 2.36 individually, adding to 6.66, and 7.19 when removed together. Layer 1 adds to 15.41 individually and costs 16.21 together. Roughly additive in both cases, which is what you want to check rather than assume: when a group costs far less together than the sum of its parts, the components are redundant and are covering for each other, and single-head ablation will under-report all of them.

One caveat on how the ablation was done. Setting a head's output to zero moves the model off the distribution it was trained on, and some of the damage you measure is the model being confused by an impossible input rather than the head's job being missing. The stricter version replaces the head's output with its mean over a batch of inputs, which keeps the scale and removes only the input-dependent part. On this model the two give the same ordering; on larger models they sometimes do not, and a paper that does not say which it used has left out something load-bearing.

What to carry, and what to skip#

• Carry: residual stream shapes and attention pattern shapes. Nearly every bug in a first week of this work is a transposed axis.
• Carry: report logit differences, not accuracy. The table above is the reason, and it will come up in an interview as 'why not just measure accuracy'.
• Carry: an attention pattern is a hypothesis about a head, and an ablation is a test of it.
• Skip: hand-tuning heatmap plots. A printed table of eight numbers settles the question faster than a picture of it.
• Skip: trying to interpret attention in the first layer of a large model before you have replicated anything. Early layers are mostly doing token-level bookkeeping and are a bad place to start.