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()
Finding 2 — retrieval is by ordering ID, not by token position¶
The plot separates the two hypotheses cleanly.
Through layer 10 the reversed swap does nothing at all (IIA 0.00, clean answer surviving
at 0.94–0.98). That is not a weak effect, it is a necessary one, and it is worth pausing
on: at layer 0 the residual at position 155 is the embedding of tea in both runs, so
copying the counterfactual's tea embedding into the clean run's tea position is writing a
value onto itself. The reversed swap is literally a no-op until the model has written
something position-dependent into those residual streams. This is the patch-onto-itself
anchor obtained a second way, and for free.
From layer 11 the answer starts to move, and layers 12–16 flip most of it: IIA 0.34 at layer 11, then 0.73, 0.73, 0.76, 0.78 (the peak, at layer 15) and 0.73, while the clean answer collapses to 0.04. The drink at each token position is unchanged; the only thing that moved is whatever the model had written into those residual streams about where in the story they sit. The positional hypothesis predicts no change here. It is wrong. Retrieval is keyed on information carried in the residual stream, which is what "the state token holds an address" means.
The same-position swap sits at 0.87–0.91 for every layer up to 16, which is exactly what it should be: it swaps the drink identities, so it rewrites the story rather than the bindings, and any mechanism at all would follow it. It is the ceiling, not evidence. The interesting comparison is that the reversed swap climbs to within a few points of it at layers 14–16 — moving the ordering information alone is worth almost as much as rewriting the drinks.
The opaque control is exactly 0.000 at all 32 layers. Same layer, same number of
positions, same pair of runs, a component the causal model says is irrelevant — and the
answer never once becomes the swapped-binding target. So the effect above is not "swapping
any two residual streams between these prompts perturbs the answer".
After layer 17 both curves collapse together — 0.22, 0.18, 0.35, 0.24 … and near zero from layer 24 — and by layer 31 the clean answer is back at 1.00. Note that the two curves collapse in step: past layer 17 it no longer matters whether you swap the ordering or the drinks, because nothing downstream is reading the drink tokens any more. That is the signature of a lookback that has already fired, and it is where §6 picks up. In the middle of that decay (layers 17–23) the clean answer only partly returns, so those patches are disruptive without being informative — another reason to read the shape rather than any single layer.
Relative to the paper, the depth is strikingly consistent. Prakash et al. locate the state token's address and payload at layers 33–38 of Llama-3-70B's 80 — a relative depth of 0.41–0.48. Our band, layers 12–16 of 32, is 0.38–0.50. We would not read much into that on its own, but it is not nothing.
The plot separates the two hypotheses cleanly.
Below layer ~10 the reversed swap does nothing at all. That is not a weak effect, it is
a necessary one, and it is worth pausing on: at layer 0 the residual at position 155 is
the embedding of beer in both runs, so copying the counterfactual's beer embedding into
the clean run's beer position is writing a value onto itself. The reversed swap is
literally a no-op until the model has written something position-dependent into those
residual streams. The clean answer survives at ~97%, which is the same anchor as the
patch-onto-itself check, obtained a second way.
From layer ~11 the reversed swap starts flipping the answer, and by layers 14–16 it flips most of them. The drink at each token position is unchanged; the only thing that moved is whatever the model had written into those residual streams about where in the story they sit. The positional hypothesis predicts no change here. It is wrong. Retrieval is keyed on information carried in the residual stream, which is what "the state token holds an address" means.
The same-position swap is high everywhere early, which is exactly what it should be: it swaps the drink identities, so it rewrites the story rather than the bindings, and any mechanism at all would follow it. It is the ceiling, not evidence. The interesting fact is that the reversed swap climbs to nearly the same height in a narrow band of layers — the ordering information alone is worth almost as much as rewriting the drinks.
The opaque control is flat at zero across all 32 layers, and the clean answer survives
at 98–100%. Same layer, same number of positions, same pair of runs, a component the causal
model says is irrelevant — and nothing moves. So the effect above is not "swapping any two
residual streams between these prompts perturbs the answer".
After layer ~17 the effect collapses to zero and the clean answer comes back. By then the dereference has already happened: the information has been read out of the state tokens and is sitting somewhere downstream, so editing the state tokens is editing a value nobody reads any more. That is the signature of a lookback that has already fired, and it is where §6 picks up.
Relative to the paper, the depth is strikingly consistent. Prakash et al. locate the state token's address and payload at layers 33–38 of Llama-3-70B's 80 — a relative depth of 0.41–0.48. Our peak sits at the same relative depth in a 32-layer model. We would not read much into that on its own, but it is not nothing.
5. Where the Ordering ID comes from¶
§4 shows that something ordering-shaped rides in the drink token's residual stream and that retrieval keys off it. It does not show where that something is computed. The paper's claim is that the source of the OI is the character and object tokens — the entities whose textual order defines "first" and "second" — and that attention copies their OIs into the drink token to form the address.
The test is a second interchange, on a different counterfactual. This time we reverse the sentences and change both drinks to two drinks that appear in neither prompt. Then:
- patch the character and object tokens (reversed) from the counterfactual into the
clean run, at every layer up to some depth
L; - freeze the drink tokens at their clean values at every layer, so the patch cannot reach them and the payloads stay clean.
If the OI is sourced at the character and object tokens, swapping theirs re-points the query at the other sentence, and — with the drinks frozen clean — the answer becomes the other clean drink. The freeze is what makes this a statement about the source rather than about the address: any effect that survives it has travelled through the query side.
The swap has the same self-cancelling property as §4. Because the counterfactual reverses the two sentences, its token at position 158 is the same character as the clean run's at 146, so writing one into the other changes no identity — only ordering. At layer 0 it is therefore an exact no-op, which is the built-in anchor for this experiment.
def build_source(n, seed=11):
rng, rows = random.Random(seed), []
for _ in range(n):
c, o, s = sample_entities(rng)
t = rng.sample([x for x in STATES if x not in s], 2) # drinks in neither prompt
r = rng.choice([0, 1])
clean_p, clean_a = prompt(c, o, s, r, r)
cf_p, cf_a = prompt(c[::-1], o[::-1], t, 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
src_rows = build_source(120)
pc = predict([r["clean"] for r in src_rows]); pf = predict([r["cf"] for r in src_rows])
src_keep = [r for r, a, b in zip(src_rows, pc, pf)
if a == r["clean_ans"] and b == r["cf_ans"]]
print(f"kept {len(src_keep)}/{len(src_rows)}")
print("clean :", src_rows[0]["clean"].split("To complete an order,")[1].split("\n")[0].strip())
print("cf :", src_rows[0]["cf"].split("To complete an order,")[1].split("\n")[0].strip())
print("target:", src_rows[0]["target"])
kept 91/120 clean : Pete grabs an opaque drum and fills it with cocoa. Then Joe grabs another opaque urn and fills it with bourbon. cf : Joe grabs an opaque urn and fills it with porter. Then Pete grabs another opaque drum and fills it with coffee. target: bourbon
This one needs two forward passes rather than one, because it reads from both prompts at many layers before it writes anything: a read-only trace with two invokes collects the counterfactual's character/object residuals and the clean run's drink residuals, and a second trace replays the clean prompt with both sets written in.
⚠️ Note where cf_src and cl_state are created. Names bound inside a trace body do not
survive the block, so a dict comprehension written inside tracer.invoke leaves you with a
NameError one cell later. Create the container outside the trace and fill it in; the
container survives, and the saved values inside it are there when the trace returns.
SRC_POS = CHAR_POS + OBJ_POS # [131, 133, 146, 158, 150, 162]
SRC_POS_REV = [133, 131, 158, 146, 162, 150] # each entity takes the other's slot
@torch.no_grad()
def source_patch(rows, upto, freeze=True, bs=8):
out = []
for i in range(0, len(rows), bs):
chunk = rows[i:i + bs]
cf_src, cl_state = {}, {} # containers live OUTSIDE the trace
with model.trace() as tracer: # read-only
with tracer.invoke([r["cf"] for r in chunk]):
for l in range(upto + 1):
cf_src[l] = layers[l].output[:, SRC_POS].save()
with tracer.invoke([r["clean"] for r in chunk]):
for l in range(N_LAYERS):
cl_state[l] = layers[l].output[:, STATE_POS].save()
with model.trace([r["clean"] for r in chunk]): # write
for l in range(N_LAYERS): # layers must be touched in order
pos, val = [], []
if l <= upto:
pos = pos + SRC_POS_REV; val.append(cf_src[l])
if freeze:
pos = pos + STATE_POS; val.append(cl_state[l])
if pos:
layers[l].output[:, pos] = torch.cat(val, dim=1)
p = model.lm_head.output[:, -1].argmax(-1).save()
out += [tok.decode([t]).strip().lower() for t in p]
return out
source_sweep = {l: iia(source_patch(src_keep, l, freeze=True), src_keep) for l in range(0, N_LAYERS, 2)}
source_nofreeze = {l: iia(source_patch(src_keep, l, freeze=False), src_keep) for l in range(0, N_LAYERS, 4)}
for l, v in source_sweep.items():
print(f"layers 0..{l:<2} IIA {v[0]:.3f} unchanged {v[1]:.3f}")
layers 0..0 IIA 0.000 unchanged 1.000 layers 0..2 IIA 0.000 unchanged 0.989 layers 0..4 IIA 0.000 unchanged 0.978 layers 0..6 IIA 0.011 unchanged 0.945 layers 0..8 IIA 0.275 unchanged 0.659 layers 0..10 IIA 0.473 unchanged 0.473 layers 0..12 IIA 0.473 unchanged 0.473 layers 0..14 IIA 0.473 unchanged 0.462 layers 0..16 IIA 0.473 unchanged 0.462 layers 0..18 IIA 0.473 unchanged 0.473 layers 0..20 IIA 0.473 unchanged 0.473 layers 0..22 IIA 0.473 unchanged 0.473 layers 0..24 IIA 0.473 unchanged 0.462 layers 0..26 IIA 0.473 unchanged 0.462 layers 0..28 IIA 0.484 unchanged 0.473 layers 0..30 IIA 0.473 unchanged 0.473
The no-freeze variant is the same patch with the drink tokens left alone. Comparing the two tells us whether the effect needs the payloads held clean, or whether the character and object patch simply propagates into the drink tokens and does the work there.
fig = go.Figure()
fig.add_trace(go.Scatter(x=list(source_sweep), y=[v[0] for v in source_sweep.values()],
mode="lines+markers", name="character + object OIs swapped, drinks frozen",
line=dict(color="#1a6ea8", width=2.5)))
fig.add_trace(go.Scatter(x=list(source_sweep), y=[v[1] for v in source_sweep.values()],
mode="lines", name="clean answer survives",
line=dict(color="#1a6ea8", width=1.4, dash="dash"), opacity=0.6))
fig.add_trace(go.Scatter(x=list(source_nofreeze), y=[v[0] for v in source_nofreeze.values()],
mode="lines+markers", name="same, drinks not frozen",
line=dict(color="#c8781a", width=2.2, dash="dot")))
fig.update_layout(title="Binding lookback: swapping the source Ordering IDs, layers 0..L",
xaxis_title="patched at every layer up to L",
yaxis_title="fraction of samples",
yaxis_range=[-0.03, 1.03], width=880, height=430,
legend=dict(orientation="h", y=-0.24))
fig.show()
Finding 3 — the Ordering ID is sourced at the character and object tokens¶
The curve is flat at zero through layers 0..6, jumps to 0.275 at 0..8, reaches 0.473 at 0..10, and then does not move again for the remaining twenty layers. Read it as a statement about when the source is finished: patching layers 0..L only helps once L is deep enough to include the layers where the character and object tokens have acquired their Ordering IDs, and once it does, adding more layers adds nothing. There is a window — roughly layers 7–10 — in which the OI is computed, and it closes.
That window sits below the layers where the address is readable in the drink tokens (12–16, §4), which is the ordering the mechanism requires: the source has to be finished before the copy that forms the address can carry it.
Two things make this a claim about the source rather than a second measurement of the address:
- the drink tokens are frozen at their clean values at every layer, so no part of this effect can be the patch leaking into the addresses we already tested in §4;
- the counterfactual's drinks appear in neither prompt, so if the effect were the payload travelling rather than the ordering, the model would answer a drink we could recognise as foreign. It does not — the answers are the clean story's other drink.
The no-freeze control settles the second point empirically, and it is the sharpest number in this section. Run the same patch without the freeze and IIA never exceeds 0.033. The freeze is not decoration: without it the patched character and object OIs propagate into the drink tokens and move the addresses along with the query, so the pointer and the address shift together and the dereference lands exactly where it did before. Holding the addresses still is what converts an intervention that changes everything into an intervention that changes one thing.
The honest caveat is the height. 0.473 is well short of the 0.775 the direct address swap reaches, so swapping the source OIs redirects the query on about half the samples rather than almost all of them. Some of that is the extra difficulty of the counterfactual (a four-way-different story rather than a re-ordered one), and some of it is presumably that the source is not confined to the six token positions we patch. We report the number rather than the interpretation we would prefer.
The curve is flat at zero for the first several layers and then rises to a plateau. Read it as a statement about when the source is finished: patching layers 0..L only helps once L is deep enough to include the layers where the character and object tokens have actually acquired their Ordering IDs. Before that we are copying representations that do not yet encode order, and — because of the self-cancelling construction — copying them is copying nothing.
Two things make this a claim about the source rather than a second measurement of the address:
- the drink tokens are frozen at their clean values at every layer, so no part of this effect can be the patch leaking into the addresses we already tested in §4;
- the counterfactual's drinks appear in neither prompt, so if the effect were the payload travelling rather than the ordering, the model would answer a drink we could recognise as foreign. It does not — the answers are the clean story's other drink.
The plateau matters as much as the rise. Adding more layers past the point where the OI is formed does not add IIA, which is what you expect if there is a specific window in which the source is computed and copied, rather than a diffuse accumulation.
Compare with the no-freeze variant. Without the freeze the patch is free to propagate into the drink tokens, which is a different and larger intervention, and it behaves differently. The freeze is doing real work: it is the difference between "we moved the ordering information" and "we moved everything downstream of the ordering information".
Put §4 and §5 together and the chain is legible in the layer indices. The character and object tokens have usable OIs first; the drink tokens become addressable a few layers later; and by the high teens editing the drink tokens no longer changes anything, because the lookback has already fired. That is the binding lookback, in a model 1/9th the size of the smallest one the paper examines.
6. The answer lookback: pointer, then payload¶
The binding lookback does not deliver a word. Its payload is the state OI — "the answer is whichever drink is second" — deposited at the final token. A second lookback turns that into the drink itself: the answer lookback, whose pointer is that state OI, whose addresses are the state OIs sitting at each drink token, and whose payload is the drink token.
That predicts something specific and testable about the final token's residual stream. Early in the stack it should hold a pointer — an OI that means nothing on its own. Late in the stack it should hold the payload — the drink itself. So take a counterfactual that shares the clean prompt's characters and containers but has different drinks and asks about the other character, and copy its final-token residual into the clean run:
| what the final token holds at layer ℓ | what the clean run answers |
|---|---|
| nothing yet | the clean answer, s_r |
| a pointer (state OI), dereferenced against the clean run's drinks | the clean story's other drink, s_{1-r} |
| the payload (the drink token itself) | a drink from the counterfactual story, t_{1-r} |
The three outcomes use disjoint vocabulary, so a single sweep separates them.
def build_answer(n, seed=13):
rng, rows = random.Random(seed), []
for _ in range(n):
c, o, s = sample_entities(rng)
t = rng.sample([x for x in STATES if x not in s], 2) # drinks in neither prompt
r = rng.choice([0, 1])
clean_p, clean_a = prompt(c, o, s, r, r)
cf_p, cf_a = prompt(c, o, t, 1 - r, 1 - r)
rows.append(dict(clean=clean_p, clean_ans=clean_a, cf=cf_p, cf_ans=cf_a,
pointer_target=s[1 - r], payload_target=t[1 - r]))
return rows
ans_rows = build_answer(140)
pc = predict([r["clean"] for r in ans_rows]); pf = predict([r["cf"] for r in ans_rows])
ans_keep = [r for r, a, b in zip(ans_rows, pc, pf)
if a == r["clean_ans"] and b == r["cf_ans"]]
print(f"kept {len(ans_keep)}/{len(ans_rows)}")
r0 = ans_keep[0]
print(f"\nclean answer {r0['clean_ans']!r}")
print(f"pointer target {r0['pointer_target']!r} (the clean story's other drink)")
print(f"payload target {r0['payload_target']!r} (a drink only the counterfactual mentions)")
kept 102/140 clean answer 'soda' pointer target 'wine' (the clean story's other drink) payload target 'juice' (a drink only the counterfactual mentions)
@torch.no_grad()
def patch_final_token(rows, layer, bs=16):
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["cf"] for r in chunk]):
box["acts"] = layers[layer].output[:, -1]
barrier()
with tracer.invoke([r["clean"] for r in chunk]):
barrier()
layers[layer].output[:, -1] = box["acts"]
p = model.lm_head.output[:, -1].argmax(-1).save()
out += [tok.decode([t]).strip().lower() for t in p]
return out
ans_sweep = {}
for l in range(N_LAYERS):
p = patch_final_token(ans_keep, l)
ans_sweep[l] = dict(
clean=sum(x == r["clean_ans"] for x, r in zip(p, ans_keep)) / len(ans_keep),
pointer=sum(x == r["pointer_target"] for x, r in zip(p, ans_keep)) / len(ans_keep),
payload=sum(x == r["payload_target"] for x, r in zip(p, ans_keep)) / len(ans_keep),
)
print(f"layer {l:>2} clean {ans_sweep[l]['clean']:.2f} "
f"pointer {ans_sweep[l]['pointer']:.2f} payload {ans_sweep[l]['payload']:.2f}")
layer 0 clean 0.99 pointer 0.00 payload 0.00
layer 1 clean 0.99 pointer 0.00 payload 0.00
layer 2 clean 1.00 pointer 0.00 payload 0.00
layer 3 clean 1.00 pointer 0.00 payload 0.00
layer 4 clean 0.99 pointer 0.00 payload 0.00
layer 5 clean 0.99 pointer 0.00 payload 0.00
layer 6 clean 1.00 pointer 0.00 payload 0.00
layer 7 clean 0.99 pointer 0.00 payload 0.00
layer 8 clean 0.99 pointer 0.00 payload 0.00
layer 9 clean 0.99 pointer 0.00 payload 0.00
layer 10 clean 0.99 pointer 0.00 payload 0.00
layer 11 clean 0.99 pointer 0.00 payload 0.00
layer 12 clean 0.70 pointer 0.23 payload 0.00
layer 13 clean 0.60 pointer 0.29 payload 0.00
layer 14 clean 0.42 pointer 0.41 payload 0.00
layer 15 clean 0.03 pointer 0.82 payload 0.00
layer 16 clean 0.01 pointer 0.90 payload 0.00
layer 17 clean 0.07 pointer 0.77 payload 0.00
layer 18 clean 0.06 pointer 0.73 payload 0.00
layer 19 clean 0.01 pointer 0.81 payload 0.00
layer 20 clean 0.02 pointer 0.76 payload 0.00
layer 21 clean 0.03 pointer 0.66 payload 0.00
layer 22 clean 0.01 pointer 0.60 payload 0.00
layer 23 clean 0.02 pointer 0.43 payload 0.09
layer 24 clean 0.02 pointer 0.05 payload 0.23
layer 25 clean 0.01 pointer 0.03 payload 0.48
layer 26 clean 0.00 pointer 0.01 payload 0.64
layer 27 clean 0.00 pointer 0.00 payload 0.95
layer 28 clean 0.00 pointer 0.00 payload 0.94
layer 29 clean 0.00 pointer 0.00 payload 0.94
layer 30 clean 0.00 pointer 0.00 payload 1.00
layer 31 clean 0.00 pointer 0.00 payload 1.00
fig = go.Figure()
for key, colour, label in [("clean", "#888888", "clean answer (no effect)"),
("pointer", "#1a6ea8", "pointer: clean story's other drink"),
("payload", "#c8781a", "payload: a counterfactual-only drink")]:
fig.add_trace(go.Scatter(x=list(ans_sweep), y=[ans_sweep[l][key] for l in ans_sweep],
mode="lines+markers", name=label,
line=dict(color=colour, width=2.5,
dash="dot" if key == "clean" else "solid")))
fig.update_layout(title="Answer lookback: patching the final token's residual stream",
xaxis_title="layer patched", yaxis_title="fraction of samples",
yaxis_range=[-0.03, 1.03], width=860, height=420,
legend=dict(orientation="h", y=-0.22))
fig.show()
Finding 4 — the final token holds a pointer before it holds an answer¶
The three curves separate almost perfectly, and in the right order.
Through layer 11, patching the final token does nothing — the clean answer survives at 0.99. Then the pointer curve rises: 0.23 at layer 12, 0.41 at 14, 0.82 at 15, peaking at 0.90 at layer 16, and holding above 0.6 through layer 22. Across that whole band the payload curve is flat at 0.000.
That is the claim, and it is worth stating plainly. In layers 12–22, copying the final token's residual stream from the counterfactual makes the model answer a drink that appears nowhere in the counterfactual prompt. Whatever was copied cannot have been the word, because the word was not there to copy. It was a reference, and it was resolved against the clean run's drink tokens after being copied in. That is a dereference, observed by handing the model a pointer from one world and the addresses of another.
The payload curve only starts at layer 23 (0.09), crosses the pointer curve at 24, and reaches 1.00 by layer 30. By then the residual stream at that position simply is the answer, and nothing about a lookback is needed to predict that half of the figure.
The ordering is what makes it a chain rather than two facts:
| layers | what is happening |
|---|---|
| 7–10 | character and object Ordering IDs are computed (§5) |
| 12–16 | the drink tokens' addresses are readable and swappable (§4) |
| 12–22 | the final token holds a state OI pointer |
| 23–31 | the final token holds the drink token itself |
The binding lookback fires and deposits a state OI at the final token; the answer lookback dereferences that OI against the drink tokens and brings the word forward. Note that layer 17 — where §4's effect collapses because nothing reads the drink tokens any more — sits inside the pointer band, exactly as it should if the read has already happened.
This is the paper's Figure 4 at one ninth the parameter count, with the same shape and the same ordering.
The three curves are the point of this section, and the middle one is the one that matters.
The payload curve — the model answering a drink that appears only in the counterfactual story — is the unsurprising half. Patch the final token's residual stream late enough and the answer comes with it; by the last layers the residual stream at that position simply is the answer. Nothing about a lookback is needed to predict that.
The pointer curve is the claim. There is a band of layers where patching the final token makes the model answer the clean story's other drink — a word that appears nowhere in the counterfactual prompt. Whatever was copied cannot have been the word, because the word was not there to copy. It was a reference, and it was resolved against the clean run's drink tokens after being copied in. That is a dereference, observed by giving the model a pointer from one world and the addresses of another.
The ordering is what makes it a chain rather than two facts. The pointer band comes first, the payload band comes later, and they barely overlap. Read together with §4 — where editing the drink tokens stops mattering at about the layer where the pointer band begins — the picture is: the binding lookback fires and deposits a state OI at the final token; a few layers later the answer lookback dereferences that OI against the drink tokens and brings the word forward.
This is the paper's Figure 4 at one ninth the parameter count, with the same shape and the same ordering.
7. Is the ordering information low-rank?¶
The paper's sharpest structural claim is that these lookbacks live in low-rank subspaces
of the state token's residual stream — a handful of directions out of thousands carry the
Ordering ID, and the rest of the residual is doing other work. That claim needs its own
control, and the control is not optional: a rank-k subspace chosen to be good will beat a
rank-k subspace chosen at random only if the structure is real.
We do the cheap version of the paper's analysis. At the peak layer from §4, collect the
swap direction for every sample — the difference between what the reversed interchange
writes and what was already there — fit an orthonormal basis to the top k principal
directions of those differences on half the samples, and then, on the held-out
half, perform the interchange only inside that subspace:
$$h_{\text{dst}} \;\leftarrow\; h_{\text{dst}} + U_k U_k^{\top}\,(h^{\text{cf}}_{\text{src}} - h_{\text{dst}})$$
The anchor is the unrestricted swap on the same held-out samples, which must reproduce the
§4 number at that layer. The control is a random orthonormal subspace of the same rank,
averaged over three draws — the number that says whether a rank-k subspace is special or
whether any k directions would do.
@torch.no_grad()
def collect(rows, layer, positions, key, bs=16):
out = []
for i in range(0, len(rows), bs):
with model.trace([r[key] for r in rows[i:i + bs]]):
a = layers[layer].output[:, positions].save()
out.append(a.float().cpu())
return torch.cat(out) # [N, len(positions), d_model]
@torch.no_grad()
def assign_and_predict(rows, layer, positions, values, bs=16):
out = []
for i in range(0, len(rows), bs):
v = values[i:i + bs].to(DEVICE, dtype=DTYPE)
with model.trace([r["clean"] for r in rows[i:i + bs]]):
layers[layer].output[:, positions] = v
p = model.lm_head.output[:, -1].argmax(-1).save()
out += [tok.decode([t]).strip().lower() for t in p]
return out
PEAK = max(range(N_LAYERS), key=lambda l: sweep["reversed"][l][0])
print(f"peak binding layer: {PEAK}")
h_dst = collect(keep, PEAK, STATE_POS, "clean") # what is there now
h_src = collect(keep, PEAK, STATE_POS_REV, "cf") # what the reversed swap writes
delta = h_src - h_dst # [N, 4, d_model]
d_model = delta.shape[-1]
split = len(keep) // 2
test_rows = keep[split:]
# principal directions of the swap, fitted on the training half only
D = delta[:split].reshape(-1, d_model)
U_full = torch.linalg.svd(D - D.mean(0, keepdim=True), full_matrices=False).Vh # [r, d_model]
print(f"fitted {U_full.shape[0]} directions from {D.shape[0]} difference vectors "
f"({split} held-in samples x {len(STATE_POS)} positions)")
peak binding layer: 15
fitted 204 directions from 204 difference vectors (51 held-in samples x 4 positions)
The anchor. Writing the full source activation on the held-out half must reproduce the
layer-PEAK number from §4, restricted to those samples. If it does not, the subspace
arithmetic below is being done on the wrong tensors.
full_swap = iia(assign_and_predict(test_rows, PEAK, STATE_POS, h_src[split:]), test_rows)[0]
print(f"full swap on the held-out half : {full_swap:.3f}")
print(f"same thing measured in section 4: {sweep['reversed'][PEAK][0]:.3f} (all samples)")
full swap on the held-out half : 0.745 same thing measured in section 4: 0.775 (all samples)
def project_swap(k, basis):
"""The interchange restricted to the span of `basis` (k x d_model, orthonormal rows)."""
P = basis[:k]
proj = torch.einsum("nkd,rd->nkr", delta[split:], P)
return h_dst[split:] + torch.einsum("nkr,rd->nkd", proj, P)
ranks = [1, 2, 4, 8, 16, 32, 64, 128]
low_rank, random_rank = {}, {}
for k in ranks:
low_rank[k] = iia(assign_and_predict(test_rows, PEAK, STATE_POS,
project_swap(k, U_full)), test_rows)[0]
scores = []
for seed in range(3):
g = torch.Generator().manual_seed(seed)
R = torch.linalg.qr(torch.randn(d_model, k, generator=g))[0].T # [k, d_model]
scores.append(iia(assign_and_predict(test_rows, PEAK, STATE_POS,
project_swap(k, R)), test_rows)[0])
random_rank[k] = sum(scores) / len(scores)
print(f"rank {k:>4} fitted subspace {low_rank[k]:.3f} "
f"random subspace {random_rank[k]:.3f}")
rank 1 fitted subspace 0.000 random subspace 0.000
rank 2 fitted subspace 0.725 random subspace 0.000
rank 4 fitted subspace 0.706 random subspace 0.000
rank 8 fitted subspace 0.725 random subspace 0.000
rank 16 fitted subspace 0.725 random subspace 0.000
rank 32 fitted subspace 0.706 random subspace 0.000
rank 64 fitted subspace 0.745 random subspace 0.000
rank 128 fitted subspace 0.765 random subspace 0.000
fig = go.Figure()
fig.add_trace(go.Scatter(x=ranks, y=[low_rank[k] for k in ranks], mode="lines+markers",
name="fitted subspace (held-out)",
line=dict(color="#1a6ea8", width=2.5)))
fig.add_trace(go.Scatter(x=ranks, y=[random_rank[k] for k in ranks], mode="lines+markers",
name="random subspace, matched rank",
line=dict(color="#888888", width=2.5, dash="dot")))
fig.add_hline(y=full_swap, line=dict(color="#c8781a", width=1.6, dash="dash"),
annotation_text="full swap", annotation_position="bottom right")
fig.update_layout(title=f"Rank of the swapped subspace at layer {PEAK}",
xaxis_title="rank k", yaxis_title="IIA on held-out samples",
xaxis_type="log", yaxis_range=[-0.03, 1.03], width=860, height=420,
legend=dict(orientation="h", y=-0.22))
fig.show()
Finding 5 — two directions carry the ordering information¶
The anchor first: the unrestricted swap on the held-out half gives 0.745, against 0.775 measured over all samples in §4. Same intervention, same layer, a different half of the data — the projection arithmetic below is operating on the right tensors.
Then the comparison, and it is stark:
| rank | fitted subspace | random subspace |
|---|---|---|
| 1 | 0.000 | 0.000 |
| 2 | 0.725 | 0.000 |
| 4 | 0.706 | 0.000 |
| 8 | 0.725 | 0.000 |
| 16 | 0.725 | 0.000 |
| 32 | 0.706 | 0.000 |
| 64 | 0.745 | 0.000 |
| 128 | 0.765 | 0.000 |
Rank 2 out of 4096 recovers 0.725 of the 0.745 full-swap effect, and every higher rank we tried recovers the same thing — the curve is a step, not a ramp. One direction is not enough; two are almost all of it. Meanwhile a random orthonormal subspace of the same rank recovers exactly nothing, at every rank up to 128. That gap is the whole content of the claim. Without the random control, "a rank-16 patch flips the answer" would be uninformative: a rank-16 slice of a 4096-dimensional residual could be doing anything. With it, we can say the ordering information occupies a small, identifiable set of directions rather than being spread through the residual stream.
Two honest qualifications. The subspace is fitted on held-in samples and evaluated on held-out ones, so it is not memorising the test set — but it is fitted to the difference vectors of this exact intervention, which is a much easier target than a subspace that has to support many interventions at once. And principal directions of a difference are a blunt instrument next to a trained rotation; DAS would plausibly find a cleaner rank-1 or rank-2 solution, so "two directions" is an upper bound on how compact the representation is, not a measurement of it.
The anchor comes first: the unrestricted swap on the held-out half reproduces the §4 number at that layer, so the projection arithmetic is operating on the right tensors.
Then the comparison. A fitted subspace of modest rank recovers much of the full effect, while a random subspace of the same rank recovers essentially none of it. That gap is the whole content of the claim. Without the random control, "a rank-16 patch flips the answer" would be uninformative — a rank-16 slice of a 4096-dimensional residual could be doing anything. With it, we can say that the ordering information occupies a small, identifiable set of directions rather than being spread evenly through the residual stream.
Two honest qualifications. The subspace is fitted on held-in samples and evaluated on held-out ones, so it is not memorising the test set — but it is fitted to the difference vectors of this exact intervention, which is a much easier target than a subspace that has to support many interventions at once. And principal directions of a difference are a blunt instrument next to a trained rotation; DAS would likely find a lower rank that works, so the rank at which our curve saturates is an upper bound on how compact the representation is, not a measurement of it.
Caveats¶
Scale. This is not the paper's experiment. Prakash et al. study Llama-3-70B-Instruct and Llama-3.1-405B-Instruct; we study an 8B model because it is what fits on one GPU. The paper explicitly declines to examine models this size on the grounds that they cannot coherently solve CausalToM, and §2 shows why that is a defensible call: 0.745 against 0.95. Every layer index quoted here is an index into 32 layers, not 80, and none of them should be mapped onto the paper's — the relative depths line up, which is suggestive and not more than that.
One model family. Everything here is Llama-3.1. The paper's own evaluation covers Qwen-2.5, OLMo-2 and Gemma-3 behaviourally, and we have not checked whether the mechanism we find transfers to any of them. A lookback found in one 8B model is a fact about that model.
Dataset size. Each intervention runs on ~100 error-filtered samples. At that size the standard error on an IIA of 0.5 is about 0.05, so differences of a few points between adjacent layers are noise and only the shape of the curve is meaningful. We report the curves rather than single numbers for that reason.
Prompt monoculture. Every prompt is the same restaurant template with the same instruction block, differing only in entities. That is what makes token-position interventions possible, and it also means we have shown a mechanism for this template. The paper checks generalisation on BigToM; we do not.
The visibility lookback is out of reach here. It is the one of the three that requires the model to first get the visibility condition right, and at 0.295 paired accuracy it does not. An interchange intervention on it would be measuring which of two wrong answers the model happens to give. We report the behavioural number and stop.
Subspaces are fitted, not learned. §7 uses the principal directions of the swap difference, which is a much blunter instrument than the paper's learned rotations (see the DAS tutorial for the trained version). A fitted subspace can only under-state how low-rank the true feature is, so this is a conservative test — but it is not the paper's test.
Conclusion¶
The paper's central abstraction survives the drop from 70B to 8B, and it survives in a form a reader can execute on one GPU.
Llama-3.1-8B-Instruct does not solve CausalToM the way Llama-3-70B-Instruct does — 0.745
against 0.95, with the unknown half carrying most of the loss and the visibility condition
(0.295 paired) beyond it entirely. But 0.745 is 25 points above the reality-strategy
baseline, which means there is a belief computation to interrogate, and when we interrogate
it we find the binding lookback and the answer lookback, chained, in the layer order the
paper's causal model predicts:
- the source Ordering IDs are computed at the character and object tokens by layer 10 (IIA 0.473, and 0.033 without the address freeze);
- the drink tokens carry an address that can be swapped independently of the drinks themselves, peaking at IIA 0.775 at layer 15, against 0.000 for a matched irrelevant-component control and a bit-exact patch-onto-itself;
- that address lives in a rank-2 subspace, which recovers 0.725 where a random subspace of the same rank recovers 0.000;
- and the final token holds a pointer to a state OI in layers 12–22 (peak 0.90) before it holds the payload in layers 23–31 (reaching 1.00) — the two never overlapping.
What we could not reach: the visibility lookback, because the model fails the behaviour it presupposes; the paper's learned low-rank subspaces, for which we substituted a fitted approximation; and any claim about generalisation beyond this one template and one model family.
The negative parts are as much of the result as the positive ones. The paper's stated reason for not looking at smaller models — that they cannot coherently solve the task — is vindicated for the visibility setting and too strong for the no-visibility one. A mechanism can be present, legible, and low-rank in a model that gets the task wrong a quarter of the time.
The paper's central abstraction survives the drop from 70B to 8B, and it survives in a form a reader can execute on one GPU.
Llama-3.1-8B-Instruct does not solve CausalToM the way Llama-3-70B-Instruct does — its
belief accuracy is well short, its unknown half is the weak half, and the visibility
condition is beyond it. But it is above the reality-strategy baseline, which means there is
a belief computation to interrogate, and when we interrogate it we find the binding
lookback: an Ordering ID sourced at the character and object tokens, written into the drink
tokens' residual streams as an address, and dereferenced by the query. The interchange that
moves only that ordering information — leaving the drink at every token position exactly
as it was — flips the model's answer in a narrow band of middle layers, while the same
interchange on a causally irrelevant component does nothing at all and patching a layer
onto itself is bit-exact.
What we could not reach: the visibility lookback, because the model fails the behaviour it presupposes; the paper's learned low-rank subspaces, for which we substituted a fitted approximation; and any claim about generalisation beyond this one template and one model family.
The negative parts are as much of the result as the positive ones. The paper's stated reason for not looking at smaller models — that they cannot coherently solve the task — is vindicated for the visibility setting and too strong for the no-visibility one. A mechanism can be present and legible in a model that gets the task wrong a third of the time.
References¶
- Prakash, N., Shapira, N., Sen Sharma, A., Riedl, C., Belinkov, Y., Rott Shaham, T., Bau, D., & Geiger, A. (2026). Language Models use Lookbacks to Track Beliefs. ICLR 2026. Project page: belief.baulab.info. Code: github.com/Nix07/belief_tracking.
- Geiger, A., Wu, Z., Potts, C., Icard, T., & Goodman, N. D. (2024). Finding Alignments Between Interpretable Causal Variables and Distributed Neural Representations. CLeaR 2024. The source of interchange interventions and IIA; the nnsight DAS tutorial implements its distributed version.
- Feng, J., & Steinhardt, J. (2024). How do Language Models Bind Entities in Context? ICLR 2024. The binding-ID picture that Ordering IDs generalise.
- Kim, N., & Schuster, S. (2023). Entity Tracking in Language Models. ACL 2023.
- Gandhi, K., Fränken, J.-P., Gerstenberg, T., & Goodman, N. D. (2023). Understanding Social Reasoning in Language Models with Language Models. NeurIPS 2023 Datasets and Benchmarks. The BigToM benchmark the paper generalises to.