💻 computer science

DeepImagine: Learning Biomedical Reasoning via Successive Counterfactual Imagining

DeepImagine is a framework that teaches large language models biomedical reasoning by training them to predict how clinical trial outcomes would change under various experimental perturbations through a combination of supervised fine-tuning on counterfactual pairs, reinforcement learning, and synthetic reasoning traces.

Original authors: Youze Zheng, Jianyou Wang, Yuhan Chen, Matthew Feng, Longtian Bao, Hanyuan Zhang, Maxim Khan, Aditya K. Sehgal, Christopher D. Rosin, Umber Dube, Ramamohan Paturi

Published 2026-04-28

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Youze Zheng, Jianyou Wang, Yuhan Chen, Matthew Feng, Longtian Bao, Hanyuan Zhang, Maxim Khan, Aditya K. Sehgal, Christopher D. Rosin, Umber Dube, Ramamohan Paturi

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 master chef trying to predict if a brand-new, experimental recipe will be a hit at a food critic's table.

You could just look at the ingredients and guess (that’s what current AI does—it looks at "correlations"). But a true master chef doesn't just guess; they imagine. They think: "If I used sea salt instead of table salt, how would the flavor change? If I turned the heat up by 10 degrees, would the crust brown faster?" By mentally running these "what-if" scenarios, the chef builds a deep, intuitive understanding of how heat, salt, and time interact.

This paper, DeepImagine, proposes doing exactly that for medical science.

The Problem: The "Guessing" Problem

Right now, when Artificial Intelligence tries to predict if a new medical drug trial will succeed, it acts like a person looking at a spreadsheet. It sees that "Trial A had 100 people and succeeded" and "Trial B had 50 people and failed," and it tries to find a pattern. But medicine is incredibly complex. A drug might fail not because it's bad, but because the dosage was slightly off, or it was tested in a different country, or the patients were slightly older. Current AI struggles because it doesn't understand the why—the "hidden laws" of biology.

The Solution: The "What-If" Machine (DeepImagine)

The researchers created a way to teach AI to "imagine" through a process they call Successive Counterfactual Imagining.

Instead of asking the AI, "Will this new trial succeed?" (which is a huge, scary question), they teach it to take small, logical steps. It’s like teaching a child to walk by first teaching them to balance, then to take one step, then two.

Here is how the "Imagination Training" works:

The "Natural" Steps (The Easy Wins):
The researchers find real medical trials where scientists already did "what-if" experiments. For example, a trial that tested a drug at 10mg, 20mg, and 30mg. The AI is shown: "Here is what happened at 10mg. Now, imagine we change only the dose to 20mg. What happens next?" Because the answer is already in the real data, the AI can learn the "rule" of how dosage affects results.
The "Approximate" Steps (The Creative Leap):
Sometimes, there isn't a perfect "what-if" scenario in history. To fix this, the researchers use a clever retrieval system to find "Twin Trials"—studies that are almost identical but have one tiny difference (like being conducted in Germany instead of Japan). The AI learns to bridge the gap between these twins.
The "Reasoning Trace" (The Inner Monologue):
To make sure the AI isn't just memorizing, they give it a "script" to follow. Instead of just jumping from Result A to Result B, the AI is trained to write out its thoughts: "Because we increased the dosage, the blood pressure should drop more significantly, which leads to..." This is like teaching a student to "show their work" in a math problem.
The "Reward" (The Taste Test):
Finally, they use a technique called Reinforcement Learning. If the AI's "imagination" leads it to a correct prediction about a real trial, it gets a "digital gold star" (a reward). If it imagines something biologically impossible, it gets corrected.

Why does this matter?

By training the AI this way, it stops being a simple calculator and starts acting more like a scientific detective.

Instead of just saying, "I think this drug will work," the AI can provide a path: "Based on how similar drugs reacted to higher doses in previous studies, I imagine this trial will succeed because..."

This makes the AI interpretable. Doctors don't want a "black box" that gives them a random answer; they want a partner that can explain its reasoning. DeepImagine is a step toward turning AI from a lucky guesser into a scientific thinker.

Technical Summary: DeepImagine: Learning Biomedical Reasoning via Successive Counterfactual Imagining

1. Problem Statement

Predicting the outcomes of prospective clinical trials is a critical but notoriously difficult task for Large Language Models (LLMs) and traditional machine learning models. Current approaches typically rely on correlational prediction—treating trial attributes (e.g., enrollment size, phase, geography, sponsor) as numerical features to predict success. However, these methods fail to capture the underlying causal mechanisms that govern how changes in experimental conditions (e.g., dosage, outcome measures, or patient demographics) influence trial results. The authors hypothesize that while the "true" causal law of clinical trials is hidden, it can be approximated by reasoning through how results change under controlled perturbations.

2. Methodology

The authors propose DeepImagine, a framework that shifts the LLM's role from a correlational predictor to a mechanistic reasoner through Successive Counterfactual Imagining.

Core Concept (Successive Counterfactual Imagination): Instead of predicting a target trial outcome $R(T)$ directly from a target configuration $x(T)$ , the model traverses a chain of local "what-if" transformations. Starting from a known historical trial $x(0)$ with result $R(0)$ , the model iteratively predicts how the result changes as it moves through intermediate states $x(1), x(2), \dots, x(T)$ , where each step involves changing only one experimental variable at a time.
Training Pipeline:
- Natural Counterfactuals (Supervised Fine-Tuning - SFT): The model is trained on real-world data where "natural" pairs exist. This includes:
  - Outcome-measure pairs: Predicting one reported outcome from another within the same trial.
  - Dose-ranging pairs: Predicting results for one dosage arm from another within the same trial.
- Approximate Counterfactuals (Reinforcement Learning - RL): For variables where natural pairs don't exist (e.g., geography or drug mechanism), the authors use an information retrieval pipeline to find highly similar trials. Since intermediate results are unobservable, they use Group Relative Policy Optimization (GRPO). The model is rewarded only for the terminal correctness of its final prediction, forcing it to discover plausible latent reasoning trajectories.
- Synthetic Reasoning Traces: To improve the quality of the SFT, the authors use a powerful teacher model (GPT-5.4) to generate "causally plausible" explanations for why a perturbation causes a specific change in outcome. This teaches the model how to reason, not just what to predict.
Data Retrieval: They use a graph-based approach using NV-Embed-v2 embeddings to identify "M-approximate counterfactual pairs"—trials that share a high number of similar attributes.

3. Key Contributions

A New Paradigm for Biomedical LLMs: Moves beyond simple pattern matching toward approximating structural causal models (SCMs) through sequential local transitions.
Hybrid Training Framework: Combines SFT (for high-certainty natural pairs) with RL (for low-certainty approximate pairs) to maximize the utility of available clinical data.
Synthetic Causal Supervision: Demonstrates that augmenting training with synthetic, expert-validated reasoning traces significantly boosts the model's ability to perform multi-step counterfactual reasoning.
Interpretability: The framework produces reasoning trajectories that can be analyzed to recover structured, interpretable causal models of clinical trial mechanisms.

4. Results

The authors evaluated a fine-tuned Qwen3.5-9B model against a strong baseline (o3-mini-2025-01-31) on two specialized evaluation sets:

Outcome-measure Perturbation: Predicting a new outcome measure's result based on a known one.
Study-arm Perturbation: Predicting a different dosage/arm result based on a known one.

Performance Gains:

Outcome Perturbation: The fine-tuned Qwen model achieved a Macro F1 of 75.50 and Weighted Accuracy of 74.81, significantly outperforming o3-mini (Macro F1: 67.10; Weighted Acc: 64.71).
Arm Perturbation: The fine-tuned Qwen model achieved a Macro F1 of 72.25 and Weighted Accuracy of 74.72, outperforming o3-mini (Macro F1: 64.46; Weighted Acc: 63.40).

5. Significance

DeepImagine demonstrates that even relatively small LLMs (under 10B parameters) can outperform much larger, state-of-the-art models in specialized biomedical tasks if they are trained to reason through causal transitions rather than direct correlations. This provides a practical roadmap for developing mechanistic biomedical language models that are not only more accurate in forecasting but also more scientifically useful and interpretable for researchers.