Source-linked AI summary
Kformer: Knowledge Injection in Transformer Feed-Forward Layers
Yunzhi Yao, Shaohan Huang, Li Dong, Furu Wei, Huajun Chen, Ningyu Zhang
TL;DR
Knowledge-driven tasks depend on external knowledge, while previous work often overlooks the implicit knowledge stored in pre-trained models. Kformer injects knowledge through Transformer FFN layers, and reports absolute improvements on Social IQA and Med-USMLE. The conclusion presents FFN knowledge injection as complementary to knowledge already in pre-trained models.
Problem
Knowledge-driven tasks depend on external knowledge, but direct fine-tuning of pre-trained language models yields suboptimal results and prior work often neglects their implicit stored knowledge.
Source: §1 · pp. 1–2Method
Kformer is a pluggable model that injects external knowledge into Transformer FFN layers while using knowledge stored in pre-trained models.
Source: §1 · p. 2Results
Kformer achieves absolute improvements on the two multi-choice tasks Social IQA and Med-USMLE.
Source: §7 · p. 11Takeaways & Limitations
The paper concludes that knowledge injection through FFN can complement the useful knowledge already present in pre-trained language models.
Source: §1 · p. 2The evaluation covers two tasks; future work proposes exploring more downstream tasks to improve Kformer’s generalization ability.
Source: §7 · p. 11
Abstract
from arXiv · showhide
Recent days have witnessed a diverse set of knowledge injection models for pre-trained language models (PTMs); however, most previous studies neglect the PTMs' own ability with quantities of implicit knowledge stored in parameters. A recent study has observed knowledge neurons in the Feed Forward Network (FFN), which are responsible for expressing factual knowledge. In this work, we propose a simple model, Kformer, which takes advantage of the knowledge stored in PTMs and external knowledge via knowledge injection in Transformer FFN layers. Empirically results on two knowledge-intensive tasks, commonsense reasoning (i.e., SocialIQA) and medical question answering (i.e., MedQA-USMLE), demonstrate that Kformer can yield better performance than other knowledge injection technologies such as concatenation or attention-based injection. We think the proposed simple model and empirical findings may be helpful for the community to develop more powerful knowledge injection methods. Code available in https://github.com/zjunlp/Kformer.
1 Introduction
Prior knowledge-injection approaches often overlook implicit knowledge stored in pre-trained models, even though knowledge-driven tasks depend on external knowledge. Kformer injects external knowledge in Transformer FFN layers and reports absolute improvements on two benchmark tasks.
1 Introduction
- Prior approaches concatenate retrieved knowledge or use attention-based interaction, while most neglect implicit knowledge stored in pre-trained model parameters.Source: §1 · pp. 1–2
- Knowledge neurons in the FFN have been observed to express factual knowledge, motivating investigation of FFN-based external knowledge injection.Source: §1 · p. 2
- Kformer injects external knowledge into Transformer FFN layers while taking advantage of knowledge already stored in pre-trained models.It is presented as a simple, pluggable model adaptable to different knowledge types.Source: §1 · p. 2
- Kformer achieved absolute improvements on SocialIQA and MedQA, with knowledge injection through FFN described as complementary to PLM knowledge.The section reports results on these two benchmarks but gives no numerical values here.Source: §1 · p. 2
2 Knowledge Neurons in the FFN
The section describes the Transformer FFN as a two-linear-layer module whose internal activations and memory vectors are associated with factual knowledge. This motivates injecting external knowledge through the FFN to draw on both external and PTM-stored knowledge.
2 Knowledge Neurons in the FFN
- An FFN uses two linear networks: its input produces coefficients that are activated and used to compute a weighted sum of the second layer's memory vectors.The input is the Transformer layer's final attention output.Source: §2 · p. 3
- The FFN is described as an unnormalized key-value memory, and its hidden state activates knowledge neurons before the second layer integrates corresponding memory vectors.Prior work describes the memory view; Dai introduced knowledge neurons to study factual knowledge storage in pre-trained Transformers.Source: §2 · p. 3
- Because FFN knowledge neurons express factual knowledge, directly modifying corresponding FFN parameters is presented as a way to use both PTM-stored and external knowledge.Source: §2 · p. 3
3 Kformer: Knowledge Injection in FFN
Kformer retrieves potential knowledge for each question, embeds it, and fuses it into pretrained Transformer FFN layers. Its pipeline combines hybrid retrieval, averaged token embeddings, and layer-specific knowledge projections.
3 Kformer: Knowledge Injection in FFN
- Kformer retrieves the top N knowledge items per question, embeds them, and fuses them into the pretrained model through Transformer FFN layers.These are the method’s three main stages: retrieval, knowledge embedding, and FFN fusion.Source: §3 · p. 3
3.1 Knowledge Retrieval
- A hybrid retrieval system ranks candidates with sparse search, then reranks the top M using dense inner-product scores to select the top N evidence items.The paper describes retrieval as replaceable across different retrieval scenarios.Source: §3.1 · pp. 3–4
3.2 Knowledge Embedding
- Each knowledge sequence is tokenized, embedded, and represented by the average of its token embeddings; the embedding matrix starts from the PTM’s input embeddings and updates during training.The knowledge embedding layer is trained simultaneously with the model.Source: §3.2 · p. 4
3.3 Knowledge Injection
- For each layer, two linear projections map knowledge vectors into the corresponding FFN matrices, which are expanded by concatenating the projected knowledge.The projection weights are randomly initialized and updated during fine-tuning.Source: §3.3 · pp. 4–5
- The injected knowledge is activated in relation to the input sequence, and the model uses it in the top three Transformer layers.Subsequent Transformer layers process and aggregate the collected information.Source: §3.3 · p. 5
4 Experiments
The experiments evaluate Kformer on Social IQA and MedQA-USMLE using knowledge from triples and textbooks. On Social IQA, feed-forward-layer injection yields reported gains over other knowledge-injection methods.
4 Experiments
- The experiments use knowledge-graph triples and textbook documents to examine Kformer on Social IQA and MedQA-USMLE.Source: §4 · p. 5
4.1 Dataset
- Social IQA uses paraphrased ATOMIC triples as knowledge candidates, while MedQA-USMLE uses English questions from US medical board exams and textbooks as its knowledge base.Source: §4.1 · p. 5
4.2 Experiment Setting
- The two multiple-choice tasks are evaluated with accuracy; the compared injection approaches concatenate retrieved knowledge at the input or inject it into self-attention's K and V parts.Source: §4.2 · pp. 6–7
4.3 Experiments Results
- For the large model, adding knowledge in the feed-forward layer gains 0.41 on the valid set and 1.56 on the test set; it outperforms RoBERTa+MCQueen by 0.68 on the test set.Source: Table 2 · p. 6, §4.3 · p. 7
- Figure 2 reports model results across top-N retrieved-knowledge settings of 1, 5, 10, 15, 20, and 25 candidates.Source: Figure 2 · p. 7
5 Analysis
Across Social IQA and MedQA, Kformer outperforms RoBERTa+ATT, while retrieval volume and injection layer affect outcomes. An interpretability case shows retrieved knowledge supporting a correct divorce-related answer.
5 Analysis
- On both tasks, Kformer outperforms RoBERTa+ATT; for long Textbook knowledge, it beats input concatenation by 1.3 in the valid set and 0.3 in the test set.The authors suggest FFN injection can focus on more evidence than input concatenation, which may include only about ten potential evidence pieces.Source: §5 · p. 8
5.2 Impact of Layers
- Top-three-layer infusion performs best on Social IQA (72.82, +0.56 Acc) and scores 33.03 in MedQA dev and 30.01 in test; layers 4–6 degrade both tasks.Infusing knowledge into all layers helps, but scores below using the top three layers.Source: Table 4 · p. 8, §5.2 · p. 9
5.1 Impact of Top N Knowledge
- Social IQA performs best with the top 15 sentences and MedQA with the top 20 documents; retrieving more reduces performance on both tasks.Too little retrieval can leave insufficient evidence, while too much can introduce knowledge noise.Source: §5.1 · pp. 8–9
5.3 Interpretability
- In one Social IQA case, activated divorce cues imply separation without stating it explicitly, letting Kformer correct RoBERTa’s lawyer choice.The activation display maps sequence tokens to retrieved knowledge; the model focuses on retrieved items 1, 7, 11, 12, and 14.Source: §5.3 · pp. 9–10, Figure 3 · p. 10
6 Related Work
Prior knowledge-fusion approaches append retrieved knowledge to model inputs or encode it separately for attention-based interaction. Kformer instead filters and incorporates external knowledge through Transformer FFN layers.
Knowledge Fusion
- External knowledge includes structured triples or paths and retrieved open-domain text; Kformer is not limited to one knowledge type and is tested with both.Prior work also retrieves open-domain contexts with information retrieval or text matching [21].Source: §6 · p. 10
Knowledge Fusion in Pre-trained Models
- Input-fusion approaches include K-BERT and concatenation [21], while other methods encode context and knowledge separately and interact them through attention.Input concatenation is limited by sequence length; the attention-based approach predicts from concatenated context and knowledge representations.Source: §6 · pp. 10–11
- Prior knowledge-embedding approaches and may encounter knowledge noise; Kformer filters and incorporates external knowledge through Transformer FFN layers.Knowledge noise is described as arising when too much knowledge is incorporated into a statement, causing it to lose its intended meaning.Source: §6 · p. 11
7 Conclusion and Future Work
Kformer injects external knowledge through Transformer feed-forward layers and achieves absolute improvements on Social IQA and Med-USMLE. The authors identify pre-training integration and evaluation on more downstream tasks as future directions.
7 Conclusion and Future Work
- Kformer is a simple, pluggable knowledge-injection model for pre-trained language models that adapts to different knowledge sources.It injects external knowledge through Transformer feed-forward layers.Source: §7 · p. 11
- Kformer achieves absolute improvements on the two evaluated multiple-choice tasks, Social IQA and Med-USMLE.Source: §7 · p. 11
- Future work includes integrating the fusion method during pre-training and exploring more downstream tasks to improve Kformer’s generalization ability.The authors also propose investigating knowledge storage mechanisms and efficient models for knowledge-intensive tasks.Source: §7 · p. 11