Extracted main text — title through conclusion, appendix excluded. This is what our citation measures are computed over, published so the extraction can be checked by eye.
56,968 characters · 19 sections · 98 citation commands
Towards Principled Causal Effect Estimation by Deep Identifiable Models
Causal inference imbens2015causal, pearl2009causality, i.e, estimating causal effects of interventions, is a fundamental task across many domains. We address the estimation of treatment effects (TEs), such as effects of public policies or new drugs, based on a set of observations consisting of binary labels for treatment/control (non-treated), outcome, and other covariates. The fundamental difficulty of causal inference is that we never observe counterfactual outcomes, which would have been if we had made another decision (treatment or control). While the ideal protocol for causal inference is randomized controlled trials (RCTs), they often have ethical and practical issues, or suffer from expensive costs. Thus, causal inference from observational data is important. It introduces other challenges, however. The most crucial one is confounding: there may be variables (called confounders) that causally affect both the treatment and the outcome, and spurious correlation follows.
A majority of works in causal inference rely on the unconfoundedness, which means that appropriate covariates are collected so that the confounding can be controlled by conditioning on those variables. This is still challenging, due to systematic imbalance (difference) of the distributions of the covariates between the treatment and control groups. Classical ways to deal with imbalance are matching and re-weighting Stuart2010matching,rosenbaum2020modern. There are semi-parametric methods van2018targeted, which have better finite sample performance, and also non-parametric methods wager2018estimation. Notably, there is a recent rise of interest in learning balanced representation of covariates, which is independent of treatment groups, starting from johansson2016learning.
There are a few lines of works that address the difficult but important problem of unobserved confounding. Without covariates to adjust for, the naive regression with observed variables introduces bias, if the decision of treatment and the outcome are confounded, as explained in Sec. (ref). Instead, many methods assume special structures among the variables, such as instrumental variables (IVs) angrist1996identification, proxy variables tchetgen2020introduction, network structure veitch2019using, and multiple causes wang2019blessings. Among them, instrumental and proxy variables are most commonly exploited. Instrumental variables are not affected by unobserved confounders, influencing the outcome only through the treatment. On the other hand, proxy variables are causally connected to unobserved confounders, but are not confounding the treatment and outcome by themselves. Other methods use restrictive parametric models allman2009identifiability, or only give interval estimation manski2009identification, kallus2019interval.
In this work, we challenge the problem of estimating TEs under unobserved confounding. We in particular discuss the individual-level TE, which measures the TE conditioned on the covariate, for example, on a patient's personal data. We highlight the natural VAE architecture following from modeling sufficient scores and the promising experimental results, under unconfounded, IV, proxy, and networked confounding settings. On the theoretical side, we show identification of TEs using our generative model under unconfoundedness, but also discuss a parallel work wu2021beta addressing limited overlap and future work(s) under unobserved confounding.
Our method exploits the important concepts of sufficient scores for TE estimation hansen2008prognostic,rosenbaum1983central and also the recent advance of VAE with identifiable latent variable, which is determined by the true observational distribution khemakhem2020variational. VAEs kingma2019introduction are suitable for causal estimation thanks to its probabilistic nature. However, most VAE methods for TEs, e.g., louizos2017causal,zhang2020treatment, are ad hoc and thus not identifiable. Instead, our goal is to build a VAE that can identify and recover from observational data a sufficient score via the latent variable, which can be seen as a causal representation scholkopf2021toward; recovering the true confounder is not necessary. The code is uploaded to OpenReview, and the proofs are in Appendix (ref). Our main contributions are:
An early version of this work, which proposed the same VAE architecture, is in wu2021identifying.
As detailed in Sec. (ref), current VAE methods for TE estimation are more heuristic than “causal”. Our work endeavors to remedy this situation. Below, we focus on other related works.
Identifiability of representation learning. The hallmark of deep neural networks (NNs) is that they can learn representations of data. A principled approach to interpretable representations is identifiability, that is, when optimizing our learning objective w.r.t. the representation function, only a unique optimum, which represents the true latent structure, will be returned. With recent advances in nonlinear ICA, identifiability of representations is proved under a number of settings, e.g., auxiliary task for representation learning hyvarinen2016unsupervised, hyvarinen2019nonlinear and VAE khemakhem2020variational. The results are exploited in causal discovery pmlr-v108-wu20b and causal representation learning shen2020disentangled. To the best of our knowledge, this work is the first to explore this identifiability in TE estimation.
Causal inference with auxiliary structures. CEVAE louizos2017causal relies on the strong assumption that the true confounder distribution can be recovered from proxies. Our method is quite different in motivation, applicability, architecture. Detailed comparisons are given in Appendix (ref). Also with proxies, kallus2018causal use matrix factorization to infer the confounders, and mastouri2021proximal use kernel methods to solve the underlying Fredholm integral equation. IVs are also exploited in machine learning, there are methods using deep NNs hartford2017deep and kernels singh2019kernel,muandet2019dual.
Representation learning for causal inference. Recently, researchers start to design representation learning methods for causal inference, but mostly limited to unconfounded settings. Some methods focus on learning a balanced representation of covariates, e.g., BLR/BNN johansson2016learning, and TARnet/CFR shalit2017estimating. shi2019adapting use similar architecture to TARnet, considering the importance of treatment probability. There are also methods using GAN yoon2018ganite and Gaussian process alaa2017bayesian. Our method shares the idea of balanced representation learning.
\begingroup
Following imbens2015causal, we begin by defining potential outcomes (or counterfactual outcomes) ${\mathbf{y}}(t), t \in \{0,1\}$, which are the outcomes that would have been observed, if we applied treatment value ${\textnormal{t}}=t$. Note that, for a unit under research, we can observe only one of ${\mathbf{y}}(0)$ or ${\mathbf{y}}(1)$, corresponding to the factual treatment applied. This is the fundamental problem of causal inference. We also observe relevant covariate ${\mathbf{x}}$, which is associated with individuals, and the observation $\mathcal{D}\coloneqq({\mathbf{x}},{\mathbf{y}},{\textnormal{t}})$ is a random variable with underlying probability distribution.
The expected potential outcome is $\mu_t({\bm{x}}) = \mathbb{E}({\mathbf{y}}(t)|{\mathbf{x}}={\bm{x}})$, conditioned on ${\mathbf{x}}={\bm{x}}$. The estimands in this work are Average TE (ATE) $\nu$ and Conditional ATE (CATE) $\tau$, defined by
CATE can be seen as an individual-level TE, if conditioned on highly discriminative covariates.
{8pt} {5pt}
Identification of TEs means that, the true observational distribution uniquely determines and gives the ATE or, better, CATE. Adapting standard identification results rubin2005causalhernanCausalInferenceWhat2020, we start with the following conditions, denoted by {\bf (A)}: there exists a (possibly unobserved) variable ${\mathbf{u}} \in \mathbb{R}^n$ such that together with ${\mathbf{x}}$, it gives (i) (Exchangeability) ${\mathbf{y}}(t) \raisebox{0.05em}{\rotatebox[origin=c]{90}{$\models$}} {\textnormal{t}}| {\mathbf{u}}, {\mathbf{x}}$ and (ii) (Overlap, or Positivity) $p({\textnormal{t}}|{\mathbf{u}}, {\mathbf{x}}) > 0$; and (iii) (Consistency of counterfactuals) ${\mathbf{y}} = {\mathbf{y}}(t)$ if ${\textnormal{t}}=t$. All of them are satisfied for both $t$, which is our convention when $t$ appears in a statement without quantification. Intuitively, exchangeability means all confounders are in essence contained in $({\mathbf{u}},{\mathbf{x}})$, and overlap means each possible value of $({\mathbf{u}},{\mathbf{x}})$ is observed for both treatment groups. Note that, joint exchangeability ${\mathbf{y}}(0),{\mathbf{y}}(1) \raisebox{0.05em}{\rotatebox[origin=c]{90}{$\models$}} {\textnormal{t}}| {\mathbf{x}},{\mathbf{u}}$ is stronger than exchangeability and is not necessary for identification hernanCausalInferenceWhat2020.
A general example of causal structure that satisfies the three conditions is shown in Figure (ref), although further structural constraints might be necessary for theoretical developments (see Sec. (ref)).. Here, ${\mathbf{x}}_{c}$,${\mathbf{x}}_{iv}$,${\mathbf{x}}_{p_a}$,${\mathbf{x}}_{p_d}$,${\mathbf{x}}_{y}$ are covariates that are: (observed) confounder, IV, antecedent proxy (that is antecedent of ${\mathbf{z}}$), descendant proxy, and antecedent of ${\mathbf{y}}$, respectively. The covariate(s) ${\mathbf{x}}$ may not have subsets in any categories in the graph. ${\mathbf{e}}$ is unobserved exogenous noise on ${\mathbf{y}}$. Assumption {\bf (A)} may hold otherwise, e.g., ${\mathbf{x}}$ is a child of ${\textnormal{t}}$.
CATE can be given by (ref), using assumption (A) in the second equality.
If the variable ${\mathbf{u}}$ is observed, then (ref) identifies CATE. However, if ${\mathbf{u}}$ is an unobserved confounder, the naive regression $\mathbb{E}[{\mathbf{y}}| {\mathbf{x}}={\bm{x}}, {\mathbf{t}}=t]$ based on observable variables is not equal to $\mu_t({\bm{x}})$. In fact, if an unknown factor correlates with ${\mathbf{t}}$ positively and tends to give higher value for ${\mathbf{y}}$, the naive regression $\mathbb{E}[{\mathbf{y}}| {\mathbf{x}}={\bm{x}}, {\mathbf{t}}=1]$ should be higher than $\mathbb{E}[{\mathbf{y}}(1)|{\mathbf{x}}={\bm{x}}]$.
Our method models prognostic scores hansen2008prognostic, adapted as Pt-scores in this paper, closely related to the important concept of balancing score rosenbaum1983central. Both are sufficient scores for identification; prognostic scores are sufficient statistics of outcome predictors and balancing score is for the treatment (see Appendix (ref) for details).
The identity function is a trivial PS. If the true data generating process (DGP) satisfies additive noise model, i.e., ${\mathbf{y}}={\bm{f}}^*({\mathbf{u}},{\mathbf{x}},{\textnormal{t}})+{\mathbf{e}}$, then ${\bm{f}}^*_t$ is a PtS hansen2008prognostic; and it is a causal representation scholkopf2021toward of the direct cause on ${\mathbf{y}}$, summarizing the effects of $({\mathbf{u}},{\mathbf{x}})$. The independence property of PtS (Lemma (ref) in Appendix),
is used in second equality of (ref) in Theorem (ref) which extends Proposition 5 in hansen2008prognostic.
Compared to (ref), $P=\mathbb{P}_t({\bm{u}},{\bm{x}})$ plays the role of ${\bm{u}}$, and $p_{{\mathbf{y}}|\mathbb{P}_{t},{\textnormal{t}}}$ conditions on $P$ instead of $({\bm{u}},{\bm{x}})$. In general, information from ${\mathbf{u}}$ is needed to determine $p_{{\mathbf{y}}|\mathbb{P}_{t},{\textnormal{t}}}$ and $p_{\mathbb{P}_t|{\mathbf{x}}}$. In Sec. (ref), we first show that our generative model can identify an equivalent PS if ${\mathbf{u}}$ is observed (Sec. (ref)), and then discuss how our model is connected to and might learn relaxations of PtS when ${\mathbf{u}}$ is unobserved (Sec. (ref)).
In this section, we first introduce our generative model and VAE architecture, then prove the identifiability of the model, and finally use our model to estimate TEs. Our method is expected to learn a latent representation sufficient for TE identification/estimation, as we will see in the next section.
\begingroup
We build the generative model (ref) of our VAE with the direct cause $\mathbb{P}_t({\mathbf{u}},{\mathbf{x}})$ as the latent variable ${\mathbf{z}}$ and connect the model to iVAE. The outcome distribution $p_{{\mathbf{y}}|\mathbb{P}_{t},{\textnormal{t}}}$ in (ref) is modeled by $p_{{\bm{f}}}({\bm{y}}|{\bm{z}},t)$, and the score distribution $p_{\mathbb{P}_{{\textnormal{t}}}|{\mathbf{x}}}$ in (ref) is modeled by $p_{\bm\lambda}({\bm{z}}|{\bm{x}},t)$. The condition on ${\mathbf{x}}$ in the joint model $p({\mathbf{y}},{\mathbf{z}}|{\mathbf{x}},{\textnormal{t}})$ reflects that our estimand is CATE given ${\mathbf{x}}$.
The outcome assumes an additive noise model such that $\bm\epsilon \sim p_{\bm\epsilon}$ denotes the exogenous noise; and $p_{\bm\lambda}$ is a factorized Gaussian, where $\bm\lambda_{{\textnormal{t}}}({\mathbf{x}})\coloneqq \operatorname{diag}^{-1}({\bm{k}}_{{\textnormal{t}}}({\mathbf{x}}))({\bm{h}}_{{\textnormal{t}}}({\mathbf{x}}), -\frac{1}{2})^T$ is the natural parameter as in the exponential family. ${\bm{\theta}}\coloneqq({\bm{f}},\bm\lambda)=({\bm{f}}, {\bm{h}}, {\bm{k}})$ contains the functional parameters.
The model is learned by the evidence lower bound (ELBO) which estimates the variational lower bound (See Appendix (ref) for the basics of VAEs):
with $p_{{\bm{f}}}$ as the decoder, $q$ as the encoder, and $p_{\bm\lambda}$ as the conditional prior.
{8pt} {5pt}
We name this architecture Intact-VAE (Identifiable treatment-condi\textit{t}ional VAE). Note that (ref) has a similar factorization with the generative model of iVAE, $p_{\bm\theta}({\bm{y}},{\bm{z}}|{\bm{x}}) = p_{{\bm{f}}}({\bm{y}}|{\bm{z}})p_{\bm\lambda}({\bm{z}}|{\bm{x}})$; the first factor (decoder) does not condition on ${\mathbf{x}}$. Similarly, our decoder $p_{{\bm{f}}}({\bm{y}}|{\bm{z}},t)$ conditions on ${\mathbf{z}}$ as a PtS which satisfies ${\mathbf{y}}\raisebox{0.05em}{\rotatebox[origin=c]{90}{$\models$}} {\mathbf{x}}|{\mathbf{z}},{\textnormal{t}}$. On the other hand, the condition on ${\textnormal{t}}$ is in both the factors of (ref). Thus, our VAE architecture can be seen as a combination of iVAE and conditional VAE (CVAE) sohn2015learning,kingma2014semi, with ${\textnormal{t}}$ as the conditioning variable. See Figure (ref) for the comparison in terms of graphical models. Our encoder $q$, which conditions on all the observables and builds the approximate posterior, is similar to other VAEs. See Appendix (ref) for summary of CVAE and iVAE.
Though our theory allows more general distributions, for tractable inference and easy implementation, we have used factorized Gaussian for the prior, and now also use it for the decoder $p_{{\bm{f}},{\bm{g}}}({\bm{y}}|{\bm{z}},t)$ and encoder $q_{\bm\phi}({\bm{z}}|{\bm{x}},{\bm{y}},t)$; i.e., they are products of univariate Gaussian distributions:
$\bm\phi\coloneqq({\bm{r}},{\bm{s}})$, like $\bm\theta$, contains functional parameters given by NNs which take $({\mathbf{x}},{\mathbf{y}},{\textnormal{t}})$ as inputs. We often write $t$ of the function argument in subscripts.
Our model identifiability lays the foundation for principled causal effect estimation, as we will discuss in Sec. (ref). The main theoretical result of iVAE is extended in our Theorem (ref) which combines the techniques from khemakhem2020variational and sorrenson2019disentanglement. Essentially the same results can be proved for other exponential family priors.
\endgroup
The essence of Theorem (ref) is ${\bm{f}}'_t = {\bm{f}}_t \circ \mathcal{A}_t$, that is, ${\bm{f}}_t$ can be identified (learned) up to an affine transformation defined by $\bm\lambda_t$. This identifiability of parameters does not directly imply TE identification; other assumptions are needed, as we show in Sec. (ref). The assumptions in Theorem (ref) are inherited from iVAE. Injectivity in i) is related to TE identification (Theorem (ref)), which requires injectivity in the true DGP and identifies a PtS up to an injective mapping. Intuitively, if ii) does not hold, then the support of $\bm\lambda({\mathbf{x}})$ degenerates to a $(2n-1)$-dimensional space; thus, ii) holds easily in practice (see B.2.3 in khemakhem2020variational). Note that, to have (ref), we only need the same observational distribution $p({\mathbf{y}}|{\mathbf{x}},{\textnormal{t}}=t)$, but this leaves room for different latent distributions.
Once the VAE is learned by the ELBO (ref), the estimate of $\mu_t$ is given by
where $q({\bm{z}}|{\bm{x}})\coloneqq \mathbb{E}_{p({\bm{y}},t|{\bm{x}})}q_{\bm\phi}({\bm{z}}|{\bm{x}},{\bm{y}},t)$ is the aggregated posterior. Our overall algorithm steps should be clear. After training Intact-VAE by (ref), we feed $({\bm{x}},{\bm{y}},t)$ into the encoder and draw posterior sample from it. Then, setting ${\textnormal{t}}=\hat{t}$ in the decoder, feed the posterior sample into it, we get counterfactual sample as outputs of the decoder. Finally, we infer CATE $\hat{\tau}({\bm{x}})=\hat{\mu}_1({\bm{x}})-\hat{\mu}_0({\bm{x}})$.
Our method also addresses the problem of imbalance--if $p({\mathbf{x}}|0),p({\mathbf{x}}|1)$ are very different for some ${\bm{x}}$, then we have few data points for one of $\mathcal{D}|{\bm{x}},t\coloneqq\{({\bm{x}},{\mathbf{y}},t)\}$, resulting in poor estimation. We set $\bm\lambda_0=\bm\lambda_1$ in the prior, and thus $p_{\bm\lambda}({\mathbf{z}}|{\mathbf{x}},{\textnormal{t}})=p_{\bm\lambda}({\mathbf{z}}|{\mathbf{x}})$ is independent of ${\textnormal{t}}$ given ${\mathbf{x}}$. The same prior for the treatment groups, i.e., the balanced prior of latent representation, is similar to balanced representation learning johansson2016learning, shalit2017estimating, where balanced representation is favored by ad hoc regularization. This is also related to the fact that, when building CVAE, unconditional prior can achieve better performance kingma2014semi.
By taking ${\mathbf{y}}$, the posterior model (the encoder) is better than the prior. On the other hand, sampling posterior requires post-treatment observation ${\bm{y}}$. Often, it is desirable that we can also have pre-treatment prediction for a new subject, with only the observation of its covariate ${\mathbf{x}}={\bm{x}}$. To this end, in (ref), we use conditional prior $p({\mathbf{z}}|{\mathbf{x}})$ as a pre-treatment predictor for ${\mathbf{z}}$: draw sample from $p({\mathbf{z}}|{\bm{x}})$ instead of $q$ and get rid of the outer average taken on $\mathcal{D}$; all the others remain the same.
As usual, we expect the variational inference and optimization procedure to be (near) optimal; that is, the consistency of VAE. Consistent estimation using the prior is a direct corollary of TE identification and the consistent VAE. See Appendix (ref) for formal statements and proofs. It is possible to prove the consistency of the posterior estimation, as shown in bonhomme2021posterior,liao2011posterior, and we leave this for future work, see Sec. (ref).
In this section, we present a critical review of existing VAE-based methods in Sec. (ref) and then discuss theoretical developments that can lead to principled causal estimation, under different settings.
Most VAE methods for TEs, e.g., louizos2017causal,zhang2020treatment,vowels2020targeted,lu2020reconsidering, add ad hoc heuristics into the VAEs, and thus break down probabilistic modeling, not to mention model identifiability. Moreover, the methods learn representations from proxy variables, leading to either impractical assumptions or conceptual inconsistency, in TE identification.
On identification. First, as to TE identification, CEVAE assumes unobserved confounder can be recovered, which is rarely possible even under further structural assumptions tchetgen2020introduction. Indeed, rissanen2021critical recently give evidence that the method often fails. Other methods ,zhang2020treatment,vowels2020targeted,lu2020reconsidering assume unconfoundedness but still rely on proxy at least intuitively; for example, lu2020reconsidering factorize the decoder as if in the proxy setting. However, unconfoundedness and proxy should not be put together. The conceptual inconsistency is that, by definition, unconfoundedness means covariates fully control confounding, while the motivation for proxy is that unconfoundedness is often not satisfied in practice and covariates are at best proxies of confounding, which are non-confounders causally connected to confounders tchetgen2020introduction. Second, without model identifiability, the empirical results of the methods lack solid ground; under settings not covered by their experiments, the methods would silently fail to learn proper representations, as we show in Sec. (ref).
On ad hoc heuristics. Ad hoc heuristics break down probabilistic modeling and/or give ELBOs that do not optimally estimate the models. For example, in CEVAE, $q({\textnormal{t}}|{\mathbf{x}})$ and $q({\mathbf{y}}|{\mathbf{x}},{\textnormal{t}})$ are added into the encoder to have pre-treatment estimation, and the ELBO has two additional likelihood terms respectively. The VAE in zhang2020treatment is even more ad hoc; it splits the latent variable ${\mathbf{z}}$ into three components, and applies the ad hoc tricks of CEVAE to each of the component. Particularly, when constructing the encoder, they implicitly assume the three components of ${\mathbf{z}}$ are conditional independent give ${\mathbf{x}}$, which violates the intended graphical model.
Compared to the above methods, our Intact-VAE is simpler and more principled, and often has better performance. It models a prognostic score as the latent variable and is based on the identification equation (ref), while not compromised by ad hoc heuristics. Our ELBO is derived by standard variational lower bound (ref). Moreover, our pre-treatment prediction is given naturally by the prior, thanks to the correspondence between our model and (ref). We show in the following subsections how our model and its identifiability inspire theoretical developments in TE identification.
As a first step, we take up theoretical analysis of Intact-VAE under assumption (A) when ${\mathbf{u}}$ is observed (i.e., ${\mathbf{u}}$ is components of ${\mathbf{x}}$). Then, we have ${\mathbf{y}}(t) \raisebox{0.05em}{\rotatebox[origin=c]{90}{$\models$}} {\textnormal{t}}|{\mathbf{x}}$ and $p({\textnormal{t}}|{\mathbf{x}})>0$ (Exchangeability and Overlap given ${\mathbf{x}}$). This is the standard unconfoundedness setting. Regarding to the PtS, $p_{\mathbb{P}_{{\textnormal{t}}}|{\mathbf{x}}}$ in (ref) degenerates to $\mathbb{P}_{{\textnormal{t}}}({\mathbf{x}})$. Theorem (ref) is an identification under shape restriction chetverikov2018econometrics, because injectivity in assumption i) is monotonicity if ${\bm{j}}_t$ is on $\mathbb{R}$.
In essence, $\mu_t$ is identified up to an invertible mapping ${\bm{v}}$, such that ${\bm{f}}_t={\bm{j}}_t\circ{\bm{v}}^{-1}$ and ${\bm{h}}_t={\bm{v}}\circ\mathbb{O}$, with same ${\bm{v}}$ for both $t$. The connection to PS\footnote{Note that, some call $\mu_0$ the prognostic score (e.g, schuler2020increasing,tarr2021estimating), even without additive noise models. In this alternative terminology, we can also say $\mathbb{O}$ is a PS, without requiring additive noise models.} is clear; if additive noise model is the ground truth, then $\mathbb{O}$ is a PS because $\mu_t$ is is a PtS. Then, because $\mathbb{O}$ is recovered up to ${\bm{v}}$, the independence ${\mathbf{y}}(t)\raisebox{0.05em}{\rotatebox[origin=c]{90}{$\models$}}{\mathbf{x}}|\mathbb{O}({\mathbf{x}})$ is preserved. Assumption i) says that the treatment affects $\mathbb{E}({\mathbf{y}}|{\bm{x}})$ only through injective ${\bm{j}}_t$ (which is identified up to ${\bm{v}}$), and CATEs are given by $\mu_0$ and an invertible function ${\bm{i}}\coloneqq{\bm{j}}_1\circ{\bm{j}}_0^{-1}$. See Appendix (ref) for real-world examples satisfying i) and the connection to “independent causal mechanisms” janzing2010causal. With the existence of $\mathbb{O}$, Assumption ii) simply matches the model to the truth. Note that, with ${\bm{k}}({\mathbf{x}}) = \bm0$, the prior ${\mathbf{z}}_{\bm\lambda,t} \sim p_{\bm\lambda}({\bm{z}}|{\bm{x}},t)$ degenerates to function ${\bm{h}}_{t}({\mathbf{x}})$.
Interestingly, Theorem (ref) does not depend on Theorem (ref), mainly because i) is strong in the sense that it indicates a PS. To work under general additive noise models with only PtSs, we use our model identifiability (ref) in wu2021beta, the main ideas of which is briefly summarized below.
The advantage of PtS can be more clearly seen when ${\mathbf{x}}$ does not satisfy overlap. Now, a straightforward estimation of TEs is not possible at a non-overlapping value ${\bm{x}}$ due to lack of data. However, $\mathbb{P}_t({\mathbf{x}})$ can map some non-overlapping values to an overlapping value, and overlap of $\mathbb{P}_t({\mathbf{x}})$ implies but is not implied by overlap of ${\mathbf{x}}$ d2020overlap. In a parallel submission wu2021beta, we assume the overlap of a PtS instead and extend Theorem (ref) to this limited overlap setting. This is a natural next step, because we already model a PtS as the latent variable.
A main result of wu2021beta is that the latent variable of our VAE recovers a PtS and identifies the CATE through the model, under limited overlap, similar to the conclusion 1) and 2) in Theorem (ref). To recover the PtS, we derive a condition and strengthen (ref) so that $\mathcal{A}_0=\mathcal{A}_1$, which compensates for $\mathbb{P}_0\neq\mathbb{P}_1$. Thus, we fully exploit the probabilistic nature of VAE in modeling, and give principled causal estimation based on the identifiability of VAE.
Still, with ${\mathbf{u}}$ observed, it would seem unnecessary to model $\mathbb{P}_{{\textnormal{t}}}({\mathbf{x}})$ by the distribution $p_{\bm\lambda}({\bm{z}}|{\bm{x}},t)$ (the prior). However, the prior, together with the encoder, quantifies the uncertainty of scores--we are uncertain how likely a PS (not only a PtS) can be recovered wu2021beta.
The positive experimental results motivate us to consider the theory under unobserved confounding. Moreover, the prior in (ref) is even more natural with ${\mathbf{u}}$ unobserved, since $p_{\mathbb{P}_{{\textnormal{t}}}|{\mathbf{x}}}$ is not degenerated due to the uncertainty of ${\mathbf{u}}$. Thus, we conjecture that, in our VAE framework, unobserved confounding is treated as a source of uncertainty of scores and is handled in a Bayesian way. We give more considerations for future theoretical work below.
Identification.\space Auxiliary structures (e.g., IVs) can give TE identifications via control functions $\mathbb{C}({\textnormal{t}},{\mathbf{x}})$, conditioning on which the treatment becomes exogenous, that is, ${\mathbf{y}}(t) \raisebox{0.05em}{\rotatebox[origin=c]{90}{$\models$}} {\textnormal{t}}|\mathbb{C}(t,{\mathbf{x}})$ matzkin2007nonparametric,wooldridge2015control. Control functions can be stochastic, as in puli2020general. Consistent TE estimation can be given by a regression of outcome on the treatment and a control function. Our model (ref) can be seen as a two-stage procedure: first, $p_{\bm\lambda}({\bm{z}}|{\bm{x}},t)$ gives a stochastic control function; second, $p_{{\bm{f}}}({\bm{y}}|{\bm{z}},t)$ regresses the outcome. We need to specify the control function learned by Intact-VAE and the required structural constraints. Control functions are recently found under the proxy setting nagasawa2018identification, or in the presence of both proxies and IVs tien2021instrumental.
Estimation.\space In causal inference, many models, including nonparametric IV regression (NPIV), are stated as conditional moment restrictions (CMRs) newey199316. Optimizing the ELBO of our VAE, given by (ref), can be seen as finding functions ${\bm{f}}$ and $\mathbb{C}$, subject to the CMR $\mathbb{E}_{p_{\bm\theta}}({\mathbf{y}}|{\mathbf{x}},{\textnormal{t}})=\mathbb{E}({\mathbf{y}}|{\mathbf{x}},{\textnormal{t}})$. We believe our Intact-VAE framework, possibly with modifications, can be shown to give optimal estimation under the CMR. There are formal connections between CMRs and quasi-Bayesian analysis using KL divergence zhang2006eps,jiang2008gibbs,kim2002limited. For example, kato2013quasi uses a quasi-likelihood from the CMR of NPIV to set the prior, and the Gibbs posterior zhang2006eps,jiang2008gibbs is a minimizer of an information complexity which has a variational characterization similar to an ELBO. For general CMR models, liao2011posterior extend kim2002limited and give the best approximation to the true likelihood function under the CMR by minimizing a KL divergence. Very recently, wang2021scalable employ quasi-Bayesian analysis to kernel-based IV methods, but only consider unconditional moments.
We use the proposed Intact-VAE for three types of data, and compare it with existing methods.
Unless otherwise indicated, for each function ${\bm{f}},{\bm{g}},{\bm{h}},{\bm{k}},{\bm{r}},{\bm{s}}$ in our VAE, we use a multilayer perceptron (MLP) that has 3*200 hidden units with ReLU activation, and $\bm\lambda=({\bm{h}},{\bm{k}})$ depends only on ${\mathbf{x}}$. The Adam optimizer with initial learning rate $10^{-4}$ and batch size 100 is employed. All experiments use early-stopping of training by evaluating the ELBO on a validation set. We test post-treatment results on training and validation set jointly. The treatment and (factual) outcome should not be observed for pre-treatment predictions, so we report them on a testing set. More details on hyper-parameters and settings are given in each experiment and Appendix.
As in previous works shalit2017estimating, louizos2017causal, we report the absolute error of ATE $\epsilon_{ate}\coloneqq |\mathbb{E}_{\mathcal{D}}(y(1)-y(0)) - \mathbb{E}_{\mathcal{D}}\hat{\tau}({\bm{x}})|$, and the square root of empirical PEHE hill2011bayesian $\epsilon_{pehe}\coloneqq \mathbb{E}_{\mathcal{D}}((y(1)-y(0))-\hat{\tau}({\bm{x}}))^2$ for individual-level TEs.
\begingroup
We generate synthetic datasets by (ref). The parameters are different between DGPs: $\mu_i$ and $\sigma_i$ are randomly generated; the functions $h,k,l$ are linear with random coefficients; and $f_0,f_1$ is built by separated randomly initialized (then fixed) NNs. We generate two kinds of outcome models, depending on the type of $f$: linear and nonlinear outcome models use random linear functions and NNs with invertible activations and random weights, respectively. We set the outcome and proxy noise level by $\alpha,\beta$ respectively. See Appendix (ref) for more details.
We experiment on three different causal structures as shown in graphical models of Figure (ref), by variation on (ref). Instead of taking inputs ${\mathbf{x}}, {\textnormal{z}}$ in $l$, we consider two special cases: $l\coloneqq l({\mathbf{x}})$, then ${\mathbf{x}}$ fully adjusts for confounding, we are in fact unconfounded; and $l\coloneqq l({\textnormal{z}})$, then we have unobserved confounder ${\textnormal{z}}$ and proxy ${\mathbf{x}}$ of ${\textnormal{z}}$. To introduce ${\mathbf{x}}$ as instrumental variable, we generate another 1-dimensional random source ${\textnormal{w}}$ in the same way as ${\mathbf{x}}$, and use ${\textnormal{w}}$ instead of ${\mathbf{x}}$ to generate ${\textnormal{z}}|{\textnormal{w}} \sim \mathcal{N}(h({\textnormal{w}}), \beta k({\textnormal{w}}))$; except indicated above, other aspects of the models are specified by (ref).
{8pt} {5pt}
For each causal structure, and with the same kind of outcome models, and the same noise levels ($\alpha,\beta$), we evaluate Intact-VAE and CEVAE on 100 random DGPs, with different sets of parameters in (ref). For each DGP, we sample 1500 data points, and split them into 3 equal sets for training, validation, and testing. Both the methods use 1-dimensional latent variable in VAE. For fair comparison, all the hyper-parameters, including type and size of NNs, learning rate, and batch size, are the same for both the methods.
Figure (ref) shows our method significantly outperforms CEVAE on all cases Both methods work the best under unconfoundedness (“ig.”), as expected. The performances of our method on IV (“inst.”) and proxy (“conf.”) settings match that of CEVAE under unconfoundedness, showing the effective deconfounding. See Appendix (ref) for results on linear outcome. Results for ATE and post-treatment are similar.
{8pt} {5pt}
Here, the true latent ${\textnormal{z}}$ is a PS, and there are no better candidate PSs than ${\textnormal{z}}$, because $f_t$ is invertible and no information can be dropped from ${\textnormal{z}}$. Thus, as shown in Figure (ref), our method learns representation as an approximate affine transformation of the true latent value, as a result of our model identifiability. More latent plots are in Appendix (ref) (the end of the paper). As expected, both recovery and estimation are better with unconditional prior $p({\textnormal{z}}|{\mathbf{x}})$, and we see examples of bad recovery using conditional $p({\textnormal{z}}|{\mathbf{x}},{\textnormal{t}})$ in Appendix Figure (ref). CEVAE shows much lower quality of recovery, particularly with large noises. Under the IV setting, while TEs are estimated as well as for the proxy setting, the relationship to the true latent is significantly obscured, because the true latent is correlated to IV ${\mathbf{x}}$ only given ${\textnormal{t}}$, while we model it by $p({\textnormal{z}}|{\mathbf{x}})$. This confirms that our method does not need to recover the true confounder distribution.
We see our method is robust to the unknown noise level. This indicates that noises are learned by our VAE. Appendix (ref) shows that the noise level affects how well we recover the latent variable.
\endgroup
This experiment shows our balanced estimator matches the state-of-the-art methods specialized for standard unconfoundedness. The IHDP dataset hill2011bayesian is widely used to evaluate machine learning based causal inference methods, e.g. shalit2017estimating, shi2019adapting.
The generating process is as following hill2011bayesian.
where ${\bm{b}}$ is constant with all elements equal to $0.5$, ${\bm{a}}$ and $o$ are a random parameters. Thus, it is not necessary to condition on ${\mathbf{x}}$, but the linear PS ${\bm{a}}^T{\mathbf{x}}$ is sufficient. See Appendix (ref) for more details.
To see our balancing property clearly, we add two components specialized for balancing from shalit2017estimating into our method. First, we use two separate NNs to build the two outcome functions ${\bm{f}}_t({\mathbf{z}}),t=0,1$ in our model (ref). Second, we add to our ELBO (ref) a regularization term, which is the Wasserstein distance cuturi2013sinkhorn between $\mathbb{E}_{\mathcal{D}\sim p({\mathbf{x}}|{\textnormal{t}}=t)}p_{\bm\lambda}({\mathbf{z}}|{\mathbf{x}}),t\in\{0,1\}$.
As shown in Table (ref), Intact-VAE outperforms or matches the state-of-the-art methods. In particular, our method has the best ATE estimation; and it has the best individual-level estimation, adding the two components from shalit2017estimating. We can see in the caption of Table (ref), the specialized additions do not really improve our method, only causing a tradeoff between CATE and ATE estimation, and this may due to the tradeoff between fitting and balancing. And notably, our method outperforms other generative models (CEVAE and GANITE) by large margins.
We find higher than 1-dimensional ${\mathbf{z}}$ in Intact-VAE gives better results, because we have discrete true PS due to the existence of discrete covariates. We report results with 10-dimensional latent variable. The robustness of VAE under model misspecification was also observed by louizos2017causal, where they used 5-dimensional Gaussian latent variable to model a binary ground truth.
We show our method is the best compared with the methods specialized for networked deconfounding, a challenging problem in its own right. Pokec leskovec2014snap is a real world social network dataset. We experiment on a semi-synthetic dataset based on Pokec, introduced in veitch2019using, and use exactly the same pre-processing and generating procedure. The pre-processed network has about 79,000 vertexes (users) connected by 1.3 $\times 10^6$ undirected edges. The subset of users used here are restricted to three living districts that are within the same region. The network structure is expressed by binary adjacency matrix ${\bm{G}}$.
Each user has 12 attributes, among which district, age, or join date is used as a confounder ${\textnormal{z}}$ to build 3 different datasets, with remaining 11 attributes used as covariate ${\mathbf{x}}$. Treatment ${\textnormal{t}}$ and outcome ${\mathbf{y}}$ are synthesised as following:
Note that district is of 3 categories; age and join date are also discretized into three bins. There is a PS that is $g({\textnormal{z}})$, which maps the three categories and values to $\{0.15, 0.5, 0.85\}$.
As in veitch2019using, we split the users into 10 folds, test on each fold and report the mean and std of pre-treatment ATE predictions. We further separate the rest of users (in the other 9 folds) by 6:3, for training and validation. Table (ref) shows the results. In addition, the pre-treatment $\sqrt{\epsilon_{pehe}}$ for Age, District, and Join date confounders are 1.085, 0.686, and 0.699 respectively, practically the same as the ATE errors. veitch2019using do not give individual-level prediction.
Intact-VAE is expected to learn a PS from ${\bm{G}}, {\mathbf{x}}$, if we can exploit the network structure effectively. Given the huge network structure, most users can practically be identified by their attributes and neighborhood structure, which means ${\textnormal{z}}$ can be roughly seen as a deterministic function of ${\bm{G}}, {\mathbf{x}}$. This idea is comparable to Assumptions 2 and 4 in veitch2019using, which postulate directly that a balancing score can be learned in the limit of infinite large network.
To extract information from the network structure, we use Graph Convolutional Network (GCN) DBLP:conf/iclr/KipfW17 in the prior and encoder of Intact-VAE. Note that GCN cannot be trained by mini-batch, instead, we perform batch gradient decent using all data for each iteration, with initial learning rate $10^{-2}$. We use dropout srivastava2014dropout with rate 0.1 to prevent overfitting.
GCN need to take as inputs the network matrix ${\bm{G}}$ and the covariates matrix ${\bm{X}} \coloneqq ({\bm{x}}_1^T,\dotsc,{\bm{x}}_M^T)^T$ of all users, where $M$ is user number, regardless of whether it is in training, validation, or testing phase; and it outputs a representation matrix ${\bm{R}}$, again for all users. To enable sample separation, we need to make sure the treatment and outcome are used only in the respective phase, e.g., $(y_m,t_m)$ of a testing user $m$ is only used in testing. During training, we select the rows in ${\bm{R}}$ that correspond to users in training set. Then, treat this training representation matrix as if it is the covariate matrix for a non-networked dataset, that is, the downstream networks in conditional prior and encoder are the same as in the other two experiments, except that they take $({\bm{R}}_{m,:})^T$ where ${\bm{x}}_m$ was expected as input. And we have respective selection operations for validation and testing. We can still train Intact-VAE including GCN by Adam, simply setting the gradients of non-seleted rows of ${\bm{R}}$ to 0.
In this work, we proposed Intact-VAE for CATE estimation. Our generative model is identifiable and has a sufficient score as its latent variable. Our method outperforms or matches state-of-the-art methods under diverse settings including unobserved confounding. In Sec. (ref), we explained why the current VAE methods are unsatisfactory from a more “causal” viewpoint. We gave theoretical analysis of Intact-VAE under unconfoundedness and discussed parallel and future theoretical work--TE identification without overlap and approaches to identification and optimal estimation under unobserved confounding. We believe this series of work will also pave the way towards principled causal effect estimation by other deep architectures, given the fast advances in deep identifiable models. For example, recently, khemakhem2020ice provide identifiability to deep energy models, and roeder2020linear extend the result to a wide class of state-of-the-art deep discriminative models. We hope this work will inspire other methods based on deep identifiable models.