PaperlayerView exact revision on arXiv

Source-linked AI summary

Incentivizing Tool-Use Reinforcement Learning for Post-Arrival Emergency Care Prediction

Pengfei Hu, Chenwei Wu, Qiao Jin, Jun Deng, Zhiyong Lu, Liyue Shen, Yue Ning

arXiv:2610.21449v1cs.AI

TL;DR

LLMs underperform conventional models on EHR prediction, with poor feature selection and unreliable numerical interpretation identified as key weaknesses. TIR-EC uses tree-guided tool selection and reverse-curriculum skill distillation to address both. Its 4B agent outperforms conventional ML baselines and is competitive with much larger frontier LLMs.

  • Problem

    LLMs underperform conventional models on EHR prediction, with poor feature selection and unreliable interpretation of numerical values contributing to the gap.

    Source: §1 · pp. 1–2

  • Method

    TIR-EC treats EHR feature groups as RL-optimized tools and distills task-specific clinical skills into model weights through a reverse curriculum.

    Source: §1 · p. 2

  • Results

    A 4B TIR-EC agent outperforms conventional ML baselines and is competitive with much larger frontier LLMs.

    Source: §6 · p. 9

  • Takeaways & Limitations

    Feature selection and interpretation improve together in TIR-EC's analysis.

    Source: §6 · p. 9

    The study covers three post-arrival tasks, and disposition and revisit labels may also reflect local practice, resource availability, and downstream care.

    Source: §7 · p. 9

Abstract

from arXiv · show

Large language models (LLMs) can reason over both structured measurements and free-text records, which makes them attractive for emergency department (ED) prediction. However, most existing LLMs still fall short of conventional ML methods like trees on EHR-based prediction tasks. Our controlled analysis traces this gap to two weaknesses: LLMs do not select informative features, and they misjudge what numerical values imply for risk even when given the right features. We propose TIR-EC, a tool-integrated reinforcement learning framework that addresses both. TIR-EC exposes EHR feature groups as tools the LLM can call. It trains the LLM with rewards from a gradient-boosted tree teacher, which provides a reference set of informative tools and a graded risk target. TIR-EC also distills task-specific clinical skills into the model weights through a reverse curriculum that gradually removes the skills from the prompt during RL process. Evaluation on three clinical prediction tasks shows that our 4B agent achieves the highest macro F1 of all compared models, outperforms conventional ML baselines, and is competitive with much larger frontier LLMs. Further analysis shows that TIR-EC selects more useful evidence with fewer tool calls and combines measurements more effectively.

1 Introduction

The introduction identifies poor feature selection and unreliable numerical interpretation as key weaknesses behind LLM underperformance on EHR prediction. It proposes TIR-EC, using tool selection and skill distillation to address both.

1 Introduction

  • The ED is a high-stakes, time-critical setting where EHR-based predictions of admission, deterioration, and return visits can support timely decisions.LLMs can reason over structured EHR data and free-text clinical notes.Source: §1 · p. 1
  • More EHR features do not consistently improve LLM predictions, while supplying GBM-ranked groups improves performance.The introduction identifies poor feature selection as a bottleneck in leveraging EHR features.Source: §1 · p. 1
  • Even with GBM-selected features, LLM risk estimates can overreact, underreact, or point in the wrong direction.In a 100-case error analysis, most errors involved mis-weighting evidence the LLM had already retrieved.Source: §1 · p. 2
  • TIR-EC makes feature selection an RL-optimized tool-calling action and rewards selections against a tree-importance reference set.The framework also distills task-specific rules about which evidence to inspect and how its values shift risk, gradually removing those rules from the prompt during training.Source: §1 · p. 2
  • TIR-EC achieves state-of-the-art performance and improves both feature selection and feature understanding, which the paper reports are mutually beneficial.The paper also describes these improvements across disposition, decompensation, and revisit prediction.Source: §1 · pp. 2–3

2 Related Work

Prior emergency-care prediction work includes arrival-time and post-arrival settings, with this paper focusing on post-arrival prediction from data gathered during care. Existing EHR agents and tool-integrated methods leave evidence selection and interpretation insufficiently addressed.

2.1 Emergency Care Prediction

  • Post-arrival prediction updates risk as repeated vital signs, labs, orders, imaging, and physiological signals become available during care.These heterogeneous sources and numerical and temporal interpretation make the setting more demanding for LLMs than compact triage inputs.Source: §2.1 · p. 3

2.2 LLMs and Agents for EHR

  • Existing EHR agents and tool-integrated RL methods largely overlook teaching LLMs to identify useful evidence and interpret retrieved information.The paper's stated goal is to use tree models to supervise evidence selection and anomaly identification.Source: §2.2 · p. 3

3 Methodology

TIR-EC trains an agent to retrieve clinical evidence and predict outcomes, using tree-based supervision for evidence choices and risk signals while gradually fading clinical guidance during training.

3.1 Task Definition

  • The policy observes an initial patient view, selectively calls clinical functions for additional evidence, and returns an outcome probability after at most K calls.Training maximizes expected terminal reward over the resulting trajectory.Source: §3.1 · p. 3

3.2 Reward Design and RL Training

  • Tree feature importance defines task-level reference functions, and matching rewards compare those functions with the agent’s calls without requiring a retrieval sequence.Functions already represented in the initial view are removed; functions returning detailed values or text are excluded because the tree does not observe their contents.Source: §3.2 · p. 4
  • The prediction reward retains tree confidence only on the observed label’s side of the decision boundary, moving an estimate on the opposite side to 0.5.The paper motivates this class-safe target as avoiding the risk of teaching the wrong class from an erroneous tree probability.Source: §3.2 · p. 4
  • The combined reward includes prediction and function-selection signals plus output-format and execution checks, and the policy is optimized with GRPO.Execution checks whether calls are complete, parseable, within budget, and error-free.Source: §3.2 · pp. 4–5

3.3 Skill Curriculum

  • A reverse curriculum starts with complete clinical skills and removes guidance in stages, while grouped trajectories for each case share the same retained skills.Validation and test prompts contain no skills, so gradual removal is intended to encourage internalization rather than prompt reliance.Source: §3.3 · p. 5

4 Experiments

Across three emergency-care prediction tasks, TIR-EC achieves the highest macro F1 among compared models and outperforms ML baselines, while matching or competing with much larger frontier LLMs.

4.1 Experimental Setups

  • Experiments use MC-MED’s 118,385 adult ED visits for three binary tasks: disposition, decompensation, and revisit.The study uses a patient-disjoint chronological 78%/11%/11% training, validation, and test split; the shared Qwen3-4B policy selects up to 12 calls from 18 feature groups.Source: §4.1 · p. 5
  • Comparisons include ML and deep-learning models, small and domain-specific LLMs, frontier LLMs, and Qwen3-4B backbone controls.All LLMs use the same agentic feature-selection setup except Qwen’s Full Features rows, which take all features.Source: §4.1 · p. 5, Table 1 · p. 6

4.2 Main Results

  • TIR-EC achieves the highest macro F1 of all models and AUROC second to Opus 5, with the best F1 on decompensation and revisit and best AUPRC on decompensation.It beats ML at all metrics and improves F1, AUPRC, and AUROC on every task over its own backbone.Source: §4.2 · p. 5, Table 1 · p. 6
  • General-purpose LLMs tend to miss positive cases, while many medical and reasoning-oriented LLMs flag nearly everyone; TIR-EC alleviates both patterns.Removing either tool-matching reward or skill curriculum lowers performance across all three tasks, while removing the GBM probability target mainly hurts ranking quality.Source: §4.2 · p. 6

5 Analysis

Across five analyses, TIR-EC selects useful evidence more effectively, combines multiple measurements better, and adapts its requests beyond the GBM teacher’s choices. Case studies illustrate these behaviors in individual predictions.

5.1 Does TIR-EC better select features?

  • TIR-EC achieves the highest AUROC while making fewer tool calls than most baselines, and its selection beats random selection at every budget.The selection advantage is largest when evidence is scarce; within a few picks, performance approaches its full-budget level.Source: §5.1 · p. 6

5.2 Does TIR-EC better interpret features?

  • With shared randomly drawn groups, TIR-EC improves steeply as groups are added and nears its ceiling with about half, while the base model gains slowly.The advantage lies in combining many measurements rather than reading any single measurement better; medical-domain models do not show this ability.Source: §5.2 · pp. 6–7, Figure 3 · p. 7

5.3 Are feature selection and interpretation independent abilities?

  • After training, TIR-EC’s requests track usable value: it requests orders routinely as they become more useful, unlike the base model’s near-avoidance.Nearly every feature group contributes more to TIR-EC than to its base model, and requests shift toward groups that gained most.Source: §5.3 · p. 7, Figure 4 · p. 7

5.4 Is TIR-EC merely mimicking the GBM teacher?

  • TIR-EC requests teacher-favored groups but also opens detailed tools the GBM cannot access, adapting evidence to task and patient.For disposition, both rely on orders, while TIR-EC also reads the specific orders and diagnoses behind that group.Source: §5.4 · p. 8, Figure 6 · p. 8

5.5 Case Study

  • In two illustrative cases, TIR-EC weighs the whole record in an admission prediction and stops after reassuring vital signs in a no-deterioration prediction.It predicts discharge in the first case, matching the outcome; in the second, both models correctly predict no deterioration, but GPT-5.6 Sol reviews more evidence.Source: §5.5 · p. 8, Table 2 · p. 8

6 Conclusion

The paper attributes LLM underperformance on EHR-based emergency-care prediction to poor feature selection and unreliable numerical interpretation. It proposes TIR-EC, a 4B agent that outperforms conventional ML baselines, is competitive with much larger frontier LLMs, and improves selection and interpretation together.

6 Conclusion

  • TIR-EC learns to select EHR features as tools through tree-guided reinforcement learning and distills clinical skills into the model via a reverse curriculum.A 4B agent outperforms conventional ML baselines, is competitive with much larger frontier LLMs, and improves feature selection and interpretation together.Source: §6 · pp. 8–9

7 Limitations

The evaluation is bounded by one dataset, three post-arrival tasks, tree-guided supervision, and a predefined tool interface. Results are retrospective, and numerical reasoning errors can remain at the individual-trajectory level.

7 Limitations

  • The study uses MC-MED alone, so the learned evidence-selection policy and risk interpretation may not transfer directly to other institutions with different populations, measurement practices, documentation, or data sources.External validation across health systems is identified as an important next step.Source: §7 · p. 9
  • The three tasks do not cover all clinical decisions, and disposition and revisit labels may reflect local practice, resource availability, and downstream care as well as physiology.The authors propose testing other outcomes and prediction horizons to assess generality.Source: §7 · p. 9
  • Tree-guided supervision can inherit dataset-specific teacher correlations, while its task-level evidence targets do not specify the optimal retrieval path for each patient.The authors suggest more adaptive or patient-specific supervision as future work.Source: §7 · p. 9
  • The predefined clinical-function interface may need remapping or redesign in another system because EHR schemas, tool granularity, and access constraints differ.Source: §7 · p. 9
  • Training improves population-level numerical evidence combination but does not prevent occasional per-trajectory errors, including comparing a correctly retrieved value against a threshold in the wrong direction.The authors suggest additional verification mechanisms for learned numerical interpretation.Source: §7 · p. 9
  • Retrospective predictive and tool-use results do not establish improved clinician decisions, workflow efficiency, or patient outcomes in prospective use.Prospective evaluation would need to account for calibration, latency, alert burden, and human interaction with model outputs.Source: §7 · p. 9

8 Ethical Considerations

The authors frame TIR-EC as a research framework rather than a clinical-care system. Ethical concerns include individual prediction failures, bias and distribution shift, and the limits of generated rationales.

8 Ethical Considerations

  • TIR-EC is a research framework, and the results are not evidence that it is ready for direct patient care.Source: §8 · p. 9
  • Incorrect predictions could cause inappropriate escalation or false reassurance, so clinical use would require independent validation, calibrated thresholds, human oversight, and checks for unreliable outputs.Strong aggregate performance does not eliminate individual failures, including numerical interpretation errors.Source: §8 · p. 9
  • Because records and labels reflect healthcare processes, learned patterns may reproduce differences in access, treatment, documentation, or resource use across groups and institutions.Tree-guided supervision does not remove this issue and may transfer such correlations into the trained policy; the authors call for subgroup evaluation and continued monitoring.Source: §8 · pp. 9–10
  • Observable tool calls and short reasoning can aid inspection, but generated rationales should not be treated as faithful explanations; evidence provenance and verifiable checks are preferred.Source: §8 · p. 10
  • The authors do not recommend TIR-EC as a substitute for clinician judgment or as a standalone basis for admission or discharge decisions.Source: §8 · p. 10

A.1 Benchmark Datasets.

The benchmark uses MC-MED, a multimodal emergency-care dataset, to evaluate three binary post-arrival prediction tasks under time-aware evidence cutoffs.

A.1 Benchmark Datasets

  • MC-MED contains 118,385 adult emergency-department visits from 70,545 patients, split chronologically by patient into training, validation, and test sets containing 78%, 11%, and 11% of patients, respectively.Source: §A.1 · p. 13
  • The benchmark covers disposition at 1, 2, or 4 hours, decompensation within 60, 90, or 120 minutes after a 15-minute assessment, and revisit within 3, 7, or 14 days.Each prediction instance uses only evidence available by its task-specific cutoff.Source: §A.1 · p. 13

A.2 Implementation Details.

Implementation trains a shared Qwen3-4B policy across three tasks with multi-turn clinical-function calls and reinforcement learning. Clinical-skill prompts are gradually removed during training and omitted during validation and testing.

A.2 Implementation Details

  • Qwen3-4B without thinking mode is the backbone for one shared policy across the three prediction tasks, with at most 12 clinical-function calls per trajectory.Training samples 10,000 instances per task and runs for 500 optimizer steps.Source: §A.2 · p. 13
  • Training uses batches of 64 prompts, generates 8 trajectories per prompt, and uses a learning rate of 1 × 10−6; reward computation combines prediction and function-matching rewards with α = 0.7.Source: §A.2 · p. 13
  • Clinical-skill sentences are fully retained for the first 20 steps, then their retention probability drops by 0.05 every 20 steps until reaching zero at step 400.Validation and test prompts omit the skill sentences, and task-specific checkpoints are selected by validation F1.Source: §A.2 · p. 13

B Tool-Use Training Details

The appendix specifies GRPO training, multi-turn evidence gathering, and the prompts, clinical skills, and cutoff-aware tool interface used to operationalize TIR-EC.

B.1 RL Training with GRPO

  • GRPO trains on groups of sampled trajectories per case, standardizes each trajectory’s reward relative to that group, and compares alternative evidence paths without a separate value model.The token-level objective uses asymmetric clipping, masks environment observations and overlong trajectories, and includes KL and entropy controls.Source: §B.1 · p. 13

B.2 TIR Rollout

  • The rollout lets the model plan, call a tool, inspect the returned observation, and continue until it predicts or reaches its call budget.Tools offer summaries for broad signals and detail access to measured values or clinical text; empty observations are retained, and the model is prompted to stop when further evidence is unlikely to change its prediction.Source: §B.2 · p. 14, Figure 7 · p. 14

C Prompts and Clinical Tools

  • The appendix separates shared instructions, task definitions, training-time skills, and a tool catalog, with bracketed prompt fields filled for each case.The prompt cards consolidate repeated instructions while preserving task definitions, numerical thresholds, and the output contract.Source: C Prompts and Clinical Tools · p. 15

C.1 Shared Prediction Prompt

  • The shared prompt supplies the task and cutoff, requires evidence-grounded tool use, and asks the model to return one probability in the final answer block.The model is instructed to prefer a few targeted calls, treat empty responses as informative, and reserve a turn for its answer.Source: §C.1 · p. 15

C.2 Task-Specific Prompts

  • Task prompts define admission, post-discharge revisit, and acute decompensation prediction by their respective cutoffs, outcomes, and prediction windows.Revisit prediction uses index-visit evidence and excludes post-discharge events; decompensation is defined by future vital-sign thresholds, not instability at the cutoff.Source: §C.2 · p. 16

C.3 Training-Time Skill Prompts

  • Training-time skill cards teach explicit vital-sign comparisons, risk combinations, order-count interpretation, and guarded use of disposition signals.They caution against inferring risk or disposition from race or ethnicity and state that the numerical rules are benchmark-specific training guidance, not clinical recommendations or outcome definitions.Source: §C.3 · p. 17

C.4 Clinical Tool Catalog and Responses

  • The interface provides 18 functions—12 structured summaries and 6 detailed-value or text tools—with empty arguments and evidence filtered by cutoff availability.The function-matching reward excludes detailed-value and text tools; timestamp rules vary by record type, and final disposition and subsequent-visit outcomes are not returned.Source: §C.4 · p. 18

Related papers

Loading 2610.21449v1…