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.
54,002 characters · 15 sections · 73 citation commands
$$-Intact-VAE: Identifying and Estimating Causal Effects under Limited Overlap
Causal inference imbens2015causal, pearl2009causality, i.e, inferring causal effects of interventions, is a fundamental field of research. In this work, we focus on treatment effects (TEs) based on a set of observations comprising binary labels ${T}$ for treatment/control (non-treated), outcome ${Y}$, and other covariates ${X}$. Typical examples include estimating the effects of public policies or new drugs based on the personal records of the subjects. The fundamental difficulty of causal inference is that we never observe counterfactual outcomes that would have been if we had made the other decision (treatment or control). While randomized controlled trials (RCTs) control biases through randomization and are ideal protocols for causal inference, they often have ethical and practical issues, or suffer from expensive costs. Thus, causal inference from observational data is important.
Causal inference from observational data has other challenges as well. One is confounding: there may be variables, called confounders, that causally affect both the treatment and the outcome, and spurious correlation/bias follows. The other is the systematic imbalance (difference) of the distributions of the covariates between the treatment and control groups––that is, ${X}$ depends on ${T}$, which introduces bias in estimation. A majority of studies on causal inference, including the current work, have relied on unconfoundedness; this means that appropriate covariates are collected so that the confounding can be controlled by conditioning on the covariates. However, such high-dimensional covariates tend to introduce a stronger imbalance between treatment and control.
The current work studies the issue of imbalance in estimating individualized TEs conditioned on ${X}$. Classical approaches aim for covariate balance, ${X}$ independent of ${T}$, by matching and re-weighting Stuart2010matching,rosenbaum2020modern. Machine learning methods have also been exploited; there are semi-parametric methods––e.g., van2018targeted––which improve finite sample performance, as well as non-parametric methods––e.g., wager2018estimation. Notably, from johansson2016learning, there has been a recent increase in interest in balanced representation learning (BRL) to learn representations ${Z}$ of the covariates, such that ${Z}$ independent of ${T}$.
The most serious form of imbalance is the limited (or weak) overlap of covariates, which means that sample points with certain covariate values belong to a single treatment group. In this case, a straightforward estimation of TEs is not possible at non-overlapping covariate values due to lack of data. Some works have focused on providing robustness to limited overlap armstrong2021finite, trimming non-overlapping sample points 10.1093/biomet/asy008, or studying convergence rates based on overlap hong2020inference. Limited overlap is particularly relevant to machine learning methods that exploit high-dimensional covariates. This is because, with higher-dimensional covariates, overlap is harder to satisfy and verify d2020overlap.
To address imbalance and limited overlap, we use a prognostic score hansen2008prognostic; it is a sufficient statistic of outcome predictors and is among the key concepts of sufficient scores for TE estimation. As a function of covariates, it can map some non-overlapping values to an overlapping value in a space of lower-dimensions. For individualized TEs, we consider conditionally balanced representation ${Z}$, such that ${Z}$ is independent of ${T}$ given ${X}$––which, as we will see, is a necessary condition for a balanced prognostic score. Moreover, prognostic score modeling can benefit from methods in predictive analytics and exploit rich literature, particularly in medicine and health hajage2017estimation. Thus, it is promising to combine the predictive power of prognostic modeling and machine learning. With this idea, our method builds on a generative prognostic model that models the prognostic score as a latent variable and factorizes to the score distribution and outcome distribution.
As we consider latent variables and causal inference, identification is an issue that must be discussed before estimation is considered. “Identification” means that the parameters of interest (in our case, representation function and TEs) are uniquely determined and expressed using the true observational distribution. Without identification, a consistent estimator is impossible to obtain, and a model would fail silently; in other words, the model may fit perfectly but will return an estimator that converges to a wrong one, or does not converge at all lewbel2019identification. Identification is even more important for causal inference; because, unlike usual (non-causal) model misspecification, causal assumptions are often unverifiable through observables white2013identification. Thus, it is critical to specify the theoretical conditions for identification, and then the applicability of the methods can be judged by knowledge of an application domain.
A major strength of our generative model is that the latent variable is identifiable. This is because the factorization of our model is naturally realized as a combination of identifiable VAE khemakhem2020variational and conditional VAE sohn2015learning. Based on model identifiability, we develop two identification results for individualized TEs under limited overlap. A similar VAE architecture was proposed in wu2021identifying,wu2021intact; the current study is different in setting, theory, learning objective, and experiments. The previous work studies unobserved confounding but not limited overlap, with different set of assumptions and identification theories. The current study further provides bounds on individualized TE error, and the bounds justify a conditionally balancing term controlled by hyperparameter $\beta$, as an interpolation between the two identifications.
In summary, we study the identification (Sec. (ref)) and estimation (Sec. (ref)) of individualized TEs under limited overlap. Our approach is based on recovering prognostic scores from observed variables. To this end, our method exploits recent advances in identifiable representation––particularly iVAE. The code is in Supplementary Material, and the proofs are in Sec. (ref). Our main contributions are:
Limited overlap. Under limited overlap, luo2017estimating estimate the average TE (ATE) by reducing covariates to a linear prognostic score. farrell2015robust estimates a constant TE under a partial linear outcome model. d2021deconfounding study the identification of ATE by a general class of scores, given the (linear) propensity score and prognostic score. Machine learning studies on this topic have focused on finding overlapping regions oberst2020characterization,dai2020quantifying, or indicating possible failure under limited overlap jesson2020identifying, but not remedies. An exception is johansson2020generalization, which provides bounds under limited overlap. To the best of our knowledge, our method is the first machine learning method that provides identification under limited overlap.
Prognostic scores\space have been recently combined with machine learning approaches, mainly in the biostatistics community. For example, huang2017joint estimate individualized TE by reducing covariates to a linear score which is a joint propensity-prognostic score. tarr2021estimating use SVM to minimize the worst-case bias due to prognostic score imbalance. However, in the machine learning community, few methods consider prognostic scores; zhang2020treatment and hassanpour2019learning learn outcome predictors, without mentioning prognostic score––while johansson2020generalization conceptually, but not formally, connects BRL to prognostic score. Our work is the first to formally connect generative learning and prognostic scores for TE estimation.
Identifiable representation. Recently, independent component analysis (ICA) and representation learning––both ill-posed inverse problems––meet together to yield nonlinear ICA and identifiable representation; for example, using VAEs khemakhem2020variational, and energy models khemakhem2020ice. The results are exploited in causal discovery pmlr-v108-wu20b and out-of-distribution generalization sun2020latent. This study is the first to explore identifiable representations in TE identification.
BRL and related methods\space amount to a major direction. Early BRL methods include BLR/BNN johansson2016learning and TARnet/CFR shalit2017estimating. In addition, yao2018representation exploit the local similarity between data points. shi2019adapting use similar architecture to TARnet, considering the importance of treatment probability. There are also methods that use GAN yoon2018ganite and Gaussian processes alaa2017bayesian. Our method shares the idea of BRL, and further extends to conditional balance––which is natural for individualized TE.
Others. Our work can lay conceptual and theoretical foundations of VAE methods for TEs louizos2017causal,lu2020reconsidering. In addition, some studies consider monotonicity, which is injectivity on $\mathbb{R}$, together with overlap johansson2020generalization,zhang2020learning. See Sec. (ref).
\begingroup
Following imbens2015causal, we assume there exist potential outcomes ${Y}(t) \in \mathbb{R}^d, t \in \{0,1\}$. ${Y}(t)$ is the outcome that would have been observed if the treatment value ${T}=t$ was applied. We see ${Y}(t)$ as the hidden variables that give the factual outcome ${Y}$ under factual assignment ${T}=t$. Formally, ${Y}(t)$ is defined by the consistency of counterfactuals: ${Y} = {Y}(t)$ if ${T}=t$; or simply ${Y}={Y}({T})$. The fundamental problem of causal inference is that, for a unit under research, we can observe only one of ${Y}(0)$ or ${Y}(1)$––w.r.t. the treatment value applied. That is, “factual” refers to ${Y}$ or ${T}$, which is observable; or estimators built on the observables. We also observe relevant covariate(s) ${X} \in \mathcal{X} \subseteq \mathbb{R}^m$, which is associated with individuals, with distribution $\mathcal{D} \coloneqq ({X},{Y},{T}) \sim p({\bm{x}},{\bm{y}},t)$. We use upper-case (e.g. ${T}$) to denote random variables, and lower-case (e.g. $t$) for realizations.
The expected potential outcome is denoted by $\mu_t({\bm{x}}) = \mathbb{E}({Y}(t)|{X}={\bm{x}})$ conditioned on ${X}={\bm{x}}$. The estimands in this work are the conditional ATE (CATE) and ATE, defined, respectively, by:
CATE is seen as an individual-level, personalized, treatment effect, given highly discriminative ${X}$.
Standard results rubin2005causalhernanCausalInferenceWhat2020 show sufficient conditions for TE identification in general settings. They are Exchangeability: ${Y}(t) \raisebox{0.05em}{\rotatebox[origin=c]{90}{$\models$}} {T}| {X}$, and Overlap: $p(t|{\bm{x}}) > 0$ for any ${\bm{x}}\in\mathcal{X}$. Both are required for $t\in\{0,1\}$. When $t$ appears in statements without quantification, we always mean “for both $t$”. Often, Consistency is also listed; however, as mentioned, it is better known as the well-definedness of counterfactuals. Exchangeability means, just as in RCTs, but additionally given ${X}$, that there is no correlation between factual ${T}$ and potential ${Y}(t)$. Note that the popular assumption ${Y}(0),{Y}(1) \raisebox{0.05em}{\rotatebox[origin=c]{90}{$\models$}} {T}| {X}$ is stronger than ${Y}(t) \raisebox{0.05em}{\rotatebox[origin=c]{90}{$\models$}} {T}| {X}$ and is not necessary for identification hernanCausalInferenceWhat2020. Overlap means that the supports of $p({\bm{x}}|t=0)$ and $p({\bm{x}}|t=1)$ should be the same, and this ensures that there are data for $\mu_t({\bm{x}})$ on any $({\bm{x}}, t)$.
We rely on consistency and exchangeability, but in Sec. (ref), will relax the condition of the overlapping covariate to allow some non-overlapping values ${\bm{x}}$––that is, covariate ${X}$ is limited-overlapping. In this paper, we also discuss overlapping variables other than ${X}$ (e.g., prognostic scores), and provide a definition for any random variable ${V}$ with support $\mathcal{V}$ as follows:
Our method aims to recover a prognostic score hansen2008prognostic, adapted as a Pt-score in Definition (ref). On the other hand, balancing scores rosenbaum1983central ${\bm{b}}({X})$ are defined by ${T}\raisebox{0.05em}{\rotatebox[origin=c]{90}{$\models$}}{X}| {\bm{b}}({X})$, of which the propensity score $p(t=1|{X})$ is a special case. See Sec. (ref) for detail.
Note that a PtS is by definition two functions; thus, overlapping $\mathbb{P}_t({X})$ means both $\mathbb{P}_0({X})$ and $\mathbb{P}_1({X})$ are overlapping. Why not balancing scores? While balancing scores ${\bm{b}}({X})$ have been widely used in causal inference, PtSs are more suitable for discussing overlap. Our purpose is to recover an overlapping score for limited-overlapping ${X}$. It is known that overlapping ${\bm{b}}({X})$ implies overlapping ${X}$ d2020overlap, which counters our purpose. In contrast, overlapping PS does not imply overlapping ${\bm{b}}({X})$ (see Sec. (ref) for an example). Moreover, with theoretical and experimental evidence, it is recently conjectured that PtSs maximize overlap among a class of sufficient scores, including ${\bm{b}}({X})$ d2021deconfounding. In general, hajage2017estimation show that prognostic score methods perform better––or as well as––propensity score methods.
Below is a corollary of Proposition 5 in hansen2008prognostic; note that $\mathbb{P}_t({X})$ satisfies exchangeability.
With the knowledge of $\mathbb{P}_t$ and $p_{{Y}|\mathbb{P}_{\hat t},T}$, we choose one of $\mathbb{P}_0,\mathbb{P}_1$ and set $t=\hat{t}$ in the density function, w.r.t the $\mu_{\hat t}$ of interest. This counterfactual assignment resolves the problem of non-overlap at ${\bm{x}}$. Note that a sample point with ${X}={\bm{x}}$ may not have ${T}=\hat{t}$.
We mainly consider additive noise models for ${Y}(t)$, which ensures the existence of PtSs.
The potential outcomes can be defined by the DGP as ${Y}(t):={\bm{f}}^*(\mathbb{M}({X}),t) + {{\mathbf{e}}}$. The DGP also specifies how other variables causally affect ${Y}$. For example, ${X}$ affects ${Y}$ through $\mathbb{M}$; and thus $\mathbb{M}({X})$ is the effect modifier hansen2008prognostic––which is often components of ${X}$ affecting ${Y}$ directly. Additive noise models are used in nonparametric regression methods for TEs caron2020estimating.
Under (ref), we can find natural examples of PS and PtS. 1) $\mathbb{P}_t({X})\coloneqq{\bm{f}}^*_t(\mathbb{M}({X}))=\mu_t({X})$\footnote{We often write $t$ of the function argument in subscripts, indicating possible counterfactual assignments.} is a PtS and not PS; 2) $\mathbb{M}$ is a PS (${X}$ is a trivial PS); and 3) $\mathbb{P}({X})\coloneqq(\mu_0({X}), \mu_1({X}))$ is a PS.
We use PS and PtS to construct a representation for CATE estimation, and their balance is important. Obviously, a PS $\mathbb{P}({X})$ is a conditionally balanced representation (defined as ${Z}\raisebox{0.05em}{\rotatebox[origin=c]{90}{$\models$}}{T}|{X}$ in Introduction), because $\mathbb{P}(X)$ does not depend on ${T}$ given ${X}$. Note that a PtS is conditionally balanced if and only if it is a PS. Thus, we introduce the notion of balanced PtS in a non-rigorous way: a PtS $\mathbb{P}_t$ is called balanced if the value of a measure for the conditional independence $\mathbb{P}_{{T}}({X}) \raisebox{0.05em}{\rotatebox[origin=c]{90}{$\models$}} {T} | {X}$ is small.
\begingroup
{8pt} {5pt}
In Sec. (ref), we specify the generative prognostic model $p({\bm{y}},{\bm{z}}|{\bm{x}},t)$, and show its identifiability. In Sec. (ref), we prove the identification of CATEs, which is one of our main contributions. The theoretical analysis involves only our generative model (i.e., prior and decoder), but not the encoder. The encoder is not part of the generative model and is involved as an approximate posterior in the estimation, which is studied in Sec. (ref).
We present the necessary definitions and results, and defer more explanations to Sec. (ref). The contents of this subsection are essentially taken from wu2021principled, but included here for completeness.
Our goal is to build a model that can be learned by VAE from observational data to obtain a PtS, or better, a PS, via the latent variable ${Z}$. The generative prognostic model of the proposed method is as follows:
The factor $p_{{\bm{f}}}({\bm{y}}|{\bm{z}},t)$ is our decoder, which models $p_{{Y}|\mathbb{P}_t,T}({\bm{y}}|P,t)$ in (ref); and $p_{\bm\lambda}({\bm{z}}|{\bm{x}},t)$ is the conditional prior, which models $\mathbb{P}_{{T}}({X})$. The outcome assumes an additive noise model such that $\bm\epsilon \sim p_{\bm\epsilon}$ denotes the noise model. The prior is a factorized Gaussian, where $\bm\lambda_{{T}}({X})\coloneqq \operatorname{diag}^{-1}({\bm{k}}_{{T}}({X}))({\bm{h}}_{{T}}({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. We denote $n\coloneqq\dim({Z})$.
For inference, the standard argument derives the ELBO:
Note that the encoder $q$ conditions on all the observables $({X},{Y},{T})$; this fact plays an important role in Section (ref). This architecture is called Intact-VAE (Identifiable treatment-condi\textit{t}ional VAE). See Figure (ref) for comparison in terms of graphical models. See Sec. (ref) for basics of VAEs.
\endgroup
Our model identifiability extends the theory of iVAE, and the following conditions are inherited.
Under (M1) and (D1), we obtain the following identifiability of the parameters in the model: if $p_{{\bm{\theta}}}({\bm{y}}|{\bm{x}},t)=p_{{\bm{\theta}}'}({\bm{y}}|{\bm{x}},t)$, we have, for any ${\bm{y}}_t$ in the image of ${\bm{f}}_t$:
where $\operatorname{diag}({\bm{a}})$ is an invertible $n$-diagonal matrix and ${\bm{b}}$ is an $n$-vector, both of which depend on $\bm\lambda_t({\bm{x}})$ and $\bm\lambda'_t({\bm{x}})$. The essence of the result is that ${\bm{f}}'_t = {\bm{f}}_t \circ \mathcal{A}_t$; that is, ${\bm{f}}_t$ can be identified (learned) up to an affine transformation $\mathcal{A}_t$. See Sec. (ref) for the proof and a relaxation of (ref). In this paper, symbol $'$ (prime) always indicates another parameter (variable, etc.): ${\bm{\theta}}'=({\bm{f}}',\bm\lambda')$.
In this subsection, we present two results of CATE identification based on the recovery of equivalent PS and PtS, respectively. Since PtSs are functions of ${X}$, the theory assumes a noiseless prior for simplification, i.e., ${\bm{k}}({X})=\bm0$; the prior ${Z}_{\bm\lambda,t} \sim p_{\bm\lambda}({\bm{z}}|{\bm{x}},t)$ degenerates to function ${\bm{h}}_{t}({\bm{x}})$.
PtSs with dimensionality lower than or equal to $d=\dim({Y})$ are essential to address limited overlapping, as shown below. We set $n=d$ because $\mu_t$ is a PtS of the same dimension as ${Y}$ under (ref). In practice, $n=d$ means that we seek a low-dimensional representation of ${X}$. (ref) makes the dimensionality explicit and reduces to (ref) if the only possibility is that $\mathbb{P}_t=\mu_t$ and ${\bm{j}}_t$ is identity.
We use (ref) instead of (ref) hereafter. Clearly, $\mathbb{P}_t$ in (ref) is PtS. In addition, injectivity and $n=d$ ensure that $n=\dim({Y}) \geq \dim(\mathbb{P}_t)$. Similarly, (ref) ensures that $n \geq \dim(\mathbb{P})$ for PS.
(ref) means that CATEs are given by $\mu_0$ and an invertible function ${\bm{i}}\coloneqq{\bm{j}}_1\circ{\bm{j}}_0^{-1}$. See Sec. (ref) for more discussions and real-world examples.
With (ref) or (ref), overlapping ${X}$ can be relaxed to overlapping PS or PtS plus the following:
Note that (ref) is only required for optimal ${\bm{h}}$ that satisfies $\mathbb{E}_{p_{\bm\theta}}({Y}|{X},{T})=\mathbb{E}({Y}|{X},{T})$ in Proposition (ref), or $p_{{\bm{\theta}}}({\bm{y}}|{\bm{x}},t)=p({\bm{y}}|{\bm{x}},t)$ in Theorem (ref). The intuition is that $\mathbb{P}_t$ maps non-overlapping ${\bm{x}}$ to an overlapping value, and ${\bm{h}}_t$ preserves this property through learning. Linear $\mathbb{P}_t$ and ${\bm{h}}_t$ imply (ref) and are often assumed, e.g., in huang2017joint,luo2017estimating,d2021deconfounding. Linear outcome models farrell2015robust,schuler2020increasing are also common.
Our first identification, Proposition (ref), relies on (ref) and our generative model, without model identifiability (so differentiable ${\bm{f}}_t$ is not needed).
In essence, i) the true DGP 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{P}_t$; and ii) $\mathbb{P}_t$ is recovered up to ${\bm{v}}$, and ${Y}(t)\raisebox{0.05em}{\rotatebox[origin=c]{90}{$\models$}}{X}|\mathbb{P}_t({X})$ is preserved––with same ${\bm{v}}$ for both $t$. Theorem (ref) below also achieves the essence i) and ii), under $\mathbb{P}_0\neq\mathbb{P}_1$.
The existence of PS is more preferred, because it satisfies overlap and (ref) more easily than PtS which requires the conditions for each of the two functions of PtS. However, the existence of low-dimensional PS is uncertain in practice when our knowledge of the DGP is limited. Thus, we depend on Theorem (ref) based on the model identifiability to work under PtS which generally exists.
Theorem (ref) implies that, without PS, we need to know or learn the distribution of hidden noise $\bm\epsilon$ to have $p_{{{\mathbf{e}}}}=p_{\bm\epsilon}$. Proposition (ref) and Theorem (ref) achieve recovery and identification in a complementary manner; the former starts from the prior by $\mathbb{P}_0=\mathbb{P}_1$ and ${\bm{h}}_0={\bm{h}}_1$, while the latter starts from the decoder by $\mathcal{A}_0=\mathcal{A}_1$ and $p_{{{\mathbf{e}}}}=p_{\bm\epsilon}$. We see that $\mathcal{A}_0=\mathcal{A}_1$ acts as a kind of balance because it replaces $\mathbb{P}_0=\mathbb{P}_1$ (balanced PtS) in Proposition (ref). We show in Sec. (ref) a sufficient and necessary condition (ref) on data that ensures $\mathcal{A}_0=\mathcal{A}_1$. Note that the singularities due to $k\to 0$ (e.g., $\bm\lambda \to \bm 0$) cancel out in (ref). See Sec. (ref) for more on the complementarity between the two identifications.
In Sec. (ref), we see that the existence of PS (Proposition (ref)) is preferable in identifying the true DGP up to an equivalent expression––while Theorem (ref) allows us to deal with PtS by adding other conditions. In learning our model with data, we assume that there is a PtS, and the decomposition of (ref) holds. However, such decompositions are not unique in general, and they are equivalent for CATE identification; and those with more balanced equivalent PtS are preferable. In this sense, we want to not only recover PtS, but also discover equivalent PS, if possible. This idea is common in practice. For example, in a real-world nutrition study huang2017joint, a reduction of 11 covariates discovers a 1-dimensional linear PS.
We consider two ways to discover and recover a equivalent PS (or balanced PtS) by a VAE. One is to use a prior which does not depend on $t$, indicating a preference for PS. Namely, we set $\bm\lambda_0=\bm\lambda_1=:\bm\Lambda$ and have $p_{\bm\Lambda}({\bm{z}}|{\bm{x}})$ as the prior in (ref). The decoder and encoder are factorized Gaussians:
where $\bm\phi =({\bm{r}},{\bm{s}})$. The other is to introduce a hyperparameter $\beta$ in the ELBO as in $\beta$-VAE higgins2016beta. The modified ELBO with $\beta$, up to the additive constant, is derived as:
For convenience, here and in $\mathcal{L}_{{\bm{f}}}$ in Sec. (ref), we omit the summation as if ${Y}$ is univariate. The encoder $q_{\bm\phi}$ depends on $t$ and can realize a PtS. With $\beta$, we control the trade-off between the first and second terms: the former is the divergence of the posterior from the balanced prior, and the latter is the reconstruction of the outcome. Note that a larger $\beta$ encourages the conditional balance ${Z}\raisebox{0.05em}{\rotatebox[origin=c]{90}{$\models$}}{T}|{X}$ on the posterior. By choosing $\beta$ appropriately, e.g., by validation, the ELBO can recover a balanced PtS while fitting the outcome well. In summary, we base the estimation on Proposition (ref) and PS as much as possible, but step into Theorem (ref) and noise modeling required by $p_{{{\mathbf{e}}}}=p_{\bm\epsilon}$ when necessary.
Note also that the parameters ${\bm{g}}$ and ${\bm{k}}$, which model the outcome noise and express the uncertainty of the prior, respectively, are both learned by the ELBO. This deviates from the theoretical conditions described in Sec. (ref), but it is more practical and yields better results in our experiments. See Sec. (ref) for more ideas and connections behind the ELBO.
Once the VAE is learned\footnote{As usual, we expect the variational inference and optimization procedure to be (near) optimal; that is, consistency of VAE. Consistent estimation using the prior is a direct corollary of the consistent VAE. see Sec. (ref) for formal statements and proofs. Under Gaussian models, it is possible to prove the consistency of the posterior estimation, as shown in bonhomme2019posterior.} by the ELBO, the estimate of the expected potential outcomes 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. We mainly consider the case where ${\bm{x}}$ is observed in the data, and the sample of $({Y},{T})$ is taken from the data given ${X}={\bm{x}}$. When ${\bm{x}}$ is not in the data, we replace $q_{\bm\phi}$ with $p_{\bm\Lambda}$ in (\theequation) (see Sec. (ref) for details and (ref) for results). Note that ${\hat{t}}$ in (ref) indicates a counterfactual assignment that may not be the same as the factual ${T}=t$ in the data. That is, we set ${T}=\hat{t}$ in the decoder. The assignment is not applied to the encoder which is learned from factual ${X},{Y},{T}$ (see also the explanation of $\epsilon_{CF,t}$ in Sec. (ref)). The overall algorithm steps are i) train the VAE using (ref), and ii) infer CATE $\hat{\tau}({\bm{x}})=\hat{\mu}_1({\bm{x}})-\hat{\mu}_0({\bm{x}})$ by (ref).
We formally justify our ELBO (ref) from the BRL viewpoint. We show that the conditional BRL via the first term of the ELBO results from bounding a CATE error; particularly, the error due to the imprecise recovery of ${\bm{j}}_t$ in (ref) is controlled by the ELBO. Previous works shalit2017estimating,lu2020reconsidering instead focus on unconditional balance and bound PEHE which is marginalized on ${X}$. Sec. (ref) experimentally shows the advantage of our bounds and ELBO. Further, we connect the bounds to identification and consider noise modeling through ${\bm{g}}_t({\bm{z}})$. Sec Sec. (ref) for detail.
We introduce the objective that we bound. Using (ref) to estimate CATE, $\hat{\tau}_{{\bm{f}}}({\bm{z}})\coloneqq{\bm{f}}_1({\bm{z}})-{\bm{f}}_0({\bm{z}})$ is marginalized on $q({\bm{z}}|{\bm{x}})$. On the other hand, the true CATE, given the covariate ${\bm{x}}$ or score ${\bm{z}}$, is:
where ${\bm{j}}_t$ is associated with a balanced PtS $\mathbb{P}_t$ discovered as the target of recovery by our VAE. Accordingly, given ${\bm{x}}$, the error of posterior CATE, with or without knowing $\mathbb{P}_t$, is defined as
We bound $\epsilon_{{\bm{f}}}$ instead of $\epsilon_{{\bm{f}}}^{*}$ because the error between $\tau({X})$ and $\tau_{{\bm{j}}}({Z})$ is small––if the balanced $\mathbb{P}_t$ is recovered, then ${\bm{z}}\approx\mathbb{P}_0({\bm{x}})\approx\mathbb{P}_1({\bm{x}})$ in (ref). We consider the error between $\hat{\tau}_{{\bm{f}}}$ and $\tau_{{\bm{j}}}$ below. We define the risks of outcome regression, into which $\epsilon_{{\bm{f}}}$ is decomposed.
With ${Y}(t)$ involved, $\mathcal{L}_{{\bm{f}}}$ is a potential outcome loss on ${\bm{f}}$, weighted by ${\bm{g}}$. The factual and counterfactual counterparts, $\epsilon_{F,t}$ and $\epsilon_{CF,t}$, are defined accordingly. In $\epsilon_{F,t}$, unit ${\bm{u}}=({\bm{x}}, {\bm{y}}, t)$ is involved in the learning of $q_t({\bm{z}}|{\bm{x}})$, as well as in $\mathcal{L}_{{\bm{f}}}({\bm{z}},t)$ since ${Y}(t)={\bm{y}}$ for the unit. In $\epsilon_{CF,t}$, however, unit ${\bm{u}}'=({\bm{x}}, {\bm{y}}', 1-t)$ is involved in $q_{1-t}({\bm{z}}|{\bm{x}})$, but not in $\mathcal{L}_{{\bm{f}}}({\bm{z}},t)$ since ${Y}(t)\neq{\bm{y}}'={Y}(1-t)$.
Thus, the regression error (second) term in ELBO (ref) controls $\epsilon_{F,t}$ via factual data. On the other hand, $\epsilon_{CF,t}$ is not estimable due to the unobservable ${Y}(1-{T})$, but is bounded by $\epsilon_{F,t}$ plus $M\mathbb{D}({\bm{x}})$ in Theorem (ref) below––which, in turn, bounds $\epsilon_{{\bm{f}}}$ by decomposing it to $\epsilon_{F,t}$, $\epsilon_{CF,t}$, and $\mathbb{V}_{{Y}}$.
$\mathbb{D}({\bm{x}})$ measures the imbalance between $q_t({Z}|{\bm{x}})$ and is symmetric for $t$. Correspondingly, the KL term in ELBO (ref) is also symmetric for $t$ and balances $q_t({\bm{z}}|{\bm{x}})$ by encouraging ${Z}\raisebox{0.05em}{\rotatebox[origin=c]{90}{$\models$}}{T}|{X}$ for the posterior. $\mathbb{V}_{{Y}}({\bm{x}})$ reflects the intrinsic variance in the DGP and can not be controlled.
Estimating $G,M$ is nontrivial. Instead, we rely on $\beta$ in the ELBO to weight the terms in (\theequation). We do not need two hyperparameters since $G$ is implicitly controlled by the third term in ELBO (ref), which is a norm constraint. $\beta$ is a trade-off between the conditional balance of learned PtS (affected by ${\bm{f}}_t$), and precision/effective sample size of outcome regression––and can be seen as the probabilistic counterpart of tarr2021estimating and kallus2018more.
We compare our method with existing methods on three types of datasets. Here, we present two experiments; the remaining one on the Pokec dataset is deferred to Sec. (ref). 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, as a surrogate of square CATE error $\epsilon_{cate}({\bm{x}})=\mathbb{E}_{\mathcal{D}|{\bm{x}}}[(y(1)-y(0))-\hat{\tau}({\bm{x}})]^2$, the empirical PEHE $\epsilon_{pehe}\coloneqq \mathbb{E}_{\mathcal{D}}\epsilon_{cate}({\bm{x}})$ hill2011bayesian, which is the average square CATE error.
Unless otherwise indicated, for each function ${\bm{f}},{\bm{g}},{\bm{h}},{\bm{k}},{\bm{r}},{\bm{s}}$ in ELBO (ref), we use a multilayer perceptron, with $3 * 200$ hidden units, and ReLU activations. Further, $\bm\Lambda=({\bm{h}},{\bm{k}})$ depends only on ${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. More details on hyper-parameters and settings are given in each experiment.
\begingroup
We generate synthetic datasets following (ref). Both ${X}\sim \mathcal{N}(\bm\mu, \bm \sigma)$ and ${W}$ are factorized Gaussians. $\bm\mu, \bm \sigma$ are randomly sampled. The functions ${\bm{h}},{\bm{k}},l$ are linear. Outcome models $f_0,f_1$ are built by NNs with invertible activations. ${Y}$ is univariate, $\dim({X})=30$, and $\dim({W})$ ranges from 1 to 5. ${W}$ is a PS, but the dimensionality is not low enough to satisfy the injectivity in (ref), when $\dim({W}) > 1$. We have 5 different overlap levels controlled by $\omega$ that multiplies the logit value. See Sec. (ref) for details and more results on synthetic datasets.
{8pt} {5pt}
With the same $(\dim({W}),\omega)$, we evaluate our method and CFR on 10 random DGPs, with different sets of functions $f,g,{\bm{h}},{\bm{k}},l$ in (ref). For each DGP, we sample 1500 data points, and split them into 3 equal sets for training, validation, and testing. We show our results for different hyperparameter $\beta$. For CFR, we try different balancing parameters and present the best results (see the Appendix for detail).
In each panel of Figure (ref), we adjust one of $\omega, \dim({W})$, with the other fixed to the lowest. As implied by our theory, our method, with only 1-dimensional ${Z}$, performs much better in the left panel (where $\dim({W}) = 1$ satisfies $\ref{ass:p_inj}$) than in the right panel (when $\dim({W}) > 1$). Although CFR uses 200-dimensional representation, in the left panel our method performs much better than CFR; moreover, in the right panel CFR is not much better than ours. Further, our method is much more robust against different DGPs than CFR (see the error bars). Thus, the results indicate the power of identification and recovery of scores. (see Figure (ref) also).
{8pt} {5pt}
Under the lowest overlap level ($\omega=22$), large $\beta(=2.5,3)$ shows the best results, which accords with the intuition and bounds in Sec. (ref). When $\dim({W}) > 1$, $f_t$ in (\theequation) is non-injecitve and learning of PtS is necessary, and thus, larger $\beta$ has a negative effect. In fact, $\beta=1$ is significantly better than $\beta=3$ when $\dim({W})>2$. We note that our method, with a higher-dimensional ${Z}$, outperforms or matches CFR also under $\dim({W})>1$ (see Appendix Figure (ref)). Thus, the performance gap under $\dim({W})>1$ in Figure (ref) should be due to the capacity of NNs in $\beta$-Intact-VAE. In Appendix Figure (ref) for ATE error, CFR drops performance w.r.t overlap levels. This is evidence that CFR and its unconditional balance overly focus on PEHE (see Sec. (ref) for detail).
When $\dim({W})=1$, there are no better PSs than ${W}$, because $f_t$ is invertible and no information can be dropped from ${W}$. Thus, our method stably learns ${Z}$ as an approximate affine transformation of the true ${W}$, showing identification. An example is shown in Figure (ref), and more plots are in Appendix Figure (ref). For comparison, we run CEVAE louizos2017causal, which is also based on VAE but without identification; CEVAE shows much lower quality of recovery. As expected, both recovery and estimation are better with the balanced prior $p_{\bm\Lambda}({\bm{z}}|{\bm{x}})$, and we can see examples of bad recovery using $p_{\bm\lambda}({\bm{z}}|{\bm{x}},t)$ in Appendix Figure (ref).
\endgroup
This experiment shows our conditional BRL matches state-of-the-art BRL methods and does not overly focus on PEHE. The IHDP hill2011bayesian is a widely used benchmark dataset; while it is less known, its covariates are limited-overlapping, and thus it is used in johansson2020generalization which considers limited overlap. The dataset is based on an RCT, but Race is artificially introduced as a confounder by removing all treated babies with nonwhite mothers in the data. Thus, Race is highly limited-overlapping, and other covariates that have high correlation to Race, e.g, Birth weight kelly2009does, are also limited-overlapping. See Sec. (ref) for detail and more results.
There is a linear PS (linear combination of the covariates). However, most of the covariates are binary, so the support of the PS is often on small and separated intervals. Thus, the Gaussian latent ${Z}$ in our model is misspecified. We use 10-dimensional ${Z}$ to address this, similar to louizos2017causal. We set $\beta=1$ since it works well on synthetic datasets with limited overlap.
As shown in Table (ref), $\beta$-Intact-VAE outperforms or matches the state-of-the-art methods. Notably, our method outperforms other generative models (CEVAE and GANITE) by large margins. Our method has the best ATE estimation and is only slightly worse than CFR for $\epsilon_{pehe}$. This fact reflects that $\epsilon_{pehe}$ is not a good criterion for CATE estimation because it is the marginalized CATE error––one expects less ATE error with overall less CATE error on individual-level, while PEHE focuses on those ${X}={\bm{x}}$ with high probability and/or large $\epsilon_{cate}({\bm{x}})$. Indeed, the unconditional balance in shalit2017estimating is based on bounding PEHE, thus results in sub-optimal ATE estimation (see also Appendix Figure (ref) where CFR gives larger ATE errors with less overlap).
To examine conditional v.s unconditional balance clearly, we modify our method and add two components for unconditional balance from shalit2017estimating (see the Appendix), and compare the modified version to the original. The over-focus on PEHE of the unconditional balance from CFR is seen more clearly in the modified version. With different values of the hyperparameter, unconditional balance does not improve (and barely affects) ATE estimation; it does affect PEHE more significantly, but often gives worse PEHE unless the hyperparameter is fine-tuned (with value 0.1).
We proposed a method for CATE estimation under limited overlap. Our method exploits identifiable VAE, a recent advance in generative models, and is fully motivated and theoretically justified by causal considerations: identification, prognostic score, and balance. Experiments show evidence that the injectivity of ${\bm{f}}_t$ in our model is possibly unnecessary because $\dim({Z})>\dim({Y})$ yields better results. A theoretical study of this is an interesting future direction. We believe that VAEs are suitable for principled causal inference owing to their probabilistic nature, if not compromised by ad hoc heuristics (see Sec. (ref)). wu2021principled introduce some newest ideas of this project.