DiffPO: A causal diffusion model for learning distributions of potential outcomes
Authors: Yuchen Ma
Valentyn Melnychuk
Jonas Schweisthal
Stefan Feuerriegel
Venue: NeurIPS 2024 poster (2024)
Keywords: Causal inference
Treatment effect estimation
Diffusion models
Machine Learning for healthcare
CATE
Primary area: causal_inference

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.

DiffPO: A causal diffusion model for learning distributions of potential outcomes
Yuchen Ma, Valentyn Melnychuk, Jonas Schweisthal Stefan Feuerriegel
LMU Munich
Munich Center for Machine Learning
yuchen. mww.mu.mmmmmmmmmmmmmmmmmmmmmmm.. de
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 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 novel causal diffusion model called DiffPO, which is carefully designed for reliable inferences in medicine by learning the distribution of potential outcomes. In our DifIPO, we leverage tailored conditional denoising diffusion model to learn complex distributions, where we address the selection bias through novel orthogonal diffusion loss. Another strength of our DiffPO method is that it is highly flexible (e.g., can also be used to estimate different causal quantities such CATE). Across wide range of experiments, show that our method achieves state-of-the-art performance.
Introduction
Predicting potential outcomes (POs) for patients from observational data is crucial for decision- making in medicine (16). For example, in cancer care, is interested in individualized predictions of survival under different treatment plans, which can then help medical practitioners choose a treatment plan that promises the largest chance of survival 48
Predicting POs of interventions from observational data challenging due to the fundamental problem of causal inference, i.e., the fact that one cannot obverse all POs for each individual (27). Further, in observational data, treatments are generally not assigned completely randomly, and, as a result, the distributions of covariates in the treatment and control groups differ (54). If this covariate shift is not taken into account, the distributions of POs given the covariates would be learned sub-optimally [13
11I
In the causal inference literature, many existing methods are aimed conditional average treatment effect (CATE) estimation where POs are used auxiliary quantities (e.g..(1 2
TO
12
13
25
28
30 32 37 45
64 66) ). In principle, some CATE methods could be used for predicting POs
these not tailored for POs and thus tend to underperform at predicting POs. The underlying reason that the estimator with the best performance in estimating CATE might not do best at predicting POs in finite samples inductive bias 13 31] i.e., they assume that POs should share similar structures. As result, they assume that the CATE follows much simpler function than the POs, that it is easier learn the CATE than each PO separately. Moreover, most of these methods only focus on point estimates instead of learning the distributions of POs. The latter is more difficult but very crucial for reliable macision-making in medical applications.
[ 13). Additionally, many state-of-the-arr CATE estimators leverage an
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

Table Overview of key methods for CATE estimation (and predicting POs). UQ refers to whether the methods allow for uncertainty quantification of POs. Bias addressing refers to whether the methods address the selection bias. Orthogonal refers to whether it is robustness (i.e., Neyman- orthogonality wrt. nuisance functions). POs are the target refers to whether the methods are originally designed for the task of predicting POs.
Estimator type
ug
Blas aldiessing
Orihogonal
POs the tuget
Key limitatinss
Seleeeer
37
step (plug model aguostic
Selection blas; estimatoc
Takarse
5
Quwstap (plug-in) modd--uu -aganniciiic
Sclezion
Di-leamd
Two-step. model agnostic
inductive point extmancr
RA-lemen
Twi-ttep, medelaggnnicc
8
CATE: inductive hase point extmoner
TARNet
One step (plug in), model specitic
X
Selection bias; point estimatoc Point extimatio
CFR R0]
Ourstipp (plug-in), medl--pecilie
X
GANTTE
64
Tan-step, specits
09
09
Selection biss; unstable training
TEDVAF
Outsstop (plug-in), mestl-opeciiic
Solection hias; ren-deemifshility
DIIPO (purs)
Duzstipp (plug-in), moddlsperiiie
Sumpling lime longer
(4) in the column UQ relers to that methods are not originally designed for uncemaiityy guaniiication of PUS, bur can be udapted for this purpoee.
In this paper, we propose a causal diffusion model called DifiPO for predicting POs. Specifically, we leverage a tailored conditional diffusion model to learn complex distributions of POs. We further propose novel orthogonal diffusion loss to adjust for the covariate shift problem in finite samples. Our orthogonal diffusion loss ensures properties. In particular, it makes our method more robust to misspecification. Further, our DiffPO is carefully designed for reliable inferences medicine and thus allows for uncertainty quantification We thereby follow recent calls in medicine to move beyond point estimates and offer distributional knowledge [23 35 which can inform how likely a certain outcome is and what a probable range of POs is.
which offers favorable theoretical
Our method is highly flexible. can learn distributions of POs for uncertainty quantification and give point estimates of POs. It can also handle different causal quantities, such as the CATE. This unlike many state-of-the thrt methods from CATE estimation which are limited to point estimates
12 13, 31 32 37 45 Hence, for different settings, custom methods have to be designed, while our causal diffusion model offers flexible and reliable approach.
Overall, our main contributions are the following: (1) We propose novel causal diffusion model for learning the distributions of POs, which can give both point estimation but also allows for uncertainty quantification. (2) We propose a novel orthogonal diffusion loss that ensures Neyman-orthooonnllty (3) We conduct a wide range of experiments in different settings demonstrate the flexibility and effectiveness of our DifiPO
2 Related Work
Our work aims to learn the distributions of POs for patients. The task related to CATE estimation in that the PO framework conceptualizes the CATE as the expected difference between the POs of a patient with and without treatment. Hence, we include also methods for CATE estimation in our literature review below. While these could, in principle, be used for predicting POs. we emphasize that CATE estimation and predicting POs will have a different umite-sample performance.
2.1
Predicting POs and CATE estimation
There have been many works designed for CATE estimation (e.g..
6,
10 12 13
25| 28
30 31 32
37, 45 64 (66) ). Existing CATE methods fall into two categories: model-agnose -agnostic and modell-seeiine eatoma We only give concise overview in this section (see Table D] A detailed literature review is in Appendix(A.3
Model-agnostic estimators. Model-agnostic learning strategies are also known as meta-learners 37, 45).. These can be split into (a) one-step (plug-in) learners that output regression functions from the observational data and then compute the CATE as the difference between the POs and (b) two-step learners that first estimate nuisance functions to build a pseudo- outcome and then obtain CATE directly by regressing the input covariates on the pseudo-outcomes in the second step. (Note that psecolosuttcomes are not potential outcomes).
(e.g.[[12 B
31 32
'Code is available at htattp::/githuu..cmm//

---

The problem with one-step plug-in in learners (e.g., S-learner [37], T-leamer (37) is that they often do not address selection bias. The problem with two-step learners (e.g. DR-learner 31 39.. RA-learner (12) is that these commonly leverage an inductive bias specific to CATE. Thus, two-step learners often have benefits for CATE estimation, but they make restrictive assumptions that hurt the finite- sample performance in predicting POs. Hence, in finite samples, such an inductive bias in CATE estimation can even lead to bias in PO prediction.
Model-specific estimators. Many model-specific estimators provide instantiations the general model- agnostic methods: Here, standard machine learning models are adapted to CATE estimation POs 61 64 66). Recently, neural networks have been used for model-specine learners [30 54| 64 66 Yet, model-specific estimators based on neural networks are designed for CATE estimation but are not directly aimed at POs.
prediction ...g......... 10
30 54
Uncertainty quantification. Most state-of-the-arr estimators fail offer uncertainty quantification of POs (see Table In particular, these methods typically only offer point estimates, yet without distributional information and thus do not allow for uncertainty quantification. Yet, distributional information is crucial for reliable decision-making in medicine and presents the focus of our method.
Implications for benchmarking, Most of the existing works are designed for CATE estimation, but not for learning the distributions of POs. For benchmarking PO prediction, many methods targeting the CATE do not include estimation of POs, and, thus, they are not applicable for our task. We show this in the column POs are the target in Table [ Instead, we can only compare with those methods that are applicable our task: that is, where the model can directly output POs predictions
37
54
64
2.2
Diffusion models for causal inference
Diffusion models were introduced for learning from complex distribution for given dataset from which one can then sample high-quality data points (c.g.. images) [26 55 577. Diffusion models have achieved state-of-the afrtthe-rrt performance, outperforming other generative models on various tasks
in the computer vision field (e.g.. [4 56 ). We give brief technical overview of diffusion models in Sec. 3
Diffusion models were previously used for different causal inference tasks but in different setting from ours, for example, generating the counterfactual of given image [36] 51), answering causal queries [8) or causal discovery [52]. We emphasize that the aforementioned tasks are different from
ours. For example, the task in [36|5]] is similar to minimizing the changes that one needs to make to an image in order for the classifier to categorize the image into different class. The task of answering causal queries (8] 33 53) builds upon structural causal models (SCMs), while we are using the PO framework [50]. The assumptions for fitting SCMs are usually much stronger than the latter. In causal discovery [52]. the output the causal graph, while our output is causal quantity. Hence, these works were all developed for different tasks (and not designed for learning the distributions of POs). Because of this, the works are not applicable as baselines.
3 Preliminaries on diffusion models
Diffusion models 26 55 57 are likelihood- based generative models that use (1) forward and
(2) reverse Markov processes. The (1) forward process 'disturbs" the data distribution g (2o) into tractable prior (0, I). For this, it gradually adds noise to an initial sample xp ~ q (20) across T steps with variance schedules {Bi,... ... }r). The forward process can be written g (Zie (r)) II q (Ite | II-1) with Gaussian distribution q (Zt (;-1) := N (Irv 3, BeI). It admits
closed form of q (3rg mg) that is also Gaussian distribution A (Vort (1 ā.) I), where ag
I (1 -.).
The (2) reverse process p (Zo.r) = Ih= (TE1P(Zt-1) xx) gradually denoises the latent variable FT I) and further allows for generating data samples from q (xo). The distribu- tions p (xt-1 1t) are usually unknown and approximated by neural network with parameters A. Thus, a parameterized Markov chain {no (J =r))
[= is trained in the reverse process. It can be parameterized as po (xt-1 xx)) := ' (Ix-15 150 (Me, (), En (z); (26) suggested using diagonal Ee (zz, t) with a constant d and computing μe (æt, t) as a function of 2, and ea (Ix, t). This yields
3

---

Me (zt, vor It -f-- (t), where ag := 1 3, := nlig ag and where eg (Te, t) predicts a 'ground truth' noise component for the noisy data sample
vT-ae
Later, we develop novel causal diffusion model for predicting POs. We further explain why the standard loss from above fails in our task because it does not account for the underlying causal structure. As remedy, develop an orthogonal diffusion loss that is tailored learn complex distributions of POs.
4 Problem Formulation
Setup: We consider observational dataset DD with i.i.d. patient data. The dataset consists of: an outcome of interest Y E y S R. dx-dimensional covariates (also called confounders) X X C [ and treatment A E {0, 13. For example, in critical care, the patient covariates X are different risk factors (e.g., age, gender, prior diseases), the treatment is whether a ventilator applied, and the outcome is the patient survival. For notation, let π<)) = p(A
X x) denote the propensity score, which gives the probability of a patient receiving the treatment. Let p(y a) p(Y y X x, a) be the probability density function of the conditional distribution P A).
We build upon the standard setting of Neyman- potential outcomes framework (0). Hence, let (a) denote the potential outcome after intervening on the treatment by setting it We have two potential outcomes for each individual: (11) if treatment is administered (i.e., = 1), and Y (0) if not treated (i.e., A = 0). However, due to the fundamental problem of causal inference (27). only one of the POs observed. Hence, Y (Y (1 (0).
Potential outcomes:
Identifiability: To ensure that POs are identifiable, we follow previous literature (e.g.. and make the following standard assumptions:
12 31
64
Assumption 1. (1) Consistency: If an individual is assigned trealment o, we observe the associated potential outcome Y = Y (a). (2) Unconfoundeddnes:: there are unobserved confounders, so that (()), Y (1) IL A X. (3) Overlap: treatment assignment is non-deterministic, i.e., 0 < ((x) <
L2 E X (f)())
Under Assumption I the distributions of POs can be identified from observational data via p(Y (a) X) = xYY X, a-)
Objective: Our main interest lies predicting POs in medical settings. Formally, (a) X] are the expected POs corresponding to a treatment assignment (intervention) for an individual with covariates X. Predicting POs is crucial for decision support in medicine [16]. For example, in critical care, doctors aim to predict the survival of each patient under different treatments (e.g., mechanical ventilation), which can then guide their decision-making.
However, decision-making in medicine needs to be reliable. Because of this, doctors are not only interested in the point estimate but need to quantify the uncertainty the POs (23]35). For example,
(35).
uncertainty quantification is needed in medical practice to decide whether to apply either a treatment that has small but certain benefit or a treatment with large benefit but potentially a large variability
in the outcomes and thus large uncertainty of whether the treatment is actually beneficial for specific patient. Hence, we estimate the distribution of the POs after assigning treatment a to an individual, ie., (a)(a) X). Learning this distribution allows us to sample from it and obtain predictive intervals for uncertainty quantification.
5 DiffPO for predicting potential outcomes
Overview: In the following, we introduce our DifiPO: a diffusion-based model for predicting POs with key components (Fig- I): ① forward diffussion process, 2 reverse diffusion process, and 3 novel orthogonal diffusion loss. Finally, we introduce our training and sample procedure.

---

5.1 Forward and reverse diffusion process
Our DifiPO builds upon diffusion models [26 55 57 We distinguish the distributions learned in the forward process and the reverse process via and p. respectively, to make the notation straightforward in this section.
① Forward process: Given a data point (y, x; a) sampled from the observational data distribution p(Y, X, A), in the forward diffusion process, gradually add Gaussian noise the initial y (denoted as yo) across 7 time steps. This thus produces a sequence of noisy samples V1,-.-. Yr in the same sample space yo- The forward process follows a Markov chain
T
(yLT 30) = II (ye Bt-1), (yle (t-1)) =- =- NV VI B 1,
(I)
==1
where 3r is variance schedule: {3, E (0, 1)}r= is small positive constant that represents a noise level and, thus, essentially controls the step sizes. By using the reparameterizatioo trick (34), sampling yt at arbitrary timestep has closed form
q (Bte yo) N (Mi Vāryo, (1 a;)
where at == Be and āi = II'= ag.. Thus, Y can be expressed as
Ye (yo. c) Vālyo + 11-aa ajC,
where c E Rd ~ N.O,,) I). As result, the data sample yo gradually 'looses' its distinguishable features for later steps Eventually, when T + oc, yr is equivalent to an isotropic Gaussian distribution.
2 Reverse process: We aim to learn the true distributions of POs, which can be identified as the conditional distribution Sec. 4 Using the notation above, it can be rewritten as q (yo xr a) 9 (30.T] I, a) dy As g (3t-1 Mr, I, a) is intractable, need to learn model distribution pe parameterized by to approximate the data distribution.
The reverse process starts at a known prior distribution with density p (vr) = N (yr; 0, I] and then proceeds backward to obtain the data distribution step process is also Markov chain, and the density of its joint distribution (yo.r 2, a) can be written as
Formally, the reverse diffusion
T Pe (yorr a) == p (8z) II pe (ye-1 Ve. x. a). yr N(0.1.
We employ conditional diffusion model with the reverse process in Eq. (4) The conditional density po (yt-1 Ve, x. a) is also Gaussian, and we parameterize it as
pa (yr-1 Vr, x. a) N (21-1; 1: (ye, t x, a) o),
(5)
The reverse process gradually denoises yr ~ No I) through the learned Gaussian transitions. Once we have computed the latter, approximate the original data distribution and can even sample from it.
Variational inference: Directly computing and maximizing the likelihood (0 a) is difficult because it involves intractable integration. As remedy, we frame the learning objective for the above diffusion processes through variational inference. Thus, the parameters @ can then be optimized by maximizing the evidence lower bound (ELBO) via
log Pe (yo:r x, a)
log (yo I. a) 2 W
vier (Virr13))))
(VLT Mo)
(6)
The right side of Eq. 6 can be rewritten as (see Appendix for the derivation)
E Pgg0 mo)
[log po (yo V1, 2., a)]
DKL (g (ur yo) llp (vr)) prior matching torm
reconstruction Lerm
(7)
(X [KLL (q (3o-1 ye, vo) pe (3r-1 II, a))]
denoising matching term
5

---

-200
PIX
68
DARLNI
DALALE
24
DATASK
Feeecinn
DA
Figure I: Overview of our causal diffusion model DifiPO. Our method involves a forward and reverse diffusion process to learn the distributions potential outcomes. Additionally, we address selection bias through our orthogonal diffusion loss.
In the denoising matching term, we try to minimize the KL-divergence between distributions, one is tractable, ground truth denoising transition step (yt-1 1/+, Mo), and other is denoising transition step po ( yt, x. a).
Conditional denoising function. In practice, directly optimizing the learning objective in Eq. 7 inefficient. Recently, (26) has shown way turn the optimization of the ELBO into a simpler problem for unconditional diffusion models. We thus follow a similar way to obtain a simpler problem for learning objective, but for our conditional diffusion models.
We employ a trainable conditional denoising function fe : (Y x R A) >. The conditional denoising function is function approximation, which we to predict the 'ground truth' noise component e for the noisy data sample ye conditioned on x and a. We consider the following parameterization that computes 1so (ye, 2; a) as a function of y/, x, and fo. This gives
Po (yi, 22
Vai yt
1
Bt.
fo (Vr, t 2, 12,0)
(8)
VI a
The simplified training objective fort Unif(L..... T} then
E (,po |||e [|l--fo (Vāryo vI āze, | v,a))l
We provide the full derivation of Eq. and Eq. (9) in Appendix
5.2 Orthogonal diffusion loss for addressing selection bias
Why we need to address selection bias through our orthogonal diffusion loss: observational data, treatments
randomized but are administered according to some (unknown) behavioral policy. This can result in a selection bias, especially in medical practice. For example, patients with a more severe health state will also receive more aggressive treatment. As result, when treatment A is selected based on the covariates X, the propensity score π (X) is not constant, implying that the distributions of covariates in treatment and control groups differ. Formally,
have distribution shift in the covariates, i.e,, p(X AA 0) * p(X A= 1). not adjusted for, this distribution shift may result in selection bias and thus inflated variance in the estimates of POs (and thus the causal effects), especially in low-sample settings [1 29 54]. To address this, we thus introduce a novel orthogonal diffusion in the following.
3 Orthogonal diffuuion loss. Our proposed orthogonal diffusion loss is inspired by orthogonal learning theory [18, 60). Orthogonal learning theory provides general toolbox for inference with semi-parametric models to provide quasi-oracle rates for statistical learning with a nuisance component. Informally, orthogonal losses first-order insensitive misspecification of the nuisance functions, which introduces many favorable properties such as double robustness 63 Below, we construct orthogonal diffusion loss for method.
start by noting that π (xx) is not constant and is unknown in observational datasets. Because of that, one cannot simply use the propensity scores to reweight the data. Instead, we need trainable approach to reweight the data and thus adjust for the distribution shift from above. To this end, the propensity score T Tzx) is nuisance function in method, which estimate through trainable

---

function Jo parameterized by Φ. Let us denote the estimated propensity score by f(x) = (z). We
then assign weights W (2, a) to each sample via some function wa (X {0. 1}) + R, More
specifically, we define the function w (X x {0. 1)) + R2+ as
1-a a
we (2, a)
+ 1-#((z)
(10)
*(x)
Then, our orthogonal diffission loss is given by
(x ||c - fo (Vāryo + v1 - d,e,t 2, 11 (11)
C(0, +) =E -) A1.
The loss in Eq. (9) fits the conditional distributions p(Y | Xx A) based on the data samples from the joint observational distribution, i.e., (Jo: xr a) ~ p(,, X, A). Due to the selection bias, this implies that p(Y X, A 1) is learned better for the treated population and p(Y A 0) for the untreated population [60]. Yet, our target to learn the potential outcome distributions (aY | X) p(Y X, A a) for both 0 and
1 equally well in all the populations. This would be equivalent to minimizing the target loss from Eq. (9 the samples are from the joint distributions of the POs, i.e.. (yo, x) (a), X). However, we do not have samples from the distributions of POs, thus we need the target loss in Eq. (TT) to be estimated via the propensity score (x) to learn the distributions of POs. Hence, the following remark holds.
The orthogonal diffusion loss in Eq. 11 evaluated with the ground truth nuisance functions matches the following target loss:
Remark 1.
C(0, x)
Σ aE{0,1}
E ((noa))-
[||e fo (Vānyo + vI a =.a) (12)
Proof. The orthogonal diffusion loss in Eq. (I) is an inverse propensity-weighte estimator of the target loss in Eq. (12).
Theorem 1 (Neyman-orthogonnlityy. The orthogonal diffusion loss in Eq. wrt. its nuisance functions.
is Neyman-orthogondl
Proof. See Appendix D
As result of Neyman-orthogonallin our orthogonal diffusion loss has clear practical advantage: offers robustness against errors in nuisance function estimation. Specifically, it is first-order insensitive wrt. errors in the propensity score (()) and, thus, is more accurate objective for our PO predictive model.
5.3 Training and sampling
Training. Our training algorithm proceeds along the following steps. We first train the function Jo on the dataset to estimate the propensity score. After this, the parameters of are frozen. We then compute each sample weight for the orthogonal diffusion loss through the weight function Wg- We finally sample from the observational data distribution and run the forward and reverse process as described in Sec 5.1 The corresponding training loss is given in Eq. (II).
Sampling. Once our diffusion model has learned the distribution po (Y--1 1 Me', J. a), we can generate samples from it. The sampling process starts by sampling yr from the prior distribution and then denoising it. To obtain y/-1 from pa (yt-1 Ye, x, a). this requires the mean j from the previous step, which can be computed by Eq. (8). Thus, V/-1 can be computed via ÿ o22, where (O, I). This process continues for T.... until we arrive at yoo
Implementation. Our model architecture of denoising function is based on [59] and uses U-Net 49 as the underlying backbone. We use 4 residual blocks where each is built with MLP layers. The diffusion embedding dimension is 128. The number of diffusion sampling steps is 100. Training is conducted with batch size of 256 and a learning rate of 0.0005. We follow (42) the way how we learn propensity scores and thus use fully connected neural networks with softmax activation. During training, we can only observe one of the two POs due to the fundamental problem of causal inference
27 (as explained in Sec. To guide the model, need to identify which PO should the loss be
7

---

computed on. Hence, we introduce causal masks as input to our model: an observational mask Mo, targeted mask m, and a conditional mask me We only compute the loss the place where the value of the targeted mask is 1. Further details about our implementation are in Appendix E.T
6 Experiments
6.1 Learning distributions of POs
Performance metrics. We use the Wasserstein distance to evaluate the model performance of learning the distributions of POs. The k-Wasserstein distance (for any ks 2 1) between two distributions 11 and D2 is
W (u1, = (fo - dl)
1/4
(13)
where (() and (1) are the quantile functions (inverse cumulative distribution functions) of 1/1 and 1/2 for quantile L, respectively. Specifically, we compute the empirical Wasserstein distance based A E ||x - Yllll 1/k where Xi...., Xx
on two sets of finite samples, i.e., (*(11, are the samples from 111 and Yi.... Yn are the samples from 1/2 In our experiments, we report the empirical Wasserstein distance for 1. Lower values of W are preferred.
Dataset. Due to the fundamental problem of causal inference, the counterfactual outcomes are never observed in real-world data. We thus follow prior literature (e.g.. 22, 38) and bench- mark our model using synthetic datasets. Fur-
Table 2: Results showing in- & out-of-sample em- pirical Wasserstein distance (i.e., W and W oug) for two potential outcomes (i.e., a = and = 1) on the synthetic dataset. Reported: mean + standard deviation over ten-fold train-test splits.
ther details about the synthetic datasets are in the Appendix(F F
11
W:1
Baselines. Many methods are targeting CATE estimation and, therefore, are not suitable for PO prediction. We thus compare against those methods that are applicable to our task (i.e., the model can directly output POs; see Table I]: (1) S-learner [37]: is a model-agnostic learner that fits single regression model by concatenat- ing the covariate and the treatment as input; (2) T-learner 37 is model-agnostic learner that fits two separate regression models, one for treated and for controls; (3) TARNet [54]: is based on representation learning to share information about the outcome across treated and controls with regularization; (4) CFR [54:: is based on representation learning used in variants of balaneing with TARNet; (5) GANITE (64] : uses a generative adversarial network to generate POs and then uses another generative adversarial network to generate CATE. We adapt GANITE to our task by removing the second stage, so that we directly sample POs from the learned distributions in the first stage. However, methods (I) to (4) above are only able to give point estimates POs. Therefore, we follow (24) to equip them with Monte Carlo (MC) dropout to make these methods comparable. Implementation details are in Appendix E.2
S-learner"
1.007 am
1.005
1754
1,781
1,055 em
Tleamme
27
1.004
113
1,002
2: k3m
L,054 2: c0
1,453
412
TARNer
54
0:334
CFR
54
0.921
1 kn1
0.943
1: 1.4
0.909 1 6112
1,078 1: 400
0.759 532
1.524
DiIPO (ours)
0.043
0.125
0,032
0.091
better (hest modified method comparable.
Results. The results are in Table 2 We find that our method gives the lowest empirical Wasserstein distance out of all methods, which desirable. Hence, the experiments show that our method outperforms the baselines by a clear margin.
6.2 Learning predictive intervals of POs
Performance metrics. To evaluate the ability of our DiffPO to allow for uncertainty estimation, we follow (24) and examine the predictive intervals (PIs) of the POs generated by different methods. In medical practice, the treatment effectiveness often reported based on the 95% and 99% Pls. We thus compute the individual equal-tailed (1 a)) Pls with [0.01, 0.05). We then calculate the faithfulness of the estimated Pls to evaluate the empirical coverage by computing the frequency with which the Pls contain the outcomes in the test data.

---

We evaluate the empirical coverage across different quantiles (1 1--o) a). The results are in Table It shows that our DiffPO is gen- erally faithful, while this not the case for the baselines. This can be expected: MC dropout relies upon mixtures of Dirac distributions in parameter space, which leads to approximations of the true posterior which are questionable and
Results:
Table 3: Results for uncertainty estimation of the two potential outcomes (i.e., a 0 and = 1). Reported: mean + standard deviation over ten-fold train-test splits.
95% Pl
99%
95% PL
99% Pl
S-leamer'
37
0.801 23m
0.891 amm
0.879 enes
0.898 2*0
THearner
0.832 230e
0.854
0.843
ate
TARNe
0.855 g3re
0.864
2000
0.845 gn1x
0.878 are
not faithful 17 24
0.861 2208
0.954 2008
4008
GANITE
64
0.892 1116
0,902
10.3
0.873
cous
0.885 1216
DiIPO (Ours)
0.981 2am
0,985
0.908 gore
0.926 2mm
6.3 Point estimates of POs
Higher better (best in bold).
modified for comparahility.
Performance metries for POs: To evaluate point estimates of POs, we compute the dif- ference between the predicted PO ÿi and the ground truth PO y, via the root mean squared N error RMSE = N 1= (0+- =))²
Table 4: Results for point estimation of POs bench- marked using the in- & out-of-sample RMSE for the two potential outcomes (i.e., a = 0 and on the synthetic dataset. Reported: mean + stan- dard deviation over ten-fold train-test splits.
1)
Σ
We compare against those methods that are applicable to this task (i.e.. the model can directly output POs; see Table I
Baselines.
RMSE
RMSH az
RMSK.
RMS
S-learner T-leamer
37
0,401 a
0.391 cmm
0.379 gm
0.398 gr
0.432
0.454 -06
0.443 gm
0.421 sea:
TARNe
0,364
0:388
0.319 g01e
Results. The results are in Table 4 We find that our DiffPO gives the best point estimates of POs.
CFR
54
0.263
0.278
0.261
0.293 18
0.378
0.390
-015
can
14.3
IMPPO (Ours)
0.143 2014
0.156 2011
0.162 2014
0.187 211
Lower better (best in bold).
6.4 Flexibility to handle other causal quantities
A strength of our method is that it is flexible and can not only be used for POs but also for other causal quantities. For example, even though we are aiming at learning the distributions of POs, our method is also capable of estimating the CATE (xx) = (1) (()) X = z/.. which is the expected difference of POs for an individual with covariate values X = x. For this, one can use
our method to first predict pa (x)) E X, A a] for E (0, 1} and then simply leverage that 7 ((r) = Mi (x) - (z)
Datasets. We estimate the CATE across ACIC
Table 5: Results for benchmarking CATE estima- tion on ACIC 2016 and ACIC 2018 for both in- & out-of-sample, respectively. Reported: % of runs with the best performance.
2016 & ACIC 2018, which are widely used
dataset collections for CATE benchmarking [66
12 42.. The ACIC2016 [46) contains 77 dif-
ferent benchmark datasets, and ACIC2018 (41) contains 24. These datasets include a wide range of data-generating mechanisms (see Appendix F for more details). We use five random train/test splits (80% / 20%) for each dataset, tune hyper- parameters on the first split, and evaluate the average out-sample on every split.
ACIC 2016 (77 datasets) bestes % best,
ACIC 2018 (24 datasets) best,
bestuu
S-learner
37
3.76
3.41
1.16
1.02
T-leamer
37
3,71
3.8
0.74
1.13
TARNet
8.92
9.21
CFR
54
10.52
11.15
7.91
7.71
GANITE 64
6.78
6.44
5.22
3.98
MWk-leanerri Tv
14.54
23.64
22.13
RA-learner
14.68
15.73
15.01
TEDVAE
[0
9.33
9.52
7.28
5.45
Performance metrics. We compare the
DinPO
26.73
25.82
31.34
36.86
<PEHE V N E.' (f (
Higher hetter (hest in hold).
Baselines. We consider a broad array of -he---rrt methods for treatment effect estimation from the literature. Considering the rich methods in the literature, we thus compare with the most
popular CATE estimation approaches in this field [T0] 30 37 [5
61
64
66). as listed in the Table
Specifically, we use the following baselines: (1) to from Sec. 6.1 (6) DR-learner [39
generates speudo-outcomer based on the doubly-robust AIPW estimator; (7) RA-
[1] uses
a regression-adjusted pseudo-outcome in the second stage; (8) TEDVAE (66]: uses a variational
autoencoder to differentiate confounding factor and learns CATE. We instantiate the meta-learnrrs
(i.e., S-, T-, DR, and RA-learner) with neural networks, similar to our DiffPO. To ensure a fair comparison across all methods and all experiments, we tune hyperparameters across all methods

---

separately for each experimental dataset. We discuss implementation details of baselines in the Appendix
Results. The results for CATE estimation are in Table We find that our method achieves a good performance similar to existing methods, though our method is not tailored for CATE estimation. Across a wide range of experiments in different settings, we even observe that our method often achieves Mate-of-the-at performance.
6.5 Visualizing the learned distributions of POs
clear advantage of our DifiPO is that we not only return point estimates but also the distribution of POs. This is crucial in medical practice [23 35 in order understand the expected probability of whether a treatment is beneficial and thus to better assess the reliability of the estimates. We show the learned distributions for real-world dataset from medicine.
IHDP dataset. This is a semi-synthetic dataset from the Infant Health and Development Program (IHDP) [25), which is based the extracted features and treatment assignments from real-world clinical trial. The dataset comprises 747 patients, with 25 features for each patient. Further details for IHDP are in Appendix F3
Insights. Our DifiPO is capable of capturing the full information about the distributions of POs. Current state-of-the-art methods usually estimate quantities expressed via the mean of POs, as these methods focus purely on point es- However, distributional knowledge of POs is important to account for un- certainty, because it informs how likely certain outcome is and gives the probability that the PO lies in desired range, especially in medicine.
303
timation 31 37 54).
Figure 2: Empirical distributions of the conditional POs. Left: p(Y (0) 2)). Right: p(Y(1) æ).
Fig. 2 shows the empirical distributions of POs. i... (0) X) and p X), given a certain patient profile X x. We can see the distributions of the two POs are different, which can provide extra information for medical decision-maie makingg. Here, the learned distributions of POs can help medical practitioners in choosing treatment plan that promises not only a large benefit for the patients but also that the benefit is highly probable.
Discussion
Conclusion. Our method is carefully designed for predicting POs in medical practice. To this end, our method not only predicts point estimates but also computes the distributions of POs to allow for uncertainty quantification and thus for reliable decision-making
Limitations. (1) As with other methods causal inference, ours rests on mathematical assumptions, yet these are standard in the literature [12 31 64). (2) The efficiency of the sampling process in our DifiPO but also diffusion models more generally could be improved further. There is already
ongoing research, such as different solvers [15 40 65 and one-step sampling (58). However, have shown above, our method can successfully seale to real-world datasets from medical practice.
Broader impact. We expect our method to have significant impact in medicine where better decision support needed to personalize treatment decisions patient profiles (16). Another strength of our method is that it is flexible. This is unlike many other methods in causal inference which are designed for highly restrictive settings. Hence, we expect our method be first step toward developing more generalizable approaches for a variety of causal inference settings.
10

---

Acknowledgements
This paper is supported by the DAAD program "Konrad Zuse Schools of Excellence in Artificial Intelligence", sponsored by the Federal Ministry of Education and Research. This work has been supported by the German Federal Ministry of Education and Research (Grant: 01IS24082).
References
[1] Ahmed M. Alaa and Mihaela van der Schaar. "Bayesian inference of individualized treatment effects using multi-task Gaussian processes". In: Advances in Neural Information Processing Systems. 2017.
[2] Ahmed M. Alaa and Mihaela van der Schaar. "Bayesian nonparametric causal inference: Information rates and learning algorithms". In: IEEE Journal of Selected Topics in Signal Processing 12 (2018), pp. 1031-1046.
[3] Ahmed M. Alaa and Mihaela van der Schaar. "Limits of estimating heterogeneous treatment effects: Guidelines for practical algorithm design". In: International Conference Machine
Learning, 2018.
[4] Serge Assaad et al. "Counterfactual representation learning with balancing weights". In: International Conference on Artificial Intelligence and Statistics. 2021.
[5] Onur Atan, William R. Zame, and Mihaela van der Schaar. "Counterfactual policy optimization using Jomain--dversarre neural networks". In: ICML CausalML workshop. 2018.
[6] Ioana Bica, James Jordon, and Mihacla van der Schaar. "Estimating the effects of continuous- valued interventions using generative adversarial networks". In: Advances in Neural Informa- tion Processing Systems. 2020.
[7] Ioana Bica "Estimating counterfactual treatment outcomes over time through adversarially balanced representations". In: International Conference on Learning Representations. 2020.
[8] Patrick Chao, Patrick Blöbaum, and Shiva Prasad Kasiviswanathan. "Interventional and coun- terfactual inference with diffusion models". In: arXiv preprint. 2023.
[9] Victor Chemozhukov et al. "Double/éeiiasedNeeyma machine learning of treatment effects".
In: American Economic Review 107.5 .5 (2017), pp. 261-265.
[10] Hugh A. Chipman, Edward I. George, and Robert E. McCulloch. "BART: Bayesian additive
regression trees". In: The Annals of Applied Statistics 4.1 (2010).
[11] Alicia Curth and Mihacla van der Schaar. "In search of insights, not magic bullets: Towards demystification of the model selection dilemma in heterogeneous treatment effect estimation".
In: International Conference on Machine Learning. 2023.
[12] Alicia Curth and Mihaela van der Schaar. "Nonparametric estimation of heterogeneous treat- ment effects: From theory to learning algorithms". In: International Conference Artificial Intelligence and Statistics. 2021.
[13] Alicia Curth and Mihaela van der Schaar. "On inductive biases for heterogeneous treatment effect estimation". In: Advances in Neural Information Processing Systems (2021).
[14] Prafulla Dhariwal and Alexander Nichol. "Diffusion models beat gans on image synthesis". In:
Advances in Neural Information Processing Systems. 2021.
[15] Tim Dockhom, Arash Vahdat, and Karsten Kreis. "Genie: Higher-order denoising diffusion solvers". In: Advances in Neural Information Processing Systems (2022).
[16] Stefan Feuerriegel et al. "Causal machine learning for predicting treatment outcomes". In:
Nature Medicine 30.4 (2024), pp. 958-968.
[17] Loic Le Folgoc et al. "Is MC dropout bayesian?" In: arXiv preprint. 2021.
[18] Dylan Foster and Vasilis Syrgkanis. "Orthogonal statistical learning". In: The Annals of Statistics 51.3 (2023), pp. 879-908.
[19] Anna Gilbert et al. "Towards understanding the invertibility of convolutional neural net- works". In: International Joint Conference on Artificial Intelligence. 2017.
[20] Negar Hassanpour and Russell Greiner. "Counterfactual regression with importance sampling
weights". In: International Joint Conference on Artificial Intelligence. 2019.
[21] Negar Hassanpour and Russell Greiner. "Learning disentangled representations for counterfac-
taal regression". In: International Conference on Learning Representations. 2019.

---

[22] Tobias Hatt and Stefan Feuerriegel. "Estimating average treatment effects via orthogonal regularization". In: Proceedings of the 30th ACM International Conference on Information & Knowledge Management. 2021.
[23] James J. Heckman, Jeffrey Smith, and Nancy Clements. "Making the most out of programme evaluations and social experiments: Accounting for heterogeneity in programme impacts". In: The Review of Economic Studies 64.4 (1997), pp. 487-535.
[24]
Konstantin Hess et al. "Bayesian neural controlled differential equations for treatment effect estimation". In: International Conference on Learning Representations. 2023.
[25] Jennifer L. Hill. "Bayesian nonparametric modeling for causal inference". In: Journal of Computational and Graphical Statistics 20.1 (201 1), pp. 217-240.
[26] Jonathan Ho. Ajay Jain, and Pieter Abbeel. "Denoising Diffusion Probabilistic Models". In:
Advances in Neural Information Processing Systems. 2020.
[27] Paul W Holland. "Statistics and causal inference". In: Journal of the American statistical Association (1986).
[28] Fredrik D. Johansson, Uri Shalit, and David Sontag. "Learning representations for counterfac- tual inference". . In: International Conference on Machine Learning. 2016.
[29] Fredrik D. Johansson et al. "Generalization bounds and representation learning for estimation of potential outcomes and causal effects". In: Journal of Machine Learning Research 23 (2022), pp. 7489-7538.
[30] Fredrik D. Johansson et al. "Learning weighted representations for generalization across designs". In: arXiv preprint (2018).
[31] Edward H. Kennedy. "Towards optimal doubly robust estimation of heterogeneous causal
effects". In: Electronic Journal of Statistics 17.2 (2023), pp. 3008-3049.
[32] Edward H. Kennedy et al. "Minimax rates for heterogeneous causal effect estimation". In:
arXiv preprint (2022).
[33] Ilyes Khemakhem et al. "Causal autoregressive flows". In: International Conference Artificial Intelligence and Statistics. 2021.
[34] Diederik P Kingma and Max Welling. "Auto-encoding variational Bayes". In: arXiv preprint.
2013.
[35] Thomas Kneib, Alexander Silbersdorff, and Benjamin Säfken. "Rage against the mean-a review of distributional regression approaches". In: Econometrics and Statistics 26 (2023),
pp. 99-123,
[36] Aneesh Komanduri et al. "Causal Diffusion Autoencoders: Toward counterfactual generation via diffusion probabilistic models". In: arXiv preprint. 2024.
[37] Sören R. Künnel et al. "Metalearners for estimating heterogeneous treatment effects using machine learning". In: Proceedings of the National Academy of Sciences 116.10 (2019),
pp.
[38] Milan Kuzmanovic et al. "Causal machine learning for cost-effective allocation of development aid". In: Proceedings of the 30th ACM SIGKDD Conference on Knowledge Discovery and Data Mining. 2024.
[39] Mark J. van der Laan. "Statistical inference for variable importance". In: The International Journal of Biostatisties 2.1 (2006).
[40] Cheng Lu et al. "DPM-Solver: A fast ODE solver for diffusion probabilistic model sampling in around 10 steps". In: Advances in Neural Information Processing Systems. 2022.
[41] M.F. MacDorman and J.O. Atkinson. "Infant mortality statistics from the linked birth/infant death data set-1995 period data". In: Monthly Vital Statistics Report 46.6 Suppl (1998), pp. 1-22.
[42] Divyat Mahajan et al. "Empirical Analysis of Model Selection for Heterogeneous Causal Effect Estimation".. In: arXiv preprint (2022).
[43] Valentyn Melnychuk, Dennis Frauen, and Stefan Feuerriegel. "Bounds representation- induced confounding bias for treatment effect estimation". In: International Conference on Learning Representations. 2024.
[44] Pawel Morzywolek, Johan Decruyenaere. and Stijn Vansteelandt. "On general class of orthogonal learners for the estimation of heterogeneous treatment effects". In: arXiv preprint.
2023.
12

---

[45] Xinkun Nie and Stefan Wager. "Quasi-oracle estimation of heterogeneous treatment effects". In: Biometrika 108 (2021), pp. 299-319.
[46]
Kenneth R. Niswander. "The collaborative perinatal study of the National Institute of Neuro- logical Diseases and Stroke". In: The Woman and Their Pregnancies (1972).
[47] Kazusato Oko, Shunta Akiyama, and Taiji Suzuki. "Diffusion models are minimax optimal
distribution estimators". In: International Conference on Machine Learning. 2023.
[48] Lucia C Petito et al. "Estimates of overall survival in patients with cancer receiving different treatment regimens: emulating hypothetical target trials in the Surveillance, Epidemiology, and End Results (SEER)-Medicar Linked Database". In: JAMA Network Open (2020).
[49] Olaf Ronneberger, Philipp Fischer, and Thomas Brox. "U-net: Convolutional networks for biomedical image segmentation". In: Medical image computing and computer-assisted inter-
vention. 2015.
[50] Donald Rubin. "Causal inference using potential outcomes: Design, modeling, decisions". In: 100.469 (2005), pp. 322-331.
[51] Pedro Sanchez and Sotirios A Tsaftaris. "Diffusion causal models for counterfactual estima- tion". In: Causal Learning and Reasoning. 2022.
[52] Pedro Sanchez et al. "Diffusion models for causal discovery via topological ordering". In: International Conference on Learning Representations. 2023.
[53] Pablo Sanchez-Martin, Miriam Rateike, and Isabel Valera. "Vaca: Design of variational graph autoencoders for interventional and counterfactual queries". In: arXiv preprint. 2021.
[54] Uri Shalit, Fredrik D. Johansson, and David Sontag. "Estimating individual treatment effect:
Generalization bounds and algorithms". In: International Conference on Machine Learning. 2017.
[55] Jascha Sohl-Dickstein et al. "Deep unsupervised learning using nonequilibrium thermodynam- ics". In: International Conference on Machine Learning. 2015.
[56] Jiaming Song, Chenlin Meng, and Stefano Ermon. "Denoising diffusion implicit models". In: arXiv preprint (2020).
[57] Yang Song and Stefano Ermon. "Generative modeling by estimating gradients of the data distribution". In: Advances in Neural Information Processing Systems. 2019.
[58] Yang Song et al. "Consistency models". In: arXiv preprint (2023).
[59] Yusuke Tashiro et al. "CSDI: Conditional score-based diffusion models for probabilistic time
series imputation". In: Advances in Neural Information Processing Systems. 2021.
[60] Stijn Vansteelandt and Pawel Morzywotek. "Orthogonal prediction of counterfactual out- comes". In: arXiv preprint. 2023.
[61] Stefan Wager and Susan Athey. "Estimation and inference of heterogeneous treatment effects
using random forests". In: Journal of the American Statistical Association 113.523 (2018), pp. 1228-12 1242.
[62] Anpeng Wu et al. "Learning decomposed representations for treatment effect estimation". In: IEEE Transactions on Knowledge and Data Engineering 35.5 (2022), pp. 4989-5001.
[63] Andrew Ying. "A geometric perspective on double robustness by semiparametric theory and information geometry". In: arXiv preprint (2024).
[64] Jinsung Yoon, James Jordon, and Mihaela van der Schaar. "GANITE: Estimation of individu- alized treatment effects using generative adversarial nets". In: International Conference Learning Representations. 2018
[65] Qinsheng Zhang and Yongxin Chen. "Fast sampling of diffusson models with exponential integrator". In: arXiv preprint (2022).
[66] Weijia Zhang, Lin Liu, and Jiuyong Li. "Treatment effect estimation with disentangled latent factors". In: AAAI. 2021.
[67] Yao Zhang, Alexis Bellot, and Mihaela der Schaar. "Learning overlapping representa- tions for the estimation of individualized treatment effects". In: International Conference on Artificial Intelligence and Statistics. 2020.
13

---

A Extended related work
A.1 Diffusion models for causal inference
Diffusion models were originally introduced for learning a complex distribution for a given dataset from which we can then sample high-quality images [55] 57, 26). Diffusion models have achieved state-of-the-art performance, outperforming other generative models on various tasks in the computer vision field (e.g..(56) [6) [4 59). Motivated by this, diffusion models were previously used for different
causal inference tasks but in a different setting from ours [5T] 52 8).
[5]] proposed model called Diff-SCM that uses diffusion models and anti-causal predictors to
generate the counterfactual of given image. It leverages a classifier guidance diffusion model. (8)
proposed using a diffusion model called DCM for modeling structural causal models (SCMs). It
focuses on approximating SCMs given observational data and underlying causal DAG. ( focused
on similar task as well as [53 36). with the aim answer causal queries by employing variational
autoencoder and normalizing flow, respectively. [36] focused on the same task as [S]] aiming for
generating high-quality counterfactual images and proposed model called CausalDiffAE. (36) follows the idea of [8] to view the diffusion model as an encoder-decodee framework and employ so-called denoising diffusion implicit models (DDIM). Diffusion models have also been used in the causal discovery field. DiffAN (52) is algorithm based on denoising diffusion training for causal discovery via topological ordering. In this work, we use diffusion models for predicting POs.
A.2
Neymantorth
Neyman -oroon-ottoggonality of a functional refers to the mean zero property of its directional derivatives along one- me-dimensional paths that change nuisance functions [18]. loss function is called Neyman- orthogonal when this property holds for all its directional derivatives along mie-dimeniioal paths that change the (infnitt-dimmenioonal) parameter of interest. This allows for a debiased machine learning approach to construct estimators and inference procedures that are robust to small mistakes in nuisance functions [9. Specific orthogonal learners for estimating CATE are given the DR-learner
39
31. R-learner 455 and i-learner (60].
A.3
Conditional average treatment effect (CATE)
Our work is about predicting potential outcomes (POs) for individuals; however, it strongly related in practical applications to conditional average treatment effect (CATE) estimation. The PO framework conceptualizes the CATE estimation problem as estimating the expected difference between individual's expected potential outcome with and without treatment. Below, we give brief overview of the CATE literature.
The estimation of the CATE has received lot of attention in the causal inference literature (c.g 10
25
28
[
54
37
2 30
64
12
66 31 45 32 13). Among those CATE estimation methods.
many have been proposed to estimate the effects of binary treatments (e.g.. [T] 25 28
54
37 2
30
3 64
6
12
66
31
45 32 ).). Based on the categorization in [13), we roughly categorize them into
three categories based on their most salient characteristics: (i) model-agnostic learning strategies for
CATE estimation (also known as meta-learners), and (ii) model-specific MI-bbaed CATE estimators (including representation learning-based and other neural network-based estimators).
A.3.1 Meta-learners
Now we consider popular approaches that involve using model-agnostic learning strategics, also known meta-learners, which can be implemented using any ML method (e.g.. (37 31
45 32
B
12). The idea behind "meta-leareer" strategies was originally introduced 37 and expanded in
45 31 12].
Existing CATE meta-learners can be categorized into: (a) one-step (plug-in) learners (indirect meta- leamner) that output two regression functions from the observational data and then compute CATE as the difference in the POs (this the strategy underlying the and T-learners); and (b) two-step learners (direct meta-learners/multi-stee functions to build a pseudo-outcome, the second step, they obtain the CATE directly by regressing the input covariates the pseudo-outcome. (Note that pseudo-outcomes not potential outcomes).
direct estimators). These learners first compute nuisance
14

---

In terms of (b), existing methods fall largely into three broad classes: regression adjustment (RA), propensity weighting (PW), doubly robust (DR) strategies.
A.3.2 Model-specific ML-based CATE estimators
Many methods proposed in the related work do not fall within the meta-learner class because they rely on the properties of a specific MI. method. We roughly categorize them into two branches: one neural network-basee (NN-based) estimators and the other does not use neural networks, for example, some using causal Forest [6T], Bayesian regression [25]. or Gaussian processes (1].
As for the NN-based estimators, some are representation learning-basse based estimators. Much work of this track focused handling selection bias by learning shared and balanced feature representations for the two PO functions or incorporating weighting strategies. Examples are BNN (28), TARNet
(54). and CFR 30 [44 used representation learning share information about the outcome across
treated and controls, while regularizing representation distributional distance between the groups. [67]
30) recognized that such regularization may result violation of ignorability in the representation and proposed to learn representations in which context information is preserved but where treatment groups overlap. Later, many works proposed further extensions (e.g.. 20
62 67 5 7
However, a wrongly chosen dimensionality of the representation or too large balancing weight can induce confounding bias [43), and, therefore, the representation learning-based methods fall outside
the scope of this paper.
15

---

B Evidence lower bound of DifPP
Directly computing and maximizing the likelihood pyy z, a) is difficult because involves inte- grating all latent variables, which is intractable for complex models. Thus, we derive the evidence lower bound (ELBO) similar the bound the variational autoencoder. Our original objective can be optimized by maximizing the ELBO, like its analog in the ELBO of a vanilla VAE (34. which is given by
log p(y 2, a) = log
(yo:T
2, a) dy1:T
p (30.T
xr a) q (V1:T Mo0) dy/.T
log
q (V1.T v10
(14)
p (youT x,
log E lo))
q (Vi.r y/0
log
p (VorT x a)
Byluerlbol lvo)
9 (Vi.r y/)
The ELBO in Eq. (14) can be rewritten as follows:
16

---

p (yo.r x.a)
log p(y) 2 2go
log
q (31rr 3/0)
(yr) [Tp= (3h-1
Mc. a)
log
II (Ble V1-1)
log (H/T) ) pe (yo 91, a)
q (31 Mo0 Il (=2 (ye- yt. 2, a)
-egg mo)
(3/e 3/--1)
log (UT) po (yo U1, 26, a) IT (yt-1 20, ())
-Reg mo)
(1/1 9o) IT 9 (Vi 31-1,
log p (UT) po (yo 21, log II q (111 Uo)
(Vr y/s Ir, a)
-Ry(surr vo)
q (y/ Vt- Mo)
t=
- wo)
log p (UT) (wo 21,
log
II
Po (1+
3/ 22, (2)
q (y1 yo)
(
M 20 (80))
t=2
alve-
(Vo)
log
(yr) po (yo y1, a)
log
9 (vn
yo)
log II
po (3e-1
y/. 2, 3/. 2/
a (y1 3/0)
① (yT
yo)
@ (We-1
r=2
log p (yr) po (yo 81, x. a)
T
Po (1/-
yes x, a)
log
(ST 300
(3t-11
3/t, wo)
EEter mo)
[log pe (yo V1 a)]
By(aer (lo) 1oo8 ((UT)
T
1ox po (31--1
M/e x. a)
E Bal
(vr vo)
(we-1 1/: 2o)
2=
EEgoerl 120) [log pa (yo 31s Ir, a)] +
Belyriun) (uo) [1og ((TT
]
Σ 1=
Jog po (Ve-1
ye. x. a) ye. yu)
E Eglon un- 1l100)))
q (yr yo)
9 (31-1
= Bylar vo) [log Po (yo y/, (n, ] reconstruction term
DRL. (a (UT yo) ll (ur))
priar matching
Boluoloo loo) [DL (g (Bt-1 Ve, vo) llpe (Me-1 ye. x, a))]
denaising matching term
(15)
17

---

C Simplified training objective for DifIPO
With the reverse process defined above, the learning objective in Eq. 7 is clearly differentiable with respect to and is ready to be employed for training. However, in practice, directly optimizing this objective is inefficient. Recall that aim is match the approximate denoising transition step Po (ye-1 Be, क a) to ground truth denoising transition step (Ve-1 yt. yo) as closely as possible, which can also model Gaussian. Thus, optimizing the KI. divergence term reduces to minimizing the difference between the means of the two distributions:
arg argmin min DKL. (g (Vt-1 ye. 30) po (Ve-1 yr, (,,,,a)) a))
=arg min PKL (N (M6-1i (1)( || N (M1-15/10; (1))
1
(t)
d + tr (Zn (1) (t) + (No ma) Rt [t) -1 (ug mai
-arg min
log Ex
2
1
log 1 + d d (μL m₂) (t) -2 (wo mnp)
aarg min
2
(16)
-arg min
1
[w1 -m1) x (1) -1 (mO - mn)]
2
arr min
1
- (Ho--py)
(o)(t) -1 (mo - [aa]
2
pargmin (t) [1100-1011 m=]
6
By using the reparameterization trick, we have yo y/
6,50
(17)
Plugging it into the previously derived true denoising transition mean My (o), we yield Py (Ye, yo) = Vas (1 a1-1) 1) V1 vā- (1 ax) 3/0
1 Oz
Vai (1 -a--1)
1) V/ + vā (1 - ar) T-
1- āi
Vou (1 (4-1) 1) Bt
(1 11-00 (4))
Vi-via
vas
Ct
Vai (1
a1-1) V/
(1 - az) 1/1
ac) VI Cgen
1 āt
(1 att) Var
(1 5x)) /
Vai
11
x1-1)
ar
(1 aa) v1 ā.
(18)
+
y/1
eo
(1 - σt) Vot
(1 xit) vatt
01 (1 8
1 0t
0t. 60
+
1/1
ai) / 0+
(1 σx)
VI arver
a āi + at
1 a₁
y1
60
(1 an) Vot
VI
σi Ba
āi
at:
y11
6u
(I āt)
Vag
v1 0iv/0
1
1 ai ED. y/
Va:
v1 - āi az 0z
Therefore, can set our approximate denoising transition mean μe (Ve, t x,, a)
1
1 aq
He (ye,
x, 0)
V1 āiya fo (ye,t xr, a)
(19)
Var
18

---

and the corresponding optimization problem becomes arg min DEL (q (Vt-1 1 yt, yo) |I po (Ve-1 ve x2 a)) aarg min DEL (N (Vt-1: Mg: (t) I (y 10; (t))
aarg min 1
y
at v1 dayan
fo (ye, t x. vol y/
1 at
(t)
at
v1 cis att
2
1
1 1-
(20)
=arg min (t)
fo (Mert x, a)
v σn
v1 0sv
aarg min
1
aq
(e- fo (ye, xz
(1)
vI āiv diva
(1 az) 2
1
[|c- fo (vc, 22,
arg min
(t) (1 az)a cz
19

---

D Proofs
Theorem 2 (Neyman-orthogonaltty)) The orthogonal diffusion loss
[(0, +)
( 1.w. wa (2, ||e So (Vorso + v1 āse, (x, (21)
Neyman-orthogonue wrt. muisance functions.
Proof. As it was shown in the Sec. 5.1 the minimization of the orthogonal diffusion loss is equivalent to the maximization following weighted ELBO wrt. to 0:
log Po (yo-r 2,
Eps, #) () wx a) Pyg
(22)
(VI.T yo)
We denote a maximizer the weighted ELBO with the ground truth nuisance functions by Po- = argmax,, E(pa, π). Given flexible enough diffusion model, Pe- (3o 2x, a) coincides with the ground truth posterior and, thus matches the ground truth conditional distribution, namely, p(Y X, A) 47
To demonstrate Neyman--orthooonality we need to show that pathwise cross-derivative is equal to zero 18 60
ic..
D Dp pg- Eppo- po-, T] = 0 for every pe and ft
(23)
We start by taking the pathwise derivative wrt. the optimal diffusion model pa-
pgp E(po-, Tpp p
(24)
d
Po- (yo.r x, a) + tt(po (Vo.r xr, a) - Po- (yo.r xr, a)
E, dl (x)0)))))))))xxx1)
W+ (I, a) Ey
log
q (VI.T 3o)
(25)
pa (yorT x, a) Pa. (yo.r x.a)
-E -B(x,...))=p(Xxxx.)
Wr (x, a) (y
Po- (yo:r
2r,, ()) + (po (yo.r x., (1) Po- (yor x, a)
(26)
po (yorr
xr a) Po (yorr x, 2.a)
=B (.)
W (x, a) Ey (uor
(27)
Po- (yorT xr, a)
20

---

Then, we take a derivative wit. the propensity score л.
pa-, ft x| D2Dro- E(po-, T)
(28)
d E (000 Wx Wx+e(+-n) dt
Po (Vo.T
Ir, a) - po- (yo.r
2. a)
(z, Pogrlll
Po- (yorr 2r a)
(=0
(29)
a -B 20)0220))-((xx.. (((T(2))2 + (1 11-T(c2 -xx)2) E Eyrroo(Girrro
po (yoer x. () - po- (UO:T [ a)
(yoer x, a)
(30)
p(A=1 xx)
po (yoer xr, 1) po- (yo:r x2, 1)
=Exwp(X)
By Byo B (Yrrll))
(T(x)))
Po- (yorr I; 1)
(31)
p(A = x))
Eg (P Bu (o po (Uorr Ir, 0) po- (yo:r 2;, 0)
+
(1 - π()))²
Pg- (Vo:T xz, 0)
(32)
1 x(x)
xx 1) / po (yorT 1) po- (yorT
-xxxp(X)
p(yo I 1) q(ver yo) dyo dynr
pe- (yorr x, 1)
(33)
+ (xx) / po (yo.r x. 0) - Po- (yo.r pø- (yOrT xr, 0)
22, 0)
p(yo xz q(31:r))
yo) dyo dy
(34)
E [x) (xx) / po (yOTT 1) - Dg- (MOeT
2x, 1) ((31eT yo) dye dy
(35)
Po- (V1.r you x 1)
+ - x(x) ʃ Po (woer x., 0) Po- (yorr xr, 0) q(vner yo) dyo dyner
po- (y1.r yu,
(36)
(40)
1 / (po (vo.T xx, 1) - Pe- (yo.r I. 1)) dyo dy1.r
Ex)
(37)
(xx)
+ 1- (xx) / (po (VO.T xx, 0) (yorr | xz 0)) dyo dy
(38)
1
-Exxp(X)
(x:) (1 1) + 1- T() (1 1) = 0,
(39)
where the equality (+) holds due p(yo 1z, a) po (yo x. a); and the equality (**) holds due to that the forward process independent of and a, namely, pe- (y1rr yo, x, a) pe- (3/17 yo)) (31er yo).
21

---

E Implementation details
In the following, we summarize the implementation details of our DifIPO and baselines.
E.1 Implementation details of DiffPO
We implemented our DifiPO in PyTorch. (Code is available at htttps///githut com/ com//ccm// DiffPO).. Experiments were carried on 1x NVIDIA A 4100-PCIE-40GBB We report the default settings of our model below but note that hyperparameters may require slight adjustments depending on the dataset.
Our model architecture of the denoising function fe is based on the architecture of [59]. and we use U-Net [49] as the basic backbone. We use 4 diffusion residual blocks, where each is built with MLP layers. The diffusion embedding dimension 128. The B starts at 0.0001 and ends at 0.5, and the schedule is the "quadratic" version. The number of diffusion sampling steps is set 100. The training batch size is set to 256 with a learning rate of 0.0005. The training epoch is set to 500.
During training. we can only observe one of the two POs due to the fundamental problem of causal inference [27) (as explained in Sec. To guide the model, we need to identify which of two POs should be generated. Hence, we introduce causal masks as input mo. a targeted mask m,, and a conditional mask mc For the observational mask ma, it is at the place where have the observational data and for the opposite case. Since condition æ and a, the conditional mask m. is for the element æ and a, and 0 otherwise. For the target mask, the observed outcomes y in the original dataset are 1, while all the other elements are 0. In this way, can compute the loss only at the place where the value of the targeted mask is 1, i.e. where we have the ground truth outcomes. We follow (42) in the way how learn propensity scores and use fully connected neural networks with softmax activation. Thus we add weight to each sample via a learned function when computing the orthogonal loss.
model: observational mask
E.2 Implementation details of baselines
We follow the implementation from for most of the CATE estimators, including S-learnr [37), T-learner (37). DR- 31). RA- learner [12), TARNet (54). For GANITE, we follow the implementation //ttps://gitiiiie com/ vanderschaarlss aaciiaarllbbbaaa /
the implementation of wwt.a//waiai.aaaaaappaeeeeaaaaaa ganite For TEDVAE, we follow
We performed hyperparameters tuning of the nuisance functions models for all the baselines based on five-fold rros-vvalidation using the training subset. For each baseline, we performed a grid search with respect different tuning criteria, evaluated on the validation subsets. We aimed for fair comparison and thus kept the number parameters, network structures, and grid size similar across models. For evaluating the uncertainty estimation, for both training and testing, the dropout probabilities were set to p = 0.1.
22

---

Dataset details
F.1
Synthetic dataset
We use one of the datasets from ACIC2018 with 177 covariates to generate the synthetic dataset with coefficient matrix W, i.e..
X ~ Real-World (-).
Heal-World (X)
(40)
Y == Uy + sin W X + AX + + Uy ~ N(..
We use a ten-fold split for train/test samples (80%
F.2 ACIC 2016 & 2018 datasets
The 2016 Atlantic Causal Inference Challenge (ACIC2016) (46) contains 77 different settings of benchmark datasets, and ACIC2018 (41) contains 24, respectively. They are designed to benchmark causal inference algorithms with various data-generating mechanisms. Covariates of ACIC 2016 are taken from a large study of developmental disorders, and covariates of ACIC 2018 are derived from the linked birth and infant death data. ACIC 2016 and ACIC 2018 differ in the number of true confounders, the varying levels of overlap, and the form of conditional outcome distributions. ACIC 2016 has 77 different tata-genera mechanisms with 100 equal- samples for each mechanism (n 4802, dx 82). ACIC 2018 provides 63 distinct data-generating mechanisms with around 40 non-equal-sized samples each mechanism (n ranges from 1, 000 to 50, 000, dx = 177). Notably, ACIC 2018 has constant CATE for most of the datasets, but heterogeneous propensity scores.
F.3
IHDP dataset
The Infant Health and Development Program (IHDP) 25 has features and treatment assignments from real-world clinical trial. The dataset comprised 747 subjects, with 25 features for each subject. Out of the 25 features, both treated and untreated, synthetic outcomes of IHDP are sampled from different conditional normal distributions. These distributions are homoscedastic (a² = 1) but have substantially different conditional means. The potential outcomes are simulated according to the standard non-linear "Response Surface B" setting in (25) with the following tata-generating mechanism:
continuous, and 19 are binary with binary treatment. For
X ~ Rerl-World(.).
A Real- (orldWo)ld
(41)
Y (X w) A) (exp((X W)))),11)
where W. constant parameters of the simulation. For further details, we refer to
25
23

---

G Additional experiments
G.1 Simulation experiments
We examine the Neyman -othooooothhgonalitty property of our orthogonal diffusion loss through simulation experiments. We conduct experiments show that the orthogonal diffusion loss can address the problem of model misspecifications when estimating nuisance functions (c.g., the propensity score (2)).
To this end, we perturb the propensity score manually to assess the robustness of the loss when the nuisance functions estimated with varying errors. During training, we consistently increase the sample size and evaluate the CATE estimation error on the synthetic dataset. As shown in Fig. 3 the CATE estimation error decreases as the sample size grows and eventually converges zero. This demonstrates that, even when the propensity score is misspecified, the orthogonal loss remains robust, enabling our DifiPO converge the correct objective. This experiment supports Theorem by illustrating how the stability and consistency of the orthogonal loss contribute to convergence. In sum, shows the theoretical benefits of our proposed orthogonal diffusion loss.
Eprme
0.3
0.2
200
400
600
800
1000
Number of samples
Figure 3: We manually perturb the propensity score during training on the synthetic data. We replace the estimated propensity score (x) with a randomly sampled value #(x) from the interval (0, 1). (a( (z. a) for each sample in the orthogonal diffusson loss C(0, f) is thus replaced by
The weight wa
weight wa (x, a). The CATE estimation error gradually converges as the sample size increases. This aligns with our expectation, the loss remains robust even with varying errors in the estimation of nuisance functions.
24

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification:
The main claims made in the abstract and introduction accurately reflect the paper's contributions and scope.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
is fine to include aspirational goals as motivation as long as is clear that these goals not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification:
We discuss the limitations of the work in Sec 7
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged judgment and recognize that individual actions in favor of transparency play impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
the paper. The authors should use their best
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
25

---

Justification:
We provide the full set of assumptions in Sec. and proof in the Appendix.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Question:
Answer: [Yes]
Justification: We fully disclose the information needed to reproduce the main experimen- tal results in the Appendix.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how reproduce that algorithm. (b) If the contribution is primarily
model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the particular way they provide for reproducibility, In the case of closed-source models, it may be that access to the model is limited in some way (e.g., registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
26

---

Answer: [Yes]
Justification: We provide open access to the data and code and provide an anonymous link to our code: inttss//annonyoout wwwpan Apppn..cciinnc///IIFFF
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips.cr ec/
for more details.
While encourage the release of code and data, understand that this might not be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https for more details.
//nips.cc/ppp cc/public/guu httpggidddeee/eebbbiiinnneee
The authors should provide instructions data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification:
We specify all the training and test details in the Sec. Sec. aand in the Appendix E
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment Statistical Significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [Yes]
Justification: we report standard deviation in our experiment results.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the enror bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
27

---

It should be clear whether the error bar is the standard deviation or the standard error of the mean.
is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out range (e.g. negative еπхо rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: Our experiments do not require any specific hardware/resourcees. We admit the runtime of our method is longer compared baselines.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary failed experiments that didn't make into the paper).
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics
Answer: [Yes]
Justification: Our research conducted in the paper conform, in every respect, with the NeurlPS Code of Ethics.
Guidelines:
The answer NA means that the authors have not reviewed the NeurlPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (c.g., there a special consid- eration due to laws or regulations in their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification:
We discuss the broader impacts of our method in Sec 7
Guidelines:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
28

---

The conference expects that many papers will be foundational research and not tied particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms monitor how system learns from feedback over time, improving the efficiency and accessibility of MI..
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: Our work does not pose such risks.
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Lieenses for existing assets
Are the creators or original owners of assets (e.g., code, data, models), used in the paper. properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification: All the datasets used in our paper are cited accordingly.
Guidelines:
The answer NA means that the paper does not existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (c.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (c.g., website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thcode. com com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
29

---

If this information is not available online, the authors are encouraged to reach out to the asset's 's creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification:
We do not provide any assets in our work.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/moodl as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification: Our work does not involve crowdsourcing or research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According the NeurIPS Code of Ethies, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: Our work does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
30


Reviews:
[{"review_id": "uWfVVI4VjV", "reviewer": "Reviewer_vD2m", "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.", "questions": "Does this approach lend itself to sensitivity analyses (e.g., of the assumption that there are no latent confounders?)\n\nCould you please define best_in and best_out in Tables 3 and 4?\n\nI believe there are usually annual ACIC challenges.  Why were acic 2016 and acic 2018 chosen over other years benchmarks?\n\nThe 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?", "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.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 4, "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.\n\n- 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.\n\n- 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.\n\n- 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?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your positive review and your helpful comments! We appreciate that you find our paper novel, useful, and with strong experiment performance. We are very happy to answer your questions and improve our paper as a result.\n\n\n## Response to “Weaknesses”\n\n**(W1) Additional performance merics for ACIC2016 & ACIC2018**\n\nThank you for your suggestions! We now report the **average performance for PEHE metrics** on ACIC2016 (77 datasets) and ACIC 2018 (24 datasets) (see **Table 2 / Rebuttal PDF**). Here, we compute the average of PEHE across all the datasets (lower values are favored). Overall, **our proposed method is better by a clear margin**. For example, on ACIC2016, the best baseline has an out-of-sample average PEHE of 0.65, while our method achieves 0.49. On ACIC2018, the best baseline has an out-of-sample average PEHE of 0.28, while our method achieves 0.14. This corresponds to improvements of 24.6% and 50.0%, respectively.\n\nWe also manually inspected where our method was superior to offer qualitative insights. Generally, we found that our method was very strong and overall best (see our \"% of runs with best performance\" in the original paper). Nevertheless, we observed a tendency that, for datasets with really simple data-generating mechanisms (for example, all linear), our method may not be best, and other, simpler models are better. \n\n**Action**: We will add the above results to our revised paper. We will further report a breakdown by dataset to allow for further qualitative insights.  \n\n**(W2) Weighting / clipping**\n\nThank you for asking this important question. Yes, you are absolutely right. For the datasets with poor overlap, extreme propensity scores (near 0 or 1) can lead to high variance. As written in your comment, a common approach is the clipping technique based on the propensity score to avoid numerical instability during training. We also use clipping in our implementation and set a minimum ($\\epsilon$) and maximum  ($1- \\epsilon$) threshold with $\\epsilon = 0.05$. \n\n**Action:** We will add details about how we follow best-practice and use clipping to the implementation details of our revised paper.\n\n## Answer to “Questions”\n\n1. Thank you for asking this interesting question. That could be a very interesting extension of our current work, as, generally, sensitivity models require fitting the conditional distribution models as nuisance functions [5]. Here, a diffusion model like ours can serve as an estimator of the conditional distribution. While this is an interesting extension for future research, our current paper focuses on the setting, where unconfoundedness holds. **Action:** We will state in our revised paper that extensions of causal diffusion models to sensitivity analysis are an interesting opportunity for future research.  \n\n\n2. The \"In-sample\" and \"out-of-sample\" refer to the portion of the dataset that is used to train and evaluate the model, respectively. The “best_in” and “best_out” mean the “% of runs with the best performance” on the training and test dataset, respectively. **Action:** We will clarify the definition in the experiment section of our revised paper. \n\n\n3. Yes, you are right that there are usually annual ACIC challenges. However, _these were mostly designed for other tasks/purposes rather than CATE estimation. Only ACIC2016 and ACIC 2018 are designed for benchmarking CATE, which is the reason why the two datasets are commonly used by previous works for CATE benchmarking [1,2,3,4].\n\n\n4. Thank you. Our method can be adapted to continuous treatment settings but, upon closer inspection, the performance gains are not as clear as for the binary treatment. **Action:** We decided to remove that claim from our paper and, instead, focus on the true strength of our method: DiffPO is carefully designed for learning the _conditional distribution_ of potential outcomes. \n\n## Response to “Limitations”:\nThank you for your honest feedback! **Action:**  We will revise our statements in the camera-ready version of our paper to be more nuanced and toned down. \n\n\n\nReference:\n\n[1] Mahajan, D., Mitliagkas, I., Neal, B. and Syrgkanis, V., 2022. Empirical Analysis of Model Selection for Heterogeneous Causal Effect Estimation. ICLR.\n\n[2] Zhang, J., Jennings, J., Zhang, C. and Ma, C., 2023. Towards causal foundation model: on duality between causal inference and attention. ICML.\n\n[3] Zhang W, Liu L, Li J, 2021. Treatment effect estimation with disentangled latent factors. AAAI. \n\n[4] Cousineau, M., Verter, V., Murphy, S. A., & Pineau, J. 2023. Estimating causal effects with optimization-based methods: A review and empirical comparison. European Journal of Operational Research.\n\n[5] Frauen, Dennis, et al, 2024. A neural framework for generalized causal sensitivity analysis. ICLR.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "ZbfloY6Ex3", "reviewer": "Reviewer_KUb5", "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.", "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? \n2. 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.", "limitations": "The authors have clearly listed the limitation of the paper, and claim the potential solution to address these limitations.", "rating": 4, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your review and your helpful comments! We appreciate that you find our proposed method important, sound, and effective. We further would like to emphasize that we perform CATE experiments to show that our method is flexible. Nevertheless, the primary objective of our method is to _learn the distribution of potential outcomes_. We are happy to answer your questions and improve our paper as a result.\n\n## Response to Weaknesses:\n\nWe thank you for your suggestion to add more datasets. **We additionally conducted experiments for the three datasets that you suggested: IHDP, Jobs, and Twins** (see **Table 3 / Rebuttal PDF**). We find that our method is comparable to existing CATE methods and achieves state-of-the-art performance.\n\n**Action**: We will add the above experiments to our revised paper.\n\n## Answer to “Questions”:\n\n**(Q1) Question about whether lower/higher PEHE is preferred**\n\nWe thank the reviewer for the important question. Yes, you are right that a **lower** value is favored for the PEHE evaluation metric (Table 1 in our original submission). In contrast, Table 3 in our original submission does not report PEHE but  “% of runs with the best performance”, but where a **higher** value is preferred. The reason for the choice of the evaluation metric is that ACIC2016 and ACIC2018 are ‘suites’ with multiple benchmark datasets; that is, they contain 77 different and 24 different benchmark datasets, respectively.  \n\nUpon reading your question, we realized that we should evaluate ACIC2016 and ACIC2018 not only in terms of “% of runs with the best performance” (where higher is better) but also in terms of average PEHE (where lower is better). Hence, **we repeated our experiments and now report the average PEHE** for these experiments (see **Table 2 / Rebuttal PDF**). Here, lower values in terms of average PEHE are favored (consistent with Table 1 from our original submission). Overall, **our proposed method is better by a clear margin**. For example, on ACIC2016, the best baseline has an out-of-sample average PEHE of 0.65, while our method achieves 0.49. On ACIC2018, the best baseline has an out-of-sample average PEHE of 0.28, while our method achieves 0.14. This corresponds to improvements of 24.6% and 50.0%, respectively.\n\nWe also manually inspected where our method was superior to offer qualitative insights. Generally, we found that our method was very strong and overall best (see our \"% of runs with best performance\" in the original paper). Nevertheless, we observed a tendency that, for datasets with really simple data-generating mechanisms (for example, all linear), our method may not be best, and other, simpler models are better. \n\n**Action:** We will add the above results to our revised paper.\n\n**(Q2) New experiments to evaluate the learned distribution**\n\nThank you for asking this important question. **We performed new experiments where we explicitly assessed the ability of our method to learn the distribution of potential outcomes.**  We thus compare the distributions using the empirical Wasserstein distance. We find that we outperform the baselines by a clear margin (see **Table 1 / Rebuttal PDF**). \n\n\n**Action:** We will add the above experiments to our revised paper.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "snXACfMP22", "reviewer": "Reviewer_PmLo", "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.", "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?", "limitations": "See weakness", "rating": 5, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 2, "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.\n- The main novelty of this work is the orthogonal diffusion loss. However, it is simply using the idea of inverse propensity scores.\n- 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.\n\n----\ndetails\n- 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.\n- Eq(11) It is unclear what the expectation is taken over.\n- L189 it is not a sentence.\n- L118 I don't think those papers are totally unrelated from this work.\n- 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your review and your helpful comments! \n\n\n## Weaknesses\n\n**(W1) Relevance of potential outcomes**\n- We would like to emphasize that potential outcomes are highly relevant in many decision-making settings such as medicine [1]. Example: predicting the survival probability under treatment A and the survival probability under treatment B, for a patient _with no previous treatment_. Here, predicting potential outcomes is directly relevant for clinical decision support** and informs which treatment is preferred (in terms of the health benefit but also side effects and quality of life) [1].  \n\n- We would like to clarify that _one potential outcome is only observed in the training data, but **not** upon deployment_. Instead, **both potential outcomes must be predicted upon deployment** (aka in the training data). \n\n[1] Feuerriegel S, Frauen D, Melnychuk V, et al. Causal machine learning for predicting treatment outcomes. Nature Medicine, 2024.\n\n\n**(W2) IPS & Neyman-orthogonality**\n\nWe admit that we leveraged the idea of inverse propensity score to address the selection bias. Nevertheless, we would like to clarify that our orthogonal diffusion loss for diffusion models has a clear advantage over existing literature in models for learning potential outcome distributions. In particular, **our orthogonal diffusion loss is Neyman-orthogonal wrt. the nuisance functions**. As such, we offer a theoretical guarantee that our loss is robust against errors in nuisance function estimation. This goes beyond the standard inverse propensity score from the literature, and the corresponding proof of Neyman-orthognality is _non-trivial_. \n\nTo improve the rigor of our work, **we added theory and new experiments** for our novel orthogonal diffusion loss:\n\n1. **Theoretical proof**:  **We provide the formal proof that our proposed orthogonal diffusion loss is Neyman-orthogonal wrt. its nuisance functions** (see **Theorem and Proof in the Rebuttal PDF**). \n\n\n2. **New experiment**: **We include the ablation study of our orthogonal diffusion loss** where we compare against a vanilla diffusion model without our orthogonal diffusion loss. The experiment results allow us to assess the contribution of our novel orthogonal diffusion loss to the overall performance (see **Table 4 / Rebuttal PDF**). We observe a clear performance gain from including our novel loss.\n\n**Action**: We will add our theorem and our detailed proof of Neyman-orthogonality to our revised paper. We will also add the new experiments with our ablation to our revised paper in order to show the performance gain from our novel orthogonal diffusion loss.\n\n\n**(W3) Presentation of background and our novel loss**\n\nThank you for your comment. It is true that methods based on conditional diffusion models always contain the forward and reverse process etc, regardless of whether the methods come from computer vision or other domains. Analogously, our method also adopts the forward and reverse process as in conditional diffusion models. However, our objective was **not** to change the forward and reverse process. Rather, **our objective was to develop a novel method for causal tasks**, and, for this reason, **we propose a novel orthogonal diffusion loss with theoretical properties** (see **Theorem and Proof in the Rebuttal PDF**).\n\nWe realized that some parts of our method are familiar to readers with a background in diffusion models. However, readers from causal inference may also appreciate the background (see, e.g., Reviewer Cvxf). However, we understand that our presentation could be improved. **Action:** We will revise our paper and split the method section into two sections presenting ‘what is known’ (=the diffusion model background) and ‘our innovation’ (=our novel orthogonal diffusion loss). \n\n\n### Details\n\n\n**(Q1)** \n\n_Difference / variational inference methods_\n\nIn principle, yes, any variational inference-based causal inference method can learn the distribution; however, none of the previous work has done so. Importantly, many of them use the variational-based method simply to learn latent variables but _then predict a point estimation of the CATE_ (e.g., TEDVAE).  In contrast, **we are the first variational inference-based causal inference method designed to learn the distribution of potential outcomes**, and **our method offers favorable theoretical properties (Neyman-orthogonality)**.\n\n_Difference / conformal prediction_\n\nIt is true that conformal prediction (CP)-based causal inference methods can predict a predictive interval. However, we would like to point out the main difference: **Our method can learn the whole distribution but CP can not**. We also would like to clarify the **two main advantages of our method** over CP:\n\n1. **More flexible.** CP targets the total uncertainty and needs to be calibrated each time for different levels of uncertainty. However, our method targets the aleatoric uncertainty and yields the full distribution of the potential outcomes at once. \n\n2. **More general/useful.** There are cases in which the mean and the variance are the same for the outcomes, but the distribution of the outcomes is different. This kind of information is very useful and can be obtained from our method. However, CP can not provide such information.\n\n**(Q2)** Thank you. The expectation in Eq (11) is taken over $(y_0, x, a) sampled from observational distribution and $\\epsilon$ sampled from normal distribution. We will add this to Eq (11) in our paper.\n\n**(Q3)** Thanks. We will fix the sentence.\n\n**(Q4)** Thanks. We will remove the term “unrelated” and use a more nuanced wording to say that the papers study a _different task_. \n\n**(Q5)** Thanks. We will rearrange Table 1 as per your suggestion.\n\n\n## Questions\n\nThank you for the suggestion. **We now provide the formal proof that our proposed orthogonal diffusion loss is Neyman-orthogonal wrt. its nuisance functions** (see **Theorem and Proof in the Rebuttal PDF**).", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "VlNqQ2RuWx", "reviewer": "Reviewer_Cvxf", "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.", "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?\n\n- How sensitive is the proposed estimator to hyperparameter, choice of model architecture or other training details?", "limitations": "The paper discusses limitations due to assumptions and efficiency of the method.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "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.\n\nQuality: Experimental evaluation of the method is quite diverse and extensive, demonstrating superior performance on various datasets compared to baselines.\n\nClarity: The motivation, method and overall paper are well-written and easy to follow. \n\nSignificance: The paper highlights and addresses significant aspects of causal inference in medical applications, like individual effects, and uncertainty quantification.", "weaknesses": "Clarity:\n\n- 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.\n\n- 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your positive review and your helpful comments! We appreciate that you find our paper interesting, important, and comprehensive. We are very delighted to answer your questions and improve our paper as a result.\n\n## Response to “Weaknesses”\n\n* Thank you very much for your suggestions! Yes, we compressed the preliminaries to save space. **Action:** We will add more details about our method in order to make it more easy to apply. We will also use the extra space in the camera-ready version to expand our section with implementation details to provide more information about how we train the diffusion models. \n\n* Thank you! **Action:** We will fix the typos and make sure that all tables and references are properly cross-referenced in the camera-ready version.\n\n\n## Answer to “Questions”:\n\n**(Q1) Choice of datasets**\n\nThank you for asking this question! Actually, there is no specific reason for splitting the estimation of these causal quantities on different datasets. Our rationale was simply that we wanted to empirically demonstrate the flexibility and generalization power of our model across a broad range of datasets.\n\nHowever, some commonly used datasets indeed can **not** be used for evaluating all three causal quantities. For example,  ACIC2016 (with 77 datasets) and ACIC2018 (with 24 datasets) datasets, which we also used in our paper. These datasets are designed for CATE benchmarking. It is impossible to use it for some other tasks, such as evaluating distribution. This is because these datasets only provide a single sample from the distribution instead of the full distribution knowledge.  Thus the ground truth is unavailable. \n\nTherefore, for CATE estimation, we benchmarked on ACIC2016 and ACIC2018. We additionally report the average performance on ACIC2016 and ACIC2018   (see **Table 2/Rebuttal PDF**). For evaluating distribution, we report the experiment results using the empirical Wasserstein distance on the semi-synthetic dataset where the ground truth distribution is available (see Table 1 / Rebuttal PDF). We find that we outperform the baselines by a clear margin for both tasks. The results also demonstrate our method is highly effective and robust across various datasets. \n\n**Action:** Previously, we have presented only a selection of our experiments in the main paper due to space constraints. In the camera-ready version, we will report additional experiments to streamline the presentation. \n\n**(Q2) Robustness to hyperparameters** \n\nThank you for asking this question! We found that our proposed model is robust and not very sensitive to the hyperparameters and model architecture. We performed a grid search of hyperparameters in a reasonable range. For example, even when changing various parameters (e.g., number of neurons), the overall performance was state-of-the-art and stable. One factor that, however, was relevant was the use of the clipping propensity score technique to avoid the numerical instability of the propensity score during training, yet this technique is commonly also applied to the baselines in order to improve their robustness.\n\n**Action:** We will report our tuning grid during hyperparameter tuning in our revised paper, and we will further discuss that our method is largely robust to different choices of the hyperparameters.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
Thank you very much for the constructive and positive evaluation of our paper and your helpful comments! We addressed all of them in the comments below and uploaded **additional results as a PDF file**.

Our **main improvements** are the following:

* **We provide a theoretical guarantee:**  
We provide **a formal proof that our proposed orthogonal diffusion loss is Neyman-orthogonal wrt. its nuisance functions** (see **Theorem and Proof in the Rebuttal PDF**). As such, we offer a theoretical guarantee that our loss is robust against errors in nuisance function estimation. This thus gives a proper justification of why our loss for predicting the distribution of potential outcomes is effective and thus beneficial in practice. 

* **We provide further extensive experiments:** 

(1) We evaluate **the learned distributions of the potential outcomes** (see **Table 1 / Rebuttal PDF**). We find that we outperform the baselines by a clear margin. Thereby, we show that our method is not only effective for providing point estimates but also for learning distributions. 

(2) We include the **ablation study of our orthogonal diffusion loss** where we compare against a vanilla diffusion loss without our orthogonal diffusion loss. (see **Table 4 / Rebuttal PDF**). We observe a clear performance gain from including our novel loss (over a vanilla diffusion model). The results also support our theory (see **Theorem and Proof in the Rebuttal PDF**).

(3) We include **three additional CATE benchmarking datasets** and find our method gives state-of-the-art performance (see **Table 3 / Rebuttal PDF**).

(4) We report the **average performance** based on the PEHE evaluation metric **across ACIC2016 and ACIC 2018 datasets** (see **Table 2 / Rebuttal PDF**). The experiment results again demonstrate that our method is superior and robust.

As a summary, we would like to emphasize our contribution again: We develop a conditional denoising diffusion model to learn complex **distributions of potential outcomes**, where we address the selection bias through a **novel orthogonal diffusion loss**. We also offer a non-trivial, **theoretical proof that our orthogonal diffusion loss is Neyman-orthogonal** wrt. its nuisance functions. This gives theoretical support that our method is effective and why it outperforms the baselines. 


We will incorporate all changes (labeled with **Action**) into the camera-ready version of our paper. Given these improvements, we are confident that our paper will be a valuable contribution to the machine learning for healthcare literature and a good fit for NeurIPS 2024.

Meta-review: Overall, the reviewers appreciated the topic of the paper which formulates potential outcomes estimation as a diffusion problem and found the initial experiments convincing. Most of them appreciated the clarity of presentation. Multiple reviewers expressed concerns about the unclear connection to IPS and insufficient ablation studies and datasets. The authors addressed these concerns in the rebuttal and should update the manuscript accordingly.
Decision: Accept (poster)