[论文解读] Biomedical knowledge graph-optimized prompt generation for large language models

Karthik Soman, Peter W. Rose+12 more|arXiv (Cornell University)|Nov 29, 2023

Topic Modeling被引用 6

一句话总结

本论文提出了一种基于知识图谱的检索增强生成（KG-RAG）框架，通过优化标记（token）使用，利用 SPOKE 知识图谱增强生物医学大语言模型。通过最小化图谱模式并采用基于嵌入的上下文剪枝方法，该框架将标记消耗减少超过 50%，同时提升性能——在 Llama-2-13b 上，于一项具有挑战性的多项选择生物医学数据集上实现 71% 的性能提升，证明了其在开源和专有大语言模型上具备成本效益高、准确且具备溯源能力的生成能力。

ABSTRACT

Large Language Models (LLMs) are being adopted at an unprecedented rate, yet still face challenges in knowledge-intensive domains like biomedicine. Solutions such as pre-training and domain-specific fine-tuning add substantial computational overhead, requiring further domain expertise. Here, we introduce a token-optimized and robust Knowledge Graph-based Retrieval Augmented Generation (KG-RAG) framework by leveraging a massive biomedical KG (SPOKE) with LLMs such as Llama-2-13b, GPT-3.5-Turbo and GPT-4, to generate meaningful biomedical text rooted in established knowledge. Compared to the existing RAG technique for Knowledge Graphs, the proposed method utilizes minimal graph schema for context extraction and uses embedding methods for context pruning. This optimization in context extraction results in more than 50% reduction in token consumption without compromising the accuracy, making a cost-effective and robust RAG implementation on proprietary LLMs. KG-RAG consistently enhanced the performance of LLMs across diverse biomedical prompts by generating responses rooted in established knowledge, accompanied by accurate provenance and statistical evidence (if available) to substantiate the claims. Further benchmarking on human curated datasets, such as biomedical true/false and multiple-choice questions (MCQ), showed a remarkable 71% boost in the performance of the Llama-2 model on the challenging MCQ dataset, demonstrating the framework's capacity to empower open-source models with fewer parameters for domain specific questions. Furthermore, KG-RAG enhanced the performance of proprietary GPT models, such as GPT-3.5 and GPT-4. In summary, the proposed framework combines explicit and implicit knowledge of KG and LLM in a token optimized fashion, thus enhancing the adaptability of general-purpose LLMs to tackle domain-specific questions in a cost-effective fashion.

研究动机与目标

• 解决在生物医学领域微调大语言模型所面临的高计算成本和领域专业知识成本问题。
• 在不牺牲生物医学推理事实准确性的前提下，减少 RAG 管道中的标记消耗。
• 使开源大语言模型（如 Llama-2-13b）在特定生物医学任务上达到或超越专有模型（如 GPT-3.5、GPT-4）的性能。
• 以稳健且成本效益高的方式，将生物医学知识图谱中的显性与隐性知识整合到大语言模型中。
• 为生物医学语境下大语言模型生成的主张提供准确的溯源信息和统计证据。

提出的方法

• 该框架将 SPOKE 生物医学知识图谱作为事实性支撑的检索源。
• 采用最小图谱模式提取上下文，以减少无关或冗余信息。
• 使用基于嵌入的方法剪枝低相关性上下文，以优化标记使用。
• 将检索到的知识以结构化格式整合到提示中，以引导大语言模型的推理过程。
• 该方法支持开源（Llama-2-13b）和专有大语言模型（GPT-3.5-Turbo、GPT-4）。
• 自动将溯源信息和统计证据与生成响应中的主张关联。

实验结果

研究问题

• RQ1 基于知识图谱优化的 RAG 框架是否能在不牺牲准确性的前提下，减少生物医学大语言模型提示中的标记消耗？
• RQ2 此类框架在多大程度上能够提升开源大语言模型在特定生物医学基准测试中的性能？
• RQ3 将生物医学知识图谱中的显性和隐性知识整合到大语言模型中，如何提升输出的事实一致性和溯源能力？
• RQ4 采用最小图谱模式与基于嵌入的剪枝方法，是否能在降低计算成本的同时保持或提升性能？
• RQ5 在多种生物医学提示下，该框架与标准 RAG 方法相比，在鲁棒性和效率方面表现如何？

主要发现

• 与标准 RAG 技术相比，KG-RAG 框架将标记消耗减少了超过 50%，且未影响准确性。
• 在一项具有挑战性的生物医学多项选择题（MCQ）数据集上，使用 KG-RAG 的 Llama-2-13b 实现了 71% 的性能提升。
• 该框架在多种生物医学提示下，显著提升了专有模型（包括 GPT-3.5-Turbo 和 GPT-4）的性能。
• 生成的响应中包含了准确的溯源信息和可用的统计证据，从而提升了可信度。
• 该方法在不同大语言模型上表现出鲁棒性，支持在专有和开源模型上实现成本效益高的部署。
• 将 SPOKE 中的显性和隐性知识整合后，显著提升了大语言模型输出的事实性基础和推理能力。