Toward Better EHR Reasoning in LLMs: Reinforcement Learning with Expert Attention Guidance

2025-08-19

AI Summary

To address the poor generalization and difficulty in modeling high-dimensional structured data of large language models (LLMs) for temporal reasoning on electronic health records (EHR), this paper proposes EAG-RL, a two-stage training framework. In Stage I, an expert model guides Monte Carlo Tree Search to generate high-quality clinical reasoning paths. In Stage II, attention mechanisms align the expert’s decision logic with the LLM’s internal representations, enabling interpretable end-to-end reinforcement learning. Unlike conventional hybrid paradigms that treat LLMs merely as feature retrievers, EAG-RL enhances the LLM’s intrinsic reasoning capability. Evaluated on two real-world EHR datasets, EAG-RL achieves an average 14.62% improvement in predictive performance, demonstrates enhanced robustness to feature perturbations, and exhibits strong cross-clinical-domain generalization.

Technology Category

• Machine Learning: Large Multimodal Models (LMMs)
• Multiagent Systems: Multiagent Learning
• Natural Language Processing: (Large) Language Models

Application Category

• Semantics and Knowledge: Data modeling to support human-machine intelligence, including LLMs agents, intelligent system behavior, explanations, and user-friendly interactions
• Search and Retrieval-Augmented AI: Agentic search
• User Modeling, Personalization and Recommendation: Large Language Models (LLM) for user modeling and recommendation

Abstract

Improving large language models (LLMs) for electronic health record (EHR) reasoning is essential for enabling accurate and generalizable clinical predictions. While LLMs excel at medical text understanding, they underperform on EHR-based prediction tasks due to challenges in modeling temporally structured, high-dimensional data. Existing approaches often rely on hybrid paradigms, where LLMs serve merely as frozen prior retrievers while downstream deep learning (DL) models handle prediction, failing to improve the LLM's intrinsic reasoning capacity and inheriting the generalization limitations of DL models. To this end, we propose EAG-RL, a novel two-stage training framework designed to intrinsically enhance LLMs' EHR reasoning ability through expert attention guidance, where expert EHR models refer to task-specific DL models trained on EHR data. Concretely, EAG-RL first constructs high-quality, stepwise reasoning trajectories using expert-guided Monte Carlo Tree Search to effectively initialize the LLM's policy. Then, EAG-RL further optimizes the policy via reinforcement learning by aligning the LLM's attention with clinically salient features identified by expert EHR models. Extensive experiments on two real-world EHR datasets show that EAG-RL improves the intrinsic EHR reasoning ability of LLMs by an average of 14.62%, while also enhancing robustness to feature perturbations and generalization to unseen clinical domains. These results demonstrate the practical potential of EAG-RL for real-world deployment in clinical prediction tasks. Our code have been available at https://github.com/devilran6/EAG-RL.

Problem

• Improving LLMs' EHR reasoning for clinical predictions
• Addressing underperformance on EHR-based prediction tasks
• Enhancing intrinsic reasoning capacity beyond hybrid paradigms

Innovation

• Reinforcement learning with expert attention guidance
• Monte Carlo Tree Search for stepwise reasoning
• Aligning LLM attention with clinical features