CaMeWM: A Causal Medical World Model for Generative Tumor Evolution and Treatment Planning

Huanlong Gao; Xuelei He; Fengjun Zhao; Xiaowei He

cs.CVOct 17, 2026

CaMeWM: A Causal Medical World Model for Generative Tumor Evolution and Treatment Planning

Authors: Huanlong Gao, Xuelei He, Fengjun Zhao, Xiaowei He

Organizations: School of Electronic Information, Northwest University, No. 1 Xuefu Avenue, Xi’an, 710127, China · Xi’an Key Lab of Radiomics and Intelligent Perception, Northwest University, No. 1 Xuefu Avenue, Xi’an, 710127, China

Abstract

Treatment planning is intrinsically counterfactual: clinicians must compare how a patient's tumor and survival would evolve under competing regimens, whereas a retrospective cohort records only the regimen received and one follow-up scan. A treatment-conditioned generator trained on such data can therefore reproduce historical treatment allocation rather than simulate the interventional disease dynamics needed to rank regimens. We introduce CaMeWM, a causal medical world model for supported post-TACE tumor evolution and treatment planning. CaMeWM implements one intervention-to-decision loop rather than three independent prediction tasks: it first admits a queried regimen only for clinically comparable histories with empirical treatment support; it then generates a distribution of lesion-local follow-up worlds through an action-conditioned response descriptor, deterministic anchor, and residual 3D diffusion transition; and it finally maps sampled worlds to restricted mean survival time (RMST) with a response-Cox utility model. Under pre-treatment adjustment and positivity, this loop estimates a prespecified, support-restricted g-formula target; all selection and nuisance estimation are fold-local, and the induced policy is evaluated by cross-fitted, censoring-adjusted doubly robust estimation. Across 79 paired out-of-fold episodes, CaMeWM achieves 0.618 SSIM and 0.528 change-map SSIM for follow-up CT transition, while its generated-world survival pathway attains a 0.615 C-index; factual-only AIPW estimates a 9.60-week RMST gain (95% CI 0.61--18.13) for the frozen policy. CaMeWM provides an inspectable basis for retrospective, support-restricted treatment planning.

Figures & tables

Figure 1: CaMeWM replaces observational assignment by a supported $\operatorname{do}(A=a)$ query. Under consistency, conditional exchangeability, and positivity, the g-formula identifies the regimen-specific population world law and its RMST value on empirical support (Theorem 1 ; full assumptions and proof in the supplement). The response code, deterministic anchor, residual DDPM, and survival head estimate components of this identified functional; their architectures neither establish exchangeability nor identify a unique patient’s counterfactual CT. The lower rail performs leakage-free IPCW/DR policy evaluation. CT thumbnails are schematic, not patient data or results.

Method	Tumor HU-MAE $\downarrow$	$\Delta$ SSIM $\uparrow$	PSNR $\uparrow$	SSIM $\uparrow$
Persistence	57.192 [52.069, 63.049]	0.157 [0.135, 0.180]	16.877 [16.148, 17.661]	0.469 [0.443, 0.497]
3D U-Net	26.944 [23.224, 31.104]	0.523 [0.481, 0.566]	19.157 [18.359, 19.974]	0.610 [0.584, 0.637]
CVAE-3D	32.300 [28.065, 37.105]	0.508 [0.464, 0.549]	18.815 [18.027, 19.670]	0.587 [0.562, 0.612]
Mask-DDPM	63.721 [58.270, 69.962]	0.051 [0.026, 0.080]	16.665 [15.966, 17.362]	0.451 [0.425, 0.477]
MeWM ∗	62.910 [57.562, 68.555]	0.074 [0.048, 0.102]	16.684 [15.993, 17.373]	0.438 [0.412, 0.465]
CaMeWM	27.145 [23.139, 31.534]	0.528 [0.484, 0.569]	19.193 [18.353, 20.112]	0.618 [0.591, 0.645]

Table 1: Factual post-TACE CT generation on 79 paired OOF episodes. Brackets denote patient-bootstrap 95% confidence intervals. Tumor HU-MAE is reported after inverse transformation to Hounsfield units. ∗ Task-matched reimplementation of MeWM’s action-conditioned generative-dynamics principle, adapted to the two regimens and public HCC-TACE-Seg outputs available in this study. All methods use the same input state, preprocessing, and patient-level OOF folds.

Planning method	Observed-regimen agreement $\uparrow$	Macro-F1 $\uparrow$	Balanced accuracy $\uparrow$	CaMeWM-only support † $\uparrow$	CaMeWM-only abstention † $\downarrow$
CoxPH	0.367 [0.266, 0.468]	0.269 [0.210, 0.319]	0.500 [0.500, 0.500]	–	–
DeepSurv	0.342 [0.241, 0.443]	0.310 [0.223, 0.409]	0.429 [0.337, 0.516]	–	–
IPTW-Cox	0.367 [0.266, 0.468]	0.269 [0.210, 0.319]	0.500 [0.500, 0.500]	–	–
MeWM ∗	0.367 [0.266, 0.468]	0.269 [0.210, 0.319]	0.500 [0.500, 0.500]	–	–
CaMeWM	0.506 [0.392, 0.620]	0.503 [0.388, 0.610]	0.523 [0.401, 0.631]	0.597 [0.561, 0.634]	0.000 [0.000, 0.000]
CaMeWM factual-only AIPW: Observed RMST $=103.62$ weeks; AIPW policy value $=113.22$ [102.34, 124.84] weeks.

Table 2: Treatment-planning behavior and factual-only causal evaluation over five patient-level out-of-fold folds. Main estimates are shown above patient-bootstrap 95% confidence intervals. Agreement metrics compare recommendations with observed historical regimens and are behavioral references, not counterfactual treatment truth. † Defined only for CaMeWM, which applies a propensity-support guardrail; “–” denotes not applicable. ∗ Task-matched reimplementation inspired by MeWM’s action-conditioned generative-dynamics principle.

Variant	Normalized tumor MAE $\downarrow$	SSIM $\uparrow$	Global HU-MAE $\downarrow$	$\Delta$ SSIM $\uparrow$
Full CaMeWM	0.070 [0.059, 0.082]	0.618 [0.591, 0.644]	61.102 [54.384, 68.512]	0.528 [0.483, 0.570]
Causal anchor only	0.070 [0.060, 0.081]	0.612 [0.585, 0.637]	61.139 [54.508, 68.499]	0.526 [0.482, 0.569]
w/o response-code conditioning	0.071 [0.060, 0.082]	0.618 [0.591, 0.643]	61.198 [54.274, 68.508]	0.528 [0.485, 0.571]
w/o direct action path	0.075 [0.064, 0.087]	0.617 [0.591, 0.644]	62.239 [55.086, 69.746]	0.524 [0.480, 0.566]
w/o confidence weighting	0.070 [0.060, 0.082]	0.618 [0.592, 0.644]	61.197 [54.610, 68.239]	0.527 [0.482, 0.570]

Table 3: Controlled generation ablation over the same five patient-level out-of-fold folds. Normalized tumor MAE is computed on the model’s

[-1,1]

CT scale and is not numerically comparable to the HU-MAE in Table 1 . “Deterministic anchor only” removes residual-DDPM refinement; “uniform pseudo-mask weighting” sets all pseudo-mask confidence weights to one. Brackets denote patient-bootstrap 95% confidence intervals.

Method	Uno C $\uparrow$	IBS $\downarrow$	RMST MAE $\downarrow$	C-index $\uparrow$ 95% CI
CoxPH	0.580	0.202	39.28	0.581 [0.500, 0.661]
DeepSurv	0.592	0.200	39.17	0.597 [0.525, 0.670]
IPTW-Cox	0.603	0.200	38.73	0.609 [0.531, 0.689]
MeWM ∗	0.461	0.242	44.53	0.462 [0.390, 0.540]
CaMeWM factual	0.589	0.211	40.80	0.592 [0.521, 0.662]
CaMeWM generated worlds	0.612	0.202	39.23	0.615 [0.546, 0.688]

Table 4: Patient-level out-of-fold survival prediction over five folds. The generated-world row evaluates the complete action-conditioned transition-and-utility path. Brackets denote patient-bootstrap 95% confidence intervals. ∗ Task-matched reimplementation inspired by MeWM’s action-conditioned generative-dynamics principle.

Figure 2: Out-of-fold transition fidelity and paired comparison. CaMeWM synthesizes action-conditioned follow-up CT volumes in registered image space, rather than tumor masks. Each row shows pre-treatment CT, observed post-treatment CT, and predictions from 3D U-Net and the mean of eight CaMeWM worlds. Cyan dashed contours mark the pre-treatment lesion only for orientation; they are neither follow-up annotations nor model outputs. Tumor-region MAE uses a fold-local pseudo post-treatment mask, whereas SSIM is computed over the complete follow-up CT. Rows are deterministically selected at the 20th, 50th, and 80th MAE percentiles.

Main-paper contribution	Formal support	Required implementation and evidence
Support-restricted interventional world distribution	Theorem S1 and its complete proof	Joint-regimen propensity, explicit support set, response-factorized g-formula, overlap diagnostics, and semi-synthetic identification recovery.
Response-factorized residual transition	Theorem S2 and Corollary S3	Operational response proxies, lesion-local residual DDPM, controlled component ablations, factor-wise errors, and empirical bound coverage.
RMST world integration and supported policy	Lemma S4 and Theorem S7	Response-Cox RMST, $K=8$ worlds per action, sampling-convergence analysis, and support abstention. The current score is not claimed to be a calibrated finite-sample confidence bound.
Cross-fitted factual policy evaluation	Theorem S6 , its proof, and the censoring identity	Outer-fold predictions, IPCW RMST pseudo-outcomes, factual-only AIPW value, overlap/balance/ESS diagnostics, and patient-level bootstrap intervals.
Claim boundaries	Proposition S8 and Remark S2	Hidden-confounding sensitivity and explicit exclusion of unique individual CT and prospective-benefit claims.

Table S1: Claim-to-evidence map. “Required evidence” denotes evidence needed to activate the corresponding main-paper claim, not optional analysis.

Symbol	Meaning
$H=(I^{-},S^{-},C)$	Pre-treatment history: baseline CT, baseline lesion mask, and clinical covariates.
$A$ , $a$	Observed joint regimen and a candidate intervention, respectively.
$M$	Treatment-conditioned response descriptor predicted from $(H,a)$ .
$X^{+}$ , $X^{+}(a)$	Observed registered follow-up CT and its potential follow-up state under $a$ .
$e_{a}(H)$ , $\mathcal{A}_{\eta}(H)$	Joint-regimen propensity and the support-restricted candidate set satisfying $e_{a}(H)\geq\eta$ .
$P_{M}$ , $P_{X}$	Response-descriptor and follow-up-state transition distributions in the g-formula.

Table S2: Principal notation used by CaMeWM.

Module	Checkpoint criterion	Safeguard
Follow-up segmenter	Maximum validation Dice	Minimum 40 epochs; patience 20
Response predictor	Minimum standardized validation MAE	Minimum 40 epochs; patience 25
Joint propensity	Minimum validation cross-entropy	Minimum 30 epochs; patience 20
Residual DDPM	Minimum sampled tumor-region MAE, checked every 10 epochs	Minimum 60 epochs; patience 5 checks
Response-Cox survival	Maximum validation C-index; validation loss only when C-index is undefined	Minimum 50 epochs; patience 30
Planner	No test-set tuning; $K=8$ , support threshold 0.05	Same locked values for every fold and method

Table S3: Validation criteria used to select fold-local checkpoints. Selection metrics are fixed before examining outer-test results.

Analytical stage	Reproducible input or operation	Required evidence
Image-state construction	label semantics, scan pairing, registration, common physical crop	cohort flow, exclusions, label and registration QC
Fold-local model fitting	training/validation partitions, fixed selection criteria, frozen outer test fold	selection records and fold-level predictions
Response and transition estimation	response descriptor, causal anchor, residual transition, world sampling	module evaluations, sampled worlds, and qualitative cases
Survival and planning	factual-state survival fit, supported action scoring, factual-only AIPW	survival, planning, support, and policy-value summaries
OOF aggregation	patient-level pooling and paired bootstrap resampling	reported tables, figures, and confidence intervals

Table S4: Reproducibility materials and audit trail. The table intentionally describes analytical stages rather than release-specific script or archive names.

Stage	Episodes
Public HCC-TACE-Seg release	105
Longitudinal local manifest with required imaging and clinical records	99
Two predeclared supported complete joint regimens	80
Empty released source label-2 tumor mask excluded	1
Final patient-level OOF cohort	79

Table S5: Outcome-blind cohort construction for the fixed five-fold OOF protocol.

Item	Final specification
Tumor semantics	Source label 2 (Mass/HCC tumor); empty source masks excluded
Spatial preprocessing	Rigid physical-grid registration; 1-mm isotropic spacing; shared $96^{3}$ baseline-tumor physical crop
Image scale	HU window $[-175,600]$ followed by the common normalization used for all methods
Follow-up lesion state	Fold-local segmenter, validation Dice selection, four-view test-time aggregation, recorded fallback QC
Treatment support	Two complete joint regimens; support threshold $\eta=0.05$
World integration	$K=8$ diffusion worlds per candidate regimen

Table S6: Locked five-fold OOF evaluation decisions.

Covariate	Release representation	Fold-local processing
Age	Continuous baseline age	Training median-centering and SD scaling
Sex	Binary release code	Training median-centering and SD scaling
AFP	Continuous reported baseline value	$\log(1+\mathrm{AFP})$ , then training median-centering and SD scaling
Child–Pugh class	Ordinal release code	Training median-centering and SD scaling
BCLC stage	Ordinal release code	Training median-centering and SD scaling
Tumor count	Binary release code	Training median-centering and SD scaling

Table S7: Pre-treatment adjustment covariates and fold-local missing-data handling in the final 79-episode cohort. “Missing” is assessed before any outer-fold preprocessing. All covariates enter the joint propensity and the response-Cox utility model together with their binary missingness indicators.

Figure S1: Observed-action propensity diagnostics. (a) Distribution of OOF observed-regimen propensities; the dashed line is the predeclared support threshold. (b) Action-wise OOF calibration. (c) Standardized mean differences before weighting, after IPTW, and after overlap weighting for the non-missing clinical covariates. These diagnostics characterize observed support and measured-covariate balance; they do not establish exchangeability.

Figure S2: Factual-to-generated survival transport diagnostic. (a–c) IPCW reliability at 12, 24, and 36 months for factual and generated states. (d) Patient-level agreement between factual-state and generated-world predicted RMST. The diagonal is identity. The paired C-index change is reported with a patient-bootstrap 95% interval and is interpreted only as a diagnostic for generated-state shift.

Figure S3: Stability with the number of sampled worlds. (a) OOF recommendation agreement with the $K=32$ reference over $K\in\{2,4,8,16,32\}$ . At the locked $K=8$ , 78/79 (98.7%) recommendations match $K=32$ . (b) Corresponding recommended-regimen composition and abstention fraction. This supports the computational choice of $K=8$ ; it is not a finite-sample confidence-calibration result.

Figure S4: Pseudo follow-up mask confidence audit. (a) Spearman associations, with patient-bootstrap 95% intervals, between three fold-local confidence signals and OOF tumor-region MAE. (b) Direct confidence-versus-error scatter. The pseudo-label confidence association is weak ( $\rho=0.20$ , 95% CI $[-0.03,0.41]$ ), so these quantities are treated as audits rather than calibrated error predictors.

Figure S5: Known-ground-truth semi-synthetic validation. (a) ATE and policy-value recovery with increasing simulated sample size. (b) Double-robustness identity under correct or misspecified propensity and outcome nuisances. (c) Positivity stress. (d) Factor-wise causal-value bound. (e) Monte Carlo convergence with sampled-world count. (f) Sensitivity to an unmeasured common cause. The experiment tests the behavior of the estimators and theoretical bounds under the simulation SCM, not clinical efficacy.

Quantity	OOF estimate
Patients	79
Factual observed-regimen RMST (weeks)	103.62
Factual-only AIPW policy value (weeks)	113.22 [102.34, 124.84]
Factual-only AIPW gain (weeks)	9.60 [0.61, 18.13]
Policy effective sample size	37.4
Mean recommended-action support	0.597

Table S8: Primary factual-only AIPW evaluation of the frozen CaMeWM policy. The outcome nuisance uses pre-treatment clinical history, factual regimen, and observed survival only.

Figure S6: Support-threshold and IPW-cap sensitivity of a historical generated-state plug-in diagnostic. Each cell uses the fixed five-fold OOF planning and propensity predictions. Rows vary the joint observed/recommended support threshold $\eta$ and columns vary the IPW cap. Panel (a) reports the descriptive plug-in RMST value gain, while panel (b) reports effective sample size. The locked primary setting is $\eta=0.05$ with cap 10. This analysis assesses overlap and weight sensitivity; it is not the primary factual-only AIPW estimand and does not convert an observational interval into evidence of clinical benefit.

Causal World Models

Causal world models represent how entities and environments change under actions, enabling prediction, planning, and counterfactual reasoning.

See all 46 papers

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 +2Disease Progression Modeling Healthcare
Institute of Artificial Intelligence, University of Central Florida
A Unifying Perspective on Causal World Models: From Observations to Representations to Structure
Aug 13, 2026Avinash Kori, Fabrizio RussoCausal Representation Learning World Model Learning
Department of Computing, Imperial College London, London, UK
PatientWM: Towards Generalizable Clinical World Models via Unified Patient States
Oct 12, 2026Ziwen Kan, Jie Peng, Bo Yan +5Action-Conditioned World Models Clinical Prediction
University of Science and Technology of China · University of Central Florida · Princeton University +2

Dec 8, 2025J/K move · Enter open · S save

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

Tianxingjian Ding, Yuanhao Zou, Chen Chen, Mubarak Shah+1

Institute of Artificial Intelligence, University of Central Florida

Disease Progression Modeling Healthcare Clinical Decision Support Clinical Decision-Making Oncology Latent World Models

Open paper AskPDF arXiv

Clinical decision-making in oncology requires predicting dynamic disease evolution, a task current static AI predictors cannot perform. While world models (WMs) offer a paradigm for generative prediction, existing medical applications remain limited. Existing methods often rely on stochastic diffusion models, focusing on visual reconstruction rather than causal, physiological transitions. Furthermore, in medical domain, models like MeWM typically ignore patient-specific temporal and clinical contexts and lack a feedback mechanism to link predictions to treatment decisions. To address these gaps, we introduce CLARITY, a medical world model that forecasts disease evolution directly within a structured latent space. It explicitly integrates time intervals (temporal context) and patient-specific data (clinical context) to model treatment-conditioned progression as a smooth, interpretable trajectory, and thus generate physiologically faithful, individualized treatment plans. Finally, CLARITY introduces a novel prediction-to-decision framework, translating latent rollouts into transparent, actionable recommendations. CLARITY demonstrates state-of-the-art performance in treatment planning. On the MU-Glioma-Post dataset, our approach outperforms recent MeWM by 12%, and significantly surpasses all other medical-specific large language models.

Abstract

Causal World Models

Explore similar work

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

A Unifying Perspective on Causal World Models: From Observations to Representations to Structure

PatientWM: Towards Generalizable Clinical World Models via Unified Patient States