Info
Last Execution: 2026-09-23
| Package | Version |
|---|---|
| nnsight | 0.8.0 |
| Python | 3.12.13 |
| torch | 2.13.0+cu126 |
| transformers | 5.15.0.dev0 |
Matryoshka Attribution¶
Introduction¶
🪆 Every attribution method ends with a ranking: which heads, MLPs or neurons matter most for a behaviour. The usual ways of getting one all pay for something. Causal interventions patch one component at a time and cost a forward pass per component, and single-component effects do not add up. Gradient methods cost one backward pass but rank by a first-order guess. Mask learning optimises a whole circuit at once, but every run commits to one sparsity through a penalty coefficient nobody knows how to set, and two runs at different sparsities need not agree on anything.
Arora et al. (2026), Matryoshka attribution: Learning to attribute language model outputs to representations and weights, reframe the problem. An attribution is a nested family of circuits — the best 1 component should sit inside the best 5, which should sit inside the best 50. Their method, MAttr, learns one score per component and turns the scores into a mask with a sigmoid top-$k$ operator: a soft mask whose entries sum to exactly $k$. The trick that gives the nesting is to sample a fresh $k$ at every training step, so the same score vector is supervised at every sparsity at once (the Matryoshka technique of Kusupati et al., 2022). There is no sparsity penalty, no straight-through estimator, and the output is a ranking rather than a circuit at one budget.
We reproduce the paper's headline recipe on its headline cell — GPT-2 small on the Indirect Object Identification task from the Mechanistic Interpretability Benchmark (MIB, Mueller et al., 2025), at the node level: 144 attention heads, 12 MLPs and the input embedding, 157 components in all — and check four claims:
- The operator. A sigmoid top-$k$ mask sums to $k$ exactly, and its implicit-differentiation gradient is correct, with a coupling term that makes only relative scores matter.
- One run, every sparsity. Five hundred steps of Adam on 157 scores gives a ranking whose faithfulness curve matches the paper's reported CPR for this cell, and beats attribution patching, integrated gradients and per-node interchange interventions on the same harness.
- Nested beats fixed. A mask trained at one budget is faithful at that budget and nowhere else; the randomised-$k$ mask is faithful everywhere.
- The ranking is the IOI circuit. The heads Wang et al. (2023) named come out on top, the two
negative name movers come out at the very bottom, and the paper's two surprises about
a9.h6anda11.h2reproduce.
Along the way we hit one property of MIB's node-level intervention that the metric hides: the input-embedding node is a master switch, and where it lands in the ranking decides whether the smallest circuits score anything at all.
📗 Paper: Matryoshka attribution: Learning to attribute language model outputs to
representations and weights (Arora, Acharya, Hu, Zhang,
Goodman, Jurafsky, Potts, 2026). Code:
aryamanarora/matryoshka-attribution.
The intervention semantics, the loss and the CPR metric below follow the paper's eval_mib.py
and MIB's evaluator; the mask is applied through nnsight instead of PyTorch hooks.
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 datasets
import math
import random
import time
import numpy as np
import torch
import nnsight
from nnsight import TransformersModel
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"
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
torch.manual_seed(0)
random.seed(0)
The model is frozen: the only parameters that will ever receive a gradient are the 157 scores.
requires_grad_(False) on the envoy forwards to the underlying module. Gradients still reach the
scores, because the mask is injected at the input embedding and everything downstream of an
injected tensor is in the graph.
model = TransformersModel("openai-community/gpt2", device_map=DEVICE, dispatch=True)
model.requires_grad_(False)
blocks = model.transformer.h
n_layers, n_heads, d_model = model.config.n_layer, model.config.n_head, model.config.n_embd
d_head = d_model // n_heads
tok = model.tokenizer
print(f"gpt2: {n_layers} layers x {n_heads} heads, d_head={d_head}, d_model={d_model}")
gpt2: 12 layers x 12 heads, d_head=64, d_model=768
The task¶
MIB's IOI data comes with several counterfactuals per prompt. The one the benchmark and the paper
use for this task is s2_io_flip: the second mention of the subject is replaced by the indirect
object, so "… Carl and Maria …, Carl gave" becomes "… Carl and Maria …, Maria gave",
and the answer flips from Maria to Carl. Clean and corrupted prompts differ at exactly one token
and have the same length. As MIB does, we keep only examples whose two names are single tokens.
Prompts are batched by token length. Every batch is one length bucket, so there is no padding anywhere in this notebook, and "the last position" means the same thing in every row.
from datasets import load_dataset
def prepare(row):
io, s = row["metadata"]["indirect_object"], row["metadata"]["subject"]
io_ids, s_ids = tok(" " + io).input_ids, tok(" " + s).input_ids
if len(io_ids) != 1 or len(s_ids) != 1:
return None
clean = tok(row["prompt"]).input_ids
corrupt = tok(row["s2_io_flip_counterfactual"]["prompt"]).input_ids
if len(clean) != len(corrupt):
return None
return dict(clean=clean, corrupt=corrupt, io=io_ids[0], s=s_ids[0],
text=row["prompt"], text_corrupt=row["s2_io_flip_counterfactual"]["prompt"])
def load_split(split, n):
examples = []
for row in load_dataset("mib-bench/ioi", split=split):
example = prepare(row)
if example is not None:
examples.append(example)
if len(examples) == n:
break
buckets = {}
for example in examples:
buckets.setdefault(len(example["clean"]), []).append(example)
return examples, buckets
train_examples, train_buckets = load_split("train", 2000) # scores are learned here
val_examples, val_buckets = load_split("validation", 200) # and evaluated here
def to_tensors(examples):
column = lambda key: torch.tensor([e[key] for e in examples], device=DEVICE)
return column("clean"), column("corrupt"), column("io"), column("s")
def sample_batch(buckets, size):
lengths = list(buckets)
length = random.choices(lengths, weights=[len(buckets[l]) for l in lengths])[0]
return random.sample(buckets[length], min(size, len(buckets[length])))
example = train_examples[0]
print("clean :", example["text"], "->", repr(tok.decode([example["io"]])))
print("corrupted :", example["text_corrupt"], "->", repr(tok.decode([example["s"]])))
print(f"\ntrain: {len(train_examples)} examples in {len(train_buckets)} length buckets"
f" validation: {len(val_examples)} examples in {len(val_buckets)} buckets")
clean : As Carl and Maria left the consulate, Carl gave a fridge to -> ' Maria' corrupted : As Carl and Maria left the consulate, Maria gave a fridge to -> ' Carl' train: 2000 examples in 14 length buckets validation: 200 examples in 10 buckets
1. The substrate: 157 nodes, one intervention¶
MIB's node level has one variable per attention head, one per MLP block and one for the input
embedding. In GPT-2 the per-head outputs are concatenated and fed to attn.c_proj, so head $h$
of layer $\ell$ is channels [h*d_head:(h+1)*d_head] of blocks[ℓ].attn.c_proj.input; the MLP's
variable is blocks[ℓ].mlp.c_proj.input, the 3072 post-GELU activations; the input node is
wte.output. All 25 sites are read in forward order, which nnsight requires.
The intervention is MIB's denoising one, which the paper calls iso. A mask $m \in [0,1]^{157}$ blends each node between its own live value and the value it took on the corrupted prompt:
$$h \leftarrow m_H\, h(\text{live}) + (1 - m_H)\, h(\text{corrupted})$$
At $m = 1$ that is the clean model; at $m = 0$ every node is patched and it is the corrupted run. Two things are worth noticing about "live". A kept node is not set to its clean value; it is computed from whatever reaches it, which may already be partly corrupted. And a soft $m$ makes the whole thing differentiable in $m$, which is what MAttr optimises through.
The corrupted values come from one no-grad trace per batch; the mask is applied in a second trace
on the clean prompts. The masked value is written back as a replacement (.input = …) rather
than an in-place edit, so autograd never sees a tensor modified after it was saved.
N_NODES = 1 + n_layers * n_heads + n_layers # input, then heads, then MLPs
def node_name(i):
if i == 0:
return "input"
i -= 1
return f"a{i // n_heads}.h{i % n_heads}" if i < n_layers * n_heads else f"m{i - n_layers * n_heads}"
def head_slice(layer):
return slice(1 + layer * n_heads, 1 + (layer + 1) * n_heads)
def mlp_index(layer):
return 1 + n_layers * n_heads + layer
def corrupt_sites(corrupt_ids):
"""The 25 node sites on the corrupted prompts, in forward order."""
with torch.no_grad(), model.trace(corrupt_ids):
sites = nnsight.save([])
sites.append(model.transformer.wte.output)
for layer in range(n_layers):
sites.append(blocks[layer].attn.c_proj.input)
sites.append(blocks[layer].mlp.c_proj.input)
return sites
def apply_mask(mask, cf):
"""Blend every node between its live value (mask=1) and its corrupted value (mask=0)."""
m = mask[0]
model.transformer.wte.output = model.transformer.wte.output * m + cf[0] * (1 - m)
for layer in range(n_layers):
m_heads = mask[head_slice(layer)].view(1, 1, n_heads, 1)
z = blocks[layer].attn.c_proj.input
z_heads = z.view(*z.shape[:2], n_heads, d_head)
cf_heads = cf[1 + 2 * layer].view(*z.shape[:2], n_heads, d_head)
blocks[layer].attn.c_proj.input = (z_heads * m_heads + cf_heads * (1 - m_heads)).reshape(z.shape)
m_mlp = mask[mlp_index(layer)]
blocks[layer].mlp.c_proj.input = blocks[layer].mlp.c_proj.input * m_mlp + cf[2 + 2 * layer] * (1 - m_mlp)
def logit_diff(io, s):
"""logit(IO) - logit(S) at the last position, one value per row. Positive = correct."""
logits = model.output.logits[:, -1]
rows = torch.arange(logits.shape[0], device=logits.device)
return logits[rows, io] - logits[rows, s]
print(f"{N_NODES} nodes: {node_name(0)}, {node_name(1)} ... {node_name(n_layers * n_heads)}, "
f"{node_name(mlp_index(0))} ... {node_name(N_NODES - 1)}")
157 nodes: input, a0.h0 ... a11.h11, m0 ... m11
Before learning anything, the two endpoints have to be right: an all-ones mask must reproduce the clean run and an all-zeros mask the corrupted run, bit-exactly, or every number below would be partly tooling.
clean, corrupt, io, s = to_tensors(sample_batch(val_buckets, 8))
cf = corrupt_sites(corrupt)
with torch.no_grad(), model.trace(clean):
ld_clean = logit_diff(io, s).save()
with torch.no_grad(), model.trace(corrupt):
ld_corrupt = logit_diff(io, s).save()
with torch.no_grad(), model.trace(clean):
apply_mask(torch.ones(N_NODES, device=DEVICE), cf)
ld_ones = logit_diff(io, s).save()
with torch.no_grad(), model.trace(clean):
apply_mask(torch.zeros(N_NODES, device=DEVICE), cf)
ld_zeros = logit_diff(io, s).save()
print(f"clean run {ld_clean.mean():+.4f} mask = 1 {ld_ones.mean():+.4f} "
f"max |diff| {(ld_ones - ld_clean).abs().max():.1e}")
print(f"corrupt run {ld_corrupt.mean():+.4f} mask = 0 {ld_zeros.mean():+.4f} "
f"max |diff| {(ld_zeros - ld_corrupt).abs().max():.1e}")
clean run +2.8334 mask = 1 +2.8334 max |diff| 0.0e+00 corrupt run -1.9905 mask = 0 -1.9905 max |diff| 0.0e+00
2. The sigmoid top-$k$ operator¶
Given scores $s \in \mathbb{R}^N$ and a budget $k$, the operator finds the threshold $\tau$ at which
$$\sum_i \sigma\!\left(\frac{s_i - \tau}{T}\right) = k$$
and returns $m_i = \sigma((s_i - \tau)/T)$. The left side is monotone in $\tau$, so fifty steps of bisection pin it down; the mask is soft, sums to $k$, and is a step function at the $k$-th score when $T \to 0$. The backward pass is where the design lives. Differentiating the constraint gives $\partial\tau/\partial s_j = \sigma'_j / \sum_l \sigma'_l$, and so
$$\frac{\partial \mathcal{L}}{\partial s_j} = \frac{\sigma'_j}{T}\left(g_j - \frac{\sum_i g_i\,\sigma'_i}{\sum_i \sigma'_i}\right), \qquad \sigma'_i = m_i(1 - m_i),$$
where $g$ is the upstream gradient. The second term centres the gradient: raising every score together changes nothing, because $\tau$ moves with them. Only relative order is learnable, which is exactly what a ranking method wants. (Wijk et al., 2025, derive the operator; the paper's Appendix A has this Jacobian.)
The implementation is the repo's, condensed. Because the forward re-solves for $\tau$, the right
check of the backward is against finite differences, not against autograd through the bisection —
the bisection's branch choices are piecewise constant, so autograd would see $\tau$ as a function
of scores.min() and scores.max() alone.
class SigmoidTopK(torch.autograd.Function):
@staticmethod
def forward(ctx, scores, k, T, n_iters):
lo, hi = scores.min() - 10 * T, scores.max() + 10 * T
with torch.no_grad():
for _ in range(n_iters): # bisection on the threshold
mid = (lo + hi) / 2
if torch.sigmoid((scores - mid) / T).sum() > k:
lo = mid
else:
hi = mid
mask = torch.sigmoid((scores - (lo + hi) / 2) / T)
ctx.save_for_backward(mask)
ctx.T = T
return mask
@staticmethod
def backward(ctx, grad_output):
(mask,) = ctx.saved_tensors
slope = mask * (1 - mask) # sigma'
centred = grad_output - (grad_output * slope).sum() / slope.sum().clamp_min(1e-8)
return slope / ctx.T * centred, None, None, None
def sigmoid_topk(scores, k, T=0.5, n_iters=50):
return SigmoidTopK.apply(scores, k, T, n_iters)
demo_scores = torch.randn(N_NODES, device=DEVICE)
for k in [1, 5, 20, 80]:
mask = sigmoid_topk(demo_scores, k)
print(f"k = {k:3d} sum(mask) = {mask.sum():8.4f} entries > 0.5: {(mask > 0.5).sum().item():3d}"
f" max entry {mask.max():.3f}")
# finite-difference check of the implicit backward (float64 for the numerics)
probe = torch.randn(12, dtype=torch.float64, requires_grad=True)
assert torch.autograd.gradcheck(lambda x: sigmoid_topk(x, 3.0, 0.5, 100), (probe,), eps=1e-6, atol=1e-5)
print("\ngradcheck against finite differences: passed")
# the centring term: a uniform shift of all scores has zero gradient
sigmoid_topk(probe, 3.0).sum().backward()
print(f"gradient of sum(mask): max |g| = {probe.grad.abs().max():.1e} (sum(mask) is pinned to k)")
k = 1 sum(mask) = 1.0000 entries > 0.5: 0 max entry 0.112 k = 5 sum(mask) = 5.0000 entries > 0.5: 0 max entry 0.430 k = 20 sum(mask) = 20.0000 entries > 0.5: 13 max entry 0.842 k = 80 sum(mask) = 80.0000 entries > 0.5: 79 max entry 0.991 gradcheck against finite differences: passed gradient of sum(mask): max |g| = 0.0e+00 (sum(mask) is pinned to k)
The mass is always exactly $k$, but on scores this close together — a standard normal draw at $T = 0.5$ — it is spread thin: the $k = 1$ mask's largest entry is 0.11, and the $k = 5$ mask has no entry above one half. The operator sharpens as the scores separate, and separating them is what training does; the learned scores below span about $\pm 7$. The finite-difference check passes, and a uniform shift of every score has a gradient of exactly zero.
order = demo_scores.argsort(descending=True)
fig = go.Figure()
for k in [5, 20, 80]:
mask = sigmoid_topk(demo_scores, k)
fig.add_scatter(y=mask[order].cpu().numpy(), mode="lines", name=f"k = {k}")
fig.update_layout(title="Sigmoid top-k on one random score vector (T = 0.5), nodes sorted by score",
xaxis_title="rank", yaxis_title="mask value", height=380)
fig.show()
Same scores, three budgets, three nested masks: the top-5 mask is the top-20 mask with a lower threshold. Nestedness is not something the method learns; it is a property of the operator. What training has to supply is the order.
3. Learning the ranking¶
The training loop is short enough to read whole. Each step draws a batch and a budget
$k \sim \mathrm{Uniform}(1, N)$, builds the soft mask from the current scores, applies it in one
trace, and minimises minus the logit difference — the iso direction: the kept circuit should
retain the clean behaviour. The backward runs inside the trace with with loss.backward():, the
gradient lands on scores outside it, and Adam does the rest. The paper's headline settings are
500 steps, learning rate 0.05, temperature 0.5. We use batches of 16 where the paper's script
defaults to 1; it makes the loss curve readable and costs nothing at this scale.
There is no sparsity term in the loss and no schedule to anneal. The randomised $k$ is the regulariser: every step supervises the ranking at a different cut.
def sample_k(schedule):
u = torch.rand(1).item()
if schedule == "log": # k log-uniform in [1, N]: 1-10 as likely as 10-100
return math.exp(math.log(N_NODES) * u)
return 1 + (N_NODES - 1) * u # k uniform in [1, N]: the paper's headline
def learn_scores(steps=500, batch_size=16, lr=0.05, T=0.5, schedule="uniform", fixed_k=None):
scores = torch.zeros(N_NODES, device=DEVICE, requires_grad=True)
optimizer = torch.optim.Adam([scores], lr=lr)
log = []
for step in range(steps):
clean, corrupt, io, s = to_tensors(sample_batch(train_buckets, batch_size))
cf = corrupt_sites(corrupt)
k = fixed_k if fixed_k is not None else sample_k(schedule)
mask = sigmoid_topk(scores, k, T)
with model.trace(clean):
apply_mask(mask, cf)
loss = -logit_diff(io, s).mean()
with loss.backward():
pass
loss_value = nnsight.save(loss.item())
optimizer.step()
optimizer.zero_grad()
log.append((step, k, loss_value))
return scores.detach(), log
t0 = time.time()
mattr_scores, mattr_log = learn_scores()
print(f"500 steps in {time.time() - t0:.0f}s")
ranking = mattr_scores.argsort(descending=True)
print("top 12:", [node_name(i) for i in ranking[:12].tolist()])
print("bottom 5:", [node_name(i) for i in ranking[-5:].tolist()])
500 steps in 14s top 12: ['a8.h10', 'a5.h5', 'a7.h9', 'a9.h9', 'a3.h0', 'm0', 'a8.h6', 'input', 'm4', 'a6.h9', 'm1', 'm2'] bottom 5: ['a5.h8', 'm6', 'm7', 'a11.h10', 'a10.h7']
steps, ks, losses = map(np.array, zip(*mattr_log))
train_ld = -losses
print(f"train logit diff, first 50 steps (all k): {train_ld[:50].mean():+.2f}")
for lo, hi in [(1, 20), (20, 80), (80, 158)]:
late = train_ld[-150:][(ks[-150:] >= lo) & (ks[-150:] < hi)]
print(f"last 150 steps, {lo:3d} <= k < {hi:3d}: mean {late.mean():+.2f} (n = {len(late)})")
print(f"clean model on the same split, roughly: {ld_clean.mean():+.2f}")
fig = px.scatter(
x=ks, y=-losses, color=steps, log_x=True,
color_continuous_scale="Viridis",
labels=dict(x="sampled k (nodes kept clean)", y="train logit diff", color="step"),
title="Every step supervises a different sparsity: the loss surface fills in across k",
height=420,
)
fig.update_traces(marker=dict(size=5))
fig.show()
train logit diff, first 50 steps (all k): +1.29 last 150 steps, 1 <= k < 20: mean +0.20 (n = 19) last 150 steps, 20 <= k < 80: mean +9.19 (n = 49) last 150 steps, 80 <= k < 158: mean +9.49 (n = 82) clean model on the same split, roughly: +2.83
Each point is one step, at the $k$ it happened to draw, on that step's training batch. The first fifty steps (dark) average +1.3 whatever $k$ was — a zero-initialised score vector is a uniform mask, the same fixed blend of clean and corrupted at every node. By the last 150 steps (yellow) the picture has split: budgets of 20 and above yield about +9, three times the clean model's +2.8, while budgets below 20 still average +0.2. That overshoot past the clean model is real, and it returns in the next section.
4. Faithfulness at every sparsity¶
MIB scores a ranking by cutting it at ten proportions $p \in \{0.001, \dots, 0.5, 1\}$, keeping the top $\lfloor pN \rfloor$ nodes as a hard mask, and measuring how much of the clean behaviour the circuit retains, normalised between the empty circuit and the full model:
$$\mathsf{Faith}(k) = \frac{\overline{\mathrm{LD}}(k) - \overline{\mathrm{LD}}(0)}{\overline{\mathrm{LD}}(N) - \overline{\mathrm{LD}}(0)}$$
CPR is the trapezoidal area under $\mathsf{Faith}$ against $p$, and the paper's second metric, Compactness, is the area under the circuit's accuracy (fraction of examples with a positive logit difference) against $\log p$, normalised to $[0, 1]$, which rewards being right at few nodes. Both on the 200 held-out validation examples, whose corrupted activations are cached once.
val_batches = []
for examples in val_buckets.values():
for start in range(0, len(examples), 32):
clean, corrupt, io, s = to_tensors(examples[start:start + 32])
val_batches.append((clean, io, s, corrupt_sites(corrupt)))
def per_example_ld(mask):
out = []
for clean, io, s, cf in val_batches:
with torch.no_grad(), model.trace(clean):
if mask is not None:
apply_mask(mask, cf)
out.append(logit_diff(io, s).save())
return torch.cat(out)
ld_full = per_example_ld(None).mean()
ld_empty = per_example_ld(torch.zeros(N_NODES, device=DEVICE)).mean()
print(f"full model: mean logit diff {ld_full:+.3f} empty circuit: {ld_empty:+.3f}")
PERCENTAGES = (0.001, 0.002, 0.005, 0.01, 0.02, 0.05, 0.1, 0.2, 0.5, 1.0)
LOG_P = [math.log(p) for p in PERCENTAGES]
def evaluate(scores):
order = scores.argsort(descending=True)
faith, acc = [], []
for p in PERCENTAGES:
hard = torch.zeros(N_NODES, device=DEVICE)
hard[order[:int(p * N_NODES)]] = 1
ld = per_example_ld(hard)
faith.append(((ld.mean() - ld_empty) / (ld_full - ld_empty)).item())
acc.append((ld > 0).float().mean().item())
trapezoid = lambda xs, ys: sum((xs[i + 1] - xs[i]) * (ys[i] + ys[i + 1]) / 2 for i in range(len(xs) - 1))
return dict(faith=faith, acc=acc, cpr=trapezoid(PERCENTAGES, faith),
compactness=trapezoid(LOG_P, acc) / (LOG_P[-1] - LOG_P[0]))
results = {}
results["MAttr (uniform k)"] = evaluate(mattr_scores)
results["random order"] = evaluate(torch.rand(N_NODES, device=DEVICE))
def show(names):
print(f"{'':>26} {'CPR':>6} {'Compact.':>9} faithfulness at " + " ".join(f"{p:g}" for p in PERCENTAGES))
for name in names:
r = results[name]
print(f"{name:>26} {r['cpr']:6.3f} {r['compactness']:9.3f} " + " ".join(f"{f:5.2f}" for f in r["faith"]))
show(results)
full model: mean logit diff +2.902 empty circuit: -2.606
CPR Compact. faithfulness at 0.001 0.002 0.005 0.01 0.02 0.05 0.1 0.2 0.5 1
MAttr (uniform k) 1.723 0.409 0.00 0.00 0.00 0.00 0.00 0.00 0.95 2.15 2.43 1.00
random order 0.247 0.109 0.00 0.00 0.00 0.00 -0.00 -0.00 -0.00 -0.00 -0.01 1.00
def plot_curves(names, title):
fig = go.Figure()
for name in names:
fig.add_scatter(x=PERCENTAGES, y=results[name]["faith"], mode="lines+markers", name=name)
fig.add_hline(y=1, line_dash="dot", annotation_text="full model")
fig.update_layout(title=title, xaxis_title="fraction of nodes kept (log)", xaxis_type="log",
yaxis_title="faithfulness", height=430, legend=dict(x=0.02, y=0.98))
fig.show()
plot_curves(results, "Faithfulness of the top-k circuit, GPT-2 / IOI, 200 validation examples")
Reading the MAttr row left to right: below 10% the ranking recovers nothing; the top 10% of nodes (15 of 157) recover 95% of the clean logit difference; at 20% and 50% the circuit is more confident than the full model, faithfulness 2.15 and 2.43. A random order stays at zero until every node is in.
The area under that curve, CPR 1.72, is the paper's number for this cell: 1.70 on the test split at seed 42, with a seed standard deviation of 0.009. Compactness 0.41 against the paper's 0.40. Two features of the curve need explaining before any comparison means anything.
Faithfulness above 1 is a known flaw of CPR — the paper discusses it, and it is the reason Compactness exists. The logit difference is unbounded, and the half of the model that the ranking leaves out includes components that push toward the subject: the negative name movers, and as section 6 shows, several late MLPs. Patching those to the corrupted run — where their job is to suppress the corrupted answer, which is the subject — removes that push. A 50% circuit is not a better IOI solver than GPT-2; it is a more confident one.
The zeros are not a failure of the ranking. They are a property of the intervention, and the cleanest way to see it is to take the input node out by hand.
4.1 The input node is a master switch¶
MIB's node-level intervention keeps a node live, not clean. The input embedding is the first thing every other node reads, so once it is patched to the corrupted prompt, every kept node computes its live value from corrupted embeddings and lands exactly on its corrupted value. A circuit that leaves the input node out is the empty circuit whatever else it contains — not approximately, but bit for bit. Under the uniform schedule MAttr learns a high score for the input node but not the highest, and every cut above its rank scores exactly zero.
input_rank = (mattr_scores > mattr_scores[0]).sum().item() + 1
print(f"MAttr ranks the input node {input_rank}th of {N_NODES}")
for k in [3, 7, 15]:
top = ranking[:k]
without = torch.zeros(N_NODES, device=DEVICE); without[top] = 1; without[0] = 0
with_input = without.clone(); with_input[0] = 1
print(f"top-{k:<2d} without the input node: LD {per_example_ld(without).mean():+.3f}"
f" (empty circuit {ld_empty:+.3f}) with it: {per_example_ld(with_input).mean():+.3f}")
MAttr ranks the input node 8th of 157
top-3 without the input node: LD -2.606 (empty circuit -2.606) with it: -2.590
top-7 without the input node: LD -2.606 (empty circuit -2.606) with it: -1.540
top-15 without the input node: LD -2.606 (empty circuit -2.606) with it: +2.638
Without the input node, the top-3, top-7 and top-15 circuits all sit at −2.606 — the empty circuit to the last digit. Put the one node back and the top-15 circuit recovers 95% of the behaviour, the top-7 about a fifth. The ranking is fine; the cliff in the metric is entirely about where one node lands.
Why does training not put it first? Under a uniform $k$ the input node only decides the loss on the steps where $k$ is drawn below its rank — a few percent of them. Everywhere else its mask entry saturates at 1, its gradient is zero, and the order among the top handful goes almost unsupervised. The paper's Appendix D compares the uniform schedule against a log-uniform one, which spends as many steps on $k \in [1, 10]$ as on $[10, 100]$, and reports it trades a little CPR for a little Compactness. That is the schedule to try.
t0 = time.time()
mattr_log_scores, _ = learn_scores(schedule="log")
print(f"log-k schedule: 500 steps in {time.time() - t0:.0f}s, "
f"input node ranked {(mattr_log_scores > mattr_log_scores[0]).sum().item() + 1}th")
print("top 12:", [node_name(i) for i in mattr_log_scores.argsort(descending=True)[:12].tolist()])
results["MAttr (log k)"] = evaluate(mattr_log_scores)
show(["MAttr (uniform k)", "MAttr (log k)"])
plot_curves(["MAttr (uniform k)", "MAttr (log k)", "random order"], "Uniform vs log-uniform k-schedule")
log-k schedule: 500 steps in 14s, input node ranked 5th top 12: ['a8.h10', 'm0', 'a5.h5', 'a7.h9', 'input', 'a3.h0', 'a9.h9', 'a8.h6', 'm1', 'a9.h6', 'm2', 'm4']
CPR Compact. faithfulness at 0.001 0.002 0.005 0.01 0.02 0.05 0.1 0.2 0.5 1
MAttr (uniform k) 1.723 0.409 0.00 0.00 0.00 0.00 0.00 0.00 0.95 2.15 2.43 1.00
MAttr (log k) 1.728 0.412 0.00 0.00 0.00 0.00 0.00 0.15 0.95 2.15 2.43 1.00
The log schedule moves the input node from 8th to 5th, so the 5% cut (7 nodes) goes from 0 to 0.15; the rest of the curve is identical and both areas move by less than 0.005 — outside the top few, the two schedules learn the same ranking. Neither puts the input node first after 500 steps. The paper's own evaluation script sidesteps the cliff in the other configuration: when the input node is not learned, it is assigned the maximum score so that it is in every circuit. We keep the learned rank in what follows, since learning it is what the headline recipe does.
5. Baselines on the same harness¶
A ranking method is only as good as what it is compared against, and comparisons across
implementations are where attribution results go wrong. Everything below reuses apply_mask,
the same corrupted caches and the same evaluator, so the only thing that varies is the ranking.
Baseline scores are computed on 320 training examples, the same split MAttr learned from.
Gradients in mask space. With the blend $h = m\,h(\text{live}) + (1-m)\,h(\text{corrupted})$,
the derivative of the metric with respect to $m_H$ at $m = 1$ is
$(h(\text{clean}) - h(\text{corrupted})) \cdot \partial \mathcal{L}/\partial H$ — exactly
input × gradient, the attribution patching estimate of a node's effect, for all 157 nodes
in one backward pass. Averaging the same gradient at $m = t$ for $t$ along $[0, 1]$ is
integrated gradients along the mask path (the paper's Appendix A shows this is also what
MAttr's very first step computes in expectation). Both are two lines here: build a mask tensor
that requires grad, read its .grad after the backward.
Interchange interventions. The classical baseline: patch one node to its corrupted value,
keep everything else clean, score the drop. That is one forward pass per node, and it is the
sweep pattern — one tracer.invoke per node, forty to a trace.
attr_examples, attr_buckets = load_split("train", 320)
attr_batches = []
for examples in attr_buckets.values():
for start in range(0, len(examples), 16):
clean, corrupt, io, s = to_tensors(examples[start:start + 16])
attr_batches.append((clean, io, s, corrupt_sites(corrupt)))
n_attr = sum(batch[0].shape[0] for batch in attr_batches)
def mask_gradient(t):
"""d(logit diff)/d(mask) at mask = t for every node, summed over examples."""
total = torch.zeros(N_NODES, device=DEVICE)
for clean, io, s, cf in attr_batches:
mask = torch.full((N_NODES,), t, device=DEVICE, requires_grad=True)
with model.trace(clean):
apply_mask(mask, cf)
with logit_diff(io, s).sum().backward():
pass
total += mask.grad
return total / n_attr
t0 = time.time()
ixg_scores = mask_gradient(1.0)
print(f"input x gradient : {len(attr_batches)} backward passes, {time.time() - t0:.0f}s")
t0 = time.time()
ig_scores = sum(mask_gradient(t) for t in torch.linspace(0.1, 1.0, 10)) / 10
print(f"integrated gradients : {10 * len(attr_batches)} backward passes, {time.time() - t0:.0f}s")
input x gradient : 30 backward passes, 1s
integrated gradients : 300 backward passes, 5s
def interchange_scores(chunk=40):
"""Drop in logit diff when one node is patched to its corrupted value, per node."""
total = torch.zeros(N_NODES, device=DEVICE)
for clean, io, s, cf in attr_batches:
with torch.no_grad(), model.trace(clean):
ld_clean = logit_diff(io, s).sum().save()
for start in range(0, N_NODES, chunk):
nodes = range(start, min(start + chunk, N_NODES))
with torch.no_grad(), model.trace() as tracer:
patched = nnsight.save([])
for node in nodes:
with tracer.invoke(clean):
mask = torch.ones(N_NODES, device=DEVICE)
mask[node] = 0
apply_mask(mask, cf)
patched.append(logit_diff(io, s).sum())
for node, ld in zip(nodes, patched):
total[node] += ld_clean - ld
return total / n_attr
t0 = time.time()
intinv_scores = interchange_scores()
print(f"interchange interventions: {N_NODES * len(attr_batches)} forward passes "
f"(one per node per batch), {time.time() - t0:.0f}s")
interchange interventions: 4710 forward passes (one per node per batch), 36s
results["input x gradient"] = evaluate(ixg_scores)
results["integrated gradients"] = evaluate(ig_scores)
results["interchange intervention"] = evaluate(intinv_scores)
for name, scores in [("input x gradient", ixg_scores), ("integrated gradients", ig_scores),
("interchange intervention", intinv_scores)]:
print(f"{name:>26} input node at rank #{(scores > scores[0]).sum().item() + 1:<3d} "
f"top 8: {[node_name(i) for i in scores.argsort(descending=True)[:8].tolist()]}")
print()
show(["MAttr (uniform k)", "MAttr (log k)", "integrated gradients", "interchange intervention",
"input x gradient", "random order"])
plot_curves(["MAttr (uniform k)", "MAttr (log k)", "integrated gradients", "interchange intervention",
"input x gradient", "random order"], "All rankings, one evaluator")
input x gradient input node at rank #152 top 8: ['a8.h6', 'a5.h5', 'a8.h10', 'a9.h9', 'a7.h9', 'a6.h9', 'a7.h3', 'a3.h0']
integrated gradients input node at rank #1 top 8: ['input', 'm0', 'a5.h5', 'a8.h6', 'a8.h10', 'a7.h9', 'a9.h9', 'a3.h0']
interchange intervention input node at rank #1 top 8: ['input', 'm0', 'a5.h5', 'a8.h6', 'a8.h10', 'a7.h9', 'a9.h9', 'a6.h9']
CPR Compact. faithfulness at 0.001 0.002 0.005 0.01 0.02 0.05 0.1 0.2 0.5 1
MAttr (uniform k) 1.723 0.409 0.00 0.00 0.00 0.00 0.00 0.00 0.95 2.15 2.43 1.00
MAttr (log k) 1.728 0.412 0.00 0.00 0.00 0.00 0.00 0.15 0.95 2.15 2.43 1.00
integrated gradients 1.493 0.395 0.00 0.00 0.00 0.00 -0.00 0.10 0.67 1.90 2.02 1.00
interchange intervention 1.470 0.377 0.00 0.00 0.00 0.00 -0.00 0.10 0.52 1.81 2.04 1.00
input x gradient 0.250 0.110 0.00 0.00 0.00 0.00 0.00 0.00 0.00 0.00 0.00 1.00
random order 0.247 0.109 0.00 0.00 0.00 0.00 -0.00 -0.00 -0.00 -0.00 -0.01 1.00
Six rankings, one evaluator.
- MAttr, CPR 1.72 from 500 backward passes on batches of 16.
- Integrated gradients, 1.49 and interchange interventions, 1.47. Both put the input node
and
m0first, so both score at the 5% cut where uniform-$k$ MAttr does not — and then fall behind: 0.67 and 0.52 against 0.95 at 10%, 1.90 and 1.81 against 2.15 at 20%. The top-8 lists are nearly the same heads in a different order; the difference is which fifteen go in together. Single-node effects, exact or first-order, do not compose into a good set, which is the whole reason the paper frames attribution as nested subsets rather than as per-component scores. - Input × gradient, 0.25, indistinguishable from random. The gradient at the clean point ranks S-inhibition and induction heads first and the input node in the bottom ten, so every cut below 100% is the empty circuit. The paper's Table 2 has I×G behind MAttr on every MIB cell.
Cost, in the currency the paper uses: MAttr spent 500 backward passes; integrated gradients 300; input × gradient 30; interchange interventions 4,710 forward passes, a number that scales with the node count while the others do not.
5.1 A mask learned at one budget¶
The cleanest test of the Matryoshka idea is to switch it off. Same operator, same optimiser, same 500 steps, but $k$ fixed at 16 (a tenth of the nodes) on every step — the way a conventional mask-learning run commits to one sparsity. The scores that come out are still a full ranking, so the evaluator can cut them anywhere.
fixed_scores, _ = learn_scores(fixed_k=16)
results["fixed k = 16"] = evaluate(fixed_scores)
print("top 12:", [node_name(i) for i in fixed_scores.argsort(descending=True)[:12].tolist()])
print()
show(["MAttr (uniform k)", "fixed k = 16"])
plot_curves(["MAttr (uniform k)", "fixed k = 16"], "Randomised k vs one fixed budget")
top 12: ['a7.h9', 'a8.h10', 'm0', 'a5.h5', 'input', 'a9.h9', 'a3.h0', 'a8.h6', 'a9.h6', 'm4', 'a6.h9', 'm1']
CPR Compact. faithfulness at 0.001 0.002 0.005 0.01 0.02 0.05 0.1 0.2 0.5 1
MAttr (uniform k) 1.723 0.409 0.00 0.00 0.00 0.00 0.00 0.00 0.95 2.15 2.43 1.00
fixed k = 16 0.881 0.399 0.00 0.00 0.00 0.00 0.00 0.15 0.97 0.97 0.89 1.00
Fixed at 16, the mask learns a circuit that is faithful at its own budget — 0.97 at 10%, fifteen nodes — and nowhere else: 0.97 at 20% and 0.89 at 50%, below the full model, where the randomised run scores 2.15 and 2.43. CPR 0.88 against 1.72, on nearly the same top twelve. What the fixed-budget run never learned is the order of everything else, because no step ever asked: nodes 17 through 157 are ranked by noise, so growing the circuit adds harmful components — the negative name movers among them — as readily as helpful ones. A conventional mask-learning method needs one run per budget, and the paper measures (its Table 1) that those runs are not nested in each other either. Randomising $k$ buys the whole curve for the price of one run.
6. What the ranking says about IOI¶
A faithfulness curve says the ranking works; it does not say what is in it. Wang et al. (2023) named the heads of the IOI circuit by function, and the paper's Appendix H checks its rankings against that list. We do the same with the uniform-$k$ scores.
head_scores = mattr_scores[1:1 + n_layers * n_heads].view(n_layers, n_heads).cpu()
fig = px.imshow(
head_scores.numpy(), color_continuous_scale="RdBu", color_continuous_midpoint=0,
labels=dict(x="head", y="layer", color="score"),
title="MAttr scores per attention head (input node and MLPs not shown)", height=480,
)
fig.show()
mlp_scores = mattr_scores[mlp_index(0):].cpu()
print("MLP scores by layer:", " ".join(f"{v:+.1f}" for v in mlp_scores.tolist()))
print(f"input node score: {mattr_scores[0]:+.2f}")
MLP scores by layer: +4.5 +3.5 +3.4 +2.6 +3.6 +2.5 -6.1 -6.3 -3.4 +0.9 +1.5 -5.5 input node score: +4.42
IOI_CIRCUIT = {
"name movers": [(9, 9), (9, 6), (10, 0)],
"S-inhibition": [(7, 3), (7, 9), (8, 6), (8, 10)],
"induction": [(5, 5), (5, 8), (5, 9), (6, 9)],
"duplicate token": [(0, 1), (0, 10), (3, 0)],
"previous token": [(2, 2), (4, 11)],
"negative name movers": [(10, 7), (11, 10)],
"backup name movers": [(9, 0), (9, 7), (10, 1), (10, 2), (10, 6), (10, 10), (11, 2), (11, 9)],
}
rank_of = torch.empty(N_NODES, dtype=torch.long)
rank_of[ranking.cpu()] = torch.arange(1, N_NODES + 1)
for role, heads in IOI_CIRCUIT.items():
cells = []
for layer, head in heads:
index = 1 + layer * n_heads + head
cells.append(f"a{layer}.h{head}: #{rank_of[index].item():<3d} ({mattr_scores[index]:+.1f})")
print(f"{role:>22} " + " ".join(cells))
print()
for label, scores in [("MAttr", mattr_scores), ("integrated gradients", ig_scores),
("interchange intervention", intinv_scores)]:
r = torch.empty(N_NODES, dtype=torch.long)
r[scores.argsort(descending=True).cpu()] = torch.arange(1, N_NODES + 1)
print(f"{label:>26}: a9.h6 at rank #{r[1 + 9 * n_heads + 6].item():<3d} a11.h2 at rank #{r[1 + 11 * n_heads + 2].item()}")
name movers a9.h9: #4 (+4.8) a9.h6: #13 (+3.1) a10.h0: #17 (+2.2)
S-inhibition a7.h3: #19 (+2.0) a7.h9: #3 (+4.8) a8.h6: #7 (+4.5) a8.h10: #1 (+5.1)
induction a5.h5: #2 (+4.9) a5.h8: #153 (-5.9) a5.h9: #23 (+1.5) a6.h9: #10 (+3.6)
duplicate token a0.h1: #14 (+2.9) a0.h10: #130 (-3.1) a3.h0: #5 (+4.5)
previous token a2.h2: #56 (-1.1) a4.h11: #51 (-0.8)
negative name movers a10.h7: #157 (-6.7) a11.h10: #156 (-6.5)
backup name movers a9.h0: #38 (-0.3) a9.h7: #24 (+1.5) a10.h1: #21 (+1.6) a10.h2: #126 (-2.7) a10.h6: #25 (+1.3) a10.h10: #18 (+2.1) a11.h2: #152 (-5.7) a11.h9: #34 (+0.2)
MAttr: a9.h6 at rank #13 a11.h2 at rank #152
integrated gradients: a9.h6 at rank #150 a11.h2 at rank #155
interchange intervention: a9.h6 at rank #152 a11.h2 at rank #155
The heatmap has the IOI circuit's geography: S-inhibition in layers 7–8, name movers in 9–10,
the induction and duplicate-token heads at a5.h5 and a3.h0, and two deep-blue cells at
a10.h7 and a11.h10. Against Wang et al.'s list:
- Ranks 1–5 are the circuit's spine: S-inhibition
a8.h10anda7.h9, inductiona5.h5, name movera9.h9, duplicate-tokena3.h0.a8.h6is 7th,a6.h910th,a9.h613th. - The negative name movers are last.
a10.h7is 157th anda11.h10156th, with the most negative scores of any node. Under the iso intervention a negative score means this node is better corrupted than clean, and a head whose job is to suppress the IO answer is exactly that. The ranking is signed, and both ends of it are informative. - The paper's two observations reproduce. Name mover
a9.h6scores clearly positive under MAttr, rank 13, while integrated gradients and interchange interventions put it at 150 and 152 — the paper reports exactly this split, MAttr and AttnLRP on one side and IG and causal interventions on the other. And backup name movera11.h2is 152nd of 157, "consistently found to be highly disfavoured", as the paper puts it. - Previous-token heads and most backup name movers sit near zero, which is the right answer: backup heads only act once a name mover is ablated, and the denoising sweep never ablates one.
- Two things the published list does not predict:
a5.h8, an induction head in Wang et al., is 153rd with a strongly negative score, and MLPs 6, 7 and 11 are as negative as the negative name movers. We did not chase either.
Caveats¶
- One cell. GPT-2 / IOI at the node level is the paper's smallest configuration and the one with a published circuit to check against. The CPR matched; nothing here speaks to the other eleven MIB cells, the edge level, the neuron and SAE bases, or the parameter-attribution results.
- Our evaluator, not MIB's. Faithfulness, the grid and both areas follow MIB's definitions, but MIB's script runs the circuit through TransformerLens with its own tokenisation, on the full split rather than 200 examples. Agreement at the second decimal with the paper's table is partly luck; the ranking-level conclusions are not.
- The input node. Every zero in the tables above is the master-switch effect. It is MIB's semantics, not ours, but a different intervention — setting kept nodes to their clean values rather than leaving them live — would not have it, and would give different curves.
- Batch 16, one seed, no sweep. The paper's script defaults to batch 1 and its learning rate was tuned on the validation split; we used its reported values as given. The paper's seed variance on this cell is 0.009 CPR.
- Simplest baselines. Input × gradient at the clean point, integrated gradients along the mask path rather than EAP-IG's activation path, and single-node interchange interventions. The paper's stronger gradient methods (EAP-IG, AttnLRP, RelP) and its mask-learning baselines are not here.
Conclusion¶
🎉 Matryoshka attribution is short enough to fit in three cells: an autograd Function for the
sigmoid top-$k$ mask, one function that blends 25 sites between live and corrupted values, and a
loop that draws a budget, applies the mask and steps Adam. The 157 scores that come out are a
ranking that is faithful at every cut, match the paper's CPR for this cell, beat gradient and
interchange baselines evaluated on identical code, and read as the IOI circuit with the negative
heads at the bottom.
In nnsight the mask is applied by replacing c_proj.input and wte.output with a blend, so
autograd sees a clean graph; the backward runs inside the trace with with loss.backward(): and
the gradient arrives on a tensor that never entered the model; and the interchange baseline is the
same apply_mask under one tracer.invoke per node. There are no hooks to register or remove,
and the corrupted activations are a saved list from a second trace.
Related: Activation Patching
for the intervention this builds on, Attribution Patching
for the input × gradient baseline done the usual way, Gradients
for the with loss.backward(): form, and Batching for the
invoke pattern used in the interchange sweep.
References¶
- Arora, Acharya, Hu, Zhang, Goodman, Jurafsky, Potts, Matryoshka attribution: Learning to attribute language model outputs to representations and weights, 2026. Code: aryamanarora/matryoshka-attribution
- Mueller et al., MIB: A Mechanistic Interpretability Benchmark, 2025. Data: mib-bench/ioi
- Wijk, Vinuesa, Azizpour, Differentiable top-k: from one-hot to k-hot, Workshop on Differentiable Systems and Scientific Machine Learning @ EurIPS 2025
- Kusupati et al., Matryoshka Representation Learning, NeurIPS 2022
- Wang, Variengien, Conmy, Shlegeris, Steinhardt, Interpretability in the Wild: a Circuit for Indirect Object Identification in GPT-2 small, ICLR 2023