AtP*: An efficient and scalable method for localizing LLM behaviour to components

TL;DR

The paper tackles the challenge of efficiently attributing LLM behavior to internal components by analyzing Activation Patching (AtP) and proposing AtP*, a gradient-based, scalable alternative. It identifies two failure modes—attention-saturation and cancellation—that cause false negatives and offers two remedies: QK fixes for queries/keys and GradDrop to disentangle direct and indirect effects. Across extensive experiments on Pythia-scale models and both single-prompt and distributional settings, AtP* (and AtP with fixes) consistently outperforms baselines, with a diagnostic method to bound remaining false negatives. The work also discusses extensions to edge patching, coarser node definitions, and practical recommendations for applying causal attribution in real-world mechanistic interpretability tasks. Overall, AtP* provides a practical, scalable framework for reliable circuit localization in large transformers, enabling more robust mechanistic insights and steering capabilities.

Abstract

Activation Patching is a method of directly computing causal attributions of behavior to model components. However, applying it exhaustively requires a sweep with cost scaling linearly in the number of model components, which can be prohibitively expensive for SoTA Large Language Models (LLMs). We investigate Attribution Patching (AtP), a fast gradient-based approximation to Activation Patching and find two classes of failure modes of AtP which lead to significant false negatives. We propose a variant of AtP called AtP*, with two changes to address these failure modes while retaining scalability. We present the first systematic study of AtP and alternative methods for faster activation patching and show that AtP significantly outperforms all other investigated methods, with AtP* providing further significant improvement. Finally, we provide a method to bound the probability of remaining false negatives of AtP* estimates.

Figures (17)

• Figure 1: Costs of finding the most causally-important nodes in Pythia-12B using different methods, on sample prompt pairs (see \ref{['tab:single_prompt_pairs']}). The shading indicates geometric standard deviation. Cost is measured in forward passes, thus each point's y-coordinate gives the number of forward passes required to find the top $x$ nodes. Note that each node must be verified, thus $y\geq x$, so all lines are above the diagonal, and an oracle for the verification order would produce the diagonal line. For a detailed description see \ref{['sec:clean_prompt_pairs']}.
• Figure 2: Relative costs of methods across models, on sample prompt pairs. The costs are relative to having an oracle, which would verify nodes in decreasing order of true contribution size. Costs are aggregated using an inverse-rank-weighted geometric mean. This means they correspond to the area above the diagonal for each curve in \ref{['fig:main_result']} and are relative to the area under the dotted (oracle) line. See \ref{['sec:irwrgm']} for more details on this metric. Note that GradDrop (difference between AtP+QKfix and AtP$^*$) comes with a noticeable upfront cost and so looks worse in this comparison while still helping avoid false negatives as shown in\ref{['fig:main_result']}.
• Figure 3: A linear approximation to the attention probability is a particularly poor approximation in cases where one or both of the endpoints are in a saturated region of the softmax. Note that when varying only a single key, the softmax becomes a sigmoid of the dot product of that key and the query.
• Figure 4: Ranks of $c(n)$ against ranks of $\hat{c}_\text{AtP}(n)$, on Pythia-12B on CITY-PP. Both improvements to AtP reduce the number of false negatives (bottom right triangle area), where in this case most improvements come from the QK fix. Coloration indicates the maximum absolute difference in attention probability when comparing $x^\text{clean}$ and patching a given query or key. Many false negatives are keys and queries with significant maximum difference in attention probability, suggesting they are due to attention saturation as illustrated in \ref{['fig:saturation_cartoon']}. Output and value nodes are colored in grey as they do not contribute to the attention probability.
• Figure 5: True rank and rank of AtP estimates with and without GradDrop, using Pythia-12B on the CITY-PP distribution with NeuronNodes. GradDrop provides a significant improvement to the largest neuron false negatives (red circles) relative to Default AtP (orange crosses).