Matryoshka attribution: Learning to attribute language model outputs to representations and weights
Abstract Overview
This paper frames attribution in language models as learning an ordering over nested subsets of internal components that minimizes a downstream loss. The authors introduce Matryoshka Attribution (MAttr), a mask-learning method built around a differentiable sigmoid top-k operator and training over randomly sampled sparsity budgets so that one model learns attributions across sparsity levels simultaneously. They evaluate MAttr for attribution to internal representations on the Mechanistic Interpretability Benchmark (MIB) and an extended MIB+ suite, and also adapt it to parameter attribution by interpolating between model checkpoints. As a practical case study, they train MAttr with reinforcement learning to identify finetuning-induced weight changes responsible for refusal behavior while preserving general capabilities.
Novelty
The main novelty is the formulation of attribution as a nested-subset optimization problem, implemented with joint training across sparsities via random budget sampling and a sigmoid top-k mask. The work is also distinctive in extending this framework from representation attribution to parameter attribution between checkpoints, including reinforcement-learning-based training for non-differentiable behavioral targets.
Results
On representation attribution, MAttr is reported to set state-of-the-art performance on MIB and MIB+, including the top official MIB leaderboard score and statistically significant gains over tested baselines on node- and edge-level localization. The method also yields sparse, transferable circuits across tasks and bases. For parameter attribution, the authors show that restoring 2% of parameters in Llama-3.2-1B-Instruct and 1% in Llama-3.1-8B-Instruct to their base-model values is sufficient to remove refusals while largely maintaining capability metrics.
Key Points
- MAttr learns a single attribution ordering over components by optimizing nested masks across multiple sparsity budgets rather than training separate masks for each sparsity.
- The method outperforms gradient-based, intervention-based, and prior mask-learning baselines on MIB-style circuit localization, while also performing well on sparsity-sensitive compactness analyses.
- MAttr can be trained with reinforcement learning for parameter-level attribution, enabling localization of finetuning-induced behaviors such as refusal to a small fraction of weight changes.