Summary
This paper proposes a method for finding subnetworks (called "circuits") in transformers that calculate intermediate variables in a hypothesized causal model. The method is evaluated in terms of two goals. First, if a variable is part of the true causal graph, the method should find a subnetwork that calculates it reliably, and if the variable is not part of the causal graph, the method should return an empty subnetwork. Second, if a subnetwork is found, ablating the subnetwork should harm the performance of the model. The method is compared to probing methods (for the first goal) and causal analysis methods (for the second goal) and generally agrees with existing methods or compares favorably.
Reasons to accept
- This paper brings together ideas from pruning, circuit discovery, probing, and causal analysis. Causal analysis methods are designed to test whether neural network representations can be abstracted as high-level variables in a causal model, whereas circuit discovery methods can help characterize the mechanisms the model uses to compute these variables. This kind of combined approach could be a promising tool for understanding transformers at different levels of abstraction.
- The paper includes experiments in a variety of different settings and compares circuit probing to a number of sensible baselines. The experiments suggest that circuit probing can help answer a variety of different interpretability questions.
- The experiments indicate that circuit probing generally agrees with existing methods and in some cases performs better, in the sense that circuit probing agrees more with our expectations. For example, in section 4.3, circuit probing seems to be work than other probing methods as an explanation for "grokking" training dynamics.
Reasons to reject
- One of the proposed benefits of this approach is that it enables a form of causal analysis, but it is not clear to me how much the method actually reveals about the causal model implicit in the network. In particular, the causal analysis involves ablating the identified circuit, and the paper states that "this ablation is equivalent to asking the counterfactual question 'How does the model’s output change if it does not compute a particular intermediate variable in a particular block?'" However, I do not see how these are equivalent. The question assumes that the chosen subnetwork (a) computes a particular intermediate variable, and (b) does not do anything else. That is, ablating the subnetwork could have other effects on the model that are not related to the intermediate variable. More generally, the paper would be stronger if claims about causality could be stated more formally. For example, existing work on causal analysis (e.g. [1]) uses a more more formal definitions of the sense in which the method can support claims about the causal structure of the neural network. It would be helpful if this paper could be more precise in explaining the type of causal claims that could be supported by the method.
- The paper argues that circuit probing is better than existing probing methods because it is more "faithful" to the original model but it does not provide a clear definition of faithfulness, and it is not clear to me how the experiments demonstrate that circuit probing is more faithful than other methods. In the introduction, faithfulness is defined parenthetically as "it only provides evidence in support of causal variables that are actually represented by the model," but I do not see how there is any ground-truth way of knowing what causal variables the model actually represents. (One possible approach would be to include experiments with networks where the ground truth causal structure is known, e.g. using [Tracr](https://arxiv.org/abs/2301.05062) [2].)
- I do not always find the experiments to be very compelling. Three of the four experiments study one-layer networks trained on very simple tasks, which do not clearly demonstrate that this is an appropriate method for analyzing transformer language models. The fourth experiment studies language modeling tasks, but I am not entirely convinced by the conclusions. Section 4.4 states that "the dependency is computed in layer 6's attention block", but this seems to overstate the result. From Appendix Figure 15, it seems that ablating a subnetwork in this block leads to a larger drop in performance relative to other layers. However, the drop still isn't especially large (e.g. ~95% to ~80% on SV agreement), and ablating circuits in other components also induces a drop in accuracy, albeit not as large. Even if the effect could be entirely localized to this component, it is not clear to me how significant this type of insight is.
- I find the description of the method in section 2 to be somewhat vague. It would be helpful to start with a more formal statement of the problem, and describe the method more formally. One particular point that confused me: When the method is applied to some component C, does this mean that the mask is learned only to the parameters in C? Or is the mask learned for all components prior to C in the computation graph, but using the representations from C to calculate the loss? Is there any justification for one or the other? Without explaining these details, it is difficult to understand how the method works, or even the exact problem it is meant to solve.
- The method involves two ideas: one is to use pruning/subnetwork analysis as a probing method, and the other is to find these subnetworks by optimizing a nearest neighbor loss. These ideas seem to be distinct to me, and I think it would make sense to evaluate these two components of the method individually, that is: (a) the same approach but replacing the nearest neighbor objective with a standard linear probing method; and (b) training a linear probe on the network representations, but using the nearest neighbor objective.
Overall, I think the paper introduces an interesting approach to interpretability, but I think a number of the key points need to be formulated more precisely to be able to determine the types of questions this method is capable of answering.
*References:*
[1] Geiger et al., 2021. Causal Abstractions of Neural Networks.
[2] Lindner et al., 2023. Tracr: Compiled Transformers as a Laboratory for Interpretability.
*Update:* After reading the rebuttal, I still have some doubts about the significance of the method and findings, but the authors have promised to address some of my concerns about clarity and presentation, so I am increasing my score from 4 to 5.