GraphWalker: Graph-Guided In-Context Learning for Clinical Reasoning on Electronic Health Records

TL;DR

GraphWalker introduces a novel demonstration selection framework that combines data-driven and model-driven insights, cohort discovery, and an efficient search algorithm to enhance clinical reasoning with large language models on electronic health records.

Contribution

It presents a new method for selecting demonstrations in in-context learning that addresses key limitations in existing approaches for EHR reasoning.

Findings

GraphWalker outperforms state-of-the-art ICL methods on multiple EHR benchmarks.

The framework significantly improves clinical reasoning accuracy.

The code is publicly available at the provided GitHub URL.

Abstract

Clinical Reasoning on Electronic Health Records (EHRs) is a fundamental yet challenging task in modern healthcare. While in-context learning (ICL) offers a promising inference-time adaptation paradigm for large language models (LLMs) in EHR reasoning, existing methods face three fundamental challenges: (1) Perspective Limitation, where data-driven similarity fails to align with LLM reasoning needs and model-driven signals are constrained by limited clinical competence; (2) Cohort Awareness, as demonstrations are selected independently without modeling population-level structure; and (3) Information Aggregation, where redundancy and interaction effects among demonstrations are ignored, leading to diminishing marginal gains. To address these challenges, we propose GraphWalker, a principled demonstration selection framework for EHR-oriented ICL. GraphWalker (i) jointly models patient clinical information and LLM-estimated information gain by integrating data-driven and model-driven perspectives, (ii) incorporates Cohort Discovery to avoid noisy local optima, and (iii) employs a Lazy Greedy Search with Frontier Expansion algorithm to mitigate diminishing marginal returns in information aggregation. Extensive experiments on multiple real-world EHR benchmarks demonstrate that GraphWalker consistently outperforms state-of-the-art ICL baselines, yielding substantial improvements in clinical reasoning performance. Our code is open-sourced at https://github.com/PuppyKnightUniversity/GraphWalker