[Paper Review] KG4Diagnosis: A Hierarchical Multi-Agent LLM Framework with Knowledge Graph Enhancement for Medical Diagnosis

Kaiwen Zuo, Yirui Jiang+2 more|arXiv (Cornell University)|Dec 22, 2024

Artificial Intelligence in Healthcare9 citations

TL;DR

KG4Diagnosis presents a hierarchical multi-agent framework that couples LLMs with automated knowledge graph construction to support medical diagnosis across 362 diseases, featuring end-to-end KG generation and human-guided reasoning to curb hallucinations.

ABSTRACT

Integrating Large Language Models (LLMs) in healthcare diagnosis demands systematic frameworks that can handle complex medical scenarios while maintaining specialized expertise. We present KG4Diagnosis, a novel hierarchical multi-agent framework that combines LLMs with automated knowledge graph construction, encompassing 362 common diseases across medical specialties. Our framework mirrors real-world medical systems through a two-tier architecture: a general practitioner (GP) agent for initial assessment and triage, coordinating with specialized agents for in-depth diagnosis in specific domains. The core innovation lies in our end-to-end knowledge graph generation methodology, incorporating: (1) semantic-driven entity and relation extraction optimized for medical terminology, (2) multi-dimensional decision relationship reconstruction from unstructured medical texts, and (3) human-guided reasoning for knowledge expansion. KG4Diagnosis serves as an extensible foundation for specialized medical diagnosis systems, with capabilities to incorporate new diseases and medical knowledge. The framework's modular design enables seamless integration of domain-specific enhancements, making it valuable for developing targeted medical diagnosis systems. We provide architectural guidelines and protocols to facilitate adoption across medical contexts.

Motivation & Objective

• Motivate the need for systematic frameworks that handle complex medical diagnosis with specialized expertise.
• Develop an end-to-end knowledge graph construction pipeline for medical knowledge.
• Integrate a two-tier MAS (GP agent + specialized consultants) with LLMs for robust diagnostic reasoning.
• Address LLM hallucinations via multi-agent verification and KG constraints.
• Provide modular architecture and implementation guidelines for adapting to new medical domains.

Proposed method

• Propose a two-tier architecture with a General Practitioner (GP) agent and domain-specific Consultant LLMs for hierarchical diagnosis.
• Implement a three-stage knowledge graph construction pipeline: data chunking, semantic-driven entity/relationship extraction using BioBERT and ontologies (SNOMED-CT/UMLS), and graph construction with LLM-augmented expansion.
• Incorporate human-guided reasoning to validate and expand the KG through expert evaluation and reinforcement learning.
• Model diagnostic reasoning with probabilistic confidence functions for GP-LLM and Consultant-LLMs, and compute a final diagnosis via agent fusion with weights.
• Use an inter-agent communication protocol to transfer queries and update GP knowledge base from Consultant feedback.
• Implement modular architecture with PyTorch for neural components and Neo4j for KG management.

Experimental results

Research questions

• RQ1 How can a hierarchical multi-agent system improve diagnostic accuracy while limiting hallucinations in medical AI?
• RQ2 Can automated knowledge graphs constructed from unstructured medical data enhance LLM-based diagnosis across multiple specialties?
• RQ3 What mechanisms best integrate GP-level triage with domain-specific consultants for efficient and accurate diagnosis?
• RQ4 How does expert-guided KG expansion influence consistency and可靠性 of medical reasoning?
• RQ5 What is the scalability and adaptability of KG4Diagnosis when adding new diseases and domains?

Key findings

• The framework addresses hallucination through multi-layer verification and KG constraints, improving diagnostic reliability over standalone LLMs.
• It covers 362 diseases across multiple specialties, demonstrating practical scope for hierarchical medical diagnosis.
• Semantic-driven extraction with BioBERT and ontology alignment enables precise entity and relation extraction for KG construction.
• LLM augmentation extends the KG with broader contextual entities and relations beyond domain-specific extractions.
• A modular architecture supports seamless integration of new domains and knowledge, with planned benchmarks for evaluation.
• The GP agent efficiently triages and coordinates with specialist agents, potentially reducing unnecessary full-domain consultations.