Finding Transformer Circuits With Edge Pruning
Authors: Adithya Bhaskar
Alexander Wettig
Dan Friedman
Danqi Chen
Venue: NeurIPS 2024 spotlight (2024)
Keywords: interpretability
circuits
pruning
optimization
Primary area: interpretability_and_explainability
TL;DR: We propose a circuit finding method that prunes connections between model components, and show it performs and scales well.

Abstract: The path to interpreting a language model often proceeds via analysis of circuits---sparse computational subgraphs of the model that capture specific aspects of its behavior. Recent work has automated the task of discovering circuits. Yet, these methods have practical limitations, as they either rely on inefficient search algorithms or inaccurate approximations. In this paper, we frame circuit discovery as an optimization problem and propose _Edge Pruning_ as an effective and scalable solution. Edge Pruning leverages gradient-based pruning techniques, but instead of removing neurons or components, prunes the _edges_ between components. Our method finds circuits in GPT-2 that use less than half the number of edges than circuits found by previous methods while being equally faithful to the full model predictions on standard circuit-finding tasks. Edge Pruning is efficient on tasks involving up to 100,000 examples, outperforming previous methods in speed and producing substantially better circuits. It also perfectly recovers the ground-truth circuits in two models compiled with Tracr. Thanks to its efficiency, we scale Edge Pruning to CodeLlama-13B, a model over 100x the size of GPT-2.
We use this setting for a case study, where we compare the mechanisms behind instruction prompting and in-context learning.
We find two circuits with more than 99.96% sparsity that match the performance of the full model. Further analysis reveals that the mechanisms in the two settings overlap substantially. This shows that Edge Pruning is a practical and scalable tool for interpretability, 
which can shed light on behaviors that only emerge in large models.

Finding Transformer Circuits with Edge Pruning
Adithya Bhaskar
Alexander Wettig
Dan Friedman
Danqi Chen
Princeton Language and Intelligence (PLI), Princeton University
adithhab@pprinceton. edu fawettig, dfriedman, danqic)es bangiches.&princeetor. @princeton.edi edu
Abstract
The path to interpreting a language model often proceeds via analysis of circuits- sparse computational subgraphs of the model that capture specific aspects of its behavior. Recent work has automated the task of discovering circuits. Yet, they rely either inefficient search
these methods have practical limitations, algorithms or inaccurate approximations. In this paper, we frame automated circuit effective
optimization problem and propose Edge Pruning and scalable solution. Edge Pruning leverages gradient-based pruning techniques, but instead of removing neurons components, it prunes the edges between components. Our method finds circuits in GPT-2 that use less than half the number of edges compared
discovery
circuits found by previous methods while being equally faithful to the full model predictions on standard circuit- cireuit-fidding tasks. Edge Pruning is efficient even with many as 100K examples, outperforming previous methods in speed and producing substantially better circuits. It also perfectly recovers the ground-truth circuits two models compiled with Tracr. Thanks to its efficiency, we scale Edge Pruning to Codel. laeeellaaa13333 methods operate on. We use this setting for case study comparing the mechanisms behind instruction prompting and in-context learning. We find two circuits with more than 99.96% sparsity that match the performance the full model and reveal that the mechanisms in the two settings overlap substantially. Our case study shows that Edge Pruning a practical and scalable tool for interpretability and sheds light behaviors that only emerge in large models.'
model over 100x the scale that prior
Introduction
Mechanistic interpretability strives to understand models via bottom-up descriptions of their com- ponents (e.g., attention heads and MLPs Transformers [Vaswani et al., 2017]). This typically proceeds via the identification and analysis of circuit [Olah et al., 2020, Elhage et al., 2021] sparse computational subgraph of the model that captures the aspects of its behavior we wish to study. The arduous process of identifying circuits (e.g., Wang et al. [2023]) was recently automated by ACDC [Conmy et al., 2023] and EAP [Syed et al., 2023]. However, ACDC uses an expensive greedy search that ablates each edge to estimate its importance. It cannot scale to datasets beyond few hundred examples or to parrameee models. EAP, on the other hand, uses gradient-based linear approximations of activation patching to estimate the importance of all edges simultaneously. While fast, these first-order approximations often sacrifice faithfulness to the full model. Besides, this approach ignores the impact of the presence/absence of other edges on the score.
In this paper, we frame circuit discovery as an optimization problem and tackle it via gradient-based pruning, rather than discrete search or first-order approximations. As such, we adapt pruning for the goal of circuit discovery instead of model compression. Rather than components, we prune the
We release our code and data publicly at
Pruning.
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

LM head
LM head
LM head
MLP
MLP
MLP
k/m|]
A(m].K
M/x].V
A[]] Q
AL12 K
KKK
V
Leanaalle hinary masks
nodes
(a) Regular Transformer
(b) Optimize edge masks
(c) Obtain sparse circuit
Figure 1: Edge Pruning disentangles the residual stream and optimizes continuous masks on the read operations via gradient descent. Discretizing the masks to {0, yields the final circuit. The full model corresponds to the case where all masks equal
edges between components and replace missing edges with counterfactual activations from corrupted examples. We enable this by replacing the residual stream of a Transformer (Figure la) with disentangled residual stream [Lindner et al., 2023, Friedman et al., 2023], which retains list of all previous activations. This allows us to introduce edge masks that determine from which components to read. We then leverage discrete optimization techniques such as Lo regularization [Louizos et al., 2018] to optimize these edge masks and produce sparse circuits (Figure lc).
We evaluate our approach, Edge Pruning, on four fronts: (I) we measure how faithfully the discovered circuits describe the behavior of the full model, (2) we verify it can recover ground-truth circuits in Tracr models [Lindner et al., 2023] compiled from known program descriptions, (3) we evaluate how the method scales more examples and (4) we assess its ability to find extremely sparse circuits in multi-billion parameter models. On four standard circuit- findiin-iiigg tasks, Edge Pruning finds circuits GPT-2 Small [Radford et al., 2019] that are consistently more faithful to the full model and have better task performance than circuits found by prior methods. The gap is especially pronounced on more complex tasks like multi-template IOI [Wang et al., 2023], where we find circuits that have 2.65x fewer edges but describe model outputs just as faithfully as the circuit found by the next-best method. We show that Edge Pruning scales effectively to a version IOI with 100K examples, where it outperforms prior methods in terms of speed and performance. Edge Pruning also perfectly recovers ground-truth circuits in two models compiled from known program descriptions with Tracr.
Finally, we establish that Edge Pruning scales to CodeLlama- 13B [Rozière et al., 2024]-100: -10x the size of models typically tackled by automated circuit discovery methods in case study. Specifically, we compare the mechanisms behind instruction instruction-prompiing and in-context learning [Brown et al., 2020] on Boolean Expressions- a task adapted from the BBH [Suzgun et al., 2022] benchmark. Edge Pruning finds circuits with just 0.04% of model edges that match the model's 's performance in either setting. Interestingly, the few-shot circuit performs well when instruction -rrpriion---ppdedd..-eoo and vice versa. The two circuits also have substantial overlap (62.7% edges of the sparser circuit), and the circuit formed by this intersection also performs significantly above chance on the task. We infer that the model relies shared mechanisms in the two settings. This case study demonstrates how Edge Pruning can inform the analysis phenomena that only emerge in large models.
In summary, our contributions are as follows:
We propose Edge Pruning, an effective and scalable method for automated circuit finding.
2. We demonstrate that Edge Pruning is competitive with or better than state-of-the-art methods simple tasks, and significantly superior on more complex ones, in terms of faithfulness and performance. Edge Pruning also scales well with more examples. Further, perfectly recovers ground-truth circuits in two Transformers compiled by Tracr.

---

3. We scale Edge Pruning to CodeLlama- 13B model over 100x larger than GPT-2 Small on a task adapted from BBH. Our case study finds that mechanisms underlying in-context learning and instruction-prompaeee
in CodeI for this task overlap significantly.
2 Background: Circuit Discovery
The goal of circuit discovery is to facilitate mechanistic understanding of Transformers by identify- ing the subset of a model's ' computational graph that is most relevant to a particular model behavior. In this section, define the computational graph of Transformer, formalize the objective for circuit discovery, and discuss the approaches of previous work.
The computational graph of Transformers. The Transformer architecture consists of a sequence of layers, namely attention layers and MLPs, which operate on the residual stream (Figure la) [Elhage et al., 2021]. The r'th layer fi reads the current state of the residual stream hs, computes its activations y/ = fi(h.), and applies it as an additive update to the residual stream hi+1 h2 + y1. We can expand this recurrence to make the dependence on prior outputs explicit:
yi = fi
Yo
+
Vj
(1)
where yo is the initialization of the residual stream with the input embeddings. We can represent the dependencies between layers as directed edges in a computational graph, where the edge -> i denotes the connection between the output of layer j to the input of layer i. Note that the computational graph may be defined at a more granular level. For instance, Conmy et al. [2023] split attention layers into multiple parallel attention heads, and represents each head by four interconnected nodes. The query/key/valle nodes receive separate input edges from previous layers, and the output node has outbound edges to downstream layers. We also follow this convention.
Circuits as subgraphs. A circuit a computational subgraph C C C. where and G denote the set of edges in the circuit and full model, respectively [Olah et al., 2020]. How do model Transformer with a missing edge j + i? Instead of simply removing the term y from the sum of inputs to node i, we adopt the approach of interchange ablation [Geiger et al., 2020, Zhang and Nanda, 2024]. For each example æ, the user provides a corrupted example I. which should consist of small change model to compute corrupted activations y, for all nodes. When an edge j - i is removed from a circuit, we replace the contribution of yj at the input of node with the corrupted activation Új- This ensures that the summed activations remain in-distribution [Zhang and Nanda, 2024] and it frames the decision to remove an edge counterfactual intervention [Vig et al., 2020].
that would result in different label in the task. We f input to the full
Circuit discovery. The goal of circuit discovery [Olah et al., 2020] to find a sparse subgraph that describes the behavior of the full model on a particular task. We use pc(y x, .) to denote the output of the Transformer circuit given original and corrupted examples x, H, and denote the output of the full model as pa (vy z) as the output of the full model. Formally, circuit discovery has the objective, arg min (ET [D(p(y)
z) Pc(y xx) subject to 1 - |CI//G 2
(2)
where the constraint enforces target sparsity of the circuit. T denotes the task distribution of interest, for which the user curates pairs of clean and corrupted examples (ar, .) that differ in crucial task features. The loss function D should capture the discrepancy between the outputs of the full model and the circuit; for language models, natural choice is the KL divergence between token predictions.
Previous approaches. We now discuss how previous methods approximate this combinatorial optimization problem and the limitations of their approaches.
1. ACDC [Conmy et al., 2023] proposes to solve the above objective using greedy search at each iteration, ACDC evaluates the effect of removing each edge individually, and removes any edge whose effect on the target metric less than specified threshold. This fails to capture the relative importance of edges and their interaction. Furthermore, the number of steps the algorithm scales linearly with the number of edges, which is prohibitive at larger model sizes (c.g., CodeLlama- 13B with 3.88M edges).
3

---

2. Edge Attribution Patching (EAP) [Syed et al., 2023] makes a linear (frrtt-oddr)) approx- imation of activation patching to assign an importance score each edge. This defines a ranking over edges, from which the top-k edges are used to form circuit of specific sparsity. While the linear approximation can compute the importance scores efficiently in single step, it is likely to find suboptimal solutions to the circuit discovery problem.
3. Conmy et al. [2023] compare to two mouning-based approaches. These either (1) prune attention heads based on estimated importance scores [Michel et al., 2019], or (2) perform structured pruning of nodes to identify the most important nodes [Cao et al., 2021]. These approaches perform worse than ACDC [Conmy et al., 2023]. Our approach differs in that we prune edges instead of neurons or nodes. This allows us to optimize at finer granularity but introduces an additional challenge as we will discuss in Section 3.
Method: Edge Pruning
In structured pruming [Wang et al., 2020, Xia et al., 2022], components such as layers and attention heads are removed increase the inference efficiency of models. The removal of a component can be modeled by a binary mask, which is relaxed to continuous parameter to be trainable with gradient-based optimization. While structured pruning produces subgraphs with fewer nodes, they are typically too coarse-grained to help with the mechanistic interpretability of model's 's computations.
We propose Edge Pruning, where define masks not nodes but over the edges connecting them. Specifically, we freeze the original model weights and introduce new trainable parameters
x E [0, 11 where [G] is the number edges in the Transformer, and the parameter 25z is a relaxed binary mask for the edge + i. In other words, the pruning mask indicates whether edge is included (2jz = 1) or removed (Zj: = 0) from the computational graph of a circuit. This formulation allows find subgraphs with greater granularity and precision compared to structured pruning, as the number of edges scales quadratically with the number of nodes in a model's computational graph.
While structured pruning discards pruned nodes by setting their activation to O, the application to interpretability calls for more careful treatment of missing nodes and edges. Specifically, the activation of a removed edge
should be replaced by the interchange activation obtained from the corrupted version of the example (Section 2). To allow gradient- optimization, we model the process as the masks continuously interpolating between the clean and corrupted activation. Specifically, we parameterize the i' 'th component as,
V/ fi 20:40 + (1 Zos
+
(2 (1 2zx) 01)
upstream of
where {ūg} denote the corrupted activations corresponding
Our formulation has a key challenge. Each node sees a different combination of activations depending on incoming edges, and thus different residual stream. Thus, can longer add the activations ya immediately to the residual stream, i.e. h = h2 + y/s as shown in Figure 1a. Instead, we modify the Transformer architecture to retain so-called disentangled residual stream [Friedman et al., 2023], in which the activations y/ are concatenated to a list of all previous activations (yo, .l........................... .).-...9i-1>> Then, dynamically aggregate these activations at the input of each node (Equation and Figure 1b).
In practice, concatenation increases the GPU memory footprint during training compared to regular edges between nodes that
structured pruning (Appendix A), but it is necessary for optimizing are separated by many layers. Despite the memory overhead, we demonstrate in Section 5 that we can scale our method to large models by parallelizing training over multiple GPUs.
We directly optimize the objective in (2) by performing stochastic gradient descent with respect to the edge weights z. The target sparsity is enforced via Lo regularization with a Lagrangian term. We leverage the formulation of Louizos et al. [2018] to model the masks as hard concrete parameters and to circumvent the non-differentiabiliyy of the L0 term. At the end of the training, the edge weights are converted to binary masks based on a threshold (e.g., 0.5), which uniquely determines the produced circuit (Figure lc). We now describe this process in more detail.

---

Details of the Edge Pruning process Our formulation of pruning is based on that used by CoFi Pruning [Xia et al., 2022]. Specifically, we model the masks x based on the hard concrete distribution as done by Louizos et al. [2018]:
Uniform(e, 1
- log
x (r-1) +1
z min(1 1, max(0, S))
where a refers the sigmoid function, 10 and log indicates that the logarithm is applied
element -iie.. We fix the temperature 8 = 3. The last two lines stretch the distribution to [4,
[-0.1, 1.1] and accumulate the "excess" probability either side to 0 and 1, respectively. The log alphas log a are the learnable parameters in this formulation.
Following, Wang et al. [2020], target sparsity is enforced via Lagrangian term [Louizos et al., 2018]. the current sparsity is 8, the term, parametrized by a reference value is
f. =1 ()=s) s) + A2 s)=
A and A2 are also updated during training via gradient ascent keep the regularization tight. We vary the value of t throughout training, linearly increasing it from O to a target value, as outlined in Appendix A. Although it may be useful to think of as "target" sparsity, it is only number. The rums usually converge to a value slightly below t, so it is prudent to set it to value greater than alatthough s can then never reach the target value, it will be pushed higher sparsities.
We have two sets of masks 2. The first set associates a 0 1 value ze with each edge = (m1, n2) in the computational graph. The second tags each node of the graph with 0 1 value 2n. The latter specifies whether a node is "active", i.e., producing output. In effect, the presence of an edge (n1, n2) is determined by the binary mask
( = m()). X 2n
We initially only used edge masks but found that the method would have difficulty converging to high sparsities (i.e., end up at low sparsities). Introducing a second set of masks allows the process to eliminate many edges quickly, accelerating the removal of unimportant components. However, the lagrangian above only applies to the edge masks. This is fine since the node masks can only remove further edges, not introduce ones top of those chosen by the edge masks. The final loss is
L LKL + Ledge,
Experiments
4.1 Experimental Setup
Methods. We compare Edge Pruning with a KL loss to ACDC and EAP in our experiments. Both outlined in Section 2. We do not compare to other pruning-based methods, as Commy et al. [2023] found them to perform much worse than ACDC. We list the hyperparameters used in Appendix A. The experiments in this section are all performed on GPT-2 Small (117M).
Tasks. Prior works evaluate their methods on the same examples used to find circuits. In a departure from this convention, we separate each dataset into train, validas dationn,, and test splits, to avoid artifacts caused by overfitting. We use the following tasks.
Indirect Object Identification (IOI-tI and ION) [Wang al., 2023] is task with instances of the format "Friends Juana and Kristi found a mango at the bar. Kristi gave it to -> Juana". Conmy al. [2023]
version with single template, which we refer as 101-t1-thh version has 50 examples in each split. We also compare the methods on variant (IOI) with 30 templates found HuggingFace?. We randomly select 200 examples each for the train and val idation splits, and 36, 084 examples for the test split.
co cteeeeecocccaaaaeeeeeeeeeeeeeeee datasets / www.aaaaaaaaaaaaaaaaaaaaaaaaaaaa /ioi/: / an example template is "Then, and A had
long argument. Afterwards said A".
5

---

ACDC
EAP
Edge Pruning
KL (Model Circult
93
90
96
96
97
58
99
93
94
95
97
98
99
Edge Sparsity (96)
Edge Spansity (56
(a) 10I-11 (IOI, I template)
(b) IOI (Indirect Object Identification)
K0. (Model Crrcutd
0.8
0.4
0.15
0.4
0.2
0.2
0.1
98.5
494
97.5
99.5
96.5
96.0
98.5
97.0
97.5
98.0
98.5
99.0
Edge Sparsity (1)
Edge Sparsity (N)
(c) GT (Greater Than)
(d) GP (Gendered Pronoun)
Figure 2: The faithfulness of the methods, given the KL divergence between the model and obtained circuits (lower is better). On IOI-tI and GP, Edge Pruning is competitive low sparsities and better at high sparsities. It outperforms both ACDC and EAP by significant margin on IOI and GT.
Greater Than (GT) [Hanna et al., 2023) consists of examples of the format "The war lasted from the year 1743 17 zy"" The objective of the task is place a greater probability on the continuations 44, 45,..., 999 than 00, ...... 42. Our dataset spans 5 templates, 120 choices for nouns, and the years 1100 through 2199. It has 150 examples in the train and validation splits, and 12, 240 examples in the test split.
Gendered Pronoun (GP) [Athwin et al., 2023] consists of statements of the form "So Evan is really great friend, isn't he". We use the templates from the original Colab notebook used by Athwin et al. [2023], but generate more examples as they only work with 5. We use the top 1, 000 most popular baby names for boys and girls each in the year 2000 to generate dataset with 150 train and validation examples each, and 378 test examples.
Tracr [Lindner et al., 2023] compiles programs written in the RASP [Weiss et al., 2021] programming language into few-layer Transformers. We evaluate Edge Pruning how well it recovers ground-truth circuits for two Tracr pngraass-xprooorte x's in the prefix) and reverse (reversing list). Both tasks were discussed in Weiss et al. [2021] and used by Conmy et al. [2023] in their evaluation.
(proportion of
Evaluation. A circuit faithful to model behavior on a task if we can corrupt all model edges outside the circuit while retaining the model's outputs [Hanna et al., 2024]. We corrupt non-cireuit edges with interchange ablation and evaluate the methods' faithfulness as the KL divergence between model and circuit outputs. Specifically, we corrupt an example by swapping the placeholder value in the same template with random example from the dataset. We appraise the circuits' performance on 10I-t1, IOI, and GP via the Logit Difference log P(correct) - log P(misleading) between the correct and misleading name/pronoun. For GT, we evaluate the Probability Difference Plyy + 99) P(00 yy 1) between the correct and incorrect ranges. All metrics on GT work with predictions restricted to the set (00, 01,..., ... 99). We always take unrestricted predictions over the entire model vocabulary on other tasks. All non-Traer experiments use GPT-2 Small model. Appendix B evaluates additional metrics-includde including circuit overlap with manually found circuits.

---

ACDC
EAP
Edge Pruning
Logit Difference
Edge Sparsity (%)
Edge Sparsity (%)
101-11 (IOI, template)
(b) IOI (Indirect Object Identification)
Probabiity Difference
3.0
0.6
Dificrence
0.5
2.6
2.4
2.2
96.0
96.5
97.0
97.5
98.0
98.5
99.0
99.5
95.5
96.0
96.5
97.0
97.5
98.0
98.5
99.0
Edge Sparsity (%6
Edge Sparsity (%6
(c) GT (Greater Than)
(d) GP (Gendered Pronoun)
Figure 3: Comparison of circuit performance between methods. We report the Logit Difference log P(correct) - log P P(misleading) for IOI-tI, IOI and GP, and the probability difference P + 99) - P(0 yy 1) for GT. Higher better for all plots. Edge Pruning finds perter-performing circuits all four tasks. The dashed line indicates the performance of the full model.
Table 1: Scaling to larger IOI dataset: ACDC improves with more examples but its runtime scales prohibitively. EAP is fast but cannot perform as well. Edge Pruning seales effectively to 100K examples, where it is the fastest and most faithful method. All runs use one NVIDIA H100 GPU.
200 examples
400 examples
100K examples
Method
Sparsity (%)
KL +
Time (s) |
KL L
Time (s)
KL 4
Time (s) |
ACDC EAP
96.6 + 0.1
0.92
18,783
0.88
40,759
96.6 1 0.1
3.47
21
3.66
43
3.78
12,260
Edge Pruning
96.6 1 0.1
0.25
2,756
0.22
2,931
0.20
3,042
4.2 Results
This section compares the three methods on our primary faithfulness and performance metrics. We report additional metrics in Appendix B, and Appendix F shows some circuits found by Edge Pruning.
Edge Pruning outperforms prior methods on more complex tasks. Edge Pruning is competitive on IOI-t1 and GP in terms of faithfulness at low sparsities, and slightly better at higher sparsities (Figure 2). is considerably more faithful on IOI and GT than both ACDC and EAP, especially at higher sparsities. In particular, ACDC does worse than randomly choosing between the two names (KL divergence 0.69) at high sparsities on IOI, whereas Edge Pruning remains better. We hypothesize that the relative simplicity of 10I-t1 and GP-one template small output space (he/she) renders local (ACDC) or first-order (EAP) approximations good proxies, potentially explaining the edge of Edge Pruning IOI and GT. A similar trend is seen in performance (Figure 3): Edge Pruning finds better-performing circuits on all four tasks. Specifically, on IOI, Edge Pruning finds circuit of 98.8% sparsity that is as faithful and performs well as the one found by ACDC at 96.8% sparsity- over 2.65x fewer edges. Interestingly, EAP scales better to higher sparsities than ACDC on GT, delivering respectable performance even 99.5% sparsity.
Edge Pruning can scale to 100K examples. We investigate how the methods scale to more examples at representative sparsities. this end, we create large version the IOI dataset's train
httapp////gggghooooooo/////////a baa--baabaeeeee names/

---

Positin
Token
Heal 3.VVV
gribedding
Boadoo
emsbedcings
Ponitur crbholdingo
Head D
[lead 1.0.Q
Ilead L.O.
MLPO
Ph
Head
MILP!
Ostput
Oxtput
(a) xproportion
(b) reverse
Figure 4: The canonical ground-truth circuits for the Tracr-compiled xproportion and reverse programs. Edge Pruning recovers both circuits perfectly.
split with 100K examples. We hold the number of gradient descent steps for Edge Pruning fixed (Appendix A). Although its runtime would scale linearly with more epochs, at 100K examples all approaches see almost all examples once.* Thus, the time reported in Table represents the relative overhead of each method. ACDC shows clear improvements with more examples, but cannot scale well due to prohibitive runtime. EAP, on the other hand, fast even with more examples. However, it underperforms the other methods significantly. Edge Pruning efficiently uses more examples and demonstrates both the least runtime and the highest faithfulness by far with 100k examples. We therefore conclude that Edge Pruning good fit for complex or mixture distributions where more examples may be needed to specify model behavior.
Edge Pruning finds ground-truth circuits in Traer programs. To check if Edge Pruning can find the ground-trute truth circuits, we use Tracr [Lindner et al., 2023] to compile two example programs xproportion and reverse as Transformers. The former yields 2-layer Transformer that outputs, at each position, the fraction of x's seen so far. The latter yields 3-layer Transformer that can reverse lists. We use zero ablation following Conmy et [2023] (more details in Appendix A). Edge Pruning achieves perfect reconstruction of both circuits (Figure 4).
Edge Pruning is robust to variance in random initialization Appendix finds that both the resulting sparsity and the faithfulness of the circuits found by Edge Pruming are remarkably consistent across different random initializations of masks. We also investigate there the question of whether multiple different circuits exist for given task, and if Edge Pruning find them.
5 Case Study: Scaling to 13B Parameters
We have seen that Edge Pruning can scale efficiently with more examples. We next investigate if it can scale with model size. This is increasingly important, given the recent interest in interpreting multi-billion parameter models [Lieberum et al., 2023, Prakash et al., 2024]. Current methods used to interpret such models, while undeniably indispensable, have limitations: path patching [Goldowsky- Dill et al., 2023] identifies important subsets of components but falls short of producing edge-level circuits. Distributed Aligmment Search [Geiger et al., 2024, Wu et al., 2023] can verify proposed symbolic execution graphs and align them with the model but requires prior knowledge of the correct symbolic graph, which is nontrivial to obtain.
On the other hand, pruning can scale to large models using model parallelism [Xia et al., 2024]. We thus apply Edge Pruning to a case study on CodeLlama- 100x larger than GPT-2-witde with real task. We are inspired by Prakash al. [2024], who compare base and fine-tuned LMs and find that finetuning enhances existing mechanisms. Instead of comparing base and fine-tuned models, compare mechanisms in the same model with different prompting schemes. Specifically, we ask whether the same mechanisms underlie (zoro-shot) instruction prompted
[Rozière al., 2024] a model over
"Wit our hyperparameters, Edge Pruning 96k unique examples (can be higher with more GD steps).

---

Table 2: Edge pruning finds circuits with 0.03-0.04% of the edges in CodeLlama-133 that match the performance of the full model. The circuits perform well in cross-evaluation and overlap highly, hinting that the same mechanisms explain large parts of instruction-proomee prompted and few-shot behavior.
Accuracy (%) f
Exact Match (%) 1
Circuit
Num. edges
Instr. prompted
Fewshot
Instr. prompted
Fewshot
Full model
3872820
82.00
89.25
100.00
100.00
Instruction prompt (IP)
1041
79.25
74.50
90.00
79.00
Fewshot (FS)
1464
75.75
87.25
84.50
91.25
IP FS
653
72.50
68.25
79.75
72.50
and few-shot behavior for the task-model pair we study. This case study serves a dual purpose. It demonstrates the scalability of Edge Pruning as a method. It also illustrates how circuit-fnding methods may fit into the interpretability arsenal. We are interested in three research questions: (RQ1)
Can Edge Pruning find edge-sparse circuits in a 13B model? (RQ2) To what extent do the circuits for instruction and few-shot prompting share the same edges? (RQ3) Does the instruction-prompted circuit perform well when used in few-shot manner, and vice versa?
Task and model setup. We work with the task Boolean Expressions from the BBH [Suzgun et al., 2022] benchmark suite. This task consists of instances of the form "((not False) and False) or (False and True) is - False". The original dataset only has 250 examples, so we programmatically generate an in-house version of the task. Our dataset has 3840, 767, and 3070 examples in the train, validation, and test splits respectively. Each instance has between 3 and 6 literals, with a maximum nesting depth of 3 and at most 2 consecutive nots. We use demonstrations for the few-shot setting. The prompts used for the instruction instruction-prompped and few-shot settings are provided in Appendix E. Our model is the instruction-finetunee version of CodeLlama-133P It achieves accuracies of 82% and 89. 99.25% in the instruction-prompted (IP) and few-shot (FS) settings, respectively.
Expresm
(RQ1) Edge Pruning produces extremely sparse circuits. We next apply Edge Pruning to the described settings. We isolate one circuit when instruction prompting and one with the few-shot prompt (hyperparameters in Appendix A, which also highlights other optimizations like distributed training and gradient checkpointing). The circuit discovered in the IP setting has 1, 041 edges, corresponding to a 99.97% edge sparsity. That discovered in the FS setting has 1, 464 edges, equivalent to 99.96% edge sparsity. The discovered circuits are evaluated in Table 2. Despite using less than 0.04% of the edges, the circuits closely match the performance of the full model the few-shot circuit achieves accuracy of 87 37.25% and performs within 2% of the full model (when prompted few-shot). The instruction-promptae circuit is accurate within 2.75% of the full model.
(RQ2) The circuits have a high overlap, and their intersection performs well. We appraise the intersection of the IP and FS circuits next. The two circuits share 653 edges, accounting for 62.7% of the edges of the sparser (instruction prompted) circuit- corresponds to an intersection over 1, 700x larger than expected by random chance. We further evaluate the cireuit formed by this intersection in the instruction prompted and few-shot settings (Table 2). It performs well in the instruction prompted setting, and worse than the model (but still significantly above chance) when prompted few-shot.
(RQ3) The circuits demonstrate strong performance in cross-evaluation. We note from Table that the circuit found with few-shot prompting shows strong performance even when instruction prompted. Analogously, the instructiin-promplle circuit also performs well in the fewshot setting.
Our case study suggests that the same mechanism (as represented by the intersection above) explains a large part of the performance both settings-i.e. they do not proceed via disjoint mechanisms. However, the performance gap between the FS and IP n FS circuits is still sizable. Further, we see modest drops in cross-evaluation eegti.nnnee......... from 87.25% when evaluating the FS circuit few-shot to 75.75% in the instruction prompted setting. This suggests that additional components are needed to
'nttps:/hhuginggface co//codellam//ooeeeeeeelll 13b- Instruct-hf

---

complete the picture. A complete mechanistic description of the components in the two circuits is an exciting avenue for future work, but beyond the scope of this case study.
Manual analysis of the CodeLLlmm--13h circuit. Interpreting a circuit in such large model even if very sparse remains a challenging task. We isolate small region of the circuit and identify curious behavior in in Appendix F, leading an intriguing conjecture. Nonetheless, we believe that thorough study requires more analysis, which is beyond the scope of this paper (but makes for exciting future work).
6 Related Work
Circuits. By reducing large model to sparse subgraph, circuits help interpret internal model computations [Olah et al., 2020, Elhage et al., 2021], and several visualization tools have been developed to aid this process [Sakarvadia et al., 2023, Katz and Belinkov, 2023, Tufanov et al., 2024]. Circuits were originally found manually [Hanna et al., 2023, Athwin et al., 2023], but this has recently been automated by tools like ACDC [Conmy al., 2023]. ACDC uses activation patching [Vig et al., 2020] to knock out unimportant edges. Other approaches instead estimate the importance of each edge via attribution scores [Nanda, 2022]; this approach was used by EAP [Syed et al., 2023]. Ferrando and Voita [2024] use attribution patching to identify domain-specifer model components in Llama-2-7B. Kramár et al. [2024] note that attribution patching may lead to incorrect approximations, and propose a variant with reduced error. In concurrent work, Hanna et al. [2024] argue that faithfulness metrics are better for evaluating circuits than measuring overlap with manual circuits. Recent work has explored other notions of a circuit. Inspired by the fact that Sparse Autoencoders (SAEs) can find cumman-ntterpretbbl features in LM activations [Cunningham et al.. 2023], Marks et al. [2024] find circuits over these features. Wu et al. [2023] align computation in Alpaca [TTorri et al., 2023] with a proposed symbolic algorithm [Geiger al., 2024]. Our method orthogonal to these developments.
Pruning. Pruning [LeCun et al., 1989] drops parameters or layers of language model for space efficiency and potential speedups. Structured pruning [Wang et al., 2020, Xia et al., 2022] imposes some regularity on the resulting subnetworks, such as an equal fraction of preserved parameters in GPU hardware at the cost of lower
each layer. Doing allows it achieve substantial speedups compression. In contrast, unstructured pruning [LeCun et al., 1989, Hassibi and Stork, 1992] does not impose such constraints. Channel pruning [He et al., 2017] is form of structured pruning that prunes input channels in vision models, which has been adapted for neural architecture search [e.g. Li al., 2022]. Pruning has occasionally been used as part of an interpretability effort, but mostly the level of neurons [Michel et al., 2019, Jain et al., 2023], or less commonly, attention heads/MLPs [Cao al., 2021]. Our work finds circuits by pruning the edges between components instead.
7 Conclusions
In this paper, we introduce Edge Pruning to find circuits by pruning edges between components. We find that it discovers sparse, faithful circuits, and we demonstrate its scalability to large datasets and large models. We close by discussing its limitations, and how future work may address them.
Limitations. We acknowledge that with small datasets, ipproximation-based approaches like EAP are faster than Edge Pruning. Circuit discovery with Edge Pruning may also require more GPU memory than these (Appendix A). Future work may precede Edge Pruning with a fast, approximate method like EAP to balance efficiency and performance. We note that even at very high sparsities, circuits for large models can still have hundreds of edges, and their full interpretation remains challenging. Further automating interpretability [Bills et al., 2023] is compelling avenue for future research. Finally, we note that even with perfect faithfulness to the model outputs, a circuit can misrepresent the necessary computations in the full model, thus leading interpretability illusion [Makelov et al., 2024]. Better metrics are needed to reveal these possibilities in practice.
eeciiida-sspeciallly at scale-where we use 32 H100 GPUs for Codee_lama-133
Societal and ethical impact. Our work aims to facilitate the process of understanding and explaining large foundation models, which is crucial for their continued safe development and deployment. We do not foresee Edge Pruning being used towards adverse socictal or ethical ends.
10

---

Acknowledgements
We are thankful to Tianyu Gao, Zirui Wang, and Mengzhou Xia for their helpful discussions regarding the experiments. Input by Abhishek Panigrahi and Carlos Jiminez was also instrumental to this project. We also thank Howard Chen and Tianyu Gao for their help in proofreading and improving the writing in this paper. AB gratefully acknowledges the support of Hisashi and Masae Kobayashi #67 Fellowship. This research also funded by the National Science Foundation (SS-2211779) and Sloan Research Fellowship.
References
Chris Athwin, Guillaume Corlouer, Esben Kran, Fazl Barez, and Neel Nanda. Identifying preliminary circuit for predicting gendered pronouns in GPT-2 small, 2023. URL https: //cmathw. itch. io/ii///ei mo/identifying sared-pronouna in-inn-nn iiw--iiinnng---e smal/.
greliminaay-cirouute wor--rdiictp--gg------- gen
Steven Bills, Nick Cammarata, Dan Mossing, Henk Tillman, Leo Gao, Gabriel Goh, Ilya Sutskever, Jan Leike, Jeff Wu, and William Saunders. Language models can explain neurons language models. https: / openaappbliic. blob. core. windows net net/neuron wxplainer/papeeeee.ee ndex. html, 2023.
Tom Brown, Benjamin Mann, Nick Ryder, Mclanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert- Voss, Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel Ziegler, Jeffrey Wu, Clemens Winter, Chris Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, Ilya Sutskever, and Dario Amodei. Language models few-shot learners. In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin, editors, Advances in Neural Information Processing Systems, volume 33, pages 1877-1901. Curran Associates, Inc., 2020. URL https::/ppooceeddiiiiite cttpscccpppp....ce fi alleeeeeeee www.fffllsseeaaeeeeeeee er.
Steven Cao, Victor Sanh, and Alexander Rush. Low-complexity probing via finding subnetworks. In Kristina Toutanova, Anna Rumshisky, Luke Zettlemoyer, Dilek Hakkani-Tur, Iz Beltagy, Steven Bethard, Ryan Cotterell, Tanmoy Chakraborty, and Yichao Zhou, editors, Proceedings of the 2021 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, pages 960-966, Online, June 2021. Association for Computational Linguistics. doi: 10.18653/v 1/202/1122221mml2. .nwccllmaaee mameclmain.774 74. URL https http///aaaaaaahhoeeo naacl-main 74.
org/20211
Arthur Conmy, Augustine N. Mavor-Parke Aengus Lynch, Stefan Heimersheim, and Adrià Garriga- Alonso. Towards automated circuit discovery for mechanistic interpretability. In Thirty-seventh Conference on Neural Information Processing Systems, 2023. URL. https: htt.htttss//ppereiiieee net /forummidd-89a177222
Hoagy Cunningham, Aidan Ewart, Logan Riggs, Robert Huben, and Lee Sharkey, Sparse autoen- coders find highly interpretable features in language models, 2023.
Nelson Elhage, Neel Nanda. Catherine Olsson, Tom Henighan, Nicholas Joseph, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, Tom Conerly, Nova DasSarma, Dawn Drain, Deep Ganguli, Zac Hatfiel-Dodds, Danny Hernandez, Andy Jones, Jackson Kernion, Liane Lovitt, Kamal Ndousse, Dario Amodei, Tom Brown, Jack Clark, Jared Kaplan, Sam McCandlish, and Chris Olah. A mathematical framework for Transformer circuits. Transformer Circuits Thread, 2021. https://traasfomee-ccccuiis
Javier Ferrando and Elena Voita. Information flow routes: Automatically interpreting language
models at scale, 2024.
Dan Friedman, Alexander Wettig, and Danqi Chen. Learning Transformer programs. In Thirty- seventh Conference on Neural Information Processing Systems, 2023. URL hwthhhtt/////// :/tsppe//

---

Atticus Geiger, Kyle Richardson, and Christopher Potts. Neural natural language inference models partially embed theories of lexical entailment and negation. In Afra Alishahi, Yonatan Belinkov. Grzegorz Chrupata, Dieuwke Hupkes. Yuval Pinter, and Hassan Sajjad, editors, Proceedings of the Third BlackboxXI.P Workshop on Analyzing and Interpreting Neural Networks for NLP, pages 163-173, Online, November 2020. Association for Computational L Linguistics. doi: 10.18653/v1 2020.blackboxull-1.16 URL https: h/tttt/oola .ow.ooyooooooooooooooo.... org/2020. 2020. blackboxnlp-1.118 1.
Atticus Geiger, Zhengxuan Wu, Christopher Potts, Thomas Ieard, and Noah Goodman. Finding alignments between interpretable causal variables and distributed neural representations. In Francesco Locatello and Vanessa Didelez, editors, Proceedings of the Third Conference on Causal Learning and Reasoning, volume 236 of Proceedings of Machine Learning Research, pages 160- 187. PMLR, 01-03 Apr 2024. URL https: /ttp//eedooce...e///e////e waapeeddiinn.m........mmmmmm mlr, wax.ppress/vee /236////23/////// /e03////eeer4..... tml.
Nicholas Goldowsky-Din Chris Macl Lucas Sato, and Aryaman Arora. Localizing model behavior with path patching, 2023.
Michael Hanna, Ollie Liu, and Alexandre Variengien. How does GPT-2 compute greater-thane than?: Interpreting mathematical abilities in pre-trained language model. In Thirty-sevnnth Conference on Neural Information Processing Systems, 2023. URL https: /ttp://///e/.////eee//e///e//eea wttpeiiiiiiieeernunne d=p4PckNQR8k.
Michael Hanna, Sandro Pezzelle, and Yonatan Belinkov. Have faith in faithfulness: Going beyond circuit overlap when finding model mechanisms, 2024.
Babak Hassibi and David Stork. Second order derivatives for network pruning: Optimal brain surgeon. In S. Hanson, J. Cowan, and C. Giles, editors, Advances in Newral Information Processing Systems (NeurIPS), volume 5. Morgan- Kaufmann, 1992. URL httpppps//poooeeeengggsssssss neurips cc/p aper files wwwesf/epeeerree 1992//999//// file/ http://ppppppp.....///////a/a/aa 300044668440302094448887000--000
pdf.
Yihui He, Xiangyu Zhang, and Jian Sun. Channel pruning for accelerating very deep neural networks.
In Proceedings of the IEEE International Conference on Computer Vision (ICCV), 2017.
Samyak Jain, Robert Kirk, Ekdeep Singh Lubana, Robert P. Dick, Hidenori Tanaka, Edward Grefen- stette, Tim Rocktäschel, and David Scott Krueger. Mechanistically analyzing the effects of fine-tuning on procedurally defined tasks, 2023.
Shahar Katz and Yonatan Belinkov. VISIT: Visualizing and interpreting the semantic information flow of Transformers, 2023.
Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In Yoshua Bengio and Yann LeCun, editors, ICLR (Poster), 2015. URL. //bblp. www.ttrrrrrreeeeee conf iclr/iclr2015 html #K #KingmaB14.
János Kramár, Tom Lieberum, Rohin Shah, and Neel Nanda. AtP*: An efficient and scalable method for localizing LLM behaviour to components, 2024.
Yann LeCun, John Denker, and Sara Solla. Optimal brain damage. In D. Touretzky, editor, Advances in Neural Information Processing Systems (NeurIPS), volume 2. Morgan-Kaufmman, 1989, URL kttpp::/ sttpp::/pproceeddings neurips. wwtpppprrpppeeeeeeee http:/pperfeeee//////pppaaeeeeee http:///aaaappee//////////////// 0933425044881227668-
Yanyu Li, Pu Zhao, Geng Yuan, Xue Lin, Yanzhi Wang, and Xin Chen. Pruning-as- Pruning-as-Searhh Search: Efficient neural architecture search via channel pruning and structural reparameterization. In Proceedings of the Thirty-Fire International Joint Conference on Artificial Intelligence, IJCAI-22, 2022.
Tom Lieberum, Matthew Rahtz, János Kramár, Neel Nanda, Geoffrey Irving, Rohin Shah, and Vladimir Mikulik. Does circuit analysis interpretability scale? Evidence from multiple choice capabilities in Chinchilla, 2023.
David Lindner, Janos Kramar, Sebastian Farquhar, Matthew Rahtz, Thomas McGrath, and Vladimir Mikulik. Tracr: Compiled Transformers as a laboratory for interpretability. In Thirty-seventh Conference on Neural Information Processing Systems, 2023. URI. https://///// httpp//ppp/eeee..///// net /f serum?id=bbbdBuu70
12

---

taattsuulab/sttaaoorrre Christos Louizos, Max Welling, and Diederik P. Kingma. Learning sparse neural networks through 1 O regularization. International Conference on Learning Representations, 2018. URL https: //openreview.o net/forumm Hite-idHH-HHihhhh
Alcksandar Makelov, Georg Lange, Atticus Geiger, and Neel Nanda. Is this the subspace you are looking for? An interpretability illusion for subspace activation patching. In The Twelfth International Conference on Learning Representations, 2024. URL atttps:/oppenreviet net /forumPiddEDttJJgMit.
Samuel Marks, Can Rager, Eric Michaud, Yonatan Belinkov, David Bau, and Aaron Mueller. Sparse feature circuits: Discovering and editing interpretable causal graphs in language models, 2024.
Thomas McGrath, Matthew Rahtz, Janos Kramar, Vladimir Mikulik, and Shane Legg. The hydra effect: Emergent self-repair in language model computations, 2023.
Paul Michel, Omer Levy, and Graham Neubig. Are sixteen heads really better than one? In H. Wallach, H. Larochelle, A. Beygelzimer. F. d'Alché-Buc, E. Fox, and Garnett, editors, Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019. URL https: httt:::://prcceeeinggg. peurips. wwt..es.cccpppppe.e.ee..... 2f 2 Paperr.aar
Neel Nanda. Attribution patching: Activation patching at industrial scale, 2022. URL https: //www.
Chris Olah, Nick Cammarata, Ludwig Schubert, Gabriel Goh, Michael Petrov, and Shan Carter. Zoom in: An introduction to circuits. Distill, 2020. doi: 0223991/diitiill00020
URL
https://distill, pub/ /tppp////////////////////// ht000///rrcuuus//ooooooooo
Nikhil Prakash, Tamar Rott Shaham, Tal Haklay, Yonatan Belinkov, and David Bau. Fine-tuning enhances existing mechanisms: A case study on entity tracking. 2024.
Alec Radford, Jeff Wu, Rewon Child, David Luan, Dario Amodei, and Ilya Sutskever. Language models unsupervised multitask learners. arXiv, 2019.
Baptiste Rozière, Jonas Gehring, Fabian Gloeckle, Sten Sootla, Itai Gat, Xiaoqing Ellen Tan, Yossi Adi, Jingyu Liu, Romain Sauvestre, Tal Remez, Jérémy Rapin, Artyom Kozhevnikov, Ivan Evtimov, Joanna Bitton, Manish Bhatt, Cristian Canton Ferrer, Aaron Grattafiori, Wenhan Xiong, Alexandre Défosse7, Jade Copet, Faisal Azhar, Hugo Touvron, Louis Martin, Nicolas Usunier, Thomas Scialom, and Gabriel Synnaeve. Code Llama: Open foundation models for code, 2024.
Mansi Sakarvadia, Arham Khan, Aswathy Ajith, Daniel Grzenda, Nathaniel Hudson, André Bauer, Kyle Chard, and Ian Foster. Attention Lens: A tool for mechanistically interpreting the attention head information retrieval mechanism. In NeurIPS Workshop Attributing Model Behavior at Scale, 2023. URL
Mirac Suzgun, Nathan Scales, Nathanacl Schärli, Sebastian Gehrmann, Yi Tay, Hyung Won Chung, Aakanksha Chowdhery, Quoc V. Le, Ed H. Chi, Denny Zhou, and Jason Wei. Challenging BIG-bench tasks and whether thhougt can solve them, 2022.
Aaquib Syed, Can Rager, and Arthur Conmy. Attribution patching outperforms automated circuit discovery. In NeurIPS Workshop on Attributing Model Behavior at Scale, 2023. URL https: //openreiiew
Rohan Taori, Ishaan Gulrajani, Tianyi Zhang, Yann Dubois, Xuechen Li, Carlos Guestrin, Percy Liang, and Tatsunori B. Hashimoto. Stanford alpaca: An instraction-following llama model. 2023.
https://github.
Igor Tufanov, Karen Hambardzumyan, Javier Ferrando, and Elena Voita. LM transparency tool: Interactive tool for analyzing Transformer language models, 2024.
13

---

Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, L ukasz Kaiser, and Illia Polosukhin. Attention is all you need. In I. Guyon, U. Von Luxburg. S. Bengio. H. Wallach, R. Fergus, S. Vishwanathan, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 30. Curran Associates, Inc., 2017. URL, http:t::/pprccccddaaggg/////// ne urips. er. er.pdf.
Jesse Vig, Sebastian Gehrmann, Yonatan Belinkov, Sharon Qian, Daniel Nevo, Yaron Singer, and Stuart Shicber. Investigating gender bias in language models using causal mediation analysis. In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin, editors, Advances in Neural Information Processing Systems, volume 33, pages 12388-12401. Curran Associates, Inc., 2020. URL htt.t:tpp///pptt ptpptpeee///ooceeeinnggeee neur wacc/papr__aee files/ /ttp://///////////////22//000000 2217116fe33222ff7770662-----1 pdf.
Kevin Ro Wang. Alexandre Variengien, Arthur Conmy, Buck Shlegeris, and Jacob Steinhardt. Interpretability in the wild: A circuit for indirect object identification in GPT-2 small. In The Eleventh International Conference Learning Representations, 2023. URL https: ew / NdNNssVSN6......
Ziheng Wang, Jeremy Wohlwend, and Tao Lei. Structured pruning of large language models. In Bonnie Webber, Trevor Cohn, Yulan He, and Yang Liu, editors, Proceedings of the 2020 Conference Empirical Methods in Natural Language Processing (EMNLP), pages 6151-6162. Online, November 2020. Association for Computational Linguistics. doi: 10.1 18653/v122020.com org. 22200. emnlp-mail mao-mmain... 496.
1p-main.496. URL pttttp:://acaaathoolooo
Gail Weiss, Yoav Goldberg, and Eran Yahav. Thinking like Transformers. In Marina Meila and Tong Zhang, editors, Proceedings of the 38th International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pages 11080-11090. PMLR, 18-24 Jul 2021. URL htwph::ttpppeee htttp://pproceeding.. mml. mlr. press 1133///es2221 html.
Zhengxuan Wu, Atticus Geiger, Christopher Potts, and Noah Goodman. Interpretability at scale: Identifying causal mechanisms alpaca. arXiv, 2023.
Mengzhou Xia, Zexuan Zhong, and Danqi Chen. Structured pruning learns compact and accurate models. In Smaranda Muresan, Preslav Nakov, and Aline Villavicencio, editors, Proceedings of the 60th Annual Meeting of the Association for Computational Linguistics (Volume I: Long Papers), pages 1513-1528, Dublin, Ireland, May 2022. Association for Computational Linguisties. doi: URL https: https////acaaaahooa/oa.o org/ 2222/222..c2 acl-long. 107.
Mengzhou Xia, Tianyu Gao, Zhiyuan Zeng, and Danqi Chen. Sheared LLaMA: Accelerating language model pre-training via structured pruning. In The Twelfth International Conference Learning Representations, 2024. URL https: /tttepp:://pppeeeeeeee /ttp///eeeeeeeeeeeeeeeee/eeeeeee net wtte//oormmmiiii wttrummmddd0000dda0000ee
Fred Zhang and Neel Nanda. Towards best practices of activation patching in language models: Metrics and methods. In The Twelfth International Conference on Learning Representations, 2024.
URL https::t//// httptttp:////pprveeveeee wtw.i/ewe...eeeeeeeeeeeeeeeeee Mett/Eooumm?dd-M117yy5
Yanli Zhao, Andrew Gu, Rohan Varma, Liang Luo, Chien-Chin Huang, Min Xu, Less Wright, Hamid Shojanazeri, Myle Ott, Sam Shleifer, Alban Desmaison, Can Balioglu, Pritam Damania, Bernard Nguyen, Geeta Chauhan, Yuchen Hao, Ajit Mathews, and Shen Li. Pytorch FSDP: Experiences scaling fully sharded data parallel, 2023.
14

---

ACDC
EAP
Edge Pruning
100
Exacct Match (863
Exact Match (%)
60
40
20
Edge Sparsity (16)
Edge Sparsity
(a) 101-11 (IOI, template)
(b) IOI (Indirect Object Identification)
1.0
100
vw
Kendalls 1Bu
Exact Match (%6
80
0.6
0.4
0.2
20
0.0
96.0
97.5
98.0
98.5
96.5
97.0
97.5
98.0
98.5
Edge Sparsity
Edge Sparsity (%6
(c) GT (Greater Than)
(d) GP (Gendered Pronoun)
Figure 5: Our secondary metric for measuring faithfulness is the Exact Match percentage between the model and circuit predictions on IOI-t1, IOI, and GP. On GT, we use the Kendall's 's Tau score between the model and circuit rankings of 00, 01.... 99 as the secondary metric. Edge Pruning is the most faithful method on all four tasks, with the difference being especially pronounced for IÓI.
A Hyperparameters and Computational Details for Edge Pruning
In this appendix, we list the hyperparameters used for the various experiments in the main text of the paper. All of our runs use the Adam [Kingma and Ba, 2015] optimizer with e = 10 and (31, (2)) = (0.9, 0.999).
GPT-2 experiments For all tasks, used sequence length of 64 tokens with padding. A batch size of 32 was adopted, and the learning rate for both the edge and node masks, as well as for the lagrangians À for both, was set 0.8. 1OI-t1 was an exception: here, we set all the above learning rates to 1 for all runs. The total number of optimization steps was 3000, and the target edge and node sparsities were linearly increased starting from 0 over the first 2500 steps. Evaluation and checkpointing were performed every 64 steps but we always used the final checkpoint to report results. To produce scatterplots, we varied the edge target up 1.1 but held the node target largely fixed for each task. These values were 0.72 for IOI-t1 and IOI, 0.68 for GT and 0.69 for GP. These values were chosen based on small number of pilot runs, and we expect that grid search can improve results further.
We also wish make several remarks about implementation. We turned off dropout for all since it made the optimization noisy. Our threshold for the final rounding is not a pre-determined value. Instead, we compute the average value of all entries of all masks, and brand that the desired sparsity. Then, we perform a binary search for a threshold such that the fraction of entries rounded 1 equals this desired sparsity. The thresholds found this way usually fell between 0.2 and 0.8. This also allows the user to achieve exactly the desired sparsity by setting a different threshold. We implement all of our code by implementing modified versions of the HuggingFace model classes, as it allows to use the HuggingFace Trainer and its optimizations out of the box. Our code also natively supports Flash Attention, though none of our results use it. Finally, we note that the role of A in the lagrangian term is to allow (and indeed, encourage). "shooting past" / when optimizing due to momentum, This prevents the model sparsities from "settling into" mode where they lag behind the targets by constant but non-zero amount throughout pruning.
15

---

Tracr experiments For both programs, we fix the aY values to O and only optimize X2 as described in Section 3. For the xproportion program, we use an edge target of 0.92 and a node target of 0.4. The edge and node mask learning rates were 1, and that for the lambdas was 0.0001. A total of 720 optimization steps were performed with batch size of 16, of which 640 was used for target warmup. The learning rates were warmed up linearly over the first 96 steps. A sequence length of 5 was used.
Initially, for reverse, setting the regularization learning rate was tricky-it was easy to end up not regularizing enough or overdoing it. Thankfully, an easy remedy was to increase the number of steps to 6000 (of which the first 5900 warmed up the edge and node targets, and the first 1500 warmed up the learning rates). This allowed us to set a relatively higher learning rate for the lambdas (0.001), along with an aggressive edge target of 1.02. The node target was set to 0.1. The learning rates of the log alphas and lambdas were 0.03 and 0.001, respectively. Despite using 6000 steps, the run took under minutes on one NVIDIA A100.
CodeL 13B experiments For our CodeLlama- 13B experiments, we use a learning rate of 0.8 for both the edge masks and the node masks. In a departure from the choice of Section 3, we also include separate lagrangian term over node masks:
Loodc. = \1 mode (faode - Spods) + A2.node - (frodo - Snode) )2
The reason for this choice was that, in our preliminary runs with small Sheared Llama [Xia et al., 2024], we found that this would achieve higher sparsities. We use a learning rate of 0.4 for all of the lambdas. The target edge and node sparsities steps with a batch size of 1. The first 200 steps linearly warm up the learning rate, while the target sparsities are linearly increased over the first 5500 steps. We enable gradient checkpointing, as well as FSDP [Zhao et al., 2023] with full sharding in BF16 precision. The maximum sequence lengths for the instruction-propptr and few-shot settings were 64 and 72, respectively.
set to 1.2 and 0.7, respectively. We use 6000
We also comment here on the computational resources used for the runs.
Computational details. The Tracr experiments use one NVIDIA A100 with 80 GB of memory. The GPT-2 experiments use either one NVIDIA A100 or one H100 (both 80 GB) each. The experiments of Table 1 all use one NVIDIA H100 for a fair runtime comparison. Each CodeLlama-133 13B run utilizes 32 H100 GPUs and 600 gigabytes of CPU memory. The typical runtime of a GPT-2 pruning run was about 45 minutes, and that of a Tracr run was under 5 minutes. The Code runs each took around 35 hours. We estimate the total computational budget to be around 5000 GPU hours.
B More results
We show more results on faithfulness and performance metrics in this appendix. Specifically, we evaluate on one alternate faithfulness (agreement) metric and one additional performance metric. For the former, we choose Exact Match percentage as the agreement metric on 10I-t1, IOI and GP. For GT. we instead report the Kendall's ' Tau score over the rankings of 00, 01, 02...... ................................ .99 as induced by the output logits of the model and circuit, which is then averaged across examples. Figure plot these metries for the three approaches. We see that Edge Pruning is consistently the most faithful method on all four tasks, with the gap to the next-best method being large for I0I.
Our choice of the performance metric is Accuracy for 10I-t1, 1OI and GP. For GT, we instead compute a variant of Probability Difference called Probability Difference 10, given by P(yy + 1 : yy + 10)
P(yy- 10: yy 1). Note that the original probability difference. P(yy + 1 99) P(00 yy 1)
can be gamed by always predicting 99. The new variant overcomes this obstacle by measuring the sharpness of the cutoff. The results, shown in Figure 6, echo the results of the main text: edge pruning is competitive on GP, and outperforms the other methods in I0I-t1, IOI and GT.
We also compare Edge Pruning to ACDC in terms of circuit overlap with manually reverse-engineered circuits. Since the manual eircuits only identified important components and not the edges between them, we plot node (component) ROC curves in Figure 7, where we consider a node included in circuit if at least one edge incident to it is included. Note that the IOI manual circuit only studied attention heads, so we ignore MLP nodes in the corresponding circuits. The results show that Edge Pruning is competitive with ACDC on circuit overlap metrics. Nevertheless, we emphasize that manually reverse-engineered circuits are not guaranteed to be optimal since they also investigate one ablation at time without considering interactions between ablations. As such, we echo Hanna et al. [2024]'s suggestions of using circuit faithfulness metrics over circuit overlap.
16

---

ACDC
EAP
Edge Pruning
100
80
Accuracy (%%
Accuracy (9%)
60
40
20
Edge Sparsity
Edge Sparsity
101-11 (IOI, I template)
(b) IOI (Indirect Object Identification)
Probability Oifference (10)
0.35
(%6
80
60
0.30
40
0.25
20
96.0
96.5
97.0
97.5
98.0
98.5
99.0
99.5
05.5
96.0
96.5
97.0
07.5
98.0
98.5
90.0
Edge Sparsity
Edge Sparsity (%)
(c) GT (Greater Than)
(d) GP (Gendered Pronoun)
Figure Comparison of the various methods on our secondary performance metric-accuraa in the case of 1OI-t1, 10I and GP, and Probability Difference 10 for GT (given by P(yy yy + 10) Plyy 10 yy 1)). Once again, Edge Pruning is competitive on GP, and outperforms other methods on IOI-t1, 10I and GT. The dashed lines indicate full model performance.
(a)
(b) GT
Figure 7: ROC curves against manual circuits for Edge Pruning and ACDC. The AUC is slightly higher for Edge Pruning on IOI, and slightly lower on GT.
C Edge Faithfulness
In other sections and appendices, have taken up the output faithfulness of Edge Pruning, i.e., whether the output distribution of the circuits matches that of the model. Here, we consider another edge faithfulness- an edge important for the model should also be important for the circuit. Con- cretely, given a circuit C' of model M, we measure for each edge E C, m. XL(M)) | {e}) and c. KL(M]|C' \ {e)), i.c., how much removing the edge from the circuit or model affects its output distribution. For method to be faithful, we expect to see strong positive correlation between the two values, especially for edges where m. large. We plot the two values against each other on the four tasks for four representative circuits found by Edge Pruning in Figure 8. The figure also provides the sparsities of each circuit. On all four tasks, whenever an edge important to the model, it is also important to the circuit. Thus, studying the circuit to infer the role/importance of the components a good proxy for the full model. On the other hand, we note that some edges completely unimportant for the model, but ablating which perturbs the circuit KL by small
17

---

35
2.5
Clouu
Cioutt
1.5
1.5
1.0
0.5
9.0
1.0
2.0
2.5
3.5
2.0
Original
Original
(a) 1OI-t1 (IOI, template)
(b) IOI (Indirect Ohject Identification)
1.4
1.0
drout
Crout
0.8
9.6
0.4
0.2
200
0.2
0.2
0.4
0.6
0.8
1.0
1.2
1.4
1.2
Original
Original
(c) GT (Greater Than)
(d) GP (Gendered Pronoun)
Figure 8: The KL divergences of the model and circuit, upon ablating individual circuit edges from each, measured against the full model. We see that all components important to the model are also important to the circuits, with an almost linear correlation between the two quantities. The circuits shown here have sparsities of 97 23%, 96.44% 98.59% and 98. 787, respectively.
amount. This perturbation much smaller than the ones seen in the former case above, but still non-negligible. This is not surprising, strcuit-finding methods may miss backup components that are deemed unnecessary for performance, and therefore be more sensitive to edge ablations. Alternatively, models may display behavior such the Hydra effect [McGrath al., 2023], whereas a circuit may not. Nonetheless, we suggest that practitioners verify any insights obtained from circuits on the full models wherever possible, regardless of the method used.
D How consistent are the circuits found by Edge Pruning?
In this appendix, we evaluate Edge Pruning can consistently find (i) good circuits, and (ii) consistent circuits in terms of chosen edges across different random initializations. To this end, we choose representative target sparsities (97.5% for IOI, 99.0% for GT, and 97.0% for GP) and prune GPT-2 small model with 12 different random seeds with these targets (and other hyperparameters as in Appendix A). As Figures 9 and 10 show, the resulting sparsities and faithfulness of these eircuits are remarkably consistent across the 12 seeds, demonstrating that Edge Pruning is robust to different initializations. It is also interesting to ask whether multiple circuits exist for performing a task (and whether Edge Pruning finds them)-Figure 11 investigates this question in this same setting by plotting the distribution of all pairwise loUs pnteessectionovee--iieee across the (222 66 pairs of circuits. We observe that the IoU values are generally high (0.5-0.7), but still far from 1. This suggests that while some components may be vital, others might be redundant.
in terms of chosen edges,
18

---

23
820
53
Seintvih
Sparrin
TOI
(b) GT
(c) GP
Figure 9: The sparsities of obtained circuits are remarkably consistent across 12 seeds.
(a) 101
(b) GT
(c) GP
Figure 10: The KI. divergences of obtained circuits are consistent across 12 seeds.
This question can be further investigated in future work, especially in how should define circuits the face of redundancy.
19

---

IOI
(b) GT
(c) GP
Figure 11: The pairwise Intersection-over-Uiie over 12 seeds is usually high, but far from 1.
E
Prompt formats for Boolean Expressions
We show the prompts used for the instruction- instruction-prompted and few-shot settings in the CodeLlama-I1 1aama-13B case study in Figure 14.
F Circuits found with Edge Pruning
In this section, we show example circuit diagrams of the circuits found by Edge Pruning. However, these come with one caveat. Since the typical circuit we found still had too many edges to present in reasonably sized figure, we only provide figures here for GT and GP, where sparsities ove 99.5% still performed well. Despite this, the circuits here are among the sparsest ones we obtained for each task and therefore perform worse than those at lower sparsities (such as those reported in Figure 2).
The GT circuit is shown in Figure 15, which also reports the faithfulness and performance metrics for it. Similarly, Figure 16 shows a circuit for GP with 99. 99.79% sparsity found by Edge Pruning, Note that the latter, due to the extremely high sparsity, does not perform that well. Nonetheless, the denser circuits compared in prior plots are too unwieldy to show here.
Interpretation the CodeLlama-133B circuit. Interpreting circuits with 1000 edges remains difficult, but we have made progress in understanding parts of the circuit. For example. we have found the following sub-circuit of composed heads (refer Figure 17 for snippet of this region): L8. H16 attends from operations (and/or) to the previous token (i.e. from op in op b). L10. H24 attends from an operand to a previous operation (i.e. from to op in op b) and read the results from L8. H16. This suggests that this duo computes the value of the expression. Interestingly, the attention pattern also holds when a is not a literal like True but an arbitrarily nested subexpression-e..g attending from to ( in "( ((True hypothesis here that the model could deal with arbitrary depth expressions by guessing the value of a a-alowiin it to proceed with the second step-and and later verifying the guess. This would also allow the model to parallelize a sequential computation by doing both steps of expression resolution in parallel. Nonetheless, further study and careful interventions required to verify this hypothesis.
False) and True)
False", A
Figure 12: Intruction prompt 
 <SYS>
Few-shot prompt 
 (True and False) or (False and True) is 
 False
<>> 
 (True or (not True)) and False is 
 False

 
 (not (True and False)) (False and True) is 
 True


Figure 13:
Evaluate the following boolean expression as either 'True' or 'False'
<SYS>
((no not True) and False) or True 

((not not True). and False) or True is 

Figure 14: The prompt used to elicit responses from the CodeLlama- 13B model in the instruction prompted and sho-shot settings, respectively. The instance is underlined.
20

---

Embeddings
Head 0.1.Q
Head 0.1.K
Head 0.1.v
MLP
Head ILV
MLP3
Head 6. 6.V.V
Head 6.9.V
Head V.
Head 1.
Head 7.10.Q
MLP6
MLPS
Head 7.10.0
Head 8.11O
MLP7
Head 6.O
Head 9.1.0
MLP8
Head 6.9.0
MLP
Head S.8.O
M.P 20
Output
Figure 15: A circuit for GT with 77% sparsity, found by Edge Pruning. This circuit obtains KL divergence of 0.3987 and a Kendall's Tau of 7062. The corresponding values for Probability Difference and Probability Difference 10 are 0.4367 and 0.2478, respectively.
21

---

Lnshoddags
Head
MLP3
Hot T.K.
Fiod 90
Head 4330
Hoad Bub
Riad KIBVV
MIP
Read
Haull
7.5.0
MILL'S
Head IEE.V
Flnad M.T.V
Head 9,7.0
Head
MIP1I
Osger
Figure 16: A circuit for GP with 99.79% sparsity, found by Edge Pruning. It obtains KL divergence of 0.4920, an accuracy of 55.03%, a Logit Difference of 0.9701, and an Exact Match of 64.02%. Note that this circuit does not perform well the less sparse (see Figure 6). However, choose to show this circuit here as the denser ones have more edges and are unwieldy to plot.
STATL
#20la
BEEAS
012333
#2233+
Figure 17: A snippet of the CodeL -13B few-shot circuit. The entire eircuit is too unwieldy to plot, but this snippet shows densely connected region. Though a bit hard to make out, a8. h16 connects a10. h24.
22

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes].
Justification: Our abstract and introduction accurately reflect the ideas, findings, and impli- cations of our work.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
It is fine to include aspirational goals as motivation long is clear that these goals are not attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes].
Justification: We acknowledge assumptions and limitations in our paper where applicable. We also discuss the limitations of our method and point future work Section
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (c.g.. independence assumptions, noiseless settings, model sellsspeciiiationn asymptotic approximations only holding locally). ). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-tex, system might not be used reliably to provide closed captions for online lectures because it fails handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and faimesss.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community, Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
23

---

Answer: [NA].
Justification:
Our paper does not include any theoretical results.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Question:
Answer: [Yes].
Justification: We provide a complete description of our method in Appendix and provide all hyperparameters and computational details in Appendix A. We also provide all prompt formats used in Appendix E.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, other means that appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm. (b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reprodueibility may be tricky in some cases, in which authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers have some path reproducing or verifying the results.
5. Open access to data and code
24

---

Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, as described in supplemental material?
Question:
Answer: [Yes].
Justification:
We will make our code and datasets publicly available.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please the NeurIPS code and data submission guidelines (uttps://nip) cc/pu blic/ Betc//guies//ae http///eeeuiisiiinnnnei
for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https: for more details.
/ /nips. wwtppblicgguiieiccccccccccocccce wwwpc//gggeee/eeee /etgdee/eoeeeuemmmiii
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental Setting/Details
Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Question:
Answer: [Yes].
We provide all details of how the data was chosen, and implementational nuances in Section and Appendices 3. We list the hyperparameters used in A.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment Statistical Significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [No].
Justification: In our comparisons, the independent variable, sparsity, can only be controlled with an approximate target sparsity and varies by model run. Therefore, we cannot measure the variance in performance of multiple circuits at exactly the same sparsity, but we run large grid of experiments using different hyperparameters and report scatterplot of distribution of circuit performance with sparsity (Figures 2, 3, 5 and 6). For our scaling study (involving no comparisons, Section 5), we run our experiments with single seed due to computational constraints.
Guidelines:
The answer NA means that the paper does not include experiments.
25

---

The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the enror bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall
with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
It should be clear whether the error bar is the standard deviation or the standard error of the mean.
is OK to report I-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that out of range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Question:
Answer: [Yes].
Justification: We provide the runtime of all three approaches compared in Table We provide other computational details, such as GPU configurations and compute budgets, in Appendix A.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount compute required for each of the individual experimental runs well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make into the paper).
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https: https://neurt:: http///eeuripppeeeee
Answer: [Yes].
Justification:
The paper strictly follows the full Code of Ethics from NeurIPS.
Guidelines:
The answer NA means that the authors have reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g.. if there is special consid- eration due to laws or regulations their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes].
Justification: We discuss possible impacts of our work in Section 7.
26

---

Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No. they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied partieular applications, let alone deployments. However, if there is direct path any negative applications, the authors should point it out, For example, it is legitimate to point that improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility of MI..
11. Safeguards
Does the paper deseribe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA].
Justification:
We do work with any high risk datasets or models in this work.
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have high risk for misuse or dual- should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]..
Justification:
All assets and related work are properly cited in the paper.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset used and, possible, include URL.
The name of the license (c.g., CC-BY 4.0) should be included for each asset.
27

---

For scraped data from a particular source (e.g., website). the copyright and terms of service of that souree should be provided.
assets are released, the license, copyright information, and terms of use the package should be provided. For popular datasets, saperswithcode com /ataaaasees has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information is available online, the authors encouraged to reach out the asset's s creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [Yes].
Justification: In experiments, we largely repurpose publicly available datasets. The in- m--ouue version of Boolean Expressions (Section 5) is generated programmatically. All details relating to its generation are discussed in Section 5.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset is used.
At submission time, remember anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification:
We do not crowdsource or research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included the main paper.
According to the NeurIPS Code of Ethics, workers involved in data collection, curation, other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based on the requirements of your country institution) were obtained?
Answer: [NA]
Justification: Our experiments do not involve crowdsourcing or research with human sub- jects.
Guidelines:
28

---

The answer NA means that the paper does not involve crowdsoureing nor research with
human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent)
may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and expect authors adhere the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
29


Reviews:
[{"review_id": "WLEHVRnKxM", "reviewer": "Reviewer_xxo6", "summary": "The paper proposes a method to automatically discover circuits within trained transformer models that are responsible for its behavior on a specific task. The automatic discovery of those circuits forms the first step to interpret the trained model. The transformer is visualized as a graph with each attention head containing 4 nodes, 3 nodes for key, query and value producing neurons with incoming edges and 1 node for the output with outgoing edges. For a specific task, the optimization problem is formulated as the difference between the output of the full model against the output of the circuit given original and corrupted inputs. The use of corrupted inputs act as a means to identify and suppress network components that are not important for the eventual circuit output. \n\nPrevious methods include greedy search and linear approximation-based edge scoring for circuit discovery. The authors propose an edge pruning based method where each edge is assigned a binary mask value relaxed to be continuous for optimization. This mask value for a specific edge is applied to the original input while the corrupted input receives the compliment and then passed on to the node. The loss used for the output is KL divergence between the circuit output and the output of the full model. After optimizing, the continuous mask values are converted into binary values by simple thresholding. Thus, a sparse graph within the overall graph is discovered that faithfully captures the behavior of the full model on the specific task. \n\nThe authors evaluate the proposed method on a wide range of tasks and metrics. The results show that the circuits uncovered are more faithful to the full model and perform better as compared to circuits discovered through previous methods. Additionally, the circuits are far sparser as compared to the previous methods. Finally, the authors show that the method can be scaled to a 13 Billion parameter model showcasing the applicability of the method.", "questions": "The experiments conducted show that circuits discovered are faithful and sparse. However, the authors should also consider reverse engineering the circuits uncovered to showcase that the circuits are not only faithful to the overall model behavior but also interpretable.\n\nThe circuits are discovered on a component level and not on a neuron level would interpreting the resulting circuits be a time-consuming avenue. How would the method work on a weight level mask ?\n\nHow do the learning rates for different component masks affect the final circuit? It is unclear what would be a good choice for the hyperparameters. \n\nHow are the mask values initialized ?", "limitations": "The authors discuss the limitations of their method in a sufficient level of detail.", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 4, "strengths": "1. The paper proposes an interesting and efficient method on an important problem in transformer interpretability.\n2. The experiments conducted and results shown are convincing of the methods superiority in terms of faithfulness and scalability.", "weaknesses": "1. The proposed method and direction are useful for discovering circuits for specific tasks. However, components of the circuits uncovered may have varied behavior across tasks or even input datapoints not considered in the dataset. This makes the overall goal of mechanistically understanding each component dependent on the specific task and the corresponding data available. \n2. The use of KL divergence is prominent in the literature. Do you also consider the possibility that multiple circuits may emerge in the model with similar output behavior? It is unclear whether the discovered circuits are consistent across models trained using different seeds. \n3. One of the major concerns is that the circuit is selected based on the sparsity level. Multiple circuits can be selected from the same optimization run depending on the sparsity and faithfulness. How do the authors decide the final circuit in terms of sparsity to be used for the next step of reverse engineering?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for providing valuable feedback and suggestions. We are happy to read that you found our method interesting and useful, and the methods convincing. We would like to respond below to some of the questions and concerns raised.\n\n> Circuits may have varied behavior on tasks or input datapoints not considered in the dataset\n\nCircuit finding is indeed sensitive to the task and which datapoints are selected. Our formulation of circuit finding as loss minimization makes this dependence explicit and we evaluate the “in-distribution” generalization of a circuit to unseen examples from the same task distribution. It would be interesting for future work to explore more intuitive ways of defining tasks and evaluate the “out-of-distribution” behavior of circuits, e.g., on task compositions.\n\n> Could there be multiple circuits with a similar KL? \n\nExcellent question!  We explore this by running Edge Pruning with 12 different seeds and a fixed sparsity target across three tasks and measure their pairwise similarity via Intersection-over-Union (IoU). Please refer to the global response, which shows that the KL values are consistent across runs. However, the IoU values are in the range of 0.52—0.64, indicating the existence of many partially overlapping circuits with similar KL. We argue that this does not represent a weakness of our method, as all these circuits are valid solutions to the problem statement. However, an analysis of their commonalities and differences could motivate future work and we will include this discussion in the next draft.\n\n> How to select the sparsity level to be used for the next steps of reverse engineering? \n\nIn circuit finding, there is a trade-off between circuit sparsity and faithfulness to the full model: larger circuits are harder to interpret, but more faithful to the full model. One option is to find circuits with a range of target sparsities and identify inflection points in the KL-sparsity curve.\n\n**Questions**\n> Interpreting some of the circuits discovered\n\nGreat suggestion! We plan to provide a detailed analysis in the next draft of the paper. We outline our ongoing efforts in the global response.\n\n> Component v/s neuron level\n\nWe discuss the challenges of extending Edge Pruning to neural circuits in the global response.\n\n> Learning rates\n\nOur method requires a high learning rate. Our experiments use a learning rate of 0.8 for all masks (log alphas) and all regularization lambdas, which we found to work well across many batch sizes and training steps.\n\n> Initialization of masks\n\nFollowing [1], we initialize our log-alphas to a normally distributed tensor with each entry having a mean of 10.0 and a standard deviation of 0.01 (corresponding to initial masks close to 1). The regularization lambdas start at 0.\n\n[1] “Structured Pruning Learns Compact and Accurate Models”, Xia et. al., ACL 2022", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Wuid95oScW", "reviewer": "Reviewer_HMkg", "summary": "The authors propose \"Edge Pruning\" as an effective and scalable method for automatic circuit discovery. Edge Pruning consists of learning a binary mask over the edges of the computation graph in a transformer neural network. Edge pruning performs favorably compared to the prior art and the authors demonstrate how it can successfully be used to find new circuits in a case study.", "questions": "- Did you attempt to use any other circuit discovery methods on the CodeLlama task?\n- You acknowledge in the limitations that it is difficult for a human to fully interpret these computational graphs with hundreds of edges, but (and this is not meant as a criticism) how hard have you tried?\n- Could Edge Pruning scale to neuron or subspace-level components?", "limitations": "The outputs of circuit discovery techniques such as Edge Pruning are large computation graphs that are not, in and of themselves, interpretable.", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 4, "strengths": "- The proposed Edge Pruning method is novel and a worthy contribution to the literature. \n- The comparisons to the prior are are thorough.\n- The paper is well-written and well-organized. The details of the Edge Pruning method are explained well.\n- The case study in which the authors find a circuit in CodeLlama-13B provides excellent evidence of the utility of Edge Pruning. The circuit that they uncovered will be able to be used as an extra \"ground truth\" circuit that all future work can use.", "weaknesses": "- It would have been useful to see whether or not any prior methods for \"automatic\" circuit discovery could have found the same or a similar circuit in CodeLlama on the same task.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for leaving a thoughtful and constructive review. It endears us to learn that you found our method novel, and our case study a good demonstration of its utility. We take this opportunity to respond to some of your questions here.\n\n> Other methods on CodeLlama\n\nWe couldn’t run ACDC on CodeLlama as it was prohibitively expensive. Although EAP can work with CodeLlama in theory, we struggled to run it as its current implementation does not interface easily with scalable frameworks like FSDP/DeepSpeed. We expect that an appropriate implementation of EAP will remain efficient for CodeLlama 13B.\n\n> Human interpretation\n\nWhile it is challenging to reverse-engineer a large circuit, we believe parts of it can still be understood with enough effort. We discuss this point further in the global response.\n\n> Can Edge Pruning scale to neuron or subspace-level components?\n\nThanks for the question! Please refer to the global response for how Edge Pruning may be used in this setting.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "U8izJwpS9t", "reviewer": "Reviewer_qgos", "summary": "In this paper, the main focus lies in finding \"transformer circuits\" to perform a variant of structured pruning on transformers. While generic structured pruning removes neurons, the authors claim that it is a too-coarse approach and propose \"edge pruning\", where one output wired to more layers can be wired just to a subset of them. The method to achieve this is borrowed from a famous paper in the pruning community [Louizos, 2018]. The empirical results show the effectiveness of the method.", "questions": "The authors are invited to provide a commentary on the differences with the aforementioned literature and to comment on the technical novelty of their work, besides providing answers for the raised points.", "limitations": "Limitations are properly discussed.", "rating": 3, "confidence": 5, "soundness": 3, "presentation": 2, "contribution": 1, "strengths": "- The approach proposed is simple and I can not find a reason why it should not work in general, even beyond transformer networks.\n\n- The presentation is overall clean and fluid, despite a very limited related works section.", "weaknesses": "- The authors miss a relevant segment of the structured pruning variant named \"channel pruning\". Essentially, instead of pruning filters, you can prune input channels to a (convolutional) neural network. Of course, in a feedforward model this results in filter pruning from the previous layer, but already in residual neural networks this is not necessarily true and discovers \"circuits\". The key factor that allows the proposed approach to work lies in the massive presence of residual branches in modern architectures (in this case, transformers). I reference [A] as one of the most representative works, but a huge effort has been conducted by the pruning community around this topic in the last 6 years. As such, the novelty claims should be massively downscaled. \n\n- Besides, it is unclear what would prevent any practitioner to re-adapt any structured pruning approach along the input dimension. As such, many more comparisons in the experiment section (with either channel pruning approaches or re-adapted structured pruning ones) should be included.\n\n- Furthermore, the proposed approach is very similar to a NAS approach where sub-networks are extracted from a supernet (in your case, the full transformer). A positioning with respect to these methods, and empirical comparison in terms of search complexity and training is also missing.\n\n- Consequently to the previous points, given that the regularizer function to undermine the subnet is borrowed by [Louizos, 2018], I fail to see the technical novelty in this work.\n\n- Besides, across the paper the way the authors address \"structured\" sparsity is too vague and generic, as well as the fact that they refer to their approach as \"Edge Pruning\" that should be \"different\" from unstructured pruning. In reality, unstructured pruning does remove edges (not in the sense of the paper), as each parameter modelizes a connection between the previous and current layer. \n\n[A] He, Yihui, Xiangyu Zhang, and Jian Sun. \"Channel pruning for accelerating very deep neural networks.\" Proceedings of the IEEE international conference on computer vision. 2017.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for pointing us to the literature on Channel Pruning and NAS! We agree this is relevant and will include a discussion of these papers in our next draft. However, we believe there has been a misunderstanding regarding the goal and key contributions of our paper, which we would like to clarify below. \n\nTo summarize, our goal is not to develop a novel pruning method. Our main contribution is to *adapt* existing pruning techniques to a problem in transformer interpretability (circuit finding). In particular, to our knowledge, we are the first to prune edges rather than nodes in the computational graph to find circuits, and to scale automated circuit discovery to 10B+ models. As per these goals, the focus of our experiments is in comparing a standard pruning approach to established baseline methods for the circuit finding problem, not in comparing different pruning methods.\n\nOur method is tailored to a particular formulation of “transformer circuits” that has been widely used in the interpretability community [1,2,3,4] and is based on cross-layer edges between interpretable components. As you noted, common pruning techniques produce other kinds of “circuits”, but we highlight below why such circuit definitions have not been as widely adapted in mechanistic interpretability.\n\n> Channel Pruning with residual connections discovers “circuits”... Any practitioner can re-adapt any structured pruning approach along the input dimension.\n\nIn the presence of residual connections, pruning the channels of the input would only allow one to either remove *all* the edges from that channel (remove channel) to future layers or *none* of them (keep channel). Edge Pruning isolates the exact contribution between any two layers, which gives interpretability practitioners a much more precise signal for understanding the internal control flow of the model. \n\nAnother difficulty with pruning the channel or hidden size dimension is that the resulting subspaces are difficult to interpret and suffer from “polysemanticity” [5]. Therefore, interpretability research tends to focus on larger components such as attention heads and retain all hidden dimensions, which can be made more interpretable with tools such as LogitLens [6, 7]. In the common response, we discuss how Edge Pruning could be extended to prune more semantic features in the hidden states of the model.\n\n> Given that the regularizer function to undermine the subnet is borrowed by [Louizos, 2018], I fail to see the technical novelty in this work.\n\nBy adapting pruning techniques to the circuit finding problem, we contribute several technical changes to prior methods. These include defining masks for edges between components across layers; disentangling the residual stream from each component to each downstream component; and replacing pruned components with counterfactual activations, rather than 0's. Note that these contributions would not make sense in the context of the standard pruning literature.\n\n> [The authors] refer to their approach as \"Edge Pruning\" that should be \"different\" from unstructured pruning. Unstructured pruning does remove edges (not in the sense of the paper), as each parameter modelizes a connection between the previous and current layer\n\nWe can see how the name “edge pruning” could technically apply to many well-known techniques commonly referred to as “unstructured pruning”. However, we hope that our method “Edge Pruning” will be primarily associated with transformer interpretability, and therefore not cause any confusion in the pruning/model compression community.\n\n\n[1] “Interpretability in the Wild: a Circuit for Indirect Object Identification in GPT-2 Small”, Wang et. al., ICLR 2023\n\n[2] “How does GPT-2 compute greater-than?: Interpreting mathematical abilities in a pre-trained language model”, Hanna et al., NeurIPS 2023\n\n[3]  “Towards Automated Circuit Discovery for Mechanistic Interpretability”, Conmy et. al., NeurIPS 2023\n\n[4] “Circuit Component Reuse Across Tasks in Transformer Language Models”, Merullo et. al., ICLR 2024\n\n[5] “Feature visualization”, Chris Olah, Alexander Mordvintsev, and Ludwig Schubert, Distill 2017\n\n[6] “Interpreting GPT: The logit lens”, nostalgebraist, LessWrong 2020\n\n[7] “Eliciting Latent Predictions from Transformers with the Tuned Lens”, Belrose et al., 2023", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "qzsPnfvlFu", "reviewer": "Reviewer_TqQ2", "summary": "The paper proposes Edge Pruning, a novel algorithm for circuit finding. They claim that it compares favourably to prior methods on GPT-2 small in terms of circuit metrics like faithfulness. They also claim it scales to the 13B model size. Finally, they apply their model to circuit-finding in a 13B model, and provide preliminary analysis of instruction-following and in-context-learning capabilities.", "questions": "1. Do you see edge pruning as a suitable technique for unsupervised circuit discovery? \n2. Were there other tasks you tried in the CodeLlama-13B model? \n3. Other than interchange interventions from a corrupt prompt, did you try other kinds of ablations? (e.g. mean ablation, random-sample ablation, zero ablation)\n4. Did circuit-discovery in the 13B model lead you to novel insights about how the model was performing computation?", "limitations": "The authors have adequately acknowledged the limitations: \n1. Requiring much more GPU memory (32 H100s for Edge Pruning)\n2. Being slower than EAP for small dataset sizes\n\nThe authors have also acknowledged that circuit metrics may be misleading. It would be good if they could address this limitation as discussed in \"Weaknesses\".", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 4, "strengths": "Originality: Excellent. The approach taken to introduce learnable parameters to determine whether to include edges in a circuit is a significant departure from prior work, and opens up exciting new avenues for future research in circuit discovery. \n\nQuality: Good. The experiments conducted are reasonable and the analysis supports the claims made. In particular, it's highly promising that Edge Pruning can recover ground-truth TracR circuits, and outperforms ACDC and EAP by a significant margin on both KL-divergence and logit difference. Some exceptions are discussed in \"Weaknesses\" below. \n\nClarity: Fair. The writing was overall clear and flowed well. Specific technical details are not present at time of review, discussed in \"Questions\" below. \n\nSignificance: Excellent. Scalable and effective circuit discovery methods open the door towards accurate interpretability for commercial-sized language models, bringing us closer towards useful applications of interpretability at large.", "weaknesses": "1. Circuits are trained to minimize KL divergence, but evaluated using the logit difference. It would be useful to discuss the differences between KL divergence and logit difference, and comment on why they are chosen for training / evaluation respectively. \n2. When evaluating baselines, this paper used KL divergence as the objective for both EAP and ACDC. However, both EAP and ACDC originally use logit difference as the optimization objective. Therefore, I am concerned that this is not an apples-to-apples comparison with prior work. \n3. The claim that Edge Pruning outperforms EAP / ACDC in GPT-2 small is based primarily on circuit metrics. However, circuit metrics may be misleading. It would be useful to directly compare the nodes and edges found by each method, and report some graph metrics such as node / edge IoU. I would be particularly interested in whether Edge Pruning recovers components of an IOI circuit previously found through manual analysis: https://arxiv.org/abs/2211.00593\n4. Similar to the above, there is little / no analysis of circuits found in the 13B model. While it is true that a smaller circuit is generally more interpretable, it is unclear to what extent we can understand the circuits found in the 13B model. It would be useful to highlight specific nodes / edges in the 13B model that seem human-interpretable. As a stretch goal, it would also be useful to formulate and verify hypotheses via standard circuit analysis techniques.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your detailed review and feedback. We are happy to find that you found our method original, our experiments convincing and the potential impact high. We respond below to some of the points raised in the review.\n\n> KL divergence v/s Logit Difference for evaluation\n\nWe actually evaluate all methods on both KL divergence and Logit Difference, as well as additional metrics. Specifically, we evaluate the methods on KL divergence in Figure 2 and Logit Difference in Figure 3. Additional faithfulness (Exact Match, Figure 5) and accuracy (Figure 6) metrics are provided in Appendix C.\n\n> Training with KL divergence instead of Logit Difference for ACDC and EAP\n\nThe ACDC algorithm actually uses KL divergence as the target metric. In Appendix C of the ACDC paper [1], they justify this choice by showing that Logit Difference (or any task-specific metric) can be over-optimized. EAP acknowledges that [1] recommends using KL divergence as the metric as well (Section 3.3 of [2]). However, as 0 is the global optimum for KL divergence, they (EAP) found that the gradient and thus their scores were the zero vector at that point, leading them to use task-specific metrics. It is therefore preferable to use KL divergence as the target when possible, as Edge Pruning does.\n\n> KL divergence v/s circuit overlap metrics\n\nOur choice of evaluation was motivated by [3], which discusses at length why faithfulness metrics are more robust and preferable as evaluation metrics as compared to circuit overlap metrics. In addition, overlap metrics are obtained against manually reverse-engineered circuits, which might not exist yet (e.g. for GP), be incomplete (e.g. IOI – only attention heads were analyzed), not be fine-grained enough (both IOI and GT manual circuits only find important nodes, not edges between them) or even inaccurate (it is difficult to evaluate second order effects like the impact of removing A on the importance of B). Nonetheless, we agree that including these metrics can offer a more rounded comparison. We show the ROC plots for Edge Pruning and ACDC on IOI and GT in the PDF attached with the global response. On the former, Edge Pruning achieves a slightly higher AUC on IOI, and a slightly lower AUC on GT as compared to ACDC. We will also include these results in the next draft.\n\n> Analysis of circuits in the 13B model / Novel insights thereof\n\nThanks for your suggestion! Please refer to our global response for our additional analysis of the circuit in the 13B model.\n\n**Questions**\n\n> Do you see Edge Pruning as a suitable technique for unsupervised circuit discovery?\n \nWe are not quite sure what kind of unsupervised circuit discovery you are referring to in this context. It would be great if you could clarify! In the paper, our method is trained with the KL loss and technically does not need labels during training.\n\n> Did you try other tasks with CodeLlama-13B?\n\nOriginally, we also performed experiments with Dyck Languages and Object Counting from BBH with CodeLlama-13B and Llama-2 13B. However, we found that (1) The instruction-prompted performance was usually very low for these tasks, and (2) Llama-2 generally performed worse than CodeLlama on the algorithmic BBH tasks. Hence, we settled upon Boolean Expressions and CodeLlama for the case study.\n\n> Did you try other ablations (mean, zero, random sampling, …)?\n\nThat is a good question. We settled on interchange ablations relatively early, and haven’t run the main experiments with other forms of ablations yet. We will include more results in a future version.\n\n[1] “Towards Automated Circuit Discovery for Mechanistic Interpretability”, Conmy et. al., NeurIPS 2023\n\n[2] “Attribution Patching Outperforms Automated Circuit Discovery”, Syed et. al., NeurIPS 2024 ATTRIB Workshop\n\n[3] “Have Faith in Faithfulness: Going Beyond Circuit Overlap When Finding Model Mechanisms”, Hanna et. al., COLM 2024", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
We thank the reviewers for their insightful comments. We are happy to see that the reviewers found our method interesting, the experiments convincing, and the potential impact high. 
At the same time, two questions were raised by multiple reviewers. We would like to address them here.

> Can Edge Pruning work with neurons? 

Individual neurons are typically polysemantic, i.e., [1] show how a single neuron responds to academic citations, English dialogue, HTTP requests and Korean text. Recent work suggests that Sparse Auto-Encoders (SAEs) trained to reconstruct the hidden states of a language model learn sparse features that correlate with distinct, interpretable concepts [2]. Recently, [3] have adapted EAP to work with “feature circuits” over such SAE features rather than heads/MLPs. The number of all feature edges is prohibitively high (quadratic in #features), and [3] used a two-stage approach where they first found important features (nodes), and then found important edges between them. Edge Pruning could be similarly adapted in this fashion: we can first model masks over features to find important features, and then learn masks over pairs of them. Since the former step only involves as many masks as the number of features, its memory overhead should be a constant factor.

> How hard is it to interpret the features of the 13B model / Did you find anything interesting? 

Interpreting circuits with >1000 edges remains difficult, but we have made progress understanding parts of the circuit. For example, we have found the following circuit of two composed heads:
L8.H16 attends from operations (and/or) to the previous token (i.e. from op to a in “a op b”)
L10.H24 attends from an operand to a previous operation (i.e. from b to op in “a op b”) and read the results from L8.H16
This suggests that this duo computes the value of the expression. Interestingly, the attention pattern also holds when a is not a literal like “True” but an arbitrarily nested subexpression. A hypothesis here is that the model could deal with arbitrary depth expressions by guessing the value of “a”—allowing it to proceed with the second step—and later verifying the guess. This would also allow the model to parallelize a sequential computation by doing them in parallel. Nonetheless, further study and careful interventions are required to verify this hypothesis. We will include these preliminary findings in the next draft and hope that it may inspire future work to study these circuits in more detail.

Finally, we have attached a PDF with this common response containing figures that we referred to in some of the responses. The first figure runs Edge Pruning with multiple seeds and finds that we discover multiple partially overlapping circuits with similar (good) faithfulness at the same sparsity. The second demonstrates that Edge Pruning is competitive on circuit overlap metrics when compared to ACDC.

[1] “Towards Monosemanticity: Decomposing Language Models With Dictionary Learning”, Bricken et al., https://transformer-circuits.pub/2023/monosemantic-features 

[2] “Sparse Autoencoders Find Highly Interpretable Features in Language Models”, Cunningham et. al., ICLR 2024

[3] “Sparse Feature Circuits: Discovering and Editing Interpretable Causal Graphs in Language Models”, Marks et. al., arXiv 2024

Meta-review: This paper proposes a novel automated method to discover circuits in transformer models. The method uses a technique from the pruning literature, adapted to this task. The paper gives good comparison with literature and is well written. The experiments show that the method scales to large models and is an important contribution in the field of interpretability of large language models.
Decision: Accept (spotlight)