PaperlayerView exact revision on arXiv

Source-linked AI summary

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

arXiv:2610.17059v1cs.AI

TL;DR

Existing medical world models often tie patient states to fixed inputs or tasks and retain history without exposing ordered state-action context, although recent treatment responses can distinguish similar current states. PatientWM combines a unified patient state with sequential latent history to model clinical evolution. Across multiple datasets, it reports superior performance in clinical prediction, multi-step forecasting, and action-planning evaluation.

  • Problem

    Existing medical world models often define patient states for fixed inputs and tasks and retain history without exposing ordered preceding states and actions, making it challenging to learn how similar current states from different treatment-response trajectories may evolve.

    Source: §1 · p. 2

  • Method

    PatientWM fuses available observations into a unified patient state, uses typed decoders for clinical readouts, and conditions transitions on current state, supplied action, and ordered latent state-action history.

    Source: §1 · p. 2

  • Results

    Experiments across MIMIC-III, MIMIC-IV, and three external clinical datasets show superior performance in clinical prediction, multi-step forecasting, and action-planning evaluation.

    Source: §1 · p. 2

  • Takeaways & Limitations

    The reported results demonstrate PatientWM’s ability to forecast clinical evolution and support planning.

    Source: §1 · p. 2

    The formulation targets the conditional predictive distribution induced by observed clinical processes, not an interventional treatment-response law.

    Source: §3.1 · p. 3

Abstract

from arXiv · show

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.

1 INTRODUCTION

Clinical decision-making requires modeling how patients evolve under treatment, but existing systems often provide task-specific outputs without a reusable state or explicit ordered treatment-response context. PatientWM combines a unified patient state with ordered latent state-action history, and experiments report superior performance across several clinical evaluation settings.

1 INTRODUCTION

  • Clinical decisions repeatedly adjust treatment based on current status, prior treatment, and observed responses, while conventional systems commonly map longitudinal observations to predefined task-specific outcomes.Such outputs do not by themselves define a shared state that can be updated after treatment and carried forward to model subsequent evolution.Source: §1 · p. 1
  • Existing medical world models often fix the clinical inputs and prediction tasks defining patient states, limiting reuse as modalities or targets change.They also retain history in observation sequences or recurrent or latent context rather than exposing ordered preceding states and actions to each transition.Source: §1 · p. 2
  • PatientWM learns a unified patient state from available observations, models its evolution under clinical actions, and uses typed decoders for task-specific outputs.The shared state is intended to support downstream predictions without defining the state around one target.Source: §1 · p. 2, Figure 1 · p. 2
  • PatientWM represents recent history as ordered latent patient-state and clinical-action pairs, conditioning each next-state prediction on that history, the current state, and the supplied action.This explicit trajectory context is intended to help distinguish similar current states reached through different treatment-response histories.Source: §1 · p. 2
  • 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.Source: §1 · p. 2

2 RELATED WORK

Related work spans clinical prediction systems, general world models, and medical trajectory models. PatientWM is positioned as integrating heterogeneous observations into a shared state while making ordered latent state-action history part of each transition.

Clinical Prediction

  • Clinical prediction systems use longitudinal EHRs and physiological signals for diagnoses, deterioration, events, and outcomes, but are typically organized around predefined targets.Those outputs do not by themselves represent how a patient’s condition evolves between observations and treatment updates.Source: §2 · p. 3

World Models

  • World models connect internal environment representations with action-conditioned future prediction, enabling trajectory simulation for planning, policy learning, or evaluation.The principle has recently been adapted to medicine to model patient states and their evolution under clinical interventions.Source: §2 · p. 3

Medical World Models

  • Medical world models vary in their evolving state: observable record, pathway, or physiological spaces; latent rollouts; or specialized imaging, postoperative, disease-specific, and signal settings.PatientWM instead integrates heterogeneous observations into a unified state with typed readouts and conditions transitions on ordered latent states and supplied clinical actions.Source: §2 · p. 3

3 METHOD

PatientWM predicts clinical evolution in a shared latent patient state, conditioning transitions on supplied actions and ordered recent history. Availability-aware state construction, typed readouts, and recursive rollout support predictions across targets and multiple steps.

3.1 PROBLEM FORMULATION

  • The next predictive state is conditioned on the current state, supplied action, and recent history; the conditional target is not an interventional treatment-response law.The implemented transition predicts toward a stopped-gradient momentum target.Source: §3.1 · p. 3

3.2 STATE REPRESENTATION

  • Available modalities are encoded and fused into a unified predictive state, which typed decoders map to clinical targets without defining the state around one target.The state supports clinical-event, outcome, and continuous laboratory readouts.Source: §3.2 · p. 4

3.3 ACTION-CONDITIONED LATENT DYNAMICS

  • Each transition uses an ordered context of recent latent state–action pairs alongside the current state and action, distinguishing similar states reached through different paths.The residual transition represents changes relative to the current predictive state, with an action- and history-conditioned update.Source: §3.3 · pp. 4–5
  • A momentum-updated observation encoder supplies a slowly evolving state target, while latent alignment and typed losses ground predictions in observable clinical quantities.Typed objectives are included when their targets are observed.Source: §3.3 · p. 5

3.4 FACTUAL ROLLOUT

  • During rollout, each predicted state becomes the next transition’s input and is appended to retained history; future observations are not re-encoded.Decoder-side priors carry across rollout steps but do not fully write generated clinical observations back into the state.Source: §3.4 · p. 5

4 EXPERIMENTS

Experiments evaluate PatientWM’s clinical readouts, transition components, recursive forecasting, bounded action planning, state reuse, and external-cohort performance. PatientWM leads the reported comparisons in clinical prediction, multi-step forecasting, action-planning agreement, and AUROC across three external datasets.

4.1 EXPERIMENTAL SETUP

  • Experiments use six-hour MIMIC-III and MIMIC-IV trajectory units and matched visible inputs, action timing where applicable, target definitions, support, and evaluators.Evaluation covers typed readouts, component ablations, recursive state reuse, bounded-candidate action planning, frozen-state reuse, adaptation, and external cohorts.Source: §4.1 · pp. 5–6

4.2 FACTUAL PREDICTION AND CLINICAL READOUT

  • PatientWM achieves the strongest overall aggregate performance across the three clinical-target panels, and outperforms all baselines in laboratory prediction at every evaluated horizon on both MIMIC cohorts.The panels cover MIMIC-III Core-7 outcomes and Common Event-7 readouts on MIMIC-III and MIMIC-IV; laboratory horizons span 6–96 hours.Source: §4.2 · p. 6, Table 1 · p. 6, Table 2 · p. 7

4.3 ABLATIONS AND MECHANISM ANALYSIS

  • Removing ordered history degrades Core-7 performance more than removing action conditioning, while restricting state construction to any single modality also lowers performance.The ablations use the same MIMIC-III split, readout, optimization budget, checkpoint rule, and evaluator as the full configuration.Source: §4.3 · pp. 6–7, Table 3 · p. 6

4.4 MULTI-STEP PATIENT-STATE ROLLOUT

  • PatientWM achieves the highest Event-7 macro-AUPRC and lowest observed-laboratory nMAE at every recursive rollout horizon, although laboratory error rises and event ranking varies more with horizon.Rollouts start from the same observed states and use recorded actions; future events and laboratory values are reserved for evaluation. Figure 4 shows the trends.Source: §4.4 · p. 7, Figure 4 · p. 8

4.5 ACTION-PLANNING EVALUATION

  • PatientWM achieves the highest covered top-1 and top-2 agreement when ranking bounded candidate actions against clinicians’ recorded actions represented in the candidate set.Methods predict future event and laboratory burdens without access to observed future values, then combine them using the same fixed scoring rule.Source: §4.5 · p. 8, Table 4 · p. 8

4.6 FROZEN-STATE REUSE

  • Without Common Event-7 supervision during world-model training, frozen PatientWM ranks third in AUPRC and second in AUROC and Brier against four baselines.Only a new typed decoder is trained for the downstream labels; the comparison uses the same target definitions, per-target support, and evaluation.Source: Figure 5 · p. 8, §4.6 · p. 9

4.8 EXTERNAL FRAMEWORK EVALUATION

  • PatientWM achieves the highest AUROC on ALOTT, VitalDB, and MC-MED, spanning sequential next-state and pointwise prediction settings.Within each dataset, methods use the same visible inputs, targets, splits, training budget, and evaluator.Source: §4.8 · p. 9, Figure 6 · p. 9

5 CONCLUSION

PatientWM is presented as a reusable predictive patient state with ordered state–action history, supporting clinical prediction, forecasting, and bounded action planning. The paper also specifies its AI-use review, data-use boundaries, and reproducibility materials.

5 CONCLUSION

  • PatientWM learns a reusable patient state and ordered state–action history, with reported strong prediction and rollout performance and higher recorded-action agreement within fixed candidate sets.The conclusion frames these results as support for clinical forecasting and bounded action planning.Source: §5 · p. 10

AI USE STATEMENT

  • The authors used generative AI for literature work, ideation, experiments, coding, result organization, figures, and manuscript drafting, and reviewed outputs against sources and records.The authors take responsibility for the final text, claims, code, and artifacts.Source: AI USE STATEMENT · p. 11

ETHICS STATEMENT

  • The retrospective study uses deidentified, access-controlled data and does not provide treatment recommendations or estimate causal treatment effects.Row-level records remain in approved environments; public artifacts are limited to code, schemas, hashes, and aggregate statistics.Source: ETHICS STATEMENT · p. 11

REPRODUCIBILITY STATEMENT

  • The main text presents evaluation roles, protocols, and primary results, while the appendix documents data construction, implementation, controls, supplementary evaluations, parameters, and splits.Source: REPRODUCIBILITY STATEMENT · p. 11

A DATASETS AND INFORMATION BOUNDARY

The appendix describes six-hour MIMIC-III and MIMIC-IV transition datasets, their action-information boundaries, and external-cohort evaluations. It also defines shared event readouts and distinguishes prediction targets from recursive state labels.

A.1 MIMIC-III STATE, OUTCOMES, AND SPLITS

  • The MIMIC-III cohort contains 53,329 ICU episodes, 911,164 windows, and 857,835 adjacent transitions, with subject-level train/validation/test splits and seven core binary outcomes.No subject crosses splits; the outcomes use non-treatment-derived target definitions.Source: §A.1 · p. 14

A.2 STRUCTURED ACTION LEDGER

  • Only boundary-known, chart-available treatment information enters the current action input; later records may enter factual replay as supplied actions, not causal intervention assignments.Event-ledger coverage counts transitions with at least one executed event and does not measure boundary-known action availability.Source: §A.2 · p. 14, Table 6 · p. 14

A.3 MIMIC-IV PORT

  • MIMIC-IV supports Common Event-7 readout evaluation, transition training for rollout and action planning, typed laboratory decoding, and transfer evaluation from a frozen MIMIC-III predictor.The native-cohort roles share six-hour transitions, patient-disjoint splits, availability masks, and leakage boundaries.Source: §A.3 · p. 14

A.4 COMMON EVENT-7 PANEL

  • The Common Event-7 sidecar enables direct event comparisons across MIMIC-III and MIMIC-IV, while remaining outside the recursive state-label matrix.Its six shared heads are next-window events; dynamic ICU mortality retains a six-hour gap and 24-hour target horizon.Source: §A.4 · pp. 14–15

A.5 EXTERNAL COHORTS

  • External evaluations use ALOTT and VitalDB adjacent-state forecasts and MC-MED pointwise prediction, with specified task exclusions and pre-anchor context.VitalDB aggregates 22 forecasting tasks and excludes three concurrent-state controls; one MC-MED task is excluded after waveform quality control.Source: §A.5 · p. 15

A.6 DATA GOVERNANCE AND INTENDED USE

The appendix describes governed data use and an architecture that encodes available observations into a shared patient state, predicts transitions, and supplies typed clinical readouts. Recursive evaluation carries the predicted state and causal readout caches forward.

A.6 DATA GOVERNANCE AND INTENDED USE

  • PatientWM is intended for retrospective representation and factual forecasting, not bedside monitoring, treatment selection, or autonomous simulation.Clinical use would require external and prospective validation, subgroup and calibration audits, and a causal design suited to the intended decision.Source: §A.6 · p. 15

B.1 ARCHITECTURE

  • The observation encoder fuses available modality-specific representations, while the action encoder represents chart-available active treatments separately from the observation state.The transition model uses a causal Transformer and predicts a gated residual around the current latent state.Source: §B.1 · pp. 15–16

B.2 TYPED DECODERS AND RECURSIVE STATE

  • Typed decoders map the predicted patient state to requested clinical targets, keeping target-specific outputs separate from the shared recursively reused state.Clinical-event and outcome readouts share the patient state but use separate output parameters and losses.Source: §B.1 · p. 17, §B.2 · p. 17
  • Continuous readouts use gated residual updates around causal persistence priors, including the last observed laboratory value or training-set mean for never-observed items.Waveform priors use the most recent quality-qualified embedding, with a learned fallback when none is available.Source: §B.2 · p. 17
  • The accumulated laboratory decoder combines the rolled state, queried item, source-visible history, and cross-laboratory context to predict observed laboratory values.Its source-only history features include running statistics and trends; future target-ledger values do not refresh those features.Source: §B.2 · p. 18
  • Recursive evaluation reuses the predicted patient state and causal continuous-readout caches, while clinical-event, outcome, and availability probabilities are decoded when requested.Those caches support causal priors in later predictions; the latent state remains the shared transition input.Source: §B.2 · p. 18

B.3 TRAINING DETAILS

Training combines latent-state prediction with available typed targets and a generated-history branch for multi-step use. The appendix also specifies recursive settings and tensor alignment, masking, and history limits.

B.3 TRAINING DETAILS

  • The direct transition objective jointly trains latent prediction and available typed targets, with masked losses for observed clinical facts, laboratory availability, and numeric laboratory values.The objective also includes waveform targets when a quality-qualified next-window embedding is available.Source: §B.3 · pp. 18–19
  • The reported transition configurations assign weight 1.0 to each enabled top-level objective term, while unavailable targets are removed by masks.Source: §B.3 · p. 18
  • Multi-step transition training adds a generated-history branch that advances predicted state and causal caches while logged actions remain externally supplied.Future observations and availability masks are used only as supervised targets; the rollout coefficient is warmed up from zero in the MIMIC-III protocol.Source: §B.3 · p. 19

B.4 HYPERPARAMETERS

  • Table 8 summarizes recursive-training settings and Table 9 lists implementation settings for the MIMIC-III causal-Transformer transition and laboratory readout.Source: Table 8 · p. 19, Table 9 · p. 20

B.5 TENSOR CONTRACT

  • The tensor contract aligns source state and action at transition index c with targets at c + 1, applying padding and availability masks within sparse sets and across transitions.Long histories retain the most recent C valid transitions; recursive evaluation consumes predicted states and causal caches after the initial window.Source: §B.5 · pp. 20–21

C COMPARISON METHODS AND CONTROLS

The comparisons adapt prior methods to a shared clinical interface while preserving their defining computational cores and matching key data, training, and evaluation conditions. The evaluation separates clinical readouts, forecasting, and action-planning questions; PatientWM leads the listed adapted baselines on external-cohort metrics.

C COMPARISON METHODS AND CONTROLS

  • Comparison methods retain their source-inspired computational cores while sharing clinical input and output interfaces and matched splits, windows, inputs, actions, targets, losses, budgets, and evaluation.Parameter counts and internal recurrence remain properties of the retained method structures.Source: C COMPARISON METHODS AND CONTROLS · p. 21, Table 11 · p. 22
  • The Core-7 ablations isolate transition mechanisms and modality composition using the same split, outcomes, readout, optimization budget, selection rule, and evaluator as the full configuration.Ablations remove action conditioning or ordered state–action history, or restrict observations to one modality.Source: C COMPARISON METHODS AND CONTROLS · p. 22

D.1 EVALUATION PROTOCOL

  • Evaluation distinguishes factual next-state prediction, typed clinical readouts, recursive rollouts, and candidate-action agreement rather than combining different target metrics into one system score.The bounded action probe measures tie-aware top-k recovery of recorded actions; the rollout supplies recorded actions at each transition and evaluates predictions at 1, 4, 8, and 16 six-hour steps.Source: §D.1 · pp. 22–23

D.2 FACTUAL READOUTS AND PREDICTION CONTROLS

  • Clinical outcome macro-metrics average per-target results across seven events within each dataset-specific panel, while external comparisons report per-target test results on shared dataset support.The panels retain their dataset-specific target horizons.Source: §D.1 · p. 23, §D.2 · p. 24, Table 12 · p. 24
  • PatientWM leads the listed adapted clinical world-model baselines on AUPRC, AUROC, and Brier score across all three external datasets.The comparisons use the same per-target observed test support within each dataset.Source: §D.2 · p. 24, Table 15 · p. 25

D.3 ROLLOUT AND ACTION-PLANNING DETAILS

The rollout and action-planning evaluations compare methods on shared trajectories and fixed scoring conditions. PatientWM achieves the highest covered top-1 and top-2 action agreement, while historical-action agreement is treated as a behavioral surrogate rather than evidence of treatment benefit.

D.3 ROLLOUT AND ACTION-PLANNING DETAILS

  • The multi-method recursive rollout compares Event-7 macro-AUPRC and observed-laboratory nMAE on the same test trajectories, horizons, recorded actions, targets, and evaluator.Table 16 reports the exact values summarized in Figure 4.Source: §D.3 · p. 26, Table 16 · p. 26
  • PatientWM obtains the highest covered top-1 and top-2 agreement under the fixed action-planning score.The evaluation includes 15,951 test trajectories per method and 13,143 decisions with non-empty recorded treatment updates.Source: §D.3 · p. 26
  • Recorded-action agreement is a clinician-behavior surrogate, not evidence of treatment optimality, counterfactual correctness, or patient-level action quality.Alternative candidate outcomes are unobserved in retrospective records; stronger claims require additional causal identification, support diagnostics, or randomized-like validation.Source: §D.3 · p. 26

D.4 FROZEN-STATE REUSE

  • The frozen-state transfer experiment removes clinical-event label loss during source training and evaluates a newly trained Common Event-7 readout on the same seven targets and test examples.The full PatientWM row serves as an in-family reference for the cost of removing event supervision.Source: §D.4 · p. 26, Table 17 · p. 27

D.5 LABEL-EFFICIENT CROSS-DATASET ADAPTATION

The cross-dataset adaptation experiment tests label-efficient reuse by transferring the Common Event-7 interface from MIMIC-III to MIMIC-IV while keeping each source predictor frozen.

D.5 LABEL-EFFICIENT CROSS-DATASET ADAPTATION

  • Only a shared 56-parameter residual logit adapter is learned on 328 MIMIC-IV examples, or 1% of fixed training support, while source representations and prediction interfaces remain fixed.All methods use the same target examples, seven outcome definitions, and evaluator.Source: §D.5 · p. 27

E REPRODUCIBILITY AND ARTIFACTS

The reproducibility details specify the principal experiment splits, model configurations, and training setups across MIMIC and external datasets.

E REPRODUCIBILITY AND ARTIFACTS

  • The MIMIC-III transition model uses all 32 training and test shards, with a test set of 3,827 subjects, 5,337 ICU episodes, and 14,317 transition trajectories.It uses a closed-loop objective advancing up to 16 six-hour windows from one observed initial state.Source: E REPRODUCIBILITY AND ARTIFACTS · p. 27
  • The MIMIC-IV rollout and action-planning model uses 15,995 validation and 15,951 common test trajectories, a two-layer width-128 GRU, and context length 16 with ordered state–action history.Training includes a 16-window closed-loop term weighted by 0.50.Source: E REPRODUCIBILITY AND ARTIFACTS · p. 27
  • The MIMIC-IV laboratory decoder has 36,483 trainable parameters and is trained with the transition and GapShrink prior frozen.Its validation selection criterion is all-target laboratory nMAE.Source: E REPRODUCIBILITY AND ARTIFACTS · p. 27
  • The MIMIC-III Core-7 clinical readout uses 2,511 inputs, including 228 appended source-history features when available.Each outcome uses 120 histogram-gradient-boosting iterations with 31 leaves and balanced class weights.Source: E REPRODUCIBILITY AND ARTIFACTS · p. 27
  • External-dataset evaluations use context length eight, shared host width 192, and dataset-specific patient-group splits; PatientWM uses three layers and six attention heads.Learned rows share balanced binary outcome loss and, where adjacent targets exist, weighted next-EHR and next-waveform losses.Source: E REPRODUCIBILITY AND ARTIFACTS · pp. 27–28
Loading 2610.17059v1…