Attention: summarize every head, then test one

Note

Generated from notebooks/applications/attention.ipynb, which you can run yourself. The outputs below are the ones it produced.

GPT-2 small has 144 attention heads. Staring at 144 attention matrices is not analysis. The useful move is to reduce each head’s pattern to a single number, scan the resulting 12x12 grid for something unusual, and only then look at the individual head.

Then, crucially, check whether the interesting-looking head actually does anything.

Key questions

• Which heads are diffuse, and which focus on one token?
• Which heads attend to the very first token (an “attention sink”)?
• Does an interesting-looking head matter causally, and what does it write?

Structure

1. Set up
2. Record and reduce every head
3. Inspect one head
4. Ablate a head and measure the damage
5. Read what the head writes with the OV circuit
6. Save

Model and data. GPT-2 small in float32. It is small enough that bfloat16 rounding visibly degrades its predictions, so the notebooks load it explicitly in float32. Eight indirect-object prompt pairs from murano.tasks.ioi. Loaded with eager attention so the per-head weights are exposed.

Requirements. pip install murano-interp[plot]

1. Set up

Attention weights are not exposed by the fused attention kernels, so the model has to be loaded with enable_attention_probs=True. It is off by default because eager attention is slower.

import torch

from murano import MuranoModel, NodeSet, Pipeline, keys
from murano.nodes import Node
from murano.plotting import plot_attention_pattern, plot_head_matrix, save_figure
from murano.steps import (
    AblateAttention,
    LoadPrompts,
    LogitDiffStep,
    Logits,
    RecordAttention,
    Save,
    Sweep,
    ov_circuit,
)

OUTPUT_DIR = "murano_outputs/attention"

# GPT-2 is small enough that bfloat16 rounding degrades its predictions.
model = MuranoModel("gpt2", dtype=torch.float32, enable_attention_probs=True)
print(f"{model.model_id}: {model.n_layers} layers, {model.n_heads} heads")

gpt2: 12 layers, 12 heads

This notebook uses the indirect-object-identification task from murano.tasks. Each clean prompt names two people and then has the subject give a drink, so the next token should be the other name, the indirect object. The corrupt prompt swaps who gives, which flips the answer.

clean  : When Mary and John went to the store, John gave a drink to ___  ->  " Mary"
corrupt: When Mary and John went to the store, Mary gave a drink to ___  ->  " John"

The metric throughout is the logit difference: the correct name’s logit minus the incorrect name’s. Positive means the model prefers the right answer.

from murano.tasks import ioi

task = ioi(n=8)
print("clean  :", task.clean[0], "->", repr(task.correct[0]))
print("corrupt:", task.corrupt[0], "->", repr(task.incorrect[0]))
print(f"{len(task.clean)} prompt pairs")

clean  : When Mary and John went to the store, John gave a drink to -> ' Mary'
corrupt: When Mary and John went to the store, Mary gave a drink to -> ' John'
8 prompt pairs

2. Record and reduce every head

RecordAttention captures the post-softmax weights for every head: [batch, head, query, key] per layer.

AttentionResult then offers reductions that collapse each head’s pattern to one number, giving a [n_layers, n_heads] grid:

• entropy() is high for diffuse heads, low for focused ones.
• sink() is the mass placed on the first token, a known transformer quirk.
• distance() is how far back a head typically looks.

results = Pipeline([LoadPrompts(task.clean), RecordAttention(model)]).run()

attention = results[keys.ATTENTION_PATTERN]
print("one number per head:", tuple(attention.entropy().shape))

one number per head: (12, 12)

Entropy: which heads are focused?

fig = plot_head_matrix(
    attention.entropy(), title="Attention entropy", value_label="entropy"
)
save_figure(fig, f"{OUTPUT_DIR}/plots/entropy.png")
fig

Sink: which heads park on the first token?

fig = plot_head_matrix(
    attention.sink(), title="Attention to the first token", value_label="mass"
)
save_figure(fig, f"{OUTPUT_DIR}/plots/sink.png")
fig

Distance: which heads look far back?

fig = plot_head_matrix(
    attention.distance(), title="Mean attention distance", value_label="tokens"
)
save_figure(fig, f"{OUTPUT_DIR}/plots/distance.png")
fig

3. Inspect one head

Head 9.9 is one of GPT-2’s name movers. Rows are query positions, columns are key positions, and the lower triangle is all a causal model can see.

The figure is dominated by a dark first column. That is the attention sink from the grid above: nearly every query position dumps most of its mass on token 0, which carries no information here. Ignore it and read the bottom row, the position that actually predicts the answer.

fig = plot_attention_pattern(attention, layer=9, head=9)
save_figure(fig, f"{OUTPUT_DIR}/plots/head_9_9.png")
fig

pattern = attention.head_pattern(layer=9, head=9)
print("query x key:", tuple(pattern.shape))

tokens = model.tokenizer.convert_ids_to_tokens(
    model.tokenizer(task.clean[0])["input_ids"]
)

top = torch.topk(pattern[-1], 4)
print("\nfrom the final position, head 9.9 attends most to:")
for value, index in zip(top.values, top.indices):
    token = tokens[int(index)].replace("\u0120", " ")
    note = "  <- attention sink" if int(index) == 0 else ""
    print(f"  position {int(index):>2}  {token!r:<10} {float(value):.3f}{note}")

query x key: (14, 14)

from the final position, head 9.9 attends most to:
  position  1  ' Mary'    0.362
  position  0  'When'     0.359  <- attention sink
  position  3  ' John'    0.132
  position  9  ' John'    0.129

The bottom row is the exception to the sink: head 9.9 splits its attention between token 0 and the indirect object’s name, and gives the subject’s name noticeably less. That is the name-mover signature, and it is a hypothesis about mechanism.

4. Ablate a head and measure the damage

A suggestive pattern is a hypothesis, not a result. The test is to delete the head and see whether the behavior degrades. AblateAttention overwrites a head’s attention weights, here with zeros, and reads the resulting logits.

baseline = Pipeline(
    [
        LoadPrompts(task.clean),
        Logits(model),
        LogitDiffStep(task.correct, task.incorrect, model=model),
    ]
).run()

clean_score = baseline[keys.LOGIT_DIFF].value
print(f"clean logit difference: {clean_score:+.4f}")

clean logit difference: +2.3971

Sweep runs that ablate-and-score chain once per head, forking the baseline each time. The candidates are the heads the entropy and distance grids flagged, plus one early head, L0H0, which shows none of the name-mover signature.

candidates = NodeSet(
    [
        Node(10, "self_attn", head=0),
        Node(9, "self_attn", head=9),
        Node(9, "self_attn", head=6),
        Node(10, "self_attn", head=7),
        Node(0, "self_attn", head=0),
    ]
)

damage = Pipeline(
    [
        Sweep(
            over=candidates,
            steps=lambda head: [
                AblateAttention(model, targets=[head], method="zero"),
                LogitDiffStep(
                    task.correct,
                    task.incorrect,
                    logits_key=keys.ATTN_ABLATED_LOGITS,
                    mask_key=keys.ATTN_ABLATED_MASK,
                    model=model,
                ),
            ],
            read=keys.LOGIT_DIFF,
        )
    ]
).run(baseline.copy())

print(f"{'head':<10} {'logit diff':>12} {'change':>10}")
print("-" * 34)
for node, score in damage[keys.SWEEP].scores.items():
    label = f"L{node.layer}H{node.head}"
    print(f"{label:<10} {score:>+12.4f} {score - clean_score:>+10.4f}")

head         logit diff     change
----------------------------------
L10H0           +1.6410    -0.7560
L9H9            +2.3897    -0.0073
L9H6            +2.5959    +0.1989
L10H7           +3.5912    +1.1941
L0H0            +2.0301    -0.3670

Two things stand out.

L10H0 is load-bearing. Removing it drops the logit difference: it was pushing toward the correct name.

L10H7 is a negative name mover. Removing it raises the logit difference. The head was actively suppressing the correct answer, and deleting it makes the model better at this task. This is the same head direct logit attribution flagged with a negative contribution.

Note also that zeroing L9H9 alone barely moves the metric, even though its attention pattern looks like a textbook name mover. Other heads compensate, and zeroing is a blunt instrument: ablation.ipynb shows that a better-chosen ablation reveals L9H9 clearly.

And L0H0 is not free. It has no role in this task, yet zeroing it costs more logit difference than zeroing L9H9 does. Nothing was deleted from the circuit; the model was simply handed a layer-0 attention pattern no forward pass ever produces, and every later layer read the damage. Treat a nonzero ablation effect as evidence that something changed, never as evidence that the head you deleted was doing the thing you were looking for.

5. Read what the head writes with the OV circuit

Attention weights say where a head looks. They say nothing about what it writes. That is the OV circuit: the linear map a head applies to whatever it reads before adding it back to the residual stream.

ov_circuit(model, layer, head) returns that map. Compose it with the unembedding and you can ask a sharp question: if this head reads the token " Mary", does it make the model more likely to say " Mary"? A head for which the answer is yes is a copying head.

One subtlety decides whether this works. GPT-2’s layer-0 MLP acts as an extended embedding, so the vector a layer-9 head actually reads is not the raw token embedding but the residual stream after layer 0. Use that “effective embedding”.

NAMES = [" Mary", " John", " Alice", " Bob", " Sarah", " Tom", " Emma", " James"]

raw_embedding = model.hf_model.wte.weight.detach().float().cpu()
effective = {
    name: model.record(name, layers=[0]).positive[(0, "residual")][0].float().cpu()
    for name in NAMES
}


def copying_score(layer: int, head: int, embeddings) -> int:
    """How many names the head promotes into its own top-5 output tokens."""
    ov = ov_circuit(model, layer, head)
    hits = 0
    for name in NAMES:
        token_id = model.tokenizer(name)["input_ids"][0]
        vocab = model.project_on_vocab(ov @ embeddings[name]).detach().float().cpu()
        hits += int(token_id in torch.topk(vocab, 5).indices.tolist())
    return hits


raw = {name: raw_embedding[model.tokenizer(name)["input_ids"][0]] for name in NAMES}

print(f"{'head':<10} {'raw embedding':>15} {'effective embedding':>21}")
print("-" * 48)
for layer, head in [(9, 9), (9, 6), (10, 0), (10, 7), (11, 10), (0, 0), (5, 5)]:
    label = f"L{layer}H{head}"
    print(
        f"{label:<10} {copying_score(layer, head, raw):>12}/{len(NAMES)} "
        f"{copying_score(layer, head, effective):>18}/{len(NAMES)}"
    )

head         raw embedding   effective embedding
------------------------------------------------
L9H9                  0/8                  7/8
L9H6                  0/8                  1/8
L10H0                 3/8                  3/8
L10H7                 0/8                  0/8
L11H10                7/8                  0/8
L0H0                  0/8                  0/8
L5H5                  0/8                  0/8

Read the right-hand column. L9H9 promotes almost every name it reads into its own top-5 output tokens. Whatever name it attends to, it pushes that name toward the logits. That is what “name mover” means, and it is a statement about the head’s weights, independent of any prompt.

L10H7, the negative name mover, copies nothing. It attends to the name and writes something that does not promote it. The two columns also show why the effective embedding matters: with the raw token embedding, L11H10 looks like a strong copier and L9H9 like none at all. Both readings are artifacts.

6. Save

run_dir = Save(
    output_dir="murano_outputs", run_name="attention", model_id=model.model_id
)(results)[keys.OUTPUT_DIR]

print("saved to:", run_dir)

saved to: murano_outputs/attention

What next

• ablation.ipynb shows that the choice of ablation decides the answer.
• activation_patching.ipynb measures every head’s causal contribution at once.
• circuit_discovery.ipynb shows how the heads found here feed one another.