causalab-mini

Use case

Attention pattern

Patching the attention pattern: does where a head looks carry the answer?

Head patching says one head of Llama-3.2-1B carries "today" to the last token: L12 H28. A head does two things there — it decides where to look, and it moves what it finds. A patch of the head's result changes both at once, so it cannot tell you which.

The attention pattern is the first half on its own. Swap it from the counterfactual prompt and the head looks where it would have looked on that prompt, while the values it mixes are still this prompt's. If "today" travels by the head's aim, the answer moves. This page runs that, per head of layer 12, on the real model.

At a glance #

Nothing moves. All 32 heads of layer 12, mean logit difference over the 42 pairs, counterfactual answer minus base answer:

headlogit difference, measured from −2.409
L12 H28 −2.396
L12 H18 −2.398
median head −2.403
L12 H15 −2.409

The bars are measured from −2.409, the lowest of the 32, because measured from zero all 32 are the same bar. The whole spread across every head of the layer is 0.013 logits, and the interchange accuracy is 0.00 on every head and every row. The mover head's pattern is not what makes it the mover head.

The document #

documents/real/weekdays_pattern_patching.json — the document this page runs, verbatim
{
  "header": {
    "description": "Is it WHERE the mover head looks? Head patching found one head, L12 H28, that carries 'today' to the last token. This swaps only that head's attention PATTERN at the last position from the counterfactual prompt \u2014 its values stay the base's \u2014 for every head of layer 12. All weekday names are one token here, so every pair is laid out identically, which the compiler checks. Needs eager attention."
  },
  "model": {
    "key": "meta-llama/Llama-3.2-1B",
    "revision": "4e20de362430cd3b72f300e6b0f18e50e7166e08",
    "dtype": "fp32",
    "attn_implementation": "eager"
  },
  "roles": {
    "base": {
      "field": "input"
    },
    "counterfactual": {
      "field": "counterfactual_inputs[0]"
    }
  },
  "sites": {
    "target": {
      "component": "attention_probs",
      "layers": [
        12
      ],
      "heads": {
        "sweep": [
          [
            0
          ],
          [
            1
          ],
          [
            2
          ],
          [
            3
          ],
          [
            4
          ],
          [
            5
          ],
          [
            6
          ],
          [
            7
          ],
          [
            8
          ],
          [
            9
          ],
          [
            10
          ],
          [
            11
          ],
          [
            12
          ],
          [
            13
          ],
          [
            14
          ],
          [
            15
          ],
          [
            16
          ],
          [
            17
          ],
          [
            18
          ],
          [
            19
          ],
          [
            20
          ],
          [
            21
          ],
          [
            22
          ],
          [
            23
          ],
          [
            24
          ],
          [
            25
          ],
          [
            26
          ],
          [
            27
          ],
          [
            28
          ],
          [
            29
          ],
          [
            30
          ],
          [
            31
          ]
        ]
      }
    },
    "head": {
      "component": "lm_head"
    }
  },
  "interventions": {
    "patching": {
      "reads": {
        "v_cf": {
          "site": "target",
          "pos": -1,
          "model": "original",
          "input": "counterfactual"
        },
        "logits": {
          "site": "head",
          "pos": -1,
          "model": "patched",
          "input": "base"
        }
      },
      "writes": {
        "patch": {
          "site": "target",
          "pos": -1,
          "mechanism": "swap",
          "operand": "v_cf"
        }
      },
      "models": {
        "patched": {
          "input": "base",
          "writes": [
            "patch"
          ]
        }
      },
      "metrics": {
        "iia": {
          "kind": "match",
          "of": "logits",
          "token_form": "space_prefixed",
          "expected": "cf_answer"
        },
        "logit_diff": {
          "kind": "logit_diff",
          "of": "logits",
          "token_form": "space_prefixed",
          "a": "cf_answer",
          "b": "base_answer"
        }
      }
    }
  },
  "steps": {
    "score": {
      "kind": "observe",
      "interventions": "patching",
      "rows": {
        "base": "weekdays_full/data",
        "counterfactual": "weekdays_full/data"
      },
      "saves": [
        {
          "value": "iia",
          "file_path": "iia.json"
        },
        {
          "value": "logit_diff",
          "file_path": "logit_diff.json"
        }
      ]
    }
  }
}

Most of its length is the head sweep: 32 single-element lists, [[0], [1], … [31]], one point per head. Apart from that and the model block, it is the layer sweep's document with a different site — read from the unmodified model on the counterfactual prompt, swap into the base run, score the head's logits against the two answers.

Two things this document says that others do not #

"attn_implementation": "eager" #

It sits in the model block, beside the checkpoint and the revision, because it is part of the experiment and not of the run. Two reasons it belongs there rather than on a command line:

  • eager is the only implementation under which attention_probs and attention_scores exist. sdpa and flash attention never materialize the pattern as a tensor — the whole point of them is that they do not — so there is nothing to address.
  • The implementations differ in the last bits. A result produced under eager is not bitwise a result produced under sdpa, and a document that did not record which one ran would not be reproducible.

Leave it out and the address does not resolve. model asks the question without running anything:

$ uv run causalab-mini model meta-llama/Llama-3.2-1B
meta-llama/Llama-3.2-1B via nnterp: 16 layers, padding_side=left
  attention_query    ok   width=2048  op=attention_interface_1
  attention_key      ok   width=512  op=attention_interface_1
  attention_scores   --   component 'attention_scores' exists only under 'eager' attention, and this model runs 'sdpa': say "attn_implementation": "eager" in the document's model block
  attention_probs    --   component 'attention_probs' exists only under 'eager' attention, and this model runs 'sdpa': say "attn_implementation": "eager" in the document's model block
  attention_z        ok   width=2048  op=attention_interface_1
…
  block_output       ok   width=2048
  ln_final           ok   width=2048
  lm_head            ok   width=128256
  logits             ok   width=128256

Every other component on the list resolves under the default. Two do not, and the refusal names the field to set.

The pattern is over keys #

attention_probs is (batch, head, query, key). The site's heads picks the head axis; pos cuts the query axis, so -1 is "what this head attends to when producing the last token"; and the key axis is left whole, because a pattern is a distribution over keys and half of one is not a pattern.

key j lands on key j

A swapped-in pattern is a vector of weights over the padded batch's key positions, and key j of the operand lands on key j here. That is only the same token if the two prompts are laid out identically: the same padded width and, row by row, the same number of real tokens. Otherwise attention mass meant for one word lands on another, or on a pad, with every shape correct and nothing to notice.

Both attention masks are in hand before any forward, so this is refused at compile time rather than discovered later — or not discovered:

$ uv run causalab-mini explain pattern_on_unaligned_rows.json
…
causalab_mini.plan.plan.PlanError: write 'look_there' swaps in the attention pattern 'their_pattern', read from prompts laid out differently: rows [0, 2] have [9, 11] tokens there and [11, 9] here. A pattern is over keys, so the two prompts must tokenize to the same length, row by row

That is why documents/v2/ has a weekdays_aligned corpus at all: on the tiny tokenizer Thursday is three tokens and Friday is one, so only some pairs line up, and the aligned corpus is the six that do. On Llama-3.2-1B every weekday name is one token, so the whole 42-pair corpus is aligned and the real document needs no subset.

What it compiles to #

$ uv run causalab-mini explain documents/real/weekdays_pattern_patching.json
<root>: Plan
  heads=0: Plan
    score: Observe  metrics=['iia/match', 'logit_diff/logit_diff']  saves=['iia.json', 'logit_diff.json']
        forward 'original' on 'counterfactual'  (42 rows x 8 tokens)
          read  'v_cf' at attention_probs[12] pos={index:-1} features=[0]/32 via 'identity'
        forward 'patched' on 'base'  (42 rows x 8 tokens)
          write 'patch' at attention_probs[12] pos={index:-1} features=[0]/32 swap(v_cf) via 'identity'
          read  'logits' at lm_head pos={index:-1} via 'identity'
  heads=1: Plan
…

features=[0]/32 is the head selection: head 0 of the 32 this layer has. A site with heads narrows the tensor to those heads, and it is the same field whether the component is a pattern, a query or a head's pre-projection result.

Running it #

$ time uv run causalab-mini run documents/real/weekdays_pattern_patching.json \
      --device-map cpu --out out/pattern_real
out/pattern_real/document.json
out/pattern_real/run.json
out/pattern_real/heads=0/document.json
out/pattern_real/heads=0/iia.json
out/pattern_real/heads=0/logit_diff.json
…
out/pattern_real/heads=31/logit_diff.json

real	1m50.005s

32 points, 64 forwards over 42 rows, under two minutes on twelve CPU cores. The recorded A100 run of the neighbouring head-patching document — three times the points — took 33 seconds (documents/real/README.md).

The result #

Mean over the 42 pairs, per head. Every one of the 32:

headlogit diffheadlogit diffheadlogit diffheadlogit diff
0−2.4038−2.40216−2.40424−2.403
1−2.4039−2.40317−2.40425−2.403
2−2.40310−2.40418−2.39826−2.403
3−2.40411−2.40419−2.40327−2.403
4−2.40312−2.40520−2.40128−2.396
5−2.40513−2.40321−2.40429−2.406
6−2.40414−2.40522−2.40330−2.400
7−2.40415−2.40923−2.40431−2.403

Median −2.403, lowest −2.409, highest −2.396, and the highest is H28. The iia.json tables are 0.0 on every row of every point: across all 1344 patched runs, the counterfactual's answer is never the top token.

For scale, the layer sweep's table has both ends of this ruler. A patch at the last token before layer 12 does nothing and reads −2.40; the same patch at layer 14 or 15 moves every row and reads +2.3 to +2.4. The pattern swap moves the number by about a quarter of one percent of the distance between them.

What the head does instead #

The contrast is the point. The same head, the same position, the same corpus, with the swap moved from its pattern to its result:

what is swappedcomponentL12 H28median headIIA at H28
where the head looksattention_probs−2.396−2.4030.00
what the head movesattention_premix+0.054−2.4030.50

The second row is documents/real/weekdays_head_patching.json — 96 points, every head of layers 10, 11 and 12 — run the same way, 8 m 09 s on CPU. One head of the 96 stands out, and it is this one: swapping its result moves the mean logit difference from −2.403 to +0.054 and flips half the 42 answers. The same head's pattern, from the same prompts, moves nothing. documents/real/README.md records +0.05 and "flips half the answers" from the A100 run.

So L12 H28 attends the same way whatever the day is. Its aim is not a function of the answer, and importing the counterfactual's aim imports nothing. What it carries is in the values it mixes — the head is a pipe whose position is fixed and whose contents are the fact.

what a null result here does not say

"The pattern carries nothing" is a claim about this task, this position and this layer. The prompts differ in exactly one word at a fixed offset, so the two patterns being interchangeable is unsurprising in hindsight — a task whose counterfactual moves a word, or changes the number of items to attend to, would not look like this.

It is also a claim about a swap at the last query position only. A head whose pattern matters earlier in the sequence would not show up: nothing was swapped at any other query.

And a pattern swap is not a knockout. Zeroing a head's pattern asks a different question — whether the head contributes at all — and documents/v2/attention_knockout.json is that document, with "operand": 0.0 in place of a read.

The tiny version #

The same experiment at a size that runs anywhere, for checking a change to the library rather than learning about a model:

documents/v2/attention_pattern_patching.json — four heads, six aligned pairs, verbatim
{
  "header": {
    "description": "Attention-pattern patching: make one head attend, at the last position, the way it did on the COUNTERFACTUAL prompt, and leave its values alone \u2014 so the question is 'does where this head looks carry the answer', separately from what it moves. `attention_probs` needs eager attention, and its last axis is the keys of the padded batch, so the two prompts must tokenize to the same length row by row; the compiler checks that before any forward (documents/data/weekdays_aligned is the subset of pairs that do, under this tokenizer)."
  },
  "model": {
    "key": "hf-internal-testing/tiny-random-LlamaForCausalLM",
    "revision": "9fb191250dd56d0ba7ec9785a025ed29c03d5998",
    "dtype": "fp32",
    "attn_implementation": "eager"
  },
  "roles": {
    "base": {
      "field": "input"
    },
    "counterfactual": {
      "field": "counterfactual_inputs[0]"
    }
  },
  "sites": {
    "one_head": {
      "component": "attention_probs",
      "layers": [
        0
      ],
      "heads": {
        "sweep": [
          [
            0
          ],
          [
            1
          ],
          [
            2
          ],
          [
            3
          ]
        ]
      }
    },
    "head": {
      "component": "lm_head"
    }
  },
  "interventions": {
    "pattern_patching": {
      "reads": {
        "their_pattern": {
          "site": "one_head",
          "pos": -1,
          "model": "original",
          "input": "counterfactual"
        },
        "logits": {
          "site": "head",
          "pos": -1,
          "model": "patched",
          "input": "base"
        }
      },
      "writes": {
        "look_there": {
          "site": "one_head",
          "pos": -1,
          "mechanism": "swap",
          "operand": "their_pattern"
        }
      },
      "models": {
        "patched": {
          "input": "base",
          "writes": [
            "look_there"
          ]
        }
      },
      "metrics": {
        "iia": {
          "kind": "match",
          "of": "logits",
          "expected": "cf_answer",
          "token_form": "space_prefixed"
        },
        "logit_diff": {
          "kind": "logit_diff",
          "of": "logits",
          "a": "cf_answer",
          "b": "base_answer",
          "token_form": "space_prefixed"
        }
      }
    }
  },
  "steps": {
    "score": {
      "kind": "observe",
      "interventions": "pattern_patching",
      "rows": {
        "base": "weekdays_aligned/data",
        "counterfactual": "weekdays_aligned/data"
      },
      "saves": [
        {
          "value": "iia",
          "file_path": "iia.json"
        },
        {
          "value": "logit_diff",
          "file_path": "logit_diff.json"
        }
      ]
    }
  }
}
$ time uv run causalab-mini run documents/v2/attention_pattern_patching.json --out out/pattern
out/pattern/document.json
out/pattern/run.json
out/pattern/heads=0/document.json
out/pattern/heads=0/iia.json
out/pattern/heads=0/logit_diff.json
…
out/pattern/heads=3/logit_diff.json

real	0m6.481s

Six and a half seconds. Its numbers are meaningless — the model has random weights — but the addresses, the head selection and the layout check are the same code.

See also #

  • Locations — the whole component vocabulary, which components need eager attention, and what heads means at each of them
  • Layer sweep — how layer 12 was chosen, and the crossover this head is the mechanism of
  • Sweeps — the head sweep, and why 32 single-element lists rather than a range
  • Reads, writes, metrics — swap with a literal operand, which is what a knockout is