The exact contents of citations.db main_text.text for this paper — one flattened LaTeX string, title through conclusion, appendix excluded, unmodified except for removing email addresses. This is what our citation measures are computed over.
56,968 characters
Towards Principled Causal Effect Estimation by Deep Identifiable Models
\maketitle
\begin{abstract}
As an important problem in causal inference, we discuss the estimation of treatment effects (TEs). Representing the confounder as a latent variable, we propose Intact-VAE, a new variant of variational autoencoder (VAE), motivated by the prognostic score that is sufficient for identifying TEs. Our VAE also naturally gives representations balanced for treatment groups, using its prior. Experiments on (semi-)synthetic datasets show state-of-the-art performance under diverse settings, including unobserved confounding. Based on the identifiability of our model, we prove identification of TEs under unconfoundedness, and also discuss (possible) extensions to harder settings.
\end{abstract}
\section{Introduction}
Causal inference \citep{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 \textit{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 \textit{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 \textit{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 \textit{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 \citep{Stuart2010matching,rosenbaum2020modern}. There are semi-parametric methods \citep[e.g. TMLE,][]{van2018targeted}, which have better finite sample performance, and also non-parametric methods \citep[e.g.~CF,][]{wager2018estimation}. Notably, there is a recent rise of interest in learning balanced representation of covariates, which is independent of treatment groups, starting from \cite{johansson2016learning}.
There are a few lines of works that address the difficult but important problem of \textit{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{sec:setup}.
Instead, many methods assume special structures among the variables, such as instrumental variables (IVs) \citep{angrist1996identification}, proxy variables \citep{tchetgen2020introduction}, network structure \citep{veitch2019using}, and multiple causes \citep{wang2019blessings}. Among them, instrumental and proxy variables are most commonly exploited. \textit{Instrumental variables} are not affected by unobserved confounders, influencing the outcome only through the treatment. On the other hand, \textit{proxy variables} are causally connected to unobserved confounders, but are not confounding the treatment and outcome by themselves. Other methods use restrictive parametric models \citep{allman2009identifiability}, or only give interval estimation \citep{manski2009identification, kallus2019interval}.
In this work, we challenge the problem of estimating TEs under unobserved confounding. We in particular discuss the \textit{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 \citep{wu2021beta} addressing limited overlap and future work(s) under unobserved confounding.
Our method exploits the important concepts of sufficient scores for TE estimation \citep{hansen2008prognostic,rosenbaum1983central} and also the recent advance of VAE with \textit{identifiable} latent variable, which is determined by the true observational distribution \citep[iVAE]{khemakhem2020variational}.
VAEs \citep{kingma2019introduction} are suitable for causal estimation thanks to its probabilistic nature. However, most VAE methods for TEs, e.g., \cite{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 \textit{causal representation} \citep{scholkopf2021toward};
recovering the true confounder is not necessary.
The code is uploaded to OpenReview, and the proofs are in Appendix~\ref{sec:proofs}. Our main contributions are:
\begin{enumerate}[topsep=0pt, partopsep=0pt, itemsep=0pt, parsep=0pt, leftmargin=13pt]
\item[1)] A new identifiable VAE, Intact-VAE, as a balanced estimator for individualized TEs;
\item[2)]
Experimental comparison to state-of-the-art methods under diverse settings;
\item[3)] Proof of TE identification via recovery of sufficient scores, under unconfoundedness;
\item[4)] Discussions of further theoretical developments and principled TE estimation using VAEs.
\end{enumerate}
An early version of this work, which proposed the same VAE architecture, is in \cite{wu2021identifying}.
\subsection{Related work}
As detailed in Sec.~\ref{sec:te_vae}, 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.
\textbf{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 \citep{hyvarinen2016unsupervised, hyvarinen2019nonlinear} and VAE \citep{khemakhem2020variational}. The results are exploited in causal discovery \citep{pmlr-v108-wu20b} and causal representation learning \citep{shen2020disentangled}. To the best of our knowledge, this work is the \textit{first} to explore this
identifiability in TE estimation.
\textbf{Causal inference with auxiliary structures.}
CEVAE \citep{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{sec:cevae}. Also with proxies, \cite{kallus2018causal} use matrix factorization to infer the confounders, and
\cite{mastouri2021proximal} use kernel methods to solve the underlying Fredholm integral equation.
IVs are also exploited in machine learning, there are methods using deep NNs \citep{hartford2017deep} and kernels \citep{singh2019kernel,muandet2019dual}.
\textbf{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 \citep{johansson2016learning}, and TARnet/CFR \citep{shalit2017estimating}.
\cite{shi2019adapting} use similar architecture to TARnet, considering the importance of treatment probability. There are also methods using GAN \citep[GANITE]{yoon2018ganite} and Gaussian process \citep{alaa2017bayesian}. Our method shares the idea of balanced representation learning.
\begingroup
\section{Preliminaries}
\label{sec:setup}
\subsection{Treatment effects and identification}
Following \cite{imbens2015causal}, we begin by defining \textit{potential outcomes} (or \textit{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 \textit{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
\vspace{-.05in}
\begin{equation}
\label{eq:ce}
\tau({\bm{x}}) = \mu_1({\bm{x}}) - \mu_0({\bm{x}}),\quad \nu = \mathbb{E}(\tau({\bm{x}})).
\end{equation}
CATE can be seen as an \textit{individual-level} TE, if conditioned on
highly discriminative covariates.
\setlength{\columnsep}{8pt}
\setlength{\intextsep}{5pt}
\begin{wrapfigure}{r}{0.18\textwidth}
\vspace{-.2in}
\begin{center}
\includegraphics[width=0.18\textwidth]{setting-4.pdf}
\end{center}
\vspace{-.2in}
\caption{
\footnotesize{A possible causal graph of unobserved confounding.}
}
\label{fig:setting}
\end{wrapfigure}
\textit{Identification of TEs} means that, the true observational distribution uniquely determines and gives the ATE or, better, CATE.
Adapting standard identification results \citep{rubin2005causal}\citep[Ch.~3]{hernanCausalInferenceWhat2020}, 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 \textit{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 \citep[pp.~15]{hernanCausalInferenceWhat2020}.
A general example of causal structure that satisfies the three conditions is shown in Figure \ref{fig:setting}, although further structural constraints might be necessary for theoretical developments (see Sec.~\ref{sec:conf})..
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 \textit{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 \eqref{eq:id}, using assumption \textbf{(A)} in the second equality.
\begin{equation}
\label{eq:id}
\mu_t({\bm{x}}) = \mathbb{E}(\mathbb{E}({\mathbf{y}}(t)|{\mathbf{u}},{\bm{x}})) = \mathbb{E}(\mathbb{E}({\mathbf{y}}|{\mathbf{u}},{\bm{x}},{\textnormal{t}}=t)) =\textstyle \int (\int p(y|{\bm{u}},{\bm{x}},t)ydy)p({\bm{u}}|{\bm{x}})d{\bm{u}}.
\end{equation}
If the variable ${\mathbf{u}}$ is observed, then \eqref{eq:id} identifies CATE. However, if ${\mathbf{u}}$ is an \textit{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}}]$.
\subsection{Prognostic score}
Our method models prognostic scores \citep{hansen2008prognostic}, adapted as \textit{Pt-scores} in this paper,
closely related to the important concept of balancing score \citep{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{sec:scores} for details).
\begin{definition}
\label{def:scores}
A \textit{Pt-score} (PtS) is two functions $\mathbb{P}_t({\mathbf{u}},{\mathbf{x}})$ ($t=0,1$) such that ${\mathbf{y}}(t)\raisebox{0.05em}{\rotatebox[origin=c]{90}{$\models$}}{\mathbf{u}},{\mathbf{x}}|\mathbb{P}_t({\mathbf{u}},{\mathbf{x}})$.
A PtS is called a \textit{P-score} (PS) if $\mathbb{P}_0=\mathbb{P}_1$.
\end{definition}
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 \citep{hansen2008prognostic}; and it is a causal representation \citep{scholkopf2021toward} of the direct cause on ${\mathbf{y}}$, summarizing the effects of $({\mathbf{u}},{\mathbf{x}})$.
The independence property of PtS (Lemma \ref{prop:pscore} in Appendix),
\begin{equation}
\label{eq:ps_indp}
{\mathbf{y}}(t)\raisebox{0.05em}{\rotatebox[origin=c]{90}{$\models$}}{\textnormal{t}},{\mathbf{u}},{\mathbf{x}}|\mathbb{P}_t({\mathbf{u}},{\mathbf{x}}),
\end{equation}
is used in second equality of \eqref{eq:cate_by_bts} in Theorem \ref{cate_by_bts} which extends Proposition 5 in \cite{hansen2008prognostic}.
\begin{theorem}[CATE by PtS]
\label{cate_by_bts}
If $\mathbb{P}_t$ is a PtS,
then CATE can be given by
\begin{equation}
\label{eq:cate_by_bts}
\begin{split}
\mu_{t}({\bm{x}}) &= \mathbb{E}(\mathbb{E}({\mathbf{y}}({t})|\mathbb{P}_{t}({\mathbf{u}},{\bm{x}}),{\bm{x}})) = \mathbb{E}(\mathbb{E}({\mathbf{y}}|\mathbb{P}_{t}({\mathbf{u}},{\bm{x}}),{t})) \\ &=\textstyle \int (\int p_{{\mathbf{y}}|\mathbb{P}_{ t},{\textnormal{t}}}(y|P,{t})ydy)p_{\mathbb{P}_t|{\mathbf{x}}}(P|{\bm{x}})dP,
\end{split}
\end{equation}
where ${\mathbf{y}}|\mathbb{P}_{t}({\mathbf{u}},{\mathbf{x}}),{\textnormal{t}} \sim p_{{\mathbf{y}}|\mathbb{P}_{t},{\textnormal{t}}}({\bm{y}}|P,t)$ and $\mathbb{P}_{t}({\mathbf{u}},{\mathbf{x}})|{\mathbf{x}}\sim p_{\mathbb{P}_t|{\mathbf{x}}}(P|{\bm{x}})$.
\end{theorem}
Compared to \eqref{eq:id},
$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{sec:principled}, we first show that our generative model can identify an equivalent PS if ${\mathbf{u}}$ is observed (Sec.~\ref{sec:unconf}), and then discuss how our model is connected to and might learn relaxations of PtS when ${\mathbf{u}}$ is unobserved (Sec.~\ref{sec:conf}).
\section{Intact-VAE}
\label{sec:intact}
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.
\subsection{Model and architecture}
\label{sec:arch}
\begingroup
We build the generative model \eqref{model_indep} 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 \eqref{eq:cate_by_bts} is modeled by $p_{{\bm{f}}}({\bm{y}}|{\bm{z}},t)$, and the score distribution $p_{\mathbb{P}_{{\textnormal{t}}}|{\mathbf{x}}}$ in \eqref{eq:cate_by_bts} 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}}$.
\begin{equation}
\label{model_indep}
\begin{gathered}
p_{{\bm{\theta}}}({\bm{y}},{\bm{z}}|{\bm{x}},t) = p_{{\bm{f}}}({\bm{y}}|{\bm{z}},t)p_{\bm\lambda}({\bm{z}}|{\bm{x}},t), \\
p_{{\bm{f}}}({\bm{y}}|{\bm{z}},t)=p_{\bm\epsilon}({\bm{y}}-{\bm{f}}_{t}({\bm{z}})),
\quad
p_{\bm\lambda}({\bm{z}}|{\bm{x}},t) \sim \mathcal{N} ( {\bm{z}} ; {\bm{h}}_{t}({\bm{x}}),\operatorname{diag}({\bm{k}}_{t}({\bm{x}}))),
\end{gathered}
\end{equation}
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{vaes} for the basics of VAEs):
\begin{equation}
\label{elbo}
\begin{split}
\log p({\bm{y}}|{\bm{x}},t) \geq \log p({\bm{y}}|{\bm{x}},t) - D_{\mathrm{KL}}(q({\bm{z}}|{\bm{x}},{\bm{y}},t)\Vert p({\bm{z}}|{\bm{x}},{\bm{y}},t)) \\
= \mathbb{E}_{{\bm{z}} \sim q}\log p_{{\bm{f}}}({\bm{y}}|{\bm{z}},t) - D_{\mathrm{KL}}(q({\bm{z}}|{\bm{x}},{\bm{y}},t)\Vert p_{\bm\lambda}({\bm{z}}|{\bm{x}},t)) .
\end{split}
\end{equation}
with $p_{{\bm{f}}}$ as the decoder, $q$ as the encoder, and $p_{\bm\lambda}$ as the \textit{conditional} prior.
\setlength{\columnsep}{8pt}
\setlength{\intextsep}{5pt}
\begin{wrapfigure}{r}{0.3\textwidth}
\vspace{-.2in}
\begin{center}
\includegraphics[width=0.3\textwidth]{VAE.pdf}
\end{center}
\vspace{-.2in}
\caption{Decoders of CVAE, iVAE, and Intact-VAE.}
\vspace{-.1in}
\label{f:vae}
\end{wrapfigure}
We name this architecture \textit{Intact-VAE} (\textit{I}de\textit{n}tifiable \textit{t}re\textit{a}tment-\textit{c}ondi\textit{t}ional VAE).
Note that \eqref{model_indep} 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 \eqref{model_indep}. Thus, our VAE architecture can be seen as a combination of iVAE and conditional VAE (CVAE) \citep{sohn2015learning,kingma2014semi}, with ${\textnormal{t}}$ as the conditioning variable.
See Figure \ref{f:vae} 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{vaes} 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:
\begin{equation}
\label{more_param}
\textstyle
{\mathbf{y}}|{\mathbf{z}},{\textnormal{t}} \sim \prod_{j}^d \mathcal{N}(y_j;f_j, g_j),
\quad
{\mathbf{z}}|{\mathbf{x}},{\mathbf{y}},{\textnormal{t}} \sim \prod_{i=1}^n \mathcal{N}(z_i ; r_i , s_i).
\end{equation}
$\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.
\subsection{Identifiability of model parameters}
Our model identifiability lays the foundation for principled causal effect estimation, as we will discuss in Sec.~\ref{sec:principled}. The main theoretical result of iVAE is extended in our Theorem \ref{idmodel} which combines the techniques from \cite{khemakhem2020variational} and \cite{sorrenson2019disentanglement}. Essentially the same results can be proved for other exponential family priors.
\endgroup
\begin{theorem}[Model identifiability]
\label{idmodel}
Given $p_{{\bm{\theta}}}({\bm{y}},{\bm{z}}|{\bm{x}},t)$ specified by \eqref{model_indep}
, for ${\textnormal{t}}=t$, assume
\begin{enumerate}
\item ${\bm{f}}_t({\mathbf{z}})$
is injective and differentiable;
\item \label{ass:inv_jac} there exist $2n+1$ points ${\bm{x}}_0,...,{\bm{x}}_{2n}$ such that the $2n$-square matrix ${\bm{L}}_t \coloneqq [\bm\gamma_{t,1},...,\bm\gamma_{t,2n}]$ is invertible, where $\bm\gamma_{t,k}\coloneqq\bm\lambda_t({\bm{x}}_k)-\bm\lambda_t({\bm{x}}_0)$.
\end{enumerate}
Then, given ${\textnormal{t}}=t$, the family is \emph{identifiable} up to an equivalence class. That is, if $p_{{\bm{\theta}}}({\mathbf{y}}|{\mathbf{x}},{\textnormal{t}}=t)=p_{{\bm{\theta}}'}({\mathbf{y}}|{\mathbf{x}},{\textnormal{t}}=t)$\footnote{${\bm{\theta}}'=({\bm{f}}',{\bm{h}}',{\bm{k}}')$ is another parameter giving the same distribution. In this paper, symbol $'$ (prime) always indicates another parameter (variable, etc.) in the equivalence class.
}, we have the relation between parameters: for any ${\bm{y}}_t$ in the image of ${\bm{f}}_t$,
\begin{equation}
\label{eq:class}
{\bm{f}}_t^{-1}({\bm{y}}_t) = \operatorname{diag}({\bm{a}}){{\bm{f}}'_t}^{-1}({\bm{y}}_t) + {\bm{b}}
=:\mathcal{A}_t({{\bm{f}}'_t}^{-1}({\bm{y}}_t))
\end{equation}
where $\operatorname{diag}({\bm{a}})$ is an invertible $n$-diagonal matrix and ${\bm{b}}$ is a $n$-vector, both depend on $\bm\lambda_t$ and $\bm\lambda'_t$.
\end{theorem}
The essence of Theorem \ref{idmodel} 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{sec:principled}.
The assumptions in Theorem \ref{idmodel} are inherited from iVAE.
Injectivity in i) is related to TE identification (Theorem \ref{th:id_os}), 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 \cite{khemakhem2020variational}). Note that, to have \eqref{eq:class}, we only need the same observational distribution $p({\mathbf{y}}|{\mathbf{x}},{\textnormal{t}}=t)$, but this leaves room for different latent distributions.
\subsection{Causal effect estimation}
Once the VAE is learned by the ELBO \eqref{elbo},
the estimate of $\mu_t$ is given by
\begin{equation}
\label{eq:cate_est}
\hat{\mu}_{\hat{t}}({\bm{x}})=\mathbb{E}_{q({\bm{z}}|{\bm{x}})}{\bm{f}}_{\hat{t}}({\bm{z}})=\mathbb{E}_{\mathcal{D}|{\bm{x}} \sim p({\bm{y}},t|{\bm{x}})}\mathbb{E}_{{\bm{z}}\sim q_{\bm\phi}}{\bm{f}}_{\hat{t}}({\bm{z}}), {\hat{t}} \in \{0,1\},
\end{equation}
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 \textbf{algorithm} steps should be clear. After training Intact-VAE by \eqref{elbo}, 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 \citep{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 \citep{kingma2014semi}.
By taking ${\mathbf{y}}$, the posterior model (the encoder) is better than the prior.
On the other hand, sampling posterior requires \textit{post-treatment} observation ${\bm{y}}$. Often, it is desirable that we can also have \textit{pre-treatment} prediction for a new subject, with only the observation of its covariate ${\mathbf{x}}={\bm{x}}$. To this end, in \eqref{eq:cate_est}, 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{sec:proofs} for formal statements and proofs. It is possible to prove the consistency of the posterior estimation, as shown in \cite{bonhomme2021posterior,liao2011posterior}, and we leave this for future work, see Sec.~\ref{sec:conf}.
\section{Principled causal effect estimation via VAEs: retrospect and prospect}
\label{sec:principled}
In this section, we present a critical review of existing VAE-based methods in Sec.~\ref{sec:te_vae} and then discuss theoretical developments that can lead to principled causal estimation, under different settings.
\subsection{VAEs for TE estimation}
\label{sec:te_vae}
Most VAE methods for TEs, e.g., \cite{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.
\textbf{On identification.} First, as to TE identification, CEVAE assumes unobserved confounder can be recovered, which is rarely possible even under further structural assumptions \citep{tchetgen2020introduction}. Indeed, \cite{rissanen2021critical} recently give evidence that the method often fails. Other methods \citep{,zhang2020treatment,vowels2020targeted,lu2020reconsidering} assume unconfoundedness but still rely on proxy at least intuitively; for example, \cite{lu2020reconsidering} factorize the decoder as if in the proxy setting. However, \textit{unconfoundedness and proxy should not be put together}. The conceptual inconsistency is that, by definition, unconfoundedness means covariates \textit{fully} control confounding, while the motivation for proxy is that unconfoundedness is often \textit{not} satisfied in practice and covariates are at best proxies of confounding, which are non-confounders causally connected to confounders \citep{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{sec:exp_syn}.
\textbf{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 \cite{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 \eqref{eq:cate_by_bts}, while not compromised by ad hoc heuristics. Our ELBO is derived by standard variational lower bound \eqref{elbo}. Moreover, our pre-treatment prediction is given naturally by the prior, thanks to the correspondence between our model and \eqref{eq:cate_by_bts}.
We show in the following subsections how our model and its identifiability inspire theoretical developments in TE identification.
\subsection{Identification under unconfoundedness}
\label{sec:unconf}
As a first step, we take up theoretical analysis of Intact-VAE under assumption \textbf{(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 \eqref{eq:cate_by_bts} degenerates to $\mathbb{P}_{{\textnormal{t}}}({\mathbf{x}})$.
Theorem \ref{th:id_os} is an identification under shape restriction \citep{chetverikov2018econometrics}, because injectivity in assumption i) is monotonicity if ${\bm{j}}_t$ is on $\mathbb{R}$.
\begin{theorem}[Identification via PS]
\label{th:id_os}
Use model \eqref{model_indep} under unconfoundedness, further assume
\begin{enumerate}
\item (Injective separation) $\mu_{t}({\mathbf{x}})={\bm{j}}_t(\mathbb{O}({\mathbf{x}}))$
for some function $\mathbb{O}$ and injective function ${\bm{j}}_t$;
\item (Score matching) in the model, $n=d$, ${\bm{f}}_t$ is injective, ${\bm{h}}_0({\mathbf{x}})={\bm{h}}_1({\mathbf{x}})$, and ${\bm{k}}({\mathbf{x}}) = \bm0$.
\end{enumerate}
Then, if $\mathbb{E}_{p_{\bm\theta}}({\mathbf{y}}|{\mathbf{x}},{\textnormal{t}})=\mathbb{E}({\mathbf{y}}|{\mathbf{x}},{\textnormal{t}})$, we have
\begin{enumerate}
\item (Recovery of score) ${\mathbf{z}}_{\bm\lambda,t}={\bm{h}}_t({\mathbf{x}})={\bm{v}}(\mathbb{O}({\mathbf{x}}))$ where ${\bm{v}}$ is an injective function;
\item (CATE Identification)
$\mu_{{t}}({\mathbf{x}}) = \hat{\mu}_{{t}}({\mathbf{x}}) \coloneqq \mathbb{E}_{p_{\bm\lambda}({\mathbf{z}}| {\mathbf{x}},t)} \mathbb{E}_{p_{{\bm{f}}}}({\mathbf{y}}|{\mathbf{z}},{t}) = {\bm{f}}_t({\bm{h}}_t({\mathbf{x}}))$.
\end{enumerate}
\end{theorem}
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 \textit{same} ${\bm{v}}$ for both $t$.
The connection to PS\footnote{Note that, some call $\mu_0$ the prognostic score (e.g, \cite{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{sec:p_inj} for real-world examples satisfying i) and the connection to ``independent causal mechanisms'' \citep{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{th:id_os} does not depend on Theorem \ref{idmodel}, 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 \eqref{eq:class} in \cite{wu2021beta}, the main ideas of which is briefly summarized below.
\subsection{Identification without overlap of ${\mathbf{x}}$}
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}}$ \citep{d2020overlap}. In a parallel submission \citep{wu2021beta}, we assume the overlap of a PtS instead and extend Theorem \ref{th:id_os} to this \textit{limited overlap} setting.
This is a natural next step, because we already model a PtS as the latent variable.
A main result of \cite{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{th:id_os}. To recover the PtS, we derive a condition and strengthen \eqref{eq:class} 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 \textit{uncertainty of scores}--we are uncertain how likely a PS (not only a PtS) can be recovered \citep[Sec.~4.1 \& C.5]{wu2021beta}.
\subsection{Preliminary thoughts under unobserved confounding}
\label{sec:conf}
The positive experimental results motivate us to consider the theory under unobserved confounding.
Moreover, the prior in \eqref{model_indep} 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.
\textbf{Identification.\space} Auxiliary structures (e.g., IVs) can give TE identifications via \textit{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}})$ \citep{matzkin2007nonparametric,wooldridge2015control}. Control functions can be stochastic, as in \cite{puli2020general}. Consistent TE estimation can be given by a regression of outcome on the treatment and a control function. Our model \eqref{model_indep} 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 \citep{nagasawa2018identification}, or in the presence of both proxies and IVs \citep{tien2021instrumental}.
\textbf{Estimation.\space} In causal inference, many models, including nonparametric IV regression (NPIV), are stated as \textit{conditional moment restrictions} (CMRs) \citep{newey199316}. Optimizing the ELBO of our VAE, given by \eqref{elbo}, 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 \textit{quasi-Bayesian} analysis using KL divergence \citep{zhang2006eps,jiang2008gibbs,kim2002limited}.
For example, \cite{kato2013quasi} uses a quasi-likelihood from the CMR of NPIV to set the prior, and the Gibbs posterior \citep{zhang2006eps,jiang2008gibbs} is a minimizer of an information complexity which has a variational characterization similar to an ELBO. For general CMR models,
\cite{liao2011posterior} extend \cite{kim2002limited} and give the best approximation to the true likelihood function under the CMR by minimizing a KL divergence. Very recently, \cite{wang2021scalable} employ quasi-Bayesian analysis to kernel-based IV methods, but only consider unconditional moments.
\section{Experiments}
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 \citep{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 \citep{hill2011bayesian} $\epsilon_{pehe}\coloneqq \mathbb{E}_{\mathcal{D}}((y(1)-y(0))-\hat{\tau}({\bm{x}}))^2$ for individual-level TEs.
\subsection{Synthetic dataset}
\label{sec:exp_syn}
\begin{figure*}[h]
\centering
\includegraphics[width=.9\textwidth]{art.pdf}
\vspace{-.1in}
\caption{\footnotesize{Graphical models for generating synthetic datasets. From left: IV ${\mathbf{x}}$, ignorability given ${\mathbf{x}}$, and proxy ${\mathbf{x}}$. Note that in the latter two cases, reversing the arrow between ${\mathbf{x}}, {\mathbf{z}}$ does not change any independence relationships, and causal interpretations of the graphs remain the same.}}
\label{fig:art}
\end{figure*}
\vspace{-.2in}
\begin{equation}
\label{art_model}
\begin{split}
\textstyle
{\mathbf{x}} \sim \mathcal{N}(\bm\mu, \bm\sigma);
{\textnormal{z}}|{\mathbf{x}} \sim \mathcal{N}(h({\mathbf{x}}), \beta k({\mathbf{x}}));
{\textnormal{t}}|{\mathbf{x}},{\textnormal{z}} \sim \operatorname{Bern}(\operatorname{Logi}(l({\mathbf{x}},{\textnormal{z}})));
{\mathbf{y}}|{\textnormal{z}},{\textnormal{t}} \sim \mathcal{N}(f({\textnormal{z}},{\textnormal{t}}), \alpha).
\end{split}
\end{equation}
\vspace{-.2in}
\begingroup
We generate synthetic datasets by \eqref{art_model}. 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{sec:exp_syn_app} for more details.
We experiment on three different causal structures as shown in graphical models of Figure \ref{fig:art}, by variation on \eqref{art_model}.
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 \textit{unconfounded}; and $l\coloneqq l({\textnormal{z}})$, then we have unobserved confounder ${\textnormal{z}}$ and \textit{proxy} ${\mathbf{x}}$ of ${\textnormal{z}}$. To introduce ${\mathbf{x}}$ as \textit{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 \eqref{art_model}.
\setlength{\columnsep}{8pt}
\setlength{\intextsep}{5pt}
\begin{wrapfigure}{r}{0.36\textwidth}
\vspace{-.2in}
\begin{center}
\includegraphics[width=0.36\textwidth]{fig_nonl.pdf}
\end{center}
\vspace{-.2in}
\caption{Pre-treatment \footnotesize{$\sqrt{\epsilon_{pehe}}$ on nonlinear synthetic dataset. Error bar on 100 random DGPs. We adjust one of the noise levels $\alpha,\beta$ in each panel, with another fixed to 0.2. }}
\vspace{-.2in}
\label{nonl_art}
\end{wrapfigure}
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 \eqref{art_model}. 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{nonl_art} 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{sec:exp_syn_app} for results on linear outcome. Results for ATE and post-treatment are similar.
\setlength{\columnsep}{8pt}
\setlength{\intextsep}{5pt}
\begin{wrapfigure}{r}{0.4\textwidth}
\vspace{-.2in}
\begin{center}
\includegraphics[width=0.4\textwidth]{latent.png}
\end{center}
\vspace{-.2in}
\caption{\footnotesize{Plots of recovered
- true
latent on the nonlinear outcome. Blue: $t=0$, Orange: $t=1$. $\alpha,\beta= 0.4$. ``no.'' indicates index among the 100 random models.}}
\label{recover}
\end{wrapfigure}
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{recover}, 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{sec:plots} (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{fig:bad_recover}. 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{sec:plots} shows that the noise level affects how well we recover the latent variable.
\endgroup
\subsection{IHDP benchmark dataset}
This experiment shows our balanced estimator matches the state-of-the-art methods specialized for standard unconfoundedness. The IHDP dataset \citep{hill2011bayesian} is widely used to evaluate machine learning based causal inference methods, e.g. \cite{shalit2017estimating, shi2019adapting}.
The generating process is as following \citep[Sec.~4.1]{hill2011bayesian}.
\begin{equation}
{\textnormal{y}}(0) \sim \mathcal{N}(e^{{\bm{a}}^T({\mathbf{x}}+{\bm{b}})},1),\quad {\textnormal{y}}(1) \sim \mathcal{N}({\bm{a}}^T{\mathbf{x}}-o,1),
\end{equation}
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{sec:ihdp_app} for more details.
To see our balancing property clearly, we add two components specialized for balancing from \cite{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 \eqref{model_indep}. Second, we add to our ELBO \eqref{elbo} a regularization term, which is the Wasserstein distance \citep{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{t:ihdp}, Intact-VAE outperforms or matches the state-of-the-art methods. In particular, our method has the \textit{best} ATE estimation; and it has the \textit{best} individual-level estimation, adding the two components from \cite{shalit2017estimating}.
We can see in the caption of Table \ref{t:ihdp}, the specialized additions do \textit{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 \textit{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 \cite{louizos2017causal}, where they used 5-dimensional Gaussian latent variable to model a binary ground truth.
\begin{table}[h]
\centering
\scriptsize
\caption{Errors on IHDP. The mean and std are calculated over 1000 random data generating models. *Results with modifications are $\epsilon_{ate}=.31_{\pm .01}/.30_{\pm .01}$ and $\sqrt{\epsilon_{pehe}}=\textbf{.77}_{\pm .02}/\textbf{.69}_{\pm .02}$. \textbf{Bold} indicates method(s) that are \textit{significantly} better than all the others. The results of the others are taken from \cite{shalit2017estimating}, except GANITE \citep{yoon2018ganite} and CEVAE \citep{louizos2017causal}.}
\vspace{-.05in}
\begin{tabular}{p{.8cm}p{1.cm}p{1.45cm}p{1.45cm}p{1.45cm}p{1.45cm}p{1.45cm}p{1.6cm}}
\toprule
Method &\textbf{TMLE} &\textbf{BNN} &\textbf{CFR} &\textbf{CF} &\textbf{CEVAE} &\textbf{GANITE} &\textbf{Ours*}\\
\midrule
$\epsilon_{ate}$ &NA/.30$_{\pm .01}$ &.42$_{\pm .03}$/.37$_{\pm .03}$
&.27$_{\pm.01}$/.25$_{\pm.01}$ &.40$_{\pm.03}$/\textbf{.18}$_{\pm.01}$ &.46$_{\pm.02}$/.34$_{\pm.01}$ &.49$_{\pm.05}$/.43$_{\pm.05}$ &$.31_{\pm .01}$/$.30_{\pm .01}$
\\
\midrule
$\sqrt{\epsilon_{pehe}}$ &NA/5.0$_{\pm .2}$ &2.1$_{\pm .1}$/2.2$_{\pm .1}$
&\textbf{.76}$_{\pm.02}$/\textbf{.71}$_{\pm.02}$ &3.8$_{\pm.2}$/3.8$_{\pm.2}$ &2.6$_{\pm.1}$/2.7$_{\pm.1 }$ &2.4$_{\pm .4}$/1.9$_{\pm .4}$ &\textbf{.77}$_{\pm .02}$/\textbf{.69}$_{\pm .02}$
\\
\bottomrule
\end{tabular}
\label{t:ihdp}
\end{table}
\subsection{Pokec social network dataset}
We show our method is the best compared with the methods specialized for networked deconfounding, a challenging problem in its own right. Pokec \citep{leskovec2014snap} is a real world social network dataset. We experiment on a semi-synthetic dataset based on Pokec, introduced in \cite{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 \texttt{district}, \texttt{age}, or \texttt{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:
\begin{equation}
\label{pokec}
{\textnormal{t}} \sim \operatorname{Bern}(g({\textnormal{z}})), \quad {\mathbf{y}} = {\textnormal{t}} + 10(g({\textnormal{z}})-0.5) + \epsilon, \text{ where $\epsilon$ is standard normal. }
\end{equation}
Note that \texttt{district} is of 3 categories; \texttt{age} and \texttt{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 \cite{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{tbl:Pokec} shows the results. In addition, the pre-treatment $\sqrt{\epsilon_{pehe}}$ for \texttt{Age}, \texttt{District}, and \texttt{Join date} confounders are 1.085, 0.686, and 0.699 respectively, practically the same as the ATE errors. \cite{veitch2019using} do not give individual-level prediction.
\begin{table}[h]
\centering
\scriptsize
\caption{Pre-treatment ATE on Pokec. Ground truth ATE is 1, as we can see in \eqref{pokec}. ``Unadjusted'' estimates ATE by $\mathbb{E}_{\mathcal{D}}(y_1)-\mathbb{E}_{\mathcal{D}}(y_0)$. ``Parametric'' is a stochastic block model for networked data \citep{gopalan2013efficient}. ``Embed-'' denotes the best alternatives given by \citep{veitch2019using}. \textbf{Bold} indicates method(s) that are \textit{significantly} better than all the others. 20-dimensional latent variable in Intact-VAE works better, and its result is reported. The results of the other methods are taken from \citep{veitch2019using}.} \label{tbl:Pokec}
\vspace{-.05in}
\begin{tabular}{cccccc}
\toprule
{} & {Unadjusted} & {Parametric} & {Embed-Reg.} & {Embed-IPW} & {Ours} \\
\midrule
\texttt{Age} & 4.34 $\pm$ 0.05 & 4.06 $\pm$ 0.01 & 2.77 $\pm$ 0.35 & 3.12 $\pm$ 0.06 & \textbf{2.08} $\pm$ 0.32 \\
\texttt{District} & 4.51 $\pm$ 0.05 & 3.22 $\pm$ 0.01 & \textbf{1.75} $\pm$ 0.20 & \textbf{1.66} $\pm$ 0.07 & \textbf{1.68} $\pm$ 0.10 \\
\texttt{Join Date} & 4.03 $\pm$ 0.06 & 3.73 $\pm$ 0.01 & 2.41 $\pm$ 0.45 & 3.10 $\pm$ 0.07 & \textbf{1.70} $\pm$ 0.13 \\
\bottomrule
\end{tabular}
\end{table}
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 \cite{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) \citep{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 \citep{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 \textit{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 \textit{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.
\section{Conclusion}
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{sec:principled}, 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, \cite{khemakhem2020ice} provide identifiability to deep energy models, and \cite{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.
\bibliography{example_paper}
\clearpage