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

发表时间: 2026-04 · arXiv:2510.17934 (ICLR 2026)

原文: https://arxiv.org/abs/2510.17934

Haoyu Huang, Hong Ting Tsang, Jiaxin Bai, Xi Peng, Gong Zhang, Yangqiu Song / The Hong Kong University of Science and Technology, Theory Lab, Huawei

速读

一句话结论 本文提出了一种名为 AtlasKV 的参数化知识注入方法，能在单张显卡不到 20GB 显存的极低消耗下，将 10亿 规模的知识图谱三元组直接注入大语言模型的注意力层，无需外部检索器或长上下文即可实现高效准确的知识增强。

要解决什么问题 现有大模型知识增强方案主要卡在推理延迟与扩展性瓶颈上。以 RAG 为代表的非参数化方法强依赖外部检索器，且必须将检索文本拼接在提示词中，长上下文前置在面对超大规模知识库时会带来极高的搜索成本与推理延迟。传统参数化方法在引入新知识时又必须重新训练模型。近期以 KBLaM 为代表的新型参数化方法试图通过将外部知识编码为键值对（Key-Value）并直接注入模型的自注意力层来免除长上下文依赖，但面临两个致命卡点：一是训练数据质量差，用固定模板合成的数据缺乏多样性，导致模型在面对分布外（OOD）提问时泛化能力极差；二是扩展性极差，其计算和显存开销随知识库大小 $M$ 呈线性增长，即 $\mathcal{O}(M)$，当知识图谱规模达到 10万 级别时就会耗尽 40GB 显存，无法扩展到真实世界动辄 10亿 规模的场景中。

怎么做的 AtlasKV 的核心思路是将知识图谱三元组转化为高质量的查询-键-值（Q-K-V）格式，并通过层次化剪枝算法将其以亚线性复杂度动态注入大模型的注意力机制中。该方法由两个关键部件构成。首先是数据层面的 KG2KV 模块，它巧妙利用了三元组 $(h, r, t)$ 的天然结构：将其中一个实体掩码作为需要预测的 Value，利用大模型将关系 $r$ 改写为名词短语作为 Key 的属性描述，并加上多样化的提问前缀构成 Query。这种转换极大丰富了查询属性的语义分布，解决了模型泛化性差的问题。其次是算法层面的层次化键值剪枝（HiKVP）模块，专门解决扩展性卡点。它事先用 UMAP 降维和高斯混合模型（GMM）将所有知识的 Key 聚类成根节点、中间层和叶子节点的三层树状结构。推理时，模型仅将根节点的 Key 驻留在 GPU 显存中，计算当前 Query 与根节点的注意力得分并选出 Top-K 个最相关分支后，再从 CPU 内存加载对应的下一层 Key 到 GPU 中，逐层剪枝直到叶子节点。最终，大模型在每一层的注意力输出变为知识图谱部分和原始序列部分的动态加权求和：

$$ \tilde{\mathbf{y}}_n^{(l)} = \lambda_{kg} \cdot \mathrm{Softmax}(\mathrm{logits}_{kg}) \cdot \bar{\tilde{\mathbf{v}}}^{(l)} + \lambda_{seq} \cdot \mathrm{Softmax}(\mathrm{logits}_{seq}) \cdot \mathbf{v}^{(l)} $$

其中 $\lambda_{kg}$ 和 $\lambda_{seq}$ 是根据两者 logits 动态计算的归一化权重。通过分层剪枝与按需加载，HiKVP 成功将时间复杂度降至 $\mathcal{O}((C_t \sqrt[3]{M} + N) \cdot N \cdot D)$，显存复杂度降至 $\mathcal{O}((C_m \sqrt[3]{M} + N) \cdot (N + D))$，彻底绕开了线性增长的显存墙。

效果如何 实验基于 LLaMA3.1-8B-Instruct 模型搭建，使用 all-MiniLM-L6-v2 和 text-embedding-3-large 作为句子编码器，在单张 48GB 显卡上仅用 2万 条三元组数据训练了 3000步。对比基线主要包括代表非参数化检索路线的 ICL（上下文学习，即标准 RAG 范式）和代表参数化注意力注入路线的 KBLaM。在极具挑战的 ATLAS-Pes2o-QKV 分布外测试集（1万 条三元组规模）下，AtlasKV 的 Top-1 知识命中率达到了 49.0%，而 KBLaM 因泛化失效命中率仅为 0.0%。扩展性方面，AtlasKV 注入 10亿 条三元组仅需不到 20GB 显存，而 KBLaM 在 10万 条时就已发生内存溢出。此外，GPT-4o 的量化评分显示，AtlasKV 生成回答的准确度远超 KBLaM，且逼近直接拼接全部知识的 ICL 上限。作者也坦诚了该方法的局限与代价：KG2KV 模块将图谱展平为独立事实，破坏了原有的图结构，阻碍了模型的多跳推理能力；同时 HiKVP 算法在各层级间频繁的 CPU 与 GPU 数据搬运会带来一定的硬件 I/O 延迟，需要在准确率与扩展性之间做出权衡。

主要贡献

检索增强生成（RAG）在利用外部知识增强大型语言模型（LLMs）方面取得了一定成功。然而，作为LLMs的非参数化知识集成范式，RAG方法严重依赖于外部检索模块和检索到的文本上下文先验。特别是在超大规模知识增强场景下，由于昂贵的搜索开销和极长的相关上下文，它们会引入巨大的推理延迟。与此相对，传统的参数化方法在适应新知识时需要重新训练模型，成本高昂。KBLaM等新型参数化方法虽然解决了免训练适应的问题，但面临缺乏高质量训练数据和极大规模下扩展性差的两大挑战。

为了解决上述问题，本文提出了AtlasKV，这是一种可扩展、有效且通用的参数化知识集成方法，能够使用极少的GPU内存（例如不到20GB VRAM）将十亿级知识图谱（KGs，例如10亿个三元组）增强到LLMs中。在AtlasKV中，引入了KG2KV和HiKVP两个核心模块，以亚线性时间和内存复杂度将KG三元组大规模集成到LLMs中。该方法利用LLMs固有的注意力机制保持了强大的知识基础和泛化性能，在适应新知识时不需要外部检索器、长上下文先验或重新训练。

背景知识与设计原则

知识图谱的定义
正如现有大多数基于图的RAG系统所做的那样，本文使用文本三元组 $(h, r, t)$ 作为知识图谱（KGs）的基本知识单元，这些三元组可以通过任何现有的KG提取方法从非结构化文本中提取。随后，KGs可以定义为 $\mathcal{G} = \{(h, r, t) | h, t \in \mathcal{E}, r \in \mathcal{R}\}$，其中 $\mathcal{E}$ 是实体集，$\mathcal{R}$ 是关系集。需要注意的是，$h, t$ 可以是命名实体，也可以是其他类型的实体，如概念、事件等。在本文中，$\mathcal{G}$ 将作为LLMs的外部事实知识被集成，以用于事实溯源和回答问题，而从文档中提取KGs的过程并非本文的重点。

自注意力层机制
在每个注意力层中，输入一个长度为 $N$ 的查询 $\boldsymbol{x} \in \mathbb{R}^{N \times 1}$，其嵌入向量可以表示为 $\mathbf{\bar{x}}^{(l)} \in \mathbb{R}^{N \times \bar{D}}$。$D$ 是注意力层的嵌入维度，$l \in \{1, .., L\}$，其中 $L$ 是注意力层的数量。同时存在三个注意力头 $\mathbf{W}_Q^{(l)}, \mathbf{W}_K^{(l)}, \mathbf{W}_V^{(l)} \in \mathbb{R}^{D \times D}$，旨在将每个输入token投影为Q-K-V嵌入 $\mathbf{q}^{(l)}, \mathbf{k}^{(l)}, \mathbf{v}^{(l)} \in \mathbb{R}^{N \times D}$。接着，第 $l$ 层第 $n$ 个token的输出计算如下：

$$ \mathbf{y}_n^{(l)} = \frac{\sum_{i=1}^n \exp(\langle \mathbf{q}_n^{(l)}, \mathbf{k}_i^{(l)} \rangle / \sqrt{D}) \mathbf{v}_i^{(l)}}{\sum_{i=1}^n \exp(\langle \mathbf{q}_n^{(l)}, \mathbf{k}_i^{(l)} \rangle / \sqrt{D})} $$

其中 $\langle \cdot, \cdot \rangle$ 表示两个向量的内积。这种标准的自注意力实现具有 $\mathcal{O}(N^2 \cdot D)$ 的时间复杂度和 $\mathcal{O}(N \cdot (N + D))$ 的内存复杂度，随着输入长度 $N$ 的增加，会导致显著的计算开销和时间延迟。

KBLaM中的矩形注意力
KBLaM中的知识库（KB）是一个键值对集合，在第 $l$ 个注意力层表示为 $\mathcal{M} = \{(\pmb{k}^{(l)m}, \pmb{v}^{(l)m})\}_{m=1}^M$，其中 $\pmb{k}^{(l)m}, \pmb{v}^{(l)m} \in \mathbb{R}^{M \times D_E}$ 是第 $m$ 个键值对的基础嵌入向量。$M$ 是KB的大小，$D_E$ 是句子编码器的输出维度。随后，第 $l$ 个注意力层和第 $n$ 个token的知识增强输出计算如下：

$$ \tilde{\mathbf{y}}_n^{(l)} = \frac{\sum_{m=1}^M \exp(\langle \tilde{\mathbf{q}}_n^{(l)}, \tilde{\mathbf{k}}^{(l)m} \rangle / \sqrt{D}) \tilde{\mathbf{v}}^{(l)m} + \sum_{i=1}^n \exp(\langle \mathbf{q}_n^{(l)}, \mathbf{k}^{(l)i} \rangle / \sqrt{D}) \mathbf{v}^{(l)i}}{\sum_{m=1}^M \exp(\langle \tilde{\mathbf{q}}_n^{(l)}, \tilde{\mathbf{k}}^{(l)m} \rangle / \sqrt{D}) + \sum_{i=1}^n \exp(\langle \mathbf{q}_n^{(l)}, \mathbf{k}^{(l)i} \rangle / \sqrt{D})} $$

其中 $\tilde{\mathbf{q}}_n^{(l)} = \tilde{\mathbf{W}}_Q^{(l)} \mathbf{x}_n^{(l)}$, $\tilde{\mathbf{k}}^{(l)m} = \tilde{\mathbf{W}}_K^{(l)} \mathbf{k}^{(l)m}$, $\tilde{\mathbf{v}}^{(l)m} = \tilde{\mathbf{W}}_V^{(l)} \mathbf{v}^{(l)m}$ 表示KB部分投影后的Q-K-V嵌入向量。$\tilde{\mathbf{W}}_Q^{(l)} \in \mathbb{R}^{D \times D}, \tilde{\mathbf{W}}_K^{(l)}, \tilde{\mathbf{W}}_V^{(l)} \in \mathbb{R}^{D \times D_E}$ 是KB的Q-K-V向量的特定投影头。这种矩形注意力具有 $\mathcal{O}((M + N) \cdot N \cdot D)$ 的时间复杂度和 $\mathcal{O}((M + N) \cdot (N + D))$ 的内存复杂度。其计算开销随 $M$ 线性增长，比标准自注意力更高效。然而，当KB大小 $M$ 剧烈扩展时，线性增长的时间和内存复杂度仍然是一个关键问题。AtlasKV通过将时间复杂度降低至 $\mathcal{O}((C_t \sqrt[3]{M} + N) \cdot N \cdot D)$ 和内存复杂度降低至 $\mathcal{O}((C_m \sqrt[3]{M} + N) \cdot (N + D))$ 进一步提高了KG增强LLM的可扩展性，其中 $C_t$ 和 $C_m$ 是远小于 $M$ 的常数。

方法细节

AtlasKV架构概览
将包含数十亿三元组的通用KGs等超大且复杂的外部知识增强到LLMs中，通常面临泛化能力低下以及难以承受的计算和内存开销问题。为了克服这些根本性挑战，本文提出了AtlasKV，这是一种通过两项关键创新将海量KGs集成到LLMs中的可扩展、有效且通用的方法：(1) KG2KV，一种新的KG集成范式，它将KG三元组自然转换为Q-K-V数据，使LLMs能够同时实现增强的泛化性能和高效的知识集成；(2) HiKVP，一种分层键值剪枝算法，在推理期间显著降低了计算和内存开销，同时保持了较高的知识基础准确性。

从KG三元组到Q-K-V数据的转换 (KG2KV)
基于观察到KG中的每个三元组都可以自然地分解为Q-K-V字符串【文献：Facts as experts: Adaptable and interpretable neural memory over symbolic knowledge, 2020, arXiv】，这与LLMs中自注意力网络的Q-K-V向量结构非常相似，本文引入了KGKV的概念，并采用了KG2KV管道将每个KG三元组转换为Q-K-V字符串及其对应的句子嵌入向量。

KG2KV的处理流程
对于给定的KG三元组 $(h, r, t)$，首先掩码其头部 $h$ 或尾部实体 $t$（可以是命名实体、事件或概念实体）。此时被掩码的实体即为该三元组中所需的“值（value）”。接着，根据掩码的位置，通过LLMs将关系 $r$ 重写为名词词汇，该名词可被视为未掩码实体的属性。例如，如果掩码三元组中的尾部实体，关系“because”可以直接通过LLMs重写为其名词形式“cause”。这个尾部掩码三元组的键（key）字符串可以表示为“the cause of John founded http://StockLemon.com”。如果掩码头部实体，则需要将关系重写为其反向名词词汇“result”。这个头部掩码三元组的键字符串可以表示为“the result of John has made profits ...”。在这个KG2KV管道中，被掩码的实体作为value数据，而未掩码的实体和关系作为key数据，从而构成完整的KGKV数据。需要注意的是，对于训练数据，通常选择命名实体作为掩码的key，并选择事件实体和关系作为value。随后通过句子编码器，KGKVs可以被压缩并编码为句子嵌入 $\mathbf{k}^m, \mathbf{v}^m$，以集成到LLMs的注意力层中。

训练查询数据的构建
同时还需要每个三元组的查询（query）句子作为训练数据。查询字符串可以通过在键字符串前添加各种提问前缀来获得。例如，提问前缀可以是“What is ...”，“Tell me ...”，或“Provide details on ..”。这种设计可以确保模型不会对特定的提问方式产生过拟合。

KG2KV的多样性与成本优势
与使用有限的人工定义模式直接从文档中合成Q-K-V数据相比，KG中海量的关系保证了使用KG2KV方法构建的训练数据中查询属性的多样性。此外，由于KG2KV中需要输入LLMs的字符串仅为掩码位置和关系，因此KG2KV比直接合成方法需要更少的token开销。KG2KV具有7.864%的显著更高的多样性比率（唯一查询属性的数量除以三元组的总数），以及低于合成方法的平均token成本（165.7）。

KGKVs的层次化聚类
受多项通过在文本块上组织层次化知识结构取得成功的工作的启发，本文采用层次聚类将KGKVs的键聚类成层次结构。这种设计旨在分担推理期间每一层层次化知识键的计算和内存负担。具体而言，首先采用均匀流形近似和投影（UMAP）对知识键进行降维，然后采用高斯混合模型（GMMs）将知识键聚类成层次结构。较高层中的每个键向量是较低层中键向量的池化。在AtlasKV中，将层数设置为3，这也可以根据实际情况设置得更大。选择3层是因为这是包含AtlasKV中所需所有定义的最小层数。为了平均分担计算和内存负担，将每层中簇的大小设置为相同，即 $S = \lceil \sqrt[3]{M} \rceil$。由此可以得到三层知识键的基础嵌入 $\pmb{k}_L^m \in \mathbb{R}^{M_L \times D_E}$, $\pmb{k}_I^m \in \mathbb{R}^{M_I \times D_E}$, $\pmb{k}_R^m \in \mathbb{R}^{M_R \times D_E}$，其中 $M_L = M, M_I = [M^{\frac{2}{3}}], M_R = \lceil \lceil M^{\frac{2}{3}} \rceil M^{-\frac{1}{3}} \rceil$。$D_E$ 是句子编码器的嵌入维度。

等效注意力机制的知识增强
在AtlasKV的微调过程中，不需要修剪KGKV对。因为AtlasKV具有泛化能力，在微调过程中不需要如此大规模的KGKVs。此处使用等效注意力方法来替代KBLaM中的矩形注意力进行知识增强。在第 $l$ 个注意力层和第 $n$ 个token处，计算公式为：

$$ \tilde{\mathbf{y}}_n^{(l)} = \lambda_{kg} \cdot \mathrm{Softmax}(\mathrm{logits}_{kg_L}) \cdot \tilde{\mathbf{v}}^{(l)} + \lambda_{seq} \cdot \mathrm{Softmax}(\mathrm{logits}_{seq}) \cdot \mathbf{v}^{(l)} $$

其中 $\mathrm{logits}_{kg_L}$ 是注意力输出的KG部分，$\mathrm{logits}_{seq}$ 是注意力输出的序列部分。对于这两个softmax结果的权重 $\lambda_{kg}$ 和 $\lambda_{seq}$，有：

并且有： 其中 $\tilde{\mathbf{k}}_L^{(l)i} = \tilde{\mathbf{W}}_K^{(l)} \mathbf{k}_L^i$ 且 $\mathbf{k}^{(l)i} = \mathbf{W}_K^{(l)} \mathbf{k}^i$。唯一的可学习变量 $\theta$ 是KG特定的查询头 $\tilde{\mathbf{W}}_Q$ 和KG投影头 $\tilde{\mathbf{W}}_K, \tilde{\mathbf{W}}_V$。然后使用LLMs原始的自回归训练目标优化 $\theta$：