Source-linked AI summary

Knowledge Circuits in Pretrained Transformers

Yunzhi Yao, Ningyu Zhang, Zekun Xi, Mengru Wang, Ziwen Xu, Shumin Deng, Huajun Chen

arXiv:2405.17969v4cs.CLcs.AIcs.CVcs.IRcs.LG

TL;DR

Knowledge storage in language models remains difficult to explain, and prior work often isolates individual components. This paper maps knowledge circuits in Transformer computation graphs and reports that circuits capture stored knowledge and help interpret editing and behaviors such as hallucination and in-context learning.

  • Problem

    Knowledge storage remains difficult to explain: prior work often isolates components, while editing methods face poor generalization, severe side effects, and ineffective use of edits.

    Source: §1 · p. 2
  • Method

    The paper constructs knowledge circuits in Transformer computation graphs for factual-recall tasks across knowledge domains, using GPT-2 and TinyLLAMA.

    Source: §1 · p. 2
  • Results

    Circuits can express stored knowledge and guide behavior; isolated circuits recall related knowledge, while hallucinations coincide with failed or incorrect transfer to the final token.

    Source: §1 · pp. 2–3, §8 · p. 10
  • Takeaways & Limitations

    The authors present circuit analysis as a preliminary perspective on language-model knowledge mechanisms, with potential insights for editing, reasoning, factuality, and hallucinations.

    Source: §8 · p. 10

    Circuit-discovery-based patching is time-intensive, although more efficient information-flow and circuit-discovery methods are noted.

    Source: Limitations and Broader Impacts · p. 10

Abstract

from arXiv · show

The remarkable capabilities of modern large language models are rooted in their vast repositories of knowledge encoded within their parameters, enabling them to perceive the world and engage in reasoning. The inner workings of how these models store knowledge have long been a subject of intense interest and investigation among researchers. To date, most studies have concentrated on isolated components within these models, such as the Multilayer Perceptrons and attention head. In this paper, we delve into the computation graph of the language model to uncover the knowledge circuits that are instrumental in articulating specific knowledge. The experiments, conducted with GPT2 and TinyLLAMA, have allowed us to observe how certain information heads, relation heads, and Multilayer Perceptrons collaboratively encode knowledge within the model. Moreover, we evaluate the impact of current knowledge editing techniques on these knowledge circuits, providing deeper insights into the functioning and constraints of these editing methodologies. Finally, we utilize knowledge circuits to analyze and interpret language model behaviors such as hallucinations and in-context learning. We believe the knowledge circuits hold potential for advancing our understanding of Transformers and guiding the improved design of knowledge editing. Code and data are available in https://github.com/zjunlp/KnowledgeCircuits.

1 Introduction

The paper shifts from locating knowledge in isolated components to tracing circuits of cooperating Transformer components, then uses those circuits to examine editing and model behavior. It reports findings about knowledge flow, editing, hallucination, and in-context learning.

Motivation

  • Prior work emphasized MLP knowledge neurons and key-value memory, while knowledge-editing methods have reported poor generalization, severe side effects, and ineffective use of edited knowledge [20] [21].Source: §1 · pp. 1–2

Knowledge circuits

  • The study constructs knowledge circuits from factual-recall tasks across factual, social-bias, linguistic, and commonsense domains using GPT-2 and TinyLLAMA.The circuits trace cooperation among attention heads, MLPs, and embeddings rather than isolating small knowledge-localization regions.Source: §1 · p. 2

Knowledge circuits unveil implicit neural knowledge representations

  • Independent knowledge circuits retain a significant portion of overall recall performance; information accumulates in early-to-middle layers, while mover and relation heads transfer and capture relational information.The authors report that information is further enhanced in later layers.Source: §1 · pp. 2–3

Knowledge circuits elucidate internal mechanisms for knowledge editing

  • ROME tends to place edited information primarily at the edited layer, after which mover heads transport it to the final token’s residual stream; fine-tuning directly integrates the edited token.The directly integrated token exerts a dominant influence on subsequent predictions.Source: §1 · p. 3

Knowledge circuits facilitate interpreting language model behaviors

  • During hallucinations, knowledge may fail to reach the final token because the circuit lacks an effective mover head or that head selects incorrect information; in-context learning brings new attention heads into the circuit.Source: §1 · p. 3

2 Background: Circuit Theory

The background frames a Transformer as a computation graph whose components exchange information through the residual stream. It describes circuits as selected subgraphs and introduces ablation as a way to identify critical connections.

2.1 Preliminaries

  • A Transformer’s residual stream accumulates input embeddings and outputs from attention heads and MLPs; each component reads from and writes to that stream.The output node operates on the embeddings and the outputs of attention heads and MLPs.Source: §2.1 · p. 3
  • A circuit is a selected set of graph nodes and edges considered critical to a specific behavior within the Transformer computation graph.The graph includes input embeddings, attention heads, MLPs, and an output node.Source: §2.1 · p. 3

2.2 Circuit Discovery

  • Circuit discovery uses causal mediation: ablate a node-to-node edge, run the altered model, and compare its output with the original using a chosen metric and threshold.Ablation replaces the edge contribution with zero or the source head’s mean value; sparse auto-encoders are also noted as an alternative discovery approach.Source: §2.2 · p. 4

3 Knowledge Circuits Discovery in Transformers

The section constructs knowledge circuits by pruning non-critical computation-graph edges, then analyzes component information and evaluates circuit completeness across stored knowledge types.

3.1 Knowledge Circuits Construction

  • Edges scoring below threshold τ are removed during topological traversal to derive a circuit of components essential for predicting the target entity from a subject-relation pair.The edge score measures the impact of ablation on MatchNLL loss.Source: §3.1 · p. 4

3.2 Knowledge Circuits Information Analysis

  • Layer-normalized node outputs are mapped into embedding space with the unembedding matrix to inspect component-written information and its influence on later computation.Source: §3.2 · pp. 4–5

3.3 Knowledge Circuits Experimental Settings

  • Circuit completeness is tested by constructing circuits on validation data and measuring isolated test performance with Hit@10 across linguistic, commonsense, factual, and bias knowledge.Hit@10 measures whether the target entity ranks among the top ten predictions; the dataset is LRE.Source: §3.3 · p. 5

4 Knowledge Circuits Unveil Implicit Neural Knowledge Representations

The discovered knowledge circuits retain substantial performance in sparse subgraphs and expose how attention heads and MLPs cooperate to represent and refine knowledge across layers.

Knowledge Circuits Evaluation

  • With less than 10% of the original knowledge circuit’s subgraph, GPT2-Medium maintains over 70% of its original performance.Random circuits of the same size fail to maintain the model’s performance.Source: §4 · pp. 5–6
  • The Landmark-country relation metric rises from 0.16 to 0.36, while attention and MLPs are more active in lower layers.The authors suggest noise from other components may hinder performance on some tasks.Source: §4 · p. 6, Figure 2 · p. 6

Special Components in Knowledge Circuits

  • GPT-2 distributes knowledge across attention heads and MLP matrices in several layers, whereas TinyLLAMA’s special components are more concentrated.The authors hypothesize this difference may reflect knowledge capacity and say it warrants further investigation.Source: §4 · p. 6, Figure 3 · p. 7

A Running Example of Knowledge Circuit

  • After MLP17, attention heads L18H14 (relation) and L20H6 (mover) collaborate to further enhance target-entity prediction; L7H14 may signal activation of the associated mover, though this requires further investigation.L14H7 is predominantly activated by L7H14.Source: §4 · p. 7

5 Knowledge Circuits Elucidate Internal Mechanisms for Knowledge Editing

The section examines how factual edits alter knowledge circuits. ROME routes edited information through mover heads, while fine-tuning writes it into components; edits can also dominate multi-hop reasoning.

Single Factual Knowledge Editing

  • ROME shifts the mover head from copying to extracting edited information from the subject position, while FT-M writes knowledge directly into MLP-0, where Intel’s logits exceed 10 and emerge top-ranked in the residual stream.Source: §5 · p. 7
  • For the unrelated “Windows server” query, the edited model still tends to answer “Intel,” illustrating an overfitting risk.The authors report that this behavior can affect unrelated knowledge.Source: §5 · p. 7, Figure 4 · p. 8

Multi-hop Factual Knowledge Editing

  • After changing Thierry Mugle’s home country from France to China, the mover head selects “China” during later-hop questions, saturating the reasoning circuit.The tested questions ask for the edited home country’s official currency and capital city.Source: §5 · p. 8
  • Even without first-hop context, the original model directly answers the multi-hop questions, consistent with reliance on relational and subject-related information regardless of grammatical adherence.Source: §5 · p. 8

6 Knowledge Circuits Facilitate Interpreting Language Model Behaviors

The authors use knowledge circuits to examine factual hallucination and in-context learning. Hallucinations involve failures to transfer correct knowledge to the final token, while demonstrations bring new attention heads into the circuit.

Factual Hallucination

  • During factual hallucinations, the model fails to move correct knowledge to the final token in earlier layers, either lacking an effective mover head or selecting incorrect information.In the Malaysia currency example, the mover at layer 16 selects “Malaysian” despite “Ringgit” also being accumulated before that layer.Source: §6 · pp. 8–9
  • In the Malaysia currency example, mover head L15H10 extracts the erroneous answer at layer 16, and a later rank drop for the true answer does not correct the mistake.Source: §6 · p. 9

In-Context Learning

  • With demonstrations, several new attention heads appear in the circuit, focus on demonstration context, and behave as Induction Heads that look up what followed a prior instance of the current token.Source: §6 · p. 9

7 Related Work

Prior work has examined knowledge in MLPs, residual streams, and model manipulation, while this paper considers knowledge circuits spanning attention and MLP components.

Knowledge Mechanism of Transformers

  • Prior studies describe Transformer MLPs as key-value memories storing knowledge, and report that knowledge accumulates across layers.Source: §7 · p. 9
  • Bayazit et al. use pruning across several layers to discover critical knowledge, but this may overlook connections between components.Source: §7 · pp. 9–10

Manipulate Language Models

  • Knowledge-editing and related manipulation methods often modify MLPs, while other findings highlight attention’s role and question whether within-layer editing locates facts reliably.Source: §7 · p. 10
  • The paper approaches knowledge manipulation through circuits that include MLP and attention components across different layers.Source: §7 · p. 10

8 Conclusion

The paper presents a circuit-theoretic perspective on knowledge storage and a preliminary analysis of its effectiveness. It reports that knowledge circuits express knowledge and guide behavior across settings.

Conclusion

  • Knowledge circuits are presented as a new perspective on knowledge storage, with preliminary analysis demonstrating their effectiveness.Source: §8 · p. 10
  • The authors report that knowledge circuits are responsible for expressing knowledge and can guide model behavior in different settings.Source: §8 · p. 10
  • The authors hope these findings inform knowledge-mechanism research and the design and editing of language models, including efforts to enhance knowledge, reasoning, and factuality and alleviate hallucinations.Source: §8 · p. 10

Limitations and Broader Impacts

The authors use causal mediation to construct knowledge circuits tailored to domains, but describe this discovery-based patching approach as time-intensive. They point to alternative circuit and information-flow methods and broader safety and privacy possibilities.

Limitations and Broader Impacts

A Implementation Details

The implementation selects a circuit-size threshold and samples zero-shot knowledge data, while layerwise target ranks show knowledge aggregation and later probability enhancement.

A Implementation Details

  • Circuit construction tests τ values of 0.02, 0.01, and 0.005 because high thresholds can omit nodes and low thresholds can add unnecessary nodes.The appropriate circuit size is chosen for different knowledge types.Source: A Implementation Details · p. 18
  • The circuit analysis uses ACDC2; for TinyLLAMA3, key and value pairs are interleaved and repeated to inspect individual attention heads.GPT2-medium circuit computation took about 1–2 days on an NVIDIA-A800 (40GB).Source: A Implementation Details · p. 18
  • The dataset is sampled for zero-shot knowledge storage using Hit@10, with test examples balanced 1:1 against validation examples.The source data is Hernandez et al. [42], originally used for few-shot settings.Source: A Implementation Details · p. 18

B.1 Rank Change Across Layers

  • At successive layer outputs, the target entity rises toward the vocabulary’s top ranks, then later layers continue increasing its probability mass.The analysis computes the target’s rank in the vocabulary at each layer.Source: §B.1 · p. 18

B.2 Special Components in Knowledge Circuit

The analyses characterize knowledge circuits as collaborations among components and trace their reuse and behavior across factual recall, multi-hop questions, hallucinations, and reversed relations. They also report that some models can answer a multi-hop question even after first-hop knowledge is removed.

B.2 Special Components in Knowledge Circuit

  • Mover-head ablation raises unrelated entity-token probabilities, while relation-head ablation raises meaningless tokens; subject probability at the final token is usually uniformly low.The paper cautions that prior subject-to-object findings came from a limited set of tasks.Source: §B.2 · pp. 18–19, Figure 6 · p. 19

B.3 A Case on TinyLLAMA

  • A TinyLLAMA circuit example traces the knowledge that Thierry Mugler’s mother tongue is French.Source: Figure 8 · p. 20

B.2 Special Components in Knowledge Circuit

  • Related relations reuse heads such as L21H12, and L7H14 appears in both official-language and official-currency circuits; GPT2 knowledge emerges at middle-to-later layers, and TinyLLAMA later.The authors suggest reused heads may be topic heads, a distinction they say warrants further investigation.Source: Figure 7 · p. 19, §B.2 · p. 20

C More Circuit Utilization Analysis

C.1 Multi-hop Factual Knowledge Editing

  • Single-hop and multi-hop circuits are compared by measuring shared nodes, while integrated reasoning is assessed through relations created by combining information from different hops.The analysis focuses on cases where predictions are correct for both single-hop and multi-hop questions.Source: §C.1 · p. 21
  • GPT2 and TinyLLAMA can still answer multi-hop questions after first-hop knowledge is removed, and their multi-hop circuits overlap with single-hop circuits, especially at mover heads.The reported case removes the first-hop question context while retaining the second-hop relation and first-hop subject.Source: §C.1 · p. 21, Figure 10 · p. 22

C.2 Hallucination

  • Hallucination cases accumulate correct and incorrect answers in earlier layers, but specific layers select the wrong answer; the authors hypothesize circuit competition.The competition explanation is proposed as a hypothesis.Source: §C.2 · p. 22

C.3 Reverse Relation

  • The superhero_person forward and reverse-relation circuits overlap by about 70%, and their activated mover heads are identical.Source: §C.3 · p. 22

D.1 Method Implementation

The paper compares two editing approaches through their effects on knowledge circuits, situates circuit analysis among related work, and notes limits in circuit granularity and interpretation.

D.1 Method Implementation

  • ROME applies a rank-one MLP weight modification to write a new subject–knowledge key-value pair; FT-M trains the MLP on the target answer while masking the original text.Both editing implementations use EasyEdit with its default parameters.Source: §D.1 · p. 22

D.2 Edit Cases on FT-M and ROME

E.1 More Related Work

  • Prior circuit work includes automated and task-specific circuit discovery, while this paper examines circuits containing both MLP and attention components and evaluates editing methods on them.The cited discussion notes that earlier studies often focused on a single component or on IOI and Color Objects tasks.Source: §E.1 · pp. 25–26

E.2 Limitation and Future Discussion

  • The circuits are relatively coarse, and logit-lens discrepancies can hinder concrete analysis of early-layer components; finer MLP-neuron contributions and head-activation conditions remain unresolved.Source: §E.2 · p. 26
Loading 2405.17969v4…