Variational Auto-Encoder Architectures that Excel at Causal Inference
Authors: Negar Hassanpour
Russell Greiner
Venue:  (2021)
Keywords: Causal Inference
Generative Modelling
Distributional Shift

Abstract: This paper provides a generative approach for causal inference using data from observational studies. Inspired by the work of Kingma et al. (2014), we propose a sequence of three architectures (namely Series, Parallel, and Hybrid) that each incorporate their M1 and M2 models as building blocks. Each architecture is an improvement over the previous one in terms of estimating causal effect, culminating in the Hybrid model. The Hybrid model is designed to encourage decomposing the underlying factors of any observational dataset; this in turn, helps to accurately estimate all treatment outcomes. Our empirical results demonstrate the superiority of all three proposed architectures compared to both state-of-the-art discriminative as well as other generative approaches in the literature. 

Under review as a conference paper at ICLR 2021
VARIATIONAL AUTO-ENCODER ARCHITECTURES THAT EXCEL AT CAUSAL INFERENCE
Anonymous authors Paper under double-blind review
ABSTRACT
This paper provides a generative approach for causal inference using data from observational studies. Inspired by the work of Kingma al. 2014), we propose a sequence of three architectures (namely Series, Parallel, and Hybrid) that each incorporate their MI and M2 models as building blocks. Each architecture is an improvement over the previous one in terms of estimating causal effect, culminating in the Hybrid model. The Hybrid model is designed to encourage decomposing the underlying factors of any observational dataset; this in turn, helps to accurately estimate all treatment outcomes. Our empirical results demonstrate the superiority of all three proposed architectures compared to both state-of-the-art discriminative as well as other generative approaches in the literature.
INTRODUCTION
As one of the main tasks in studying causality (Peters et al. 2017 Guo al. 2018), the goal of Causal Inference is to figure out how much the value of a certain variable would change (i.e., the effect) had another certain variable (i.e., the cause) changed its value. A prominent example the counterfactual question (Rubin 1974, Pearl, 2009) "Would this patient have lived longer [and by how much], had she received an alternative treatment?". Such question often asked in the context of precision medicine, which attempts to identify which medical procedure t E 'T will benefit a certain patient x the most, terms of the treatment outcome E R (e.g., survival time).
A fundamental problem in causal inference is the unobservablity of the counterfactual outcomes (Hol- land 1986). That is, for each subject i, any real-world dataset can only contain the outcome of the administered treatment (aka the observed outcome: 34), but not the outcome(s) of the alternative treatment(s) (aka the counterfactual outcome(s)) - i.e., y for t E T {tt}. In other words, the causal effect is never observed (i.e., missing in any training data) and cannot be used train predictive models, nor can it be used to evaluated a proposed model. This makes estimating causal effects a more difficult problem than that of generalization in the supervised learning paradigm.
In general, we can categorize most machine learning algorithms into two general approaches, which differ in how the input features x and their target values y modeled (Ng & Jordan 2002):
Discriminative methods focus solely on modeling the conditional distribution p(y[z;) with the goal of direct prediction of y for each instance æ. For prediction tasks, discriminative approaches are often more accurate since they use the model parameters more efficiently than generative approaches. Most of the current causal inference methods are discriminative, including the Balancing Neural Network (BNN) (Jhhanson et al. 2016), CounterFactual Regression Network (CFR-Net) (Shalit et al. (017). and Cet-Net's extensions cf.. (Yao et al. 2018 Hassanpour & Greiner 2019 2020) as well as
Dragon-Net (Shi et al. 2019).
Generative methods, on the other hand, describe the relationship between x and y by their joint probability distribution p(x, v). This, in turn, would allow the generative model to answer arbitrary queries, including coping with missing features r using the marginal distribution p(x) or [similar to discriminative models] predicting the unknown target values y via pzy(z)) A promising direction forward for causal inference is developing generative models, using either Generative Adverserial Network (GAN) (Goodfellow et al. (0044) or Variational Auto-Encoder (VAE) (Kingma & Welling 2014 Rezende et al.] 2014). This has led to two generative approaches for causal inference: GANs for inference of Individualised Treatment Effects (GANITE) (Yoon et al. 2018) and Causal Effect

---

Under review as a conference paper at ICLR 2021
VAE (CEVAE) Louizos et al. (2017). However, neither of the two achieve competitive performance in terms of treatment effect estimation compared to the discriminative approaches.
Although discriminative models have excellent predictive performance, they suffer from two draw- backs: (i) overfitting, and (ii) making highlyccofideet predictions, even for instances that are "far" from the observed training data. Generative models based on Bayesian inference, on the other hand, can handle both of these drawbacks: issue (i) can be minimized by taking an average over the poste- rior distribution of model parameters; and issue (ii) can be addressed by explicitly providing model uncertainty via the posterior (Gordon & Hernández-Lobato 2020). Although the exact inference is often intractable, efficient approximations to the parameter posterior distribution is possible through variational methods. Here, we use the Variational uuto-EEcoddrr (VAE) (Kingma & Welling 2014 Rezende al. 2014) for the Bayesian inference component of our causal inference method.
Contribution: In this paper, we propose three interrelated Bayesian model architectures (namely Series, Parallel, and Hybrid) that employ the VAE framework to address the task of causal inference for binary treatments. We find that the best performing architecture is the Hybrid model, that is [partially] successful in decomposing the underlying factors of any observational dataset. This is a valuable property, as that means it can accurately estimate all all treatment outcomes. We demonstrate that these models significantly outperform the state-of of-the-art in terms of treatment effect estimation performance two publicly available benchmarks, as well fully synthetic dataset that allows for detailed performance analyses.
2 RELATED WORKS
CFR-Net Shalit et al. (017)) considered the binary treatment task and attempted to learn a representation space & that reduces selection bias by making Pr( (bx) (4=0) and Pr( (z:) [|=]|]]] as close to each other as possible, provided that ( retains enough information that the learned regressors (+(()) E (0, 1}} can generalize well on the observed outcomes. Their objective function includes L ye. htt (e(((Zz)) (Ø(z.))]. which is the loss of predicting the observed outcome for sample i (described as æz), weighted by wi
& + 2 where u = Pr( (r(t=1) 1). This is effectively setting 2n
w/ 2P(4) where Pr( t, is the probability of selecting treatment t, over the entire population.
DR-CFR Hassanpour & Greiner 2020)) argued against the standard implicit assumption that all of the covariates X are confounders (i.e., contribute both treatment assignment and outcome determination). Instead, they proposed a graphical model similar to that in Figure I and designed a discriminative causal inference approach accordingly built on top of the CFR-Net. Specifically, their model, named Disentangled Representations for CFR (DR-CFR), includes three representation networks, each trained with constraints to insure that each component corresponds its respective underlying factor. While the idea behind DR-CFR provides an interesting intuition, it is known that only generative models (and not discriminative ones) can truly identify the underlying data generating mechanism. This paper is a step in this direction.
Dragon-Net Shi al. (019)'s 's main objective was to estimate the Average Treatment Effect (ATE), which they explain requires two stage procedure: (i) fit models that predict the outcomes for both treatments; and (ii) find a downstream estimator of the effect. Their method is based on classic result from strong ignorability - i.e., Theorem 3 in (Rosenbaum & Rubin 1983)
that states:
(y' v0) L (v'20 r 1 D(x)
& &
Pr(t 1|x) 10 Pr(t = 1 b(b))) E (0.1)
where 6(x:) a balancing score" They consider propensity score balancing score and argue that only the parts of X relevant for predicting T'' are required for the estimation of the causal effect This theorem only provides way to match treated and control instances though - i.e., it helps finding potential counterfactuals from the alternative group to calculate ATE, (Sh et al. (201)), however, used this theorem derive minimal representations on which to regress estimate the outcomes.
That is, 1. T b(X) (Rosenbaum Rubin 1983).
Thh authors acknowledge that this would hurt the predictive performance for individual outcomes. As result, this yields inaccurate estimation of Individual Treatment Effects (ITEs).
2

---

Under review as a conference paper at ICLR 2021
X
Z
Y!
yo
Underlying factors of any observa- tional dataset (Hassanpour Greiner 2020)
Figure 1:
Figure 2: Graphical model of CEVAE [Louizos al. 2017)
GANITE Yoon al. (2018) proposed the counterfactual GAN, whose generator G, given {z, t.y² y²) estimates the counterfactual outcomes (ú ());); and whose discriminator tries to identify which of {[x, v, [xr, 1, 1, } is the factual outcome. It is, however, unclear why this requires that G must produce samples that are indistinguishable from the factual outcomes, especially as D can just learn the treamment selection mechanism instead of distinguishing the factual outcomes from counterfactuals. Although this work is among the few generative approaches for causal inference, our empirical results (in Section 4)
show that it does effectively estimate counterfactual outcomes.
CEVAE Louizos et al. 2017) used VAE to extract latent confounders from their observed proxies in X. While this an interesting step in the right direction, empirical results show that does not always accurately estimate treatment effect (see Section 4. The authors note that this may be because CEVAE is not able address the problem of selection bias. Another reason that we think contributes to CEVAE's sub-optimal performance is its assumed graphical model of the underlying data generating mechanism (depicted in Figure 2) This model assumes that there is only one latent variable Z (confounding 7 and Y) that generates the entire observational data; however, we know from (Kaagg et al. 2017) and (Hassanpour & Greiner 2020) that there must more (see Fig
R2: M1 and M2 VAEs In an attempt to enhance the conventional representation learning with VABs - referred as the MI model (Kingma Welling (214]) Rezende et al. 2014)) a semi- supervised manner, Kingma et al. (2014) proposed the M2 VAE. While the M1 model helps learning latent representations from the covariate matrix alone, the M2 model allows the target information also to guide the representation learning process. In work, the target information includes the treatment bit as well as the observed outcome Y This additional information helps learning more expressive representations, that was possible with the unsupervised M1 model. Appendix A. presents a more detailed overview of the MI and M2 VAEs.
3 METHOD
Following (Hassanpour & Greiner 2020) and without loss of generality, assume that the random variable X follows an unknown joint probability distribution Pr( X A, T,E), where r, A, T, and E are non-overlapping independent factors. Moreover, assume that treatment T follows Pr( ((Z'IT A) (i.e., and A are the responsible factors for selection bias) and outcome Y" follows Pr, (Y") A, )) see Figure Observe that the factor N (resp., Y) partially determines only T (resp., Y), but not Y (resp., T); and A includes the confounding factors between 7 and Y.
Our goal is to design generative model architectures that encourage learning disentangled repre- sentations of these four underlying latent factors (see Figure [ In other words, is an attempt to decompose and separately learn the underlying factors that are responsible for determining T and
Y. To achieve this, we propose three architectures (as illustrated in Figures 3(a) 3(b) and 3 each employing VAE (Kingma & Welling 2014] Rezende et al. 2014) that each include a decoder (generative model) and an encoder (variational posterior). Specifically, we use the MI and M2 models from (Kingma et al., (2014) as our building blocks, leading to Series architecture, a Parallel architecture, and Hybrid one. Each component is parametrized deep neural network.
3.1
THE VARIATIONAL. AUTO-ENCODER COMPONENT
3.1.1
THE SERIES ARCHITECTURE
The architecture of the Series model is illustrated in Figure 3(a) Louizos et al. (2015) proposed a similar architecture address faimesss in machine learning, but using binary sensitive variable S
3

---

Under review as a conference paper at ICLR 2021
Z
y
Y1
ZE
Z
yo
Y!
Z
Z
Z
T
Z
Z
Z
Z
X
(a) The Series model
(b)
The Parallel model
(c) The Hybrid model
Figure 3: Belief nets of the proposed architectures.
(e.g., gender, race, etc.) rather than the treatment T. Here, we employ this architecture for causal inference and explain why it should work. 8s We hypothesize that this structure functions as distillation tower: the bottom M2 VAE attempts to decompose (guided by T`) from À and T (cap- tured by 71): and the top M2 VAE attempts to learn A and T (guided by Y). Decoder and encoder components of the Series model (parametrized by 02 and Φ, respectively) involve the following distributions:
Priors
Likelihood
Posteriors
Pe. (#2)
Po. (xc (1)
qo. (2i (21 1)
Pe, (#1 22)
qs, (u|z1)
4s. (22 21)
The goal is to maximize the conditional log ligeliiel of the observed data (left-hand-sie of the following inequality) by maximizing the Evidence Lower BOund (ELBO; right-hand- hand-side) - iee.
Σ log ((.))t)).
2
Egs, (#1 (n,t) [log Po. (2:1211, t,]]
KL go, (z) po. (2)1,, KL gø. (z2ly, 21) po, (Ez) (2)
where KL denotes the Kullback-Leible divergence, Po, (22) is the unit multivariate Gaussian (i.e., ((O.I) I)), and the other distributions are parameterized deep neural networks.
3.1.2 THE PARALLEL ARCHITECTURE
84 The Series model is composed of two M2 stacked models. However, Kingma et al. 2014 showed that an MI+M2 stacked architecture learns better representations than an M2 model alone for downstream prediction task. This motivated us to design a double MI+M Parallel model; where one arm for the outcome to guide the representation learning via Zi and another for the treatment to guide that via Zz. This architecture is illustrated in Figure 3(b) We hypothesize that Zi would learn and Y, and Zy would learn I' (and perhaps partially A). Decoder and encoder components of Parallel model (parametrized by 0p and Φp respectively) involve the following distributions:
Priors
Likelihood
Posteriors
Pa [2) Pa, [21 W, x2)
pa, (24) Pa,, (2 24)
Po, [2] 23)
4eb (2) p. t) thh (y|21)
Lo, (2g x. v) Ie (1[23)
4o (22 21)
Lo (z1lt, (3)
Here, the conditional log- likelihoodd can be upper bounded by:
N
lo p(:: (.)
2
Beg (ar val w.t,)) log (22/ 23:]]
(3)
KL qo, (21 po, (21 22) ) KL g0p (2219, 21) pe, (22)
(4)
KI. go, (2 (2/ po, (2a(1, KI. go, (zul1, (3) po, (24)
(5)
4

---

Under review as a conference paper at ICLR 2021
3.1.3 THE HYBRID ARCHITECTURE
The final architecture, Hybrid, attempts to get the best capabilities of the previous two architectures. The backbone of the Hybrid model has Series architecture, that separates I (factors related to the treatment T: captured by the right module with Z3 as its head), from A and Y (factors related the outcome Y; captured by the left module with Z7 as its head). The left module, itself, consists of Parallel model that attempts to proceed one step further and decompose A from Y. This is done with the help of a discrepancy penalty (see Section 3.3). b. Figure 3(c) illustrates our designed architecture for the Hybrid model. Decoder and encoder components of the Hybrid model (parametrized by 0z
and Φb respectively) involve the following distributions:
Priors
Likelihood
Posteriors
po. (2)) po. [1 V. 22)
po, (26) Po, (2alt, z4) Po, (2s U, 26)
pdp (24)
po,
(x) 23, 27)
qeh. 27 x.,t
Qs (21
Joh (25 (7
qsh. 22g I. V)
Goa. (y)21, 25)
qd. (t|23)
po,, (77 21, 25)
4dh. (32 y, z1)
Gon. (zoly, 26)
qás. (z4 z3)
Here, the conditional log- likelikiiiood can be upper bounded by:
log ((2))t.)))))
Bag (53, (m,t,))) log po,, (2il 27:)
(6)
KL gon (21 Pe, (z1ly,, 22) ) - KL gon (22ly, 21) || pa,, (22)
(7)
KL( go. (2g (2, y) || Po. (zalt, 24) KKL gon (zalt, 23) || pos (21)
(8)
KL( gov (25|27) Paa (2sly- #c)) - KLL Cox (ze 25) || pa,, (zo)
(9)
KL go. (27122 || Pa.
(10)
The first term in the ELBO (i.e., right-hand-side of Equations ( (3), or (6) is called the Recon-
struction Loss (RecL) and the next term(s) (i.e.. Equation ( summation of Equations 4 and (
or summation of Equations (7 9 and (10) is referred to as the KL Rs Divergence (KLD).
Concisely, the ELBO can be written RecL KLD, which is to be maximized.
3.2 DISENTANGLEMER WITH 8-VAE
As mentioned earlier, we want the learned latent variables to be disentangled, to match our assumption of non-overlapping factors I, A, and Y. To ensure this, we employ the 3-VAE (Higgins al. 2017) which adds hyperparameter /8 as a multiplier of the KLD part of the ELBO. This adjustable hyperparameter facilitates a trade-off that helps balance the latent channel capacity and independence constraints (handled by the KL terms) with the reconstruction accuracy - i.e., including the B hyperarameter grants a better control over the level of disentanglement in the learned representations (Burgess et al. 2018). Therefore, the generative objective to be minimized becomes:
Lvae
ReeLL + 3- KLD
(11)
Although Higgins al. (017)) suggest the 33 to be set greater than in most applications, Hoffman ct al. (017)) show that having a B < weight on the KL term can be interpreted as optimizing the ELBO under an alternative prior, which functions as regularization term prevent degeneracy.
3.3 DISCREPANCY
Although all the three proposed graphical models suggest that T and 21 are statistically independent (see, for example, the collider structure (at X): T + X + Z₁ in Figure ((a]). an information leak is quite possible due to the correlation between the outcome y and treatment in the data. We therefore require an extra regularization term on qo (z in order to penalize the discrepancy (denoted by disc) between the conditional distributions of 21 given t= 0 versus given 1. To achieve this regularization, we calculate the disc using an Integral Probability Metric (IPM) (Mansour et al. 2009 that measures the distance between the two above-mentioned distributions:
Laiec
IPM {21 hai, {x1ha,,11
(12)
In this work, we use the Maximum Mean Discrepancy (MMD) (Gretton al. 2012
5

---

Under review as a conference paper at ICLR 2021
3.4 PREDICTIVE LOss
Note, however, that neither the VAE nor the disc losses contribute to training a predictive model for outcomes. To remedy this, extend the objective function include discriminative term for the regression loss of predicting y: GT
Cpee
Σ w C y/. De] N
1
(13)
where the predicted outcome Oi is derived as the mean of the o3 (3:(21e) posterior trained for the respective treatment t;; [[ Vs. 0t] is the factual loss (f.e., L2 loss for real-valued outcomes and log loss for binary-valued outcomes); and wi represent the weights #) that attempt account for selection bias. We consider two approaches in the literature to derive the weights: (i) the Population Based (PB) weights proposed in CFR-Net (Shalit et al. 2017| ; and (ii) the Context-Aware (CA) weights pro- posed by Hassanpour & Greiner (2019). Note that disentangling A from T is only beneficial when
using the CA weights, since we need just the A factors derive them (Hassanpour & Greiner
2020
3.5 FINAL MODEL(S)
Putting everything together, the overall objective function to be minimized is then:
J Spedd + Laisc + Y LVAE + A Reg
(14)
where Reg penalizes the model complexity.
This objective function is motivated by the work of |MCalllum et al. (006)), which suggested optimizing convex combination of discriminative and generative losses would indeed improve predictive performance. As an empirical verification, note that for = 0, the Series and Parallel models effectively reduce to -et.Net. However, our empirical results (cf., Section 4 suggest that the generative term in the objective function helps learning representations that embed more relevant information for estimating outcomes than that of & in CFR-Net.
We refer to the family of our proposed methods as VAE-CI (Variational Inference); specifically: {S, P. HJ--AE- for Series, Parallel, and Hybrid respectively.
for Causal
EXPERIMENTS, RESULTS, AND DISCUSSION
4.1 BENCHMARKS
Infant Health and Development Program (IHDP) The original IHDP randomized controlled trial was designed evaluate the effect of specialist home visits on future cognitive test scores of premature infants. Hill (011) induced selection bias by removing a non-random subset of the treated population. The dataset contains 747 instances (608 control and 139 treated) with 25 covariates. We use the same benchmark (with 100 realizations of outcomes) provided by and used in (Johansson jet al. 2016) and (Shalit et al. (017).
Atlantic Causal Inference Conference 2018 (ACIC'18) ACIC' 18 is collection of binary- treatment datasets released for a data challenge. Following (Shi et al. 2019), we use subset of the datasets with instances N 51,5, 10} 10x (four datasets in each category). The covariates matrix for each dataset involves 177 features and is sub-sampled from table of medical measurements taken from the Linked Birth and Infant Death Data (LBIDD) (MacDorman & Atkinson. 1998), that contains information corresponding to 100,000 subjects.
Fully Synthetic Datasets We generated set of synthetic datasets according to the procedure described in (Hassanpour & Greiner 2020); see Section A.2 for an overview. We considered all the viable datasets in mesh generated by various sets of variables, of sizes mr, m., my {0, 4, 8} and mg 1. This creates 24 scenarios that consider all possible relative sizes of the factors T, A,
This similar to the way (Kingma et al. 2014] included classification loss in their Equation (9). "There are 33 = 27 combinations in total; however, removed three of these combinations that generate pure noise outcomes- iee. A A Y= 0: (0, 0). (4, 0), and (8, 0, 0).

---

Under review as a conference paper at ICLR 2021
S-VAE-CI
21
22
Z3
24
25
26
27
21
22
0.2599
1,2512
1.1979
0.3109
0.2890
0.2958
0.3407
0.6433
0.7704
0.5202
0.5662
0.8229
0.7073
0.7170
0.7925
0. 0.7828
0.8957
0.7929
0.2622
0.2641
0.5679
0.7106
0.8204
P-VAE-CI
DR-CFR
[_rep
Y.rep
0,2920
0.2922
0.6374
0.6676
0.5228
0..173
0,8412
0.7799
0.7563
0.1882
0. 0.1888
0.5429
0.7038
0.8402
0.7686
0.7735
0.4738
0.3737
0,0544
0.6749
1.0008
Figure 4: Performance analysis for decomposition of the underlying factors on the synthetic dataset with A,,, =8, m===
Table 1: IHDP (100 realizations) benchmark
Table ACIC'18 (N S 101) benchmark
METHOD
PEHE
CATE
METHOD
PEHE
CATE
CFR
0.75 (0.57)
0.08 (0.10)
CFR
5.13 (5.59)
1.21 (1.81)
DR-CFR
0.65 (0.37)
0.03 (0.04)
DR-CFR
(3.39).33)
0.80 (1.41)
DRAGON
NA
0.14 (0.15)
DRAGON
NA
0.48 (0.77)
GANITE
2.81 (2.30)
0.24 (0.46)
GANITE
3.55 (2.27)
0.69 (0.65)
CEVAE
(3.47)
0.18 (0.25)
CEVAR
5.30 (5.52)
3.29 (3.50)
S-VAE-CI
0.51 (0.37)
0.00 (0.02)
S-VAE-CI
2.73 (2.39)
0.51 (0.82)
P-VAE-CI
0.52 (0.36)
0.01 (0.03)
P.VAE- CI
2.62 (2.26)
0.37 (0.75)
H-VAE-CI (PB)
0.49 (0.36)
0.01 (0.02)
H-VAE- CI (PB)
1.78 (1.27)
0.44 (0.77)
H-VAF-CI (CA)
0.48 (0.35)
0.01 (0.01)
H-VAE-CI (CA)
1.66 (1.30)
0.39 (0.75 (0.75)
PEHE and <AE measures (lower is better) represented in the form of "mean (standard deviation)".
and Y. For each scenario, we synthesized multiple datasets with various initial random seeds in order allow for statistical significance testing of the performance comparisons between various methods.
4.2 EVALUATING IDENTIFICATION OF THE UNDERLYING FACTORS
To evaluate the identification performance of the underlying factors, we use fully synthetic dataset with mr -ur=m0 ma = my = and mg m==1. 1. We set x to be one of the four dummy vectors Vi.4 and input to each trained representation network Zj. Three of these vectors had "1" in the 8 positions associated with r, A, and Y respectively, and the remaining 17 positions of each vector was filled with "O". The fourth vector all "1" except for the last position (the noise) which was "0". This helps measure the maximum amount of information that is passed to the final layer of each representation network.
We let Oz be the elu output (here, E 22200) of the encoder network Zj when =. The average of the 200 values of Oi, (Aug(Oi.)) represents the power of signal that was produced by the Zj channel on the input V.. The values shown in Figure 4 tables Acg(O2.3) and Avg(Os.5) divided by Aug(O+.)) for each of the learned representation networks. Note that, larger ratio indicates that the respective representation network Z, has allowed more of the input signal Vi to pass through. Section A.3 includes more details on this procedure
the ratios of Aug(O
As expected, Z3 and Z4 capture I (e.g., the Zı ratios for r in the (P, H)- tables are largest), and Z1, Zz. Zo. Zo. and Z7 capture Á and Y. Note that decomposition of A from T has not been achieved by any of the methods except for H-VAE-CI, which captures T by Z1 and A by 2n, se. (note the ratios are largest for Z1 and Zz). This decomposition is vital for deriving context-awan importance sampling weights because they must be calculated from A only (Hassanpour & Greiner 2020). Also observe that both (P. HJ-VAE- CI are able to separate I from A. However, DR-CFR, which tried to disentangle all factors, failed not only to disentangle A from Y, but also I from A.
4.3 EVALUATING TREATMENT EFFECT ESTIMATION
Evaluation of treatment effect estimation is often done with semi- or fully- synthetic datasets that include both factual and counterfactual outcomes. There
categories of performance measures:
Unlike the evaluation strategy presented (Hassanpour & Greiner 2020) that only looked at the first layer's weights of each representation network, we propagate the values through the entire network and check how much of each factor is exhibited in the final layer of every representation network. RP: Yet, the proposed procedure still crudely evaluates the quality of disentanglement of the underlying factors. We did explore using the Mutual Information (Belghazi et al. 2013 8 for this task (not shown here); however, it appears that it does not work for high-dimensional data such as ours. in all, more research is needed to address this task.
7

---

Under review as a conference paper at ICLR 2021
Individual-basedd baase:: "Precision in Estimation of Heterogeneous Effect" PEHE
Σ N
(é, (é.-c;)²
uses è; =/ oن as the estimated effect and e; = y/ yL as the true effect (Hill
2011 and
baalat:oo--bsed: "Bias of the Average Treatment Effect" 6ATE |ATE- ATE, where ATE
1 N
1- V. - ビ vf and ATE is calculated based on the estimated outcomes.
In this paper, we compare performances of our proposed methods (S, P. H.- versus the following treatment effect estimation methods: CFR-Net (Shalit et al. 2017). DR-CER (Hassanpour
& Greiner 2020), Dragon-Net (Shi et al. 2019), GANITE (Yoon et al. 2018), and CEVAE (Louizos et all]]_201 2017), 7). The basic search grid for hyperparameters of the CFR-Net based algorithms (including our proposed methods) is available in Section A.4] For the other algorithms, we searched around their default hyperparameters' setting. We ran the experiments for the contender methods using their publicly available code-bases; note the following points:
Since Dragon-Net is designed estimate ATE only, we did not report its performance results for the PEHE measure (which, as expected, were significantly inaccurate).
The original GANITE bode-base was implemented for binary outcomes only. We modified the code (losses, etc.) such that it could process real-valued outcomes also.
We were surprised that CEVAF diverged when running on the ACIC' 18 datasets. To avoid this, we had to run the ACIC' 18 experiments on the binary covariates only.
Tables 1 and 3 summarize the mean and standard deviation of the PEHE and EATE measures (lower is better) on the IHDP, 18, and Synthetic benchmarks respectively. VAE-CI achieves the best performance among the contending methods. These results are statistically significant (in bold; based on Welch's unpaired t-test with 0.05) for the IHDP and Synthetic benchmarks. Although VAE-CI also achieves the best performance on the ACIC 18 benchmark, the results are not statistically significant due to the high standard deviation of the contending methods' performances.
Figure F visualizes the PEHE measures on the entire synthetic datasets with sample size of N= 10.000.
We observe that both plots corresponding to H- L-VAE-CI method (PB well CA) are inscribed by plots of all other methods, showcasing H-VAE-Cl's superior performance under every possible
selection bias scenario. Note that for scenarios where ma 0 (i.e., the ones of form my _0_ 0_m on perimeter of the radar chart in Figure 5). the performances of H- HAE-AICI (PB) and H- VAE CI (CA) are almost identical. This is expected, since for these scenarios, the learned representation for A would be degenerate, and therefore, the context- aware weights would reduce to population- ones. On the other hand, for scenarios where mA + O, the H-VAE-CI (CA) often performs better than H-VAE-CI (PB). This can be attributed to the fact that H- VAE-CI has correctly disentangled A from Y. This facilitates learning good context-aware weights that better account for selection bias, which in turn, results in better causal effect estimation performance.
We also performed hyperparameters' sensitivity analyses in terms of PEHE (see Figure 6). We discuss the results in the following:
For the a hyperparameter (i.e., coefficient of the discrepancy penalty), Figure 6(a) suggests that DR-CFR and H-VAE-CI methods have the most robust performance ki &t. throughout various values of a. This is expected, because, unlike CFR and (S, P)-VAE-CI, DR-CFR and H-VAE-CI possess independent node for representing A. This helps them still capture A grows; since for them, a only affects learning a representation of Y.
Comparing H-VAE-CI (PB) with (CA), we observe that for all a 0.01, H-VAE-CI (CA) outper- forms H- VAE-CI (PB). This is because the discrepancy penalty would force Z1 to only capture T and Zg to only capture A. This results in deriving better CA weights (that should be learned from A; here, from its learned representation Zg). H-VAE-CI (PB), on the other hand, cannot take advantage of this disentanglement, which explains its sub-optimal performance.
According to Figure 6(b) various & values (i.e,, coefficient of KL divergence penalty) R/ & R3 Re
do not make much difference for H- VAE-CII (except for 3 > 1; since this large value means the learned representations will be close to Gaussian noise). We initially thought using V-E might
further disentangle the underlying factors. However, Figure ((b) suggests that close-to-zer zero or even zero 8s also work effectively. We now think that the H- VA-V Cl's architecture itself sufficiently decomposes the I, A, and Y factors, without needing the help of a KLD penalty. Appendix A.5 includes more evidence and a detailed discussion on why this interpretation should hold.
8

---

Under review as a conference paper at ICLR 2021
0.0.4
Table 3: PEHE and EXTE measures (lower is better) represented in the form of "mean (ston- dard deviation)" on the entire synthetic bench- mark (average performance of 24 x datasets, each with sample size of 10,000).
RURL
C_0_8
R_8,4
045
0,0
RAA
0.2
0,1
-VAAECI
--VA-
HIVE (19))
8.4.0
0,8.4
METHOD
PEHE
CATE
8,0,8
0
-I--TT
CFR
0.39 (0.08)
0.027 (0.020)
8,0.1
4,0.4
DR-CFR
0.26 (0.07)
0.007 (0.004)
4,0,
4,0_8
DRAGON
NA
0.007 (0.005)
4,8.4
4
GANITE
1.28 (0.43)
0.036 (0.015)
4.8
4,4
4
CEVAE
1.39 (0.32)
0.287 (0.217)
Radar graphs of PEHE (on the radii; lower is better) for the entire synthetic benchmark (24 x with N = 10,000; each vertex denotes the respective dataset). Figure is best viewed in color.
Figure 5:
S-VAE-CI
0.28 (0.05)
0.004 (0.003)
P-VAE-CI
0.28 (0.05)
0.004 (0.003)
H-VAE- CI (PB)
0.20 (0.03)
0.003 (0.002)
H-VAE- CI (CA)
0.18 (0.02) 0.003 (0.002)
05
64
64
04
83
62
62
64
Lovee
Le
0.011
0.001 0.01
(a) C
(b) /8
Figure 6: Hyperparameters' (x-axis) sensitivity analysis based on PEHE (y-axis) on the synthetic dataset with mr, A -8, ml 1. Legend is the same Figure 5 Plots are best viewed in color.
& For hyperparameter y (i.e. coefficient of the generative loss penalty), H-VAE-CI achieves the most stable performance compared to the (S. P|- VAE-CI models - see Figure 6 Of particular interest is the superior performance of VAE-CI for y S 0.01 compared to that of [J. VAE-CI. This means that having the generative loss term (i.e. Σ is more important for (S, PJ-VAE-CI than for H-VAE-CI to perform well note extreme case happens 7 O. where the latter performs significantly better than the former. We hypothesize that this is be- cause H-VAE-CI already learns expressive representations Z3 and Z7, meaning the optimization no longer really requires the LVA term to impose that. This is in contrast to Z1 in S- S-VAE-CI and Z1 and Zg in P-V-ECC
5 FUTURE WORKS AND CONCLUSION
Despite the success of the proposed methods, especially the Hybrid model, in addressing causal inference, no known algorithms can yet learn to perfectly disentangle factors A and T. RI: This goal is important because we know isolating A, and learning context-award weights from it, does enhance the quality of the causal effect estimation performance note the superior performance of H-VAE-CI (CA). RI 83: The results of our ablation study (in Figure 6(b). however, revealed that the currently used 8- VAE does not help with disentanglement of the underlying factors. This shows that we can attribute all the decomposition we get to the proposed architectures and objective function.
A future direction of this research is to explore the use of better disentangling constraints e.g.. works Chen et al. (2018) and Lopez et al. (2018) to see if they would yield sharper results.
The goal of this paper was to estimate treatment effects (either for individuals or the entire population) from observational data. We designed three VAE-based (Kingma & Welling [0144] Rezende et al.. 2014) architectures (namely Series, Parallel, and Hybrid), that employed (Kingma al. 2014)'s MI and M2 models. The Hybrid model, as the best performing architecture, partially succeeded at decomposing the underlying factors r, A. and T; which helped in accurate estimation of treatment outcomes. Our empirical results demonstrated the superiority of the proposed methods, compared to both state-of-the-ar discriminative well as generative approaches in the literature.

---

Under review as a conference paper at ICLR 2021
REFERENCES
Mohamed Ishmael Belghazi, Aristide Baratin, Sai Rajeshwar, Sherjil Ozair, Yoshua Bengio, Aaron Courville, and Devon Hjelm. MINE: Mutual information neural estimation. JCML, 2018.
Christopher Burgess, Irina Higgins, Arka Pal, Loic Matthey, Nick Watters, Guillaume Desjardins, and Alexander Lerchner. Understanding disentangling in 3-VAR. arXiv preprint: 0008.004409999 2018.
Ricky TQ Chen, Xuechen Li, Roger Grosse, and David Duvenaud. Isolating sources of disentanglement in variational autoencoders. In NeurIPS, pp. 2610-2620 2018.
Ian Goodfellow, Jean Pouget-Abadie, Mehdi Mirza, Bing Xu, David Warde-Farley Sherjil Ozair, Aaron Courville, and Yoshua Bengio. Generative adversarial nets. In NeurIPS, 2014.
Jonathan Gordon and José Miguel Hernández-Lobato. Combining deep generative and discriminative models for bayesian semi-supervised learning. Pullern Recognition, 100:107156, 2020.
Arthur Gretton, Karsten M Borgwardt, Malte J Rasch, Bernhard Schölkopf, and Alexander Smola. A kernel two-sample test. JMLR, 1 13(Mrchh), 2012.
Ruocheng Guo, Lu Cheng, Jundong Li, P. Richard Hahn, and Huan Liu. A survey of learning causality with data: Problems and methods. arXiv preprint: 8809,09337, 2018.
Negar Hassampour and Russell Greiner. Counterfactual regression with importance sampling weights. In IJCAI, 2019.
Negar Hassanpour and Russell Greiner. Learning disentangled representations for counterfactual regression. In ICLR, 2020.
Irina Higgins, Loic Matthey, Arka Pal, Christopher Burgess, Xavier Glorot, Matthew Botvinick, Shakir Mohamed, and Alexander Lerchner. 3-VAE: Learning basic visual concepts with constrained variational framework. ICLR, 2017.
Jennifer L Hill. Bayesian nonparametric modeling for causal inference. Journal of Computational and Graphical Statistics, 20(1). 2011.
Matt Hoffman, Carlos Riquelme, and Matthew Johnson.
The 3 V-E''ss implicit prior. 2017. URL
http:/ /ttt:::///ayest/ andeepl ea rn ng. http://////////a//////// 666 .dff]
Paul Holland. Statistics and causal inference. Journal of the American statistical Association, 81 (396):945-900, 1986.
Fredrik Johansson, Uri Shalit, and David Sontag. Learning representations for counterfactual inference. In ICML, 2016.
Diederik P Kingma and Jimmy Lei Ba. Adam: A method for stochastic optimization. In ICLR, 2015.
Diederik Kingma and Max Welling. Auto-encoding variational bayes. In ICLR, 2014.
Diederik P Kingma, Shakir Mohamed, Danilo Jimenez Rezende, and Max Welling, Semi-supervised learning with deep generative models. In NeurIPS, 2014.
Kun Kuang, Peng Cui, Bo Li, Meng Jiang, Shiqiang Yang. and Fei Wang. Treatment effect estimation with data-driven variable decomposition. In AAAI, 2017.
Francesco Locatello, Stefan Bauer, Mario Lucie, Gunnar Raetsch, Sylvain Gelly, Bernhard Schölkopf, and Olivier Bachem. Challenging common assumptions in the unsupervised learning of disentan- gled representations. In ICML, 2019.
Romain Lopez, Jeffrey Regier, Michael I Jordan, and Nir Yosef. Information constraints on auto- encoding variational bayes. In NeurIPS, pp. 6114-6125, 2018.
Christos Louizos, Kevin Swersky, Yujia Li, Max Welling, and Richard Zemel. The variational fair autoencoder. arXiv 151.000::15110008900 2015.
10

---

Under review as a conference paper at ICLR 2021
Christos Louizos, Uri Shalit, Joris M Mooij, David Sontag, Richard Zemel, and Max Welling. Causal effect inference with deep latent-variable models. In NeurIPS. 2017.
Marian F MacDorman and Jonnae O Atkinson. Infant mortality statistics from the 1996 period linked birth/infant death dataset. Monthly Vital Statistics Report, 46(12), 1998.
Yishay Mansour, Mehryar Mohri, and Afshin Rostamizadeh. Domain adaptation: Learning bounds and algorithms. arXiv preprint:0902.3440 2009.
Andrew McCallum, Chris Pal, Gregory Druck, and Xuerui Wang. Multi-conditional learning: Generativedissriminaatie training for clustering and classification. In AAAI, pp. 433-439, 2006.
Andrew Y Ng and Michael Jordan. On discriminative vs. generative classifiers: A comparison of logistic regression and naive bayes. In NeurlPS, pp. 841-848, 2002.
Judea Pearl. Causality. Cambridge University Press, 2009.
Jonas Peters, Dominik Janzing, and Bernhard Schölkopf. Elements of causal inference: foundations and learning algorithms. MIT press, 2017.
Danilo Jimenez Rezende, Shakir Mohamed, and Daan Wierstra. Stochastic backpropagation and approximate inference deep generative models. JCMI., 2014.
Paul R Rosenbaum and Donald B Rubin. The central role of the propensity score in observational studies for causal effects. Biometrika, 1983.
Donald B Rubin. Estimating causal effects of treatments in randomized and nonrandomized studies. Journal of Educational Psychology, 66(5), 1974.
Uri Shalit, Fredrik D. Johansson, and David Sontag. Estimating individual treatment effect: General- ization bounds and algorithms. In JCML, 2017.
Claudia Shi, David Blei, and Victor Veitch. Adapting neural networks for the estimation of treatment effects. NeurlPS, 2019.
Liuyi Yao, Sheng Li, Yaliang Li, Mengdi Huai, Jing Gao, and Aidong Zhang. Representation learning for treatment effect estimation from observational data. In NeurlPS, 2018.
Jinsung Yoon, James Jordon, and Mihaela van der Schaar. GANITE: Estimation of individualized treatment effects using generative adversarial nets. In ICLR, 2018.

---

Under review as a conference paper at ICLR 2021
A APPENDIX
A.1 MI AND M2 VARIATIONAL AUTO-ENCODERS
As the first proposed model, the MI VAE the conventional model that is used learn representations of data (Kingma & Welling 2014 Rezende et al. 2014). These features are learned from the covariate matrix X only. Figure 7(a) illustrates the encoder and decoder of the MI VAE. Note the graphical model on the left depicts the encoder; and the one the right depict the decoder, which has arrows going the other direction.
Proposed by Kingma al. (2014), the M2 model was attempt to incorporate the information in target Y into the representation learning procedure. This results in learning representations that separate specifications of individual targets from general properties shared between various targets. In case of digit generation, this translates into separating specifications that distinguish each digit from writing style or lighting condition. Figure 7(b) illustrates the encoder and decoder of the M2 VAE.
We stack the M1 and M2 models as shown in Figure 7 to get the best results. This way, can first learn a representation Z from raw covariates, then find second representation Zz, now learning from Z1 instead of the raw data.
9
Z
Z
Z
X
X
X
(a) MI model
(b) M2 model
(c) MI+M3 model
Figure 7: Decoders (parametrized by 8) and encoders (parametrized by Φ) of the M1, M2, and
M1+M2 VAEs.
A.2 PROCEDURE OF GENERATING THE SYNTHETIC DATASETS
Given as input the sample size N; dimensionalities [mrr, ma, my] E 2+)3), for each factor L {T, A, }. the means and covariance matrices (H4,, Σr.); and a scalar that determines the slope of the logistic curve.
For each latent factor L E T. form by drawing instances (each of size my) from N(HL Σz). The covariates matrix X is the result of concatenating r, A, and Y. Refer to the concatenation of I and A as and that of A and Tas & (for later use).
For treatment T' sample mr+ma tuple of coefficients 0 from Wo0, 1) Detine the logging policy To (
1-exp(-(2)* where = 0. For each instance 2,, sample treatment f; from the Bernoulli distribution with parameter To 1=1
For outcomes y and Y sample m. +mr tuple of coefficients 80 and Ø from from.N(O. (0, mm Define y" = ( o o t + 0.5) 80 /m/s +mrr + è and y = (Φ $). /0/mmm//////////mmmm + + e, where white noise sampled from (00 0.1) and is the symbol for element-wise product.
A.3 EVALUATING IDENTIFICATION OF THE UNDERLYING FACTORS
Here, elaborate the procedure followed to evaluate identification performance of the underlying factors. We produced four dummy vectors V; E t +y + as depicted on the left-side of Figure The first third vectors had ones (constant) in the positions associated with T,
12

---

Under review as a conference paper at ICLR 2021
A, and Y respectively, and the remainder of them were filled with zeroes. The fourth vector was all ones, so we can measure the maximum amount of information that is passed to the final layer each representation network.
4
g
V1: 11 1 00 0
.....
Z2: 00 ... 0 11 ...
V,
O
Vz: 00 ... 0 00
11...
44 11.. 1 11 11...
Figure 8: The four dummy z-like vectors (left); and the input/output vectors of the representation networks (right).
In the next step, each vector Vi is fed to each trained network Zj. and the output O; is recorded (see the right-side of FigureBB, The average of O; represents the power of signal that communicated from V; and passed through the Zj channel. The values reported the tables illustrated Figure 4 the ratios of (average of O₁, O2, O₂} divided by (average of O4 for all the learned representation networks.
A.4 HYPERPARAMETERS
For all CFR, DR-CFR, and VAE-CI methods, we trained the neural networks with 3 layers (each consisting 200 hidden neurons] non-lineaa activation function elu, regularization coefficient of A=1E-4, Adam optimizer (Kingma & Ba 2015) with learning rate of 1E-3, batch size of 300, and maximum number of iterations of 10, 000. See Table Al for our hyperparameter search space.
Table 4: Hyperparameters and ranges
Hyperparameter
Range
Discrepancy coefficient a {0, -2,-3,-2,-110,
KIDD coefficient 3
{0, 0,,
Generative coefficient y
{0, 1,,,,,11
A.5 A DETAILED ANALYSIS OF THE EFFECT OF
Our initial hypothesis in using 3-VAE was that it might further disentangle the under- lying factors, in addition to the other constraint already in place (lee. the architecture as well as the discrepancy penalty). However, Figure 6(b) suggests that close-to-zero or even zero 3s also work effectively. To further explore this hypothesis, we examined the decomposition tables (similar to Figure 4 of H-VAE-CI for extreme configurations with B = 0 and observed that they were all effective decomposing the underlying factors I', A, and T (similar to the performance reported in green table Figure Figure 9 shows several these tables.
Our interpretation of this observation is that the H-VAE-CIS architecture already takes care of decomposing the I, A, and T factors, without needing the help of a KLD penalty. This means either of the following happening: (i) 3-VAE is not the best performing disentangling method and other disentangling constraints should be used instead - e,g., works Chen et al. (2018) and Lopez et al. 22018) or (ii) it is thcoretically impossible to achieve disentanglement without some
5RR1 & R2: In addition to this basic configuration, we also perform our grid search with an updated number of layers and/or number of neurons in each layer. This makes sure that all methods enjoy similar model complexity.
13

---

Under review as a conference paper at ICLR 2021
H-VAE-CI
71
22
23
24
25
76
27
0.2185
1.7711
1.771
0.2164
0.2190
00.099
0.6041
006051
0.6142
0.6308
0.8688
0.8315
0.6138
Y
0.8523
0.8552
0.3321
10 0834
0.7552
0.7859
0.7384
H-VAE-I
r
0.2242
0.2182
0.7439
0.6770
0.2373
0.2583
0.2189
0.22612
Y
0.5385
0.5430
0.3211
0.3333
0.434
0.4412
0.4571
Z3
24
25
26
0.7090
A
0.6438
0.6461
00.2750
00.312
0.7452
0.7569
0.8237
H-VAE-C
23
24
25
r
0.3646
0.3643
1 1457
0.8659
0.3942
0.4069
0.3166
0.5127
0.5307
0.6463
0.7016
0.6717
Y
0.5565
0.5780
0.4260
0 3964
0.4119
0.4234
0.4534
23
24
25
0.8821
0.8752
5.05
0.3326
0.3309
0.4006
4
1.2542
1.2480
0..8843
0.3488
0.8533
0.8568
0.9392
1.923
0.4997
H-VAE-C
23
24
25
0.0850
0.0875
1.5791
1.711
0.1464
0.1459
0.1833
0.6107
0.60885
0.6142
0.6388
0.7851
0.7937
0.6877
Y
0.8349
0.8242
0.3321
0 0.3834
0.5177
0.5073
0.6832
Figure 9: Decomposition tables for H-VAE-CI with 3=0.
supervision Locatello et al. 2019 which might not be possible to provide in this task, Exploring these options is of the scope of this paper and is left to future work.
14


Reviews:
[{"review_id": "_nf3URHCD5", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "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": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "The authors provide some novel VAE architectures for causal inference. They compare different architectures and assess their respective strengths and weaknesses. They empirically test the models on synthetic and real-world data and compare them to some baselines.\n\nMajor comments:\n- The Hybrid model seems to perform best, but it also has the largest number of parameters and latent variables. What would happen if one would introduce even more layers of latent variables? Would the performance get even better? Is there a tradeoff between performance and complexity?\n- In Figure 6 (a,b), it seems like the best method (H-VAE-CI) works very well with alpha=0 and beta=0. How can that be explained? If those loss terms are not needed, what is it that actually makes the model better?\n- Moreover, in Fig. 6 (c), the H-VAE-CI also still has a considerable gap to the other models at gamma=0, which means completely without VAE loss. How can that be explained?\n\nMinor comments:\n- The KLD is the Kullback-Leibler *divergence*, which is crucially *not a distance* (because it's not symmetric)\n- Could a more performant disentangling VAE be used? The beta-VAE is generally not the best [1]\n- Generally, how do these disentanglement results relate to the impossibility theorem in [1]?\n- How is the IPM chosen for L_disc? If it's the MMD, how is the kernel chosen? Does it depend on the task?\n- In Table 3, the DR-CFR seems better than the S-VAE-CI and P-VAE-CI. Why are they boldened?\n\nSummary:\nThe idea is very interesting, but the experiments are not fully convincing. Especially the ablation studies (Fig. 6) seem to suggest that most of the loss terms are not actually needed for the proposed model (H-VAE-CI) to perform well. I think further experimentation is needed to fully understand what really makes the proposed model work better than the baselines. For this, it could also be interesting to construct even \"deeper\" models than the (H-VAE-CI), which would be less informed by prior assumptions, and see if they might not work even better.\n\nUpdate: Thanks to the additional experiment and discussions, I have increased my score.\n\n\n[1] Locatello, F., Bauer, S., Lucic, M., Raetsch, G., Gelly, S., Schölkopf, B., & Bachem, O. (2019, May). Challenging common assumptions in the unsupervised learning of disentangled representations. In international conference on machine learning (pp. 4114-4124).\n", "extra_scores": {}, "extra_text": {}}, {"review_id": "Tf8nQ5s5YNY", "reviewer": "AnonReviewer4", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "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": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "The paper proposes three VAE architectures for ATE estimation.\n\nThe approach is motivated by CEVAE which uses a VAE to learn a single latent representation of confounding between the treatment, target and covariates, but attempts to disentangle the confounding between the covariates and the treatment, confounding between the covariates, treatment and target, confounding between the covariates and the target and covariates-only information. \n\nTo attempt to to this, 3 different hierarchical VAE architectures are proposed which involve stacking M1 and M2 units. They also include an MMD term in the loss to encourage independence between one of the latent representations and the target.\n\nThe proposed method is evaluated using synthetic data and on two real datasets. In both cases, the proposed method outperforms the SoTA methods, but the results are only significant in one case.\n\nThe paper is generally well written. The problem and general motivation are clearly stated. However, the specific motivations for the 3 specific architectures tested could be more rigorously explained and backed up experimentally.\n\nGiven that this is more of an architecture paper, the evaluation of the specific architecture choices could be more thorough. The attempt to account for the contributions of the different latent factors is appreciated, but the results are not strongly convincing. It seems clear that Z3/Z4 captures Gamma, but it’s not clear that any of the other architecture choices make a difference. The performance results do not seem to support that they are (they do not appear to be significantly different).\n\nIt’s also not clear what to conclude regarding the hyper parameter sensitivity. It would be interesting to see whether alpha affects the ATE estimation, i.e. whether the MMD term is making a difference.\n\nIn summary, the paper proposes 3 architectures for ATE estimation. It is well written and presents reasonable ideas, but the novelty is somewhat limited (it combines existing ideas and architectures), the performance advantage over the existing SoTA is not strongly convincing and the specific architecture choices could be better motivated and evaluated experimentally.\n\n--Post-rebuttal--\n\nI have increased my score in response to the authors' clarifications and additional experiments and updates added to the paper.", "extra_scores": {}, "extra_text": {}}, {"review_id": "1heHh-sxYd", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 7, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "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": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "Summary: Some generative models have been proposed for causal effect estimation but they often do not have a competitive performance. Recent work suggested that a combination of generative and discriminative model may improve treatment estimation with observational data, and further suggests a generic latent variable model for factorizing selection bias, as well as outcome. The author(s) build on this work and propose a set of deep generative models, with a hybrid objective function (generative + discriminative), that outperforms current approaches for ATE. \n\nThe proposed family of models is appealing. The proposed objective function is motivated from previous literature. I think it is always a shame (but alright) that one must add those beta constraints on VAEs in order to get sensible results (see my question below). Still, the experimental section makes a strong point that the method is effective and outperform previous approaches.\n\nQuestion:\n1. What are the results of the hyperparameter search for beta? In particular, if the result of beta is small (or zero), can the model still be called a VAE, as it should degenerate as an autoencoder. In that case, the approximate posterior may have near zero variance and the model barely capture any uncertainty. I wish to understand (a) whether the resulting model is still able to really generate data, from the prior predictive density and (b) what would be the performance of the model with a fixed beta, or a beta annealing as an ablation study. \n2. Would it be more advantageous to use other disentanglement constraints instead of the beta-VAE? [1, 2]\n\n[1] https://arxiv.org/pdf/1802.04942.pdf\n[2] https://arxiv.org/abs/1805.08672\n\n ", "extra_scores": {}, "extra_text": {}}, {"review_id": "1IUHSnQWqOh", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 7, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "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": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "Summary:\n\nThis paper introduces a new VAE architecture for performing\ncausal inference. It shows superior precision on estimating\nthe heterogeneous effect and lower bias in estimating the\ntreatment effect.\n\nClarity:\n\nThe paper was fairly clearly written and straightforward to follow.\nI think a small amount of explanation of Kingma's M1 and M2 models\nwould have be improved the flow as your model does build on these\nsemi-supervised architectures.\n\nTechnical Quality:\n\nThe method and experiments are well-thought out and very clearly\nevaluate the work.\n\nWhat's not clear to me is how much the extra performance is from just\nhaving a larger model and more parameters that can be fit. The\nappendix says there are 3 layers of 200 variables each but the hybrid\narchitecture clearly has more parameters. Are those distributed\nappropriately? It's also unclear how much the other architectures were\ndesigned to estimate the heterogeneous effect.\n\nSignificance and Originality:\n\nThe work does feel more incremental than anything but it is an\nimportant problem and there is value in estimate treatment\neffect with higher accuracy.\n\n\nMinor notes:\n\n* First sentence of Section 3. \"follows a(n unknown)\" -> \"follows an unknown\"\n\n* Some of the figures don't look like they are vector graphics\n\n* The code doesn't seem intended to be executed. Hopefully if a full implementation can be made available later", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: The authors suggest a VAE model for causal inference. The approach is motivated by CEVAE (Louizos et al., 2017) which uses a VAE to learn a latent representation of confounding between the treatment, target, and covariates. This paper goes beyond this approach and tries to design generative model architectures that encourage learning disentangled representations between different underlying factors of variation inspired by Hassanpour & Greiner (2019). 

The reviewers agreed that the topic will be of interest to a large group of readers. While the first version of the papers raised questions about the experimental design, several questions on the architecture design were addressed during the rebuttal period (e.g., deeper architectures). Other improvements were suggested and not adopted (e.g., alternative methods to achieve better disentanglement). The ablation studies seem to suggest that some of the loss terms are not actually needed and that non-probabilistic autoencoders (beta=0) also work well. We recommend aiming at improving the writing quality and coverage of more background material on the proposed architectures and causal factors. 

Decision: Reject