Latent g-Computation for Potential Outcomes Distributional Estimation under Time-Varying Treatments

ICLR 2026 Conference Submission22562 Authors

Keywords: Treatment Effect Estimation, time-varying treatments, g-formula, g-computation, potential outcome distribution

TL;DR: We estimate full distributions of individualized potential outcomes under time-varying treatments via a latent space implementation of the g-formula.

Abstract: Estimating individualized potential outcomes (POs) under time-varying treatments is central to fields like medicine, marketing, or public policy. However, most methods yield only point estimates, offering limited guidance for risk-sensitive decisions. In this work, we propose G-Latent, a novel method to estimate full PO distributions via the g-formula. Unlike prior work such as G-Net (Li et al., 2021), we learn per-step outcome distributions directly through variational autoencoders (VAEs) rather than relying on global residual pools. Our core contribution is the introduction of a latent-space rollout in which each time-step embedding is updated from latent representations instead of observed samples. At inference time, this enables Monte Carlo (MC) sampling without data-space autoregression, reducing error accumulation and increasing sampling speed. To enhance expressivity, we adapt a VAE parameterization based on infinite mixtures of asymmetric Laplace distributions from An & Jeon (2023) to our time-series setting. We also decouple sequence processing: a transformer encodes the history up to a given time, while a lightweight GRU advances the forecast horizon, avoiding repeated transformer passes across steps and MC samples at inference. We validate our approach on synthetic, semi-synthetic, and real-world datasets, and provide theoretical guarantees.

Primary Area: causal reasoning

Submission Number: 22562