Source-linked AI summary

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

Kaiwen Zuo, Yirui Jiang, Fan Mo, Pietro Lio

arXiv:2412.16833v4cs.AIcs.LG

TL;DR

Medical knowledge-graph construction and reasoning remain difficult for unstructured medical data, while LLM diagnosis requires specialized expertise and hallucination control. KG4Diagnosis addresses these needs with a hierarchical GP-and-consultant architecture integrated with an end-to-end, human-guided knowledge-graph pipeline. The framework is presented as a modular foundation for specialized medical diagnosis systems, although its performance depends on high-quality and comprehensive medical data and remains challenged by rare or rapidly evolving conditions.

  • Problem

    Medical knowledge graphs face unresolved construction and reasoning challenges on unstructured and multimodal data, while LLM-based methods face hallucination and accuracy concerns.

  • Method

    KG4Diagnosis combines a GP agent, domain-specific consultant agents, semantic extraction, decision-relation reconstruction, and human-guided knowledge expansion.

  • Results

    The framework provides a hierarchical, modular architecture integrating automated knowledge-graph construction with multi-agent diagnostic reasoning across 362 common diseases.

  • Takeaways & Limitations

    KG4Diagnosis is presented as an extensible foundation for incorporating new medical domains and knowledge into specialized diagnosis systems.

  • Takeaways & Limitations

    Performance depends on the quality and comprehensiveness of the knowledge graph and medical data, with remaining challenges for rare disease combinations and rapidly evolving knowledge.

Abstract

from arXiv · show

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.

Introduction

KG4Diagnosis addresses challenges in constructing and reasoning over medical knowledge graphs from unstructured data by combining hierarchical multi-agent diagnosis with an end-to-end knowledge-graph pipeline. Its stated contributions include semantic extraction, decision-relation reconstruction, human-guided expansion, hallucination controls, real-world scenarios, and modular extensibility.

  • Medical knowledge graphs support clinical decision-making, but constructing and reasoning over them from unstructured and multimodal data remains challenging.
  • Existing approaches trade off reliability, scalability, unstructured-data handling, accuracy, hallucination control, or computational complexity.
  • KG4Diagnosis combines a general-practitioner agent for initial assessment with specialized agents for domain-specific diagnosis across 362 common diseases.
  • Its end-to-end knowledge-graph pipeline uses semantic-driven extraction, multi-dimensional decision-relation reconstruction, and human-guided reasoning for scalable knowledge expansion.
  • The framework also proposes multi-agent verification, knowledge-graph constraints, real-world healthcare scenarios, and modular integration of new domains and knowledge.

System Architecture Overview

KG4Diagnosis separates knowledge-graph construction from hierarchical medical decision-making while connecting both through a Camel-based multi-agent system. The architecture mirrors clinical practice by pairing general practitioners with specialists.

  • The architecture has two primary components: a medical knowledge-graph construction pipeline and a Camel-based multi-agent system for hierarchical decision-making.
  • The multi-agent design mirrors real-world medical practice, with general practitioners collaborating with specialists for comprehensive patient care.

Knowledge Graph Construction Pipeline

The knowledge-graph pipeline segments medical documents, extracts ontology-linked entities and relationships, constructs a graph, and prepares structured knowledge for later reasoning. BioBERT provides biomedical extraction while contextual segmentation and standardized ontologies organize the inputs.

  • Knowledge Graph Construction Pipeline: The pipeline first segments medical documents into context-constrained chunks before extracting entities and relationships.
  • Knowledge Graph Construction Pipeline: The framework presents automated graph construction as a staged process that ends with a structured representation suitable for downstream medical reasoning.
  • Knowledge Graph Construction Pipeline: BioBERT and medical ontologies such as SNOMED-CT and UMLS support semantic-driven extraction of medical entities and relationships.
  • Knowledge Graph Construction Pipeline: Extracted entities and relationships are stored in dedicated databases and used to construct the medical knowledge graph for diagnostic reasoning.
  • Knowledge Graph Construction Pipeline: The constructed graph represents medical entities as nodes and relationships between those entities as edges.

Stage 4: LLM-Augmented Knowledge Graph

KG4Diagnosis augments biomedical extraction with LLM-based context-aware knowledge expansion and uses expert validation to improve graph reliability. The validated graph then supports hierarchical diagnosis, with a GP agent referring low-confidence or specialized cases to domain consultants.

  • LLM-Augmented Knowledge Graph: LLMs extend BioBERT by identifying broader, context-aware entities and relations in complex or ambiguous medical text.
  • LLM-Augmented Knowledge Graph: The enriched knowledge is stored in databases and integrated into the graph to support multi-agent diagnostic reasoning and decision-making.
  • Human-Guided Reasoning: Experts validate extracted relationships, correcting or removing invalid or uncertain links before incorporating verified entities and relationships into the graph.
  • Human-Guided Reasoning: The expanded graph includes previously extracted and expert-verified entities and relationships, improving reliability and comprehensiveness.
  • Hierarchical Multi-Agent Framework: Consultant-LLMs are optimized for specific medical domains, including cardiology, neurology, endocrinology, and rheumatology.
  • GP-LLM: Primary Diagnostic Agent: The GP-LLM refers cases when confidence falls below 0.7 or when the diagnosis requires specialized expertise.

Consultant-LLMs: Specialized Diagnostic Agents

Consultant-LLMs provide domain-specific diagnostic confidence, collaborate on complex cases, and support referral-driven refinement within the hierarchical framework.

  • Consultant-LLMs: Four Consultant-LLMs specialize in cardiology, neurology, endocrinology, and rheumatology, using domain-specific training datasets and clinical guidelines.
  • Advanced Diagnosis with Multi-Agent Collaboration: Collaborative cases combine multiple agents’ weighted diagnostic contributions into a final confidence estimate for complex diagnoses.
  • Inter-Agent Communication Protocol: Queries are transferred between agents through a compatibility transformation, enabling referral and collaborative refinement across the hierarchy.
  • Inter-Agent Communication Protocol: Consultant feedback contributes incremental knowledge to the GP-LLM, and the framework sets the referral threshold at τ = 0.7.
  • Referral Decision Threshold: The GP-LLM refers cases when its confidence falls below τ or when the diagnosis belongs to the specialist set, while otherwise retaining the query.

Future Training and Evaluation Work

The framework covers 362 diseases, represents complex clinical relationships through knowledge graphs, and is being extended with continuous learning and a future benchmark.

  • Training Scope: Targeted fine-tuning gives specialist agents deep disease-specific knowledge while preserving integration with broader medical knowledge.
  • Knowledge Graph Scope: The knowledge graph links diseases, symptoms, diagnostic patterns, medications, patient status, surgery, drugs, and BMI to organize obesity-related clinical information.
  • Continuous Learning: The current knowledge base encompasses all 362 diseases and uses agent interactions and feedback loops to refine diagnostic capabilities over time.
  • Future Evaluation: A forthcoming benchmark will evaluate diagnostic accuracy, hallucination prevention, and multi-agent coordination efficiency using standardized metrics.

Discussion

KG4Diagnosis combines automated knowledge graphs with hierarchical multi-agent diagnosis across 362 diseases, showing strengths in accuracy, hallucination prevention, clinical coordination, and modular expansion while retaining scalability and data-coverage challenges.

  • Overall findings: KG4Diagnosis covers 362 common diseases and combines knowledge-graph construction with hierarchical multi-agent diagnosis, demonstrating promising capabilities alongside important challenges.The framework is evaluated across multiple medical specialties.
  • Knowledge graph construction: Higher precision in semantic entity extraction and relationship reconstruction supports handling complex medical terminology and relationships.These modules are described as outperforming conventional approaches in precision.
  • Hierarchical diagnosis: Multiple validation layers and knowledge-graph constraints significantly reduce incorrect diagnoses compared with standalone LLM implementations.The GP agent also triages cases and coordinates specialist consultation for complex cases.
  • Knowledge organization: The knowledge graph organizes interconnected diseases, symptoms, and diagnostic patterns, supporting efficient navigation, hierarchical decisions, and incorporation of new medical domains.Its modularity is positioned as useful for changing medical knowledge.
  • Scalability: As medical domains expand, coordinating more specialist agents increases system complexity and motivates more sophisticated coordination mechanisms.The tiered structure still manages computational resources efficiently.
  • Limitations: Performance depends on knowledge-graph quality and coverage, with rare diseases, unusual combinations, and rapidly evolving knowledge remaining difficult cases.Future work also plans MedQA experiments and comparison with prominent models, while deployment may be constrained by limited medical data resources.

Related Work

Related work spans rule-based, deep-learning, hybrid symbolic-neural, multimodal, medical-LLM, and hierarchical multi-agent approaches. These methods improve extraction, reasoning, or clinical specialization but remain constrained by scalability, domain nuance, hallucination, verification, and integration challenges.

  • Rule-Based and Ontology-Driven Approaches: Rule-based and ontology-driven methods offer interpretable, structurally consistent medical knowledge but struggle with scalability and unstructured data.The cited approaches include SNOMED-CT and UMLS-based systems.
  • Deep Learning and Pre-Trained Models: Pre-trained biomedical models improve clinical information extraction but can miss domain-specific nuances and require substantial computational resources.LLMs extend processing of unstructured medical data but introduce accuracy and hallucination concerns.
  • Hybrid Symbolic-Neural Approaches: Hybrid symbolic-neural systems balance interpretability with adaptability, yet require curated inputs and face computational scalability challenges.Multimodal knowledge graphs add clinical notes, imaging, and laboratory data while retaining standardization and fusion difficulties.
  • Advancements in Medical LLMs: Medical LLMs such as ESM-1b, Med-PaLM, MediTron, BioGPT, and GPT-4medprompt target biomedical representation, clinical reasoning, multimodal integration, or extraction.The passage distinguishes their application emphases rather than reporting a common benchmark result.
  • Hierarchical Multi-Agent Architectures: Hierarchical multi-agent architectures assign initial assessment to general agents and domain diagnosis to specialists, improving diagnostic accuracy, scalability, and reliability.This organization is described as mirroring real-world medical systems.
  • Open Challenges: The field still needs better extraction from unstructured data, stronger hallucination verification, and improved multimodal integration for medical applications.These challenges motivate continued methodological development.

Conclusion

KG4Diagnosis integrates automated knowledge-graph construction with specialized LLMs and a hierarchical agent structure for diagnosis across 362 common diseases. Its validation layers and targeted specialist consultation support hallucination prevention and resource management, while standardized benchmarks remain under development.

  • Conclusion: KG4Diagnosis integrates a three-stage knowledge-graph pipeline with hierarchical agents covering 362 common diseases.The pipeline uses semantic processing and human-guided reasoning, while agents mirror real-world medical practices.
  • Conclusion: Multiple validation layers help prevent hallucinations, and targeted specialist consultation manages computational resources.The conclusion presents these as significant advantages of the framework.
  • Conclusion: Comprehensive standardized benchmarks are still being developed to provide shared evaluation metrics for the research community.The current implementation is described as promising but not yet fully benchmarked.
Loading 2412.16833v4…