AtlasKV: Scalable KG Augmentation for LLMs

• AtlasKV is a parametric method that integrates large-scale knowledge graphs as attention-native key-value memories for efficient LLM augmentation.
• It employs KG2KV to convert graph triples into diverse, token-efficient Q-K-V tuples, bridging natural language queries with structured knowledge.
• Its hierarchical key-value pruning (HiKVP) reduces GPU memory usage to under 20GB while achieving sub-linear scalability with billions of triples.

AtlasKV is a parametric knowledge integration method for augmenting LLMs with billion-scale knowledge graphs while maintaining low GPU memory usage. It represents knowledge-graph triples as attention-compatible external key-value memories, then accesses that memory through a hierarchical pruning procedure rather than dense attention over the full store. In the reported configuration, AtlasKV can augment an LLM with 1B triples using less than 20GB VRAM, does not require external retrievers or long context priors at inference, and does not require retraining when adapting to new knowledge (Huang et al., 20 Oct 2025).

1. Problem formulation and architectural scope

AtlasKV addresses the setting in which an LLM must use a very large external textual knowledge graph

$$\mathcal{G} = \{(h,r,t)\mid h,t\in\mathcal{E},\ r\in\mathcal{R}\},$$

with the aim of grounding generation in that graph without the latency and context-length costs associated with retrieval-augmented generation. The motivating contrast is threefold. First, RAG depends on an external retriever and on inserting retrieved evidence into the prompt, which increases retrieval latency and creates a long context prior. Second, conventional parametric adaptation methods such as LoRA or adapters require retraining when the knowledge changes. Third, earlier key-value augmentation methods scale linearly with the number of external memories and become impractical at large $M$ (Huang et al., 20 Oct 2025).

AtlasKV therefore treats the knowledge graph as an external attention memory rather than as textual context to be concatenated into the prompt. Its design has two principal components: KG2KV, which converts knowledge-graph triples into natural-language query-key-value supervision and offline memory embeddings, and HiKVP, or Hierarchical Key-Value Pruning, which restricts inference-time access to a small candidate subset of the external memory. The method follows a lightweight adaptation regime: the backbone model is not retrained in full, and knowledge updates are handled by rebuilding external memory rather than modifying the base LLM weights (Huang et al., 20 Oct 2025).

A concise description of the system components is as follows:

Component	Function	Role in scaling
KG2KV	Converts triples into Q-K-V memories	Improves training data quality
HiKVP	Hierarchical pruning over external memory	Reduces linear memory access to sub-linear
KG-specific heads	Project external memories into attention space	Enables attention-native knowledge grounding

2. KG2KV: transforming triples into attention-native memory

KG2KV is the mechanism that converts a knowledge graph into supervision compatible with transformer attention. For each triple $(h,r,t)$, AtlasKV constructs textual memories by masking one entity and rewriting the relation into an attribute-like noun phrase. In the tail-masked case, the masked tail becomes the value and the head-plus-relation phrase becomes the key; in the head-masked case, the relation is rewritten in reverse noun form so that the masked head becomes the value. The paper presents examples such as “the cause of John founded StockLemon.com” and “the result of John has made profits ...”, reflecting the two masking directions (Huang et al., 20 Oct 2025).

A corresponding query sentence is then constructed by prepending question prefixes such as “What is …”, “Tell me …”, and “Provide details on …”. The resulting Q-K-V tuples serve two purposes. They create a direct alignment between natural-language questions and KG-derived key-value pairs, and they avoid overfitting to a single rigid template. AtlasKV explicitly argues that this is superior to fully synthetic schema-based QKV generation because it inherits the diversity of relations already present in real knowledge graphs (Huang et al., 20 Oct 2025).

The reported quantitative comparison is specific. The Synthetic diversity ratio is 0.003%, whereas the KG2KV diversity ratio is 7.864%. The Synthetic avg token cost is 349.9, whereas the KG2KV avg token cost is 165.7. The intended interpretation is that KG2KV is both more diverse and more token-efficient than the synthetic alternative (Huang et al., 20 Oct 2025).

After text construction, each key and value string is encoded offline into base embeddings $\bm{k}^m$ and $\bm{v}^m$. AtlasKV reports two sentence encoders: all-MiniLM-L6-v2 with $D_E=384$ and text-embedding-3-large with $D_E=3072$. These embeddings are then projected into the LLM attention space by KG-specific heads $\tilde{W}_Q^{(l)}$, $\tilde{W}_K^{(l)}$, and $\tilde{W}_V^{(l)}$, so that the model can attend jointly to sequence memory and external KG memory through the same attention mechanism (Huang et al., 20 Oct 2025).

3. HiKVP: hierarchical pruning and sub-linear external attention

The central systems contribution of AtlasKV is HiKVP, which reduces the cost of external-memory attention from linear in the number of memory entries to sub-linear. The method begins by reducing key dimensions with UMAP and then clustering with Gaussian Mixture Models into a 3-layer hierarchy: a leaf layer containing original keys, an intermediate layer containing pooled clusters of leaf keys, and a root layer containing pooled clusters of intermediate keys. To balance the hierarchy, the cluster size is set to

$M$0

The corresponding layer sizes are described as $M$1, $M$2, and root size roughly $M$3 (Huang et al., 20 Oct 2025).

Inference proceeds in three pruning stages. First, root-layer projected keys are uploaded to GPU and scored against the token query; only top-$M$4 root nodes are kept. Second, the corresponding intermediate keys are uploaded, scored, and pruned to top-$M$5. Third, the corresponding leaf keys are uploaded, scored, and pruned to top-$M$6; only the final selected values are fetched to GPU for the KG branch of attention. AtlasKV reports default pruning settings of $M$7, $M$8, and $M$9 (Huang et al., 20 Oct 2025).

This design is tightly coupled to memory placement. Root-layer projected keys are initially uploaded to GPU; intermediate and leaf keys and values remain in CPU memory. After each pruning stage, previously used higher-level keys are offloaded back to CPU, and only the currently selected candidates are transferred upward. This staging is the practical basis of the claim that the full KG is never resident on GPU at once (Huang et al., 20 Oct 2025).

The method reformulates rectangular attention so that the sequence branch and KG branch can be separated. At inference, only the pruned leaf set participates in the KG contribution. AtlasKV reports the following overall complexity: $(h,r,t)$0 for time and

$(h,r,t)$1

for memory, with

$(h,r,t)$2

This is the formal expression behind the claim that AtlasKV replaces the linear $(h,r,t)$3-dependence of earlier attention-based external-memory methods with a sub-linear $(h,r,t)$4-dependence (Huang et al., 20 Oct 2025).

4. Training regime and adaptation to new knowledge

AtlasKV uses a lightweight adaptation regime in which the only learnable variables are the KG-specific query and projection heads. The backbone is LLaMA-3.1-8B-Instruct. The paper states that $(h,r,t)$5 is initialized from the ordinary attention query head $(h,r,t)$6, while $(h,r,t)$7 and $(h,r,t)$8 are initialized randomly (Huang et al., 20 Oct 2025).

Training data are derived from ATLAS-Wiki, described as a KG with 900M+ nodes and 5.9B edges. The appendix also mentions sampling from ATLAS-Wiki containing 1.492B triples, from which only 20K triples are sampled to construct the QKV training set. This small supervised set is emphasized as part of AtlasKV’s generalization claim (Huang et al., 20 Oct 2025).

The reported training configuration is specific: AdamW, initial learning rate $(h,r,t)$9, cosine decay to $\bm{k}^m$0, 3K iterations, batch size 10, and injection of KGKVs into attention every 3 layers. During training, the KG size increases by 4 every 100 iterations. All experiments are run on a single 48GB GPU in bfloat16 (Huang et al., 20 Oct 2025).

HiKVP is not required during training, because the training-time KGKV sets are relatively small; pruning is an inference scalability mechanism. Knowledge updates therefore follow an external-memory path rather than a retraining path: add new triples, convert them with KG2KV, encode them offline, rebuild or update the hierarchical index, and reuse the trained KG-specific heads. This is the operational meaning of the paper’s claim that AtlasKV adapts to new knowledge without retraining (Huang et al., 20 Oct 2025).

5. Empirical results and evaluation profile

AtlasKV is evaluated on out-of-distribution grounding tasks using Enron, ATLAS-CC-QKV, and ATLAS-Pes2o-QKV, with ACC@1, ACC@5, and a GPT-4o-based relevance score for generated answers. The grounding metric is extracted from the 15th attention layer and averaged over heads. The principal baselines are Zero-shot, ICL, and KBLaM (Huang et al., 20 Oct 2025).

The headline systems result is the memory comparison: AtlasKV requires less than 20GB VRAM for 1B triples, whereas KBLaM requires over 40GB VRAM for even 100K triples. This claim is paired with the sub-linear complexity analysis and with the CPU/GPU staged memory design (Huang et al., 20 Oct 2025).

Grounding results are reported at several candidate-set sizes. On Enron at $\bm{k}^m$1 triples, KBLaM (3e3 steps) obtains ACC@1 = 9.1 and ACC@5 = 20.0, whereas AtlasKV reports 21.8 / 32.7, and AtlasKV w/o HiKVP reports 27.3 / 47.3. On ATLAS-Pes2o-QA at $\bm{k}^m$2 triples, KBLaM (2e4 steps) reports 0.0 / 5.5, AtlasKV reports 16.4 / 49.0, and AtlasKV w/o HiKVP reports 47.3 / 67.2. At $\bm{k}^m$3 triples on the same dataset, the numbers are 25.5 / 52.7 for KBLaM, 87.3 / 92.7 for AtlasKV, and 92.7 / 100.0 for AtlasKV w/o HiKVP (Huang et al., 20 Oct 2025).

On ATLAS-CC-QA at $\bm{k}^m$4 triples, the reported values are 3.6 / 10.9 for KBLaM, 40.0 / 54.5 for AtlasKV, and 61.8 / 81.8 for AtlasKV w/o HiKVP. At $\bm{k}^m$5 triples, they are 23.6 / 56.4, 89.1 / 90.9, and 96.4 / 100.0, respectively. These comparisons establish two consistent facts: KG2KV substantially improves out-of-distribution grounding over the baseline, and HiKVP introduces some accuracy loss relative to the unpruned variant, although the pruned system still remains substantially stronger than KBLaM in the reported settings (Huang et al., 20 Oct 2025).

The evaluation also reports that AtlasKV reaches strong performance with only 20K KGKV samples and 3K training steps, whereas the paper contrasts this with 20K steps reported for KBLaM. In generation relevance scoring, ICL can produce answers often above 0.9, but the paper characterizes it as memory- and latency-intensive; it states that with more than 100 triples, ICL can require over 48GB VRAM (Huang et al., 20 Oct 2025).

6. Interpretation, scope boundaries, and limitations

AtlasKV belongs to the class of attention-based external memory systems for knowledge graphs, not to the class of KV-cache compression systems or storage-engine key-value stores. This distinction matters because the abbreviation “KV” appears across several unrelated research areas. Retrieval-based KV-cache reduction methods such as A$\bm{k}^m$6ATS target long-context inference by offloading or retrieving past token states (He et al., 18 Feb 2025). Runtime KV-cache management methods such as ARKV and AnTKV target eviction, quantization, or tri-state memory control for transformer caches (Lei et al., 19 Feb 2026, Li et al., 24 Jun 2025). Persistent or scalable LLM KV-cache systems such as SGLang-LSM and PiKV focus on storage layout, cache-serving, or MoE-specific cache placement (Yu et al., 20 Nov 2025, Liu et al., 2 Aug 2025). ArceKV, by contrast, is an LSM-based key-value store for dynamic compaction control under changing workloads (Liu et al., 5 Aug 2025). AtlasKV addresses a different problem: grounding LLM generation in a very large external knowledge graph through attention-native key-value memory (Huang et al., 20 Oct 2025).

The paper also makes its trade-offs explicit. AtlasKV depends on the quality of the underlying textual KG; KG extraction errors are out of scope and would propagate into the memory. HiKVP depends on UMAP + GMM hierarchy quality, so poor clustering may reduce retrieval accuracy. AtlasKV w/o HiKVP is usually more accurate than AtlasKV with pruning, which shows that sub-linear scaling is not free. The evaluation focuses primarily on grounding-style QA and does not extensively test complex multi-hop reasoning over huge KGs. Finally, the training objective is not introduced as a novel loss in the main text; the paper states that the training settings follow the earlier KBLaM style, so reproducing exact behavior may require those inherited details (Huang et al., 20 Oct 2025).

Within those limits, AtlasKV’s contribution is specific and technically narrow: it makes external knowledge-graph augmentation attention-native, scales that mechanism through hierarchical pruning, and separates knowledge updates from backbone retraining. The resulting system occupies a distinct position between retriever-based RAG and conventional fine-tuning, with the external KG retained as an explicit memory rather than absorbed into model weights or injected as long textual context (Huang et al., 20 Oct 2025).