Knowledge Circuits in Pretrained Transformers

Dear reviewer 8n8x,

Many thanks for your detailed and constructive comments. We hope that this reply answers all your questions, and we look forward to further discussion.

Addressing Weaknesses and Questions:

WK1&WK2: The findings in this paper are not particularly surprising, and the introduced methodology is not very technically involved.

• Factual recall may be a simple task for humans, but for current LMs, it’s still challenging, and we still do not know the mechanism of current LMs. In the previous work aiming to understand the behavior of LM, they mainly focused on the identified objection identification (IOI), which doesn't require memory of the knowledge in the model. Instead, in our work, we focus on knowledge recall, which is the basis of many complex tasks like multi-hop reasoning and chain-of-thought. As you mentioned in the strengths, our work could inspire further applications of circuit discovery.

• Second, previous findings about the special attention head are intriguing but do not apply to the factual recall task. Also, even if we anticipate these mechanisms in a transformer model in the first place, our work provides a new perspective to find them and investigate their behavior change. The analysis through the knowledge editing task enables us to see the mechanism more clearly and understand the drawbacks of the current LM and editing methods, which are also our main contributions.

WK3: include more statistics to help convince the reader why things are the way they are described in the main text. For example, how do quantitatively define what a "mover head" is, and how did you get this from the statistical evidence.

Thanks for your constructive advice. We provide the definition of these heads in Appendix and we will refine it in the main text. Here, we show more concrete definitions and statistics.

• Mover Head: Attention Heads that attend strongly to the subject token and extract attributes of the subject. From Figure 3 and Figure 4, we can see that the rank of the token is similar to the subject position and the last token position. And in the knowledge editing task, we can view the changes in the behaviors.

• Relation Head: Relation heads extract many relation attributes by attending from the end token to the relation tokens.

We select these attention heads by the attention score of the specific direction in which they put more than half of their influence on the attention pattern. To better see the function of these components, we corrupt them in the computation to see the change in the model’s output. The amount of the mover head and relation head we removed is 10.

From the Table, we can find the model's performance drop when we remove these special attention heads. Specially, mover head contribute more to the subject-specific information, as ablating them would increase other relation-related tokens' probability. Moreover, if we ablate the relation head, we can find the model tends to generate some meaningless tokens instead of the relation information. We also show one case of the output in Figure 3 in the attached pdf. From the figure, we can find ablating the mover head would increase the probability of "Italian", "English" and "Spanish", which are not subject-related. While ablating the relation head would lead to the increase of some meaningless word "a", "that", which are not relation-related.

WK4: Are there issues with the current LM's mechanisms on factual recall, and how to address them? Take hallucinations for example, is it always the case that "the model fails to move the correct knowledge to the final token in the earlier layers"? If this is the case, why does this happen, and is there a general method to resolve the issue besides checking every error example and fixing the problematic components?

A good question, and that is exactly our ultimate goal. Current LMs still show some disadvantages, and understanding the mechanism is important for us to solve them. In our work, as mentioned by reviewer jmbU and in Figure 8 in the paper, we found the failure in the multi-hop questions scenario can be attributed to saturated information and the skip of the first hop. Also, regarding the hallucination scenarios when the knowledge is stored in the knowledge model, we can see the mover's head select the wrong information, and until the last layer, the true answer moves to the top token.

For the possible solutions of these issues, knowledge editing currently shows the potential to address hallucinations and errors in large language model. In our analysis, we can see the ROME methods can alter the behavior of the mover head to help it gain the right information. In the future we can conduct some regularization here to make the information here at a proper level. As you mentioned, our work could inspire further applications of circuit discovery, We acknowledge that there are still many mysteries that we haven’t uncovered, but we believe our work is a good beginning.

WK5: The writing overall could be improved.

Thanks for your advice. We will polish our paper to make it more clear and easy to understand.

Thank you for the response, which addresses my concerns to a certain degree. Regarding "previous findings about the special attention head are intriguing but do not apply to the factual recall task", I believe this is not accurate. For example, "Summing Up the Facts: Additive Mechanisms Behind Factual Recall in LLMs" by Chughtai et al (also cited as [15] in the paper here) investigates factual recall and finds "subject heads", "relation heads", 'mixed heads', etc. which overlap a lot with the findings here.

Overall, I believe the major weaknesses I mentioned (lack of novelty in techniques and findings, lack of deeper insights) still hold, but I'll slightly raise my evaluation due to the extra positive signals and promises from the response.

Dear reviewer 8w5y,

Many thanks for your detailed and constructive comments. We appreciate that you are optimistic about the impact of the work.

Response to weakness and questions:

WK: It would be better if the authors analyzed more knowledge editing methods.

Thanks for your advice, we agree with your opinion. In fact, we previously evaluated the Knowledge Neuron methods to edit and we found utilizing knowledge neurons always failed to overturn the model’s prediction. We would consider other hyper-network editing methods in the future.

Question: circuit size and distribution.

• In Table 1, we show the number of edges of the discovered circuit. Take GPT2-Medium as an example; the average number of edges in the final identified circuit is around 400, and the average number of nodes is about 80. This proves that with only less than 5% of the original computing subgraph edges, the model can maintain over 70% of its original performance, demonstrating the effectiveness of our methods.
• We list the distribution of the circuit in Figure 2. Consistent with our expectations, the initial layers of the network, where the model processes inputs and extracts foundational information, are dominated by the activity of attention and feed-forward blocks. As the model progresses to undertake more nuanced and specialized tasks, the engagement of individual attention heads and FFN blocks diminishes correspondingly.

Dear reviewer 7LFR,

Many thanks for your detailed and constructive comments. We recognize and agree with the limitations of our work and address the specific comments and improvements here.

Response to weakness and questions:

WK1: Relation with reference [36]

Thank you for bringing this up! We would modify our paper and add this missing discussion. The aim of the VISIT aligned with our paper. However, they mainly build the graph by logit lens and by pruning the unactivated neuron, while our methods measure the contribution to the prediction. Also, they mainly focus on the behavior of the specific neuron or block in the Transformer model, ignoring the cooperation of these different components, while our work provides a new perspective. Finally, they simply conduct several case analyses and do not consider further analysis of the background mechanism under knowledge editing and hallucinations. We believe our two papers supplement each other, and employing the fine-grained analysis could help us understand our identified circuit more in the future.

WK2: Robustness of the methods.

Very good point. Our method is computed based on the reverse-topological order (so the nodes are sorted from output to input) to traverse the computation graph.

• We put the different computed circuits based on the topological order and our reverse-topological order in Figure 1 in the attached pdf in the global response. We can see that if we use the original topological orders, the resulting circuits tend to retain redundant nodes in the initial layers. This occurs because the Transformer processes information layer-wise. If nodes from an earlier layer are ablated first, it disrupts the input to subsequent layers, potentially causing a substantial decline in performance. In contrast, when we apply our reverse-topological order, the ablation of nodes in a later layer does not impact the computations of the preceding layers. Also, the topological order would omit the later layers as in our experiments in the main text, the model would give us the answer before the last layer, leading to an incomplete circuit.
• Meanwhile, in our experiments, we found that our reverse-topological order also takes less time than the topological order. By ablating nodes in the later layers at the outset, our method significantly reduces the overall computation time. In comparison, the topological order necessitates computing through all layers until the final stage, which is inherently more time-consuming due to the denser circuits in the later layers.

These results demonstrate the robustness of our methods, and from our view, this point deserves a new paper to discuss the casual effect order when computing the graph and is not in the main scope of our work. Thanks again for your very good point.

WK3: lacks clarity.

Thanks for your advice. As to the ICL problem you mentioned, our work supports the previous ICL mechanism [1]. They found the special attention heads’ output acts like the function vector for ICL demonstrations, and if we add the function vector output to the original zero-shot input, the model can do the same task. However, they didn’t analyze the behavior of these attention heads. Based on our proposed methods, we can discover these special attention heads demonstrate the behavior in Figure 5(b).

To clarify our findings, we conduct experiments by ablating the newly appeared attention head in the ICL circuit. We can find that compared to the randomly selected attention head by ablating this attention, the probability drops significantly in the prediction, demonstrating the importance of the component our method identified.

Moreover, for the editing task, we provide another observation of the mover head’s behavior change in Figure 2 in the attached pdf in the global response, which also make our findings more clear. We can find that before editing, the mover head retrieves the correct information for the Windows Server 2003 “Microsoft”, but after the edit, the mover head retrieves the error information ‘Intel’, which is actually the edited information at the subject place by the ROME method.

WK4: formal definitions of mover head, relation head

Thanks for your advice. In our previous version, we defined these concepts in the Appendix due to the space limit, and we would mention them in the main paper when modifying it.

Question: the performance drop of the circuit varies across different datasets.

Our experiments test the Hit@10 to measure whether the model knows the knowledge. As we ablate everything but the circuit, we anticipate a little drop-down in the prediction. In our analysis, we found that for some kind of knowledge like "fruit inside color", the original rank of the target token is around 6-8, and after ablating other things, the rank is slightly higher than 10. Hence, this drop-down here is a bit more. The most interesting finding is that we can see an increase in several kinds of knowledge, suggesting that the discovered knowledge circuits may encapsulate the relevant knowledge and the model’s performance on these tasks could have been hindered by noise from other components.

---

[1] Eric Todd, Millicent L. Li, Arnab Sen Sharma, Aaron Mueller, Byron C. Wallace, David Bau. Function Vectors in Large Language Models (2024)

Dear reviewer jmbU,

We are very happy that you are optimistic about the impact of the work. In this comment, we respond to the question and the suggested improvements.

Response to weakness & questions:

WK1&Q1: About the editing layer we select in our experiments.

The layer in current editing methods is an important hyper-parameter; currently, there is no conclusion on which layer is the best. In ROME[1] paper, they edit layer 0 for the FT method and in our experiments, we follow their setting as the case to compare the behaviors of ROME and FT under the same setting.

Moreover, we also conduct editing experiments under different layers in Appendix, Figure 11 and 12. We found the same phenomenon across these different layers, which enhances our previous findings. Actually, despite the casual tracing results pointing out that the middle layer stores the knowledge, the identified layer is not always the best layer to edit. In Figure 12, editing ROME using the mid-layer 6 failed to overturn the model's results for the ‘Super-hero’ knowledge. This phenomenon is also found in previous work[2]. Our work provides a new view of what happened when we edited the knowledge.

WK2&Q2: inconsistent effects of ablating everything but the circuit.

Our experiments test the Hit@10 to measure whether the model knows the knowledge. As we ablate everything but the circuit, we anticipate a little drop-down in the prediction. In our analysis, we found that for "fruit inside color", the original rank of the target token is around 6-8, and after ablating other things, the rank is slightly higher than 10. Commonly, the drop-down is acceptable and the most interesting finding is that we can see an increase in several kinds of knowledge, suggesting that the discovered knowledge circuits may encapsulate the relevant knowledge and the model’s performance on these tasks could have been hindered by noise from other components.

---

[1] Kevin Meng, David Bau, Alex Andonian, Yonatan Belinkov. "Locating and Editing Factual Associations in GPT." (2023).

[2] Peter Hase, Mohit Bansal, Been Kim, Asma Ghandeharioun. “Does Localization Inform Editing? Surprising Differences in Causality-Based Localization vs. Knowledge Editing in Language Models.” (2023)

Dear reviewer Njw6,

Thank you very much for your well-considered review! We are very happy that you appreciate the work we put into exploring a novel perspective on interpreting and understanding the subject. Thank you also for the insightful questions, and please let us know if you have any more.

Response to weakness:

WK1: absence of a concrete, successful application of circuits.

We agree that our current work still has room for improvement in the proposed application. However, as you mentioned, the main contribution of our work is proving the existence of a connection between circuits and these applications, including model editing, hallucination, and ICL. We are also investigating to address these issues and we believe our work is a good start for future work in this field.

WK2: no baseline performance for comparison. Thanks for your advice. We have conducted experiments that compare the performance of randomly constructed circuits. The results can be seen as follows:

Here, we compute the random circuits by randomly deciding whether the edge should be removed and making sure the graph is connected. The random circuit is the same size as the circuit we discovered by our method. From the table, we can see that the random circuit failed to maintain the model’s performance, which further enhanced the robustness and efficacy of our methods.

Meanwhile, Regarding some odd performances like the "capital city”, we conduct more analysis here. We found the test data here may not be stored in the language model. In this case, even if we find the circuit for the knowledge, the circuit itself still cannot provide us with the correct answer as the LM itself does not store the knowledge in the model.

WK3: related work to be improved.

Thanks for your advice. Due to the limited pages of the submission, we put some of the related work in Appendix E. Your insights are greatly appreciated and will be taken into consideration. We would polish our introduction to make it more convincing.