One Loss to Rule Them All: Marked Time-to-Event for Structured EHR Foundation Models

TL;DR

EHR data are irregularly sampled and include numerical measurements, which challenges standard next-token pretraining. The authors propose ORA, a discretized, code-specific marked time-to-event objective that jointly models event timing and values, enabling a non-sparse, multi-code predictive signal. Across two large datasets and two architectures, ORA yields improved generalizability and performance across classification, time-to-event, and regression tasks, outperforming next-token and prior time-to-event losses. The work argues that pretraining objectives should align with EHR structure to expand downstream capabilities and suggests directions for scaling and robustness in healthcare foundation models.

Abstract

Clinical events captured in Electronic Health Records (EHR) are irregularly sampled and may consist of a mixture of discrete events and numerical measurements, such as laboratory values or treatment dosages. The sequential nature of EHR, analogous to natural language, has motivated the use of next-token prediction to train prior EHR Foundation Models (FMs) over events. However, this training fails to capture the full structure of EHR. We propose ORA, a marked time-to-event pretraining objective that jointly models event timing and associated measurements. Across multiple datasets, downstream tasks, and model architectures, this objective consistently yields more generalizable representations than next-token prediction and pretraining losses that ignore continuous measurements. Importantly, the proposed objective yields improvements beyond traditional classification evaluation, including better regression and time-to-event prediction. Beyond introducing a new family of FMs, our results suggest a broader takeaway: pretraining objectives that account for EHR structure are critical for expanding downstream capabilities and generalizability

Figures (5)

• Figure 1: Our work introduces ORA, a marked time-to-event pretraining loss that accounts for the value and irregular timing of EHR, demonstrating the importance of the loss design for improved downstream capabilities and generalizability.
• Figure 2: Relative improvement (%) over the strongest Next-Token Prediction (NTP) baseline on each task for the Transformer architecture. For each dataset (MIMIC-IV and CUMC), we compare ORA (green circles), the best Temporal Point Process (TPP) pretraining model (blue squares), and task-specific supervised baselines (red triangles) across classification, time-to-event, and regression tasks.
• Figure 3: Relative improvement (%) over the strongest Next-Token Prediction (NTP) baseline on each task for the Mamba architecture. For each dataset (MIMIC-IV and CUMC), we compare ORA (green circles), the best Temporal Point Process (TPP) pretraining model (blue squares), and task-specific supervised baselines (red triangles) across classification, time-to-event, and regression tasks.
• Figure 4: Visualization of Cohort Definition. The detailed definition of at-risk, prediction time, and case events can be found in Appendix B of pang2025fomoh.
• Figure 5: Efficient prediction head: We share the same projection layer for both nonnumerical and numerical codes and then use two different projection heads to output the separate probability matrices.