CLARITY: Medical World Model for Guiding Treatment Decisions by Simulating Context-Aware Disease Trajectories

Abstract

Clinical decision-making in oncology requires forecasting how disease evolves under treatment, yet most AI systems remain static predictors that cannot model longitudinal, treatment-conditioned progression. Although generative and world models have demonstrated strong capabilities in general domains, their adaptation to medicine remains limited and insufficient for capturing complex, treatment-induced physiological dynamics across temporal scales. To address these gaps, we introduce CLARITY, a medical world model that enables counterfactual simulation of treatment-conditioned disease trajectories for clinical decision-making. By jointly encoding imaging-derived latent states, temporal intervals that capture irregular follow-ups, and patient-specific clinical context, CLARITY learns smooth and interpretable representations of disease progression, allowing the model to simulate how alternative treatments reshape future disease dynamics. Because treatment optimization is inherently sequential and uncertain, requiring evaluation of long-term outcomes across multiple possible interventions, we further propose an entropy-regularized, computationally efficient long-horizon prediction-to-decision framework that plans treatment strategies over imagined disease trajectories and iteratively refines therapy proposals through 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 prediction across three cancer datasets, including two brain tumor cohorts (MU-Glioma-Post and zero-shot on UCSF-ALPTDG) and one breast cancer dataset (ISPY-2), demonstrating strong generalization across cancer types while consistently outperforming prior generative methods and medical-domain large language model baselines.