AtlasKV: Augmenting LLMs with Billion-Scale Knowledge Graphs in 20GB VRAM
Authors: Haoyu Huang
Hong Ting Tsang
Jiaxin Bai
Xi Peng
Gong Zhang
Yangqiu Song
Venue: ICLR 2026 Poster (2026)
Keywords: Large language model; Knowledge augmentation; Knowledge graph;
Primary area: foundation or frontier models, including LLMs
TL;DR: This paper proposes AtlasKV, a scalable, effective, and general way to augment LLMs with billion-scale KGs in less than 20GB GPU VRAM, where KG2KV and HiKVP are introduced to integrate KG triples at scale with sub-linear time and memory complexity.

Abstract: Retrieval-augmented generation (RAG) has shown some success in augmenting large language models (LLMs) with external knowledge. However, as a non-parametric knowledge integration paradigm for LLMs, RAG methods heavily rely on external retrieval modules and the retrieved textual context prior. Especially for very large scale knowledge augmentation, they would introduce substantial inference latency due to expensive searches and much longer relevant context. In this paper, we propose a parametric knowledge integration method, called $\textbf{AtlasKV}$, a scalable, effective, and general way to augment LLMs with billion-scale knowledge graphs (KGs) (e.g. 1B triples) using very little GPU memory cost (e.g. less than 20GB VRAM). In AtlasKV, we introduce KG2KV and HiKVP to integrate KG triples into LLMs at scale with sub-linear time and memory complexity. It maintains strong knowledge grounding and generalization performance using the LLMs' inherent attention mechanism, and requires no external retrievers, long context priors, or retraining when adapting to new knowledge.

Published as a conference paper at ICLR 2026
ATLASKV: AUGMENTING LLMS WITH BILLION-SCALE KNOWLEDGE GRAPHS IN 20GB VRAM
Haoyu Huang', Hong Ting Tsang', Jiaxin Bai
+ Xi Peng", Gong Zhang", Yangqiu Song'
1 The Hong Kong University of Science and Technology
Theory Lab, Huawei
(hhuangcp, httsangaj, jbai Jêconnect.ust hk (pancy pengxi, nicholas. yasongêcse.use .st
Source Code
Data and Models
ABSTRACT
Retrieval-augmented generation (RAG) has shown some success in augmenting large language models (LLMs) with external knowledge. However, as non- parametric knowledge integration paradigm for LLMs, RAG methods heavily rely on external retrieval modules and the retrieved textual context prior. Especially for very large scale knowledge augmentation, they would introduce substantial inference latency due to expensive searches and much longer relevant context. In this paper, propose parametric knowledge integration method, called AtlasKV, a scalable, effective, and general way to augment LLMs with billion-scale knowledge graphs (KGs) (e.g. 1B triples) using very little GPU memory cost (e.g. less than 20GB VRAM). In AtlasK we introduce KG2KV and HiK to integrate KG triples into LLMs scale with sub-linear time and memory complexity. It maintains strong knowledge grounding and generalization performance using the LLMs' inherent attention mechanism, and requires no external retrievers, long context priors, or retraining when adapting to new knowledge.
Long Context
Adaptor
Encoder
attention
LoRA
LLMs
attention
Retriever
L.Mss
QLoRA
LLMs
(a) Non-Parametric Methods
(b) Traditional Parametric Methods
(c) AtlasKV
Figure The simple illustrations of two kinds of popular knowledge augmentation paradigms for LLMs and new parametric knowledge augmentation paradigm adopted by AtlasK (a) Non- parametric methods usually rely on external retrievers and long context prior, which have retriever- limited performance and substantial inference latency. (b) Traditional parametric methods require re- re-training the model when adapting to new knowledge, which are also expensive. achieve injecting external knowledge efficiently at scale without need for external retrievers long context prior with strong generalization ability.
(c) AtlasKV can
INTRODUCTION
Large language models (LLMs) have demonstrated impressive generation abilities in various down- stream tasks (Brown et al., 2020; Touvron et al., 2023; Kung et al., 2023; et al., 2024; Dagdelen
"Corresponding author.

---

Published as a conference paper at ICLR 2026
et al., 2024), where their expanding parameter scales enable them to function comprehensive knowledge stores by encoding factual information directly into their parameters (Petroni et al., 2019; Jiang et al., 2020; Rae et al., 2021: Bubeck et al., 2023; Morris et al., 2025). generation (RAG) (Gao et al., 2023; Fan et al., 2024) is more cost-efficient solution to enhance the capabilities of LLMSs in knowledge-intensiee tasks, which may require vast amount of external knowledge, without altering an LLM's parametric representation. RAG methods usually retrieve text chunks (Sarthi et al., 2024; Jimenez Gutierrez et al., 2024) or subgraphs (Edge et al., 2024; Huang et al., 2025) that are relevant to query from an external textual knowledge base (KB) or knowledge graph (KG), which serve as the context prior for LLMs to generate responses.
Although RAG methods have achieved some success in efficiently augmenting LLMs with external knowledge, they still face some critical limitations. As shown in (a) of Figure 1, these con-parametric methods heavily rely on external retrieval modules and the textual context prior, which introduce substantial inference latency due to expensive searches (e.g. nearest-neighhor searches) (Khandelwal et al., 2019; He et al., 202 la; Shi et al., 2023) and longer context (Ram et al., 2023; Cao et al., 2025) specially for very large scale knowledge augmentation.
In contrast, parametric approaches (Gururangan et al., 2020; Wang et al., 2024; Cao et al., 2025) can achieve the integration of external knowledge into LLMs without need for external retrievers or long context prior. Traditional knowledge adaptation techniques (Hu et al., 2022; Diao et al., 2023), as shown in (b) of Figure 1, require retraining the model when adapting to a new distribution of knowledge, which significantly limits their application scenarios. A new parametric knowledge augmentation paradigm introduced by KBLaM (Wang et al., 2024) like (c) of Figure 1, well addresses this issue by encoding external knowledge into series of key-value parametric representations and seamlessly injecting them into the self-attention layers of an LLM, which well preserves the advantages of parametric methods and can also adapt to any new KB in a training- free manner.
Nevertheless, we found that there are two critial challenges of this novel knowledge augmentation (1) Lack of high quality training data. It requires query-key-var vvaue (Q-K-V) sentences of the external knowledge as the training data. Directly synthesizing Q-K-V training data from unformatted documents with fixed pre-defined schemas suffers from limited query diversity, which could result to poor generalization performance in out- of-distribution (OOD) scenarios. (2) Poor scalability. When augmenting LLMs with very large scale external KB or KGs, the computational and memory overhead of this knowledge augmentation paradigm becomes prohibitively high, even with linear time and memory complexity.
paradigm that restrict its real-world applications:
To this end, propose AtlasKV, scalable method that enables end-to-end knowledge augmentation of LLMSs with billion-scale KGs (e.g. 1B triples) using very little GPU memory cost (e.g. less than 20GB VRAM) while achieving superior knowledge grounding and generalization performance in OOD scenarios. We achieve this by two innovative designs from both data and algorithm perspectives. Specifically, to address the training data quality issue, we observe that each triple in KGs can be naturally converted into Q-K-V data, which shares very similar structure with the Q-K-V vectors of self celf-attention networks in LLMs. So we introduce the concept of KGKV and propose the V pipeline that naturally converts each KG triple into high- sigh-quality Q-K-V data for both training and inference, enabling a better injection of KGs into L.LMs. To solve the scalability challenge, we propose a hierarchical kay-value pruning (HiKVP) algorithm that can dramatically reduce computational and memory overhead while maintaining high knowledge grounding accuracy during inference time.
In summary, we make the following main contributions:
We propose AtlasKV. a scalable method that enables end-to-end augmentation of LLMs with billion-scale KGs (e.g. 1B triples) using very little GPU memory (e.g. less than 20GB VRAM) while achieving superior knowledge grounding performance and strong generalization abilities.
We introduce KG2KV and HiKVP complementary designs to address data and algo- rithmic challenges respectively: KG2KV naturally transforms KG triples into high-quality Q-K-V data enhance generalization, while VP enables scalable integration through hierarchical pruning that dramatically reduces computational and memory overhead during inference.
2

---

Published as a conference paper at ICLR 2026
Extensive experiments and analysis demonstrate the superior effectiveness and scalability V compared to ICL. KBLaM, and RAG methods, with comprehensive ablation studies validating the contribution of each component.
2
RELATED WORK
Non-parametric Knowledge Augmentation Methods For LLMs. The most popular non-parametric knowledge augmentation methods for LLMs is RAG (Lewis et al., 2020; Gao et al., 2023; Fan al., 2024), which significantly enhances LLMs by incorporating an external retriever that fetches relevant context from external KBs (Zhang et al., 2024a; Sarthi et al., 2024) or KGs (Mavromatis & Karypis, 2024; Edge al., 2024; Jimenez Gutierrez et al., 2024; Huang et al., 2025). Their retrievers usually heavily rely either the separately pretrained sentence transformers (Wang et al., 2020a; Izacard et al., 2021; Zhang et al., 2025a), or a well-designed tuning process with the LLM's output as the feedback signal (Shi et al., 2023; Chang et al., 2025). However, the performance of this knowledge augmentation paradigm could significantly limited by the capabilities of the retrievers. There are also some graph-based RAG advancements
scale to large knowledge bases more efficiently. E GrGraphRAG (Zhao et al., 2025) constructs bidirectional indexes between entities and chunks enable fast lookup during both local and global retrieval. LinearRAGi (Zhuang al., 2025) constructs a relation-fin free hicrarchical graph to let graph construction scale linearly with corpus size. Nevertheless, they still follow the ICL-based RAG paradigm. And the long context retrieved from large KBs or KGs could still introduce substantial inference latency.
Parametric Knowledge Augmentation Methods For LLMs. Parametric knowledge augmentation approaches are much more native to LLMs. Because the inherent memory of LLMs is integrated by pretraining and supervised fine-tuning (Geva et al., 2020; Petroni et al., 2019; Morris et al., 2025), which are also parametic methods. Some early attempts to efficiently integrate new external knowledge into L.MS such as LoRA (Hu al., 2022) and adapter (He et al., 2021b; Diao et al., 2023) still suffer from retraining the model when integrating new external knowledge into LLMs. Cache- augmented generation (CAG) (Chan et al., 2025) is a new caching- based knowledge augmentation paradigm that preloads the LLM with all relevant documents in advance and precomputing the KV cache. MemDec (Cao et al., 2025) provides a domain-specinf memory module that enhances various frozen LLMs without parameter modifications. KBLaM (Wang et al., 2024) introduces new knowledge augmentation paradigm that augments knowledge into the attention layers of LLMs, achieving linear computational complexity while enabling training-free adaptation to new KBs after initial training.
Knowledge Graph Augmentation Methods For LLMs. There are many works augmenting KGs into LLMS in both parametric and non- con-aarmettric ways. Except for the graph-based RAG methods mentioned above, some LLM-based KGQA methods like RAR (Shen et al., 2025), KnowGPT (Zhang et al., 2024b), and the work done by Ji et al. (2024) also depend on the training of their retrievers or path aligner to find the most relevant knowledge from KGs as context. KELP (Liu et al., 2024) explores latent semantic matching to improve path-level knowledge selection from KGs. They are still limited by the performance of the retrievers and also face inference latency issues.
BACKGROUNDS AND DEFINITIONS
3.1 KNOWLEDGE GRAPHI
As most of the existing works on graph-based RAG systems did, we use the textual triples (h, r, f) as the basic knowledge unit of KGs in our work, which could be extracted from unstructured text with any existing KG extraction method (Angeli et al., 2015; Huang et al., 2024; Mo et al., 2025). Then the KGs can be defined as G {(h, r, 1)) E E, RJ. where E is the set of entities and R is the set of relations. Note that h, t could be either named entities, or other types of entities, such as concepts, events, etc (Zhang al., 2020; Tan al., 2024; Bai et al., 2025). In our work, G will be integrated as external factual knowledge for LLMs to ground facts and answer questions. And the process of extracting KGs from documents is not the focus of work.
3

---

Published as a conference paper at ICLR 2026
3.2 ATTENTION NETWORKS
In this section, we will first give the definitions of self-attention layers, which are the key components of the transformer (Vaswani et al., 2017) backbone. Then we will describe the definitions of the rectangular attention in KBLaM (Wang et al., 2024), which is a new knowledge augmentation paradigm adopted by AtlasKV.
Self-Attention Layers. In each attention layer, we input a query 2 E IR with N token length, which embedding vector can be denoted as x0 E RNDD, D is the embedding dimension of attention layers and E (1, ... L}, where is the number of attention layers. There also three attention
heads W W W
Q K V E RDDD, which are designed to project each input token into Q-K-V embeddings (1) k()), (1(1 IRDD Then the output at the /-th layer and n-th token computed as
((0) E exp( (q K (O/,,K/)/DDD //DD/!
E exp( ( K. //VD)
where (+--) denotes the inner product of two vectors. This standard implementation of self-attention can have a time complexity of O(N2. ((N2.D) D) and memory complexity of O(N . (N + D)), which could lead to significant computational overheads and time delay as the input length N gets larger.
Rectangular Attention in KBLaM. The KB in KBLaM is a set of key-value pairs, which is denoted as M = {[k(Um v(0m } J at the (-th attention layer, where k k()m, z(0m E IR Mx are the base embedding vectors of the m-th key-value pair. M is the size of the KB and DE is the output dimension of the sentence encoder. Then the knowledge augmented output at the /-th attention layer and n-th token is computed as
Em exp((i"), //0mm/// //DD)/DD + E exp( 1 k0O /D)// EM exp((q!)) k(0m) //D) + E exp( 1 kD///DD
y
(2)
where q. w Q Inn k01m WT k0m, 7(1m -W m(lm denote the Q-K-V embedding V E Dx are the specific projection heads for the Q-K-V vectors of the KB. This rectangular attention have a time complexity of O ((M + N)- -N D) and memory complexity of O ((M + N)- (N+ + D)). Its computational overhead would grow linearly with M, which is more efficient than standard self-attention.
(1)
vectors after projection of the KB part. w(l
Q E RDDD w w
However, as the size of KB M scales up heavily, the linearly growing time and memory complexity would still be critical problem. AtlasKV can further improve the scalability of the KG augmented
I.M with O (C.VM ory complexity, where C, and C. are constants that much smaller than M.
N)- N D) time complexity and O (Cnm VM + N). (N + P)) mem-
4 METHODOLOGY
Overview of AtlasKV. Augmenting LLMs with super large and complex external knowledge, e.g. general KGs with billions of triples, often struggles with low generalization abilities and unbearable computational and memory overhead (Wang et al., 2024; Zhang et al., 2024b; Jin et al., 2024). To overcome these fundamental challenges, method to integrate massive KGs into LLMs through two key innovations: (1) KG2KV, a new KG integration paradigm that naturally converts KG triples into Q-K-V data, enabling LLMs achieve both enhanced generalization performance and efficient knowledge integration, and (2) HiKVP, a hierarchical key-value pruning algorithm that dramatically reduces computational and memory overhead while maintaining high knowledge grounding accuracy during inference time.
propose AtlasKV, scalable, effective and general
4.1 KG TO KV
Building on the observation that every triple in KGs can be naturally decomposed into Q-K-V strings (Verga et al., 2020), which shares very similar structure with the Q-K-V vectors of self- attention networks in LLMs, we introduce the concept of KGKV and employ the KG2KV pipeline to transform each KG triple into Q-K-V strings and their corresponding sentence embedding vectors.
4

---

Published as a conference paper at ICLR 2026
John founded Stock1 Lemonue om
John made profits every year since started short selling
John founded Brocklamoonee om
John has made profits evey your sinen started short selling
bocause
because
What the Tell me the
What is the Tell me the
John has made profits every year
chase
of
John founded StockLenonn.com comn
result
of
?
started short selling
Tail Masked Triple Transformation
Head-Masker Masked Triple Transformation
Figure 2: An example of how we transform the KG triples to Q-K-V data.
For given KG triple (h, or concept entity). Then the masked entity can be the value we need in this triple. And then we will masked, which could be considered as the
firstly mask its head h or tail entity t (could be named entity, event
rewrite into noun words according the position attribute of the entity that is not masked. For example, as shown in Figure 2, if we mask the tail entity be directly rewritten into its noun word "cause" through LLMs. And the key string of this tail-masked triple can be represented as "the cause of John founded StockLemon.co commeemoooomm..... If we mask the head entity in the triple, we need rewrite the relation into its reversed noun word "result". And the key string of this head-masked triple can be represented as "the result of John has made profits ...", In this KG2KV pipline, we consider the masked entity as the value data, and the other entity as well as the relation the key data, which complete KGKV data. Note that for the training data, we usually select named entities as the key to mask, and select event entities and relations as the value due to the reasons elaborated in Section 5.3. Then through a sentence encoder, KGKVs can be compressed and encoded into sentence embeddings k" v to integrate to the attention layers of LLMs.
in the triple, the relation "because"
We also need the query sentence of each triple to serve as the training data. The query string can be obtained by adding various questioning prefix to the key strings. For example, the questioning prefix can be "What
"Tell me... ...", or " "Provide details on. ...." This design can ensure the model would not overfit to specific type of questioning way.
Compared to directly synthesizing Q-K-V data from documents with limited human-defined schemas, the massive relations from KGs can guarantee the diversity of the enquiry attributes in the constructed training data with KG2KV method. Besides, due to the only strings we need input into the LLMs in KG2KV is the masked position and the relation, KG2KV needs fewer token cost than directly synthetic method. We also provide the prompt template in Appendix H. As shown in Table 1, KG2KV holds a significantly higer diversity ratio (number of unique enquiry attributes divided by the total number of triples) of 7.864% and lower average token cost of 165.7 than the synthetic method. We also provide some samples of the training data constructed by synthetic and KG2KV methods respectively in Appendix G. And the influence of relation rewriting process, which is the only part that consumes tokens in KG2KV, is also analyzed in Appendix B.2. As verified, more accurate key strings with relation rewriting in KG2KV are more beneficial than the simple combinations of relation and unmasked entity strings without relation rewriting.
Table 1: Comparison of the data diversity ratio and average token cost between KG2KV and synthetic method.
Method |iieessity Ratio 1/vgg. Token Cost +
Synthetic/ KG2KV 7.864% %
0.003%
349.9
165.7
4.2 AUGMENTING LLM WITH KGKVS
In this section, we will introduce how AtlasKV integrates the KGKVs constructed in the previous section into the L.MSs with the help of hierarchical key-value pruning (HiKVP), which makes more scalable than other methods.
Hierarchical Clustering on KGKVs. Inspired by various works (Sarthi et al., 2024; Zhang et al., 2024a; Huang et al., 2025; Zhang et al., 2025b) that achieve some success by organizing hierarchical knowledge structure on textual chunks, which is an intuitive way to organize the world's d knowledge, we employ the hierarchical clustering to cluster the keys of KGKVs Vs into a hierarchical structure. This design aims to share the computational and memory burden during inference time on each layer of the hierarchical knowledge keys. Specifically, as previous works (Sarthi et al., 2024; Huang et al.,
5

---

Published as a conference paper at ICLR 2026
2025) did, we first employ Uniform Manifold Approximation and Projection (UMAP) to reduce the dimension of the knowledge keys, and then employ Gaussian Mixture Models (GMMs) to cluster the knowledge keys into hierarchical structure. Each key vector in a higer layer is the pooling of the key vectors in the lower layer. In Atlas we set the number of layers to be 3, which can also be larger according to the actual situation. We select 3 layers because that is the minimum number of layers to include all of the definitions we need in At To share the computational and memory burden equally, we set the size of clusters in each layer to be the same, which is
S
M Then can have the base embeddings of three layer knowledge keys ky RM,xDDD
KT E R M, kR E M&MxxD where M,.- M, Mr DE is the embedding dimension of the sentence encoder.
M MR
[
M
M
And
Knowledge Augmentation. During the tuning process of AtlasKV, do not need prune the KGKV pairs. Because due to the generalization capability of AtlasKV as verified in Section 5, we do not need such large scale KGKVs in the tuning process. And we use an equivalent attention method to replace the rectangular attention in KBLaM for knowledge augmentation. At the (-th attention layer and n-th token, it is computed
y >kg Softmax (logitSkgz) (0)
Aseg Softmax (logitS (()),
where logits is the KG part of the attention output, logits seg is the sequence part of the attention output. For the weights Akg and Aseq of these two softmax results, we have
ΣM exp(logits
Akg
(4)
exp(logits ) + En exp(logits seg) Σ =1
E exp(logits
Aseg
(5)
E M exp(logitsk ) + Em= exp(logits
And we also have
logits (
k ((6))))) /VD, logits seq (q ( k kW///V
(6)
where 1/ W ky and k()- / k. The only learnable variables @ are the KG-specific query heads We q and KG projection heads Wx, Wv. Then we optimize using LLMs' original auto-regressive training objective:
L p(ul, 9) = II po gen, Wes),
(7)
where L is the total length of the query q and the answer And denotes the i-th token in the combination of and v. q<i and U'a are the query and answer tokens in all turns hefore the current prediction token F4. It is also very intuitive that the attention output at both the prefilling process and each decoding step is a weighted combination of the KG part and the sequence part values. And the weights are determined dynamically according to the logits of the KG and sequence part. The equivalence to the rectangular attention can be proven in Appendix C.
Hierarchical Key- Value Pruning. During inference time, usually there are only few relevant KG values needed for the query. And the injection of irrelevant knowledge would also introduce noise. So we use a hierarchical key-value pruning (HiKVP) pipeline as shown in Figure to efficiently and scalably find the most relevant KGKVs for the query. And the attention will be computed as
5 = Xe. Soottmx (logits;,)) 50) + Aseg Softmax (logits.eq)
(²))
(8)
where the logits and 500 are the pruned KG logits and values in the leaf-layer. Aig is the weight based on the pruned KG logits. They can be obtained by the following steps.
Step 1. Initially, only the root-layer projected key vectors K are uploaded in the GPU memory,
while other layer key and value vectors (k("), k(() and (())) remain in the CPU memory. We first compute attention weights between the query and root-layer keys:
logits
(2(0) (U) \/D.

---

Published as a conference paper at ICLR 2026
VM
GPU
kg
Rt
Top-ka id
Root Layer kangs
Mapping
ky
|ll
x
Top-Ky its
Inter Layer Keys
Selected Keys
kg
Top-ky
Loqe-Layer Kays
Sefested Xeys
Salect
!!!
Xay Sodtumne Softtmon(kgiinn (eginn))
-1lI
KGKVS
Luger
Figure 3: An overview of hierarchical key-value pruning (HiKVP) with three layers of knowledge keys at the I-th attention layer. The gray background indicate that the part stored and computed in the GPU memory.
After calculating the softmax of it, we prune root keys by selecting the keys with top-kæ highest scores and use a mapping to obtain the included inter-layer keys k ((*xS)xD And the root-layer keys will be offloaded back to the CPU memory.
Step 2.
We conduct the similar process to the step 1, but with the selected inter-layer keys. We upload the selected inter-layer keys the GPU memory, and compute the attention weights:
logitskgr (ān (() /VD.
(10)
Then we can prune the selected inter-layer keys and obtain the selected leaf-layer keys k 8)x the same way. And the selected inter-layer keys will be offloaded the CPU memory.
Step 3. Finally, we upload the selected leaf-layer keys from the CPU to GPU memory. Then we compute the attention weights after the softmax calculation and prune the leaf-layer keys by directly selecting the corresponding logits with top-kı highest softmax scores:
logits = Lopp__ooitsSSoftmal (logits ) .L) E AL x1 logits (q k0 //VD. (11)
And through the mapping indices, we can also obtain the pruned values 5(0) R and upload them the GPU memory. Then logits (ll) and Akg can be obtained and the attention output of Equation 8 during the inference time can be finally computed.
these steps of HiKVP can be done o ( NN) time complexity and O ((Cm M + N) (N + D)) memory complexity, respectively, where C2 and Cm are constants that are much smaller than M. We also provide the detailed derivation in Appendix D.
EXPERIMENTS
In this section, report the performance of AtlasKV from GPU memory cost, knowledge grounding accuracy and the relevance the generation results perspectives. And we also compare the perfor- mance of AtlasKV with other knowledge integration baseline methods demonstrate the superiority of AtlasKV, We also report the training and evaluation details of AtlasKV in Appendix A. and Appendix A.2.
5.1 EXPERIMENTAL SETTINGS
Baselines. Following the settings in KBLaM (Wang al., 2024), compare AtlasKV with both non- parametric and parametric methods. For non parammeeiii methods, we include in-context learning (ICL), which is the basic knowledge augmentation paradigm used in RAG methods. For parametric methods, we include KBLaM, which is an advanced parametric knowledge augmentation paradigm. We also include zero-shot learning provide boundaries of our experimental results.
7

---

Published as a conference paper at ICLR 2026
Training Datasets. In AtlasKV, construct the Q-K-V V training dataset ATLAS-WiKi-QKS with ATLAS-Wiki, which is one of ATLAS (Bai et al., 2025) family KGs with 900+ million nodes and 5.9 billion edges containing both event and named entities. Because it offer sufficiently large KGs that enable us to train the model and comprehensively assess the performance of AtlasK and other baseline methods across various KG sizes. And we also use the Synthetic training dataset used in KBLaM (Wang et al., 2024) to compare with.
Evaluation Datasets. To comprehensively assess the performance of AtlasKV and other baseline methods in a scenario closer to the real world, we test all methods in the OOD settings. We not only include the Enron dataset (Klimt& & Yang, 2004), which is an OOD dataset used in KBLaM (Wang et al., 2024), but also introduce the ATLAS-CC-QKY dataset (Bai et al., 2025) and ATLAS-Pes20 QKV dataset (Bai al., 2025), which are also constructed from ATLAS-CC and ATL.AS-Pe220 in ATLAS (Bai et al., 2025) with the KG2KV method, respectively, to evaluate the performance of different methods in more comprehensive way. Note that ATLAS-CC-QKS and WTLAS-Pes20.. QKV are much harder than Enron dataset because they are constructed from more complex KGs and include much more unique enquiry attributes that are closer to the real world scenarios. With these OOD datasets, we can better evaluate the generalization capabilities of various methods.
5.2
EXPERIMENTAL RESULTS
Method
Time Complexity
Memory Complexity
ICL
O ((MT NY2.
((MT+ N)-(MT + D))
(MT
RAG
O
(M + RT + (RT + N2 D)
O ((RT+ + N.NNNRRRT + +N+
CAG
0
(RT N22. N22.
O ((RT+N] (N +D)
KBLaM
0
(M NNNNNNNN
O ((M+N) N)- DD}
AtlasKV
O
(C.VM + N)- D))
0 (CmVM NN) (NV
Table 2: Comparison of the time and memory complexity of AtlasKV, KBLaM, RAG, CAG (Chan
et al., 2025) and ICL methods, where the parts marked in teal color represent they could be very large.
AtlasKV is more scalable with HiKVP. To verify the scalabil- ity of AtlasKV with HiKVP, we
GPU VRAM Cost(GB)
ATLAS-CC-QKV
ATLAS- Pes2o-QK
In-context learning
compare the GPU memory us-
40.0
40.0
KBLaM
age at inference time of AtlasKV and other methods across wide
30.0
30.0
AtlasKV
Zero-shot
20.0
20.0
Total GPU Memory
range of KG sizes from 1 to IB triples. As shown in Figure 4, the colored areas represent how much VRAM is saved compared with the other method. It demon- strates that with the help of HiKVP, AtlasKV can save a huge amount of VRAM compared with ICL as well as KBLaM and achieve much lower GPU memory cost. With the increasing of KG scale, the VRAM usage of AtlasKV even just little bit higher than the zero-shot generation. And in Atlask V, less than 20GB VRAM is required augment L.MMs with 1B triples. However, in KBLaM,, over 40GB VRAM is required to deal with even 100K triples. The key reason why AtlasKV can achieve this is that HiKVP significantly reduces the time and memory complexity of KBLaM from linear to sub-linear, as demonstrated in Table 2, where T denotes the average token length of the triples. In ICL-based RAG methods, when the size of relevant triples R scales up, the inference latency and VRAM usage caused by long context prior dependence would grow heavily. And their performance will also influenced by the lost-in-ther the--ddddddeee dilemma (Liu et al., 2023), which would not exist in AtlasKV. And compared with the increased cost with every additional retrieved R in CAG (Chan et al., 2025), which is a new caching-bass knowledge augmentation paradigm, AtlasK V trades the sub-linear attention-basse based retrieval cost for the ability to perform efficient inference independent of the scale of the retrieved relevant knowledge.
10
10a
10
10
# of KG Triples
# of KG Triples
Figure 4: GPU memory usage comparison of AtlasKV and other methods across various KG sizes from to 1B triples.
AtlasKV is more accurate and generalizable with KGKVs. Not only can A save huge amount of inference cost with strong scalability, but it can also maintain a higher knowledge grounding performance with strong generalization ability. To quantitatively assess the knowledge grounding performance of AtlasKV, we extract the seraged--vver-hule KG part sost-sof attention scores at the 15th layer due to the reason described in Appendix A.2, and then obtain the Top-1
8

---

Published as a conference paper at ICLR 2026
10' Triples
10' Tripls
10" Triples
10t Triples
Method
Steps
ACC
ACCB5
ACC01
ACC
ACC#1
ACC
ACCIR1
ACC
ou Eazon
KBLaM
2e4
3e3
100.0
90.0
100.0
100.0
50,9
50.9
83.6
76.4
29.1
25.5
47.3
56.4
7.3
9.1
23.6
20.0
AtlssY (128-64-16)
3c3
100.0 (10.0)
(40.0)
(+166)))
90.9 (47.3)
418 (1127)
50.9
218 (412.7)
32.7 (412.7
AdlasY weo HIKVP
3e3
100.0 (e0.0)
100.0
76.4 (+25.5)
92.7 (+9.1)
56.4 (22.3)
80.0 (+23.))
27.3 (x18.2)
47.3 (+27.)
Eval on A7Z-0-
KBLaM
3e3
40.0
80.0
16.4
45.5
5.5
14.5
3.6
2e4
50.0
80.0
25.5
52.7
3.6
14.5
0.0
5.5
AtlasY (128-64-161
363
90.0 (+40.0)
100L
(+2000)
823 (+6.8) By
92.2 (+400.0)
52.2
(+47.2) 2222)
20.9 (+56.4)
16.4 (+16.4)
(+4.5)
AtlaKY w/o IDKYP
3e3
100.0
(-50.0)
100.0
(+20.0)
92.7
(+67.2)
100.0 (+17.3)
72.7
(+67.2)
90,9
(+76.4)
47.3 (+47.3)
67.2 (+61.7)
Eval ATLAS: CC-QA
KBLaM
Se3
60.0
90.0
21.8
38.2
12.7
23.6
3.6
10.9
2o4
50.0
100.0
23.6
56.4
10.9
21.8
3.6
10.9
AtlacV (128 64-161 AllaKV w/o IGKVP
3e3
100,0 (+0.0)
100.0
(+0.0)
S9 (+65.5)
90.9
(+34.5)
61.8
(+18.1) 1)
24.5
(+50.9)
10.0
(+36.4)
$1.5 (+13.6)
3e3
100.0
(+0.0)
100.0
(+0.0)
96.4 (+72.8)
100.0
(+13.6)
83.6
(+70.9)
96.4
(+72.8)
6188
(+58.2)
8188 (+70.9)
Table 3: The knowledge grounding performance of AtlasKV against KBLaM with all-MiniLM-L6--22 as the sentence encoder on three OOD evaluation datasets across various tuning steps and KG sizes. We defaultly set the top-k in HiKVP to 128, 64, and 16 for the kR, K1, ky respectively.
and Top-5 accuracy of the knowledge grounding performance. As shown in Table 3, across all three OOD datasets and a wide range of KG sizes, A achieves significantly higher Top-1 and Top-5 accuracy than KBLaM with the KGKVs as the training data. Especially in ATLAS- Pes2o-QKV and ATLAS-CC-QKV datasets, which are much harder due to their complex and diverse enquiry attributes, KBLaM performs very bad because there are too limited enquiry attributes in Synthetic training data to make it more generalizable. However, only 20K KGKV samples as the training data needed to make AtalsKV much more accurate and generalizable. It also suggests KGK can make the training process more efficient with only 3K steps, compared with the 20K training steps reported in KBLaM. We also experiment with different top-k settings in VP in Appendix B.4.1. Another interesting observation that, even with HiKVP, there is not a big performance drop of AtlasKV and it still performs better than KBLaM. This is mainly because the specific heads trained in AtlasK have the capabilities to conduct fuzzy retrieval at different layers of semantic granularity. Besides, we observe the training dynamics of AtlasK in Appendix E, which is that from a specific training step, the model regularly start learning to retrieve relevant knowledge from the external KG triples, instead of brute force over-fitting We also report the results with a larger model as the sentence encoder in Appendix B.
Enron
ATLAS-CC-QRV
ATLAS-Pes2o-QKA
1.0
1.0
1.0
In-context learning
GPTScore
0.8
KBLaM
0.5
0.5
0.5
AtlasKV w/o HiKVP
Zero-shat
0.0
10a
10a
10
102
10
# of KG Triples
# of KG Triples # of KG Triples
Figure 5: Scored by GPT-40 between O and 1, the shaded area exhibits the standard error over 5 random seeds. The score of each random seed is also the average of generation results.
We also use GPT-4o (Hurst et al., 2024) to score the relevance between the ground truth and generated answers. As shown in Figure 5, AtlasKV achieves significantly higher GPTScores than KBLaM. Although ICL can generate very accurate result with over 0.9 scores, it is super time-consuming to put all external knowledge into the context of LLMs. Expecially when there are more than 100 triples in a KG. over 48GB VRAM is required and can not be run on the limited GPU memory. And KBLaM also performs poorly on the two difficult datasets, which is also due to the reasons we explained above. Remarkably, despite having only limited training samples with enquiry attributes similar to Enron in ATLAS- -IL--Wiii--KK. AtlasKV still outperforms KBLaM on both knowledge grounding accuracy and answer relevance metries, even though KBLaM's training data contains exactly the same enquiry attributes as Enron. This is mainly due the the diversity of the enquiry attributes in ATLAS- Wiki-QKV, which is constructed by KG2KV module. We have compared that with fully synthetic method in Table I and it makes AtlasKV have the capability to be generalized to more unseen enquiry attributes in complex scenarios. We also provide some samples of ATLAS- Wiki-QKt and Synthetic dataset in Appendix G.1.

---

Published as a conference paper at ICLR 2026
5.3
ABLATION STUDY
10 Triples
10F Triples
10 Triples
10t Triples
Method
Steps
ACC#1
Acces
ACCe1
ACCØ5
ACC
AcCes
ACC
AcCes
Evcll ATLAS ATLAS--PesZ-OAA
AtlasKY wio HiKVP
363
100.0
100.0
92.7
100.0
72.7
90.9
47.3
67.2
AllasKV wio HiKVP& & Event
3c3
90.0
100.0
80.0
89.1
31.5
63.6
9.1
36.4
AllasVV w/o HiXVP& Emitty
3x3
100.0
100.0
49.0
67.3
20.0
30.9
3.6
5.5
Eval oz Enzon
AtlasKV w/o HiKVP AtlasKV HiKVP Event
3x3
100.0
100.0
76.4
92.7
56.4
80.0
27.3
47.3
3e3
80.0
100.0
73.6
10.9
38.2
AllasKV w/o HiKVP& & Eatity
363
40.0
100.0
40.0
74.5
16.4
27.3
1.8
9.1
Table 4:
The knowledge grounding performance of different variants of AtlasK with all-MiniLM- L6-v2 as the sentence encoder on three OOD evaluation datasets across various tuning steps and KG sizes.
Cooperating named and event entities together KG2KV process helps with the model's learning. To analyze the reasonability of the design described in Section 4.1, we conduct experiments on the variant with only KG2KV component, which are denoted as "AtalsKV w/o HiKVP". Because we need focus the ablations of training data and avoid the influence of pruming process. We compare the training data containing both named and event entities with the variants containing only named or event entities, which are denoted "AtalsKV w/o HiKVP & Event" and "AtalsKV w/o HiKVP & Entity". Here is the analysis based on the results in Table 4: (1) Without any one kind of entities, there will be drop of the knowledge grounding performance. (2) Especially when there are only event entities, due to the high complex semantics in both key and value strings, it becomes very hard for the specific heads learn from scratch, which leads to huge performance drop. (3) When only named entities are employed in KG2KV. the performance drop is smaller because the key and value strings of them are shorter and the semantics are simpler, which makes easier for the specific heads to learn from scratch. But the worse performance with only named entities suggests that we still need some complex semantics in event entities to help the model to learn better.
6 CONCLUSION
In this paper, we presented AtlasK scalable, effective, and general framework augment LLMs with billion-scale knowledge graphs under very low GPU memory budgets. Compared with non- parametric methods, AtlasK requires external retrievers and does not depend on long context prior, which could lead to substantial inference latency. Compared with traditional parametric methods, AtlasKV can be adapted to new knowledge a training-free manner. We achieve that by (1) KG2KV which naturally converts KG triples into Q-K-V data, enabling LLMs to achieve both enhanced generalization performance and efficient knowledge integration, and (2) HiKVP which conducts hierarchical key-value pruning to dramatically reduces computational and memory costs while maintaining high performance during inference time.
7 ETHICS STATEMENT
We affirm adherence to the ICLR Code of Ethics. This work develops AtlasKV does not involve human subjects or interventions. We use publicly available datasets (e.g., ATLAS family KGs, Synthetic, and Enron datasets) under their licenses. We construct Q-K-V data via KG2KV without collecting new personal data or attempting deanonymization. We also comply with model or API providers' terms (c.g.. LLLMA3.1-88 1-B.A...---BBB Instrrcct GPT-4o, and 4i--mmiii without uploading proprietary or sensitive information. No undisclosed conflicts of interest exist. All experiments were performed on a single 48GB GPU, and we encourage energy- configurations.
8 REPRODUCIBILITY STATEMENT
LLMs for Training and Evaluation. Three LLMs used in our AtlasKV experiments: LLaMA3.1- Instruurtt,, GPT-4o, and GPT-4o-mini. We use LLaMA3.1 I-8B-Inntruct as the backbone of AtlasKV and it can be obtained at Hugging Face. GPT-4o and GPT-4o-mini are used to score the generated
10

---

Published as a conference paper at ICLR 2026
answers and rewrite the relations in KG2KV process, respectively. They can be accessed via OpenAI API calls.
Training and Evaluation Details. We provide comprchensive descriptions about our training and evaluation settings in Appendix A, including the hyper- parameter settings and detailed processing steps. All of the datasets used in work can be obtained through public resources (Bai et al., 2025; Wang et al., 2024; Klimt & Yang, 2004). We will also release our source code soon after the submission.
ACKNOWLEDGMENT
The authors of this paper were supported by the ITSP Platform Research Project (ITS/189/23FP) from ITC of Hong Kong, SAR, China, and the AoE from RGC of Hong Kong, SAR, China. We also thank the support from Huawei.
and the GRF (16205322)
REFERENCES
Gabor Angeli, Melvin Jose Johnson Premkumar, and Christopher D Manning. Leveraging linguistic structure for open domain information extraction. In Proceedings of the 53rd Annual Meeting of the Association for Computational Linguistics and the 7th International Joint Conference Natural Language Processing (Volume 1: Long Papers). pp. 344-354, 2015.
Jiaxin Bai, Wei Fan, Qi Hu, Qing Zong, Chunyang Li, Hong Ting Tsang, Hongyu Luo, Yauwai Yim, Haoyu Huang, Xiao Zhou, al. Autoschemakg: Autonomous knowledge graph construction through dynamic schema induction from web-scale corpora. arXiv preprint arX1v:2505.23628, 2025.
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot leamerr. Advances in neural information processing systems, 33:1877-1901, 202.. 2020.
Sébastien Bubeck, Varun Chandrasekaran, Ronen Eldan, Johannes Gehrke, Eric Horvitz, Ece Kamar, Peter Lee, Yin Tat Lee, Yuanzhi Li, Scott Lundberg, et al. Sparks of artificial general intelligence: Early experiments with gpt-4. arXiv preprint 2.03. 12712 2023.
Jiaqi Cao, Jiarui Wang, Rubin Wei, Qipeng Guo, Kai Chen, Bowen Zhou, and Zhouhan Lin. Mem- ory decoder: A pretrained, plug-and-play memory for large language models. arXiv preprint arXiv:2508. 098774 2025.
Brian Chan, Chao-Ting Chen, Jui-Hung Cheng, and Hen-Hsen Huang. Don't do rag: When augme-e generation is all you need for knowledge tasks. In Companion Proceedings of the ACM Web Conference 2025, pp. 893-897. 2025.
Ge Chang, Jinbo Su, Jiacheng Liu, Pengfei Yang, Yuhao Shang, Huiwen Zheng, Hongli Ma, Yan Liang, Yuanchun Li, and Yunxin Liu. Grail: Learning to interact with large knowledge graphs for retrieval augmented reasoning. arXiv preprint arXiv:2508.n5 055998 2025.
John Dagdelen, Alexander Dunn, Sanghoon Lee, Nicholas Walker, Andrew S Rosen, Gerbrand Ceder, Kristin Persson, and Anubhav Jain. Structured information extraction from scientific text with large language models. Nature communications, 15(1):1418, 2024.
Shizhe Diao, Tianyang Xu, Ruijia Xu, Jiawei Wang, and Tong Zhang. fixture-f-dommai-addee Decoupling and injecting domain knowledge to pre-trained language models memories. arXiv preprint arXiv:2306.0544 2023.
Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishck Kadian, Ahmad Al-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Amy Yang, Angela Fan, et al. The llama 3 herd of models. arXiv e-prints, pp. arXiv-2407, 2024.
Darren Edge. Ha Trinh, Newman Cheng, Joshua Bradley, Alex Chao, Apurva Mody, Steven Truitt, Dasha Metropolitansky, Robert Osazuwa Ness, and Jonathan Larson. From local to global: A graph rag approach focusedd summarization. arXiv preprint 4rXiv:2404.161.33 2024.
11

---

Published as a conference paper at ICLR 2026
Shahul Es, Jithin James, Luis Espinosa Anke, and Steven Schockaert. Ragas: Automated evaluation of retrieval augmented generation. In Proceedings of the 18th Conference of the European Chapter of the Association for Computational Linguistics: System Demonstrations, pp. 150-158, 2024.
Wenqi Fan, Yujuan Ding, Liangbo Ning, Shijie Wang, Hengyun Li, Dawei Yin, Tat-Seng Chua, and
Qing Li. A survey on rag meeting Ilms: Towards retrieval-augmennn large language models. In Proceedings of the 30th ACM SIGKDD conference on knowledge discovery and data mining, pp- 6491-6501, 2024.
Yunfan Gao, Yun Xiong, Xinyu Gao, Kangxiang Jia, Jinliu Pan, Yuxi Bi, Yixin Dai, Jiawei Sun, Haofen Wang. and Haofen Wang. Retrieval-augmented generation for large language models: A survey, arXiv preprint arXiv:2312.10997, 2(1), 2023.
Mor Geva, Roei Schuster, Jonathan Berant, and Omer Levy. Transformer feed-forward layers are key-value memories. arXiv preprint arXiv:2012.14913, 2020.
Zirui Guo, Lianghao Xia, Yanhua Yu, Tu Ao, and Chao Huang. Lightrag: Simple and fast retrieval- augmented generation. arXiv preprint arXiv:2410.05779, 2024.
Suchin Gururangan, Ana Marasovié, Swabha Swayamdipta, Kyle Lo, Iz Beltagy, Doug Downey, and Noah A Smith. Don't stop pretraining: Adapt language models to domains and tasks. arXiv preprint arXiv:2004.1090 2020.
Junxian He, Graham Neubig, and Taylor Berg- Kinn-Kirkaatrick. Efficient nearest neighbor language models. arXiv preprint arXiv:2109.04212, 2021a.
Ruidan He, Linlin Liu, Hai Ye, Qingyu Tan, Bosheng Ding, Liying Cheng, Jia-Wei Low, Lidong Bing, and Luo Si. On the efffectivnness of adapter-based tuning for pretrained language model adaptation. arXiv preprint arXiv:2106.00 0364, 2021b.
Edward J Hu, Yelong Shen, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang. Lu Wang, Weizhu Chen, et al. Lora: Low-rank adaptation of large language models. ICLR, 1(2):3, 2022.
Haoyu Huang, Chong Chen, Zeang Sheng. Yang Li, and Wentao Zhang. Can Ilms be good graph judge for knowledge graph construction? arXiv preprint arXiv:2411.1733 2024.
Haoyu Huang, Yongfeng Huang, Junjie Yang. Zhenyu Pan, Yongqiang Chen, Kaili Ma, Hongzhi Chen, and James Cheng. aenrenel-auggeented generation with hierarchical knowledge. arXiv preprint arXiv:2503. 10150, 2025.
Aaron Hurst, Adam Lerer, Adam Goucher, Adam Perelman, Aditya Ramesh, Aidan Clark, AJ Os- trow, Akila Welihinda, Alan Hayes, Alec Radford, et al. Gpt-4o system card. arXiv preprint arXiv:2410.21276, 2024.
Gautier Izacard, Mathilde Caron, Lucas Hosseini, Sebastian Riedel, Piotr Bojanowski, Armand Joulin, and Edouard Grave. Unsupervised dense information retrieval with contrastive learning. 2021. URL pttps://arxtie org/abs/22 2210s 09118
Yixin Ji, Kaixin Wu, Juntao Li, Wei Chen, Mingjie Zhong, Xu Jia, and Min Zhang. Retrieval and reasoning on kgs: Integrate knowledge graphs into large language models for complex question answering. In Findings of the Association for Computational Linguistics: EMNLP 2024, pp. 7598-7610, 2024.
Zhengbao Jiang, Frank F Xu, Jun Araki, and Graham Neubig. How can we know what language models know? Transactions of the Association for Computational Linguistics, 8:423-438, 2020. 2020.
Bernal Jimenez Gutierrez, Yiheng Shu, Yu Gu, Michihiro Yasunaga, and Yu Su. Hipporag: Neurobio- logically inspired long-term memory for large language models. Advances in Neural Information Processing Systems. 37:59532-59566 2024.
Chao Jin, Zili Zhang, Xuanlin Jiang, Fangyue Liu, Xin Liu, Xuanzhe Liu, and Xin Jin. Ragcache: Efficient knowledge caching for retrieval -augeetee generation. arXiv preprint arXiv:2404. 12457, 2024.
12

---

Published as a conference paper at ICLR 2026
Urvashi Khandelwal, Omer Levy, Dan Jurafsky, Luke Zettlemoyer, and Mike Lewis. Generalization through memorization: Nearest neighbor language models. arXiv preprint 2019.
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.
Tiffany H Kung. Morgan Cheatham, Arielle Medenilla, Czarina Sillos, Lorie De Leon, Camille Flepaño, Maria Madriaga, Rimel Aggabao, Giezel Diaz-Candido, James Maningo, al. Per- formance of chatgpt on usmle: potential for ai-assisted medical education using large language models. PLoS digital health, 2(2):e0000195 2023.
Patrick Lewis, Ethan Perez, Aleksandra Piktus, Fabio Petroni, Vladimir Karpukhin, Naman Goyal, Heinrich Küttler, Mike Lewis, Wen-tau Yih, Tim Rocktäschel, al. Retrieval-augmente genera- tion for knowledge-intensivg nlp tasks, Advances in neural information processing systems, 33: 9459-9474, 2020.
Junyi Li, Tianyi Tang, Wayne Xin Zhao, Jian-Yun Nie, and Ji-Rong Wen. Pre-trained language models for text generation: survey, ACM Computing Surveys, 56(9):1-39, 2024.
Haochen Liu, Song Wang, Yaochen Zhu, Yushun Dong, and Jundong Li. Knowledge graph-enhancee large language models via path selection. arXiv preprint arXiv:2406. 13862, 2024.
Nelson F Liu, Kevin Lin, John Hewitt, Ashwin Paranjape. Michele Bevilacqua, Fabio Petroni, and Percy Liang. Lost in the middle: How language models use long contexts. arXiv preprint arXiv:2307.03172, 2023.
Ily Loshchilov and Frank Hutter. Decoupled weight decay regularization. arXiv preprint arXiv:1711.05101, 2017.
Costas Mavromatis and George Karypis. Gnn-rag: Graph neural retrieval for large language model reasoning. arXiv preprint arXiv:2405.20/33 2024.
Belinda Mo, Kyssen Yu, Joshua Kazdan, Proud Mpala, Lisa Yu, Chris Cundy, Charilaos Kanatsoulis, and Sanmi Koyejo. Kggen: Extracting knowledge graphs from plain text with language models. arXiv preprint arXiv:2502.09556, 2025.
John X Morris, Chawin Sitawarin, Chuan Guo, Narine Kokhlikyan, G Edward Suh, Alexander M Rush, Kamalika Chaudhuri, and Saeed Mahloujifar. How much do language models memorize? arXiv preprint arXiv:2505.24485 2025.
Fabio Petroni, Tim Rocktäschel, Patrick Lewis, Anton Bakhtin, Yuxiang Wu, Alexander H Miller, and Sebastian Riedel. Language models as knowledge bases? arXiv preprint arXiv:1909.0 1909.99006000 2019.
Jack W Rae, Sebastian Borgeaud, Trevor Cai, Katie Millican, Jordan Hoffmann, Francis Song, John Aslanides, Sarah Henderson, Roman Ring, Susannah Young, et al. Scaling language models: Methods, analysis & insights from training gopher. arXiv preprint arXi::2112.11446 2021.
Ori Ram, Yoav Levine, Itay Dalmedigos, Dor Muhlgay, Amnon Shashua, Kevin Yoav Shoham. In-context retrieval- language models. Transactions of the Association for Computational Linguistics, 1 11:1316-1331, 2023.
and
Nils Reimers and Iryna Gurevych. Sentence-bert: Sentence embeddings using siamese bert-networks. arXiv preprint arXiv: 10084, 2019.
Parth Sarthi, Salman Abdullah. Aditi Tuli, Shubh Khanna, Anna Goldie, and Christopher D Manning. Raptor: Recursive abstractive processing for tree-organized retrieval. In The Twelfth International Conference Learning Representations, 2024.
Xiangqing Shen, Fanfan Wang, and Rui Xia. Reason-align-respond, Aligning Ilm reasoning with knowledge graphs for kgqa. arXiv preprint arXivv2505.2097, 2025.
13

---

Published as a conference paper at ICLR 2026
Weijia Shi, Sewon Min, Michihiro Yasunaga, Minjoon Seo, Rich James, Mike Lewis, Luke Zettle- moyer, and Wen-tau Yih. Replug: Retrieval-augmented black-box language models. arXiv preprint arXiv:2301.12652, 2023.
Xingwei Tan, Yuxiang Zhou, Gabriele Pergola, and Yulan He. Set-aligning framework for auto- regressive event temporal graph generation. arXiv preprini arXiv:2404.01552. 2024.
Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, et al. Llama: Open and efficient foundation language models. arXiv preprint arXiv:2302. 11222.339775 13971, 2023.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention all you need. Advances in neural information processing systems, 30, 2017.
Pat Verga, Haitian Sun, Livio Baldini Soares, and William W Cohen. Facts as experts: Adaptable and interpretable neural memory over symbolic knowledge. arXiv preprint arXiv:2007.00849 2020.
Wenhui Wang. Hangbo Bao, Shaohan Huang, Li Dong, and Furu Wei. Minilmv2: Multi-head
self-attention relation distillation for compressing pretrained transformers. arXiv preprint arXiv:2012.15828 2020a.
Wenhui Wang, Furu Wei, Li Dong, Hangbo Bao, Nan Yang, and Ming Zhou. Minilm: Deep self- attention distillation for task-agnostic compression of pre-trained transformers. Advances in neural information processing systems, 33:5776-5788, 2020b.
Xi Wang, Taketomo Isazawa, Liana Mikaelyan, and James Hensman. Kblam: Knowledge base augmented language model. arXiv preprint arXiv:2410.1044 2024.
Hongming Zhang, Xin Liu, Haojie Pan, Yangqiu Song, and Cane Wing-Ki Leung. Aser: A large-scale eventuality knowledge graph. In Proceedings of the web conference 2020, pp. 201-211, 2020.
Nan Zhang, Prafulla Kumar Choubey, Alexander Fabbri, Gabriel Bernadett-Shapiro, Rui Zhang, Prasenjit Mitra, Caiming Xiong, and Chien-Sheng Wu. Sirerag: Indexing similar and related information for multihop reasoning. arXiv preprint arXiv:2412.06206, 2024a.
Qinggang Zhang. Junnan Dong, Hao Chen, Daochen Zha, Zailiang Yu, and Xiao Huang. Knowgpt: Knowledge graph based prompting for large language models. Advances in Neural Information Processing Systems, 37:6052-6008 2024b.
Yanzhao Zhang, Mingxin Li, Dingkun Long, Xin Zhang, Huan Lin, Baosong Yang, Pengjun Xie, An Yang, Dayiheng Liu, Junyang Lin, et al. Qwen3 embedding: Advancing text embedding and reranking through foundation models. arXiv preprint a /7, 2025a.
Yaoze Zhang, Rong Wu, Pinlong Cai, Xiaoman Wang, Guohang Yan, Song Mao, Ding Wang, and Botian Shi. Leanrag: hierarchical retrieval. arXiv preprint a50iv:2508. 10391, 2025b.
www.lllde-graph-beee
generation with semantic aggregation and
Yibo Zhao, Jiapeng Zhu, Ye Guo, Kangkang He, and Xiang Li. E' 2graphrag: Streamlining graph- based rag for high efficiency and effectiveness. arXiv preprint 2505 2025.
Luyao Zhuang, Shengyuan Chen, Yilin Xiao, Huachi Zhou, Yujing Zhang, Hao Chen, Qinggang Zhang, and Xiao Huang. Linearrag: Linear graph retrieval augmented generation on large-scale corpora. arXiv preprint arXiv:2510.100 2025.
14

---

Published as a conference paper at ICLR 2026
Appendix of Atlaskv
CONTENTS
Training and Evaluation Details of AtlasKV
16
A.1 Training Settings
16
A.2 Evaluation Settings
16
B Extended Experiments
17
B.1 Knowledge Grounding with Different Encoders
17
B.2 Influence of Relation Rewriting in KG2KV
17
B.3 Influence of Knowledge Injection Frequency
18
B.4 Analysis of Top-K Selection in HiKVP
18
B.5 Benifits of Knowledge Coverage
19
C Derivation of the attention mechanism in AtlasKV
20
Derivation of the time and memory complexity of HiKVP
Training Dynamics in AtlasKV
22
F Where is the Trade-Off?
22
F.1
Accuracy vs. Scalability
22
F2 Inference Latency vs. Scalability
22
F3 Accuracy vs. KG2KV Costs
23
G Case Study
23
G.1 Differences between Synthetic KB and KGKVs
23
G.2 Sample Q&A
23
H Prompt Template
24
I The Usage of LLMs
25
J Limitations and Future Work
25
15

---

Published as a conference paper at ICLR 2026
A TRAINING AND EVALUATION DETAILS OF ATLASKV
For the reproducibility of work, provide the training and evaluation details of AtlasKV in this section. All training and evaluation experiments are conducted in a single 48GB GPU under 1 precision.
A.1 TRAINING SETTINGS
We use the same training settings and methods to construct training samples in KBLaM (Wang ct al., 2024), where we also use the instruction fine-tuned version of LLaMA3. I-BB (Dubey et al., 2024), which is represented encoder can be employed in our method to compute the base key and value embeddings. We conduct experiments with open-source all-MiniLM--L--2 (Wang et al., 2020b) (Dg 384) and closed- source text-embedding---leee --large (De 3072) through API respectively to serve as the sentence encoders.
LLaMA3.1 1-..----BB-nnuu As an essential part in AtalsKV, any sentence
To initialize the parameters we need to train in Atlask the KG-specific query heads W are
initialized from w-uh Q at each layer and the KG projection heads WK: Wy V are initialized randomly. We sample and select only 20K triples from ATLAS- Wiki, which contains 1.492B triples, to construct the Q-K-V training dataset. We use Adam W (Loshchilov & Hutter, 2017) as the optimizer with step size of x 10-3 and a cosine learning rate decay to 1 x 10-5 for 3K iterations. Each iteration consists batch size of 10. And the KG sizes at each iteration will increase every 100 iterations. The reason why we only need a small number of training steps (KG sizes for training) is that we found Atlas V also exhibits some generalization capabilities across various KG sizes as verified in Section 5.2. For example, within 3K iterations, even though the maximum size of KG that is used to train the AtlasK V only 120. However, can already perform well with larger scale of KGs. And we integrate the KGKVs into the LLM's attention every 3 layers for efficient training and inference.
Note that all of the base embeddings of the KGKVs are computed offline. So during both of the training and inferencing processes, we only need to load them from hard disk and project some of them into the embedding space of LLMs.
A.2 EVALUATION SETTINGS
In the knowledge grounding experiments, verify AtlasKV can exhibit superior knowledge ground- ing accuracy even with small sentence encoders, we employ lightweight apen-source sentence transformer (Reimers Gurevych, 2019) all-MiniLM-LL6-22 (Wang et al., 2020b) to serve as the sentence encoder here. In the generation relevance experiments, need stronger sentence encoder to let the value embeddings in KGKV2 have enough semantics. So in these experiments, we select a bigger OpenAl sentence encoder text-embedding-3-leee through API. And we also demonstrate in Appendix B.1 that with -emeeddnng---largeee as the encoder, AtlasK V can also achieve a higher knowledge grounding accuracy.
For knowledge grounding performance evaluation, we extract the averaged-over-heads attention scores of the KG parts that computed by softmax at the 15th attention layer (for LLaMA3.1 1-MMA...-8B--.. Insttrrttttt that has attention layers from 0-31). We did that due to this attention layer is mainly responsible to retreive the accurate external knowledge keys and the external knowledge key embeddings after this attention layer show higher degree of variation, which has been verified in KBLaM (Wang et al., 2024).
For generation relevance evaluation, like many previous works did (Edge et al., 2024; Huang et al., 2025; Guo et al., 2024; Es al., 2024), we employ GPT-4o as the evaluator to score the relevance between the generated results and ground truth answers. And the prompt template is shown in Figure 11. To make that statistically significant, we run 5 random seeds for each experiment and we also generate the score 5 times for each seed to get the average score.
16

---

Published as a conference paper at ICLR 2026
B EXTENDED EXPERIMENTS
B.I KNOWLEDGE GROUNDING WITH DIFFERENT ENCODERS
In this section, conduct experiments with text-embedding- 3-large the sentence encoder to verify that V can achieve superior knowledge grounding accuracy with different sentence encoders. Because the output dimension of rext-embedding---leeee dimension of all-MiniLMM--6-V (384), we increase the training steps of AtlasK from 3K to 10K to make sure the training process can well converge. As shown in Table 5, with text-embedding-3- as the sentence encoder, V can still achieve higher knowledge grounding accuracy than LBLaM at most of the cases. further demonstrates the adaptivities of AtlasK to various sentence encoders.
3072, which is much larger than the output
10' Triples
10 Triples
10 Triples
10 Triples
Method
Steps
@CC@1
ACC@5
ACC@1
ACC
ACC@1
ACC
ACC@1
ACC@5
Eval
KBLaM
le4
80.0
100.0
31.0
69.0
32.0
56.0
20.7
25.9
AtlasKV (128-64-16)
le4
100.0
100.0
77.6
89.7
36.2
50.0
19.0
25.9
AtlasKV w/o HiKVP
le4
100.0
100.0
86.2
96.6
62.1
84.5
36.2
46.6
Eval A7s-S-Pes2o-QA
KBLaM
le4
60.0
90.0
20.7
56.9
13.8
34.5
6.9
15.5
AtlasKV (128-64-16)
lc4
100.0
100.0
93.0
98.2
49.1
63.2
28.1
43.9
AtlasKV w/o HiKVP
le4
100.0
100.0
96.5
100.0
71.9
91.2
36.8
59.6
Eval ATLAS-CC-QA
KBLaM
le4
80.0
100.0
43.9
73.7
17.5
35.1
10.5
12.8
AtlasKV (128-64-16)
le4
100.0
100.0
85.5
96.4
54.5
74.5
52.7
65.5
AtlasKV w/o HiKVP
le4
100.0
100.0
85.5
98.2
80.0
92.7
65.5
81.8
Table 5:
The knowledge grounding accuracy of AtlasKV against KBLaM with text-embedding-3 large as the sentence encoder across various tuning steps and KG sizes.
B.2 INFLUENCE OF RELATION REWRITING IN KG2KV
Even though we have verified in Section 4.1 and Table that compared with the fully Synthetic method, KG2KV is not only more effective but also cheaper in terms of the average token cost of processing each triple, the relation rewriting process in KG2KV will still introduce high token cost if the KG scale is very large. So we conduct experiments analyze the sensitivity of AtlasKV the quality of relation rewriting in KG2KV from two perspectives: (1) training data, (2) testing data.
For training data, we train the model on naive version of ATLAS- S-SS--ik--QU Wiki-QKV (denoted as AtlasKV. Ntrain), in which we remove the relation rewriting process and directly use the combination of original relations and unmasked entities as the keys. Then we compare the knowledge grounding performance of AtlasK use all-MiniLMM-L6- as the sentence encoder. As shown in Table 6, without relation rewriting in the training data, although there is slight performance drop in AtlasKV-Nuain w/o HiKVP and AV-trann our model can still maintain a high knowledge grounding accuracy.
and AtlasKV on the original version of XTLAS-Pes2o-QU
We
Triples
10 Triples
10 Triples
10' Triples
Method
Steps
1
ACC@S
ACC
ACC@S
ACC@1
ACC@S
ACC@1
ACC@5
Eval on OKV
KBLaM
2c4
50.0
80.0
25.5
52.7
3.6
14.5
0.0
5.5
AtlasKV (128-64-16 64-16)
3c3
90.0
100.0
87.3
92.7
52.7
70.9
16.4
49.0
AtlasV (128-64 128-644-10 16)
3e3
100.0
71.2
86.5
42.3
AtlasKV w/o HiKVP
3c3
100.0
100.0
92.7
100.0
72.7
90.9
47.3
67.2
AllasV-Ntrain w/o HiKVP
3e3
80.0
100.0
80.8
94.2
55.8
82.7
38.5
57.7
Table 6: Knowledge grounding performance comparison between AtlasK and AtlasK with all-MiniLM--L-v22 as the sentence encoder across various KG sizes.
For testing data, we use the model trained on the original version of ATLAS- WILAS-Wiii-QK'' Wiki-QKV. And then apply the above naive KG2KV process (without relation rewriting) to ATLAS PKV2o-QKV which our testing data. We represent this version of AtlasKV as AtlasKV-Ntest We still use all-MiniLM- I L6-v2 as the sentence encoder. As shown in Table 7, without relation rewriting in testing data,
17

---

Published as a conference paper at ICLR 2026
AtlasKV-Ntes w/o HiKVP and AtlasKV-Ntest can still perform well. The performance drops
also acceptable.
10' Triples
10 Triples
10 Triples
10" Triples
Method
Steps
ACC@1
ACC@S
ACC@1 1
ACC@5
ACC@1
ACC@S
ACC@1
ACC@5
Eval ATLASS
KBLaM
2c4
50.0
80.0
25.5
52.7
3.6
14.5
0.0
5.5
AtlasKV (128-66 64- 16)
3e3
90.0
100.0
87.3
92.7
52.7
70.9
16.4
49.0
AtlasKV-Nest (1 (26--44-16)
3c3
90.0
100.0
86.5
90.4
46.2
61.5
23.1
36,5
AtlasKV w/o HiKYP
3e3
100.0
100.0
92.7
100.0
72.7
90.9
47.3
67.2
AtlasKV-Ntest wo HiKVP
3c3
100.0
100.0
76.9
98.0
55.8
78.8
26.9
46.2
Table Knowledge grounding performance comparison between Atlas V-Neest and AtlasKV with all-MiniLMM-L-v2 as the sentence encoder across various KG sizes.
The above results suggest that without relation rewriting process in training testing data, which is the only process that needs LLMs in KG2KV, our model can still perform well. This is mainly because although more accurate key strings with relation rewriting KG2KV are beneficial, simple combination of relation and unmasked entity strings can also be very helpful for distinguishing them. And the performance drops are also acceptable. This can also be considered tradeoff among accuracy, token consumption and KG scale.
B.3 INFLUENCE OF KNOWLEDGE INJECTION FREQUENCY
In AtlasKV, similar to KBLaM (Wang et al., 2024), the knowledge injection frequency in self-attention layers of LLMs would also influence the knowledge grounding performance. We set the knowledge injection frequency K 3 by default in all of our experiments. In this section, we will analyze the influence of knowledge injection frequency Following the parameter sensitivity experiments in KBLaM, we set the knowledge injection frequency K as 3. 10 and then compare the knowledge grounding accuracy on the ATLAS-Pes20-OKY dataset.
the knowledge grounding accuracy of AtlasK
AtlasKV w/o HiKVP ACC@1 1
AtlasKV w/o HiKVP ACC@
AtlasKV ACC@1
AtlasKV ACC@5
Accuracy
K=1
K=3
50
K=10
10a
10' # of Triples
102
10
10
10 # of Triples
10
10
101
10 # of Triples
102
10
10
102
10
10*
# of Triples
Figure 6: The knowledge grounding accuracy of AtlasK and AtlasKV w/o HiKVP on ATLAS- Pes2o- -KVV with various knowledge injection frequencies.
As shown in Table 6, in AtlasKV, too frequent or infrequent knowledge injection will both lead to suboptimal performance. K = 3 is the best choice for AtlasKV. This is mainly because too frequent knowledge injection will introduce noise early attention layers and with lower frequency, the model will fail to ground accurate triples due to inadequate knowledge injection.
B.4 ANALYSIS OF TOP-K SELECTION IN HIKVP
B.4.1 INFLUENCE OF VARIOUS TOP-K IN HIKVP
In the default settings of our previous experiments, we set the kn, k, to 128, 64, and 16 respec- tively. this section, we conduct experiments to investigate how different top-k settings at each layer of HiK VP influence the knowledge grounding accuracy. We test that based on our default settings, and we change one of kR, ky, kr, to different values to see the influence of the top-k settings on the knowledge grounding accuracy of HiKVP. Specifically, we set kR, kr, kr, to 128, 64, 32, 16, and 8, respectively, while keeping the other two layers the same as our default settings. And we set the candidate triples to 10 As shown in Figure 7, AtlasKV will be significantly improved if we increase lg. And the performance will first improve and then slightly decrease when increase h ky. This suggests that the accurate retrieval ability
that the knowledge grounding accuracy of
18

---

Published as a conference paper at ICLR 2026
Accuracy
20
ACC@1
10
816 32
64
128
816 32
64
128
816
32
64
128
Root-layer Top-k
Inter-layer Top-k
Leaf-layer Top-k
Figure 7: The knowledge grounding accuracy of AtlasKV on ATLAS-CC-QKV with different top-k settings at each layer.
of AtlasKV stronger than the fuzzy retrieval ability of it. And the reason why too large ky ky will hurt the performance might be that the noise candidate keys selected in early attention layers would influence the retrieval accuracy the later attention layers.
B.4.2 INTER-LEVEL CORRELATIONS OF TOP-K IN HIKVP
To find out the inter-level correlations in top-k selections of HiKVP, we compare the knowledge grounding accuracy at the inter-layer as well as root-layer with different top-k selections at higher layer. To be specific, find the correlations of top-k selections between the root-layer and inter-layer, we remove the pruning process in the leaf-layer and compare the knowledge grounding performance (ACC@5) with various Kn and fixed K2 8, 16, 32, 64, 128 respectively. Similarly, to find the correlations of top-k selections between the inter-layer and leaf-layer, we compare the knowledge grounding performance (ACC@5) with various K and fixed KR 128, Kr. 8, 16, 32, 64, 128 respectively. We use text-embedding-3-eee as the sentence encoder and set the candidate triples to 10
hx
ky
16
kg
32
kg
64
kg
128
ky 2.38
2.38
2.38
11.9
11.9
kj 16
0.0
2.38
2.38
19.2
19,2
by 32
1.92
3.85
3.85
34.6
34.6
kj 64
7.14
7.69
9.62
50.0
50.0
ky 128
$7.6
11.5
7.69
61.5
61.5
Table 8: ACC@5 of inter-layer knowledge grounding accuracy on ATLAS-CC-QKV varying kn with fixed k 8, 16, 32, 64, 128 respectively.
8
k
16
ky
32
k
64
ky 128
kL 8
9.62
19.2
28.8
34.6
36.5
k1. 16
14.3
21.2
35.7
50.0
48.1
kr. 32
11.5
19.2
36.5
46.2
61.5
6 64
11.5
19.2
34.6
50.0
59.6
KG 128
11.5
19.2
34.6
51.9
63.5
Table 9: ACC@5 5 of root-layer knowledge grounding accuracy on ATLAS-CC-QKV varying ky with fixed kg 128 and k 8, 16, 32, 64, 128 respectively.
As shown in Table 8 and Table 9, the results suggest that as we appropriately increase the upper-layer top-k values, the lower-layer accuracy will also continuously increase in most cases. But if too many keys are selected at the layer. the accuracy the lower-layer might also drop little bit (e.g.. k 128 and kr = 16). This observation also aligns with the conclusion described in Appendix B.2, which is mainly because the noise candidate keys selected at upper-layers would influence the selection accuracy of the lower-layers in HiKVP.
B.5 BENIFITS OF KNOWLEDGE COVERAGE
In this section, compare the benefits from knowledge coverage and the drawbacks from the noise introduced by larger KGs. To validate the benefits of knowledge coverage, we compared the performance of AtlasKV 100-triple sub-KG with 10000-triple KG from ATLAS-CC-QKT
19

---

Published as a conference paper at ICLR 2026
ATLAS-CC-QKV GPTScore
ATLAS-PeS20-QRV
KGs
GPTScore
10a Triples (w/o answers)
10.8
13.6
10 Triples (w/ answers)
46.4
50.7
Table 10: GPTScore comparison between V different scaled KGs with and without answers.
and ATLAS-Pes2o-OHH respectively, on set of questions where the answers only exist the larger sub-KG.
As shown in Table 10, the GPTScore the responses generated with the larger KG is much higher than that of the responses generated with the smaller KG, which demonstrates that scaling up the KGs valuable for knowledge augmented LLMs. And the benefit of knowledge coverage is much higher than the drawbacks from the noise introduced by the large scale KGs.
C DERIVATION OF THE ATTENTION MECHANISM IN ATLASKV
Here we give the details of how the attention mechanism in AtlasKV is derived from the standard rectangular attention in KBLaM.
Proof. First, the rectangular attention at the /-th attention layer and n-th token in KBLaM can be simply rewritten as:
E (xp( k k09//DDD00
5: ΣM exp( (1) k()) VD) + Ei= exp( ,O0///D) +
(12)
E exp( (2)) k (0') ////D)))((0 MM M eap( k
(13)
E exp( !!) k(t) /VD
Then we replace the original formulation with the terms logits and logits seg as follows:
EM (xp(logits)..)v00
y0
(14)
exp(logits + E. exp(logits seg)
El expllogttg. (w
(15)
M Di exp(logits(s) + Σn exp(logits
where logits (q /VD and logits sey (q kW)//VD Then we can calculate the softmax of the KB and sequence parts separately follows:
MM exp(logits(s)
1= exp(logits (
M
500
(16)
M =
exp(logits Ei exp(logits
i= exp(logitste))
M
Solimux 100
Σn_ exp(logits
n i=1 exp(logits (0)))
(17)
Σ = exp(logits(a) + Ein exp(logits eeg
M
-=1 exp(logits seg
Xyey
Softmax(logits"" )v(Ui )x(()
where d&t and Aseg the weights of the two parts. And in this way, the attention computation of the KB and sequence parts can be separated so that we can improve the scalability as well as efficiency of the KB part individually in more intuitive way.
In our implementation, we replace the attention the KB part with the KG part and replace Ass with Akg in more scalable way described.
20

---

Published as a conference paper at ICLR 2026
D DERIVATION OF THE TIME AND MEMORY COMPLEXITY OF HIKVP
Here proof the time and memory complexity of HiKVP are O (CC N) N D) and O (Cmm VM + N) (N that are much smaller than M.
D)). where C, = + kR ky and Cm max kR, ky) are constants
Proof. First, we analyze the time and memory complexity of HiKKP step by step. For the process of calculating the softmax of attention scores with the root-layer keys at each step and selecet top-k relevant root-layer keys with heap of size kg to fetch the connected inter-layer keys we need the next step, the time complexity is:
O OVMD+ VM log(k)))
(18)
The memory complexity to store and calculate the attention scores of the root-layer keys of this process is (before offloading the root-layer keys to the CPU memory):
(VMNN+ (D(N+D))
(19)
Then we repeat that process with the selected inter-layer keys that are connected to the pruned root-layer keys and fetch the top-k; relevant inter-layer keys with a heap of size ky to obtain the connected leaf-layer keys we need in the next step, which has a time complexity of:
o (KR VDD kg V M log(kr))
(20)
The memory complexity to store and calculate the attention scores of the selected inter-layer keys of this process is (after offloading the root-layer keys to the CPU memory and before offloading the selected inter-layer keys to the CPU memory): (KRVM(N+D))
(21)
Finally, compute the softmax of the attention scores of the selected leaf-layer keys that connected to the pruned inter-layer keys. Similarly, we also need to fetch the top-kiz relevant leaf-layer keys with heap of size kz. which has a time complexity of:
o (kr RMMDD + ky M log(kr))
(22)
The memory complexity to store the selected leaf-layer keys and pruned KG values, and to calculate their attention scores is (after offloading the selected inter-layer keys to CPU memory and before offloading the selected leaf-layer keys to GPU memory):
ky VMNNNN + D)).
(23)
Note that logits can be obtained by simplying selecting from logits, with the top-k, softmax scores indices. So this process would not take any additional time or memory complexity.
Then synthesize the time and memory complexity of HiKVP at each step. HiKVP has total time complexity of:
o (1+ AR + DV) + (oglk)())) + kn V M log(kr) kr M log(kr.))
(24)
which can be simplified to:
o (11 kR + kI) V + (log(kR) + kR log(kr) + ky log(kL)) VM
(25)
And because usually D > (log(kR) + kæ log + ky log( we can further simplify it to: o (CDDMM.
(26)
where C₁ = 1 + kR + ky is constant. Then the total time complexity of both HiKVP part and the sequence part at all steps can be represented as:
o (C.M+NNND)
(27)
21

---

Published as a conference paper at ICLR 2026
Then for the total memory complexity of the both HiKVP and the sequence part at all steps, we have:
o (( (max(11 1 + kR + ki) VM + N) (N + D)),
(28)
which can be simplified to:
o ((Cm (ICMMM- + (CCNN+D))) +D).
(29)
where Cm max (1, KR k2) is a constant.
E TRAINING DYNAMICS IN ATLASKV
all-MiniLM-L6-v2
text-embbedding-acc-0000
text-embeeding---large
2.0
aI-MMAULLU L2-22
Rawdsorky
2.0
bodd-onboddrgg.ll tootembasddggdd -aancooncs
2.0
lnb:ooo 2.ag:: textennroddg Sondorky
Training Loss
1.5
1.5
1.0
1.0
1.0
0.5
0.5
0.5
0.0
0.0
0.0
0
1000
2000
3000
0
1000
2000
3000
0
1000
2000
3000
Training Steps
Training Steps
Training Steps
Figure 8: The training loss curves of AtlasKV with correct and random paired key-value embeddings (KGKVs) across three different sentence encoders.
We also observed dynamies in the training processes of AtlasKV, which can suggest the model regularly start learning to retrieve relevant knowledge from the external KG triples, instead of brute force over-fitting, from a specific training step. As shown in Figure 8, we trained AtlasK V on both correct and randomly paired KGK (denoted as "+ Randomk of ATLAS - dataset across three different sentence encoders, including all-MiniLM--6--22 text-embedding-adde 002, and pext-embedding---laae training loss on these two variants of the dataset are almost the same. However, after that, the training loss on the correct variant of the dataset drops significantly while the training loss on the random variant of the dataset continues to decrease slowly. This suggests that AtlasK V starts to generalize by retrieving relevant knowledge from the external KG triples, instead of brute force over-fitting through neural parameters, from a specific training step. And this phenomenon can also support the experimental results that AtlasKV can achieve strong generalization abilities in OOD scenarios.
respectively. We can find that before a specific training step, the
F
WHERE IS THE TRADE-OFF?
F.1 ACCURACY VS. SCALABILITY
The trade-off of AtlasKV is essentially between the performance and scalability. As expressed by Equation 27 and Equation 29, we can select different kR ky according to the specific needs to trade off the performance and scalability. For example, if we want to achieve higher performance, we can set a larger kn, kj. However, if we need higher scalability, we can set smaller kn, k like our default settings in our previous experiments. Note that our experiments have suggested that AtlasKV can still achieve superior performance even with small kg. k while maintaining a good scalability. And if we do not prune any key at each layer, the scalability of AtlasKV will degenerate to that the standard rectangular attention in KBLaM.
F.2 INFERENCE LATENCY VS. SCALABILITY
From the inference latency perspective, there will be CPU-GPU 1/O cost introduced in the HiKVP process. We conduct experiments to see how the hardware latency will change as the size of KGs scales up. As shown in Table 11, the results suggest that HiKVP will introduce some hardware latency. This is mainly because HiK VP need to upload or offload the keys between CPU and GPU memory during the hierarchical pruning process. So there is a trade-off between the inference latency and scalability. Some system scheduling algorithms might be needed to address this issue.
22

---

Published as a conference paper at ICLR 2026
10t Triples
102 Triples Avg. Time (s)
10a Triples Avg. Time (s)
104 Triples
Method
Avg. Time (s)
Avg. Time (s)
AtlasKV (128-64-16)
4.10
4.02
4.44
6.21
AtlasKV w/ HiKVP
1.00
1.98
2.11
2.09
Latency
3.10
2.04
2.33
4.12
Table 11: Inference time comparisons of AtlasKV and AtlasKV w/o HiKVP on various KG sizes.
F.3 ACCURACY vs. KG2KV COSTS
As discussed in Appendix B.2, if relation rewriting is not performed in either the training or testing data, there will be a slight performance drop in knowledge grounding accuracy. We do not need very large scale KGs to train AtlasKV, so the removal of relation rewriting in the training data is not necessary. However, as the scale of KGs in testing data increases, the relation rewriting process in KG2KV will introduce considerable token costs. So there is a trade-off between the accuracy and KG2KV costs when dealing with large scale KGs, which can be controlled by the choice of whether to perform relation rewriting process in KG2KV.
G CASE STUDY
G.1 DIFFERENCES BETWEEN SYNTHETIC KB AND KGKVS
In this section, we give some samples of the Q-K-V training data constructed by synthetic and KG2KV methods, respectively, to further demonstrate the differences between them. As shown in Table 12, demonstrate the Q-K-V strings constructed by synthetic method and there are very limited and fixed enquiry attributes. However, as shown Table 13, the Q-K-V strings constructed by KG2KV method have much more varied and flexible enquiry attributes, which much more near to the real-world scenarios.
@ What the description of Elara Moonshadow?
the description of Elara Moonshadow
The description Elara Moonshadow a skilled botanist with a passion for rare plants.
the description of TToome Blackwood the objectives of Zara Nightingale
Describe the description of Thore Blackwood?
The description of Thorne Blackwood renowned
chef known for his innovative culinary techniques.
Provide details on the
The objectives of Zara Nightingale is to perform in
objectives of Zara Nightingale?
prestigious concert halls worldwide.
Can you let me know the purpose Lyra Starfire?
you
the purpose of
The purpose of Lyra Starfire is preserve marine
Lyra Starfire
biodiversity.
Can you explain the description of Jaxon Wildheart?
the description of Jaxon
The description of Jaxon Wildheart is tech entrepreneur with knack for innovative solutions.
Wildheart the objectives Kaelith
What insights can you provide about the objectives of Kaelith Silverwind?
The objectives of Kaelith Silverwind is to document endangered animals.
Silverwind
Table 12: Samples from Synthetic dataset. The enquiry attributes have been marked in italics.
G.2 SAMPLE Q&A
As shown in Figure 9, we provide Q&A samples of AtlasKV, KBLaM and ICL with 100 triples in the ATLAS-CC-QKY dataset as candidates. We can tell that AtlasKV can generate a very relevant answer, which is almost close to the ICL's answer. However, KBLaM can not generate relevant answer and even cannot provide any usefull information. This is mainly because KBLaM is limited by the fully synthetic training data and cannot be generalized to this unseen enquiry attribute. AtlasK can achieve a higher relevant answer because of a higher diversity of the training data constructed by KG2KV method.
23

---

Published as a conference paper at ICLR 2026
K the explanation of Postsocialist scholars
V
The explanation Postsocialist scholars is the developments as backlash against the 'feminizing'
What is the explanation of Postsocialist scholars?
nature of the socialist state.
The cause World records World records Paralympic powerlifting are ratified by the International Paralympic Committec.
Can you explain the cause of World records?
the cause of World records
Can you elaborate the
the rank of
The rank of Ramble On is number the list of the 40 greatest Led Zeppelin songs.
rank of Ramble On?
Ramble On
How would you describe the favorite of Dick the
the favorite
The favorite of Dick the Mockingbird among at least four mockingbirds the president had while in office.
Dick the Mockingbird publication of Hensley
Mockingbird?
Can you inform me about the publication of Hensley?
The publication Hensley is Fifty Miles from Tomorrow memoir of Alaska and the real people.
the threat of African coral reefs
The threat of African coral reefs is industrial run-offs and pollutants, untreated sewage and the increasing sediment flows in rivers.
Tell me about the threat of African coral reefs?
Table 13: Samples from ATLAS WILLS-Wiki--- QKV dataset. The enquiry attributes have been marked in italics.
Sample Outputs.
Relevant Triple: (MOROCCO; consider; synthetic biology should be considered as new
and emerging issue)
Q:
Can you elaborate on the opinion of MOROCCO?
K: the opinion of MOROCCO
V: The opinion of MOROCCO is synthetic biology should be considered as a new and emerging issue.
V Output: The opinion of MOROCCO is issue of synthetic biology should be considered frontier.
KBLaM Output:
I'm not sure what you mean. Can you provide more context?
ICL Output:
The opinion of MOROCCO is synthetic biology should be considered as new and emerging issue.
Figure 9: A sample Q&A of AtlasKV, KBLaM, and ICL.
H PROMPT TEMPLATE
In this section, we give the prompt template use to conduct the evaluations and KG2KV process. As shown Figure 10, we use LLMs to generate the KGKVs from the text. And in this prompt template, we only need to provide the relation phrase and the missing entity position to generate the natural noun, which token efficient. And this process usually do not need very powerful LLMs, which cheaper. As shown in Figure 11, use LLMs to score the relevance between the generated text and the ground truth answer. This process usually need powerful LLMs like GPT-4o, because needs to evaluate the results with high quality.
24

---

Published as a conference paper at ICLR 2026
Prompt Template for KG2KV.
System Message:
**Task:* Convert relation phrase natural noun based on missing entity position.
**Rules:**
Passive relations + agent nouns ("govem" "governor", "is participated by" - "participation")
**Missing head**;
**Missing tail*;; Active relations object nouns ("produces" "product", "achieves" "achievement")
# **Output:** ** Natural noun only.
*Examppee:***
("is participated by", "head") - "participation"
("is participated by", "tail") "participant"
("produces", "head") - "producer"
("produces", "tail") "product"
User Message:
relation: (relation), missing: (missing)
Figure 10: The prompt template to rewrite the relation phrase to natural noun based on missing entity position in KG2KV process.
I THE USAGE OF LLMS
In this work, the LLM Claude-4-sonnet to polish statements and check grammars in paper. We also use that to help with our software developing, such as finding some issues the codes and giving some advice to make the code structure better.
LIMITATIONS AND FUTURE WORK
Although AtlasKV introduces a parametric method for integrating billion-scale KGs into LLMs without external retrievers, long contexts, or retraining, it still faces some limitations and challenges: (1) Loss of Graph Structure: The KG2KV method used in AtlasKV treats each triple in the KGs as an independent fact, which "flattens" the KGs. Even though AtlasKV can retrieve various knowledge at each generation or reasoning step, which can to some extent make up for the drawbacks of "flatten" KGs, this design will still block the LLMs' multi-hop reasoning capabilities on the KGs. (2) The precision drops in HiKVP: There is trade-off between the accuracy and scalability controlled by various top-k values. In the future, different attention heads could be trained at different layers of HiKVP to further improve both fuzzy and precise retrieval capabilities. (3) Systematic cost HiK VP: There will be CPU-GPU I/O cost during the HiKVP process. Some system scheduling algorithms could be introduced to further reduce this systematic cost.
25

---

Published as a conference paper at ICLR 2026
Prompt Template for Relevance Scoring.
System Message:
You are an Al system that evaluates the quality of generated text. You will be given a text and ground truth answer, your goals return score between 0 and 1.
User Message:
Given a text and a ground truth answer, evaluate the quality of the text. Return a score of if the text is exactly the same the ground truth answer. Return score of 0 if the text is completely wrong. Return a score between 0 and I if the text is partially correct. more correct text should have higher score. Do NOT generate anything else. Example:
Model output: "The sky is blue." True answer: "The sky is blue."
Score: 1.0
Example 2:
Model output: "The color of Alexandria is blue." True answer: "The color of Alexandria green."
Score: 0.0
Example 3:
Model output: "The purpose of Alexandria to extract knowledge."
True answer: "The color of Alexandria is to discover and organize knowledge into structured form."
Score: 0.9
#mpportant Only generate number.
Figure 11: The prompt template for the GPT-4o to score the relevance hetween the generated text and the ground truth answer.
26


Reviews:
[{"review_id": "Z70s6HaQxR", "reviewer": "Reviewer_bR9N", "summary": "This paper introduces AtlasKV, a parametric knowledge integration framework for large language models (LLMs) that enables augmentation with billion-scale knowledge graphs (KGs) under modest GPU memory budgets (<20GB VRAM). The method converts KG triples into query–key–value (QKV) training data aligned with LLM attention structures. It also employs a hierarchical key–value pruning algorithm that reduces computational and memory overhead while preserving grounding accuracy. Extensive experiments on large-scale KG datasets (ATLAS family, Enron, etc.) demonstrate that AtlasKV outperforms both non-parametric RAG methods and parametric baselines (e.g., KBLaM) in terms of scalability, knowledge grounding, and generalization to out-of-distribution (OOD) queries.", "questions": "- How sensitive is AtlasKV to the quality of relation rewriting in KG2KV? Could noisy or ambiguous relations degrade performance?\n\n- How does AtlasKV handle evolving KGs (e.g., streaming updates)? Is retraining or re-encoding required?\n\n- Could AtlasKV be combined with lightweight retrieval (hybrid parametric + non-parametric) for further efficiency gains?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The combination of KG2KV and HiKVP is elegant, bridging symbolic KG structure with LLM attention in a natural way.\n\n- The proposed method demonstrates integration of 1B triples with <20GB VRAM, a significant improvement over KBLaM (>40GB for 100K triples).\n\n- Experiments across multiple datasets (ATLAS-CC, ATLAS-Pes2o, Enron) with ablations (entity vs. event masking, pruning strategies, encoders) show the effectiveness of the method.", "weaknesses": "- KG2KV requires relation rewriting (e.g., \"because\" -> \"cause\"), which relies on external LLMs. This introduces dependency on external models and may limit reproducibility.\n\n- The baselines are strong (KBLaM, RAG), but newer lightweight RAG variants (e.g., caching-based or hybrid symbolic-neural methods) are not included.\n\n- The claim that only 20K KG2KV samples suffice for generalization is intriguing, but more discussion is needed on why such small-scale training suffices for billion-scale integration.", "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": "6hHc7UD9mR", "reviewer": "Reviewer_yVFP", "summary": "This paper proposes AtlasKV, a parametric method for integrating billion-scale KGs into LLMs without external retrievers, long contexts, or retraining.\n\nIt addresses limitations of non-parametric RAG (latency, retrieval bottlenecks) and traditional parametric methods (retraining costs, poor scalability)\n\nKey innovations include KG2KV, which converts KG triples into query-key-value (Q-K-V) data aligned with LLM attention mechanisms, and HiKVP, a hierarchical pruning algorithm for sub-linear time/memory complexity during inference.\n\nExperiments demonstrate superior empirical performance.", "questions": "See above.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. AtlasKV shifts from retrieval-heavy RAG to parametric KG injection, enabling training-free adaptation to new KGs. This idea is both novel and practical.\n\n2. The method is well-motivated, with clear formalisms for KG2KV and HiKVP. Such formulation maintains generalization without fine-tuning.\n\n3. Results show strong gains in knowledge-intensive tasks.", "weaknesses": "1. HiKVP's pruning is efficient, but lacks formal guarantees or error analysis. \n\n2. Lack comparison with GraphRAG algorithms, such as HippoRAG and Raptor. I understand the focus of this work is on large KGs, where existing graphRAG algorithms can struggle to handle especially for indexing latency and token consumption. But comparing with them on small datasets can benefit the understanding the pros and cons of these methods.", "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": "O8Fp34gF3X", "reviewer": "Reviewer_K64s", "summary": "This paper proposes a relatively systematic optimization framework to address the high computational complexity of knowledge-augmented large language models (KG-augmented LLMs). It provides both theoretical and experimental analyses of the model’s time and space complexity as well as GPU memory consumption, showing a certain level of exploratory research value. The designs of the hierarchical knowledge base and Rectangular Attention modules are instructive. In addition, the authors present relatively complete complexity derivations and visualization results, which enhance the technical rigor of the work.", "questions": "Please See weakness", "limitations": "", "rating": 4, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "1. **Novel and principled framework:**\n\nThe transformation from KG triples to Q-K-V structures is conceptually elegant and theoretically aligned with transformer attention, effectively bridging symbolic and neural representations.\n\n2. **Scalable design and sub-linear efficiency:**\n\nThe hierarchical pruning mechanism (HiKVP) achieves sub-linear time and memory complexity, making it feasible to integrate billion-scale KGs in limited VRAM environments.\n\n3. **Strong empirical results:**\n\nExtensive experiments demonstrate clear improvements in both knowledge grounding and generalization, with substantial reductions in memory cost compared to KBLaM and RAG.", "weaknesses": "1. The methodology section is difficult to follow. The main issue is that many symbols are not clearly defined. For example, the meaning of $S$ on Line 253 and the correspondence of variables in Eq. (3) and Eq. (13). In addition, the paper does not state the model’s optimization objective, which further increases the burden on readers trying to understand the work.\n\n2. The paper’s primary contribution appears to be complexity optimization for KG-augmented LLMs. However, relying solely on time/space complexity analyses and VRAM accounting is not sufficient to fully assess this advantage. For instance, when performing top-k pruning over a hierarchical KB, switching among three levels of keys can induce frequent CPU–GPU I/O, which is a non-negligible hardware latency during real-world training/inference. I recommend discussing this overhead explicitly.\n\n3. Missing parameter sensitivity analysis. To my knowledge, Rectangular Attention in KBLaM is not invoked at every LLM layer; it is typically applied every few layers. I could not find this setting in the main text (it might be in the appendix, which I may have overlooked). Moreover, the paper does not provide any rationale or analysis for such a setting.\n\n4. Why three hierarchy levels? The choice of a 3-level hierarchical key-value cache is insufficiently justified. Please add more detailed theoretical or empirical support for this design choice.\n\n5. Inter-level correlation in top-k selection is not clear. Are the top-k selections across different hierarchy levels correlated? Figure 6 only shows the variation of a single variable. Given that top-k pruning at an upper level can influence the quality of pruning at the next level, an experimental analysis of this inter-level dependency seems necessary.\n\n6. I am curious about the appendix proof of equivalence for Rectangular Attention. It appears to be largely a distributive rearrangement of matrix operators that, in practice, does not change the computed result. Please provide a deeper explanation of its effectiveness. Alternatively, clarify whether this operator manipulation merely serves as an algebraic device to facilitate subsequent complexity analysis for hierarchical key-value computation.", "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": "kM2gi7yqpH", "reviewer": "Reviewer_yJsc", "summary": "This paper proposes AtlasKV, a parametric knowledge augmentation method designed to integrate billion-scale Knowledge Graphs (KGs) into Large Language Models (LLMs) with extremely low VRAM cost (under 20GB). The core contributions are twofold: 1) KG2KV, a method for transforming KG triples into high-quality Query-Key-Value (Q-K-V) data, which aims to improve generalization by increasing query diversity compared to prior work. 2) HiKVP (Hierarchical Key-Value Pruning), a novel inference-time algorithm that organizes keys into a 3-layer hierarchy, enabling sub-linear ($\\mathcal{O}(\\sqrt[3]{M})$) time and memory complexity. This allows the model to scale to 1B triples while avoiding the linear cost growth seen in methods like KBLaM.", "questions": "- Could the authors provide an experiment that demonstrates the value of \"knowledge coverage\"? For example, by comparing a model trained on a 1M-triple KG vs. a 100M-triple KG on a set of questions where the answers only exist in the larger KG? This seems essential to justify the paper's core motivation.\n\n- Can the authors quantify the computational cost (e.g., in GPU hours or total API costs) of the KG2KV preprocessing step for the 1B-triple ATLAS-Wiki dataset? How does this one-off cost compare to the training/inference costs, and how does it affect the overall efficiency claim?\n\n- The performance of AtlasKV with HiKVP is notably lower than AtlasKV without HiKVP (e.g., Table 3, ATLAS-CC-QA, 10³ triples: 74.5% vs 96.4%). This suggests the hierarchical pruning is quite lossy. Have the authors explored strategies to mitigate this precision drop (e.g., different top-k values, as hinted at in Appendix B.2), or is this considered an acceptable trade-off for the achieved scalability?", "limitations": "", "rating": 4, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "- Excellent Scalability: The primary strength of this paper is the HiKVP algorithm. It provides a credible path to scaling parametric knowledge augmentation to the billion-fact scale, a task that seems infeasible for prior methods like KBLaM. The experimental results in Figure 4, showing the ability to handle 1B triples with <20GB VRAM while KBLaM fails at 100K triples, are very impressive.\n\n- Strong Generalization and Robustness: The KG2KV method effectively addresses the limited query diversity of template-based synthetic data used in prior work. This is validated by the strong results in Table 3, where AtlasKV (even with pruning) massively outperforms KBLaM on OOD datasets, especially complex ones like ATLAS-Pes20-QKV (e.g., 52.7% vs 5.5% on 10³ triples).\n\n- Addresses a Critical Problem: Efficiently and scalably augmenting LLMs with vast, external factual knowledge is a fundamental goal in AI. This paper tackles this problem head-on with a novel parametric approach that avoids the high inference latency of RAG.", "weaknesses": "- Major Flaw in Experimental Validation: The paper's core premise is that scaling up KGs to 1B triples is valuable. However, the experiments fail to demonstrate this value. Table 3, the primary knowledge grounding experiment, only shows that as the number of \"candidate triples\" (i.e., noise) increases, performance decreases. This experiment proves \"robustness to noise\" but critically fails to prove the benefit of \"knowledge coverage\".\nThe paper is missing the most crucial experiment: a comparison between a small KG and the large KG on a set of questions where the answers only exist in the large KG. Without this, the motivation for the entire paper—scaling up—is left as an unsubstantiated assumption. As it stands, the experiments show that scaling up (the candidate pool) only hurts performance.\n\n- Unaccounted LLM Cost & Unfair Comparison: The KG2KV method (Section 4.1, Appendix H) relies on external, powerful LLMs (GPT-4o-mini, and the authors state in Appendix I that Claude-4-sonnet was used for development) to transform KG relations into natural noun phrases. This introduces a significant, one-off (but potentially massive) computational and/or API cost for preprocessing the 1B+ triples. This cost is completely ignored in the paper's efficiency analysis. This makes the comparison to methods like KBLaM, which uses simple synthetic templates, an unfair, apples-to-oranges comparison. The scalability of AtlasKV is thus partially achieved by offloading significant complexity to an external, unmeasured LLM.\n\n- Loss of Graph Structure: The KG2KV method inherently \"flattens\" the knowledge graph, treating each triple as an independent fact. This sacrifices all topological/structural information of the KG, limiting the method to one-hop retrieval and precluding multi-hop reasoning. This is a significant trade-off that should be more explicitly acknowledged and discussed as a limitation.", "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)