Biomedical knowledge graph-enhanced prompt generation for large language models

1. 生物医学知识图谱增强大语言模型提示生成

论文地址：[2311.17330] Biomedical knowledge graph-enhanced prompt generation for large language models (arxiv.org)

源码地址：https://github.com/BaranziniLab/KG_RAG 

2. 摘要

大语言模型（LLM）一直在以前所未有的速度推动人工智能的进步，但在生物医学等知识密集型领

域仍然面临挑战。预训练和特定领域微调等解决方案增加了大量的计算开销，后者也需要领域专业

知识。外部知识注入是特定于任务的，需要模型训练。在这里，我们引入了一种与任务无关的基于

知识图谱的检索增强生成（KG-RAG）框架，通过利用 Llama-2-13b、GPT-3.5-Turbo 和 GPT-4

等大语言模型和大规模生物医学知识图谱 SPOKE 来生成植根于既定知识的有意义的生物医学文

本。 KG-RAG 持续增强了 LLM 在各种提示类型上的表现，包括一跳和两跳提示、药物再利用查

询、生物医学真/假问题和多项选择题 (MCQ)。值得注意的是，KG-RAG 在具有挑战性的 MCQ 数

据集上使 Llama-2 模型的性能显着提升了 71%，证明了该框架能够以更少的参数为特定领域的问

题赋能开源模型。此外，KG-RAG 还增强了专有 GPT 模型的性能，例如 GPT-3.5，它在 MCQ 数

据的上下文利用率方面比 GPT-4 有所改进。我们的方法还能够解决药物再利用问题，返回有意义

的再利用建议。总之，所提出的框架以优化的方式分别结合了KG和LLM的显性和隐性知识，从而

增强了通用LLM在统一框架中解决特定领域问题的适应性。

3. 简介

大语言模型 (LLM) 在解决涉及语言建模和处理各个领域的复杂任务方面表现出了令人印象深刻的

性能 [1]。 LLM 在自监督学习框架中对大量文本数据进行预训练，该框架可以是掩码语言建模（例

如 BERT 等模型 [2,3]）或自回归框架（GPT 等模型 [4,5]）。这种预训练将有关语言的知识编码到

模型参数中。与深度神经网络中常用的迁移学习方法类似，这种隐性知识可以通过监督训练进行细

化，以在一系列特定领域的任务中表现出色[5,6]。尽管如此，作为对输入提示的响应[7-9]，给出的

答案在语言上是连贯的（即：“幻觉”），但LLM中知识的“隐式表示”也被证明会产生非事实信息。

这个问题对在生物医学等精度要求严格的领域具有重大挑战。

人们引入了各种策略来解决LLM的幻觉。其中一种解决方案涉及利用特定领域的数据来预训练

LLM，而不是依赖通用文本语料库。这种方法产生了PubMedBERT[10]、BioBERT[11]、

BlueBERT[11,12]、SciBERT[13]、ClinicalBERT[14]、BioGPT[5]、Med-PaLM[15]等模型的创建，

和BioMedGPT[16]。然而，考虑到当前LLM实现类人性能所需的巨大计算成本和时间，这种从头开

始对LLM进行预训练是不切实际的。提出了一种称为提示调整的替代方法，作为增强 LLM 性能的

方法，例如通过使用零样本[17]、少样本[4] 和思想链[18] 提示策略。

尽管提示调整方法已被证明是有效的，但其性能仅限于知识密集型任务，这些任务需要提供有关世

界的最新知识来解决用户提示问题。为了解决此类知识密集型任务，提出了一种将KG与LLM相结

合的替代方法[19-25]。这种方法主要在问答任务中实现，其中知识图谱中包含的结构化信息用于提

供预测问题答案的上下文。虽然这种多模式综合方法显示出希望，但其下游监督培训是针对特定任

务量身定制的，限制了其多功能性和广泛适用性，从而限制了其充分利用LLM“新兴能力”[26]。为

了增强LLM的情境意识，最近引入了一种称为检索增强生成（RAG）的方法[27]。该方法涉及增强

参数预训练的LLM，使其能够访问包含有关世界的新知识的非参数存储器（例如维基百科）[27]。

在本文中，我们提出了一种将 KG 与 RAG 框架内预先训练的 LLM 相结合的方法，从而实现两全​​

其美。为了实现这一目标，我们利用了一种称为可扩展精准医学开放知识引擎（SPOKE）的生物

医学知识图谱[28]。

SPOKE 集成了 40 多个不同领域的公开生物医学知识源，其中每个源都围绕一个生物医学概念，

例如基因、蛋白质、药物、化合物、疾病以及它们的一种或多种已知关系。因为这些概念是形成定

义集的循环实体（例如，所有命名的人类基因、所有 FDA 批准的药物等），所以将这些概念集成

到单个图中揭示了连接知识源并提供生物学和本体论的新颖的多跳事实关系和每个概念的上下文。

我们提出的框架巧妙地检索了生物医学来自 SPOKE 的上下文，优化其提取以充分满足提示，从而

创建我们所说的“提示感知上下文”。这种丰富的提示进一步用作 RAG 框架中 LLM 的输入，以生成

有意义的生物医学文本。我们使用各种预训练的 LLM（包括 Llama-2-13b、GPT-3.5-Turbo 和 GP

T-4）评估了这种方法。通过改变LLM的系统提示，我们能够利用相同的框架来执行各种生物医学

任务，例如真/假和多项选择回答、药物重新利用建议以及解决与疾病相关的一般生物医学查询。

这种适应性凸显了我们的方法与任务无关的本质。进一步进行了系统的比较分析，以评估 Llama-

2-13b、GPT-3.5-Turbo 和 GPT-4 在此框架内的性能。因此，我们的方法展示了处理各种特定领域

生物医学任务的能力，而无需对通用LLM进行额外的预训练或微调。这种方法不仅简化了流程，而

且具有成本效益和时间效率。

4. 框架

我们开发了一个框架，为大语言模型构建准确且值得信赖的生物医学提示。该 KG-RAG 框架涉及

根据用户提示进行实体识别、从 KG 中提取生物医学概念、提示感知上下文生成、转换为自然语

言、提示组装和答案检索。使用不同的场景对该方法的性能进行了广泛的测试。

图 1 显示了使用两种方法作为 GPT-4 模型输入的一跳提示（图 1A 黄色框）和两跳提示（图 1B 黄

色框）：（i）基于提示，即没有 KGRAG（蓝色框）和 (ii) 有 KG-RAG（绿色框）。我们观察到，

只有 KG-RAG 能够解决该问题，并为一跳（图 1A 黄色框）和两跳提示（图 1B 黄色框）提供了简

短而直接的答案。 

基于知识图谱的检索增强生成 (KGRAG) 框架的架构。箭头方向表示该框架中管道的流向。

疾病实体识别这是KG-RAG的第一步。这一步的目标是从输入的文本提示中提取疾病概念（一个实

体），然后在KG中找到对应的匹配疾病节点（一个SPOKE概念）。这是通过两步过程实现的：i)

从提示中提取实体，ii) 与 SPOKE 进行实体匹配。实体提取识别并提取输入文本提示中提到的疾病

实体，也称为“提示疾病提取”。为了实现这一目标，在 GPT-3.5-Turbo 模型上使用了零样本提示

[17] 方法。具体来说，系统提示被设计为从输入文本中提取疾病实体并以 JSON 格式返回它们。

接下来，使用实体匹配来获取 KG 中表示的疾病的概念名称。为此，SPOKE 中所有疾病概念（即

节点）的嵌入均使用“all-MiniLM-L6-v2”句子转换器模型[46]。该过程将疾病概念的名称转换为 384

维密集向量空间，使其适合语义搜索。接下来，这些新创建的疾病概念嵌入被存储在“Chroma”向

量数据库中[47]。选择与提取的实体具有最高向量相似度的疾病概念用于后续上下文检索。如果零

样本方法无法识别提示中的疾病实体，则会从向量数据库中选择与整个输入文本提示具有最显着向

量相似性的五个疾病概念。

虚线框显示了知识图中“疾病实体识别”和“提示感知上下文提取”的详细信息 。

首先，使用转换器模型（使用超参数调整完成模型选择）将输入提示和所有提取的上下文关联嵌入

到同一向量空间（上图中的上下文潜在空间）。接下来，仅选择与输入提示向​​量具有最高余弦相似

度的上下文关联（上图）。对于上下文选择，提示上下文余弦相似度应该满足两个条件：(i) 大于

包含与所选疾病节点相关的所有上下文的相似性分布的 75%，以及 (ii) 具有 0.5 的最小相似性值。

5. 超参数分析

优化了“上下文体积”和“上下文嵌入模型”等两个超参数。

上下文量定义了允许从 KG 流向 LLM 的图形连接数量的上限。上下文嵌入模型是一种在从 KG 检

索准确的生物医学背景以响应输入提示方面表现出更高相关度的方法。为了优化这些超参数，我们

使用了两种上下文嵌入模型（基于 MiniLM 和 PubMedBert 的模型），并且上下文体积的大小不断

增加（图 2A）。对于一跳提示，在所有上下文卷设置中，基于 PubMedBert 的模型的平均性能

（Jaccard 相似度）比 MiniLM 模型高出约 10%（PubMedBert 的平均性能 = 0.67，MiniLM 的平

均性能 = 0.61）。对于两跳提示，在所有上下文设置中，PubMedBert 的性能比 MiniLM 高出

8.1%（PubMedBert 模型的平均性能 = 0.4，MiniLM 的性能 = 0.37）。

就上下文量而言，图 2A 显示单跳提示的性能曲线趋于稳定，并且双跳提示的性能曲线呈现趋于稳

定的趋势（对于两种模型）。

基于这些发现，我们选择基于 PubMedBert 的模型作为上下文嵌入模型，并将上下文量设置为

100-200 之间的值（对于大多数下游任务，我们选择上下文量为 150，对于 True/False 问题，我

们选择选择上下文体积 100）。

图 2. 超参数分析和药物再利用性能 (A) 使用一跳（上）和两跳（下）提示的超参数分析性能曲

线。 x 轴表示“上下文量”（来自 KG 的关联数量），y 轴表示提示中的平均表现（杰卡德相似

度）。红色曲线表示“S-PubMedBert-MSMARCO”，蓝色曲线表示“all-MiniLM-L6-v2”上下文嵌入模

型。 (B) LLM在药物再利用问题上的平均表现（杰卡德相似度）（上）。误差条表示 95% CI。底

部面板显示了关于 COVID-19 的药物再利用提示示例（黄色框）以及使用 KG-RAG 框架的 GPT-4

模型的相应响应（绿色框）。注意：由于空间限制，绿色框中的响应被截断。 

在本研究中，控制 LLM 输出随机性水平的“温度”参数 [4] 对于所有 LLM 均设置为 0。

6. 药物再利用

图 2B 显示了三种 LLM 模型（Llama-2-13b、GPT-3.5-Turbo 和 GPT-4）的药物再利用性能分析。

我们观察到性能指标持续增长（Jaccard 相似度）从 Llama-2-13b 到 GPT4，有趣的是这个顺序与

这些模型的参数大小的顺序一致（图 2B，表 1）。 KG-RAG 框架能够从 SPOKE 提供必要的背景

信息，以回答与 COVID-19 相关的药物再利用问题（图 2B 下图黄色和绿色框）。除了提供可行的

再利用建议外，它还对临床试验评估这些建议有效性的必要性发出了谨慎的警告。这表明一种负责

任且具有安全意识的文本生成方法（图 2B 下图绿色框）。

表 1. 使用 KG-RAG 框架在药物再利用提示上的模型性能（Jaccard 与 KG 的基本事实相似度）

7. True/False 和 MCQ 

为了评估LLM在 True/False 和 MCQ 数据集上的表现，我们随机抽取了 150 个问题并替换了

1000 次（使用 bootstrapping）。然后计算每次采样迭代的准确度指标，从而得出性能分布。在药

物再利用数据集的情况下，杰卡德相似度被用作性能指标。

图 3 显示了三个 LLM 在 True/False（图 3A）和 MCQ（图 3B）数据集上使用基于提示和 KG-

RAG 框架的性能（准确性）的引导分布。 

表 2 总结了三个LLM在这些数据集上的表现。我们观察到 KG-RAG 框架下的 LLM 模型在

True/False 和 MCQ 数据集上都有一致的性能增强（表 2）。KG-RAG 在更多方面将 Llama-2 的

性能从初始水平（0.31±0.03 至 0.53±0.03）显着提高了约 71%，具有挑战性的 MCQ 数据集（表

2）。有趣的是，我们还观察到，在不基于提示的方法中，与使用 KG-RAG 框架的 MCQ 数据集上

的 GPT-3.5-Turbo 模型 (0.79±0.02) 相比，GPT-4 模型 (0.74±0.03) 的性能有小幅下降（T 检验，p

-value < 0.0001，t-statistic = -47.7，N = 1000）。

图 3. True/False 和 MCQ 数据集上的 LLM 表现 LLM 在 (A) True/False 和 (B) MCQ 数据集上的表

现（准确度）分布。蓝色分布表示使用基于提示的方法的性能，红色分布表示使用基于 KG-RAG

的方法的性能。黑色垂直虚线表示分布的平均值。 

表 2. True/False 和 MCQ 数据集上的 LLM 表现（准确度：平均值±标准差）