Back to all posts
AI & Data

BioGraphRAG — Biomedical Knowledge Graph Retrieval Augmented Generation

October 25, 2024 · 28 min read · Originally published on Kunj's Substack

This is a summary. Read the full article on Substack.

BioGraphRAG combines retrieval-augmented generation with an expert-curated biomedical knowledge graph to answer biomedical questions with explainable, source-cited responses. LLMs alone suffer from outdated training data, no access to proprietary datasets, and hallucination; vanilla vector RAG helps but struggles with the intricate entity relationships that dominate biomedical knowledge. Graph-based RAG addresses both by retrieving structured subgraphs instead of flat text chunks.

System architecture

The system is a five-stage pipeline: users submit queries through a Chainlit frontend; a FastAPI backend extracts the relevant subgraph context from the knowledge graph; an LLM generates an initial response grounded in that subgraph; the response is enriched with external metadata from UniProt, AlphaFold, and RXNav; and the final enriched answer is rendered back in Chainlit with a visualization of the supporting subgraph.

Dataset and graph database

The knowledge graph is built on BioKG — constructed from 34M+ PubMed abstracts, with 11,479,285 unique entities across 12 types and 42,504,077 relations across 52 types, drawn from 13 reputable data sources. BioKG was chosen over the graphs used by prior systems (KG-RAG's SPOKE, KRAGEN) for its data variety and quality. NebulaGraph hosts the graph for its open-source license, distributed architecture, performance at billions of nodes and edges, and the nGQL query language. Ingestion was a practical bottleneck: serial entity upserts took roughly 33 hours, which multiprocessing cut to about 3 hours.

Retrieval and enrichment

LlamaIndex handles indexing and retrieval of knowledge-graph triplets: a natural-language query is translated into a NebulaGraph lookup, the relevant entity–relationship–entity sets are retrieved, and the LLM composes an answer from them. A multi-stage enrichment pipeline then augments the raw graph answer — which is often just one or two lines — with gene information, protein structures, drug data, and links to external resources, with an LLM fallback when the graph query returns insufficient context.

Performance analysis

Performance was profiled across low-degree (~2 connections), mid-degree (31–51), and high-degree (8,151–13,939) nodes. Retrieval time dominates for high-degree nodes — up to ~58 seconds — a direct consequence of the graph's scale-free topology, while generation time stays negligible throughout. Proposed mitigations include caching, parallel processing, better indexing, and precomputation for hot high-degree entities.

Future work

Planned directions include high-degree node optimization, integrating clinical-trial and genomic data, domain-specific model fine-tuning, multimodal responses incorporating biomedical imaging, and HIPAA/GDPR compliance work.