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
Quick verdict
MAttr learns one nested attribution ordering via sigmoid top-k; tops MIB (5.6 vs 1.95) and ablates refusal with 1% of weights; rigorous, partly self-metric'd.
Summary
The paper introduces Matryoshka Attribution (MAttr), a mask-learning method that attributes language-model outputs to internal components. Rather than training one mask per sparsity with L0 penalties, MAttr learns a single score per variable: each step samples a random budget k, builds a differentiable sigmoid top-k mask, softly intervenes on non-top-k variables, and minimises the task loss across all sparsities simultaneously. The authors derive that the first SGD step recovers path-reweighted Integrated Gradients, unifying mask learning with gradient attribution. On the MIB Circuit Localization track (12 task–model pairs, 4 models), edge-level MAttr reports the top official leaderboard score (5.6 vs 1.95 for the runner-up); the paper extends MIB to MIB+ with neuron- and SAE-level bases and introduces a Compactness metric to fix CPR's flaws. Significance is assessed with paired Wilcoxon tests and Holm–Bonferroni correction. Extended to weight deltas via GRPO, MAttr restores 1% (8B) or 2% (1B) of Llama Instruct weights to base, removing refusals while preserving GSM8K, IFEval and MMLU, versus >27.6% parameter changes for baselines. Limitations include a self-defined headline metric, sensitivity to Adam's ε (needing 10⁻² at granular bases), higher cost than gradient baselines (about 3k vs 0.1–1k backward passes), and judge-based refusal evaluation on a small StrongREJECT subset.
Review complete
Assessment at a glance
5
risk flags identified
● high 0
● medium 3
● low 2
5 strengths · 0 weaknesses · 5 author questions
AI text signal = Binoculars zero-shot detector score over 20 sampled passages (Fast-DetectGPT co-reported) — absence does not prove human authorship, and a flag does not imply misconduct.
Levels reflect verification complexity — not a quality score. Evolve never rates your paper.
Key claims under review
- Edge-level MAttr ranks number 1 on the official MIB leaderboard with an average secret-test-set score of 5.6 vs 1.95 for the runner-up.
- The first SGD step of MAttr equals path-reweighted centred Integrated Gradients, unifying mask learning and gradient-based attribution.
- MAttr identifies sparse, task-transferable circuits across node, edge, MLP-neuron and SAE bases.
- Restoring 1% (Llama-3.1-8B) or 2% (1B) of weights to base removes refusals while maintaining capabilities, versus >27.6% parameter changes for baselines.
Strengths
Theoretically grounded unification of mask learning and gradient attribution
Appendix A derives that the first-step gradient is centred I×G at the masked forward and that one SGD step yields path-reweighted Integrated Gradients; §A.4 explains why the Bilodeau et al. (2024) impossibility theorem does not apply (MAttr is neither complete nor linear).
Unusually honest baseline treatment for a mask-learning paper
I×G, IG (m=5/10/30), EG, RelP variants, GIM, AttnLRP, IntInv, DBM and Node/Edge Pruning are re-tuned with sweeps in §E.1, with cost accounting (Fig. 2 bottom) and ≈50×-cost multi-sparsity sweep variants of DBM/NP that still lose (§E.2, Fig. 12).
Extensive ablations of every design knob
Practically consequential weight-level result
Self-critical metric analysis
Full Review Report
$1.99 · one-timeUnlock the complete preliminary analysis Evolve produced for this paper.
Report preview
Scope & Claims
The paper introduces Matryoshka Attribution (MAttr, §3), a mask-learning method for attributing LLM outputs to internal components. Two design choices define it: (1) one score per variable trained across all sparsity budgets simultaneously by sampling k ~ Uniform(1, |S|−1) each step (the Matryoshka trick), and (2) a differentiable sigmoid top-k mask (Def. 2, after Wijk et al. 2025) that removes the explicit sparsity penalty and straight-through estimators. Central claims: (a) state-of-the-art circuit localization — number 1 on the official MIB leaderboard at an average secret-test-set score of 5.6 vs 1.95 for the runner-up, with public-test increases of 45.9%–343.0% per task–model pair (§4, Fig. 2); (b) theoretical unification — the first SGD step of MAttr equals path-reweighted centred Integrated Gradients (§A.1–A.2), and the Bilodeau et al. (2024) impossibility theorem does not bind because MAttr is neither complete nor linear (§A.4); (c) parameter-level attribution — trained with GRPO on refusal judge scores, restoring 1% (Llama-3.1-8B-Instruct) or 2% (1B) of weights to base removes refusals while maintaining GSM8K, IFEval and MMLU, versus >27.6% parameter changes for baselines (§5.1, Fig. 5a).
The remaining sections cover Methodology Assessment, Statistical Evidence, Results & Reproducibility, Limitations & Risks, and Verdict & Recommendations.
- All 5 strengths & 0 weaknesses, analyzed in detail
- 5 risk flags with severity, section & reviewer questions
- Methods & statistical evidence assessment
- 5 questions a peer reviewer would ask
- Full six-dimension evaluation report & recommendations
One-time purchase · this paper only
One-time purchase, this paper only. AI-generated preliminary analysis — not peer review.
Frequently asked questions
What is Matryoshka Attribution (MAttr)?
An interpretability method that learns a single importance score for every internal component (attention head, neuron, or weight row) of a language model by training across all circuit sizes at once: each step samples a random budget k, builds a differentiable sigmoid top-k mask, replaces the non-top components with their counterfactual values, and tunes the scores so base behaviour is preserved. The result is a ranking from which the smallest sufficient circuit for a task can be read off.
How well does MAttr score on the MIB benchmark?
The authors report the number-1 spot on the official MIB Circuit Localization leaderboard with an average score of 5.6 versus 1.95 for the runner-up (secret test set). On the public test set, node-level MAttr scores about 2.06 CPR versus 1.85 for the best swept mask-learning baseline, and edge-level about 6.44 versus 2.49 for Edge Pruning.
What does restoring 1% of weights to remove refusals mean?
Training MAttr with reinforcement learning on a refusal judge score identified which rows of the fine-tuned Llama weight delta cause refusal. Setting just 1% of weights back to base-model values (2% for the 1B model) eliminated refusals on harmful-request benchmarks while preserving GSM8K, IFEval and MMLU — showing safety training is highly localized and easily reversed, relevant to both auditing and misuse.
How is MAttr different from Integrated Gradients?
Its first SGD step is mathematically equivalent to path-reweighted Integrated Gradients, but continued training optimizes the actual task loss across all sparsities simultaneously. The authors argue MAttr escapes the Bilodeau et al. impossibility theorem because it is neither complete nor linear, and can learn to beat gradient-based attribution on causal faithfulness.
About this review
This is an AI-generated preliminary review by Evolve.Pub, produced from the openly available text of the paper. It is not peer review and does not represent the authors' views. The canonical record of the paper is at the publisher link above.
- Figures and tables extracted with garbled label-value alignment (e.g., Fig. 2); per-variant numbers cited carry alignment uncertainty.
- Portions of §4, §5 (Fig. 5a values) and appendices H–L are truncated in the provided text.
- Repository handle (aryamanarora/matryoshka-attribution) is stated in the paper but its contents could not be verified offline.
- The headline 5.6 leaderboard score refers to a secret test set; only public-set numbers are auditable here.
Go deeper on this paper
Wrote a paper? Stress-test it before reviewers do.
Upload your manuscript for the same structured scrutiny — reasoning, statistics, evidence, figures, and reproducibility.