The heatmap trap

Multi-head attention was introduced so a model could attend to information from different representation subspaces and positions in parallel.[1] That architectural invitation sounds modular: perhaps one head tracks syntax, another reference, another position, and another topic. Early analyses gave the intuition real support. In BERT, individual heads often showed stable positional patterns, delimiter attraction and striking alignment with some dependency or coreference relations.[2] In machine translation, some important heads were confident, consistent and linguistically interpretable.[3]

Yet another set of observations pulled in the opposite direction. Large fractions of heads could be removed from trained models with little loss on the evaluated task, sometimes leaving a layer with one head.[4] BERT analyses found a small repertoire of repeated attention patterns across many heads, and disabling certain heads could even improve task performance.[5] The natural conclusion was that a few heads specialise while the rest are redundant.

That conclusion is useful but too quick. It assumes that a head’s visible routing pattern and its causal contribution are the same object. They are not. A heatmap answers, “Which source positions received large scalar weights for this query?” The model’s prediction depends on additional questions: “Which vectors were read from those positions? What did the output projection write? How did later components use that write? Could another pathway supply the same information?”

This distinction also explains why attention-as-explanation debates became difficult. Different attention distributions can support similar predictions, and attention weights may not track other importance measures.[6] More fundamentally, for some sequence and head dimensions, attention weights are not identifiable from the output: multiple weight patterns can produce the same contextual result.[7] Even when the weights are stable, they are only coefficients. The norm and direction of the transformed value vectors also determine the head output.[8]

A useful head claim therefore needs coordinates. At minimum, record the model state, the task slice, the candidate relation and the intervention. “Layer 7, head 3 is syntactic” omits all four. “In this checkpoint, on these relative-clause examples, its routing aligns with subject-to-verb dependency edges and replacing its output removes 62 per cent of the agreement margin” is less memorable but scientifically usable. The second statement can be challenged, repeated and narrowed. The first is merely a nickname.

Scope matters because the same head can change role across token positions and input regimes. A route that is useful for short declaratives may collide with another relation in nested clauses. A head that appears redundant under clean text may become the only surviving path under corruption. A head that helps one token’s logit may suppress another. Specialisation is not an essence inside a head; it is a conditional causal relation between a component, a configured system and a defined behaviour.

The practical implication is immediate. A head name should carry a typed claim. “Attends to direct objects on this corpus” is a routing claim. “Is necessary for direct-object agreement on this intervention set” is a causal necessity claim. “Implements direct-object handling” is a mechanism claim that additionally requires a write, a downstream use and a boundary around the relevant circuit. Most interpretability errors begin when these claim types are compressed into one friendly noun.

flowchart LR
  A["The heatmap trap"] --> B["A head routes and writes"]
  B --> C["Three causal organisations"]
  C --> D["From maps to causal profiles"]
  D -. "re-test" .-> A
  classDef first fill:#0a1424,stroke:#7B4CD8,stroke-width:2px,color:#e6f3ff
  classDef step fill:#0a1424,stroke:#3f7da0,stroke-width:1px,color:#e6f3ff
  class A first
  class B,C,D step

A head routes and writes

Consider the residual vector at token position i. A head first constructs a query for that position and keys for possible source positions. Their scaled dot products become attention logits. Softmax converts those logits into routing weights. The head then forms a weighted sum of value vectors and projects the result back into the model dimension.

Here, X contains the current token representations, d_k is the key dimension, and h indexes the head. The product of query and key projections controls where information is read. The value and output projections control which feature is moved and how it is written. Later layers decide whether that write is amplified, cancelled, transformed or ignored. The head is therefore not simply a soft pointer. It is a low-dimensional routing-and-writing channel.

The source representation is itself the product of earlier computation. A head that appears to route from a pronoun to a name may actually be reading a marker written at the name position by an earlier head or feed-forward unit. Likewise, the output is added to a residual stream already carrying many features. Another component may amplify the write, cancel it, or rotate it into a useful direction. A head’s semantic description can therefore be inherited from upstream preparation and downstream use rather than contained in its own parameters.[9]

This is why the familiar phrase “the head stores relation R” is usually too strong. The query-key path may detect a condition, the value-output path may move a feature, and a later circuit may convert that feature into behaviour. The relation belongs to the smallest causal path that preserves all three operations, not automatically to the component with the clearest picture.

The distinction becomes even more important across heads. If two heads produce similar attention maps, they may still write different features. If they produce different maps, their writes may land in the same downstream subspace. A head can also be polysemantic: one routing-and-writing operation may behave differently across input regimes. Successor heads, for example, have been found across several model families and sizes, incrementing ordered tokens such as weekdays or numbers, while at least one analysed head also showed interpretable polysemantic behaviour on natural text.[12]

A recent theoretical preprint gives one reason multiple heads may help even when each query needs a single target. Under an idealised relational task with compressed embeddings, splitting a fixed key-query dimension across heads can reduce interference from superposed relations and increase recoverable capacity.[16] This does not prove that real language-model heads become neat semantic modules. It does suggest a deeper design role for head multiplicity: several narrower routing spaces may deconflict relations that collide in one wider space.

The causal unit is therefore not given by the architecture’s head boundary. A head is a convenient intervention site. The function may be narrower than a head, duplicated across heads, or spread across a circuit that includes several heads and other components.

Three causal organisations

The terms specialised, redundant and distributed are often treated as informal descriptions. They become sharper when defined by interventions. Let the full model achieve a baseline performance P. Let P_−h be performance after ablating head h. Let P_{h} be performance when only that head is restored from a common inactive baseline. The exact performance measure may be accuracy, loss, logit difference or another task-specific quantity, but it must be fixed before comparing interventions.

A causally specialised head has a selective route-write pattern, a meaningful single-head ablation effect and low replaceability by plausible alternatives. A redundant head belongs to a set where single ablations are small but group ablation is large, and one member can rescue much of another’s function. A distributed head belongs to a coalition where individual components are insufficient: the group rescues a behaviour that its members cannot produce alone.

A toy system where the difference is visible

Why real circuits often exceed one head

Induction heads provide a clean example. In small attention-only models, a lower-layer previous-token head can write information that lets a later induction head match the current token against an earlier context and copy the token that followed that earlier match. The familiar pattern [A][B] ... [A] → [B] is therefore not the work of one isolated lens. It is implemented by a cross-layer circuit. Evidence is strong and causal in small attention-only models, while the broader claim that induction heads explain most in-context learning in large models remains more indirect.[10]

Larger natural-language behaviours can involve many more heads. A detailed analysis of indirect-object identification in GPT-2 small described a circuit of 26 attention heads grouped into seven functional classes, including name movers, induction-related components, inhibition and backup mechanisms.[11] The important word is not “26”. It is “circuit”. Some components only become intelligible through the paths connecting them.

Recent causal head-gating work on Llama 3 models likewise reports multiple sparse task-sufficient subcircuits and low modularity, meaning an individual head’s role often depends on interactions with other heads.[15] This is evidence against a simple atlas in which every head receives one stable cognitive noun. It supports a conditional atlas: for this task slice, under this intervention and baseline, these heads facilitate, interfere with or do not measurably affect the behaviour.

The classification also depends on grain. Head B in the toy model is specialised relative to the candidate set of heads, but its effect still requires a downstream sign rule. An induction head is specialised for a routing operation yet distributed with the previous-token head that prepares its key. A useful report therefore names both the local role and the circuit boundary: “specialised writer within a distributed relay” is coherent; “the induction algorithm head” is not.

stateDiagram-v2
  [*] --> Claim
  state "A toy system where the<br/>difference is visible" as Claim
  state "Why real circuits often exceed<br/>one head" as Test
  state "1. Route: which relation does<br/>the head select?" as Revision
  Claim --> Test: expose mechanism
  Test --> Decision: evidence survives
  Test --> Revision: rival survives
  Revision --> Claim: narrow or rebuild
  Decision --> [*]

From maps to causal profiles

The practical alternative to one-word head labels is a route-write-replace-coalition profile. It does not promise a complete mechanistic explanation. It forces four questions to remain separate long enough for the hidden organisation to appear.

Begin with a behaviour, not a component. Define the exact token positions, contrast cases and metric whose mechanism matters. Then use routing observations to generate candidate relations. Only after that should interventions test writes, substitutes and coalitions. Starting from a visually distinctive head and searching for a story invites selection bias: the analyst chooses the relation after seeing the map, then evaluates it on examples that display the pattern.

The profile is a versioned object rather than a permanent label. It should change when the checkpoint, prompt distribution, tokenisation, decoding rule or intervention baseline changes. That discipline may seem administrative, but it prevents a result about one configured system from becoming folklore about an architecture family.

1. Route: which relation does the head select?

Measure whether the attention matrix aligns with a candidate relation beyond simple baselines. If a head appears to track syntactic parents, compare it with fixed offsets, delimiter frequency, token identity and local-distance rules. Evaluate across paraphrases, lengths and task slices. Routing selectivity is useful evidence because it narrows the hypothesis space. It is not sufficient evidence because the head may write little, write the wrong feature or be one of several substitutes.

A compact route score can be written as S_h = R(A_h, relation) − R(baseline, relation), where R is an alignment measure chosen before seeing the result. The equation is valid only for the named relation and corpus. Remove that assumption and S_h says nothing about other inputs or causal use.

2. Write: what information reaches the shared computation?

Inspect the value-output contribution, not just the attention coefficient. Useful probes include the norm and direction of each source contribution, the subspace written by the OV product, direct logit attribution where appropriate, and activation replacement that changes the write while holding other variables fixed. Norm alone is not a mechanism claim. A large vector can be orthogonal to the downstream readout, while a small vector can cross a sensitive decision boundary.

For a query position i, the source contribution from token j is approximately a_ij,h(x_jW^V_h)W^O_h. This makes the smallest useful test clear: preserve a_ij,h, alter the transformed value, and observe whether the targeted behaviour changes as predicted.

3. Replace: which alternative pathway can substitute?

Single-head ablation measures current necessity only when the replacement baseline is meaningful. To test redundancy, attempt rescue. Replace the ablated head with another head’s output, a learned sparse combination of heads, a mean activation or a matched counterfactual activation. Then repeat on a shifted slice. A substitute that works only on the training-like benchmark may share a shortcut rather than the intended function.

Define the necessity cost N_h = L(M_−h) − L(M), where larger loss means worse performance. For a candidate substitute g, define remaining damage D_h←g = L(M_h←g) − L(M). If N_h is large but some substitute leaves damage near zero across the declared distribution, the head is used but replaceable. The conclusion depends on the candidate substitute class. Failure to find a substitute is not proof that none exists.

4. Coalition: which group is necessary and sufficient?

Run group ablations and restoration tests. Redundancy appears when individual removals are safe but removing the group is harmful. Distributed complementarity appears when individual restorations are weak but a group restoration is strong. For a set S, let coalition rescue be C(S) = P(M_S) − P(M_∅). In the synthetic lab, C({C}) = 0 and C({D}) = 0, while C({C,D}) = 0.167. The relation belongs to the pair.

Coalitions need not stop at heads. An attention head may depend on an MLP feature, a preceding head, a positional signal or a later readout. The intervention boundary should expand until it contains the causal path relevant to the claim. That is why automated circuit discovery methods search edges or component sets rather than assuming one head is the natural module.[17]

A positive profile permits only a scoped claim. It can justify preserving a function, targeting a circuit for editing or testing a pruning candidate. A negative result also has content: failed route alignment rejects the proposed relation; a silent write rejects causal use through that path; successful substitution rejects uniqueness; failed coalition rescue says the chosen boundary is incomplete or the hypothesis is wrong. None of these outcomes establishes that the head has no role on untested behaviours.

How to run and interpret the lab

The Python script uses no model framework. It implements scaled dot-product routing with fixed one-hot queries and keys, then applies explicit downstream rules. Run python paper-42-head-ablation-lab.py. The script prints baseline accuracy, single-head ablation cost, rescue from an all-off baseline and selected pair interventions. Assertions protect the two core invariants: the intact model must score 1.0 and the all-off model 0.5.

A positive single-ablation result permits the narrow conclusion that the component is used by the frozen configured system on the tested cases under that ablation. A positive single-rescue result shows sufficiency relative to the all-off baseline for a task slice. A negative single-ablation result does not establish uselessness because another path may compensate. A negative rescue result does not establish irrelevance because the head may need a partner. The lab cannot establish how naturally trained models organise their heads, whether the synthetic categories are exhaustive, or whether a similar pattern transfers to another architecture.

from __future__ import annotations

from dataclasses import dataclass
from itertools import product
from math import exp, sqrt
from typing import Dict, Iterable, List, Mapping, Sequence, Tuple

HEADS = (
    "A_twin_1",
    "A_twin_2",
    "B_specialist",
    "C_relay",
    "D_relay",
    "silent_delimiter",
)
MODES = ("copy_A", "copy_B", "xor_CD")
KEYS = ("A", "B", "C", "D", "SEP")


def softmax(xs: Sequence[float]) -> List[float]:
    maximum = max(xs)
    exponentials = [exp(x - maximum) for x in xs]
    normaliser = sum(exponentials)
    return [x / normaliser for x in exponentials]


@dataclass(frozen=True)
class Case:
    """One synthetic sequence, its query mode and its signed target."""

    bits: Mapping[str, int]
    mode: str

    @property
    def target(self) -> int:
        if self.mode == "copy_A":
            return self.bits["A"]
        if self.mode == "copy_B":
            return self.bits["B"]
        if self.mode == "xor_CD":
            # Signed XNOR: +1 when C and D agree, -1 when they differ.
            return self.bits["C"] * self.bits["D"]
        raise ValueError(f"Unknown mode: {self.mode}")


class FixedAttentionLab:
    """A deterministic six-head attention system with known causal organisation.

    Each head uses a fixed query to attend to one keyed record. Five heads write
    the routed bit. The delimiter head has a sharp attention pattern but a zero
    output gain. A small downstream rule creates redundant, specialised and
    distributed cases that can be distinguished by ablation and rescue.
    """

    target_key = {
        "A_twin_1": "A",
        "A_twin_2": "A",
        "B_specialist": "B",
        "C_relay": "C",
        "D_relay": "D",
        "silent_delimiter": "SEP",
    }

    write_gain = {
        "A_twin_1": 1.0,
        "A_twin_2": 1.0,
        "B_specialist": 1.0,
        "C_relay": 1.0,
        "D_relay": 1.0,
        "silent_delimiter": 0.0,
    }

    def __init__(self, logit_scale: float = 8.0) -> None:
        self.logit_scale = logit_scale

    def attention(self, case: Case, head: str) -> Tuple[List[float], float]:
        """Return attention weights and the scalar written by one head.

        Q is a head-specific one-hot target, K is each token's one-hot key, and
        V is the signed record value. The implementation is intentionally small
        so that the intervention logic is visible rather than hidden in a model
        framework.
        """

        values = [case.bits.get(key, 1 if key == "SEP" else 0) for key in KEYS]
        target = self.target_key[head]
        scaled_match = self.logit_scale / sqrt(len(KEYS))
        logits = [scaled_match if key == target else 0.0 for key in KEYS]
        weights = softmax(logits)
        routed_value = sum(weight * value for weight, value in zip(weights, values))
        return weights, self.write_gain[head] * routed_value

    @staticmethod
    def sign(score: float) -> int:
        return 1 if score >= 0.0 else -1

    def predict(self, case: Case, active: Iterable[str] = HEADS) -> int:
        active_set = set(active)
        outputs: Dict[str, float] = {}
        for head in HEADS:
            _, value = self.attention(case, head)
            outputs[head] = value if head in active_set else 0.0

        if case.mode == "copy_A":
            score = outputs["A_twin_1"] + outputs["A_twin_2"]
        elif case.mode == "copy_B":
            score = outputs["B_specialist"]
        elif case.mode == "xor_CD":
            # A downstream non-linearity makes the two retrieved values jointly necessary.
            score = outputs["C_relay"] * outputs["D_relay"]
        else:
            raise ValueError(f"Unknown mode: {case.mode}")
        return self.sign(score)


def dataset() -> List[Case]:
    """Return all 48 combinations: 16 bit assignments across three query modes."""

    cases: List[Case] = []
    for raw_bits in product((-1, 1), repeat=4):
        bits = dict(zip(("A", "B", "C", "D"), raw_bits))
        for mode in MODES:
            cases.append(Case(bits=bits, mode=mode))
    return cases


def accuracy(model: FixedAttentionLab, cases: Sequence[Case], active: Iterable[str]) -> float:
    active_tuple = tuple(active)
    correct = sum(model.predict(case, active_tuple) == case.target for case in cases)
    return correct / len(cases)


def evaluate() -> None:
    model = FixedAttentionLab()
    cases = dataset()
    baseline = accuracy(model, cases, HEADS)
    no_heads = accuracy(model, cases, ())

    # Invariants make accidental changes visible.
    assert baseline == 1.0
    assert no_heads == 0.5

    print(f"baseline accuracy: {baseline:.3f}")
    print(f"all heads off:     {no_heads:.3f}\n")
    print("single-head interventions")
    print("head                 ablated   loss_delta   restored_from_zero")
    for head in HEADS:
        ablated = accuracy(model, cases, (h for h in HEADS if h != head))
        restored = accuracy(model, cases, (head,))
        print(
            f"{head:20s} {ablated:8.3f} "
            f"{baseline - ablated:11.3f} {restored - no_heads:18.3f}"
        )

    pairs = [
        ("A_twin_1", "A_twin_2"),
        ("C_relay", "D_relay"),
        ("B_specialist", "silent_delimiter"),
    ]
    print("\npair interventions")
    print("pair                                ablated   restored_from_zero")
    for left, right in pairs:
        without_pair = tuple(h for h in HEADS if h not in {left, right})
        ablated = accuracy(model, cases, without_pair)
        restored = accuracy(model, cases, (left, right))
        print(f"{left + ' + ' + right:35s} {ablated:8.3f} {restored - no_heads:18.3f}")

    print("\nInterpretation")
    print("- A_twin_1 and A_twin_2 are redundant: either one rescues copy_A; only joint ablation hurts.")
    print("- B_specialist is specialised: its single ablation hurts and it rescues performance alone.")
    print("- C_relay and D_relay are distributed: each ablation hurts, neither rescues alone, the pair does.")
    print("- silent_delimiter has a sharp attention pattern but zero write gain, so its map is not a causal role.")


if __name__ == "__main__":
    evaluate()

The script makes the counterfactual logic inspectable. The synthetic organisation is engineered deliberately, so the experiment validates the diagnostic method rather than claiming an emergent property of trained language models.

The profile is intentionally conservative. It preserves useful head-level observations without pretending the head boundary is ontologically privileged. It also yields predictions. A specialised head should fail targeted ablation, resist plausible replacement and rescue its task slice alone. A redundant set should survive single removals but fail group removal. A distributed set should show coalition rescue and path-dependent effects. When those predictions fail, revise the causal organisation before renaming the head.

sequenceDiagram
  participant Q as Question
  participant M as Mechanism
  participant E as Evidence
  participant D as Decision
  Q->>M: The heatmap trap
  M->>E: A toy system where the difference is visible
  E-->>M: compare rival explanations
  E->>D: Classify the coalition before naming the head
  D-->>Q: revise scope or act

Classify the coalition before naming the head

When you inspect, prune, steer or compare attention heads, stop asking “What does this head do?” as the first question. Ask four narrower questions: Where does it route? What does it write? What can replace it? Which coalition makes the behaviour possible?

The answers change engineering decisions. Preserve a low-replaceability specialised function even if its average activation is small. Prune a redundant set only after group ablation, rescue and shifted-task tests. Treat distributed circuits as paths or coalitions; compressing one head at a time may destroy the interface between components while preserving individually impressive maps. Treat sharp but silent heads as hypotheses about routing, not mechanisms. Treat heads whose removal improves performance as interfering on the declared slice, then test whether the effect persists elsewhere.

For interpretability, replace the static atlas with scoped role records. A record should say “routes relation R, writes feature F, is replaceable by set G, and participates in coalition S under intervention I”. This makes disagreements productive: one study may contest the route, another the write, and a third the boundary without pretending they disagree about one vague label.

For pruning, rank candidate sets, not just heads. Measure saved computation against shifted performance, calibration and recovery after faults. A head that is redundant in normal operation may be useful as a backup under distribution shift. Conversely, two individually harmless removals may cross a group threshold. Compression should preserve causal functions, rather than only high-scoring components.

For steering or editing, intervene where the target information is written or read. Changing an attention map may not change the write; amplifying a write may also affect unrelated behaviours that share the same output subspace. Coalition evidence tells you whether a local edit can work or whether the surrounding path must be changed and retested.

For safety and release decisions, low replaceability is a concentration risk. A single fragile pathway deserves fault injection and monitoring. Multiple substitute pathways reduce one failure mode but complicate attribution. Distributed circuits require recovery tests at their interfaces. The same profile that improves explanation therefore informs assurance: it reveals where failure can be contained, where compensation exists and where a local diagnostic will miss a system-level break.

The smallest useful implementation is modest. Choose one behaviour, ten candidate heads, two replacement baselines and a held-out shift. Produce a route-write-replace-coalition record for each candidate, then inspect the smallest rescuing sets. This will not reveal the whole model, but it will replace attractive speculation with a falsifiable local mechanism and expose which additional boundary must be investigated.

This decision is less satisfying than a colourful head atlas, but it is more useful. It turns an interpretability picture into an experimental claim, a pruning heuristic into a reliability test, and a semantic label into a versioned statement about a configured system. Many heads do not imply many neat relations. They create a space in which relations may separate, duplicate, interfere and compose. The job of analysis is to discover which causal organisation is actually present.