Llama-3.1-8B, cut open
Eight slabs, left to right, stand for eight bands of four layers. Rows are token positions. The heavy row at the bottom of each slab is END, where the next token is read out. The heads that write into END sit on the diagonal rising from it.
Moving dots are attention. They leave a source position, enter a head, and the head's write joins END. Dots moving right along a row are the residual stream.
Same query, six prefixes
The query never changes. Only the prefix does. Two of these supply no evidence and propose no capital, yet still lower the Berlin logit. Gold-logit shifts from Table 1:
Germany is still enriched
Subject-position MLP writes make Germany's attributes available at the Germany position. Under the irrelevant prefix that signal is effectively unchanged (App. B.1.1).
The MLPs enrich the irrelevant tokens too. “Computers” picks up its own associations (B.1.2). The two tokens differ in relevance, not in whether they carry associated information.
Prefix tokens get copied to END
Parroting heads carry prefix-token identity into END without a matching induction cue. On a multi-token prefix their attention spreads over several content words, and the write is approximately an attention-weighted mixture of those tokens' unembedding directions.
Swapping the value at one source position with attention held fixed shifts these heads' writes toward the swapped token (B.2.1). The effect is in OV, not QK routing.
The readers look away from the filler
Mid-layer readers, selected on golden context, attend to candidates of the expected answer type and assign little attention to unrelated filler (B.4).
Effective, not selective
Replace every middle-block head's END write with its context-free counterpart while the prefix stays in the input. This restores most of the gold-logit loss. Masking the same heads' END→prefix reads barely helps.
A sparse, late population
Score each head by how much its direct gold contribution falls, weighted by its attention to the prefix.
In Llama-3.1-8B the population sits mostly in late layers, and its attention to the prefix rises when one is present (Table C.1). Distractor writing is not part of the score, and it still rises.
What the mask does inside one head
Split each write into oh = Ph + Uh. P comes from prefix positions and carries most of the extra distractor content, with near-zero mean gold attribution. The gold drop sits in a weaker U. Masking END→prefix attention removes P, and softmax redistribution scales U up.
Click any orange square in the model to load it here. Fig. 5's example head has a = 0.47, so U is scaled ×1.90.
Two populations, mostly apart
Retrieval heads read the answer span in golden context and write toward it. Ranked independently, the top heads of the two populations differ, with a small overlap (Fig. 6). That separation is what makes a class-conditional mask possible.
The oracle and the deployable mask
The oracle swaps the END residual at the intrusion-band layer for its clean counterpart, so it needs a matched clean pass. The mask needs nothing at inference.
The λ slider under the model interpolates the intrusion heads' writes toward their context-free values: o(λ) = oirr + λ(onctx − oirr).
Every control is live
Click heads in the model to mask them one at a time. Combine regimes, masks, λ and the oracle patch.