Info
Last Execution: 2026-08-19
| Package | Version |
|---|---|
| nnsight | 0.8 |
| Python | 3.12.13 |
| torch | 2.13.0+cu126 |
| transformers | 5.x |
Language Models Use Lookbacks to Track Beliefs¶
Introduction¶
🧠 If Bob fills a bottle with beer and then Carla, on the other side of the kitchen, fills a cup with coffee without either of them watching the other, then what is in the cup and what Bob believes is in the cup are different questions with different answers. Answering the second one requires the model to keep a representation of the world indexed by whose head it is in — a Theory-of-Mind problem, and one that a language model has to solve with attention and residual streams rather than with a database.
Prakash et al. (2026), Language Models use Lookbacks to Track Beliefs, ask how that representation is built. Their answer is a reusable circuit motif they call a lookback. A piece of reference information is computed once at a source token and copied to two places: an address that sits with the information it labels, and a pointer that sits at the token where the model will later need it. When the model needs the information, it dereferences the pointer — attends from the pointer back to the matching address — and brings the payload stored beside that address forward. The paper finds three of these chained together in belief tracking: a binding lookback that ties each character to the object and state they acted on, an answer lookback that fetches the state token once the right state has been identified, and a visibility lookback that lets one character's beliefs be updated from another's when the story says they can observe each other.
The models. The paper studies Llama-3-70B-Instruct and Llama-3.1-405B-Instruct, and says explicitly that it does not examine smaller ones, "as they are unable to coherently solve the CausalToM task." We do not have either of those here — no 70B weights, no remote execution — so this notebook cannot reproduce the paper's setting, and nothing below should be read as having done so.
What we can do is ask the question the paper left open: does the lookback mechanism appear
in a model small enough to run on one GPU? We use meta-llama/Llama-3.1-8B-Instruct —
about 16 GB in bfloat16, one A6000 — and treat "the mechanism is absent" and "the model
cannot do the task" as two different outcomes that have to be told apart. The behavioural
baseline in §2 is what tells them apart, and it comes before any internals.
What we find. The 8B model solves the no-visibility task at 0.745 against the 70B's 0.95, and fails the visibility condition outright. In the setting it can do, the binding lookback and the answer lookback are both present and measurable: an Ordering ID computed at the character and object tokens by layer 10, carried as a rank-2 address in the drink tokens' residual streams and swappable there with IIA 0.775 at layer 15, dereferenced from a pointer at the final token in layers 12–22, and resolved into the answer token in layers 23–31 — against 0.000 for every control we ran. The visibility lookback we could not reach, and we say why rather than reporting a number for it.
Setup¶
If using Colab, install NNsight:
!pip install -U nnsight
try:
import google.colab
is_colab = True
except ImportError:
is_colab = False
if is_colab:
!pip install -U nnsight
import json
import random
from collections import Counter
from IPython.display import clear_output
import torch
import numpy as np
import nnsight
from nnsight import TransformersModel
import matplotlib.pyplot as plt
import plotly.express as px
import plotly.graph_objects as go
import plotly.io as pio
pio.renderers.default = "colab" if is_colab else "plotly_mimetype+notebook_connected+colab+notebook"
/home/localjadenfk/miniconda3/envs/ndif2/lib/python3.12/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html from .autonotebook import tqdm as notebook_tqdm
Llama-3.1-8B-Instruct in bfloat16 needs roughly 16 GB of GPU memory. Every
measurement in this notebook is a forward pass over a ~180-token prompt, batched, so the
whole notebook runs in about 40 minutes on one A6000.
We load with an explicit task= and dispatch immediately, so the weights are on the GPU
before the first trace.
model = TransformersModel(
"meta-llama/Llama-3.1-8B-Instruct",
task="text-generation",
dtype=torch.bfloat16,
device_map="auto",
dispatch=True,
)
clear_output()
tok = model.tokenizer
layers = model.model.layers
N_LAYERS = model.config.num_hidden_layers
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
DTYPE = torch.bfloat16
print(f"{N_LAYERS} layers, d_model={model.config.hidden_size}")
32 layers, d_model=4096
1. CausalToM: the task¶
The paper's dataset is called CausalToM, and it is deliberately dull. Two characters each act on a distinct object, in two sentences of identical shape, and a question asks what one of them believes about one of the objects. Everything that could vary — who, what, which — varies over a fixed pool of entities, so the only thing that distinguishes one sample from another is the assignment. That is what makes it usable for causal analysis: two prompts can be made to differ in exactly one variable.
The story template, following Prakash et al.:
{char_1}and{char_2}are working in a busy restaurant. To complete an order,{char_1}grabs an opaque{obj_1}and fills it with{state_1}. Then{char_2}grabs another opaque{obj_2}and fills it with{state_2}.
The containers are opaque and the two actions are independent, so in this
no-visibility condition each character knows the contents of their own container and
has no belief at all about the other's. The question is
"What does {char} believe the {obj} contains?", and the answer is the drink if the
queried character is the one who filled that container, and unknown otherwise. Half
the questions are of each kind, which is what stops "just read out the contents" from
being a winning strategy.
The paper also uses variants that append explicit visibility sentences — "{char_2}
cannot observe {char_1}'s actions", or can observe — which is the substrate for the
visibility lookback. We build those too and use them in §2.
CHARACTERS = ["Dean", "Beth", "Jake", "Josh", "Karen", "Carl", "Lee", "Pam", "Donna",
"Frank", "Diane", "Bob", "Ellen", "Ivy", "Pete", "Anne", "Greg", "Nina",
"Mark", "Judy", "Paul", "Rose", "Sam", "Tara", "Phil", "Olivia", "Jean",
"Joe", "Kyle", "Ruth"]
OBJECTS = ["jar", "cup", "mug", "glass", "flute", "pitcher", "jug", "bottle", "can",
"flask", "pint", "quart", "horn", "tun", "urn", "vat", "tank", "drum"]
STATES = ["water", "milk", "tea", "beer", "soda", "juice", "coffee", "wine", "gin",
"rum", "champagne", "cocktail", "punch", "espresso", "cocoa", "sprite",
"bourbon", "port", "stout", "ale", "porter"]
INSTRUCTION = (
"1. Track the belief of each character as described in the story. "
"2. A character's belief is formed only when they perform an action themselves or can "
"observe the action taking place. "
"3. A character does not have any beliefs about the container and its contents which "
"they cannot observe. "
"4. To answer the question, predict only what is inside the queried container, strictly "
"based on the belief of the character, mentioned in the question. "
"5. If the queried character has no belief about the container in question, then predict "
"'unknown'. "
"6. Do not predict container or character as the final output."
)
⚠️ Every entity above is a single token under the Llama-3 tokenizer when preceded by a space. That is not cosmetic. All of the interventions below copy activations between two prompts at fixed token positions, and that is only meaningful if the two prompts tokenize to the same length with the same layout. Filtering the entity pool to single tokens buys that for free; we assert it rather than trust it.
for pool in (CHARACTERS, OBJECTS, STATES):
for w in pool:
assert len(tok.encode(" " + w, add_special_tokens=False)) == 1, w
print(f"{len(CHARACTERS)} characters, {len(OBJECTS)} containers, {len(STATES)} drinks, "
f"all single-token")
30 characters, 18 containers, 21 drinks, all single-token
def story(c, o, s, visibility=None):
"""visibility: None = no visibility sentences (the paper's no-visibility condition);
'none' = explicit mutual non-observation; 0/1 = that character can observe the other."""
t = (f"{c[0]} and {c[1]} are working in a busy restaurant. To complete an order, "
f"{c[0]} grabs an opaque {o[0]} and fills it with {s[0]}. "
f"Then {c[1]} grabs another opaque {o[1]} and fills it with {s[1]}.")
if visibility == "none":
t += f" {c[1]} cannot observe {c[0]}'s actions. {c[0]} cannot observe {c[1]}'s actions."
elif visibility == 0:
t += f" {c[1]} cannot observe {c[0]}'s actions. {c[0]} can observe {c[1]}'s actions."
elif visibility == 1:
t += f" {c[0]} cannot observe {c[1]}'s actions. {c[1]} can observe {c[0]}'s actions."
return t
def belief(c_idx, o_idx, s, visibility=None):
"""Ground-truth belief of character c_idx about container o_idx."""
if c_idx == o_idx: # they filled it themselves
return s[o_idx]
if visibility == c_idx: # they can observe the other character
return s[o_idx]
return "unknown"
def prompt(c, o, s, c_idx, o_idx, visibility=None, ask_reality=False):
q = (f"What does the {o[o_idx]} contain?" if ask_reality
else f"What does {c[c_idx]} believe the {o[o_idx]} contains?")
p = (f"Instruction: {INSTRUCTION}\n\n"
f"Story: {story(c, o, s, visibility)}\n"
f"Question: {q}\n"
f"Answer:")
ans = s[o_idx] if ask_reality else belief(c_idx, o_idx, s, visibility)
return p, ans
def sample_entities(rng):
return rng.sample(CHARACTERS, 2), rng.sample(OBJECTS, 2), rng.sample(STATES, 2)
A worked example. The story is the same in both cells below; only the question changes.
rng = random.Random(0)
c, o, s = sample_entities(rng)
p_own, a_own = prompt(c, o, s, 0, 0) # ask char 0 about char 0's container
print(p_own)
print(f"\n>>> correct answer: {a_own!r}")
Instruction: 1. Track the belief of each character as described in the story. 2. A character's belief is formed only when they perform an action themselves or can observe the action taking place. 3. A character does not have any beliefs about the container and its contents which they cannot observe. 4. To answer the question, predict only what is inside the queried container, strictly based on the belief of the character, mentioned in the question. 5. If the queried character has no belief about the container in question, then predict 'unknown'. 6. Do not predict container or character as the final output. Story: Joe and Ellen are working in a busy restaurant. To complete an order, Joe grabs an opaque tun and fills it with gin. Then Ellen grabs another opaque cup and fills it with bourbon. Question: What does Joe believe the tun contains? Answer: >>> correct answer: 'gin'
p_other, a_other = prompt(c, o, s, 0, 1) # ask char 0 about char 1's container
p_real, a_real = prompt(c, o, s, 0, 1, ask_reality=True)
print(p_other.split("Question:")[1].strip(), "->", repr(a_other))
print(p_real.split("Question:")[1].strip(), "->", repr(a_real))
What does Joe believe the cup contains? Answer: -> 'unknown' What does the cup contain? Answer: -> 'bourbon'
The last pair is the control that makes the whole thing legible. Same story, same
container, and the two correct answers are different: reality is the drink, belief is
unknown. A model that answers both the same way is not tracking beliefs at all,
whatever its accuracy on the belief question alone happens to be.
2. The behavioural baseline¶
Before any internals. If Llama-3.1-8B-Instruct cannot do this task, then any null
result from an intervention tells us nothing — we would be probing a mechanism that is
not there to be probed. So we measure four things on 200 samples each:
- belief accuracy in the no-visibility condition, split by whether the queried character filled the container themselves;
- reality accuracy — the same stories, asking what is actually in the container;
- belief accuracy with the explicit "cannot observe" sentences appended;
- belief accuracy in the visibility condition, where one character can observe the other, scored separately for the observer and the non-observer.
The model is scored on its top-1 next token after Answer:, which is how the paper's
own evaluation script scores it. All drinks and unknown are single tokens, so this is
exact rather than a prefix match.
tok.padding_side = "left"
if tok.pad_token is None:
tok.pad_token = tok.eos_token
@torch.no_grad()
def predict(prompts, bs=16):
"""Greedy next token after 'Answer:', batched."""
out = []
for i in range(0, len(prompts), bs):
with model.trace(prompts[i:i + bs]):
p = model.lm_head.output[:, -1].argmax(-1).save()
out += [tok.decode([t]).strip().lower() for t in p]
return out
We wrap every collection pass in torch.no_grad(). nnsight traces run with autograd on
by default, which for read-only work costs memory and time we do not need.
def build_eval(n, visibility=None, ask_reality=False, seed=0):
rng, rows = random.Random(seed), []
for _ in range(n):
c, o, s = sample_entities(rng)
ci, oi = rng.choice([0, 1]), rng.choice([0, 1])
p, a = prompt(c, o, s, ci, oi, visibility, ask_reality)
rows.append(dict(prompt=p, ans=a, own=(ci == oi)))
return rows
def report(name, rows):
preds = predict([r["prompt"] for r in rows])
ok = [p == r["ans"] for p, r in zip(preds, rows)]
own = [k for k, r in zip(ok, rows) if r["own"]]
oth = [k for k, r in zip(ok, rows) if not r["own"]]
print(f"{name:<28} overall {sum(ok)/len(ok):.3f} "
f"own container {sum(own)/len(own):.3f} (n={len(own)}) "
f"other container {sum(oth)/len(oth):.3f} (n={len(oth)})")
return sum(ok) / len(ok), preds
acc_belief, preds_belief = report("no-visibility, belief", build_eval(200, seed=0))
acc_reality, preds_reality = report("no-visibility, reality", build_eval(200, seed=0, ask_reality=True))
acc_explicit, _ = report("explicit 'cannot observe'", build_eval(200, seed=0, visibility="none"))
[transformers] Ignoring clean_up_tokenization_spaces=True for BPE tokenizer TokenizersBackend. The clean_up_tokenization post-processing step is designed for WordPiece tokenizers and is destructive for BPE (it strips spaces before punctuation). Set clean_up_tokenization_spaces=False to suppress this warning, or set clean_up_tokenization_spaces_for_bpe_even_though_it_will_corrupt_output=True to force cleanup anyway.
no-visibility, belief overall 0.745 own container 0.856 (n=104) other container 0.625 (n=96)
no-visibility, reality overall 0.940 own container 0.942 (n=104) other container 0.938 (n=96)
explicit 'cannot observe' overall 0.565 own container 0.625 (n=104) other container 0.500 (n=96)
Two baselines make these numbers readable.
Chance. The model's plausible outputs are the two drinks in the story and the word
unknown; guessing among them is 33%.
The reality strategy. A model that ignores the belief framing entirely and always reports the actual contents of the queried container scores 100% on the "own container" half and 0% on the other half — 50% overall. This is the number that matters, because it is what a model with no belief representation but perfect reading comprehension would get.
print(f"belief accuracy {acc_belief:.3f}")
print(f"reality-strategy ceiling 0.500 (always answer the container's actual contents)")
print(f"chance 0.333")
print()
print("prediction distribution, belief question:", Counter(preds_belief).most_common(5))
belief accuracy 0.745
reality-strategy ceiling 0.500 (always answer the container's actual contents)
chance 0.333
prediction distribution, belief question: [('unknown', 61), ('porter', 11), ('sprite', 10), ('coffee', 9), ('milk', 9)]
The visibility condition. Same story, plus one sentence saying that character A can
observe character B. Now A's belief about B's container is B's drink, while B's belief
about A's container is still unknown. We score the observer and the non-observer
separately, and also paired — both questions right for the same story — which is how
the paper's evaluation script scores this condition.
rng = random.Random(1)
obs_rows, non_rows = [], []
for _ in range(200):
c, o, s = sample_entities(rng)
obs = rng.choice([0, 1]) # this one can observe
obs_rows.append(prompt(c, o, s, obs, 1 - obs, visibility=obs)) # ask the observer
non_rows.append(prompt(c, o, s, 1 - obs, obs, visibility=obs)) # ask the non-observer
p_obs = predict([p for p, a in obs_rows])
p_non = predict([p for p, a in non_rows])
ok_obs = [p == a for p, (_, a) in zip(p_obs, obs_rows)]
ok_non = [p == a for p, (_, a) in zip(p_non, non_rows)]
print(f"visibility, observer {sum(ok_obs)/200:.3f} (correct answer: the other's drink)")
print(f"visibility, non-observer {sum(ok_non)/200:.3f} (correct answer: 'unknown')")
print(f"visibility, paired {sum(a and b for a, b in zip(ok_obs, ok_non))/200:.3f}")
visibility, observer 0.535 (correct answer: the other's drink) visibility, non-observer 0.515 (correct answer: 'unknown') visibility, paired 0.295
Finding 1 — the model can do the no-visibility task, and cannot do the visibility one¶
Three things to take from those numbers.
The belief representation is real but weak. Belief accuracy is 0.745, above the reality-strategy ceiling of 0.50 — a model that read the story perfectly and ignored the belief framing entirely would score exactly 0.50. So it is doing something with the belief question. But it is far below the 0.95 Prakash et al. report for Llama-3-70B-Instruct on the same task, and it matches the authors' own released evaluation of this checkpoint (0.72 ± 0.04 over ten runs of 100 samples). The paper's decision not to study models this size is defensible; the mechanism question is nevertheless still askable, because the accuracy is not at floor.
The failure is asymmetric, and it is the unknown half that fails. 0.856 when the
queried character filled the container themselves, 0.625 when the correct answer is
unknown. That is the shape you would expect from a model with a working binding mechanism
and a weak ignorance mechanism: it can find the drink bound to a character, and it is
worse at concluding that no such binding exists.
Reality accuracy is 0.940, so the failures are not reading-comprehension failures. Same stories, same containers: asking what is in the container is answered 20 points better than asking what someone believes is in it. Belief and reality are answered differently, which is the minimum evidence that there is a belief computation to look for.
Adding the explicit "cannot observe" sentences makes things worse (0.565), not better. The extra clause is redundant with the no-visibility setting and the model is distracted by it — a reminder that these templates are load-bearing and that "more explicit" is not automatically "easier".
The visibility condition is a different story. Paired accuracy — both the observer's and the non-observer's question right for the same story — is 0.295, and each question separately is near the 0.5 you would get by choosing between "the drink" and "unknown" at random. Any intervention on the visibility lookback would be measuring noise rather than mechanism. We report it and stop there. The visibility lookback is out of reach at this scale, and saying so is the honest version of the result.
So: not "the model cannot do the task", and not "the model does the task". It does the no-visibility task 25 points above the strategy baseline and fails the visibility one, which is exactly the situation in which asking about mechanism is worth doing and worth qualifying.
Three things to take from those numbers.
The belief representation is real but weak. Belief accuracy sits above the reality-strategy ceiling of 0.50 — a model that read the story perfectly and ignored the belief framing entirely would score exactly 0.50, and ours scores meaningfully more than that. So it is doing something with the belief question. But it is far below the ~0.95 that Prakash et al. report for Llama-3-70B-Instruct on the same task, and the authors' own released evaluation of this checkpoint (0.72 ± 0.04 over ten runs of 100 samples) agrees with what we measure. The paper's decision not to study models this size is defensible; the mechanism question is nevertheless still askable, because the accuracy is not at floor.
The failure is asymmetric, and it is the unknown half that fails. The model is
reliable when the queried character filled the container themselves and much less reliable
when the correct answer is unknown. That is the shape you would expect from a model with
a working binding mechanism and a weak ignorance mechanism: it can find the drink bound
to a character, and it is worse at concluding that no such binding exists.
Reality accuracy is high, so the failures are not reading-comprehension failures. Same stories, same containers, and simply asking what is in the container is answered far more accurately than asking what someone believes is in it. Belief and reality are answered differently, which is the minimum evidence that there is a belief computation to look for.
The visibility condition is a different story. Paired accuracy — both the observer's and the non-observer's question right for the same story — is low enough that any intervention on the visibility lookback would be measuring noise rather than mechanism. We report it and stop there. The visibility lookback is out of reach at this scale, and saying so is the honest version of the result.
So: not "the model cannot do the task", and not "the model does the task". It does the no-visibility task above the strategy baseline and fails the visibility one, which is exactly the situation in which asking about mechanism is worth doing and worth qualifying.
3. What a lookback is¶
Before measuring anything, the abstraction, because it is easy to lose in prose.
A lookback has four parts:
| part | where it lives | what it is |
|---|---|---|
| source | an early token | the token(s) where the reference information is first computed |
| address | the residual stream of the token to be recalled | a copy of the reference information, sitting next to the thing it labels |
| pointer | the residual stream of the token doing the recalling | the other copy of the same reference information |
| payload | alongside the address | the information that will be retrieved |
The mechanism is: attention copies the source's reference information into two different residual streams — one becomes the address, the other the pointer. Later, an attention head at the pointer token queries against the addresses; the head whose address matches wins, and its payload is written into the pointer token's residual stream. The model has performed a dereference.
In CausalToM the reference information is an Ordering ID (OI): a representation of whether this entity was mentioned first or second. That is the crucial claim, and it is what makes the binding lookback testable. If the model bound Bob to beer by literal co-location — "read the drink in the same sentence as the name" — then moving activations between token positions would move the answer with the positions. If instead it binds through an OI carried in the residual stream, then moving the OI moves the answer even when the token positions and the drink identities at those positions are unchanged.
The three lookbacks the paper identifies chain like this:
- Binding lookback. Source: the character and object tokens of each sentence, whose OIs are computed from their order. Address: a copy of those OIs in the state (drink) token's residual stream. Payload: the state's own OI. Pointer: the query character and object tokens at the end of the prompt. Dereferencing it yields which state is the answer, as an OI — not yet the word.
- Answer lookback. Source: the state tokens. Address: the state OI at each state token. Pointer: the state OI that the binding lookback just deposited at the final token. Payload: the state token itself. Dereferencing it yields the word.
- Visibility lookback. Only in the visibility condition: a visibility ID generated at the visibility sentence lets the observing character's query reach the observed character's binding, updating their beliefs.
def draw_lookback():
fig, ax = plt.subplots(figsize=(11.5, 5.0), facecolor="white")
ax.set_facecolor("white")
ax.set_xlim(0, 100); ax.set_ylim(0, 62); ax.axis("off")
def box(x, y, w, h, label, fc):
ax.add_patch(plt.Rectangle((x, y), w, h, facecolor=fc, edgecolor="#333333",
lw=1.2, zorder=2))
ax.text(x + w / 2, y + h / 2, label, ha="center", va="center",
fontsize=9.5, color="#111111", zorder=3)
def arrow(p, q, color, rad, lw=1.7, ls="-"):
ax.annotate("", xy=q, xytext=p, zorder=4,
arrowprops=dict(arrowstyle="-|>", color=color, lw=lw, ls=ls,
shrinkA=2, shrinkB=2,
connectionstyle=f"arc3,rad={rad}"))
SRC, ADDR, PTR, OTH = "#f6d6a8", "#a8cfe8", "#c9b3e6", "#ededed"
y, h = 30, 7
toks = [("Bob", 2, 9, SRC), ("bottle", 12, 11, SRC), ("beer", 24, 9, ADDR),
("Carla", 35, 10, SRC), ("cup", 46, 8, SRC), ("coffee", 55, 11, ADDR),
("Bob", 69, 9, PTR), ("bottle", 79, 11, PTR), ("Answer:", 91, 8, OTH)]
for label, x, w, fc in toks:
box(x, y, w, h, label, fc)
ax.text(17, y - 3.4, "sentence 1 (OI = 1)", fontsize=8.5, color="#555555", ha="center")
ax.text(52, y - 3.4, "sentence 2 (OI = 2)", fontsize=8.5, color="#555555", ha="center")
ax.text(84, y - 3.4, "question", fontsize=8.5, color="#555555", ha="center")
# source -> address (one per sentence)
arrow((17.5, y + h), (28.5, y + h), "#c8781a", -0.6)
arrow((50, y + h), (60.5, y + h), "#c8781a", -0.6)
ax.text(23, y + h + 6.0, "OI copied → $address$", fontsize=8.5, color="#c8781a",
ha="center", va="bottom")
# source -> pointer
arrow((6.5, y + h), (73.5, y + h), "#7a52ad", -0.26)
ax.text(40, y + h + 15.0, "OI copied → $pointer$", fontsize=8.5, color="#7a52ad",
ha="center", va="bottom")
def elbow(x0, x1, ylev, color, ls="-"):
ax.plot([x0, x0], [y, ylev], color=color, ls=ls, lw=2.0, zorder=4)
ax.plot([x0, x1], [ylev, ylev], color=color, ls=ls, lw=2.0, zorder=4)
ax.annotate("", xy=(x1, y), xytext=(x1, ylev), zorder=4,
arrowprops=dict(arrowstyle="-|>", color=color, lw=2.0, ls=ls))
# dereference
elbow(73.5, 28.5, 19, "#1a6ea8")
ax.text(51, 15.0, "$dereference$: attend from pointer to the matching address",
fontsize=8.5, color="#1a6ea8", ha="center", va="center")
# payload forward
elbow(28.5, 95, 8, "#1a6ea8", ls="--")
ax.text(58, 4.0, '$payload$ carried forward to the final token → "beer"',
fontsize=8.5, color="#1a6ea8", ha="center", va="center")
handles = [plt.Rectangle((0, 0), 1, 1, fc=c, ec="#333333") for c in (SRC, ADDR, PTR)]
ax.legend(handles, ["source (OI computed here)", "address + payload", "pointer"],
loc="upper right", fontsize=8.5, frameon=False,
bbox_to_anchor=(1.02, 1.03))
ax.set_title("The binding lookback in CausalToM", fontsize=11.5, color="#111111", pad=0)
plt.tight_layout()
plt.show()
draw_lookback()
The dashed blue arrow is a simplification: in the paper the binding lookback's payload is the state OI, and a second lookback (the answer lookback, §6) turns that OI into the word. We measure them separately below.
4. The binding lookback, measured¶
The paper's test for the binding lookback is an interchange intervention: run the model on a clean prompt, run it on a counterfactual prompt, copy some activations from the second into the first, and see whether the answer becomes the one the counterfactual's binding implies. The metric is interchange intervention accuracy (IIA) — the fraction of samples where the model's top-1 answer equals what the hypothesised causal model predicts after the swap.
The counterfactual has to differ in exactly one thing, and here the construction is elegant. Take a clean story
Bob grabs an opaque bottle and fills it with beer. Then Carla grabs another opaque cup and fills it with coffee.
and build the counterfactual by reversing the order of the two sentences, carrying characters, containers and drinks along with them:
Carla grabs an opaque cup and fills it with coffee. Then Bob grabs another opaque bottle and fills it with beer.
Every binding is identical. Bob still filled the bottle with beer. The only thing that changed is which sentence slot each triple occupies — the Ordering ID. And because every entity is one token and the two sentences have the same shape, the two prompts line up token for token.
Now look at the drink tokens. In the clean prompt, position 155 is beer and position
167 is coffee. In the counterfactual, position 155 is coffee and position 167 is
beer. So if we copy the counterfactual's drink-token residual from 167 into the
clean run's 155, and from 155 into 167, then:
- the drink at each position is unchanged —
beeris still at 155,coffeestill at 167; - only the ordering information carried in those residual streams has swapped.
That makes the test sharp. Two hypotheses give different predictions:
| hypothesis | prediction after the reversed swap |
|---|---|
| the model reads the drink from the same sentence as the queried character (positional) | answer unchanged — beer |
| the model retrieves by an ordering ID carried in the residual stream (lookback) | answer flips to coffee |
We ask both questions of the same 100-odd samples.
STATE_POS = [155, 156, 167, 168] # the two drink tokens and the '.' after each
STATE_POS_REV = [167, 168, 155, 156]
OPAQUE_POS = [149, 161] # the two 'opaque' tokens - a matched-size control
OPAQUE_REV = [161, 149]
CHAR_POS = [131, 133, 146, 158] # both names in sentence 0, then each subject
OBJ_POS = [150, 162]
⚠️ Those are literal token indices, which is only safe because every prompt in this notebook tokenizes to the same length with the same layout. We verify that rather than assume it, and print the tokens at the positions we are about to intervene on.
def build_binding(n, seed=7):
"""Clean vs counterfactual: same bindings, reversed sentence order."""
rng, rows = random.Random(seed), []
for _ in range(n):
c, o, s = sample_entities(rng)
r = rng.choice([0, 1]) # which character/container is queried
clean_p, clean_a = prompt(c, o, s, r, r)
cc, oo, ss = c[::-1], o[::-1], s[::-1] # reverse the two sentences
cf_p, cf_a = prompt(cc, oo, ss, 1 - r, 1 - r)
rows.append(dict(clean=clean_p, clean_ans=clean_a,
cf=cf_p, cf_ans=cf_a, target=s[1 - r]))
return rows
rows = build_binding(120)
lens = {len(tok.encode(r["clean"])) for r in rows} | {len(tok.encode(r["cf"])) for r in rows}
assert len(lens) == 1, lens
print(f"every prompt is {lens.pop()} tokens\n")
ids = tok.encode(rows[0]["clean"])
cf_ids = tok.encode(rows[0]["cf"])
for p in CHAR_POS + OBJ_POS + STATE_POS + OPAQUE_POS:
print(f" {p:>4} clean {tok.decode([ids[p]])!r:<12} counterfactual {tok.decode([cf_ids[p]])!r}")
every prompt is 181 tokens 131 clean ' Diane' counterfactual ' Karen' 133 clean ' Karen' counterfactual ' Diane' 146 clean ' Diane' counterfactual ' Karen' 158 clean ' Karen' counterfactual ' Diane' 150 clean ' horn' counterfactual ' cup' 162 clean ' cup' counterfactual ' horn' 155 clean ' tea' counterfactual ' port' 156 clean '.' counterfactual '.' 167 clean ' port' counterfactual ' tea' 168 clean '.\n' counterfactual '.\n' 149 clean ' opaque' counterfactual ' opaque' 161 clean ' opaque' counterfactual ' opaque'
print("CLEAN")
print(rows[0]["clean"].split("Story:")[1].strip())
print(f" -> {rows[0]['clean_ans']!r}\n")
print("COUNTERFACTUAL")
print(rows[0]["cf"].split("Story:")[1].strip())
print(f" -> {rows[0]['cf_ans']!r}\n")
print(f"TARGET after the reversed swap: {rows[0]['target']!r}")
CLEAN Diane and Karen are working in a busy restaurant. To complete an order, Diane grabs an opaque horn and fills it with tea. Then Karen grabs another opaque cup and fills it with port. Question: What does Diane believe the horn contains? Answer: -> 'tea' COUNTERFACTUAL Karen and Diane are working in a busy restaurant. To complete an order, Karen grabs an opaque cup and fills it with port. Then Diane grabs another opaque horn and fills it with tea. Question: What does Diane believe the horn contains? Answer: -> 'tea' TARGET after the reversed swap: 'port'
Error filtering. An intervention on a sample the model already gets wrong is uninterpretable, so — as the paper does — we keep only samples where the model answers both the clean and the counterfactual prompt correctly with no intervention.
pred_clean = predict([r["clean"] for r in rows])
pred_cf = predict([r["cf"] for r in rows])
keep = [r for r, a, b in zip(rows, pred_clean, pred_cf)
if a == r["clean_ans"] and b == r["cf_ans"]]
print(f"kept {len(keep)}/{len(rows)} samples answered correctly on both prompts")
kept 102/120 samples answered correctly on both prompts
The intervention. Two tracer.invoke blocks inside one model.trace: the first runs
the counterfactual and reads the residual stream at layers[l], the second runs the clean
prompt and writes those values in at the destination positions. Batching the two runs into
a single forward is the main win over hooks here — one pass, not two, and no bookkeeping
to keep the two aligned.
⚠️ Two nnsight details. A transformers-5 decoder block's .output is a bare tensor, not
a tuple, so it is indexed directly — no [0]. And a name bound inside one invoke is only
visible to a later one after that block has run; invokes resume in model-reached order,
not definition order, so the reader and the writer touching the same module need a
tracer.barrier between them.
@torch.no_grad()
def interchange(rows, layer, src_pos, dst_pos, bs=8, src_key="cf"):
"""Copy layer-`layer` residuals at `src_pos` of one run into `dst_pos` of the clean run."""
out = []
for i in range(0, len(rows), bs):
chunk = rows[i:i + bs]
box = {}
with model.trace() as tracer:
barrier = tracer.barrier(2)
with tracer.invoke([r[src_key] for r in chunk]):
box["acts"] = layers[layer].output[:, src_pos]
barrier()
with tracer.invoke([r["clean"] for r in chunk]):
barrier()
layers[layer].output[:, dst_pos] = box["acts"]
p = model.lm_head.output[:, -1].argmax(-1).save()
out += [tok.decode([t]).strip().lower() for t in p]
return out
def iia(preds, rows):
"""(IIA against the swapped-binding target, rate of the unchanged clean answer)."""
return (sum(p == r["target"] for p, r in zip(preds, rows)) / len(rows),
sum(p == r["clean_ans"] for p, r in zip(preds, rows)) / len(rows))
Anchor first: patch-onto-itself must be an exact no-op. Before believing any effect, we run the identical machinery with the clean prompt as both source and destination and the identity permutation of positions. If the plumbing is right, the logits come back bit-identical to an untouched forward pass. If they do not, everything downstream is measuring our own bug.
@torch.no_grad()
def selfpatch_logits(rows, layer, bs=8):
got, ref = [], []
for i in range(0, len(rows), bs):
chunk = rows[i:i + bs]
box = {}
with model.trace() as tracer:
barrier = tracer.barrier(2)
with tracer.invoke([r["clean"] for r in chunk]):
box["acts"] = layers[layer].output[:, STATE_POS]
barrier()
base = model.lm_head.output[:, -1].save()
with tracer.invoke([r["clean"] for r in chunk]):
barrier()
layers[layer].output[:, STATE_POS] = box["acts"]
patched = model.lm_head.output[:, -1].save()
got.append(patched.float().cpu()); ref.append(base.float().cpu())
return torch.cat(got), torch.cat(ref)
for l in (0, 8, 16, 24, 31):
g, r = selfpatch_logits(keep[:24], l)
print(f"layer {l:>2} identical logits: {torch.equal(g, r)} "
f"max |Δ| = {(g - r).abs().max().item():.2e}")
layer 0 identical logits: True max |Δ| = 0.00e+00
layer 8 identical logits: True max |Δ| = 0.00e+00
layer 16 identical logits: True max |Δ| = 0.00e+00
layer 24 identical logits: True max |Δ| = 0.00e+00
layer 31 identical logits: True max |Δ| = 0.00e+00
Exactly zero at every layer. The read/write path is doing what we think it is.
The sweep. Three interventions at every layer, on the same samples:
- reversed — the discriminating test above: counterfactual drink-token residuals copied in with positions swapped, which leaves the drink at each position unchanged and moves only the ordering information;
- same position — the same residuals copied straight across. This does change the drink at each position, so it should move the answer under either hypothesis; it is the upper bound on what a state-token patch can do;
opaquecontrol — the same reversed swap applied to the twoopaquetokens instead. Same layer, same number of positions, same two runs, a component the causal model says is irrelevant. This is what tells us the effect is not "any large activation swap between these two prompts flips the answer".
sweep = {}
for name, src, dst in [("reversed", STATE_POS, STATE_POS_REV),
("same position", STATE_POS, STATE_POS),
("opaque control", OPAQUE_POS, OPAQUE_REV)]:
sweep[name] = [iia(interchange(keep, l, src, dst), keep) for l in range(N_LAYERS)]
peak = max(range(N_LAYERS), key=lambda l: sweep[name][l][0])
print(f"{name:<16} peak IIA {sweep[name][peak][0]:.3f} at layer {peak}")
reversed peak IIA 0.775 at layer 15
same position peak IIA 0.912 at layer 5
opaque control peak IIA 0.000 at layer 0
fig = go.Figure()
style = {"reversed": dict(color="#1a6ea8", dash="solid"),
"same position": dict(color="#c8781a", dash="solid"),
"opaque control": dict(color="#888888", dash="dot")}
for name, series in sweep.items():
fig.add_trace(go.Scatter(x=list(range(N_LAYERS)), y=[a for a, _ in series],
mode="lines+markers", name=name,
line=dict(width=2.5, **style[name])))
fig.add_trace(go.Scatter(x=list(range(N_LAYERS)), y=[u for _, u in sweep["reversed"]],
mode="lines", name="reversed: clean answer survives",
line=dict(color="#1a6ea8", width=1.4, dash="dash"), opacity=0.6))
fig.update_layout(title="Binding lookback: interchanging the drink tokens' residual streams",
xaxis_title="layer patched",
yaxis_title="fraction of samples",
yaxis_range=[-0.03, 1.03], width=880, height=430,
legend=dict(orientation="h", y=-0.24))
fig.show()