Research map / Large language models / Retrieval and knowledge
KBLaM: Knowledge Base augmented Language Model
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes KBLaM, which encodes knowledge-base triples as continuous key-value "knowledge tokens" via a frozen sentence encoder plus learned linear adapters and injects them into a frozen LLM through rectangular attention, trained by instruction tuning on synthetic KB data.
Problem
Augmenting pre-trained LLMs with large external knowledge bases without retrieval modules, quadratic in-context cost, or per-update fine-tuning.
Approach
Each KB triple becomes a fixed-length key-value pair (one knowledge token) from a pre-trained sentence encoder with linear adapters; prompt tokens attend to all knowledge tokens via a rectangular attention pattern in which knowledge tokens do not attend to each other, so cost scales linearly with KB size and the KB can be updated per triple without fine-tuning the LLM.
Key results
- KBLaM with an 8B Llama-3 backbone integrates over 10K triples on a single 80GB A100, with performance comparable to in-context learning on synthetic data at much lower memory.
- Attention at layer 15 gives top-5 retrieval accuracy close to BM25 on the synthetic KB, and outperforms BM25 on perturbed entity names.
- On Enron (out-of-distribution), KBLaM degrades relative to synthetic data but still beats the zero-shot baseline.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use KBLaM instead of in-context learning | Llama3 8B · Synthetic KB simple and two-entity Q&A | BERT score (F1) | No clear effect |
| Use KBLaM instead of zero-shot learning | Llama3 8B · Enron Q&A | BERT score (F1) | Helped |
| Evaluate on Enron rather than synthetic KB data | KBLaM · Enron Q&A | BERT score (F1) | Hurt |
| Use KBLaM instead of BM25 | KBLaM · Synthetic KB retrieval | Top-5 accuracy | No clear effect |
| Use KBLaM instead of BM25 | KBLaM · Perturbed synthetic KB retrieval | Retrieval accuracy | Helped |
| Use an open-source sentence encoder instead of OpenAI ada-002 (ablation) | KBLaM · Q&A | Answer quality | Hurt |
| Use a lower-dimensional sentence encoder (ablation) | KBLaM · Q&A | Answer quality | Hurt |
| Add knowledge tokens every 3 layers instead of every layer (ablation) | KBLaM · Simple and two-entity Q&A | BERT score (F1) | No clear effect |
| Add knowledge tokens every 3 layers instead of every layer (ablation) | KBLaM · Unanswerable questions | Precision | Hurt |
| Add knowledge tokens every 10 layers instead of every layer (ablation) | KBLaM · Unanswerable questions | Refusal behavior | Hurt |
| Add knowledge tokens every 10 layers instead of every layer (ablation) | KBLaM · Q&A | Answer quality | Hurt |
| Add knowledge tokens every 10 layers instead of every layer (ablation) | KBLaM · Q&A | Efficiency | Helped |
Surprising findings
- Using only linear adapters trained on fully synthetic GPT-generated data generalizes to real Enron emails, despite the value text distribution differing substantially.
- With knowledge tokens added only every 10 layers, the model fails to produce refusal answers, suggesting earlier-layer knowledge tokens act as soft instruction prompts.
Limitations
- Information loss: each triple is compressed into a fixed-length vector, so exact names and numbers may not be reproduced.
- Performance degrades on out-of-distribution data (Enron) relative to the synthetic KB.
- Non-trivial one-time training cost (about 24-48 hours on a single 80GB GPU).
Experimental setup
Datasets: synthetic KB generated by GPT (about 135K triples) Enron email-derived KB
Benchmarks: synthetic KB validation set Enron KB
Metrics: BERTScore F1 GPT-4 score (0-5) top-1 and top-5 retrieval accuracy precision and recall for unanswerable questions
Baselines: in-context learning zero-shot learning BM25 retrieval-augmented generation (5 triples retrieved)
Builds on: retrieval-augmented generation in-context learning KV cache
Closest papers by meaning
- MKGL: Mastery of a Three-Word Language (NeurIPS 2024 · similarity 0.88)
- A Systematic Exploration of Knowledge Graph Alignment with Large Language Models in Retrieval Augmented Generation (AAAI 2025 · similarity 0.87)
- Harnessing Large Language Models for Knowledge Graph Question Answering via Adaptive Multi-Aspect Retrieval-Augmentation (AAAI 2025 · similarity 0.87)
- Knowledge Card: Filling LLMs' Knowledge Gaps with Plug-in Specialized Language Models (ICLR 2024 · similarity 0.87)
- Simple is Effective: The Roles of Graphs and Large Language Models in Knowledge-Graph-Based Retrieval-Augmented Generation (ICLR 2025 · similarity 0.87)
- RuAG: Learned-rule-augmented Generation for Large Language Models (ICLR 2025 · similarity 0.86)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.