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

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.

Paper

References (67)

Scroll for more · 38 remaining

Similar papers

Peer review

Reviewer Cvxf6/10 · confidence 4/52024-07-11

Summary

This paper proposes DiffPO, a causal diffusion model, to predict individual potential outcomes. The authors motivate why predicting potential outcomes can be a more complex problem than (conditional) average treatment effect (CATE) prediction in practical settings. They adapt a conditional denoising diffusion model for this problem setting and augment its training with an orthogonal diffusion loss to correct for confounding. The proposed method is applicable to diverse settings and empirical evaluation demonstrates performance superior to comparable baselines across datasets. Further, the method is deployed to measure CATE, which is direct function of the predicted potential outcomes. The diffusion-model based instantiation of this method lends itself readily to uncertainty quantification, which indicates greater reliability of the method in high-stakes domains like medicine.

Strengths

Originality: The motivation of the paper is particularly interesting. It is intuitive that a method designed for good ATE/CATE prediction will not necessarily perform well in potential outcome prediction. The paper motivates the need to separately focus on the latter setting well. Quality: Experimental evaluation of the method is quite diverse and extensive, demonstrating superior performance on various datasets compared to baselines. Clarity: The motivation, method and overall paper are well-written and easy to follow. Significance: The paper highlights and addresses significant aspects of causal inference in medical applications, like individual effects, and uncertainty quantification.

Weaknesses

Clarity: - I found the preliminaries about diffusion models highly compressed (which is understandable due to space constraints). Since the proposed method relies heavily on diffusion modeling, the authors might consider using the extra space allowed in camera ready versions for expanding on the relevant details about training diffusion models. - Authors should also make sure to reference all tables and figures in the main text to better guide the reader to results relevant to a given section. Nit: line 309 does have what was probably intended as a reference but is broken.

Questions

- The three different kind of experiments (PO estimation, CATE estimation, PO distribution estimation) all use different datasets. Is there a reason all these quantities cannot be estimated for all datasets? - How sensitive is the proposed estimator to hyperparameter, choice of model architecture or other training details?

Rating

6

Confidence

4

Soundness

3

Presentation

3

Contribution

3

Limitations

The paper discusses limitations due to assumptions and efficiency of the method.

Reviewer PmLo5/10 · confidence 4/52024-07-11

Summary

This paper aims to develop a model which focuses on potential outcomes (POs) instead of CATE. It focuses on estimation of POs and proposed an estimation method which is an extension of the well know diffusion model with a different loss function.

Strengths

- This work proposes a diffusion model based estimation method for potential outcome estimation. The method is a simple extension of the standard diffusion model.

Weaknesses

- The motivation behind this is not clear as in most of the cases people are interested in the causal effect instead of potential outcomes, especially given one potential outcome is observed already. - The main novelty of this work is the orthogonal diffusion loss. However, it is simply using the idea of inverse propensity scores. - There is nothing specially designed for potential outcomes in the forward and reverse process, the variational inference and the conditional denoising function, which should not be clarified as background instead of methodology. ---- details - L350 In fact any variational inference based causal inference method can predict the distribution of potential outcomes. In addition, conformal prediction based causal inference methods can predict a predictive interval for each potential outcome. - Eq(11) It is unclear what the expectation is taken over. - L189 it is not a sentence. - L118 I don't think those papers are totally unrelated from this work. - Table 1. I believe for the purpose of this work, the comparison should be between different estimation models. The meta learners should be applicable along with the proposed model and other baselines, so it is not that crucial to list them in Table 1.

Questions

- Remark 1 may need more discussion. How to prove the loss is Neyman-orthogonal? What is the benefit to make the loss Neyman-orthogonal? Is there any experiment showing the advantage of this loss?

Rating

5

Confidence

4

Soundness

2

Presentation

2

Contribution

2

Limitations

See weakness

Reviewer KUb54/10 · confidence 4/52024-07-13

Summary

This paper applies the technology of diffusion models to estimation of potential outcome and treatment effect. Owing to the capability of modelling distribution in diffusion model, this paper can not only give point estimation but also uncertainty. Specifically, it proposes to use variational inference to learn the model and applies sample re-weighting to remove confounding bias. The estimation is proved to satisfiying Neyman-orthogonality property, and empirical validation reveal the effectiveness of the method.

Strengths

This paper exploits the power of diffusion models and achieve more capability than previous works on CATE estimation. The proposed method is technologically sound and achieve Neyman-orthogonality, which is an important theoretical property. The experimental results demonstrate the effectiveness of the method.

Weaknesses

Although the paper is technologically sound and well-written, the empirical examination is somewhat lacking. There have been a lot of dataset, such IHDP, Jobs, Twins in CATE estimation. I suggest the authors to complement the experimental results on these datasets to enhance the persuasive appeal of the experiments. There are additional questions about the experiments, see the details in Questions.

Questions

1. The results in Table 3 (CATE estimation) is confusing. The authors claim the metric is PEHE, which is favored to be lower. However, the proposed method (DiffPO) achieves largest value in Table 3. What does the value in Table 3 mean? 2. The authors examines the inferred potential outcome distribution in Section 6.3. However, how to validate the rationality of the inferred distribution? I think it is a very important problem. Otherwise, the inferred results are not trustworthy.

Rating

4

Confidence

4

Soundness

3

Presentation

3

Contribution

3

Limitations

The authors have clearly listed the limitation of the paper, and claim the potential solution to address these limitations.

Reviewer vD2m6/10 · confidence 3/52024-07-20

Summary

This paper describes DiffPO, a causal diffusion model for predicting distributions of potential outcomes, as well as related causal quantities such as CATE estimates. DiffPO accounts for the confounding between covariates and outcomes through a simple weighting procedure, under standard potential outcomes assumptions. The paper evaluates its performance in comparison to a variety of PO and CATE estimators over several standard benchmarks, including the ACIC 2016 and 2018 datasets and TCGA.

Strengths

- The application of a diffusion approach to potential outcome and CATE estimation is useful, especially given the flexibility to handle a variety of problem settings. - The paper's theoretical formulations of causal diffusion is novel, including the orthogonal diffusion loss and the implications for the DiffPO ELBO and simplified training objective. - The experiments show strong performance compared to baseline approaches for CATE estimation on several standard causal inference benchmark datasets.

Weaknesses

- The high level evaluation results (table 2 and table 3) are impressive. However, it would be good include additional details on benchmark performance, perhaps in the appendix. For example, these results tables show the "% of runs with best performance". It would be insightful to see the underlying performance metrics as well, to understand whether the approach was "just a little better" or a lot better. This would also help evaluate whether there are contexts where DiffPO is consistently stronger or weaker than alternatives. - The weighting approach in equation 10 will be unstable when samples have propensities approaching 0 or 1. While this is covered by standard causal assumptions related to population overlap, in practice it is common to have to clip propensity scores or drop high or low propensity samples. Was that unnecessary here, or what approach was chosen?

Questions

Does this approach lend itself to sensitivity analyses (e.g., of the assumption that there are no latent confounders?) Could you please define best_in and best_out in Tables 3 and 4? I believe there are usually annual ACIC challenges. Why were acic 2016 and acic 2018 chosen over other years benchmarks? The paper claims DiffPO can adapt to continuous treatments --- but the experiments seem to all be binary treatments? Did I miss something about the ACIC datasets or were continuous treatments not evaluated?

Rating

6

Confidence

3

Soundness

3

Presentation

4

Contribution

4

Limitations

Limitations were adequately discussed though, I think the claim (in the limitations paragraph, line 366) that the method scales to real-world datasets based on ACIC, IHDP, and TCGA is perhaps a little much to claim given the breadth of real-world scenarios and datasets, and the fact that these are datasets created or collated for academic purposes.

Reviewer Cvxf2024-08-10

Thank you for the clarifications and additional results! I will keep my score.

Reviewer KUb52024-08-13

Thanks for your response and effort in revise paper. The response address almost all my concerns. Still, I have another question. What does the empirical Wasserstein distance means in evaluating the learned distribution?

Authorsrebuttal2024-08-13

Response to Reviewer KUb5

Thank you for your reply! We are happy that we have addressed your concerns sufficiently! We also welcome the opportunity to answer your remaining question about the empirical Wasserstein distance. The Wasserstein distance is a commonly used metric in machine learning [1,2,3] for quantifying the distance between two probability distributions based on the concept of optimal transport. In our work, we use it to evaluate the model performance in learning the potential outcome distributions. Thus, this allows us to move _beyond evaluating point estimates_ to the _full distribution_. The $k$-Wasserstein distance (for any $k \geq 1$) between two distributions $\mu_1$ and $\mu_2$ is $W^k\left(\mu_1, \mu_2\right)=\left(\int_0^1\left|\mathbb{F}_1^{-1}(l)-\mathbb{F}_2^{-1}(l)\right|^k \mathrm{~d} l\right)^{1 / k}$ where $\mathbb{F}_1^{-1}(l)$ and $\mathbb{F}_2^{-1}(l)$ are the quantile functions (inverse cumulative distribution functions) of $\mu_1$ and $\mu_2$ for quantile $l$, respectively. In practice, Wasserstein distance is estimated by the empirical Wasserstein distance based on two sets of finite samples, i.e., $\hat{W}^k\left(\mu_1, \mu_2\right)$ $=\left(\frac{1}{n} \sum_{i=1}^n\left\|X_i-Y_i\right\|^k\right)^{1 / k}$, where $X_1, \ldots, X_n$ are the samples from $\mu_1$ and $Y_1, \ldots, Y_n$ are the samples from $\mu_2$, and we set $k=1$ in our experiments, following [1]. (3) We use empirical Wasserstein distance to evaluate our model performance on learning the potential outcome distributions, where _lower_ distances imply _better_ performance. Here, our method gives the **best** results and outperforms the baselines by a clear margin (see **Table 1 / Rebuttal PDF**). This demonstrates the strong performance of our model not only on point estimates, but also on estimating the _full distribution_. Thank you again for reading our rebuttal! We hope that the above results answer all your questions. We would greatly appreciate it if you would raise your score accordingly. Thank you. And do let us know if you have further questions -- we will do our best to answer them promptly. Thanks again for reviewing our submission! Reference: [1] Maeda T, Ukita N. Fast inference and update of probabilistic density estimation on trajectory prediction. International Conference on Computer Vision. 2023 [2] Melnychuk V, Frauen D, Feuerriegel S. Normalizing flows for interventional density estimation. International Conference on Machine Learning. 2023. [3] Letzelter V, Perera D, Rommel C, et al. Winner-takes-all learners are geometry-aware conditional density estimators. International Conference on Machine Learning. 2024.

Reviewer PmLo2024-08-13

Thanks for the very helpful response. I raised my score to 5.

Authorsrebuttal2024-08-13

Response to Reviewer PmLo

Thank you so much for your response and for raising the score! We will include all suggested points in the paper. Thank you again for your effort in reviewing this paper and your constructive comments!

Program Chairsdecision2024-09-25

Decision

Accept (poster)

© 2026 NYSGPT2525 LLC