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()