一、从KG-RAG的任务及现有方案说起

KG-RAG检索与查询相关的最相关的k条推理路径，为LLMs推理提供简洁准确的上下文信息，可以先回顾下KG-RAG的执行阶段，涉及两个主要阶段：检索和生成。

检索阶段，给定一个查询q，第一阶段通过将q与知识图谱中的实体和关系进行匹配，构建一组候选的“推理路径”，通过从知识图谱(G)中搜索相关的三元组(s,r,e)来完成的，检索(q,G)→ {P}，然后按它们与q的相关性对检索到的推理路径集合{P}进行排序。

推理路径Pi可以正式定义为：Pi=(s,r1,m1,r2,…,rk−1,mk−1,rl,e)，其中s是起始实体，e是答案实体，mj是中间实体，rj是连接这些实体的关系。路径中的跳数，等于关系的数量，决定了路径的长度l。对于给定的查询q，正确的推理路径包括与q对应的正确答案实体。

在推理阶段，排名靠前的k条路径增强查询，形成一个丰富的查询q′。然后将这个q′输入到大型语言模型（LLM）中以生成最终输出：Generate(q′,LLM)→output.

再看现有针对问题都有哪些具体方案。

KG-RAG框架根据知识图谱（KG）的信息是否微调到大模型（LLMs）中，被归类为模块化或耦合。

如模块化方法，ToG，通过用LLMs替换传统的排名模型来提高检索准确性。然而，由于缺乏先验的KG知识和频繁调用LLMs的需求，ToG落地比较难讲；

又如耦合框架，RoG，使用KG信息对LLMs进行微调，使它们能够生成作为查询模板的“关系路径”，直接从KGs检索正确的推理路径。也有gnn-rag利用图神经网络（GNN）提取有用的推理路径。

二、引入分类思想的FFRAG-KGRAG方案

最近的工作《FRAGRAG: A Flexible Modular Framework for Retrieval-Augmented Generation based on Knowledge Graphs》(https://arxiv.org/pdf/2501.09957)，核心思路是通过动态调整检索过程以适应查询的复杂性，里面用到分类治之的思想，虽然没有开源代码，但也可以看看一些思路。

核心包括推理感知和检索两个主要核心点。

1、推理感知

推理感知根据查询的复杂性将推理任务分类为简单或复杂。正确推理路径的跳数与查询中的特定统计数据密切相关，例如实体数量、关系数量和子句数量。所以可以通过预测推理路径的跳数范围，使用二分类器对查询进行分类。

一条路径P包含两种类型的信息：语义（即实体和关系）和结构性（即跳数）。语义信息主要来源于知识图谱，难以提前感知和利用，结构信息与知识图谱和q都有关联，查询上下文q越复杂，路径P中的跳数就越多（即表示更难的推理任务）。

因此，在一个可容忍的误差范围内，可以仅基于查询上下文q预测路径P中的跳数。预测出的跳数可以作为提升非特定检索过程的关键因素。例如，设定跳数阈值为2，如果推理路径的最小跳数小于等于2，则任务被分类为简单；否则，被分类为复杂。

可以看下具体细节：

使用一组公共的知识图谱问答数据集（Freebase和Wiki-Movies）训练一个二分类器，每个数据集包括一个基础知识图谱和大量配对的查询Q及其对应的答案A。

在训练之前，对于每个查询q∈Q，识别从查询实体Entq到答案实体Enta∈A的所有最短推理路径。该查询所有推理路径中的最小跳数H决定了查询标签Y:Y=1（复杂），如果H≥δ；Y=0（简单），如果H<δ，阈值δ设为2。

但要进行建模的时候，需要进行编码：

其中查询编码器可以是任何编码机制，如语言模型BERT、词嵌入Word2Vec或TF-IDF。这样，在训练期间二分类器的分类损失为：

这里，p(yq∣q)=解码器(hq)表示q被归类为简单或复杂的概率。

2、检索生成

检索生成模块旨在准确识别与查询相关的推理路径，将KG-RAG检索过程细化为“预处理-检索-后处理”pipeline，然后最后交给LLM生成结果。

首先是预处理阶段，从原始KG中提取包含重要实体和关系的子图，目的是缩小检索范围，提高检索效率。

这里用到常用的思路，基于重要性评估移除较不相关的实体和边。如基于实体的子图剪枝中采用广义排名机制（GRM），例如随机游走重启（RWR）、个性化PageRank（PPR）以及PageRank微调（PRN），然后选择前n个实体。基于边的子图剪枝中，应用边排名模型（ERM）作为检索器（例如BM25、SentenceTransformer）计算它们与查询q的语义相似性。

然后是检索阶段，根据查询的复杂性，采用不同的检索策略。

对于简单推理任务，涉及较短推理路径的简单查询，采用更广泛的检索方法至关重要，以最小化信息丢失，所以采用广度优先搜索（BFS）算法，允许高效遍历查询实体与实体集之间的所有推理路径。

对于复杂推理任务，对于具有较长推理路径的复杂查询，增加检索路径不仅会呈指数级增加计算成本，还会引入大量冗余信息，所以采用最短路径检索算法（如Dijkstra算法），以最小化计算开销并减少噪声。

最后是后处理阶段，检索过程主要关注于寻找从查询实体到潜在答案实体的路径，忽略中间实体的语义及其与查询的相关性，这可能导致推理路径的收集是无序且冗余的，即几个明显问题，包含不相关或罕见中间实体的推理路径可能会误导大型语言模型（LLM）的推理；路径增加了提示的长度，这不仅增加了推理成本，还冒着超出上下文长度限制的风险；理性能受到这些路径在提示内位置的影响，位于开头的路径影响更为显著；

所以，使用路径排名模型（如DPR、ColBERT等）过滤冗余推理路径，确保最终输入LLMs的推理路径是最相关且有益的。

最后，设计一个提示模板，用过滤后的推理路径P来增强问题q。

总结

本文主要介绍了KGRAG的一个进展工作，核心还是对路径剪枝思路的优化，分儿治之。每个流程其实都是之前方案中常用的