CLARITY: Medical World Model for Guiding Treatment Decisions by Simulating Context-Aware Disease Trajectories in Latent Space
Abstract
Clinical decision-making in oncology requires forecasting howdisease evolves under treatment, yet most AI systems remain static pre-dictors that cannot model longitudinal, treatment-conditioned progres-sion. Although generative and world models have demonstrated strongcapabilities in general domains, their adaptation to medicine remainslimited and insufficient for capturing complex, treatment-induced phys-iological dynamics across temporal scales. To address these gaps, weintroduce CLARITY, a medical world model that enables counterfac-tual simulation of treatment-conditioned disease trajectories for clini-cal decision-making. By jointly encoding imaging-derived latent states,temporal intervals that capture irregular follow-ups, and patient-specificclinical context, CLARITY learns smooth and interpretable representa-tions of disease progression, allowing the model to simulate how alter-native treatments reshape future disease dynamics. Because treatmentoptimization is inherently sequential and uncertain, requiring evaluationof long-term outcomes across multiple possible interventions, we furtherpropose an entropy-regularized, computationally efficient long-horizonprediction-to-decision framework that plans treatment strategies overimagined disease trajectories and iteratively refines therapy proposalsthrough survival-aware feedback, forming a closed-loop simulation-to-decision framework for treatment planning. CLARITY achieves state-of-the-art performance in treatment planning and survival predictionacross three cancer datasets, including two brain tumor cohorts (MU-Glioma-Post and zero-shot on UCSF-ALPTDG) and one breast can-cer dataset (ISPY-2), demonstrating strong generalization across can-cer types while consistently outperforming prior generative methods andmedical-domain large language model baselines.