SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Overall, we found that KL divergence was the least flawed of all metrics we tried in this project, and think more practitioners should use it given i) the problems with other metrics listed, and ii) how empirically it can recover several circuits that were found by researchers using other metrics.
D Details of Subnetwork Probing and Head Importance Score for Pruning
D.1 Subnetwork Probing
There are 3 modifications we made to Subnetwork Probing (Cao, Sanh, and Rush, 2021, SP) in our work. In this Appendix we provide technical detail and motivation for these modifications:
1. We do not train a probe. ACDC does not use a probe. Cao, Sanh, and Rush (2021) train a linear probe after learning a mask for every component. The component mask can be optimized without the probe, so we just omit the linear probing step.
2. We change the objective of the SP process to match ACDC's. ACDC uses a task-specific metric, or the KL divergence to the model's outputs (Algorithm 1). In order to compare the techniques in equivalent settings we use the same metric (be it KL divergence or task-specific) in SP. Cao, Sanh, and Rush (2021) use negative log probability loss.
3. We generalize the masking technique so we can replace activations with both zero activations and corrupted activations. Replacing activations with zero activations<sup>6</sup> is useful for pruning (as they improve the efficiency of networks) but are not as commonly used in mechanistic interpretability (Goldowsky-Dill et al., 2023), so we adapt SP to use corrupted activations. SP learns a mask  $Z$  and then sets the weights  $\phi$  of the neural network equal to  $\phi * Z$  (elementwise multiplication), and locks the attention mask to be binary at the end of optimization (Jang, Gu, and Poole, 2017). This means that outputs from attention heads and MLPs in models are scaled closer to 0 as mask weights are decreased. To allow comparison with ACDC, we can linearly interpolate between a clean activation when the mask weight is 1 and a corrupted activation (i.e a component's output on the datapoint  $x_{i}^{\prime}$ , in the notation of Section 3) when the mask weight is 0. We do this by editing activations rather than weights of the model.
Additionally, we used a constant learning rate rather than the learning rate scheduling used in Cao, Sanh, and Rush (2021).
The regularization coefficients  $\lambda$  (in the notation of Cao, Sanh, and Rush (2021)) we used in Figure 4 were 0.01, 0.0158, 0.0251, 0.0398, 0.0631, 0.1, 0.158, 0.251, 0.398, 0.631, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 30, 50, 70, 90, 110, 130, 150, 170, 190, 210, 230, 250.
The number of edges for subgraphs found with Subnetwork Probing are computed by counting the number of edges between pairs of unmasked nodes.
D.2 Head Importance Score for Pruning
In Section 4-4.2 we compared ACDC with the Head Importance Score for Pruning (Michel, Levy, and Neubig, 2019, HISP). We borrow the author's notation throughout this section, (particularly from Section 4.1) and adapt it so that this Appendix can be read after Section 2.
The authors use masking parameters  $\xi_{h}$  for all heads, i.e. scaling the output of each attention head by  $\xi_{h}$ , similar to the approach in Subnetwork Probing (Appendix D.1), so that each head  $\mathrm{Att}_h$  's output is  $\xi_{h}\mathrm{Att}_{h}(x_{i})$  on an input  $x_{i}$ . The authors keep MLPs present in all networks that they prune. We generalize their method so that it can be used for any component inside a network for which we can take gradients with respect to its output.
The authors define head importance as
$$
I _ {h} := \frac {1}{n} \sum_ {i = 1} ^ {n} \left| \frac {\partial \mathcal {L} \left(x _ {i}\right)}{\partial \xi_ {h}} \right| = \frac {1}{n} \sum_ {i = 1} ^ {n} \left| \operatorname {A t t} _ {h} \left(x _ {i}\right) ^ {T} \frac {\partial \mathcal {L} \left(x _ {i}\right)}{\partial \operatorname {A t t} _ {h} \left(x _ {i}\right)} \right|. \tag {3}
$$
<sup>6</sup>Which is generally equivalent to setting weight parameters equal to 0.
(a) Corrupted activations
(b) Zero activations
Figure 6: Examples of subgraphs recovered by ACDC on the induction task with different types of activations and threshold  $\tau = 0.5623$ . These find the two different induction heads (1.5, 1.6) and a previous token head (0.0) as identified by Goldowsky-Dill et al. (2023). (a) shows the result with corrupted activations, while (b) shows the result with zero activations.
where the equivalence of expressions is the result of the chain rule. We make three changes to this setup to allow more direct comparison to ACDC: i) we use a metric rather than loss, ii) we consider corrupted activations rather than just zero activations and iii) we use the 'head importance' metric for more internal components than merely attention head outputs.
Since our work uses in general uses a metric  $F$  rather than loss Appendix C, we instead use the derivative of  $F$  rather than the derivative of the loss. The HISP authors only consider interpolation between clean and zero activations, so in order to compare with corrupted activations, we can generalize  $\xi_h$  to be the interpolation factor between the clean head output  $\mathrm{Att}_h(x)$  (when  $\xi_h = 1$ ) and the corrupted head output  $\mathrm{Att}_h(x')$  (when  $\xi_h = 0$ ). Finally, this same approach works for any internal differentiable component of the neural network. Therefore we study the HISP applied to the query, key and value vectors of the model and the MLPs outputs.
In practice, this means that we compute component importance scores
$$
I _ {C} := \frac {1}{n} \sum_ {i = 1} ^ {n} \left| \left(C \left(x _ {i}\right) - C \left(x _ {i} ^ {\prime}\right)\right) ^ {T} \frac {\partial F \left(x _ {i}\right)}{\partial C \left(x _ {i}\right)} \right|. \tag {4}
$$
Where  $C(x_{i})$  is the output of an internal component  $C$  of the transformer, which is equivalent to 'attribution patching' (Nanda, 2023) up to the absolute value sign.
To compute the importance for zero activations, we adjust Equation (4) so it just has a  $C(x_{i})$  term, without the  $-C_h(x_i')$  term. We also normalize all scores for different layers as in Michel, Levy, and Neubig (2019). The identical setup to Equation (4) works for outputs of the query, key and value calculations for a given head, as well as the MLP output of a layer. In Section 4 we use query, key and value components for each head within the network, as well as the output of all MLPs.
The number of edges for subgraphs found with HISP is also computed by counting the number of edges between pairs of unmasked nodes, like Subnetwork Probing (Appendix D.1).
E Experimental study of algorithm design
This section evaluates design choices for ACDC and SP, by re-doing the experiments in Section 4. We explore two axes of variation.
In theory. In practice, components need be torch(nn). Modules such that we can calculate the gradient of  $F$  with respect to the components' outputs.
Table 2: AUCs for corrupted activations, Random Ablation. (E)  $=$  Edge, (N)  $=$  Node.
Metric	Task	ACDC(E)	HISP(E)	SP(E)	ACDC(N)	HISP(N)	SP(N)
KL	Docstring	0.982	0.805	0.937	0.950	0.881	0.928
Greaterthan	0.853	0.693	0.806	0.890	0.642	0.827
IOI	0.869	0.789	0.823	0.880	0.668	0.842
Loss	Docstring	0.972	0.821	0.942	0.938	0.889	0.941
Greaterthan	0.461	0.706	0.812	0.766	0.631	0.811
IOI	0.589	0.836	0.707	0.777	0.728	0.797
Tracr-Proportion	0.679	0.679	0.525	0.750	0.909	0.818
Tracr-Reverse	0.200	0.577	0.193	0.312	0.750	0.375
Table 3: AUCs for corrupted activations, Zero Ablation. (E)  $=$  Edge, (N)  $=$  Node.
Metric	Task	ACDC(E)	HISP(E)	SP(E)	ACDC(N)	HISP(N)	SP(N)
KL	Docstring	0.906	0.805	0.428	0.837	0.881	0.420
Greaterthan	0.701	0.693	0.163	0.887	0.642	0.134
IOI	0.539	0.792	0.486	0.458	0.671	0.605
Loss	Docstring	0.929	0.821	0.482	0.825	0.889	0.398
Greaterthan	0.491	0.706	0.639	0.783	0.631	0.522
IOI	0.447	0.836	0.393	0.424	0.728	0.479
Tracr-Proportion	1.000	0.679	0.829	1.000	0.909	1.000
Tracr-Reverse	1.000	0.577	0.801	1.000	0.750	1.000
- Minimizing the task-specific metric, rather than the KL divergence.
- Patching activations with zeros, rather than with the result on a corrupted input (interchange intervention).
- Looking at node-level TPR and FPR for the ROC curves, rather than edge-level.
The results paint a mixed picture of whether ACDC or SP is better overall, but reinforce the choices we implicitly made in the main text. A stand-out result is that ACDC with zero-patching is able to perfectly detect the tracr circuits (Figs. 9 and 10).
A numerical summary of the results is in Tables 2 and 3, which display the areas under the ROC curve (AUC) for all the design choices we consider.
E.1 Minimizing the task-specific metric, rather than the KL divergence
We ran ACDC, SP and HISP with the task-specific metric from Table 1, instead of KL divergence. The exact modification is described in Appendix C. The ROC result is in Fig. 7. Compared to minimizing KL divergence (Fig. 3), ACDC works better for Docstring, but worse for Greater-Than and IOI, indicating that it is not a very robust method.
We prefer using the KL divergence instead of the task-specific metric, because the task-specific metric can be over-optimized (Appendix C). This means that the recovered circuit ends up performing the task more than the original model, and is thus not accurate. We can observe this effect by comparing the task-specific metric achieved by the methods in Fig. 8.
E.2 Activation patching with zeros, instead of corrupted input
In the main text experiments that compared using corrupted activations and zero activations (Figure 4), all three methods recovered subgraphs with generally lower loss when doing activation patching with zeros, in both the experiments with the normal model and with permuted weights. It is unclear why the methods achieve better results with corruptions that are likely to be more destructive. A possible explanation is that there are 'negative' components in models (Appendix F.3) that are detrimental to the tasks, and the zero activations are more disruptive to these components. A discussion of how methods could be adjusted to deal with this difficulty can be found in Alternative 2 in Appendix C.
Figure 7: Edge-wise ROC curves generated by minimizing the task-specific metric in Table 1, rather than KL divergence.
Figure 8: Optimizing the task-specific metric of the subject model, on trained and reset networks. For each recovered circuit, we plot its task-specific metric (Table 1) against its number of edges. The reset networks metrics don't change much with the number of edges, which is good. We found that for IOI that extremely large logit differences could be achieved (over 15) but this didn't happen when the network had a large number of edges.
Figure 9: Edge-wise ROC curves generated by minimizing the KL divergence, but using zero activations.
Figure 10: Edge-wise ROC curves generated by minimizing the task-specific metric in Table 1, using zero activations.
E.3 Node-level ROC curve, rather than edge-level ROC curve
We compute the FPR and TPR of classifying whether a node belongs to the circuit. We consider this alternative task because SP and HISP operate at the node-level, whereas ACDC operates at the edge-level, so this is fairer to HISP and SP. The results are broadly similar to edge-level ROCs, and are described in Figs. 11 and 12 and Tables 2 and 3.
Figure 11: Node-level ROC when optimizing the KL divergence, with corrupted activations.
F IOI task: details and qualitative evidence
F.1 Further details on the IOI experiments
In the ACDC run in Figure 1, we used a threshold of  $\tau = 0.0575$ . We also removed all edges which did not lie on a directed path from the input (which is equivalent in computation since we use corrupted activations). Our library now only supports splitting query, key and input, rather than merely looking at the connections between heads. Additionally, For ease of visualization, in the diagram on the left of Figure 1 we removed all edges between grey nodes more than 2 layers apart, and  $90\%$  of the edges between grey and red nodes.
Our IOI experiments were conducted with a dataset of  $N = 50$  text examples from one template the IOI paper used (' When John and Mary went to the store, Mary gave a bottle of milk to'). The
Figure 12: Node-level ROC when optimizing the task-specific metric in Table 1, with corrupted activations.
Figure 13: The IOI circuit (Figure 2 from Wang et al. (2023)). We include edges between all pairs of heads connected here, through the Q or the  $\mathrm{K + V}$  as indicated. We also include all connections between MLPs and these heads, the inputs and outputs. See https://github.com/ArthurConmy/Automatic-Circuit-Discovery/blob/main/acdc/ioi/util.py#L205. The full circuit is in Fig. 14.
corrupted dataset was examples from the ABC dataset (Wang et al., 2023) — for example ‘When Alice and Bob went to the store, Charlie gave a bottle of milk to’.
In the IOI experiment in Figure 1, we did not split the computational graph into the query, key and value calculations for each head. This enabled the ACDC run to complete in 8 minutes on an NVIDIA A100 GPU. However, the larger experiments that kept  $>10\%$  of the edges of the original edges in the computational graph sometimes took several hours. On one hand we don't expect these cases to be very important for circuit discovery, but they make up the majority of the points of the pareto frontier of curves in this paper.
F.2 The IOI circuit
Wang et al. (2023) find a circuit 'in the wild' in GPT-2 small (Radford et al., 2019). The circuit identifies indirect objects (see for example Table 1) by using several classes of attention heads. In this subsection we analyze how successful ACDC's circuit recovery (Figure 1) is. All nine heads found in Figure 1 belong to the IOI circuit, which is a subset of 26 heads out of a total of 144 heads in GPT-2 small. Additionally, these 9 heads include heads from three different classes (Previous Token Heads,
Figure 14: Our low-level implementation of the IOI circuit (Wang et al., 2023), in terms of heads split by query, key, value; and MLPs. It has 1041 edges. For edge between groups  $A$ ,  $B$  in Fig. 13, we connect each member of group  $A$  with each member of group  $B$ . The colors of groups in this figure correspond to the group colors in Fig. 13.
S-Inhibition Heads and Name Mover Heads) and are sufficient to complete the IOI task, showing that ACDC indeed can recover circuits rather than just subgraphs.
For our ROC plots, we considered the computational graph of IOI described in Figure 13.
The ground-truth circuit gets a logit difference of 3.24 compared to the model's 4.11 logit difference. It has a KL divergence of 0.44 from the original model.
F.3 Limitations of ACDC in recovering the IOI circuit
The main figure from the paper Figure 1 shares several features with circuits recovered with similar thresholds, even when logit difference rather than KL divergence is minimized. The figure does not include heads from all the head classes that Wang et al. (2023) found, as it does not include the Negative Name Mover Heads or the Previous Token Heads. In Figure 15 we run ACDC with a lower threshold and find that it does recover Previous Token Heads and Negative Name Mover Heads, but also many other heads not documented in the IOI paper. This is a case where KL divergence performs better than logit difference maximisation (which does not find Negative Name Movers at any threshold), but still is far from optimal (many extraneous heads are found). Ideally automated circuit discovery algorithms would find negative components even at higher thresholds, and hence we invite future empirical and theoretical work to understand negative components and build interpretability algorithms capable of finding negative components. An early case study gaining wide understanding of a negative head in GPT-2 Small can be found in McDougall et al. (2023).
Figure 15: A subset of the 443/32923 edges of GPT-2 Small that ACDC recovered when optimizing for KL divergence at threshold  $\tau = 0.00398$ . This subset includes edges between Negative Heads (10.7 and 11.10). A number of heads not found by the IOI work (9.2 and 11.11 for example) were also found.
Figure 16: The ground-truth Greater-Than circuit (16a) and a circuit that ACDC recovers (16b).
G Greater-Than task: details and qualitative evidence
We use a random sample of 100 datapoints from the dataset provided by Hanna, Liu, and Variengien (2023).
We use the circuit from Figure 13 from their paper, including connections between MLPs that are in the same group (e.g. MLPs 8, 9, 10 and 11) but not including connections connections between attention heads in the same group. We also include all Q and K and V connections between attention heads present. Their circuit includes all earlier layer connections to the queries of the mid-layer attention heads that cause a logit spike (Figure 16a). This would account for more than half of the edges in the circuit were we to include all such edges that compute these query vectors, and hence we compromised by just adding all early layer MLPs as connections to these query vectors. Full details on the circuit can be found in our codebase<sup>8</sup> and in Fig. 17. This circuit gets a probability difference score of  $72\%$  on a subset of their dataset for which GPT-2 Small had an  $84\%$  probability difference and a KL divergence from the original model of 0.078.
<sup>8</sup>https://github.com/ArthurConmy/Automatic-Circuit-Discovery/blob/main/acdc/greaterthan/util.py#L231
Figure 17: Our low-level implementation of the Greater-Than circuit (Hanna, Liu, and Variengien, 2023), in terms of heads split by query, key, value; and MLPs. It has 262 edges. For edge between groups  $A$ ,  $B$  in Fig. 16a, we connect each member of group  $A$  with each member of group  $B$ .
An example subgraph ACDC recovered, including a path through MLP 0, mid-layer attention heads and late-layer MLPs, is shown in Figure 16b. This run used the Greater-Than probability difference metric and a threshold of 0.01585 and recovered the results in the abstract.
H Docstring task: details and qualitative evidence
H.1 The docstring circuit
Heimersheim and Janiak (2023) find a circuit in a small language model that is responsible for completing Python docstrings. The model is a 4-layer attention-only transformer model, trained on natural language and Python code. This circuit controls which variable name the model predicts in each docstring line, e.g for the prompt in Table 1 it chooses shape over the other variable names files, obj, state, size, or option.
The circuit is specified on the level of attention heads, consisting of 8 main heads (0.2, 0.4, 0.5, 1.4, 2.0, 3.0, and 3.6) that compose over four layers, although it only makes use of three levels of composition. It consists of 37 edges between inputs, output, and the attention heads, as we show in Figure 18a. We discuss below why we exclude 0.2 and 0.4.
We apply the ACDC algorithm (Section 2) to this dataset, using the prompts from Heimersheim and Janiak (2023) found in their accompanying Colab notebook. For our corrupted dataset we use their random_random dataset which randomizes both the variable names in the function definition as well as in the docstring of prompts.
ACDC generates the subgraph shown in Figure 18b. We now compare this to the original 8-head circuit from Heimersheim and Janiak (2023), which was the most specific circuit evaluated in that work. We refer to this circuit as the 'manual' circuit to distinguish it from the ground truth, which
${}^{9}$  Available at https://colab.research.google.com/drive/17CoA1yARaWHvV14zQGcI3ISz1bIRZKS5 as of 8th April 2023
(a) Canonical Docstring circuit (37 edges)
(b) ACDC circuit (KL-divergence,  $\tau = 0.095$ )
Figure 18: Our implementation of the Docstring circuit (Heimersheim and Janiak, 2023), compared to an ACDC-generated circuit.
Metric	Full model	ACDC KL τ = 0.005	ACDC KL τ = 0.095 (Fig. 18b)	ACDC LD τ = 0.067 (Fig. 19)	Manual 8 heads, all connections	Ground-truth circuit (Fig. 18a)
KL-divergence	0	0.33	1.2	0.67	0.83	1.1
Mean logit diff.	0.48	0.58	-1.7	0.32	-0.62	-1.6
Num. of edges	1377	258	34	98	464	37
Table 4: Comparing our ACDC docstring results to the ground-truth from Heimersheim and Janiak (2023) using their metrics. We compare (from left to right) the full model, the subgraph from ACDC runs optimizing for KL divergence ( $\tau = 0.005$  and 0.095) and logit difference ( $\tau = 0.067$ ), as well as the two subgraphs made manually from Heimersheim and Janiak (2023): One including all connections between the given attention heads, and one using only the given circuit. The metrics used are KL divergence between full-circuit outputs and resemble-ablated output (lower is better), mean logit difference between correct and wrong completions (higher is better), and the number of edges in the circuit (lower is better).
includes the edge connections that the authors speculated were most important but did not evaluate due to a lack of software for edge-editing. We find (a) overlapping heads, (b) heads found by ACDC only, and (c) heads found in the manual interpretation only. In the first class (a) we find heads 0.5, 1.4, 2.0, 3.0, and 3.6. All these manually identified heads are recovered by ACDC. In class (b) we find head 1.0 which the authors later add to their circuit to improve performance; ACDC shows for the first time where this head belongs in the circuit. In class (c) we find heads 0.2, 0.4 and 1.2. However, the first two of these are not actually relevant under the docstring distribution and only added by the authors manually. Head 1.2 is considered a non-essential but supporting head by the authors and not identified by ACDC at the chosen threshold of  $\tau = 0.095$  (for KL divergence). This might be because 1.2 is less important than the other heads, and indeed we recover this head in larger subgraphs (such as the subgraph in Figure 19).
We compare the numerical results between the ACDC circuits and the circuit described in Heimer-sheim and Janiak (2023) in Table 4. In addition to the  $\tau = 0.095$  run (Figure 18b) we perform a run with lower KL divergence threshold of  $\tau = 0.005$  recovering a larger circuit (258 edges) containing also head 1.2 that was missing earlier.
Figure 19: ACDC-found subgraph for docstring task minimizing logit difference  $(\tau = 0.067)$  instead of KL divergence.
Since Heimersheim and Janiak (2023) use logit difference as their metric, we add an ACDC run that optimizes logit difference rather than KL divergence (see Appendix C for details on this adjustment) with threshold  $\tau = 0.067$ . This circuit (Figure 19) recovers the relevant manual-interpretation heads (including 1.2) as well.[10] It is even more specific, containing  $93\%$  less edges than the full circuit. This is also  $79\%$  less edges than the head-based circuit from Heimersheim and Janiak (2023) while achieving a better score on all metrics.
Note that there are two versions of the manual circuit we consider. There is (i) the set of 8 heads given in Heimersheim and Janiak (2023) that the authors test with a simple methods (not specifying edges), and (ii) the circuit of 39 edges as suggested by the authors that they were not able to test due to not having software to implement editable transformer computational graphs in PyTorch. We reconstruct this circuit, shown in Figure 18a, from their descriptions and perform tests (Table 4).
In case (i) the ACDC run (threshold  $\tau = 0.005$ ) achieves better performance in both metrics, Logit Difference and KL divergence, while being more specific (258 edges) when compared to the set of heads found by Heimersheim and Janiak (2023). In the more specific case (ii) the ACDC run (with threshold  $\tau = 0.095$ ) closely matches the manual interpretation, with a very similar circuit recovered (Figure 18). The ACDC run is slightly more specific but has slightly worse KL divergence and Logit Difference.