SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Arthur Conmy* Independent
Augustine N. Mavor-Parker* UCL
Aengus Lynch* UCL
Stefan Heimersheim  
University of Cambridge
Adrià Garriga-Alonso*  
FAR AI
Through considerable effort and intuition, several recent works have reverse-engineered nontrivial behaviors of transformer models. This paper systematically the mechanistic interpretability process they followed. First, researchers choose a metric and dataset that elicit the desired model behavior. Then, they apply activation patching to find which abstract neural network units are involved in the behavior. By varying the dataset, metric, and units under investigation, researchers can understand the functionality of each component.
We automate one of the process' steps: finding the connections between the abstract neural network units that form a circuit. We propose several algorithms and reproduce previous interpretability results to validate them. For example, the ACDC algorithm rediscovered 5/5 of the component types in a circuit in GPT-2 Small that computes the Greater-Than operation. ACDC selected 68 of the 32,000 edges in GPT-2 Small, all of which were manually found by previous work. Our code is available at https://github.com/ArthurConmy/Automatic-Circuit-Discovery.
Rapid progress in transformer language modelling (Vaswani et al., 2017; Devlin et al., 2019; OpenAI, 2023, inter alia) has directed attention towards understanding the causes of new capabilities (Wei et al., 2022) in these models. Researchers have identified precise high-level predictors of model performance (Kaplan et al., 2020), but transformers are still widely considered 'black-boxes' (Alishahi, Chrupa, and Linzen, 2019) like almost all other neural network models (Fong and Vedaldi, 2017; Buhrmester, Münch, and Arens, 2021). Interpretability research aims to demystify machine learning models, for example by explaining model outputs in terms of domain-relevant concepts (Zhang et al., 2021).
Mechanistic interpretability focuses on reverse-engineering model components into human-understandable algorithms (Olah, 2022). Much research in mechanistic interpretability views models as a computational graph (Geiger et al., 2021), and circuits are subgraphs with distinct functionality (Wang et al., 2023). The current approach to extracting circuits from neural networks relies on a lot of manual inspection by humans (Räuker et al., 2022). This is a major obstacle to scaling up mechanistic interpretability to larger models, more behaviors, and complicated behaviors composed of many sub-circuits. This work identifies a workflow for circuit research, and automates part of it by presenting several methods to extract computational graphs from neural networks.
Our main contributions are as follows. First, we systematize the common workflow prevalent in many existing mechanistic interpretability works, outlining the essential components of this process
*Work partially done at Redwood Research. Correspondence to arthurconmy@gmail.com
Though this perspective is not universal (Lipton, 2016).
Figure 1: Automatically discovering circuits with ACDC. Left: a computational graph for GPT-2 Small, with a recovered circuit for the IOI task highlighted in red. Only edges between adjacent layers are shown. Right: the recovered circuit with labelled nodes. All heads recovered were identified as part of the IOI circuit by Wang et al. (2023). Edge thickness is proportional to importance.
(Section 2). One of its steps is to find a subgraph of the model which implements the behavior of interest, which is a step possible to automate. We introduce Automatic Circuit Discovery (ACDC), a novel algorithm that follows the way in which researchers identify circuits (Section 3), and adapt Subnetwork Probing (SP; Cao, Sanh, and Rush, 2021) and Head Importance Score for Pruning (HISP; Michel, Levy, and Neubig, 2019) for the same task. Finally, we introduce quantitative metrics to evaluate the success of circuit extraction algorithms (Sections 4 and 4.2). We present a detailed ablation study of design choices in Appendix E and qualitative studies in Appendices F, G, H, I and J.
2 The Mechanistic Interpretability Workflow
Mechanistic interpretability attempts to explain and predict neural network behaviors by understanding the underlying algorithms implemented by models. In the related work section we discuss the mechanistic interpretability field and its relationship to 'circuits' research (Section 5). Neural network behaviors are implemented by algorithms within the model's computational graph, and prior work has identified subgraphs (circuits, following Wang et al. (2023)'s definition) that capture the majority of particular behaviors. In this section, we describe a workflow that several prior works have followed that has been fruitful for finding circuits in models.
As a concrete example of an approach taken to finding a circuit, Hanna, Liu, and Variengien (2023) prompt GPT-2 Small with a dataset of sentences like "The war lasted from 1517 to 15". GPT-2 Small completes this sentence with "18" or "19" or any larger two digit number, but not with any two digit number that is at most "17" (from here, we refer to prompt completions like this as the "Greater-Than" task). This behavior can be measured by the difference in probability the model places on a completion "18" or "19" or larger and the probability the model places on a completion "17" or smaller. Note that we use the term 'dataset' to refer to a collection of prompts that elicit some behavior in a model: we do not train models on these examples, as in this paper we focus on post-hoc interpretability.
The researchers then create a corrupted dataset of sentences that do not have any bias against particular two digit completions (the '01-dataset' (Hanna, Liu, and Variengien, 2023)). The researchers attribute the greater-than operation to late layer MLPs and then find earlier components that identify the numerical values of years, including attention heads in the model. Finally, Hanna, Liu, and Variengien (2023) interpret the role of each set of components. For example, they identify early model components that respond to the "17" token, and later model components that boost the importance of logits for years greater than 17.
There are equivalent steps taken in a growing number of additional works (Heimersheim and Janiak, 2023, the "Docstring" task; Goldowsky-Dill et al., 2023, the "Induction" task; Wang et al., 2023,
the "IOI" task), described in brief in Table 1 and in detail in Appendices F, H and J. We identify the workflow that eventually finds a circuit as following three steps. Researchers:
1. Observe a behavior (or task $^3$ ) that a neural network displays, create a dataset that reproduces the behavior in question, and choose a metric to measure the extent to which the model performs the task.
2. Define the scope of the interpretation, i.e. decide to what level of granularity (e.g. attention heads and MLP layers, individual neurons, whether these are split by token position) at which one wants to analyze the network. This results in a computational graph of interconnected model units.
3. Perform an extensive and iterative series of patching experiments with the goal of removing as many unnecessary components and connections from the model as possible.
Researchers repeat the previous three steps with a slightly different dataset or granularity, until they are satisfied with the explanation of the circuit components.
This work (ACDC) presents a tool to fully automate Step 3. Before we dive into the details of ACDC, we expand on what Steps 1-3 involve, and review examples from previous work that we use to evaluate ACDC.
2.1 Step 1: Select a behavior, dataset, and metric
The first step of the general mechanistic interpretability workflow is to choose a neural network behavior to analyze. Most commonly researchers choose a clearly defined behavior to isolate only the algorithm for one particular task, and curate a dataset which elicits the behavior from the model. Choosing a clearly defined behavior means that the circuit will be easier to interpret than a mix of circuits corresponding to a vague behavior. Some prior work has reverse-engineered the algorithm behind a small model's behavior on all inputs in its training distribution (Nanda et al., 2023; Chughtai, Chan, and Nanda, 2023), though for language models this is currently intractable, hence the focus on individual tasks.
We identified a list of interesting behaviors that we used to test our method, summarized in Table 1. These include previously analyzed transformer models (1 and 3 on GPT-2 Small, 2 and 6 on smaller language transformers) where researchers followed a workflow similar to the one we described above. Tasks 4 and 5 involve the full behavior of tiny transformers that implement a known algorithm, compiled with tracr (Lindner et al., 2023). For each task, we mention the metric used in previous work to measure the extent to which the model performs the task on the corresponding dataset.
2.2 Step 2: Divide the neural network into a graph of smaller units
To find circuits for the behavior of interest, one must represent the internals of the model as a computational directed acyclic graph (DAG, e.g. Figure 2a). Current work chooses the abstraction level of the computational graph depending on the level of detail of their explanations of model behavior. For example, at a coarse level, computational graphs can represent interactions between attention heads and MLPs. At a more granular level they could include separate query, key and value activations, the interactions between individual neurons (see Appendix I), or have a node for each token position (Wang et al., 2023).
Node connectivity has to be faithful to the model's computation, but that does not fully specify its definition. For example, following Elhage et al. (2021), many works consider the connections between model components in non-adjacent layers due to the additivity of the residual stream, even though these are computed with dynamic programming in the actual model implementation. Connectivity defines what is considered a direct or a mediated interaction (Pearl, 2009; Vig et al., 2020). See for example Figure 2a, where component B has both a direct effect on the output node O and an indirect effect on the output through component A.
<sup>3</sup>Section 3 formally defines "task". We use "behavior" and "task" interchangeably.
Task	Example Prompt	Output	Metric
1: IOI
(Appendix F.2)	“When John and Mary went to the store, Mary gave a bottle of milk to”	“_John”	Logit
difference
2: Docstring
(Appendix H.1)	def f(self, files, obj, state, size, shape, option):
"\"document string example
:param state: performance analysis
:param size: pattern design
:param	“_shape”	Logit
difference
3: Greater-Than
(Appendix G)	“The war lasted from 1517 to 15”	“18” or “19” or ...or “99”	Probability
difference
4: trunc-xproportion
(Appendix I.1)	["a", "x", "b", "x"]	[0, 0.5, 0.33, 0.5]	Mean Squared
Error
5: trunc-reverse
(Appendix I.2)	[0, 3, 2, 1]	[1, 2, 3, 0]	Mean Squared
Error
6: Induction
(Section 4.2)	“Vernon Dursley and Petunia Durs”	“ley”	Negative log-
probability
Table 1: Five behaviors for which we have an end-to-end circuit from previous mechanistic interpretability work, plus Induction. We automatically rediscover the circuits for behaviors 1-5 in Section 4. Tokens beginning with space have a “_” preponded for clarity.
2.3 Step 3: Patch model activations to isolate the relevant subgraph
With the computational DAG specified, one can search for the edges that form the circuit. We test edges for their importance by using recursive activation patching: i) overwrite the activation value of a node or edge with a corrupted activation, ii) run a forward pass through the model, and iii) compare the output values of the new model with the original model, using the chosen metric (Section 2.1). One typically starts at the output node, determines the important incoming edges, and then investigates all the parent nodes through these edges in the same way. It is this procedure that ACDC follows and automates in Algorithm 1.
Patching with zeros and patching with different activations Activation patching methodology varies between mechanistic interpretability projects. Some projects overwrite activation values with zeros (Olsson et al., 2022; Cammarata et al., 2021), while others erase activations' informational content using the mean activation on the dataset (Wang et al., 2023). Geiger et al. (2021) prescribe interchange interventions instead: to overwrite a node's activation value on one data point with its value on another data point. Chan et al. (2022) justify this by arguing that both zero and mean activations take the model too far away from actually possible activation distributions. Interchange interventions have been used in more interpretability projects (Hanna, Liu, and Variengien, 2023; Heimersheim and Janiak, 2023; Wang et al., 2023), so we prefer it. However we also compare all our experiments to replacing activations with zeros (Section 4.2, Appendix E.2).
2.4 Explaining the circuit components
After successfully isolating a subgraph, one has found a circuit (Section 1). The researcher then can formulate and test hypotheses about the functions implemented by each node in the subgraph. There is early evidence that ACDC is helpful for making novel observations about how language models complete tasks, such as the importance of surprising token positions that help GPT-2 Small predict correctly gendered pronouns (Appendix K). In our work we focus on automating the time-consuming step 3 that precedes functional interpretation of internal model components, though we think that automating the functional interpretation of model components is an exciting further research direction.
(a) Choose computational graph, task, and threshold  $\tau$
(b) At each head, prune unimportant connections.
(c) Recurse until the full circuit is recovered.
Figure 2: How ACDC works (Steps 2a-2c). Step 2a: a practitioner specifies a computational graph of the model, the task they want to investigate, and a threshold under which to remove connections. Step 2b: ACDC iterates over nodes in the computational graph, replacing activations of connections between a node and its children, and measuring the effect on the output metric. Connections are removed if their measured effect on the metric under corruption is below the threshold  $\tau$ . Step 2c: recursively apply Step 2b to the remaining nodes. The ACDC procedure returns a subgraph of the original computational graph.
3 Automating circuit discovery (Step 3)
This section describes algorithms to automate Step 3 of the mechanistic interpretability workflow (Section 2.3). In all three cases, we assume that the 'task' being studied is defined by a set of prompts  $(x_{i})_{i = 1}^{n}$  on which the model's predictions have a noticeable pattern (see Table 1 for examples) and a set of prompts  $(x_{i}^{\prime})_{i = 1}^{n}$  where this task is not present. We then use the activations of the models on a forward pass on the points  $x_{i}^{\prime}$  as corrupted activations (Section 2.3).
Automatic Circuit Discovery (ACDC). Informally, a run of ACDC iterates from outputs to inputs through the computational graph, starting at the output node, to build a subgraph. At every node it attempts to remove as many edges that enter this node as possible, without reducing the model's performance on a selected metric. Finally, once all nodes are iterated over, the algorithm (when successful) finds a graph that i) is far sparser than the original graph and ii) recovers good performance on the task.
To formalize the ACDC process, we let  $G$  be a computational graph of the model of interest, at a desired level of granularity (Section 2.2), with nodes topologically sorted then reversed (so the nodes are sorted from output to input). Let  $H \subseteq G$  be the computational subgraph that is iteratively pruned, and  $\tau > 0$  a threshold that determines the sparsity of the final state of  $H$ .
We now define how we evaluate a subgraph  $H$ . We let  $H(x_{i},x_{i}^{\prime})$  be the result of the model when  $x_{i}$  is the input to the network, but we overwrite all edges in  $G$  that are not present in  $H$  to their activation on  $x_{i}^{\prime}$  (the corrupted input). This defines  $H(x_{i},x_{i}^{\prime})$ , the output probability distribution of the subgraph under such an experiment. Finally we evaluate  $H$  by computing the KL divergence  $D_{\mathrm{KL}}(G(x_i)||H(x_i,x_i'))$  between the model and the subgraph's predictions. We let  $D_{\mathrm{KL}}(G||H)$  denote the average KL divergence over a set of datapoints. Appendix C discusses alternatives to the KL divergence, and Appendix E.1 explores the consequences of optimizing the task-specific metrics from Table 1 instead.
Algorithm 1 describes ACDC. The order in which we iterate over the parents  $w$  of  $v$  is a hyperparameter. In our experiments the order is lexicographically from later-layer MLPs and heads to earlier-layer MLPs and heads, and from higher- to lower-indexed heads. We note that in one case in our work, the order of the parents affected experimental results (Appendix J).
Subnetwork Probing (SP; Cao, Sanh, and Rush, 2021). SP learns a mask over the internal model components (such as attention heads and MLPs), using an objective that combines accuracy and
4To implement the computation of  $H(x_{i},x_{i}^{\prime})$  , we initially run a forward pass with the unmodified model on the input  $x_{i}^{\prime}$  and cache all activations.
Algorithm 1: The ACDC algorithm.
Data: Computational graph  $G$  dataset  $(x_{i})_{i = 1}^{n}$  corrupted datapoints  $(x_i^{\prime})_{i = 1}^{n}$  and threshold  $\tau >0$  Result: Subgraph  $H\subseteq G$    
1  $H\gets G$  // Initialize H to the full computational graph   
2  $H\gets H.$  reverse_topological_sort() // Sort H so output first   
3 for  $v\in H$  do   
4 for w parent of v do   
5  $\begin{array}{rlr}{H_{\mathrm{new}}\leftarrow H\setminus \{w\to v\}} & {} & {\mathrm{/~Temporarily~remove~candidate~edge}}\\ {\mathrm{if}D_{KL}(G||H_{\mathrm{new}}) - D_{KL}(G||H) <   \tau \mathrm{then}}\\ {|H\leftarrow H_{\mathrm{new}}\qquad \mathrm{/~Edge~is~unimportant,~remove~permanently}}\\ {\mathrm{end}} \end{array}$    
6   
7   
8   
9 end   
10 end   
11 return H
sparsity (Louizos, Welling, and Kingma, 2018), with a regularization parameter  $\lambda$ . At the end of training, we round the mask to 0 or 1 for each entry, so the masked computation corresponds exactly to a subnetwork of a transformer. SP aims to retain enough information that a linear probe can still extract linguistic information from the model's hidden states. In order to use it to automate circuit discovery, we make three modifications. We i) remove the linear probe, ii) change the training metric to KL divergence as in Section 2, and iii) use the mask to interpolate between corrupted activations and clean activations (Section 3) rather than zero activations and clean activations. Appendix D.1 explains the details of these changes.
Head Importance Score for Pruning (HISP; Michel, Levy, and Neubig, 2019). HISP ranks the heads by importance scores (Appendix D.2) and prunes all the heads except those with the top  $k$  scores. Keeping only the top  $k$  heads corresponds to a subnetwork that we can compare to ACDC. We plot the ROC obtained from the full possible range of  $k$ . Like SP, this method only considers replacing head activations with zero activations, and therefore we once more generalize it to replace heads and other model components with corrupted activations (for details, see Appendix D.2).
4 Evaluating Subgraph Recovery Algorithms
To compare methods for identifying circuits, we seek empirical answers to the following questions.
- Q1: Does the method identify the subgraph corresponding to the underlying algorithm implemented by the neural network?
- Q2: Does the method avoid including components which do not participate in the elicited behavior?
We attempt to measure Q1 and Q2 using two kinds of imperfect metrics: some grounded in previous work (Section 4.1), and some that correspond to stand-alone properties of the model and discovered subgraph (Section 4.2).
4.1 Grounded in previous work: area under ROC curves
The receiver operating characteristic (ROC) curve is useful because a high true-positive rate (TPR) and a low false-positive rate (FPR) conceptually correspond to affirming Q1 and Q2, respectively.
We consider canonical circuits taken from previous works which found an end-to-end circuit explaining behavior for tasks in Table 1. We formulate circuit discovery as a binary classification problem, where edges are classified as positive (in the circuit) or negative (not in the circuit). Appendices F, G, H, I and J describe and depict the canonical circuits for each task. Appendix E.3 considers the node classification problem instead, which is less appropriate for ACDC but more appropriate for other methods.
Figure 3: ROC curves of ACDC, SP and HISP identifying model components from previous work, across 5 circuits in transformers. The points on the plot are cases where SP and ACDC return subgraphs that are not on the Pareto frontier. The corresponding AUCs are in Table 2.
We sweep over a range of ACDC thresholds  $\tau$ , SP regularization parameters  $\lambda$ , or number of HISP elements pruned  $k$ . We plot pessimistic segments between points on the Pareto frontier of TPR and FPR, over this range of thresholds (Fawcett, 2006). ACDC and SP optimize the KL divergence for tasks where this makes sense (all but tracr tasks, which use the L2 distance). All methods employ activations with corrupted data. Appendix C describes and Appendix E experiments with different design choices for the metric and activation patching methodology.
Figure 3 shows the results of studying how well existing methods recover circuits in transformers. We find that i) methods are very sensitive to the corrupted distribution, ii) ACDC has competitive performance (as measured by AUC) with gradient-descent based methods iii) ACDC is not robust, and it fails at some settings.
Several of the tasks appeared to require specific distributions and metrics for the areas under the curves to be large. For example, ACDC achieved poor performance on both tracr tasks in Fig. 3, but the circuit was perfectly recovered by ACDC at any threshold  $\tau >0$  when patching activations with zeros (Appendix I). Furthermore, ACDC achieves a greater AUC on the IOI and Greater-Than and tracr-reverse tasks than both of the other methods, and hence overall is the optimal algorithm. As an example of the variable performance of circuit recovery algorithms, on the Docstring task we achieve the high performance when using the ACDC algorithm with the docstring metric (Appendix H). However in other tasks such as the IOI task, ACDC performance was worse when optimizing for logit difference.
Further research in automated interpretability will likely yield further improvements to the FPR and TPR of circuit discovery. We outline limitations with all current methods, but also gesture at likely fundamental limitations of the false positive and true positive measures. A limitation with all existing methods is that they optimize a single metric. This means they systematically miss internal model components such as the "negative" components found in previous work (IOI, Docstring) that are actively harmful for performance. The IOI recovery runs were not able to recover negative heads when optimizing for logit difference. Even when optimizing for low KL divergence, the negative components were only recovered when very small thresholds were used (Figure 15).