KG-RAG：将知识图谱注入大模型「记忆」

痛点：LLM 在生物医学领域的「幻觉」困境

想象一下这样的场景：一位研究罕见病的医生向 GPT-4 询问某种基因突变与特定疾病的关联，模型给出了一个听起来专业、流畅，但实际上在最新医学文献中早已被推翻的答案。这并非模型「笨」，而是它所依赖的「记忆」——训练数据——天然存在滞后性。GPT-4 的知识截止日期停留在某个过去的时间点，而生物医学领域的知识更新速度是出了名的快：新药获批、基因功能新发现、疾病机制新认识，每一天都在发生。

这正是 KG-RAG（Knowledge Graph-based Retrieval Augmented Generation）要解决的核心问题。

背景：UCSF BaranziniLab 的医学 AI 探索

KG-RAG 由加州大学旧金山分校（UCSF）的 BaranziniLab 开发。这是一个专注于神经遗传学和计算生物学的实验室，团队成员横跨生物信息学、机器学习和临床研究。2023年11月，他们在 arXiv 上发布了论文《Biomedical knowledge graph-enhanced prompt generation for large language models》（arXiv:2311.17330），提出了将结构化知识图谱注入 LLM 推理流程的方法。

项目使用的知识图谱名为 SPOKE（Scalable Precision Medicine Oriented Knowledge Engine），这是一个专注于生物医学领域的大规模异构图谱。SPOKE 整合了超过 40 个生物医学数据源，涵盖基因、蛋白质、药物、化合物、疾病等 21 种节点类型，总计超过 2700 万个节点和 5300 万条边。这些数据来自文献挖掘、临床试验、基因表达分析等多个渠道，并以结构化的关系形式组织起来。

核心创新：「提示词感知上下文」提取

传统 RAG（Retrieval-Augmented Generation）的做法是：从大量文本语料中检索与问题相关的文本片段，注入到 prompt 中。这种方法有效，但有一个显著缺陷——检索结果往往是「语义相似」而非「因果相关」。

举个具体例子。当用户询问「setmelanotide 对 Bardet-Biedl 综合征的疗效」时，传统 RAG 可能返回所有包含这些关键词的新闻文章。但 KG-RAG 的做法是：先通过 SPOKE 图谱找出「setmelanotide → 适应症 → Bardet-Biedl 综合征」这条因果路径，然后只提取这条路径上的上下文。这被称为「提示词感知上下文」（Prompt-Aware Context）——即「足以回答当前问题的最小上下文」。

这种设计带来两个关键优势：减少了输入 token 数量（降低成本和延迟），以及保证了上下文的因果可验证性（不是统计相关，而是真实生物医学关系）。

技术架构：三层解耦的 RAG Pipeline

KG-RAG 的代码架构清晰，分为三个核心模块：

第一层：向量数据库构建（vectorDB/create_vectordb.py）。系统使用 Chroma 向量数据库存储疾病节点的嵌入向量。嵌入模型采用 sentence-transformers/all-MiniLM-L6-v2（用于节点检索）和 pritamdeka/S-PubMedBert-MS-MARCO（用于上下文检索）。配置文件中定义的分块策略为 chunk_size=650、chunk_overlap=200，批量大小为 200。

第二层：上下文检索（utility.py）。这是核心逻辑所在。检索流程为：（1）用户问题编码为向量；（2）在 Chroma 中检索最相似的疾病节点；（3）通过 SPOKE REST API 获取该疾病节点的子图上下文；（4）使用第二个嵌入模型对上下文进行重排序，过滤掉与问题相似度低于 0.5（75 百分位阈值）的候选上下文。

第三层：LLM 生成（rag_based_generation/GPT/text_generation.py）。支持的模型包括 GPT-4、GPT-3.5-Turbo（通过 OpenAI API 或 Azure OpenAI）和 Llama-2-13b-chat（本地运行）。生成时将系统 prompt（定义在 system_prompts.yaml 中）与检索到的上下文拼接，temperature 固定为 0 以保证确定性输出。

整个 Pipeline 还提供了一个 Cypher 查询变体（notebooks/cypher_rag_using_langchain.ipynb），允许用户直接用 Cypher 查询 SPOKE 图谱的 Neo4j 实例，实现更精细的图路径控制。

Benchmark 数据集 BiomixQA

团队还发布了一个专门的评测数据集 BiomixQA，托管在 HuggingFace（kg-rag/BiomixQA），包含两类问题：多项选择题（MCQ）和判断题（True/False）。题目覆盖多种生物医学概念，用于系统评估不同 LLM 在生物医学 QA 任务上的表现。这个数据集可以通过三行 Python 代码加载，是评估 KG-RAG 效果的标准化工具。

上手门槛与局限性

门槛确实不低。 项目以 Jupyter Notebook 为主要使用界面，核心运行脚本是命令行 Python 模块。具体来说：需要 Python 3.10.9 环境、CUDA 支持（如果本地运行 Llama）、GPT API Key 或 Azure 账户、以及 HuggingFace 账户用于下载 Llama 模型。首次运行需要执行 python -m kg_rag.run_setup 构建向量数据库（从已提供的疾病-基因关系 pickle 文件），并可选下载 Llama 模型。整个过程没有 Docker 支持，不提供 Web UI，配置全靠手改 config.yaml。

当前版本的重要局限： 系统目前仅支持疾病相关的查询。README 中明确写道「At the moment, KG-RAG is specifically designed for running prompts related to Diseases. We are actively working on improving its versatility.」也就是说，如果你询问的是某类药物的研发进展、某种蛋白质的结构分析，当前版本可能无法提供有意义的上下文。这是一个需要持续关注的发展方向。

行业意义：结构化知识与概率模型的桥梁

KG-RAG 代表了一个重要趋势：在 LLM 的「统计智慧」与专家系统的「规则严谨」之间寻找结合点。SPOKE 图谱提供了可验证的因果关系，LLM 负责理解和生成自然语言，两者互补。与其争论「大模型能否替代专业知识」，不如问「如何让大模型安全地调用专业知识」——KG-RAG 正是对这个问题的务实回答。

此外，项目将生物医学数据与 RAG 技术结合的方式，为其他垂直领域（法律、金融、工程）提供了参考范式：只要有高质量的结构化知识图谱，就可以复制这套「图谱 → 向量检索 → LLM 生成」的 Pipeline。