ReToP: Learning to Rewrite Electronic Health Records for Clinical Prediction

TL;DR

ReToP tackles clinical prediction from heterogeneous EHRs by learning to rewrite patient records with an LLM-based EHR rewriter that is trained on synthetic rewrites and aligned to downstream prediction via a CSC-guided end-to-end objective. The framework jointly tunes the rewriter and a task-specific predictor to maximize predictive likelihood across rewrites, while preserving fidelity to the original EHR. Empirical results on MIMIC-IV and transfer tests show substantial gains across mortality, readmission, LOS, and congenital malformation prediction, with particular strength on imbalanced tasks and under domain shift. The work highlights the feasibility of task-aware rewrite mechanisms for healthcare AI, offering a lightweight adaptation path for unseen datasets and tasks along with insights into feature emphasis and explainability tradeoffs.

Abstract

Electronic Health Records (EHRs) provide crucial information for clinical decision-making. However, their high-dimensionality, heterogeneity, and sparsity make clinical prediction challenging. Large Language Models (LLMs) allowed progress towards addressing this challenge by leveraging parametric medical knowledge to enhance EHR data for clinical prediction tasks. Despite the significant achievements made so far, most of the existing approaches are fundamentally task-agnostic in the sense that they deploy LLMs as EHR encoders or EHR completion modules without fully integrating signals from the prediction tasks. This naturally hinders task performance accuracy. In this work, we propose Rewrite-To-Predict (ReToP), an LLM-based framework that addresses this limitation through an end-to-end training of an EHR rewriter and a clinical predictor. To cope with the lack of EHR rewrite training data, we generate synthetic pseudo-labels using clinical-driven feature selection strategies to create diverse patient rewrites for fine-tuning the EHR rewriter. ReToP aligns the rewriter with prediction objectives using a novel Classifier Supervised Contribution (CSC) score that enables the EHR rewriter to generate clinically relevant rewrites that directly enhance prediction. Our ReToP framework surpasses strong baseline models across three clinical tasks on MIMIC-IV. Moreover, the analysis of ReToP shows its generalizability to unseen datasets and tasks with minimal fine-tuning while preserving faithful rewrites and emphasizing task-relevant predictive features.

Figures (5)

• Figure 1: Comparison of three LLM-based approaches for clinical prediction from EHRs. Approach (a) relies on fine-tuned LLMs as EHR encoders. Approach (b) relies on LLMs to complement EHRs. Our approach, ReToP (c), uses a trainable LLM-based EHR rewriter aligned with clinical tasks.
• Figure 2: Overview of ReToP training and inference. We train the EHR rewriter on $\mathcal{D}_{Rw}$ (step 1). Then, jointly end-to-end train the rewriter and predictor on the target task (steps 2-3).
• Figure 3: ROC scores across $\alpha$ values for inference level (§ Eq. 2) using ReToP$_{Llama}$ on MIMIC-IV, shown for MOR (left), RA (center), and LOS (right) tasks. Curves are stratified by EHR length: overall (blue), short EHR $<2048$ (green), medium EHR in $[2048,4096]$ (yellow), and long $>4096$ (red).
• Figure 4: Effect of $\lambda$ on KL training (up to 4k steps, eval. every 1000 steps) for MOR, RA, LOS using ReToP$_{Llama}$.
• Figure 5: Example patient EHR for the MOR task. Red text shows main content-based differences across rewrites, while highlighted text shows generated unfaithful information.