KBLaM: Knowledge Base augmented Language Model

Authors: Xi Wang, Taketomo Isazawa, Liana Mikaelyan, James Hensman
Primary area: foundation or frontier models, including LLMs
Keywords: Large language model; Knowledge augmentation;

Abstract:
In this paper, we propose Knowledge Base augmented Language Model (KBLAM), a new method for augmenting Large Language Models (LLMs) with external knowledge. KBLAM works with a knowledge base (KB) constructed from a corpus of documents, transforming each piece of knowledge in the KB into continuous key-value vector pairs via pre-trained sentence encoders with linear adapters and
integrating them into pre-trained LLMs via a specialized rectangular attention mechanism. Unlike Retrieval-Augmented Generation, KBLAM eliminates external retrieval modules, and unlike in-context learning, its computational overhead scales linearly with KB size rather than quadratically. Our approach enables integrating a large KB of more than 10K triples into an 8B pre-trained LLM of only 8K context window on one single A100 80GB GPU and allows for dynamic updates without model fine-tuning or retraining. Experiments demonstrate KBLAM’s effectiveness in various tasks, including question-answering and open-ended reasoning, while providing interpretable insights into its use of the augmented knowledge. Code and datasets are available at https://github.com/microsoft/KBLaM/

Decision: accept (poster)
Reviewer scores: 8, 5, 5, 8, 3

Reviews and discussion:
<thread>
creation date:
1729071815699

thread_type:
official_review

<thread_object>
creator: reviewer_oum4

<content>
summary:
This paper introduces a new method, KBLaM, for enhancing language models with external knowledge. The main idea can be interpreted as an information compression method of in-context learning (directly putting the entire external corpus in context). Specifically, KBLaM first transforms each piece of knowledge in the knowledge base into continuous key-value vector pairs using a pre-trained sentence encoder. These vector pairs are then mapped into the language model's hidden space via linear adapters. By employing a rectangular attention mechanism, KBLaM enables the integration of external knowledge into the model. Experimental results show that KBLaM achieves performance comparable to in-context learning while offering significantly better scalability.

strengths:
1. Compared to in-context learning, KBLAM introduces linear attention computation scaling relative to the size of the knowledge base (KB). Additionally, it preserves the key advantage of in-context learning by enabling dynamic updates to the knowledge without requiring model fine-tuning.

2. The proposed method also enhances interpretability by potentially offering insights into how the model leverages knowledge through the attention matrix.

3. This paper is clearly-written and easy to follow.

weaknesses:
### 1. Limitations of the method:
- The proposed method encodes each piece of knowledge in the knowledge base separately, without leveraging the graph structure or interconnections between knowledge entities. This could limit the method's ability to capture richer, contextual relationships within the knowledge base.
- The method introduces significant computational overhead by requiring fine-tuning of multiple adapters: a key adapter ($\in \mathbb{R}^{D \times P}$), a value adapter ($\in \mathbb{R}^{D \times P}$), and a query transformation ($\in \mathbb{R}^{D \times D}$) in each transformer layer. For example, using the experimental setup described in the paper, where the hidden sizes of Llama3-8B and the pre-trained OpenAI ada-002 encoder are 4096 and 1536, respectively, the total number of trainable parameters amounts to around **1 billion**—roughly 1/8 of the base model. Compared to in-context learning (training-free), this method represents a substantial increase in resource requirements.
- The performance of the method appears to be heavily dependent on the choice of pre-trained encoder. As shown in Figure 7, using a weaker encoder can lead to a significant drop in performance, in some cases falling below that of the zero-shot baseline.

###  2. Limitations of the experiments:
- The experiments are only conducted on a single base LLM (Llama3-8B). Testing on a wider range of model architectures and sizes would help assess the generalizability of the proposed method.
- The evaluation is based on a small number of datasets, specifically a synthetic knowledge base and Enron. A broader evaluation on more diverse and real-world datasets would provide stronger evidence of the method's effectiveness.
- Missing comparison with Retrieval Augmentation Generation (RAG).

questions:
1. Could you elaborate on the design rationale behind Eq. (9)? Specifically, it seems that the output of each attention layer, $\Tilde{\bf{y}}_n^l$, could alternatively be derived by aggregating the information from both $\Tilde{\bf{v}}_m^l$ and $\bf{v}_i^l$ independently. What specific advantages does the chosen design in Eq. (9) offer over this potential alternative formulation?

2. Could you provide insights on how another competitive baseline, such as RAG, might perform under your experimental settings? Additionally, it would be interesting to know if KBLaM could offer reduced memory overhead compared to RAG.

3. How would KBLaM perform in more challenging KBQA scenarios, such as multi-hop reasoning tasks? Additionally, can KBLaM maintain its advantage over in-context learning in such complex environments?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730489563371

thread_type:
official_review

<thread_object>
creator: reviewer_9hrd

<content>
summary:
The authors implement a scheme that encodes Knowledge Base triplets into Key+Value pairs via a pretrained embedding encoder.  The information from the triplets can then be continually 'looked up' by a pretrained local LLM through the training of adapter matrices that (essentially) convert the local LLM representations into embedding space.  Experiments are performed to show that this is effective, including the training of the adapter matrices using synthetic data - which points to greater applicability of the method.

strengths:
The paper presents an interesting new KB embedding lookup technique, with some small-scale experiments.  Also included was some visualisation of the attention across the KB - which was particularly striking for the multistep case (KB items 4+6).

Oeverall, it seems like the ideas in the paper are of general interest, even though they may not be fully fleshed out as presented here.  Even if "KBLaM" isn't the end-point of this line of research, this paper points towards a line of research that could be quite fruitful : Loosely, that weight matrix adapters can be trained to reverse the embedding process using synthetic data.

It also seems that the top-k choices among the Keys are all that are necessary (particularly since the facts are likely stand-alone, rather than mutually reliant), which could mean that much cheaper approximate search could be used to select out relevant entries before doing a full dot-product.

weaknesses:
At several points, the 'linear rather than quadratic scaling' of the method is emphasised.  This seems a weak point, since the size of the tuples $M$ already exceeds the context length $N$, so $MN>N^2$ in this regime (as obliquely mentioned in L284).

The constant $C$ introduced in Eqn 11 feels like something that should have a more principled origin.

small edit: L103 "training samples is not of interests." -> "training samples is not important."

questions:
Looking at the triplet examples (in the Appendix), and the test questions, it seems that there is an extremely high character-overlap between the name-property text and the Questions being asked.  To what extent do the queries have to 'hit' the keys exactly (for instance : Asking about 'Maggie Thatcher' rather than 'Margaret Thatcher')?  Is there an identifiable effect whereby inexact matches get honed-in on subsequent layers/tokens (for instance)?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730645792931

thread_type:
official_review

<thread_object>
creator: reviewer_w1by

<content>
summary:
This paper proposes a new regime called KBLaM to integrate the entire KB into LLM with a complexity that grows linearly. The paper primarily introduces Rectangular Attention to achieve this. Experimental data shows that the method can be scaled to handle data with up to 10K triples.

strengths:
-	The paper proposes a new method for integrating KB into LLM, and theoretically, this method is one order of magnitude more efficient than the ICL method.
-	It is a neat and interesting way to map each triple into a fixed-length key-value vector pair
-	The main body of the paper is written well, and the explanations for each step of the calculation are precise

weaknesses:
-	The experimental comparison setup is not well-designed:
	No comparisons in the first experiment. How is the retrieval accuracy compares to that of BM25?
	In the last two experiments, although the authors mentioned that ICL encountered memory constraints, it also resulted in fewer data points.
-	The conclusion of the experiment is relatively ordinary. In the last two experiments, the author claims that KBLaM can reason about knowledge and KBLaM knows when it cannot answer a question. However, these two features seem to be the result of the newly proposed structure working in conjunction with the original parameters and structure in LLMs.
-	The novelty of this paper is limited. As the author points out in the appendix, the overall approach is very close to structured prompting[1] and structured attention[2]. If each triple is considered as a demonstration, could the method proposed in this article be seen as transferring structured prompting from text to triples? Additionally, although the author cites work related to structured attention in the appendix, I believe the degree of relevance to these methods is much greater than that to the work on the key-value (KV) cache mechanism and multi-modal language models, and should be cited in the main body.
[1]Hao, Yaru, et al. "Structured prompting: Scaling in-context learning to 1,000 examples." arXiv preprint arXiv:2212.06713 (2022).
[2]Cai, Tianle, et al. "Scaling In-Context Demonstrations with Structured Attention." arXiv preprint arXiv:2307.02690 (2023).
-	There are some repetitive words in the article:
	There is no need to use a large section in Chapter 3 background to introduce Self-attention layer.
	Image label and paper content overlap. For example, “uses the attention score at the 15th layer, averaged over all attention heads, as a classification score for each triple and measures the top-1 and top-5 accuracy” in Fig5 overlaps with the content introduced in the first experiment.
-	The accuracy of ICL on the Enron dataset does not seem to be demonstrated in Fig6(a).

questions:
-	As seen in Fig6, within the range where the ICL method can operate, i.e., less than 200 triples, the ICL method is generally not inferior to KBLaM in terms of accuracy. Can these gaps be narrowed?
-	How does this method compare to retrieval + ICL?
-	How does KBLaM perform on more popular datasets and different types of problems?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730719667729

thread_type:
official_review

<thread_object>
creator: reviewer_q4dg

<content>
summary:
This paper presents KBLAM, a novel approach to augment Large Language Models (LLMs) with external knowledge from a structured knowledge base (KB). The method utilizes pre-trained sentence encoders with linear adapters to transform knowledge into continuous key-value vector pairs, which are then integrated into the LLM via a specialized rectangular attention mechanism. This approach brings the following benefits: improvements in computational efficiency, dynamic knowledge updates, and interpretability compared to existing methods such as Retrieval-Augmented Generation (RAG) and in-context learning.

strengths:
1. The KBLAM method offers a significant advantage in memory efficiency by compressing the kv cache, which enables the model to scale to longer contexts.

2. KBLAM eliminates the need for separate retrieval steps, simplifying the architecture and making the model more efficient compared to RAG.

3. KBLAM allows for runtime KB updates without the need to re-pretrain the model, providing a level of flexibility in deploying LLMs. Additionally, the ability to refuse to answer when the required information is not present in the KB adds a layer of reliability, an essential trait for LLMs used in real-world applications.

4. The paper exhibits a robust methodology with meticulous experimentation. It particularly emphasizes KBLAM’s scalability, efficiency, and interpretability. The use of a rectangular attention mechanism enables the model to scale linearly with the size of the KB. Furthermore, the authors provide comprehensive ablation studies to substantiate their design choices.

weaknesses:
1. I think that the most relevant baseline for comparison with the proposed method should be prompt caching (e.g., as described in the following link). This approach converts the knowledge base into a key-value cache for subsequent use. Compared to this method, I think the advantages of the KBLAM method maybe diminished. For instance: a) it requires instruction tuning; b) the accuracy of retrieval might be substantially lower compared to prompt caching; c) the key-value cache generated by a pretrained encoder is inherently inferior to that obtained from LLMs.
the prompt caching method:
a) https://platform.openai.com/docs/guides/prompt-caching
b) https://ai.google.dev/gemini-api/docs/caching
c) https://proceedings.mlsys.org/paper_files/paper/2024/file/a66caa1703fe34705a4368c3014c1966-Paper-Conference.pdf


2. In the KBLAM method, there is no relative order among different triples, which means there are no position embeddings. In scenarios involving extremely long contexts, this could render the attention mechanism ineffective (the over-smoothing phenomenon), greatly reducing the precision of retrieval. Although the authors propose an attention score scaling method to mitigate this, it would be beneficial to conduct some analytical experiments, such as a needle-in-a-haystack experiment (https://github.com/gkamradt/LLMTest_NeedleInAHaystack), for further analysis.

3. The KBLAM method requires KB instruction tuning. Does this imply freezing the pretrained LLM and only training the adapters? According to the paper, PEFT methods were not used. Consequently, updating the KB kv cache without re-training the LLMs might result in a performance decline, potentially compromising the plug-and-play nature of the method.

4. The performance of the adapter is highly dependent on the data used for instruction tuning. A non-homogeneous kv cache encoder can lead to a significant drop in performance. Therefore, in ablation experiments comparing the pretrained encoder from other sources, it would be best to compare it with the kv cache generated by the prompt caching method, i.e. the LLM itself. This would help readers better understand the performance of the proposed method.

questions:
1. Given that the KBLAM method inevitably leads to information loss (due to the compression of the kv cache and over-smoothing in scenarios with large kb sizes), have the authors considered incorporating a retrieval module to mitigate this issue? Additionally, it might be beneficial to use retrieval scores as a position bias (for example, the kv cache closer to the prompt corresponds to a higher relevance of the triple, which aligns with the inherent characteristics of the attention mechanism and should effectively enhance the method's performance).

2. Have the authors experimented with making the method more modular, such as by using PEFT (Parameter-Efficient Fine-Tuning) methods and training only the adapters part? If so, what were the results of this approach?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730791694621

thread_type:
official_review

<thread_object>
creator: reviewer_6e7t

<content>
summary:
This paper focuses on the open-domain question-answering task and proposes the Knowledge Base Augmented Language Model (KBLaM) framework. Specifically, KBLaM augments a large language model (LLM) with an external knowledge base (KB) at inference time by first transforming the knowledge in the KB into continuous key-value vectors using pre-trained encoders, then integrating them into the LLM through a specifically designed rectangular attention mechanism. Unlike Retrieval-Augmented Generation (RAG), the proposed KBLaM does not require costly retrieval, and, unlike in-context learning (ICL), KBLaM achieves comparable performance with significantly lower computational overhead. Experiments on both a synthesized KB and a real KB dataset demonstrate the effectiveness of the proposed method.

strengths:
1. The motivation of the paper is strong. It addresses scenarios where external knowledge is required for LLMs to solve tasks and explores how to integrate that information into LLMs efficiently and effectively. This paper identifies the limitations of current methods, including long in-context learning (ICL) and retrieval-augmented generation (RAG), and proposes a new solution that mitigates these limitations.
2. The proposed KBLaM, to my knowledge, is novel and offers the following advantages:
- It transforms external unstructured knowledge into a structured KB with knowledge triplets, which compresses lengthy context information and facilitates convenient integration of knowledge into LLMs.
- It encodes knowledge triplets as key-value vectors of uniform size, enabling straightforward incorporation into LLMs. The structured nature of the triplets also allows efficient retrieval of relevant knowledge within the KB.
- Each knowledge triplet can be encoded and incorporated into LLMs independently, providing greater scalability and lower complexity than ICL. This approach also offers flexibility for updating the knowledge within the KB.
- It also allows high interpretability via viewing the attention matrix.
3. Overall, this paper provides solid evidence through extensive experiments and ablation studies to demonstrate the effectiveness of the proposed method. The writing is also clear and easy to follow.

weaknesses:
1. While this paper includes experiments with both synthetic and real-world KB datasets, the baselines are limited to only ICL and zero-shot approaches. Recent state-of-the-art methods, such as RAG-based and long-context models, are not included as baselines, making it difficult to assess the effectiveness of the proposed method relative to these approaches. Including results from these comparisons would strengthen the evaluation.
2. The generalization of KBLaM is less well-known. While it has been tested on an OOD dataset, other commonly used QA datasets have not been explored, such as [1][2][3][4]. Please see the questions below.

---
[1] Natural Questions: A Benchmark for Question Answering Research, Kwiatkowski et al., 2019

[2] TriviaQA: A Large Scale Distantly Supervised Challenge Dataset for Reading Comprehension, Joshi et al., 2019

[3] MuSiQue: Multihop Questions via Single-hop Question Composition, Trivedi et al., 2022

[4] When Not to Trust Language Models: Investigating Effectiveness of Parametric and Non-Parametric Memories, Mallen et al., 2023

questions:
1. In Figure 6(a), why are the in-context learning results (orange line) missing from the Enron dataset?
2. I’m interested in the generalization of the proposed method. For other QA datasets and benchmarks commonly used by the community, such as [1][2][3][4], is it possible to extend KBLaM to address these problems? How can a KB be constructed for them? Do you need to build a KB for each dataset? What is the size (number of knowledge triplets) of the KB that is sufficient or suitable to solve these problems?
</content>
</thread_object>

</thread>
<thread>
creation date:
1732159434318

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
General response regarding KBLaM v.s. RAG

comment:
We would like to thank again to all the reviewers for their helpful and constructive comments!

One common question raised by the reviewers is: *How does KBLaM compare with RAG*?

Before diving into the comparison, we would to emphasize that RAG is a solid and robust method powered by more than 50 years of information retrieval research.
KBLaM on the other hand, "isn't the end-point of this line of research" (pointed out by reviewer 9hRD), but is just *the very first paper* exploring a new possibility of knowledge augmentation.

As such, "KBLaM replaces RAG" is NOT a claim or message of our submission, instead, we aim to answer two fundamental questions:
- Can we augment knowledge as continuous vectors (i.e. a modality) rather than as discrete strings?
- Can we perform retrieval *without an external retriever* but using LLM's attention only?

The answer given by KBLaM is: Yes. Under the settings considered, we successfully encode external knowledge triples as "pre-trained LLM readable" fixed length vectors, and the LLM's attention indeed demonstrates highly interpretable retriever behavior (Fig. 4)

Nevertheless, we agree with reviewer 6e7T on the importance of having a comprehensive comparison between the two approaches, which we analyze from the following four aspects:

**System optimization**

All systems require tuning for optimal performance, same do RAG and KBLaM. 

- RAG involves tuning *hyperparameters* such as the chunk size, and the value of K when retrieving the top-K candidates. Notice that these parameters are not trivially optimizable through backpropagation in that: 1. The retrieving process is a discrete operation; 2. The retrievers are external to the LLM and potentially heterogeneous to LLM or non-differentiable.

- KBLaM's "retriever" is essentially the LLM's attention, which can be optimized in a fully differentiable and end-to-end way through backpropagating from any type of LLM's output loss, e.g. cross-entropy loss or RLHF rewards, to the weights.

**Inference latency**

Both RAG and KBLaM require data preprocessing: RAG requires pre-chunking text and storing the index; KBLaM requires pre-computing and caching the key and value.

Assuming the preprocessing is done, given a KB composed of $M$ triples, where each triple takes $T$ tokens, and a user query that is $N$ tokens long, the asymptotic inference time complexity for the two methods is

RAG: 
- Retrieve the top-K relevant documents ($O(M)$).
- Tokenize the retrieved documents ($O(KT)$)
- LLM synthesis generation ($O((KT + N)^2)$ from attention, $O(KT + N)$ from FFNs)

KBLaM:
- LLM synthesis generation ($O((M+N)N)$ from attention, $O(N)$ from FFNs.)

*FFN --> Feed-forward networks, also referred to as MLPs*

While RAG's actual latency depends heavily on retriever implementation, we consider an idealized scenario with **highly optimized** retrieval systems and **negligible IO/network/API overhead**, where RAG's overhead only comes from the LLM. Assuming K=5 retrieved triples, our measurements averaged over 40 queries show RAG has "time to first token" of 0.049 seconds, 12% slower than KBLaM's 0.042 seconds with 4,096 triples in the context (full results in Fig.11, appendix).

Interestingly, KBLaM's latency *stays almost constant* across KB sizes. We believe this is because high-FLOP operations in an LLM have input sizes *constant* with respect to KB size in KBLaM:

1. KBLaM's QKV transformation overhead is constant with regard to KB size: All the keys and values for the KB are pre-computed (as we discussed earlier in the paragraph), unlike RAG, which requires computing keys and values for the $K$ triples on the fly.

2. KBLaM's knowledge tokens (i.e. the KB parts) are **not** passed through the FFNs, so the overhead of FFNs is fixed, unlike RAG, whose input to the FFNs scales with the value of $K$.

**Memory overhead**

Similiar to the inference latency, RAG's memory overhead depended on the value of $K$, so it will be fixed as $O((KT + N)^2)$. 

KBLaM's memory overhead, $O((M+N)N)$, grows linearly with the KB size.

In our empirical evaluation (Fig.11 in the appendix), we find that given $K=5$, RAG has the same memory overhead as KBLaM with $M=512$.

**Question answering quality**

On the settings we considered, BM25 as a retriever shows almost perfect top-5 accuracy in retrieving relevant triples to the context, therefore RAG would give highly accurate answers similar to in-context learning.

However, in case the question does not have an answer in the KB, RAG's behavior becomes undeterminable: If there is no explicit "unfound detection" mechanism built into the retriever, the pipeline would retrieve wrong triples into the context and LLM may use its own knowledge to synthesize the answer, potentially causing hallucination.
On the other hand, thanks to its favorable memory and latency characteristics, KBLaM can fit the whole KB into the context and is explicitly instruction-tuned to perform refusal when facing such questions.
</content>
</thread_object>

</thread>
<thread>
creation date:
1733260419695

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Final author general response

comment:
# General response

Dear reviewers and AC

We would like to thank again to all the reviewers for their detailed responses and constructive discussion!

To summarize, our submission proposes KBLaM, a new method of augmenting external knowledge represented as a knowledge base (KB) into a pre-trained LLM utilizing the key-value structure in the KB.
KBLaM has the following contributions:
- We show that each triple in the KB can be presented to an LLM as a fixed-length continuous vector rather than discrete tokens.
- We show that LLM's attention can demonstrate an **interpretable** retriever behavior, i.e. it attends accurately over the relevant vector to answer a query (Figure 4 and Figure 5).
- KBLaM allows for augmenting more than 10K triples (~200K tokens of information) into the context of a pre-trained LLM on a single A100 GPU.

The novelty and value of our contributions was recognized by the reviewers as

"...a promising step forward in knowledge augmentation..." (9e7T)

"...this paper points towards a line of research that could be quite fruitful..." (9hRD)

The high-level takeaway message we would like to convey through this work is that:

*If you are attempting to augment a pre-trained LLM with **structured knowledge**, making use of the structure brings you scalability, interpretability, and a principled model that follows the inherent properties of the data.*

KBLaM's design is motivated by the goal of developing a knowledge augmentation strategy that *respects and utilizes* the structure in the knowledge base and the triples, as we discuss in more detail in the introduction section.
</content>
</thread_object>

</thread>


Metareview:
metareview:
This paper focuses on the open-domain question-answering task and proposes the Knowledge Base (KB) Augmented Language Model (KBLaM) framework. Specifically, KBLaM transforms the KB  into continuous key-value vectors using pre-trained encoders, then integrating them into the LLM through an attention mechanism. Unlike Retrieval-Augmented Generation (RAG), the proposed KBLaM does not require retrieval, and, unlike in-context learning (ICL), KBLaM achieves comparable performance with significantly lower computational overhead. Experiments on both a synthesized KB and a real KB dataset demonstrate the effectiveness of the proposed method.

Reviewers generally liked the proposed methods both in motivation and in its potential advantages  (6e7T, q4dg, W1bY, 9hRD, ouM4), and selecting an illuminating set of experiments (6e7T, q4dg). They raised questions about baselines, especially RAG baselines (6e7T, W1bY, ouM4) and that evaluation task on synthetic data / Enron is not the most standard or competitive (6e7T, ouM4) or that evaluating on 1 model is not sufficient (ouM4). Overall the authors addressed most reviewer concerns, in particular convincing comparisons with RAG. The authors also addressed several concerns from the main reject reviewer (ouM4) about resource consumption.

Overall the paper should be accepted. For a method paper, addressing the remaining concerns of reviewer ouM4 could include significant trade-offs in the choice of experiments. Still, for a stronger paper, it could benefit from some experiments on standard competitive tasks (6e7T, ouM4).