GFM-RAG: 用于大语言模型检索增强的图基础模型

原标题：GFM-RAG: Graph Foundation Model for Retrieval Augmented Generation

论文链接：https://arxiv.org/pdf/2502.01113

项目主页：https://rmanluo.github.io/gfm-rag/

TL;DR: GFM-RAG是第一个适用于大语言模型检索增强的图基础模型。

摘要

检索增强生成（RAG）在将知识集成到大语言模型（LLM）方面已被证明是有效的。然而，传统的RAG难以捕捉知识片段之间的复杂关系，这限制了它们在需要集成多来源知识的复杂推理中的表现。最近，图增强检索增强生成（GraphRAG）构建图结构来显式地对这些关系建模，从而实现更有效和高效的检索器。然而，它的性能仍然受到图结构内的噪声和不完整性的阻碍。为了解决这个问题，我们引入了一种用于检索增强生成的图基础模型（GFM）。GFM-RAG由图神经网络作为基本架构，该网络可以对图结构进行推理，以捕获复杂的查询-知识关系。具有8M参数的GFM在包含60个知识图、超过14M个三元组和700k个文档的大规模数据集上经历了两个阶段的训练过程。这为GFM-RAG带来了更强的性能和通用性，使其成为第一个适用于不可见数据集的图基础模型，无需任何微调即可进行检索。

主要贡献

1. 文章提出了一种适用于大语言模型检索增强的图基础模型（GFM-RAG）,该模型由一种与查询依赖GNN提供支持，可以在单个步骤内实现高效的多步检索。
2. 文章训练了一个参数量为8M的模型，标志着第一个可以直接应用于各种未知数据集的检索增强生成的图基础模型的诞生。
3. 我们在三个多跳QA数据集和七个特定领域RAG数据集上评估了GFM-RAG，它在所有数据集上实现了最先进的性能，证明了它的有效、高效、通用性，以及成为进一步增强生成研究的基础模型的潜力。

模型

GFM-RAG由三个主要部分组成：A. 构建知识图谱索引，从文档语料库中构建知识图谱索引；B. 图基础模型检索器（GFM retriever），在大规模数据集上预先训练，可根据任何用户查询和知识图谱索引检索文档；C. 文档排序和答案生成，对检索到的文档进行排序并生成最终答案。具体的GFM-RAG的框架以及图基础模型的训练过程如图1所示。

图1 GFM-RAG的框架以及图基础模型的训练过程

KG-index构建

给定一组文档$D$，首先从文档中提取实体E和关系$R$以形成三元组$T$。然后，构造到文档的实体倒排索引$M\in{0,1}^{|E|\times|D|}$来记录每个文档中提到的实体。为了更好地捕捉知识之间的联系，需要进一步进行实体解析(entity resolution) ，在语义相似的实体之间添加额外的边$T^+$。

图基础模型（GFM）检索器

查询依赖GNNs在捕获查询特定信息和对未见图的可泛化性方面表现出色，这对于GFM检索器至关重要，它允许GFM检索器根据用户查询动态调整消息传递过程，并在图中找到最相关的信息。可以表述为：$H_q^L=GNN_q(q,G,H^0)$，其中$H^0\in\mathbb{R}^{|E|\times d}$是初始实体特征，$G$是知识图谱，$q$是查询，$H_q^L$表示在查询$q$条件下经过$L$层查询以来的消息传递后更新的实体表示。

查询初始化

给定一个查询$q$，先用文本嵌入模型将其编码成一个查询向量：$\mathbf{q}=\text{SentenceEmb}(q),\mathbf{q}\in\mathbb{R}^{d}$。然后用其初始化实体表示：

$$H^0 = \begin{cases} \mathbf{q}, & e\in E_q, \\ \mathbf{0}, & \text{otherwise}. \end{cases}$$

其中$E_q$是查询$q$中的提到实体集合。

查询依赖信息传播

查询依赖的消息传递将从问题实体向知识图中的其他实体传播信息，以捕捉它们与查询的相关性。

$$三元组信息传递：m_e^{l+1} = \text{Msg}(h_e^l,g^{l+1}(h_r^l),h_{e'}^l), (e,r,e')\in G,\\ 实体信息聚合：h_e^{l+1} = \text{Update}(h_e^l,\text{Agg}(\{m_{e'}^{l+1} | e'\in N_r(e),r\in R\})).$$

经过$L$层消息传递后，最后的MLP层结合sigmoid函数将实体嵌入映射到与查询的相关性得分。

$$P_q = \sigma(\text{MLP}(H_q^L)),~P_q\in\mathbb{R}^{|E|\times 1}$$

训练目标

GFM检索器的训练目标是最大化相关实体与查询的似然，可以通过最小化二进制交叉熵（binarycross-entropy, BCE）损失来优化：

$$L_{\text{BCE}} = -\frac{1}{|A_q|}\sum_{e\in A_q} \log P_q(e) - \frac{1}{|E^{\texttt{-}}|}\sum_{|E^{\texttt{-}}|} \log (1-P_q(e)),$$

其中$A_q$表示与查询q相关的目标实体集合，$E^{\texttt{-}}$是负样本实体集合。

然而，由于目标实体的稀疏性，BCE损失可能存在梯度消失问题。为了解决这个问题，我们进一步引入了排序损失来最大化正负实体之间的差值：

$$L_{\text{RANK}} = - \frac{1}{|A_q|}\sum_{e\in A_q} \frac{P_q(e)}{\sum_{e'\in E^{\texttt{-}}} P_q(e')}$$

最终的训练目标是BCE损失和排名损失的加权组合：

$$L = \alpha L_{\text{BCE}} + (1-\alpha) L_{\text{RANK}}$$

训练任务

GFM-RAG的训练分为两个阶段：1）自监督预训练阶段，通过在大规模知识图谱上采样的三元组生成人造查询进行训练，增强其图推理的能力；2）微调阶段，在真实标注的数据集上进行微调，以提高模型在下游任务的性能和泛化能力。

文档排序和答案生成

给定GFM检索器预测的实体相关性得分$P_q\in\mathbb{R}^{|E|\times 1}$，我们首先检索相关度得分最高的top-T实体$E_q^{T}$如下：

$$E_q^{T} = \arg\text{top-}T(P_q),~E_q^{T}=\{e_1,\ldots,e_T\}$$

这些检索到的实体被文档排序器用来获得最终的文档。为了减少热门实体的影响，我们通过它们作为文档倒排索引M\in{0,1}^{|E|\times|D|}中提到的实体的频率的倒数来对实体进行加权，并通过将权重相加来计算最终的文档相关性分数：

$$F_e = \begin{cases} \frac{1}{\sum_{d\in D} M[e,d]}, & e\in E_q^{T}, \\ 0, & \text{otherwise}, \end{cases} \\ P_d = M^{\top}F_e,~P_d\in\mathbb{R}^{|D|\times 1}$$

根据文档相关性得分$P_d$检索top-K文档，并以检索增强生成方式馈送到LLM的上下文中，以生成最终答案：

$$D^K = \arg\text{top-}K(P_d),~D^K=\{D_1,\ldots,D_K\}, \\ a = \text{LLM}(q,D^K).$$

实验设置

数据集

首先评估了GFM-RAG在三种广泛使用的多跳QA数据集上的有效性，包括Hot-potQA、MuSiQue、2WikiMultiHopQA (2Wiki)。

此外，文章还评估了GFM-RAG在来自三个领域的七个RAG数据集上的性能，以此来证明GFM-RAG作为基础模型的通用性。(1) biomedical: PubMedQ; (2) customer support: DelucionQ, TechQ, ExpertQA, EManual; (3) general knowledge: MS Marco, HAGRID。

表1 用于测试的数据集与知识图谱索引的统计

对比方法

在实验中，文章与三个类别下几个广泛使用的检索方法进行了比较：

（1）单步朴素方法：BM2、Contriever、GTR、ColBERTv2、RAPTOR、Proposition；

（2）图增强方法：LightRAG、HippoRAG；

（3）多步骤方法：IRCoT。

评价指标

对于检索性能，文章使用recall@2 (R@2)和recall@5 (R@5)作为评估指标。

对于最终的QA性能，文章使用之前研究中的EM分数和F1分数。

实现细节

GFM检索器由6个与查询相关的消息传递层实现，隐藏维度设置为512。采用预训练的all-mpnet-v2作为句子嵌入模型。GFM检索器的总参数为8M，在8台NVIDIA A100(80G)上进行训练，batch size为4，learning rate为5e-4, loss weight α= 0.3。训练数据包含60公斤数和超过14M个三元组，这些三元组是从训练集中提取的700k个文档中构建的。

表2 用于训练的查询文档与知识图谱的统计

实验结果

检索性能

文章首先对三个多跳QA数据集的基线进行了GFM-RAG检索性能的评估。如表3所示，GFM-RAG在所有数据集上都取得了最好的性能，在HotpotQA、MuSiQue和2Wiki上R@2分别比SOTA IRCoT + HippoRAG高出16.8%、8.3%和19.8%。实验结果证明了GFM-RAG在多跳检索中的有效性。从结果可以看出，图增强的HippoRAG优于单纯的单步检索器（如BM25、RAPTOR），这凸显了图结构在多跳检索中的重要性。

表3 检索性能比较

问答性能

第二，文章评估了GFM-RAG的QA性能，这直接影响了用户的检索质量。文章选用GPT-4o-mini作为大语言模型，并使用检索到的前5个文档生成答案。从表4所示的结果来看，单步GFM-RAG已经在所有其他基线上取得了最先进的性能。同时，我们还将GFM-RAG与IRCoT相结合，进行多步检索和推理，进一步提高了EM在三个数据集上的性能，分别提高了8.5%、21.2%和3.9%。结果表明，GFM-RAG与任意多步框架在多跳推理任务中的有效性和良好的兼容性。

表4 问答性能对比

效率分析

GFM-RAG在单步进行多步推理方面取得了很高的效率。如表5所示，朴素的单步推理方法效率最高，但性能并不令人满意。诚然，多步骤框架IRCoT可以提高性能，但由于要使用LLMs进行迭代检索和推理，其计算成本较高。相比之下，GFM-RAG在单步GNN推理中进行多跳推理，比单步方法更有效，比多步方法更高效。

表5 检索效率和性能比较

消融实验

文章进行了消融研究，来考察GFM-RAG中不同组成部分的有效性，包括不同的句子嵌入模型（表6）、预训练策略（表7）和损失加权策略（表8）。结果表明，GFM-RAG对不同的句子嵌入模型并不敏感，预训练策略和损失加权策略对GFM-RAG的性能都至关重要。

表6 GFM-RAG中使用的不同句子嵌入模型的比较

表7 GFM-RAG中预训练和微调的有效性

表8 GFM-RAG中损失加权策略的影响

模型的可泛化性

为了证明GFM-RAG作为基础模型的通用性，文章在不做任何微调的情况下，在七个特定领域的RAG数据集上测试了GFM-RAG的性能（R@5）。首先，根据每个数据集中的文档建立知识图谱索引。接着，在给定查询的情况下，我们使用预先训练好的GFM Retriever，在相应KG-index的帮助下检索出Top-K文档。如图2所示，GFM-RAG在所有数据集上都取得了最佳性能，平均比SOTA HippoRAG高出18.9%，证明了GFM-RAG作为基础模型的通用性，无需任何微调即可直接应用于各种未见数据集。此外，表9中的结果表明，GFM-RAG具有很强的可移植性，可在特定领域数据集上进行微调，从而进一步提高性能。

图2 模型性能和通用性比较。

表9 模型性能（R@5）和可迁移性比较

模型性能缩放定律

我们进一步研究了GFM-RAG的神经缩放定律，它量化了模型性能如何随训练数据规模和模型参数大小而增长。如图3所示，GFM-RAG的性能（MRR：z）与训练数据（x）和模型大小（y）呈良好的比例关系：$z \propto 0.24x^{0.05} + 0.11y^{0.03}$，可以用幂律缩放定律来拟合。结果表明了GFM-RAG作为基础模型的可扩展性和进一步提升的潜力。

图3 GFM-RAG的模型性能缩放定律。

模型的可解释性

GFM-RAG由于多层的GNN结构，使其具有很好的多跳推理能力。我们可以通过展现GNN捕捉到的重要路径信息，来解释其推理过程。路径对最终预测的重要性可以通过最终预测分对每一层（跳）的三元组的偏导数来量化。

如图4所示，GFM-RAG在多跳推理中捕捉到了重要的路径信息，这有助于解释模型的推理流程。

$$s_1,s_2,\ldots,s_L=\arg\mathop{\text{top-}}k\frac{\partial p_e(q)}{\partial s_l}$$

表10 GFM捕获的用于多跳推理的重要路径。

结论

这篇文章介绍了首个用于检索增强生成的图基础模型。通过利用知识图索引，GFM-RAG对知识与文档之间的复杂关系进行了明确建模，从而促进了更有效、更高效的检索过程。依赖于在大规模数据集上进行预训练的GNN，GFM-RAG可以有效地对图结构执行多跳推理，从而在一个步骤中找到相关的知识边。在三个基准数据集和七个特定领域数据集上进行的广泛实验表明，GFM-RAG在有效性、效率和通用性方面明显优于最先进的方法。同时，GFM-RAG符合扩展规律，这也表明它有潜力扩展到更大的数据集。