PatientWM: Towards Generalizable Clinical World Models via Unified Patient States
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
| MIMIC-III: Core-7 | MIMIC-III: Common Event-7 | MIMIC-IV: Common Event-7 | |||||||
|---|---|---|---|---|---|---|---|---|---|
| Method | AUPRC | AUROC | Brier | AUPRC | AUROC | Brier | AUPRC | AUROC | Brier |
| 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 |
| MIMIC-III Observed Laboratory Values | |||||
|---|---|---|---|---|---|
| Method | 6 h nMAE | 24 h nMAE | 48 h nMAE | 96 h nMAE | Mean nMAE |
| 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 |
| MIMIC-III Core-7 | |||
|---|---|---|---|
| Method | AUPRC | AUROC | Brier |
| 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 |
| 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 |
| Method | AUPRC | AUROC | Brier |
|---|---|---|---|
| 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 |
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 |
| 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 |
| Experiment | Final | 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 |
| Setting | Value |
|---|---|
| Width | 256 |
| Layers | 4 |
| Attention heads | 8 |
| Feed-forward hidden width | 1,024 |
| Activation | GELU |
| Dropout | 0.1 |
| Tensor role | Shape | Meaning |
| laboratory item identities | sparse item indices for each state window | |
| laboratory numeric features | value summaries, counts, timing, and availability features | |
| laboratory token mask | valid, non-padding laboratory positions | |
| clinical-fact values | current-window binary clinical facts | |
| clinical-fact mask | availability of each clinical fact | |
| waveform embedding | frozen waveform-derived representation |
| 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 |
| AUPRC | ||||||||
|---|---|---|---|---|---|---|---|---|
| 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 |
| 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 | ||||||||
| 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 |
| 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 | ||||||||
| 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 |
| Dataset | Method | Macro AUPRC | Macro AUROC | Brier |
|---|---|---|---|---|
| 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 |
| Event-7 macro-AUPRC | Observed-laboratory nMAE | |||||||
|---|---|---|---|---|---|---|---|---|
| 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 |
| Method | Macro AUPRC | Macro AUROC | Macro Brier |
|---|---|---|---|
| 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 |