【KG+RAG 论文】医学知识图谱检索增强 LLM 的框架 —— KG-RAG

论文：Biomedical knowledge graph-enhanced prompt generation for large language models

⭐⭐⭐

Code：github.com/BaranziniLab/KG_RAG

文章目录

• • 论文速读
  • 模型效果
  • 总结

论文速读

这篇论文提出了 KG-RAG 的框架，使用医学知识图谱（SPOKE）来对 LLM 进行检索增强。

该框架的运行效果如下图：

上图中，黄色部分是用户问题，蓝色部分是 GPT-4 的原生回答，绿色部分是经过 KG-RAG 框架处理后生成的回答。左边的 (A) 是一个关于一跳推理的问题，右边的 (B) 是一个关于两条推理的问题。

可以观察到，KG-RAG 可以解决这个单跳和双跳的问题，并且相比于 GPT-4，可以提供更加简单明了的答案。

工作过程：KG-RAG 框架的基本工作原理如下：

1. 实体识别与实体链接：根据用户的问题，使用 LLM 做问句中的疾病实体识别，再对识别的结果对 KG 进行实体链接的检索，得到 KG 中相应的节点（即疾病的节点）
2. 上下文提取（Context pruning）：从 KG 中召回与这个实体相关联子图，再基于 embedding 计算语义相似度从子图中过滤出有用的三元组，之后再将这些三元组将其转换为自然语言
3. 提示组装与文本生成：把上一步得到的自然语言，与 question 拼在一起，组合为 prompt，再加上 SYSTEM_PROMPT，送给 LLM 来回答，从而获得最终答案

模型效果

可以看到，在 KG-RAG 框架下，各 LLM 的表现都有提升。

总结

这篇文章提出的框架是一个结合 KG 来做 RAG 的有效方案，但当用于工业落地时，仍会存在很多问题：

• 实体识别使用了 LLM，之后又做了 entity link，这样的效率肯定不太高。
• 为了从召回子图过滤出有用的三元组，这里需要专门的 embedding 模型去做
• 从关联子图 -> 自然语言这一步，也存在很多坑

这篇文章的工作主要是在医学领域结合 KG 来实现 RAG，但在其他领域，需要结合实际的场景去定制具体的策略。