Source-linked AI summary
Knowledge Neurons in Pretrained Transformers
Damai Dai, Li Dong, Yaru Hao, Zhifang Sui, Baobao Chang, Furu Wei
TL;DR
Pretrained Transformers recall substantial factual knowledge, but how that knowledge is stored remains unclear. The paper introduces knowledge neurons and a knowledge attribution method for BERT’s cloze task, finding that their activation tracks corresponding knowledge expression and can support preliminary fact editing without fine-tuning.
Problem
The paper asks how factual knowledge is stored in pretrained Transformers rather than only measuring whether models can recall it.
Method
The authors use BERT’s fill-in-the-blank cloze task and knowledge attribution to identify neurons that express relational facts.
Results
Knowledge-neuron activation is positively correlated with corresponding knowledge expression across suppression, amplification, and prompt-activation analyses.
Takeaways & Limitations
Knowledge neurons provide a way to analyze factual storage and support preliminary attempts to update or erase facts without fine-tuning.
Takeaways & Limitations
The studies focus on single-word cloze expressions of factual knowledge, leaving generalized reasoning, neuron interactions, and multi-word extensions open.
Abstract
from arXiv · showhide
Large-scale pretrained language models are surprisingly good at recalling factual knowledge presented in the training corpus. In this paper, we present preliminary studies on how factual knowledge is stored in pretrained Transformers by introducing the concept of knowledge neurons. Specifically, we examine the fill-in-the-blank cloze task for BERT. Given a relational fact, we propose a knowledge attribution method to identify the neurons that express the fact. We find that the activation of such knowledge neurons is positively correlated to the expression of their corresponding facts. In our case studies, we attempt to leverage knowledge neurons to edit (such as update, and erase) specific factual knowledge without fine-tuning. Our results shed light on understanding the storage of knowledge within pretrained Transformers. The code is available at https://github.com/Hunter-DDM/knowledge-neurons.
1 Introduction
The paper investigates how factual knowledge is stored in pretrained Transformers by identifying knowledge neurons and testing their relationship to knowledge expression. It also presents preliminary attempts to edit factual knowledge without fine-tuning.
- Findings: Knowledge-neuron activation is positively correlated with expression of the corresponding factual knowledge.Suppressing or amplifying these neurons affects knowledge expression, and they activate more strongly for prompts expressing the associated fact.
- Findings: Top-activating prompts from open-domain texts usually express the corresponding fact, whereas bottom-activating prompts usually do not express the correct relation.This provides qualitative evidence connecting neuron activation with factual expression beyond the original cloze prompts.
- Editing studies: The authors conduct preliminary studies that update or erase factual knowledge by directly modifying feed-forward-network parameters without fine-tuning.The reported knowledge surgery shows promising results while retaining a moderate influence on other knowledge.
- Approach: The paper introduces knowledge neurons and a knowledge attribution method to identify neurons expressing specific relational facts in BERT’s cloze task.The method treats feed-forward networks as key-value memories and computes each neuron’s contribution to knowledge prediction.
2 Background: Transformer
A Transformer encoder consists of repeated blocks containing self-attention and feed-forward-network modules. The paper relates the FFN’s two linear layers to keys and values in a key-value memory.
- Transformer structure: A Transformer block mainly contains a self-attention module and a feed-forward network module.The encoder is described as a stack of L identical blocks.
- Key-value interpretation: The FFN and self-attention formulas are structurally similar, differing mainly in their activation functions: GELU for FFNs and softmax for self-attention.This similarity motivates interpreting FFN input as a query vector.
- Key-value interpretation: The two FFN linear layers can be viewed as keys and values, respectively, analogous to query-key-value mechanisms in self-attention.This interpretation supports modeling FFNs as key-value memories.
3 Identifying Knowledge Neurons
The paper identifies knowledge neurons by attributing factual predictions to FFN intermediate neurons, then refining candidates across diverse prompts for the same fact.
- 3.1 Knowledge Assessing Task: Fill-in-the-blank cloze prompts represent relational facts as a head entity and relation with the tail entity masked for prediction.The task treats a model as knowing a fact when it predicts the correct answer.
- 3.2 Knowledge Attribution: Integrated gradients estimate each intermediate neuron's contribution by tracking the correct-answer probability as its activation changes from zero to its original value.The attribution score accumulates output-probability change along this interpolation.
- 3.2 Knowledge Attribution: Neurons with attribution scores above threshold t form a coarse knowledge-neuron set for each prompt.The model output is the probability of the correct answer, and the neuron is an intermediate unit in an FFN.
- 3.3 Knowledge Neuron Refining: The refining strategy removes false positives by retaining neurons shared across more than p% of diverse prompts expressing the same fact.It assumes true-positive neurons are shared across varied prompts, whereas syntactic or lexical neurons are not consistently shared.
4 Experiments
Experiments evaluate knowledge-neuron identification in BERT using cloze prompts, neuron statistics, interventions, and open-domain prompts. The identified neurons concentrate in upper layers, affect correct-answer probabilities when manipulated, and respond more strongly to prompts expressing their corresponding facts.
- Experimental Setup: Experiments use BERT-base-cased on the PARAREL fill-in-the-blank cloze task, retaining 34 relations with diverse prompt templates.The model has 12 Transformer blocks, a hidden size of 768, and FFN inner hidden size of 3,072.
- Knowledge-Neuron Statistics: Most fact-related neurons occur in the topmost Transformer layers, consistent with prior findings cited by the authors.Figure 3 reports the layer distribution of neurons identified by the proposed attribution method.
- Knowledge-Neuron Statistics: 4.13 knowledge neurons are identified per relational fact on average, compared with 3.96 for the activation-based baseline.The similar scale supports fair subsequent comparisons.
- Knowledge-Neuron Statistics: Proposed-method neurons are more relation-specific than baseline neurons: same-relation facts share 1.23 neurons on average, whereas different-relation facts share almost none.The baseline shares many neurons across intra-relation pairs and a substantial portion across inter-relation pairs.
- Knowledge-Neuron Interventions: 29.03% average suppression decreases correct probability, while baseline suppression decreases it by only 1.47%.Knowledge neurons are suppressed by setting activations to 0; baseline-identified neurons have negligible average influence.
- Knowledge-Neuron Interventions: 31.17% average amplification increases correct probability, whereas baseline amplification decreases average correct probability by 1.27%.Amplification doubles the selected neurons’ activations.
- Knowledge-Neuron Interventions: The intervention results are presented as a proof of concept because neuron distribution affects control precision, which the authors leave for future work.Wider neuron distributions require manipulating more top-k neurons for better control.
- Prompt Activation: On BINGREL, proposed-method neurons activate more for knowledge-expressing prompts (T1 = 0.485) than controls (T2 = 0.019; T3 = −0.018).Top-activating prompts express the corresponding relation, while bottom-activating prompts can contain the same entities without expressing that relation.
5 Case Studies
The paper uses knowledge neurons to update and erase factual relations in pretrained Transformers without fine-tuning. These preliminary surgeries affect targeted knowledge while generally preserving other predictions.
- 5.1 Updating Facts: Knowledge-neuron surgery updates a learned relation ⟨h, r, t⟩ to ⟨h, r, t′⟩ by modifying corresponding FFN value slots.The procedure identifies fact-specific neurons, removes those shared by more than 10% of intra-relation facts, and changes the remaining slots.
- 5.1 Updating Facts: About four top knowledge neurons produce a nontrivial fact-update success rate, whereas the same number of random neurons is insufficient.Including more top knowledge neurons can further improve the success rate.
- 5.1 Updating Facts: Fact updating has little negative influence on other knowledge predictions, including knowledge expressed by same- and different-relation prompts.The evaluation measures change rate, success rate, and perplexity changes within and across relations.
- 5.2 Erasing Relations: The erasing study targets four relations involving sensitive personal information by zeroing value slots of 20 frequently occurring knowledge neurons.The relations are place_of_birth, country_of_citizenship, occupation, and work_location.
- 5.2 Erasing Relations: Erasing increases perplexity for the removed knowledge while leaving perplexity for other relations similar, suggesting targeted knowledge removal with minimal effort.The authors characterize this result as a promising way to erase undesired knowledge.
6 Related Work
Related work probes factual knowledge in pretrained models, develops attribution methods for interpreting model outputs, and analyzes Transformer components beyond self-attention. This paper connects these strands through knowledge attribution over FFN neurons.
- Probing Knowledge in Pretrained Models: Cloze-query studies show that pretrained models such as BERT can recall factual knowledge without fine-tuning.Later work improves cloze queries through mining and paraphrasing methods.
- Attribution Methods: Attribution methods assign model outputs to input features using gradient-based, back-propagation, and integrated-gradient approaches.Integrated gradients were proposed with sensitivity and implementation invariance as guiding axioms.
- Analysis of Transformer: Transformer analysis has increasingly examined feed-forward networks alongside self-attention, including their interpretation as key-value memories.This key-value perspective motivates the paper’s focus on FFN-based knowledge storage.
7 Conclusion and Future Directions
The paper identifies neurons associated with factual knowledge, shows that manipulating them changes corresponding knowledge expression, and demonstrates preliminary updating and erasing studies. Its conclusions remain bounded by cloze-based factual evaluations and unresolved questions about generalized use and neuron interactions.
- Conclusion: The proposed attribution method identifies knowledge neurons that express factual knowledge in pretrained Transformers.The paper examines these neurons in the fill-in-the-blank cloze setting.
- Conclusion: Suppressing or amplifying knowledge-neuron activation accordingly affects the strength of knowledge expression.Open-domain analyses also find that these neurons tend to activate on prompts expressing their corresponding knowledge.
- Conclusion: Two preliminary case studies attempt to use knowledge neurons to update or erase knowledge without fine-tuning.These studies are presented as initial applications of the identified neurons.
- Future Directions: The studies are limited to factual knowledge expressed through single-word cloze blanks, leaving generalized reasoning, neuron interactions, multi-word blanks, and multilingual models open.The authors also note that knowledge can be expressed more implicitly and that the method’s applicability to other knowledge types remains to be explored.