Time-Series Causal Normalizing Flows
A framework that extends causal normalizing flows to time-series, enabling simulation-based interventional density estimation over time.
Estimating Interventional Outcomes over Time with Causal Normalizing Flow Accepted at UAI 2026 · joint work with LG AI Research
Abstract
Estimating outcome distributions under time-varying treatments is an essential task for personalized decision-making, particularly in domains such as healthcare. Most prior work focuses on point predictions, which fail to capture the inherent variability in outcomes. Recent efforts have begun integrating generative models to estimate interventional distributions — but existing approaches, including causal normalizing flows, are restricted to static settings and are not well suited to sequential data.
We propose a framework that extends causal normalizing flows to time-series, enabling simulation-based interventional density estimation over time. The method learns representations of treatment and covariate history that capture temporal dependencies; conditioned on these and guided by a causal graph, the flow generates interventional samples, allowing simulation of outcome trajectories under alternative treatment strategies.
Contributions
- Extend causal normalizing flows to the time-series setting.
- Design a learning and inference framework for multi-step interventional distribution estimation.
- Generalize across query types without requiring query-specific design.
Problem setup
We consider a multivariate time-series \(\mathcal{X}=(\mathbf{X}_1,\ldots,\mathbf{X}_T)\in\mathbb{R}^{T\times d}\) under two assumptions:
- Its data-generating process follows a Markovian SCM, with jointly independent exogenous variables.
- It is causally stationary, so its causal structure is invariant over time.
Window causal graph. We assume that the SCM-induced graph \(\mathcal{G}\) is given. For each variable \(X^j_t\), it distinguishes:
- \(\mathrm{Pa}^{<}_{\mathcal{G}}(X^j_t)\) — time-lagged parents
- \(\mathrm{Pa}^{=}_{\mathcal{G}}(X^j_t)\) — contemporaneous parents
Target estimand. Given an observed history \(\bar{\mathbf{x}}_m\) and a subsequent treatment sequence \(\underline{a}_{m+1}\), we estimate the interventional distribution of a future outcome:
\[P\!\left(Y_t(\underline{a}_{m+1}) \mid \bar{\mathbf{x}}_m\right), \qquad t \geq m+1\]This formulation encompasses conditional queries given pre-intervention information and marginal queries that integrate it out. The experiments below consider both marginal and history-adjusted queries.
Method
Model architecture
Causal History Encoder — encodes only the time-lagged parents \(\mathrm{Pa}^{<}_{\mathcal{G}}(X^j_t)\) into an embedding \(h^j_{\phi,t}\), using a feature-wise causal (parent) mask \(\mathbf{M}^j\):
\[h^j_{\phi,t} = E^j_{\phi}\!\left(\mathbf{M}^j \odot \mathbf{X}_{t-\tau:t-1}\right)\]Causal Conditional Normalizing Flows — model contemporaneous causal dependencies, with each \(X^j_t\) conditioned on its own embedding:
\[x^j_t = f^j_\theta\!\left(z^j_t;\ \mathrm{Pa}^{=}_{\mathcal{G}}(X^j_t),\ h^j_{\phi,t}\right), \qquad z^j_t \overset{\text{i.i.d.}}{\sim} P_z \quad \forall j\]Training — maximize the observational log-likelihood:
\[\max_{\theta,\,\phi}\ \sum_{t=1}^{T}\sum_{j=1}^{d} \log P_\theta\!\left(X^j_t \mid \mathrm{Pa}^{=}_{\mathcal{G}}(X^j_t),\ h^j_{\phi,t}\right)\]Estimation of interventional distribution
The target estimand admits the following form:
\[P\!\left(Y_t(\underline{a}_{m+1}) \mid \bar{\mathbf{x}}_m\right) = \int_{\mathbf{x}'_{m+1:t-1}} P\!\left(Y_t \mid do(a_t), \bar{\mathbf{x}}_{t-1}\right) \prod_{k=m+1}^{t-1} P\!\left(\mathbf{x}'_k \mid do(a_k), \bar{\mathbf{x}}_{k-1}\right)\]This is analytically intractable, so we approximate it by Monte Carlo simulation with TSCNF.
Theoretical guarantee
Does a TSCNF trained on observational data recover the exogenous variables needed to implement interventions? We show that matching the observational distribution isolates each exogenous variable up to a single invertible, component-wise transformation that is invariant across time and observed histories.
Theorem (Identifiability of TSCNF). Suppose we fit TSCNF \(f_\theta(\cdot \mid \mathbf{h}_{\phi,t})\) with fixed base distribution \(P_z\). If
- at \(t=1\) the conditioning history is empty, so TSCNF reduces to a static causal NF that induces the same initial distribution as the SCM of the initial time step, and
- for almost every history \(\bar{\mathbf{x}}_{t-1}\) sufficiently encoded by \(\mathbf{h}_{\phi,t}\), TSCNF induces the same conditional distribution as the history-conditioned contemporaneous SCM,
then there exists a component-wise invertible mapping
\[\mathbf{a}(\mathbf{z}_t)=\bigl(a_1(z^1_t),\ldots,a_d(z^d_t)\bigr),\]where \(\mathbf{a}\) is independent of time and history, such that
\[f_\theta^{-1}\bigl(f_{\bar{\mathbf{x}}}(\mathbf{z}_t)\mid \mathbf{h}_{\phi,t}\bigr)=\mathbf{a}(\mathbf{z}_t)\]for almost every history, for all \(t=2,\ldots,T\).
Why it matters. The mapping \(\mathbf{a}(\cdot)\) links each latent of TSCNF to a single exogenous variable, so conditional interventions can be implemented by modifying only the latents of the intervened variables. Since the mapping is invariant across time and history, this relation stays stable during multi-step rollouts — TSCNF can therefore estimate sequential interventional distributions and recover the target estimand at the model level.
Experiments
Setup. Fully synthetic time-series from linear and non-linear DGPs, plus a simulated PK–PD tumor-growth benchmark. Metrics: MMD and 1-D Wasserstein for distributional similarity at a given step, and RMSE for multi-step accuracy. Two query types throughout:
- Marginal — population-level effects under a hypothetical treatment sequence
- History-adjusted — conditioned on each subject’s covariate and treatment history
Ablation study
The full model consistently performs best, and the two mechanisms are complementary:
- Causal graph masking alone gives a notable improvement — filtering the history to causally relevant variables avoids spurious signals.
- Feature-wise conditioning on top of it improves further — each flow dimension conditions solely on its causal parents.
| Metric | Method | \(do(X^2_t)\) | \(do(X^4_t)\) | \(do(X^2_t, X^4_t)\) |
|---|---|---|---|---|
| Mean Dist. | w/o both | 1.98 (1.32) | 2.03 (1.10) | 1.61 (1.16) |
| w/o feature-wise cond. | 1.48 (1.05) | 1.30 (0.88) | 1.18 (0.86) | |
| TSCNF | 0.72 (0.57) | 0.66 (0.61) | 0.51 (0.47) | |
| MMD | w/o both | 0.53 (0.29) | 0.46 (0.23) | 0.48 (0.25) |
| w/o feature-wise cond. | 0.38 (0.25) | 0.21 (0.18) | 0.40 (0.23) | |
| TSCNF | 0.11 (0.14) | 0.08 (0.12) | 0.13 (0.21) | |
| Wass. | w/o both | 2.19 (1.26) | 2.39 (1.03) | 1.79 (1.10) |
| w/o feature-wise cond. | 1.57 (1.03) | 1.42 (0.87) | 1.26 (0.85) | |
| TSCNF | 0.74 (0.57) | 0.74 (0.63) | 0.53 (0.47) |
A. Synthetic datasets — non-linear DGPs
Marginal query. TSCNF and G-Net consistently outperform MSCVAE. MSCVAE relies on inverse propensity weighting, multiplying treatment probabilities across time — numerically unstable over long horizons. TSCNF and G-Net instead follow a g-computation-style approach that avoids explicit propensity modeling.
History-adjusted query. TSCNF is comparable or superior to G-Net across linear and non-linear settings. The gap widens under sparser time-lagged structures: G-Net incorporates all past covariates and treatments without filtering, which injects noise, whereas TSCNF selectively attends to the relevant time-lagged parents.
B. Tumor growth — PK–PD simulation
Point accuracy holds as the prediction horizon \(k\) grows. Normalized RMSE (%) at each horizon (\(\gamma_c = 5,\ \gamma_d = 5\); best in bold, second best underlined):
| Model | \(k=1\) | \(k=2\) | \(k=3\) | \(k=4\) | \(k=5\) |
|---|---|---|---|---|---|
| RMSN | 0.63 (0.11) | 0.74 (0.12) | 0.79 (0.11) | 0.83 (0.11) | 0.88 (0.11) |
| G-Net | 0.52 (0.03) | 0.64 (0.04) | 0.72 (0.06) | 0.79 (0.09) | 0.92 (0.24) |
| Causal CPC | 0.70 (0.10) | 0.77 (0.10) | 0.81 (0.10) | 0.84 (0.12) | 0.88 (0.14) |
| DoFlow | 1.01 (0.05) | 1.54 (0.06) | 1.90 (0.06) | 2.19 (0.07) | 2.40 (0.09) |
| TSCNF | 0.55 (0.09) | 0.58 (0.10) | 0.63 (0.10) | 0.67 (0.10) | 0.71 (0.11) |
Resources
- Paper link — to be added once UAI 2026 proceedings are published
- Code — release planned