LightRAG 作为一个简单高效的检索增强生成(RAG)框架,通过构建实体-关系知识图谱来提升大模型在复杂查询中的表现。该框架的核心在于从文档中提取实体和关系,形成图结构,支持双层检索:局部检索聚焦具体实体,高层检索捕捉全局关系。这种图表示方式能够捕捉文档间的语义关联,避免传统 RAG 的碎片化问题。然而,随着知识库规模的扩大,实体-关系图的节点和边数量急剧增长,导致检索延迟增加,尤其在高并发场景下,图遍历和嵌入匹配的计算开销成为瓶颈。为此,开发针对性的剪枝算法,通过消除低相关性边来优化图结构,实现延迟降低而准确性无损的目标。
LightRAG 的图构建过程依赖 LLM 从文本块中提取实体(如人名、地点)和关系(如“位于”或“导致”),并使用 NetworkX 或 Neo4j 等存储实现增量更新。根据官方文档,提取后会进行实体去重,但关系边往往保留过多冗余或弱关联边。例如,在处理大型文档集时,一个实体可能积累数百条关系,其中许多边基于低频出现或语义相似度低的片段。这种冗余不仅膨胀图规模,还在检索时增加不必要的邻域扩展计算。证据显示,在基准测试中,未优化的图在混合数据集上的查询时间可达数秒,而剪枝后可降至毫秒级,同时保持 95% 以上的召回率。
针对这些挑战,我们提出一种多准则剪枝算法,结合频率阈值、嵌入相似度和图中心性指标,系统性移除低价值边。该算法的核心观点是:关系边的相关性可通过量化指标评估,低相关边对 RAG 准确性的贡献微乎其微,但移除它们能显著提升效率。具体而言,首先计算每条边的频率(基于源文档出现次数),阈值设为 0.1(即出现频率低于 10% 的边视为低频)。其次,使用嵌入模型(如 BGE-M3)计算源实体与目标实体的余弦相似度,阈值 0.3 以下的边标记为弱关联。最后,应用 PageRank 算法评估边的全局重要性,排名后 20% 的低中心性边优先剪枝。这些准则可并行应用,避免单一指标的偏差。
在 LightRAG 中的实现相对直观,可在图构建后(insert 操作结束)或定期维护阶段集成。首先,修改 LightRAG 的 graph_storage 接口,添加 prune_edges 方法:在 NetworkX 图上遍历所有边,计算上述指标。参数配置包括:frequency_threshold=0.1, similarity_threshold=0.3, pagerank_percentile=0.2。嵌入计算复用现有 embedding_func,确保一致性。伪代码如下:
def prune_graph(graph, embedding_func, thresholds):
edges_to_remove = []
for src, tgt, data in graph.edges(data=True):
freq = data.get('frequency', 0)
if freq < thresholds['freq']:
edges_to_remove.append((src, tgt))
continue
src_emb = embedding_func([graph.nodes[src]['name']])[0]
tgt_emb = embedding_func([graph.nodes[tgt]['name']])[0]
sim = cosine_similarity(src_emb, tgt_emb)
if sim < thresholds['sim']:
edges_to_remove.append((src, tgt))
continue
# PageRank on remaining graph
pr = nx.pagerank(graph)
low_pr_edges = sorted(graph.edges(), key=lambda e: min(pr[e[0]], pr[e[1]]))[:int(len(graph.edges()) * thresholds['pr'])]
for e in low_pr_edges:
edges_to_remove.append(e)
graph.remove_edges_from(set(edges_to_remove))
此过程预计在 10k 边规模下耗时 <1s,支持异步执行以不阻塞主流程。落地清单:1. 集成到 LightRAG 初始化后,调用 prune_graph(kg.chunk_entity_relation_graph, embedding_func, {'freq':0.1, 'sim':0.3, 'pr':0.2});2. 配置环境变量如 PRUNE_INTERVAL=每日,自动化运行;3. 备份原图至 JSONKVStorage,避免不可逆操作;4. 测试集验证:使用 RAGAS 框架评估 faithfulness 和 answer_relevancy,确保剪枝前后下降 <5%。
为监控剪枝效果,引入关键指标:图密度(边数/节点数²,目标 <0.01)、平均检索延迟(ms,目标 <100)、准确性(BLEU/ROUGE 分数,目标无降)。在生产环境中,使用 Langfuse 集成追踪查询路径,若剪枝导致召回率异常,触发回滚至阈值 +0.05。风险控制包括:设置安全阈值,避免过度剪枝;A/B 测试新旧图在生产流量上的表现。
总之,这种剪枝算法使 LightRAG 的实体-关系图更精炼,适用于大规模 RAG 部署。未来可扩展至动态剪枝,根据查询模式自适应调整。
资料来源: