Treatment-Aware Spatio-Temporal Diffusion Models

• The paper presents two formulations of treatment-aware diffusion models, one using causal patch selection and the other leveraging treatment embeddings for clinical forecasting.
• It employs selective inpainting and diffusion sampling to disentangle causal regions from environmental noise, thereby reducing computational complexity and mitigating spurious correlations.
• Empirical results on diverse benchmarks and MS datasets validate the models’ capability in both generic prediction tasks and personalized counterfactual medical forecasting.

Searching arXiv for the specified papers and closely related work on treatment-aware spatio-temporal diffusion. Treatment-aware spatio-temporal diffusion models are generative frameworks that integrate spatio-temporal forecasting with explicit conditioning on interventions, treatments, or causally salient regions. In the literature represented here, the term covers two distinct but related formulations. One formulation treats selected spatio-temporal patches as “treatment” or causal patches and models the remaining regions as environmental or non-causal variables to be inpainted with a diffusion prior, as in CaPaint for spatio-temporal prediction (Duan et al., 2024). The other formulation conditions a spatio-temporal diffusion process on clinical treatment assignments and longitudinal imaging features in order to generate future lesion masks under specific therapies, as in forecasting future multiple sclerosis lesion masks conditioned on treatments (Favero et al., 9 Aug 2025). Taken together, these works place diffusion modeling inside a causal or intervention-aware setup, with treatment-awareness arising either from causal partitioning of input regions or from explicit therapy embeddings.

1. Conceptual scope and problem setting

In CaPaint, the observed object is a historical spatio-temporal sequence $X_1,\dots,X_T$ with $X_t\in\mathbb{R}^{H\times W\times C_{in}}$, and the task is to predict the next $K$ frames $Y:=\{X_{T+1},\dots,X_{T+K}\}$ (Duan et al., 2024). The central causal assumption is that each frame can be decomposed into two disjoint sets of patches, $X_t=[X_t^c;X_t^e]$, where $X_t^c$ denotes “treatment” or causal patches that genuinely drive $Y$, and $X_t^e$ denotes “environmental” or non-causal patches that induce spurious correlations. The simplified structural causal model is stated as $X_t \to \{X_t^c,X_t^e\}\to Y$ with $X_t^e \nrightarrow Y$ in the sense that its contribution is only spurious (Duan et al., 2024).

In the multiple sclerosis forecasting model, the target is a future NET2 lesion mask at a target visit $X_t\in\mathbb{R}^{H\times W\times C_{in}}$0, conditioned on MRI visits at pre-treatment $X_t\in\mathbb{R}^{H\times W\times C_{in}}$1 and interim $X_t\in\mathbb{R}^{H\times W\times C_{in}}$2, together with treatment information (Favero et al., 9 Aug 2025). The model follows the Variational Diffusion Model formalism with a forward noising process $X_t\in\mathbb{R}^{H\times W\times C_{in}}$3 and a learned reverse denoising process $X_t\in\mathbb{R}^{H\times W\times C_{in}}$4. The clean target is a NET2 label volume $X_t\in\mathbb{R}^{H\times W\times C_{in}}$5, where $X_t\in\mathbb{R}^{H\times W\times C_{in}}$6 axial slices, converted into a pseudo-2D tensor by folding the slice dimension into the channel axis (Favero et al., 9 Aug 2025).

These two formulations instantiate treatment-awareness differently. In CaPaint, “treatment” refers to causal spatio-temporal regions selected from the observed sequence. In the multiple sclerosis model, treatment refers to the patient’s treatment arm among six categories plus a null category for classifier-free dropout (Favero et al., 9 Aug 2025). A plausible implication is that “treatment-aware spatio-temporal diffusion model” is not a single canonical architecture, but a family of models in which the diffusion process is conditioned on intervention-relevant structure.

2. Causal formulation and treatment semantics

CaPaint is explicitly framed in causal terms. After identifying $X_t\in\mathbb{R}^{H\times W\times C_{in}}$7 and $X_t\in\mathbb{R}^{H\times W\times C_{in}}$8, the conceptual starting point is the back-door formula

$X_t\in\mathbb{R}^{H\times W\times C_{in}}$9

Operationally, this is approximated by masking out $K$0, generating multiple samples $K$1 from its marginal, and averaging: $K$2 The paper notes that NuwaDynamics enumerates exponentially many environmental configurations, whereas CaPaint uses front-door adjustment to avoid that cost, although the back-door formula remains the conceptual starting point (Duan et al., 2024).

The multiple sclerosis model is treatment-aware through direct conditioning rather than through a patchwise causal partition. Its conditioning variable is

$K$3

where $K$4 is a learned embedding of the patient’s treatment arm and $K$5 is the feature map output of a frozen encoder processing the multi-modal MRI volumes $K$6 comprising FLAIR, T2, and Gd-enhanced data (Favero et al., 9 Aug 2025). Counterfactual generation is performed by varying $K$7 while keeping $K$8 fixed, thereby producing future NET2 masks under different therapies (Favero et al., 9 Aug 2025).

The contrast between these formulations is methodologically important. In one case, treatment is an internal latent partition of the observed field into causal and environmental regions; in the other, treatment is an externally assigned intervention represented as a discrete arm. This suggests that treatment-awareness in spatio-temporal diffusion can be grounded either in causal discovery over the input domain or in interventional covariates supplied by the application domain.

3. Identification and conditioning mechanisms

CaPaint identifies causal versus non-causal regions with a Vision Transformer backbone trained on a self-supervised spatio-temporal reconstruction task: $K$9 Per-patch multi-head attention is extracted at each layer,

$Y:=\{X_{T+1},\dots,X_{T+K}\}$0

and aggregated across heads and query positions into an importance score

$Y:=\{X_{T+1},\dots,X_{T+K}\}$1

The $Y:=\{X_{T+1},\dots,X_{T+K}\}$2 patches are ranked by $Y:=\{X_{T+1},\dots,X_{T+K}\}$3, and the top $Y:=\{X_{T+1},\dots,X_{T+K}\}$4 are deemed causal $Y:=\{X_{T+1},\dots,X_{T+K}\}$5, while the remainder $Y:=\{X_{T+1},\dots,X_{T+K}\}$6 are treated as non-causal (Duan et al., 2024). The resulting binary mask $Y:=\{X_{T+1},\dots,X_{T+K}\}$7 is then used during diffusion inpainting and later data augmentation.

The multiple sclerosis model uses two conditioning channels. First, treatment embedding is defined by a standard lookup

$Y:=\{X_{T+1},\dots,X_{T+K}\}$8

with $Y:=\{X_{T+1},\dots,X_{T+K}\}$9 indexing placebo, NE, LE, MiE, MoE, and HE (Favero et al., 9 Aug 2025). This embedding is summed with a time-step embedding $X_t=[X_t^c;X_t^e]$0, and the result $X_t=[X_t^c;X_t^e]$1 is injected into every ResBlock of the U-Net via learned linear layers (Favero et al., 9 Aug 2025).

Second, image conditioning is introduced through ControlNet fusion. The ControlNet is described as a frozen copy of the Stage-1 diffusion encoder, fine-tuned to process the concatenated pseudo-2D slices of $X_t=[X_t^c;X_t^e]$2. At each U-Net resolution $X_t=[X_t^c;X_t^e]$3, the ControlNet feature map $X_t=[X_t^c;X_t^e]$4 is merged into the main U-Net feature map $X_t=[X_t^c;X_t^e]$5 via a zero-initialized $X_t=[X_t^c;X_t^e]$6 convolution: $X_t=[X_t^c;X_t^e]$7 Because the convolution is zero-initialized, the network can gradually learn to attend to image conditioning without disturbing the pre-trained diffusion prior (Favero et al., 9 Aug 2025).

A common element across both systems is selective conditioning. CaPaint conditions on a mask that distinguishes causal from environmental patches. The multiple sclerosis model conditions on both therapy assignment and longitudinal imaging. In both cases, the diffusion model is not merely unconditional generation over spatio-temporal data; it is constrained by variables intended to encode intervention-relevant structure.

4. Diffusion architecture and sampling procedures

CaPaint adopts an unconditional DDPM $X_t=[X_t^c;X_t^e]$8 pretrained on ImageNet or spatio-temporal snapshots and fine-tunes it on spatio-temporal frames $X_t=[X_t^c;X_t^e]$9 (Duan et al., 2024). The denoising objective is

$X_t^c$0

where $X_t^c$1 is the causal mask encoding $X_t^c$2 location (Duan et al., 2024). At sampling time, two sub-images are maintained: $X_t^c$3, targeting the causal patches and kept anchored to the original $X_t^c$4, and $X_t^c$5, targeting the environmental patches and re-sampled via the learned denoiser (Duan et al., 2024). For each diffusion step $X_t^c$6, the causal branch is updated by

$X_t^c$7

the environmental branch by

$X_t^c$8

and the result is merged using the binary mask $X_t^c$9: $Y$0 The stated purpose is to inpaint only the spurious region while respecting the treatment region (Duan et al., 2024).

The multiple sclerosis model uses a forward diffusion process

$Y$1

or equivalently

$Y$2

with $Y$3 and $Y$4 (Favero et al., 9 Aug 2025). The reverse process is

$Y$5

where the mean is produced by a U-Net denoiser (Favero et al., 9 Aug 2025).

The U-Net backbone uses 6 down-sampling stages and 6 up-sampling stages, with optional self-attention at spatial resolutions $Y$6 and $Y$7 (Favero et al., 9 Aug 2025). Input NET2 masks and MRI volumes of shape $Y$8 are reshaped to $Y$9, with $X_t^e$0 axial slices folded into channels (Favero et al., 9 Aug 2025). The final output predicts the noise $X_t^e$1 with shape $X_t^e$2 (Favero et al., 9 Aug 2025).

At inference, classifier-free guidance is applied: $X_t^e$3 where $X_t^e$4 denotes both embeddings dropped out (Favero et al., 9 Aug 2025). The sampling loop is a standard DDPM sampler initialized with $X_t^e$5, and multiple draws, for example 100, yield an ensemble of plausible futures whose mean or heatmap summarizes uncertainty (Favero et al., 9 Aug 2025).

5. Training strategy, augmentation, and computational properties

CaPaint is organized as a unified causal-diffusion pipeline. Historical frames $X_t^e$6 are divided into patches; ViT self-attention produces importance vectors; causal and non-causal indices $X_t^e$7 are extracted to build masks $X_t^e$8; a DDPM is fine-tuned on $X_t^e$9 with mask side information; diffusion inpainting is run for each sequence replicate $X_t \to \{X_t^c,X_t^e\}\to Y$0 to produce augmented sequences $X_t \to \{X_t^c,X_t^e\}\to Y$1; original and augmented sequences are mixed with sampling probability $X_t \to \{X_t^c,X_t^e\}\to Y$2 per time step into the training set; any spatio-temporal predictor $X_t \to \{X_t^c,X_t^e\}\to Y$3 is trained on this mixture; and at test time only $X_t \to \{X_t^c,X_t^e\}\to Y$4 on the original $X_t \to \{X_t^c,X_t^e\}\to Y$5 is used (Duan et al., 2024). The model is therefore not solely a forecaster but also an augmentation mechanism for downstream predictors.

A central claim of CaPaint concerns complexity. NuwaDynamics’ back-door procedure is stated to require exploring $X_t \to \{X_t^c,X_t^e\}\to Y$6 combinations of environmental patches, which is exponential in $X_t \to \{X_t^c,X_t^e\}\to Y$7, whereas CaPaint runs one diffusion pass per time step with $X_t \to \{X_t^c,X_t^e\}\to Y$8 arithmetic operations, or $X_t \to \{X_t^c,X_t^e\}\to Y$9 with $X_t^e \nrightarrow Y$0 replicates; because $X_t^e \nrightarrow Y$1, the method remains quasi-linear (Duan et al., 2024). The paper characterizes this as reducing data generation complexity from exponential to quasi-linear levels (Duan et al., 2024).

The multiple sclerosis model uses a two-stage training strategy. Stage 1 pre-trains a diffusion prior on NET2 labels with

$X_t^e \nrightarrow Y$2

with $X_t^e \nrightarrow Y$3 sampled from an interpolated-shifted cosine schedule and min-SNR weighting $X_t^e \nrightarrow Y$4 (Favero et al., 9 Aug 2025). Stage 2 freezes the U-Net prior and trains only the ControlNet zero-convs using

$X_t^e \nrightarrow Y$5

(Favero et al., 9 Aug 2025).

Because NET2 lesions occupy $X_t^e \nrightarrow Y$6 of voxels, lesion voxels are up-weighted by factor 10 in the MSE using $X_t^e \nrightarrow Y$7 (Favero et al., 9 Aug 2025). During Stage 1, $X_t^e \nrightarrow Y$8 is dropped 10% of the time to enable classifier-free guidance; during Stage 2, $X_t^e \nrightarrow Y$9 and/or $X_t\in\mathbb{R}^{H\times W\times C_{in}}$00 are each dropped 10% of the time (Favero et al., 9 Aug 2025).

6. Empirical results and application domains

CaPaint is evaluated on five real-world spatio-temporal benchmarks: TaxiBJ$X_t\in\mathbb{R}^{H\times W\times C_{in}}$01, KTH, SEVIR, DRS, and FireSys (Duan et al., 2024). The reported backbones are ConvLSTM, PredRNN-V2, ViT, MAU, SimVP, Mmvp, and EarthFarsser, and the metrics are MAE, MSE, and SSIM (Duan et al., 2024). The reported gains across combinations are 3.7%–77.3% relative improvement, with examples including an MSE drop from $X_t\in\mathbb{R}^{H\times W\times C_{in}}$02 on DRS and an MAE drop from $X_t\in\mathbb{R}^{H\times W\times C_{in}}$03 (Duan et al., 2024). In data-scarce regimes using 10%, 25%, and 50% of the data, the model yields the largest SSIM gains when data are few; in comparisons against Flip, Rotate, Crop, and Nuwa mix-up, only CaPaint improves MAE consistently on all five benchmarks (Duan et al., 2024).

The multiple sclerosis model is evaluated on a multi-centre dataset of 2131 patient 3D MRIs from randomized clinical trials for relapsing-remitting MS across six different treatments (Favero et al., 9 Aug 2025). It is described as the first treatment-aware spatio-temporal diffusion model able to generate future masks demonstrating lesion evolution in MS (Favero et al., 9 Aug 2025). The model is assessed on regional spatial accuracy, downstream lesion-activity classification, and log-count regression.

The regional metrics on the held-out test set are as follows:

Region	Proposed model	Monte Carlo baseline
Cerebral WM BA / Precision / Recall / F1	0.684 / 0.537 / 0.629 / 0.579	0.497 / 0.321 / 0.322 / 0.320
Cerebral CX BA / Precision / Recall / F1	0.697 / 0.367 / 0.562 / 0.444	0.503 / 0.154 / 0.135 / 0.143

For binary lesion-activity classification at week 96, the reported AP/AUC by treatment arm are: Placebo $X_t\in\mathbb{R}^{H\times W\times C_{in}}$04 versus $X_t\in\mathbb{R}^{H\times W\times C_{in}}$05; NE $X_t\in\mathbb{R}^{H\times W\times C_{in}}$06 versus $X_t\in\mathbb{R}^{H\times W\times C_{in}}$07; LE $X_t\in\mathbb{R}^{H\times W\times C_{in}}$08 versus $X_t\in\mathbb{R}^{H\times W\times C_{in}}$09; MiE $X_t\in\mathbb{R}^{H\times W\times C_{in}}$10 versus $X_t\in\mathbb{R}^{H\times W\times C_{in}}$11; MoE $X_t\in\mathbb{R}^{H\times W\times C_{in}}$12 versus $X_t\in\mathbb{R}^{H\times W\times C_{in}}$13; and HE $X_t\in\mathbb{R}^{H\times W\times C_{in}}$14 versus $X_t\in\mathbb{R}^{H\times W\times C_{in}}$15 (Favero et al., 9 Aug 2025). For log-count regression, the reported MSE values are Placebo $X_t\in\mathbb{R}^{H\times W\times C_{in}}$16 versus $X_t\in\mathbb{R}^{H\times W\times C_{in}}$17, NE $X_t\in\mathbb{R}^{H\times W\times C_{in}}$18 versus $X_t\in\mathbb{R}^{H\times W\times C_{in}}$19, LE $X_t\in\mathbb{R}^{H\times W\times C_{in}}$20 versus $X_t\in\mathbb{R}^{H\times W\times C_{in}}$21, MiE $X_t\in\mathbb{R}^{H\times W\times C_{in}}$22 versus $X_t\in\mathbb{R}^{H\times W\times C_{in}}$23, MoE $X_t\in\mathbb{R}^{H\times W\times C_{in}}$24 versus $X_t\in\mathbb{R}^{H\times W\times C_{in}}$25, and HE $X_t\in\mathbb{R}^{H\times W\times C_{in}}$26 versus $X_t\in\mathbb{R}^{H\times W\times C_{in}}$27 (Favero et al., 9 Aug 2025). Qualitative figures are said to show variability in lesion placement around periventricular white matter and systematic lesion reduction under higher-efficacy treatments (Favero et al., 9 Aug 2025).

These results illustrate two application regimes. CaPaint addresses generic spatio-temporal prediction and augmentation in settings such as meteorology, human actions, reaction-diffusion PDEs, taxi flow, and wildfire spread (Duan et al., 2024). The multiple sclerosis model addresses personalized medical forecasting and counterfactual therapy analysis (Favero et al., 9 Aug 2025).

7. Interpretive significance, limitations, and common points of confusion

A common source of confusion is the meaning of “treatment-aware.” In CaPaint, the “treatment” is not a clinical intervention but the subset of causal patches $X_t\in\mathbb{R}^{H\times W\times C_{in}}$28 that genuinely drive future outcomes, while the environmental subset $X_t\in\mathbb{R}^{H\times W\times C_{in}}$29 is treated as a source of spurious correlation to be regenerated by inpainting (Duan et al., 2024). In the multiple sclerosis model, treatment is literal treatment-arm information embedded as a conditioning variable $X_t\in\mathbb{R}^{H\times W\times C_{in}}$30 (Favero et al., 9 Aug 2025). These are distinct semantics that share a causal orientation.

Another potential confusion concerns whether the causal claims are purely empirical or formally justified. CaPaint explicitly invokes back-door adjustment and states that front-door logic is used to justify avoiding exponential enumeration of environmental configurations (Duan et al., 2024). The multiple sclerosis model emphasizes causal, image-based generative modeling and enables counterfactual future NET2 masks under different efficacies by varying treatment embeddings while keeping MRI conditioning fixed (Favero et al., 9 Aug 2025). This suggests an intervention-aware use of diffusion models, although the paper also notes that more sophisticated counterfactual analyses and causal metrics remain future work (Favero et al., 9 Aug 2025).

The limitations in the multiple sclerosis setting are explicit. MS lesion formation varies widely across individuals and plausible trajectories, so single-outcome metrics such as Dice or MSE can penalize clinically valid but different forecasts; multi-centre MRI differences and NET2 segmentation noise can divert model capacity; and while the Monte Carlo baselines reflect real-world population-level prognosis, more sophisticated counterfactual analyses and causal metrics remain future work (Favero et al., 9 Aug 2025). In CaPaint, the main stated problem setting is data scarcity, data imbalance, and the lack of causal connections in overly customized models, which undermine generalizability and interpretability (Duan et al., 2024).

Taken together, these studies define an emerging research direction in which diffusion models are coupled to intervention-aware structure in spatio-temporal data. One branch uses causal region identification and selective inpainting to regularize generic spatio-temporal forecasting (Duan et al., 2024). Another branch uses explicit treatment embeddings and longitudinal imaging conditioning to generate future disease states and counterfactual outcomes (Favero et al., 9 Aug 2025). A plausible implication is that future treatment-aware spatio-temporal diffusion models may increasingly unify causal adjustment, structured conditioning, uncertainty-aware generation, and downstream decision support within a single framework.