Matryoshka attribution: Learning to attribute language model outputs to representations and weights
Aryaman Arora, Kirill Acharya, Nathan Hu, Yanzhe Zhang, Noah Goodman, Dan Jurafsky, Christopher Potts
cs.CL, cs.LG
2026-09-22
MAttr learns a nested component ranking, scores 5.6 on MIB's secret test vs 1.95, and undoes Llama 3.1 8B refusals by restoring 1% of weights while GSM8K and MMLU stay close.
Attributing a language model's output to its internals still splits into three imperfect tools. Causal interventions swap one component at a time and need a forward pass per variable; covering combinations is 2^n. Gradient methods scan everything in one backward pass and often miss computations that actually matter. Learnable masks can be trained, but they usually add a sparsity penalty plus a straight-through estimator or Gumbel noise, and they are brittle.
The Stanford group restates the job as learning an order over internal variables. Add variables back to the unintervened set in that order so a downstream loss stays small. Nested subsets at every sparsity is the Matryoshka property.
MAttr stores a scalar score per variable. Each step samples a budget k uniformly, then turns the scores into a sigmoid top-k mask whose entries sum to k. High-score variables keep their original computation. The rest are softly interpolated between the base-run activation and the source-run activation. That interpolation is a soft interchange intervention.
Two design choices do most of the work.
Because mask values are monotone in k, the learned object is a global ranking. After one gradient step the scores recover centered Input×Gradient; after one SGD step they recover a path-reweighted Integrated Gradients. Further training can leave IG's axioms behind. Bilodeau et al. showed that any complete and linear attribution method can fail to beat chance on local counterfactuals. MAttr learns an order that is invariant to constant shifts, and it directly optimizes the task loss.
Representation attribution uses paired base/counterfactual inputs and a logit-difference loss. Parameter attribution interpolates in weight space instead: θ = θb + α ⊙ (θs − θb). The loss becomes unclipped GRPO, with a judge scoring sampled completions and advantages flowing back into S. The basis is each dmodel-shaped row of a weight matrix, plus layernorm vectors.
On the Mechanistic Interpretability Benchmark circuit-localization track, edge-level MAttr scores 5.6 on the official secret test set against 1.95 for second place. Per model-task pair the lift over the runner-up ranges from 45.9% to 343.0%. On the public test set it beats retuned I×G, IG, Expected Gradients, AttnLRP, interchange interventions, DBM, and Node Pruning; a paired Wilcoxon test with Holm–Bonferroni correction is p < 0.05 against every baseline. Sweeping sparsity for DBM and Node Pruning costs about 50× MAttr and still loses.
They also report compactness, the area under interchange-intervention accuracy on a log sparsity grid, which rewards small circuits that already preserve the behavior. MAttr has the highest CPR at every basis they tried. Compactness matches or beats IG in most settings. Sampling k log-uniformly can push compactness past IG at a small CPR cost.
Transfer on Llama is concrete: a ranking trained on subtraction still works on addition and month arithmetic, repeatedly putting heads a15.h13 and a16.h21 at the top, in line with prior circuit work. Rankings disagree with weak gradient methods mainly on MLP blocks, not attention heads.
The parameter experiment is sharper. Restoring 1% of Llama 3.1 8B Instruct rows to the base checkpoint moves StrongREJECT from 2.6 to 84.0 and SORRY-Bench from 27.1 to 99.3, while GSM8K goes 81.0 to 79.5, IFEval 74.7 to 71.7, and MMLU 69.1 to 69.5. Abliteration changes 27.6% of parameters to get a similar refusal drop. Expected Gradients at the same 1% budget tanks IFEval to 47.5. On the 1B model, MAttr changes 2% of parameters: StrongREJECT 75.8, MMLU still 43.6.
| method | params changed | StrongREJECT | GSM8K | MMLU |
| Instruct (8B) | 0 | 2.6 | 81.0 | 69.1 |
| Abliteration | 27.6% | 68.6 | 81.5 | 68.6 |
| MAttr | 1.0% | 84.0 | 79.5 | 69.5 |
| Expected Gradients | 1.0% | 59.5 | 76.5 | 67.8 |
An appendix on emergent misalignment recovers almost all held-out loss with 0.1% of the weight delta. On ViT-B/16, Faith log-AUC is 5.2, and 2% of patches already make basenji the argmax class.
Interpretability becomes a backprop-compatible objective. For circuit localization, MAttr is currently first on the official MIB board, and one ranking covers every sparsity. For people who ship a base checkpoint next to an instruct checkpoint, the 8B result is a measurement: refusal can sit in about 1% of the weight delta, while GSM8K, IFEval, and MMLU mostly do not. That is an auditing tool and a refusal-removal recipe at the same time.
The algorithm is incremental. The contribution is collapsing three attribution families into one trainable objective, then running that objective on weight deltas with an RL reward.
The authors only learn a dataset-level scalar score S. They do not train a featurizer that would make scores example- or checkpoint-specific. Representation attribution still needs paired counterfactuals; parameter attribution needs a reward model. CPR is a flawed headline number: logit difference is unbounded, the evaluation grid integrates linearly on a log-spaced sparsity axis, and a random order already picks up about 0.25 CPR between p=0.5 and 1.0. The 5.6 vs 1.95 gap should not be read as a 2.9× circuit-quality gap.
Refusal training uses the StrongREJECT held-out split and evaluates on the same distribution plus SORRY-Bench. There is no evidence that 1% still suffices on a disjoint harm distribution. GRPO here has no clip and no KL. Row-level sparsity is shown at 1B and 8B, not at 70B.