AtP${}^{*}$: An efficient and scalable method for localizing LLM behaviour to components

We are given a model $\mathcal{M}:X\rightarrow\mathbb{R}^{V}$ that maps a prompt (token sequence) $x\in X:=\{1,\ldots,V\}^{T}$ to output logits over a set of $V$ tokens, aiming to predict the next token in the sequence. We will view the model $\mathcal{M}$ as a computational graph $(N,E)$ where the node set $N$ is the set of model components, and a directed edge $e=(n_{1},n_{2})\in E$ is present iff the output of $n_{1}$ is a direct input into the computation of $n_{2}$. We will use $n(x)$ to represent the activation (intermediate computation result) of $n$ when computing $\mathcal{M}(x)$.

The choice of $N$ determines how fine-grained the attribution will be. For example, for transformer models, we could have a relatively coarse-grained attribution where each layer is considered a single node. In this paper we will primarily consider more fine-grained attributions that are more expensive to compute (see Section 4 for details); we revisit this issue in Section 5.

Following past work (Geiger et al., 2022; Chan et al., 2022; Wang et al., 2022), we assume a distribution $\mathcal{D}$ over pairs of inputs $x^{\text{clean}},x^{\text{noise}}$, where $x^{\text{clean}}$ is a prompt on which the behaviour occurs, and $x^{\text{noise}}$ is a reference prompt which we use as a source of noise to intervene with¹¹1This precludes interventions which use activation values that are never actually realized, such as zero-ablation or mean ablation. An alternative formulation via distributions of activation values is also possible.. We are also given a metric²²2Common metrics in language models are next token prediction loss, difference in log prob between a correct and incorrect next token, probability of the correct next token, etc. $\mathcal{L}:\mathbb{R}^{V}\rightarrow\mathbb{R}$, which quantifies the behaviour of interest.

Similarly to the work referenced above we define the contribution $c(n)$ of a node $n$ to the model’s behaviour as the counterfactual absolute³³3The sign of the impact may be of interest, but in this work we’ll focus on the magnitude, as a measure of causal importance. expected impact of replacing that node on the clean prompt with its value on the reference prompt $x^{\text{noise}}$.

Using do-calculus notation (Pearl, 2000) this can be expressed as $c(n):=|\mathcal{I}(n)|$, where

Note that the need to average the effect across a distribution adds a potentially large multiplicative factor to the cost of computing $c(n)$, further motivating this work.

We can also intervene on a set of nodes $\eta=\{n_{i}\}$. To do so, we overwrite the values of all nodes in $\eta$ with their values from a reference prompt. Abusing notation, we write $\eta(x)$ as the set of activations of the nodes in $\eta$, when computing $\mathcal{M}(x)$.

We note that it is also valid to define contribution as the expected impact of replacing a node on the reference prompt with its value on the clean prompt, also known as denoising or knock-in. We follow Chan et al. (2022); Wang et al. (2022) in using noising, however denoising is also widely used in the literature (Meng et al., 2023; Lieberum et al., 2023). We briefly consider how this choice affects AtP in Section 5.2.

The most straightforward method is to directly do Activation Patching to find the true effect $c(n)$ of each node, in some uninformed random order. This is necessarily inefficient.

However, if we are scaling to a distribution, it is possible to improve on this, by alternating between phases of (i) for each unverified node, picking a not-yet-measured prompt pair on which to patch it, (ii) ranking the not-yet-verified nodes by the average observed patch effect magnitudes, taking the top $|N|/|\mathcal{D}|$ nodes, and verifying them. This balances the computational expenditure on the two tasks, and allows us to find large nodes sooner, at least as long as their large effect shows up on many prompt pairs.

Our remaining baseline methods rely on an approximate node additivity assumption: that when intervening on a set of nodes $\eta$, the measured effect $\mathcal{I}(\eta;x^{\text{clean}},x^{\text{noise}})$ is approximately equal to $\sum_{n\in\eta}\mathcal{I}(n;x^{\text{clean}},x^{\text{noise}})$.

Under the approximate node additivity assumption, we can construct an approximately unbiased estimator of $c(n)$. We select the sets $\eta_{k}$ to contain each node independently with some probability $p$, and additionally sample prompt pairs $x^{\text{clean}}_{k},x^{\text{noise}}_{k}\sim\mathcal{D}$. For any node $n$, and sets of nodes $\eta_{k}\subset N$, let $\eta^{+}(n)$ be the collection of all those that contain $n$, and $\eta^{-}(n)$ be the collection of those that don’t contain $n$; we’ll write these node sets as $\eta^{+}_{k}(n)$ and $\eta^{-}_{k}(n)$, and the corresponding prompt pairs as ${x^{\text{clean}}_{k}}^{+}(n),{x^{\text{noise}}_{k}}^{+}(n)$ and ${x^{\text{clean}}_{k}}^{-}(n),{x^{\text{noise}}_{k}}^{-}(n)$. The subsampling (or subset sampling) estimator is then given by

The estimator $\hat{\mathcal{I}}_{\text{SS}}(n)$ is unbiased if there are no interaction effects, and has a small bias proportional to $p$ under a simple interaction model (see Section A.1.1 for proof).

In practice, we compute all the estimates $\hat{c}_{\text{SS}}(n)$ by sampling a binary mask over all nodes from i.i.d. Bernoulli${}^{|N|}(p)$ – each binary mask can be identified with a node set $\eta$. In Algorithm 1, we describe how to compute summary statistics related to Equation 13 efficiently for all nodes $n\in N$. The means $\bar{i}^{\pm}$ are enough to compute $\hat{c}_{\text{SS}}(n)$, while other summary statistics are involved in bounding the magnitude of a false negative (cf. Section 3.2). (Note, $\text{count}^{\pm}_{n}$ is just an alternate notation for $|\eta^{\pm}(n)|$.)

Instead of sampling each $\eta$ independently, we can group nodes into fixed “blocks” $\eta$ of some size, and patch each block to find its aggregated contribution $c(\eta)$; we can then traverse the nodes, starting with high-contribution blocks and proceeding from there.

There is a tradeoff in terms of the block size: using large blocks increases the compute required to traverse a high-contribution block, but using small blocks increases the compute required to finish traversing all of the blocks. We refer to the fixed block size setting as Blocks. Another way to handle this tradeoff is to add recursion: the blocks can be grouped into higher-level blocks, and so forth. We call this method Hierarchical.

We present results from both methods in our comparison plots, but relegate details to Section A.1.2. Relative to subsampling, these grou**-based methods have the disadvantage that on distributions, their cost scales linearly with size of $\mathcal{D}$’s support, in addition to scaling with the number of nodes⁷⁷7AtP* also scales linearly in the same way, but with far fewer forward passes per prompt pair..

When attributing model behavior to components, an important choice is the partition of the model’s computational graph into units of analysis or ‘nodes’ $N\ni n$ (cf. Section 2.1). We investigate two settings for the choice of $N$, AttentionNodes and NeuronNodes. For NeuronNodes, each MLP neuron⁸⁸8We use the neuron post-activation for the node; this makes no difference when causally intervening, but for AtP it’s beneficial, because it makes the $n\mapsto\mathcal{L}(n)$ function more linear. is a separate node. For AttentionNodes, we consider the query, key, and value vector for each head as distinct nodes, as well as the pre-linear per-head attention output⁹⁹9We include the output node because it provides additional information about what function an attention head is serving, particularly in the case where its queries have negligible patch effects relative to its keys and/or values. This may happen as a result of choosing $x^{\text{clean}},\,x^{\text{noise}}$ such that the query does not differ across the prompts.. We also refer to these units as ‘sites’. For each site, we consider each copy of that site at different token positions as a separate node. As a result, we can identify each node $n\in N$ with a pair $(T,S)$ from the product TokenPosition $\times$ Site. Since our two settings for $N$ are using a different level of granularity and are expected to have different per-node effect magnitudes, we present results on them separately.

We investigate transformer language models from the Pythia suite (Biderman et al., 2023) of sizes between 410M and 12B parameters. This allows us to demonstrate that our methods are applicable across scale. Our cost-of-verified-recall plots in Figures 1, 7 and 8 refer to Pythia-12B. Results for other model sizes are presented via the relative-cost (cf. Section 4.2) plots in the main body Figure 9 and disaggregated via cost-of-verified recall in Section B.3.

All reported results use the negative log probability¹⁰¹⁰10Another popular metric is the difference in logits between the clean and noise target. As opposed to the negative logprob, the logit difference is linear in the final logits and thus might favor AtP. A downside of logit difference is that it is sensitive to the noise target, which may not be meaningful if there are multiple plausible completions, such as in IOI. as their loss function $\mathcal{L}$. We compute $\mathcal{L}$ relative to targets from the clean prompt $x^{\text{clean}}$. We briefly explore other metrics in Section B.4.

As mentioned in the introduction, we’re primarily interested in finding the largest-effect nodes – see Appendix D for the distribution of $c(n)$ across models and distributions. Once we have obtained node estimates via a given method, it is relatively cheap to directly measure true effects of top nodes one at a time; we refer to this as “verification”. Incorporating this into our methodology, we find that false positives are typically not a big issue; they are simply revealed during verification. In contrast, false negatives are not so easy to remedy without verifying all nodes, which is what we were trying to avoid.

We compare methods on the basis of total compute cost (in # of forward passes) to verify the $K$ nodes with biggest true effect magnitude, for varying $K$. The procedure being measured is to first compute estimates (incurring an estimation cost), and then sweep through nodes in decreasing order of estimated magnitude, measuring their individual effects $c(n)$ (i.e. verifying them), and incurring a verification cost. Then the total cost is the sum of these two costs.

Sometimes we find it useful to summarize the method performance with a scalar; this is useful for comparing methods at a glance across different settings (e.g. model sizes, as in Figure 2), or for selecting hyperparameters (cf. Section B.5). The cost of verified recall of the top $K$ nodes is of interest for $K$ at varying orders of magnitude. In order to avoid the performance metric being dominated by small or large $K$, we assign similar total weight to different orders of magnitude: we use a weighted average with weight $1/K$ for the cost of the top $K$ nodes. Similarly, since the costs themselves may have different orders of magnitude, we average them on a log scale – i.e., we take a geometric mean.

This metric is also proportional to the area under the curve in plots like Figure 1. To produce a more understandable result, we always report it relative to (i.e. divided by) the oracle verification cost on the same metric; the diagonal line is the oracle, with relative cost 1. We refer to this as the IRWRGM (inverse-rank-weighted relative geometric mean) cost, or the relative cost.

Note that the preference of the individual practitioner may be different such that this metric is no longer accurately measuring the important rank regime. For example, AtP* pays a notable upfront cost relative to AtP or AtP+QKfix, which sets it at a disadvantage when it doesn’t manage to find additional false negatives; but this may or may not be practically significant. To understand the performance in more detail we advise to refer to the cost of verified recall plots, like Figure 1 (or many more in Section B.3).

As a starting point we report results on single prompt pairs which we expect to have relatively clean circuitry¹¹¹¹11Formally, these represent prompt distributions via the delta distribution $p(x^{\text{clean}},x^{\text{noise}})=\delta_{x^{\text{clean}}_{1},x^{\text{noise}}_{1}}(x^{\text{clean}},x^{\text{noise}})$ where $x^{\text{clean}}_{1},x^{\text{noise}}_{1}$ is the singular prompt pair.. All singular prompt pairs are shown in Table 1. IOI-PP is chosen to resemble an instance from the indirect object identification (IOI) task (Wang et al., 2022), a task predominantly involving attention heads. CITY-PP is chosen to elicit factual recall which previous research suggests involves early MLPs and a small number of late attention heads (Meng et al., 2023; Geva et al., 2023; Nanda et al., 2023). The country/city combinations were chosen such that Pythia-410M achieved low loss on both $x^{\text{clean}}$ and $x^{\text{noise}}$ and such that all places were represented by a single token.

We show the cost of verified 100% recall for various methods in Figure 1, where we focus on NeuronNodes for CITY-PP and AttentionNodes for IOI-PP. Exhaustive results for smaller Pythia models are shown in Section B.3. Figure 2 shows the aggregated relative costs for all models on CITY-PP and IOI-PP.

Instead of applying the strict criterion of recalling all important nodes, we can also relax this constraint. In Figure 7, we show the cost of verified 90% recall in the two clean prompt pair settings.

The previous prompt pairs may in fact be the best-case scenarios: the interventions they create will be fairly localized to a specific circuit, and this may make it easy for AtP to approximate the contributions. It may thus be informative to see how the methods generalize to settings where the interventions are less surgical. To do this, we also report results in Figure 8 (top) and Figure 9 on a random prompt pair chosen from a non-copyright-protected section of The Pile (Gao et al., 2020) which we refer to as RAND-PP. The prompt pair was chosen such that Pythia-410M still achieved low loss on both prompts.

We find that AtP/AtP* is only somewhat less effective here; this provides tentative evidence that the strong performance of AtP/AtP* isn’t reliant on the clean prompt using a particularly crisp circuit, or on the noise prompt being a precise control.

Causal attribution is often of most interest when evaluated across a distribution, as laid out in Section 2. Of the methods, AtP, AtP*, and Subsampling scale reasonably to distributions; the former 2 because they’re inexpensive so running them $|\mathcal{D}|$ times is not prohibitive, and Subsampling because it intrinsically averages across the distribution and thus becomes proportionally cheaper relative to the verification via activation patching. In addition, having a distribution enables a more performant Iterative method, as described in Section 3.3.

We present a comparison of these methods on 2 distributional settings. The first is a reduced version of IOI (Wang et al., 2022) on 6 names, resulting in $6\times 5\times 4=120$ prompt pairs, where we evaluate AttentionNodes. The other distribution prompts the model to output an indefinite article ‘ a’ or ‘ an’, where we evaluate NeuronNodes. See Section B.1 for details on constructing these distributions. Results are shown in Figure 8 for Pythia 12B, and in Figure 9 across models. The results show that AtP continues to perform well, especially with the QK fix; in addition, the cancellation failure mode tends to be sensitive to the particular input prompt pair, and as a result, averaging across a distribution diminishes the benefit of GradDrops.

An implication of Subsampling scaling well to this setting is that diagnostics may give reasonable confidence in not missing false negatives with much less overhead than in the single-prompt-pair case; this is illustrated in Figure 6.

We only considered a small set of prompt pair distributions, which often were limited to a single prompt pair, since evaluating the ground truth can be quite costly. While we aimed to evaluate on distributions that are reasonably representative, our results may not generalize to other distributions.

In the NeuronNodes setting, we took MLP neurons as our fundamental unit of analysis. However, there is mounting evidence (Bricken et al., 2023) that the decomposition of signals into neuron contributions does not correspond directly to a semantically meaningful decomposition. Instead, achieving such a decomposition seems to require finding the right set of directions in neuron activation space (Bricken et al., 2023; Gurnee et al., 2023) – which we viewed as being out of scope for this paper. In Section 5.2 we further discuss the applicability of AtP to sparse autoencoders, a method of finding these decompositions.

More generally, we only considered relatively fine-grained nodes, because this is a case where very exhaustive verification is prohibitively expensive, justifying the need for an approximate, fast method. Nanda (2022) speculate that AtP may perform worse on coarser components like full layers or entire residual streams, as a larger change may have more of a non-linear effect. There may still be benefit in speeding up such an analysis, particularly if the context length is long – our alternative methods may have something to offer here, though we leave investigation of this to future work.

It is popular in the literature to do Activation Patching with these larger components, with short contexts – this doesn’t pose a performance issue, and so our work would not provide any benefit here.

In this work we took the ground truth of activation patching, as defined in Equation 1, as our evaluation target. As discussed by McGrath et al. (2023), Equation 1 often significantly disagrees with a different evaluation target, the “direct effect”, by putting lower weight on some contributions when later components would shift their behaviour to compensate for the earlier patched component. In the worst case this could be seen as producing additional false negatives not accounted for by our metrics. To some degree this is likely to be mitigated by the GradDrop formula in Eq. 11, which will include a term drop** out the effect of that downstream shift.

However, it is also questionable whether we need to concern ourselves with finding high-direct-effect nodes. For example, direct effect is easy to efficiently compute for all nodes, as explored by nostalgebraist (2020) – so there is no need for fast approximations like AtP if direct effect is the quantity of interest. This ease of computation is no free lunch, though, because direct effect is also more limited as a tool for finding causally important nodes: it would not be able to locate any nodes that contribute only instrumentally to the circuit rather than producing its output. For example, there is no direct effect from nodes at non-final token positions. We discuss the direct effect further in Section 3.1.2 and Section A.2.2.

Another nuance of our ground–truth definition occurs in the distributional setting. Some nodes may have a real and significant effect, but only on a single clean prompt (e.g. they only respond to a particular name in IOI¹²¹²12We did observe this particular behavior in a few instances. or object in A-AN). Since the effect is averaged over the distribution, the ground truth will not assign these nodes large causal importance. Depending on the goal of the practitioner this may or may not be desirable.

When evaluating the performance of various estimators, we focused on evaluating the relative rank of estimates, since our main goal was to identify important components (with effect size only instrumentally useful to this end), and we assumed a further verification step of the nodes with highest estimated effects one at a time, in contexts where knowing effect size is important. Thus, we do not present evidence about how closely the estimated effect magnitudes from AtP or AtP* match the ground truth. Similarly, we did not assess the prevalence of false positives in our analysis, because they can be filtered out via the verification process. Finally, we did not compare to past manual interpretability work to check whether our methods find the same nodes to be causally important as discovered by human researchers, as done in prior work (Conmy et al., 2023; Syed et al., 2023).

While we think it likely that our results on the Pythia model family (Biderman et al., 2023) will transfer to other LLM families, we cannot rule out qualitatively different behavior without further evidence, especially on SotA–scale models or models that significantly deviate from the standard decoder-only transformer architecture.

While we focus on computing the effects of individual nodes, edge activation patching can give more fine-grained information about which paths in the computational graph matter. However, it suffers from an even larger blowup in number of forward passes if done naively. Fortunately, AtP is easy to generalize to estimating the effects of edges between nodes (Nanda, 2022; Syed et al., 2023), while AtP* may provide further improvement. We discuss edge-AtP, and how to efficiently carry over the insights from AtP*, in Section C.2.

We focused on fine-grained attribution, rather than full layers or sliding windows (Meng et al., 2023; Geva et al., 2023). In the latter case there’s less computational blowup to resolve, but for long contexts there may still be benefit in considering speedups like ours; on the other hand, they may be less linear, thus favouring other methods over AtP*. We leave investigation of this to future work.

Nanda (2022) observed that AtP’s approximation to layer normalization may be a worse approximation when it comes to patching larger/coarser nodes: on average the patched and clean activations are likely to have similar norm, but may not have high cosine-similarity. They recommend treating the denominator in layer normalization as fixed, e.g. using a stop-gradient operator in the implementation. In Section C.1 we explore the effect of this, and illustrate the behaviour of this alternative form of AtP. It seems likely that this variant would indeed produce better results particularly when patching residual-stream nodes – but we leave empirical investigation of this to future work.

Denoising (Meng et al., 2023; Lieberum et al., 2023) is a different use case for patching, which may produce moderately different results: the difference is that each forward pass is run on $x^{\text{noise}}$ with the activation to patch taken from $x^{\text{clean}}$ — colloquially, this tests whether the patched activation is sufficient to recover model performance on $x^{\text{clean}}$, rather than necessary. We provide some preliminary evidence to the effect of this choice in Section B.4 but leave a more thorough investigation to future work.

Further, in some settings it may be of interest to do mean-ablation, or even zero-ablation, and our tweaks remain applicable there; the random-prompt-pair result suggests AtP* isn’t overly sensitive to the noise distribution, so we speculate the results are likely to carry over.

A natural application of the methods we discussed in this work is the automatic identification and localization of sparse subgraphs or ‘circuits’ (Cammarata et al., 2020). A variant of this was already discussed in concurrent work by Syed et al. (2023) who combined edge attribution patching with the ACDC algorithm (Conmy et al., 2023). As we mentioned in the edge patching discussion, AtP* can be generalized to edge attribution patching, which may bring additional benefit for automated circuit discovery.

Another approach is to learn a (probabilistic) mask over nodes, similar to Louizos et al. (2018); Cao et al. (2021), where the probability scales with the currently estimated node contribution $c(n)$. For that approach, a fast method to estimate all node effects given the current mask probabilities could prove vital.

Recently there has been increased interest by the community in using sparse autoencoders (SAEs) to construct disentangled sparse representations with potentially more semantic coherence than transformer-native units such as neurons (Cunningham et al., 2023; Bricken et al., 2023). SAEs usually have a lot more nodes than the corresponding transformer block they are applied to. This could pose a larger problem in terms of the activation patching effects, making the speedup of AtP* more valuable. However, due to the sparseness of the SAE, on a given forward pass the effect of most features will be zero. For example, some successful SAEs by Bricken et al. (2023) have 10-20 active features for 500 neurons for a given token position, which reduces the number of nodes by 20-50x relative to the MLP setting, increasing the scale at which existing iterative methods remain practical. It is still an open research question, however, what degree of sparsity is feasible with tolerable reconstruction error for practically relevant or SOTA–scale models, where the methods discussed in this work may become more important again.

AtP* could be used to discover single nodes in the model that can be leveraged for targeted inference time interventions to control the model’s behavior. In contrast to previous work (Li et al., 2023; Turner et al., 2023; Zou et al., 2023) it might provide more localized interventions with less impact on the rest of the model’s computation. One potential exciting direction would be to use AtP* (or other gradient-based approximations) to see which sparse autoencoder features, if activated, would have a significant effect.