Finding Transformer Circuits with Edge Pruning

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.

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. Quality: 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. Clarity: Fair. The writing was overall clear and flowed well. Specific technical details are not present at time of review, discussed in "Questions" below. Significance: 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. 2. 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. 3. 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 4. 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.

Questions

1. Do you see edge pruning as a suitable technique for unsupervised circuit discovery? 2. Were there other tasks you tried in the CodeLlama-13B model? 3. Other than interchange interventions from a corrupt prompt, did you try other kinds of ablations? (e.g. mean ablation, random-sample ablation, zero ablation) 4. Did circuit-discovery in the 13B model lead you to novel insights about how the model was performing computation?

Rating

7

Confidence

4

Soundness

3

Presentation

3

Contribution

4

Limitations

The authors have adequately acknowledged the limitations: 1. Requiring much more GPU memory (32 H100s for Edge Pruning) 2. Being slower than EAP for small dataset sizes The authors have also acknowledged that circuit metrics may be misleading. It would be good if they could address this limitation as discussed in "Weaknesses".

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.

Strengths

- The approach proposed is simple and I can not find a reason why it should not work in general, even beyond transformer networks. - 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. - 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. - 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. - 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. - 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. [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.

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.

Rating

3

Confidence

5

Soundness

3

Presentation

2

Contribution

1

Limitations

Limitations are properly discussed.

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.

Strengths

- The proposed Edge Pruning method is novel and a worthy contribution to the literature. - The comparisons to the prior are are thorough. - The paper is well-written and well-organized. The details of the Edge Pruning method are explained well. - 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.

Questions

- Did you attempt to use any other circuit discovery methods on the CodeLlama task? - 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? - Could Edge Pruning scale to neuron or subspace-level components?

Rating

8

Confidence

3

Soundness

3

Presentation

3

Contribution

4

Limitations

The outputs of circuit discovery techniques such as Edge Pruning are large computation graphs that are not, in and of themselves, interpretable.

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. Previous 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. The 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.

Strengths

1. The paper proposes an interesting and efficient method on an important problem in transformer interpretability. 2. 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. 2. 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. 3. 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?

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. The 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 ? How do the learning rates for different component masks affect the final circuit? It is unclear what would be a good choice for the hyperparameters. How are the mask values initialized ?

Rating

8

Confidence

4

Soundness

3

Presentation

3

Contribution

4

Limitations

The authors discuss the limitations of their method in a sufficient level of detail.