Targeted VAE: Structured Inference and Targeted Learning for Causal Parameter Estimation
Authors: Matthew James Vowels
Necati Cihan Camgoz
Richard Bowden
Venue:  (2021)
Keywords: causal inference
variational inference
disentanglement
variational autoencoder

Abstract: Undertaking causal inference with observational data is extremely useful across a wide range of domains including the development of medical treatments, advertisements and marketing, and policy making. There are two main challenges associated with undertaking causal inference using observational data: treatment assignment heterogeneity (i.e., differences between the treated and untreated groups), and an absence of counterfactual data (i.e., not knowing what would have happened if an individual who did get treatment, were instead to have not been treated). We address these two challenges by combining structured inference and targeted learning. To our knowledge, Targeted Variational AutoEncoder (TVAE) is the first method to incorporate targeted learning into deep latent variable models. Results demonstrate competitive and state of the art performance.


Under review as a conference paper at ICLR 2021
TARGETED VAE: STRUCTURED INFERENCE AND TAR- GETED LEARNING FOR CAUSAL PARAMETER ESTIMA- TION
Anonymous authors Paper under double-bling blind review
ABSTRACT
Undertaking causal inference with observational data is extremely useful across a wide range of domains including the development of medical treatments, ad- vertisements and marketing, and policy making. There are two main challenges associated with undertaking causal inference using observational data: treat- ment assignment heterogeneity (i.e., differences between the treated and untreated groups), and absence of counterfactual data (i.e. not knowing what would have happened an individual who did get treatment, were instead to have not been treated). We address these two challenges by combining structured inference and targeted learning. To our knowledge, Targeted Variational AutoEncoder (TVAE) is the first method incorporate targeted learning into deep latent variable mod- els. Results demonstrate competitive and state of the art performance.
INTRODUCTION
The estimation of the causal effects of interventions or treatments on outcomes is of the upmost importance across range of decision making processes and scientific endeavours, such as policy making (Kreif & DiazOrdaz (019). advertisement (Bottou et al. 2013). the development of medical treatments (Petersen al. 2017), the evaluation of evidence within legal frameworks (Pearl 2009
Siegerink et al. 2016) and social science (Vowels, 2020: Hernan 2018 Grosz et al.] (020). Despite the common preference for Randomized Controlled Trial (RCT) data over observational data, this preference is not always justified. Besides the lower cost and fewer ethical concerns, observational data may provide a number statistical advantages including greater statistical power and increased generalizability (Deaton & Cartwright 2018). However, there two main challenges when deal- ing with observational data. Firstly, the group that receives treatment is usually not equivalent to the group that does not (treatment assignment heterogeneity), resulting in selection bias and con- founding due to associated covariates. For example, young people may prefer surgery, older people may prefer medication. Secondly, we are unable directly estimate the causal effect of treatment, because only the factual outcome for a given treatment assignment is available. In other words, we do not have the counterfactual associated with the outcome for different treatment assignment to that which was given. Treatment effect inference with observational data is concerned with find- ing ways to estimate the causal effect by considering the expected differences between factual and counterfactual outcomes.
We seek to address the two challenges by proposing a method that incorporates targeted learning techniques into disentangled variational latent model, trained according to the approximate max- imum likelihood paradigm. Doing so enables us to estimate the expected treatment effects, as well as individual-level treatment effects. Estimating the latter is especially important for treatments that interact with patient attributes, whilst also being crucial for enabling individualized treatment assignment. Thus, propose the Targeted Variational AutoEncoder (TVAE), undertake abla- tion study, and compare our method's performance against current alternatives on two benchmark datasets.

---

Under review as a conference paper at ICLR 2021
Pro-rreatment Covariates
Zk
Zu
Assigmment
Outcome
y
(8)
(b)
(c)
Figure 1: Directed Acyelic Graphs (DAGs) for (a) the problem of estimating the effect of treatment t on outcome y with confounders x. DAG (b) reflects an RCT. DAG (c) illustrates TVAE and an extension of the DAG by Zhang et al. (222)), where the structure is priori assumed to factorize into into risk zy. instrumental 2,, and confounding factors Ze. We extend their model with z, to account for the scenario whereby not all covariates will be related to treatment and/or outcome.
2 BACKGROUND
Problem Formulation: A characterization of the problem of causal inference with no unob- served confounders is depicted in the Directed Acyclic Graphs (DAGss shown in Figs. Ta) and [b). Fig. Ta) is characteristic of observational data, where the assignment of treatment related to the covariates. Fig. Db) is characteristic of the ideal RCT, where the treatment is unrelated to the covariates. Here, X, p(x) R represents the m-dimensional, pre-treatment covariates for individual i assigned factual treatment & ~ p(t[x) resulting in factual outcome yf ply|x, t). Together, these constitute dataset D = {[ye, L, x:] where N is the sample size.
The conditional average treatment effect for individual with covariates x; may be estimated as (x,) = Elys xx, do = 1) Y/ |x; do(t t = 0)], where the expectation accounts for the non- determinism of the outcome (Jesson et al. 2020). Alternatively, by comparing the post-intervention distributions when we intervene on treatment t. the Average Treatment Effect (ATE) is f(x) =
Ex [ do(t 1)] E[y] |x, do(t O)]] Here, do(t) indicates the intervention on t. setting all instances to a static value, dynamic value, or distribution and therefore removing any dependencies originally had (Pearl) 2209 van der Laan Rose, 2018 2011). This scenario corresponds with the DAG in Fig. [b), where treatment t is no longer function of the covariates x.
Using an estimator for the conditional mean Qtt, x) t, t, x), we can calculate the Average Treatment Effect (ATE) and the empirical error for estimation of the ATE (eATE) In order to estimate cATE we assume access to the ground truth treatment effect т, which is only possible with synthetic semi-synthetic datasets. The Conditional Average Treatment Effect (CATE) may also be calculated and the Precision in Estimating Heterogeneous Effect (PEHE) is one way to evaluate model's ' efficacy in estimating this quantity. See the appendix for the complete definitions of these terms.
The Naive Approach: The DAG in Fig. aa) highlights the problem with taking naive approach to modeling the joint distribution p(y, t. x). The structural relationship t. < x - indicates both the covariates x, and that backdoor path exists
that the assignment of treatment is dependent through x to y. In addition to our previous assumptions, if we also assume linearity, adjusting for this backdoor path a simple matter of adjusting for x by including it in logistic regression. The naive method is an example of the uppermost methods depicted in Fig. 2, and leads to the largest bias. The problem with the approach is (a) that the graph is likely misspecifted such that the true relationships between covariates well as the relationships between covariates and the outcome may be more complex. There also problem (b), that linearity is not sufficient to 'let the data speak' (van der Laan & Rose 2011) or to avoid biased parameter estimates. Using powerful nonparametric models (e.g., neural networks) may solve the limitations associated with linearity and interactions to yield a consistent estimator for p(y]x), and such a model is example of the middlemost methods depicted in Fig. 2 However, this estimator is not targeted to the estimation of the causal effect parameter т, only predicting the outcome, and require means to reduce residual bias.
Targeted Learning: Targeted Maximum Likelihood Estimation (TMLE) (Schuler Rose 2016 van der Laan & Rose 2011 2018 van der Laan & Starmans 2014 falls under the lowermost
For a binary outcome variahle y (0, 1). (R x) is the same as the conditional probability distribution x))
2

---

Under review as a conference paper at ICLR 2021
Non-targeted misspecified parametric
estimator (e.g. logistic regression)
Non-targetod misspecified non-/semi- paramelni estinator (e.g. DNN)
Datu
Targeted non- pammm--a estimator
Figure 2: Different methods for estimating the causal parameter T yield different levels of bias. Adapted from vaa der Laan & Rose 2011).
methods depicted in Fig. 2 and follows an approach involving three main steps: (1) estimation of the conditional mean Ely it. x) with estimator Q (t, x). (2) estimation of the propensity scores with estimator g(|)x), and (3) updating the conditional mean estimator Q" to get Q* using the propensity scores to attain an estimate for the causal parameter т.
The propensity score for individual i is defined as the conditional probability of being assigned treatment g(t, x,) = p(t = f,(x x:), E [o, 1] (Rosenbaum Rubin, 1983). The scores can be used to compensate for the relationship between the covariates and the treatment assignment using Inverse Probability of Treatment Weights (IPTWs), reweighting each sample according to its propen- sity score. Step (3) is undertaken using 'clever covariates' which are similar to the IPTWs. They form an additional covariate variable H(1, x,) = [1{x)) for individual i assigned treatment, and (0,, x;) -g(1 | for individual i not assigned treatment. Note that when condition single numeric value we imply an intervention (e.g. g(1 \x:) = g(do(t = )),))) A logistic regres- sion is then undertaken as y
is the logit/inverse sigmoid function, Q" x) is set to be a constant, suppressed offset and represents fluctuation parameter which to be estimated from the regression. Once has been estimated, we acquire an updated estimator:
(tQ x)] + eH (t, x) where
('o(t = x),x) = 10 + e/H
(1)
This equation tells us that our new estimator Q² equal the old estimator balanced by the correc- tive (t, term. This term adjusts for the bias associated with the propensity scores. When the parameter zero, means that there no longer any influence from the 'clever covariates' H() The updated estimator Q' can then be plugged into the estimator for + +(Q x). When the optimal solution is reached (i.e. when e 0), the estimator also satisfies what is known as the efficient
Influence Curve (IC), or canonical gradient equation (Hampel 1974 van der Laan & Rose 2011 Kennedy 2016):
N Σ 1C (v,t)xx = Σ [((t,,,,)() ((t,.x))) -( +(t,xx))
-(0,x,,) =(Q:x)] (2)
where TC(y,, 14, x,) represents the IC, and TC TC*(y,, 44, x.) represents the efficient IC for consistent Q and g. It can be seen from the right hand side Eq. 2 that at convergence, the estimator and its estimand are equal: y/ (t,, xx) and Q(1, x;) Q0, x,) 7(Q;* x). Over the whole dataset, all terms in Eq. 2
'cancel' resulting in the mean IC = 0. As such, the logistic regression in Eq. represents a solution to the IC via a parametric submodel.
The TMLE method provides a doubly robust, asymptotically efficient estimate of the causal or 'target' parameter, and these theoretical guarantees make it attractive for adaptation into neural networks for causal effect estimation.
3 METHODOLOGY
In this section we present the Targeted Variational AutoEncoder (TVAE), a deep generative latent variable model that enables estimation of the average and conditional average treatment effects (ATE and CATE resp.) via the use amortized variational inference techniques and Targeted Maximum Likelihood Estimation (TMLE). For review the relevant VAE theory, the appendix. A top-
3

---

Under review as a conference paper at ICLR 2021
0 (Z)
p(t 22. 2.)
p(2.)
p(x Zop Zen Bay 200
((2]
p(2))
g(2, x)
plyft -0. 2,)
g(2,/x)
p(z.)
= p(t la., 2,)
4.
plyle 1.
M it 1, 24y)
p(z,)
p(y 0, 2ny)
Figure 3: The block-diagram for Targeted VAE. Dashed boxes indicate the variationally inferred latent variables Zo: Ze. Ze, and Zy. Arrows indicate functions, and colors distinguish treatment (green), outcome (blue), covariates (grey), and confounder (purple) related entities.
level diagram for TVAE is shown in Fig. 3 and follows the structure implied by the DAG in Fig. Ic). A more detailed architectural block-diagram is also presented in the appendix
Assumptions: As is common (aao et al. 2020 Guo et al. 2020 Rubin 2005 Imbens & Rubin
2013) when undertaking causal inference with observational data, we make three assumptions: (I) Stable Unit Treatment Value Assumption (SUTVA): the potential outcomes for each individual or data unit are independent of the treatments assigned to all other individuals, such that there are no interactions between individuals.
(2) Positivity: the assignment of treatment probabilities are all non-zero and non-deterministic p(t = fx|x =x,) > O, V t and x. (3) Ignorability: all confounders are observed such that the likelihood of treatment for two individuals with the same covariates is equal, and the potential outcomes for two individuals with the same covariates are also equal
s.t. (y" yy' IL t/x and TT (y'm1." .y
If one had knowledge of the true causal DAG underlying a set of data, one could undertake causal inference without being concerned for issues relating to structural misspecification. Unfortu- nately, and this is particularly the case with observational data, we rarely have access to this knowl- edge. Quite often an observed set of covariates x are modelled as a group of confounding variables (as per the DAG in Figure (p). Furthermore, and as noted by Zhang et al. (2020), researchers may be encouraged to incorporate as many covariates into their model as possible, in an attempt to reduce the severity of the ignorability assumption. However, including more covariates than is necessary leads to other problems relating to the curse of dimensionality and (in)efficiency of estimation.
TVAE:
A large set of covariates may be separable into subsets of factors such instrumental, risk, and confounding factors. Doing so helps us to match our model more closely to the true data generating process, well to improve estimation efficiency by 'distilling' Prior work has explored the potential to discover the relevant confounding covariates via Bayesian networks (Haggstrom) 2017), regularized regression (Kuang et al.) (017). and deep latent variable models based on Variational Autoencoders (VAEs) (Zhang et al. 2020: Louizos et al., (017). The first two methods identify variables (and are variable selection algorithms), whereas VAES infer them, and learn compact, disentangled representations of the observations. The benefit of the latter approach is that it (a) infers latent variables on a subsets from population aggregates), (b) under additional assumptions, VAEs have been shown to infer hidden confounders in the presence of noisy proxy variables, thereby potentially reducing the reliance on ignorability) (Louizos et al. 2017), and (c) makes no assumptions about the functional form used to map between covariate and latent space.
covariate adjustment set.
basis (rather than deriving
We seek to infer and disentangle the latent distribution into subsets of latent factors using VAEs. These latent subsets are (Z,, Zy, Ze, }.}. which represent the instrumental factors on t, the risk factors on y. the confounders on both and y, and factors solely related to x, respectively. Without inductive bias, consistently disentangling the latent variables into these factors would be impossible (Locatello et al. 2019). In TVAE this inductive bias is incorporated in number of ways: firstly, by incorporating supervision and constraining Z and n, to be predictive of t and respectively; secondly, by constraining ze to be predictive of both and y; and finally, by employing diagonal-
"Source code will also made available in supplementary material upon acceptance.
'Taken together, assumptions (2) and (3) constitute strong ignorability.
4

---

Under review as a conference paper at ICLR 2021
covariance priors (isotropic Gaussians) to encourage disentanglement and independence between latent variables. The structural inductive bias on the model is such that Zy, and 2,, and 2, learn factors relevant to outcome and treatment, for which we provide explicit supervision, thereby leaving Zo for all remaining factors.
In general, it is impossible isolate the effect of t - y due unobserved confounding (D'Amour 2019), and this is why we make the assumption of ignorability. However, it is worth noting that, under additional assumptions, deep latent variable techniques have been shown be able to infer hidden confounders from what are known as noisy proxy variables present in the dataset (see e.g., Montgomery al. 2000; Louizos et al. 2017). The assumption of ignorabilityy then shifts from 'all confounders are observed' to "all unobserved confounders have been inferred from proxies'. Whilst the capability to infer confounders from proxies represents an additional motivation for the use of VAEs, the focus of this work is not to explore whether and by how much we are able to do so, and we therefore maintain the assumption of ignorability.
Inference:
Drr II
Day
(22 xx
N(ua fid(x), fea(x)); ((zy)x) II N(aa = faa(x), =fad(x)) d=1
d=1
D.c
D...
<(zclx)= II
N N(a faa(x), f(a((x)): q(oolx>)
II N N fra(x), 02 foa(x))
p(olze no) = Been(Boooz k.)) Bern(gg(Zzs, 2c))
p(ý Byy Ze t) Bern(t Sio(Zy, k.) + (1-t) fii 2c)) = Bern Bern(Q,())
Generation:
((soutieiv)
p(z(otea)) II Me 1); p(t|z,, ze) Beem(hn ze)) (ern(gp(.))
p(ylay: Ze: i)) = Bern(t h2(2y, 2.) + (1 - i) 113(2y 2c)) = Bern(Qp(.)
(Xbin /hee Zoo Zz, Bny Bern(ho Zo; byy)
Decant
p(Rcont Bey 2o. 21, 2ny = II d=1
M(Zemttalla = ha(Zor Zn, 2,, 2y), o h5 Z,, Zz, Zy))
(3)
The proof for identifiability under the assumption of ignorability (or, alternatively, under the sumption that all unobserved confounders have been inferred from proxy variables) has been de- rived previously by Louizos al. (2017) and Zndng al. (2020). The factor 2, is d-separated from t and given x, and does not affect the identification of the causal effect. i.c.. plyldo(t), dot), 2ft,o.g.ct) p(y]t, Zy, ze) (see Zhang et al. 2020 and Louizos et 2017p. We impose the priors and parameterizations denoted in Equatron 3 where D(.) is the number of dimensions in the respective variable (latent or otherwise), and f1-11 and h1-6 represent fully connected neural network functions. The parameters for these neural networks are learnt via variational Bayesian approximate inference (Kingma Welling, 2014) according to the following objective:
CE.OO
Σ
F Baaaaa [logr (&: (2,, Zo 20) + logg (L./2r, 2.) + -log (Wil/i Zy: 2cc]
(4)
- [DKL (q (Z)(xa) 1p (z.)) + DKL (q (ze x, llp (zc))
+ Di (q (z,lx,) lp (2,)) + D L (g (z, x;) ||p (z.)))
Note that all Gaussian variance parameterizations are diagonal. cases where prior knowledge dictates a discrete rather than continuous outcome, equivalent parameterizations to those in Eqs. may be employed. For example, in the IHDP dataset, the outcome data are standardized to have variance of 1, and the outcome generation model becomes Gaussian with variance also equal to 1.
5

---

Under review as a conference paper at ICLR 2021
Note that separate treatment and outcome classifiers are used both during inference and generation (Qg. 94 and Qp. 9p resp.). The classifiers for inference have separate parameters to those use during generation. Predictors or classifiers of outcome incorporate the two-headed approach of (Shalit) et al. 2017). and use ground-truth t are used for Q2 whereas samples t are used for Qp. For unseen test cases, either the ground-truth used simulate an outcome. During training is used.
an sampled treatment from treatment classifier 9p may be
We now introduce the targeted regularization, the purpose of which is to encourage the outcome to we define the fluctuation sub-model be independent of the treatment assignment. Following Eq. [
and corresponding logistic loss for finding e as:
Qg. 2., 25 =a 0 (Q 22,
14
T(1: 0)
(5)
= 1: git: = 0; 2 2i) E (((:c)== y, log(Q( 1, 2 2,, -
(1 (1- log(1 Qg, 1g 2(.
(6)
In Eq. 5 I is the indicator function. For unbounded regression loss, mean squared error loss may be used (see appendix). Note that the logistic loss is suitable for continuous outcomes bounded between 0 and I (see van der Laan & Rose 2011, pp. pp.121:133 121:132 for proof). Putting all together, we then optimize to find generative parameters for functions h1-6. inference parameters for functions f1-12, and fluctuation parameter c as follows:
L min Σ
C CELBO + +xruS((...0)):
θ = IC" =0
de
===
Where XrL represents hyperparameter loss weight for the targeted regularization. At convergence, = 0 when Q and g become consistent estimators, satisfying the conditions for the EIC (see Eq. and reference van der Laan & Rose 2011 pp125:128).
A further element that differentiates our work from one other recent contribution (Shi et 2019) that uses targeted regularization is that the gradients resulting from are not taken with respect to Ip or By (which are the propensity score arms which we assume to be consistent and unbiased). Tar- geted learning is concerned with de-biasing the outcome classifier Q using propensity scores from g. In other words, assuming the propensity scores are well estimated, the targeted learning regularizer intended to affect the outcome classifier only, and nor the propensity score estimator. It is there- fore more theoretically aligned (with the targeted learning literature) to apply regularization to the outcome estimator Q. and not to g. As per Eq. D in TMLE, y is assumed be a consistent estimator, forming part of the de-biasing update process for Q. but it is not subject to update itself. In order to prevent the regularization from affecting the propensity arms, the gradients from the regularizer are only taken with respect to all parameters that influence this outcome classifier (which include upstream parameters for Q, as well as the more direct parameters Qp). We use Pytorch's 's 'detach' method on the propensity scores when calculating the targeted regularization. This method decou- ples the propensity score arm from backpropagation relating to the computation of the regularization value. In contrast, Dragonnet (Shi et al. 2019) applies regularization to the entire network.
In summary, the notable aspects our model are as follows: the introduction of new latent vari- able z, for factors unrelated to outcome and/or treatment to aid the recovery of the true underlying structure; the ability to estimate both individual level and average treatment effects; and, as far as we are aware, the first incorporation of targeted learning in a deep latent variable approach, and one which backpropagates the regularization gradients to specific outcome-related parameters in the network.
4 RELATED WORK
There are number of ways to mitigate the problems associated with the confounding between the covariates and the treatment. For a review on such methods, readers are pointed to the recent survey by Yao al. 2020). Here consider methods that utilize neural networks part of their models,
but note that many non-neural network methods exist (Chernozhukov et al., 2017 van der Laan &
Rose 2011 van der Laan & Starmans 2014 Rubin 2005 Hill 2011 Athey & Imbens 2016).
Perhaps the most similar works to ours are those of Dragonnet (Shi et al. 2019) and TEDVAE /hhng et al. 2020). We discuss the differences between these and TVAR in turn. Dragonnet

---

Under review as a conference paper at ICLR 2021
incorporates the same targeted learning regularization process which allows for the simultaneous optimization of Q and c. However, the method sacrifices the ability to estimate individual treatment effect in preference to achieving good estimation of the average treatment effect across the sample. Indeed, they do not report PEHE. Finally, Dragonnet applies regularization to the entire network, whereas we specifically 'target' the regularization the outcome prediction arm by restricting the backpropagation of gradients.
TEDVAE, on the other hand, builds on CEVAE (Louizos et al. 2017) and seeks inference and disentanglement of the latent instrumental, risk, and confounding factors from proxy variables with a variational approach. However, it has no means to allocate latent variables that are unrelated to treatment and/or outcome (i.e.. TVAE's s Zo). The advantage of including factors Zo with variational penalty is that the model has the option to use them, or not to use them, depending on whether they are necessary (i.e. KI. is pushed to zero). It is important not to force factors unrelated to treatment and outcome into Zi [c....)) because doing so restricts the overlap between the class of models that can be represented using TEDVAE, and the class of models describing the true data generating process,
Other methods include GANITE (Yoon et al. 2018) which requires adversarial training, and may therefore be more difficult to optimise. PM (Schwab et al.] 2019), SITE (Yao et al. 2018)) and Mul- tiMBNN (Sharma et al., 2020) incorporate propensity score matching. TARNET (Shalit et al. (017) inspired the two-headed outcome arm in our TVAE, well the three-header architecture in (Shi ett al. 2019). RSB (Zhang et al. 2019) incorporates a regularization penalty, based on the Pearson Correlation Coefficient, intended to reduce the association between latent variables predictive of treatment assignment and those predictive of outcome.
5 EXPERIMENTS
We perform an ablation study, beginning with (a) TVAE (base) which is equivalent to TEDVAE (b) TVAE with z.. and (c) TVAE with both with 2, and targeted regularization & during training. In order to fairly evaluate the benefits of introducing Zo, we ensure that the total number of latent dimensions remains constant. We undertake the ablation study
synthetic dataset which we call
TVAEsynth, before comparing against methods on both the IHDP dataset (Hill 2011 Gross 1993) and the Jobs dataset (LaLonde 1986 Smith & Todd 20055 Dehejia & Wahba 2002). In particular, the synthetic dataset was intentionally created such that not all covariates are exogenous and so that there exist some latent factors not related outcome or treatment. Thus, we should expect a significant improvement in performance to occur with the introduction of 2,, demonstrating the importance of incorporating inductive bias that closely matches the true structure of the data. Note that while these datasets vary in whether the outcome variable is continuous (IHDP, TVAESynth) or binary (Jobs), the treatment variable is always binary. Whilst it is possible to undertake Targeted Learning continuous treatment effects, we leave this to future work.
For the IHDP dataset, we evaluate our network on the Average Treatment Effect estimation error (cATE), and the Precision in Estimation of Heretogeneous Effect (PEHE). As per (Louizos et al. 2017 Shalit et al.] 2017] Yao et al.] 2018), for the Jobs dataset (for which we have only partial evaluate
network the Average Treatment effect on the Treated error (eATT) to approximate the eATE, and the policy risk Reot to approximate the error on the CATE. See the appendix for definitions of these metrics. When estimating treatment effects, 100 samples are drawn for each set of input covariates X;. We compare against Dragonnet (Shi et al.) 2019).
effect supervision)
TEDVAE (Zhang al. 2020), CEVAE (Louizos et 2017)) GANITE (Yoon et al. | (2018). Targeted Maximum Likelihood Estimation (TMLE) (van der Laan & Rose 2018), and TARNET + variants (Shalit et al. 2017). We provide results for both within sample and out-of-sample performance. is worth noting that within sample and out-of-sample results are equally valid for treatment effect
estimation, because the network never supervised on treatment effect (Shi et al. 2019).
For model selection we follow Louizos et al. (017)) and Zhang et al. (2020) and use the minimum validation loss on the total objective function. Whilst some model selection heuristics exist that serve as surrogates for the eATE itself (e.g., see Hassanpour & Greiner 2019 or Athhe & Imbens 2016) we take the same view as Zhang et al. (2020), namely that the development of our model "should be self-sufficien and not rely others". Furthermore, we undertake minimal hyperparameter tuning for the simple reason that, in real- world applications, the supervision required for effective tuning would not be available. For all experiments, we undertake 100 replications and provide mean and
7

---

Under review as a conference paper at ICLR 2021
standard error. See the appendix for details on these datasets and the architecture, well training and testing details.
Table 1: Means and standard errors for the ablation study using the TVAESynth dataset. Here, 'oos' means out-of-sample, 'ws' means within sample. "+z." indicates the introduction of the miscella- neous factors, "tz"' indicates the introduction of miscellaneous factors but without changing the dimensionality of Ze thereby increasing total latent capacity, '+5" indicates our targeted regulariza- tion (with selected backpropagation), "+5" indicatest targeted regularization equivalent to the one used by [hi et al. (2019) (i.e. with gradients applied all upstream parameters), the ç subscript indicates its weight in the loss.
Method
VEPEHE
VEPERE
EATÆ
EA7Е
TVAE (base/TEDVAE)
.179:E.003
.1781.003
.128:1.005
.128-.005
TVAE 2
.1744.003
.173:.003
.1214.005
.120+1.05
TVAE + Zo
.166-t.003
.166::003
.0691.004
.06911004
TVAE {x=0.1
.171#.003
.1704.003
.122..004
.121
TVAE + Z. + Ex=0.1
.111:.002
.150=.003
.0481.003
.088-.033
TVAE Zo + CA=0.1 (full model)
.150+t.002
.150+.003
.048+-0033
.048
Ablation Study TVAESyuth: The results for the ablation study on the synthetic dataset are shown in Table T The results demonstrate that both eATE and PEHE significantly improved by the incorporation of 2, or targeted regularization, with combination of the two yielding the best results for both within sample and out of sample testing. The fact that TVAE +2, outperforms TVAE +2 despite the latter having a larger latent capacity, suggests that reducing the capacity of the latent space has beneficial, regularizing effect. Based on the results of this ablation, the benefits of this regularizing effect appear to be distinct from the benefits that derive from the addition of miscellancous factors. Finally, the results indicate negligible empirical benefits to restricting the backgpropagation of the regularizer to non- proppenssitty related parameters. However, the restriction of the backpropagation (according to our implementation) more closely aligns with the original TMLE and efficient influence curve theory, and we therefore retain this feature for the remaining experiments.
IHDP Results: Results on IHDP are shown in Table 2 and indicate state of the art performance both within sample and out-of-sample eATE and PEHE. The results corroborate the ablation results in Table in that the incorporation of z, and targeted regularization result in monotonic improvement above TEDVAE. TVAE outperformed only by Dragonnet on within-sample cATE performance. However, this method does not provide estimations for indvidual CATE, and limited to the estimation of average treatment effects.
Table 2: Means and standard erTors for evaluation on the semi semi-syntheiis IHDP dataset Hill 2011
Results from: (Louizos al. 2017 Shalit al. 2017; Zhang al. 2020 Yoon al. 2018 Here, "oos" means oft-of-sample and 'ws' means within sample. "+ zo indicates the introduction of the miscellaneous factors, and indicates targeted regularization with subscript indicating weight in the loss.
Method
EPEEE WS
EEEEE 00S
CATE Ws
CΛTE 0OS
TMLE (van der Laan & Rose
2018
5.04.20
330
CEVAE (Louizos et al. 2207)
2.74.10
2.64.10
.34.001
46-t.220
TARNet (Shalit al. (017)
.88.000
.95+.00
.26+.01
.28..01
CFR-MMD Shaht et al. 2017
.734.00
.78+.00
.30.001
31+.01
CFR-Wass (Shalit
2017)
.71+.00
.764.00
.55.001
.24-001
TEDVAE /Zhng et al.
2020)
.62+-11
.63+.11
.00--05
GANITE Yoon et 2018)
1.94.40
2.44.40
.31.05
.494.05
Dragonnet w/ t-reg (Shi et al.
2019
.14.001
.20-.01
TVAE (w/ 2o. (x=0.0)
.574.03
.57+.03
.16.001
.16..01
TVAE (w/ Zp. Ex=0.4)
.52+.02
.54..02
.15.01
.16+.01
8

---

Under review as a conference paper at ICLR 2021
Table 3: Means and standard errors for evaluation the real-world Jobs dataset (Lal.onde) 1986
Smith & Todd 2005 Shalit et al. 2017 Yoon et al. 2018). Results taken from: (Louizos et al. 2017 Zhang et 2020). Here, "oos' means out-of-sample and 'ws' means within sample. '+ Zo
indicates the introduction of the miscellaneous factors, and '+ 5 indicates targeted regularization with subseript indicating its weight in the loss.
Method
Rpot W8
Ryw 00s
EA77 Ws
EATT OOS
TMLE van der Laan & Rose 2018)
.22+.00
.02::.01
CEVAE (ouuizos et al. 2017
.15+.00
.26+-00
.02+-01
00=001
TARNet (Shalit et al. 2017)
.17+.00
.21+000
.05+-02
.11 11=004 04
CFR-MMD (Shalit et al. 2017
.18+.00
.214.00
.04.001
.08=.03
CFR- Wass (Shalit et 2017
.17=:00
.21..00
.044.01
.09-.03
GANITE Yoon et al.
2018)
.13:.00
.14..00
.01..01
.06=.03
TEDVAE Zhang et
2020
.06+.00
.06=000
TVAE (w/ Zo: (x=o)
.16+.00
.164.00
.01..00
.01=.00
TVAE (w/ Zo. 1x=1)
.164.00
.164.00
.01
Jobs Results: The results for the Jobs data are shown in Table GANITE was found to perform the best across most metrics, although this method has been argued to be more reliant on larger sample sizes than others, on the basis that it performs relatively poorly on the smaller IHDP dataset Yoon et al.] 2018). Furthermore, GANITE relies potentially unstable/unreliable adversarial training (Moyer et al., 2018 Lezama 2019 Gabbay & Hosen 2019). Finally, TVAE outperforms GANITE on eATT, is consistent (beyond 2 decimal places) across out-of-sample and within-sample evaluations and has a lower standard err, and is competitive across all metrics. On this dataset, the concomitant improvements associated with the additional latent factors and targeted learning were negligible.
6 CONCLUSION
In this work we aimed to improve existing latent variable models for causal parameter estimation in ways: Firstly, by introducing a latent variable model factors unrelated to treatment and outcome, thereby enabling the model to more closely reflect the data structure; and secondly, by incorporating targeted learning regularizer with selected backpropagation to further de-bias out- come predictions. Our experiments demonstrated concomitant improvements in performance, and our comparison against other methods demonstrated TVAE's ability to compete with and/or exceed state of the art for both individual as well as average treatment effect estimation. For future work, we plan to explore the application of TVAE to longitudinal data with continuous or categorical treat- ment. There is also opportunity to explore the use of TVAE in inferring hidden confounders from proxies in the dataset, as well as interpreting the model to explore and validate what information inferred in the model's 's latent factors. Additionally, it was noted from the ablation study results that the restriction of regularization gradients did not yield a significant change in performance when compared with applying the regularization to the entire network. As well undertaking further experiments to understand this behavior, we propose to explore alternative ways to integrate the targeted learning update procedure the learning procedure. Finally, the ablation results indicated that part of the improvement associated with the introduction of z, is associated with regularizing effect relating the reduction the dimensionality of ze. This aspect of the model's behavior also lends itself to future exploration.
REFERENCES
A. A. Alemi, I. Fischer, J. V. Dillon, and K. Murphy. Deep variational information bottleneck. 1XXY:16612.0041007 2017.
S. Athey and G.W. Imbens. Recursive partitioning for heterogeneous causal effects. Proc. of the National Academy of Sciences of the USA, 113(27), 2016.

---

Under review as a conference paper at ICLR 2021
Eli Bingham, Jonathan P. Chen, Martin Jankowiak, Fritz Obermeyer, Neeraj Pradhan, Theofanis Karaletsos, Rohit Singh, Paul A. Szerlip. Paul Horsfall, and Noah D. Goodman. Pyro: Deep universal probabilistic programming. J. Mach. Learn. Res., 20, 2019.
D. M. Blei, A. Kucukelbir, and D. McAuliffe. Variational inference: review for statisticians. arXhv:16010006700 2018.
L. Bottou, J. Peters, Quinonero-Candelll J., D. X. Charles, D. M. Chickering, E. Portugaly, D. Ray, P. Simard, and E. Snelson. Counterfactual reasoning and learning systems: the example com- putational advertising. Journal of Machine Learning Research, 14, 2013.
C. Burgess, Higgins, A. Pal, L. Matthey, N. Watters, G. Desjardins, and A. Lerchner. Under- standing disentangling in Beta- VAE. Urxiv::804.0339900 2018.
V. Cheroozuukov, D. Chetverikov, M. Demirer, E. Dullo, C. Hansen, and W. Newey. Dou- ble/debiaseedNeeyman machine learning of treatment effects. American Economic Review, 5, 2017.
B. Dai, Z. Wang, and Wipf. The usual suspects? reassessing blame for VAE posterior collapse. arXiv:/912. 10702211, 2019.
A. D'Amour. On multi-cause causal inference with unobserved confounding: Counterexamples, impossibilitty and alternatives. Proceedings of the 22nd International Conference on Artificial Intelligence and Statisfics, 89, 2019.
A. Deaton and N. Cartwright. Understanding and misundertstanding randomized controlled trials. Social Science and Medicine, 210:2-21, 2018. doi: 10.101 socscimed..011..11
R. H. Dehejia and S. Wahba. Propensity score-matching methods for nonexperimental causal stud- ies. Review of Economics and Statistics, 84(1):151-661, 2002.
V. Dorie. Non-parametrics for causal inference. http:://giluub..oo//vrrrr
2016.
A. Gabbay and Y. Hosen. Demystifying inter-class disentanglement. 1XXi:1906.11796/1. 2019.
I. J. Goodfellow, J. Pouget-Abadie, M. Mirza, B. Xu, D. Warde-Farley, S. Ozair, A. Courville, and
Y. Bengio. Generative adversarial nets. arXiv: 4406. 2661, 2014.
R. T. Gross. Infant health and development program (IHDP): enhancing the outcomes of low birth weight, premature infants in the United States. MI: Inter-university Confortium for Political and Social Research, Ann Arbor, 1993.
M.P. Grosz, J.M. Rohrer, and F. Thoemmes. The taboo against explicit causal inference in nonexperimental psychology. Perspectives on Psychological Science, pp. 1-13, 2020. doi: (41177/174561160022155
R. Guo, Li, and H. Liu. Learning individual causal effects from networked observational data. Association for Computing Machinery, 2020.
J. Haggstrom. Data-driven confounder selection via Markov and Bayesian networks. Biometriks, 74(2):389-398, 2017.
F. R. Hampel. The influence curve and its role in robust estimation. Journal of the American Statistical Association, 69(346))383-393 1974.
N. Hassanpour and R. Greiner. Counterfactual regression with importance sampling weights. Pro- ceedings of the 28th International Joint Conference Artificial Intellgence (IJCAI-19), 2019.
N. Hassanpour and R. Greiner. Learning disentangled representations for counterfactual regression. ICLRR 2020.
M. Herman. The c-word: scientific euphemisms do not improve causal inference from observational data. American Journal of Public Health, 108(5):625 626, 2018. doi: 10.2105/AJPH.2113 304337.
10

---

Under review as a conference paper at ICLR 2021
1. Higgins, L. Matthey, A. Pal, C. Burgess, X. Glorot, M. Botvinick, S. Mohamed, and A. Lerchner. Beta- VAE: Learning basic visual concepts with a constrained variational framework. ICLR, 2017.
J. L. Hill. Bayesian nonparametric modeling for causal inference. Journal of Computational and Graphical Statistics, 20(1), 2011.
G.W. Imbens and D.B. Rubin. Causal inference for statistics, social, and biomedical sciences. An Introduction. Cambridge University Press, New York, 2015.
A. Jesson, S. Mindermann, U. Shalit, and Y. Gal. Identifying causal effect inference failure with uncertainty-aware models. wrXh::2007.001633 2020.
E. H. Kennedy. Statistical causal inferences and their applications in public health research, chapter Semiparametric theory and empirical processes in causal inference. Spinger, 2016.
D. Kingma and J. L. Ba. Adam: method for stochastic optimization. arXiv: / 141.:698029. 2017.
D. Kingma and M. Welling. Auto-encoding variational Bayes. urXiv:/312.61 2014.
N. Kreif and K. DiazOrdaz. Machine learning in policy evaluation: new tools for causal inference. arXiv: 1oXiv:1903.000 004000, 2019.
K. Kuang, P. Cui, B. Li, M. Jiang. S. Yang, and F. Wang. Treatment effect estimation with data- driven variable decomposition. AAAI, 2017.
R.J. LaLonde. Evaluating the econometric evaluations training programs with experimental data. The American Economic Review, pp. 604 604-620, 1986.
J. Lezama. Overcoming the disentanglement reconstruction trade-off via Jacobian supervision. ICLR, 2019.
F. Locatello, S. Bauer, M. Lucic, G. Ratsch, S. Gelly, B. Scholkopf, and Bachem O. Chal- lenging common assumptions in the unsupervised learning of disentangled representations. arXiv:1811.1235993 2019.
C. Louizos, U. Shalit, Mooij. D. Sontag, R. Zemel, and M. Welling. Causal effect inference with deep latent-variable models. 31st Conference on Neural Information Processing Systems, 2017.
J. Lucas, G. Tucker, R. Grosse, and M Norouzi. Understanding posterior collapse in generative latent variable models. ICI., 2019a.
J. Lucas, G. Tucker, R. Grosse, and M. Norouzi. Don't blame the ELBO! a linear VAE perspective on posterior collapse. arXiv: 1. 1.02469v1, 2019b.
M. Maier, K. Marazopoulou, D. Arbour, and D. David. A sound and complete algorithm for learning causal models from relational data. Proceedings of the 29th Conf. on Uncertainty Artificial Intelligence, 2013.
I. Mayer, J. Josse, and J. P. Raimundo, Fand Vert. MissDeepCausal. causal inference from incom- plete data using deep latent variable models. arXiv:2002.10837/1 2020.
M.R. Montgomery, M. Gragnolati, K.A. Burke, and E. Paredes. Measuring living standards with proxy variables. Demography, 37(2):155-774, 2000.
J. M. Mooij, O. Stegle, D. Janzing, K. Zhang, and B. Scholkopf. Probabilistic latent variable models for distinguishing between cause and effect. NIPS, pp. 1687-1695, 2010.
D. Moyer, S. Gao, R. Brekelmans, G. Steeg, and A. Galstyan. Invariant representations without adversarial training. NeurlPS, 2018.
H. Oktay. Using latent variable models improve causal estimation. PhD thesis, Doctoral Disser- tation 1205, 2018
J. Pearl. Causality. Cambridge University Press, Cambridge, 2009.
11

---

Under review as a conference paper at ICLR 2021
M. Petersen, L. Balzer, D. Kwarsimma, N. Sang, G. Chamie, J. Ayieko. J. Kabami, A. Owaraganise, T. Liegler, F. Mwangwa, and K. Kadede. Association of implementation of a universal testing and treatment intervention with HIV diagnosis, receipt of antiretroviral therapy, and viral suppression in East Africa. Journal of American Medical Association, 317(21)2219-2200 2017. doi: 10. (001/jama.001..5705
M. Rolinek, D. Zietlow, and G. Martius, Variational autoencoders pursue PCA directions (by acci- dent). arXiv: / /2.66755v, 2019.
P. R. Rosenbaum and D. B. Rubin. The central role of the propensity score in observational studies for causal effects. Biometrika, 70(1):41-55, 1983.
D. B. Rubin. Causal inference using potential outcomes: Design, modeling, decisions. Journal of the American Statistical Association, 100(469):322 2-212, 2005. doi: 10.1 10.1198/01221450400000 880.
M. S. Schuler and S. Rose. Targeted maximum likelihood estimation for causal inference in obser- vational studies. American Journal of Epidemiology. 185(1):65-73, 2016.
Schwab, Linhardt, and W. Karlen. Perfect match: a simple method for learning representations for counterfactual inference with neural networks. wwXiv:///0.00666:5 2019.
U. Shalit, F. D. Johansson, and D. Sontag. Estimating individual treatment effect: generalization bounds and algorithms. arxiv:1606. 03976v5, 2017.
A. Sharma, G. Gupta, A. Prasad, R.and Chatterjee, L. Vig, and G. Shroff. MultMBNN: matched and balanced causal inference with neural networks. arXiv:2004. 13446v2, 2020.
C. Shi, D. M. Blei, and V. Veitch. Adapting neural networks for the estimation of treatment effects. 33rd Conference Neural Information Processing Systems, 2019.
B. Siegerink, W. den Hollander, M. Zeegers, and Middelburg. Causal inference in law: an epidemiological perspective. European Journal of Risk Regulation, 7(1):155-186, 2016. doi: 10.1017/S1867799X0000
R. Silva, R. Scheine, G. Clark, and Spirtes. Learning the structure of linear latent variable models. Journal of Machine Learning Research, 7:191-246, 2006.
J. A. Smith and P. E. Todd. Does matching overcome LaLonde's critique of nonexperimental esti- mators? Journal of Econometrics, 125(1):305-353, 2005.
N. Tishby and N. Zaslavsky. arXiv:1503.02406111 2015.
Deep learning and the information bottleneck principle.
M. J. van der Laan and S. Rose. Targeted Learning Causal Inference for Observational and Experimental Data. Springer International, New York, 2011.
M. van der Laan and S. Rose. Targeted Learning in Data Science. Springer International, Switzer- land, 2018
M. J. van der Laan and R. J. C. M. Starmans. Entering the of data science: targeted learning and the integration of statistics and computational data analysis. Advances in Statistics, 2014.
M. J. Vowels. Limited functional form, misspecification, and unreliable interpretations in psychol- ogy and social science. arXiv:2009. 10025, 2020.
L. Yao, S. Li, Y. Li, M. Huai, Gao, and A. Zhang. Representation learning for treatment effect estimation from observational data. 32nd Conference Neural Information Processing Systems (NeurIPS), 2018.
I.. Yao, 7. Chu, S. Li, L., Gao, and A. Zhang. A survey on causal inference. arXiv:2002.02770, 2020.
Yoon, Jordan, and M. van der Schaar. GANITE: Estimation of individualized treatment effects using generative adversarial nets. ICLR, 2018
12

---

N Under review as a conference paper at ICLR 2021
W. Zhang, L. Liu, and J. Li. Treatment effect estimation with disentangled latent factors. arXiv:2001.06552, 2020.
Z. Zhang, Q. Lan, Y. Wang, N. Hassampour, and R. Greiner. Reducing selectin bias in counterfactual reasoning for individual treatment effects estimation. 33rd Conference Neural Information Processing Systems, 2019.
A APPENDIX
This appendix includes details the metrics used for evaluating treatment effect estimation; back- ground on Variational AutoEncoders; a formulation of the targeted learning regularizer for contin- uous, unbounded outcomes (as used with the IHDP dataset evaluations); details about the datasets; and training. testing, hyperparameters. architecture, and hardware details. Source code and data (including IHDP (Hill) 2011; Gross 1993). Jobs (LaLonde) 1986] [mith & Todd (005). and our own EVAESy will also be attached as supplementary material for reproducibility purposes upon publication.
A.1 METRICS
This section presents the metrics used for the experiments (Section 5 in the main paper).
The Average Treatment Effect (ATE) and error on the estimation of ATE (eATE) is given in Eq.
N Σ (f(Q:x))-((x)
+(Q;x)
E (11,x,,) Q(0,x,)) CATE = N
(8)
N
1
To estimate the error the model's capacity model the Conditional Average Treatment effect (CATE), the Precision in Estimation of Heteogenous Effects (PEHE) is given in Eq g
1 N Σ (+(Q:x)) - -(xx))²
CPEHE
(9)
It be seen from Equations 8 that the ATE is essentially the expectation of the conditional treat- ment effect (conditioned on the covariates for each individual) over the data set (Jesson et al. 2020) For scenarios when proportion of the dataset is from Randomized Controlled Trial (RCT), ), as is the case for the Jobs dataset, we may use the error on the estimation of Average Treatment effect on the Treated (ATT), which is given in Eq. T0 Shalit et al. 2017 Louizos et al. 2017:
CATT
U1
Uj JET
Σ FE
(Q(1,x) - Q(0,
(10)
T1
To
Ti
ET
where 7 = Tj U To constitutes all individuals in the RCT, and the subscripts denote whether or not those individuals were in treatment (subscript 1) or control groups (subscript O). The first two terms in Eq. comppise the true ATT, and the third term the estimated ATT. In datasets where the ground-truth for the CATE not available (as the case with the Jobs dataset) we may use the policy risk as proxy for PEHE:
Ryol = [xl) (xp(=(()) = + (=) = (x((=(x) 0))
(11)
where Txx) = is the policy to treat when / 1 of > c, and π (xx) = 0 is the policy not to treat otherwise Yao et al. [018 Shalit al. 2017). is treatment threshold. This threshold can be varied to understand how treatment inclusion rates affect the policy risk. For our experiments we set to zero, per (Shalit al. 2017 Louizos al. 2017).
13

---

Under review as a conference paper at ICLR 2021
A.2 VARIATIONAL AUTOENCODERS
This part of the appendix is referenced in Section 3 (Methodology) of the main paper.
We now consider the theory behind popular and powerful latent variable representation learning and density estimattion method known as Variational AutoEncoders (VAEs) (Kingma & Welling 2014). Although adversarial methods (Goodfellow et al. 2014) have been shown to be effective for density estimation, they are also troublesome to train. Hence we choose VAEs which have more stable training dynamics (Moyer et al. 2018). A further motivation for the use of VAEs relates to their ability to infer latent variables. The Ignorability assumption holds that all confounders are observed, and much causal inference is undertaken with this assumption (Mooij et al., 2010 Maier er al. 2013; Oktay 2018 Silva et al. 2006). The use of a latent variable model allows to infer unobserved/hiddenn confounders (Louizos et al., 2017). although the use of VAES means that there is some uncertainty concerning the guarantees of the learned model (Rolinek et al.) 2019 Dai al., 2019 Lucas et al.] 20.9a(b). This uncertainty notwithstanding, previous implementations for latent variable models with causal inference show promising results in comparison with other methods Louizos et al. 2017 Mayer et al. 2020. Hassanpour & Greiner 2020).
manifold via stochastic mapping p(z]x), where latent variables z provide a compact representation. The distribution po x) also serves as the posterior the generative model po (x) = J po (x (z) ())p())))) (zzzz having parameters 0. Marginalizing out z is usually intractable, and the true posterior p( [ (x) is unknown, so a simpler approximating posterior qo (z)x) having parameters 6 is introduced (Blei et al., 2018) such that:
We wish to encode the m-dimensional covariates
Pa (x(x)))((z) + log (z/xxx
go [0(z)x)
log (x) = E (x(e)) log
(12)
2 (2 x)
Taken together with the expectation, the second term on the right hand side Eq. 12 represents the Kullback- Liebler Divergence (KLD) between the approximating posterior and the true posterior. The KLD is always positive, and by ignoring this term, we are left with a lower-bound on the log- likelihood known as the variational lower bound, or Evidence Lower BOund (ELBO). The VAE provides the means to scale this amortized variational inference to intractable, high-dimensiional problems, and minimizes the negative log likelihood over a dataset of N samples by adjusting the parameters of neural networks {0, }} according to the ELBO.
N Σ
(x.
S N
(-Bgalelxs) lloge (x, (x:/z)] + 3DL [go (xx())x))l)x()))
(13)
where B 1 is used for the standard variational approximation procedure, but may be set empiri- cally (Higgins et al., 2017). annealed (Burgess et al. 2018) or optimized according to the Informa- tion Bottleneck principle (Alemi et al.) 2017: Tishby & Zaslavsky 2015). The first term in Eq. [3] is the negative log- likeliiol and is calculated in the form of a reconstruction error. The second term is the KLD between the approximating posterior and the prior, and therefore acts as prior regularizer. Typically, the family of isotropic Gaussian distributions is chosen for the posterior (.). and an isotropic Gaussian with unit variance for the prior p(z), which helps to encourage disentanglement Higgins et al. (017).
A.3 TARGETED REGULARIZATION FOR BOUNDED AND UNBOUNDED OUTCOMES
This part of the appendix is referenced in Section 3 (Methodology) of the main paper. In contrast with Eq. 6 in the main paper, this formulation of targeted regularization is for an unbounded. contin- uous outcome (as is the case in the IHDP experiments). A mean-squared error version of &, similar to the one found in (Shi et al. 2019), is given as follows:
T(4=11 1)
T(I =
Q(t,, 25, Q(t,, 24, +
(14)
g(t; = 1: 1; z()
g(t: = 0: 21, 2,)
(( 9: &otno 6) (yi Q(4,, 2f)))
(15)
14

---

Under review as a conference paper at ICLR 2021
Zo
X
X,
X
X
X
X
Ze
Zy
Figure 4: The DAG for TVAESynth dataset.
However, we stress that this formulation only appropriate for unbounded outcomes, and not con- tinuous outcomes in general (van der Laan & Rose [201T) For continuous, bounded outcomes, as well as binary outcomes, the NLL formulation (as in the main paper) has been shown be more appropriate, and theoretically sound.
A.4 DATASETS
This part of the appendix referenced in Section (Experiments) of the main paper.
We utilize 100 replications of the semi-synthetic Infant Health and Development Program (IIDP))
dataset (Hill 2011 Gross 1993 The linked version (see footnote) corresponds with setting of
the NPCI data generating package (Dorie 2016), which is the version that is used most frequently in other comparisons (e.g., see Shi et al. 2019; Shalit al. [017; Yao al. 2018) and comprises 608 untreated and 139 treated samples (747 in total). There are 25 covariates, 19 of which are discrete/binary, and the rest are continuous. The outcome for the IHDP data is continuous and
unbounded. Similarly to (Louizos et al. 2017 Shalit et al., 2017) and others, we utilize 60/30/10 rrain/validation/esst split.
We also utilize the job outcomes dataset which we refer to as Jobs (LaLonde 1986 Smith & Todd 2005 Unlike the IHDP datset, Jobs real-world data with a binary outcome. We follow a similar procedure to (Shalit et al. 2017) who indicate that they used the Dehejia and Wahba (Dehejia & Wahba 2002) and PSID comaprison samples. The Dehejia-Wahba Wahba sample comprises 260 treated samples and 185 control samples, along with the PSID comparison group comprising 2490 samples'. The dataset contains a mixture of observational and Randomized Controlled Trial data. Similarly to (Louizos al.) 2017 Shalit et al.] 2017) and others, utilize 56/24/20 vrain/validttiontesst split, and undertake 100 runs with varying random split allocations in order to acquire an estimate of average performance and standard error. Note that, between models, the same random seed used both for intialization as well as dataset splitting, and therefore the variance due to these factors is equivalent across experiments.
Finally, we introduce new synthetic dataset named TVAESyuth which follows the structure shown in Figure 4 While the weightings are chosen relatively arbitrarily, the structure is intentionally designed such that there are a mixture of exogenous and endogenous covariates. This enables us to compare the performance tent dimensions constant). These data comprised 1000 samples, continuous outcome and binary treatment, 8 covariates, and generative/strrucuuall equations as follows:
TVAE with and without z, (while keeping the total number of
U U ~ 1(0, U U ~ Bernoulli(0.5)
U U U2a
N
(16)
https://github. com Jiahangg22) TEDVAE/ / and elsewhere, and will be included in supplementary folder with source code upon acceptance.
AAaillabl from https:: /www. fredjo. com/
AAvailabee from htt.t::pp/////tt.ttt //tpsss/////es.seeee nber meeeeeorrrrrreeeeeeeaeeaeeeeeeee
htlpaaattt//swwaaaaaaaa 2.w.ata2.. html
15

---

Under review as a conference paper at ICLR 2021
X
Inference
Pred y block
Encoder Block
ELU
PC + ELU
Encoder Block
Encoder Block
Encoder Block
Encoder Block
ELU
ELU
g(z.
((2
((2
9(2,
PC
FC
FC
FC
Var
Pred t block
2., 2,) ptll 0. (.))
pt 2z)
block
Predy y block 2
Pred block
Generation
pfS 2. 2,))
FC ELU
FC ELU
p(z.) via.)
p(z.)
((2y))
FC - ELU
FC ELU
Pred block
FC
FC ELU
FC + FLU
Pred t block
FC
FC
Pred x block
p(t/z 2.)
Pred y block 2
ply 1,
píy 0,
FC
FC
FC Logils
(8 Re: 24, Fly)
p( 2 Ryy
Figure 5: Block architectural diagram for TVAE's inference and generation models (number of layers in Pred 1 & 2 and Pred blocks varies by experiment, as does the number of neurons in each layer (see details in text).
2, U2. Bn U2y
UU.
Z. =U=+
(17)
X1 Bemouli((ozz, + 0. 0.5))) x2 No0..20 + 0..3ze + 0.5z, 0.10 22, 0.2)
(18)
X3 ~ (00.2... 0.22. + 1.22, +0.10 0.2) X4 (erouuli(((66xxx+011(1; -0.5)) (19)
X5 ~ N(0.6z + 0.11U,, 0.1) X6 ~ + 2..1
(20)
X7 ~ NN0.52..+ + 0.100, 0.1) Xs ~ ~(0.52, + 0.1U
(21)
tp = 000.2%. + 0.82, + 0.1U,) t Bernoulli (t.p)
(22)
= 0.2%. + 0.5z,, 0.2t. + 0.1U,
(23)
Interventional distributions were generated by setting equal and 0 thereby yielding ground truth ATE (+s 0.8) and individual effects for evaluation purposes. The number treated individuals 70%. Note that for every replication, the 1000 sample evaluation set regenerated, and may have an empirical statistics that vary. For experiments, we use 80/20 train/test split.
A.5 MODEL AND HYPERPARAMETERS
This part of the appendix is referenced in Section 3 (Methodology), as well as in Section (Ex- periments) of the main paper. A block diagram of the TVAE architecture is shown in Fig. 5 For continuous outcomes (as in the IHDP dataset) we standardize the values and model as a Gaussian with fixed variance of 1, and mean determined by the outcome arm. All binary outcomes in the model (e.g. treatment or the relevant covariates) are modelled as Bernoulli distributed with a probability determined by the associated neural network function.
We now list the hyperparameters that were explored as part of model training. There may be room improve figures with further hyperparameter tuning. However, given that the tuning of
16

---

Under review as a conference paper at ICLR 2021
hyperparameters in causal inference paradigm is problematic in general, intentionally limited the space of hyperparameters (similarly to Zhang et al. 2020). Bold font indicates the settings that were used in the presented results.
Hyperparameter settings for IHDP dataset experiments were: hidden layers: 3; the weight on tar- geted regularization \ {0.0, 0.1, 0.2, 0.4, 0.6, 0.8, 1.0]; Adam (Kingma & Ba 2017) opti- mizer with learning rate LR = (1e-3, le-4, 5e-5); number hidden neurons was 300; number of layers- 4: dimensionality of the latent factors was Dz, Dz, 10, Dz. 15, D.. 5; batch size of 200; number of epochs 200; weight regularization 1e-4; and learning rate decay 5e-4.
Hyperparameter settings for the Jobs dataset experiments were: hidden layers: the weight
targeted regularization ÀTL = {0.0, 0.1, 0.2, 0.4, 0.6. 0.8, 1.0]: an Adam (Kingma & Ba 2017) optimizer with learning rate LR (5e-5, 1e-53: number of hidden neurons was 200; number of layers 2; dimensionality of the latent factors was Dz, Dz, 6, D=, 8, D=, 4: batch size
of 200; number of epochs 200; weight regularization 1e-4; and learning rate decay 5e-4=3.
Hyperparameter settings for the TVAESynth dataset experiments were: hidden layers: 2: the weight on targeted regularization ÀTr. = (0.0, 0.1, 0.2, 0.4, 0.6, 0.8, 1.0); an Adam (Kingma Ba 2017) optimizer with learning rate LR = 5e-5; number of hidden neurons was 20; number of layers 2; dimensionality of the latent factors was D2, D2 D= 2, Dz. 1; batch size of 200;
number of epochs 40; weight regularization 1e-4; and learning rate decay 5e-3.
As described in the main paper, wherever range of hyperaparameters was explored, the validation loss on the total objective function was used as the model selection criterion (i.e., not the causal effect estimation performance, which is not available in real-world scenarios).
A.6 SOFTWARE AND HARDWARE
The network is coded using Pyro (Bingham al. 2019) and builds base code by (Zhang al. 2020). We train on GPU (e.g. NVIDIA 20807T driven by 3.6GHz Intel 19-9900K CPU running Obuntu 18.04. Training 200 epochs of the IHDP dataset (training split of 450 samples) takes approx. 35 seconds (0. 175s per epoch).
17


Reviews:
[{"review_id": "VPXzujAmV4D", "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 Targeted Variational Autoencoders for ATE and CATE estimation. \n\nIt is assumed that all relevant confounding variables can be measured through proxy variables and then a VAE architecture is used to disentangle the latent variables into 4 sets: z_t - confounding between the treatment and covariates, z_c - confounding between the treatment, covariates and target, z_y - confounding between the covariates and target, z_0 - solely related to the covariates. \n\nThe proposed architecture builds on an existing approach, TEDVAE. The main difference is that TEDVAE did not include the z_0 term. \n\nA targeted regularization approach is also used which is similar to Shi et al. 2019\n\nThe proposed method outperforms existing approaches on to real world datasets.\n\nIn general the paper is well written and the approach appears sound and reasonable. The novelty is somewhat limited as both the architecture and regularization build incrementally on previous approaches. However, the proposed method appears to lead to a decent increase in performance over SoTA approaches. ", "extra_scores": {}, "extra_text": {}}, {"review_id": "K1jvRWaJrNw", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 5, "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": "This work provides an improved method for individual and average causal parameter estimation. The main idea is to apply targeted learning to deep latent variable models. The proposed approach is based on two existing work: \n- Disentangled variational latent model [Zhang et al., 2020].\n- Implementing targeted learning in neural nets via regularization [Shi et al., 2019].\n\nConsidering the above two papers, the contribution of this work is not very significant. However, it provides a nice combination of those works and adds the missing element of z_0 to model factors unrelated to treatment and outcome and to give more flexibility to the design of the latent variables, which all together makes a complete picture. The authors compare the proposed method with some alternatives on two datasets and the results show the better performance of the approach in most of the cases.\n\nIn general, the presentation of the paper is subpar. In many parts only a high level explanation is provided. For instance, \n- A motivation for using latent variable models is needed, similar to the one given in [Zhang et al., 2020].\n- I do not believe that the explanation for equation (1) will clarify the matter. \n- \"The combination of KL regularization and supervision helps prevent z_o from learning information in z_{t,c,y}\" This is an important point, yet it is very vaguely stated and should be elaborated.\n- The last paragraph of page 5 regarding the difference between the training in this work and [Shi et al., 2019] is not clear.\n\nOn page 4, the authors mention that: \"the use of deep latent variable techniques enables us to attempt to infer these hidden confounders from what are known as noisy proxy variables present in the dataset\" This is in general not true. We are not able to infer the latent confounders, unless if we assume very strong assumptions on them. The aforementioned sentence is also confusing because in this paper it is assumed that we do not have any latent confounders, i.e., the ignorability assumption. If the authors are indeed assuming that we have latent confounders and they are estimated using VAE, then nothing can be said about the correctness of the outputs of the method. Therefore, the sentence \"Doing so enables us to infer hidden confounders from proxy variables in the dataset, and to estimate the expected treatment effects, as well as individual-level treatment effects\" in the introduction should also be clarified. The only way that the approach is justifiable is by assuming ignobility and assuming that VAE is only used to get a concise representation of the observed confounders to improve the prediction.", "extra_scores": {}, "extra_text": {}}, {"review_id": "CMezdu2qyoO", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 5, "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": "This paper proposes a method for estimating conditional and average treatment effects under unconfoundedness. There are two main ideas: (1) train a VAE with latent space aimed at the adjustment-relevant information, and (2) incorporate the targeted regularization of Shi et al into the training.\n\nThe problem of estimating causal effects using deep learning is important, and I think paper has a promising direction. It is closely related to a growing literature on this subject (as noted in the paper itself), and does a reasonable job of explaining the innovations.\n\nHowever, there are some significant issues.\n\nThe exposition is generally unclear. The paper needs a complete rewrite, with particular attention to clarity about the (salient parts of) TMLE, and the claimed advantages of this approach over closely related methods. There also needs to be substantial improvement around the VAE component of the model, and precisely what the identification assumptions are.\n\nThe empirical results are inadequate. The IHDP and LaLonde datasets are not challenging enough to separate methods in the bakeoff, nor rich enough to give substantial insight into their performance. I suggest you use ACIC competition data. Further, the ablation study doesn't seem to test the right aspects of the model. E.g., there should be some experiment showing that not applying targeted regularization to the propensity score part of the model leads to better performance, since that's the core distinction with Shi et al.\n\nThis paper isn't ready for publication in its current form. However, I reiterate that I think the basic ideas are interesting, and worth developing.\n\nSome further comments and questions:\n\n1. missing expectation symbol in 2nd paragraph of background (in def of \\tau)\n2. equation 1 is meant to be describing an observational quantity; the do(T=t) should just be t\n3. the discussion about optimal epsilon=0 is confusing (I understand that you mean that running more than 1 round of the TMLE update doesn't make a difference, but this is not clear in your writing)\n4. indeed, all of the prose around equation 2 should be rewritten. It is not clear in the current version.\n5. under DAG (c), x doesn't block backdoor paths and strong ignorability is not satisfied\n6. the related work should be clearer that targeted regularization was proposed in Shi et al \n7. what models are you using for Q and g for the TMLE reported in the table? In particular, are you comparing training with targeted regularization to training without targeted regularization and then plugging in to TMLE?\n8. when you compare to 'Dragonnet' are you comparing to Dragonnet + Targeted Regularization, or just vanilla dragonnet?\n9. what exactly is done to prevent the propensity score from being modified by the targeted regularization? I presume that at least the shared latent representation is still affected by the inclusion of the targeted regularization term?", "extra_scores": {}, "extra_text": {}}, {"review_id": "3LasS-5xXPz", "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 proposed contribution of this work is to build of the existing literature which uses variational autoencoders for causal inference by (1) allowing an explicit mechanism for modeling irrelevant covariates and (2) incorporating targeted regularization into the latent variable nnet framework. Optimization of the model is done by minimizing the ELBO subject to a penalty term which the authors refer to as “targeted regularization”. This term is essentially an application of the TMLE model. Both of these proposals appear to improve the performance on what are now pretty standard benchmark datasets (jobs and ihdp). \n\nOverall, I think this work presents two very sensible additions to the latent variable formulation of causal neural networks. My largest concern is with novelty–each contribution borders on incremental, and neither necessarily open doors for substantial amounts of follow on work. With that being said, I think that this work does provide value to the community given (a) the sensible, simple, model changes proposed and, (b) the pretty compelling empirical evidence.\n\nComments / Questions:\n\n* You note that the gradients are taken with respect to zeta and not with respect to g_p or g_q in the paper. This language is kind of confusing at first read. As I read the difference is that you are more explicitly reproducing the machinery of TMLE by adding what is essentially a logistic regression (for bounded outcomes). This is very sensible, but would benefit from more explicit discussion relating back to the TMLE literature. \n* It appears that the change mentioned above does not provide a substantial benefit for estimating the ATE over Dragonnet, can the authors provide intuition around that? \n* The oblation study to tease apart the contribution of the two proposed changes is very interesting. Given that the introduction of zeta provides substantial benefit on its own, it would be interesting to see a variant that does not use z_0 but uses zeta.\n\nSmall edit notes:\n\n* In the main text you make reference to equation 14 which is in the supplement. I believe you meant to make reference to equation 5. \n", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: The paper introduces some good ideas, but I don't think it is quite there in terms of a method to be recommended for publications. I think it is mostly reasonably written (I do not agree with the comment of a 'complete rewrite') but there are indeed some passages for improvement (for instance, an equation as y = σ−1[Q0(t, x)] + εH(t, x)], Section 2, needs comments, as the left hand side is discrete and the right hand side is continuous, unbounded).

My main concern is the disregard for identification. Some citations are unclear (the second-to-last paragraph in Section 4 cites a few papers in identification that have little to do with the problem here, which is proxy-based. The papers cited don't even mention latent variables at all). As stated, the split in three sets of variables as suggested by Figure 1 is just an idealization: there is no reason at all they can be identified, and actually the theory where just Zc is considered impose a lot of restrictions on when we can possible identify Zc (see e.g., Miao et al. 2018, Biometrika, https://arxiv.org/pdf/1609.08816.pdf ). I know that some papers like Louizos et al. play fast and loose with identification too, but at least their Z_c structure they aim at has been studied elsewhere (like the Miao et al. paper), while here, like the Zhang et al. paper cited, may be leading researchers to an unfruitful path. This, combined with the relative modesty of the novelty, is the primary reason for my recommendation. I do think the paper can be improved in a productive way by investigating it from the point of view of either i) the theoretical justification for identification; ii) or from a more empirical direction with much experimentation on the different ways the structured latent space is capturing confounding (the target learning aspect of it is pretty much orthogonal to this).
Decision: Reject