Biomedical knowledge graph-optimized prompt generation for large language models (2311.17330v2)

Abstract: 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.

Overview of Biomedical Knowledge Graph-Optimized Prompt Generation for LLMs

This paper addresses the challenge of integrating domain-specific knowledge into LLMs for improved accuracy and relevance within the biomedical field. LLMs like GPT-3.5, GPT-4, and Llama-2, while powerful, often experience difficulties when navigating the complexities of knowledge-intensive domains such as biomedicine. The authors propose a novel framework, Knowledge Graph-based Retrieval Augmented Generation (KG-RAG), utilizing a comprehensive knowledge graph named SPOKE to efficiently supplement LLMs with authoritative biomedical context.

KG-RAG Framework

KG-RAG employs a series of steps to enhance LLM outputs using SPOKE, which aggregates data from over 40 biomedical sources. The framework involves recognizing entities from input prompts, extracting relevant biomedical concepts from SPOKE, embedding these concepts, and optimizing context retrieval. This integration results in a significant reduction of token usage—over 50% less compared to conventional RAG methods—without compromising accuracy.

Numerical Results

The authors report substantial improvements in model performance, notably achieving a 71% increase in accuracy on a challenging multiple-choice question (MCQ) dataset for Llama-2. Additionally, both GPT-3.5 and GPT-4 showed enhanced prompt responses when supported by KG-RAG. These results underscore the framework's efficacy in reinforcing open-source models and proprietary systems alike.

Comparative Analysis

KG-RAG's robustness and efficiency were compared against the existing Cypher-RAG technique. While Cypher-RAG struggled, particularly with slight input perturbations leading to drastic retrieval failures (0% accuracy), KG-RAG maintained a consistent retrieval accuracy of 97%. This confirms KG-RAG's resilience and cost-effectiveness due to its optimized context retrieval.

Implications and Future Directions

This research demonstrates how explicit knowledge from structured databases, like SPOKE, can be efficiently integrated with LLMs, enhancing their capability to handle specialized queries in biomedicine. The practice of minimizing token consumption while maximizing contextual relevancy is critical, especially given the token limitations inherent in many LLMs.

Future research could expand this framework's applicability to broader sets of biomedical entities or to other specialized domains by incorporating additional knowledge graphs. This would likely enhance the versatility of the model further. Moreover, as LLMs evolve, the integration of knowledge graphs could continue to mitigate the risks of “hallucinations” by grounding outputs in verifiable data.

Conclusion

The KG-RAG framework represents a significant step in the adaptation of LLMs to domain-specific tasks, enabling them to deliver well-substantiated, accurate, and relevant responses. By merging the implicit learning of LLMs with the explicit, structured knowledge of KGs, this approach paves the way for more reliable applications in fields that demand high accuracy, such as biomedicine.