DiffPO: A causal diffusion model for learning distributions of potential outcomes

TL;DR

This paper proposes a novel causal diffusion model called DiffPO, which is carefully designed for reliable inferences in medicine by learning the distribution of potential outcomes by leveraging a tailored conditional denoising diffusion model to learn complex distributions.

Abstract

Predicting potential outcomes of interventions from observational data is crucial for decision-making in medicine, but the task is challenging due to the fundamental problem of causal inference. Existing methods are largely limited to point estimates of potential outcomes with no uncertain quantification; thus, the full information about the distributions of potential outcomes is typically ignored. In this paper, we propose a novel causal diffusion model called DiffPO, which is carefully designed for reliable inferences in medicine by learning the distribution of potential outcomes. In our DiffPO, we leverage a tailored conditional denoising diffusion model to learn complex distributions, where we address the selection bias through a novel orthogonal diffusion loss. Another strength of our DiffPO method is that it is highly flexible (e.g., it can also be used to estimate different causal quantities such as CATE). Across a wide range of experiments, we show that our method achieves state-of-the-art performance.

Figures (3)

• Figure 1: Overview of our causal diffusion model DiffPO. Our method involves a forward and reverse diffusion process to learn the distributions of potential outcomes. Additionally, we address selection bias through our orthogonal diffusion loss.
• Figure 2: Empirical distributions of the conditional POs. Left: $p(Y(0) \mid x)$. Right: $p(Y(1) \mid x)$.
• Figure 3: We manually perturb the propensity score during training on the synthetic data. We replace the estimated propensity score $\hat{\pi}(x)$ with a randomly sampled value $\tilde{\pi}(x)$ from the interval $(0, 1)$. The weight $w_{\hat{\pi}}(x, a)$ for each sample in the orthogonal diffusion loss $\mathcal{L}(\theta, \hat{\pi})$ is thus replaced by weight $w_{\tilde{\pi}}(x, a)$. The CATE estimation error gradually converges as the sample size increases. This aligns with our expectation, as the loss remains robust even with varying errors in the estimation of nuisance functions.

Theorems & Definitions (6)

• Remark 1
• proof
• Theorem 1: Neyman-orthogonality
• proof
• Theorem 2: Neyman-orthogonality
• proof