AGE 的核心洞察就一句:图节点不平等,关键节点掩不得。用可学习采样器识别关键节点保护起来,只掩辅助节点做 JEPA,图嵌入就能和 frozen LLM 对齐。掩码 SSL 在图上一直不灵,根因是掩错了节点。今天分享一篇 GraphRAG 嵌入对齐的工作 AGE:用自适应掩码替代随机掩码,把 frozen LLM 的图嵌入问题给解了。GraphRAG 把图作为外部知识喂给 LLM,但图嵌入和文本嵌入一直对不齐——尤其 frozen LLM。常规做法是掩码 SSL(像 MAE/JEPA 那样随机掩码节点再预测),但在图上跑效果差。AGE 判断干脆:图不是文本,有关键节点信息密集难预测,随机掩到它们就全塌。旧做法哪里不够随机掩码在文本上好用,在图上失灵,两个原因:图是简洁表示:节点少,每个节点信息密度高关键节点难预测:有些节点 hold 主导上下文,从邻居推不出来随机掩码一掩到关键节点,SSL 就学不到东西——预测目标本身就没法从周围重建。四模块架构AGE 用 Transformer 掩码 SSL,架构设计成类似文本编码器(解决图-文 latent 对齐)。四个模块配合:图编码器把检索到的子图 S* 用 GNN 编码成节点特征 h_in GNN_GE(S*) ∈ R^{N×d_g}。这是 AGE 的输入。节点采样器核心创新。用 MHA Linear Softmax 给每个节点打分,得到概率分布 p_NS。按采样率 ρ0.3 抽 N_key ⌈ρN⌉ 个关键节点 I_key,剩下的 I_aux V* \ I_key 是辅助节点。采样器可学习——它学会识别哪些节点信息密集、难从邻居预测,把这些保护起来不掩。概念编解码把关键节点特征 h_key 编码成潜在 z_key MHA_CE(h_key PE),和辅助节点的占位向量 z_aux(直接拷 h_in 的值)拼回 z,保持原节点位置。然后解码 h_out MHA_CD(z PE)。h_out 要预测的是辅助节点的目标嵌入——也就是说,用关键节点去重建辅助节点,反过来学。目标编码器JEPA 的预测目标。MHA_TE 把 h_in 投到语义空间得 h_target。AGE 训练 h_out 逼近 h_target(只对辅助节点算 loss)。语义空间预测比输入空间重建鲁棒(JEPA 核心思想)。图聚合器GNN_GSA 把 h_out 聚合成单个 token,MLP_Proj 投到 LLM 输入维度,喂给 frozen LLM。三损失优化三个 loss 各管各的模块,不重叠,不用调权重:L_PT(提示调优):下游任务准确率,优化 target encoder 聚合器 projectorL_target(JEPA重建):辅助节点的 h_out 逼近 h_target,优化图编码器 概念编解码器。用知识蒸馏:关键节点表示 → 重建辅助节点L_NS(RL采样):采样操作不可微,用 RL 启发的监督。辅助节点重建误差大的,增大其被选为关键节点的概率——把难预测的节点推到 I_key总 loss L_PT L_target L_NS,1:1:1 等权最优(消融验证)。实验全面领先四个 benchmark 全面 SOTA:ExplaGraphs:Llama3.2-1B 26.72pp(0.5595→0.8267),Llama3.1-8B 9.07pp(0.8705→0.9612)SceneGraphs:Llama3.2-3B 0.9486WebQSP:AGEAMAR Llama2-7B 86.5(vs AMAR 84.3)CWQ:AGEAMAR 85.2,超过 GPT-4 的 ReKnoSfrozen LLM 非参数检索器,能打过 GPT-4 的 LLM 检索器——图嵌入对齐做对了,小模型也能赢。消融层层递进:G-Retriever 0.5595 → GA随机掩码 0.6532(9.37%)→ JEPA随机掩码 0.7141(15.46%)→ GA采样器 0.7870(22.75%)→ JEPA采样器 0.8267(26.72%)。JEPA 比 GA 强,采样器比随机掩码强,两者叠加效果最大。采样率 ρ0.3 最优——掩太多(ρ大)关键节点被掩进去了,掩太少(ρ小)学不到东西。小扬总结AGE 的核心洞察就一句:图节点不平等,关键节点掩不得。用可学习采样器识别关键节点保护起来,只掩辅助节点做 JEPA,图嵌入就能和 frozen LLM 对齐。掩码 SSL 在图上一直不灵,根因是掩错了节点。论文标题: AGE: Adaptive-masking for Graph Embedding in Graph Retrieval-Augmented Generation 论文链接: https://arxiv.org/abs/2607.00052