Latent Dynamics · 2026-10-20

Forecasting Patient Trajectories from the Electronic Health Records with Action-Conditioned World Models

Simon A. Lee, Jiahang Hank Sha, Soomin Chung, Kai Akamatsu, Eran Halperin, Jeffrey N. Chiang

AI-assisted summary based on the paper. Consult the linked source for methods and limitations. Benchmark numbers appear after editor verification.

TL;DR

The authors introduce an action-conditioned latent world model for electronic health records that separates patient states from observed clinical actions. Using a state-dependent linearization predictor, the model propagates local latent changes for time-resolved prediction, multi-step rollouts, and action-conditioned simulation, then decodes rollouts for forecasting evaluation. They report competitive, task-dependent downstream performance vs larger EHR foundation models with a ~30× smaller parameter footprint.

Why it matters

This work proposes a controlled latent-dynamics formulation for longitudinal EHR: learned transitions are explicitly conditioned on observed actions rather than only autoregressively predicting future events. The authors evaluate fidelity in latent space and in decoded clinical feature space, test downstream predictive utility, and show that perturbing action sequences changes generated trajectories. This frames EHR modeling as forecasting and simulation in an explicit time-indexed representation.

Method

Models EHR as controlled trajectories (state x_t, action a_t) with discrete-time windows (Δ=6 hours) from MIMIC-IV/ MEDS ICU setting.
Uses a latent encoder f_φ(x_t) to get z_t (D=128) and a state- and action-dependent linearization transition F_θ(z_t,a_t) to propagate latent increments for multi-step rollout.
Trains with latent prediction loss comparing predicted vs encoder-derived future latent states, regularized with SIGReg; evaluates rollout fidelity and decoded forecasting.

Limitation

The authors caution that intervention-conditioned trajectories should not be interpreted as causal estimates of treatment effects; they characterize how learned observational dynamics respond to alternative action sequences under appropriate causal assumptions.

Abstract (from arXiv)

We introduce a predictive world model for electronic health records (EHR) that represents patient histories as action conditioned latent trajectories. Rather than modeling clinical events solely through autoregressive token prediction, our framework explicitly separates patient states from observed clinical actions and learns their evolution through a state and action conditioned transition operator. At its core is a state dependent linearization predictor that propagates local changes in latent patient state, enabling time resolved prediction, multi step rollout, and action conditioned simulation. Empirically, we show that generated latent trajectories preserve downstream predictive signal, temporal structure, and subgroup level statistical relationships present in encoder derived representations. We further decode free running rollouts into the original clinical feature space and evaluate future clinical measurements against observed values using task matched forecasting baselines, demonstrating that the learned dynamics retain predictive information beyond latent space agreement. Across binary classification and time to event tasks, EHR WM achieves competitive, though task dependent, performance relative to substantially larger EHR foundation models while using a 30x smaller parameter footprint. Finally, perturbing observed action sequences produces corresponding changes in generated trajectories, demonstrating the action sensitivity of the learned dynamics. Together, these results support world modeling as a complementary framework for learning explicit, time indexed representations of patient evolution from longitudinal EHR.

Publication record

Bibliographic evidence describes where and when a record appeared. It does not validate scientific claims.

First public release: 2026-08-16 · Indexed by Pub-AI: 2026-10-20 UTC

preprint