Matryoshka attribution: Learning to attribute language model outputs to representations and weights
cs.CL, cs.LG
Submitted: 2026-09-22
Updated: 2026-09-22
Comments: 10 pages main text, 58 pages total; preprint
Code: https://github.com/tatsu-lab/stanford_alpaca
License: http://creativecommons.org/licenses/by/4.0/
The gist: Attributing language model outputs to their internal computations is an open problem in interpretability.
Terminology
Abstract
Attributing language model outputs to their internal computations is an open problem in interpretability. Existing methods, which use causal interventions, gradients, or learnable masks, either are infeasibly expensive or struggle to identify actual causally-important internal computations. We propose framing attribution as the problem of identifying nested subsets of internal components which minimise a downstream loss. To learn this task, we introduce Matryoshka Attribution (MAttr), a mask learning method that parametrises the mask with a simple differentiable sigmoid top- k operator. We supervise training over all sparsities simultaneously by randomising k over training, resulting in a learned ordering of components by attribution score. MAttr achieves number 1 on the official leaderboard of the Mechanistic Interpretability Benchmark (Mueller et al., 2025); our method identifies sparse and task-transferrable circuits across varying circuit bases. As a practical application, we show that MAttr can be trained with reinforcement learning to identify weight changes responsible for downstream behaviours in LLM finetuning. We train MAttr on refusal judge scores and find that restoring 1% of Llama 3.1 8B Instruct's weights to their base model state is sufficient to remove refusals while maintaining capabilities. We view MAttr as a successful formulation of interpretability into a learnable objective that we can tackle with gradient descent, and encourage future work along these lines.
Sources
- ADAG: Automatically Describing Attribution Graphs
- Language Model Circuits Are Sparse in the Neuron Basis
- Estimating or Propagating Gradients Through Stochastic Neurons for Conditional Computation
- Weird Generalization and Inductive Backdoors: New Ways to Corrupt LLMs
- Stochastic Parameter Decomposition
- Reasoning Models Don't Always Say What They Think
- Think you have Solved Question Answering? Try ARC, the AI2 Reasoning Challenge
- Discovering Variable Binding Circuitry with Desiderata
- Correcting Gradient-Based Circuit Localization via Interaction-Aware Backpropagation
- Arithmetic in the Wild: Llama uses Base-10 Addition to Reason About Cyclic Concepts
- The State of Sparsity in Deep Neural Networks
- Localizing Model Behavior with Path Patching
- Llama Scope: Extracting Millions of Features from Llama-3.1-8B with Sparse Autoencoders
- Predictive Concept Decoders: Training Scalable End-to-End Interpretability Assistants
- RelP: Faithful and Efficient Circuit Discovery in Language Models via Relevance Patching
- Activation Oracles: Training and Evaluating LLMs as General-Purpose Activation Explainers
- Bactrian-X: Multilingual Replicable Instruction-Following Models with Low-Rank Adaptation
- The Hydra Effect: Emergent Self-repair in Language Model Computations
- Tail-Likelihood Reinforcement Learning
- GRP-Obliteration: Unaligning LLMs With a Single Unlabeled Prompt
Related papers
- Exploring Solution Divergence and Its Effect on Large Language Model Problem Solving
- Ishigaki-IDS-Bench: A Benchmark for Generating Information Delivery Specification from BIM Information Requirements
- Subliminal Steering: Stronger Encoding of Hidden Signals
- MedStruct-S: A Benchmark for Key Discovery, Key-Conditioned QA and Semi-Structured Extraction from OCR Clinical Reports
- The End of Transformers? On Challenging Attention and the Rise of Sub-Quadratic Architectures
- Untangling the Mechanisms of Misleading Context in Medical Question Answering