MAttr turns attribution into a training objective, hitting MIB SOTA in just 500 steps

aryaman2020 · x · 2026-09-24

Aryaman introduces MAttr, a new interpretability method that frames attribution itself as a training objective:

Method: A differentiable sigmoid top-k operator (no sparsity loss, no straight-through tricks) parametrizes causal interventions; attribution scores are trained with RL using a judge score as reward, and randomizing the top-k budget during training induces an attribution ranking. Unlike prior work on representations, the mask is applied to model weights across two checkpoints, tracing qualitative LLM behaviours to the subset of parameters updated during training.
Results: On the Mechanistic Interpretability Benchmark (MIB), just 500 training steps per subtask sets state-of-the-art among node-level methods, driving logit difference far above baselines with a compact set of units.
Transfer: Rankings aren't overfitting — the MIB subtraction ranking recovers 86% of performance on the Month modular addition task.
Appendices: Tracing emergent misalignment to parameter deltas, virtual-weight toy models, ViT experiments, and a theoretical link to one-step Integrated Gradients, answering the call of the "Impossibility Theorems" paper.

More from Research