Admissible Diffusion for Multimodal Interventional Trajectories

Xing Han; Shravan Chaudhari; Jiarui Shao; Paul Pu Liang; Suchi Saria

cs.LGSep 28, 2026

Admissible Diffusion for Multimodal Interventional Trajectories

Authors: Xing Han, Shravan Chaudhari, Jiarui Shao, Paul Pu Liang, Suchi Saria

Organizations: Johns Hopkins University · Massachusetts Institute of Technology · Bayesian Health

Abstract

Generating a plausible clinical trajectory does not establish what would happen under a different treatment. We present ADMIT, a framework combining irregular multimodal representations, treatment-conditioned latent diffusion and explicit constraints on generated states or actions. We formulate its interventional target through sequential g-computation and distinguish causal assumptions from constraint satisfaction. Its admissibility mechanism translates physiological prior knowledge into explicit constraints on generated states and proposed actions. Treatment-exposure dynamics condition latent transitions, while state projection or action gating applies the constraints during rollout so that they influence subsequent trajectory generation. In our preliminary experiments, multimodal inputs improved supervised hidden-state recovery and reduced treatment-contrast error. In a simulated dosing-schedule experiment with leak-free history encoding, ADMIT predicted most of the tumor-volume change caused by redistributing a fixed total dose. An exposure input improved these predictions around a temporary dose reduction whether or not the assumed clearance rate was correct, but reduced the predicted size of a dose effect, and a deterministic recurrent baseline matched ADMIT's average predictions. Exposure projection reduced constraint violations, although enforcement remained incomplete. Semi-synthetic experiments using eICU context illustrated treatment-response generation under fixed and adaptive policies. Observational examples further characterize model treatment sensitivity. ADMIT provides a framework for testing whether complementary observations and physiological restrictions improve intervention trajectories, with representation recovery, effect accuracy and rule enforcement assessed separately.

Figures & tables

Figure 1: ADMIT framework. Aligned multimodal observations form latent states. Treatment-conditioned diffusion generates successive blocks, with optional constraints before updating the history. Modality-specific heads reconstruct inputs and decode selected targets. The experiments generate numerical outputs; the notes channel uses categorical tokens. The physiology guard approximately enforces a specified rule.

$H$ visibility	Inputs	$H$ MSE	Contrast RMSE	Contrast correlation
1	Multimodal	$0.130\pm 0.007$	$0.755\pm 0.022$	$0.838\pm 0.010$
1	Target only	$0.159\pm 0.016$	$0.783\pm 0.014$	$0.838\pm 0.004$
0	Multimodal	$0.203\pm 0.012$	$0.783\pm 0.008$	$0.843\pm 0.016$
0	Target only	$0.784\pm 0.045$	$0.913\pm 0.034$	$0.827\pm 0.004$

Table 1: Synthetic representation and treatment-contrast diagnostics. Mean

\pm

sample standard deviation across three seeds. Visibility zero removes the direct contribution of

H

to target observations. The

H

readout is jointly supervised. Treatment contrasts compare always treating with never treating, pooled over 20 future steps and three standardized target channels.

Figure 2: Synthetic tumor trajectories. Four examples span severity at intervention after 20 steps. Solid lines and bands show generated means and sample standard deviations; dashed lines show paired simulator trajectories. Treatment plans never treat, always treat, or treat when tumor volume exceeds 12. The model updates the adaptive plan once per generated block; the simulator updates it every step. Black history points denote simulator states. The model encodes the full observed record, including future observations, so these examples do not establish forecasting accuracy.

	Schedule-difference RMSE			Trajectory	Dose-effect
Model	Front-loaded	Break	Break, steps 5–8	RMSE	ratio
Zero-effect reference	0.465	0.207	0.197	—	0
ADMIT, no exposure	$0.104\pm 0.010$	$0.081\pm 0.009$	$0.089\pm 0.030$	$0.62\pm 0.16$	$0.96\pm 0.03$
ADMIT, $\lambda=0.60$	$0.096\pm 0.003$	$0.061\pm 0.002$	$0.064\pm 0.014$	$0.56\pm 0.16$	$0.89\pm 0.03$
ADMIT, $\lambda=0.78$ (true)	$0.102\pm 0.012$	$0.064\pm 0.007$	$0.064\pm 0.010$	$0.50\pm 0.15$	$0.84\pm 0.03$
ADMIT, $\lambda=0.90$	$0.107\pm 0.012$	$0.064\pm 0.006$	$0.065\pm 0.007$	$0.48\pm 0.12$	$0.80\pm 0.02$

Table 2: Predicting the effects of dose timing and dose size. Mean

\pm

sample standard deviation across three training seeds, each scoring the same 300 test patients. The first three columns give the RMSE of predicted schedule differences and the fourth the RMSE of predicted tumor volume, averaged over the three schedules; lower is better for all four. Every model has higher trajectory error in seed 2, which dominates that column’s standard deviation. The dose-effect ratio compares predicted and simulated effects of sustained dosing at 0.6 versus 0.2; one is correct.

\lambda

is the retention used to compute the exposure input; the simulator uses 0.78.

Policy	Outcome RMSE	Effect RMSE	Effect correlation
Never treat	0.535	—	—
Always treat	0.601	0.488	0.679
Adaptive	0.498	0.357	0.724

Table 3: Treatment comparisons with real clinical context and simulated outcomes. Single training seed, 1,712 validation stays, 16 future hours and two standardized synthetic targets. Effects are relative to never treating after the forecast begins; the reference policy has zero effect by definition. The adaptive plan uses separate thresholds to start and stop treatment. Values describe the latest completed configuration and do not estimate variability across training runs.

Target	Policy contrast	Factual RMSE	Contrast/RMSE
Lactate	0.067	1.170	0.058
Urine output	0.120	0.466	0.256
Creatinine	0.018	0.398	0.045

Table 4: Sensitivity to treatment changes in observational clinical records. One-seed normalized results for 1,687 eICU stays. Policy contrast is the mean absolute difference between sustained-treatment and no-treatment predictions. Factual RMSE scores observed entries; the final column is a descriptive scale ratio.

Appendix figures & tables3 assets

Supplementary material from the paper’s appendix.

Appendix

Simulator	Tumor model of Section 4 with one continuous drug and the hidden comorbidity off; on the logit scale, each dose is centered at 0.9 times the previous dose with noise standard deviation 0.8; retention $\lambda=e^{-1/4}$ ; 12 history and 12 forecast steps; block size four.
Patients	1,000 training, 200 validation and 300 test patients, fixed for all conditions and training seeds.
Stage 1	Encoder of Section 3 with time attention and a local kernel; 16-dimensional latent state; 60 epochs, batch size 128, learning rate $10^{-3}$ ; no exposure head; also supervised with the simulator’s noise-free target signal.
Transition	40 epochs, batch size 128, learning rate $10^{-3}$ , hidden size 64, 100 diffusion steps; rollout-loss weight three over four free-running steps; no condition dropout; guidance scale one.
GRU baseline	Deterministic recurrent network with hidden size 128 whose input is the latent state, dose and exposure; 200 epochs; trained on true previous states and on its own previous predictions.
Evaluation	64 ADMIT samples with sampling noise shared across schedules; simulator references average 32 runs with process noise shared across schedules.

Appendix

Table 5: Dosing-schedule experiment settings.

	History change			Probe $R^{2}$
Seed	Leak-free	Full record (max / mean)	Exposure error	None	$\lambda=0.60$	$\lambda=0.78$	$\lambda=0.90$
0	0	6.75 / 0.51	$<3\times 10^{-7}$	0.962	0.969	0.971	0.973
1	0	4.98 / 0.54	$<3\times 10^{-7}$	0.961	0.971	0.969	0.971
2	0	5.09 / 0.48	$<3\times 10^{-7}$	0.961	0.969	0.973	0.969

Appendix

Table 6: Validity checks per training seed. History change: largest absolute change in the encoded history after all observations following the last history step are replaced. Exposure error: largest absolute difference between the model’s exposure (

\lambda=0.78

) and the simulator’s drug level over factual doses and all planned schedules. Probe

R^{2}

: test-set fraction of variance in the simulator’s exposure explained by a linear regression on the history summary

R_{t}

, for each exposure input.

Comparison	Difference	Seed 0	Seed 1	Seed 2
No exposure $-$ true $\lambda$	Front-loaded	$-$ 0.019 [ $-$ 0.023, $-$ 0.014]	0.012 [0.006, 0.017]	0.014 [0.009, 0.018]
	Break	0.003 [0.001, 0.005]	0.016 [0.013, 0.018]	0.033 [0.029, 0.036]
$\lambda=0.60$ $-$ true $\lambda$	Front-loaded	$-$ 0.018 [ $-$ 0.021, $-$ 0.016]	0.003 [ $-$ 0.000, 0.005]	$-$ 0.002 [ $-$ 0.005, 0.001]
	Break	$-$ 0.012 [ $-$ 0.014, $-$ 0.011]	0.002 [0.001, 0.003]	0.001 [ $-$ 0.001, 0.003]
$\lambda=0.90$ $-$ true $\lambda$	Front-loaded	$-$ 0.004 [ $-$ 0.006, $-$ 0.001]	0.005 [0.003, 0.007]	0.016 [0.013, 0.019]
	Break	$-$ 0.012 [ $-$ 0.014, $-$ 0.011]	$-$ 0.001 [ $-$ 0.003, 0.000]	0.012 [0.010, 0.013]

Appendix

Table 7: Paired patient-bootstrap differences in schedule-difference RMSE. For each training seed, the 300 test patients are resampled with replacement 4,000 times, with the same resample for both models. Entries give the first model’s RMSE minus the second’s, with a 95% percentile interval in brackets; negative values favor the first model. ADMIT and GRU rows without a stated input use the true

\lambda=0.78

Explore similar work

May 21, 2026cs.LG

ChronoMedicalWorld: A Medical World Model for Learning Patient Trajectories from Longitudinal Care Data

Long-horizon clinical simulation -- predicting how a patient's physiology evolves over years under specified interventions -- is central to chronic-disease care, yet existing electronic health record (EHR) models are predominantly discriminative, and general-purpose large language models drift under repeated interventions. We propose the \textbf{ChronoMedicalWorld Model (CMWM)}, an action-conditioned latent world-model framework for learning patient trajectories from longitudinal care data. CMWM couples a joint-embedding state encoder with a wide action encoder that admits both structured intervention indicators and free-text communication embeddings, and trains a recurrent latent transition module under a six-term objective: next-observation supervision, next-latent prediction, SIGReg latent regularisation, and three physiology-aware shape priors (slope, continuity, large-jump penalty). A closed-loop rollout-prefix protocol matches training to deployment, so the model is optimised against the same multi-step error it exhibits at inference. As a concrete case study, we instantiate CMWM for annual estimated glomerular filtration rate (eGFR) trajectory forecasting in chronic kidney disease (CKD). On a 2{,}232-patient nephrology cohort, the CKD instantiation achieves a dynamic-50% history rollout test mean absolute error (MAE) of 7.384 and root-mean-square error (RMSE) of 10.256, against 7.964 and 11.069 for a tuned GPT-5.5 structured-prompting baseline (

-7.28\%

MAE,

-7.35\%

RMSE), with the gain dominated by the dialogue portion of patient--health-coach communication. The framework is not CKD-specific: its architecture, loss design, and training protocol apply to any chronic condition that can be cast as periodic clinical state interleaved with structured and conversational interventions.

Jiangyuan Wang, Xuyong Chen, Junwei He +3

Beijing KidneyTec Medical Technology Co., Ltd., Beijing, China

Patient Trajectories Medical World Model Simulated Patients Chronic Kidney Disease Longitudinal Data

Ask PDF

Jul 28, 2026cs.LG

DRIFT: Direct-Recursive Intervention-Conditioned Forecasting of ICU Physiological Trajectories

Many time-series forecasts depend not only on prior observations but also on actions specified during the forecast period. In intensive care units (ICUs), future vital signs and laboratory values are influenced by treatments such as vasopressors. However, models that predict the full future sequence all at once make little use of these treatments, whereas autoregressive models can accumulate errors. We introduce DRIFT, a hybrid framework in which a direct model produces the primary forecast and a recursive, action-conditioned model contributes constrained corrections. We evaluate DRIFT on 6,046 admissions from MIMIC-IV and 8,345 admissions from eICU-CRD. Averaged across the 8-, 24-, and 48-hour forecast endpoints, DRIFT reduces mean absolute error for mean arterial pressure (MAP) by 0.673% relative to an action-conditioned Temporal Fusion Transformer (TFT-action) on MIMIC-IV and achieves the lowest corresponding error among the compared models on eICU-CRD. Although the overall accuracy improvement is modest, a MIMIC-IV audit restricted to windows in which the supplied treatment sequence was altered showed that DRIFT achieved lower observed-target MAP error than TFT-action at 8 and 24 hours. Treatment-sequence alteration increased DRIFT's MAP error by 0.21-0.26 mmHg more than it increased TFT-action's error, with prediction changes occurring primarily after the supplied paths diverged. In a separate robustness experiment, the MAP advantage persisted under three shared checkpoint-selection rules emphasizing overall endpoint error, MAP error, or both equally.

Weixin Liu, Juming Xiong, Congning Ni +4

Vanderbilt University, Nashville, Tennessee, USA · Vanderbilt University Medical Center, Nashville, Tennessee, USA

Intensive Care Unit Mimic-Iv Datasets Multi-Stream Temporal Fusion Treatment

Ask PDF

May 16, 2026cs.AI

From Static Risk to Dynamic Trajectories: Toward World-Model-Inspired Clinical Prediction

Clinical decision-making is a feedback system where risk estimates influence treatment, which in turn changes disease trajectories, and both shape clinicians' measurement practices. Static prediction often fails clinically: models trained on observational care logs conflate disease biology with clinician behavior, particularly under treatment confounder feedback and irregular or informative observation. This Review focuses on intervention-aware disease trajectory modeling in clinical AI--methods estimating patient-specific longitudinal disease evolution and assessing trajectory changes under alternative treatments. We organize the field around six linked components: three decision tasks (factual forecasting, counterfactual estimation, policy evaluation) and three data-generating mechanisms (disease evolution, treatment assignment, observation process) that determine identifiability. We present the first unified framework bridging forecasting, counterfactual trajectories, and policy evaluation across discrete/continuous time, explicitly addressing treatment assignment, time-varying confounding, and observation bias. We synthesize key method families (multistate/joint models, temporal point-process, deep sequence architectures, longitudinal causal inference), map them to relevant components, and align evaluation with claim strength via overlap diagnostics, uncertainty quantification, off-policy robustness, and target-trial validation. This synthesis advances benchmark prediction to decision-grade clinical evidence, enabling treatment-sensitive individualized futures, pre-deployment policy stress-testing, and safer closed-loop learning health systems that adapt/abstain when evidence is insufficient.

Pujun Feng, Xiaoyu Guo, Seyed Ehsan Saffari +10

Faculty of Applied Sciences, Macao Polytechnic University, Macao SAR 999078, China. · Medin.ai, Beijing 100871, China. · School of Software & Microelectronics, Peking University, Beijing 100871, China. +8

Readmission Prediction Clinical Prediction Clinical Decision Support Treatment Latent Confounders

Ask PDF