Knowledge Circuits in Pretrained Transformers

Summary

This paper tries to study how language models (LMs) recall factual knowledge by finding circuits within the model internals in a mechanistic way, on a dataset of (subject, relation object)-like sentences. The discovered circuits help understand the mechanisms of LM's knowledge recall, and also enable some explorations on the impact of knowledge editing methods and analyze hallucinations and in-context learning (ICL).

Strengths

- Discovering the internal mechanisms of LMs' factual recall is an important problem for better understanding and improving LM's knowledge-related abilities. The methodologies and the discovered circuits presented in the paper are interesting and may be useful for the community. - The analysis based on the discovered circuits is informative and could inspire further applications of circuit discovery.

Weaknesses

- Overall, the findings in this paper are not particularly surprising, and mostly verify known mechanisms studied by prior investigations on knowledge recall in LMs. This could also be seen from the paper writing - in most (if not all) places where a discovered circuit component has a dedicated functionality, it is something already found in the literature (e.g., mover head, relation head, mixture head, etc.). Also given the simplicity of the task (recall the object entity given the subject and relation), it is not hard to anticipate these mechanisms in a transformer model in the first place. - The introduced methodology is not very technically involved. The adopted circuit discovery method is a direct application of standard causal mediation analysis in the literature, and very much follows the framework of Conmy et al. as cited. - The interpretations of the mechanisms are somewhat subjective and not rigorously justified. For example, there are usually direct claims about a certain head or a certain MLP layer functioning as such and such, without enough supporting evidence. It would be better to include more statistics to help convince the reader why things are the way they are described in the main text. For example, how to quantitatively define what a "mover head" is, and how did you get this from the statistical evidence? - There are few in-depth discussions, potentially due to the simplicity of the task. But even within the task itself, the overall findings are quite fragile and it is very hard to get the bigger picture from all the presented findings. Ultimately, the important question is: are there issues with the current LM's mechanisms on factual recall, and how to address them? Take hallucination 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? - The writing overall could be improved. For example, "zero ablation" and "MatchNLL" are repeatedly mentioned in Section 3.1 and 3.3. The caption around Figure 2 is not informative enough to understand the figure content (e.g., are these attention/mlp nodes?) and is too close to the main text.

Questions

See weaknesses.

Rating

5

Confidence

3

Soundness

3

Presentation

2

Contribution

2

Limitations

Yes.

Summary

The paper introduce a critical subgraph named Knowledge Circuit in the language model to understand how LLMs store and express the knowledge. Through extensive experiments and cases analysis, the authors show the significance of this proposed Knowledge Circuit, which can unveil implicit neural knowledge representations, elucidate internal mechanisms for knowledge editing, and facilitate interpreting language model behaviors.

Strengths

1. The authors approach the knowledge mechanism of LLMs from the Knowledge Circuit perspective, which has not been explored in this field and has good interpretability. 2. The authors conduct extensive experimental analysis, and find that the Knowledge Circuit is beneficial to unveiling implicit neural knowledge representations, understanding knowledge editing mechanism, and interpreting language model behaviors. 3. The paper is written-well, with a direct motivation and a well organized structure.

Weaknesses

1. It would be better if the authors analyze more knowledge editing methods. 2. What is the average size of this subgraph? And what is the distribution of nodes?

Questions

Please see above Weaknesses.

Rating

7

Confidence

4

Soundness

4

Presentation

4

Contribution

4

Limitations

None

Summary

This paper focus on the problem of interpreting the knowledge storage mechanism of the large language models. The authors proposed a new perspective, which uses the knowledge circuit to understand how the language model stores and expresses the knowledge. Knowledge circuit is a subgraph of the computation graph of the model. Specifically, the knowledge circuit of a given factual knowledge is defined as the smallest subgraph to recall it "from the knowledge stored in the model". They proposed a knowledge circuit construction method, where the nodes are ablated one-by-one to measure the impact of removing it from the computation graph. The resulted knowledge circuit will then be a necessary computation graph for maintaining the original inference result. Experiments are conducted across a variety of domains of factual recall tasks. To intuitively interpret the constructed knowledge circuits, they probes the information flow by projecting the latent representations into the space of token embeddings. Under this interpretation, several conclusions are drawn as follows. - language model tends to aggregate knowledge in the earlier to middle layers and further enhances this information in the later layers. - the knowledge editing method ROME tends to incorporate edited information primarily at the edited layer. - when hallucination occurs, the language model fails to correctly transfer knowledge to the final token in the earlier layers. - several new attention heads emerge in the knowledge circuit during in-context learning.

Strengths

- It is in general well formatted and easy to follow. - Experimental details are well included.

Weaknesses

- Regarding the novelty of this work, reference [36] appears to be highly relevant, as it involves the construction of a flow graph and the presentation of semantic information flow. However, this reference is not discussed in the related work section. - The robustness of the knowledge circuit construction method should be addressed. For instance, the results of the construction method should be unique or equivalent under different traversal orders. If not, the solution may not be well-defined. - The interpretation based on the knowledge circuit lacks clarity. Although several conclusions are drawn in the paper, it is unclear if these conclusions align with existing objective observations on related phenomena, such as in-context learning, and whether they could aid in predicting the model's behavior. The effectiveness of an explanation method should be validated by deriving a conclusion or guidance that can be subsequently verified through observation. This aspect is not clearly demonstrated in the paper. - It would be beneficial to include formal definitions of mover head, relation head, and mixture head, as these are crucial concepts in this paper.

Questions

- Please refer to the weakness part. - According to Table 1, the performance drop of Circuit varies across different datasets. Is there any analysis of this phenomenon?

Rating

5

Confidence

4

Soundness

2

Presentation

3

Contribution

3

Limitations

Yes

Summary

The authors propose utilizing techniques from mechanistic interpretability to explore the connections between model components involved in factual recall. They systematically ablate connections between model components in a reverse topological order, maintaining a list of the connections that most significantly recover loss on factual recall tasks. Several properties of the resultant "knowledge circuit" are measured, including its sufficiency for recovering model performance on factual recall. The authors then demonstrate the roles of various components of the knowledge circuit, such as relation heads, mover heads, extract heads, and MLPs. They interpret the roles of these components by projecting their outputs directly into the logit space of the model. This approach is applied to understanding hallucinations and model editing failures. The authors demonstrate how knowledge editing methods like ROME switch the role of the subject mover head to instead extract the final answer directly from the subject token rather than moving information about the subject. The authors also manage to explain failures in multi-hop knowledge editing - the edited knowledge "dominantly saturates" the subject mover head preventing it from performing its natural function. They also show how factual hallucination failures can be attributed to a mistake on the part of the subject mover head.

Strengths

- As far as I am aware, this paper is the first to provide a mechanistic interpretation of the impact of knowledge editing techniques on a model. The authors demonstrate how knowledge editing methods like ROME switch the role of the subject mover head to instead extract the final answer directly from the subject token rather than moving information about the subject. The authors also manage to explain failures in multi-hop knowledge editing - the edited knowledge "dominantly saturates" the subject mover head preventing it from performing its natural function. This work will be quite relevant to other researchers thinking about improving model editing techniques. - The authors provide a new perspective on the origins of hallucinations in the model. They attribute hallucination failures to the inability of the subject mover head to move the correct answer to the final token position. This mechanistic understanding offers valuable insights into why hallucinations occur and how they might be mitigated. - The authors contextualize their findings amongst other papers that attempt to mechanistically interpret factual recall in models. The authors also reproduce a claim mentioned in several previous works, involving components in the early layers of a model being more important for knowledge retrieval (Figure 2).

Weaknesses

- The authors choose to perform ROME [1] edits at layer 0 of the model. In contrast, the original ROME paper performs edits in the middle of the model. The precise layer is chosen using a causal tracing algorithm. This choice is not justified in the paper. It is unclear what difference this makes in the results provided in this paper. - In table 1, it is unclear why ablating everything but the circuit improves performance on some tasks, but decreases others. For example, restricting the model to just the circuit increases the accuracy of "Country Language" but decreases the accuracy of "fruit inside color". [1] Kevin Meng, David Bau, Alex Andonian, Yonatan Belinkov. "Locating and Editing Factual Associations in GPT." (2023).

Questions

- Why did the authors choose to perform ROME edits at layer 0 of the model, in contrast to the original ROME paper's approach of editing in the middle layers using a causal tracing algorithm, and what impact might this choice have on the results presented in the paper? - In Table 1, what explains the inconsistent effects of ablating everything but the circuit on task performance, such as improving accuracy for "Country Language" while decreasing it for "fruit inside color"?

Rating

8

Confidence

4

Soundness

4

Presentation

4

Contribution

4

Limitations

The authors discuss the limitations of their work, and how it can contribute to ensuring safety and privacy information to promote trustworthy AI.

Summary

This paper presents a circuit-based analysis of factual recall and model editing in language models. The authors make two key claims about knowledge storage and retrieval: 1) Factual recall is compatible with circuit-based frameworks. 2) Traditional model editing methods influence behavior change within the model's "knowledge circuit." The paper compellingly advocates for studying knowledge storage and recall using the circuit framework, presenting preliminary but promising results.

Strengths

This paper is well-written, and the experiments are presented with clarity. While some of the results are still preliminary, the paper provides a novel perspective on interpreting and understanding the subject.

Weaknesses

## General Comments Despite the strengths and contributions highlighted, a significant shortcoming of this paper is the absence of a concrete, successful application of circuits for the tasks mentioned in the introduction and Section 6, such as model editing, hallucination detection, and in-context learning (ICL). This is why I described the review as "preliminary," even though the paper provides concrete experiments and evidence on the existence of a connection between factual knowledge recall and circuits. Despite its shortcomings, this paper offers an intriguing and promising new approach to addressing the problem by building on Niu et al.'s [21] call to "...examine the entire decision-making circuit" when studying how knowledge is stored and recalled in language models (LMs) with supporting evidence and experiments. The authors highlight the limitations of the previous knowledge neuron theory of LMs, noting its inadequacies as a fundamental theory for interpreting the factual recall process in LMs. The community needs a new theoretical framework, and I believe this paper points us in the right direction. ### Knowledge Circuit Information Analysis This section of the experiment aims to demonstrate Claim #1: that circuits can adequately explain the process of factual storage and recall. While the experimental results provide some preliminary support for this claim, the evidence is not as compelling as suggested in the introduction. Although the model shows initial success in using circuit-based interpretations for factual recall, some of the results are underwhelming. Notably, the relation "capital city" has a test set performance of 0%. Despite the authors' assertion that the standalone performance of the circuit is good, there is no baseline performance for comparison. Perhaps using random circuits as a baseline could provide a clearer benchmark. ### Related Work The paper provides an adequate summary of prior work in Section 7 and the introduction; however, there is room for improvement. The introduction initially highlights the issue with previous knowledge editing methods at lines 40-42, noting that "different types of knowledge are often located in the same areas." While this is a valid issue, it is not the most critical one. The relevant work section (lines 336-339) offers a more comprehensive summary of the issues. Therefore, it would be beneficial to expand the introduction by incorporating a condensed summary from lines 336-339 to provide a clearer and more thorough context for the readers.

Questions

Please see the weaknesses section.

Rating

8

Confidence

5

Soundness

3

Presentation

3

Contribution

3

Limitations

The authors have adequately addressed the limitations.

Acknowledgement to the rebuttal

Thank you for providing a response to my comments and questions. It is very helpful. I maintain my positive review to the paper.

Response to rebuttal

Thank you for taking the time to respond. I appreciate your thoughtful feedback, which has helped strengthen my analysis of the comparison between ROME editing and FT. This has increased my confidence in the high quality of your work.