← 返回 2026-07-06

图检索增强生成中的自适应掩码图嵌入方法 AGE: Adaptive-masking for Graph Embedding in Graph Retrieval-Augmented Generation

Bao Long Nguyen Huu, Atsushi Hashimoto 📅 2026-06-30 👍 6 2026-07-13 08:37

Graph Embedding GraphRAG Joint-Embedding Predictive Architecture Knowledge Graph QA Self-Supervised Learning

通过自适应节点采样和JEPA改进GraphRAG图嵌入表示

前置知识

GraphRAG

Graph Retrieval-Augmented Generation是检索增强生成的扩展，通过引用图结构数据作为外部知识来支持大语言模型。与基于文本的RAG不同，GraphRAG使用节点（实体）和边（关系）来表示知识，能够清晰地呈现复杂关系，提供数据整合、提高搜索准确性、增强推理能力以及减少幻觉等优势。GraphRAG将检索到的子图转换为文本表示或图嵌入后输入到LLM中生成答案。

本文的核心就是改进GraphRAG中的图嵌入模块，理解GraphRAG的基本框架和工作流程是读懂本文的前提。

Joint-Embedding Predictive Architecture (JEPA)

JEPA是一种自监督学习架构，由编码器Eθ、预测器Pφ和目标编码器Tθ组成。它通过stop-gradient操作防止目标编码器的表示坍塌。与传统的生成式架构（如MAE）需要重建像素或token级别的细节不同，JEPA消除了对不必要输入特征细节的重建，专注于学习更抽象的语义表示。这种方法在认知神经科学中有理论依据，模仿人类自上而下的模式推理能力。

JEPA是本文方法的核心创新之一，AGE首次在GraphRAG框架中应用JEPA来改进图表示学习，理解JEPA的设计动机和原理对理解本文贡献至关重要。

非参数检索器

非参数检索器是指在检索过程中不使用可训练参数的检索方法，例如基于k近邻（k-NN）的静态检索。与基于LLM的可训练检索器相比，非参数检索器计算开销小、成本低，但可能包含冗余节点或遗漏关键节点，导致缺乏显式的结构约束。本文使用SentenceBert对知识数据进行索引，结合静态k近邻检索和Prize-Collecting Steiner Tree子图构建方法。

本文选择非参数检索器作为基础，以保持实用性，然后通过改进图嵌入模块来提升性能。理解这一选择有助于理解本文的设计理念和实验设置。

关键节点

在知识图谱中，某些节点包含主导性的上下文信息，对于维护图的完整性和语义表示至关重要。这些节点被称为关键节点。与自然语言文本不同，图是简洁的逻辑结构，冗余度最小，因此关键节点很难从周围节点预测。掩码这些关键节点会导致自监督学习过程的效率降低。本文通过强化学习训练的节点采样器来自动识别和选择关键节点。

关键节点的概念是本文自适应掩码策略的核心动机，理解为什么需要区分关键节点和辅助节点，以及如何利用这种区分来改进图嵌入，是理解本文方法的关键。

研究动机

现有的GraphRAG方法在处理冻结LLM时面临图表示与文本表示的潜在特征不对齐问题。具体场景包括：当使用非参数检索器时，检索到的子图可能包含冗余节点或遗漏关键节点，导致缺乏显式的结构约束；传统的随机节点掩码方法在处理非判别性节点时效果不佳，因为这些节点很难从周围节点预测，导致表示质量下降；图嵌入需要嵌入元素之间的关系而非其各自内容，但现有的图编码器设计未能很好地对齐到LLM的文本编码器空间。实验数据显示，G-Retriever在ExplaGraphs数据集上的准确率仅为0.5595（Llama3.2-1B），说明还有很大的改进空间。

本文的目标是本文提出自适应掩码图嵌入（AGE）方法，旨在解决GraphRAG中图嵌入与LLM文本编码器的潜在特征不对齐问题。具体目标包括：设计一个模仿文本嵌入编码器的架构，解决潜在特征不对齐问题；引入基于强化学习的节点采样器，用自适应方法替代传统的随机节点掩码；应用JEPA自监督学习架构，消除不必要的细节重建，专注于学习更抽象的语义表示；在保持非参数检索器和冻结LLM的低计算成本前提下，显著提升GraphQA任务的性能。

与已有工作不同的是，本文的独特切入角度在于将认知神经科学的洞察应用到GraphRAG中。从认知科学出发，人类大脑的颞叶区域参与自下而上的联想学习，能够选择关键知识并将其与相关数据链接。本文试图在GraphRAG中重现这一功能，提出了新颖的节点采样器模块。与现有的两阶段训练方法（第一阶段SSL，第二阶段prompt tuning）不同，本文采用一阶段训练过程，将SSL与prompt tuning集成，避免了复杂的训练流程。此外，本文首次在GraphRAG框架中应用JEPA，通过分离潜在识别和表示学习，实现了更高效的图表示学习。

核心方法

AGE方法的整体思路是模仿文本嵌入的自监督学习过程，同时结合JEPA改进表示质量。框架包含四个主要步骤：检索、图预处理、嵌入和推理。检索阶段使用SentenceBert对知识数据进行索引，结合静态k近邻检索和Prize-Collecting Steiner Tree子图构建。嵌入阶段是AGE的核心，将检索到的子图S*首先通过图编码器得到节点特征hin，然后通过节点采样器将节点分为关键节点和辅助节点。关键节点通过概念编码器-解码器处理，目标编码器生成预测目标，形成JEPA架构。最后通过图结构聚合器将输出聚合为单一token输入LLM。推理阶段使用任意LLM生成答案。

核心创新点在于引入了基于强化学习的节点采样器和JEPA架构的组合。节点采样器通过MHA网络、线性层和softmax激活计算节点概率分数pNS，自适应地选择关键节点而非随机掩码。这一策略源于图的简洁逻辑结构特性——某些关键节点对维护图完整性至关重要，掩码它们会导致SSL效率低下。JEPA架构通过stop-gradient操作防止表示坍塌，概念编码器-解码器学习从关键节点重建辅助节点表示，目标编码器直接学习对下游任务有用的表示。两者协同工作：节点采样器识别潜在的关键节点，编码器-解码器通过ELBO优化学习重建这些表示，实现了高效的图表示学习。本质上，AGE将图表示学习分解为潜在识别和表示学习两个独立过程，每个过程有明确的优化目标。

方法步骤详情

AGE方法的完整步骤如下：第一步，检索阶段，使用SentenceBert对知识数据进行索引，给定查询q，通过静态k近邻检索找到相关图元素，然后使用Prize-Collecting Steiner Tree扩展检索到的图元素及其邻接关系，构建子图S*=(V*, E*)。第二步，文本嵌入，将检索到的子图S*转换为文本格式，与输入查询q拼接，使用与冻结LLM兼容的预训练TextEmbedding函数得到htext = TextEmbedding([textualize(S*); q]) ∈ ℝ^{L×dl}。第三步，图嵌入，使用图编码器GNN_GE从子图S*得到hin = GNN_GE(S*; θ_GE) ∈ ℝ^{N×dg}，其中N = |V*|，dg是节点特征维度。第四步，节点采样，节点采样器通过MHA处理hin，得到zMS，然后通过线性层和softmax得到节点概率分数pNS = Softmax(Linear(zMS; θ^{Linear}_{MS})) ∈ [0,1]^N，根据采样率ρ选择Nkey = ⌈ρN⌉个关键节点Ikey和辅助节点Iaux。第五步，概念编码-解码，从hin中提取关键节点特征hkey，通过概念编码器MHA_CE编码为zkey = MHA_CE(hkey + PE(hkey); θ_CE)，将zkey与辅助节点的占位向量组合为z，通过概念解码器MHA_CD解码得到hout = MHA_CD(z + PE(z); θ_CD)。第六步，目标编码，使用目标编码器MHA_TE从hin生成目标嵌入htarget = MHA_TE(hin + PE(hin); θ_TE)。第七步，聚合，通过图结构聚合器GNN_GSA聚合hout，使用mean pooling和MLP投影得到适合LLM输入维度的token。第八步，优化，使用三个损失函数优化不同模块：提示调优损失LPT优化θ_TE、θ_GSA和θ_Proj；目标损失Ltarget优化θ_GE、θ_CE和θ_CD；采样损失LNS通过REINFORCE算法优化θ_NS。总损失为LPT + Ltarget + LNS，由于每个损失优化不同的参数且无重叠，无需调整权重。

技术新颖性

技术新颖性体现在多个方面：首先，AGE首次将JEPA架构应用到GraphRAG框架中，这是JEPA在图结构数据表示学习中的首次应用，通过消除不必要的细节重建，专注于学习更抽象的语义表示。其次，引入了基于强化学习的节点采样器，用自适应方法替代传统的随机节点掩码，这是图SSL中的创新设计，能够识别图中的关键节点。第三，采用一阶段训练过程，将SSL与prompt tuning集成，避免了现有两阶段方法的复杂性。第四，从认知神经科学汲取灵感，模仿人类颞叶的自下而上联想学习机制，这种跨学科的应用是新颖的。第五，设计了分离的优化目标，每个损失函数优化不同的参数集，无需权重调整，这是工程上的创新。实验显示，使用JEPA与可学习节点采样的组合相比基线G-Retriever在ExplaGraphs数据集上提升了26.72个百分点（从0.5595到0.8267），证明了这些创新的有效性。

**Fig. 1**: Overview of GraphRAG with the proposed Adaptive-masking for Graph Embedding (AGE) embedding

**Fig. 2**: Architecture for Adaptive-masking for Graph Embedding

实验结果

核心发现包括：AGE方法在四个基准数据集上一致性地提升了非参数搜索组件的性能，取得了最先进的准确率。在ExplaGraphs数据集上，AGE G-Retriever（Llama3.2-1B）相比基线G-Retriever提升了26.72个百分点（从0.5595到0.8267），这是最显著的提升。在SceneGraphs数据集上，相同的配置提升了6.44个百分点（从0.7540到0.8184）。在WebQSP数据集上，AGE G-Retriever（Llama3.2-1B）提升了2.4个百分点（从60.1到62.5）。消融研究揭示了三个关键发现：JEPA架构比生成式架构（GA）更有效，AGE使用GA和随机掩码比基线提升9.37%，而使用JEPA和随机掩码提升15.46%；可学习节点采样器相比随机掩码带来额外7.27个百分点的提升（从15.46%到22.75%）；JEPA与可学习节点采样的组合实现26.72%的总体提升。采样率ρ的分析显示，ρ=0.3在ExplaGraphs上达到最佳性能（Llama3.2-1B为81.4%，Llama3.2-3B为92.6%），在WebQSP上Llama3.2-1B也达到最佳（62.5%），Llama3.2-3B为第二佳（72.2%，最佳为73.5%在ρ=0.35）。损失权重消融显示，等权重（1:1:1）始终表现最佳，偏离等权重会降低性能，这支持了三个损失操作于不相交参数、无需显式加权的设计。定性分析显示，节点采样器能够成功从特定领域采样实体作为关键节点，例如从'saving souls'推断'missionaries'和'Christians'比反向推断更容易。与其他方法集成时，AGE与AMAR结合在WebQSP上达到86.5%（Llama2-7B），相比AMAR的84.3%提升2.2个百分点。

**Table 1**: Performance comparison across ExplaGraphs, SceneGraphs, and WebQSP datasets under the five settings

**Table 2**: Performance improvements (Llama3.2 1B, ExplaGraphs; % points)

**Fig. 4**: Performance against sampling rate (WebQSP)

**Fig. 5**: Node embedding of G-Retriever and AGE G-Retriever with sampling rate ρ = 0.3 using t-SNE

查看结构化数据

任务	指标	本文	基线	提升
ExplaGraphs	Accuracy	0.8267	0.5595	+26.72个百分点 (Llama3.2-1B, AGE G-Retriever vs G-Retriever)
SceneGraphs	Accuracy	0.8184	0.7540	+6.44个百分点 (Llama3.2-1B, AGE G-Retriever vs G-Retriever)
WebQSP	Hit@1	62.5	60.1	+2.4个百分点 (Llama3.2-1B, AGE G-Retriever vs G-Retriever)
WebQSP	Hit@1	86.5	84.3	+2.2个百分点 (Llama2-7B, AGE AMAR vs AMAR)
CWQ	Hit@1	85.2	82.9	+2.3个百分点 (Llama2-7B, AGE AMAR vs AMAR)

局限与改进

作者承认的局限性包括：使用了固定的采样率ρ=0.3，尽管不同图中的关键节点密度存在变化；只在GraphRAG任务上测试AGE，尽管它适用于其他模态；主要针对小规模模型，由于计算限制，对大规模LLM的有效性尚未探索；专注于表示检索到的结构数据供LLM使用，而非直接解决图学习任务（如节点分类或链接预测）。观察到的局限性包括：方法在WebQSP等大型知识图谱数据集上提升相对较小（2.4个百分点），可能是因为这些数据集包含预训练时未出现的文本知识，非参数检索器难以提供关键信息；节点采样器的训练依赖强化学习，可能存在训练不稳定或收敛慢的问题；方法假设图结构中存在明确的关键节点，但对于某些均匀分布的图结构可能不适用；由于使用冻结LLM，方法的有效性可能依赖于LLM的预训练质量和文本编码器的特性，对不同LLM的泛化性需要进一步验证。

独立分析的弱点

独立分析的弱点包括：固定采样率的设计忽略了不同图结构的关键节点密度差异。改进方向可以是设计自适应采样率机制，根据图的统计特征（如节点度分布、聚类系数等）动态调整采样率。在WebQSP等大型知识图谱上提升有限，可能是因为非参数检索器的固有限制。改进方向是结合轻量级的可训练检索组件，或者在子图构建阶段引入更复杂的相关性评分机制。节点采样器的RL训练可能不稳定。改进方向是引入更稳定的RL算法（如Actor-Critic）或设计辅助监督信号（如基于图拓扑特征的先验知识）。方法专注于GraphRAG任务，泛化到其他图学习任务的能力未知。改进方向是在节点分类、链接预测等经典图学习任务上验证AGE的有效性，或设计任务无关的图表示学习目标。使用冻结LLM限制了方法的潜力。改进方向是探索与轻量级LLM微调方法（如LoRA）的协同优化，或者设计适配不同LLM架构的图嵌入模块。对大规模LLM的有效性未探索。改进方向是在更大规模的LLM（如GPT-4级别的模型）上验证AGE的有效性，虽然计算成本较高，但可以揭示方法在更强基线上的相对优势。

未来方向

作者提出的未来工作方向包括：探索动态采样率机制，根据不同图的关键节点密度自适应调整；将AGE应用到其他模态，如视觉-语言任务中的图结构表示；在大规模LLM上验证AGE的有效性；探索AGE在图学习任务（如节点分类、链接预测）中的应用。基于成果可延伸的研究方向包括：设计更复杂的节点采样器，例如基于图注意力网络的分层采样或基于图生成模型的采样策略；将JEPA扩展到多模态图表示学习，同时处理文本、图像、结构化数据等多种模态；探索AGE在时序图或动态图中的应用，处理图结构随时间演化的场景；研究AGE与轻量级LLM微调方法的协同优化，例如通过联合训练图嵌入模块和LoRA适配器；开发针对特定领域（如医疗、法律、金融）的AGE变体，利用领域知识设计更好的节点采样策略；探索AGE在联邦学习或隐私保护设置中的应用，例如通过图嵌入的差分隐私保护用户数据。

复现评估

复现评估：论文未明确提及代码开源情况，这是一个可复现性方面的不足。数据方面，使用了四个公开基准数据集：ExplaGraphs（生成式常识推理）、SceneGraphs（视觉问答）、WebQSP和CWQ（从Web问题派生的大型问答数据集，使用Freebase知识图谱）。这些数据集都是公开可获取的，有利于复现。算力方面，使用了多个LLM作为冻结组件：Llama3.2（1B和3B）、Llama2（7B和13B）、Llama3.1（8B）。这些模型规模相对适中，但训练图嵌入模块仍需要可观的GPU资源，估计需要至少24GB显存的GPU。实验细节相对完整，包括超参数设置（如采样率ρ=0.3）、损失函数定义、优化器配置等，但在某些方面仍有不足：未提供训练的具体时长或收敛曲线；未消融图编码器架构的选择；未提供不同随机种子的实验结果以评估方差；未提供训练集和验证集的划分细节。总体而言，复现难度中等，主要挑战在于计算资源和缺乏开源代码。如果作者能够开源代码并提供更详细的训练日志，复现将更容易。

论文图表

Fig. 3: Performance against sampling rate (ExplaGraphs)

该图展示了采样率ρ对AGE G-Retriever在ExplaGraphs数据集上性能的影响。x轴是采样率ρ从0.3到0.7，y轴是准确率。图中包含两条曲线，分别对应Llama3.2-1B和Llama3.2-3B。结果显示，ρ=0.3时两个模型都达到最佳性能：Llama3.2-1B为81.4%，Llama3.2-3B为92.6%。随着采样率增加，性能下降，说明过高的采样率（即选择过多的关键节点）会降低模型性能。

这张图对理解AGE的超参数选择至关重要。它提供了采样率ρ的消融研究结果，帮助读者理解为什么作者选择ρ=0.3作为默认设置。图中的结果还暗示了关键节点的最佳比例——大约30%的节点应该被识别为关键节点，这为理解图结构中关键节点的分布提供了洞察。此外，Llama3.2-3B在所有采样率上都优于Llama3.2-1B，验证了更大模型的优势。