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

Authors: yuchen  ma
valentyn  melnychuk
jonas  schweisthal
stefan  feuerriegel
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-14 21:22:14

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.

Main paper: 
 introduction: 
 Predicting potential outcomes (POs) for patients from observational data is crucial for decisionmaking in medicine [16] . For example, in cancer care, one 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 is 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, 11] .In the causal inference literature, many existing methods are aimed at conditional average treatment effect (CATE) estimation where POs are used as auxiliary quantities (e.g., [1, 2, 3, 6, 10, 12, 13, 25, 28, 30, 31, 32, 37, 45, 64, 66] ). In principle, some CATE methods could be used for predicting POs but these are not tailored for POs and thus tend to underperform at predicting POs. The underlying reason is that the estimator with the best performance in estimating CATE might not do best at predicting POs in finite samples [11, 13] . Additionally, many state-of-the-art CATE estimators leverage an inductive bias [13, 31] , i.e., they assume that POs should share similar structures. As a result, they assume that the CATE follows a much simpler function than the POs, so that it is easier to 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 decision-making in medical applications.Table 1 : 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., Neymanorthogonality wrt. nuisance functions). POs are the target refers to whether the methods are originally designed for the task of predicting POs.Estimator type UQ Bias-addressing Orthogonal POs are the target Key limitations S-learner [37] One-step (plug-in), model-agnostic ✗ ✗ ✗ ✓ Selection bias; point estimator T-learner [37] One-step (plug-in), model-agnostic ✗ ✗ ✗ ✓ Selection bias; point estimator DR-learner [31, 39] Two-step, model-agnostic ✗ ✓ ✓ ✗ CATE inductive bias; point estimator RA-learner [12] Two-step, model-agnostic ✗ ✓ ✗ ✗ CATE inductive bias; point estimator TARNet [54] One-step (plug-in), model-specific ✗ ✗ ✗ ✓ Selection bias; point estimator CFR [30] One-step (plug-in), model-specific ✗ ✓ ✗ ✓ Point estimator GANITE [64] Two-step, model-specific(✗) ✗ ✗ (✗)Selection bias; unstable training TEDVAE [66] One-step (plug-in), model-specific ✗ ✗ ✗ ✗ Selection bias; non-identifiability DiffPO (ours) One-step (plug-in), model-specific ✓ ✓ ✓ ✓ Sampling time longer (✗) in the column UQ refers to that methods are not originally designed for uncertainty quantification of POs, but can be adapted for this purpose.In this paper, we propose a causal diffusion model called DiffPO for predicting POs. Specifically, we leverage a tailored conditional diffusion model to learn complex distributions of POs. We further propose a novel orthogonal diffusion loss to adjust for the covariate shift problem in finite samples.Our orthogonal diffusion loss ensures Neyman-orthogonality, which offers favorable theoretical properties. In particular, it makes our method more robust to misspecification. Further, our DiffPO is carefully designed for reliable inferences in 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.Our method is highly flexible. It 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 is unlike many state-of-the-art 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 a flexible and reliable approach.Overall, our main contributions are the following: (1) We propose a 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-orthogonality.(3) We conduct a wide range of experiments in different settings to demonstrate the flexibility and effectiveness of our DiffPO. 12 Related WorkOur work aims to learn the distributions of POs for patients. The task is 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 finite-sample performance. 
predicting pos and cate estimation: 
 There have been many works designed for CATE estimation (e.g., [1, 2, 3, 6, 10, 12, 13, 25, 28, 30, 31, 32, 37, 45, 64, 66] ). Existing CATE methods fall into two categories: model-agnostic and model-specific estimators. We only give a concise overview in this section (see Table 1 ). A detailed literature review is in Appendix A.3.Model-agnostic estimators. Model-agnostic learning strategies are also known as meta-learners (e.g., [12, 13, 31, 32, 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 pseudooutcome and then obtain CATE directly by regressing the input covariates on the pseudo-outcomes in the second step. (Note that pseudo-outcomes are not potential outcomes).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 an estimation of POs, and, thus, they are not applicable for our task. We show this in the column POs are the target in Table 1 . Instead, we can only compare with those methods that are applicable to our task: that is, where the model can directly output POs predictions [37, 54, 64] .Diffusion models were introduced for learning from complex distribution for a given dataset from which one can then sample high-quality data points (e.g., images) [26, 55, 57] . Diffusion models have achieved state-of-the-art performance, outperforming other generative models on various tasks in the computer vision field (e.g., [14, 56, 59] ). We give a brief technical overview of diffusion models in Sec. 3.Diffusion models were previously used for different causal inference tasks but in a different setting from ours, for example, generating the counterfactual of a 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, 51] is similar to minimizing the changes that one needs to make to an image in order for the classifier to categorize the image into a 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 is the causal graph, while our output is a 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. 
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 q (x 0 ) into a tractable prior N (0, I). For this, it gradually adds noise to an initial sample x 0 ∼ q (x 0 ) across T steps with variance schedules {β 1 , . . . , β T }. The forward process can be written as q (x 1:T | x 0 ) = T t=1 q (x t | x t-1 ) with Gaussian distribution q (x t | x t-1 ) := N x t ; √ 1 -β t x t-1 , β t I . It admits a closed form of q (x t | x 0 ) that is also a Gaussian distribution N ( √ ᾱt x 0 , (1 -ᾱt ) I), where ᾱt = t i=1 (1 -β i ). The (2) reverse process p (x 0:T ) = T t=1 p (x t-1 | x t )gradually denoises the latent variable x T ∼ N (0, I) and further allows for generating new data samples from q (x 0 ). The distributions p (x t-1 | x t ) are usually unknown and approximated by a neural network with parameters θ. Thus, a parameterized Markov chain {p θ (x t-1 | x t )} T t=1 is trained in the reverse process. It can be parameterized as p θ (x t-1 | x t ) := N (x t-1 ; µ θ (x t , t) , Σ θ (x t , t)). [26] suggested using diagonal Σ θ (x t , t) with a constant σ t and computing µ θ (x t , t) as a function of x t and ϵ θ (x t , t). This yieldsµ θ (x t , t) = 1 √ αt x t -βt √ 1-ᾱt ϵ θ (x t , t), where α t := 1β t , ᾱt := i≤t α i and where ϵ θ (x t , t) predicts a 'ground truth' noise component ϵ for the noisy data sample x t .Later, we develop a 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 a remedy, we develop an orthogonal diffusion loss that is tailored to learn complex distributions of POs. 
problem formulation: 
 Setup: We consider an observational dataset D with i.i.d. patient data. The dataset consists of: Potential outcomes: We build upon the standard setting of Neyman-Rubin potential outcomes framework [50] . Hence, let Y (a) denote the potential outcome after intervening on the treatment by setting it to a. We have two potential outcomes for each individual: Y (1) if treatment is administered (i.e., A = 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 is observed. Hence, Y = AY (1) + (1 -A) Y (0).an outcome of interest Y ∈ Y ⊆ R, d X -dimensional covariates (also called confounders) X ∈ X ⊆ R d X ,Identifiability: To ensure that POs are identifiable, we follow previous literature (e.g., [12, 31, 64] ) and make the following standard assumptions: Objective: Our main interest lies in predicting POs in medical settings. Formally, E[Y (a) | X] are the expected POs corresponding to a treatment assignment (intervention) a 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 in the POs [23, 35] . For example, uncertainty quantification is needed in medical practice to decide whether to apply either a treatment that has a small but certain benefit or a treatment with a large benefit but potentially a large variability in the outcomes and thus large uncertainty of whether the treatment is actually beneficial for a specific patient. Hence, we estimate the distribution of the POs after assigning treatment a to an individual, i.e., p(Y (a) | X). Learning this distribution allows us to sample from it and obtain predictive intervals for uncertainty quantification. 
diffpo for predicting potential outcomes: 
 Overview: In the following, we introduce our DiffPO: a diffusion-based model for predicting POs with 3 key components (Fig. 1 ): 1 a forward diffusion process, 2 a reverse diffusion process, and 3 a novel orthogonal diffusion loss. Finally, we introduce our training and sample procedure.Our DiffPO builds upon diffusion models [26, 55, 57] . We distinguish the distributions learned in the forward process and the reverse process via q and p, respectively, to make the notation straightforward in this section. 1 Forward process: Given a data point (y, x, a) sampled from the observational data distribution p(Y, X, A), in the forward diffusion process, we gradually add Gaussian noise to the initial y (denoted as y 0 ) across T time steps. This thus produces a sequence of noisy samples y 1 , . . . , y T in the same sample space as y 0 . The forward process follows a Markov chainEQUATIONwhere β t is a variance schedule: {β t ∈ (0, 1)} T t=1 is a small positive constant that represents a noise level and, thus, essentially controls the step sizes. By using the reparameterization trick [34] , sampling y t at an arbitrary timestep t has a closed formEQUATIONwhere α t := 1β t and ᾱt := t s=1 α s . Thus, y t can be expressed asEQUATIONwhere ϵ ∈ R d Y ∼ N (0, I). As a result, the data sample y 0 gradually 'looses' its distinguishable features for later steps t. Eventually, when T → ∞, y T is equivalent to an isotropic Gaussian distribution.The reverse process starts at a known prior distribution with density p (y T ) = N (y T ; 0, I) and then proceeds backward to obtain the data distribution at step t = 0. Formally, the reverse diffusion process is also a Markov chain, and the density of its joint distribution p θ (y 0:T | x, a) can be written as We employ a conditional diffusion model with the reverse process in Eq. ( 4). The conditional density p θ (y t-1 | y t , x, a) is also Gaussian, and we parameterize it asEQUATIONThe reverse process gradually denoises y T ∼ N (0, I) through the learned Gaussian transitions. Once we have computed the latter, we approximate the original data distribution and can even sample from it.Variational inference: Directly computing and maximizing the likelihood p θ (y 0 | x, a) is difficult because it involves intractable integration. As a 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) viaEQUATIONThe right side of Eq. ( 6) can be rewritten as (see Appendix B for the derivation)E y1∼q(Y1|y0) [log p θ (y 0 | y 1 , x, a)] reconstruction term -D KL (q (y T | y 0 ) ∥p (y T ))prior matching term- T t=2 E yt∼q(Yt|y0) [D KL (q (y t-1 | y t , y 0 ) ∥ p θ (y t-1 | y t , x, a))]denoising matching term .Forward ProcessReverse ProcessFigure 1 : Overview of our causal diffusion model DiffPO. Our method involves a forward and reverse diffusion process to learn the distributions of potential outcomes. Additionally, we address selection bias through our orthogonal diffusion loss.In the denoising matching term, we try to minimize the KL-divergence between two distributions, one is tractable, ground truth denoising transition step q (y t-1 | y t , y 0 ), and other is denoising transition step p θ (y t-1 | y t , x, a).Conditional denoising function. In practice, directly optimizing the learning objective in Eq. ( 7) is inefficient. Recently, [26] has shown a way to 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 our learning objective, but for our conditional diffusion models.We employ a trainable conditional denoising functionf θ : (Y × R | X , A) → Y.The conditional denoising function is a function approximation, which we use to predict the 'ground truth' noise component ϵ for the noisy data sample y t conditioned on x and a. We consider the following parameterization that computes µ θ (y t , t | x, a) as a function of y t , x, a and f θ . This givesEQUATIONThe simplified training objective for t ∼ Unif{1, . . . , T } is thenEQUATION)We provide the full derivation of Eq. ( 8) and Eq. ( 9) in Appendix C.Why we need to address selection bias through our orthogonal diffusion loss: In observational data, treatments are not 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 a more aggressive treatment. As a 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, we have a distribution shift in the covariates, i.e., p(X| A = 0) ̸ = p(X | A = 1).If not adjusted for, this distribution shift may result in a selection bias and thus an inflated variance in the estimates of POs (and thus the causal effects), especially in low-sample settings [12, 29, 54] . To address this, we thus introduce a novel orthogonal diffusion in the following.function g ϕ parameterized by ϕ. Let us denote the estimated propensity score by π(x) = g ϕ (x). We then assign weights w π (x, a) to each sample via some function w π : (X × {0, 1}) → R + . More specifically, we define the function w π : (X × {0, 1}) → R + asEQUATIONThen, our orthogonal diffusion loss is given byEQUATION)The loss in Eq. ( 9) fits the conditional distributions p(Y | X, A) based on the data samples from the joint observational distribution, i.e., (y 0 , x, a) ∼ p(Y, 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 | X, A = 0) for the untreated population [60] . Yet, our target is to learn the potential outcome distributionsp(Y (a) | X) = p(Y | X, A = a)for both a = 0 and a = 1 equally well in all the populations. This would be equivalent to minimizing the target loss from Eq. ( 9) if the samples are from the joint distributions of the POs, i.e., (y 0 , x) ∼ p(Y (a), X). However, we do not have samples from the distributions of POs, thus we need the target loss in Eq. ( 11) to be estimated via the propensity score π(x) to learn the distributions of POs. Hence, the following remark holds.Remark 1. The orthogonal diffusion loss in Eq. ( 11) evaluated with the ground truth nuisance functions matches the following target loss:EQUATION)Proof. The orthogonal diffusion loss in Eq. ( 11) is an inverse propensity-weighted estimator of the target loss in Eq. ( 12).Theorem 1 (Neyman-orthogonality). The orthogonal diffusion loss in Eq. ( 11) is Neyman-orthogonal wrt. its nuisance functions.Proof. See Appendix D.As a result of Neyman-orthogonality, our orthogonal diffusion loss has a clear practical advantage: it offers robustness against errors in nuisance function estimation. Specifically, it is first-order insensitive wrt. errors in the propensity score π(x) and, thus, is a more accurate objective for our PO predictive model.Training. Our training algorithm proceeds along the following steps. We first train the function g ϕ on the dataset to estimate the propensity score. After this, the parameters of g ϕ are frozen. We then compute each sample weight for the orthogonal diffusion loss through the weight function w π . 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. (11) .Sampling. Once our diffusion model has learned the distribution p θ (Y t-1 | y t , x, a), we can generate samples from it. The sampling process starts by sampling y T from the prior distribution and then denoising it. To obtain y t-1 from p θ (y t-1 | y t , x, a), this requires the mean μ from the previous step, which can be computed by Eq. ( 8). Thus, y t-1 can be computed via μ + σ t z, where z ∼ N (0, I). This process continues for t = T, . . . until we arrive at y 0 .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 a batch size of 256 and a learning rate of 0.0005. We follow [42] in 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. 4). To guide the model, we need to identify which PO should the loss be computed on. Hence, we introduce causal masks as input to our model: an observational mask m o , a targeted mask m t , and a conditional mask m c . We only compute the loss at the place where the value of the targeted mask is 1. Further details about our implementation are in Appendix E.1.6 Experiments 
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 k ≥ 1) between two distributions ν 1 and ν 2 isEQUATION)where F -1 1 (l) and F -1 2 (l) are the quantile functions (inverse cumulative distribution functions) of ν 1 and ν 2 for quantile l, respectively. Specifically, we compute the empirical Wasserstein distance based on two sets of finite samples, i.e.,W k (ν 1 , ν 2 ) = 1 n n i=1 X (i) -Y (i) k 1/k, where X 1 , . . . , X n are the samples from ν 1 and Y 1 , . . . , Y n are the samples from ν 2 . In our experiments, we report the empirical Wasserstein distance for k = 1. Lower values of W k are preferred. ) for two potential outcomes (i.e., a = 0 and a = 1) on the synthetic dataset. Reported: mean ± standard deviation over ten-fold train-test splits. 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 benchmark our model using synthetic datasets. Further details about the synthetic datasets are in the Appendix F.a = 0 a = 1 Ŵ 1 in Ŵ 1 out Ŵ 1 in Ŵ 1 out S-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 1 ):(1) S-learner [37] : is a model-agnostic learner that fits a single regression model by concatenating the covariate and the treatment as input; (2) T-learner [37] : is a 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 balancing 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 (1) to (4) above are only able to give point estimates of 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.Results. The results are in Table 2 . We find that our method gives the lowest empirical Wasserstein distance out of all methods, which is desirable. Hence, the experiments show that our method outperforms the baselines by a clear margin.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 is often reported based on the 95% and 99% PIs. We thus compute the individual equal-tailed (1α) PIs with α ∈ [0.01, 0.05]. We then calculate the faithfulness of the estimated PIs to evaluate the empirical coverage by computing the frequency with which the PIs contain the outcomes in the test data. Results: We evaluate the empirical coverage across different quantiles (1α). The results are in Table 3 . It shows that our DiffPO is generally faithful, while this is 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 not faithful [17, 24] . Performance metrics for POs: To evaluate point estimates of POs, we compute the difference between the predicted PO ŷi and the ground truth PO y i via the root mean squarederror RMSE = 1 N N i=1 (ŷ i -y i ) 2 .Baselines. We compare against those methods that are applicable to this task (i.e., the model can directly output POs; see Table 1 ).Results. The results are in Table 4 . We find that our DiffPO gives the best point estimates of POs.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 τ(x) = E[Y (1) -Y (0) | X = x],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 µ a (x) = E[Y | X, A = a] for a ∈ {0, 1} and then simply leverage that τ (x) = µ 1 (x)µ 0 (x). Datasets. We estimate the CATE across ACIC 2016 & ACIC 2018, which are widely used dataset collections for CATE benchmarking [66, 12, 42] . The ACIC2016 [46] contains 77 different 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 hyperparameters on the first split, and evaluate the average out-sample on every split.We compare theϵ PEHE = 1 N N i=1 (τ (x i ) -τ (x i )) 2 .Baselines. We consider a broad array of state-of-the-art 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 [10, 30, 37, 54, 61, 64, 66] , as listed in the Table 1 . Specifically, we use the following baselines: (1) to (5) from Sec. 6.1. ( 6) DR-learner [39, 31] : generates pseudo-outcomes based on the doubly-robust AIPW estimator; (7) RA-learner [12] : 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-learners (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 E.Results. The results for CATE estimation are in Table 5 . We find that our method achieves a good performance similar to existing methods, even 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 state-of-the-art performance.A clear advantage of our DiffPO is that we not only return point estimates but also the distribution of POs. This is crucial in medical practice [23, 35] in order to 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 a 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 on the extracted features and treatment assignments from a real-world clinical trial. The dataset comprises 747 patients, with 25 features for each patient. Further details for IHDP are in Appendix F.3. Insights. Our DiffPO 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 estimation [31, 37, 54] . However, distributional knowledge of POs is important to account for uncertainty, because it informs how likely a certain outcome is and gives the probability that the PO lies in a desired range, especially in medicine. Fig. 2 shows the empirical distributions of POs, i.e., p(Y (0) | X) and p(Y (1) | 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-making. Here, the learned distributions of POs can help medical practitioners in choosing a treatment plan that promises not only a large benefit for the patients but also that the benefit is highly probable.Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, 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. 6 and 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 to a 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, or as supplemental material. 
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.As with other methods in 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 DiffPO -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, as we have shown above, our method can successfully scale to real-world datasets from medical practice.Broader impact. We expect our method to have a significant impact in medicine where better decision support is needed to personalize treatment decisions to 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 to be a first step toward developing more generalizable approaches for a variety of causal inference settings.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, 14, 59] ). Motivated by this, diffusion models were previously used for different causal inference tasks but in a different setting from ours [51, 52, 8] .[51] proposed a model called Diff-SCM that uses diffusion models and anti-causal predictors to generate the counterfactual of a 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. [8] focused on a similar task as well as [53, 36] , with the aim to answer causal queries by employing variational autoencoder and normalizing flow, respectively. [36] focused on the same task as [51] , aiming for generating high-quality counterfactual images and proposed a model called CausalDiffAE. [36] follows the idea of [8] to view the diffusion model as an encoder-decoder framework and employ so-called denoising diffusion implicit models (DDIM). Diffusion models have also been used in the causal discovery field. DiffAN [52] is an algorithm based on denoising diffusion training for causal discovery via topological ordering. In this work, we use diffusion models for predicting POs.Neyman-orthogonality of a functional refers to the mean zero property of its directional derivatives along one-dimensional paths that change nuisance functions [18] . A loss function is called Neymanorthogonal when this property holds for all its directional derivatives along one-dimensional paths that change the (infinite-dimensional) 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 in the DR-learner [39, 31] , R-learner [45] , and i-learner [60] .Our work is about predicting potential outcomes (POs) for individuals; however, it is 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 an individual's expected potential outcome with and without treatment. Below, we give a brief overview of the CATE literature.The estimation of the CATE has received a lot of attention in the causal inference literature (e.g., [10, 25, 28, 1, 54, 37, 2, 30, 3, 64, 61, 12, 66, 31, 45, 32, 13] ). Among those CATE estimation methods, many have been proposed to estimate the effects of binary treatments (e.g., [10, 25, 28, 1, 54, 37, 2, 30, 3, 64, 6, 12, 66, 31, 45, 32, 13] ). 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 ML-based CATE estimators (including representation learning-based and other neural network-based estimators).Now we consider popular approaches that involve using model-agnostic learning strategies, also known as meta-learners, which can be implemented using any ML method (e.g., [37, 31, 45, 32, 13, 12] ). The idea behind "meta-learner" strategies was originally introduced in [37] and expanded in [45, 31, 12] .Existing CATE meta-learners can be categorized into: (a) one-step (plug-in) learners (indirect metalearners) that output two regression functions from the observational data and then compute CATE as the difference in the POs (this is the strategy underlying the S-and T-learners); and (b) two-step learners (direct meta-learners/multi-stage direct estimators). These learners first compute nuisance functions to build a pseudo-outcome. In the second step, they obtain the CATE directly by regressing the input covariates on the pseudo-outcome. (Note that pseudo-outcomes are not potential outcomes).In terms of (b), existing methods fall largely into three broad classes: regression adjustment (RA), propensity weighting (PW), or doubly robust (DR) strategies.Many methods proposed in the related work do not fall within the meta-learner class because they rely on the properties of a specific ML method. We roughly categorize them into two branches: one is neural network-based (NN-based) estimators and the other does not use neural networks, for example, some using causal Forest [61] , Bayesian regression [25] , or Gaussian processes [1] .As for the NN-based estimators, some are representation learning-based estimators. Much work of this track focused on 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] . [54] used representation learning to 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 in a 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., [29, 20, 4, 21, 62, 67, 5, 7] ). However, a wrongly chosen dimensionality of the representation or a too large balancing weight can induce confounding bias [43] , and, therefore, the representation learning-based methods fall outside the scope of this paper.Directly computing and maximizing the likelihood p(y | x, a) is difficult because it involves integrating all latent variables, which is intractable for complex models. Thus, we derive the evidence lower bound (ELBO) similar to the bound of 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 | x, a) = log p (y 0:T | x, a) dy 1:T = log p (y 0:T | x, a) q (y 1:T | y 0 ) q (y 1:T | y 0 )EQUATION)The ELBO in Eq. ( 14) can be rewritten as follows:log p(y) ≥E q(y 1:T |y0) log p (y 0:T | x, a) q (y 1:T | y 0 ) =E q(y 1:T |y0) log p (y T ) T t=1 p θ (y t-1 | y t , x, a) T t=1 q (y t | y t-1 ) =E q(y 1:T |y0) log p (y T ) p θ (y 0 | y 1 , x, a) T t=2 p θ (y t-1 | y t , x, a) q (y 1 | y 0 ) T t=2 q (y t | y t-1 ) =E q(y 1:T |y0) log p (y T ) p θ (y 0 | y 1 , x, a) T t=2 p θ (y t-1 | y t , x, a) q (y 1 | y 0 ) T t=2 q (y t | y t-1 , y 0 ) =E q(y 1:T |y0) log p (y T ) p θ (y 0 | y 1 , x, a) q (y 1 | y 0 ) + log T t=2 p θ (y t-1 | y t , x, a) q (y t | y t-1 , y 0 ) =E q(y 1:T |y0)   log p (y T ) p θ (y 0 | y 1 , x, a) q (y 1 | y 0 ) + log T t=2 p θ (y t-1 | y t , x, a) q(yt-1|yt,y0)q(yt|y0) q(yt-1|y0)   =E q(y 1:T |y0) log p (y T ) p θ (y 0 | y 1 , x, a) q (y 1 | y 0 ) + log q (y 1 | y 0 ) q (y T | y 0 ) + log T t=2 p θ (y t-1 | y t , x, a) q (y t-1 | y t , y 0 ) =E q(y 1:T |y0) log p (y T ) p θ (y 0 | y 1 , x, a) q (y T | y 0 ) + T t=2 log p θ (y t-1 | y t , x, a) q (y t-1 | y t , y 0 ) =E q(y 1:T |y0) [log p θ (y 0 | y 1 , x, a)] + E q(y 1:T |y0) log p (y T ) q (y T | y 0 ) + T t=2 E q(y 1:T |y0) log p θ (y t-1 | y t , x, a) q (y t-1 | y t , y 0 ) =E q(y1|y0) [log p θ (y 0 | y 1 , x, a)] + E q(y T |y0) log p (y T ) q (y T | y 0 ) + T t=2 E q(yt,yt-1|y0) log p θ (y t-1 | y t , x, a) q (y t-1 | y t , y 0 ) = E q(y1|y0) [log p θ (y 0 | y 1 , x, a)] reconstruction term -D KL (q (y T | y 0 ) ∥p (y T ))prior matching termEQUATIONWith 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 our aim is to match the approximate denoising transition step p θ (y t-1 | y t , x, a) to ground truth denoising transition step q (y t-1 | y t , y 0 ) as closely as possible, which we can also model as a Gaussian. Thus, optimizing the KL divergence term reduces to minimizing the difference between the means of the two distributions:EQUATIONBy using the reparameterization trick, we haveEQUATIONPlugging it into the previously derived true denoising transition mean µ q (y t , y 0 ), we yieldEQUATIONTherefore, we can set our approximate denoising transition mean µ θ (y t , t | x, a) asEQUATIONand the corresponding optimization problem becomesEQUATIONTheorem 2 (Neyman-orthogonality). The orthogonal diffusion lossEQUATIONis Neyman-orthogonal wrt. its nuisance functions.Proof. As it was shown in the Sec. 5.1, the minimization of the orthogonal diffusion loss is equivalent to the maximization of the following weighted ELBO wrt. to θ:EQUATIONWe denote a maximizer of the weighted ELBO with the ground truth nuisance functions by p θ * = argmax p θ E(p θ , π). Given a flexible enough diffusion model, p θ * (y 0 | x, a) coincides with the ground truth posterior and, thus matches the ground truth conditional distribution, namely, p(Y | X, A) [47] .To demonstrate Neyman-orthogonality, we need to show that a pathwise cross-derivative is equal to zero [18, 60, 44] , i.e.,EQUATIONWe start by taking the pathwise derivative wrt. the optimal diffusion model p θ * : EQUATIONEQUATIONIn the following, we summarize the implementation details of our DiffPO and baselines.We implemented our DiffPO in PyTorch. (Code is available at https://github.com/yccm/ DiffPO). Experiments were carried out on 1× NVIDIA A100-PCIE-40GB. 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 f θ 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 is 128. The β starts at 0.0001 and ends at 0.5, and the schedule is the "quadratic" version. The number of diffusion sampling steps is set to 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. 4). To guide the model, we need to identify which of two POs should be generated. Hence, we introduce causal masks as input to our model: an observational mask m o , a targeted mask m t , and a conditional mask m c . For the observational mask m o , it is 1 at the place where we have the observational data and 0 for the opposite case. Since we condition on x and a, the conditional mask m c is 1 for the element x 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, we 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 we 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.We follow the implementation from https://github.com/AliciaCurth/CATENets/tree/main for most of the CATE estimators, including S-learner [37] , T-learner [37] , DR-learner [39, 31] , RAlearner [12] , TARNet [54] . For GANITE, we follow the implementation of https://github.com/ vanderschaarlab/mlforhealthlabpub/tree/main/alg/ganite. For TEDVAE, we follow the implementation of https://github.com/WeijiaZhang24/TEDVAE.We performed hyperparameters tuning of the nuisance functions models for all the baselines based on five-fold cross-validation using the training subset. For each baseline, we performed a grid search with respect to different tuning criteria, evaluated on the validation subsets. We aimed for a fair comparison and thus kept the number of 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.We use one of the datasets from ACIC2018 with 177 covariates to generate the synthetic dataset with a coefficient matrix W , i.e.,EQUATIONWe use a ten-fold split for train/test samples (80%/20%).The 2016 Atlantic Causal Inference Challenge (ACIC2016) [46] The Infant Health and Development Program (IHDP) [25] has features and treatment assignments from a real-world clinical trial. The dataset comprised 747 subjects, with 25 features for each subject. Out of the 25 features, 6 are continuous, and 19 are binary with a binary treatment. For both treated and untreated, synthetic outcomes of IHDP are sampled from different conditional normal distributions. These distributions are homoscedastic (σ 2 = 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 data-generating mechanism:EQUATIONwhere β, W , ω are constant parameters of the simulation. For further details, we refer to [25] .We examine the Neyman-orthogonality property of our orthogonal diffusion loss through simulation experiments. We conduct experiments to show that the orthogonal diffusion loss can address the problem of model misspecifications when estimating nuisance functions (e.g., the propensity score π(x)).To this end, we perturb the propensity score manually to assess the robustness of the loss when the nuisance functions are 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 to zero. This demonstrates that, even when the propensity score is misspecified, the orthogonal loss remains robust, enabling our DiffPO to converge to the correct objective. This experiment supports Theorem 1 by illustrating how the stability and consistency of the orthogonal loss contribute to convergence. In sum, it shows the theoretical benefits of our proposed orthogonal diffusion loss. The weight w π (x, a) for each sample in the orthogonal diffusion loss L(θ, π) is thus replaced by weight w π (x, a). The CATE estimation error gradually converges as the sample size increases. This aligns with our expectation, as the loss remains robust even with varying errors in the estimation of nuisance functions.Justification: We provide open access to the data and code and provide an anonymous link to our code: https://anonymous.4open.science/r/DIFFPO Guidelines:• The answer NA means that paper does not include experiments requiring code. • Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?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, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper. • The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run 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).• It should be clear whether the error bar is the standard deviation or the standard error of the mean. • If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics. • The authors should make sure to preserve anonymity (e.g., if there is a special consideration 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 is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or 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.• The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a 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 is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology. • If there are 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 to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?Answer: [NA] Justification: Our work does not pose such risks.Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a 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 to 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 a best faith effort.12. Licenses for existing assets Question: 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?Answer: [Yes] Justification: All the datasets used in our paper are cited accordingly.Guidelines:• The answer NA means that the paper does not use 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 a URL. • The name of the license (e.g., CC-BY 4.0) should be included for each asset.• For scraped data from a particular source (e.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, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a 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.Code is available at https://github.com/yccm/DiffPO.Reverse process: We aim to learn the true distributions of POs, which can be identified as the conditional distribution in Sec. 4. Using the notation above, it can be rewritten as q (y 0 | x, a) = q (y 0:T | x, a) dy 1:T . As q (y t-1 | y t , x, a) is intractable, we need to learn a model distribution p θ parameterized by θ to approximate the data distribution.Orthogonal diffusion loss. Our proposed orthogonal diffusion loss is inspired by orthogonal learning theory[18,60]. Orthogonal learning theory provides a general toolbox for inference with semi-parametric models to provide quasi-oracle rates for statistical learning with a nuisance component. Informally, orthogonal losses are first-order insensitive to misspecification of the nuisance functions, which introduces many favorable properties such as double robustness[63]. Below, we construct an orthogonal diffusion loss for our method.We start by noting that π(x) 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 a trainable approach to reweight the data and thus adjust for the distribution shift from above. To this end, the propensity score π(x) is a nuisance function in our method, which we estimate through a trainable 
 
 Appendix:  


Decision: 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.

Reviews: review: paper_summary: paper_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.
strength_weakness: strength_weakness: - The high level evaluation results (table 2 and table 3) are impressive.  However, it would be good include additional details on benchmark performance, perhaps in the appendix.  For example, these results tables show the "% of runs with best performance".  It would be insightful to see the underlying performance metrics as well, to understand whether the approach was "just a little better" or a lot better.  This would also help evaluate whether there are contexts where DiffPO is consistently stronger or weaker than alternatives.

- The weighting approach in equation 10 will be unstable when samples have propensities approaching 0 or 1. While this is covered by standard causal assumptions related to population overlap, in practice it is common to have to clip propensity scores or drop high or low propensity samples. Was that unnecessary here, or what approach was chosen?
questions: questions: Does this approach lend itself to sensitivity analyses (e.g., of the assumption that there are no latent confounders?)

Could you please define best_in and best_out in Tables 3 and 4?

I believe there are usually annual ACIC challenges.  Why were acic 2016 and acic 2018 chosen over other years benchmarks?

The paper claims DiffPO can adapt to continuous treatments --- but the experiments seem to all be binary treatments? Did I miss something about the ACIC datasets or were continuous treatments not evaluated?
limitations: 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.
paper_summary: paper_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.
strength_weakness: strength_weakness: Although the paper is  technologically sound and well-written, the empirical examination is somewhat lacking. There have been a lot of dataset, such IHDP, Jobs, Twins in CATE estimation. I suggest the authors to complement the experimental results on these datasets to enhance the persuasive appeal of the experiments. There are additional questions about the experiments, see the details in Questions.
questions: questions: 1. The results in Table 3 (CATE estimation) is confusing. The authors claim the metric is PEHE, which is favored to be lower. However, the proposed method (DiffPO) achieves largest value in Table 3. What does the value in Table 3 mean? 
2. The authors examines the inferred potential outcome distribution in Section 6.3. However, how to validate the rationality of the inferred distribution? I think it is a very important problem. Otherwise, the inferred results are not trustworthy.
limitations: limitations: The authors have clearly listed the limitation of the paper, and claim the potential solution to address these limitations.
paper_summary: paper_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.
strength_weakness: strength_weakness: - The motivation behind this is not clear as in most of the cases people are interested in the causal effect instead of potential outcomes, especially given one potential outcome is observed already.
- The main novelty of this work is the orthogonal diffusion loss. However, it is simply using the idea of inverse propensity scores.
- There is nothing specially designed for potential outcomes in the forward and reverse process, the variational inference and the conditional denoising function, which should not be clarified as background instead of methodology.

----
details
- L350 In fact any variational inference based causal inference method can predict the distribution of potential outcomes. In addition, conformal prediction based causal inference methods can predict a predictive interval for each potential outcome.
- Eq(11) It is unclear what the expectation is taken over.
- L189 it is not a sentence.
- L118 I don't think those papers are totally unrelated from this work.
- Table 1. I believe for the purpose of this work, the comparison should be between different estimation models. The meta learners should be applicable along with the proposed model and other baselines, so it is not that crucial to list them in Table 1.
questions: 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: limitations: See weakness
paper_summary: paper_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.
strength_weakness: strength_weakness: Clarity:

- I found the preliminaries about diffusion models highly compressed (which is understandable due to space constraints). Since the proposed method relies heavily on diffusion modeling, the authors might consider using the extra space allowed in camera ready versions for expanding on the relevant details about training diffusion models.

- Authors should also make sure to reference all tables and figures in the main text to better guide the reader to results relevant to a given section. Nit: line 309 does have what was probably intended as a reference but is broken.
questions: questions: - The three different kind of experiments (PO estimation, CATE estimation, PO distribution estimation) all use different datasets. Is there a reason all these quantities cannot be estimated for all datasets?

- How sensitive is the proposed estimator to hyperparameter, choice of model architecture or other training details?
limitations: limitations: The paper discusses limitations due to assumptions and efficiency of the method.

Comments: title: Response to Reviewer PmLo
Response to Reviewer KUb5
comment: Thank you so much for your response and for raising the score! 

We will include all suggested points in the paper.
Thank you again for your effort in reviewing this paper and your constructive comments!
Thank you for your reply! We are happy that we have addressed your concerns sufficiently!

We also welcome the opportunity to answer your remaining question about the empirical Wasserstein distance.

The Wasserstein distance is a commonly used metric in machine learning [1,2,3] for quantifying the distance between two probability distributions based on the concept of optimal transport. In our work, we use it to evaluate the model performance in learning the potential outcome distributions. Thus, this allows us to move _beyond evaluating point estimates_ to the _full distribution_.


The $k$-Wasserstein distance (for any $k \geq 1$) between two distributions $\mu_1$ and $\mu_2$ is

$W^k\left(\mu_1, \mu_2\right)=\left(\int_0^1\left|\mathbb{F}_1^{-1}(l)-\mathbb{F}_2^{-1}(l)\right|^k \mathrm{~d} l\right)^{1 / k}$

where $\mathbb{F}_1^{-1}(l)$ and $\mathbb{F}_2^{-1}(l)$ are the quantile functions (inverse cumulative distribution functions) of $\mu_1$ and $\mu_2$ for quantile $l$, respectively.

In practice, Wasserstein distance is estimated by the empirical Wasserstein distance based on two sets of finite samples, i.e., 

$\hat{W}^k\left(\mu_1, \mu_2\right)$ $=\left(\frac{1}{n} \sum_{i=1}^n\left\|X_i-Y_i\right\|^k\right)^{1 / k}$,

where $X_1, \ldots, X_n$ are the samples from $\mu_1$ and $Y_1, \ldots, Y_n$ are the samples from $\mu_2$, and we set $k=1$ in our experiments, following [1]. 

(3) We use empirical Wasserstein distance to evaluate our model performance on learning the potential outcome distributions, where _lower_ distances imply _better_ performance. Here, our method gives the **best** results and outperforms the baselines by a clear margin (see **Table 1 / Rebuttal PDF**). This demonstrates the strong performance of our model not only on point estimates, but also on estimating the _full distribution_.


Thank you again for reading our rebuttal! We hope that the above results answer all your questions. 

We would greatly appreciate it if you would raise your score accordingly. Thank you. And do let us know if you have further questions -- we will do our best to answer them promptly.

Thanks again for reviewing our submission!


Reference:

[1]  Maeda T, Ukita N. Fast inference and update of probabilistic density estimation on trajectory prediction. International Conference on Computer Vision. 2023

[2] Melnychuk V, Frauen D, Feuerriegel S. Normalizing flows for interventional density estimation. International Conference on Machine Learning. 2023.

[3] Letzelter V, Perera D, Rommel C, et al. Winner-takes-all learners are geometry-aware conditional density estimators. International Conference on Machine Learning. 2024.
Thanks for the very helpful response. I raised my score to 5.
Thanks for your response and effort in revise paper. The response address almost all my concerns. Still, I have another question. What does the empirical Wasserstein distance means in evaluating the learned distribution?
Thank you for answering my questions.
Thank you for the clarifications and additional results! I will keep my score.