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

AI summary

Overview

• Research area: Medical machine learning / clinical decision support — specifically world models applied to oncology treatment planning, combining latent dynamics modeling, survival analysis, and multimodal large language models.
• Technical level: Advanced (requires familiarity with world models, latent state-space modeling, survival analysis losses such as Cox partial likelihood and Brier score, LoRA adaptation, contrastive learning, and entropy-regularized planning).
• Scope in one sentence: The paper introduces CLARITY, a medical world model that simulates treatment-conditioned disease trajectories in a latent space and uses survival-aware feedback to iteratively propose and refine cancer therapy regimens, evaluated on two brain tumor cohorts (MU-Glioma-Post, UCSF-ALPTDG) and one breast cancer cohort (ISPY-2).

What This Paper Is About

Most clinical AI systems in oncology are static predictors: they map an image or a patient record to a single outcome and cannot simulate how a disease would evolve under a treatment that was never given. CLARITY instead treats treatment planning as a sequential, counterfactual problem — it learns disease progression directly in a latent representation space, conditioned on imaging, elapsed time between follow-ups, and patient-specific clinical context, so that alternative "what-if" therapies can be simulated before committing to one. The goal is a closed-loop system that imagines disease trajectories, scores their survival implications, and iteratively refines therapy proposals toward lower predicted risk.

Key Contributions

1. Treatment-conditioned disease modeling in latent space. Longitudinal progression is modeled as continuous physiological transitions inside a visual representation space rather than through pixel-level image synthesis, enabling counterfactual modeling of treatment-induced disease dynamics.
2. Temporal and clinical conditioning. Continuous time intervals are encoded as embeddings (sinusoidal encoding of the gap between pre- and post-treatment observations) and multimodal patient context — genomics, demographics, prior therapies — is integrated to produce longitudinally coherent, individualized predictions across irregular follow-up intervals.
3. A prediction-to-decision integration mechanism. The paper introduces an Inverse Survival Evaluation framework: predicted latent rollouts are optimized under an entropy-regularized long-horizon objective and fed back to a therapy policy agent, producing adaptive and interpretable treatment recommendations.
4. A demonstrated simulation-to-decision loop. The framework reaches state-of-the-art results in treatment planning and survival prediction across three cancer datasets, including zero-shot transfer to an external brain tumor cohort.

Main Findings

• Treatment planning performance on MU-Glioma-Post: CLARITY achieves an F1-score of 57.1% (precision 61.3%, recall 53.5%, Jaccard 39.9%), a 9.2% absolute improvement over the strongest medical-specific baseline, Huatuo-Vision (46.4% F1).
• External zero-shot validation on UCSF-ALPTDG: The paper's text reports 48.9% F1, outperforming the second-best method at 44.1%. Note that Table 1 lists 50.2% F1 (precision 52.6%, recall 48.0%, Jaccard 33.5%) for the same row, so the text and table figures differ within the paper.
• Breast cancer generalization on ISPY-2: F1 of 53.9% (precision 56.2%, recall 51.8%, Jaccard 36.8%), trained independently on that cohort.
• MLLM prompting alone is weak, but works inside the loop: Claude-4.5-Sonnet prompted directly achieves 41.6% F1, whereas CLARITY uses the MLLM as a component of a simulation-to-decision loop and substantially improves effective performance.
• Comparison against a re-implemented diffusion baseline: Direct comparison with the original MeWM was not feasible because its in-house data is not public and it targets CT data, so the authors re-implemented a MeWM-style diffusion predictor on their MRI data (MeWM*) and report 43.6% F1 on MU-Glioma-Post versus CLARITY's 57.1%.
• Survival ranking: CLARITY reaches a state-of-the-art C-index of 0.794 ± 0.02, compared with SurvTRACE (0.713), MeWM* (0.701), RSF (0.672), and DeepSurv (0.664).
• Kaplan-Meier stratification: CLARITY produces clearer separation between risk strata with a log-rank p-value of 0.0017 versus 0.0763 for the diffusion-based MeWM baseline, and a C-index of 0.7943 in that analysis.
• Blinded expert assessment (40 held-out MU-Glioma-Post cases): CLARITY was judged equivalent/non-inferior in 60.0% of cases, preferred over the physician-administered regimen in 35.0%, and the physician regimen was preferred in 5.0% — meaning 95.0% were rated equivalent/non-inferior or preferable.
• Latent beats diffusion in the ablation: Replacing diffusion synthesis with latent dynamics produced the largest gain, +8.8% in F1 (43.6% to 52.4%). Adding clinical context raised F1 to 52.8%, feedback alone (K=3) reached 54.0%, and both together peaked at 57.1%.
• Iteration count matters: Performance peaks at K=3 with an F1 of 57.1%; a slight decline at K=4 (55.4% F1) indicates diminishing returns and possible noise introduction, so K=3 was adopted.
• Loss design: Adding the Brier score improves over the baseline by enforcing probability calibration (+0.7%); the soft-label contrastive variant yields the highest C-index (79.4%), outperforming hard labels because it better captures nuanced treatment similarities.
• Encoder choice matters: Natural-image pretraining (DINOv2) gives the lowest C-index (0.665 ± 0.020, Brier 0.198 ± 0.008); medical-domain pretraining (MedSigLIP) improves to 0.705 ± 0.016 (Brier 0.176 ± 0.006); MRI-specific encoders do better still, with brainIAC at 0.786 ± 0.012 (Brier 0.168 ± 0.004) and MRI-CORE best at 0.794 ± 0.010 (Brier 0.162 ± 0.003).
• Efficiency: A diffusion-based approach requires up to 61.3 TFLOPs and 38.6 seconds for a single simulation (1000 steps; 39.5 TFLOPs and 19.7 seconds at 500 steps), whereas the latent-based predictor uses 4.21 TFLOPs and 0.341 seconds — described as an approximate 9–15× reduction with sub-second inference — which makes multi-iteration Inverse Survival Evaluation feasible at K=3.
• Representation drift: Diffusion-based reconstruction yields a cosine distance of 0.0174 to the ground-truth post-treatment latent versus 0.0084 for direct latent prediction, roughly 1.9× higher drift.
• Human-level context: The paper cites prior studies in which neurosurgeons achieve roughly 40% accuracy predicting outcomes for high-grade glioma patients and multidisciplinary oncology teams reach approximately 55% accuracy in advanced cancer cohorts.
• Simulated decision trajectory example: The model constructs a temporally consistent treatment lane — RT+TMZ → TMZ → CCNU → Avastin+Brachy — that minimizes cumulative risk, with boxed numbers in the figure denoting 30-day risk scores.

Methodology in Plain English

CLARITY is built from four pieces. A frozen 3D MRI encoder (MRI-CORE, a 12-layer ViT-B producing 768-dimensional latents) is adapted with LoRA so it can specialize to longitudinal disease modeling while preserving its general visual representations. A therapy policy agent — GPT-5 — reads the patient's clinical profile and proposes multiple clinically valid drug combinations, constrained by safety rules (for example, avoiding co-administration of Bevacizumab and Temozolomide). Each proposed regimen is turned into a dense embedding by a medical text encoder (MedGemma, fine-tuned with 4-bit quantization and LoRA).

The core is the "Actor," a disease evolution model with two parts: a Post-Treatment Latent Predictor (a 4-layer self-attention Transformer) that forecasts the post-treatment latent state, and a Survival Predictor (a 4-layer two-way cross-attention Transformer) that maps pre- and predicted post-treatment latents to a one-year survival probability and a continuous risk score. The Actor concatenates four inputs — the pre-treatment latent, the clinical embedding, the sinusoidal time-gap embedding, and the drug embedding — and predicts the next latent state plus survival outputs.

Training uses paired pre- and post-treatment MRIs with an L1 latent consistency loss (predicted latent versus ground-truth latent) and a soft-label contrastive regularization in which pairwise similarities between treatment text embeddings act as soft targets, so that mechanistically similar therapies (such as alkylating chemotherapy agents) produce more similar disease transitions than fundamentally different ones like radiation. Survival supervision combines a Brier loss for calibration with a Cox partial likelihood loss for risk ordering; the ablation shows combining both is more reliable than either alone. The optimizer is AdamW.

At inference, two modes operate. In Direct Survival Evaluation, the agent proposes candidates and the Actor scores each one. In Inverse Survival Evaluation, those scores are fed back to the agent as structured feedback and it proposes refined therapies, repeating for K iterations (K=3 empirically) before selecting the lowest-risk proposal. Planning is framed as entropy-regularized trajectory optimization: rather than greedily minimizing the discounted cumulative risk J(a_{1:H}), the framework maximizes negative expected risk plus a temperature-scaled entropy term, keeping exploration alive while shifting probability mass toward lower-risk trajectories. Because the model can roll out latent states autoregressively under arbitrary policies without needing paired imaging at intermediate steps, it supports "what-if" simulation of treatment schedules at arbitrary time points.

Why This Matters

Impact on research. The paper argues that medical generative models should not spend compute on pixel-level image synthesis when clinical decisions depend on robust visual features and their longitudinal transitions. Its ablation quantifies this: switching from diffusion synthesis to latent dynamics gives the single largest improvement (+8.8% F1), and diffusion reconstruction nearly doubles representation drift (1.9×). It also positions world models — previously successful in robotics and control — as a viable paradigm for medicine, while naming the obstacles: irregular time intervals, high inter-patient heterogeneity, multimodal conditioning, and stringent interpretability requirements.

Real-world applications.

• Treatment planning support in neuro-oncology, where clinicians must choose among options such as continuing temozolomide, introducing targeted agents, or escalating to salvage treatments.
• Counterfactual "what-if" simulation of therapy sequences that were never observed in the patient's record, including multi-stage regimens such as RT+TMZ → TMZ → CCNU → Avastin+Brachy.
• Survival stratification and risk-stratified monitoring using predicted one-year survival probabilities and continuous risk scores.
• Cross-cancer generalization as a general pipeline — demonstrated by independent training on the ISPY-2 breast cancer cohort under neoadjuvant chemotherapy, beyond the brain tumor setting.

Industry relevance. The efficiency numbers matter for deployment: 4.21 TFLOPs and 0.341-second inference versus 61.3 TFLOPs and 38.6 seconds for diffusion makes iterative, multi-candidate planning computationally practical. The design also composes off-the-shelf components — a frontier MLLM for candidate generation, a medical text encoder, and an MRI foundation model adapted with LoRA — which maps onto how clinical AI products are assembled from pretrained building blocks.

Future Directions

• Closing the gap between imitation and optimization. The paper notes that the physician-administered regimen is an observed clinical decision rather than a unique optimum, and that CLARITY was preferred in 35.0% of blinded cases against it — raising the open question of how to validate recommendations that deviate from historical practice.
• Scaling beyond three cohorts and two cancer types. Generalization is demonstrated on two brain tumor datasets plus ISPY-2 breast cancer; extension to other tumor types and to zero-shot settings with less overlap in imaging protocol and treatment vocabulary remains untested.
• Choosing the planning horizon principledly. Performance peaks at K=3 and declines at K=4, so the trade-off between iterative refinement and accumulated noise in the long-horizon objective needs a more systematic treatment than the fixed setting adopted here.
• Making the latent space more interpretable. The paper's stated goal is smooth, interpretable representations of disease progression, but interpretability of the latent states themselves and of the survival feedback signal is not quantitatively evaluated.

Target Audience

This paper is most useful to machine learning researchers working on world models, latent dynamics, and sequential decision-making who are interested in medical applications; to clinical AI and computational oncology researchers building treatment-planning or survival-prediction systems; and to practitioners assembling multimodal pipelines from foundation models (MRI encoders, medical text encoders, and LLM-based agents). Clinicians and medical affairs readers may find the blinded expert assessment and the Kaplan-Meier analysis most directly relevant, while readers without a background in survival analysis losses or latent state-space models will find the methodology sections demanding.

Authors’ abstract

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.