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.
54,002 characters
$$-Intact-VAE: Identifying and Estimating Causal Effects under Limited Overlap
\maketitle
\begin{abstract}
As an important problem in causal inference, we discuss the identification and estimation of treatment effects (TEs) under limited overlap; that is, when subjects with certain features belong to a single treatment group. We use a latent variable to model a prognostic score which is widely used in biostatistics and sufficient for TEs; i.e., we build a generative prognostic model. We prove that the latent variable recovers a prognostic score, and the model identifies individualized treatment effects. The model is then learned as $\beta$-Intact-VAE––a new type of variational autoencoder (VAE). We derive the TE error bounds that enable representations balanced for treatment groups conditioned on individualized features. The proposed method is compared with recent methods using (semi-)synthetic datasets.
\end{abstract}
\section{Introduction}
Causal inference \citep{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 \textit{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 \textit{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 \textit{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
\textit{covariate balance}, ${X}$ independent of ${T}$,
by matching and re-weighting \citep{Stuart2010matching,rosenbaum2020modern}. Machine learning methods have also been exploited; there are semi-parametric methods––e.g., \citet[TMLE]{van2018targeted}––which improve finite sample performance, as well as non-parametric methods––e.g., \citet[CF]{wager2018estimation}. Notably, from \citet{johansson2016learning}, there has been a recent increase in interest in \textit{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 \textit{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 \citep{armstrong2021finite}, trimming non-overlapping sample points \citep{10.1093/biomet/asy008}, or studying convergence rates based on overlap \citep{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 \citep{d2020overlap}.
To address imbalance and limited overlap, we use a prognostic score \citep{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 \textit{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 \citep{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, \textit{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 \citep[particularly Sec.~8]{lewbel2019identification}. Identification is even more important for causal inference; because, unlike usual (non-causal) model misspecification, causal assumptions are often unverifiable through observables \citep{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 \citep[iVAE]{khemakhem2020variational} and conditional VAE \citep[CVAE]{sohn2015learning}. Based on model identifiability, we develop two identification results for individualized TEs under limited overlap.
A similar VAE architecture was proposed in \citet{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{sec:model&id}) and estimation (Sec.~\ref{sec:estimation}) 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{sec:proofs}. Our main contributions are:
\vspace{-5pt}
\begin{enumerate}[topsep=0pt, partopsep=0pt, itemsep=0pt, parsep=0pt, leftmargin=13pt]
\item[1)] TE identification under limited overlap of ${X}$, via prognostic scores and an identifiable model;
\item[2)] bounds on individualized TE error, which justify our conditional BRL;
\item[3)]
a new regularized VAE, $\beta$-Intact-VAE,
realizing the identification and conditional balance;
\item[4)] experimental comparison to the state-of-the-art methods on (semi-)synthetic datasets.
\end{enumerate}
\subsection{Related work}
\textbf{Limited overlap.}~Under limited overlap, \citet{luo2017estimating} estimate the average TE (ATE) by reducing covariates to a linear prognostic score. \citet{farrell2015robust} estimates a constant TE under a partial linear outcome model. \citet{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 \citep{oberst2020characterization,dai2020quantifying}, or indicating possible failure under limited overlap \citep{jesson2020identifying}, but not remedies. An exception is \citet{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.
\textbf{Prognostic scores\space} have been recently combined with machine learning approaches, mainly in the biostatistics community.
For example,
\citet{huang2017joint} estimate individualized TE by reducing covariates to a linear score which is a joint propensity-prognostic score. \citet{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; \citet{zhang2020treatment} and \citet{hassanpour2019learning} learn outcome predictors, without mentioning prognostic score––while \citet{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.
\textbf{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 \citep{khemakhem2020variational}, and energy models \citep{khemakhem2020ice}.
The results are exploited in causal discovery \citep{pmlr-v108-wu20b} and out-of-distribution generalization \citep{sun2020latent}. This study is the first to explore identifiable representations in TE identification.
\textbf{BRL and related methods\space} amount to a major direction. Early BRL methods include BLR/BNN \citep{johansson2016learning} and TARnet/CFR \citep{shalit2017estimating}. In addition, \cite{yao2018representation} exploit the local similarity between data points. \cite{shi2019adapting} use similar architecture to TARnet, considering the importance of treatment probability. There are also methods that use GAN \citep[GANITE]{yoon2018ganite} and Gaussian processes \citep{alaa2017bayesian}. Our method shares the idea of BRL, and further extends to conditional balance––which is natural for individualized TE.
\textbf{Others.}~Our work can lay conceptual and theoretical foundations of VAE methods for TEs \citep[e.g.,][]{louizos2017causal,lu2020reconsidering}. In addition, some studies consider monotonicity, which is injectivity on $\mathbb{R}$, together with overlap \citep{johansson2020generalization,zhang2020learning}. See Sec.~\ref{sec:other_related}.
\begingroup
\section{Setup and preliminaries}
\label{sec:setup}
\subsection{Counterfactuals, treatment effects, and identification}
Following \citet{imbens2015causal}, we assume there exist \textit{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 \textit{factual outcome} ${Y}$ under \textit{factual assignment} ${T}=t$. Formally, ${Y}(t)$ is defined by the \textit{consistency of counterfactuals}: ${Y} = {Y}(t)$ if ${T}=t$; or simply ${Y}={Y}({T})$. The \textit{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 \textit{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:
\begin{equation}
\label{eq:ce}
\tau({\bm{x}}) = \mu_1({\bm{x}}) - \mu_0({\bm{x}}),\quad \nu = \mathbb{E}(\tau({X})).
\end{equation}
CATE is seen as an \textit{individual-level}, personalized, treatment effect, given
highly discriminative ${X}$.
Standard results \citep{rubin2005causal}\citep[Ch.~3]{hernanCausalInferenceWhat2020} show sufficient conditions for TE identification in general settings. They are \emph{ Exchangeability}: ${Y}(t) \raisebox{0.05em}{\rotatebox[origin=c]{90}{$\models$}} {T}| {X}$, and \emph{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, \textit{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 \citep[pp.~15]{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{sec:identification}, will relax the condition of the overlapping covariate to allow some non-overlapping values ${\bm{x}}$––that is, covariate ${X}$ is \textit{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:
\begin{definition}
${V}$ is \textit{Overlapping} if $p(t|V={\bm{v}}) > 0$ for any $t\in\{0,1\},{\bm{v}} \in \mathcal{V}$. If the condition is violated at some value ${\bm{v}}$, then ${\bm{v}}$ is \textit{non-overlapping} and ${V}$ is \textit{limited-overlapping}.
\end{definition}
\subsection{Prognostic scores}
\label{sec:pgs}
Our method aims to recover a prognostic score \citep{hansen2008prognostic}, adapted as a Pt-score in Definition \ref{def:scores}. On the other hand, balancing scores \citep{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{sec:scores} for detail.
\begin{definition}
\label{def:scores}
A \textit{Pt-score} (PtS) is two functions $\mathbb{P}_t({X})$ ($t=0,1$) such that ${Y}(t)\raisebox{0.05em}{\rotatebox[origin=c]{90}{$\models$}}{X}|\mathbb{P}_t({X})$.
A PtS is called a \textit{P-score} (PS) if $\mathbb{P}_0=\mathbb{P}_1$.
\end{definition}
Note that a PtS is by definition two functions; thus, overlapping $\mathbb{P}_t({X})$ means \textit{both} $\mathbb{P}_0({X})$ and $\mathbb{P}_1({X})$ are overlapping.
\textbf{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}$ \citep{d2020overlap}, which counters our purpose.
In contrast, overlapping PS does not imply overlapping ${\bm{b}}({X})$ (see Sec.~\ref{sec:pgs_vs_bs} 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})$ \citep{d2021deconfounding}. In general, \cite{hajage2017estimation} show that prognostic score methods perform better––or as well as––propensity score methods.
Below is a corollary of Proposition 5 in \cite{hansen2008prognostic}; note that $\mathbb{P}_t({X})$ satisfies exchangeability.
\begin{proposition}[Identification via PtS]
\label{cate_by_bts}
If $\mathbb{P}_t$ is a PtS and ${Y}|\mathbb{P}_{\hat t}({X}),T \sim p_{{Y}|\mathbb{P}_{\hat t},T}({\bm{y}}|P,t)$ where $\hat{t}\in \{0,1\}$ is a counterfactual assignment, then CATE and ATE are identified, using \eqref{eq:ce} and
\begin{equation}
\label{eq:cate_by_bts}
\begin{split}
\mu_{\hat{t}}({\bm{x}}) &= \mathbb{E}({Y}({\hat{t}})|\mathbb{P}_{\hat{t}}({X}),{X}={\bm{x}}) = \mathbb{E}({Y}|\mathbb{P}_{\hat{t}}({\bm{x}}),{T}=\hat{t})=\textstyle \int p_{{Y}|\mathbb{P}_{\hat t},T}({\bm{y}}|\mathbb{P}_{\hat{t}}({\bm{x}}),\hat{t}){\bm{y}} d{\bm{y}}
\end{split}
\end{equation}
\end{proposition}
\vspace{-.1in}
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.
\begin{enumerate}
\item\hspace{-2mm}\footnote{The symbols \textbf{G}, \textbf{M}, and \textbf{D} in the labels of conditions stand for Generating process, Model, and Data.} \label{ass:anm} (Additive noise model) the data generating process (DGP) for ${Y}$ is
${Y} = {\bm{f}}^*(\mathbb{M}({X}),{{T}}) + {{\mathbf{e}}}$
where ${\bm{f}}^*,\mathbb{M}$ are functions, and ${{\mathbf{e}}}$ denotes a zero-mean exogenous (external) noise.
\end{enumerate}
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 \citep{hansen2008prognostic}––which is often components of ${X}$ affecting ${Y}$ directly.
Additive noise models are used in nonparametric regression methods for TEs \citep{caron2020estimating}.
Under \ref{ass:anm}, 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 \textit{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.
\section{Identification under generative prognostic model}
\label{sec:model&id}
\begingroup
\setlength{\columnsep}{8pt}
\setlength{\intextsep}{5pt}
\begin{wrapfigure}{r}{0.26\columnwidth}
\vspace{-.2in}
\begin{center}
\includegraphics[width=0.26\columnwidth]{VAE.pdf}
\end{center}
\vspace{-.15in}
\caption{CVAE, iVAE, and Intact-VAE: Graphical models of the decoders.
}
\vspace{-.5in}
\label{f:vae}
\end{wrapfigure}
In Sec.~\ref{sec:arch}, we specify the generative prognostic model $p({\bm{y}},{\bm{z}}|{\bm{x}},t)$, and show its identifiability. In Sec.~\ref{sec:identification}, 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{sec:estimation}.
\subsection{Model, architecture, and identifiability}
\label{sec:arch}
We present the necessary definitions and results, and defer more explanations to Sec.~\ref{sec:arch_details}.
The contents of this subsection are essentially taken from \cite{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:
\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 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 \eqref{eq:cate_by_bts}; 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:
\begin{equation}
\label{elbo}
\log p({\bm{y}}|{\bm{x}},t)
\geq \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{equation}
Note that the encoder $q$ conditions on all the observables $({X},{Y},{T})$; this fact plays an important role in Section \ref{sec:regular}. This architecture is called \textit{Intact-VAE} (\textit{I}de\textit{n}tifiable \textit{t}re\textit{a}tment-\textit{c}ondi\textit{t}ional VAE).
See Figure \ref{f:vae} for comparison in terms of graphical models. See Sec.~\ref{vaes} for basics of VAEs.
\endgroup
Our model identifiability extends the theory of iVAE,
and the following conditions are inherited.
\begin{enumerate}
\item \label{ass:model} i) ${\bm{f}}_t$ is injective, and ii) ${\bm{f}}_t$ is differentiable.
\end{enumerate}
\begin{enumerate}
\item \label{ass:inv_jac_smp} $\bm\lambda_t({X})$ is non-degenerate, i.e., the linear hull of its support is $2n$-dimensional.
\end{enumerate}
Under \textbf{(M1)} and \textbf{(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$:
\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 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{sec:proofs} for the proof and a relaxation of \ref{ass:inv_jac_smp}.
In this paper, symbol $'$ (prime) always indicates another parameter (variable, etc.): ${\bm{\theta}}'=({\bm{f}}',\bm\lambda')$.
\subsection{Identifications under limited-overlapping covariate}
\label{sec:identification}
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{ass:anm}. In practice, $n=d$ means that we seek a low-dimensional representation of ${X}$.
\ref{ass:pt_inj} makes the dimensionality explicit and reduces to \ref{ass:anm} if the only possibility is that $\mathbb{P}_t=\mu_t$ and ${\bm{j}}_t$ is identity.
\begin{enumerate}
\item \label{ass:pt_inj} (Low-dimensional PtS) Under \ref{ass:anm}, $\mu_{t}({X})={\bm{j}}_{t}(\mathbb{P}_t({X}))$
for some $\mathbb{P}_t$ and injective ${\bm{j}}_t$.
\end{enumerate}
We use \ref{ass:pt_inj} instead of \ref{ass:anm} hereafter. Clearly, $\mathbb{P}_t$ in \ref{ass:pt_inj} is PtS. In addition, injectivity and $n=d$ ensure that $n=\dim({Y}) \geq \dim(\mathbb{P}_t)$. Similarly, \ref{ass:p_inj} ensures that $n \geq \dim(\mathbb{P})$ for PS.
\begin{enumerate}
\item \label{ass:p_inj} (Low-dimensional PS) Under \ref{ass:anm}, $\mu_{t}({X})={\bm{j}}_{t}(\mathbb{P}({X}))$
for some $\mathbb{P}$ and injective ${\bm{j}}_t$.
\end{enumerate}
\ref{ass:p_inj} 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{sec:p_inj} for more discussions and real-world examples.
With \ref{ass:pt_inj} or \ref{ass:p_inj}, overlapping ${X}$ can be relaxed to overlapping PS or PtS plus the following:
\begin{enumerate}
\item \label{ass:prepart} \text{(Score partition preserving) For any ${\bm{x}},{\bm{x}}' \in \mathcal{X}$, if $\mathbb{P}_t({\bm{x}})=\mathbb{P}_t({\bm{x}}')$, then ${\bm{h}}_t({\bm{x}})={\bm{h}}_t({\bm{x}}')$.}
\end{enumerate}
Note that \ref{ass:prepart} 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{th:id_nop_pscore}, or $p_{{\bm{\theta}}}({\bm{y}}|{\bm{x}},t)=p({\bm{y}}|{\bm{x}},t)$ in Theorem \ref{th:id_nop_ptscore}. 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{ass:prepart} and are often assumed, e.g., in \cite{huang2017joint,luo2017estimating,d2021deconfounding}. Linear outcome models \citep{farrell2015robust,schuler2020increasing} are also common.
Our first identification, Proposition \ref{th:id_nop_pscore}, relies on \ref{ass:p_inj} and our generative model, \textit{without} model identifiability (so differentiable ${\bm{f}}_t$ is not needed).
\begin{proposition}[Identification via recovery of PS]
\label{th:id_nop_pscore}
Suppose we have DGP \ref{ass:p_inj} and model \eqref{model_indep} with $n=d$. Assume \ref{ass:model}-i) and\begin{inparaenum}\item \label{ass:match_pscore} \emph{(PS matching) ${\bm{h}}_0({X})={\bm{h}}_1({X})$ and ${\bm{k}}({X}) = \bm0$.}
\end{inparaenum}
Then, if $\mathbb{E}_{p_{\bm\theta}}({Y}|{X},{T})=\mathbb{E}({Y}|{X},{T})$, we have
\begin{enumerate}
\item (Recovery of PS) ${\bm{z}}_{\bm\lambda,t}={\bm{h}}_t({\bm{x}})={\bm{v}}(\mathbb{P}({\bm{x}}))$
on overlapping ${\bm{x}}$, \\ where ${\bm{v}}: \mathcal{P} \to \mathbb{R}^n$ is an injective function, and $\mathcal{P}\coloneqq\{\mathbb{P}({\bm{x}})|\text{overlapping }{\bm{x}}\}$.
\item (CATE Identification) if $\mathbb{P}$ in \ref{ass:p_inj} is overlapping, and \ref{ass:prepart} is satisfied, then
$\mu_{{t}}({\bm{x}}) = \hat{\mu}_{{t}}({\bm{x}})$ for any ${t} \in \{0,1\}$ and ${\bm{x}} \in \mathcal{X}$, where $\hat{\mu}_{{t}}({\bm{x}}) \coloneqq \mathbb{E}_{p_{\bm\lambda}({Z}| {\bm{x}},t)} \mathbb{E}_{p_{{\bm{f}}}}({Y}|{Z},{t}) = {\bm{f}}_t({\bm{h}}_t({\bm{x}}))$.
\end{enumerate}
\end{proposition}
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 \textit{same} ${\bm{v}}$ for both $t$. Theorem \ref{th:id_nop_ptscore} 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{ass:prepart} 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{th:id_nop_ptscore} based on the model identifiability to work under PtS which
generally exists.
\begin{theorem}[Identification via recovery of PtS]
\label{th:id_nop_ptscore}
Suppose we have DGP \ref{ass:pt_inj} and model \eqref{model_indep} with $n=d$.
For the model, assume \ref{ass:model}
and\begin{inparaenum}
\item \label{ass:match_noise} \emph{(Noise matching) $p_{{{\mathbf{e}}}}=p_{\bm\epsilon}$ and ${\bm{k}}({X})=k{\bm{k}}'({X}),k \to 0$}.
\end{inparaenum}
Assume further that \ref{ass:inv_jac_smp} and\begin{inparaenum}
\item \label{ass:g_bcovar_smp} \emph{(Balance from data) $\mathcal{A}_0=\mathcal{A}_1$ in \eqref{eq:class}}.
\end{inparaenum}
Then, if $p_{{\bm{\theta}}}({\bm{y}}|{\bm{x}},t)=p({\bm{y}}|{\bm{x}},t)$; conclusions 1) and 2) in Proposition \ref{th:id_nop_pscore} hold with $\mathbb{P}$ replaced with $\mathbb{P}_t$ in \ref{ass:pt_inj}; and the domain of ${\bm{v}}$ becomes
$\mathcal{P}\coloneqq\{\mathbb{P}_t({\bm{x}})|p(t,{\bm{x}}) > 0\}$.
\end{theorem}
Theorem \ref{th:id_nop_ptscore} 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{th:id_nop_pscore} and Theorem \ref{th:id_nop_ptscore} 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{th:id_nop_pscore}. We show in Sec.~\ref{sec:proofs} a sufficient and necessary condition \ref{ass:g_bcovar} 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 \eqref{eq:class}.
See Sec.~\ref{sec:complementarity} for more on the complementarity between the two identifications.
\section{Estimation by $\beta$-Intact-VAE}
\label{sec:estimation}
\subsection{Prior as PS, posterior as PtS, and $\beta$ as regularization strength}
\label{sec:regular}
In Sec.~\ref{sec:identification}, we see that the existence of PS (Proposition \ref{th:id_nop_pscore}) is preferable in identifying the true DGP up to an equivalent expression––while Theorem \ref{th:id_nop_ptscore} 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{ass:pt_inj} 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 \textit{discover} equivalent PS, if possible. This idea is common in practice. For example, in a real-world nutrition study \citep{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 \eqref{model_indep}. The decoder and encoder are factorized Gaussians:
\begin{equation}
\label{eq:enc}
p_{{\bm{f}},{\bm{g}}}({\bm{y}}|{\bm{z}},t) =\mathcal{N}({\bm{y}} ; {\bm{f}}_{t}({\bm{z}}) , \operatorname{diag}({\bm{g}}_{t}({\bm{z}}))),
\mkern9mu
q_{\bm\phi}({\bm{z}}|{\bm{x}},{\bm{y}},t)= \mathcal{N} ( {\bm{z}} ; {\bm{r}}_{t}({\bm{x}},{\bm{y}}) , \operatorname{diag}({\bm{s}}_{t}({\bm{x}},{\bm{y}}))),
\end{equation}
where $\bm\phi =({\bm{r}},{\bm{s}})$.
The other is to introduce a hyperparameter $\beta$ in the ELBO
as in $\beta$-VAE \citep{higgins2016beta}. The modified ELBO with $\beta$, up to the additive constant, is derived as:
\begin{equation}
\label{eq:elbo_imp}
\mathbb{E}_{\mathcal{D}} \{- \betaD_{\mathrm{KL}}(q_{\bm\phi}\Vert p_{\bm\Lambda}) -\mathbb{E}_{{\bm{z}}\sim q_{\bm\phi}}[({\bm{y}}-{\bm{f}}_{t}({\bm{z}}))^2/2{\bm{g}}^2_{t}({\bm{z}})] - \mathbb{E}_{{\bm{z}}\sim q_{\bm\phi}}\log|{\bm{g}}_{t}({\bm{z}})| \}.
\end{equation}
For convenience, here and in $\mathcal{L}_{{\bm{f}}}$ in Sec.~\ref{sec:balance}, 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{th:id_nop_pscore} and PS as much as possible, but step into Theorem \ref{th:id_nop_ptscore} 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{sec:identification}, but it is more practical and yields better results in our experiments.
See Sec.~\ref{sec:elbo_disc} 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. \textit{Consistent estimation} using the prior is a direct corollary of the consistent VAE. see Sec.~\ref{sec:consist_vae} for formal statements and proofs. Under Gaussian models, it is possible to prove the consistency of the posterior estimation, as shown in \citet{bonhomme2019posterior}.} by the ELBO, the estimate of the expected potential outcomes 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. 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{sec:prepost} for details and \ref{sec:add_exp} for results). Note that ${\hat{t}}$ in \eqref{eq:cate_est} 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{sec:balance}).
The overall \textbf{algorithm} steps are i) train the VAE using \eqref{eq:elbo_imp}, and ii) infer CATE $\hat{\tau}({\bm{x}})=\hat{\mu}_1({\bm{x}})-\hat{\mu}_0({\bm{x}})$ by \eqref{eq:cate_est}.
\subsection{Conditionally balanced representation learning}
\label{sec:balance}
We formally justify our ELBO \eqref{eq:elbo_imp} 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{ass:pt_inj} is controlled by the ELBO. Previous works \citep{shalit2017estimating,lu2020reconsidering} instead focus on unconditional balance and bound PEHE which is marginalized on ${X}$. Sec.~\ref{sec:ihdp} 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{sec:novel_bound} for detail.
We introduce the objective that we bound. Using \eqref{eq:cate_est} 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 \textit{true CATE}, given the covariate ${\bm{x}}$ or score ${\bm{z}}$, is:
\begin{equation}
\label{eq:taus}
\tau({\bm{x}})={\bm{j}}_1(\mathbb{P}_1({\bm{x}}))-{\bm{j}}_0(\mathbb{P}_0({\bm{x}})),\quad \tau_{{\bm{j}}}({\bm{z}})={\bm{j}}_1({\bm{z}})-{\bm{j}}_0({\bm{z}}),
\end{equation}
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 \textit{error of posterior CATE}, with or without knowing $\mathbb{P}_t$, is defined as
\begin{equation}
\epsilon_{{\bm{f}}}^{*}({\bm{x}}) \coloneqq \mathbb{E}_{q({\bm{z}}|{\bm{x}})}(\hat{\tau}_{{\bm{f}}}({\bm{z}})-\tau({\bm{x}}))^2;\quad
\epsilon_{{\bm{f}}}({\bm{x}}) \coloneqq \mathbb{E}_{q({\bm{z}}|{\bm{x}})}(\hat{\tau}_{{\bm{f}}}({\bm{z}})-\tau_{{\bm{j}}}({\bm{z}}))^2.
\end{equation}
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 \eqref{eq:taus}. 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.
\begin{definition}[CATE risks]
Let
${Y}({\hat t})|\mathbb{P}_{\hat t}({X}) \sim p_{{Y}({\hat t})|\mathbb{P}_{\hat t}}({\bm{y}}|P)$ and
$q_t({\bm{z}}|{\bm{x}})\coloneqq q({\bm{z}}|{\bm{x}},t)=\mathbb{E}_{p({\bm{y}}|{\bm{x}},t)} q_{\bm\phi}$. The \textit{potential outcome loss} at $({\bm{z}}, t)$, \textit{factual risk}, and \textit{counterfactual risk} are:
\begin{equation*}
\begin{gathered}
\mathcal{L}_{{\bm{f}}}({\bm{z}},\hat{t})\coloneqq \mathbb{E}_{p_{{Y}({\hat t})|\mathbb{P}_{\hat t}}({\bm{y}}|P={\bm{z}})}({\bm{y}}-{\bm{f}}_{\hat t}({\bm{z}}))^2/{\bm{g}}_{\hat t}({\bm{z}})^{2}
={\bm{g}}_{\hat t}({\bm{z}})^{-2}\textstyle\int ({\bm{y}}-{\bm{f}}_{\hat t}({\bm{z}}))^2 p_{{Y}({\hat t})|\mathbb{P}_{\hat t}}({\bm{y}}|{\bm{z}})d{\bm{y}};
\\
\epsilon_{F,t}({\bm{x}}) \coloneqq \mathbb{E}_{q_t({\bm{z}}|{\bm{x}})}\mathcal{L}_{{\bm{f}}}({\bm{z}},t);\quad \epsilon_{CF,t}({\bm{x}}) \coloneqq \mathbb{E}_{q_{1-t}({\bm{z}}|{\bm{x}})}\mathcal{L}_{{\bm{f}}}({\bm{z}},t).
\end{gathered}
\end{equation*}
\end{definition}
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, \textit{the regression error (second) term in ELBO \eqref{eq:elbo_imp} 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{th:gen_bound} below––which, in turn, bounds $\epsilon_{{\bm{f}}}$ by decomposing it to $\epsilon_{F,t}$, $\epsilon_{CF,t}$, and $\mathbb{V}_{{Y}}$.
\begin{theorem}[CATE error bound]
\label{th:gen_bound}
Assume $|\mathcal{L}_{{\bm{f}}}({\bm{z}},t)| \leq M$ and $|{\bm{g}}_t({\bm{z}})| \leq G$, then:
\begin{equation}
\begin{split}
\textstyle
\epsilon_{{\bm{f}}}({\bm{x}}) \leq 2[G^2(\epsilon_{F,0}({\bm{x}}) + \epsilon_{F,1}({\bm{x}}) +M\mathbb{D}({\bm{x}})) - \mathbb{V}_{{Y}}({\bm{x}})]
\end{split}
\end{equation}
where $\mathbb{D}({\bm{x}})\coloneqq \sum_t\sqrt{D_{\mathrm{KL}}(q_t \Vert q_{1-t})/2}$,
and $\mathbb{V}_{{Y}}({\bm{x}})\coloneqq\mathbb{E}_{q({\bm{z}}|{\bm{x}})}\sum_t\mathbb{E}_{p_{{Y}({t})|\mathbb{P}_{t}}({\bm{y}}|{\bm{z}})}({\bm{y}}-{\bm{j}}_{t}({\bm{z}}))^2$.
\end{theorem}
$\mathbb{D}({\bm{x}})$ measures the imbalance between $q_t({Z}|{\bm{x}})$ and is symmetric for $t$.
Correspondingly, \textit{the KL term in ELBO \eqref{eq:elbo_imp} 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 \textit{$G$ is implicitly controlled by the third term in ELBO \eqref{eq:elbo_imp}}, 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 \citet{tarr2021estimating} and \citet{kallus2018more}.
\section{Experiments}
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{sec:pokec}.
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, 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}})$ \citep{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 \eqref{eq:elbo_imp}, 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.
\subsection{Synthetic dataset}
\label{sec:exp_syn}
\begingroup
\begin{equation}
\label{art_model}
\begin{split}
\textstyle
{W}|{X} \sim \mathcal{N}({\bm{h}}({X}), {\bm{k}}({X}));\;
{T}|{X} \sim \operatorname{Bern}(\operatorname{Logi}(\omega l({X})));\;
{Y}|{W},{T} \sim \mathcal{N}(f_{{T}}({W}), g_{{T}}({W})).
\end{split}
\end{equation}
\vspace{-.05in}
We generate synthetic datasets following \eqref{art_model}. 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{ass:p_inj}, when $\dim({W}) > 1$.
We have 5 different overlap levels controlled by $\omega$ that multiplies the logit value. See Sec.~\ref{sec:exp_syn_app} for details and more results on synthetic datasets.
\setlength{\columnsep}{8pt}
\setlength{\intextsep}{5pt}
\begin{wrapfigure}{r}{0.38\textwidth}
\vspace{-.35in}
\begin{center}
\includegraphics[width=0.38\textwidth]{fig_depnoise.pdf}
\end{center}
\vspace{-.2in}
\caption{\footnotesize{$\sqrt{\epsilon_{pehe}}$ on synthetic datasets. Error bar on 10 random DGPs.}}
\vspace{-.15in}
\label{fig:depn}
\end{wrapfigure}
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 \eqref{art_model}. 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{fig:depn}, we adjust one of $\omega, \dim({W})$, with the other fixed to the lowest. As implied by our theory, our method, with only \textit{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{recover} also).
\setlength{\columnsep}{8pt}
\setlength{\intextsep}{5pt}
\begin{wrapfigure}{r}{0.38\textwidth}
\vspace{-.1in}
\begin{center}
\includegraphics[width=0.38\textwidth]{latent_2plots_h.png}
\end{center}
\vspace{-.18in}
\caption{\footnotesize{Plots of recovered
- true
latent. Blue: $t=0$, Orange: $t=1$.}}
\vspace{-.1in}
\label{recover}
\end{wrapfigure}
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{sec:estimation}.
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{fig:z200}). Thus, the performance gap under $\dim({W})>1$ in Figure \ref{fig:depn} should be due to the capacity of NNs in $\beta$-Intact-VAE. In Appendix Figure \ref{fig:ate} 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{sec:ihdp} 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{recover}, and more plots are in Appendix Figure \ref{fig:recover}.
For comparison, we run CEVAE \citep{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{fig:bad_recover}.
\endgroup
\subsection{IHDP benchmark dataset}
\label{sec:ihdp}
This experiment shows our conditional BRL matches state-of-the-art BRL methods and does not overly focus on PEHE. The IHDP \citep{hill2011bayesian} is a widely used benchmark dataset; while it is less known, its covariates are limited-overlapping, and thus it is used in \citet{johansson2020generalization} which considers limited overlap. The dataset is based on an RCT, but \texttt{Race} is artificially introduced as a confounder by removing all treated babies with nonwhite mothers in the data. Thus, \texttt{Race} is highly limited-overlapping, and other covariates that have high correlation to \texttt{Race}, e.g, \texttt{Birth weight} \citep{kelly2009does}, are also limited-overlapping.
See Sec.~\ref{sec:ihdp_app} 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 \citet{louizos2017causal}.
We set $\beta=1$ since it works well on synthetic datasets with limited overlap.
As shown in Table \ref{t:ihdp}, $\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 \cite{shalit2017estimating} is based on bounding PEHE, thus results in sub-optimal ATE estimation (see also Appendix Figure \ref{fig:ate} where CFR gives larger ATE errors with less overlap).
\begin{table}[h]
\centering
\scriptsize
\vspace{-.05in}
\caption{Errors on IHDP over 1000 random DGPs.
``Mod.~*'' indicates the modified version with unconditional balancing hyperparameter of value ``*''. \textit{Italic} indicates where the modified version is
significantly worse than the original.
\textbf{Bold} indicates method(s) which is
significantly better than others. The results of other methods are taken from \cite{shalit2017estimating}, GANITE \citep{yoon2018ganite}, and CEVAE \citep{louizos2017causal}.}
\vspace{-.05in}
\begin{tabular}{p{.6cm}p{.5cm}p{.5cm}p{.5cm}p{.5cm}p{.6cm}p{.6cm}p{.8cm}p{.8cm}p{.8cm}p{.8cm}p{.8cm}p{.8cm}}
\toprule
Method &{TMLE} &{BNN} &{CFR} &{CF} &{CEVAE} &{GANITE} &{Ours} &{Mod. 1} &{Mod.~0.2} &{Mod.~0.1} &{Mod.~0.05} &{Mod.~0.01}\\
\midrule
$\epsilon_{ate}$ &.30$_{\pm .01}$ &.37$_{\pm .03}$
&.25$_{\pm.01}$ &\textbf{.18}$_{\pm.01}$ &.34$_{\pm.01}$ &.43$_{\pm.05}$ &\textbf{.177}$_{\pm .007}$ &\textit{.196}$_{\pm .008}$
&\textbf{.177}$_{\pm.007}$ &\textbf{.167}$_{\pm .005}$ &\textbf{.177}$_{\pm .006}$ &\textbf{.179}$_{\pm.006}$
\\
\midrule
$\sqrt{\epsilon_{pehe}}$ &5.0$_{\pm .2}$ &2.2$_{\pm .1}$
&\textbf{.71}$_{\pm.02}$ &3.8$_{\pm.2}$ &2.7$_{\pm.1 }$ &1.9$_{\pm .4}$ &.843$_{\pm .030}$ &\textit{1.979}$_{\pm .082}$
&\textit{1.116}$_{\pm.046}$ &$.777_{\pm .026}$ &.894$_{\pm.039}$ &.841$_{\pm .029}$
\\
\bottomrule
\end{tabular}
\label{t:ihdp}
\vspace{-.05in}
\end{table}
To examine conditional v.s unconditional balance clearly, we modify our method and add two components for unconditional balance from \citet{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).
\vspace{-1mm}
\section{Conclusion}
\vspace{-1mm}
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 \textit{principled} causal inference owing to their probabilistic nature, if not compromised by ad hoc heuristics (see Sec.~\ref{sec:vae_te}).
\citet[Sec.~4.2]{wu2021principled} introduce some newest ideas of this project.
\bibliography{main}
\bibliographystyle{iclr2022_conference}