Attribution Patching Outperforms Automated Circuit Discovery

Venue: EMNLP_2024_Workshop_BlackBoxNLP
Published: 2024-08-20 02:45:43

Abstract: Automated interpretability research has recently attracted attention as a potential research direction that could scale explanations of neural network behavior to large models. Existing automated circuit discovery work applies activation patching to identify subnetworks responsible for solving specific tasks (circuits). In this work, we show that a simple method based on attribution patching outperforms all existing methods while requiring just two forward passes and a backward pass. We apply a linear approximation to activation patching to estimate the importance of each edge in the computational subgraph. Using this approximation, we prune the least important edges of the network. We survey the performance and limitations of this method, finding that averaged over all tasks our method has greater AUC from circuit recovery than other methods.

Main paper: 
 introduction: 
 Mechanistic interpretability is a subfield of AI interpretability that focuses on attributing model behaviors to its components, thus reverse engineering the network (Olah, 2022) . This field aims to identify subnetworks (circuits) within the model which are responsible for solving specific tasks (Olah et al., 2020) . Prior attempts at finding circuits in language models have led to finding networks of attention heads and multi-layer perceptrons (MLPs) that partially or fully explain model behaviors at tasks such as indirect object identification, modular arithmetic, completion of docstrings, and predicting successive dates (Wang et al., 2023; Nanda et al., 2023; Heimersheim and Janiak, 2023; Hanna et al., 2023) . However, almost all previous work has been limited to relatively small models since manually applying mechanistic interpretability methods has not currently scaled to end-to-end circuits in larger models (Lieberum et al., 2023) .It may be important to scale interpretability to large models as these are the neural networks most widely deployed and used by a wide range of people. Currently, we have little understanding into how these models work and failure modes are not always found ahead of deployment. If successful, scaled interpretability could address a wide variety of concerns about the lack of transparency of language models (Vig et al., 2020) , in addition to speculative risks about the alignment of machine learning systems (Hubinger, 2020) . Automated Circuit Discovery (ACDC; Conmy et al. (2023) ) attempts to automate a large portion of the mechanistic interpretability workflow -the pruning of edges between attention heads and MLPs that do not affect the task being studied. ACDC begins with a computational graph, and recursively calculates the importance of an edge in the graph for a specific task. In our work, we use edges to refer to activations inside models between two components (Section 2 describes this motivation further). ACDC's pruning algorithm applies activation patching. (Note that activation patching is not attribution patching. Both are defined in full in Section 3.3.) At a high level, activation patching edits a specific activation in a model forward pass and measures a model statistic (e.g loss) under this intervention. Activation patching is inefficient for circuit discovery because getting each statistic about model activations requires another forward pass. Our work uses attribution patching to recover circuits more efficiently (Section 3.3).Our main contributions are:1. Introducing a method for using attribution patching on all computational graph edges for automated circuit discovery (Edge Attribution Patching, Section 3.3). 
benchmarking edge attribution patching vs: 
 existing circuit discovery methods (Section 4). 
finding and explaining some limitations with: 
 Edge Attribution Patching (Section 5).Automated Circuit Discovery refers to finding the important subgraph of models' computational graphs for performance on particular tasks (Conmy et al., 2023) . Existing algorithms include efficient heuristics (Michel et al., 2019) and gradientdescent based methods (Louizos et al., 2018; Cao et al., 2021) . ACDC is related to pruning (Blalock et al., 2020) and other compression techniques (Zhu et al., 2023) , but differs in how the compressed networks are reflective of the circuits that model uses to compute outputs to certain tasks and the goal of ACDC is not to speed up forward passes (all techniques studied in this work slow forward passes).Activation Patching is a technique for analyzing the role of individual components in a model. It involves targeted manipulations of activations during a forward pass (further explained in Section 3.1). Previous works applied this technique under various names, such as Interchange Interventions (Geiger et al., 2021) , Causal Mediation Analysis (Vig et al., 2020) and Causal Tracing (Meng et al., 2022) . We adapt the terminology used by Conmy et al. (2023) .Transformer Circuits. Our work builds upon the framework for understanding transformers for interpretability as introduced by Elhage et al. (2021) . The important details include how they formulate forward passes of transformer models. Individual attention heads and MLPs (collectively called nodes) read and write information to a central communication channel, also called the residual stream. In these terms we can examine dependencies of nodes with the output of earlier nodes, i.e we can measure the effect of attention heads in layer 0 on the attention heads in layer 2. In the following, we view these dependencies as edges between nodes, building on existing work using this perspective (Heimersheim and Janiak, 2023; Hanna et al., 2023; Wang et al., 2023) .We present Edge Attribution Patching (EAP) as a technique to identify relevant model components for solving a specific task. In the following, we view language models as directed, acyclic graphs. In these terms, we aim to find small subgraphs that retain good performance on narrow tasks. We determine the importance of a specific edge through targeted manipulation of activations during a forward pass. We compare two approaches, Attribu-tion Patching and Activation Patching, in order to motivate EAP.Activation patching refers to replacing the activations from one model forward pass with the activations from a different forward pass. This method is typically applied to measure the counterfactual importance of model components, i.e. to measure a statistic L(x) from model outputs under the activation patching, where x is an input prompt. For example, L often represents loss or logit difference (Wang et al., 2023) .Following existing work (Section 2), we study the effect of activation patching on specific model edges by setting these equal to activations from different forward passes. Concretely, suppose that an edge E in the computational graph has activation e corr on some corrupted prompt. In this work, we use the change in metric under activation patching|L(x clean | do(E = e corr )) -L(x clean )| (1)to measure the impact of edge E. We use donotation from causality (Pearl, 1995) to emphasise that activation patching is a causal intervention.Activation patching slows ACDC since each measurement (like Equation ( 1)) requires another forward pass. Attribution patching (Nanda, 2023 ) is a technique for estimating Equation (1) for many different edges E using only two forward passes and one backward pass.1 It linearly approximates the metric difference after corrupting a single edge in the computational graph (Figure 1 ) by expanding L as a function of the edge activation as a Taylor series with terms up to the first order:2 L(x clean | do(E = e corr )) ≈ L(x clean )+ (e corr -e clean ) ⊤ ∂ ∂e clean L(x clean | do(E = e clean ))Call this ∆eL, the attribution score.(2)A simple rearrangement implies that Equation ( 1) is approximately equal to |∆ e L| (3) which we call the absolute attribution score for the rest of this paper. In this work we always compute this score across a set of (x clean , x corr ) pairs and take the mean.In practice, all gradients needed to calculate the attribution scores come from intermediate terms computed in one ordinary backwards pass3 in Py-TorchWe can use the insights from Section 3.2 to build an automated circuit discovery algorithm. This takes two steps: i) use Equation ( 2) to obtain absolute attribution scores for the importance of all edges in the computational graph, and then ii) sort these scores and keep the top k edges in a circuit. We use Edge Attribution Patching (EAP) to refer to this algorithm. In the rest of the work we report results for all k values when we evaluate EAP (similar to HISP in Conmy et al. (2023) ). Note that one limitation of attribution patching is that it will not work when the gradient of the metric is the zero vector. Conmy et al. (2023) recommended the use of KL divergence as a metric, which is i) equal to 0 when we run the model without ablations and ii) a non-negative metric. Therefore the zero point is a global minimum and hence all gradients are the zero vector at this point. In this work we use the task-specific metrics' (not KL divergence) from Conmy et al. (2023) so avoid this issue.  
edge attribution patching vs activation: 
 Patching vs ACDCWe compare Edge Attribution Patching (EAP) and ACDC on the Indirect Object Identification (IOI), Docstring, and Greater-Than tasks. For each of these tasks, previous studies identified a subgraph (circuit) relevant for solving the task. We use their results as a ground truth for benchmarking both methods. We also compare using ACDC with the task-specific metrics (e.g logit difference) and KL Divergence (which was originally recommended).For the docstring task, we also include repeated activation patching as another point of reference for performance comparisons. We applied repeated activation patching by running the same circuit discovery method described in Section 3.3 but using Equation (1) rather than absolute attribution scores. Activation patching was not included in the other tasks as it was too computationally expensive to run on the GPT-2 small models used by IOI and Greater-Than. Subnetworks found using EAP for all three tasks are shown in Appendix A.The ROC curves in Figure 2 suggest the performance of EAP is better than ACDC overall: it has the maximal AUC in Figure 2 -2, while ACDC used with the KL Divergence metric outperforms EAP in Figure 2 .In this section, we look at the approximate metric change (attribution score) EAP assigns to each edge in the model. We aim to understand the relation between the attribution score and the function of the edge in the task being studied. First, we look at the distribution of scores for edges in the circuit compared to edges not in the circuit for each of the three tasks.Figure 3 shows the distribution of attribution scores for the IOI task. The distributions for the remaining tasks can be found in Appendix B. Qualitatively, attribution scores for edges in the circuit tend to be spread further from zero. Furthermore, there are only 6 edges outside of the interval [-0.25, 0.25] that aren't part of the IOI circuit. We further explore the attribution scores for the IOI circuit's classes of heads in Appendix E.  
limitations: 
 We introduced edge activation patching as an approximation to activation patching. However, we found that edge activation patching outperformed ACDC, a technique based on activation patching (Section 4). In this section, we investigate whether attribution patching's success is due to extremely accurate approximations (in Section 5.1 we find that the answer is no), and whether there is any further use for ACDC (in Section 5.2 we find that the answer is yes). We use the docstring task as a case study due to the small model size used.To study how faithful the approximation Equation ( 2) is, we plot the attribution patching scores (Equation ( 2)) against the activation patching scores (Equation ( 1)) in Figure 4a . Surprisingly, we find a fairly weak correlation between activation and attribution patching scores (R 2 = 0.27). Further, the line of best fit has gradient 0.531, suggesting that attribution patching estimates the effect of activation patching as twice as important as it really is. Moreover, we can gain some sense for the discrepancy between activation and attribution patching by studying the continuous transition between clean (e clean ) and corrupted (e corr ) activations in Equation (1), i.e studying the values |L(x clean | do(E = λe corr + (1 -λ)e clean )) -L(x clean )| for 0 ≤ λ ≤ 1. We can compare this to the linear approximations of Attribution Patching λ∆ e L. Figure 4b shows the result for one edge in the docstring circuit where the linear approximation to activation patching is not accurate.We find that interpolating towards the corrupted input creates a concave curve (Figure 4b ) such that the linear approximation at λ = 0 overestimates the effect of activation patching this edge. In Appendix D we show that this also holds for the other outlier edges in the ellipse in Figure 4a .In Section 5.1 above, we found that EAP overestimates activation patching in cases where the attribution score is concave. This suggests the potential to refine the result by running ACDC on the pruned subgraph returned by EAP. We ran EAP first, then ACDC on the resulting subgraph for the Docstring task, varying pruning thresholds for EAP and ACDC independently. Figure 5 compares the TPR and FPR for the combined methods with the ROC curve of EAP only. The combined methods show increased performance compared to EAP only. Finally, one further limitation of this research is that the metrics used for interpretability do not precisely capture meaningful human understanding. Recovering a subgraph that humans previously recovered is limited because i) we can't evaluate this metric for interpretability tasks that we don't yet understand and ii) human-found circuits are imperfect, increasing the noise in this measurement. 
conclusion: 
 We provide evidence that Edge Attribution Patching (EAP) outperforms ACDC in identifying circuits while being substantially faster to compute. This result is surprising, as EAP is an approximation for activation patching, the method applied by ACDC. However, running ACDC on the prepruned subnetwork found by EAP can improve the identification of relevant edges. Therefore, we suggest future circuit discovery experiments to run EAP first and then apply ACDC.  
 
 Appendix: In Figure 9 we show further cases where in the docstring task attribution patching can be misleading. These cases all involve an edge that comes from the model's embeddings (positional and tokens). Our interpretation is that weighted averages of embeddings are anomalous inputs to the model and cause the concave change in docstring logit diff which doesn't occur when edges ae between non-embedding model components. We further explore the attribution scores for the IOI circuit. The IOI circuit is comprised of different attention head classes such as Induction heads, S-Inhibition heads, etc. (Wang et al., 2023) . Figure 10 shows the distributions of scores stratified by the roles of the edges. The edge roles are defined according to the role of their origin node. While edge roles such as Previous Token, Duplicate Token, Induction, and S-Inhibition edges have attribution scores centered around zero, we see a bias in edge scores given to name mover and negative name mover edges. As the name mover edges are directly responsible for the model outputting the indirect object, the attribution scores are largely negative since ablating these edges removes the model's ability to output the indirect object, lowering the logit difference. Similarly, the negative name movers have attribution scores that are largely positive since ablating these edges improves the logit difference. This matches the intuitive function of the edges. There are only two types of edges iterated over in ACDC: i) residual edges where the result is added at its endpoint, and ii) edges between the residual stream and the query, key and value calculations. Clearly for all edges like ii) we can compute the gradient terms in Equation (2) in one backwards pass.Interestingly, for all ∆ e L terms where e is a type i) edge (i.e added at the endpoint), we only need calculate the gradient with respect to the endpoint of the edge! For example, suppose we're calculating the effect of L0H0 on L1H0Q. If we represent the input to L1H0Q as a node V in the computational graph then due to how V is just the sum of all the edges entering V . This allows efficient calculation of all the ∆ e L values since gradients with respect to nodes in computational graphs are calculated by default in backwards passes.Attribution patching (like activation patching) also applies to nodes and other model internal components that aren't edges, but we only use edges in this work.Note that L(xclean| do(E = eclean)) = L(xclean) as all activations in this equation are from clean forward passes. We highlight the eclean since we take the gradient with respect to this activation.In Appendix F we show how only one backwards pass is required. 


Decision: All the reviewers agree that this paper introduces a very valuable interpretability method that will have a positive impact on the field. As such we gladly accept it to BlackboxNLP.

The reviewers, however, do raise several salient points for improvement that the authors are strongly encouraged to engage with in the camera ready version of the paper. Reviewer KRi9 provides various suggestions for improving the writing style and presentation of the paper that should be incorporated. Furthermore, given that the paper is currently only 5 pages, the authors are encouraged to take up the suggestion of reviewer HGi3 of incorporating a more detailed background on circuit finding methods.

Reviews: review: paper_summary: paper_summary: This work addresses the problem of automated circuit discovery, i.e. of identifying a subnetwork (circuit) of a model that is responsible for solving a specific subtask. 

To this end the authors propose to prune unimportant edges in the computational graph of the neural network based on gradient-based estimation of edge importance. They call this new approach Attribution Patching (since they use a standard gradient-based XAI attribution method to estimate the edges' importance).

The authors apply their method to transformer-based language models on three tasks (completion of docstrings, Indirect Object Identification IOI, greater-than task) and show it performs better in terms of the AUC for identifying edges in the circuit than Automated Circuit Discovery (ACDC) from Conmy et al. 2023 which is based on Activation Patching.

The advantage of the proposed method is that it is less expensive to compute than previous methods based on iteratively editing the network's connections to assess each edge's importance in order to prune the network to obtain the circuit. The authors further show that this new approach can also be complementary to ACDC, indeed they show that when Attribution Patching is applied as a first step before applying ACDC, then this gives the best results.
strength_weakness: strength_weakness: - The work is almost identical to an approach proposed in an "incomplete project report" posted by Neel Nanda in February 2023: https://www.neelnanda.io/mechanistic-interpretability/attribution-patching
Although this report (blog post) is not an official peer-reviewed publication, and the authors do cite it, it shall be more clearly acknowledged by the authors due to the great similarity to the present work IMHO.

- Another recent work by Ferrando et al. goes in a similar fashion of finding circuits in an automated way using an XAI method, namely ALTI (Aggregation of Layer-Wise Token-to-Token Interactions) XAI method: https://arxiv.org/pdf/2403.00824
You might want to cite this work too.
questions: questions: This work uses less than the maximum allowed number of pages, indeed it has only 5 pages. So I recommend the authors use the full paper length of 8 pages to clarify a number of points in their camera-ready version, if accepted for publication.

- In the abstract you write that your approach requires only 2 forward passes and 1 backward pass. While I understand that 1 backward pass is required to compute gradients, it is less clear to me why 2 forward passes are needed. Could you please clarify what is done in each of theses forward passes?

- page 2 Related Work on Automated Circuit Discovery: "that model used" -> "that THE model uses", "slow forward passes" -> "USE slow forward passes" (Btw why are these forward passes slow? Is there some overhead due to storing some intermediate activations during the forward pass?)

- Please specify the tasks and models used in your experiments (no reference is given for them so far)

- Page 3 end of section 4.1: Please clarify when EAP is better than ACDC, and when the opposite is true. Does this depend on the metric used to find important edges? Why is it so?

- Section 5 first sentence: "We introduced edge activation patching" -> "We introduced edge ATTRIBUTION patching"

- Section 5.2 first sentence: "... in cases where the attribution score is concave" -> "... in cases where the CHANGE IN LOGIT DIFFERENCE METRIC is concave"

- What is the threshold you apply to obtain the ROC curves? Is it the value of the edge importance above which edges get selected for the circuit?

- In figure 4b and 9 you show examples where the linear approximation to activation patching is not accurate. Could you please also show examples where the approximation IS accurate?

- In my understanding section 4.2 is kind of superfluous since you show there that the edges selected to be in the circuit have higher absolute attribution scores. But this is exactly how you choose the edges to keep in the circuit in section 3.3. So this result seems not surprising. Maybe remove that section or justify why it is important and not redundant with your edge selection approach.
paper_summary: paper_summary: This paper examines the automated interpretability of large NNs. It compares the SOTA Automatic Circuit discovery(ACDC) algorithm with a new technique called attribution patching (Nanda, 2023). The system uses the scores from the attribution patching algorithm to keep the top K edges of an NN. This work conducts empirical experiments to compare the performance of this new technique with more rigorous but computationally expensive ACDC method.
strength_weakness: strength_weakness: 1. This paper does not introduce any new methodology or system for NN interpretability. It seems to apply a sorting algorithm to the existing attribution patching algorithm. 
2. The abstract and introduction should be rewritten to state this paper’s contributions clearly. An initial read may suggest that attribution patching is a contribution of this work, but that was proposed in Nanda, 2023.
questions: questions: NA
paper_summary: paper_summary: This paper extends an existing method for finding node-level circuits (attribution patching) to the edge domain, allowing circuits to be found at the edge level. It then conducts experiments that demonstrate that, despite taking significantly less time to run than another popular method (ACDC) the circuits it finds have equal or better overlap with known task circuits for the IOI, Greater-Than, and Docstring tasks. Surprisingly, though edge attribution patching is effective at finding circuits, it is not actually a good approximation of activation patching scores.
strength_weakness: strength_weakness: This paper has a number of weaknesses, namely that:
1. it does not evaluate against all the techniques that ACDC did
2. its evaluation is limited to a small number of known circuits, in small models
3. it doesn't measure faithfulness
4. it is lacking a bit of background on what circuits are, to make this paper more accessible to readers new to circuits

However, this paper is also almost a year old, and has been succeeded by a number of other papers / circuit-finding methods (combining this method with integrated gradients, or using differentiable masking to find circuits). These papers have addressed much of 2-3 (and 1 looks in retrospect to be less important). So, it's not clear to me that these flaws are very important now. However, I would love to see the background on circuits be explained more clearly, and to see the comments below be addressed.
questions: questions: - Say a little more about circuits. How does mech interp view the computational graph? What are your nodes V (attn heads and MLPs) and edges E between them? What does this conceptualization of circuits capture (structure) and what does it miss (semantics)?
- Don't capitalize "Interchange Interventions", etc.; they're not proper nouns
- Turn Figures 1,2, 4a, and 5 into PDF/vector graphics figures with selectable text
- You refer to "Figure 2-2", but it's not clear which one that is