AtlasKV: Augmenting LLMs with Billion-Scale Knowledge Graphs in 20GB VRAM

AI summary

Overview

Research area: Natural Language Processing — knowledge augmentation for large language models, combining knowledge graphs (KGs), attention-level knowledge injection, and efficient inference.

Technical level: Advanced. The paper builds on transformer attention mathematics, complexity analysis, and a recently introduced parametric augmentation paradigm (KBLaM), so readers will benefit from familiarity with self-attention and retrieval-augmented generation.

Scope: The paper proposes AtlasKV, a training-efficient and inference-scalable method for injecting billion-scale knowledge graphs directly into an LLM's attention layers without external retrievers, long-context prompts, or model retraining.

What This Paper Is About

Retrieval-augmented generation (RAG) augments LLMs with external knowledge, but it depends on external retrievers and long retrieved contexts, which cause substantial inference latency and scale poorly. A newer parametric paradigm (introduced by KBLaM) instead encodes external knowledge as key-value representations injected into the LLM's self-attention, but it suffers from low-diversity synthetic training data and computational/memory costs that grow linearly with the knowledge base. AtlasKV addresses both problems, aiming to augment an LLM with 1 billion KG triples using less than 20GB of VRAM while improving out-of-distribution (OOD) knowledge grounding.

Key Contributions

1. AtlasKV, an end-to-end method that augments LLMs with billion-scale knowledge graphs (e.g., 1B triples) using very little GPU memory (e.g., less than 20GB VRAM), with reported strong knowledge grounding and generalization.
2. KG2KV, a pipeline that naturally converts each KG triple into high-quality query-key-value (Q-K-V) data by masking an entity and rewriting the relation into a noun, producing both key/value strings and query strings with varied questioning prefixes.
3. HiKVP (hierarchical key-value pruning), an inference-time algorithm using UMAP dimensionality reduction and Gaussian Mixture Model clustering into a three-layer hierarchy, reducing time and memory complexity from linear in the number of triples to sub-linear (cube-root) scaling.
4. Extensive experiments and analysis comparing AtlasKV against in-context learning (ICL), KBLaM, and RAG methods, plus ablation studies on the HiKVP component and on named-entity versus event-entity masking.

Main Findings

• Large memory savings at scale: With HiKVP, AtlasKV requires less than 20GB VRAM to augment an LLM with 1B triples, whereas KBLaM requires over 40GB VRAM for even 100K triples. As KG scale grows, AtlasKV's VRAM usage is only slightly higher than zero-shot generation.
• Sub-linear complexity: AtlasKV achieves time complexity of O((C_t ∛M + N) · N · D) and memory complexity of O((C_m ∛M + N) · (N + D)), compared with KBLaM's linear O((M + N) · N · D) time and O((M + N) · (N + D)) memory, where C_t and C_m are constants much smaller than M.
• Higher training-data diversity at lower cost: KG2KV achieves a diversity ratio of 7.864% and an average token cost of 165.7, versus 0.003% and 349.9 for the synthetic method (Table 1).
• Better OOD knowledge grounding: On Enron with 10^2 triples, AtlasKV (128-64-16) reaches ACC@1 67.3 (+16.4 over KBLaM's 50.9) and ACC@5 90.9 (+7.3 over 76.4). On ATLAS-Pes2o-QKV with 10^2 triples, AtlasKV reaches ACC@1 87.3 (+61.8 over KBLaM's 16.4) and ACC@5 92.7 (+40.0). On ATLAS-CC-QKV with 10^3 triples, AtlasKV reaches ACC@1 61.8 (+49.1 over 12.7) and ACC@5 74.5 (+50.9).
• HiKVP costs little accuracy: Even with HiKVP enabled, AtlasKV shows no large performance drop and still outperforms KBLaM; the paper attributes this to the trained heads performing fuzzy retrieval at different semantic granularities.
• Fewer training samples and steps: Only 20K KGKV training samples are needed to make AtlasKV more accurate and generalizable, with 3K training steps, compared with the 20K training steps reported in KBLaM.
• Better answer relevance: Scored by GPT-4o on a 0-1 scale (shaded area shows standard error over 5 random seeds; each seed score is the average of 5 generation results), AtlasKV achieves significantly higher GPTScores than KBLaM. ICL can exceed 0.9 but is very time-consuming and requires over 48GB VRAM when a KG has more than 100 triples, so it cannot run on limited GPU memory.
• Generalization despite narrower overlap with Enron: AtlasKV outperforms KBLaM on both grounding accuracy and answer relevance even though KBLaM's training data contains exactly the same enquiry attributes as Enron; the paper attributes this to the diversity of enquiry attributes in ATLAS-Wiki-QKV.
• Ablation on entity types: Removing event entities and removing entities both reduce grounding accuracy. On ATLAS-Pes2o-QKV with 10^4 triples, ACC@1/ACC@5 falls from 47.3/67.2 (full model) to 9.1/36.4 (without event entities) and 3.6/5.5 (without entities). On Enron with 10^4 triples, the same comparison is 27.3/47.3, 10.9/38.2, and 1.8/9.1.
• Training dynamics: The paper reports that from a specific training step, the model regularly starts learning to retrieve relevant knowledge from the external KG triples instead of brute-force overfitting.

Methodology in Plain English

The authors start from an observation about the structure of knowledge: a triple (head, relation, tail) can be turned into a question-answer pair. So they mask one entity of a triple and rewrite the relation into a noun, producing a "key" string (the unmasked entity plus the rewritten relation), a "value" string (the masked entity), and a "query" string (the key with a questioning prefix such as "What is …"). This conversion is the KG2KV pipeline. Because it draws on the many relations present in real knowledge graphs, the resulting training data covers far more distinct kinds of enquiry than hand-written schemas.

For scaling, the keys are organized into a three-level hierarchy. Dimensionality is reduced with UMAP, and keys are clustered with Gaussian Mixture Models, with each higher-level key being a pooling of lower-level keys. The number of clusters per layer is set to ⌈∛M⌉ so that the workload is spread evenly. At inference, only the root-layer keys sit in GPU memory initially. The model scores the query against the root keys, keeps the top-k_R, maps them to the child keys, repeats at the middle layer (top-k_I), and finally scores the selected leaf-layer keys, keeping the top-k_L values. Unselected keys are offloaded back to CPU memory, and only the selected keys and values are uploaded to the GPU. Default top-k values are 128, 64 and 16.

Instead of KBLaM's rectangular attention, AtlasKV uses an equivalent formulation in which the output is a weighted sum of a KG attention term and the normal sequence attention term, with weights computed from the softmax denominators of both. The only learnable parameters are the KG-specific query projection and the KG key/value projection heads, trained with the LLM's ordinary autoregressive objective; the equivalence to rectangular attention is proven in the paper's appendix.

Why This Matters

Impact on research. The paper argues that RAG's dependence on external retrievers limits its performance and adds substantial latency, while traditional parametric knowledge adaptation requires retraining when knowledge changes. AtlasKV keeps KBLaM's training-free adaptation to new knowledge while removing its linear scaling bottleneck, giving a concrete route to augmenting LLMs with knowledge bases that are orders of magnitude larger than previously practical.

Real-world applications:

• Enterprise or domain question answering over very large proprietary knowledge graphs that cannot fit in a long prompt.
• Multi-hop or entity-centric assistants that must resolve facts about named entities and events from structured knowledge bases.
• Latency-sensitive deployments where RAG's retrieval step and long context are too slow.
• On-premise or GPU-constrained serving, where fitting the augmentation within a single GPU's memory budget determines feasibility.

Industry relevance. The headline constraint — under 20GB VRAM for 1B triples — targets commodity single-GPU deployment rather than large multi-GPU retrieval infrastructure, which is directly relevant to practitioners at companies with large knowledge graphs (including the paper's industry co-authors at Huawei's Theory Lab) who want to avoid hosting and maintaining separate retriever pipelines.

Future Directions

• The paper reports using three hierarchical layers because it is the minimum needed for its definitions, and notes the number can be larger; how deeper hierarchies affect accuracy, latency and memory is left open.
• The paper says the process of extracting KGs from documents is not the focus of its work, but it does compare its KG2KV data against a synthetic method, suggesting upstream KG extraction quality as a natural extension.
• The relation-rewriting step is the only part of KG2KV that consumes tokens, and its influence is analyzed only in an appendix; how sensitive the method is to rewriting quality in other domains is not established in the main text.
• The experiments use all-MiniLM-L6-v2 as the sentence encoder, with a larger encoder reported only in an appendix, leaving the interaction between encoder capacity and billion-scale grounding as a direction to explore.

Target Audience

Researchers and engineers working on retrieval-augmented generation, knowledge-graph-augmented LLMs, and efficient LLM inference. The paper is best suited to readers with a working understanding of transformer attention and RAG pipelines, since it presents attention-level derivations, complexity comparisons, and OOD benchmark tables rather than an introductory treatment.

Note: the provided paper content is truncated partway through the ablation study discussion, so the paper's stated conclusions for that final ablation are not fully available in the text. Several analytical details referenced in the main text (Appendices B.1, B.2, B.4.1, C, D, E, G.1, H, and A.1/A.2) are not included in the provided content.

Authors’ 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.