PatientWM: Towards Generalizable Clinical World Models via Unified Patient States

Ziwen Kan; Jie Peng; Bo Yan; Davis Hobley; Zheng Zhang; Yu Tian; Tianlong Chen; Song Wang

cs.AIOct 12, 2026

PatientWM: Towards Generalizable Clinical World Models via Unified Patient States

Authors: Ziwen Kan, Jie Peng, Bo Yan, Davis Hobley, Zheng Zhang, Yu Tian, Tianlong Chen, Song Wang

Organizations: University of Science and Technology of China · University of Central Florida · Princeton University · Nanyang Technological University · University of North Carolina at Chapel Hill

Abstract

Clinical decisions depend on anticipating how a patient's condition will evolve as treatment is adjusted and new evidence becomes available. Existing medical world models capture patient-state evolution, but often define the evolving state around a particular modality or task, limiting reuse as inputs or prediction targets change. These models also retain history through observation sequences or recurrent or latent summaries rather than exposing preceding patient states and clinical actions as an ordered transition context. This distinction matters because similar current measurements may follow different recent changes, treatment exposures, and observed responses. To tackle these limitations, we propose PatientWM, a clinical world-model framework built around a unified patient state and sequential latent history. Availability-aware encoders integrate observations available at each decision time into the patient state, while typed decoders produce task-specific clinical readouts without defining the state around any single target. PatientWM further conditions each next-state prediction on the current state and supplied action together with a bounded, ordered history of preceding latent patient states and clinical actions, preserving recent treatment and response context. Experiments across MIMIC-III, MIMIC-IV, and three external clinical datasets show superior performance in clinical prediction, multi-step forecasting, and action-planning evaluation. Transfer experiments further demonstrate the generalizability of PatientWM's patient state. Together, these results establish its ability to forecast clinical evolution and support planning.

Figures & tables

Figure 1: Conceptual Overview of PatientWM. Conventional clinical prediction focuses on learning task-specific predictions from heterogeneous observations. PatientWM instead learns a unified patient state from available observations and models how this state evolves under clinical actions and recent history. The learned state supports multiple clinical readouts and multi-step forecasting.

Figure 2: PatientWM encodes available observations and the supplied action into a unified patient state $z_{t}$ and action representation $\alpha_{t}$ . Conditioned on $z_{t}$ , $\alpha_{t}$ , and sequential latent history $H_{t}$ , the transition predicts $\widehat{z}_{t+1}$ . The query decoder produces clinical event and outcome predictions and query-specific readouts, including laboratory values. Black arrows show the training-time forward pass with observed history; blue arrows show recursive multi-step rollout, where predicted states update the history for subsequent transitions.

	MIMIC-III: Core-7			MIMIC-III: Common Event-7			MIMIC-IV: Common Event-7
Method	AUPRC $\uparrow$	AUROC $\uparrow$	Brier $\downarrow$	AUPRC $\uparrow$	AUROC $\uparrow$	Brier $\downarrow$	AUPRC $\uparrow$	AUROC $\uparrow$	Brier $\downarrow$
medDreamer	0.0579	0.5130	0.2192	0.0278	0.5049	0.2631	0.0341	0.5050	0.2547
SMB-Structure	0.0638	0.5412	0.2165	0.0522	0.6000	0.2278	0.0604	0.6375	0.2572
SepsisAgent	0.0772	0.6191	0.2181	0.0517	0.6631	0.2457	0.1742	0.7971	0.1875
ACWM	0.0762	0.5945	0.2235	0.0444	0.5879	0.2444	0.1768	0.7959	0.1846
PatientWM (Ours)	0.2227	0.7974	0.0761	0.0598	0.6837	0.1964	0.1784	0.7988	0.1704

Table 1: Clinical event and outcome prediction on the MIMIC-III and MIMIC-IV test sets. The three panels report MIMIC-III Core-7 outcomes and Common Event-7 readouts on both datasets. Bold indicates the best result in each column.

Figure 3: Per-target AUROC profiles for (a) MIMIC-III Core-7 outcomes, (b) MIMIC-III Common Event-7 readouts, and (c) MIMIC-IV Common Event-7 readouts. Corresponding per-target AUPRC and Brier results are reported in Appendix D.2 .

MIMIC-III Observed Laboratory Values
Method	6 h nMAE $\downarrow$	24 h nMAE $\downarrow$	48 h nMAE $\downarrow$	96 h nMAE $\downarrow$	Mean nMAE $\downarrow$
medDreamer	0.4296	0.4791	0.5387	0.6079	0.5138
SMB-Structure	0.6257	0.6236	0.6300	0.6286	0.6270
SepsisAgent	0.4295	0.4792	0.5388	0.6086	0.5140
ACWM	0.4296	0.4789	0.5386	0.6081	0.5138
PatientWM (Ours)	0.3248	0.3888	0.4430	0.4966	0.4133

Table 2: Laboratory-value prediction on the MIMIC-III and MIMIC-IV test cohorts across 6–96-hour horizons. Bold marks the best result within each panel.

MIMIC-III Core-7
Method	AUPRC $\uparrow$	AUROC $\uparrow$	Brier $\downarrow$
PatientWM (Ours)	0.2227	0.7974	0.0761
Without action	0.0861	0.6067	0.2024
Without history	0.0840	0.5507	0.2523
Laboratory only	0.0800	0.5874	0.2171
Clinical facts only	0.0777	0.5888	0.2137

Table 3: MIMIC-III Core-7 component and modality ablations.

Figure 4: MIMIC-IV multi-step rollout comparison. Panel (a) reports recursive Event-7 macro-AUPRC, and panel (b) reports observed-laboratory nMAE from the typed laboratory decoder. Detailed results are reported in Appendix D.3 .

Method	Top-1 covered	Top-2 covered
medDreamer	0.4713	0.7997
SMB-Structure	0.3616	0.7925
SepsisAgent	0.6920	0.9081
ACWM	0.3576	0.7833
PatientWM (Ours)	0.7477	0.9126

Table 4: Recorded-action agreement in the action-planning evaluation.

Figure 5: Frozen-state transfer on MIMIC-III Common Event-7. Panels (a)–(c) report AUPRC, AUROC, and Brier score, respectively, with full PatientWM and the four Table 1 baselines shown on the same test support. Higher is better for AUPRC and AUROC; lower is better for Brier.

Method	AUPRC $\uparrow$	AUROC $\uparrow$	Brier $\downarrow$
medDreamer	0.0323	0.4985	0.2204
SMB-Structure	0.0450	0.5944	0.2061
SepsisAgent	0.0515	0.5904	0.2035
ACWM	0.0386	0.5532	0.2163
PatientWM (Transfer)	0.0517	0.6418	0.2023

Table 5: MIMIC-III to MIMIC-IV adaptation with 1% target training data.

Figure 6: External AUROC comparison. Panels (a), (b), and (c) report results on ALOTT, VitalDB, and MC-MED, respectively. Higher is better in all panels.

Appendix figures & tables12 assets

Supplementary material from the paper’s appendix.

Appendix

Split	Transitions	Event windows	Coverage	Events	Episodes
Train	685,561	620,559	90.52%	11,870,214	41,172
Validation	85,295	77,490	90.85%	1,460,107	5,141
Test	86,979	78,673	90.45%	1,531,140	5,146

Appendix

Table 6: Structured event-ledger coverage. Coverage is the fraction of transitions whose corresponding event window contains at least one executed event of any supported type. This dataset statistic does not measure boundary-known action availability or prediction performance.

Decoder	Output	Role
clinical event/outcome	target-specific probability	clinical readout
waveform	512-D embedding	waveform prediction and cache update
clinical fact	six probabilities	transition supervision
lab availability	693 probabilities	observation-process supervision
joint lab value	693 normalized values	transition supervision and cache update
queried lab value	one normalized value	observed-laboratory prediction

Appendix

Table 7: Typed decoder outputs and their roles in training and evaluation.

Experiment	$H$	Final $\lambda_{r}$	Additional objective
MIMIC-III transition rollout	16	0.25	residual and waveform terms
MIMIC-IV primary rollout	16	0.50	residual term
One-step/readout training	–	0	none

Appendix

Table 8: Recursive training settings used by the reported transition experiments.

Setting	Value
Width	256
Layers	4
Attention heads	8
Feed-forward hidden width	1,024
Activation	GELU
Dropout	0.1

Appendix

Table 9: Implementation settings for the main MIMIC-III causal-Transformer transition and laboratory readout.

Tensor role	Shape	Meaning
laboratory item identities	$[B,C+1,N_{L}]$	sparse item indices for each state window
laboratory numeric features	$[B,C+1,N_{L},11]$	value summaries, counts, timing, and availability features
laboratory token mask	$[B,C+1,N_{L}]$	valid, non-padding laboratory positions
clinical-fact values	$[B,C+1,K_{y}]$	current-window binary clinical facts
clinical-fact mask	$[B,C+1,K_{y}]$	availability of each clinical fact
waveform embedding	$[B,C+1,D_{W}]$	frozen waveform-derived representation

Appendix

Table 10: Tensor contract for the main MIMIC-III transition model. The state axis contains

C+1

windows because

C

transitions connect consecutive source and target states.

Method	Retained core	Interface adaptation	Fairness controls
medDreamer	Dreamer-style recurrent state-space dynamics	Replace task-specific interfaces with shared laboratory and clinical-fact encoders, frozen waveform embeddings, action inputs, and typed readouts	Same split, windows, visible inputs, action timing, targets, masks, training budget, selection rule, support, and evaluator
SMB-Structure	Frozen predictive representation	Map the same laboratory, clinical-fact, and frozen-waveform representations into the frozen core, then fit an action-aware transition and shared typed readouts	Same split, visible inputs, action timing, clinical targets, masks, support, and evaluator
SepsisAgent	Recurrent history-to-next-state update	Replace native feature and output layers with shared laboratory and clinical-fact encoders, frozen waveform embeddings, action inputs, and typed readouts	Same split, windows, visible inputs, action timing, targets, masks, training budget, selection rule, support, and evaluator
ACWM	Latent state-transition update	Replace native signal/label interface with shared laboratory and clinical-fact encoders, frozen waveform embeddings, action inputs, and typed readouts	Same split, windows, visible inputs, action timing, targets, masks, training budget, selection rule, support, and evaluator

Appendix

Table 11: Adaptation and fairness contract for the MIMIC-IV comparisons. Each method retains its defining source-inspired core while using the stated common benchmark interface.

AUPRC $\uparrow$
Method	Hypercapnic acidemia	Severe acidemia	Creatinine AKI	Urine-output AKI	GCS decrease	Incident CAM+	24-hour mortality	Macro AUPRC
medDreamer	0.0581	0.0262	0.0226	0.2236	0.0102	0.0401	0.0244	0.0579
SMB-Structure	0.0637	0.0312	0.0155	0.2484	0.0115	0.0597	0.0170	0.0638
SepsisAgent	0.0660	0.0404	0.0270	0.2518	0.0198	0.1003	0.0354	0.0772
ACWM	0.0647	0.0389	0.0227	0.2637	0.0194	0.0982	0.0261	0.0762
PatientWM (Ours)	0.2556	0.1747	0.1192	0.5179	0.0386	0.1092	0.3440	0.2227

Appendix

Table 12: Per-target MIMIC-III Core-7 test results at the native target horizons. Macro scores are unweighted averages across the seven targets; the AUROC rows correspond to Figure 3 (a).

Method	Any acute deterioration	New invasive ventilation	New pressor	Circulation pressor	Respiratory IMV	Sepsis-3 onset	Dynamic ICU mortality	Macro
MIMIC-III Common Event-7 Test Panel: AUPRC $\uparrow$
medDreamer	0.0423	0.0335	0.0185	0.0234	0.0276	0.0317	0.0178	0.0278
SMB-Structure	0.0733	0.0725	0.0302	0.0495	0.0518	0.0694	0.0187	0.0522
SepsisAgent	0.0656	0.0530	0.0477	0.0507	0.0416	0.0786	0.0250	0.0517
ACWM	0.0517	0.0324	0.0404	0.0501	0.0437	0.0668	0.0256	0.0444
PatientWM (Ours)	0.0691	0.0620	0.0409	0.0718	0.0483	0.0758	0.0505	0.0598

Appendix

Table 13: Per-target MIMIC-III Common Event-7 test results. Macro scores are unweighted averages across the seven targets; the AUROC rows correspond to Figure 3 (b).

Method	Any acute deterioration	New invasive ventilation	New pressor	Circulation pressor	Respiratory IMV	Sepsis-3 onset	Dynamic ICU mortality	Macro
MIMIC-IV Common Event-7 Test Panel: AUPRC $\uparrow$
medDreamer	0.0653	0.0428	0.0253	0.0257	0.0346	0.0272	0.0178	0.0341
SMB-Structure	0.0901	0.0589	0.0548	0.0514	0.0528	0.0494	0.0657	0.0604
SepsisAgent	0.2733	0.1193	0.2093	0.2044	0.1228	0.0819	0.2087	0.1742
ACWM	0.2725	0.1202	0.2107	0.2042	0.1259	0.0827	0.2214	0.1768
PatientWM (Ours)	0.2816	0.1253	0.2066	0.1970	0.1265	0.0843	0.2278	0.1784

Appendix

Table 14: Per-target MIMIC-IV Common Event-7 test results. Macro scores are unweighted averages across the seven targets; the AUROC rows correspond to Figure 3 (c).

Dataset	Method	Macro AUPRC $\uparrow$	Macro AUROC $\uparrow$	Brier $\downarrow$
ALOTT	medDreamer	0.0675	0.5659	0.0474
	SMB-Structure	0.0651	0.5672	0.0525
	SepsisAgent	0.2007	0.7864	0.0451
	ACWM	0.1838	0.7757	0.0451
	PatientWM (Ours)	0.2268	0.7915	0.0437
VitalDB	medDreamer	0.0538	0.5082	0.0441

Appendix

Table 15: External test evaluation of PatientWM and available clinical-model baselines. All rows use the same per-target observed test support within each dataset. Bold marks the best value within each dataset and metric.

	Event-7 macro-AUPRC $\uparrow$				Observed-laboratory nMAE $\downarrow$
Method	6 h	24 h	48 h	96 h	6 h	24 h	48 h	96 h
medDreamer	0.0474	0.0733	0.1038	0.1052	0.4243	0.4527	0.5093	0.5728
SMB-Structure	0.1141	0.0546	0.0493	0.0425	0.6528	0.7082	0.8018	0.9533
SepsisAgent	0.1637	0.0844	0.0927	0.0939	0.4191	0.4520	0.5085	0.5690
ACWM	0.1378	0.0665	0.0761	0.0714	0.4240	0.4519	0.5084	0.5729
PatientWM (Ours)	0.1755	0.1005	0.1090	0.1308	0.4093	0.4161	0.4556	0.4958

Appendix

Table 16: MIMIC-IV multi-step rollout on the same test trajectories and targets. We report recursive Event-7 macro-AUPRC and observed-laboratory nMAE from the typed laboratory decoder. Bold marks the best value in each column.

Method	Macro AUPRC $\uparrow$	Macro AUROC $\uparrow$	Macro Brier $\downarrow$
medDreamer	0.0278	0.5049	0.2631
SMB-Structure	0.0522	0.6000	0.2278
SepsisAgent	0.0517	0.6631	0.2457
ACWM	0.0444	0.5879	0.2444
PatientWM (Ours)	0.0598	0.6837	0.1964
PatientWM (Frozen transfer)	0.0500	0.6122	0.2401

Appendix

Table 17: MIMIC-III Common Event-7 frozen-state transfer on the Table 1 test support. Bold marks the best result in each column.

Explore similar work

ChronoMedicalWorld: A Medical World Model for Learning Patient Trajectories from Longitudinal Care Data
May 21, 2026Jiangyuan Wang, Xuyong Chen, Junwei He +3Latent Dynamics Modeling Healthcare
Beijing KidneyTec Medical Technology Co., Ltd., Beijing, China
Medical world models: representing medical states, modelling clinical dynamics and guiding intervention policies
Jun 15, 2026Ke Liu, Mengxuan Li, Yanyi Bao +4World Model Learning Clinical Decision Support
College of Computer Science, Zhejiang University, Hangzhou, China. · School of Medicine, Zhejiang University, China. · Department of Biomedical Informatics, Harvard University, USA.
CLARITY: Medical World Model for Guiding Treatment Decisions by Modeling Context-Aware Disease Trajectories in Latent Space
Dec 8, 2025Tianxingjian Ding, Yuanhao Zou, Chen Chen +2Healthcare Clinical Decision Support
Institute of Artificial Intelligence, University of Central Florida

May 21, 2026J/K move · Enter open · S save

ChronoMedicalWorld: A Medical World Model for Learning Patient Trajectories from Longitudinal Care Data

Jiangyuan Wang, Xuyong Chen, Junwei He, Xu Xu+2

Beijing KidneyTec Medical Technology Co., Ltd., Beijing, China

Latent Dynamics Modeling Healthcare Action-Conditioned World Models Clinical Outcome Prediction Latent World Models Longitudinal EHR Modeling

Open paper AskPDF arXiv

Long-horizon clinical simulation -- predicting how a patient's physiology evolves over years under specified interventions -- is central to chronic-disease care, yet existing electronic health record (EHR) models are predominantly discriminative, and general-purpose large language models drift under repeated interventions. We propose the ChronoMedicalWorld Model (CMWM), an action-conditioned latent world-model framework for learning patient trajectories from longitudinal care data. CMWM couples a joint-embedding state encoder with a wide action encoder that admits both structured intervention indicators and free-text communication embeddings, and trains a recurrent latent transition module under a six-term objective: next-observation supervision, next-latent prediction, SIGReg latent regularisation, and three physiology-aware shape priors (slope, continuity, large-jump penalty). A closed-loop rollout-prefix protocol matches training to deployment, so the model is optimised against the same multi-step error it exhibits at inference. As a concrete case study, we instantiate CMWM for annual estimated glomerular filtration rate (eGFR) trajectory forecasting in chronic kidney disease (CKD). On a 2{,}232-patient nephrology cohort, the CKD instantiation achieves a dynamic-50% history rollout test mean absolute error (MAE) of 7.384 and root-mean-square error (RMSE) of 10.256, against 7.964 and 11.069 for a tuned GPT-5.5 structured-prompting baseline (

-7.28\%

MAE,

-7.35\%

RMSE), with the gain dominated by the dialogue portion of patient--health-coach communication. The framework is not CKD-specific: its architecture, loss design, and training protocol apply to any chronic condition that can be cast as periodic clinical state interleaved with structured and conversational interventions.

Abstract

Explore similar work

ChronoMedicalWorld: A Medical World Model for Learning Patient Trajectories from Longitudinal Care Data

Medical world models: representing medical states, modelling clinical dynamics and guiding intervention policies

CLARITY: Medical World Model for Guiding Treatment Decisions by Modeling Context-Aware Disease Trajectories in Latent Space