KBLaM: Knowledge Base augmented Language Model
Authors: Xi Wang
Taketomo Isazawa
Liana Mikaelyan
James Hensman
Venue: ICLR 2025 Poster (2025)
Keywords: Large language model; Knowledge augmentation;
Primary area: foundation or frontier models, including LLMs

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

Published as a conference paper at ICLR 2025
KBLAM: KNOWLEDGE BASE AUGMENTED LANGUAGE MODEL
Xi Wang ** Johns Hopkins University
Taketomo Isazawa
Microsoft Research t f--isazawattmiccroooffe
xwang4578cs. jhu. edu
Liana Mikaelyan
James Hensman
Microsoft
Microsoft Research jameshen www.ssseemsnnnnoneeooooooooeeeee
Jm.keellaaneniccroffft.co com
ABSTRACT
In this paper, we propose Knowledge Base augmented Language Model (KBLAM), a new method for augmenting Large Language Models (LLMs) with external knowledge. KBLAM works with a knowledge base (KBB) 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 specialized rectangular attention mechanism. Unlike Retrieval- Auggeeele 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 retraining. Experiments demonstrate KBLAM's effectiveness in various tasks, including question guestion-answering answering and open open-ended reasoming, while providing interpretable insights into its use of the augmented knowledge. Code and datasets are available at https: : / thub. com/mi crosoft/ /KBLaM/
INTRODUCTION
Large language models (LLMs) have demonstrated impressive knowledge and reasoning capabilities. However, in many scenarios, users need to augment L.LMs with external knowledge, particularly when concept definitions differ from or extend beyond the information stored in the L.M''s 's parameters. One straightforward approach is to perform supervised fine-tuning the external corpus (Hu et al., 2021; Liu et al., 2024c) 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.c. 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., 2020). RAG assumes external knowledge is stored unstructured documents, then given input prompt, RAG first uses a retriever module to extract relevant corpus chunks from the documents. These extractions then concatenated the input question as a prompt and fed into the LLM (Fig. top). This approach allows RAG to utilize knowledge from extensive documents while maintaining small context size, overcoming the context limitations of many early- generation LLMs.
Recently, with the emergence of long-contedt language models such as GPT4 (Achiam et al.. 2023) and Gemini (Team et al.. 2023), 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. middle). Several recent works (Lee et al., 2024; Li et al., 2024a; Yu et al., 2024) have explored whether long-context models can replace RAG in long-context reasoning tasks. particular, Lee et al. (2024) argues that in-context learning has
Equul contribution. Work done during an internship Microsoft Research.

---

Published as a conference paper at ICLR 2025
External
Retriever Module
Relevant corpus
RAG
Separate componants, Cheap
Tokenizer
Retrieve
Summarize
Prompt
Tokenizer
Concatenation
In-context learning
Single component, Expensive
Summarize
Prompt
Retrieve
KB
Knowledge
KBLAM (ours) Single component, Chesp
Construct KB Summarize
Encode
operation
Prompt
Tokenizer
Rectangulaz
attention
LLM
Retrieve
Figure 1: Overview of the KBLAM pipeline and comparison with existing approaches. KBLAM augments knowledge into pre-trained LLM in the form of knowledge tokens, a set of continuous key-value vectors, using modified rectangular attention structure. Unlike RAG, KBLAM does not rely on separate retriever module at inference time and unlike in-context learning, KBLAM computation and memory overhead scales linearly rather than quadratically with the size of the K.B.
potential advantages, such 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 new regime for augmenting pre- trained LLMs with external knowledge: Knowledge Base augmented Language Model (KBLAM, Fig. bottom). We consider a setting where the unstructured external corpus is transformed into structured knowledge base (KB) through existing tools, which summarizes the critical information in the data and generates knowledge triples (Eq. (1)) containing an entity name (<name>), property <property>> and value (<val ue>). Then, KBLAM's design fully leverages the structure of the KB, achieving efficient integration of external knowledge into LLMs (Sec. 4).
First, KBLAM maps each triple into fixed-lennthh key-value vector pair, referred to knowl- edge 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 <propert rty> into a key vector, which serves as an identifier, mimicking the key embedding of token; <value> into value vector, which provides the actual content, similar to the value embedding of a token.
Then, KBLAM uses the structure between triples augment knowledge tokens into the LLM's attention in scalable and dynamic way using a simple rectangular attention structure: Triples with different <name> and <property> can be considered 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 the number of triples, unlike in-context learning's quadratically growing overhead, giving KBL.MM much better scalability. Additionally, the independence allows us to update/remove/add triple by only modifyyng its corresponding single knowledge token without any further changes, which not achievable in, c.g. standard KV cache mechanism. Perhaps more importantly, the attention matrix under this design is highly interpretable. As show in Fig. the model's use of the knowledge tokens can be directly inspected through the attention score.
represent independent pieces of
Lastly, show that the linear adapters be learned using instruction tuning on purely synthetic data (Sec. 5) while being able to generalizes to real data. Different from supervised fine-tuning, which aims memorize knowledge into model weights, the learning process of KBI. KBLAM M aims at
2

---

Published as a conference paper at ICLR 2025
finding a projection between the mre-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 important.
Compared with RAG, KBLAM incorporates all external information in the context and answers input prompts in 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 not present in the KB, improving the reliability of the model and reducing hallucinations.
In summary, KBLAM is 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 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 K KBs, which can help future research in augmenting LLM with KB, as well as in other topics such as long-contenn language models, hallucination detection/reduution, and structured attention.
2 RELATED WORK!
Retrieval augmented generation (RAG) RAG (Lewis et al., 2020) is one of the most successful approaches for augmenting external knowledge into LLMs, and many approaches have been proposed to perform RAG, such similarity search and text2Sql (Qin et al., 2022). KBLAM can also be understood as the model doing RAG in an implicit way: given input prompts, the rectangular attention compares the sequence's '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 then added the hidden state, augmenting next-token generation with information from the KB. However, KBLAM differs from RAG in that it does not rely on separate retrieval module. Instead, everything is carried out by attention in an end-to-end fashion.
Structured attention KBLAM also shares similarities with works that use structured attention masks instead of the standard lower- weir-triangular causal attention mask (Ratner et al., 2022; Cai et al., 2023; Merth et al., 2024). These works utilize independence assumptions in the context, e.g. between different randomly-chosen sets of few-shot examples (Ratner et al., 2022) or documents from different sources (Cai et al., 2023; Merth al., 2024), and let independent contexts not attend over each other, which reduces the overhead of attention operations. In KBLAM, we also utilize similar independence assumption among inputs: We assume independence between different KB triples, which provides KBLAM with scalability and interpretability.
Key-value (KV) cache mechanism Given context string, the KV cache mechanism (Pope et al., 2023) 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.
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
{<<name>>> <property> mi <value>..))M
In the rest of the paper, we will refer to a set of triples of the format Eq. (1) as a KB and each triple inside it as a knowledge triple. Examples of such KB is shown in Table. and 2 in Appendix E.
We provide more extensive related work discussion in Appendix A
3

---

Published as a conference paper at ICLR 2025
The KB construction process summarizes information from multiple documents and organizes it in a structured way. Importantly, the structured nature of the KB enables us to augment LLM with information from documents in an efficient way, which we discuss in the next section.
is also worth noting that KB construction is not the focus of our paper: KBLAM works on KB KBss: A synthetic
generated from corpus using existing tools. In the paper, we mainly focus KB generated via GPT, and a real KB constructed from the Enron email (Klimt & Yang, 2004) dataset. Details of the two KBs will be discussed in detail in later sections.
Self-attention layer A rescoder-based transformer is mostly composed of multiple self- layers (Vaswani, 20.7). Consider a transformer of layers. Each layer has three projection heads: W R DxD Wk E RDDD W RDxD E (l,..., LJ, 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 x² [ari:-. -..-.... wy.... ... x'y]" VR input, then for the nihh token embedding, it is first transformed by the projection heads into three vectors
d = 2
k" W
where we refer Ly ky and vy as the wth token's query, key, and value embedding at the /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, y² = .u....... Va..... I E ERNDD is computed as
y/ Ei= exp (wn,s) of Σ exp (Wn,s)
where Wn,s (d., /VD
(3)
where (: denotes the inner product of two vectors. After y acquired, it often fed into a feedforward network (FFN) with output dimension of D for further transformation.
standard implementation of self-attention would have a time complexity of O(N2 and memory complexity of ((N2) for computing and storing all Wn.i. Additionally, the FFN would also introduce a significant computation overhead of magnitude O(ND?). Due to these factors, as the sequence gets longer, self-attention faces high memory consumption and slow computation.
4 AUGMENTING LLM WITH THE KB
In this section, we discuss how KBLAM augments an LLM with KB. The process has two steps (visualized in Fig. 2): 1. We convert each triple in the KB from the string form into continuous key-value vector pair, referred to as 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 KB in the format of Eq. (1), for each triple, we first adopt pre-trained sentence encoder model, denoted as f(() which converts string into P-dimension continuous embedding vector. Through the encoder, we convert each triple into base key embedding and base value embedding. In particular, for the ruth triple, we have:
km f(hhe <property>,, m of <name>.)) RR,
Vm f(<value>m) E ERP.
(4)
Next, we introduce a linear key and a linear value adapter
WK E LxDXP
W. E LDDXXP
(5)
where L denotes the number of attention layers in a model. With the adapters, we map km and vm from the sentence encoder's 's space LLM's key and value embedding space at each attention layer. In particular, for the mth knowledge triple, we transform its base key and value embedding into
Am = [ m3
K
J. m = Wak,, E ERI.
m
im [ .......
J! L T Wyvm e RLXD
(7)
Since each ie', and oy share the same size with LLM's key and value embedding (i.c. ky and v/ in Eg. (2)). they can be directly incorporated into attention's computation. Therefore, the encoding
4

---

Published as a conference paper at ICLR 2025
In: Triple ID, between and M.
Frozen weights.
Learned weights
A: Total number attention layers
Key: spropenntyhg, of Chamee,
Pro trained sentence encader
Linear key edapter
Value:-yyalion Loa
T Pro trained sentence enccdar
Linear value edapter
Pinn
Save on disk
(a) Step 1. K.B encoding For each triple from a KB (Eq. (1)) of M triples, we first construct a key and value string (left most white boxes). Then both strings are processed by pre-trained sentence encoder followed by learned linear adapter (middle
blue boxes). The acquired knowledge tokens (Eqs. (6) and (7), right- most small boxes) are then stored on disk for later usage. Note that all km and e' s have a fixed dimension D, i... identical to that of the key and value embeddings of token, regardless of input length.
Input& & output the +t attention layer KBLaM
2emee exp (( /DDD Pin Ei=; exp (1gh, KDDD) vs
C Pre-trained
P/1
exp km.)D) ET_, exp( k))/D)
It
y
where
4%
8 computed using learned query head Wi
Rectangular attention matrix (before softmax)
A:
Prompt queries
fi
Knowledge tokens
4)/vD
(oaded from memory
Pro trained QKV head
= (Masked)
X
X
KB keys
Prompt keys
(b) Step 2. Augmenting knowledge tokens into attention For layer I, given an sequence of N D-dimensional embeddings from the prompt (æij.... ...i...... axx. e.g. a question about the KB), augmented with M knowledge tokens as context (K 6'm,) 3. }m= KBLAM's attention outputs N embedding vectors with each element D E R computed via Eq. (9) (top right panel) under rectangular attention matrix (bottom right panel): Blue regions show the extra components introduced by KBLAM whereas the parts are from standard self-attention. Note that input sizes from the blue parts in the equation and the matrix, which scales as O(M).
FFN stay unchanged, the only additional overhead introduced by KBLAM comes
Figure 2: Overview of KBLAM's KB augmentation process
process is equivalent to converting a knowledge triple from a string of multiple tokens into a special token, whose key and value embeddings each layer are generated through encoder, rather than through self-attention. As such we refer to (km, 0m) pair knowledge token.
We visualize the process in Fig. 2a. This encoding process is applied to all triples in the KB, which transforms the information from KB into collection of knowledge tokens
{( (<<ame>>,, m. <property> <value>,,) mm)m=
Encode
{(ks 6m.) J m==
(8)
Injecting knowledge token into prompt tokens After converting KB into a collection of knowledge tokens we inject the information the KB into the each attention layer of a pre-trained LLM through modified attention structure?. In particular, consider an input prompt of N tokens, the /th attention layer takes in x = [-]..... ......, "........" xyy|" and the 1th dimension of the knowledge tokens { ) the output of each attention layer, ÿ' = [v(...., V Ul....., [x] E NDD given by
Rectangular Attention:
.......Th
ΣM exp (wi (i,m)) m + E exp (w of
(9)
Σ EMEn exp (ih) (m)) En= exp (wh)))
M
where
5 n. = (g., R.. ब. = W/En.
w/ = (g.. //)/
(10)
and W/ T DxD is a learnable query head initialized from W!
We use symhols with tilde to denote additional attention components introduced by KBLAM.
5

---

Published as a conference paper at ICLR 2025
Memory used (GB)
We name the attention structure of Eq. (3) rectangular attention as the attention matrix is of rectangular shape (bottom-right right panel in Fig. 2b). 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 + NN) X N,, where M the number of triples in the KBB. Additionally, if knowledge tokens are not introduced, i.e. M = 0, Eq. (9) falls back to Eq. (3) and we recover the pre-trained L.M..
75
Total GPU
memory
50
Zero-shot
KBLaM
25
In-context
leaming
102
103 10*
# #Triplets
Figure 3: Memory overhead of differ- ent methods. Given a KB of M triples (with each triple K-token long on av- erage). In-context learning's memory scales with (K)) whereas KBLAM's memory scales with M.
The memory and time complexity of the rectangular atten- tion is O((M + N)N) and ((M + D) respectively. In real settings, it is likely that M >> 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& 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) 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., 2024b). Additionally, since all knowledge tokens are encoded independently, when the KB is updated, e.g. new knowledge triples are created 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.
triple is updated, one could simply generate a
Note that KBLAM's rectangular attention differs from cross-attention (Vaswani, 2017: Chen et al., 2021), which operates on a source sequence with key and value embeddings and target sequence with only query embeddings, where tokens within each sequence do not attend to one another. In setting, if we consider the KB as the source and the input prompt as the target, cross attenttono corresponds to retaining only the summations over M in Eq. (9). However, KBLAM's attention also contains self-attention over the prompt tokens, given by the summations over in Eq. (9).
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 " m E {1..... 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) 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 constant shift to the KB part's unnormalized attention score
wh n,m log C' - log M + (T.. //DDDD
(1))
where C' E R+ is a hyperparameter. This is equivalent to multiplying the summations over M by M such that the total contribution of the KB part stays roughly unchanged as scales. In our experiments, we fix C' as 100, which is the largest number of triples (i.e. M) seen at training time.
5 KB INSTRUCTION TUNING
The learnable parameters in KBLAM comprise the weights for the weights for the linear adapters/- Inspired by recent works in multi-modal language models (Liu et al., 2023), we use instruction tuning for parameter learning. Specifically, given KB, we generate question-answer pairs about the KB (denoted as Q and A respectively) using formatted strings or GPT, then we optimize 0 using:
heads 0 {{ WK, K, Wr, W,,K (W)) }
max Po. (A KB).
(12)
where & denotes parameters of the pre-trained LLM, encompassing QKV heads, FFN parameters and embedding/outpu layer weights. Notably, KB instruction tuning preserves the base L.MM's reasoning

---

Published as a conference paper at ICLR 2025
Question
Question item
Question on KB item 4 and
Question irrelevant to the KB
item 2
0.4
12
Query index
0.3
16
0.2
20
0.1
24
0.0
2
KB index
KB index
index
KB index
Figure 4: KBLAM's attention matrix is interpretable. Consider toy KB of 10 triples, the heatmap shows the attention weights under different questions. We select the 15th attention layer and visualize the sost-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 's query.
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 al., 2023; Mireshghallah al., 2022). 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 memorize specific information, but learn projection from the pre-trained sentence encoder space to the LLM's semantic space. Consequently, the exact content of the KB less critical, provided the training data encompasses a diverse range of text.
Generation of synthetic KBB To synthesize a KB, we first use GPT generate 50 <name>s based 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 <value> of three <property>>: "description", "objectives", and "purpose" for each name, in the same conversation context. Crucially, we instruct GPT to generate <value>s uncorrelated with names, ensuring the information comes from the KB rather than the LLM's predictive ability. This yields 45K names and a KB of approximately 135K triples. We provide example triples in Table. in Appendix. E, the full list of types and prompts in Appendix. G.1 and release the resulting synthetic KB dataset for research purposes.
Generation of instructions The construction of the instruction tuning dataset crucial, as reflects our intended use of the KB and the desired behavior of KBLAM under various scenarios. We considered the following types of instructions (examples shown in Fig. 12 Appendix. F):
Simple Q& about single <name> and <property> with <va lue> the answer. Multi-entities Q&A involving multiple <name>s and <property>s, with their corresponding <value>s as answers.
Open-ended Q&A similar to simple and multi-entities Q&A. but including additional open- ended question component that requires open-ended reasoning about <value>, Reference answers for these questions are generated by GPT. The exact prompt is detailed in Appendix. G.2.
Unanswerable questions irrelevant any <name> <property> in the KB. For such cases, the standard response is "Sorry, cannot find relevant information in the KB."
For simple, multi-entities, and unanswerable questions, both questions and answers are generated using formatted strings, with templates of various styles such as "Can you inform me..." or "What is the..." " A comprehensive list of these templates is provided in Appendix. G.5.
EXPERIMENTS
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, show that for questions with answer in the KB, the model can refuse to answer questions, with "over-refusal" behavior occurring later than in in-context learning. Lastly, conduct ablation studies understand the design choices of KBLAM.
7

---

Published as a conference paper at ICLR 2025
6.1 EXPERIMENT SETTING
Model specification For all experiments, we use the instruction fine-tuned version of Llama3 8B (Dubey et al., 2024) 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)).
Optimization setting We initialize the key and value adapters: WK and Wy randomly, and we initialize additional query heads W at each layer from the pre-trained weights W Optimization is conducted using AdamW (Loshchilov, 2017) with a step size of 5 x 10-4 and a cosine learning rate decay to 5 x 10 6 for 20K iterations. Each iteration uses a mini- of 400 Q&A pairs, composed of 20 micro-batches of 20 samples. The instruction tuning is performed on single 80GB A100 GPU under 6 without any parameter-effictte 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)). These samples are derived from the first 120K triples of the synthetic KB described in Sec. 5. 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 sample-specifie KB. Depending on the instruction type (simple, multi-entities, open-en -edeeedddd 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 -atehes 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)) offline for all triples in the synthetic KB.
Evaluation dataset For evaluation, we considered the following two KB datasets*
Synthetic data The validation set of the synthetic KB, i.c. the 15000 triples not used for training. Enron A KB constructed from the Enron (Klimt & Yang, 2004) dataset, an open-sourced corporate email dataset. Following (Wu et al., 2024), constructed KB of triples from the Enron emails data by fine-tuning a small language model, which we then cluster to remove duplicates (Winn al., 2018; 2021).
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 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 Given question about the mth triple in the KB, we expect certain tokens in the question, particularly those containing the keyword <name>m. to attend more strongly to the mith knowledge token compared to others. This would be reflected in larger values /h n. = (Ā., > for these relevant tokens. We observe such patterns in practice, as visualized in Fig. 4. In particular, we consider 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 answer input questions in an interpretable way, for any query, the top-K triples with the highest attention score form supporting evidence. This behavior suggests that K KBLAM's attention mechanism functions implicitly as a retriever. We can quantitatively assess its retrieval performance by treating the attention scores, w! nm m E I,..., M, as classification score. In particular, we again extract the averaged- -ereggdd-ove--heee -ver--eedde attention score at the 15th layer, and we evaluate the top-1 and top-5
All KBs used for training and evaluation are released together with the paper.
8

---

Published as a conference paper at ICLR 2025
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 presents these results, demonstrating that KBLAM's attention remains highly accurate even with large number of triples, both for in-distribution and out-of diittii-tittittionnn (OOD) data. Notably, this performance is achieved solely through instruction tuning using paired data, without any explieit regularization or retrieval objectives. In fact, imposing regularization on the attention structure would be challenging and potentially counterproductive as it unclear, egg. which tokens, attention heads or attention layers should be regularized.
Synthetic data
Enron
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),
Top-1 Acc
Top-5 Acc
1.00
0.75
and open-ended Q&A. For the first two types, we evaluate the output quality using BERT score (Zhang* et al., 2020). For open-ended questions, we use GPT-4 to score the out- put quality between 0 and 5 (full prompt presented in Ap- pendix. G.3). Results are presented in Fig. 6a and Fig. 6b. 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) and hence better scalability. On Enron, the out oit-of-dde of-distribution (OOD) data, KBLAM shows degraded performance, but still performs better than zero-shot baseline, indicating that KBLAM using the KB information effectively. The observations from sample outputs on open-ended Q&A (Appendix. H) also align with the metries values: On synthetic KB, KBLAM gives answer almost identical to origin triple content (Appendix. H) whereas on Enron, KBLAM's s answer captures the meaning of the original text, with more pronounced difference in wording and de- tails (Appendix. H.2). Note that Enron has <va 1 ue>s that look entirely different from that of the synthetic data (see Table. 1 v.s. 2 in Appendix. E). We therefore believe such performance degradation could be alleviated by using larger and more diverse synthetic dataset or through synthetic dataset derived from real-world dataset.
0.50
0.25
10a
10a
10a
10a
#Triples
#Triples
Figure 5: Through instruction tuning, the attention shows retrieval behav- ior. 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, 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).
KBLAM knows when it cannot answer 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
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 KBLAMM with in-context learning, where we explicitly prompt the L.M 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 presented in Fig. 6c. Broadly, we notice that when the number of triples increases, both KBLAM and in-context learning show decreased precision, i.e. the models start 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: We find that KBLAM also works with open-sourced sentence encoders, with exact performance depending on the encoders' capacity (Fig. 7); The frequency of adding knowledge tokens can be set less than every layer, which further improve efficiency but could cause performance drop (Fig, 8). 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).
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
9

---

Published as a conference paper at ICLR 2025
Zero-shot
In-context learning
KBLaM
Synthetic data (one entity)
Synthetic data (two entities)
Enron (one entity)
Enron (two entities)
BERT Score (F1)
1.0
1.0
0.7
0.7
0.8
0.8
0.6
0.6
0.6
0.6
0.5
0.5
1.00
102
10
102
102
10
102
10a
10
102
10a
104
# #Triples
# Trriples
# Tripples
# Tripples
(a) Simple and two-entity Q&& Evaluated by BERT score (F1), which measures word-level cosine similarity between model outputs and reference answers by on embedding space. In-context learning does not have training process, showing identical performance Enron and synthetic data, we omitted its Enron results for graph visibility.
Synthetic data
Enron
Precision
Recall 1
1.0
GPT score
5.0
2.5
0.6
0.2
102
10a 10a
102
10a
104
0
500
1000
500 # +Triples
1000
#Triples
# #Triples
#Triples
(b) Open-ended Q&A Scored by GPT-4 between (c) Unanswerable questions Precision and recall O and 5, dot and bar shows the mean and standard for detecting questions without relevant informa- error over 5 random seeds.
tion the KB.
Figure 6: KBLAM can reason about the KB. (a) and (b): Performance on Q&A tasks evaluated by BERT and GPT-4 scores. KBLaM shows comparable quality to in-context learning on synthetic data while using less memory (Fig. 3). Some degradation occurs with Enron data, which is out-of- distribution. Zero-shot learning fails to provide sensible answers as the questions are not answerable using model's internal knowledge; (c): We consider unanswerable questions with no relevant triples in the KB as positive class (20 percent of 100 questions) and answerable questions as negative class. We then use precision and recall measure whether the model knows say that it "Cannot find relevant information" when is supposed to say so. Both in-context learning and KBLAM show "over-refusall", i.e. failing to respond to answerable questions, when the number of triples grows larger (i.e. degradation in precision), however, KBLAM degrades slower.
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 settings using raw text as knowledge representation.
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. 6a and 6b, KBLAM generalizes to duitt--disttibution 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 K.B 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 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 hyperparameter that adjusts the compression rate, to accommodate various amounts of information in triple.
More sophisticated instruction tuning Future work could explore complex instruction tuning examples for L,,,,,, such as multi-hop, multi-step, or chain-of-thought reasoning about the KB. Since KBLLM incorporates all KB information in the context, similar in-context learning, we believe it has the potential for more flexible KB reasoning compared to RAG.
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGMENTS
We thank Mathew Salvaris for his valuable contribution to polishing our work and assistance in cleaning up and open-soureing the codebase.
REFERENCES
Josh Achiam, Steven Adler, Sandhini Agarwal, Lama Ahmad, 1ige Akkaya, Florencia Leoni Aleman, Diogo Almeida, Janko Altenschmidt, Sam Altman, Shyamal Anadkat, et al. Gpt-4 technical report. arXiv preprini arXiv:2303.08 087774. 2023.
Sebastian Borgeaud, Arthur Mensch, Jordan Hoffmann, Trevor Cai, Eliza Rutherford, Katie Millican, George Bm Van Den Driessche, Jean-Baptise Lespiau, Bogdan Damoc, Aidan Clark, et al. Improving language models by retrieving from trillions of tokens. International conference on machine learning, pp. 2206 22400. PMLR, 2022.
Tianle Cai, Kaixuan Huang, Jason D Lee, and Mengdi Wang. Scaling in-context demonstrations with structured attention. arXiv preprint arXiv:2307.0269/ 2023.
Chun-Fu Richard Chen, Quanfu Fan, and Rameswar Panda. Crossvit: Cross-attention multi-scale vi- sion transformer for image classification. In Proceedings the IEEE/CVF international conference on computer vision, pp. 357-366, 2021.
Zihang Dai, Zhilin Yang, Yiming Yang, Jaime Carbonell, Quoe Le, and Ruslan Salakhutdi- Transformer-xl: Attentive language models beyond fixed-length context. arXiv preprint 19X1v29001.02600, 2019.
Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Amy Yang, Angela Fan, al. The Ilama 3 herd of models. arXiv preprint arXiv:2407.21788 2024.
Peng Gao, Jiaming Han, Renrui Zhang, Ziyi Lin, Shijie Geng, Aojun Zhou, Wei Zhang, Pan Lu, Conghui He, Xiangyu Yue, al. Llama- adaater v2: eeffiietr-efficient visual instruction model. arXiv preprint arXiv:2304. 100004.55000 15010, 2023.
Tao Ge, Hu Jing, Lei Wang. Xun Wang, Si-Qing Chen, and Furu Wei. In-context autoencoder for context compression in a large language model. ICLR, 2024.
Edouard Grave, Armand Joulin, and Nicolas Usunier. Improving neural language models with a continuous cache. arXiv preprint arXiv: 2ee.:1112..4444 2016.
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:2/06.00663 2021.
Gautier Izacard and Edouard Grave. Leveraging passage retrieval with generative models for open domain question answering. arXiv preprint wwXiv:2007.01282 2020.
Gautier Izacard, Patrick Lewis, Maria Lomeli, Lucas Hosseini, Fabio Petroni, Timo Schick, Jane Dwivedi-Yuu 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.
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.
Bryan Klimt and Yiming Yang. The enron corpus: A new dataset for email classification research. In European conference machine learning, pp. 217-226. Springer, 2004.
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. 13/2/, 2024.
11

---

Published as a conference paper at ICLR 2025
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 genera- tion for knowledge-intensive nlp tasks. Advances in Neural Information Processing Systems, 33: 9459-9474, 2020.
Zhuowan Li, Cheng Li, Mingyang Zhang, Qiaozhu Mei, and Michael Bendersky. Retrieval augmented generation long-context Ilms? comprehensive study and hybrid approach. arXiv preprini arXiv:2407.16833, 2024a.
Zhuowei Li, Zihao Xu, Ligong Han, Yunhe Gao, Song Wen, Di Liu, Hao Wang, and Dimitris N. Metaxas. Implicit co-contxt learning. In The Thirteenth International Conference on Learning Representations, 2025.
Zonggian Li, Yinhong Liu, Yixuan Su, and Nigel Collier. Prompt compression for large language models: A survey. arXiv preprint arXiv:2410. 12388, 2024b.
Haotian Liu, Chunyuan Li, Qingyang Wu, and Yong Jae Lee. Visual instruction tuning. In seirth- Conference Neural Information Processing Systems, 2023. URL https:
//
net. forum? i -iwiid022220 #0H2xGHlkw.
Haotian Liu, Chunyuan Li, Yuheng Li, and Yong Jae Lee. Improved baselines with visual instruction tuning. In Proceedings of the IEEE/CVF Conference Computer Vision and Pattern Recognition, pp. 26296-26306, 2024a.
Nelson Liu, Kevin Lin, John Hewitt, Ashwin Paranjape, Michele Bevilacqua, Fabio Petroni, and Percy Liang. Lost in the middle: How language models long contexts. Transactions of the Association for Computational Linguistics, 12:157-173, 2024b.
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 arXhv:2402.09353 2024c.
I Loshchilov. Decoupled weight decay regularization. arXiv preprint arXi:: 2017.
Thomas Merth, Qichen Fu, Mohammad Rastegari, and Mahyar Najibi. Superposition prompting: Improving and accelerating retrieval-augmennted generation. arXiv preprint arXiv:2404.066 06000, 2024.
Fatemehsadat Mireshghallah, Archit Uniyal, Tianhao Wang, David Evans, and Taylor Berg- Kirkpatrick. Memorization in nlp fine-tuning methods. arXiv preprint arXiv:2205.12506, 2022.
Jesse Mu, Xiang Li, and Noah Goodman. Learning to compress prompts with gist tokens. NeurIPS, 2024.
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.
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.
Nir Ratner, Yoav Levine, Yonatan Belinkov, Ori Ram, Inbal Magar, Omri Abend, Ehud Karpas, Amnon Shashua, Kevin Leyton- and Yoav Shoham. Parallel context windows for large language models. arXiv preprint 22..:::22/22/04444 2022.
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:: /ttp:///axiivv//i org. www./o///////oooooooo/.....aaaaa 10084.
abs /190.
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. 1805, 2023.
12

---

Published as a conference paper at ICLR 2025
A Vaswani. Attention is all you need. Advances in Neural Information Processing Systems, 2017.
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.
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.
Haolun Wu, Ye Yuan, Liana Mikaelyan, Alexander Meulemans, Xue Liu, James Hensman, and Bhaskar Mitra. Structured entity extraction using large language models. arXiv preprint 24l.::2002.044377, 2024.
Shengqiong Wu, Hao Fei, Leigang Qu, Wei Ji, and Tat-Seng Chua. Next-gpt: Any-to-any multimodal Ilm. arXiv preprint arXiv:2309.0553 2023.
Yuhuai Wu, Markus N Rabe, Delesl Hutchins, and Christian Szegedy. Memorizing transformers. arXiv preprint 2rXiv:2003.089) 2022.
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 Linguisties (Volume I: Long Pupers), pp.
2023.
Howard Yen, Tianyu Gao, and Danqi Chen. Long-context language modeling with parallel context encoding. arXiv preprint arXiv:2402.16617, 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.
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 ar arXiv:2310.06714 2023.
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: www.oeeeeeeeeeeeee net.
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 Computer Vision and Pattern Recognition, pp. 6586-6597, 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 ilm-as-a-judge with mt-bench and chatbot arena. Advances in Neural Information Processing Systems, 36:46595-46623, 2023.
Yufan Zhuang, Chandan Singh, Liyuan Liu, Jingbo Shang, and Jianfeng Gao. Vector-ICL: In-context learning with continuous vector representations. In The Thirteenth International Conference on Learning Representations, 2025.
13

---

Published as a conference paper at ICLR 2025
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., 2016; Dai al., 2019; Khandelwal et al., 2020; Wu et al., 2022). 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. (2022). which constructs external memory as key-value vector pairs, similar to KBLAM. However, they require training 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 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 long context into more continuous representations (Li et al., 2024b). Ge et al. (2024) replaces short prompts with special tokens to instruct the L.M to do some particular tasks, however they require fine-tuning of the L.M.M Mu et al. (2024) encode a long context into memory slots through fine-tuned I.M,, which a more powerful but also costly compressor, which we believe could future work be combined with KBLAM's framework as alternative to pre-trained sentence encoder.
Multi-modal language model Recent advances in Multi-modal Language Models (Liu et al., 2023; Gao et al., 2023; Liu et al., 2024a; Zhao et al.. 2023; Wu et al., 2023, MLMs), demonstrate the possibility of incorporating data of other modalities (c.g. images, videos) into pre-trained LLM through learned adapters and instruction tuning. With KBLAM, we can consider a KB 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.
Augmenting information into an LLM in a continuous form RAG and in-context learning incorporates external information into LLMs in the form of 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, 2020; Izacard et al., 2023; Ye et al., 2023) or using cross-attention inject these vectors into the decoder (Borgeaud et al., 2022). However, these works focus on training a new model from scratch, while KBLAM is applied to augment pre-trained LLM with knowledge. Concurrent works (Li et al., 2025; Zhuang et al., 2025) consider augmenting few-shot demonstration information into a pre-trained LLM in the form of continuous vectors, however, KBLAM focuses on augmenting factual knowledge from a large scale KB rather than tosk-solving ability via few-shot examples. Recent work (Yen et al., 2024) 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 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.
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 results (Fig. 7) verify this: We experimented with some open-sourced alternatives from sentence transformer (Reimers & Gurevych, 2019). Broadly we observe that the commercial OpenAl 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. 4 and 5, we noted that the attention score from the 15th layer exhibits accurate retrieval behavior. Extending this analysis to other layers (Fig. 9), 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
14

---

Published as a conference paper at ICLR 2025
all-MiniLM-L6--2
all-mpnet- base-v2
ada-002
text-embeddingglary
(Open source, P 384)
(Open source, P 768)
(API only, P 1536)
(API (API only, P 3072)
Synthetic data
Enron
Synthetic data
Enron
Top-1 Acc.
1.0
1.0
BERT Score (F1)
0.7
0.9
0.5
0.5
0.0
0.6
102
10a
102
10
102
10a
10
10a
# Triplets
#Triplets
# #Triples
# #Triples
Figure 7: KBLAM works with other types of pre-trained encoders. Similar to the setting in Fig. 5 and Fig. 6a, with sentence encoder of varying capacity.
K=3
K=10
BERT Score (F1) (Synthetic data) 1
BERT Score (F1) (Enron) f
Refusal precision
Memory used (GB)
0.9
0.8
25
0.65
0.5
0.8
20
0.60
0.2
102
103
10a
103
10a
10a
10a
10a
# #Triplets
# #Triplets
# #riiplets
# Trippeess
Figure 8: Ablation study on knowledge token injection frequency (K). we have knowledge tokens added to attention the model performance drops (K = 10), but the memory overhead comparable
decreases and efficiency increases. The BERT score of K two columns), however when K 3, the model shows lower precision at refusing answering unanswerable questions (third column), when K 10, the model fails to output refusal answers.
(first
the variation of keys and values at each layer (V [k] and V. (5/1)) 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 consider varying the frequency of adding knowledge tokens (denoted by K). That is, for certain attention layers of KBLAM, we used Eq. (3) instead of Eq. (9). particular, we considered only adding knowledge every 1, 3 and 10 layers and the results are presented in Fig. 8. Broadly, we notice with lower frequency, the model tends 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.
follow instructions: When K 10, KBLAM fails to provide
C EXTENDED EXPERIMENT RESULTS
Comparison between KBLAM'S retrieval performance and BM25 We also ran experiments with BM25, lexical search algorithm, as baseline for the top-1 and top-5 retrieval accuracy, and the results can be seen in Figure 10a. We find that, on the synthetic dataset, KBLaM achieves comparable t accuracy BM25 across all KB sizes.
As with KBLAM, BM25's performance also degrades on Enron, although the degradation less sever, leading to a larger gap between the performance of the two methods as the number of triples increases.
Perturbing the synthetic dataset dataset, so that the entity's 's name in the question no longer exactly matches the entity's 's name in the knowledge base (e.g. entity with name "Fuler's Edifice" in the knowledge base may be referred to
We also experiment with setting where we perturb the synthetic
15

---

Published as a conference paper at ICLR 2025
Key embeddings
Value embeddings
Top-1
Top-5
Variance acorss KB entities
x10
1.0
0.3
Accuracy
1.5
0.2
0.5
1.0
0.1
0.0
15
30
15
30
15
30
15
30
Layer idx
Layer idx
Layer idx
Layer idx
Figure 9: Knowledge tokens' embeddings contain various level of KB information across layers. The embeddings in carlier 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.
BM255 (Synthetic ditaa
KBLaM dyytthetic data) (Enron) Top-5 Acc.
KBLaM
BM25
BM25 (Erron)
Top-1 Acc.
Top-1 Acc.
Top-5 Acc.
1.0
1.0
0.5
0.5
0.0
4mm
0.0
102
10
10a
102
10
102
10
102
10a
10
# Trriples
#Triples
#Triples
#Triples
Retrieval accuracy on synthetic KB and Enron.
(b) Retrieval accuracy perturbed synthetic KB.
Figure 10:
KBLaM shows retrieval accuracy comparable with BM25. We consider a setting similar to that of Fig. 5, additionally including BM25 as baseline. On synthetic data, KBLAM's top-1 accuracy is slightly worse than BM25 but shows very close top-5 accuracy (left subfigure, solid purple and blue lines). In the perturbed synthetic dataset, where the keyword in the questions no longer matches that in the triple exactly, KBLAM shows superior performance to BM25 at all numbers of triples tested.
as "Euler's Building" in the question). This was done by feeding the entity names through GPT-4 and asking perturb the question described above.
Using the same KBLAM model trained on synthetic data where there is always an exact match between the question and knowledge base, we find that despite never having such queries during training. KBLAM can still retrieve correctly from the knoweldge base, as shown in Figure 10b. Compared to the baseline BM25, KBLAM degrades less and, unlike with the other datasets, outperforms BM25 at all KB sizes.
Latency comparison with RAG In Fig. 11, compare the RAG's time and memory overhead with KBLAM, where we focus on the overhead from the LLM assuming RAG uses a highly-optimized retriever with negligible overhead. We find that RAG, with 5 triples retrieved into the context, shows time latency and memory higher than KBLAM, due to the large number of characters in the context. KBLAM on the other hand shows stable latency with respect the KB size, in that KBLAM's overhead from QKV projection and FFN does change with KB size.
D CURRENT LIMITATIONS AND EXTENDED FUTURE WORK
One-time training costs of KBLAM A limitation of KBLAM lies in its non-zero one-time costs (around 24-48 hours on a single 80GB GPU) compared to some other methods for knowledge augmentation such as in-context-learning RAG, or prompt caching. Namely, KBLAM requires the fine-tuning of adapters, and the creation of triples if the input is not already organized this format. However, the former is one-time cost, with adapters generalizing across corpora. with the Enron dataset, and the latter only needs to be performed once per corpus. These one-time time costs contrast to the of fine-tuning for knowledge augmentation (Hu et al., 2021; Liu et al.,
16

---

Published as a conference paper at ICLR 2025
Time to first token
Memory usage
0.10
20
GigaBytes (GB)
Seconds (s)
0.08
19
0.06
18
0.04
2
2
2
2
212
2
2
2
2
2
#Triples
# #Triples
KBLaM
Overhead K=5
KBLAM maintains lower latency and memory usage than RAG while keeping the full KB in context. Averaged over 40 random queries from the synthetic KB, for KB with 512 triples, KBLAM matches RAG's memory footprint which retrieves 5 triples (gray dashed line, right), while providing consistently faster token generation across all KB sizes (red line, left).
Figure 11:
2024c) which have much higher one-time costs, as models need to be fine-tuned for each corpus. In addition, as a result of these one-time costs, the runtime costs of KBLAN are significantly lower than in-context learning, RAG, prompt caching, shown in Figure 11.
Utilizing more extensive structural information in the KB. One of the motivations behind KBLAM's design is that a KB can be viewed as a collection of key-value pairs, similar to a dictionary, therefore can independently encode each triple and augment them into pre-trained LLM efficiently. Future work can study designing customized attention structures for more sophisticated structures underlying KBB e.g. tree structure, if we consider <name> under same <name> but different <prope r ty>s as the leaf nodes, or graph structure, if we assume each triple as node and edges connect triples with <value> equal to the <name> of another triple.
root node and triples
17

---

Published as a conference paper at ICLR 2025
E
SAMPLE KB
Here we present some example triples from both KBs, where each triple is of format
(<name>; <property>; <value>)
We present ten examples from the synthetic KB (Table. 1) and from Enron (Table. 2) respectively,
<name>
<property>
<value>
Posh Poodle
description
A line of organic teas with health-boosting properties
Maserati Vasco da Gama
description
A smart lighting system with customizable settings
A historical strategy game in ancient Rome.
Thorn Tacticaan
description
Conquer the empire
Bohr's s Bookshelf
purpose
To promote a sense of calm and improve quality of life To raise awareness about the ethical implications of technology
Pixelated Prose
purpose
A historical library housing rare manuscripts, books, and documents
Matrix Monument
description
A secure data center offering cloud storage and
Nova Citadel
description
cybersecurity services
Test radiation shields, monitor devices, and come up
Celestial Genome
objectives
with fixes
To provide guidance, build supportive network, and promote positive parenting practices
The Mamas and the Papas
objectives
A high-stakes fashion competition featuring talented designers from around the world. It's showcase of creativity and style
The Doom That Came to Sarnath
description
Table 1: Synthetic KB
<name>
<property>
<va 1 ue>
Feature that allows traders to automatically fill any
Fill Order
description
order that is at the top of the stack.
Allen & Overy
description
A solicitors' partnership Part the PA'S list
EAP PA
description
Analyst Recommendations
Screen for attractive stocks using Wall Street analysts"
Screener
purpose
opinions
Financial Accounting
A standard that deals with accounting rules for derivatives and other financial instruments.
description
Standard 133
Vividence
description
A customer management applications firm.
LD Calculation
A methodology proposed by CDWR for consistency among contracts, viewed as more 'objective' by CDWR.
description
Methodology
Boxed Trio of Rosewood
Uniquely shaped rosewood stoppers protect opened
description
Stoppers
wines from rapid oxidation.
troubled software company sold by Mattel a
Learning Company
description
buy-out firm exchange for a share of future profits.
Refiects dollars paid to generator for DAM Contract Balancing that is Out-of- for LRR. Pre-employment screening services
Code 238
description
Validata
purpose
Table 2: Enron KB
F SAMPLE Q&A
An example of four types of questions used for instruction tuning is presented in Fig. 12.
18

---

Published as a conference paper at ICLR 2025
Example KB four triples
1. <Tempura Titon,
description, objective,
A bodybui der competing in international competitions.>
2. <Pythia's Plates,
To promote iatic health and relaxation.>
3. <Samosa Sage,
purpose,
To provoke thought and discussior.>
4. <Somosa Sage,
description, phi losopher pandering the meaning of life.>
Simple Q&
Open-ended Q&A
Q: What's the description of Samosa Sage?
What objective Pythia's Plates and what you think of its social impacts?
A: The deseription of Samosa Suge philosupher pondering the meaning of life.
objective of Pythia's Plates promote holistic health and relaxation, which could potentially lead
Multi-entities Q&A Q: Tell me the purpose of Samosa Sage and the description Tempuru Titan. A: The purpose of Samosa Sage to provoke thought and discussion and the description of Tempura Titan is bodybuilder competing international competitions.
improved physical and mental wellbeing their community,
increased health awareness, stronger connections,
Unanswerable Q&A
you tell description Pesto Paladin? Sorry cannot find relevant information in the K.B.
Examples of instructions tuning dataset. Given KB of 4 triples (top panel), we generate 4 types of Q&A for instruction tuning. The KB shown is sampled from the synthetic KB; The question and answer for open-ended Q&A are also synthesized by GPT, the rest Q&As are generated using formatted string.
Figure 12:
G PROMPT
In this section, we provide all prompts used for KBLAM's experiments.
G.1 PROMPT FOR SYNTHETIC KB GENERATION
To prompt the GPT-4 for generating synthetic KB, begin setting system prompt as
You are system that generates synthetic data examples in JSON format
Then we construct a list of idea types and a list of object types:
idoa_types [
'greck letters', 'fiction characters', 'fmoou rock bands', 'birds', 'animals', 'natural phenomens', 'physical locations', 'artist names', 'classical music', "mssical instruments', 'music genres', 'art styles', 'aacient Roman concepts', Hiiauu myths', 'Cthulhu Mythos', real-worrl company names', mythological creatures', "planets and stars', 'nistorical figures', Literary genres', "botanical nanes', famous Landmarks', scientific concepts', 'space missions', rinventions!
"ph phhicssopical terms', chemical elements', famous scientists', 'marine life', 'mythological places'
object Sepect_typee_-
'cducation company', 'tech company', car company', 'entertainment company', construction company', 'retail company', finance company', 'healtheare company', 'restaurant' 'hotel', github repo', 'project', 'meeting room', 'building', "lab', aailine', 'textbook', 'website', 'personal blog', gaming company', "consulting firm', 'biotech company', 'app', software tool', "bookstore', 'e-commerce site', social media platform', "fitness brand', "fasnion brand', non-profit organization'
Then for each combination of object type and idea type in the list, we prompt GPT with
19

---

Published as a conference paper at ICLR 2025
Please randemly generate 50 (object_typel name innovated by [idea_typel. The name should have various styles. The generated name should be of diverse style and length, for example the name could have hyphen, space in it multiple words.
which gives us 50 different names. Then the same context, we further prompt GPT with
Now for each of the objectives, and a purpose for the data.
generated, generate short desciption, short
Please ensure that the generated contents has LON*** correlation with the name.
which generates the <value> 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
Now for each of the name, descriptien, objective and purpose generated, make their text style more diverse using mixture of formal and informal language.
G.2 PROMPT FOR OPEN-ENDED Q& GENERATION
To generate open-ended Q&A instructions, we first take a triple from the KB, construct a simple Q& from it, and then feed the simple Q&A into GPT and let GPT augment it into more complex form via the following prompt
You are given question and answer pair, please extend the question to be open-ended and generate a short answer. For example, you could generate
what is the objective of and what do you think of it? Make sure the answer is oonly* based on information provided from the QA pair. In addition, please generate the format of: Q: A:
G.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 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 system prompt of
You are AI system that evaluates the quality of generated responses. Your goal
return score between and indi indicating how accurate and useful the response is. An accurate and useful response should get high score of 5.
Then, the following prompt, which encourages the model give chain-of-thought reasoning for scoring (Zheng et al., 2023):
A model is given a question about some information and evidence. The question is composed of two parts, a part that involves repeating information in the evidence and part that potentially involves open- ended thinking, Then the model generates response. Evaluate the response based on how grounded it is given the evidence and how reasonable it is. Return an integer score and step by step explanation of how you arrived at the score. Score of 5 means the response is accurate, relevant and reasonable (in that it meets common sense).
20

---

Published as a conference paper at ICLR 2025
If the response addresses the question and the evidence in relevant way, should get a high score Score means the response is inaccurate and irrelevant or model is hallucinating, Score between 0 and 5 means the respense is partially correct and relevant.
Followed by the prompt, we include examples better help GPT calibrate the scoring rule
Example 1:
Rvidence: "The purpose of Alexandria to extract knowledge."
Question: "bescribe the purpose of Alexandria and how it can benefit users."
Model output: "The purpose of Alexandria is extract knowledge, it
benefit users by providing a structured way to organize knowledge.
Score:
Reason: The model's response is accurate and relevant to the question and evidence, the open-ended part is reasonable.
Example 2:
Evidence: "The purpose of Alexandria is to extract knowledge."
Question: "Describe the purpose of Alexandria and what can it extract.
Model output: "The purpose of Alexandria is to extract knowledge, it can extract knowledge knowledge.
Score:
Reason: The model's s response is accurate and relevant to the question and evidence.
Example 3:
Evidence: "GreatTool is an app that helps users to be more productive."
Question: "Describe GreatTool and how it may affect the community.
Model output: "GreatTool is app that helps users be more productive may affect the community by helping users to sleep better."
Score:
Reason: The model's s response is accurate and relevant to the question and evidence but it is not very reascnable.
Example 4:
Evidence: "GreatTool is an app that helps users to be more productive." Question: "Describe GreatTool and how it may affect the cormunity."
Model output: "GreatTool is app that helps users be more productive may affect the community by helping users to organize their tasks and manage their time better improving their productivity."
Score:
Reason: The model's response is accurate and relevant to the question and evidence and the open ended part is sensible and reasonable.
Example 5:
Evidence: "GreatTool is an app that helps users to be more productive."
Question: "Describe GreatTool and hew it may affect the community,
Model output: "GreatTool great teal with many feature" Score:
Reason: The model's response is net accurate and doesn't answer the question.
G.4 PROMPT FOR LLAMA EVALUATION
For in-context learning, the following prompts for evaluation
Simple Q&A
Please answer questions based on the given text with format: "The { property) of (name) is (descripticn]"
Two-entity Q
21

---

Published as a conference paper at ICLR 2025
Please answer questions based the given text with format: "The { property) (namel) (descriptien); The (property) of (name2) is description); ..."
Open-ended Q&A You are previded context and question that has retrieval part and an open-ended part.
Please answer the question based en the given text.
If the information for the epen-ended part is not provided in the context please generate potential possible answer.
Unanserable questions
Please answer questions based on the given text with format: "The { property) of (name) is (descripticn]", if relevant information cannot be found in the text, please respond "I sorry cannot find relevant information in the KB".
For zero-shot learning, we use the following prompt
Simple Q&A
Please answer the question in very compact manner with format: The property) of (name) is (descripticn)
Two-entity Q6A
Please answer the question in very compact manner with format: "The property) of (namel) is (descriptien); The (property) of (name2) is descripticn);
{
Open-ended Q&A Please answer the question based on your knowledge.
G.5 QUESTION TEMPLATE
For simple Q&A, we use the following templates
what <property> does <name> have?,
What is the <property> of <name>?, Tell me about the <property> of <name>., Can you 1et know the <property> of <name>?, Can you inform me about the <property> of <name>?, Describe the <property> of <name>., What dctails can What kind of <property> does <name> have?, Provide details on the <property> of <name>.,
share about the <property> of <name>?,
What features does the <property> of <name> include?,
Can you elaborate on the <property> of <name>?,
How woula you describe the <property> of <name>?,
What you tell about the <property> characteristics of <name>?, Can you explain the <property> of <name>?, what insights can you provide about the <property> of <name>?, what should I know about the <property> of <name>?,
For multi-entities Q&A involving triples, {(<name> <property>,, <value>,)]e we use the following template
What is the (}
Tell me (h,
Can you let know (), Can you inform me 1), Describe (),
Explain (h,
Could you describe the (), what you tell about (),
22

---

the <property>_g <name> -r, and Published as a conference paper at ICLR 2025
Could you provide information (h, Please enlighten me about (h, can you clarify () for me?,
Could you give me detailed description of 1h, need more information on {),
then for each blank in the curly brackets, we fill in the string
the <property> <name> 1,
the <property>_C of <name>_
H SAMPLE OUTPUT
In this section, present sample output from the model. In particular, present generated answers from KBLAM, in-context learning, and zero-shot learning, on open-ended Q&A tasks for synthetic KB (Appendix. H.1) and Enron (Appendix. H.2), given questions about "Relevant triple" from KB of 100 triples.
that zero-shot consistently produces hallucinations, the contents of both the synthetic KB and Enron KB differ significantly from the real-world factual knowledge stored in the LLM's parameters.
Broadly,
KBLAM and in-context learning, on the other hand, utilize the information from the KB to synthesize the answer. In-context learning provides highly precise answers all questions, however incurs substantial computational overhead and struggles to scale to larger KBs On synthetic data, KBLAM accurately captures the information from the KB terms of meaning and sentence structure, though it may use slightly different wording, e.g. in the last example on "Titdal surge". KBLAM uses "renewable" instead of "sustainable" in the original triple
When tested on the out-of-distributt Enron data, KBLAM still grasps the general meaning of the relevant triple, albeit with more noticeable differences. As discussed in the limitations section of the main text, we believe this performance gap can be narrowed by using more diverse synthetic KBB with a stronger real-world foundation during training.
H.1 SYNTHETIC KB
Relevant triple:
(ReefFulse Jellyfish; objectives; to provide real-time alerts, support remote monitoring, and improve home safety)
Question:
Can you inform about the objectives of ReefPulsc Jellyfish and what do you think of its offectiveness in achieving these goals?
KBLaM output:
The objectives of ReefPulse Jellyfish are to provide real-time monitoring alerts, and improve home safety. While the effectiveness ReefFulse Jellyfish in achieving these goals would depend on various factors such as the accuracy of the monitoring technology, the reliability of the alert system, and the user experience, the outlined objectives suggest comprehensive approach to enhancing home safety.
In-context learning output:
According to the context, the objectives of ReefPulse Jellyfish are to provide real-tine alerts, support remote monitoring, and improve hone safety.
Based on these objectives, think ReefPulse Jellyfish is effective achieving its goals. The real-time alerts feature allows users to stay informed and take prompt action in case of any issues, which can help prevent potential hazards and ensure safer home envirenment. The renote monitoring feature enables
keep track of their home's
23

---

In-context learning output: Published as a conference paper at ICLR 2025
conditions from anywhere, which be particularly useful for people are away from home for extended periods. Additionally, improved home safety feature can help reduce the risk of accidents and injuries, making it a valuable tool for hemecwners. Overall, I believe ReefPulse Jellyfish is
Zero-shot learning output:
ReefFulse Jellyfish is a product designed to monitor and control the water parameters in aquariums, specitically targeting jellyfish and other delicate marine life. Based my knowledge, the objectives of ReefFulse Jellyfish are:
1. Accurate Water Parameter Moniteringa:: To provide real-time monitoring of water parameters such as på, ammonia, nitrite, nitrate, and temperature, ensuring that the aquarium's water conditions optimal for the jellyfish and other inhabitants.
2. omated water Parameter Control**; тo automatically adjust water parameters to maintain a stable and healthy enviroament, reducing the need for manual intervention and minimizing the risk of water parameter fluctuations.
3. *xEnhanced Jellyfish Care**: To provide a more precise and controlled envi ronment
Relevant triple:
(Sagan's starry Hotel; objectives; to create magical and enchanting environment for guests)
Question:
What features does the objectives of Sagan's s Starry Hotel include and how do these features contribute to the overall guest experience?
KBLaM output:
The objectives of Sagan's s starry Hetel is to create magical and enchanting environment for guests. These features likely include themed decor, special amenities, and exceptional service, all designed transport guests into fantastical world, enhancing their stay and making it memorable.
Based on the context, the objectives Sagan's starry Hotel are to create magical and enchanting environment for guests. The features that contribute to this objective include:
Creating magical and enchanting environment for guests Providing unique and memorable experience for guests Enhancing the overall guest experience through attention to detail and a focus on customer satisfaction
These features contribute to the overall guest experience by creating a sense of wonder and excitement, making guests feel special and valued, and providing memorable and enjoyable stay.
Zero-shot learning output:
I apologize, but I couldn't find information "Sagan's Starry Hotel American astrenomer and science communicator, but
", Carl Sagan couldn't find any evidence of hotel bearing his name.
If you meant to ask about different hotel or a hypothetical hotel concept, please provide context details, and I'll do my best to your question.
Relevant triple:
(Currents Coach; purpose; to help you stay active and achieve your fitness goals)
Question:
24

---

Published as a conference paper at ICLR 2025
Describe the purpose of Currents Coach and how it benefit in their daily lives.
KBLaM output:
The purpose of Currents Coach is help individuals track and improve their fitness. By using Currents Ceach, users can monitor their physical activities, set fitness goals, and track their progress over time. This can Lead to a more structured and effective fitness routine, ultimately contrisuting to better overall health and well-being.
In-context learning output:
The purpose of Currents Coach is to help individuals track and improve their fitness. This app can benefit users in their daily lives by providing personalized and comprehensive fitness tracking system. With Currents Coach,
set and achieve their fitness goals, monitor their progress, and receive guidance and support to help them stay motivated and on track.
Here
potential benefits of using Currents Coach:
Improved physical health: By tracking and monitoring their fitness progress, users can make data-driven decisions to improve their overall physical health and well-peing.
Increased motivation: Currents Ceach provides
accountability
and motivation, helping users stay committed to their fitness goals and make progress towards achieving them.
Fersonalized guidance: The app offers personalized guidance and support
Zero-shot learning output:
Currents Coach is personalized financial wellness platform designed to help users manage their finances, achieve their financial goals, and improve their overall financial well-being. The purpose of Currents Coach is to provide users with comprehensive and interactive financial planning experience that tailered their individual needs goals.
Here are some ways Currents Coach can benefit users in their daily lives:
1. **Persnnalzed Financial Planning**: Currents Coach machine learning algorithms to analyze users' financial data and provide personalized financial plans, including budgeting, saving, and investing recommencttions.
2. **Goal Setting and Tracking**: Users goals, such as saving for a down paynent on house, paying off debt, or building an emergency fund. Current
and track their financial
Relevant triple:
(Hertz Horizens; description; music production company specializing in independent artists)
Question:
Can you explain the description of Hertz Horizons and what kind of artists they specialize in?
KBLaM output:
The description of Hertz Horizons music production company specializing in independent artists. They focus on supporting and promoting independent musicians, providing them with the resources and opportunities needed to succeed in the music industry.
In-context learning output:
According to the context, Hertz Herizons is a music production company specializing in independent artists.
Zero-shot learning output:
25

---

Published as a conference paper at ICLR 2025
Bertz Horizens is music producticn and distribution соmрany that specializes in creating and releasing music film, television, and other visual media. They are known for their unique approach to music production, which involves creating custom music tracks for specific scenes or moments in film or TV show.
Hertz Horizens is particularly known for their work in the electronic and ambient genres, and they have worked with a wide range of artists and producers to create music for varicus projects. Some of the artists they have worked with include:
Tycho: A electronic music producer and composer who has worked on numerous film and TV projects, including the Netflix series "Stranger Things".
Rival Consoles: British electronic music producer and composer has worked on
Relevant triple:
(Tidal Surge; objectives; to generate clean and sustainable energy)
Question:
Provide details 1 on the objectives of Tidal Surge and explain how aims to achieve them.
KBLaM output:
The objectives of Tidal Surge is generate clean and renewable energy. It aims to achieve these objectives by harnessing the power tidal currents, which are a consistent and powerful source of energy. This involves the development of advanced technologies and infrastructure designed to capture and convert the kinetic energy of the tides into electricity.
In-context learning output: According to the context, the objectives of Tidal Surge are to "generate clean sustainable energy".
To achieve these objectives, Tidal Surge likely aims to harness the power of tidal currents, which are the reqular rise and fall of the sea level caused by the gravitational pull of the moon and sun. This renewable energy source has the potential to provide significant portion of the world's energy needs while reducing greenhouse gas emissiens and mitigating climate change.
тo generate clean and sustainable energy, Tidal Surge may employ various technologics, such as:
These are structures that harness the energy of tidal currents by creating dam-like structure across tidal estuary or narrow channel.
1. Tidal barrages:
2. Tical stream generators: These are
Zero-shot learning output:
Tidal Surge is research project fecused on developing novel tidal energy converter (TEC) that can harness the power of ocean tides to generate electricity. The primary objectives of Tidal Surge are:
1. +*Design and development of more efficient TEC**: The project ains design and develop TEC that capture larger portion of the tidal energy available in the ocean, thereby increasing the overall efficiency of the device.
Tidal Surge aims to create TEC that can be easily scaled up or down depending the speci fic tidal energy resource and the needs of the power grid. This wi11
2. IIIprovedd scalability and cost-effectiveness-
26

---

OPIC funds for offshore sil and natural gas Published as a conference paper at ICLR 2025
help reduce the tidal energy generation and make it competitive with other renewable energy
H.2 ENRON
Relevant triple:
(Sitraka; description; provider of Internet software technologies for organizations)
Question:
What kinc of description does Sitraka have and what services do they offer?
KBLaM output:
The description of Sitraka is software company providing networking solutions for businesses. They offer services that likely include developing and inplementing netwerk infrastructure, managing network security, and possibly providing technical support to ensure smooth network operations for their clients.
In-context learning output:
Based on the given text, Sitraka described as provider of Internet software technologies for IT organizations.
As for the services they offer, the text does not provide specific information. However, based on the general description, can inferred that Sitraks offers software technologies and solutions for IT organizations, possibly related to internet internet-based based services, network management, or other I1-related areas.
Zero-shot learning output:
According to my knowledge, Sitraka is a software company that provides IT service management, solutions. They offer a range of servi ces and products that help organizations manage and optimize their IT infrastructure, including:
IT Service Management (ITSM) scftware: Sitraka's flagship product is the ITSM software, which provides comprehensive platforn. for managing IT services, including incident, problem, and change management. Asset Management: sitraka's s asset management solution helps organizations track and manage their IT assets, including hardware, software, and network devices.
Service Desk: Sitraka's service desk solution provides centralized flatform for managing requests, incidents, and problems. Reporting and Analyties: Sitraka's s reporting and analytics solution provices real
Relevant triple:
(Unocal; objectives; projects)
Question:
what features does the objectives of Unocal include and how do these features impact their operations?
KBLaM output:
The objectives of Unocal include securing offshore oil and gas reserves. These features impact their operaticns by ensuring that they can and extract these natural resources efficiently and safely, which is crucial for meeting energy demands and maintaining a stable supply chain.
In-context learning output:
Hased the provided context, the ab objectives Unocal are:
27

---

Published as a conference paper at ICLR 2025
Secure OPIC funds for offshore cil natural gas projects
These dbjectives do not provide specific features, but rather outline the company's s geals and priorities. However, we can infer that securing funds for offshore oil and natural gas projects may involve features such as:
F Financial planning and oudgeting
Risk assessment and management
Project development and execution Requlatory compl iance and permitting
Collaboration with stakeholders, including government agencies and investors
These features would likely impact Unocal's operations by:
Guiding the allocation of resources and budgeting for specific projects Informing the company's strategic decisions and pricrities Influencing the company's relatienships with stakeholders and partners
Zero-shot learning output:
Unocal, multinational energy corperation, has several objectives that shape its operations. Based my knowledge, the following features incluced in Unocal's objectives and their impact on operations:
1. +*Safety and Environmental Protectionss Unocal prioritizes the safety of its emplcyees, contractors, and the public, well environmental protection. This objective is reflected in their commitment environmental impacts, implementing safety protocols, and conducting regular training programs. Impact: This objective ensures a safe working environment, reduces the risk of accidents, and minimizes environmental damage.
reducing
2. +opperational Excellence**: Unocal aims to optimize its operations by improving efficiency, reducing cests, and increasing preductivity. This objective ive is achieved through precess improvements, technclogy adopt on, and employes
Relevant triple: (San Francisco Symphony; purpose; cencert organization)
Question:
Provide details on the purpose of San Francisco Symphony and what activities they engage in to fulfill this purpose.
KBLaM output:
The purpose of San Francisco Sympheny is to organize comrunity concerts.
To fulfill this purpose, they likely engage in activities such as planning and coordinating concert events, recruiting musicians, pronoting the events to the community, and pessibly collaborating with local organizations ensure the concerts are accessible and enjoyable for a wide aucience.
In-context learning output:
Based on the given text, purpose San Francisco Symphony a concert organization. This implies that the San Francisco Symphony is a music organization that aims to promote and perform classical music, likely through concerts and other musical events.
To fulfill this purpose, the San Francisco Symphony may engage various activities such as:
Organizing and performing concerts, including classical music performances, symphony orchestras, and chamber music recitals
28

---

Published as a conference paper at ICLR 2025
Producing and presenting musical events, such opera, ballet, and dance performances
Providing music education and cutreach programs, such as music classes, workshops, and community outreach initiatives collaborating with other music organizations, schools, and community groups premote music education and appreciation Developing and maintaining a professional
Zero-shot learning output:
The San Francisco Symphony (SFS) is world-renoxned erchestra based in san Francisco, California. The purpose of the San Francisco symphony is to provide high-quality musical performances, promote music education, and foster sense of community through music.
To fulfill this purpose, the Francisco Symphony engages in variety of activities, including:
1. Concerts: The SFS performs over 200 concerts cach year, featuring a wide range of repertoire from classical masterpicces to contemporary works. These concerts take place home venue, as well as at other locations throughout the Ray Area.
Davies Symphony Hall, the orchestra's
2. Music Education: The SFS offers a range of music education programs for chilcren and adults, including the SFS Youth Orchestra
Relevant triple:
(TB/A/Chiat/Dar objectives; eliminate private office and individual cubicle, Promoee collaboration)
Promote
Question:
What insights can you provide about the objectives of TeNA/Chiat/Day and what do you think of their appreach to promoting collaboration?
KBLaM output:
The objectives of TBWA/Chiat/Day is promote face-to- face collaboration and eliminate conference rooms. Their approach to promoting collaboration by eliminating conference rooms suggests focus more informal and spontaneous interactions, which can be beneficial for fostering creativity and open communication among team members.
In-context learning output:
Based the provided context, the objectives of TeNA/Chiat/Day are:
Eliminate private office and individual cubicle Fromote collaboration
Insights that can provided about the objectives of TBWA/Chiat/Day are:
The agency is focusing on creating more collaborative work envi ronment by eliminating private offices and individus cubicles. This suggests that they value teamwork and communication among employees. By promoting collaboration, TBRA/Chiat/Day creativity, innovation, and productivity among its employees. This approach can lead to better ideas, faster probiem-solving, and improved morale.
likely trying to increase
As for their approach to promoting collaboration, it seens that they are taking bold step by eliminating private offices and
Zero-shot learning output:
TBMa/Chiat/Day is global advertising agency that was founded in 1978 by Jay Chiat and Richard Laermer. The agency is known for its creative and innovative approach to advertising, and has worked with wide range of clients across various industries.
29

---

Published as a conference paper at ICLR 2025
As for their objectives, TBMA/Chiat/Day's mission is to "Challenge the Status Que" and to "change the Game" for their clients. They aim to create advertising that is not only creative and effective, but also memorable and impactful. The agency's approach is centered around the idea of "Disruption," which invelves challenging the conventional norms and expectations of advertising and finding and innovative ways connect with consumers.
In terms of promoting collaboration,
Relevant triple:
(Go2net Mail; purpose; used for sending and receiving emails)
Question:
Describe the purpose of Go2net Mail and explain how benefits users.
KBLaM output:
The purpose of Go2net Mail is used for sending and receiving emails. It benefits users by providing reliable and efficient platform for communication, ensuring that emails are delivered promptly and securely.
In-context learning output:
The purpose of Go2net Mail is to previde service for sending and receiving emails. This service allows users communicate with others through electronic mail, which is convenient and efficient way to exchange information.
Go2net Mail benefits users in several ways. Firstly, it provides platform for users to send and receive emails, which essential for personal and professional communication. Secondly, it allows users to access their emails from anywhere, at any time, as long as they have an internet connection. This makes it easy to stay connected with others, and
regardless of their location. Thirdly, Go2net Mail provides reliable way to send and receive emails, which is important for users who need to share sensitive informaticn. Overall, Go2
Zero-shot learning output:
Go2net Mail is free web-based email service provided by Go2net, popular onl ine directory and search engine. The primary purpose of Go2net Mail is to offer users a reliable and user-friendly email service that allows them to send and receive emails, manage their contacts, and store files online.
Here are some benefits that Go2net Mail offers to its users:
1. F*Fre Email Service*:: Go2net Mail provides free email service, don't have pay any subscription fees or maintenance
which means costs.
2. e*Web-Based Access**: Users can access their email accounts from anywhere, at any time, using web browser and an internet connection.
3. *LLarge Storage Space**: Go2net Mail offers
30


Reviews:
[{"review_id": "EHDVFgTeS1", "reviewer": "Reviewer_6e7T", "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.", "questions": "1. In Figure 6(a), why are the in-context learning results (orange line) missing from the Enron dataset?\n2. 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?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 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.\n2. The proposed KBLaM, to my knowledge, is novel and offers the following advantages:\n- 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.\n- 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.\n- 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.\n- It also allows high interpretability via viewing the attention matrix.\n3. Overall, this paper provides solid evidence through extensive experiments and ablation studies to demonstrate the effectiveness of the proposed method. The writing is also clear and easy to follow.", "weaknesses": "1. While this paper includes experiments with both synthetic and real-world KB datasets, the baselines are limited to only ICL and zero-shot approaches. Recent state-of-the-art methods, such as RAG-based and long-context models, are not included as baselines, making it difficult to assess the effectiveness of the proposed method relative to these approaches. Including results from these comparisons would strengthen the evaluation.\n2. 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.\n\n---\n[1] Natural Questions: A Benchmark for Question Answering Research, Kwiatkowski et al., 2019\n\n[2] TriviaQA: A Large Scale Distantly Supervised Challenge Dataset for Reading Comprehension, Joshi et al., 2019\n\n[3] MuSiQue: Multihop Questions via Single-hop Question Composition, Trivedi et al., 2022\n\n[4] When Not to Trust Language Models: Investigating Effectiveness of Parametric and Non-Parametric Memories, Mallen et al., 2023", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "cn0bKlhzTj", "reviewer": "Reviewer_q4dg", "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.", "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).\n\n2. 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?", "limitations": "", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "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.\n\n2. KBLAM eliminates the need for separate retrieval steps, simplifying the architecture and making the model more efficient compared to RAG.\n\n3. 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.\n\n4. The paper exhibits a robust methodology with meticulous experimentation. It particularly emphasizes KBLAM’s scalability, efficiency, and interpretability. The use of a rectangular attention mechanism enables the model to scale linearly with the size of the KB. Furthermore, the authors provide comprehensive ablation studies to substantiate their design choices.", "weaknesses": "1. I think that the most relevant baseline for comparison with the proposed method should be prompt caching (e.g., as described in the following link). This approach converts the knowledge base into a key-value cache for subsequent use. Compared to this method, I think the advantages of the KBLAM method maybe diminished. For instance: a) it requires instruction tuning; b) the accuracy of retrieval might be substantially lower compared to prompt caching; c) the key-value cache generated by a pretrained encoder is inherently inferior to that obtained from LLMs.\nthe prompt caching method:\na) https://platform.openai.com/docs/guides/prompt-caching\nb) https://ai.google.dev/gemini-api/docs/caching\nc) https://proceedings.mlsys.org/paper_files/paper/2024/file/a66caa1703fe34705a4368c3014c1966-Paper-Conference.pdf\n\n\n2. 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.\n\n3. 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.\n\n4. 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "yPC3K54Ysk", "reviewer": "Reviewer_W1bY", "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.", "questions": "-\tAs 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?\n-\tHow does this method compare to retrieval + ICL?\n-\tHow does KBLaM perform on more popular datasets and different types of problems?", "limitations": "", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "-\tThe 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.\n-\tIt is a neat and interesting way to map each triple into a fixed-length key-value vector pair\n-\tThe main body of the paper is written well, and the explanations for each step of the calculation are precise", "weaknesses": "-\tThe experimental comparison setup is not well-designed:\n\tNo comparisons in the first experiment. How is the retrieval accuracy compares to that of BM25?\n\tIn the last two experiments, although the authors mentioned that ICL encountered memory constraints, it also resulted in fewer data points.\n-\tThe 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.\n-\tThe 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.\n[1]Hao, Yaru, et al. \"Structured prompting: Scaling in-context learning to 1,000 examples.\" arXiv preprint arXiv:2212.06713 (2022).\n[2]Cai, Tianle, et al. \"Scaling In-Context Demonstrations with Structured Attention.\" arXiv preprint arXiv:2307.02690 (2023).\n-\tThere are some repetitive words in the article:\n\tThere is no need to use a large section in Chapter 3 background to introduce Self-attention layer.\n\tImage 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.\n-\tThe accuracy of ICL on the Enron dataset does not seem to be demonstrated in Fig6(a).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "We6GudKN13", "reviewer": "Reviewer_9hRD", "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.", "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)?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 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).\n\nOeverall, 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.\n\nIt also seems that the top-k choices among the Keys are all that are necessary (particularly since the facts are likely stand-alone, rather than mutually reliant), which could mean that much cheaper approximate search could be used to select out relevant entries before doing a full dot-product.", "weaknesses": "At several points, the 'linear rather than quadratic scaling' of the method is emphasised.  This seems a weak point, since the size of the tuples $M$ already exceeds the context length $N$, so $MN>N^2$ in this regime (as obliquely mentioned in L284).\n\nThe constant $C$ introduced in Eqn 11 feels like something that should have a more principled origin.\n\nsmall edit: L103 \"training samples is not of interests.\" -> \"training samples is not important.\"", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "O4vlgGJBit", "reviewer": "Reviewer_ouM4", "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.", "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?\n\n2. 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.\n\n3. 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?", "limitations": "", "rating": 3, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 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.\n\n2. The proposed method also enhances interpretability by potentially offering insights into how the model leverages knowledge through the attention matrix.\n\n3. This paper is clearly-written and easy to follow.", "weaknesses": "### 1. Limitations of the method:\n- 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.\n- 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.\n- 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.\n\n###  2. Limitations of the experiments:\n- 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.\n- 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.\n- Missing comparison with Retrieval Augmentation Generation (RAG).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)