🤖 machine learning

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

CLARITY is a novel medical world model that forecasts individualized disease trajectories in a structured latent space by integrating temporal and clinical contexts, thereby enabling physiologically faithful, treatment-conditioned predictions that outperform existing methods in oncology treatment planning.

Original authors: Tianxingjian Ding, Yuanhao Zou, Chen Chen, Mubarak Shah, Yu Tian

Published 2026-03-30

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Tianxingjian Ding, Yuanhao Zou, Chen Chen, Mubarak Shah, Yu Tian

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are a doctor trying to decide the best treatment for a patient with cancer. Usually, AI tools act like a crystal ball that gives you a single, static prediction: "If we do Treatment A, the patient has a 60% chance of surviving."

But real life isn't static. Patients change, treatments have side effects, and the disease evolves differently depending on when and how you treat it. The problem is that current AI can't really "play out" different scenarios to see what happens next.

Enter CLARITY. Think of CLARITY not as a crystal ball, but as a high-tech flight simulator for the human body.

The Core Idea: The Medical Flight Simulator

In aviation, pilots don't just guess how a plane will handle a storm; they use simulators to try out different maneuvers in a safe, virtual environment before taking off.

CLARITY does the same thing for cancer treatment:

The "What-If" Engine: Instead of giving you one answer, CLARITY creates a virtual world where it can simulate the patient's disease journey under different treatments. It asks: "What if we give Drug X? What if we wait two weeks? What if we combine Radiation with Immunotherapy?"
The Virtual Patient: It builds a digital twin of the patient's disease using MRI scans, genetic data, and medical history. It doesn't just look at a picture of a tumor; it understands the story of how that tumor has been growing over time.

How It Works (The Three Magic Ingredients)

The paper describes three special "ingredients" that make this simulator so smart:

1. The Time Machine (Temporal Context)
Cancer doesn't grow in a straight line. Sometimes patients get scans every month, sometimes every six months.

The Analogy: Imagine watching a movie where the scenes are missing. A normal AI gets confused. CLARITY is like a director who knows exactly how much time passed between the missing scenes. It understands that a tumor growing over 3 days looks very different than one growing over 3 months. It fills in the gaps smoothly.

2. The Personalized GPS (Clinical Context)
Every patient is unique. Their age, genetics, and past treatments matter.

The Analogy: A standard GPS gives everyone the same route. CLARITY is like a personalized GPS that knows your car's engine type, your driving style, and the current traffic. It tailors the disease simulation specifically to this patient's biology, not just a generic "average" patient.

3. The "Reverse" Coach (Inverse Survival Evaluation)
This is the most clever part. Usually, you pick a treatment and see what happens. CLARITY works backward.

The Analogy: Imagine a chess coach who doesn't just tell you the best move. Instead, the coach simulates 100 different games where you try different moves. Then, the coach says, "Hey, the moves that lead to a win in the simulation are these three. Let's try them."
- CLARITY generates many possible treatment plans.
- It simulates the future for each one.
- It looks at the "survival score" of each future.
- It then iteratively refines the plan, discarding the bad paths and polishing the good ones until it finds the strategy with the highest chance of survival.

Why Is This a Big Deal?

1. It's Faster and Smarter than "Image Generators"
Previous medical AI tried to "draw" what the tumor would look like next month (like a generative art tool). This is slow and often creates blurry, unrealistic pictures.

CLARITY's Trick: Instead of drawing the picture, it predicts the mathematical essence of the tumor's state. It's like predicting the weather pattern rather than trying to draw every single raindrop. This makes it 15 times faster and much more accurate.

2. It Handles the "Unknown"
Doctors often have to guess what will happen with a new drug they haven't tried on a patient before.

The Superpower: CLARITY can simulate these "counterfactual" scenarios. It can tell a doctor, "We haven't tried this specific drug combo on this patient yet, but our simulation suggests it will shrink the tumor faster than the standard treatment."

3. It Actually Works
The researchers tested CLARITY on real brain tumor and breast cancer data.

The Result: It beat all other AI models, including massive general-purpose AI (like GPT-4) and specialized medical AI. It predicted survival outcomes with much higher accuracy, essentially acting as a "second opinion" that has run thousands of virtual trials before you even make a decision.

The Bottom Line

CLARITY is a tool that turns cancer treatment planning from a guessing game into a strategic simulation.

Instead of asking, "What will happen?" it allows doctors to ask, "What happens if we try this? What if we try that?" By running these simulations in a fast, virtual world, it helps doctors choose the path that gives the patient the best possible future. It's not just predicting the future; it's helping to design a better one.

1. Problem Statement

Clinical decision-making in oncology requires forecasting how a patient's disease evolves under specific treatments over time. Current AI systems face three critical limitations:

Static Prediction: Most models predict a single static outcome (e.g., survival probability) based on current data but fail to model the longitudinal trajectory of disease progression under different interventions.
Lack of Counterfactual Reasoning: Clinicians need to evaluate "what-if" scenarios (e.g., "What happens if we switch from Temozolomide to CCNU?"). Existing models cannot simulate disease states under unobserved treatment strategies.
Inefficiency and Noise: Generative approaches (like diffusion models) that synthesize future medical images are computationally expensive and often introduce pixel-level noise that obscures the underlying physiological dynamics needed for decision-making.

2. Methodology: The CLARITY Framework

CLARITY is a Medical World Model designed to simulate treatment-conditioned disease trajectories in a latent space, bridging the gap between prediction and decision-making. It operates via a Simulation-to-Decision framework.

Core Architecture

The system consists of three primary components:

MRI Encoder (Visual State Encoding):
- Uses a pre-trained 3D vision foundation model (e.g., MRI-CORE) with LoRA (Low-Rank Adaptation) to extract a latent visual state ( $z_{pre}$ ) from pre-treatment MRI scans.
- This approach preserves general visual representations while specializing for longitudinal disease modeling.
Therapy Policy Agent (MLLM):
- A Multimodal Large Language Model (e.g., GPT-5) acts as an agent that proposes candidate treatment combinations based on patient context (genomics, demographics, history).
- It generates structured therapy actions (chemotherapy, radiotherapy, immunotherapy, etc.) adhering to medical safety constraints.
Actor Module (Disease Evolution Model):
- Functions as the "World Model." It takes the pre-treatment latent state, clinical context, temporal intervals, and a specific therapy embedding to predict the post-treatment latent state ( $\hat{z}_{post}$ ).
- Post-Treatment Latent Predictor: Uses a Transformer to forecast future latent states. It is trained with a latent consistency loss (matching ground truth) and a soft-label contrastive loss (ensuring semantically similar treatments produce similar latent transitions).
- Survival Predictor: A bidirectional cross-attention module that processes both pre- and post-treatment latents to output a 1-year survival probability and a continuous risk score.

Key Mechanisms

Context-Aware Conditioning: The model explicitly encodes continuous time intervals (using sinusoidal embeddings) and multimodal clinical data (genomics, demographics) to handle irregular follow-ups and patient heterogeneity.
Inverse Survival Evaluation (The Feedback Loop):
- Instead of greedily selecting a therapy, CLARITY treats treatment planning as a long-horizon optimization problem.
- It employs an Entropy-Regularized objective. The Therapy Policy Agent proposes a set of candidate trajectories.
- The Actor simulates these trajectories to estimate cumulative survival risk.
- These risk scores are fed back to the Agent to iteratively refine therapy proposals (up to $K=3$ iterations), shifting probability mass toward lower-risk strategies while maintaining exploration diversity.

3. Key Contributions

Treatment-Conditioned Disease Modeling: CLARITY models disease progression as continuous physiological transitions within a visual representation space, enabling counterfactual simulation of how different treatments reshape future dynamics.
Temporal and Clinical Conditioning: The framework integrates continuous time encoding and heterogeneous patient data (omics, history) to generate personalized, longitudinally coherent predictions.
Prediction-to-Decision Integration: It introduces an Inverse Survival Evaluation framework that transforms latent rollouts into adaptive treatment recommendations via a closed-loop, entropy-regularized planning process.

4. Experimental Results

The model was evaluated on three datasets: MU-Glioma-Post (brain tumors), UCSF-ALPTDG (brain tumors, zero-shot), and ISPY-2 (breast cancer).

Treatment Planning Performance:
- CLARITY achieved a 57.1% F1-score on MU-Glioma-Post, outperforming the strongest medical baseline (Huatuo-Vision, 46.4%) and re-implemented diffusion models (MeWM*, 43.6%).
- It demonstrated strong zero-shot generalization on UCSF-ALPTDG (48.9% F1) and breast cancer data (53.9% F1).
Survival Analysis:
- Achieved a state-of-the-art C-index of 0.794, significantly outperforming DeepSurv (0.664) and RSF (0.672).
- Kaplan-Meier curves showed superior separation between high- and low-risk strata with a log-rank p-value of 0.0017.
Efficiency:
- CLARITY is 9–15x more computationally efficient than diffusion-based approaches (4.21 TFLOPs vs. 61.3 TFLOPs) and achieves sub-second inference (0.34s).
- It reduces representation drift (cosine distance to ground truth) by nearly 50% compared to pixel-space generation methods, ensuring higher fidelity in downstream survival evaluation.

5. Significance

CLARITY represents a paradigm shift in medical AI from static outcome prediction to dynamic treatment planning.

Clinical Utility: By simulating counterfactual trajectories, it provides clinicians with data-driven insights into how alternative therapies might alter a patient's disease course, directly addressing the uncertainty of oncology decision-making.
Technical Advancement: It proves that latent space modeling is superior to pixel-space generation for medical decision support, offering a balance of high fidelity, interpretability, and computational feasibility.
Generalizability: The framework's success across different cancer types (brain and breast) suggests a scalable approach for personalized medicine in other complex, longitudinal disease domains.