Estimating Potential Outcome Distributions with Collaborating Causal Networks
Authors: Tianhui Zhou
William E Carson IV
David Carlson
Venue: Accepted by TMLR ()

Abstract: Traditional causal inference approaches leverage observational study data to estimate the difference in observed (factual) and unobserved (counterfactual) outcomes for a potential treatment, known as the Conditional Average Treatment Effect (CATE). However, CATE corresponds to the comparison on the first moment alone, and as such may be insufficient in reflecting the full picture of treatment effects. As an alternative, estimating the full potential outcome distributions could provide greater insights. However, existing methods for estimating treatment effect potential outcome distributions often impose restrictive or overly-simplistic assumptions about these distributions. Here, we propose Collaborating Causal Networks (CCN), a novel methodology which goes beyond the estimation of CATE alone by learning the full potential outcome distributions. Estimation of outcome distributions via the CCN framework does not require restrictive assumptions of the underlying data generating process (e.g. Gaussian errors). Additionally, our proposed method facilitates estimation of the utility of each possible treatment and permits individual-specific variation through utility functions (e.g. risk tolerance variability). CCN not only extends outcome estimation beyond traditional risk difference, but also enables a more comprehensive decision making process through definition of flexible comparisons. Under assumptions commonly made in the causal inference literature, we show that CCN learns distributions that asymptotically capture the correct potential outcome distributions. Furthermore, we propose an adjustment approach that is empirically effective in alleviating sample imbalance between treatment groups in observational studies. Finally, we evaluate the performance of CCN in multiple experiments on both synthetic and semi-synthetic data. We demonstrate that CCN learns improved distribution estimates compared to existing Bayesian and deep generative methods as well as improved decisions with respects to a variety of utility functions.

Published in Transactions on Machine Learning Research (09/2022)
Estimating Potential Outcome Distributions with Collaborating Causal Networks
Tianhui Zhou
thuizhou@gmail.com com
Deportment of Biostatistics and Bioinformatics
Duke University
Durham, NC 27705, U.S.
William E Carson IV
william..crssmm@dduke edu
Deportment of Biomedical Engineering
Duke University
Durham, NC 27705, U.S.
David Carlson
david, www.saaarrnnnuuseeee edu
Department of Civil and Environmennal Engineering
Deportment of Biostatistics and Bioinformatics Deportment of Computer Science
Department of Electrical and Computer Engincering Duke Universily Durham, NC 27705, U.S.
Reviewed on OpenReview: https: openreview. net/ forun? id= FFey
Abstract
Traditional causal inference approaches leverage observational study data to estimate the difference observed (factual) and unobserved (counterfactual) outcomes for a potential treatment, known as the Conditional Average Treatment Effect (CATE). However, CATE corresponds to the comparison on the first moment alone, and as such may be insufficient in reflecting the full picture of treatment effects. As alternative, estimating the full potential outcome distributions could provide greater insights. However, existing methods for estimating treatment effect potential outcome distributions often impose restrictive or overly- simplistic assumptions about these distributions. Here, we propose Collaborating Causal Networks (CCN), a novel methodology which goes beyond the estimation of CATE alone by learning the full potential oulcome distribulions. Estimation of outcome distributions via the CCN framework does not require restrictive assumptions of the underlying data generating process (e.g. Gaussian errors). Additionally, proposed method facilitates cstimation of the utility of each possible treatment and permits individual-speciff variation through utility functions (e.g. risk tolerance variability). CCN not only extends outcome estimation beyond traditional risk difference, but also enables a more comprehensive decision making process through definition of flexible comparisons. Under assumptions commonly made in the causal inference literature, we show that CCN learns distributions that asymptotically capture the correct potential outcome distributions. Furthermore, we propose an adjustment approach that is empirically effective in alleviating sample imbalance between treatment groups in observational studies. Finally, evaluate the performance of CCN in multiple experiments on both synthetic and semi-synthetic data. We demonstrate that CCN learns improved distribution estimates compared to existing Bayesian and deep generative methods as well as improved decisions with respects to a variety of utility functions.

---

Published in Transactions on Machine Learning Research (09/2022)
1 Introduction
Personalized medicine requires estimating how an individual's intrinsic biological characteristics trigger unique and heterogencous responses to treatments (Yaddaní & Bocrwinkle 2015). Under the causal inference potential outcomes framework (Imbens & Rubin 2015) these trealment effects are characterized by the difference between the conditional expectations of potential outcomes under different treatment assignments, known as the Conditional Average Treatment Effect (CATEF). Since only the outcome for the assigned treatment is observed, estimating CATE often requires inferring the unobserved or counterfactual potential outcomes (Ding & Li 2018.. Recently, machine learning approaches have been developed for CATE, estimation including extensions of Random Forests (Wager & Athey 2018) and bespoke neural network frameworks Shalit el al. 2017 Shi et al. 2019). among other methods.
However, CATE does not necessarily align with optimal choices. In a decision theoretic framework, the optimal docision maximizes the expected utility function U(7) over the distribution of outcomes (Joyco 1999,, where a utility function U(7) quantifies an individual's preferences or tolerances to certain treatments with respects to 7, the variable over which these preferences tolerances vary, CATE represents special case of an identity utility function (()) y; however, more complex utility functions require alternative estimation approaches. One approach to learning decision maker is to make the utility function itself the objective function and optimize with respects to transformed outcomes, known as policy learning (Qian & Murphy 2011: Kallus & Zhou 2018). However, traditional policy learning methods require the utility function to be pre-specified, whereas utility functions typically vary between individuals (Pennings & Smidts 2003)- Defining proper losses with respects to complex decision criteria could be challenging from both a theoretical as well as computational standpoint. Morcover, if we also want to account for an individual's specific needs or have trained model generalize to new individuals with different utilities, traditional policy learning methods fall short. Previous efforts to estimate potential outcome distributions include Bayesian Additive Regression Trees (BART) (Chipman et al.) 2010 Hill 2011), variational methods (Louizos et al., 20170, generalized additive models with location, scale and shape (GAMLSS) (Hohberg et al. 2020), and techniques based on adversarial networks (Yoon et al. 2018 Ge et al., 2020). These techniques ofhen impose explicit or implicit assumptions about the outcome distributions (e.g. Gaussian errors) which may not align with the true data generating mechanism.
To address these issues, we propose novel neural network-based approach, Collaborating Causal Networks (CCN), to estimate the full potential outcome distributions, in turn providing flexible personalization and valuable insights with regards to specific individuals. CCN extends the structure of the Collaborating Networks (CN) framework (Zhou el al. 202T) to create a new causal framework that flexibly represents distributions. Under assumptions commonly made the causal inference literature, we show that CCN asymptotically captures conditional potential outcome distributions without having to explicitly make restrictive assumptions about the form of these distributions. We also introduce adjustment method that alleviates effects of imbalance between treatment groups, thus addressing a common confound that impairs model generalization. Empirically, this adjustment method improves point estimates, distribution estimates, and docision-making We sunmarize the contributions of our work as follows:
We propose a novel framework, Collaborating Causal Networks (CCN), to estimate full potential outcome distributions.
2. We characterize the asymptotic properties of CCN for estimation of smooth outcome distributions (i.e., distributions without point masses) by extending CN theory to the causal inference setting.
3. We propose adjustment scheme that combines -omain-invariant, propensity-speciee information and propensity stratification to alleviate the effects of treatment group imbalance.
4. We evaluate our framework with respect to different personalized utility functions in causal inference decision making, thus demonstrating the ability of CCN to address distinct user needs when such personalized utilities are made available.
5. We demonstrate empirically that CCN improves conditional decisions in a potential outcomes framework.
CCTTE is sometimes also referred the Individual Treatment Effectt or ITF Shalit al. 2017 Yao et al. 2018
2

---

Published in Transactions on Machine Learning Research (09/2022)
2 Problem Statement
2.1 Notations
Let X denote the random variable from which observed covariates I are drawn, X E X C RP. We assume a binary treatment condition with each unit or observation assigned a treatment T E (0. We choose the binary setup for clarity; however, the extension to multi-class problems is straightforward to construct under the proposed framework. We let (0) E R and (1) E R represent the continuous potential outcomes under the two treatment conditions, with (T)) representing the observed outcome. Lowercase letters with subscript. i are used to denote individual observations on each subject: {1)(1), +4, 24
F4,
A common goal in the causal inference literature is CATE estimation, T (z;; = (1) = - = xz)- Due to the presence of unmeasured features in practice, Vegetabile 2021 points out that T(x,) is an average taken over individuals with the same observed features. The main objective of this work is to address the problem of conditional causal inference and to develop a framework that addresses a broader range of objectives beyond CATE alone. Specifically, our goal is to use incomplete data to infer the full distributions over both potential outcomes in binary treatment scenario, Y(0) X and Y())X.. Successful estimation of these distributions will facilitate exploration of personalized needs and preferences through the introduction of various utility functions U()) (Dehejia 2005).
2.2 Utility Functions
While estimating potential outcome distributions and capturing uncertainties can be helpful to understand predictions in and of themselves, we also want to evaluate whether capturing these distributions leads to improved downstream decisions. We incorporate utility functions as part of our proposed CCN framework as a way to facilitate personalized decisions and recommendations. Utility functions aim to quantify the value of assigning a treatment to a patient based on traits or characteristics of said patient, while also taking into account personal preferences, tolerances, and needs. The structure of utility functions can be adapted different levels and degrees of comparisons. Below, we present different setups for quantifying utility with the following equations, each providing a different scope of comparisons:
(i) unified utility: - [=
(ii) treatment-specific utility: Ey-nt(()) (U1())) -anp)l)))_-nooo_aaa. En-ptr(o))
(iii) utility related to inherent features: B, By (U) x - E Bo [U((,
(iv) utility for personalized needs: Ex-eo(()=))Ui E Eo
Going in order from to ii the scope of comparison is broadened. A major distinction between iii and fiv) is that, for subjects with same features x, (iv) allows the utility to differ in accordance to personal preferences, which better aligns with the concept of personalized medicine. We note that setting U(r) = in ① results in the objective for CATE. Therefore, view the incorporation of utility more general framework.
Rather than estimating the potential outcomes and then calculating the utility, a model could be learned to predict the outcome the utility function in scopes ( fi.. and Mi... This is the strategy taken in "Policy Learning" (Athey & Wager 2021). However, transforming the outcomes may result in loss of information. For example, consider a threshold utility, U(7) Ly>0. Applying the transformation 1/ > CJ would transform continuous information into binary information, thus resulting in loss of information. On the contrary, estimating distributions not only enables to study personalized utilities like (iv but also retain all information in the system during training. Using estimated distributions to evaluate utility functions can be viewed as a two-step procedure or a plug-in estimator, which might be less efficient than direct optimization in some cases (Bickel & Ritov 2003). However, its flexibility and adaptability in studying various types of utilities makes it advantageous for practical purposes (Grunewalder 2018
may or may not depend on X.
3

---

Published in Transactions on Machine Learning Research (09/2022)
In short, utility functions provide way to flexibly take into account individual features, needs, and preferences when considering a possible treatment. Thus, the incorporation of utility functions with the CCN forms a framework that enables comprehensive treatment comparisons according to the above considerations.
2.3 Causal Assumptions
Like many causal methods, CCN relies on the assumptions of strong ignorability and consistency to estimate the potential outcome distributions when a single treatment outcome only observed for each datum Rooeenaamm & Rubin 1983; Hernan & Robins, 2020). Strong ignorability and consistency are characterizod via the following assumptions:
Assumption 1 (Positivity or overlap). AX E C C RP, the probability of assignment to any trealment group is bounded away from zero: 0 < Pr(T = 1|X = z) < Vx such that p(x) >
Assumption 2 (Consistency). The observed outcome given specific trealment is equal to its potential outcome: YJ,, X (T)(T X.
Assumption 3 (Ignorability unconfoundedness)) The potential oulcomes are jointly independent of the treatment assignment conditional on X: (Y()). 1 TXX.
3 Collaborating Causal Networks
The CCN approach approximates the conditional distributions of Y X0) and YL)) CCN uses two- function framework based on the Collaborating Networks (CN) method (Zhou et al. 2021). We choose to exlend CN to the causal setting because it automatically adapls to different distribution families, including non- Gan-Gansiian distributions, and we believe that this flexibility and robustness will result in more reliable estimates of the different utilities detailed in Section 2.2] We first give an overview of CN, then present the CCN framework in detail, and finally introduce our adjustment strategics. Proofs of all theoretical claims are provided in Appendix A.
3.1 Overview of Collaborating Networks
Similar to Generative Adversarial Networks (GAN) (Goodfellow et al. 2014), CN is based jointly lcarning two functions, both of which are approximated by neural networks. However, unlike GAN where the two networks have opposing objective functions that force the networks to "compete" against one another, in CN the two networks form a collaborative approach in that they work towards the same goal from different angles. Specifically, CN estimates the conditional distribution, Y|X, using two neural networks: a network g(y.z) to approximate the conditional CDF, Pr(Y < y| X), and a network f'g, x) to approximate Information sharing enforced by the fact that the CDF and its inverse are an identity mapping for any quantile qE g(ʃ(g, x)) = The networks form collaborative scheme with their respective losses,
inverse.
g-loss: : Eqn.. 9(
(1)
f-loss : Ey,, [(q x)(
x),
(2)
The quantile q is randomly sampled (e.g., ~ Unif(0, 1)). ((.,)) represents the binary cross-entropy loss. The parameters of I and g are only updated according to their respective losses. When the objective the simultaneous minimization of equations and a fixed point of the optimization is at the true conditional CDF and its inverse (Zhou et al. 2021 In this framework, g(-), the function used to approximate the conditional CDF, is considered the main function. f(-) is viewed extensively search the full outcome space to help g(-) acquire information about the distribution function over the full relevant. space. On the other hand, the optimality of f(-) depends on an optimal g(-), as acquires information solely through inverting g(-). Zhou et al. (2021) show that f(.) can be replaced by other space searching tools, including prefixed uniform distributions, which may result a minor performance loss favor of case of optimization. This case of optimization can be especially beneficial when combining CN with other regularization terms. We show depiction of the CN framework from Zhou et al. (2021) in Figure
auxiliary function whose job is to
4

---

Published in Transactions on Machine Learning Research (09/2022)
q
x)(q.x)
(q))))
fa(q.x)
f(gx)
f-loss g-loss
x
Epero/ Popul
(a) f-network.
(b) Training scheme.
Figure 1: Diagram of the CN framework. [(a) depicts training for prediction of a conditional quantile Ō(g. xx)
directly. The dashed arrow coming from the component representative of Ø(g. x) indicates that the objective function does not produce a useful gradient. [(b] shows the full CN framework, where g(.) and f()) are
trained jointly to learn the CDF and conditional CDF.
In this work, we focus on extending the function g(-) and the g-loss for causal inference. We replace flq. x) with variable general form of a space searching tool, such as uniform distribution covering the range of the observed outcomes. We thus simplify the g-loss to,
g-loss Eay, ((,
(3)
Thus, equation 8 is generalization of the g-loss introduced in Zhou et al. (2021). In practice, equation 3 replaced with an empirical approximation.
3.2 Collaborating Causal Networks Formulation
Following the taxonomy of Künzel et al. 2019). CN can be extended either as an "S-learner," " where the treatment label is included as an additional covariate and is thus more scalable for multiple treatment groups, or as a "T-learner," " where the outcome under each treatment arm is estimated using separate functions or networks. Below, we detail the T-learner extension the CN. Based on equation 3 we define two functions paramctorized by neural networks, go(-) and 91(+), which approximate the untreated group conditional CDF, Pr(Y x| 0), and the treated group conditional CDF, Pr(Y < <, = 1), respectively. Combining the losses of these two treatment groups results in the following loss function:
g-loss"
By(t=0), [<(Lp(t=0)<=: 90(2, + [)(t=)), (1
(4)
We refer to the optimization of this objective as the CCN framework. Under Assumptions 1, 2, and 3, the fixed point solution and consistency of CCN framework still hold regardless of the treatment group imbalance. To summarize, Assumption / connects the conditional distribution Y|X, X,7 T to the potential outcome distribution (T))X, in each treatment space with density functions: p(x[T = 0) and p(r(T = 1). Assumption guarantoes that despite the treatment group imbalance, p(x[T 0) * p(x(T = 1) * p(x), cach treatment space can still have sufficient coverage of the full space p(a) given enough samples. Lastly, Assumption 3 generalizes the potential outcome distributions from each space Y (T)) T to the full space (T)) removing its conditioning on the treataentt label T. Given our assumptions,
state:
Proposition 1 (Optimal solution for go and 91). When the support of the outcomes is subset of the support of 2, or as covers the whole oulcome space, the functions go and 9ı that mónimize g-loss" are oplåmal when they are equivalent to the conditional CDF of Y (0)]X = and Y (1)) = =x, Vr such that p(x) > 0.
Proposition 2 (Consistency of go and 91). Assume the ground truth CDF functions for T E {0,1} are Lipschitz continuous with respect to both the features X and the potential outcomes {Y Y()) and the support of the outcomes is a subset of the support of Denote the ground truth functions as gû and gi- As n + co, the finite sample estimators øö and of have the following consistency property: (49) >r 0: dgg gi) + under some metrics d, such as the LI norm.
Taken together, these propositions state that the CDF estimators go and 91 inherit large sample properties when estimating potential outcome distributions under the CCN framework. Proofs are provided in Appendix
5

---

Published in Transactions on Machine Learning Research (09/2022)
T Specific
Φw
D()
Wass-loss
S: (: (()))))))))))))))))))
9o Go
91 G..)
,o(..) 91 (..))
Invariant
S: s: (x), (x))()))aa())ee)))) (()))))())))))((())))))
g losspro
((U)
Y(1)
X
(a(.)
e(.)
Assign-loss
(a) Causal Diagram
(b) FCCN Training
(c) Infer Distribution
Figure 2: 2(a) depicts how two sources of information could impact Y. 2(b) visualizes the FCCN network.
[(c) depicts how the trained go and ø1 functions can be used to sketch the underlying CDFss of (0) = z and Y (1) æ.
Though consistency is only claimed for functions belonging to the Lipschitz family, we consider this to be reasonable assumption as Lipschitz functions can provide good approximations to a wide range of functions (Sohrab 2003); especially considering the boundedness of CDFs between (0, 1). In practice, our approximation functions are implemented as neural networks, which allows for additional flexibility. Enforcing Lipschitz constraint on a neural network can be done through a variety of techniques, such as by limiting the weights to finite value via weight clipping (Arjovsky et al. 2017). ). In practice, the learned functions were smooth and we did not find it necessary to enforce such a constraint.
3.3 Adjustment for Treatment Group Imbalance
One obstacle in observational studies is treatment group imbalance, where the distributions of the covariate spaces p(x[T T = 0) and p(z[T = 1) differ. This creates two major issues for causal predictions: poor generalization over different treatment spaces and confounding effects. In the asymptotic regime (Proposition 2), this imbalance is less problematic since overlap (Assumption 1) ensures all regions with positive density will eventually be densely covered with samples. However, for finite number of samples this imbalance hurts inference. Therefore, we pair our method with an adjustment scheme to address this obstacle.
We encode the covariates into a space that simultaneously facilitates generalization between spaces and adjusts for confounding effects. We represent this new space as 5 = wow (2), &A e(dx()))) Two neural networks, (-(+): : I + Raw and Φa(+)) IP + RaA, transform the input æ E IP into qw and ga-dimensional latent spaces. In Figure [[a] these spaces are graphically depicted as the two sources of information that can impact the outcome. (-) enforced to be invariant to trealment groups (domain-invarian) and the (a(-) is enforced to learn specific differences between the treatment spaces (domain-specific) (Shalit al., 2017 Ben-David et al. 2010). The invariant component ow(+) finds a more balanced representation between spaces that benefits generalization, while the component @A(-) controls for confounding effects through learning the treatment. assignment mechanism. We denote the g-loss" trained on this representation space as g-loss pro. Additionally, a neural network e(.) : 1894 + [o, 1] uses the output of (a(-) to predict the propensity of treatment assignment, Pr(T = 1|X = æ). While e()a(x)) may seemn redundant given (x)x,, including (z)) in the covariate space encourages propensity score stratification (Hahn al. 2020).
Domain-invariance is encouraged in the space ow through penalty on the Wasserstein distance (Wass-loss) between the two treatment arms, and (px()),e(()) is encouraged through cross-entropy loss on the assigned treatment labels (Assign-loss), as detailed in Sections 3.3.1 and 333.2 respectively. The representations (ow (-), (.).)) e(-)) with their respective losses are incorporated as regularization terms. This framework is depicted in Figure 2 and the full loss can be expressed as the sum of g-loss pro-
L(go 91, Φw, (W, g- loss 91, ow, &.,e) e) + (Wass-loss((ww) (Assign-loss(AA.e))
(5)
The hyperparameters a and /3 adjust the relative importance of their respective losses during learning. Empirically, we have found performance to be robust to selection of these hyperparameter values. We call this full adjustment CCN (FCCN). In summary, in FCCN two additional functions, Φw(-) and (a(+), are learned
6

---

Published in Transactions on Machine Learning Research (09/2022)
extensions of the existing CCN framework to learn domain-invariant and domain-specific representations of the treatment spaces, respectively. These functions are parameterized by neural networks in practice.
3.3.1 Wass-loss to Alleviate the Covariate Space Imbalance (oomain-Invariant)
The introduction of Wass-loss motivated by CounterFactual Regression (CFR) implemented with the Wasserstein distance (Shalit et al., 2017). CFR is a causal estimator based on representation learning Φw() R Row The goal to find latent representations where p(o = 1) and p(ów (x) 0) are more balanced domain-invarianl than the original space. We use the Wasserstein-1 distance, which represents the total "work" required to transform one distribution into another (Vallender 1974). Through the Kantorovich-Rubinstti duality (Villani 2008), this distribution distance is,
W (Paa Pb) sup 21 Ez [D(x)] Ex [D(x)]
|D|| E represents the family L-Lipschitz functions. We approximate this distance by adopting the approach of Arjovsky et al. (2017) This turns into the following regularization term,
Wass-loss max to 1-1) (ow (x)] o) (P(w
(() parameterized by a small neural network, and the Lipschitz constraint on is enforced through weight clipping. The Wass-loss penalizes differences in the latent space between the treatment and control group, which in turn improves generalization between groups.
3.3.2 Assign-loss and Propensity Stratification for Confounding Effects (Domain-Specific)
The introduction of Assign-loss is inspired by Dragonnet (Shi et al. 2019), deep learning method that learns a latent represenlation for treatment assignment mechanism to control for confounding effects. It is defined as a binary cross-entrops loss on prediction the treatment assignment label with e(oA(z)),
Assign-loss: -E,, (too(((())())) + (1 - t) log(11 - c(pa(x)))
We additionally incorporate the estimated propensity (x)) directly into our covariate space, an adjustment. term propensity stratification (PS). Using propensity scores in the predictive model is form of continuous stratification to reduce the bias of estimation by facilitating information sharing within sample strata created by (x)) (Hahn et al. 20200.. PS corresponds to learning propensily scores via a cross-entropy objective when estimated propensity scores are included in the original feature space r. As a second step, this new feature space (Ir, (x:) is used to train functions g0(.) and gl(.).
4 Related Work
CATE Estimation. A common approach to CATE estimation with machine learning is matching, which
identifies pairs of similar individuals Rubin 1973 Rosenbaum & Rubin
1983 Li & Fu 2017 Schwab et al.)
2018 This idea motivates many free-based methods that identify similar individuals within automatically
identified regions the covariate space (Liaw Wiener 2002 Zhang & Lu 2012 Athey Imbens 2016
Wager & Athey 2018). Deep learning methods have also been proposed for CATE prediction. These include
networks with additional loss terms encourage a treatment-invariiat space (Johansson et al. 2020
2016
Du et al. 2021 and networks that explicitly encode treatment propensity information (Shi et al. 2019
Representation learning can be combined with weighting strategies to enforce covariate balance (Assaad et al, 2021 Hassanpour & Greiner 2019 2020a). Despite some resemblance to previous methods in learning more balanced fealure embeddings, this manuscript performs an extensive ablation study to support and justify that each component of our proposed method leads to increased robustness. Additionally, these methods largely focus on CATE estimation only (with some exceptions noted below), which may be insufficient reflect the full picture of different treatment regimes (Park et al., (2021). In contrast, CCN estimates full risttiiuioons to assess the utility and confidence of decision.
7

---

Published in Transactions on Machine Learning Research (09/2022)
Potential Outcome Distribution Sketching. Multiple Bayesian methods have been proposed for the estimation of outcome distributions, including Gaussian Processes (Alaa & van der Schaar 2017). Bayesian dropout (Alaa et al. 2017), and Bayesian Additive Regression Trees (BART) (Chipman et al. 2010). BART has gained popularity in recent years and has been the focus of further modifications, including variations to account for regions with poor overlap (Hahn et al. 2020). However, Bayesian met mothods can suffer under model misspecification (Walker 2013), such as mismatch between the assumed and true outcome distributions. Bayesian methods have also been integrated with deep learning, such as the Causal Effect Variational Autoencoder (CEVAE) and its extensions (Louizos et al., 2017 Jesson et al., 2020]; hybrid architectures are sometimes adopted to account for certain types of missing data mechanisms Hassanpour & Greiner 2020b
Frequentist approaches can achieve flexible representations of distributions. A well-known adaptation is the Generalized Additive Model with Location, Scale and Shape (GAMLSS), which estimates the parameters for a baseline distribution with up to three transformations given a specific distribution family (Briseño Sanchez ct al. 2020: Hohberg et al.) 2020). The CDF may also be estimated nonparametrically by adapting density inti methods such as nearest. neighbors (Shen 2019), which is less reliable in areas of treatment group imbalance. GAN-inspired methods, including GANITE (Yoon et al. 2018), can also learn non- Gaussian outcome distributions. There is emerging literature on conformal prediction in treatment effect estimation
(Chernozhukov et al. 2021 Lei & Candès 2021). However, conformal prediction only learns specific level of coverage and coverage probabilities are proven for marginal rather than conditional distributions.
Quantile Regression. Quantile regression is an established method to estimate uncertainty and variability of continuous outcomes (Koenker & Hallock 2001 Meinshausen & Ridgeway 2006 Tagasovska & Lopez- 2019). Aspects of CN and the proposed CCN framework both bear resemblance to quantile regression, with the f network described in Section 3.1 learning quantile regression function. However, the learning methodology is different: canonical quantile regression uses a tilted absolute value loss function, whereas CN focuses on learning the conditional CDF, 9(-), and uses different loss functions. These modifications are necessary when using deep networks, as applying the tilted absolute value loss function in a deep network can result in drastically underest imated variability (Zhou et al. 2021). Quantile regression has been studied in the context of causal inference (Chernozhukov & Hansen, 2006; Tagasovska al. 2020 Sun et al.) 2021). However, these works looked at the population-level effects of endogenous variables treatments on potential outcomes rather than individual-level effects and do not consider incorporation of personalized utility functions. Furthermore, many of these studies are not as flexible as the proposed approach and would not work with more complex estimators, in contrast to the deep learning approach proposed here. The comparisons made in the original CN work suggest that the CN approach is quite robust compared to quantile regression losses (Zhou el al. 2021j. Regardless, quantile regression can still be beneficial when simpler estimator forms hold or when learning is data-limited.
Policy Learning and Utility Functions. A key purpose of estimating the conditional causal effect is to serve as information in decision making processes. A common strategy is policy learning, where the policy is expressed as a function of the feature space and learn a policy that optimizes a pro-defined utility [Beyellimerr & Langford 2009; Qian & Murphy 2011 Bertsimas et al.) 2017 Kallus & Zhou 2018). This involves transforming observed outcomes in accordance to the pre-defined utility as the new objective for optimization. Traditionally, utilities studied in policy learning are linear transformation of the potential outcomes, which can be described as the difference between the benefit and cost (Athey & Wagerr 2021]. p.
However, when we are presented with threshold based utilities, transforming the observed outcomes may be subject to information loss (e.g., binarization of continuous variable greatly reduces information) according to the Data Processing Incquality (Beaudry& & Renner 2012). Additionally, policy learning can only study pre-specified utilities, whereas a trained potential outcome distribution estimator can serve as a one-stop shop to explore a myriad of personal utility functions.
Off-Policy Evaluation. Estimating utility differences and choosing optimal decisions is closely related to the concept of off-policy evaluation Thomas & Brunskill 2016). We note that the policy learning strategies above could be considered as off-policy evaluations they evaluate decisions with respects to utility functions. In reinforcement learning, off-policy evaluation is frequently used to estimate the reward, which is analogous to the utilities examined in this paper. Off-policy evaluation methods have been developed to estimate distributional properties and uncertainties over the return (cumulative reward), including first estimating
8

---

Published in Transactions on Machine Learning Research (09/2022)
the confidence interval on the estimator (Thomas et al. 2015). These methods have been recently extended to develop estimators for the complete distribution of the return through Universal Off-Policy Evaluation (Chandak et al. 2021). However, these methods differ from the proposed approach in the overall goal, as off-policy evaluation estimates the overall return and is more similar to estimating the average distribution of difference of utilities, allowing off-policy evaluation methods to take advantage of repeated observations for their theoretical proofs. In contrast. CCN estimates potential outcome distributions for individual data samples, allowing for more personalized estimates and treatment recommendations.
5 Experiments
We follow established causal literature and use semi-synthetic scenarios to assess estimates of conditional causal effects. First, we evaluate causal methods using the Infant Health and Development Program (IHDP) dataset (Hill 2011), where the outcome of each subject is simulated under standard Gaussian distribution with a heterogencous treatment effect. The IHDP dataset represents an idcal scenario for methods with Gaussian assumptions, including BART and CMGP. The second dataset used is derived from data collected in a field experiment in India studying the impact of education (EDU). In this case, we synthesize each individual outcome with heterogeneous effect and variability using a non-Gaussian distribution. The semi-syuthetic procedures are briefly outlined in Section 5.3 with full details in Appendix B Additionally, we provide evaluations number of difffernt synthetic outcome distributions to compare methods under different. sccnarios, including multi-modal distributions and several other different distribution familics.
We include our base approach, CCN, and its adjusted version, FCCN, in experiments. We compare them to existing approaches that estimate potential outcome distributions including Bayesian approaches (BART
Hill 2011 CMGP (Alaa & van der Schaar 2017). CEVAE (Louizos et al., 2017)), a frequentist approach (GAMLSS [Hobberg et al. (2020]), and a GAN-based approach (GANITE (Yoou al. 2018)). Causal Forests (CF) [Wager & Athey| 2018) is benchmarked for con-distribution metrics as a popular recent CATE-only method. GAMLSS's s flexibility and strength in estimating distributions is dependent
close match to the true distribution familics from which data is generated, which is rarely known in practice. Therefore, we evaluate GAMESS by providing the closest possible distribution to the distributions underlying data generation, meaning that GAMLSS is provided more information than any other method. We also benchmark the proposed approaches against policy learning approaches on decision- telisionmmaiing metrics. To fully understand the impact of the various adjustments in FCCN compared to CCN, we run ablation studies and evaluate the performance over a suite of hyperparameters for tuning the adjustment.
Details for all model implementations, including model architectures and hyperparameter selection procedures, are provided in Appendix C We use neural kewwrrk-based architecture for CCN and FCCN, but also detail an alternative structure that enforces a monotonic constraint in Appendix D Model and experiment code is available at https: httpah/w.ooaaarooooaaallooaalaaa
-etworks..
5.1 Metrics
We evaluate mean outcome estimates via the Precision in Estimation of Heterogencous Effect (PEHE) metric and the full distribution by estimating the log-likelihood (LL) of the potential outcomes. We regard LL as the key evaluation metric since it evaluates estimation of full distributions. In addition, we evaluate how well each method makes decisions by using Area Under the Receiver Operating Characteristic Curve (AUC) for chosen utility functions to show that improved distributional estimates lead to improved decisions.
Precision in Estimation of Heterogeneous Effect (PEHE); We adopt the definition of PEHE defined in Hill 2011). For unit i and covariates I;, we have the CATE and its estimates as r(x;) = E[ = 2, E E X xr] and f(x;). PEHE quantifies the distance between CATE and estimates according to the following:
PEHE =
N Σ (f(())) - (2x)))²
9

---

Published in Transactions on Machine Learning Research (09/2022)
Table 1: Quantitative results on IHDP (Section 5.2 The mean and standard error of each metric are reported. CMGP attains top marks in all metrics due to modelling assumptions that align with the true data generating mechanism. FCCN outperforms CCN in all metrics with statistical significance, and remains competitive with CMGP on LI and both AUC metrics. GANITE is only used to estimate the CATE as it is challenging to train on this small dataset according to Yoon et al. 2018
Metrics PEНE
CCN
FOON
GAMISS
GANITE
BART
CMGP
CEVAE
CF
1.59 = 16
1.13 = .14
3.00 = 39
2.40 = 40
2.23 = 33
.714 + .087
2.60 + 10
3.52 = .57
LL
-1.78 : 02
-1.62 1 02
-2.34 = .13
-1.99 + 08
-1.51 1 .01
-2.82 1 .08
NA
AUC (linear)
.25 + 011
942 4 010
.30 - 0100
.723 + 017
.923 + 000
.957 + 0012
.23 = 008
.896 +: .009
AUC (threshold)
.933 + .01
.95 + .10
.25 + .010
.917 + .009
.955 0012
.564 + .010
NA
Log Likelihood (LL): Log likelihood (LL) mcasures how well each method models potential outcome distributions. Typically, LL is calculated by evaluating the probability density function (PDF) functions at the observed points. However, closed-form distributions are not directly available for sampling based approaches such as GANITE and variants of CCN. Instead, we approximate LL by calculating the probability on a neighborhood of the realized outcome y, By.e= = (y - c,y + c), where e is a small positive value:
N Σ 2N ==
LL
log(Pr|X(()) E Bo, =x:)).
Asymptotically, the true distribution dominates in this evaluation, and this can be shown under the crilerion of the Kullback-Leibe Leibler divergence (Kullback & Leibler 1951] as N + 0o and e e 0. We define e = 0.5 for IHDP and <= 0.2 for EDU to adjust for the scale of the outcomes.
Area Under the Curve (AUC): As mentioned in 2.2 personalized decisions can be described by the quantity: r(I;, Uo,,, U1,) = (===pl(())xxx:(UU.... =)] E ==xpp (7,))) The optimal decision is based on the sign of this contrast Ui...) which is regarded the true label. Using the estimated contrast f2,, Uo,s, a1) the decision score, can estimate the AUC by varying the decision threshold.
5.2 IHDP
The Infant Health and Development Program (IHDP) dataset is derived from results of a randomized experiment, and the data is transformed to be more observational study-like by removing non-random portion of the data belonging to the treatment group. We use response surface B as described in Hill 2001 for estimating heterogeneous trealment elfects. The study consists of 747 subjects (139 in the treated group) with 19 binary and 6 continuous variables (I; E RR25. We use 100 replications of the data for out-of-sample evaluation by following the simulation process of Shalit et al. (2017).
The quantitative results are summarized in Table T CMGP achieves top marks, as its Gaussian noise model aligns with the true data generating mechanism in this dataset resulting in a significant advantage given limited data. This contrasts to later experiments when its strong assumptions are invalid. Overall, CCN is either competitive with or outright outperforms other methods in estimating both mean and distribution metrics. The ad advantages of combining the proposed adjustment strategy is evident as FCCN improves over CCN by a clear margin. It is worth noting that LL calculated under the ground truth model is -1.41, demonstrating that FCCN is also effective in capturing the true distributions with an LL estimate of -1.62 + 0.02 that is only behind the LL estimate made by CMGP.
7 4 We evaluate two different utility functions: a linear utility - and a non-linear utility with U0, xx.) ([))))))) and U, æ;;
(X) x;;])) +4), the ATE for surface B is 4 (Hill 2011 Table D shows AUC (linear) and AUC (threshold) corresponding to results under these two utilities. These results demonstrate the improved distribution estimates made by CCN compared to many techniques contribute to more accurate decisions, despite the fact that a homoskedastic Gaussian distribution is well matched to the specifications in BART and GAMLSS, CMGP and FCCN are very similar in their decision quality, despite CMGP's assumptions being well-matched to the data generation procodure. The evaluations of all these metrics can be regarded as a plug-in process for distribution estimators. We notice that PEHES of these distribution methods may not be as competitive as some of the state-of-the-art CATE estimators
10

---

Published in Transactions on Machine Learning Research (09/2022)
Table 2: Quantitative results on the EDU dataset (Section 5.3 outperforms CCN with statistical significance.
FCCN achieves top marks on all metrics and
Metries
CCN
FOCN
GAMLSS
GANITE
BART
CMGP
CEVAE
CF
PEIEE
.392 + 049
.296
.042
3344 + .033
1,253 + 1811
.334 + 042
2,087 + .015
1.911 + .351
1,022 + .051
LL
.2.178 + 024
- .022
2.250 1: .025
5.0922 + 596
2.443 1. .063
-3.533 .008
3.538 +: .055
NA
AUC
.933 = 026
.953 .114
.911 + .100
.700 = 053
.906 + .015
.875 = .066
.62 + .039
NA
Shalit et al. 2017 Hassanpour & Greiner 2020a). as optimizing the mean is simpler and has exact match with this metric. The strategy of estimating distributions can be flexibly adapted to various comparisons with trained models.
Comparison Policy Learning. We next compare the proposed approaches to a policy learning approach, specifically policytree (Sverdrup et al. 2021), with full details in Appendix F Policy learning is limited to pre-speeified utility functions, so we set up two scenarios: one with lincar utility (Uu(r) 7 U1(7) 7 4), and one with threshold (binary) utility (Uo(7) 1 E)) U1(7) = 1 {)>E))())-- Since policytree only outputs its predicted optimal treatment, we compare on accuracy (predicted vs true optimal treatment). On the IHDP dataset, FCCN performs well on both linear and threshold utilities, achieving accuracies of 88.6% and 87. 77% accuracy, respectively. However, policytree's accuracy drops from 76.9% to 57.6% when we switch from a linear a threshold utility, signifying how much information lost by binarizing the outcomes.
5.3 EDU
The EDU dataset is based on data collected from randomized field experiment in India between 2011 and 2012 (Banerji et al. 2017: 2019). The experiment studies whether providing a mother with adult education benefits their children's 's learning. We define the binary trealment as whether a mother receives adult education and the continuuus outcome as the difference between the final and the baseline test scores. After the preprocessing described in Appendix B the sample size is 8,627 with 18 continuous covariates and 14 binary covariates, 24 R
We create semi-synthetic case over the two potential outcomes according to the following procedures. We first train two neural networks, fgo (.) and fú, (-), on the observed outcomes for the control and treatment groups. The uncertainty model for the control and treatment group are based on a Gaussian distribution and an exponential distribution, respectively. The dependence of the outcome uncertainties on different. distributions is intended to showcase the abilities of models that are able to adapt to different distribution families (c.g., CCN and FCN). We represent m, as an indicator of whether the mother has received any previous education, we hypothesize the variability is higher for the mothers not educated previously" Then, the potential outcomes synthesized as,
(0) fg (z.) + (2- (0).0), Y(1)a
(z:) + (2 m.) ex)())))
In this experiment, treatment group imbalance comes from two aspects of our data generation procedure. The first is from a treatment assignment model with propensity Pr(T; 1/x)) = =I +1+exp(--]] ()) where we assign large coefficients in 3 to create imbalance. The other is from truncation, as we remove well-balanced subjects with estimated propensitics in the range of 0.3 < Pr(T, = 1|x;) < 0.7. We keep 1,000 samples for evaluation and use the rest for training. The full procedure is repeated 10 times for variability assessment. The utility function is customized for each subject to mimic personalized decisions. For subject i, U0,(1,,z)) 1(> 0z), and 27; [(> > e + 1-11- mm) where i; ~ UO, 1.5). The interpretation of this utility is that different. mothers have different expectations of their children's 's improvements with threshold u, For the mothers without previous education, their expectations are higher by 1. This design coincides with the expectation that education should have a positive effect on outcomes in exchange for a finite cost, as we would only invest in the intervention if a positive return was expected.
The variable indicating whether the mother has received previous education is included as a predictor in X; however, models are aware how the variable relates outcome variability and must learn this relationship. We simply assign its own variable for ease of notation when describing how potential outcomes for the EDU semi-synthetic. data are synthesized.
11

---

Published in Transactions on Machine Learning Research (09/2022)
Estimated CDF
Estimated CDF
Estimated CDF
Estimga
0.6
0.6
0.6
0.6
0.4
True Y(o)
9.4
0.4
0.4
True Y(I)
-1
3
(a) OCN
(b) FCCN
(c) GAMLSS
(d) GANITE
1.0
1.0
CDF
Estimated CDF
0.6
0.6
Eatimated
0.6
0.4
0.4
0.2
0.2
-1
3
22
-1
2
-1
0
(e) BART
(l) CMGP
(x) CEVAE
Figure 3: Visualization of estimated vs. true CDFs for a random EDU dataset sample (Section 5.3 FCCN estimates are the only ones that closely follow the true control and treatment group curves with high fidelity.
Results from this experiment are summarized in Table 2 CCN-based methods demonstrate their ability to flexibly model different distributions with FCCN providing improvement over CCN and top marks in all metrics. Figure B shows that FCCN is the only method able to recover both Gaussian and exponential iont with high fidelity, which we believe contributes to its top performance. Making an optimal decision is highly dependent on how close a given method's estimated distribution aligns with the true values and all relevant heterogeneity. Thus, the AUCs follow their respective LLs with FCCN producing the best performance in both. Unlike the IHDP data, CMGP's modeling assumptions do not match the generative procedure, and its performance suffers. CEVAE makes two model misspecifications: (i) it assumes homogencous Gaussian error, whereas the real outcomes arise from heteroskedastic Gaussian exponential distributions, and (ii) it decodes the continuous covariates into a Gaussian distribution. Thus, CEVAE captures the marginal distributions well (visualized in Figure 1(h) in Appendix E but does not provide helpful personalized suggestions.
Next, we randomly draw an EDU data sample and compare the estimated CDFs against the true CDFs Figure (see Figure [S] for additional samples). We find that CCN-based approaches are capable of faithfully recovering the true conditional CDFs, whereas the other methods have gaps in their estimation. GAMLSS is accurate the control group but not the treatment group. This is partially due to gamlss package Stasinoooolos et al. 2021p not supporting the exponential distribution with location shift, so skewed normal is chosen as the closest reasonable substitute. Overall, GAMLSS is flexible but requires precise specification on a case-by-e by-case basis, whereas CCN can robustly use the same approach. In our experiments for GAMLSS, we must choose very close distributions and limit. the uncertainty to the relevant variables or the package does not converge. Similarly, CMGP estimates the control CDF well as its Gaussian assumptions align with the underlying distribution of the control group. However, is not able accurately capture the treatment CDF.
Lastly, we assess whether the heteroskedastiicty of the outcomes is captured. The combination of M 0, and T= 0, 1 produces four uncertainty models. We visualize the predictive 90% interval widths in Figure FCCN captures the bimodal nature of the interval widths. In contrast, GANITE only captures a small fraction of the difference between low and high variance cases. Both CEVAE (Louizos et al. 2017) and BART Hill 2011) fail to capture the heteroskodasticcty and are not shown. For GAMLSS, we explicitly feed its uncertainty model with only T and M for to converge effectively. It does not produce variability in interval widths. Overall, CCN and its variants produce higher quality ranges.
12

---

Published in Transactions on Machine Learning Research (09/2022)
Table 3: The estimated LL under different simulated distributions (Section 5.4.2) GAMLSS with the true family and heteroskcdastic Gaussian, respoctively.
and represent fitting
True Value
CCN
FCCN
BART
GAMLSS!
GAMLSS2
Gumbel -2.87
-3.67 + .22
-3.56 .02
-3.92 .06
-3.67 + .02
-3.90 + .05
Gammna
-3.17
-3.83 + .04
-3.74 + .06
-3.97 + .02
-3.77 + .02
-3.95 + .02
Weibull
-2.87
-3.41 + .02
-3.32 + .03
-3.86 + .12
-3.32 + .04
-3.71 + .09
Estimated Interval Width
4.0
T-O, M-1
Estimated Interval Width
4.0
T-0, M-1
Estimated Interval Width
T-0, M-1
Tw1, M=O
M-0
4.0
T-1, M=O
Tw1,
3.0
3.0
3.0
2.0
2.0
2.0
1.0
1.0
2.0 3.0
4.0
2.0
3.0
1.0
2.0
3.0
Interval
Interval
(a) CCN
(b) FCCN
(c) CANITE
Figure 4: Ability of methods to capture heteroskodastic outcomes of EDU data (Section 5.3 are demonstrated by plotting the estimated versus true 90% interval widths given four different combinations of T and M. GANITE reflects the main trend of how uncertainties change with T and M. FCCN can clearly discern the four scenarios by aligning its estimated interval widths and the true interval widths. Additional visualizations are provided in Appendix
5.4
Additional Comparisons and Properties
In this section, we include several additional experiments to reflect other properties of CCN and its adjustment to demonstrate their advantages in estimating the potential outcome distributions.
5.4.1 Estimating Multimodal Distributions
We choose to extend CN to estimate potential outcome distributions due to its ability adapt to variety of outcome familics. We demonstrate this with qualitative experiment that uses synthetically-generate multi-modal data distributions. Many of other methods, including CMGP, GAMLSS, BART, and CEVAE, cannot capture it with their existing structure, whereas CCN, FCCN, and methods like GANITE, can. To succinctly summarize this experiment, only CCN and FCCN naturally adjust to the multi-modal space, as shown briefly in Figure 5 GANITE is aware of the mixtures but does weight them well. The other algorithms are not able to capture the mixture model. The data generating process is described in Appendix H
5.4.2 Additional Distribution Tests
Next, we skotch out how diffcrent methods work on a varicty of known outcome distributions, including Gumbel, Gamma, and Weibull distributions, with full details in Appendix G The results Table 3 similar the previously presented semi-synthetic cases, where CCN straightforwardly adapts these distributions and FCCN provides additional improvements. In fact, FCCN even slightly outperforms GAMLSS even when GAMLSS is provided the true outcome distribution. GAMLSS does not come close to rivaling the performance of CCN when provided a flexible but not perfectly matched outcome distribution.
5.4.3 Sample Size and Convergence
Proposition 2 suggests CCN can asymptotically estimate the optimal value given a large sample size. We create synthetic example with the logistic distribution to visualize it. The data generating process is
13

---

Published in Transactions on Machine Learning Research (09/2022)
0.25
True Y(o)
True Y(3)
0.25
True Y(o)
True Y(3) Est Y(3)
Est Y(o)
Est Y(3)
Est Y(O)
9.15
Densit
0.15
9.30
0.30
0.65
-5
5
10
-5
10
(a) CCN
(b) FCCN
True Y(o)
True Y(U Est Y(1)
True Y(o)
True Y(1)
Y(O)
Y(0)
Est Y(1)
Density
0.20
Dengity
0.20
9.35
0.15
0.50
0.00
(c) GAMLSS
(d) GANITE
0.25
True Y(o)
True Y(3)
0.25
True YIo)
True Y(1)
Est Y(G)
Est Y(3)
Est Y(O)
Est Y(1)
Density
Dendity
0.36
0.15
9.10
0.00
0.85
0.05
10
10
(e) BART
(f) CMGP
0.25
True Y(O) Est Y(0)
True Y(1)
Est Y(1)
0.20
0.15
0.30
0.00
(g) CEVAE
Figure 5: Visualization of each method's estimated density of the potential outcomes for synthetically generatod multi-modal data (Section 5.4.1) Visually, CCN and FCCN outperform other methods in estimating the multimodal data. It appears GANITE learns that the data is non-unimodal in nature, but makes innacurate density estimates, greatly overestimating densities in some regions and underestimating density in others.
described in Appendix D We vary the input data size and compare all methods on log-likelihood in Figure G We note that CCN and its variants all asymptotically approach the optimal value. All adjusted versions of CCN present faster convergence rates with FCCN dominating the curve. Gaussian methods (BART, CMGF CEVAE) provide more stable approximations in smaller samples. Due to distributional mismatch, the optimal value can not be attained by those methods. Though GAMLSS has the correct family specification, it is restrained by the flexibility of the additive model to approach the optimal value. GANITE progresses slower and needs over 15,000 samples to generate competitive results.
5.4.4 Ablation Study
We include three components to account for treatment group imbalance. They are the Assign-loss (Assignment), the Wass-loss (Wass), and propensity stratification (PS). In this section, we inspect how CCN empirically benefits from each component. Additional details and visualizations can be found in Appendix J
*CMGP convergence was only tested up certain sample size due to computational tractability of fitting the model to larger sample sizes (n 20,000).
14

---

Published in Transactions on Machine Learning Research (09/2022)
Table 4: Quantitative results on the IHDP dataset regarding different variants of CCN (Section 5.4.4
Metrics/ Metrics/Method Method
CCN
Wass
Assign
PS
Assign+ PS
FOCN
PEHE
1.59 .16
1.32 +
1.42 .25
1.15 .10
1.22 .15
1.13 .14
LL
-1.78 = .22
-1.65 + -02
-1.64 .03
-1.67 = .15
-1.65 = .02
-1.62 + .02
AUC (Linear)
.925 + 0011
.938 + .010
.40 .000
.188 + .012
.90 = .00
.942 + .000
AUC (Threshold)
.913 .011
.932 .011
.932 .011
.911 .012
.934 .011
.935 .010
Table Quantitative results on the EDU dataset regarding different variants CCN (Section 5.4.4
Metrics/Methho
CCN
Wass
Assign
PS
Assign PS
FCCN
PEHE
.392 + .049
.324 .046
.33 + .041
400 + .052
.339 = .099
.296 + .042
LL
-2.178 + .044
-2.128 + .020
-2.132 + .033
-2171 + 024
-2.129 + .029
-2.125 + .022
AUC
.933 + .026
.951 .013
.946 + .018
.932 .022
.952 .019
.953 .014
Table 4 and summarize the results on the evaluating metrics in both the IHDP and EDU datasets. Overall, variants of CCN with adjust- ment more accurately estimate the potential outcomes. However, the aspects of how these components contribute vary according to their altributes. The propensity score stratification mainly facilitates information sharing between strata (Lunceford & Davidian 2004;; hence, it excels in the IHDP dataset where the imbalance is caused by removing a specific subset from the treatment group. However, on the condi- tional level the stratification can be regarded as form of aggregation, which might hinder the precision. Therefore, do not see much gain in distribution--aseed metrics or personalized de- cisions by solely including the propensity. The Assign-loss overcomes the confounding effect by extracting representation relevant to confounding effects, and we observe that it effectively increases the model performance all metrics. The combination of Assign-loss and propensity stratification demonstrates the merits of these two approaches. The Wass-loss finds representation that balances the treatment and control groups and improves both point and distribution estimates in the two datascts. Nevertheless, it does not account for the domain-spectee information (Shi et al. 2019 FCCN achieves the best performance in all cases by a clear margin.
-2.5
-3.0
33.5
CCN
Estimated
FCCN
4.0
GAMLSS
-4.5
GANITE
5.0
BART
CMGP
-5.5
CEVAE
-6.0
True Value
2.5
3.0
4.0
4.5
log(sample size)
Figure 6: Convergence rate of models as a function sample size (Section 5.4.3).
Moreover, we examine the proposed adjustment components in different scenarios in Figure Figure ((a) shows that CCN and its variants all asymptotically approach the optimal value with FCCN being the quickest. Figure 7(b) made by varying the values of a or 8. It suggests that parameter values that are either too large too small are more likely hurt models with only a single adjustment component, while FCCN is more robust to these changes. Figure 7(c) describes case where irrelevant dimensions with standard Gaussian distributions added to the covariate space. We observe that adding noise worsens the performance. In this case, the Assign-loss is more likely to overfit the propensity model with extra covariates containing noise only. Hence, mdjustment through Wass-loss is preferable. In Figure 7(d) we vary the propensity model by changing its coefficient. Larger values represent less balanced space and larger confounding effects. In this setup, the Assign-loss slightly better when the imbalance more extreme, as a balanced represontation becomes more challenging to obtain. Among them, FCCN is the most robust due to simultancously considering the domain-invariant and specific information.
15

---

Published in Transactions on Machine Learning Research (09/2022)
-1.60
-1.65
-3.0
-1.70
Estimated
3.5
CCN
listmmaed timdated
-1.75
-4.0
Assign
-1.80
PS
Assign(A)
-4.5
Assign+PS
-1.8
AS)
WASS
-1.90
WASS(a)
FCCN
FCCN(a)
True Value
FCCN(8)
-2.00
3.0
4.5
lsamseeel size)
-log(tuning parameter)
(a) Varying tuning parameters
(b) Varying sample size
-1.8
-2.0
Estimated LL
Ittmmated
-2.2
-2.0
-2.4
-2.5
Assign
Assign
-2.8
PS
PS
Assign+PS
-3.0
Assign+
FCCN
FCCN
-3.2
-35
2
6
10
0.2
0.4
0.6
0.8
1.0
1.2
1.4
Irrelevant Covariates
Coefficient of Propensity Model
(c) Adding irrelevant covariates
(d) Varying the coefficient of the propensity model
Figure 7: Estimated LL under different scenarios (Section 5.4.4). ((a) demonstrates differences in convergence speed of CCN and its variants as a function of sample size. ((b) depicts the performance of CCN variants given different tuning parameter values. 7(c) and 7(d) show that adding noise or treatment group imbalance hurts model performance. Under these different circumstances, FCCN remains the most robust.
6 Discussion
Here, we propose Collaborating Causal Networks (CCN): novel framework to estimate conditional potential outcome distributions, backed by theoretical proofs and paired with an adjustment method that rectifies treatment group imbalance. Our experiments demonstrate that CCN is able to adapt to a variety of outcomes, including exponential family distributions and multi-modal distributions, thus empirically demonstratting that CCN is effective in inferring full potential outcomes. Additionally, incorporation of our proposed adjustment technique in FCCN is relatively robust with regards to treatment group imbalance in synthetic and semi- synthetic experiments. We note that improving distribution estimates leads to improved decision- decision-making making even without priori access to utility functions by comparing to policy learning. In all of our evaluations, FCCN meets or exceeds the current state-of-the-art methodology for potential outcomes distribution estimation, and asymptotically approaches our thcoretical claims.
CCN and its variants are competitive with but do not outperform methods that make modeling assumptions that are perfectly aligned with the true data generating distribution, as is the case with CMGP and the IHDP dataset. Therefore, if the true data generating process is known prior to modeling then other methods may be preferred over CCN, However, we note that access to the true data generating mechanism is not common in practice and that these processes may be quite complex, in which case we would still expect CCN and its variants lo asymplotically capture outcome distributions regardless of form. The Lipschitz continuous assumption a potential limitation, albeit a small one. Lipschitz CDF entails PDF without point masses, which prohibits some distributions, such as zero-inflated probability distributions, but is mild in practice under an additive noise model. Like many causal inference models, the proposed approach is dependent key assumptions. In this case, the primary assumptions summarized by the strong ignorability
16

---

Published in Transactions on Machine Learning Research (09/2022)
assumptions. Finally, with any method with relevance to healthcare precision medicine, emphasize this framework should not be used alone to guide clinical practice, suggest. treatments, or recommend any health-related actions without consultation and close collaboration with medical professionals. When applying CCN, a practitioner should take care to evaluate the reasonableness of the assumptions, especially that of unconfoundedresss, by consulting with domain experts about the features included as predictors the model.
Acknowledgments
We acknowledge the anonymous reviewers for their helpful suggostions and comments.
Research reported in this manuscript was supported by the National Institute of Biomedical Imaging and Bioengineering and the National Institute of Mental Health through the National Institutes of Health BRAIN Initiative under Award Number R01EB026937. The contents of this manuscript are solely the responsibility of the authors and do not necessarily represent the official views of any of the funding agencies or sponsors.
References
Ahmed M Alaa and Mihaela van der Schaar. Bayesian Inference of Individualized Treatment Effects using Multi-task Gaussian Processes. Advances in Neural Information Processing Systems, pp. 3424 3432, 2017.
Ahmed M Alaa, Michael Weisz, and Mihacla Van Der Schaar. Deep Counterfactual Networks with Propensity- Dropout. Proceedings of International Conference on Machine Learning, 2017.
Martin Arjovsky, Soumith Chintala, and Léon Bottou. Wasserstein Generative Adversarial Networks.
International Conference on Machine Learning, 2017.
Serge Assaad, Shuxi Zeng, Chenyang Tao, Shounak Datta, Nikhil Mehta, Ricardo Henao, Fan Li, and Lawrence Carin. Counterfactual Representation Learning with Balancing Weights. Artificial Intelligence and Statistics, 2021.
Susan Athey and Guido Imbens. Recursive partitioning for heterogeneous causal effects. Proceedings of the National Academy of Sciences, 113((7):7353 7360, 2016.
Susan Athey and Stefan Wager. Policy Learning with Obscrvational Data. Econometrica, 89(1):333-161, 2021.
Rukmini Banerji, James Berry, and Marc Shotland. The Impact of Maternal Literacy and Participation Programs: Evidence from a Randomized Evaluation in India. American Economic Journal: Applied Economics, 9(4):303 37, 2017.
Rukmini Banerji, James Berry, and Marc Shotland. The Impact of Mother Literacy and Participation Programs on Child Learning: Evidence from a Randomize Evaluation in India, 2019.
Normand Beaudry and Renato Renner. An intuitive proof of the data processing inequality. Quantum Information & Computation, 12(5-)):432 441, 2012.
Shai Ben-David, John Blitzer, Koby Crammer, Alex Kulesza, Fernando Percira, and Jennifer Wortman Vaughan. A Theory of Learning from Different Domains. Machine Learning, 79(1-2):151 -110:111,1111111 2010.
Dimitris Bertsimas, Nathan Kallus, Alexander M Weinstein, and Ying Daisy Zhuo Personalized Diabetes Management Using Electronic Medical Records. Diabetes Care, 40(2):210-2) 2170 2017.
Alina Beygekzimer and John Langford. The Offset Tree for Learning with Partial LLabels. In Proceedings of the International Conference on Knowledge Discovery and Data Mining, pp. 129 138, 2009.
Peter Bickel and Ya' la'acov Ritov. Nonparametric estimators which can be "plugged-in". The Annols of Statistics, 31(4):1033 1053, 2003.
17

---

Published in Transactions on Machine Learning Research (09/2022)
Guillermo Briseño Sanchez, Maike Hohberg, Andreas Groll, and Thomas Kncib, Flexible instrumental variable distributional regression. Journal of the Royal Stalistical Sociely: Series A (Stalistics in Sociely), 183(4): 1553 1574, 2020.
Yash Chandak, Scott Niekum, Bruno da Silva, Erik Learned-Miller, Emma Brunskill, and Philip Thomas. Universal Off-Policy Evaluation. Advances in Neural Information Processing Systems, 34:27475-27490, 2021.
Victor Chernozhukov and Christian Hansen. Instrumental quantile regression inference for structural and treatment effect models. Journal of Economelries, 132(2))491 -525, 2006.
Victor Chernozhukov, Kaspar Wüthrich, and Yinchu Zhu. An Exact and Robust. Conformal Inference Method for Counterfactual and Synthetic Controls. Journal of the Americon Statistical Association, pp. 1 1-44, 2021.
Hugh Chipman and Robert. McCulloch. Bayes Tree: Bayesion Additive Regression Trees, 2016. URL
https://CRAN R- http://et..rrgpccaaeeeeeeeee
R package version 1.4.
Hugh Chipman, Edward George, and Robert E McCulloch. Bayesian Ensemble Learning. In Advances in
Neural Information Processing Systems, pp. 265-272, 2007.
Hugh A Chipman, Edward I George, and Robert E McCulloch. BART: Bayesian additive regression trees.
The Annals of Applied Statistics, 4(1):206 298, 2010.
Rajeev H Dehejia. Program evaluation as decision problem. Journal of Econometrics, 125(1-2):141-173, 2005.
Peng Ding and Fan Li Causal Inference: A Missing Data Perspective. Statistical Science, (3(2):214 237, 2018.
Xin Du, Lei Sun, Wouter Duivesteijn, Alexander Nikolaev, and Mykola Pechenizkiy, Adversarial balancing- based representation learning for causal effect inference with observational data. Data Mining and Knowledge Discovery, 3544::1713 1738, 2021.
Qiyang Ge, Xuelin Huang, Shenying Fang, Shicheng Guo, Yuanynan Liu, Wei Lin, and Momiao Xiong. Conditional Generative Adversarial Networks for Individualized Treatment Effect Estimation and Treatment Selection. Frontiers in Genelies, 11, 2020.
Ian Goodfellow, Jean Pouget-Abadie, Mehdi Mirza, Bing Xu, David Warde-Farley, Sherjil Ozair, Aaron Courville, and Yoshua Bengio. Generative Adversarial Bets. Advances in Neural Information Processing Systems, pp. 1872-2680, 2014.
Steffen Grunewalder. Plug-in Estimators for Conditional Expectations and Probabilities. In International Conference on Artificial Intelligence and Statistics, pp. 1513-1521, 2018.
P Richard Hahn, Jared S Murray, and Carlos Carvalho. Bayesian Regression Tree Models for Causal Inference: Regularization, Confounding, and Heterogeneous Effects (with Discussion). Bayesion Anabysis, 15(3): 965-1056, 2020.
Jun Han and Claudio Moraga. The influence of the sigmoid function parameters on the speed of back- propagation learning. In International Workshop on Artificial Neural Networks, pp. 195-201. Springer, 1995.
Negar Hassanpour and Russell Greiner. CounterFactual Regression with Importance Sampling Weights. IJCAI, pp. 5880-5887, 2019.
Negar Hassanpour and Russell Greiner. Learning disentangled representations for counterfactual regression. In International Conference on Learning Representations, 2020a.
Negar Hassanpour and Russell Greiner. Variational A Aduacces in Newral Information Processing Systems, 2020b.
Architectures that Excel at Causal Inference.
18

---

Published in Transactions on Machine Learning Research (09/2022)
Miguel A Hernan and James M Robins. Causal Inference. Boca Raton: Chapman & Hall/ 2020.
Jennifer Hill. Bayesian Nonparametric Modeling for Causal Inference. Journal of Computational and Graphical Slalistics, 20(1):217 240, 2011.
Maike Hohberg, Peter Pütz, and Thomas Kncib. Treatment effects beyond the mean using distributional regression: Methods and guidance. PIOS One, 15())=e0226514, 2020.
Guido W Imbens and Donald Rubin. Causal Inference Statistics, Social, and Biomedical Sciences. Cambridge University Press, 2015.
Andrew Jesson, Sören Mindermann, Uri Shalit, and Yarin Gal. Identifying Causal-Effoct Inference Failure with Uncertainty-Aware Models. Advances in Newral Information Processing Systems, 33, 2020.
Fredrik Johansson, Uri Shalit, and David Sontag. Learning Representations for Counterfactusl Inference. International Conference on Machine Learning, pp. 3020-3028 2016.
Fredrik D Johansson, Uri Shalit, Nathan Kallus, and David Sontag, Generalization Bounds and Representation Learning for Estimation of Potential Outcomes and Causal Effects. arXiv preprint or Xiw:2001.07420 2020.
James M Joyce. The Foundations of Causal Decision Theory. Cambridge University Press, 1999.
Nathan Kallus and Angela Zhou. Confounding- Robusst Policy Improvement. Advances in Newral Information Processing Systems, 31, 2018.
Roger Koenker and Kevin F Hallock. Quantile Regression. Journal of Economic Perspectives, 15(4):143-156, 2001.
S Kullback and RA Leibler. On Information and Sufficiency. The Armals of Mathematical Statistics, 22: 79 79-86, 1951.
Sören R Künzel, Jasject S Sekhon, Peter Bickel, and Bin Yu. Metalearners for estimating heterogencous treatment effects using machine learning. Proceedings of the National Academy of Sciences, 116(10): 4156 4165, 2019.
Lihua Lei and Emmanuel J Candès. Conformal inference of counterfactuals and individual treatment effects. Journal of the Royal Statistical Society: Series B (Statistical Methodology). 2021.
Sheng Li and Yun Fu. Matching on Balanced Nonlinear Representations for Treatment Effects Estimation. Advances in Neural Information Processing Systems, pp. 929 939, 2017.
Andy Liaw and Matthew Wiener. Classification and Regression by randomForest. R News, 2(3):18 22, 2002.
Christos Louizos, Uri Shalit, Joris M Mooij, David Sontag, Richard Zemel, and Max Welling, Causal Elfect Inference with Deep Tatent-Variabl Models. Advances in Neural Information Processing Systems, pp. 6446 6456, 2017.
Jared K Lunceford and Marie Davidian. Stratification and weighting via the propensity score in estimation of causal treatment effects: a comparative study. Statistics in Medicine, 23(19):2937-2960, 2004.
Nicolai Meinshausen and Greg Ridgeway. Quantile Regression Forests. Journal of Machine Learning Research, 7(6), 2006.
Junhyung Park, Uri Shalit, Bernhard Schölkopf, and Krikamol Muandet. Conditional distributional treatment effect with kernel conditional mean embeddings and u-statistic regression. In International Conference Machine Learning, pp.
PMLR, 2021.
Joost ME Pennings and Ale Smidts. The Shape of Utility Functions and Organizational Behavior. Management Science, 19(1):12511120 12633, 2003.
Min Qian and Susan A Murphy. Performance guarantees for individualized treatment rules. Annals of Slatistics, 39(2):1188) 2011.
19

---

Published in Transactions on Machine Learning Research (09/2022)
RA Rigby and DM Stasinopoulos. Generalized additive models for location, scale and shape. Applied Statistics, 54:507-554, 2005.
Paul Rosenbaum and Donald Rubin. The central role of the propensity score in observational studies for causal effects. Biometrika, 70(1):41-55, 1983.
Donald B Rubin. Matching to Remove Bias in Observational Studies. Biometrics, pp. 159-183, 1973.
Patrick Schwab, Lorenz Linhardt, and Walter Karlen, Perfect Match: A Simple Method for Learning Representations for Counterfactual Inference with Neural Networks. arXiv preprint arXiv:1810. 0..06000 2018.
Uri Shalit, Fredrik D Johansson, and David Sontag. Estimating Individual Treatment Effect: Generalization Bounds and Algorithms. In International Conference on Machine Learning, pp- 5076-3088 2017.
Shu Shen. Estimation and Inference of Distributional Partial Effects: Theory and Application. Journal of Business & Economic Statistics, 37(1):54 66, 2019.
Claudia Shi, David Blei, and Victor Veitch. Adapting Neural Networks for the Estimation of Treatment Effects. In Advances in Neural Information Processing Systems, pp. 2507 2517, 2019.
Houshang H Sohrab. Basic Real Analysis, volume 231. Springer, 2003.
Mikis Stasinopoulos, Bob Rigby, Vlasios Voudouris, Calliope Akantziliotou, Marco Enea, and Daniil Kiose. gomlss: Generalised Additive Models for Location Scale and Shape, 2021. URL https://CRAN R- R-project
org/packagesgamless R package version 5.3-4.
Shuo Sun, Erica EM Moodie, and Johanna G Nešlehová. Causal inference for quantile treatment effects. Enwvrrommetriss, (244)=26688 2021.
Frik Sverdrup, Ayush Kanodia, Zhengyuan Zhou, Susan Athey, and Stefan Wager. policytree: Policy learning via Doubly Robust Empirical Welfare Mazimization over Trees, 2021. URL https://CRAN R-rojject org/packkagepoolccytree R package version 1.1.1.
Natasa Tagasovska and David Lopez-Paz. Single-Model Uncertainties for Deep Learning, Advances in Newral Information Processing Systems, 32, 2019.
Natasa Tagasovska, Valérie Thavez-Demoulin and Thibault Vatter. Distinguishing Cause from Effect Using Quantiles: Bivariate Quantile Causal Discovery, In International Conference Machine Learning, pp. 9311 9323. PMLR, 2020.
Philip Thomas and Emma Brunskill. Data-Efficient Off-Policy Policy Evaluation for Reinforcement Learning. In International Conference on Machine Learning, pp. 2139-2148 PMLR, 2016.
Philip Thomas, Georgios Theocharous, and Mohammad Ghavamzadeh. High-Confidence Off-Policy Evaluation. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 29, 2015.
Julie Tibshirani, Susan Athey, and Stefan Wager. grf: Generalized Random Forests, 2020. URL https
7/CRAN. www.:rorojjecc.rr/aaaaaaa-geeeee
package version 1.2.0.
SS Vallender. Calculation of the Wasserstein Distance Between Probability Distributions on the Line. Theory of Probability & Its Applications, 18(4):784 786, 1974.
Brian G Vegetabile. On the Distinction Between "Conditional Average Treatment Effects" (CATE) and
"Individual Treatment Elfletts" (TTE) Under Ignorability Assumptions. ICML 2021 Workshop, 2021.
Cédric Villani. Oplämal Transport: Old and New, volume 338. Springer Science & Business Media, 2008
Stefan Wager and Susan Athey. Estimation and Inference of Heterogencous Treatment Effectt using Random Forests. Journal of the American Stalistical Association, 113(523): 1242, 2018.
20

---

Published in Transactions on Machine Learning Research (09/2022)
Stephen G Walker. Bayesian inference with misspecified models. Journal of Stotistical Planning and Inference, 143(10):1621-1663, 2013.
Liuyi Yao, Sheng Li, Yaliang Li, Mengdi Huai, Jing Gao, and Aidong Zhang. Representation Learning for Treatment Effect Estimation from Observational Data. Advances in Neural Information Processing Systems, 31, 2018.
A Yazdani and E Boerwinkle. Causal Inference in the Age of Decision Medicine. Journal of Data Mining in Genomics & Proteomics, 6(1), 2015.
Jinsung Yoon, James Jordon, and Mihaela van der Schaar. GANITE: Estimation of Individualized Treatment Effects using Generative Adversarial Nets. International Conference on Learning Represenlations, 2018.
Guoyi Zhang and Yan Lu. Bias-corrected random forests in regression. Journal of Applied Statistics, 39(1): 151-160, 2012.
Tianhui Zhou, Yitong Li, Yuan Wu, and David Carlson. Estimating Uncertainty Intervals from Collaborating Networks. Journal of Machine Learning Research, 22(257):1 1-47, 2021. URL http://jmlr orgepapers/ v22/20- 1100. html
21

---

Published in Transactions on Machine Learning Research (09/2022)
A Proofs of Propositions 1 and 2
Zhou et al. (2021) assume that the covariate distributions are the same for training and generalization for CN. In obscrvational studies, the training spaces p(z[T = 1) and p(z[T = 0) differ from the evaluation space p(a). Thus, the central challenge of migrating the properties of CN to CCN is demonstrating robustness of CCN with respects to this covariate space mismatch.
First, we explore the properties of CN and CCN in the presence of covariate space mismatch. Second, under assumptions commonly made in the causal inference literature, expand on how CCN can overcome covariate space mismatch. For completeness, we restate the two propositions from the main section of the paper. Note that both claims are made with respects to the full covariate space, Vx such that p(x) > 0.
Proposition S1 (Optimal solution for go and gi). When the support of the outcomes is a subset of the support of 2, or as covers the whole outcome space, the functions go and 91 that minômize g-loss' are optimal when they are equivalent to the conditional CDF of Y
and ( I, Vr. such that p(æ) > 0.
Proposition $2 (Consistency of go and 91). Assume the ground truth CDF functions for T E {0, 1] are Lipschitz continuous with respect to both the features X and the potential outcomes (Y( (()) and the support of the outcomes is a subset of the support of z. Denole the ground truth functions as Pu and 91- As + 20, the finite sample estimators 96 and o have the following consistency property: d(gō. + 0; d gi) +p O under some metrics d, such as the I, norm, and with the space searching tool z being able to cover the full outcome space.
A.1 Claims from Zhou et al. (2021) with and without Space Mismatch
We now restate two CN propositions from Zhou et al. (202)) under the non-causal sctting where there no space mismatch between the training and the evaluation spaces.
Proposition S3 (Optimal solution for g from Zhou et al. (2021)). Assume that (q,x) approcimales the conditional 4th quantile of YYX= = x (inverse CDF, not necessarily perfect). If f(g, spans R then g minimizing equation [ is optimal when it is equivalent to the conditional CDF, or Pr(Y < = g(. x), Vy E R. Vz such that p(z) > 0.
Proposition S4 (Consistency of g from Zhou et al. (2221) ) Assume the true CDF function g" satisfies Lipschilz continuily with respect to x and oulcome y. As n - oo, the finile sample eslimator g" has the following consistency property: d(g", + 0 under some metric d such as L1 norm and with f capable of searching the full outcome space.
Proposition S3 demonstrates that a fixed point solution estimates the correct distributions. Proposition states that the optimal learned function asymptotically estimates the true distributions. However, they do not answer whether these properties hold on a new space that differs from the original training space. We address this existing limitation by developing Proposition SS5 which shows that these properties can still be retained given certain type of space mismatch. For generality, define the training space as p(z)) and the new space as p' (ir).
S4
Proposition S5 (The dependency of CN on p(x)). If the conditional oulcome distribution Y|X- = x remains invoriant between X ~ p(x) and X ~ p'(x) (covariate shift), and p(z) > 0 = p'(x) > O, the solutions in Propositions S3 and the consistency in Proposition S4 also generalize to the new space where p'(x) > O.
Proof of Proposition 55 With the Propositions 53 and
S4
we observe that g estimates the conditional
distribution of Y|X x in the space where p(x) 0.
Next, we generalize it to new space p'(x). Given the condition p(z) > 0 = p'(x) > 0, for any z in
evaluation space with p'(x) > 0, it is covered in the training space where p(x) > O. From Proposition S3 and S4 we know that for such T, the optimum can be obtained. The covariate shift assumption on the invariance of outcome distributions then guarantees that the optimum of such x in the training space is also
the optimum in the new space.
Therefore, each point x in the new space with p'(x) can obtain their optimum through training the model the training space p(x).
22

---

Published in Transactions on Machine Learning Research (09/2022)
A.2 Overcoming Covariate Space Mismatch for CCN
Proposition 55 enables us to extend the optimum of CN to new spaces given two conditions: the covariate stift and the space overlap p(z)= > 0 = p'(x) > O. Our main task is to show how they hold in causal settings. First, we give a weaker version of Propositions and 2 as direct result from Proposition 53 and without accounting for the mismatch bet ween the training and generalization spaces.
Claim S1 (Potential outcome distributions each treatment space). Given all the conditions in Proposition
S3 and S4 optimal solution for go and 91 in g-loss" in equation / are the true CDF's of (O) z, T =
0, dx, such that p(z[T = 0) > O; and Y()) = I. ==1, Vx, such that p(z]T = 1) > 0. The finite sample estimalors % and of can also consistently estimate these CDFs.
Discussion on Clain S1. We only discuss the first part of Claim ST for T 0 without loss of generality, while optimizing go involves updating parameters using batches with t 0. The other group can be shown using identical steps.
A direct conclusion from Propositions S3 and S is that the optimal 90 is the fixed point solution, and finite sample estimator ø6 consistently estimates the true CDF of Y|X = x, T=0, Vx, such that p(=|T= 0) > 0. This is claimed for the full conditional distribution Y|X = I, T = rather than for the potential outcome (().. By virtue of the treatment consistency (Assumption 2 the following two outcomes are identically distributed: Y|X x, T 0 - Y(O) x; T 0. Therefore, we can successfully establish the estimators for potential outcome distributions, but currently limited to each treatment subspace.
As mentioned above, we need two conditions to generalize Claim ST back to the full space with density p(r). The covariate shift has been explicitly described in (Johansson et al. 2020). Covariate shift is a straightforward consequence of Assumption 3 and we state for clarity.
Lemma S1 (Covariate shift). The conditional distributions of potential outcomes (O) = x and ( = are independent of the covoriate space in which they are located.
Discussion on Covariate Shift. The ignorability states that [Y(0), Y(1) 1 T|X. It implies that Vy(0). x, the three following density functions are equivalent: p(y(O), 2, = 0) (()()) =, 1) = p(v(0), x(]]|x = z).
This supports that the potential outcome distributions are invariant to the treatment groups
We next show the condition of the space overlap, which is straightforward consequence of Assumption 1 Lemma S2 (Positivity relating to the space overlap). Under Assumption 1 the equivalent condition holds: p(x) 0- p(x[T = 0) > 0, and p(x) > + p(x[T'= = > O.
Proof of Lemma S2 The positivity in Assumption I claims that Vz, 0 < Pr(T = 1/x) 1. Then for each æ where p(x) > 0, we can find a constant 1 > C2 > 0 that satisfies Pr(T 1/z) > Cr.
By Bayes rule, p(=/T = 1) = Pr(T = 11)))(=)/PP(() = 1) Czp(x) > 0. Then p(x) > 0 = p 0.
From the other direction, if p(=[T = 1) = Pr(T = (lx) Pr(T = 1) > 0, then each component on the right hand side neods to be positive. Therefore, p(x(T 1) = p(x) > We can use the same procodure verify the condition for T= O.
With Lemma and 52 satisfying the conditions in Proposition S5 we acquire that the space mismatch given the standard causal assumptions is the specific lype that does not impact the large sample properties of CCN. As a result, Propositions and naturally follow from these two Lemmas (or alternatively, following from Assumption [ and 3). Thus, under the standard assumptions in causal inference, CCN will capture the full potential outcome distributions. This procedure is not limited a binary treatment condition and is extendable multiple treatment setups.
23

---

Published in Transactions on Machine Learning Research (09/2022)
B Semi-Synthetic Data Generation
B.1 IHDP
Simulated replicates for the IHDP data were downloaded directly from ltttps://githubb http:////ggtttbb..oommmmmmm.aa.a cowmmmiiicalmmmm////mcmc.../a//a cfrnet, which were used to generate the results of WASS-CFR (Shalit et al. 2017) and CEVAE (Louizos
ct al. 2017). The IHDP dataset does not contain personally identifiable information or offensive content.
B.2
Education Data
The raw education data corresponding to the EDU dataset were downloaded from the Harvard Datavers o which consist of 33,167 observations and 378 variables. The dataset does not contain personally identifiable information or offensive content.
We pre-process the data by combining repetitive information and deleting covariates with over 2,500 missing values. This results in processed dataset containing 8,627 observations. The function fyi (-) and fg (-) are learned from the observed outcomes for the treatment and control groups. They are both designed as single- hidden-layer neural networks with 32 units and sigmoid activation functions (Han & Moraga 1995). A logistic regression model with coefficient 3 = [31,..., B2s] and propensity score Pr(T = =xxx),,
14e2p(-2 are used to generate treatment labels and mimic observational study data. The coefficienss are randomly generated as Bi U(-0.8,0.8) 0.8).
C Detailed Method Implementations
The models CCN, CEVAE and GANITE were trained and evaluated on the various datasets described in this work on a machine with a single NVIDIA P100 GPU. The R-based methods BART, CF, and GAMI.SS were trained and evaluated on machine with an Intel(R) Xeon(R) Gold 6154 CPU. The CMGP model was trained and evaluated on machine with an Intel Core i7 1h generation processor and NVIDIA GeForce RTX 3090 GPU,
CCN and FCCN
The implementation of CCN and all its variants are based on the code base for CN Zhou et al. 2021 which is provided at /tttps://iithu:
under the MIT license. Functions 9u and 91 follow the structures of g in Zhou et al. (2021). We implement the full collaborating structure and find the optimization to be harder when added with regularization terms. Therefore, we fix ʃ a uniform distribution covering the range of the observed outcomes. This is referred to the "g-only" setup in Zhou ct al. 2021b, and is casier to optimize with only marginal loss in accuracy.
In FCCN, we introduce two latent representation @a(-) and Φw(-). We set their dimensions to 25. They are both parameterized through a neural network with a single hidden layer of 100 units. The Wasserstein distance in Wass-loss is learned through D(-) which is a network with two hidden layers of 100 and 60 units per layer. We adopt the weight clipping strategy with threshold: (-0.01, 0.01) to maintain Lipschitz constraint. (Arjovsky al., 2017 The hyperparameters a and /3 are tuned for FCCN by selecting candidate values that do not significantly impact the log-likelihood calculated upon the observed outcomes, with the idea being to balance fitting the observed data while also encouraging generalization through non-zero a and 3 values. We propose few candidate values for a and B as: 5e-3, 1e-3, 5e-4, 1e-4, 5e-5, 1e-5, as we do not want these values to be too large to overpower the part of the objective that learns the empirical distribution. Then, we perform a grid search to determine the hyper-parameters. Based on the results on the first few simulations, we fix a 5e-4 and 1e-5 in IHDP experiments and a 1e-5 and 3 5e-3 in EDU experiments. We find this specification to consistently improve the performance over regular CCN. To assess the potential outcome distributions, and take expectation over defined utility function and draw 3,000 samples for each test. data point with the learned go and 91- The code for CCN and ils adjustment will be public on GitHfub under the MIT license if the monuscript is accepted.
http:://aasautee...rrradddaaaaad harvard. haw.arrd.cch//aataet...... http://ww....ttooooooooo/ottaaat
24

---

Published in Transactions on Machine Learning Research (09/2022)
BART (Chipman et al,, 2010) The implementation of BART uses the R package BayesTree CChipaan & McCulloch 2016) under the GPL. (>= 2) License. We use the default setups in its model structure. Chipman et al. (2007) suggest that BART's performance with a default prior is already highly competitive and is not highly dependent on fine tuning. We set the burn-in itcration to 1,000. We draw 1,000 random samples per individual to access their posterior predicted distributions, as the package stores all the chain information and is not scalable for large data.
CEVAE (Louizos et al., 2017) CEVAE is implemented with the publicly available code from https: : //github. www.riioooeeweeeeeeeeeeee CEVAEE CEEAEEEEVeerCSSS pytorch/ We follow its default structure in defining encoders and decoders. The latent confounder size is 20. The optimizer is based on ADAM with weight decay according to Louizos et al. (2017). We use their recommended learning rate and decay rate in IHDP experiments. In EDU experiments, the learning rate is set to 1e-4, and the decay rate to 1e-3 after tuning. We draw 3,000 posterior samples to access the posterior distributions.
GANITE Yoon et al., 2018) The implementation of GANITE is based on https://github: com/ jsyoon0823/ GANITE under the MIT License. The model consists of two GANs: one for imputing missing outcomes (counterfactual block) and one for generating the potential outcome distributions (ITE block). Within each block, they have a supervised loss on the observed outcomes to augment the mean estimation of the potential outcomes. We use the recommended specifications in Yoon et al. (0018) to train the IHDP data. In the EDU dataset, the hyperparameters for the supervised loss are set to a = (counterfactual block) and B = 1e-3 (TTE block) after tuning.
CF WWager & Athey, under the GPL-3 License. are aulomatically tuned.
2018) The implementation of CF uses the package grf (Tibshirani et al. 2020) We specify the argument tune. paramaaere=aaalll that all the hyperparametens
GAMLSS (Hohberg et al. 2020) The implementation of GAMLSS uses the R package gamlss (Rigby & Stasinopoulos 2005) under the GPL-3 License. Since this method uses likelihood to estimate its paramoters and often does not converge under complex models, we feed the location, scale, and shape models with relevant variables only. In location models, we fit all continuous variables with penalized splines. In scale and shape models, we use relevant variables in their linear forms. This choice is based on consideration of the balance between representation power and model stability.
CMGP (Alaa & van der Schaar 2017) The implementation of CMGP uses the Python package cmgp under the 3-Clause BSD License. For all experiments in which CMGP was used, we use a maximum iteration count (max_gp_ gpriitrrtiiiinsss parameter) of 100 when fitting the CMGP model.
D Enforcing a Monotonicity Constraint
The Icarned CCN model should have a monotonic property that g(z, +)) 2 g(2, Vc 2 0. In our experiments, this condition learned with a neural network architecture with our training scheme. We do not see any non-trivial violations of this requirement. If required, though, this scheme can be enforced by modifying the neural network strueture. One way of accomplishing this goal is to use a neural network with the form,
g( 2) = Σ (o)tmaa(""()))) + (x)))
Here, (-) represents the sigmoid function. - dh and get are all neural networks that map from the input space to a J-dimensional vector, IP + R.. In this case, the formulation of the outcome is still highly flexible but becomes a mixture of sigmoid functions. As the multiplier on z is required to be positive, each individual sigmoid function is monotonically increasing as a function of z. Because the weight on each sigmoid is positive, this creates a full monotonic function as a function of z.
We implement this structure and find that it competitive with more standard architecture but is more difficult to optimize. As it is not empirically necessary to implement this strategy, we prefer the standard architecture in our implementations.
25

---

Published in Transactions on Machine Learning Research (09/2022)
E Additional CDF Visualizations
We provide additional visualizations for the estimated potential outcome distributions with each method in Figure Si based on another data point, which agrees with the results visualized in Figure The two variants of CCN are capable of capturing the main shape of the true CDF curves, including the asymmelry of the exponential distribution, with higher fidelity. GAMLSS less accurate in the treatment group due to using skewed normal for the exponential distribution with location shifts. GANITE's two-GAN structures reliant data richness for accurate predictions (Yoon et al., 2018), : so it falls short in cases with greater treatment group imbalance. BART provides reasonable estimates but struggles with the misspecification by its Gaussian form. Similar to GAMLSS, CMGP captures the control group CDF well but does not approximate the ment curve with high fidelity. CEVAE captures the overall marginal distribution for the potential outcomes as shown in Figure Ihh but fails to discern the heterogeneity in conditional distributions figure (g)
highly
Estimated CDF
Estimated CDF
Estiman
0.6
0.6
0.6
Estimated
0.4
True YIo)
0.4
0.4
0.4
True YI3)
Eat
0.2
(a) CCN
(b) FCCN
(c) GANITE
(d) GAMLSS
1.0
Marginal YIO)
Estimated CDF
Estimated CDR
Marginal Y(D)
Est YID)
Est YIL)
Estimars
0.6
0.6
Estimated
0.4
0.4
0.4
0.4
0.2
0.2
0.2
0.2
(e) BART
(f) CMGP
(g) CEVAR
(h) CEVAR
Figure S1: Additional visualizations of estimated vs. true CDFs for a random EDU dataset sample (Section The two variants CCN give good distribution estimates, while other methods provide less accurate estimates. By comparing the posterior distributions of CEVAE against the conditional distribution and marginal distribution of the ground truth in Figure (g) and Figure (1(h) conclude that CEVAE primarily captures the marginal distribution in this study.
5.3 Section E
F Policy Learning
In decision making, the core difference between a traditional policy learning method and a distribution learning mothod is whether the utility is determined in advance. Though a policy learning method can tailor decisions based on different utilities, computation, we discuss below another shortcoming of traditional policy
the cost of fitting new model each proposed utility.
Regardless of the inconvenienee learning methods. To train a traditional policy learning approach, the first step is often to convert the outcome to the observed utility. While this is less problematic for bijective transformations, it might result in information loss if we deal with discretized utilities.
To demonstrate, we compare FCCN to policytree with its published package (Athey & Wager 2021 Sverdrup et al. 2021 on IHDP. We propose two types of utilities. They include a linear utility, Un(7) = 7, U1(y) = y 4,
and a threshold utility, Un(r) = 1,2 (o), (?( = Ey>)))(()-- Since the policytree package only outputs the decision, we use accuracy as the evaluation metric. The results are summarized in Table SI FCCN consistently makes more correct. decisions. The information loss in the threshold utility negatively impacts
26

---

Published in Transactions on Machine Learning Research (09/2022)
Table S1: Comparing the accuracy of policy learning between FCCN and policytree (Section
Utehity/Method
FCCN %
policytree %
Linear
88.60 + 1.09
76.86 1.03
Threshold
87.72 1.19
57.64 + .71
policytree. We note that threshold utility drastically reduces the available information by converting a continuous scale to a binary scale.
Fundamentally, these two methods are different and they address similar problems from different perspectives. There might be some possibilities that we could combine their merits. Hence, we will consider exploring their interactions more in future work.
G Additional Distribution Tests
To further illustrate the potential of CCN to model different types of distributions with high fidelity, we simulate potential outcomes from three extra distributions to assess its adaptability. For comparison we include GAMLSS, BART, CCN, and FCCN. The assessment based on log likelihood (LL) to reflect the closeness to the true distributions. We simulate the covariate spaces and treatment labels with the following procedures:
140
T-1
100
40
Covariates:
K; = (I IIO,;), Tz sid. N(O, 1);
True Propensity Score Pr(T 1(x)
Treatment assignment:
The limited propensity score overlap between two treatment groups (Sec- tion This indicates a severe treatment group imbalance.
Figure S2:
Pr(7: 1/xx) exp(exx)) x[]
..., T.
Given the magnitude of we have created a covariate space with limited overlap between two treatment groups as shown in Figure With limited sample size, also helps evaluate the robustness of our method when positivity in Assumption is possibly violated. Then specify three scenarios with sufficient nonlinearity added to the potential outcome generating processes. We generate 2,000 data points in each case, and the variability assessment is based on 5-fold cross validation.
Gumbel Distribution:
10 25a) 1): Y(1), Gumbel cos Σ f=1
10 sin .
25.5) j=1
Y(0),,, Gumbel sin Σ
xyd)
5 cos Σ 1=1
2z,0)
Σ
Gamma Distribution:
10
(00)
Gamma
sin Σ
23,1) + cos
+ 0.5, 2
cos
x f,4) + sin Σ 23,1
j=6
/=1
10
5
10
xzoe) )
Y(1),a Gamma
cos
+ sin
+ 0.5, 2
sin
kg,)) + cos
Σ
j=1
j=6
j=l
j=
Weibull Distribution:
((0)= Weibull 5 sin Σ x...) ) + cos
2
cos
x..a)
+ sin Σ
+ 0.2
j=
j=6
27

---

Published in Transactions on Machine Learning Research (09/2022)
10
10
211) ~ Weibull 5
cos
Fj,i + sin Σ kz,))
2
sin
jn ) + cos
Σ
Tj,a
+ 0.2
j=6
y=1
==
H Estimating Multimodal Distributions
We use the following steps generate the data from a multimodal distribution:
Covariales: Tz ii.d
1);
Treatment assignment: Pr(T, 1) 1x,00
Mixture parameter: d ii.d, (ernoulli(0.5):
Potential outcomes: ((0) N
Y()aa (6, exp(1) +
The control group is mixture of two Gaussian distributions, and the (reatment group mixture of Gaussian and exponential distributions. The mixture information is not given to any model, and we simply use their original form to approximate the distributions. Each model is trained using 1,600 simulated samples.
Sample Size and Convergence
We create example with the logistic distribution to visualize the performance of models with respect to sample size. We simulate 40,000 samples in total and hold out 2,000 for evaluation based log likelihood (LL) with the following procedures:
IJz i.d. N 1);
Covariates: x = (01, 22,,
Treatment assignment: Pr(T; 1 [s:;) [tezp(=x[0) B 22;
Scale parameter: a, 1 + 22,4 + x3,i| + 0.5;
Location parameter: (0);; sin + =2 + sin( p(1); cos(x1, + x2, + cos(23,;7);
Potential outcomes: (0); Logistic(((()); 0,), (1), Logistic(u(1):...)
J
Ablation Study
First, we provide an additional visualization in Figure SB
reflect how well each adjustment improves the the uncertainty estimates of EDU data. FCCN not only captures the heteroskedastieity, but also reduces uncertainty by exhibiting narrower uncertainty bars in the corresponding plots.
J.1 Additional Imbalance Adjustment Study
We give another motivating example to visualize the added robustness provided by our adjustment scheme. The visualizations in Figure 7(b) 7 and 7(d) are all based the same synthetic data, described below. We the same covariate space and treatment assignment mechanism in Appendix G However, posit nonstandard distribution with its location model trigonometric function, and outcome uncertainty model as heteroskedastic Beta distribution. We generate 2,000 data points in total, with 8/2 split for training and testing. The detailed synthetic procedure for the outcomes is described below:
5
10
10
1
Σ
Σ
(0)) Beta
sin Σ zz
Y()) Beta
+
Σ Zjf
28

---

Published in Transactions on Machine Learning Research (09/2022)
Estimated Interval Width
4.0
T=0, M=1
Estimated Interval Width
Tw0, M-1
Estimated Interval Width
4.0
T+0, M-1
Tw1,
T-1, M-1
Tw1, M-1
Tw1, M-1
3.0
3.0
3.0
2.0
2.0
0
1.0
1.0
1.0
True Interval Width
True Interval Width
True Interval Width
(a) CCN
(b) Wass
Assignment
Estimated Interval Width
4.0
T-0, M-0
M=1
Estimated Interval Width
T=0,
T-0,
Estimated Interval Width
Tw0, M-0
M-1
Tw1, M-0
T=1, Ma0
Tw1, M-0
Mw1
T=1, M-1
Mw1
3.0
3.0
3.0
2.0
2.0
8
1.0
1.0
2.0
3.0
4.0
1.0
2.0
3.0
1.0
2.0
3.0
4.0
True Interval Width
True Interval Width
True Interval Width
(d) PS
(e) Assignment-+P5
(f) FCCN
The ability of different variations CCN to capture heteroskedastic outcomes of EDU data (Section 5.3) are demonstrated by plotting the estimated versus true 90% interval widths given four different
Figure S3:
combinations of T and M in an additional ablation study (Section D.. FCCN not only captures the het- eroskedasticity, but also reduces the uncertainty by exhibiting narrower regions of uncertainty.
We visualize the performance of different adjustment schemes in scatter plots where the axis corresponds the true propensity score of each point. Figure S4 depicts the absolute difference between the inferred CATES and true CATFs. Lower vertical positions represent lower errors. We observe that the performance deteriorates in each method point is close to either the two boundaries (extreme propensity scores), which are areas in which models generally struggle the most observational studics. Compared with the baseline CCN, each adjustment scheme by itself lowers the error to some extent. Among them, Wass- Wass-CCN, Assign-CCN and FCCN are able reduce the average error by over 50 %. The propensity stratification (PS) adjustment can effectively reduce bias when there is more homogeneity within each stratum. However, severe imbalance in this case only gives homogeneity in the strata where propensity is around 0.5. Hence, PS provides limited benefits. In contrast, Wass-loss and Assign-loss seek new representations to either rectify group level imbalance or exclude confounding effects. They prove to be more effective the regions of imbalance, which represent the majority of the data points.
29

---

Published in Transactions on Machine Learning Research (09/2022)
Absolute Difference in CAT
0.46
Medlan
Medlan 90% Range
90% Range
90% Range
Absolute Difference
2.0
Absoluee
6.6
6.6
True Propensity Score Pr(T= 1(x)
0.2
True Propensity Score Pr(P=11X0
True Propensity Score Pr(T=1(XX
(a) CCN
(b) Wass
(c) Assign
CATS
agez
Average:
Averages
Mediann 90% Range
Median
Median
90% Range
2.0
90% Range
Absolute Differance
2.5
2.5
Absoluts.
Absolute
0.6
6-2
Propensity Pr(T=1(X)
0.4
66
1.0
6.0
0.2
Propensity
04
046
Pr(f=1(X)
1.0
0.6
Propensity Pr(f=11X0 1.0
04 0.6
(d)
(c)
Assign+
(f) FCCN
Figure $4: Scatter plot of propensity scores (x-axis) versus the absolute difference between the true CATFs and their estimates (y-axis) (Section J.Ip. Among variants of CCN, Wass- Assign- CCN, and FCCN are able to reduce the average error by over 50%.
Extimated
Latimated
Eatmateat
-2.37
-1.69
-1.72 90% Range
Median
90% Range
90% Range
True Propensity Score Pr(T=1/X)
True Propensity Score Pr(T=1|XX
True Propensity Score Pr(T=1/X)
(a) CCN
(b) Wass
(c) Assign
Estimated LL
satimated 1.
Extmated
-2.23
Average: .1.64
Averager -1.61
Median
90% Range
90% Range
90% Range
True Propensity Score Pr(T=1/X)
True Propensity Score Pr(T=1|X)
True Propensity Score Pr(T=1/X)
(d) PS
(e) Assign+PS
FCCN
Figure $5: Scatter plot of propensity scores (x-axis) versus the estimated log- likeliieooodd (LL) (y-axis) (Section Collectively, the Assign-loss and Wass-loss contribute to FCCN making more robust distributions estimates than any one single component.
30


Reviews:
[{"review_id": "Z9eqVlblL9", "reviewer": "Reviewer_fqza", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "** Strengths **\n- I assume the TMLR audience is interested in the problem potential outcomes distribution estimation and its effect on decision making\n- The proposed framework is an interesting application of previous work (Zhou et al 2021)\n- To the best of my limited knowledge, the experimental evaluation seems extensive, including some interesting ablations and a comprehensive comparison with existing methods\n\n** Weaknesses **\n- The clarity and preciseness of the writing can be improved, including adding a bit more context and examples (more details in the requested changes)\n- Some of the claims in terms of what the algorithm can do asymptotically should be qualified by the assumptions (Lipschitz continuous CDF) and its potential implications should be discussed", "was_revised": false, "final_justification": "", "rebuttal": "", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The paper proposes Collaborating Causal Networks (CCNs) to estimate the potential outcome distributions, which allows one to estimate the utility of the treatment even in case of personalized utility functions (i.e. individual-specific). \n\nProposition 1 and 2 show that CCNs asymptotically converge to the potential outcome distribution under a set of assumptions, including that the CDFs for X and the potential outcomes are Lipschitz continuous. The theoretical results are strongly based on (Zhou et al 2021). \n\nBesides the asymptotic properties, the authors propose an adjustment for the finite sample method, FCCN, based on a regularization that encourages a domain invariant and domain specific embedding for X. \n\nThe paper provides an extensive experimental validation on some benchmark datasets that seem to validate the advantages of the methods. The ablation studies on the various components of the FCCN regularization provide some interesting insights (e.g. propensity score stratification being beneficial in IHDP, but not necessarily in general).", "claims_and_evidence_yn": "", "claims_explanation": "", "audience_yn": "", "audience_explanation": "", "requested_changes": "I have a few questions and suggestions:\n\n- In the introduction U(\\gamma) is used, but it isn’t clear what \\gamma is, so the statement U(\\gamma)=\\gamma for CATE is a bit unclear\n- Personally I found the use of CCN for both plural (Collaborating Causal Networks) and singular (a Collaborating Causal Network) a bit confusing. I would instead use CCN for the singular and CCNs for the plural. Moreover, in the singular case I would use “a CNN” \n- Some of the contributions could be stated more precisely, e.g. the asymptotic properties in 2. (under which conditions are they proven? What do these conditions, e.g. a Lipschitz CDF, entail?)  \n- Contribution 4. Is part of the motivation for contribution 1, and “practically more meaningful” when discussing personalized utilities seems a bit of an overstatement (since we might still need to elicit these utility functions, which seems even more complicated on an individual basis)\n- The ITE vs CATE discussion is more of a footnote\n- In Section 2.2. can you maybe provide for illustration and clarification purposes an example in which naively transforming the outcome for a threshold utility could result in a loss of information? (There are some results in the experimental section, but maybe you have a simpler example to show)\n- Positivity can be more precisely state as  $0 < P(T=1|X=x) <1$ for all x such that $P(X=x) \\neq 0$\n- Section 3.1 was very short and I would have appreciated some more explanation or examples, since most of the work is based on CNs. For example I didn’t understand what does it mean that f(.) is a function that searches the outcome space, since it seems that it mostly constrains the function to an invertible one. What am I missing?\n- Also I think it would be less confusing to call the network g(Y,X) in the first line as g(y,x) (similarly to later). I think this is what is intended also by the conditional CDF, which seems to be estimating P(Y < y(input) | X=x (input))\n- Section 3.2 typo “Assumption 1 guarantees that …” $p(X|T=0) = p(X|T=1)$ …\n- Proposition 1 seems sensible, but can you discuss how can one implement it in practice?\n- Proposition 2, what does a Lipschitz CDF entail, e.g. no discrete Y ? Can you comment on how limiting this assumption is?\n- In Section 3.3. Why do we need $\\phi_A$, wouldn’t it be enough to have a propensity score estimation directly from X (so $e: \\mathbb{R}^p \\to \\[0,1\\]$?\n- Equation 5 is unclear, what are $g_1, g_2$ ? I can assume from the text that there are two networks, each with a regularization term, but can you maybe expand the explanation?\n- Typo in the VIllani(2008) citation", "broader_impact_concerns": "While there isn’t any specific broader impact concern for this work, the general estimation of potential outcomes distributions might easily be misused or misunderstood by practitioners, especially if the assumptions and limitations are not completely clear. I would suggest the authors discuss the limitations of their method in the conclusions.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "jUKNJlNq87", "reviewer": "Reviewer_4C4n", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "STRENGTHS\n1. I think, for the most part, the paper is well written and easy to follow. I have a few suggestions in this regard, which I detail in the “requested changes” section. \n2. The problem makes sense. I agree with the authors’ observation that it may be more informative to estimate a potential outcome distribution rather than a point estimate. In typical scenarios, there is plenty of data to estimate these distributions accurately, yet estimating the distribution is typically not adopted in other methods.\n3. I find the empirical validation of the accuracy of (F)CCN’s estimated distributions to be quite thorough and tend to believe their superiority.\n\n\nWEAKNESSES\n1. It is not entirely clear to me what purpose the utility functions serve. Could the authors provide an example of utility functions used in practice? Furthermore, the provided categories (i - iv) all seem to be based on an expectation over the same distribution? To me this doesn’t seem to be a good argument for estimating the distributions? Can’t we just map the point estimate to the components required to complete the utility function? Don’t get me wrong, I still think that estimating distributions is a useful idea, but the utility functions may not be the best argument for this.\n2. It seems the adjustment from a standard CN to (F)CCN is to simply reinterpret it as a T-Learner. This is fine, but why not use more adopted distribution estimators? Perhaps a Gaussian Process (such as in Alaa & van der Schaar (2017)), or even a simple Kernel Density Estimator? Is there a unique property of CNs that make them more suitable for this task that I am missing?\n3. While the set of benchmarks is quite good, I am missing the GPs used in Alaa & van der Schaar (2017). As their CMGP is purely based on GPs, they also end up estimating the same distributions as (F)CCN, could you use CMGPs the same way as you used (F)CNN?\n4. IHDP and EDU are indeed standard benchmark tests for CATE estimators. However, the authors propose to estimate more than just CATE. How would (F)CCN react to more complicated distributions? Perhaps an additional benchmark test on synthetic data with varying distributions (according to some parameter you could visualise over) would be interesting.", "was_revised": false, "final_justification": "", "rebuttal": "", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The authors propose Collaborative Causal Networks (CCNs) and the extended full adjustment CCNs (FCCNs). Their idea is centred around CATE being a point-estimate, while the entire TE distribution would be more informative. Rather than estimating an expectation of the potential outcomes directly (as is the approach taken by most relevant methods in this area), the authors predict the entire potential outcome distribution using Collaborative Networks (CNs). \n\nIn their paper, the authors show how CNs can be extended as a T-Learner (2 distinct CNs, one for each treatment option), and suggest how one may build an S-Learner (where the treatment is considered a covariates in X). Their work is empirically validated on IHDP and EDU, two semi-synthetic datasets typically used for validating CATE models.", "claims_and_evidence_yn": "", "claims_explanation": "", "audience_yn": "", "audience_explanation": "", "requested_changes": "*Changes suggested above*\n1. Provide some examples and motivation for utility functions in the main text.\n2. Further motivate the use of CNs for this particular task.\n3. Add CMGP as a benchmark for comparing distributions (e.g. in fig4)\n4. Perform additional experiments on synthetic data with varying distributions (with increased/decreased complexity), i.e. provide additional tables such as tab1 and tab2, but using a new synthetic setup.\n\n*Additional changes*\nIt is not very clear to me how CNs work. A more complete (algorithmic or architectural) description of CNs, either in the main text or the appendix, would help me a lot.", "broader_impact_concerns": "I see no broader impact concerns.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "QVPW8qatTm", "reviewer": "Reviewer_YQTK", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "**Strengths**\n- The studied problem is well motivated: much attention has been paid to estimating a particular parameter of potential outcome distribution---the conditional average treatment effect (CATE)---but many other functions of the distribution could be of interest. \n\n- The overall solution to the problem is sound, and builds naturally on previous results. \n\n- The baseline models used in the empirical evaluation are well-chosen and represent qualitatively different approaches. \n\n- The proposed task \"EDU\" is a reasonable construction given the topic of the paper. However, the noise model is somewhat simplistic given that the task is to estimate the distribution. \n\n- The empirical results are examined from several different perspectives. \n\n**Weaknesses**\n\n- The structure of the paper could be improved. At a high level, the paper is about estimating p(Y(t) | X).  2.2 places large emphasis on utility functions which are only a small part of the rest of the paper. The title and Section 3 places large emphasis on extending the CN model. The main problem gets lost at times. The overall approach could be described as quantile regression with sampled target quantiles using deep neural networks. Additional losses are added to handle covariate shift. \n\n- In 2.2., the authors discuss personalisation of treatment through the expected utility over the (to-be-estimated) potential outcome distributions. This task is closely related to off-policy evaluation, and in particular, to distributional off-policy evaluation. For example, Chandak et al (Neurips 2021), studied Universal Off-policy Evaluation and \"take the first steps towards a universal off-policy estimator (UnO)—one that provides off-policy estimates and high-confidence bounds for any parameter of the return distribution\". A discussion of off-policy evaluation is absent from the current paper. Similarly, the relationship to general quantile regression is not discussed. \n\n- In the main contributions, the authors claim to \"prove the asymptotic properties of CCN\". These proofs are given in Appendix A and amount to arguing that if the conditional distribution is well-estimated asymptotically (proven in existing work), the same is true in expectation over different distributions of the conditioning variable (see proof of Proposition S3). While this is accurate, the claimed contribution is smaller than what the list in the introduction suggests because of its heavy reliance on previous work---I would argue that the proof is not new, even though the implication may be wider than what the initial work (Zhou et al., 2021) concluded. Moreover, the claim, as stated, could be misinterpreted as a claim about a practical algorithm, but no indication for how to actually find an optimiser of the g-loss is given. \n\n- The proposed FCCN objective includes two hyper parameters, \\alpha and \\beta, which are selected (as described in Appendix C) based on the empirical log-likelihood. It is not clear to me why a non-zero value of either hyper-parameter should yield a higher log-likelihood since both are aimed at improving generalisation to a new distribution, not across samples from the same distribution. Choosing hyper parameters used for this purpose has been a large point of discussion around previous work (such as the cited Shalit et al. (2017)). \n\n- The proof of Proposition S3 is very informal, ending with the sentence \"We claim that the optimum remains invariant when we generalize from p(x) to p'(x).\" Is this not what should have been proven at that point? Moreover, Lemma S1 and Lemma S2 are trivial consequences of Assumptions 1 and 3 and stating them as Lemmas is not necessary. (They are also not used to prove a Theorem). \n\n- On page 5, the authors claim that including e( ) in the covariate space encourages propensity score stratification. The empirical support for this claim is rather mixed since including PS seems to have positive effect in some cases and negative in some (see Table 4 and Table 5). This brings up another question: In the \"PS\" ablation (not PS+Assign), is the Assign-loss omitted? In this case, PS has no reason to be representative of the propensity score.  \n\n- The CCN neural network parameterization is not completely described, in that 3.2 contains no reference to  the learned representation \\phi. Does CCN use the same representation sharing as FCCN but without the adjustments to the objective of FCCN (this seems suggested by the ablation study in 5.4.4)? Or alternatively, are g0 and g1 fit completely independently of each other?", "was_revised": false, "final_justification": "", "rebuttal": "", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The paper studies estimation of potential outcome distributions using a framework dubbed Collaborating Causal Networks (CCN), extending the existing framework Collaborating Networks (CN). The problem is formalised as learning potential outcome CDFs which generalize from treatment groups to the full study  population, under assumptions often made in observational causal estimation: consistency, overlap and ignorability. It is argued that, in this setting, the asymptotic properties of CN apply also to CCN. Additionally, three heuristic adjustments to the learning objective for improving finite-sample performance of neural network estimators are proposed, grounded in the causal inference literature: (i) minimisation of representation space treatment group imbalance, (ii) propensity stratification through treatment assignment loss (iii) inclusion of fitted propensity score in representation. The resulting estimator is dubbed FCCN. The proposed framework is evaluated in two simulated datasets, IHDP and EDU, both based on randomized experiments but with synthesised outcomes. It is shown that the FCCN  adjustments to the objective of CCN  improves likelihood, error in estimated causal effects and accuracy in the estimated utility. The frameworks also compare favourably to other estimators of potential outcome distributions, including both parametric and non-parametric methods. On the EDU dataset, the authors inspect estimates more closely, and conclude that FCCN and CNN more closely fit the CDFs than baseline methods and learn at a higher rate. An ablation study concludes that each of the added components contribute to the results. ", "claims_and_evidence_yn": "", "claims_explanation": "", "audience_yn": "", "audience_explanation": "", "requested_changes": "**Necessary for recommending acceptance**\n- Adding a paragraph discussing the relation to quantile regression using neural networks and to off-policy evaluation and distributional reinforcement learning---in particular work which estimate the CDF of the value function or uses quantile regression---would be necessary to better contextualise the current work. \n\n- Describe the architectures used to fit the CCN and FCCN models in equations, including representation components. Consider adding an illustration of CCN as well. If the CCN model does not share architecture with FCCN (but without the same objective), it would be interesting for the ablation study in should include FCCN trained only with the g-loss. In Shalit (2017), it was found that representation sharing between potential outcome estimators made the biggest improvement to estimates (not changes to the learning objective). \n\n- Define what the PS abbreviation in the ablation study means already in 3.3.2 and describe what it does when the Assign-loss is not used. \n\n- Include description of hyper parameter selection in main paper and a short discussion of why this method of selection is appropriate here.  \n\n- In Figure 1, g0 and g1 appear to be functions of S, but in 3.2, they are described as functions of Z and X. In other words, in (5), it is not clear how g0 and g1 relate to S. \n\n**Would strengthen the work**\n- In the EDU experiments, it is not clear whether gains over baselines come from improvements in fitting the means (which are neural networks also in the true model) or the noise distribution. Given that the PEHE improvement is quite large compared to most baselines, I would assume that a large part is due to the mean.   This merits discussion. Additionally, please state whether the \"s\" variable is included in \"x\". \n\n- If the utility functions in 2.2. are established categories, consider adding a reference for this. Consider restructuring section 2 to better highlight what is the main problem studied in this work, what is motivation, etc. \n\n- Consider restructuring 3 to give the reader a better idea of why we end up with the loss in (4). For example, for a fixed z, this is quantile regression, and for the expectation, you optimise the loss w.r.t. every quantile (like Zhou et al., 2021). \n\n**Typos, minor corrections** \n- p,1, abstract: \"Standard causal inference assumptions\". Causal inference is a large subject, there are multiple sets of standard assumptions depending on the problem studied. \n- p.1, abstract: \"We demonstrate CCN\" -> \"We demonstrate that CCN\"\n- p.2, \"Empirically, these techniques often impose\". The assumptions aren't tied to empirical observation.\n- p.2, start of 2.1. \"We define the covariates as X \\in \\cX...\". This is not a definition per se, moreover X are random variables on \\cX, not elements of it, no?\n- p.2 \"We adopt this rigor in definition\". I'm not sure what this is referring to here. \n- p.3 Start of 3. \"Distributions Y(0)|X\". Previously, this was denoted \"p(Y(0)|X)\"\n- p.4 \"Zhou et al, (2021) shows that\" -> \"... show that\"\n- p.4 \"ground truth CDF functions for T \\in {0,1} is Lipchitz\". is -> are\n- p.9 \"respectfully\" -> \"respectively\"\n- The Figure 4 caption doesn't say which task is considered. \n- Weakening/contextualizing Claim 2 in main contributions would better represent its reliance on previous work. \n- Equation (3) is also due to Zhou et al (2021), right?", "broader_impact_concerns": "None", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: 
The authors have substantially revised the paper and well addressed the reviewers' comments and concerns. The reviewers and I are happy to have the revised paper accepted.

With the revision, some minor textual errors have crept in that the authors should address in the camera-ready version:

1. In the new paragraph on off-policy evaluation in Sec 4, it seems like textual citations (\citet) are sometimes used when parenthetical citations (\citep) would seem better choices.

2. Something seems to have gone awry in the formulation of the first sentence in the second paragraph of the discussion section:

    "CCN and its variants are competitive with but do not best methods that ..."

    
Decision: Accept (poster)