Biomedical knowledge graph-enhanced prompt generation for large language models/生物医学知识图增强的大型语言模型提示生成

摘要

大型语言模型（LLMs）以前所未有的速度推动了人工智能的进步，但在生物医学等知识密集领域仍然面临挑战。诸如预训练和领域特定微调的解决方案增加了大量的计算开销，后者需要领域专业知识。外部知识的注入是任务特定的，并且需要模型训练。在这里，我们介绍了一种基于任务的不可知的知识图检索增强生成（KG-RAG）框架，通过利用庞大的生物医学KG SPOKE与LLMs（如Llama-2-13b、GPT-3.5-Turbo和GPT-4）相结合，生成根植于已建立知识的有意义的生物医学文本。KG-RAG持续增强了LLMs在各种提示类型上的性能，包括一跳和两跳提示、药物重定位查询、生物医学真/假问题和多项选择题（MCQ）。值得注意的是，KG-RAG在具有挑战性的MCQ数据集上为Llama-2模型的性能提供了显著的71%提升，表明该框架能够为具有较少参数的开源模型赋予领域特定问题的能力。此外，KG-RAG提高了专有GPT模型的性能，例如GPT-3.5在MCQ数据上的上下文利用方面表现出对GPT-4的改进。我们的方法还能够解决药物重定位问题，提供有意义的重定位建议。总之，所提出的框架以优化的方式将KG和LLM的显式和隐式知识相结合，从而增强了通用型LLM应对领域特定问题的适应性，形成了统一的框架。

引言

大型语言模型（LLM）在解决涉及语言建模和处理的各种复杂任务中表现出令人印象深刻的性能[1]。LLMs 在自监督学习框架下对大量文本数据进行了预训练，该框架可以是掩码语言建模（例如 BERT 类模型[2,3]）或自回归框架（类似 GPT 的模型[4,5]）。这种预训练将有关语言的知识编码到模型参数中。类似于深度神经网络中常用的迁移学习方法，通过监督训练可以优化这种隐式知识，以在一系列领域特定任务中表现出色[5,6]。然而，LLM 中的知识的“隐式表示”已被证明会生成非事实信息，尽管语言上连贯（即“幻觉”），作为对输入提示的响应[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）与语言模型集成在一起[19–25]。这种方法主要在问答任务中实现，其中KG中包含的结构化信息用于提供预测答案的上下文。虽然这种多模态整合方法显示出了潜力，但其下游监督训练是针对特定任务定制的，限制了其通用性和广泛适用性，从而限制了其充分发挥LLMs的“新兴能力”的能力[26]。为了增强LLM的上下文意识，最近引入了一种称为检索增强生成（RAG）的方法[27]。这种方法涉及增强参数预训练的LLM的能力，使其能够访问一个包含关于世界的更新知识的非参数化内存（例如，维基百科）[27]。

在本文中，我们提出了一种方法，将知识图谱与预训练的LLM集成在RAG框架中，从而兼具两者的优点。为了实现这一目标，我们利用了一种名为可扩展精准医学开放知识引擎（SPOKE）的生物医学知识图谱[28]。SPOKE已整合了40多个公共生物医学知识源，涉及单独的领域，每个知识源围绕生物医学概念（如基因、蛋白质、药物、化合物、疾病等）中心，以及一个或多个已知关系。由于这些概念是形成定义集合的经常出现的实体（例如，所有命名的人类基因、所有FDA批准的药物等），将这些概念集成到一个图中，暴露了连接知识源的新型多跳事实关系，并为每个概念提供了生物和本体上下文。我们提出的框架熟练地从SPOKE中检索生物医学上下文，优化其提取以充分满足提示，从而创建我们称之为“提示感知上下文”的内容。这种丰富的提示进一步被用作RAG框架中LLM的输入，用于生成有意义的生物医学文本。我们使用了各种预训练的LLM，包括Llama-2-13b、GPT-3.5-Turbo和GPT-4来评估这种方法。通过将系统提示变化为LLM，我们能够将同一框架用于各种生物医学任务，如真/假和多项选择题回答、药物再利用建议以及与疾病相关的一般生物医学查询。这种适应性突显了我们方法的任务不可知性。进一步进行了系统的比较分析，以评估Llama-2-13b、GPT-3.5-Turbo和GPT-4在该框架内的性能。因此，我们的方法展示了处理各种领域特定的生物医学任务的能力，而无需对通用型LLM进行额外的预训练或微调。这种方法不仅简化了流程，还使我们的方法具有成本效益和时间效率。

结果

在这里，我们开发了一个框架，用于为大型语言模型构建准确可信的生物医学提示。这个KG-RAG框架涉及从用户提示中识别实体，从知识图谱中提取生物医学概念，提示感知上下文生成，转换成自然语言，提示组装和答案检索。这种方法的性能在不同的场景下进行了广泛测试。

KG-RAG框架的提示

图1显示了一个单跳提示（图1A黄色框）和两跳提示（图1B黄色框）作为输入提供给GPT-4模型的两种方法：（i）基于提示，即不使用KG-RAG（蓝色框），和（ii）使用KG-RAG（绿色框）。我们观察到只有KG-RAG能够回答问题，并为单跳提示（图1A黄色框）和两跳提示（图1B黄色框）提供简洁明了的答案。（更多提示示例请参见S1文本中的表A。）

图1. 提示示例 示例提示（黄色框）和基于提示（蓝色框）/基于KG-RAG（绿色框）的GPT-4响应。 （A）显示了用于搜索“Bardet-Biedl综合征”患者的体重管理药物的单跳提示，（B）显示了用于搜索“莱姆病”和“斑秃”共享的基因的两跳提示。 （A）和（B）显示了相关的“提示感知图”上下文。 图中的黑色方框表示提示中提到的疾病节点/节点。 注意：由于空间限制，（A）中基于提示的GPT-4响应（蓝色框）被截断。

超参数分析

我们优化了两个超参数，即“上下文容量”和“上下文嵌入模型”。 上下文容量定义了从KG到LLM允许流动的图连接数量的上限（请参见材料和方法，S1文本）。 上下文嵌入模型是一种在从KG中检索准确的生物医学上下文以回答输入提示方面表现更好的模型（请参见材料和方法，S1文本）。 为了优化这些超参数，我们使用了两个上下文嵌入模型（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，并且对于真/假问题，我们选择了上下文容量为100）。

**图2. 超参数分析和药物再利用性能（A）**使用单跳（顶部）和两跳（底部）提示的超参数分析性能曲线。 x轴表示“上下文容量”（来自KG的关联数量），y轴表示跨提示的平均性能（Jaccard相似性）。红色曲线表示“S-PubMedBert-MSMARCO”，蓝色曲线表示“all-MiniLM-L6-v2”上下文嵌入模型。 （B）LLMs在药物再利用问题上的平均性能（Jaccard相似性）（顶部）。误差棒表示95％的CI。底部面板显示了COVID-19相关的药物再利用提示的示例（黄色框）以及GPT-4模型使用KG-RAG框架的相应响应（绿色框）。 注意：由于空间限制，绿色框中的响应已被截断。

药物再利用性能

图2B显示了三种LLM模型（Llama-2-13b、GPT-3.5-Turbo和GPT-4）的药物再利用性能分析。 我们观察到，从Llama-2-13b到GPT4，性能指标（Jaccard相似性）持续增加，并且有趣的是，这种顺序与这些模型的参数大小的顺序一致（图2B，表1）。 KG-RAG框架能够从SPOKE提供必要的上下文来回答与COVID-19相关的药物再利用问题（图2B底部面板黄色和绿色框）。 除了提供可操作的再利用建议外，它还对评估这些建议的有效性进行了慎重的警告，指出了文临床试验的必要性。 这表明了一种负责任和安全意识的文本生成方法（图2B底部面板绿色框）。

表1. 使用KG-RAG框架进行药物再利用提示的模型性能（与来自KG的真实值的Jaccard相似性）

真/假和MCQ数据集上的性能

图3显示了三种LLM在真/假（图3A）和MCQ（图3B）数据集上使用基于提示和KG-RAG框架的性能（准确率）的bootstrap分布。表2总结了这些数据集上三种LLM的性能。我们观察到在KG-RAG框架下，LLM模型在真/假和MCQ数据集上都表现出了一致的性能提升（表2）。 KG-RAG显着提升了Llama-2模型的性能，从其初始水平（0.31±0.03）提高约71%到更具挑战性的MCQ数据集上的水平（0.53±0.03）（表2）。有趣的是，我们还观察到，在MCQ数据集上，相较于GPT-3.5-Turbo模型（0.79±0.02），使用KG-RAG框架的GPT-4模型的性能略有但显著下降（0.74±0.03）（T检验，p值<0.0001，t统计量= -47.7，N = 1000），但在基于提示的方法中没有观察到这种情况。

图3. LLM在真/假和MCQ数据集上的性能（准确率）图示
LLM在（A）真/假和（B）MCQ数据集上的性能（准确率）分布。蓝色分布表示基于提示的方法的性能，红色分布表示基于KG-RAG的方法的性能。黑色垂直虚线表示分布的平均值。数值越高，性能越好。

表2. LLM在真/假和MCQ数据集上的性能（准确率：均值±标准差）

讨论

在这项工作中，我们介绍了一个简单但高效的框架，将生物医学知识图谱与LLM聊天模型结合起来，这里实现为KG-RAG框架。这种整合导致了一个特定领域的生成系统，其响应牢固地基于已建立的生物医学知识。我们使用各种类别的生物医学问题评估了该框架，涵盖了单跳、更复杂的两跳查询、药物再利用查询、真/假查询和多项选择题。在所有测试场景中，所提出的KG-RAG框架始终表现出比基于提示的基线LLM模型更优越的性能，即没有使用KG-RAG。我们相信，这种性能改进源于KG的显式知识与LLM的隐式知识的融合。一个具有各种概念（在本例中为生物医学概念）的异构知识图谱，在大规模互连的情况下，有可能产生新的知识作为“新兴能力”。实际上，随着LLM在模型参数、训练数据和训练计算等各个维度的扩展，它们被认为表现出推理或“新兴能力”，尽管这一观察结果可能更多地可解释为“在上下文学习”或其他示例的方面。无论如何，在这里，我们展示了在提示级别提供特定领域地面实况背景的价值。具体来说，我们假设一个整合了KG和LLM的框架将产生协同效应，从而产生更可靠和基于知识的响应。

先前的研究已经将KG与LLM结合在一起，用于诸如问答[23]、多跳关系推理[19]、常识推理[20,24]甚至模型预训练[25,32]等知识密集型任务。此外，通过整合结构化知识来增强提示已被描述和研究过。这些有前途的方法已经增强了KG和LLM之间的积极强化。然而，值得注意的是，这些方法可能是任务特定的，并且可能需要额外的模型训练。在某些情况下，通过包含更高阶关系来增加知识可能会呈指数增长。这种方法可能会损害LLM的有限标记空间。替代方法通过直接使用诸如SPARQL之类的查询语言进行知识注入。然而，这可能会使系统受到底层KG架构的约束，可能会影响提示的灵活性和适应性。此外，随着KG的扩展和其模式的增长，它可能会占据LLM输入标记空间的重要部分。

所提出的KG-RAG方法是一个任务无关的框架，适用于广泛的生物医学提示类型，从真/假问题到药物再利用查询。它的特点是具有最少的超参数，并且在操作过程中不需要任何训练数据。通过优化的提示感知上下文提取，它有选择地包含了来自KG的相关关联，有效利用了输入标记空间。此外，我们的方法与图查询语言独立（或最多只是轻微依赖），从而增强了提示的灵活性。因此，所提出的KG-RAG框架通过具有任务无关性，针对生物医学领域（由于SPOKE KG），并针对标记使用进行了优化，解决了上述限制。

在我们的分析中，我们发现LLM性能随着模型大小（即参数数量）的增加而增强。有趣的是，使用KG-RAG框架时，尽管GPT-4的模型大小，其在MCQ数据集上的性能与GPT-3.5-Turbo相比显著下降。事实上，在KG-RAG框架下，GPT-3.5-Turbo的性能与GPT-4模型在真/假数据集上的性能相当。这些结果表明，目前而言，GPT-3.5可能比GPT-4更好地理解上下文。实际上，一项最近的研究比较了GPT-4的2023年3月版和2023年6月版，揭示了LLM性能随时间漂移的情况。该研究显示，随着时间的推移，GPT-4表现出对用户指令的遵循程度下降。相比之下，GPT-3.5的指令遵循行为随时间变化并没有保持一致。

当专有的GPT模型与开源的Llama-2-13b模型进行比较时，在生物医学真/假数据集上，它们的性能相似。然而，在更具挑战性的MCQ数据集上，Llama-2最初的表现要低于GPT模型。有趣的是，KG-RAG框架为Llama-2提供了显著的性能提升，从基线提高了约71%的性能。尽管这种提升缩小了性能差距，但Llama-2的性能仍然低于GPT模型。这表明，KG-RAG框架有潜力发挥开源预训练模型（如Llama-2）的固有上下文理解能力，使其在与GPT等专有模型的竞争中更具竞争力。这些发现突显了在改进语言模型在专业领域复杂任务上的性能方面，上下文丰富技术的重要性。

尽管我们在测试药物再利用方案时采用了简单的方法，但所获得的结果是令人鼓舞的。例如，使用KG-RAG的GPT-4能够提出先前被认为是基于机制合理性的COVID-19再利用的药物，例如氯喹、托珠单抗、阿那基那、利托那韦、头孢三嗪、克拉霉素、普拉双贝等。进一步的系统分析显示KG-RAG框架在其他疾病上表现良好。药物再利用的全面验证方法超出了本研究的范围，但我们希望强调，在RAG框架中将生物医学KG和LLM相结合，可以在其响应中提供医学上合理和可测试的见解。

虽然所提出的框架成功解决了许多挑战，但也存在改进的机会。目前，这种方法仅限于处理围绕疾病的生物医学问题，因为我们在从提示中识别生物医学实体时专注于将疾病概念嵌入SPOKE KG。未来的工作可以通过将SPOKE KG中的所有生物医学概念（节点）包括在内来扩大这个范围。由于SPOKE包含超过2500万个生物医学节点，这种扩展将使KG-RAG框架能够处理更广泛的生物医学问题，从而增强其多功能性。

材料和方法

基于知识图谱的检索增强生成（KGRAG）框架
所提出的KG-RAG框架的模式如图4所示。以下各节将解释该框架的每个组成部分。

图4. 基于知识图谱的检索增强生成（KGRAG）框架的模式。 箭头的方向表示该框架中流程的流向。

疾病实体识别

这是KG-RAG中的第一步。该步骤的目标是从输入文本提示中提取疾病概念（一个实体），然后在知识图谱中找到相应的匹配疾病节点（SPOKE概念）。这是作为一个两步过程实现的：i）从提示中提取实体，ii）将实体与SPOKE匹配。实体提取识别和提取输入文本提示中提到的疾病实体，也称为“提示疾病提取”（图5）。为了实现这一点，使用了GPT-3.5-Turbo模型上的零样本提示方法。具体来说，设计了一个系统提示来从输入文本中提取疾病实体，并以JSON格式返回它们。

接下来，使用实体匹配来获取疾病概念在知识图谱中的概念名称。为此，使用了“all-MiniLM-L6-v2”句子转换模型（即图5中的疾病潜在空间）对SPOKE中的所有疾病概念（即节点）的嵌入进行了预计算。这个过程将疾病概念的名称转换为一个384维的稠密向量空间，使其适用于语义搜索。接下来，这些新创建的疾病概念嵌入被存储在“Chroma”向量数据库中。选择与提取实体具有最高向量相似性的疾病概念进行后续的上下文检索。如果零样本方法未能在提示中识别出疾病实体，则从具有与整个输入文本提示具有最显著向量相似性的五个疾病概念中选择。

**图5. KG-RAG的详细模式。**虚线框显示了从知识图谱中的“疾病实体识别”和“提示感知上下文提取”中的详细信息。

欢迎关注：https://blog.csdn.net/hanhanwanghaha

从SPOKE中检索疾病上下文

SPOKE是一个异构的生物医学知识图谱，通过语义相关关系连接数百万个生物医学概念。SPOKE包括了21种不同类型的2700万个节点和55种不同类型的5300万条边，实现为属性图。SPOKE是通过从41个不同的生物医学数据库下载和整合信息而组装而成。值得注意的是，SPOKE的绝大部分由系统性实验测量确定的筛选信息组成，几乎没有来自文献的贡献，因此可以被视为经验性、事实性的生物医学知识的语料库。在本研究中，SPOKE被用作输入提示中提到的疾病的生物医学背景的来源。简而言之，从SPOKE KG中提取与疾病节点相关联的上下文三元组（主体、谓词、客体），并随后将其转换为英文，使其适用于输入到LLM中。

upperCase(predicateName)_<upperCase(firstLetter(subjectType)),
lowerCase(firstLetter(predicateName)),
upperCase(firstLetter(objectType))>

该模式允许使用以下规则将提取的三元组直接转换为英语（见S1文本）：

(S, P, O) → Subject lowerCase(predicateName) Object
For example:
(Disease hypertension, ASSOCIATES_DaG, Gene VHL) → Disease hypertension associates Gene VHL

上下文修剪

接下来，通过选择必要的最相关的上下文，对提取的疾病上下文进行修剪，以便回答给定的提示，如下所述。首先，使用句子转换模型将输入提示和所有提取的上下文关联嵌入到相同的向量空间中（图5中的上下文潜空间），并使用超参数调整进行模型选择。接下来，仅选择与输入提示向量具有最高余弦相似度的上下文关联（图5）。对于上下文选择，提示 - 上下文余弦相似度应满足两个条件：（i）大于包含所有与所选疾病节点相关的上下文的相似度分布的第75百分位数，并且（ii）具有最小相似度值为0.5。

大型语言模型（LLM）

当将输入提示与提示感知上下文结合在一起时，产生了一个丰富的提示，用作输入到LLM进行文本生成。为此，使用了三个预训练的聊天模型：Llama-2-13b，GPT-3.5-Turbo和GPT-4。具有130亿参数和4096个标记大小的Llama模型被下载并部署在Amazon Elastic Compute Cloud（EC2）GPU P3实例中。通过OpenAI API访问GPT模型。由于GPT模型的参数数量比Llama更高，这为我们提供了比较KG-RAG性能作为LLM大小函数的机会，以其参数数量为准。在这项研究中，对于所有LLM，设置了控制LLM输出中随机性水平的“温度”参数为0。

超参数分析和验证

评估了KG-RAG的性能，跨两组超参数，如“上下文体积”和“上下文嵌入模型”。上下文体积定义了允许从KG流向LLM的图连接数量的上限。这个超参数介绍了上下文增强和LLM输入标记空间利用之间的平衡。认识到生物医学上下文通常使用的词汇可能与一般领域的场景不同，引入了“上下文嵌入模型”超参数。该超参数确定了哪个模型在检索准确的生物医学上下文以响应输入提示方面表现更好。为此任务使用了“all-MiniLM-L6-v2”和“S-PubMedBert-MS-MARCO”模型。此外，使用单跳（75个问题）和双跳提示（90个问题）创建了两组验证数据（共165个问题）。单跳提示形成与单跳图相关的上下文，而双跳提示利用双跳图作为其上下文。这些提示使用GPT-4模型在KG-RAG框架中执行，利用特定的系统提示以JSON格式返回结果。通过解析这些JSON响应并将其与基本真相进行比较，计算了Jaccard相似度。通过这些分析，为下游任务进行了超参数的经验选择。

测试数据集

采用三种类型的测试数据集对所提出的框架的性能进行定量分析（见S1文本中的表B）：（i）真/假数据集；（ii）多项选择题（MCQ）数据集；以及（iii）药物再利用数据集。真/假数据集是从三个外部数据源创建的，包括DisGeNET、MONDO和SemMedDB。DisGeNET整合了与人类疾病相关的基因和基因变异的数据，这些数据来自筛选存储库、GWAS目录、动物模型和科学文献。MONDO提供了关于开放生物医学本体（OBO）格式中疾病实体的本体分类信息。SemMedDB包含从PubMed引文中提取的语义断言，并且我们使用该资源来制定关于药物和疾病的真/假问题。MCQ包含了每个问题有五个选项，其中有一个是正确答案，这些问题是使用来自Monarch Initiative和ROBOKOP的数据创建的。从SPOKE KG中提取了疾病 - 化合物关联来创建测试LLM在KG-RAG框架中的药物再利用能力的问题。药物再利用问题的制定采用了一个简单的策略，即如果一种疾病类似于具有已知治疗的另一种疾病，则识别可以用于治疗该疾病的化合物。因此，针对所提出的框架进行了系统定量分析，创建了322个真/假问题，306个MCQ问题和68个药物再利用的生物医学问题数据集。为了评估LLM在真/假和MCQ数据集上的性能，随机采样了150个问题，并进行了1000次（使用自助法）重复采样。然后计算了每个采样迭代的准确性指标，得出了性能分布。在药物再利用数据集的情况下，使用Jaccard相似性作为性能指标。

致谢

SEB在UCSF担任Heidrich家族和朋友资助的神经病学教授。SEB在UCSF担任神经病学一级教授。我们还感谢UCSF的Versa团队提供访问OpenAI API的基础设施。

数据可用性

SPOKE KG可通过https://spoke.rbvi.ucsf.edu/neighborhood.html进行访问。也可以使用REST-API（https://spoke.rbvi.ucsf.edu/swagger/）进行访问。KG-RAG代码可以在https://github.com/BaranziniLab/KG_RAG获取。本研究中使用的生物医学数据集（单跳问题、双跳问题、真/假问题、MCQ问题、药物再利用问题、来自SPOKE KG的疾病上下文）也可以在同一GitHub存储库中提供给研究界。

原文地址：https://arxiv.org/ftp/arxiv/papers/2311/2311.17330.pdf
开源代码：https://github.com/BaranziniLab/KG_RAG
这个项目的环境有点复杂。