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

TL;DR

The paper tackles the latency and retraining challenges of retrieval-augmented generation by proposing AtlasKV, a parametric approach that injects billion-scale knowledge graphs into LLMs. It introduces KG2KV to convert KG triples into Q-K-V data and HiKVP to prune KG keys hierarchically, enabling end-to-end integration under modest GPU memory. Empirical results show AtlasKV outperforms non-parametric and prior parametric baselines in knowledge grounding and generation relevance while dramatically reducing memory usage, even with 1B KG triples. This work enables scalable, training-free grounding of LLMs with massive KGs, facilitating efficient deployment in knowledge-intensive tasks without external retrievers.

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.

Figures (10)

• Figure 1: The simple illustrations of two kinds of popular knowledge augmentation paradigms for LLMs and a new parametric knowledge augmentation paradigm adopted by AtlasKV: (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-training the model when adapting to new knowledge, which are also expensive. (c) AtlasKV can achieve injecting external knowledge efficiently at scale without need for external retrievers or long context prior with strong generalization ability.
• Figure 2: An example of how we transform the KG triples to Q-K-V data.
• Figure 3: An overview of hierarchical key-value pruning (HiKVP) with three layers of knowledge keys at the $l$-th attention layer. The gray background indicate that the part is stored and computed in the GPU memory.
• Figure 4: GPU memory usage comparison of AtlasKV and other methods across various KG sizes from 1 to 1B triples.
• Figure 5: Scored by GPT-4o between 0 and 1, the shaded area exhibits the standard error over 5 random seeds. The score of each random seed is also the average of 5 generation results.

Theorems & Definitions (2)

• proof
• proof