KBLaM: Knowledge Base augmented Language Model

## Abstract

Abstract In this paper, we propose K nowledge B ase augmented La nguage M odel ( 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.

## 1 Introduction

Large language models (LLMs) have demonstrated impressive knowledge and reasoning capabilities. However, in many scenarios, users need to augment LLMs with external knowledge, particularly when concept definitions differ from or extend beyond the information stored in the LLM’s parameters. One straightforward approach is to perform supervised fine-tuning on the external corpus (Hu et al., ; Liu et al., ) to alter the parameter weights. However, fine-tuning can be inefficient as the weights need to be updated each time knowledge is updated. Additionally, fine-tuning can cause catastrophic forgetting, i.e. degrading performance on general tasks.
As such, alternative methods that do not modify LLM’s weights have gained more and more popularity.

One such alternative approach is Retrieval Augmentation Generation (RAG, Lewis et al., ). RAG assumes external knowledge is stored as unstructured documents, then given an input prompt, RAG first uses a retriever module to extract relevant corpus chunks from the documents. These extractions are then concatenated to the input question as a prompt and fed into the LLM (Fig. [1](#S1.F1) top). This approach allows RAG to utilize knowledge from extensive documents while maintaining a small context size, overcoming the context limitations of many early-generation LLMs.

Recently, with the emergence of long-context language models such as GPT4 (Achiam et al., ) and Gemini (Team et al., ), it has become possible to directly put the entire external corpus in the context, eliminating the need for the retriever for document selection, simplifying the pipeline. We refer to this approach as *in-context learning* (Fig. [1](#S1.F1) middle).
Several recent works (Lee et al., ; Li et al., ; Yu et al., ) have explored whether long-context models can replace RAG in long-context reasoning tasks. In particular, Lee et al. () argues that in-context learning has potential advantages, such as better compatibility with chain-of-thought reasoning and few-shot demonstrations. However, the time and memory complexities of in-context learning are quadratic with the context length. Additionally, the dependencies between tokens in the context introduce difficulties in interpreting the attention matrix and dynamically updating the KV cache.

In this paper, we propose a new regime for augmenting pre-trained LLMs with external knowledge:
Knowledge Base augmented Language Model (KBLaM, Fig. [1](#S1.F1) bottom).
We consider a setting where the *unstructured* external corpus is transformed into a *structured* knowledge base (KB) through existing tools, which summarizes the critical information in the data and generates knowledge triples ([Eq. 1](#S3.E1)) containing an entity name (<name>), a property (<property>), and a value (

).
Then, KBLaM’s design fully leverages the structure of the KB, achieving efficient integration of external knowledge into LLMs (Sec. [4](#S4)).

First, KBLaM maps each triple into a *fixed-length* key-value vector pair, referred to as a *knowledge token*, which has sizes identical to the KV cache of a single token and can be seamlessly incorporated into each attention layer of the LLM. In particular, KBLaM encodes from <name> and <property> into a key vector, which serves as an identifier, mimicking the key embedding of a token;

into a value vector, which provides the actual content, similar to the value embedding of a token.

Then, KBLaM uses the structure between triples to augment knowledge tokens into the LLM’s attention in a scalable and dynamic way using a simple rectangular attention structure:
Triples with different <name> and <property> can be considered to represent independent pieces of information, therefore knowledge token from each triple can be encoded and injected into pre-trained LLM independently. This allows the KBLaM’s complexity (memory and time) to grow linearly with respect to the number of triples, unlike in-context learning’s quadratically growing overhead, giving KBLaM much better scalability.
Additionally, the independence allows us to update/remove/add a triple by only modifying its corresponding single knowledge token without any further changes, which is not achievable in, e.g. standard KV cache mechanism.
Perhaps more importantly, the attention matrix under this design is highly interpretable. As we show in Fig. [5](#S5.SS0.SSS0.Px1), the model’s use of the knowledge tokens can be directly inspected through the attention score.

Figure: Figure 1: Overview of the KBLaM pipeline and comparison with existing approaches. KBLaM augments knowledge into a pre-trained LLM in the form of knowledge tokens, a set of continuous key-value vectors, using a modified rectangular attention structure. Unlike RAG, KBLaM does not rely on separate retriever module at inference time and unlike in-context learning, KBLaM’s computation and memory overhead scales *linearly* rather than quadratically with the size of the KB.
Refer to caption: /html/2410.10450/assets/x1.png

Lastly, we show that the linear adapters can be learned using instruction tuning on purely synthetic data (Sec. [5](#S5)) while being able to generalizes to real data. Different from supervised fine-tuning, which aims to memorize knowledge into model weights, the learning process of KBLaM aims at finding a projection between the pre-trained sentence encoder space and the LLM’s embedding space, which motiviates us the train KBLaM with fully synthetic data since the exact knowledge in the training samples is not of interests.

Compared with RAG, KBLaM incorporates all external information in the context and answers input prompts in an end-to-end way without utilizing additional external modules, similarly to in-context learning. Indeed, KBLaM also shows empirical performance comparable to in-context learning but with significantly better scalability to more than 10K triples.
In addition, through instruction tuning, KBLaM can learn to refuse to provide an answer if the information required to answer the question is not present in the KB, improving the reliability of the model and reducing hallucinations.

In summary, KBLaM is a novel method for augmenting pre-trained LLMs with external knowledge bases (KBs). We demonstrate that KB triples can be efficiently encoded into ”knowledge tokens” - continuous key-value vector pairs equivalent in size to one LLM token - using a pre-trained sentence encoder with a learned linear adapter. These tokens are incorporated into the LLM through a modified attention structure, enabling linear scaling with KB size and allowing dynamic knowledge updates without model fine-tuning. The adapters are trained end-to-end using instruction tuning on synthetic data. Lastly, we release our training and evaluation KBs, which can help future research in augmenting LLM with KB, as well as in other topics such as long-context language models, hallucination detection/reduction, and structured attention.

## 2 Related work 1 1 1 We provide more extensive related work discussion in Appendix A

Retrieval augmented generation (RAG)  RAG (Lewis et al., ) is one of the most successful approaches for augmenting external knowledge into LLMs, and many approaches have been proposed to perform RAG, such as similarity search and text2Sql (Qin et al., ). KBLaM can also be understood as the model doing RAG in an implicit way: given input prompts, the rectangular attention compares the sequence’s queries with knowledge tokens’ keys, then the weighted average is taken of all knowledge tokens’ values based on query-key similarity, similarly to a soft retrieval process. The averaged value embedding is then added to the hidden state, augmenting next-token generation with information from the KB. However, KBLaM differs from RAG in that it does not rely on a separate retrieval module. Instead, everything is carried out by attention in an end-to-end fashion.

Multi-modal language model  Recent advances in Multi-modal Language Models (Liu et al., ; Gao et al., ; Liu et al., ; Zhao et al., ; Wu et al., , MLMs), demonstrate the possibility of incorporating data of other modalities (e.g. images, videos) into a pre-trained LLM through learned adapters and instruction tuning. With KBLaM, we can consider a KB as another modality in that a KB is presented to the LLM in the form of continuous vectors via an encoder, and an encoder-adapter is learned also through instruction tuning inspired by MLM literature.

Key-value (KV) cache mechanism  Given a context string, the KV cache mechanism (Pope et al., ) caches the key and value embeddings at each layer, also known as the prefill stage. With KV cache, the complexity of generating new tokens conditioned on the context decreases from quadratic in context length to linear. KBLaM’s encoded knowledge tokens work in a way similar to KV cache, but we have the key and value acquired from an external encoder rather than through self-attention process. Additionally, if the context is modified, the standard KV cache requires re-computing the whole KV cache due to the causal attention mask, whereas in KBLaM only the corresponding knowledge token need be updated as the knowledge tokens do not attend to each other.

Augmenting information into an LLM in a continuous form  RAG and in-context learning incorporates external information into LLMs in the form of a *discrete string*, whereas KBLaM augments information into LLMs as fixed length *continuous* representations. Many works also consider first encoding external documents as continuous vectors, then concatenating the vectors as the input for the decoder model (Izacard & Grave, ; Izacard et al., ; Ye et al., ) or using cross-attention to inject these vectors into the decoder (Borgeaud et al., ).
However, these works focus on training a new model from scratch, while KBLaM is applied to augment a *pre-trained* LLM with knowledge. Recent work (Yen et al., ) also considers augmenting a *pre-trained* LLM with encoded information, similar to our setting. However, their approach requires training an encoder *from scratch*, whereas KBLaM’s only trains linear adapters after a pre-trained black-box encoder. In addition, they add extra cross-attention layers after each self-attention layer for augmentation, while KBLaM modifies existing self-attention layers without introducing new components.

## 3 Background

#### Knowledge base in the form of triples

In this paper, we assume that the external knowledge is represented as a structured knowledge base (KB) obtained from unstructured text using, in particular, a KB is composed of triples of format

$$ $\{(\texttt{<name>}_{m};\texttt{<property>}_{m};\description_m)\}_{m=1}^{M}. \end{equation}Intherestofthepaper,wewillrefertoasetoftriplesoftheformat\lx@cref{creftype~refnum}{eq:triple_def}asa\textsc{KB}andeachtripleinsideitasa\emph{knowledge triple}.ExamplesofsuchKBisshowninTable.~{}\ref{tab:syn_kb}and~{}\ref{tab:enron_kb}inAppendix~{}\ref{sec:sample_kb}.TheKBconstructionprocesssummarizesinformationfrommultipledocumentsandorganizesitinastructuredway.Importantly,thestructurednatureofthe\textsc{KB}enablesustoaugmentLLMwithinformationfromdocumentsinanefficientway,whichwediscussinthenextsection.\par Itisalsoworthnotingthat\textsc{KB}constructionis\emph{not}thefocusofourpaper:\textsc{KBLaM}workson\textsc{KB}generatedfromcorpususing\emph{existing tools}.Inthepaper,wemainlyfocusontwo{\textsc{KB}}s:Asynthetic\textsc{KB}generatedviaGPT,andareal\textsc{KB}constructedfromtheEnronemail~{}\cite[citep]{(\@@bibref{AuthorsPhrase1Year}{klimt2004enron}{\@@citephrase{, }}{})}dataset.Detailsofthetwo{\textsc{KB}}swillbediscussedindetailinlatersections.\par\par\par$ (1) $$

#### Self-attention layer

A decoder-based transformer is mostly composed of multiple self-attention layers (Vaswani, ).
Consider a transformer of $L$ layers.
Each layer has three projection heads: $\bm{W}_{Q}^{l}\in\mathbb{R}^{D\times D},\bm{W}_{K}^{l}\in\mathbb{R}^{D\times D},\bm{W}_{V}^{l}\in\mathbb{R}^{D\times D},l\in\{1,\ldots,L\}$,
where $D$ denotes *embedding dimension*.
Given a user prompt (e.g. a question) of $N$ tokens, each attention layer takes a sequence of $N$ token embeddings $\bm{x}^{l}=[\bm{x}_{1}^{l},\ldots,\bm{x}_{n}^{l},\ldots\bm{x}_{N}^{l}]^{\top}\in\mathbb{R}^{N\times D}$ as input,
then for the $n$th token embedding, it is first transformed by the projection heads into three vectors

$$ $\displaystyle\bm{q}_{n}^{l}$ $\displaystyle=\bm{W}_{Q}^{l}\bm{x}_{n}^{l},$ $\displaystyle\bm{k}_{n}^{l}$ $\displaystyle=\bm{W}_{K}^{l}\bm{x}_{n}^{l},$ $\displaystyle\bm{v}_{n}^{l}$ $\displaystyle=\bm{W}_{V}^{l}\bm{x}_{n}^{l},$ (2) $$

where we refer to $\bm{q}_{n}^{l}$, $\bm{k}_{n}^{l}$, and $\bm{v}_{n}^{l}$ as the $n$th token’s query, key, and value embedding at the $l$th layer respectively. Note that in practice, each attention layer often has multiple sets of attention heads but here we assume only a single head for notation simplicity.

The output of the attention layer, $\bm{y}^{l}=[\bm{y}_{1}^{l},\ldots,\bm{y}_{n}^{l},\ldots,\bm{y}_{N}^{l}]^{\top}\in\mathbb{R}^{N\times D}$, is computed as

$$ $\bm{y}_{n}^{l}=\frac{\sum_{i=1}^{n}\exp\left(w_{n,i}\right)\bm{v}_{i}^{l}}{\sum_{i=1}^{n}\exp\left(w_{n,i}\right)},~{}\text{where}~{}w_{n,i}=\langle\bm{q}_{n}^{l},\bm{k}_{i}^{l}\rangle/\sqrt{D}$ (3) $$

where $\langle\cdot,\cdot\rangle$ denotes the inner product of two vectors. After $\bm{y}_{n}^{l}$ is acquired, it is often fed into a feedforward network (FFN) with output dimension of $D$ for further transformation.

A standard implementation of self-attention would have a time complexity of $\mathcal{O}(N^{2}D)$ and memory complexity of $\mathcal{O}(N^{2})$ for computing and storing all $w_{n,i}$. Additionally, the FFN would also introduce a significant computation overhead of magnitude $\mathcal{O}(ND^{2})$.
Due to these factors, as the sequence gets longer, self-attention faces high memory consumption and slow computation.

Figure: (a) Step 1. KB encoding For each triple from a KB ([Eq. 1](#S3.E1)) of $M$ triples, we first construct a key and value string (left most white boxes). Then both strings are processed by a pre-trained sentence encoder followed by a learned linear adapter (middle two blue boxes). The acquired knowledge tokens ([Eqs. 6](#S4.E6) and [7](#S4.E7), right-most small boxes) are then stored on disk for later usage. Note that all $\bm{\tilde{k}}_{m}^{l}$s and $\bm{\tilde{k}}_{m}^{l}$s have a fixed dimension of $D$, i.e. identical to that of the key and value embeddings of a token, regardless of input length.
Refer to caption: /html/2410.10450/assets/x2.png

## 4 Augmenting LLM with the KB

In this section, we discuss how KBLaM augments an LLM with a KB.
The process has two steps (visualized in Fig. [2](#S3.F2)): 1. We convert each triple in the KB from the string form into a continuous key-value vector pair, referred to as a *knowledge token*, through a pre-trained sentence encoder followed by linear adapters; 2. All knowledge tokens are then injected into each attention layer of an LLM through a *rectangular attention* structure.

#### Knowledge tokens

Provided a KB in the format of [Eq. 1](#S3.E1), for each triple, we first adopt a *pre-trained* sentence encoder model, denoted as $f(\cdot)$ which converts a string into a $P$-dimension continuous embedding vector.
Through the encoder, we convert each triple into a base key embedding and a base value embedding.
In particular, for the $m$th triple, we have:

$$ (4) $$

Next, we introduce a linear key and a linear value adapter

$$ $\displaystyle\bm{\tilde{W}}_{K}\in\mathbb{R}^{L\times D\times P},$ $\displaystyle\bm{\tilde{W}}_{V}\in\mathbb{R}^{L\times D\times P},$ (5) $$

where $L$ denotes the number of attention layers in a model.
With the adapters, we map $\bm{\mathsf{{k}}}_{m}$ and $\bm{\mathsf{{v}}}_{m}$ from the sentence encoder’s space to LLM’s key and value embedding space at each attention layer.
In particular, for the $m$th knowledge triple, we transform its base key and value embedding into

$$ $\displaystyle\tilde{\bm{k}}_{m}=[\tilde{\bm{k}}_{m}^{1},\ldots,\tilde{\bm{k}}_{m}^{l},\ldots,\tilde{\bm{k}}_{m}^{L}]^{\top}=\bm{\tilde{W}}_{K}\bm{\mathsf{{k}}}_{m}\in\mathbb{R}^{L\times D},$ (6) $\displaystyle\tilde{\bm{v}}_{m}=[\tilde{\bm{v}}_{m}^{1},\ldots,\tilde{\bm{v}}_{m}^{l},\ldots,\tilde{\bm{v}}_{m}^{L}]^{\top}=\bm{\tilde{W}}_{V}\bm{\mathsf{{v}}}_{m}\in\mathbb{R}^{L\times D}.$ (7) $$

Since each $\tilde{\bm{k}}_{m}^{l}$ and $\tilde{\bm{v}}_{m}^{l}$ share the same size with LLM’s key and value embedding (i.e. $\bm{k}_{n}^{l}$ and $\bm{v}_{n}^{l}$ in [Eq. 2](#S3.E2)), they can be directly incorporated into attention’s computation. Therefore, the encoding process is equivalent to converting a knowledge triple from a string of multiple tokens into a special token, whose key and value embeddings at each layer are generated through an encoder, rather than through self-attention. As such we refer to $(\tilde{\bm{k}}_{m},\tilde{\bm{v}}_{m})$ pair as a *knowledge token*.

We visualize the process in Fig. [2(a)](#S3.F2.sf1).
This encoding process is applied to all triples in the KB, which transforms the information from a KB into a collection of knowledge tokens

$$ $\{(\texttt{<name>}_{m},\texttt{<property>}_{m},\description_m)\}_{m=1}^{M} \xrightarrow[]{\text{Encode}} \{(\tilde{\bm{k}}_m , \tilde{\bm{v}}_m)\}_{m=1}^M. \end{equation}\par$ (8) $$

#### Rectangular Attention: Injecting knowledge token into prompt tokens

After converting a KB into a collection of knowledge tokens we inject the information of the KB into the each attention layer of a pre-trained LLM through a modified attention structure(^3^33We use symbols with a tilde to denote additional attention components introduced by KBLaM.).
In particular, consider an input prompt of $N$ tokens, the $l$th attention layer takes in $\bm{x}^{l}=[\bm{x}_{1}^{l},\ldots,\bm{x}_{n}^{l},\ldots,\bm{x}_{N}^{l}]^{\top}$ and the $l$th dimension of the knowledge tokens $\{(\bm{\tilde{k}}_{m}^{l},\bm{\tilde{v}}_{m}^{l})\}_{m=1}^{M}$, the output of each attention layer, $\tilde{\bm{y}}^{l}=[\tilde{\bm{y}}_{1}^{l},\ldots,\tilde{\bm{y}}_{n}^{l},\ldots,\tilde{\bm{y}}_{N}^{l}]^{\top}\in\mathbb{R}^{N\times D}$, is given by

$$ $\displaystyle\tilde{\bm{y}}_{n}^{l}$ $\displaystyle=\frac{{\sum_{m=1}^{M}\exp\left(\tilde{w}_{n,m}^{l}\right)\bm{\tilde{v}}_{m}^{l}}+\sum_{i=1}^{n}\exp\left(w_{n,i}^{l}\right)\bm{v}_{i}^{l}}{{\sum_{m=1}^{M}\exp\left(\tilde{w}_{n,m}^{l}\right)}+\sum_{i=1}^{n}\exp\left(w_{n,i}^{l}\right)},$ (9) $$

where

$$ $\displaystyle\tilde{w}_{n,m}^{l}=\langle\tilde{\bm{q}}_{n}^{l},\bm{\tilde{k}}_{m}^{l}\rangle/\sqrt{D},~{}\tilde{\bm{q}}_{n}^{l}=\bm{\tilde{W}}_{Q}^{l}\bm{x}_{n},$ $\displaystyle w_{n,i}^{l}=\langle\bm{q}_{n}^{l},\bm{k}_{i}^{l}\rangle/\sqrt{D},$ (10) $$

and $\bm{\tilde{W}}_{Q}^{l}\in\mathbb{R}^{D\times D}$ is a *learnable* query head initialized from $\bm{W}_{Q}^{l}$.

Figure: Figure 3: Memory overhead of different methods. Given a KB of $M$ triples (with each triple $K$-token long on average). In-context learning’s memory scales with $(KM)^{2}$, whereas KBLaM’s memory scales with $M$.
Refer to caption: /html/2410.10450/assets/x4.png

We name the attention structure of [Eq. 3](#S3.E3) *rectangular attention* as the attention matrix is of rectangular shape (bottom-right panel in Fig. [2(b)](#S3.F2.sf2)).
The idea is that tokens in the prompt part can attend to, in addition to all previous tokens in the prompt, all knowledge tokens. On the other hand, knowledge tokens cannot attend to each other, i.e. they do not have self-attention nor query embeddings. This design gives an attention matrix of rectangular shape of size $(M+N)\times N$, where $M$ is the number of triples in the KB.
Additionally, if knowledge tokens are not introduced, i.e. $M=0$, [Eq. 9](#S4.E9) falls back to [Eq. 3](#S3.E3) and we recover the pre-trained LLM.

The memory and time complexity of the rectangular attention is $\mathcal{O}((M+N)N)$ and $\mathcal{O}((M+N)ND)$ respectively. In real settings, it is likely that $M>>N$, i.e. the number of external knowledge items is far larger than the size of the prompt/question, then the overhead would grow only linearly with $M$ (instead of quadratically), which allows KBLaM to scale up to very large values of $M$, enabling putting a large amount of information in the context.
Also, note that the output sequence length of rectangular attention does not vary with the number of knowledge tokens, therefore the overhead for the intermediate FFN modules stays unchanged regardless of $M$.

Lastly, it can be seen from  [Eq. 3](#S3.E3) that changing the order of the knowledge tokens does not affect the output, therefore KBLaM does not suffer from the positional bias issue of in-context learning (Liu et al., ). Additionally, since all knowledge tokens are encoded independently, when the KB is updated, e.g. new knowledge triples are created or a triple is updated, one could simply generate a new knowledge token with the encoder or update the corresponding knowledge tokens value. This highlights the difference between KBLaM and the standard KV cache mechanism in LLM, which needs re-computation when the cached content is modified in any way.

Note that KBLaM’s rectangular attention *differs* from cross-attention (Vaswani, ; Chen et al., ), which operates on a source sequence with key and value embeddings and a target sequence with only query embeddings, where tokens within each sequence do not attend to one another. In our setting, if we consider the KB as the source and the input prompt as the target, cross-attention corresponds to retaining only the summations over $M$ in [Eq. 9](#S4.E9). However, KBLaM’s attention also contains self-attention over the prompt tokens, given by the summations over $n$ in [Eq. 9](#S4.E9).

#### KB length generalization through attention score scaling

Notice that the part of the attention output corresponding to the KB would scale with $M$.
In particular, assume $\tilde{w}_{n,m}^{l},m\in\{1,\ldots,M\}$ are all of similar magnitude, then as $M$ grows, eventually the second term (the prompt/questions part) in both denominator and numerator of [Eq. 9](#S4.E9) will be overwhelmed by the first term (the KB part), and the information from the prompt/question will be lost.
To resolve this, during *inference*, we add a constant shift to the KB part’s unnormalized attention score

$$ $\tilde{w}_{n,m}^{l}=\log C-\log M+\langle\tilde{\bm{q}}_{n}^{l},\bm{\tilde{k}}_{m}^{l}\rangle/\sqrt{D},$ (11) $$

where $C\in\mathbb{R}^{+}$ is a hyperparameter which we fix as $100$ in our experiments. This is equivalent to multiplying the summations over $M$ by $\tfrac{C}{M}$, such that the total contribution of the KB part stays roughly unchanged as $M$ scales.

## 5 KB instruction tuning

The learnable parameters in KBLaM comprise the weights for the weights for the linear adapters/heads $\theta=\left\{\bm{\tilde{W}}_{K},\bm{\tilde{W}}_{V},\{\bm{\tilde{W}}_{Q}^{l}\}_{l=1}^{L}\right\}$,
Inspired by recent works in multi-modal language models (Liu et al., ), we use instruction tuning for parameter learning. Specifically, given a KB, we generate question-answer pairs about the KB (denoted as $Q$ and $A$ respectively) using formatted strings or GPT, then we optimize $\theta$ using:

$$ $\max_{\theta}\log p_{\theta,\phi}(A\mid Q,\textsc{KB}),$ (12) $$

where $\phi$ denotes parameters of the pre-trained LLM, encompassing QKV heads, FFN parameters and embedding/outpu layer weights.
Notably, KB instruction tuning preserves the base LLM’s reasoning abilities by avoiding fine-tuning the pre-trained LLM itself. Moreover, learning only the linear adapters minimizes the risk of memorizing training data, a known issue in LLM fine-tuning (Zeng et al., ; Mireshghallah et al., ).
To further mitigate this risk, we conducted instruction tuning using a *fully synthetic* KB generated by GPT. This approach is motivated by the understanding that the instruction tuning process aims *not* to memorize specific information, but to learn a projection from the pre-trained sentence encoder space to the LLM’s semantic space. Consequently, the exact content of the KB is less critical, provided the training data encompasses a diverse range of text.

#### Generation of synthetic KB

To synthesize a KB, we first use GPT to generate 50 <name>s based on combinations of 30 object types (e.g. restaurant name, software tool) and 30 idea types (e.g. natural phenomena, famous landmarks). Then for each name, we prompt GPT to generate the
Figure 4: KBLaM’s attention matrix is interpretable. Consider a toy KB of 10 triples, the heatmap shows the attention weights under different questions. We select the 15th attention layer and visualize the post-softmax attention score averaged over all attention heads. The x-axis only shows the keys corresponding the KB part and the y-axis is aligned by the end of each question’s query.

In this section, we perform empirical evaluation for KBLaM.
We begin by showing that KBLaM’s attention matrix, after instruction tuning, shows interpretable patterns and works as an accurate retriever.
We then show that KBLaM can answer questions with performance comparable to in-context learning, but with much lower memory cost, and can scale to 10K triples with little performance degradation.
Lastly, we show that for questions with no answer in the KB, the model can refuse to answer questions, with “over-refusal” behavior occurring later than in-context learning.
Lastly, we conduct ablation studies to understand the design choices of KBLaM.

### 6.1 Experiment setting

Model specification  For all experiments, we use the instruction fine-tuned version of Llama3 8B (Dubey et al., ) as the backbone LLM, and OpenAI’s ada-002 sentence embedding model ($P=1536$) as the pre-trained encoder for computing base key and value embedding ([Eq. 5](#S4.E5)).

Optimization setting
We initialize the key and value adapters: $\bm{\tilde{W}}_{K}$ and $\bm{\tilde{W}}_{V}$ randomly, and we initialize additional query heads $\bm{\tilde{W}}_{Q}^{l}$ at each layer from the pre-trained weights $\bm{W}_{Q}^{l}$.
Optimization is conducted using AdamW (Loshchilov, ) with a step size of $5\times 10^{-4}$ and a cosine learning rate decay to $5\times 10^{-6}$ for 20K iterations. Each iteration uses a mini-batch of 400 Q&A pairs, composed of 20 micro-batches of 20 samples. The instruction tuning is performed on a single 80GB A100 GPU under bfloat16 without any parameter-efficient tuning methods.

Construction of the training batches
We train KBLaM using instruction tuning, where each training sample consists of a KB, a question, and an answer ([Eq. 12](#S5.E12)). These samples are derived from the first 120K triples of the synthetic KB described in Sec. [5](#S5). To construct each training sample, we perform the following procedure:
We randomly select a subset of 10 to 100 triples from the synthetic KB to form a sample-specific KB.
Depending on the instruction type (simple, multi-entities, open-ended, or unanswerable Q&A), we designate one, multiple, or no triples as the *relevant* triple(s), with the rest serving as *distractors*.
We generate a question based on the instruction template and the relevant triple(s).
In each batch of 20 micro-batches, we include 2 micro-batches of unanswerable Q&A pairs and 6 micro-batches each for the other three Q&A types.
During training, we find limiting the KB size crucial for successful convergence. To further optimize training efficiency, we pre-compute the base key and value vectors ([Eq. 5](#S4.E5)) offline for all triples in the synthetic KB.

Evaluation dataset
For evaluation, we considered the following two KB datasets(^4^44All KBs used for training and evaluation are released together with the paper.)

- –
Synthetic data  The validation set of the synthetic KB, i.e. the 15000 triples not used for training.
- –
Enron  A KB constructed from the Enron (Klimt & Yang, ) dataset, an open-sourced corporate email dataset.
Following  (Wu et al., ), we constructed a KB of triples from the Enron emails data by fine-tuning a small language model, which we then cluster to remove duplicates  (Winn et al., ; ).

For each validation dataset, we consider <property> in “description”, “purpose” and “objectives”.

Baseline
We considered the following two methods as baselines

- –
In-context learning  Flatten all triples in a KB as strings, and attach it in front of the prompt. In-context learning’s memory overhead grows quadratically with the number of triples. Notice that we can only experiment with a maximum of 200 triples due to memory constraints.
- –
Zero-shot learning  Directly ask the LLM a particular question and provide no additional context, using the LLM’s internal knowledge for answering the questions.

Evaluation setting  For all evaluations, we repeat the experiments with 5 random seeds, each run containing 100 test samples, i.e. we randomly generate 100 KBs of various sizes and query the model with questions about the KB. The results reported are the averaged values over all 500 questions.

### 6.2 Experiment results

#### KBLaM attention is an accurate retriever

Figure: Figure 5: Through instruction tuning, the attention shows retrieval behavior. Given simple Q&A on the validation set of the synthetic data (solid line) and Enron dataset (OOD, dashed line), we use the attention score at the 15th layer, averaged over all attention heads, as a classification score for each triple and measure the top-1 and top-5 accuracy. KBLaM assigns the highest attention score to the truly relevant triple most of the time (with performance degraded but still reasonable on OOD).
Refer to caption: /html/2410.10450/assets/x5.png

Given a question about the $m$th triple in the KB, we expect certain tokens in the question, particularly those containing the keyword ${\texttt{<name>}}_{m}$, to attend more strongly to the $m$th knowledge token compared to others. This would be reflected in larger values of $\tilde{w}_{n,m}^{l}=\langle\tilde{\bm{q}}_{n}^{l},\bm{\tilde{k}}_{m}^{l}\rangle$ for these relevant tokens. We observe such patterns in practice, as visualized in Fig. [5](#S5.SS0.SSS0.Px1). In particular, we consider a KB of 10 triples and ask KBLaM questions about a single entity, two entities, or a question without an answer in the KB. We find the post-softmax attention score at the 15th layer, i.e. the middle attention layer of the 32-layer Llama3 8B, averaged over all 32 attention heads, to correlate well with the desired entity.
This property allows KBLaM to answer input questions in an interpretable way, as for any query, the top-K triples with the highest attention score form supporting evidence.
This behavior suggests that KBLaM’s attention mechanism functions implicitly as a retriever. We can quantitatively assess its retrieval performance by treating the attention scores, $\tilde{w}_{n,m}^{l},m\in{1,\ldots,M}$, as a classification score. In particular, we again extract the averaged-over-heads attention score at the 15th layer, and we evaluate the top-1 and top-5 accuracy, which represent the percentage of evaluation samples where the true triple received either the highest attention score or was among the top-5 highest scores, respectively.
Fig. [5](#S6.F5) presents these results, demonstrating that KBLaM’s attention remains highly accurate even with a large number of triples, both for in-distribution and out-of-distribution (OOD) data. Notably, this performance is achieved *solely* through instruction tuning using paired data, without any explicit regularization or retrieval objectives. In fact, imposing regularization on the attention structure would be challenging and potentially counterproductive as it is unclear, e.g. which tokens, attention heads or attention layers should be regularized.

#### KBLaM can reason about knowledge

We consider querying the model with three types of questions: Simple, two-entities (multi-entities Q&A involving two triples), and open-ended Q&A. For the first two types, we evaluate the output quality using BERT score (Zhang* et al., ). For open-ended questions, we use GPT-4 to score the output quality between 0 and 5 (full prompt presented in Appendix. [E.3](#A5.SS3)). Results are presented in Fig. [6.2](#S6.SS2.SSS0.Px4) and Fig. [6.2](#S6.SS2.SSS0.Px4).
Broadly, we can see that on the synthetic dataset, KBLaM shows performance comparable to in-context learning, but with a much smaller memory footprint (Fig. [3](#S4.F3)) and hence better scalability.
On Enron, the out-of-distribution (OOD) data, KBLaM shows degraded performance, but still performs better than zero-shot baseline, indicating that KBLaM is using the KB information effectively.
The observations from sample outputs on open-ended Q&A (Appendix. LABEL:sec:sample_output) also align with the metrics values: On synthetic KB, KBLaM gives answer almost identical to origin triple content (Appendix. LABEL:sec:syn_kb_sample_output), whereas on Enron, KBLaM’s answer captures the meaning of the original text, with more pronounced difference in wording and details (Appendix. LABEL:sec:enron_kb_sample_output).
Note that Enron has

s that look entirely different from that of the synthetic data (see Table. LABEL:tab:syn_kb v.s. LABEL:tab:enron_kb in Appendix. [C](#A3)).
We therefore believe such performance degradation could be alleviated by using a larger and more diverse synthetic dataset or through a synthetic dataset derived from a real-world dataset.

#### KBLaM knows when it cannot answer a question

If the model can refuse to answer a question in case it cannot find relevant information in the context, then the risk of model hallucinations can be largely controlled.
We consider a setting where, given a KB, we ask the model 100 questions in total, out of which 80 questions are answerable, and the other 20 are not.
We then consider the unanswerable questions as the positive class and answerable questions as negative.
To evaluate the model, we use standard binary classification metrics: precision and recall.
We compare the performance of KBLaM with in-context learning, where we explicitly prompt the LLM with the following instructions: “if relevant information cannot be found in the text, please respond I am sorry I cannot find relevant information in the KB”. The results are presented in Fig. [6.2](#S6.SS2.SSS0.Px4).
Broadly, we notice that when the number of triples increases, both KBLaM and in-context learning show decreased precision, i.e. the models start to incorrectly refuse for having not found the information, however, in-context learning shows more drastic degradation. The recall for both methods stays constant, i.e. both methods do not start to hallucinate as we add more triples.

#### Ablation study

We conduct ablation studies over the design choice of KBLaM and present the results in Appendix. [B](#A2):
We find that KBLaM also works with open-sourced sentence encoders, with exact performance depending on the encoders’ capacity (Fig. [7](#A2.F7));
The frequency of adding knowledge tokens can be set less than every layer, which further improve efficiency but could cause performance drop (Fig. [8](#A2.F8)). Lastly, we study the role of knowledge tokens at each layer, where we find that knowledge tokens in earlier layers may serve as “soft prompt” for instruction following (Fig. [9](#A2.F9)).

[Uncaptioned image]: /html/2410.10450/assets/x6.png

[Uncaptioned image]: /html/2410.10450/assets/x7.png

[Uncaptioned image]: /html/2410.10450/assets/x8.png

[Uncaptioned image]: /html/2410.10450/assets/x9.png

## 7 Conclusion

In this paper, we propose KBLaM, which efficiently augments pre-trained LLMs with external knowledge bases. KBLaM represents external knowledge as dense continuous vectors and leverages the independence structure between triples, which offers several advantages: avoiding expensive self-attention over large knowledge sources, enabling dynamic knowledge updates without fine-tuning, improving interpretability, and allowing attention to perform retrieval. These benefits could potentially also be applied to settings using raw text as knowledge representation.

## 8 Limitations and future work

Higher quality synthetic KB  In the paper, we demonstrate the feasibility of training KBLaM with fully synthetic data. However, as shown in Fig. [6.2](#S6.SS2.SSS0.Px4) and  [6.2](#S6.SS2.SSS0.Px4), KBLaM generalizes to out-of-distribution dataset (Enron) but the performance degrades. We believe this can be improved through a larger and more diverse synthetic dataset.
It is also possible to generate the synthetic data seeded from real KB data, such that the adapters can better capture the distribution of the real text.

Information loss in knowledge token  KBLaM encodes sentences in the triple as a fixed-length vector, therefore KBLaM may fail to *precisely* generate the text word by word.
This is acceptable for properties we considered such as “description” or “objectives”. However, this is problematic for cases when we need exact names or numerical values.
One future direction is to introduce a hyperparameter that adjusts the compression rate, to accommodate various amounts of information in a triple.

More sophisticated instruction tuning
Future work could explore complex instruction tuning examples for KBLaM, such as multi-hop, multi-step, or chain-of-thought reasoning about the KB. Since KBLaM incorporates all KB information in the context, similar to in-context learning, we believe it has the potential for more flexible KB reasoning compared to RAG.

## References

- Achiam et al. (2023)
Josh Achiam, Steven Adler, Sandhini Agarwal, Lama Ahmad, Ilge Akkaya, Florencia Leoni Aleman, Diogo Almeida, Janko Altenschmidt, Sam Altman, Shyamal Anadkat, et al.
Gpt-4 technical report.
*arXiv preprint arXiv:2303.08774*, 2023.
- Borgeaud et al. (2022)
Sebastian Borgeaud, Arthur Mensch, Jordan Hoffmann, Trevor Cai, Eliza Rutherford, Katie Millican, George Bm Van Den Driessche, Jean-Baptiste Lespiau, Bogdan Damoc, Aidan Clark, et al.
Improving language models by retrieving from trillions of tokens.
In *International conference on machine learning*, pp.  2206–2240. PMLR, 2022.
- Cai et al. (2023)
Tianle Cai, Kaixuan Huang, Jason D Lee, and Mengdi Wang.
Scaling in-context demonstrations with structured attention.
*arXiv preprint arXiv:2307.02690*, 2023.
- Chen et al. (2021)
Chun-Fu Richard Chen, Quanfu Fan, and Rameswar Panda.
Crossvit: Cross-attention multi-scale vision transformer for image classification.
In *Proceedings of the IEEE/CVF international conference on computer vision*, pp.  357–366, 2021.
- Dai et al. (2019)
Zihang Dai, Zhilin Yang, Yiming Yang, Jaime Carbonell, Quoc V Le, and Ruslan Salakhutdinov.
Transformer-xl: Attentive language models beyond a fixed-length context.
*arXiv preprint arXiv:1901.02860*, 2019.
- Dubey et al. (2024)
Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Amy Yang, Angela Fan, et al.
The llama 3 herd of models.
*arXiv preprint arXiv:2407.21783*, 2024.
- Gao et al. (2023)
Peng Gao, Jiaming Han, Renrui Zhang, Ziyi Lin, Shijie Geng, Aojun Zhou, Wei Zhang, Pan Lu, Conghui He, Xiangyu Yue, et al.
Llama-adapter v2: Parameter-efficient visual instruction model.
*arXiv preprint arXiv:2304.15010*, 2023.
- Ge et al. (2023)
Tao Ge, Jing Hu, Lei Wang, Xun Wang, Si-Qing Chen, and Furu Wei.
In-context autoencoder for context compression in a large language model.
*arXiv preprint arXiv:2307.06945*, 2023.
- Grave et al. (2016)
Edouard Grave, Armand Joulin, and Nicolas Usunier.
Improving neural language models with a continuous cache.
*arXiv preprint arXiv:1612.04426*, 2016.
- Hu et al. (2021)
Edward J Hu, Yelong Shen, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang, Lu Wang, and Weizhu Chen.
Lora: Low-rank adaptation of large language models.
*arXiv preprint arXiv:2106.09685*, 2021.
- Izacard & Grave (2020)
Gautier Izacard and Edouard Grave.
Leveraging passage retrieval with generative models for open domain question answering.
*arXiv preprint arXiv:2007.01282*, 2020.
- Izacard et al. (2023)
Gautier Izacard, Patrick Lewis, Maria Lomeli, Lucas Hosseini, Fabio Petroni, Timo Schick, Jane Dwivedi-Yu, Armand Joulin, Sebastian Riedel, and Edouard Grave.
Atlas: Few-shot learning with retrieval augmented language models.
*Journal of Machine Learning Research*, 24(251):1–43, 2023.
- Khandelwal et al. (2020)
Urvashi Khandelwal, Omer Levy, Dan Jurafsky, Luke Zettlemoyer, and Mike Lewis.
Generalization through memorization: Nearest neighbor language models.
In *International Conference on Learning Representations*, 2020.
- Klimt & Yang (2004)
Bryan Klimt and Yiming Yang.
The enron corpus: A new dataset for email classification research.
In *European conference on machine learning*, pp.  217–226. Springer, 2004.
- Lee et al. (2024)
Jinhyuk Lee, Anthony Chen, Zhuyun Dai, Dheeru Dua, Devendra Singh Sachan, Michael Boratko, Yi Luan, Sébastien MR Arnold, Vincent Perot, Siddharth Dalmia, et al.
Can long-context language models subsume retrieval, rag, sql, and more?
*arXiv preprint arXiv:2406.13121*, 2024.
- Lewis et al. (2020)
Patrick Lewis, Ethan Perez, Aleksandra Piktus, Fabio Petroni, Vladimir Karpukhin, Naman Goyal, Heinrich Küttler, Mike Lewis, Wen-tau Yih, Tim Rocktäschel, et al.
Retrieval-augmented generation for knowledge-intensive nlp tasks.
*Advances in Neural Information Processing Systems*, 33:9459–9474, 2020.
- Li et al. (2024)
Zhuowan Li, Cheng Li, Mingyang Zhang, Qiaozhu Mei, and Michael Bendersky.
Retrieval augmented generation or long-context llms? a comprehensive study and hybrid approach.
*arXiv preprint arXiv:2407.16833*, 2024.
- Liu et al. (2023)
Haotian Liu, Chunyuan Li, Qingyang Wu, and Yong Jae Lee.
Visual instruction tuning.
In *Thirty-seventh Conference on Neural Information Processing Systems*, 2023.
URL [https://openreview.net/forum?id=w0H2xGHlkw](https://openreview.net/forum?id=w0H2xGHlkw).
- Liu et al. (2024a)
Haotian Liu, Chunyuan Li, Yuheng Li, and Yong Jae Lee.
Improved baselines with visual instruction tuning.
In *Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition*, pp.  26296–26306, 2024a.
- Liu et al. (2024b)
Nelson F Liu, Kevin Lin, John Hewitt, Ashwin Paranjape, Michele Bevilacqua, Fabio Petroni, and Percy Liang.
Lost in the middle: How language models use long contexts.
*Transactions of the Association for Computational Linguistics*, 12:157–173, 2024b.
- Liu et al. (2024c)
Shih-Yang Liu, Chien-Yi Wang, Hongxu Yin, Pavlo Molchanov, Yu-Chiang Frank Wang, Kwang-Ting Cheng, and Min-Hung Chen.
Dora: Weight-decomposed low-rank adaptation.
*arXiv preprint arXiv:2402.09353*, 2024c.
- Loshchilov (2017)
I Loshchilov.
Decoupled weight decay regularization.
*arXiv preprint arXiv:1711.05101*, 2017.
- Merth et al. (2024)
Thomas Merth, Qichen Fu, Mohammad Rastegari, and Mahyar Najibi.
Superposition prompting: Improving and accelerating retrieval-augmented generation.
*arXiv preprint arXiv:2404.06910*, 2024.
- Mireshghallah et al. (2022)
Fatemehsadat Mireshghallah, Archit Uniyal, Tianhao Wang, David Evans, and Taylor Berg-Kirkpatrick.
Memorization in nlp fine-tuning methods.
*arXiv preprint arXiv:2205.12506*, 2022.
- Mu et al. (2024)
Jesse Mu, Xiang Li, and Noah Goodman.
Learning to compress prompts with gist tokens.
*Advances in Neural Information Processing Systems*, 36, 2024.
- Pope et al. (2023)
Reiner Pope, Sholto Douglas, Aakanksha Chowdhery, Jacob Devlin, James Bradbury, Jonathan Heek, Kefan Xiao, Shivani Agrawal, and Jeff Dean.
Efficiently scaling transformer inference.
*Proceedings of Machine Learning and Systems*, 5:606–624, 2023.
- Qin et al. (2022)
Bowen Qin, Binyuan Hui, Lihan Wang, Min Yang, Jinyang Li, Binhua Li, Ruiying Geng, Rongyu Cao, Jian Sun, Luo Si, et al.
A survey on text-to-sql parsing: Concepts, methods, and future directions.
*arXiv preprint arXiv:2208.13629*, 2022.
- Ratner et al. (2022)
Nir Ratner, Yoav Levine, Yonatan Belinkov, Ori Ram, Inbal Magar, Omri Abend, Ehud Karpas, Amnon Shashua, Kevin Leyton-Brown, and Yoav Shoham.
Parallel context windows for large language models.
*arXiv preprint arXiv:2212.10947*, 2022.
- Reimers & Gurevych (2019)
Nils Reimers and Iryna Gurevych.
Sentence-bert: Sentence embeddings using siamese bert-networks.
In *Proceedings of the 2019 Conference on Empirical Methods in Natural Language Processing*. Association for Computational Linguistics, 11 2019.
URL [https://arxiv.org/abs/1908.10084](https://arxiv.org/abs/1908.10084).
- Team et al. (2023)
Gemini Team, Rohan Anil, Sebastian Borgeaud, Yonghui Wu, Jean-Baptiste Alayrac, Jiahui Yu, Radu Soricut, Johan Schalkwyk, Andrew M Dai, Anja Hauth, et al.
Gemini: a family of highly capable multimodal models.
*arXiv preprint arXiv:2312.11805*, 2023.
- Vaswani (2017)
A Vaswani.
Attention is all you need.
*Advances in Neural Information Processing Systems*, 2017.
- Winn et al. (2018)
John Winn, John Guiver, Sam Webster, Yordan Zaykov, Martin Kukla, and Dany Fabian.
Alexandria: Unsupervised high-precision knowledge base construction using a probabilistic program.
In *Automated Knowledge Base Construction (AKBC)*, 2018.
- Winn et al. (2021)
John Winn, Matteo Venanzi, Tom Minka, Ivan Korostelev, John Guiver, Elena Pochernina, Pavel Mishkov, Alex Spengler, Denise Wilkins, Siân Lindley, et al.
Enterprise alexandria: Online high-precision enterprise knowledge base construction with typed entities.
In *3rd Conference on Automated Knowledge Base Construction*, 2021.
- Wu et al. (2024)
Haolun Wu, Ye Yuan, Liana Mikaelyan, Alexander Meulemans, Xue Liu, James Hensman, and Bhaskar Mitra.
Structured entity extraction using large language models.
*arXiv preprint arXiv:2402.04437*, 2024.
- Wu et al. (2023)
Shengqiong Wu, Hao Fei, Leigang Qu, Wei Ji, and Tat-Seng Chua.
Next-gpt: Any-to-any multimodal llm.
*arXiv preprint arXiv:2309.05519*, 2023.
- Wu et al. (2022)
Yuhuai Wu, Markus N Rabe, DeLesley Hutchins, and Christian Szegedy.
Memorizing transformers.
*arXiv preprint arXiv:2203.08913*, 2022.
- Ye et al. (2023)
Qinyuan Ye, Iz Beltagy, Matthew E Peters, Xiang Ren, and Hannaneh Hajishirzi.
Fid-icl: A fusion-in-decoder approach for efficient in-context learning.
In *Proceedings of the 61st Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers)*, pp.  8158–8185, 2023.
- Yen et al. (2024)
Howard Yen, Tianyu Gao, and Danqi Chen.
Long-context language modeling with parallel context encoding.
*arXiv preprint arXiv:2402.16617*, 2024.
- Yu et al. (2024)
Tan Yu, Anbang Xu, and Rama Akkiraju.
In defense of rag in the era of long-context language models.
*arXiv preprint arXiv:2409.01666*, 2024.
- Zeng et al. (2023)
Shenglai Zeng, Yaxin Li, Jie Ren, Yiding Liu, Han Xu, Pengfei He, Yue Xing, Shuaiqiang Wang, Jiliang Tang, and Dawei Yin.
Exploring memorization in fine-tuned language models.
*arXiv preprint arXiv:2310.06714*, 2023.
- Zhang* et al. (2020)
Tianyi Zhang*, Varsha Kishore*, Felix Wu*, Kilian Q. Weinberger, and Yoav Artzi.
Bertscore: Evaluating text generation with bert.
In *International Conference on Learning Representations*, 2020.
URL [https://openreview.net/forum?id=SkeHuCVFDr](https://openreview.net/forum?id=SkeHuCVFDr).
- Zhao et al. (2023)
Yue Zhao, Ishan Misra, Philipp Krähenbühl, and Rohit Girdhar.
Learning video representations from large language models.
In *Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition*, pp.  6586–6597, 2023.
- Zheng et al. (2023)
Lianmin Zheng, Wei-Lin Chiang, Ying Sheng, Siyuan Zhuang, Zhanghao Wu, Yonghao Zhuang, Zi Lin, Zhuohan Li, Dacheng Li, Eric Xing, et al.
Judging llm-as-a-judge with mt-bench and chatbot arena.
*Advances in Neural Information Processing Systems*, 36:46595–46623, 2023.

## Appendix A Extended related work

In this section, we present additional discussion on work related to KBLaM.

#### Memory augmented language models

Some other works have also incorporated external memory, i.e. a collection of fixed length latent representations, into language models (Grave et al., ; Dai et al., ; Khandelwal et al., ; Wu et al., ). KBLaM’s knowledge tokens can also be seen as external memory generated from the KB. The most similar approach along this line of work is Wu et al. (), which constructs external memory as key-value vector pairs, similar to KBLaM. However, they require training a transformer from scratch and the memory key-value pairs come from training tokens seen in the past, whereas KBLaM’s memory is from an external KB and is augmented into a pre-trained model at inference time.

#### Token compression

KBLaM’s encoder transforms each triple into a continuous vector pair equivalent to 1 token in length, which essentially compresses information from the multiple-token string into one single token. There have been some recent works aiming at compressing a long context into more continuous representations. Ge et al. () replaces short prompts with special tokens to instruct the LLM to do some particular tasks, however they require fine-tuning of the LLM. Mu et al. () encode a long context into memory slots through a fine-tuned LLM, which is a more powerful but also costly compressor, which we believe could in future work be combined with KBLaM’s framework as an alternative to pre-trained sentence encoder.

#### Structured attention

KBLaM also shares similarity with works that use structured attention masks instead of the standard lower-triangular causal attention mask (Ratner et al., ; Cai et al., ; Merth et al., ). These works utilize independence assumptions in the context, e.g. between different randomly-chosen sets of few-shot examples (Ratner et al., ) or documents from different sources (Cai et al., ; Merth et al., ), and let independent contexts not attend over each other, which reduces overhead of attention operations. In KBLaM, we make a similar independence assumption: We assume independence between different KB triples, which provides KBLaM with scalability and interpretability.

## Appendix B Ablation study

We conducted a number of ablation studies on the design choices of KBLaM.

#### Choice of encoders

KBLaM relies on a pre-trained sentence encoder, and the capacity of the encoder can affect the performance. Indeed our results (Fig. [7](#A2.F7)) verify this: We experimented with some open-sourced alternatives from sentence transformer (Reimers & Gurevych, ). Broadly we observe that the commercial OpenAI embedding works better than the open-sourced one, models with higher embedding dimensions perform better than those with lower dimensions.

#### Where to add knowledge tokens

In Fig. [5](#S5.SS0.SSS0.Px1) and [5](#S6.F5), we noted that the attention score from the $15$th layer exhibits accurate retrieval behavior.
Extending this analysis to other layers (Fig. [9](#A2.F9)), we find that they *do not* demonstrate such straightforward accuracy (although further probing might reveal more nuanced patterns).
We suspect that this indicates that the amount of KB information provided by the knowledge tokens vary at different layers provide.
To better understand this, we examining the variation of keys and values at each layer ($\mathbb{V}_{m}[\bm{\tilde{k}}^{l}_{m}]$ and $\mathbb{V}_{m}[\bm{\tilde{v}}^{l}_{m}]$) across triples, we observe that encoder outputs vary minimally in earlier layers.

#### Frequency of knowledge tokens

The previous observation indicates that the knowledge tokens may not be providing information about the KB at every layer, therefore we consider varying the frequency of adding knowledge tokens (denoted by $K$). That is, for certain attention layers of KBLaM, we used [Eq. 3](#S3.E3) instead of [Eq. 9](#S4.E9).
In particular, we considered only adding knowledge every 1, 3 and 10 layers and the results are presented in Fig. [8](#A2.F8).
Broadly, we notice with lower frequency, the model tends not to follow instructions: When $K=10$, KBLaM fails to provide any refusal answer. As such we suspect that knowledge tokens in earlier layers may serve as soft instruction prompts that guides the LLM how to use the KB.

Figure: Figure 7: KBLaM works with other types of pre-trained encoders. Similar to the setting in Fig. [5](#S6.F5) and Fig. [6.2](#S6.SS2.SSS0.Px4), but with sentence encoder of varying capacity.
Refer to caption: /html/2410.10450/assets/x10.png

Figure: Figure 8: Ablation study on knowledge token injection frequency (K). If we have knowledge tokens added to attention , the model performance drops ($K=10$), but the memory overhead decreases and efficiency increases. The BERT score of $K=3$ is comparable to $K=1$ (first two columns), however when $K=3$, the model shows lower precision at refusing answering unanswerable questions (third column), and when $K=10$, the model fails to output refusal answers.
Refer to caption: /html/2410.10450/assets/x13.png

Figure: Figure 9: Knowledge tokens’ embeddings contain various level of KB information across layers. The embeddings in earlier layers vary less among different entities, indicating that they may serve as instruction prompt on how to use the KB. Later layers show a higher degree of variation, indicating that they may be used to provide actual knowledge.
Refer to caption: /html/2410.10450/assets/x15.png

## Appendix C Sample KB

Here we present some example triples from both KBs, where each triple is of format

$$ $(\texttt{<name>};\texttt{<property>};$ $$

## Appendix E Prompt

In this section, we provide all prompts used for KBLaM’s experiments.

### E.1 Prompt for synthetic KB generation

To prompt the GPT-4 for generating a synthetic KB, we begin by setting a system prompt as

Then we construct a list of idea types and a list of object types:

Then for each combination of object type and idea type in the list, we prompt GPT with

which gives us 50 different names. Then in the same context, we further prompt GPT with

which generates the

for three <property>s: description, objectives and purpose. Lastly, we perform one round of polishing, to let GPT diversity the language and style of the generated KB triple

### E.2 Prompt for open-ended Q&A generation

To generate open-ended Q&A instructions, we first take a triple from the KB, construct a simple Q&A from it, and then we feed the simple Q&A into GPT and let GPT augment it into a more complex form via the following prompt

### E.3 Prompt for GPT evaluation of open-ended Q&A

To evaluate the answer quality of open-ended Q&A tasks, we again use GPT. In particular, we ask GPT to score the answer from two aspects: 1. How grounded the answer is given the question and the KB; 2. How reasonable the open-ended part of the answer is.

We begin with a system prompt of

Then, we use the following prompt, which encourages the model to give chain-of-thought reasoning for scoring (Zheng et al., ):

Followed by the prompt, we include 5 examples to better help GPT calibrate the scoring rule

### E.4 Prompt for Llama evaluation

For in-context learning, we use the following prompts for evaluation

For zero-shot learning, we use the following prompt

### E.5 Question template

For simple Q&A, we use the following templates

For multi-entities Q&A involving $G$ triples, { ( <name> g , <property> g , _g)}_g=1^G,weusethefollowingtemplateWhatisthe{}Tellme{},Canyouletmeknow{},Canyou fragments { ( <name> 𝑔 , <property> 𝑔 , _g)}_g=1^G,weusethefollowingtemplateWhatisthe{}Tellme{},Canyouletmeknow{},Canyou

## 2 Related work 1 1 1 We provide more extensive related work discussion in Appendix A

Retrieval augmented generation (RAG)  RAG (Lewis et al., ) is one of the most successful approaches for augmenting external knowledge into LLMs, and many approaches have been proposed to perform RAG, such as similarity search and text2Sql (Qin et al., ). KBLaM can also be understood as the model doing RAG in an implicit way: given input prompts, the rectangular attention compares the sequence’s queries with knowledge tokens’ keys, then the weighted average is taken of all knowledge tokens’ values based on query-key similarity, similarly to a soft retrieval process. The averaged value embedding is then added to the hidden state, augmenting next-token generation with information from the KB. However, KBLaM differs from RAG in that it does not rely on a separate retrieval module. Instead, everything is carried out by attention in an end-to-end fashion.

Multi-modal language model  Recent advances in Multi-modal Language Models (Liu et al., ; Gao et al., ; Liu et al., ; Zhao et al., ; Wu et al., , MLMs), demonstrate the possibility of incorporating data of other modalities (e.g. images, videos) into a pre-trained LLM through learned adapters and instruction tuning. With KBLaM, we can consider a KB as another modality in that a KB is presented to the LLM in the form of continuous vectors via an encoder, and an encoder-adapter is learned also through instruction tuning inspired by MLM literature.

Key-value (KV) cache mechanism  Given a context string, the KV cache mechanism (Pope et al., ) caches the key and value embeddings at each layer, also known as the prefill stage. With KV cache, the complexity of generating new tokens conditioned on the context decreases from quadratic in context length to linear. KBLaM’s encoded knowledge tokens work in a way similar to KV cache, but we have the key and value acquired from an external encoder rather than through self-attention process. Additionally, if the context is modified, the standard KV cache requires re-computing the whole KV cache due to the causal attention mask, whereas in KBLaM only the corresponding knowledge token need be updated as the knowledge tokens do not attend to each other.

Augmenting information into an LLM in a continuous form  RAG and in-context learning incorporates external information into LLMs in the form of a *discrete string*, whereas KBLaM augments information into LLMs as fixed length *continuous* representations. Many works also consider first encoding external documents as continuous vectors, then concatenating the vectors as the input for the decoder model (Izacard & Grave, ; Izacard et al., ; Ye et al., ) or using cross-attention to inject these vectors into the decoder (Borgeaud et al., ).
However, these works focus on training a new model from scratch, while KBLaM is applied to augment a *pre-trained* LLM with knowledge. Recent work (Yen et al., ) also considers augmenting a *pre-trained* LLM with encoded information, similar to our setting. However, their approach requires training an encoder *from scratch*, whereas KBLaM’s only trains linear adapters after a pre-trained black-box encoder. In addition, they add extra cross-attention layers after each self-attention layer for augmentation, while KBLaM modifies existing self-attention layers without introducing new components.

## 3 Background

#### Knowledge base in the form of triples

In this paper, we assume that the external knowledge is represented as a structured knowledge base (KB) obtained from unstructured text using, in particular, a KB is composed of triples of format

$$ $\{(\texttt{<name>}_{m};\texttt{<property>}_{m};\description_m)\}_{m=1}^{M}. \end{equation}Intherestofthepaper,wewillrefertoasetoftriplesoftheformat\lx@cref{creftype~refnum}{eq:triple_def}asa\textsc{KB}andeachtripleinsideitasa\emph{knowledge triple}.ExamplesofsuchKBisshowninTable.~{}\ref{tab:syn_kb}and~{}\ref{tab:enron_kb}inAppendix~{}\ref{sec:sample_kb}.TheKBconstructionprocesssummarizesinformationfrommultipledocumentsandorganizesitinastructuredway.Importantly,thestructurednatureofthe\textsc{KB}enablesustoaugmentLLMwithinformationfromdocumentsinanefficientway,whichwediscussinthenextsection.\par Itisalsoworthnotingthat\textsc{KB}constructionis\emph{not}thefocusofourpaper:\textsc{KBLaM}workson\textsc{KB}generatedfromcorpususing\emph{existing tools}.Inthepaper,wemainlyfocusontwo{\textsc{KB}}s:Asynthetic\textsc{KB}generatedviaGPT,andareal\textsc{KB}constructedfromtheEnronemail~{}\cite[citep]{(\@@bibref{AuthorsPhrase1Year}{klimt2004enron}{\@@citephrase{, }}{})}dataset.Detailsofthetwo{\textsc{KB}}swillbediscussedindetailinlatersections.\par\par\par$ (1) $$

#### Self-attention layer

A decoder-based transformer is mostly composed of multiple self-attention layers (Vaswani, ).
Consider a transformer of $L$ layers.
Each layer has three projection heads: $\bm{W}_{Q}^{l}\in\mathbb{R}^{D\times D},\bm{W}_{K}^{l}\in\mathbb{R}^{D\times D},\bm{W}_{V}^{l}\in\mathbb{R}^{D\times D},l\in\{1,\ldots,L\}$,
where $D$ denotes *embedding dimension*.
Given a user prompt (e.g. a question) of $N$ tokens, each attention layer takes a sequence of $N$ token embeddings $\bm{x}^{l}=[\bm{x}_{1}^{l},\ldots,\bm{x}_{n}^{l},\ldots\bm{x}_{N}^{l}]^{\top}\in\mathbb{R}^{N\times D}$ as input,
then for the $n$th token embedding, it is first transformed by the projection heads into three vectors

$$ $\displaystyle\bm{q}_{n}^{l}$ $\displaystyle=\bm{W}_{Q}^{l}\bm{x}_{n}^{l},$ $\displaystyle\bm{k}_{n}^{l}$ $\displaystyle=\bm{W}_{K}^{l}\bm{x}_{n}^{l},$ $\displaystyle\bm{v}_{n}^{l}$ $\displaystyle=\bm{W}_{V}^{l}\bm{x}_{n}^{l},$ (2) $$

where we refer to $\bm{q}_{n}^{l}$, $\bm{k}_{n}^{l}$, and $\bm{v}_{n}^{l}$ as the $n$th token’s query, key, and value embedding at the $l$th layer respectively. Note that in practice, each attention layer often has multiple sets of attention heads but here we assume only a single head for notation simplicity.

The output of the attention layer, $\bm{y}^{l}=[\bm{y}_{1}^{l},\ldots,\bm{y}_{n}^{l},\ldots,\bm{y}_{N}^{l}]^{\top}\in\mathbb{R}^{N\times D}$, is computed as

$$ $\bm{y}_{n}^{l}=\frac{\sum_{i=1}^{n}\exp\left(w_{n,i}\right)\bm{v}_{i}^{l}}{\sum_{i=1}^{n}\exp\left(w_{n,i}\right)},~{}\text{where}~{}w_{n,i}=\langle\bm{q}_{n}^{l},\bm{k}_{i}^{l}\rangle/\sqrt{D}$ (3) $$

where $\langle\cdot,\cdot\rangle$ denotes the inner product of two vectors. After $\bm{y}_{n}^{l}$ is acquired, it is often fed into a feedforward network (FFN) with output dimension of $D$ for further transformation.

A standard implementation of self-attention would have a time complexity of $\mathcal{O}(N^{2}D)$ and memory complexity of $\mathcal{O}(N^{2})$ for computing and storing all $w_{n,i}$. Additionally, the FFN would also introduce a significant computation overhead of magnitude $\mathcal{O}(ND^{2})$.
Due to these factors, as the sequence gets longer, self-attention faces high memory consumption and slow computation.

Figure: (a) Step 1. KB encoding For each triple from a KB ([Eq. 1](#S3.E1)) of $M$ triples, we first construct a key and value string (left most white boxes). Then both strings are processed by a pre-trained sentence encoder followed by a learned linear adapter (middle two blue boxes). The acquired knowledge tokens ([Eqs. 6](#S4.E6) and [7](#S4.E7), right-most small boxes) are then stored on disk for later usage. Note that all $\bm{\tilde{k}}_{m}^{l}$s and $\bm{\tilde{k}}_{m}^{l}$s have a fixed dimension of $D$, i.e. identical to that of the key and value embeddings of a token, regardless of input length.
Refer to caption: /html/2410.10450/assets/x2.png

## 4 Augmenting LLM with the KB

In this section, we discuss how KBLaM augments an LLM with a KB.
The process has two steps (visualized in Fig. [2](#S3.F2)): 1. We convert each triple in the KB from the string form into a continuous key-value vector pair, referred to as a *knowledge token*, through a pre-trained sentence encoder followed by linear adapters; 2. All knowledge tokens are then injected into each attention layer of an LLM through a *rectangular attention* structure.

#### Knowledge tokens

Provided a KB in the format of [Eq. 1](#S3.E1), for each triple, we first adopt a *pre-trained* sentence encoder model, denoted as $f(\cdot)$ which converts a string into a $P$-dimension continuous embedding vector.
Through the encoder, we convert each triple into a base key embedding and a base value embedding.
In particular, for the $m$th triple, we have:

$$ (4) $$

Next, we introduce a linear key and a linear value adapter

$$ $\displaystyle\bm{\tilde{W}}_{K}\in\mathbb{R}^{L\times D\times P},$ $\displaystyle\bm{\tilde{W}}_{V}\in\mathbb{R}^{L\times D\times P},$ (5) $$

where $L$ denotes the number of attention layers in a model.
With the adapters, we map $\bm{\mathsf{{k}}}_{m}$ and $\bm{\mathsf{{v}}}_{m}$ from the sentence encoder’s space to LLM’s key and value embedding space at each attention layer.
In particular, for the $m$th knowledge triple, we transform its base key and value embedding into

$$ $\displaystyle\tilde{\bm{k}}_{m}=[\tilde{\bm{k}}_{m}^{1},\ldots,\tilde{\bm{k}}_{m}^{l},\ldots,\tilde{\bm{k}}_{m}^{L}]^{\top}=\bm{\tilde{W}}_{K}\bm{\mathsf{{k}}}_{m}\in\mathbb{R}^{L\times D},$ (6) $\displaystyle\tilde{\bm{v}}_{m}=[\tilde{\bm{v}}_{m}^{1},\ldots,\tilde{\bm{v}}_{m}^{l},\ldots,\tilde{\bm{v}}_{m}^{L}]^{\top}=\bm{\tilde{W}}_{V}\bm{\mathsf{{v}}}_{m}\in\mathbb{R}^{L\times D}.$ (7) $$

Since each $\tilde{\bm{k}}_{m}^{l}$ and $\tilde{\bm{v}}_{m}^{l}$ share the same size with LLM’s key and value embedding (i.e. $\bm{k}_{n}^{l}$ and $\bm{v}_{n}^{l}$ in [Eq. 2](#S3.E2)), they can be directly incorporated into attention’s computation. Therefore, the encoding process is equivalent to converting a knowledge triple from a string of multiple tokens into a special token, whose key and value embeddings at each layer are generated through an encoder, rather than through self-attention. As such we refer to $(\tilde{\bm{k}}_{m},\tilde{\bm{v}}_{m})$ pair as a *knowledge token*.

We visualize the process in Fig. [2(a)](#S3.F2.sf1).
This encoding process is applied to all triples in the KB, which transforms the information from a KB into a collection of knowledge tokens

$$ $\{(\texttt{<name>}_{m},\texttt{<property>}_{m},\description_m)\}_{m=1}^{M} \xrightarrow[]{\text{Encode}} \{(\tilde{\bm{k}}_m , \tilde{\bm{v}}_m)\}_{m=1}^M. \end{equation}\par$ (8) $$

#### Rectangular Attention: Injecting knowledge token into prompt tokens

After converting a KB into a collection of knowledge tokens we inject the information of the KB into the each attention layer of a pre-trained LLM through a modified attention structure(^3^33We use symbols with a tilde to denote additional attention components introduced by KBLaM.).
In particular, consider an input prompt of $N$ tokens, the $l$th attention layer takes in $\bm{x}^{l}=[\bm{x}_{1}^{l},\ldots,\bm{x}_{n}^{l},\ldots,\bm{x}_{N}^{l}]^{\top}$ and the $l$th dimension of the knowledge tokens $\{(\bm{\tilde{k}}_{m}^{l},\bm{\tilde{v}}_{m}^{l})\}_{m=1}^{M}$, the output of each attention layer, $\tilde{\bm{y}}^{l}=[\tilde{\bm{y}}_{1}^{l},\ldots,\tilde{\bm{y}}_{n}^{l},\ldots,\tilde{\bm{y}}_{N}^{l}]^{\top}\in\mathbb{R}^{N\times D}$, is given by

$$ $\displaystyle\tilde{\bm{y}}_{n}^{l}$ $\displaystyle=\frac{{\sum_{m=1}^{M}\exp\left(\tilde{w}_{n,m}^{l}\right)\bm{\tilde{v}}_{m}^{l}}+\sum_{i=1}^{n}\exp\left(w_{n,i}^{l}\right)\bm{v}_{i}^{l}}{{\sum_{m=1}^{M}\exp\left(\tilde{w}_{n,m}^{l}\right)}+\sum_{i=1}^{n}\exp\left(w_{n,i}^{l}\right)},$ (9) $$

where

$$ $\displaystyle\tilde{w}_{n,m}^{l}=\langle\tilde{\bm{q}}_{n}^{l},\bm{\tilde{k}}_{m}^{l}\rangle/\sqrt{D},~{}\tilde{\bm{q}}_{n}^{l}=\bm{\tilde{W}}_{Q}^{l}\bm{x}_{n},$ $\displaystyle w_{n,i}^{l}=\langle\bm{q}_{n}^{l},\bm{k}_{i}^{l}\rangle/\sqrt{D},$ (10) $$

and $\bm{\tilde{W}}_{Q}^{l}\in\mathbb{R}^{D\times D}$ is a *learnable* query head initialized from $\bm{W}_{Q}^{l}$.

Figure: Figure 3: Memory overhead of different methods. Given a KB of $M$ triples (with each triple $K$-token long on average). In-context learning’s memory scales with $(KM)^{2}$, whereas KBLaM’s memory scales with $M$.
Refer to caption: /html/2410.10450/assets/x4.png

We name the attention structure of [Eq. 3](#S3.E3) *rectangular attention* as the attention matrix is of rectangular shape (bottom-right panel in Fig. [2(b)](#S3.F2.sf2)).
The idea is that tokens in the prompt part can attend to, in addition to all previous tokens in the prompt, all knowledge tokens. On the other hand, knowledge tokens cannot attend to each other, i.e. they do not have self-attention nor query embeddings. This design gives an attention matrix of rectangular shape of size $(M+N)\times N$, where $M$ is the number of triples in the KB.
Additionally, if knowledge tokens are not introduced, i.e. $M=0$, [Eq. 9](#S4.E9) falls back to [Eq. 3](#S3.E3) and we recover the pre-trained LLM.

The memory and time complexity of the rectangular attention is $\mathcal{O}((M+N)N)$ and $\mathcal{O}((M+N)ND)$ respectively. In real settings, it is likely that $M>>N$, i.e. the number of external knowledge items is far larger than the size of the prompt/question, then the overhead would grow only linearly with $M$ (instead of quadratically), which allows KBLaM to scale up to very large values of $M$, enabling putting a large amount of information in the context.
Also, note that the output sequence length of rectangular attention does not vary with the number of knowledge tokens, therefore the overhead for the intermediate FFN modules stays unchanged regardless of $M$.

Lastly, it can be seen from  [Eq. 3](#S3.E3) that changing the order of the knowledge tokens does not affect the output, therefore KBLaM does not suffer from the positional bias issue of in-context learning (Liu et al., ). Additionally, since all knowledge tokens are encoded independently, when the KB is updated, e.g. new knowledge triples are created or a triple is updated, one could simply generate a new knowledge token with the encoder or update the corresponding knowledge tokens value. This highlights the difference between KBLaM and the standard KV cache mechanism in LLM, which needs re-computation when the cached content is modified in any way.

Note that KBLaM’s rectangular attention *differs* from cross-attention (Vaswani, ; Chen et al., ), which operates on a source sequence with key and value embeddings and a target sequence with only query embeddings, where tokens within each sequence do not attend to one another. In our setting, if we consider the KB as the source and the input prompt as the target, cross-attention corresponds to retaining only the summations over $M$ in [Eq. 9](#S4.E9). However, KBLaM’s attention also contains self-attention over the prompt tokens, given by the summations over $n$ in [Eq. 9](#S4.E9).

#### KB length generalization through attention score scaling

Notice that the part of the attention output corresponding to the KB would scale with $M$.
In particular, assume $\tilde{w}_{n,m}^{l},m\in\{1,\ldots,M\}$ are all of similar magnitude, then as $M$ grows, eventually the second term (the prompt/questions part) in both denominator and numerator of [Eq. 9](#S4.E9) will be overwhelmed by the first term (the KB part), and the information from the prompt/question will be lost.
To resolve this, during *inference*, we add a constant shift to the KB part’s unnormalized attention score

$$ $\tilde{w}_{n,m}^{l}=\log C-\log M+\langle\tilde{\bm{q}}_{n}^{l},\bm{\tilde{k}}_{m}^{l}\rangle/\sqrt{D},$ (11) $$

where $C\in\mathbb{R}^{+}$ is a hyperparameter which we fix as $100$ in our experiments. This is equivalent to multiplying the summations over $M$ by $\tfrac{C}{M}$, such that the total contribution of the KB part stays roughly unchanged as $M$ scales.

## 5 KB instruction tuning

The learnable parameters in KBLaM comprise the weights for the weights for the linear adapters/heads $\theta=\left\{\bm{\tilde{W}}_{K},\bm{\tilde{W}}_{V},\{\bm{\tilde{W}}_{Q}^{l}\}_{l=1}^{L}\right\}$,
Inspired by recent works in multi-modal language models (Liu et al., ), we use instruction tuning for parameter learning. Specifically, given a KB, we generate question-answer pairs about the KB (denoted as $Q$ and $A$ respectively) using formatted strings or GPT, then we optimize $\theta$ using:

$$ $\max_{\theta}\log p_{\theta,\phi}(A\mid Q,\textsc{KB}),$ (12) $$

where $\phi$ denotes parameters of the pre-trained LLM, encompassing QKV heads, FFN parameters and embedding/outpu layer weights.
Notably, KB instruction tuning preserves the base LLM’s reasoning abilities by avoiding fine-tuning the pre-trained LLM itself. Moreover, learning only the linear adapters minimizes the risk of memorizing training data, a known issue in LLM fine-tuning (Zeng et al., ; Mireshghallah et al., ).
To further mitigate this risk, we conducted instruction tuning using a *fully synthetic* KB generated by GPT. This approach is motivated by the understanding that the instruction tuning process aims *not* to memorize specific information, but to learn a projection from the pre-trained sentence encoder space to the LLM’s semantic space. Consequently, the exact content of the KB is less critical, provided the training data encompasses a diverse range of text.

#### Generation of synthetic KB

To synthesize a KB, we first use GPT to generate 50 <name>s based on combinations of 30 object types (e.g. restaurant name, software tool) and 30 idea types (e.g. natural phenomena, famous landmarks). Then for each name, we prompt GPT to generate the
Figure 4: KBLaM’s attention matrix is interpretable. Consider a toy KB of 10 triples, the heatmap shows the attention weights under different questions. We select the 15th attention layer and visualize the post-softmax attention score averaged over all attention heads. The x-axis only shows the keys corresponding the KB part and the y-axis is aligned by the end of each question’s query.

In this section, we perform empirical evaluation for KBLaM.
We begin by showing that KBLaM’s attention matrix, after instruction tuning, shows interpretable patterns and works as an accurate retriever.
We then show that KBLaM can answer questions with performance comparable to in-context learning, but with much lower memory cost, and can scale to 10K triples with little performance degradation.
Lastly, we show that for questions with no answer in the KB, the model can refuse to answer questions, with “over-refusal” behavior occurring later than in-context learning.
Lastly, we conduct ablation studies to understand the design choices of KBLaM.

### 6.1 Experiment setting

Model specification  For all experiments, we use the instruction fine-tuned version of Llama3 8B (Dubey et al., ) as the backbone LLM, and OpenAI’s ada-002 sentence embedding model ($P=1536$) as the pre-trained encoder for computing base key and value embedding ([Eq. 5](#S4.E5)).

Optimization setting
We initialize the key and value adapters: $\bm{\tilde{W}}_{K}$ and $\bm{\tilde{W}}_{V}$ randomly, and we initialize additional query heads $\bm{\tilde{W}}_{Q}^{l}$ at each layer from the pre-trained weights $\bm{W}_{Q}^{l}$.
Optimization is conducted using AdamW (Loshchilov, ) with a step size of $5\times 10^{-4}$ and a cosine learning rate decay to $5\times 10^{-6}$ for 20K iterations. Each iteration uses a mini-batch of 400 Q&A pairs, composed of 20 micro-batches of 20 samples. The instruction tuning is performed on a single 80GB A100 GPU under bfloat16 without any parameter-efficient tuning methods.

Construction of the training batches
We train KBLaM using instruction tuning, where each training sample consists of a KB, a question, and an answer ([Eq. 12](#S5.E12)). These samples are derived from the first 120K triples of the synthetic KB described in Sec. [5](#S5). To construct each training sample, we perform the following procedure:
We randomly select a subset of 10 to 100 triples from the synthetic KB to form a sample-specific KB.
Depending on the instruction type (simple, multi-entities, open-ended, or unanswerable Q&A), we designate one, multiple, or no triples as the *relevant* triple(s), with the rest serving as *distractors*.
We generate a question based on the instruction template and the relevant triple(s).
In each batch of 20 micro-batches, we include 2 micro-batches of unanswerable Q&A pairs and 6 micro-batches each for the other three Q&A types.
During training, we find limiting the KB size crucial for successful convergence. To further optimize training efficiency, we pre-compute the base key and value vectors ([Eq. 5](#S4.E5)) offline for all triples in the synthetic KB.

Evaluation dataset
For evaluation, we considered the following two KB datasets(^4^44All KBs used for training and evaluation are released together with the paper.)

- –
Synthetic data  The validation set of the synthetic KB, i.e. the 15000 triples not used for training.
- –
Enron  A KB constructed from the Enron (Klimt & Yang, ) dataset, an open-sourced corporate email dataset.
Following  (Wu et al., ), we constructed a KB of triples from the Enron emails data by fine-tuning a small language model, which we then cluster to remove duplicates  (Winn et al., ; ).

For each validation dataset, we consider <property> in “description”, “purpose” and “objectives”.

Baseline
We considered the following two methods as baselines

- –
In-context learning  Flatten all triples in a KB as strings, and attach it in front of the prompt. In-context learning’s memory overhead grows quadratically with the number of triples. Notice that we can only experiment with a maximum of 200 triples due to memory constraints.
- –
Zero-shot learning  Directly ask the LLM a particular question and provide no additional context, using the LLM’s internal knowledge for answering the questions.

Evaluation setting  For all evaluations, we repeat the experiments with 5 random seeds, each run containing 100 test samples, i.e. we randomly generate 100 KBs of various sizes and query the model with questions about the KB. The results reported are the averaged values over all 500 questions.

### 6.2 Experiment results

#### KBLaM attention is an accurate retriever

Figure: Figure 5: Through instruction tuning, the attention shows retrieval behavior. Given simple Q&A on the validation set of the synthetic data (solid line) and Enron dataset (OOD, dashed line), we use the attention score at the 15th layer, averaged over all attention heads, as a classification score for each triple and measure the top-1 and top-5 accuracy. KBLaM assigns the highest attention score to the truly relevant triple most of the time (with performance degraded but still reasonable on OOD).
Refer to caption: /html/2410.10450/assets/x5.png

Given a question about the $m$th triple in the KB, we expect certain tokens in the question, particularly those containing the keyword ${\texttt{<name>}}_{m}$, to attend more strongly to the $m$th knowledge token compared to others. This would be reflected in larger values of $\tilde{w}_{n,m}^{l}=\langle\tilde{\bm{q}}_{n}^{l},\bm{\tilde{k}}_{m}^{l}\rangle$ for these relevant tokens. We observe such patterns in practice, as visualized in Fig. [5](#S5.SS0.SSS0.Px1). In particular, we consider a KB of 10 triples and ask KBLaM questions about a single entity, two entities, or a question without an answer in the KB. We find the post-softmax attention score at the 15th layer, i.e. the middle attention layer of the 32-layer Llama3 8B, averaged over all 32 attention heads, to correlate well with the desired entity.
This property allows KBLaM to answer input questions in an interpretable way, as for any query, the top-K triples with the highest attention score form supporting evidence.
This behavior suggests that KBLaM’s attention mechanism functions implicitly as a retriever. We can quantitatively assess its retrieval performance by treating the attention scores, $\tilde{w}_{n,m}^{l},m\in{1,\ldots,M}$, as a classification score. In particular, we again extract the averaged-over-heads attention score at the 15th layer, and we evaluate the top-1 and top-5 accuracy, which represent the percentage of evaluation samples where the true triple received either the highest attention score or was among the top-5 highest scores, respectively.
Fig. [5](#S6.F5) presents these results, demonstrating that KBLaM’s attention remains highly accurate even with a large number of triples, both for in-distribution and out-of-distribution (OOD) data. Notably, this performance is achieved *solely* through instruction tuning using paired data, without any explicit regularization or retrieval objectives. In fact, imposing regularization on the attention structure would be challenging and potentially counterproductive as it is unclear, e.g. which tokens, attention heads or attention layers should be regularized.

#### KBLaM can reason about knowledge

We consider querying the model with three types of questions: Simple, two-entities (multi-entities Q&A involving two triples), and open-ended Q&A. For the first two types, we evaluate the output quality using BERT score (Zhang* et al., ). For open-ended questions, we use GPT-4 to score the output quality between 0 and 5 (full prompt presented in Appendix. [E.3](#A5.SS3)). Results are presented in Fig. [6.2](#S6.SS2.SSS0.Px4) and Fig. [6.2](#S6.SS2.SSS0.Px4).
Broadly, we can see that on the synthetic dataset, KBLaM shows performance comparable to in-context learning, but with a much smaller memory footprint (Fig. [3](#S4.F3)) and hence better scalability.
On Enron, the out-of-distribution (OOD) data, KBLaM shows degraded performance, but still performs better than zero-shot baseline, indicating that KBLaM is using the KB information effectively.
The observations from sample outputs on open-ended Q&A (Appendix. LABEL:sec:sample_output) also align with the metrics values: On synthetic KB, KBLaM gives answer almost identical to origin triple content (Appendix. LABEL:sec:syn_kb_sample_output), whereas on Enron, KBLaM’s answer captures the meaning of the original text, with more pronounced difference in wording and details (Appendix. LABEL:sec:enron_kb_sample_output).
Note that Enron has

s that look entirely different from that of the synthetic data (see Table. LABEL:tab:syn_kb v.s. LABEL:tab:enron_kb in Appendix. [C](#A3)).
We therefore believe such performance degradation could be alleviated by using a larger and more diverse synthetic dataset or through a synthetic dataset derived from a real-world dataset.

#### KBLaM knows when it cannot answer a question

If the model can refuse to answer a question in case it cannot find relevant information in the context, then the risk of model hallucinations can be largely controlled.
We consider a setting where, given a KB, we ask the model 100 questions in total, out of which 80 questions are answerable, and the other 20 are not.
We then consider the unanswerable questions as the positive class and answerable questions as negative.
To evaluate the model, we use standard binary classification metrics: precision and recall.
We compare the performance of KBLaM with in-context learning, where we explicitly prompt the LLM with the following instructions: “if relevant information cannot be found in the text, please respond I am sorry I cannot find relevant information in the KB”. The results are presented in Fig. [6.2](#S6.SS2.SSS0.Px4).
Broadly, we notice that when the number of triples increases, both KBLaM and in-context learning show decreased precision, i.e. the models start to incorrectly refuse for having not found the information, however, in-context learning shows more drastic degradation. The recall for both methods stays constant, i.e. both methods do not start to hallucinate as we add more triples.

#### Ablation study

We conduct ablation studies over the design choice of KBLaM and present the results in Appendix. [B](#A2):
We find that KBLaM also works with open-sourced sentence encoders, with exact performance depending on the encoders’ capacity (Fig. [7](#A2.F7));
The frequency of adding knowledge tokens can be set less than every layer, which further improve efficiency but could cause performance drop (Fig. [8](#A2.F8)). Lastly, we study the role of knowledge tokens at each layer, where we find that knowledge tokens in earlier layers may serve as “soft prompt” for instruction following (Fig. [9](#A2.F9)).

[Uncaptioned image]: /html/2410.10450/assets/x6.png

[Uncaptioned image]: /html/2410.10450/assets/x7.png

[Uncaptioned image]: /html/2410.10450/assets/x8.png

[Uncaptioned image]: /html/2410.10450/assets/x9.png

## 7 Conclusion

In this paper, we propose KBLaM, which efficiently augments pre-trained LLMs with external knowledge bases. KBLaM represents external knowledge as dense continuous vectors and leverages the independence structure between triples, which offers several advantages: avoiding expensive self-attention over large knowledge sources, enabling dynamic knowledge updates without fine-tuning, improving interpretability, and allowing attention to perform retrieval. These benefits could potentially also be applied to settings using raw text as knowledge representation.

## 8 Limitations and future work

Higher quality synthetic KB  In the paper, we demonstrate the feasibility of training KBLaM with fully synthetic data. However, as shown in Fig. [6.2](#S6.SS2.SSS0.Px4) and  [6.2](#S6.SS2.SSS0.Px4), KBLaM generalizes to out-of-distribution dataset (Enron) but the performance degrades. We believe this can be improved through a larger and more diverse synthetic dataset.
It is also possible to generate the synthetic data seeded from real KB data, such that the adapters can better capture the distribution of the real text.

Information loss in knowledge token  KBLaM encodes sentences in the triple as a fixed-length vector, therefore KBLaM may fail to *precisely* generate the text word by word.
This is acceptable for properties we considered such as “description” or “objectives”. However, this is problematic for cases when we need exact names or numerical values.
One future direction is to introduce a hyperparameter that adjusts the compression rate, to accommodate various amounts of information in a triple.

More sophisticated instruction tuning
Future work could explore complex instruction tuning examples for KBLaM, such as multi-hop, multi-step, or chain-of-thought reasoning about the KB. Since KBLaM incorporates all KB information in the context, similar to in-context learning, we believe it has the potential for more flexible KB reasoning compared to RAG.

## Appendix A Extended related work

In this section, we present additional discussion on work related to KBLaM.

#### Memory augmented language models

Some other works have also incorporated external memory, i.e. a collection of fixed length latent representations, into language models (Grave et al., ; Dai et al., ; Khandelwal et al., ; Wu et al., ). KBLaM’s knowledge tokens can also be seen as external memory generated from the KB. The most similar approach along this line of work is Wu et al. (), which constructs external memory as key-value vector pairs, similar to KBLaM. However, they require training a transformer from scratch and the memory key-value pairs come from training tokens seen in the past, whereas KBLaM’s memory is from an external KB and is augmented into a pre-trained model at inference time.

#### Token compression

KBLaM’s encoder transforms each triple into a continuous vector pair equivalent to 1 token in length, which essentially compresses information from the multiple-token string into one single token. There have been some recent works aiming at compressing a long context into more continuous representations. Ge et al. () replaces short prompts with special tokens to instruct the LLM to do some particular tasks, however they require fine-tuning of the LLM. Mu et al. () encode a long context into memory slots through a fine-tuned LLM, which is a more powerful but also costly compressor, which we believe could in future work be combined with KBLaM’s framework as an alternative to pre-trained sentence encoder.

#### Structured attention

KBLaM also shares similarity with works that use structured attention masks instead of the standard lower-triangular causal attention mask (Ratner et al., ; Cai et al., ; Merth et al., ). These works utilize independence assumptions in the context, e.g. between different randomly-chosen sets of few-shot examples (Ratner et al., ) or documents from different sources (Cai et al., ; Merth et al., ), and let independent contexts not attend over each other, which reduces overhead of attention operations. In KBLaM, we make a similar independence assumption: We assume independence between different KB triples, which provides KBLaM with scalability and interpretability.

## Appendix B Ablation study

We conducted a number of ablation studies on the design choices of KBLaM.

#### Choice of encoders

KBLaM relies on a pre-trained sentence encoder, and the capacity of the encoder can affect the performance. Indeed our results (Fig. [7](#A2.F7)) verify this: We experimented with some open-sourced alternatives from sentence transformer (Reimers & Gurevych, ). Broadly we observe that the commercial OpenAI embedding works better than the open-sourced one, models with higher embedding dimensions perform better than those with lower dimensions.

#### Where to add knowledge tokens

In Fig. [5](#S5.SS0.SSS0.Px1) and [5](#S6.F5), we noted that the attention score from the $15$th layer exhibits accurate retrieval behavior.
Extending this analysis to other layers (Fig. [9](#A2.F9)), we find that they *do not* demonstrate such straightforward accuracy (although further probing might reveal more nuanced patterns).
We suspect that this indicates that the amount of KB information provided by the knowledge tokens vary at different layers provide.
To better understand this, we examining the variation of keys and values at each layer ($\mathbb{V}_{m}[\bm{\tilde{k}}^{l}_{m}]$ and $\mathbb{V}_{m}[\bm{\tilde{v}}^{l}_{m}]$) across triples, we observe that encoder outputs vary minimally in earlier layers.

#### Frequency of knowledge tokens

The previous observation indicates that the knowledge tokens may not be providing information about the KB at every layer, therefore we consider varying the frequency of adding knowledge tokens (denoted by $K$). That is, for certain attention layers of KBLaM, we used [Eq. 3](#S3.E3) instead of [Eq. 9](#S4.E9).
In particular, we considered only adding knowledge every 1, 3 and 10 layers and the results are presented in Fig. [8](#A2.F8).
Broadly, we notice with lower frequency, the model tends not to follow instructions: When $K=10$, KBLaM fails to provide any refusal answer. As such we suspect that knowledge tokens in earlier layers may serve as soft instruction prompts that guides the LLM how to use the KB.

Figure: Figure 7: KBLaM works with other types of pre-trained encoders. Similar to the setting in Fig. [5](#S6.F5) and Fig. [6.2](#S6.SS2.SSS0.Px4), but with sentence encoder of varying capacity.
Refer to caption: /html/2410.10450/assets/x10.png

Figure: Figure 8: Ablation study on knowledge token injection frequency (K). If we have knowledge tokens added to attention , the model performance drops ($K=10$), but the memory overhead decreases and efficiency increases. The BERT score of $K=3$ is comparable to $K=1$ (first two columns), however when $K=3$, the model shows lower precision at refusing answering unanswerable questions (third column), and when $K=10$, the model fails to output refusal answers.
Refer to caption: /html/2410.10450/assets/x13.png

Figure: Figure 9: Knowledge tokens’ embeddings contain various level of KB information across layers. The embeddings in earlier layers vary less among different entities, indicating that they may serve as instruction prompt on how to use the KB. Later layers show a higher degree of variation, indicating that they may be used to provide actual knowledge.
Refer to caption: /html/2410.10450/assets/x15.png

## Appendix C Sample KB

Here we present some example triples from both KBs, where each triple is of format

$$ $(\texttt{<name>};\texttt{<property>};$ $$

## Appendix E Prompt

In this section, we provide all prompts used for KBLaM’s experiments.

### E.1 Prompt for synthetic KB generation

To prompt the GPT-4 for generating a synthetic KB, we begin by setting a system prompt as

Then we construct a list of idea types and a list of object types:

Then for each combination of object type and idea type in the list, we prompt GPT with

which gives us 50 different names. Then in the same context, we further prompt GPT with

which generates the

for three <property>s: description, objectives and purpose. Lastly, we perform one round of polishing, to let GPT diversity the language and style of the generated KB triple

### E.2 Prompt for open-ended Q&A generation

To generate open-ended Q&A instructions, we first take a triple from the KB, construct a simple Q&A from it, and then we feed the simple Q&A into GPT and let GPT augment it into a more complex form via the following prompt

### E.3 Prompt for GPT evaluation of open-ended Q&A

To evaluate the answer quality of open-ended Q&A tasks, we again use GPT. In particular, we ask GPT to score the answer from two aspects: 1. How grounded the answer is given the question and the KB; 2. How reasonable the open-ended part of the answer is.

We begin with a system prompt of

Then, we use the following prompt, which encourages the model to give chain-of-thought reasoning for scoring (Zheng et al., ):

Followed by the prompt, we include 5 examples to better help GPT calibrate the scoring rule

### E.4 Prompt for Llama evaluation

For in-context learning, we use the following prompts for evaluation

For zero-shot learning, we use the following prompt

### E.5 Question template

For simple Q&A, we use the following templates

For multi-entities Q&A involving $G$ triples, { ( <name> g , <property> g , _g)}_g=1^G,weusethefollowingtemplateWhatisthe{}Tellme{},Canyouletmeknow{},Canyou fragments { ( <name> 𝑔 , <property> 𝑔 , _g)}_g=1^G,weusethefollowingtemplateWhatisthe{}Tellme{},Canyouletmeknow{},Canyou

### E.2 Prompt for open-ended Q&A generation

To generate open-ended Q&A instructions, we first take a triple from the KB, construct a simple Q&A from it, and then we feed the simple Q&A into GPT and let GPT augment it into a more complex form via the following prompt

### E.3 Prompt for GPT evaluation of open-ended Q&A

To evaluate the answer quality of open-ended Q&A tasks, we again use GPT. In particular, we ask GPT to score the answer from two aspects: 1. How grounded the answer is given the question and the KB; 2. How reasonable the open-ended part of the answer is.

We begin with a system prompt of

Then, we use the following prompt, which encourages the model to give chain-of-thought reasoning for scoring (Zheng et al., ):

Followed by the prompt, we include 5 examples to better help GPT calibrate the scoring rule

### E.4 Prompt for Llama evaluation

For in-context learning, we use the following prompts for evaluation

For zero-shot learning, we use the following prompt

### E.5 Question template

For simple Q&A, we use the following templates

For multi-entities Q&A involving $G$ triples, { ( <name> g , <property> g , _g)}_g=1^G,weusethefollowingtemplateWhatisthe{}Tellme{},Canyouletmeknow{},Canyou fragments { ( <name> 𝑔 , <property> 𝑔 , _g)}_g=1^G,weusethefollowingtemplateWhatisthe{}Tellme{},Canyouletmeknow{},Canyou

## Peer Reviews

confidence: 3
contribution: 3
initial_rating: 6
presentation: 3
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?
soundness: 3
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.
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.
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

---

confidence: 4
contribution: 2
initial_rating: 5
presentation: 3
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?
soundness: 3
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.
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.
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.

---

confidence: 4
contribution: 2
initial_rating: 5
presentation: 3
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?
soundness: 3
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
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.
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).

---

confidence: 3
contribution: 3
initial_rating: 6
presentation: 4
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)?
soundness: 3
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.
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.
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."

---

confidence: 4
contribution: 2
initial_rating: 3
presentation: 3
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?
soundness: 2
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.
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.
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).

Meta-review: 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).

Decision: Accept (Poster)