Teaching EHR models to reason in time

Electronic health records (EHRs) contain sequences of events: admissions, transfers, diagnoses and other details recorded at different times. Models trained on these records can learn to predict what comes next. But records may be incomplete, and a model’s generation limit can make it difficult to represent events unfolding over long periods. The authors ask whether reinforcement learning (RL)—training with feedback tied to a task’s outcome—can help models reason more effectively about these timelines.

The researchers treated an EHR foundation model as a generator of possible future patient histories. They tested three prediction tasks using the MIMIC-IV database: whether a patient would be readmitted to hospital within 30 days of discharge; whether someone discharged from intensive care would return within two years; and whether a patient would die in hospital within two years of admission. These are predictions from recorded histories, not recommendations about treatment.

A key design problem was deciding what to do when a generated timeline ended before it could answer the question. If a model ran out of tokens before the relevant time window had passed, the absence of a recorded event did not necessarily mean the event would not occur. The researchers therefore treated such outputs as inconclusive, rather than rewarding or penalizing them. They also filtered training examples when the available record did not cover enough time, or the outcome lay beyond the model’s supported generation length.

The team fine-tuned ETHOS, an EHR model, using two RL methods, GRPO and DAPO, and compared results with the pretrained model and several other baselines. They evaluated performance using measures that assess how well a model ranks cases, including AUROC and AUPRC. For 30-day readmission, both RL versions improved on the 9-million-parameter pretrained model: its AUROC was 0.7385, compared with 0.7455 for GRPO and 0.7515 for DAPO. The 9-million-parameter DAPO model was broadly comparable to a larger, 56-million-parameter pretrained version on this task, though the measures did not both favor the smaller model.

The authors also report improvements across the ICU readmission and inpatient mortality tasks when RL training was extended to multiple tasks. In their experiments, training on one task sometimes improved results on others, and adding tasks generally brought further gains. This is evidence of transfer within the study’s setup—not proof that the same pattern will hold in other hospitals or populations.

The researchers examined generated timelines as well as prediction scores. By several measures, RL-generated trajectories more closely resembled the recorded timelines than outputs from the pretrained model. They also provided somewhat more useful inputs to a simple decision-tree classifier for readmission prediction. These findings suggest the improvement was not limited to the models’ task scores, though resemblance to recorded data does not by itself show that a generated history is clinically correct.

There are important limits. The work used a single EHR dataset and one foundation-model family, and the authors say the evaluation did not include clinical verification. The training filters removed substantial portions of some task datasets, especially for the longer time horizons; filtering and balancing also changed which examples the model learned from. Results therefore depend on those choices and on MIMIC-IV’s data. The experiments measure retrospective predictive performance, not whether using a model improves care or patient outcomes.

The paper appears in the Proceedings of Machine Learning Research for ML4H 2026; the supplied source is its arXiv version. The paper’s stated limitations and this source text provide the basis for this account. As with any extracted paper text, formatting can obscure details, and reported model performance should not be mistaken for clinical validation. The approach is a research direction, not evidence that an EHR model is ready to guide patient care.