PaperlayerView exact revision on arXiv

Source-linked AI summary

Biomedical knowledge graph-optimized prompt generation for large language models

Karthik Soman, Peter W Rose, John H Morris, Rabia E Akbas, Brett Smith, Braian Peetoom, Catalina Villouta-Reyes, Gabriel Cerono, Yongmei Shi, Angela Rizk-Jackson, Sharat Israni, Charlotte A Nelson, Sui Huang, Sergio E Baranzini

arXiv:2311.17330v2cs.CL

TL;DR

Knowledge-intensive tasks require provenance and up-to-date knowledge, while building from scratch carries substantial computational cost and time overhead. The paper proposes KG-RAG, a robust, token-optimized framework. It reports superior performance to prompt-based baseline LLMs across human-curated benchmarks, with Llama-2 improving ~71% from baseline.

  • Problem

    Knowledge-intensive tasks require provenance and up-to-date knowledge, while building from scratch carries substantial computational cost and time overhead.

    Source: Introduction · pp. 3–4

  • Method

    KG-RAG is proposed as a robust, token-optimized framework.

    Source: Introduction · p. 4

  • Results

    KG-RAG outperformed prompt-based baseline LLMs across the evaluated human-curated benchmark datasets, with Llama-2 improving ~71% from baseline.

    Source: Discussion · pp. 12, 14

  • Takeaways & Limitations

    The framework is presented as a cost-effective way to generate from a knowledge graph.

    Source: Results · p. 5

    As the knowledge graph expands, it may affect prompt flexibility and adaptability.

    Source: Discussion · p. 13

Abstract

from arXiv · show

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.

Introduction

Biomedical LLMs can produce non-factual answers, while pretraining is costly and prompt tuning is limited on knowledge-intensive tasks requiring provenance and up-to-date knowledge. The paper introduces KG-RAG, which retrieves prompt-aware biomedical context from SPOKE for pretrained LLMs.

  • LLM hallucinations pose a challenge in biomedicine, where answers require stringent accuracy.Source: Introduction · p. 3
  • Pretraining carries substantial time and computational overhead, while prompt tuning is restricted on tasks requiring provenance and up-to-date knowledge.Source: Introduction · pp. 3–4
  • Earlier KG-language-model integration used graph context primarily for question answering, and task-specific downstream supervision limited versatility and broad applicability.Source: Introduction · p. 4
  • KG-RAG integrates a biomedical KG with a pretrained LLM and retrieves essential prompt-aware context from SPOKE to enrich prompts with provenance and statistical evidence.Source: Introduction · pp. 4–5
  • The framework was evaluated with pretrained LLMs including Llama-2-13b and GPT-3.5-Turbo.Source: Introduction · p. 5

Results

KG-RAG retrieves and prunes biomedical graph context for LLM prompts, then benchmarks show improved question-answering performance, robust retrieval, and reduced token use. The results also compare context-embedding models and illustrate answers with supporting evidence and provenance.

Results

  • The KG-RAG pipeline includes entity recognition, SPOKE concept extraction, concept embedding, prompt-aware context generation, natural-language conversion, prompt assembly, and answer retrieval.Source: Results · p. 5

Prompting KG-RAG framework

RAG comparative analysis

Performance on True/False and MCQ datasets

Discussion

KG-RAG combines biomedical graph knowledge with LLMs to generate grounded responses while using minimal tokens. It improves performance across biomedical tasks, but its current implementation is focused on disease questions and SPOKE.

Discussion

  • KG-RAG outperformed prompt-based LLM baselines on all evaluated tasks in specialized domains.The discussion also characterizes it as more robust to prompt perturbation and more efficient in token use.Source: Discussion · p. 12
  • KG-RAG used over 50% fewer tokens than Cypher-RAG, which averaged 8006 tokens and includes the full graph schema in its prompt.SPOKE contains over 40 million nodes, and the authors note Cypher-RAG requires models with a larger context window for such a graph.Source: Discussion · p. 13
  • Under KG-RAG, GPT-3.5-Turbo performed on par with GPT-4 on True/False datasets, while GPT-4’s performance on the MCQ dataset dropped significantly compared with GPT-3.5-Turbo under the same framework.Source: Discussion · p. 14
  • ~71%: KG-RAG improved Llama-2 performance from baseline on the challenging MCQ dataset, though it remained below GPT models.The reported improvement narrowed, but did not close, the performance gap.Source: Discussion · p. 14
  • The current approach handles disease-centered biomedical questions because entity recognition focuses on disease concepts in SPOKE.The authors propose including all SPOKE biomedical concepts and other knowledge graphs to broaden the question scope.Source: Discussion · p. 15
  • The framework is implemented on SPOKE, and the work evaluates the framework rather than formally evaluating the knowledge graph itself.The authors also note that retrieved context depends on graph information and do not claim SPOKE is error-free or ready for clinical use.Source: Discussion · p. 15
  • KG-RAG retrieves semantically meaningful graph context with minimal tokens and combines it with an LLM’s implicit knowledge to generate biomedical responses.The responses include provenance and statistical evidence when available.Source: Discussion · pp. 12, 15, 16

RAG) Framework

The section describes knowledge-graph retrieval for LLM text generation and its evaluation, covering KG-RAG’s disease identification, context retrieval and pruning, and Cypher-RAG’s schema-driven retrieval, datasets, and token-use and prompt-sensitivity limitations.

Disease entity recognition

  • The framework used zero-shot prompting with GPT-3.5-Turbo and matched diseases to SPOKE concepts using precomputed all-MiniLM-L6-v2 embeddings, retrieving disease nodes with 99.7% accuracy.Source: Disease entity recognition · p. 17

Disease context retrieval from SPOKE

  • KG-RAG fetches disease-neighbor triples from SPOKE, converts them to English, and appends edge provenance to the context.Optional edge evidence, such as p-values and z-scores, can also be incorporated through the -e option.Source: Disease context retrieval from SPOKE · p. 19

Context Pruning

  • Context pruning embeds the prompt and candidate associations in a shared vector space, retaining associations above both the 75th-percentile similarity and 0.5 minimum.The retained contexts have the highest cosine similarity to the input prompt.Source: Context Pruning · p. 20

Large Language Model (LLM)

  • The enriched prompt is sent to Llama-2-13b, GPT-3.5-Turbo, or GPT-4 for text generation, with temperature set to 0 for all models.Llama-2-13b was deployed on an EC2 GPU instance, while GPT models were accessed through the OpenAI API.Source: Large Language Model (LLM) · pp. 20–21

Hyperparameter analysis and validation

  • Hyperparameter validation evaluated KG-RAG on 165 questions: 75 prompts mentioned one disease and 90 mentioned two diseases, while the evaluated hyperparameters were context volume and context embedding model.Source: Hyperparameter analysis and validation · p. 21
  • MiniLM and PubMedBERT were candidate context embedding models, and context volume set the upper limit on graph connections passed to the LLM.PubMedBERT was considered because it was pretrained on biomedical text.Source: Hyperparameter analysis and validation · p. 21

Test dataset

  • Evaluation used true/false, MCQ, and RAG-comparison test sets; the first two were compiled from external resources and reviewed by domain experts.The true/false dataset drew on DisGeNET, MONDO, and SemMedDB.Source: Test dataset · p. 22
  • The evaluation included 311 True/False, 306 MCQ, and 100 graph-context questions; accuracy was estimated from 1,000 bootstrap samples of 150 questions for the first two datasets.The graph-context questions were designed from Disease-Gene associations in SPOKE to compare KG-RAG and Cypher-RAG retrieval.Source: Test dataset · pp. 22–23

Cypher-RAG

  • Cypher-RAG (Bratanic 2023) embeds graph schema in the prompt, generates a Cypher query for Neo4j retrieval, and uses returned information as LLM context.Source: Cypher-RAG · p. 23
  • Explicit schema embedding restricts prompt token space and increases token use; more complex schemas may require LLMs with longer context windows.Source: Cypher-RAG · p. 23
  • Cypher-RAG is sensitive to prompt formulation: slight perturbations can produce incorrect queries and affect downstream generation.Source: Cypher-RAG · p. 23

Data availability

SPOKE can be accessed through its neighborhood website or REST API. The KG-RAG code and biomedical benchmark datasets are available in the same GitHub repository.

Data availability

Supporting information

The supporting information describes supplementary methods and materials, including two supplementary tables.

Supporting information

Loading 2311.17330v2…