# Weighting-Based Treatment Effect Estimation via Distribution Learning

Dongcheng Zhang
Institue of Project Management
Tsinghua University
Haidian District, Beijing, China
<EMAIL_ADDRESS>
&Kunpeng Zhang
Robert H. Smith School of Business
University of Maryland, College Park
College Park, MD 20742, USA
<EMAIL_ADDRESS>

###### Abstract

Existing weighting methods for treatment effect estimation are often built
upon the idea of propensity scores or covariate balance. They usually impose
strong assumptions on treatment assignment or outcome model to obtain unbiased
estimation, such as linearity or specific functional forms, which easily leads
to the major drawback of model mis-specification. In this paper, we aim to
alleviate these issues by developing a distribution learning-based weighting
method. We first learn the true underlying distribution of covariates
conditioned on treatment assignment, then leverage the ratio of covariates’
density in the treatment group to that of the control group as the weight for
estimating treatment effects. Specifically, we propose to approximate the
distribution of covariates in both treatment and control groups through
invertible transformations via change of variables. To demonstrate the
superiority, robustness, and generalizability of our method, we conduct
extensive experiments using synthetic and real data. From the experiment
results, we find that our method for estimating average treatment effect on
treated (ATT) with observational data outperforms several cutting-edge
weighting-only benchmarking methods, and it maintains its advantage under a
doubly-robust estimation framework that combines weighting with some advanced
outcome modeling methods.

_K_ eywords Observational study $\cdot$ Weighting $\cdot$ Distribution
learning $\cdot$ Treatment effect estimation

##  1 Introduction

Randomized control experiments have become increasingly appealing as a way to
help researchers in many fields understand the causal effect of a given
treatment. For example, firms invest in A/B testing to see whether a certain
homepage layout will attract more traffic. Marketers design randomized
experiments to test the effectiveness of a specific incentive program to
retain customers and increase revenue. In the field of healthcare and
sociology, scientists desire to understand the effect of technology access and
advancement on human well-being, and they usually conduct their studies
through randomized experiments. In general, these randomized experiment
designs can guarantee that the treatment and the control groups are randomly
different from one another in regard to all background variables. This in turn
allows researchers to gauge the causal impact of a treatment on a certain
outcome variable. One major challenge, however, is the cost of field
experiments, which tend to be expensive and occasionally harmful to
participants or platforms. Additionally, randomization is impossible in many
situations. For example, if we were concerned about the negative effects of
some treatment, such as drugs, tobacco smoke, or alcohol, it would be
unethical to deliberately assign and expose a group of individuals to this
treatment. Thus, researchers must develop new techniques to estimate treatment
effects using the collected observational data.

Existing methods and challenges The alternatives to randomized field
experiments are various computational methods that use observational data,
primarily including regression, matching and weighting. In regression models,
researchers adjust confounding variables that are correlated with both
treatment and outcome variables to alleviate the confounding effect. Although
regression models are gaining widespread popularity, they rely heavily on the
assumption of a functional form, often a linear form, that does not reflect
most real-world scenarios. Matching and weighting methods are then developed
to tackle endogeneity issues in observational studies, by balancing covariates
in treatment and control groups. While matching methods do not fully use all
data, weighting overcomes this limitation by assigning all individuals
different weights so they can all be used [1]. Both matching and weighting
have been well studied to replicate randomized control experiments as closely
as possible, and their popularity is reflected in scattered research across
disciplines such as statistics, economics, political science, sociology, and
business.

Current weighting methods for treatment effect estimation are often built upon
the idea of propensity scores or covariate balance. In 1983, Rosenbaum and
Rubin introduced the propensity score, which summarizes all covariates into
one scalar to represent the probability of treatment assignment conditioned on
observed covariates [2]. A typical propensity score-based weighting method for
treatment estimation is inverse probability of treatment weighting (IPTW) [3],
which has been theoretically proven to result in an unbiased estimation [4].
Thus, the key challenge becomes estimating the propensity score in an accurate
way.

To estimate the propensity score, researchers often assume a simply specified
functional form to model the relationship of covariates to the treatment
assignment. For example, a simple logistic regression is most typically used.
However, this simple assumption does not hold in real applications, and its
failure causes model mis-specification [5]. This mis-specification in turn can
easily lead to poor propensity score prediction and consequently to biased
treatment effect estimation. Therefore, many researchers have proposed using
machine learning methods such as Lasso, boosting regression, bagged CART and
random forest in order to improve propensity score estimation by modeling the
complicated non-linearity.

In addition to developing propensity score-based weighting, the existing
literature has also introduced statistical weighting methods that aim to
directly balance covariates between treatment and control groups. These
methods include entropy balancing [6], optimization-based weighting [7],
empirical balancing calibration weighting [8], and many others [9, 10, 11].
One common drawback of these methods is that rather than balancing the joint
distribution of covariates, their objective is only balancing the first moment
(i.e., means) and sometimes the second moment or interactions of covariates.
In consequence, these methods call for some strong assumptions, such as
linearity or specific functional forms, in order to achieve unbiased
estimation of treatment effects.

Present study To address the aforementioned issue of model mis-specification
in weighting, in this paper we propose a new weighting-based method via
distribution learning for treatment effect estimation with observational data.
The fundamental idea is that when we have the true underlying distribution of
covariates conditioned on treatment assignment, we can compute the ratio of
covariates’ probability density in the treatment group to that in the control
group, use it as weighting to balance the distribution of covariates, and
thereby obtain an unbiased treatment effect estimation.

The key question then becomes how we can learn the underlying distribution of
covariates in both groups. Some parametric methods can be used. For example,
we can assume the data follows some commonly used probability distribution
(e.g., Gaussian) and estimate associated parameters using maximum likelihood
estimation or other techniques. This might be problematic, as the true
probability density function of data is often unknown, is likely far from
Gaussian or any other well-known distributions, and even an analytical form
might not exist. Therefore, to approximate the true distribution from data, we
borrow the idea of constructing flexible learnable probability distribution
via transformation [12]. Specifically, we start with a simple base
distribution and take a series of invertible transformations via change of
variables to approach a distribution that is as close as possible to the
truth. Previous studies have shown that the transformed distribution can
converge to the truth, as long as there exist an infinite amount of
observations and the transformations have enough expressive power [12, 13].
However, finding a closed-form analytical function for a probability
distribution via transformation is not easy and is sometimes impossible.
Building upon the universal approximation theory, we develop a data-driven
method by leveraging deep neural networks to learn such transformation
functions.

The transformation can be computationally expensive since it involves a
component related to computing the determinant of a Jacobian. To make our
method practical, we consider using transformations with determinants of
Jacobians that are easy to compute. To this end, we choose an autoregressive
transformation function $x_{i}=f(z_{1:i})=\tau(z_{i},c(z_{1:i-1}))$, where
$\tau$ is the invertible transformation and $c$ is the autoregressive
conditioner. For autoregressive transformations, Jacobians are triangular
matrices whose log-determinants can be easily computed in linear time, as the
sum of diagonal entries on a log scale, as we will describe more fully in the
methodology section. After approximating the distribution of covariates
conditioned on whether the individual is treated $p(x\mid W=w)$,
$w\in\\{0,1\\}$, we then use the density ratio $\frac{p(x\mid W=1)}{p(x\mid
W=0)}$ as the weight for every individual $x$ in the control group to estimate
the treatment effect.

To evaluate the performance of our proposed method, in this study we focus on
estimating the average treatment effect on the treated (ATT) with
observational data. The ATT is the average treatment effect on those subjects
who ultimately received the treatment, which is one of the most widely studied
treatment effects; however, our method can also easily be extended to estimate
the average treatment effect (ATE). In subsequent sections of this paper, we
conduct simulation studies under model specification settings that vary with
respect to treatment assignment and outcome model. In addition, to demonstrate
the effectiveness, robustness, scalability, and generalizability of our
approach, we vary the dimensions of covariates, degrees of confounding and
sample sizes in our simulations. The experiment results show that our model
outperforms several cutting-edge weighting-based baselines. For example, it
can achieve up to 81% RMSE reduction over the best weighting-based benchmark.
We will explain the details of our results in the experiments section.

Unlike statistical weighting approaches, our method does not impose strong
assumptions on the relationship between covariates and outcome, such as
linearity and additivity. Moreover, our method calculates weights based on the
learned distribution of covariates, rather than using propensity score
function fitting as existing propensity score-based weighting methods do. This
explains why our method has consistent performance across different settings
and outperforms state-of-the-art weighting-based methods. To further
demonstrate the superiority and generalizability of our method, we apply our
approach to a real-world dataset, the Twins data introduced by [14]. To the
best of our knowledge this data is the only public dataset in which we can
obtain counterfactual outcomes for all individuals and which allows us to
arbitrarily introduce confounding. The experiment results tell us that
compared to other weighting-based methods, the ATT estimated using our
proposed method is more accurate, as it is closer to the actual ATT.

As previously discussed, treatment effect estimation primarily lies in
weighting or outcome modeling, which each have their own pros and cons. One
natural extension of earlier research lies in combining these two approaches
for better treatment effect estimation. In statistics and economics,
researchers have explored along these lines, developing a doubly-robust
framework [15, 16] where a weighting method is combined with outcome modeling.
In this study, we also show the effectiveness of our weighting method by
plugging it into the doubly-robust framework, which serves as the replacement
of the weighting component. Using an advanced machine learning component in
the doubly-robust framework as alternative to regression for outcome modeling,
such as BART or random forest, our approach consistently achieves the best
performance among all baseline estimators. Finally, we compare our method with
existing state-of-the-art approaches, such as deep learning-based generative
models [14], purely statistical methods [11] and machine learning methods
[17]. Not surprisingly, our method exhibits better performance. In sum, we
make four major contributions in this study.

  1. (i)

We propose a new weighting-based method via distribution learning to
approximate the underlying distribution of covariates conditioned on the
treatment assignment for estimating treatment effects. We believe this
innovative approach could open a new perspective in studying ATT and ATE
estimation using observational data.

  2. (ii)

We theoretically demonstrate that our method can be viewed as importance
sampling and can obtain an unbiased estimation as long as the learned
distribution is close to the truth. We also empirically validate the
convergence of the learnable transformed distribution to the true
distribution.

  3. (iii)

To evaluate the performance of our weighting-based method, we conduct
extensive experiments on synthetic data across various settings, as well as
real data. The results show the superiority and robustness of our method as
compared to cutting-edge baselines.

  4. (iv)

In addition to weighting-only-based ATT estimation, we further demonstrate the
efficacy and superiority of our method under the paradigm of the doubly-robust
framework.

##  2 Literature Review

A large body of work in the literature addresses treatment effect estimation,
especially the weighting-based methods, with observational data. The methods
studied can be primarily categorized into propensity score and covariate
balance-based approaches. In this section we provided an overview of some
pertinent studies, then we review the doubly-robust framework and machine
learning methods used for treatment effect estimation.

Propensity score-based weighting estimators: The most well-known and widely
used propensity score weighting method is the inverse probability of treatment
weighting (IPTW) [18, 19, 3, 20]. The weight in IPTW is set to
$\gamma_{i}=\frac{w_{i}}{\hat{e}_{i}}+\frac{1-w_{i}}{1-\hat{e}_{i}}$, where
$\hat{e}_{i}$ is the estimated propensity score for individual $i$.
$\gamma_{i}$ is then used to weight both treatment ($w_{i}=1$) and control
($w_{i}=0$) groups to estimate ATE; alternatively
$\gamma_{i}=w_{i}+(1-w_{i})\frac{\hat{e}_{i}}{1-\hat{e}_{i}}$ is used to
estimate ATT [20, 4]. Often, $\hat{e}_{i}$ is estimated by a regression (e.g.,
logistic regression) over covariates. However, most regression models impose a
strong assumption about the functional form of the true propensity score
model, which may not be true in real applications. This consequently leads to
a common issue in causal inference, called model mis-specification, which can
easily result in bad propensity score prediction, accordingly leading to an
unbalanced covariate distribution via weighting, and thus finally yielding a
biased treatment effect estimation [21]. To overcome this issue, researchers
have developed various advanced machine learning-based methods as alternatives
to regression approaches including Lasso [22], boosting regression [23],
bagged CART and random forest [21, 24]. These regression approaches have all
been developed in the context of propensity score estimation to improve the
propensity score weighting performance. Due to the imperfect nature of
prediction models, estimations of unbiased propensity scores remain non-
guaranteed both theoretically and practically.

Covariate balance-based weighting estimators: Another stream of literature on
weighting-based treatment effect estimation focuses on balancing covariates.
For example, Hainmueller introduced entropy balancing, which relies on a
maximum entropy reweighting scheme that calibrates unit weights to achieve
pre-specified balance conditions [6]. Imai and Ratkovic proposed a covariate
balancing propensity score (CBPS) method, which combines treatment assignment
modeling and covariate balance optimization [9]. Zubizarreta learned the
stable weights of minimum variance that balance precisely the means of the
observed covariates and other features of their distributions [7]. Chan et al.
considered a wide class of calibration weights constructed to obtain balance
of the moments of observed covariates [8]. All these statistical weighting
estimators aim to achieve covariate balance for the treatment and control
groups. However, in practice they can only model the first, second, and
possibly higher moments of covariates, and they are not able to adjust
inequalities of the covariates’ joint distribution. Therefore, they rely on
some strong assumptions, such as linearity and additivity, to obtain unbiased
estimates of treatment effects.

Doubly-robust framework: Existing studies have demonstrated that the treatment
effect estimation can be tackled by either properly controlling the treatment
assignment mechanism (e.g., propensity score weighting) or modeling the
outcome function (e.g., regressions) [4, 25]. Doubly-robust estimators combine
the propensity score weighting with outcome modeling. They are “doubly robust”
in the sense that they can achieve consistency as long as either the
propensity score or the outcome model is estimated consistently, and they can
achieve efficiency if these two components are estimated at a sufficiently
fast rate [4, 11]. A simple doubly-robust method is augmented inverse
probability weighting (AIPW) [15], which estimates ATT as follows:

$\hat{E}(Y^{0}\mid
W=1)=\frac{1}{n_{1}}\sum_{\\{i:w_{i}=1\\}}\hat{h}_{0}(x_{i})+\sum_{\\{i:w_{i}=0\\}}\frac{\hat{e}_{i}}{1-\hat{e}_{i}}(y_{i}^{0}-\hat{h}_{0}(x_{i}))/\sum_{\\{i:w_{i}=0\\}}\frac{\hat{e}_{i}}{1-\hat{e}_{i}},\\\
\widehat{ATT}=\hat{E}(Y^{1}\mid W=1)-\hat{E}(Y^{0}\mid
W=1)=\frac{1}{n_{1}}\sum_{\\{i:w_{i}=1\\}}y_{i}^{1}-\hat{E}(Y^{0}\mid W=1),$
(1)

where $n_{1}$ is the size of the treatment group and $\hat{h}_{0}(x_{i})$ is
the fitted outcome function (e.g. linear regression). In fact, the propensity
score weighting component
$\frac{\hat{e}_{i}}{1-\hat{e}_{i}}/\sum_{\\{i:w_{i}=0\\}}\frac{\hat{e}_{i}}{1-\hat{e}_{i}}$
in the AIPW can be substituted by other weighting estimators. For example,
[11] proposed an approximate residual balancing method,

$\hat{E}(Y^{0}\mid
W=1)=\frac{1}{n_{1}}\sum_{\\{i:w_{i}=1\\}}\hat{h}_{0}(x_{i})+\sum_{\\{i:w_{i}=0\\}}\gamma_{i}\cdot(y_{i}^{0}-\hat{h}_{0}(x_{i})),$
(2)

where $\gamma_{i}$ is based on covariate balance optimization and
$\hat{h}_{0}(x_{i})$ is fitted by a Lasso or an elastic net. The outcome model
$\hat{h}_{0}(x_{i})$ can be fitted using highly data-adaptive machine learning
algorithms, such as random forest, Bayesian additive regression tree (BART)
and boosting [26]. Hence, a doubly-robust framework can be seen as a generic
framework in which different weighting methods (including propensity score
weighting and covariate-balance based weighting) and different outcome
modeling approaches (including linear regression, regularized regression and
other machine learning methods) can be employed.

Machine learning-based methods: As machine learning has advanced and
especially as deep neural networks have become successful in many domains,
various algorithms have been developed for causal inference with observational
data. For example, tree-based algorithms such as causal forest are extended in
the potential outcome framework to estimate treatment effects. For example,
Wager and Athey developed a nonparametric causal forest for estimating
heterogeneous treatment effects that extends the widely used random forest
algorithm. In their potential outcomes framework with unconfoundedness, they
showed that causal forests are pointwise consistent for the true treatment
effect and have an asymptotically Gaussian and centered sampling distribution
[17]. Athey and Imbens developed a recursive partition approach to construct
trees for causal effects that allow to do valid inference for the causal
effects in randomized experiments and in observational studies satisfying
unconfoundedness. Their method provides valid confidence intervals without
restrictions on the number of covariates or the complexity of the data-
generating process [27]. Yahav et al. introduced a tree-based approach
adjusting for observable self-selection bias in intervention studies [28].
Chen et al. applied causal forest to analyze patient-level treat effect and
facilitate the implementation of the targeted outreach program [29]. Deep
learning-based methods such as balancing neural network (BNN) [30],
counterfactual regression (CFR) [31], and causal effect variational
autoencoder (CEVAE) [14], are proposed for causal effect inference by
leveraging the expressive power and high flexibility of neural networks. Among
these methods, CEVAE uses a variational autoencoder to learn the latent
variable representation for observed covariates in both treatment and control
groups, and it estimates the treatment effect based on counterfactual
prediction. To the best of our knowledge, CEVAE is the best deep neural
network-based model for treatment effect estimation.

##  3 Methodology

In this section, we first formally define our research problem and illustrate
necessary assumptions. Then we propose a general framework of treatment effect
estimation via density approximation. Meanwhile, we show it allows us to
obtain an unbiased estimator in theory. Finally, we introduce our weighting
approach DLW under this framework.

###  3.1 Problem definitions and assumptions

Definition 1: Propensity score. The propensity score for an individual $i$
with covariates $x_{i}$ is the probability that the individual would be
assigned to the treatment group ($W=1$), i.e., $e(x_{i})\equiv P(W=1\mid
X=x_{i})$, where $X$ and $W$ are random variables representing individuals and
treatment/control group assignment, respectively.

Definition 2: Average treatment effect on the treated (ATT). Let
$X=(x_{1},x_{2},\cdots,x_{N})$, $x_{i}\in R^{d}$ denote covariates,
$Y^{0}=(y_{1}^{0},y_{2}^{0},\cdots,y_{N}^{0})$ be the potential outcome when
not treated, $Y^{1}=(y_{1}^{1},y_{2}^{1},\cdots,y_{N}^{1})$ be the potential
outcome when treated, and $W=(w_{1},w_{2},\cdots,w_{N})$, $w_{i}\in\\{0,1\\}$
represent a binary treatment variable. The treatment effect for an individual
$i$ is $\eta_{i}=y_{i}^{1}-y_{i}^{0}$. The average treatment effect for the
entire population is $ATE\equiv E[Y^{1}-Y^{0}]$. If focusing only on the
treatment group, the goal becomes estimating the average treatment effect on
the treated (ATT), formally defined as $ATT\equiv E[Y^{1}-Y^{0}\mid W=1]$. As
previously stated, we focus on ATT estimation in this study, but our proposed
method can easily be extended to estimate ATE.

Assumptions: When weighting methods are used to estimate treatment effects in
non-experimental studies, strongly ignorable treatment assignment and overlap
are always assumed [2, 4]. That is: (1) treatment assignment ($W$) is
independent of the potential outcomes ($Y^{0},Y^{1}$) given the covariates
($X$), i.e., $W\perp(Y^{0},Y^{1})\mid X$, and (2) there is a non-zero
probability of an individual with covariate $X$ receiving each treatment,
i.e., $0<p(W=1\mid X)<1$ for all $X$.

###  3.2 A generic framework of treatment effect estimation via density
approximation

In this section, we first demonstrate that ATT estimation can be derived under
the strongly ignorable treatment assignment assumption as follows:

$\displaystyle ATT$ $\displaystyle=E(Y^{1}-Y^{0}\mid W=1)$ (3)
$\displaystyle=E(Y^{1}\mid W=1)-E(Y^{0}\mid W=1)$ $\displaystyle=E(Y^{1}\mid
W=1)-\int E(Y^{0}\mid W=1,x)*p(x\mid W=1)dx$

By the strongly ignorable treatment assignment assumption, i.e.,
$W\perp(Y^{0},Y^{1})\mid X,E(Y^{0}\mid W=1,X)$ should be equal to $E(Y^{0}\mid
W=0,X)$, the above equation can be rewritten as:

$\displaystyle ATT$ $\displaystyle=E(Y^{1}\mid W=1)-\int E(Y^{0}\mid
W=0,x)*p(x\mid W=1)dx$ (4) $\displaystyle=E(Y^{1}\mid W=1)-\int
y^{0}*p(y^{0}\mid W=0,x)*p(x\mid W=0)*\frac{p(x\mid W=1)}{p(x\mid W=0)}dxdy$
$\displaystyle=E(Y^{1}\mid W=1)-\int y^{0}*p(y^{0},x\mid W=0)*\frac{p(x\mid
W=1)}{p(x\mid W=0)}dxdy$ $\displaystyle=E(Y^{1}\mid
W=1)-\int\left[y^{0}*\frac{p(x\mid W=1)}{p(x\mid W=0)}\right]*p(y^{0},x\mid
W=0)dxdy$ $\displaystyle=E(Y^{1}\mid W=1)-E\left[y^{0}*\frac{p(x\mid
W=1)}{p(x\mid W=0)}\mid W=0\right]$

The last step tells us that $E(Y^{0}\mid W=1)$ can be viewed as the numerical
integration by Markov chain Monte Carlo (MCMC) and specifically by importance
sampling, where the density ratio $\frac{p(x\mid W=1)}{p(x\mid W=0)}$ can be
used as the weight on the control group. Thus, ATT can be obtained as long as
$\frac{p(x\mid W=1)}{p(x\mid W=0)}$ can be estimated. There are many
parametric and nonparametric methods for probability density estimation [32,
33, 34, 35]. In this paper, we develop a data-driven method to learn this
density ratio, and several important aspects related to this density ratio are
worth noting here.

(1) $E(Y^{1}\mid W=1)$ can be estimated via a sample mean of Y in the
treatment group, and it has been proven to be a consistent nonparametric
estimator [36].

(2) In theory, using the true density ratio as weight can be viewed as the
inverse probability of treatment weighting for ATT (IPTW), based on the
Bayesian theorem:

$\displaystyle\frac{p(x\mid W=1)}{p(x\mid W=0)}$
$\displaystyle=\frac{p(W=1\mid x)*P(W=0)}{p(W=0\mid x)*P(W=1)}$ (5)
$\displaystyle=\frac{e(x)}{1-e(x)}*\frac{P(W=0)}{P(W=1)}$

(3) As mentioned, the density ration can be easily extended to ATE estimation.

$\displaystyle ATE$ $\displaystyle=E(Y^{1}-Y^{0})$ (6)
$\displaystyle=E(Y^{1})-E(Y^{0})$ $\displaystyle=\int E(Y^{1}\mid
x)*p(x)dx-\int E(Y^{0}\mid x)*p(x)dx$ $\displaystyle=\int E(Y^{1}\mid
W=1,x)*p(x)dx-\int E(Y^{0}\mid W=0,x)*p(x)dx$ $\displaystyle=\int
y^{1}*p(y^{1}\mid W=1,x)*p(x\mid W=1)*\frac{p(x)}{p(x\mid W=1)}dx$
$\displaystyle-\int y^{0}*p(y^{0}\mid W=0,x)*p(x\mid W=0)*\frac{p(x)}{p(x\mid
W=0)}dx$ $\displaystyle=\int y^{1}*p(y^{1},x\mid W=1)*\frac{p(x)}{p(x\mid
W=1)}dx-\int y^{0}*p(y^{0},x\mid W=0)*\frac{p(x)}{p(x\mid W=0)}dx$
$\displaystyle=\int\left[y^{1}*\frac{p(x)}{p(x\mid W=1)}\right]*p(y^{1},x\mid
W=1)dxdy-\int\left[y^{0}*\frac{p(x)}{p(x\mid W=0)}\right]*p(y^{0},x\mid
W=0)dxdy$

Similarly, $E(Y^{1})$ and $E(Y^{0})$ can be viewed as the MCMC integration
(importance sampling) where $\frac{p(x)}{p(x\mid W=1)}$ and
$\frac{p(x)}{p(x\mid W=0)}$ can be used as weights for the treatment group and
the control group, respectively. According to the Bayesian theorem,

$\displaystyle\frac{p(x)}{p(x\mid W=0)}$ $\displaystyle=\frac{p(x\mid
W=0)*P(W=0)+p(x\mid W=1)*P(W=1)}{p(x\mid W=0)}$ (7)
$\displaystyle=P(W=0)+P(W=1)*\frac{p(x\mid W=1)}{p(x\mid W=0)}$
$\displaystyle\frac{p(x)}{p(x\mid W=1)}$ $\displaystyle=\frac{p(x\mid
W=0)*P(W=0)+p(x\mid W=1)*P(W=1)}{p(x\mid W=1)}$ (8)
$\displaystyle=P(W=1)+P(W=0)*\frac{p(x\mid W=0)}{p(x\mid W=1)}$

As $P(W=0)$ and $P(W=1)$ can be estimated straightforwardly using the
proportion of treated units and controlled units, the estimation of the
importance weights $\frac{p(x)}{p(x\mid W=0)}$ and $\frac{p(x)}{p(x\mid W=1)}$
relies on the estimation of the density ratio $\frac{p(x\mid W=1)}{p(x\mid
W=0)}$.

###  3.3 DLW: our proposed weighting approach via distribution learning

Now we turn to probability density estimation for both treatment and control
groups using observational data. In theory, any probability distribution can
be constructed via a method of change of variables, i.e., through a sequence
of invertible mappings starting from a simple base distribution like Gaussian
or Uniform [13]. Specifically, let $X\in R^{d}$ be a random variable (e.g.
covariates) and express $X$ as an invertible mapping $f:R^{d}\rightarrow
R^{d}$ from random variable $Z$. Then the true distribution of $X$ can be
obtained by a change of variables via the following operations [37, 38]:

$\displaystyle x=f(z),\text{where }z\sim p(z),$ (9) $\displaystyle
p(x)=p(z)\cdot\mid det\frac{\partial f}{\partial z}\mid^{-1}$

Note that the computational complexity of the transformation is dependent on
the determinant of a Jacobian $\mid det\frac{\partial f}{\partial z}\mid$,
which can be very computationally expensive.

By repeatedly applying the rule of change of variables, the base distribution
$p(z^{0})$ "flows" through a sequence of invertible mappings
$(f_{1},f_{2},\cdots,f_{k})$ and finally reaches a much more complex
distribution $p(x)$ that is close to the true underlying distribution.

$\displaystyle x=f_{k}\cdots\cdots f_{1}(z^{0}),\text{where }z^{0}\sim
p(z^{0}),$ (10) $\displaystyle p(x)=p(z^{0})\cdot\prod_{k=1}^{K}\mid
det\frac{\partial f_{k}}{\partial z^{k-1}}\mid^{-1}$

Our goal is to estimate the distribution of $X\in R^{d}$ where we update
parameters involved in the mappings using maximum likelihood estimation (MLE).
As previous studies have proven, when there is an infinite number of
observations and the transformations have enough expressive power, the
transformed distribution can converge to the true distribution of $X$, where
the likelihood function converges to a constant [12, 13].

A number of transformations have been introduced in the machine learning
community, and of those we choose the autoregressive transformation for the
ease of calculating the determinant of Jacobian $\mid det\frac{\partial
f}{\partial z}\mid$. The autoregressive transformation is formalized as

$\displaystyle x_{i}=f(z_{1:i})=\tau(z_{i},c(z_{1:i-1};\theta))$ (11)

where $\tau$ is the invertible transformation and $c$ is the autoregressive
conditioner. As $x_{i}$ only depends on $z_{1:i-1}$ and $z_{i}$ in
autoregressive transformations, the Jacobians are triangular matrices with
log-determinants that can be computed in linear time as the sum of the
diagonal entries on a log scale. Previous studies have proven that this kind
of transformation is sufficiently expressive to form a universal approximator
for probability distributions if the autoregressive conditioner, $c$, is
characterized by a masked autoregressive neural network $NN_{1}$ [39, 40], and
the invertible transformation $\tau$ is characterized by another neural
network $NN_{2}$ [41, 42]. The outputs of $NN_{1}$ are weights and biases of
$NN_{2}$. Specifically, the “neural autoregressive transformation” is
formalized as follows:

$\displaystyle x_{i}$ $\displaystyle=f(z_{1:i})=\tau(z_{i},c(z_{1:i-1}))$ (12)
$\displaystyle=NN_{2}(z_{i},NN_{1}(z_{1:i-1}))\text{ for }i\text{ in }1:d,$

where $d$ is the dimension of covariates $X$ and $i$ represents the index of
every dimension.

Using the transformations described above, we can construct a universal
approximator for probability distributions by repeatedly applying the rule of
change of variables from the base distribution $p(z_{0})$. Then by maximum
likelihood estimation, we can approximate $p(x\mid W=1)$ and $p(x\mid W=0)$
using all training data in the treatment and the control groups, respectively.
To summarize, the procedure for our distribution learning-based weighting for
ATT estimation is outlined below, and the overall framework is shown in Figure
1.

  1. I)

For the treatment group, the distribution of covariates $x\in R^{d}$, $p(x\mid
W=1)$ is learned by $K$ “neural autoregressive transformations” from a simple
base distribution ($d$-dimensional Gaussian with unit variance), i.e.,
$p(x\mid W=1)=p(z_{0})*\prod_{k=1}^{K}\mid det\frac{\partial f_{k}}{\partial
z_{k-1}}\mid^{-1}$, where $z_{0}\sim N(0,I)$. All associated parameters are
learned by MLE using data in the treatment group. Specifically, we choose the
widely used stochastic gradient-based optimization method $Adam$ [40] to
optimize the likelihood function.

  2. II)

For the control group, we have a similar procedure where the distribution
$p(x\mid W=0)$ is learned by $K$ “neural autoregressive transformations” from
a $d$-dimensional Gaussian with unit variance.

  3. III)

The density ratios $\frac{p(x\mid W=1)}{p(x\mid W=0)}$ for all individuals in
the control group are calculated. Based on the derivation in Section 3.2,
$E(Y^{0}\mid W=1)$ can be estimated by taking the weighted average of $Y^{0}$
in the control group, with $\frac{p(x\mid W=1)}{p(x\mid W=0)}$ being the
weights.

  4. IV)

$E(Y^{1}\mid W=1)$ can be estimated nonparametrically via the sample mean of
$Y^{1}$ in the treatment group, and finally we estimate the ATT by:

$\displaystyle\widehat{ATT}$ $\displaystyle=\hat{E}(Y^{1}\mid
W=1)-\hat{E}\left[Y^{0}*\frac{p(x\mid W=1)}{p(x\mid W=0)}\mid W=0\right]$ (13)
$\displaystyle=\frac{1}{n_{1}}\sum_{\\{i:w_{i}=1\\}}y_{i}^{1}-\left[\sum_{\\{i:w_{i}=0\\}}\frac{p(x\mid
W=1)}{p(x\mid W=0)}*y_{i}^{0}\right]/\sum_{\\{i:w_{i}=0\\}}\frac{p(x\mid
W=1)}{p(x\mid W=0)}.$

Note that this procedure can be applied for ATE estimation as illustrated in
Section 3.2.

Figure 1: Overall framework of ATT estimation via distribution learning.

##  4 Experiments

In this section, we evaluate the performance of our proposed distribution
learning-based weighting method in comparison with existing methods for ATT
estimation, with the goal of addressing the following questions: (Q1): How
does our proposed weighting method perform compared with other weighting
methods? (Q2): Can our method consistently achieve good performance of ATT
estimation under different model specifications with respect to treatment
assignment and outcome models? (Q3): Can our method really approximate the
underlying true distribution of covariates conditioned on treatment
assignment? All experiments are conducted using Python and some R
packages111https://CRAN.R-project.org/package=WeightIt. [1] on a machine with
a RAM of 32Gb. The details of hyper-parameter tuning are described in the
Appendix. To facilitate reproducibility, we release our code and data
online222https://github.com/DLweighting/Distribution-Learning-based-
weighting..

###  4.1 Evaluation metrics

We evaluate all methods using two widely used metrics: bias and root mean
squared error (RMSE), which are defined as

$\displaystyle Bias(ATT,\widehat{ATT}_{i})$
$\displaystyle=\frac{1}{K}\sum_{i=1}^{K}\widehat{ATT}_{i}-ATT,$ (14)
$\displaystyle RMSE(ATT,\widehat{ATT}_{i})$
$\displaystyle=\sqrt{\frac{1}{K}\sum_{i=1}^{K}(\widehat{ATT}_{i}-ATT)^{2}}$

where $K$ is the number of experiments (e.g., 100 in our study),
$\widehat{ATT}_{i}$ is the estimated ATT in the $i^{th}$ experiment, and ATT
represents the true average treatment effect on the treated.

###  4.2 Baselines

We compare our method with other weighting-based estimators and several
methods commonly used in machine learning and statistics. We also include the
unadjusted estimator of ATT, denoted as $\widehat{ATT}_{base}$, which
evaluates ATT by directly taking the average outcome of the control group as
the counterfactual for the treatment group. It completely ignores the
confounding bias and is defined mathematically as
$\widehat{ATT}_{base}=\frac{1}{n_{1}}\sum_{\\{i:w_{i}=1\\}}y_{i}^{1}-\frac{1}{n_{0}}\sum_{\\{i:w_{i}=0\\}}y_{i}^{0}$,
where $n_{0}$ and $n_{1}$ are the number of units in the control and treatment
groups, respectively. The comparison baselines are organized as follows:

  1. (1)

First, in the weighting-only framework for ATT estimation, we compare our
distribution learning-based weighting estimator $\widehat{ATT}_{DLW}$ with
other weighting estimators, including propensity score-based weighting and
covariate balance-based weighting estimators.

  2. (2)

Second, in the doubly-robust framework, we replace the weighting component
with all weighting estimators, including ours and baselines. Combining
weighting with a widely used machine learning component for outcome modeling,
i.e., BART or random forest, we again show the effectiveness of our weighting
estimator.

  3. (3)

Finally, we compare our method with existing frequently used methods in
statistics and machine learning.

I. Weighting-only methods

  1. 1)

Covariate balance-based weighting methods

     * •

$\widehat{ATT}_{ebal}$: Entropy balancing estimates the ATT by using a maximum
entropy reweighting scheme that balances the covariates in treatment and
control groups [6].

     * •

$\widehat{ATT}_{cbps}$: Covariate balancing propensity score weighting
combines treatment assignment modeling and covariate balance optimization for
ATT estimation [9].

     * •

$\widehat{ATT}_{opt}$: Optimization-based weighting estimates ATT via stable
weights obtained by optimizing the balance of observed covariates [7].

     * •

$\widehat{ATT}_{ebcw}$: Empirical balancing calibration weighting estimates
the ATT by constructing a wide class of calibration weights based on the
balance of the moments of observed covariates [8].

  2. 2)

Propensity score-based weighting methods
As stated in Section 2, propensity score weighting (IPTW) sets the weight
$\gamma_{i}=w_{i}+(1-w_{i})*\frac{\hat{e}_{i}}{1-\hat{e}_{i}}$ to estimate the
ATT (denoted as $\widehat{ATT}_{ps}$), where the propensity score
$\hat{e}_{i}$ is often modeled by logistic regression [1, 20]. Advanced
machine learning models, including random forest ($\widehat{ATT}_{RF}$),
bagged cart ($\widehat{ATT}_{BCART}$), gradient boosting
($\widehat{ATT}_{XGBoost}$), and Lasso ($\widehat{ATT}_{Lasso}$), are also
proposed as alternatives to logistic regression to improve the performance of
propensity score estimation.

II. Doubly-robust framework
As previously mentioned, a doubly-robust framework is another common framework
that adopts weighting for ATT estimation:

$\displaystyle\hat{E}(Y^{0}\mid W=1)$
$\displaystyle=\frac{1}{n_{1}}\sum_{\\{i:w_{i}=1\\}}\hat{h}_{0}(x_{i})+\sum_{\\{i:w_{i}=0\\}}\gamma_{i}\cdot(Y_{i}^{0}-\hat{h}_{0}(x_{i}))$
(15) $\displaystyle\widehat{ATT}$
$\displaystyle=\frac{1}{n_{1}}\sum_{\\{i:w_{i}=1\\}}y_{i}^{1}-\hat{E}(Y^{0}\mid
W=1),$

where $n_{1}$ is the treatment group size, $\gamma_{i}$ is the weight based on
either covariate balance or propensity score, and $\hat{h}_{0}(x_{i})$ is the
outcome function which can be fitted by regression or other data-adaptive
algorithms. Various advanced machine learning approaches can be used to model
the outcome function, and in this study we choose two widely used ones: random
forest and BART. We compare the performance of different weighting methods
under the doubly-robust framework, with the outcome function modeled by random
forest and BART, respectively. Note that these estimators are denoted as
$\widehat{ATT}_{DR\\_*\\_ranf}$ and $\widehat{ATT}_{DR\\_*\\_bart}$, where ‘*’
represents the weighting component implemented by various weighting
approaches, including our distribution learning-based weighting DLW.

III. Frequently used methods
To further demonstrate the effectiveness and usefulness of our proposed
method, we also compare it against current frequently used methods in
statistics and machine learning, including:

  * •

$\widehat{ATT}_{OLS}$: Ordinary least square regresses the outcome variable on
the treatment assignment variable while controlling covariates to estimate the
ATT.

  * •

$\widehat{ATT}_{PSM}$: Propensity score matching uses logistic regression to
calculate the probability of an individual being treated and then performs 1:1
nearest neighbor matching to select for each treated individual $i$ a control
individual with the smallest distance from that individual $i$ [2].

  * •

$\widehat{ATT}_{ARB}$: Approximate residual balancing combines balancing
weights with a regularized regression adjustment [11].

  * •

$\widehat{ATT}_{CF}$: Causal forest extends the widely used random forest
algorithm to estimate treatment effect [17].

  * •

$\widehat{ATT}_{CEVAE}$: CEVAE estimates the ATT by latent variable
representation learning and counterfactual prediction employing a deep
learning-based generative model, i.e., variational autoencoder [14].

###  4.3 Experiments using synthetic data

In this section, we first introduce our generation of synthetic data and then
demonstrate the effectiveness of our distribution learning-based weighting
method for ATT estimation under various settings.

Dataset To demonstrate that our method is capable of handling data that is
complex and close to the reality, we generate synthetic data with relatively
complex distributions. We first follow prior studies and draw each covariate
from a three-component Gaussian mixture distribution with means of (-3, 0, 3),
unit variance and equal mixture proportion [11]. Then we standardize the
covariates in each dimension to ensure all covariates have zero mean and unit
variance as most in the literature do [43, 21, 11]. To show the
generalizability and robustness of our method, we simulate three different
model specification scenarios (i.e., linear, moderately non-linear and
strongly non-linear) based on existing studies [43, 7, 44, 45, 5]. We also
experiment with small and large sample sizes $N=\\{5000,10000\\}$, weak and
strong confounding strengths $s_{c}=\\{0.2,0.4\\}$, and low and high
dimensions of covariates $d=\\{8,16\\}$.

Setting 1 (Linear and additive): This very simple setting has been widely used
in many prior studies. It uses a correctly specified logistic function to
generate the binary treatment variable $W$ and a linear function to generate
the outcome $Y$ as

$\displaystyle W\sim Bernoulli(e(x)),$ (16) $\displaystyle e(x)=p(w=1\mid
x)=\frac{1}{1+exp(s_{c}*\sum_{i=1}^{d}x_{i})},$ $\displaystyle
Y_{linear}=\sum_{j=1}^{d}\beta_{1j}x_{j}+W+\epsilon$

Setting 2 (Moderately non-linear and non-additive): This setting uses a
moderately mis-specified logistic function to generate the binary treatment
variable $W$ and a moderately nonlinear function to generate the outcome $Y$
as

$\displaystyle W\sim Bernoulli(e(x)),$ (17) $\displaystyle e(x)=p(w=1\mid
x)=\frac{1}{1+exp(s_{c}*(\sum_{i=1}^{d}(x_{i}^{2}-1)+\sum_{1\leq i,j\leq
d}x_{i}x_{j}))},$ $\displaystyle
Y_{nonlin}=\sum_{j=1}^{d}\beta_{1j}x_{j}^{2}+\sum_{1\leq i,j\leq
d}\beta_{2j}x_{i}x_{j}+W+\epsilon$

Setting 3 (Strongly non-linear and non-additive): This setting uses a strongly
mis-specified logistic function to generate the binary treatment variable $W$
and a strongly nonlinear function to generate the outcome $Y$ as

$\displaystyle W\sim Bernoulli(e(x)),$ (18) $\displaystyle e(x)=p(w=1\mid
x)=\frac{1}{1+exp(s_{c}*(\sum_{i=1}^{d}(log(x_{i}^{2}+1)-0.5)+\sum_{1\leq
i,j\leq d}x_{i}x_{j}))},$ $\displaystyle
Y_{nonlin}=\sum_{j=1}^{d}\beta_{1j}log(x_{j}^{2}+1)+\sum_{1\leq i,j\leq
d}2*\beta_{2j}sin(x_{i}x_{j})+W+\epsilon$

where $\beta_{ij}\sim Uniform(0,1),\beta_{2j}\sim Uniform(0,1)$, and
$\epsilon\sim N(0,1)$.

Results Tables 1-4 summarize our results and compare our method to the
baselines in terms of ATT estimation performance under Setting 2. From these
tables, we draw the following observations and findings.

  * •

Overall, compared to other covariate balance-based and propensity score-based
weighting methods, our distribution learning-based weighting estimator
performs the best, both with and without the doubly-robust framework. It also
outperforms other frequently used state-of-the-art methods that are not purely
based on weighting. Specifically, our method can achieve up to approximately
81% RMSE reduction over the best baseline, which answers Q1. Furthermore, it
is noteworthy that doubly-robust methods outperform outcome modeling-based
methods in most cases, which explains the positive effect of combining outcome
modeling with a weighting component.

  * •

Statistical weighting estimators based on covariate balance, i.e.,
$\widehat{ATT}_{ebal}$, $\widehat{ATT}_{cpbs}$, $\widehat{ATT}_{opt}$ and
$\widehat{ATT}_{ebcw}$, all perform poorly. Their performance is similar to
that of $\widehat{ATT}_{base}$, which means they are unable to address the
confounding problem. This is not surprising because the basic assumptions are
violated under non-linear and non-additive settings. In contrast, our method
balances the distribution of covariates rather than just the moments, and thus
it does not rely on strong assumptions. In the doubly-robust framework, these
results also hold when comparing $\widehat{ATT}_{DR\\_*\\_ranf}$ (or
$\widehat{ATT}_{DR\\_*\\_bart}$) with $\widehat{ATT}_{ranf}$ (or
$\widehat{ATT}_{bart}$), where ‘*’ represents the weighting component.

  * •

As expected, traditional propensity score weighting based on logistic
regression, denoted as $\widehat{ATT}_{ps}$, fails due to the issue of model
mis-specification. Machine learning-based propensity score weighting
estimators, such as $\widehat{ATT}_{RF}$, $\widehat{ATT}_{BCART}$ and
$\widehat{ATT}_{XGBoost}$, achieve performance that is better than
$\widehat{ATT}_{ps}$ but apparently worse than our method,
$\widehat{ATT}_{DLW}$. The reason for this performance difference is that our
method calculates weights based on the learned distribution of covariates
rather than on the predicted propensity score as these machine learning
methods do. These comparative results of weighting methods also hold
consistently under the doubly-robust framework.

  * •

Furthermore, our weighting estimator can be improved when the population size
increases from 5,000 to 10,000, indicating a good asymptotic property. In
addition, we find that the estimation error increases when (i) more
confounders exist or (ii) confounding becomes stronger.

  * •

While the results reported here are for Setting 2, note that we obtain similar
results for Setting 3 (strongly non-linear and non-additive). For Setting 1
(linear and additive), which is not even practical, our distribution learning-
based weighting estimator also achieves robust and sufficiently good
performance, though not as good as the estimators specifically designed for
linear cases. Please refer to the Appendix for more details.

Based on these observations and in response to Q2, we can conclude that our
method provides consistent and robust performance across different settings,
and it outperforms several state-of-the-art weighting-based methods because it
is able to learn the underlying true distribution of covariates conditioned on
the treatment assignment. It is therefore able to alleviate the model mis-
specification issue. Note that the best performance is bold while the second
best is underlined.

Table 1: Performance comparison of our method ($\bm{\widehat{ATT}_{DLW}}$) with weighting-based methods. | d = 8 | d = 16
---|---|---
| N = 5,000 | N = 10,000 | N = 5,000 | N = 10,000
$s_{c}$ | Estimator | Bias | RMSE | Bias | RMSE | Bias | RMSE | Bias | RMSE
0.2 | $\widehat{ATT}_{base}$ | -2.674 | 2.677 | -2.673 | 2.674 | -6.073 | 6.076 | -6.057 | 6.058
$\widehat{ATT}_{ebal}$ | -2.674 | 2.677 | -2.673 | 2.674 | -6.074 | 6.077 | -6.059 | 6.060
$\widehat{ATT}_{cbps}$ | -2.674 | 2.676 | -2.673 | 2.674 | -6.074 | 6.077 | -6.059 | 6.060
$\widehat{ATT}_{opt}$ | -2.674 | 2.676 | -2.673 | 2.674 | -6.075 | 6.078 | -6.058 | 6.060
$\widehat{ATT}_{ebcw}$ | -2.674 | 2.677 | -2.673 | 2.674 | -6.074 | 6.077 | -6.059 | 6.060
$\widehat{ATT}_{ps}$ | -2.675 | 2.678 | -2.673 | 2.674 | -6.075 | 6.078 | -6.059 | 6.061
$\widehat{ATT}_{RF}$ | -1.575 | 1.576 | -1.469 | 1.470 | -4.999 | 5.001 | -4.842 | 4.843
$\widehat{ATT}_{BCART}$ | -1.456 | 1.458 | -1.324 | 1.325 | -4.866 | 4.868 | -4.663 | 4.665
$\widehat{ATT}_{XGBoost}$ | -0.587 | 0.590 | -0.427 | 0.430 | -2.359 | 2.361 | -2.148 | 2.149
$\widehat{ATT}_{Lasso}$ | -2.675 | 2.677 | -2.673 | 2.674 | -6.073 | 6.077 | -6.058 | 6.059
$\bm{\widehat{ATT}_{DLW}}$ | -0.132 | 0.195 | -0.124 | 0.177 | -0.679 | 1.197 | -0.347 | 0.416
0.4 | $\widehat{ATT}_{base}$ | -3.744 | 3.746 | -3.744 | 3.745 | -7.767 | 7.769 | -7.761 | 7.763
$\widehat{ATT}_{ebal}$ | -3.745 | 3.746 | -3.743 | 3.744 | -7.767 | 7.769 | -7.761 | 7.762
$\widehat{ATT}_{cbps}$ | -3.745 | 3.746 | -3.743 | 3.744 | -7.766 | 7.769 | -7.761 | 7.762
$\widehat{ATT}_{opt}$ | -3.745 | 3.747 | -3.743 | 3.745 | -7.769 | 7.771 | -7.761 | 7.763
$\widehat{ATT}_{ebcw}$ | -3.745 | 3.746 | -3.743 | 3.744 | -7.767 | 7.769 | -7.761 | 7.762
$\widehat{ATT}_{ps}$ | -3.745 | 3.747 | -3.744 | 3.745 | -7.769 | 7.771 | -7.761 | 7.763
$\widehat{ATT}_{RF}$ | -2.369 | 2.370 | -2.217 | 2.218 | -6.720 | 6.721 | -6.574 | 6.576
$\widehat{ATT}_{BCART}$ | -2.216 | 2.218 | -2.042 | 2.042 | -6.609 | 6.611 | -6.414 | 6.416
$\widehat{ATT}_{XGBoost}$ | -0.980 | 0.982 | -0.845 | 0.847 | -3.907 | 3.908 | -3.563 | 3.565
$\widehat{ATT}_{Lasso}$ | -3.745 | 3.746 | -3.744 | 3.745 | -7.767 | 7.769 | -7.761 | 7.763
$\bm{\widehat{ATT}_{DLW}}$ | -0.198 | 0.221 | -0.177 | 0.188 | -0.843 | 1.396 | -0.398 | 0.690

Table 2: Performance comparison of our method ($\bm{\widehat{ATT}_{DR\\_DLW\\_ranf}}$) with weighting-based methods under the doubly-robust framework where random forest is used for the outcome function component. | d = 8 | d = 16
---|---|---
| N = 5,000 | N = 10,000 | N = 5,000 | N = 10,000
$s_{c}$ | Estimator | Bias | RMSE | Bias | RMSE | Bias | RMSE | Bias | RMSE
0.2 | $\widehat{ATT}_{base}$ | -2.674 | 2.677 | -2.673 | 2.674 | -6.073 | 6.076 | -6.057 | 6.058
$\widehat{ATT}_{ranf}$ | -1.426 | 1.427 | -1.265 | 1.265 | -4.863 | 4.865 | -4.638 | 4.639
$\widehat{ATT}_{DR\\_ebal\\_ranf}$ | -1.428 | 1.429 | -1.270 | 1.271 | -4.812 | 4.814 | -4.593 | 4.594
$\widehat{ATT}_{DR\\_cbps\\_ranf}$ | -1.428 | 1.429 | -1.270 | 1.271 | -4.812 | 4.814 | -4.593 | 4.594
$\widehat{ATT}_{DR\\_opt\\_ranf}$ | -1.428 | 1.429 | -1.270 | 1.271 | -4.813 | 4.815 | -4.594 | 4.595
$\widehat{ATT}_{DR\\_ebcw\\_ranf}$ | -1.428 | 1.429 | -1.270 | 1.271 | -4.812 | 4.814 | -4.593 | 4.594
$\widehat{ATT}_{DR\\_ps\\_ranf}$ | -1.428 | 1.429 | -1.270 | 1.271 | -4.813 | 4.814 | -4.594 | 4.595
$\widehat{ATT}_{DR\\_RF\\_ranf}$ | -1.167 | 1.168 | -1.014 | 1.014 | -4.482 | 4.484 | -4.230 | 4.231
$\widehat{ATT}_{DR\\_BCART\\_ranf}$ | -1.139 | 1.140 | -0.987 | 0.987 | -4.448 | 4.449 | -4.190 | 4.191
$\widehat{ATT}_{DR\\_XGBoost\\_ranf}$ | -0.948 | 0.950 | -0.832 | 0.833 | -3.584 | 3.586 | -3.413 | 3.414
$\widehat{ATT}_{DR\\_Lasso\\_ranf}$ | -1.428 | 1.429 | -1.270 | 1.271 | -4.813 | 4.815 | -4.594 | 4.595
$\bm{\widehat{ATT}_{DR\\_DLW\\_ranf}}$ | -0.828 | 0.830 | -0.735 | 0.735 | -2.962 | 2.988 | -2.817 | 2.819
0.4 | $\widehat{ATT}_{base}$ | -3.744 | 3.746 | -3.744 | 3.745 | -7.767 | 7.769 | -7.761 | 7.763
$\widehat{ATT}_{ranf}$ | -2.142 | 2.143 | -1.921 | 1.921 | -6.528 | 6.529 | -6.286 | 6.287
$\widehat{ATT}_{DR\\_ebal\\_ranf}$ | -2.140 | 2.141 | -1.923 | 1.923 | -6.465 | 6.466 | -6.233 | 6.234
$\widehat{ATT}_{DR\\_cbps\\_ranf}$ | -2.140 | 2.141 | -1.923 | 1.923 | -6.465 | 6.466 | -6.233 | 6.234
$\widehat{ATT}_{DR\\_opt\\_ranf}$ | -2.140 | 2.141 | -1.923 | 1.923 | -6.466 | 6.468 | -6.234 | 6.235
$\widehat{ATT}_{DR\\_ebcw\\_ranf}$ | -2.140 | 2.141 | -1.923 | 1.923 | -6.465 | 6.466 | -6.233 | 6.234
$\widehat{ATT}_{DR\\_ps\\_ranf}$ | -2.140 | 2.141 | -1.923 | 1.923 | -6.466 | 6.467 | -6.233 | 6.234
$\widehat{ATT}_{DR\\_RF\\_ranf}$ | -1.785 | 1.786 | -1.566 | 1.567 | -6.124 | 6.125 | -5.856 | 5.857
$\widehat{ATT}_{DR\\_BCART\\_ranf}$ | -1.746 | 1.747 | -1.527 | 1.528 | -6.097 | 6.098 | -5.821 | 5.822
$\widehat{ATT}_{DR\\_XGBoost\\_ranf}$ | -1.421 | 1.422 | -1.267 | 1.268 | -5.102 | 5.103 | -4.837 | 4.838
$\widehat{ATT}_{DR\\_Lasso\\_ranf}$ | -2.140 | 2.141 | -1.923 | 1.923 | -6.466 | 6.467 | -6.234 | 6.235
$\bm{\widehat{ATT}_{DR\\_DLW\\_ranf}}$ | -1.237 | 1.238 | -1.112 | 1.112 | -3.946 | 3.988 | -3.746 | 3.753
Table 3: Performance comparison of our method ($\bm{\widehat{ATT}_{DR\\_DLW\\_bart}}$) with weighting-based methods under the doubly-robust framework where BART is used for the outcome function component. | d = 8 | d = 16
---|---|---
| N = 5,000 | N = 10,000 | N = 5,000 | N = 10,000
$s_{c}$ | Estimator | Bias | RMSE | Bias | RMSE | Bias | RMSE | Bias | RMSE
0.2 | $\widehat{ATT}_{base}$ | -2.674 | 2.677 | -2.673 | 2.674 | -6.073 | 6.076 | -6.057 | 6.058
$\widehat{ATT}_{bart}$ | -0.128 | 0.136 | -0.085 | 0.089 | -1.409 | 1.415 | -0.926 | 0.929
$\widehat{ATT}_{DR\\_ebal\\_bart}$ | -0.128 | 0.136 | -0.085 | 0.089 | -1.410 | 1.415 | -0.926 | 0.929
$\widehat{ATT}_{DR\\_cbps\\_bart}$ | -0.128 | 0.136 | -0.085 | 0.089 | -1.410 | 1.415 | -0.926 | 0.929
$\widehat{ATT}_{DR\\_opt\\_bart}$ | -0.128 | 0.136 | -0.085 | 0.089 | -1.410 | 1.416 | -0.926 | 0.929
$\widehat{ATT}_{DR\\_ebcw\\_bart}$ | -0.128 | 0.136 | -0.085 | 0.089 | -1.410 | 1.415 | -0.926 | 0.929
$\widehat{ATT}_{DR\\_ps\\_bart}$ | -0.128 | 0.136 | -0.085 | 0.089 | -1.410 | 1.415 | -0.926 | 0.929
$\widehat{ATT}_{DR\\_RF\\_bart}$ | -0.099 | 0.109 | -0.060 | 0.065 | -1.302 | 1.308 | -0.824 | 0.828
$\widehat{ATT}_{DR\\_BCART\\_bart}$ | -0.096 | 0.107 | -0.059 | 0.066 | -1.289 | 1.295 | -0.815 | 0.818
$\widehat{ATT}_{DR\\_XGBoost\\_bart}$ | -0.081 | 0.094 | -0.054 | 0.060 | -1.029 | 1.036 | -0.614 | 0.618
$\widehat{ATT}_{DR\\_Lasso\\_bart}$ | -0.128 | 0.136 | -0.085 | 0.089 | -1.410 | 1.416 | -0.926 | 0.929
$\bm{\widehat{ATT}_{DR\\_DLW\\_bart}}$ | -0.062 | 0.085 | -0.034 | 0.043 | -0.834 | 0.919 | -0.439 | 0.448
0.4 | $\widehat{ATT}_{base}$ | -3.744 | 3.746 | -3.744 | 3.745 | -7.767 | 7.769 | -7.761 | 7.763
$\widehat{ATT}_{bart}$ | -0.172 | 0.180 | -0.105 | 0.112 | -2.148 | 2.153 | -1.388 | 1.392
$\widehat{ATT}_{DR\\_ebal\\_bart}$ | -0.172 | 0.180 | -0.105 | 0.112 | -2.148 | 2.154 | -1.388 | 1.392
$\widehat{ATT}_{DR\\_cbps\\_bart}$ | -0.172 | 0.180 | -0.105 | 0.112 | -2.148 | 2.154 | -1.388 | 1.392
$\widehat{ATT}_{DR\\_opt\\_bart}$ | -0.172 | 0.180 | -0.105 | 0.112 | -2.149 | 2.154 | -1.388 | 1.392
$\widehat{ATT}_{DR\\_ebcw\\_bart}$ | -0.172 | 0.180 | -0.105 | 0.112 | -2.148 | 2.154 | -1.388 | 1.392
$\widehat{ATT}_{DR\\_ps\\_bart}$ | -0.172 | 0.180 | -0.105 | 0.112 | -2.148 | 2.154 | -1.388 | 1.392
$\widehat{ATT}_{DR\\_RF\\_bart}$ | -0.133 | 0.143 | -0.070 | 0.079 | -2.020 | 2.026 | -1.265 | 1.269
$\widehat{ATT}_{DR\\_BCART\\_bart}$ | -0.127 | 0.139 | -0.068 | 0.078 | -2.002 | 2.008 | -1.245 | 1.249
$\widehat{ATT}_{DR\\_XGBoost\\_bart}$ | -0.105 | 0.118 | -0.061 | 0.072 | -1.695 | 1.701 | -0.984 | 0.989
$\widehat{ATT}_{DR\\_Lasso\\_bart}$ | -0.172 | 0.180 | -0.105 | 0.112 | -2.149 | 2.155 | -1.388 | 1.392
$\bm{\widehat{ATT}_{DR\\_DLW\\_bart}}$ | -0.085 | 0.102 | -0.038 | 0.054 | -1.202 | 1.242 | -0.611 | 0.640
Table 4: Performance comparison of our methods ($\bm{\widehat{ATT}_{DLW}}$ and $\bm{\widehat{ATT}_{DR\\_DLW\\_bart}}$) with frequently used methods. | d = 8 | d = 16
---|---|---
| N = 5,000 | N = 10,000 | N = 5,000 | N = 10,000
$s_{c}$ | Estimator | Bias | RMSE | Bias | RMSE | Bias | RMSE | Bias | RMSE
0.2 | $\widehat{ATT}_{base}$ | -2.674 | 2.677 | -2.673 | 2.674 | -6.073 | 6.076 | -6.057 | 6.058
$\widehat{ATT}_{OLS}$ | -2.675 | 2.677 | -2.674 | 2.675 | -6.075 | 6.078 | -6.058 | 6.060
$\widehat{ATT}_{PSM}$ | -2.542 | 2.559 | -2.567 | 2.577 | -6.040 | 6.046 | -6.014 | 6.017
$\widehat{ATT}_{ARB}$ | -2.674 | 2.676 | -2.673 | 2.674 | -6.076 | 6.079 | -6.059 | 6.061
$\widehat{ATT}_{CF}$ | -1.242 | 1.243 | -1.007 | 1.007 | -4.903 | 4.905 | -4.597 | 4.598
$\widehat{ATT}_{CEVAE}$ | -0.576 | 0.588 | -0.282 | 0.291 | -2.619 | 2.766 | -2.145 | 2.154
$\bm{\widehat{ATT}_{DLW}}$ | -0.132 | 0.195 | -0.124 | 0.177 | -0.679 | 1.197 | -0.347 | 0.416
$\bm{\widehat{ATT}_{DR\\_DLW\\_bart}}$ | -0.062 | 0.085 | -0.034 | 0.044 | -0.834 | 0.919 | -0.926 | 0.929
0.4 | $\widehat{ATT}_{base}$ | -3.744 | 3.746 | -3.744 | 3.745 | -7.767 | 7.769 | -7.761 | 7.763
$\widehat{ATT}_{OLS}$ | -3.745 | 3.747 | -3.744 | 3.745 | -7.768 | 7.770 | -7.762 | 7.763
$\widehat{ATT}_{PSM}$ | -3.657 | 3.666 | -3.610 | 3.624 | -7.746 | 7.751 | -7.691 | 7.694
$\widehat{ATT}_{ARB}$ | -3.746 | 3.747 | -3.744 | 3.745 | -7.769 | 7.772 | -7.762 | 7.763
$\widehat{ATT}_{CF}$ | -1.786 | 1.787 | -1.446 | 1.446 | -6.453 | 6.454 | -6.114 | 6.115
$\widehat{ATT}_{CEVAE}$ | -0.824 | 0.828 | -0.456 | 0.459 | -3.318 | 3.324 | -3.564 | 3.568
$\bm{\widehat{ATT}_{DLW}}$ | -0.198 | 0.221 | -0.177 | 0.188 | -0.843 | 1.396 | -0.398 | 0.690
$\bm{\widehat{ATT}_{DR\\_DLW\\_bart}}$ | -0.085 | 0.102 | -0.038 | 0.054 | -1.202 | 1.242 | -0.611 | 0.640

The effect of parameters. As Tables 1-4 show, parameters do affect the quality
of ATT estimation. For example, larger observational data can reduce the
estimation error. Stronger confounding and increased dimensionality of data
can make ATT estimation more challenging. To further understand the
relationship between parameters and the performance, we conduct several
additional experiments in which we vary the sample size from 5,000 to 100,000
and use various model settings, confounding strengths and dimensionality of
covariates. Figure 2 shows that (i) the ATT estimation performance is not
determined by a single parameter but is jointly affected by all of them and
(ii) increasing the sample size can improve ATT estimation to a sufficient
level across all scenarios.

Figure 2: The impact of parameters on the performance of ATT estimation.

Model convergence. In theory, our model converges to the underlying true
distributions for both treatment and control groups as long as we have enough
transformations. In practice, we usually have limited observational data and
only make a few transformations, which leads to an obvious question: does the
model converge? To answer this question, we empirically calculate the negative
log-likelihood (NLL) of covariates in both groups in every iteration. From
Figure 3, we can see that our model does converge within a small number of
iterations, which implies that the estimated distribution approaches the
truth. This finding responds to Q3 (We note that this is for Setting 2 and
$d=8$. Please refer to Appendix for other settings). Blue and red denote
training and validation, respectively.

Figure 3: Convergence of NNL.

###  4.4 Experiments on real data

We now turn to evaluating our model using real data. In this section, we apply
our distribution learning-based weighting approach to a real dataset for ATT
estimation and compare the results with baselines.

Dataset We use the Twins dataset, introduced by [14]. It has 46 covariates and
71,345 pairs of twins. The treatment is being born the lighter twin, whereas
the outcome corresponds to the mortality rate of each of the twins in their
first year of life. Following previous studies [14, 46], we choose twins
weighing less than 2kg and having the same sex. In order to select appropriate
covariates to introduce selection bias, we first choose the top 10 covariates
that have high correlation coefficients with $Y^{0}$, and then we run an OLS
regression of $Y^{0}$ on these 10 covariates to obtain our final set of
covariates that have significant regression coefficients. After this process,
seven covariates remain for each twin pair. They are related to the number of
gestation weeks prior to birth, quality of care, child’s sex, parent marital
status, and some risk factors. Thus, our final dataset contains 11,984 pairs
of twins with seven covariates.

In order to simulate an observational study, we generate the binary treatment
variable conditioned on the seven covariates. Of these covariates, we
primarily use the variable ‘GESTAT10’ (the number of gestation weeks prior to
birth) to introduce selection bias, for two reasons: (1) it is highly
correlated with the outcome and (2) textcolorredmany prior studies show that
both extremely long and extremely short gestation period cause a
baby333https://www.stanfordchildrens.org/en/topic/default?id=postmaturity-in-
the-newborn-90-P02399. to lose weight, which can affect the mortality rate.
Thus, the treatment variable is generated as

$\displaystyle W\sim Bernoulli(e(z,x)),$ (19) $\displaystyle
Conf\\_term=log(1+z^{2})+0.01*\sum_{x\neq z}x,$ $e(z,x)=p(W=1\mid
z,x)=\begin{cases}0.1,&if~{}Conf\\_term>median(Conf\\_term)\\\
0.9,&if~{}Conf\\_term\leq median(Conf\\_term)\\\ \end{cases}$

where $z$ is GESTAT10 and $x$ represents other covariates.

Table 5: Performance comparison of our method ($\bm{\widehat{ATT}_{DLW}}$) with weighting-based methods in Twins dataset. Estimator | Bias | RMSE
---|---|---
$\widehat{ATT}_{base}$ | -0.108 | 0.108
$\widehat{ATT}_{ebal}$ | -0.946 | 0.946
$\widehat{ATT}_{cbps}$ | -0.101 | 0.102
$\widehat{ATT}_{opt}$ | -0.103 | 0.103
$\widehat{ATT}_{ebcw}$ | -0.102 | 0.102
$\widehat{ATT}_{ps}$ | -0.104 | 0.104
$\widehat{ATT}_{RF}$ | -0.050 | 0.050
$\widehat{ATT}_{BCART}$ | -0.042 | 0.043
$\widehat{ATT}_{XFBoost}$ | -0.039 | 0.041
$\widehat{ATT}_{Lasso}$ | -0.104 | 0.104
$\bm{\widehat{ATT}_{DLW}}$ | 0.000 | 0.011
Table 6: Performance comparison of our method ($\bm{\widehat{ATT}_{DR\\_DLW\\_ranf}}$) with weighting-based methods under the doubly-robust framework where random forest is used for the outcome function component in Twins dataset. Estimator | Bias | RMSE
---|---|---
$\widehat{ATT}_{base}$ | -0.108 | 0.108
$\widehat{ATT}_{ranf}$ | -0.013 | 0.016
$\widehat{ATT}_{DR\\_ebal\\_ranf}$ | -0.010 | 0.014
$\widehat{ATT}_{DR\\_cbps\\_ranf}$ | -0.010 | 0.014
$\widehat{ATT}_{DR\\_opt\\_ranf}$ | -0.010 | 0.014
$\widehat{ATT}_{DR\\_ebcw\\_ranf}$ | -0.010 | 0.014
$\widehat{ATT}_{DR\\_ps\\_ranf}$ | -0.010 | 0.014
$\widehat{ATT}_{DR\\_RF\\_ranf}$ | -0.008 | 0.012
$\widehat{ATT}_{DR\\_BCART\\_ranf}$ | -0.009 | 0.013
$\widehat{ATT}_{DR\\_XGBoost\\_ranf}$ | -0.010 | 0.014
$\widehat{ATT}_{DR\\_Lasso\\_ranf}$ | -0.010 | 0.014
$\bm{\widehat{ATT}_{DR\\_DLW\\_ranf}}$ | -0.007 | 0.012
Table 7: Performance comparison of our method ($\bm{\widehat{ATT}_{DR\\_DLW\\_bart}}$) with weighting-based methods under the doubly-robust framework where BART is used for the outcome function component in Twins dataset. Estimator | Bias | RMSE
---|---|---
$\widehat{ATT}_{base}$ | -0.108 | 0.108
$\widehat{ATT}_{bart}$ | -0.010 | 0.013
$\widehat{ATT}_{DR\\_ebal\\_bart}$ | -0.010 | 0.013
$\widehat{ATT}_{DR\\_cbps\\_bart}$ | -0.010 | 0.013
$\widehat{ATT}_{DR\\_opt\\_bart}$ | -0.010 | 0.013
$\widehat{ATT}_{DR\\_ebcw\\_bart}$ | -0.010 | 0.013
$\widehat{ATT}_{DR\\_ps\\_bart}$ | -0.010 | 0.013
$\widehat{ATT}_{DR\\_RF\\_bart}$ | -0.005 | 0.010
$\widehat{ATT}_{DR\\_BCART\\_bart}$ | -0.005 | 0.012
$\widehat{ATT}_{DR\\_XGBoost\\_bart}$ | -0.003 | 0.011
$\widehat{ATT}_{DR\\_lasso\\_bart}$ | -0.010 | 0.013
$\bm{\widehat{ATT}_{DR\\_DLW\\_bart}}$ | 0.000 | 0.010

Results Similar to our handling of simulation studies, we compare the
performance of our approach with baselines on two evaluation metrics regarding
bias and RMSE. The results in Tables 5-8 tell us that:

  * •

Our distribution learning-based weighting estimator performs the best and it
is very close to the true ATT, as reflected by the near zero bias and small
RMSE in Table 5.

  * •

Similar to what we see in the results of our experiments on synthetic data in
Section 4.3, statistical weighting estimators based on covariate balance
cannot address the confounding issue, as the strong assumptions do not hold in
real data. Machine learning-based propensity score weighting estimators, such
as $\widehat{ATT}_{RF}$, $\widehat{ATT}_{BCART}$ and
$\widehat{ATT}_{XGBoost}$, outperform $\widehat{ATT}_{ps}$ but perform
significantly worse than our method, $\widehat{ATT}_{DLW}$. These results also
hold when we plug the weighting component into the doubly-robust framework.

  * •

Our proposed method $\widehat{ATT}_{DLW}$ also exhibits performance superior
to other frequently used methods in statistics and machine learning. Note that
the most widely used approaches, OLS regression and propensity score matching,
do not improve either bias or RMSE over the base when used on real data, just
as they do not in the experiments on synthetic data for Setting 2 and Setting
3.

Table 8: Performance comparison of our methods ($\bm{\widehat{ATT}_{DLW}}$, $\bm{\widehat{ATT}_{DR\\_DLW\\_bart}}$) with frequently used methods in Twins dataset. Estimator | Bias | RMSE
---|---|---
$\widehat{ATT}_{base}$ | -0.108 | 0.108
$\widehat{ATT}_{OLS}$ | -0.103 | 0.103
$\widehat{ATT}_{PSM}$ | -0.102 | 0.102
$\widehat{ATT}_{ARB}$ | -0.102 | 0.102
$\widehat{ATT}_{CF}$ | -0.002 | 0.011
$\widehat{ATT}_{CEVAE}$ | -0.002 | 0.012
$\bm{\widehat{ATT}_{DLW}}$ | 0.000 | 0.011
$\bm{\widehat{ATT}_{DR\\_DLW\\_bart}}$ | 0.000 | 0.010

Model convergence As in the simulation studies in Section 4.3, the estimated
distributions of covariates in both treatment and control groups converge.
This is shown in Figure 4, where the negative log-likelihood (NLL) of
covariates stabilizes after certain number of iterations. Blue and red denote
training and validation, respectively.

Figure 4: Convergence of NNL in Twins data.

##  5 Discussions and Conclusion

Model mis-specification is a major issue in causal inference using
observational data. Prior studies have attempted to develop various propensity
score- or covariate balance-based weighting methods, yet they must still
impose many strong assumptions to obtain unbiased estimators. To fill in this
research gap, we present a new weighting method via distribution learning for
treatment effect estimation. In this paper, we approximate probability density
of covariates conditioned on treatment assignment via transformations through
change of variables. Specifically, we implement a series of neural
autoregressive transformations and leverage the ratio of covariates’ density
in the treatment group to that of the control group as the weight for
estimating treatment effects. Extensive experiments on synthetic data and real
data demonstrate the effectiveness of our proposed method.

Our study contributes to the literature and to the causal inference practices
in the following ways. First, substantively, we propose a weighting method via
distribution learning to address the model mis-specification issue for
treatment effect estimation with observational data. Theoretically, we show
that our method can achieve unbiased estimation, similar to IPTW when the
propensity score of treatment assignment is close to the truth. We believe
that this provides a new path for studying causal inference. Second,
methodologically, we employ various advanced methodologies and tools that
provide new angles of approach in studying treatment effect estimation,
particularly when it comes to leveraging the power of deep neural networks and
statistics. To demonstrate the robustness and generalizability of our method,
we not only conduct extensive simulation studies under different parameter and
model specification settings, but also conduct experiments on a real-world
dataset. Finally, managerially, our study provides actionable and practical
guidance for managers who are seeking methods for understanding causal effects
of certain events, particularly when a randomized field experiment is not
possible. It implements a data-driven weighting-based ATT estimator, which can
help scholars better sharpen the collective toolkit and harness the power of
machine learning methods in their causal inference studies. The results
indicate that our method DLW is a powerful analytical tool to tackle the model
mis-specification for ATT estimation. In addition, DLW is also efficient and
useful in handling both high and low dimensional data.

This study is not without limitations. First, our method is not designed to
address the issue of omitted variables, which commonly arises in causal
inference in observational studies since almost all such studies are based on
the ignorability assumption. We presuppose that the covariates in our model
are the most important and relevant ones and can fully determine group
assignment that is completely accounted for by the propensity score. However,
in reality some crucial variables are likely to be omitted from the model. In
this case the groups may remain unbalanced, which can lead to seriously biased
results. Second, although our method can be generalized to high-dimensional
cases as we see in our experiments, larger data is necessary for good density
approximation. Third, the neural autoregressive model is very powerful and
capable of dealing with continuous cases. Thus, developing new distribution
learning algorithms that can handle a mixture of continuous and categorical
covariates is always worth pursuing in future research.

## References

  * [1] E. Stuart.  Matching methods for causal inference: A review and a look forward.  Statistical science : a review journal of the Institute of Mathematical Statistics, 25 1:1–21, 2010.
  * [2] PAUL Rosenbaum and Donald Rubin.  The central role of the propensity score in observational studies for causal effects.  Biometrika, 70:41–55, 04 1983.
  * [3] Jared Lunceford and Marie Davidian.  Stratification and weighting via the propensity score in estimation of causal treatment effects: A comparative study.  Statistics in medicine, 23:2937–60, 10 2004.
  * [4] Guido Imbens.  Nonparametric estimation of average treatment effects under exogeneity: A review.  Review of Economics and Statistics, 2004.
  * [5] C. Drake.  Effects of misspecification of the propensity score on estimators of treatment effect.  Biometrics, 49:1231–1236, 1993.
  * [6] Jens Hainmueller.  Entropy balancing for causal effects: A multivariate reweighting method to produce balanced samples in observational studies.  Political Analysis, 20(1):25–46, 2012.
  * [7] José Zubizarreta.  Stable weights that balance covariates for estimation with incomplete outcome data.  Journal of the American Statistical Association, 110:0–0, 04 2015\.
  * [8] Kwun Chan, Sheung Yam, and Zhang Zheng.  Globally efficient nonparametric inference of average treatment effects by empirical balancing calibration weighting.  Journal of the Royal Statistical Society Series B (Statistical Methodology), 78, 04 2015.
  * [9] Kosuke Imai and Marc Ratkovic.  Covariate balancing propensity score.  Journal of the Royal Statistical Society: Series B (Statistical Methodology), 76(1):243–263, 2014.
  * [10] Fan Li, Kari Lock Morgan, and Alan M. Zaslavsky.  Balancing covariates via propensity score weighting.  Journal of the American Statistical Association, 113(521):390–400, 2018.
  * [11] Susan Athey, Guido W. Imbens, and Stefan Wager.  Approximate residual balancing: debiased inference of average treatment effects in high dimensions.  Journal of the Royal Statistical Society: Series B (Statistical Methodology), 80(4):597–623, 2018.
  * [12] E. G. Tabak and Cristina V. Turner.  A family of nonparametric density estimation algorithms.  Communications on Pure and Applied Mathematics, 66(2):145–164, 2013\.
  * [13] George Papamakarios, Eric T. Nalisnick, Danilo Jimenez Rezende, Shakir Mohamed, and Balaji Lakshminarayanan.  Normalizing flows for probabilistic modeling and inference.  ArXiv, abs/1912.02762, 2019.
  * [14] Christos Louizos, Uri Shalit, Joris Mooij, David Sontag, Richard Zemel, and Max Welling.  Causal effect inference with deep latent-variable models.  In Proceedings of the 31st International Conference on Neural Information Processing Systems, NIPS’17, page 6449–6459, Red Hook, NY, USA, 2017. Curran Associates Inc.
  * [15] James Robins, A G Rotnitzky, and Lue Zhao.  Estimation of regression coefficients when some regressors are not always observed.  Journal of The American Statistical Association - J AMER STATIST ASSN, 89:846–866, 09 1994.
  * [16] Victor Chernozhukov, Denis Chetverikov, Mert Demirer, Esther Duflo, Christian Hansen, and Whitney Newey.  Double/debiased/neyman machine learning of treatment effects.  American Economic Review, 107(5):261–65, May 2017.
  * [17] Stefan Wager and Susan Athey.  Estimation and inference of heterogeneous treatment effects using random forests.  Journal of the American Statistical Association, 113(523):1228–1242, 2018.
  * [18] J. Czajka, Sharon M. Hirabayashi, R. Little, and D. Rubin.  Projecting from advance data using propensity modeling: An application to income and tax statistics.  Journal of Business & Economic Statistics, 10:117–131, 1992.
  * [19] James Robins, Miguel Hernán, and Babette Brumback.  Marginal structural models and causal inference in epidemiology.  Epidemiology (Cambridge, Mass.), 11:550–60, 10 2000.
  * [20] Keisuke Hirano, Guido Imbens, and Geert Ridder.  Efficient estimation of average treatment effects using the estimated propensity score.  Econometrica, 71:1161–1189, 02 2003.
  * [21] Brian Lee, Justin Lessler, and Elizabeth Stuart.  Improving propensity score weighting using machine learning.  Statistics in medicine, 29:337–46, 11 2009.
  * [22] Max H. Farrell.  Robust inference on average treatment effects with possibly more covariates than observations.  Journal of Econometrics, 189(1):1 – 23, 2015.
  * [23] Daniel Mccaffrey, Greg Ridgeway, and Andrew Morral.  Propensity score estimation with boosted regression for evaluating causal effects in observational studies.  Psychological methods, 9:403–25, 01 2005.
  * [24] Funk MJ. Westreich D, Lessler J.  Propensity score estimation: neural networks, support vector machines, decision trees (cart), and meta-classifiers as alternatives to logistic regression.  J Clin Epidemiol., 63(8):826–33, 2010.
  * [25] Jennifer Hill.  Bayesian nonparametric modeling for causal inference.  Journal of Computational and Graphical Statistics, 20:217–240, 03 2011.
  * [26] Noemi Kreif and Karla DiazOrdaz.  Machine learning in policy evaluation: new tools for causal inference.  2019\.
  * [27] Susan Athey and Guido Imbens.  Recursive partitioning for heterogeneous causal effects: Table 1.  Proceedings of the National Academy of Sciences, 113:7353–7360, 07 2016.
  * [28] Inbal Yahav, Galit Shmueli, and Deepa Mani.  A tree-based approach for addressing self-selection in impact studies with big data.  MIS Quarterly, 40:819–848, 04 2016.
  * [29] Yixing Chen, Ju-Yeon Lee, Shrihari (Hari) Sridhar, Vikas Mittal, Katharine McCallister, and Amit G. Singal.  Improving cancer outreach effectiveness through targeting and economic assessments: Insights from a randomized field experiment.  Journal of Marketing, 84(3):1–27, 2020.
  * [30] Uri Shalit, Fredrik D. Johansson, and David Sontag.  Estimating individual treatment effect: generalization bounds and algorithms.  In Doina Precup and Yee Whye Teh, editors, Proceedings of the 34th International Conference on Machine Learning, volume 70 of Proceedings of Machine Learning Research, pages 3076–3085, International Convention Centre, Sydney, Australia, 06–11 Aug 2017. PMLR.
  * [31] Fredrik D. Johansson, Uri Shalit, and David Sontag.  Learning representations for counterfactual inference.  In Proceedings of the 33rd International Conference on International Conference on Machine Learning - Volume 48, ICML’16, page 3020–3029. JMLR.org, 2016.
  * [32] A. P. Dempster, N. M. Laird, and D. B. Rubin.  Maximum likelihood from incomplete data via the em algorithm.  JOURNAL OF THE ROYAL STATISTICAL SOCIETY, SERIES B, 39(1):1–38, 1977.
  * [33] Zudi Lu and Dag Tjøstheim.  Nonparametric estimation of probability density functions for irregularly observed spatial data.  Journal of the American Statistical Association, 109, 10 2014.
  * [34] Bharath Sriperumbudur, Kenji Fukumizu, Arthur Gretton, Aapo Hyvärinen, and Revant Kumar.  Density estimation in infinite dimensional exponential families.  Journal of Machine Learning Research, 18(57):1–59, 2017.
  * [35] V. Maroulas, J. Mike, and Christopher Oballe.  Nonparametric estimation of probability density functions of random persistence diagrams.  J. Mach. Learn. Res., 20:151:1–151:49, 2019.
  * [36] Heng Shu and Z. Tan.  Improved estimation of average treatment effects on the treated: Local efficiency, double robustness, and beyond.  arXiv: Methodology, 2018.
  * [37] Walter Rudin.  Real and Complex Analysis.  McGraw-Hill Science/Engineering/Math, May 1986.
  * [38] V.I. Bogachev.  Measure Theory.  Number v. 1 in Measure Theory. Springer Berlin Heidelberg, 2007.
  * [39] Mathieu Germain, Karol Gregor, Iain Murray, and Hugo Larochelle.  Made: Masked autoencoder for distribution estimation.  In Francis Bach and David Blei, editors, Proceedings of the 32nd International Conference on Machine Learning, volume 37 of Proceedings of Machine Learning Research, pages 881–889, Lille, France, 07–09 Jul 2015\. PMLR.
  * [40] Durk P Kingma, Tim Salimans, Rafal Jozefowicz, Xi Chen, Ilya Sutskever, and Max Welling.  Improved variational inference with inverse autoregressive flow.  In D. Lee, M. Sugiyama, U. Luxburg, I. Guyon, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 29, pages 4743–4751. Curran Associates, Inc., 2016.
  * [41] Chin-Wei Huang, David Krueger, Alexandre Lacoste, and Aaron Courville.  Neural autoregressive flows.  In Jennifer Dy and Andreas Krause, editors, Proceedings of the 35th International Conference on Machine Learning, volume 80 of Proceedings of Machine Learning Research, pages 2078–2087, StockholmsmÃ€ssan, Stockholm Sweden, 10–15 Jul 2018. PMLR.
  * [42] I. Kobyzev, S. Prince, and M. Brubaker.  Normalizing flows: An introduction and review of current methods.  IEEE Transactions on Pattern Analysis and Machine Intelligence, pages 1–1, 2020.
  * [43] Author Diamond.  Genetic matching for estimating causal effects: A general multivariate matching method for achieving balance in observational studies.  Review of Economics and Statistics - REV ECON STATIST, 0, 01 2005\.
  * [44] Jerome H. Friedman.  Multivariate adaptive regression splines.  Ann. Statist., 19(1):1–67, 03 1991.
  * [45] Younathan Abdia, K.B. Kulasekera, Somnath Datta, Maxwell Boakye, and Maiying Kong.  Propensity scores based methods for estimating average treatment effect and average treatment effect among treated: A comparative study: Propensity scores based methods for estimating average treatment.  Biometrical Journal, 59, 04 2017.
  * [46] Jinsung Yoon, J. Jordon, and M. V. D. Schaar.  Ganite: Estimation of individualized treatment effects using generative adversarial nets.  In ICLR, 2018.

## Appendix

### A1. Results of experiments on synthetic data for Setting 1

From Tables A1 - A4, we can clearly observe that our distribution learning-
based weighting estimator achieves robust and sufficiently good performance.
It significantly performs better than various propensity score-based
weighting, and machine learning-based methods such as $\widehat{ATT}_{CF}$ and
$\widehat{ATT}_{CEVAE}$. However, it does not outperform those covariate
balance-based weighting estimators and some frequently used methods such as
$\widehat{ATT}_{OLS}$ and $\widehat{ATT}_{ARB}$, whose estimation biases are
near zero. This is because (1) the strong assumptions of these estimators,
i.e., linearity and additivity, are designed to hold in Setting 1 and (2) the
model mis-specification issue does not exist anymore, which is not likely to
be true in practice and thus is not the focus of this study.

Table A1: Performance comparison of our method ($\bm{\widehat{ATT}_{DLW}}$) with weighting-based methods. | d = 8 | d = 16
---|---|---
| N = 5,000 | N = 10,000 | N = 5,000 | N = 10,000
$s_{c}$ | Estimator | Bias | RMSE | Bias | RMSE | Bias | RMSE | Bias | RMSE
0.2 | $\widehat{ATT}_{base}$ | -0.743 | 0.745 | -0.749 | 0.750 | -1.364 | 1.366 | -1.370 | 1.370
$\widehat{ATT}_{ebal}$ | -0.002 | 0.030 | 0.001 | 0.021 | -0.003 | 0.028 | 0.000 | 0.020
$\widehat{ATT}_{cbps}$ | -0.003 | 0.030 | -0.001 | 0.021 | -0.009 | 0.030 | -0.007 | 0.021
$\widehat{ATT}_{opt}$ | -0.041 | 0.050 | -0.039 | 0.044 | -0.081 | 0.085 | -0.077 | 0.080
$\widehat{ATT}_{ebcw}$ | -0.001 | 0.030 | 0.001 | 0.021 | -0.003 | 0.028 | 0.000 | 0.020
$\widehat{ATT}_{ps}$ | -0.003 | 0.035 | 0.001 | 0.025 | -0.008 | 0.048 | -0.002 | 0.036
$\widehat{ATT}_{RF}$ | -0.394 | 0.396 | -0.388 | 0.389 | -0.895 | 0.896 | -0.875 | 0.875
$\widehat{ATT}_{BCART}$ | -0.380 | 0.382 | -0.373 | 0.374 | -0.852 | 0.854 | -0.839 | 0.840
$\widehat{ATT}_{XGBoost}$ | -0.366 | 0.369 | -0.351 | 0.353 | -0.693 | 0.695 | -0.686 | 0.687
$\widehat{ATT}_{Lasso}$ | -0.148 | 0.189 | -0.161 | 0.205 | -0.218 | 0.273 | -0.213 | 0.288
$\bm{\widehat{ATT}_{DLW}}$ | -0.043 | 0.149 | -0.014 | 0.058 | -0.067 | 0.113 | -0.049 | 0.098
0.4 | $\widehat{ATT}_{base}$ | -1.281 | 1.282 | -1.288 | 1.288 | -2.136 | 2.137 | -2.142 | 2.143
$\widehat{ATT}_{ebal}$ | -0.002 | 0.038 | 0.001 | 0.023 | 0.003 | 0.050 | -0.009 | 0.032
$\widehat{ATT}_{cbps}$ | -0.012 | 0.044 | -0.005 | 0.024 | -0.113 | 0.137 | -0.069 | 0.083
$\widehat{ATT}_{opt}$ | -0.042 | 0.054 | -0.040 | 0.046 | -0.076 | 0.088 | -0.089 | 0.091
$\widehat{ATT}_{ebcw}$ | -0.001 | 0.038 | 0.001 | 0.023 | 0.003 | 0.050 | -0.009 | 0.032
$\widehat{ATT}_{ps}$ | 0.003 | 0.056 | 0.003 | 0.040 | -0.004 | 0.132 | -0.001 | 0.091
$\widehat{ATT}_{RF}$ | -0.725 | 0.727 | -0.709 | 0.709 | -1.537 | 1.538 | -1.506 | 1.506
$\widehat{ATT}_{BCART}$ | -0.685 | 0.686 | -0.672 | 0.673 | -1.488 | 1.489 | -1.452 | 1.452
$\widehat{ATT}_{XGBoost}$ | -0.639 | 0.641 | -0.619 | 0.620 | -1.083 | 1.084 | -1.090 | 1.090
$\widehat{ATT}_{Lasso}$ | -0.311 | 0.399 | -0.314 | 0.419 | -0.406 | 0.537 | -0.441 | 0.581
$\bm{\widehat{ATT}_{DLW}}$ | -0.053 | 0.103 | -0.042 | 0.076 | -0.131 | 0.258 | -0.139 | 0.200
Table A2: Performance comparison of our method ($\bm{\widehat{ATT}_{DR\\_DLW\\_ranf}}$) with weighting-based methods under the doubly-robust framework where random forest is used for the outcome function component. | d = 8 | d = 16
---|---|---
| N = 5,000 | N = 10,000 | N = 5,000 | N = 10,000
$s_{c}$ | Estimator | Bias | RMSE | Bias | RMSE | Bias | RMSE | Bias | RMSE
0.2 | $\widehat{ATT}_{base}$ | -0.743 | 0.745 | -0.749 | 0.750 | -1.364 | 1.366 | -1.370 | 1.370
$\widehat{ATT}_{ranf}$ | -0.206 | 0.209 | -0.178 | 0.180 | -0.633 | 0.634 | -0.580 | 0.581
$\widehat{ATT}_{DR\\_ebal\\_ranf}$ | -0.118 | 0.122 | -0.102 | 0.104 | -0.375 | 0.376 | -0.343 | 0.344
$\widehat{ATT}_{DR\\_cbps\\_ranf}$ | -0.118 | 0.122 | -0.102 | 0.105 | -0.376 | 0.377 | -0.345 | 0.345
$\widehat{ATT}_{DR\\_opt\\_ranf}$ | -0.124 | 0.128 | -0.108 | 0.110 | -0.393 | 0.394 | -0.359 | 0.360
$\widehat{ATT}_{DR\\_ebcw\\_ranf}$ | -0.118 | 0.122 | -0.102 | 0.104 | -0.375 | 0.376 | -0.343 | 0.344
$\widehat{ATT}_{DR\\_ps\\_ranf}$ | -0.118 | 0.122 | -0.102 | 0.105 | -0.377 | 0.378 | -0.344 | 0.345
$\widehat{ATT}_{DR\\_RF\\_ranf}$ | -0.166 | 0.169 | -0.143 | 0.145 | -0.546 | 0.547 | -0.496 | 0.497
$\widehat{ATT}_{DR\\_BCART\\_ranf}$ | -0.165 | 0.168 | -0.142 | 0.144 | -0.537 | 0.538 | -0.491 | 0.491
$\widehat{ATT}_{DR\\_XGBoost\\_ranf}$ | -0.162 | 0.165 | -0.140 | 0.142 | -0.506 | 0.507 | -0.464 | 0.465
$\widehat{ATT}_{DR\\_Lasso\\_ranf}$ | -0.136 | 0.140 | -0.119 | 0.121 | -0.417 | 0.419 | -0.381 | 0.383
$\bm{\widehat{ATT}_{DR\\_DLW\\_ranf}}$ | -0.126 | 0.134 | -0.102 | 0.105 | -0.388 | 0.390 | -0.351 | 0.352
0.4 | $\widehat{ATT}_{base}$ | -1.281 | 1.282 | -1.288 | 1.288 | -2.136 | 2.137 | -2.142 | 2.143
$\widehat{ATT}_{ranf}$ | -0.400 | 0.402 | -0.346 | 0.347 | -1.089 | 1.090 | -1.013 | 1.014
$\widehat{ATT}_{DR\\_ebal\\_ranf}$ | -0.228 | 0.230 | -0.197 | 0.198 | -0.638 | 0.639 | -0.599 | 0.600
$\widehat{ATT}_{DR\\_cbps\\_ranf}$ | -0.229 | 0.232 | -0.198 | 0.199 | -0.664 | 0.665 | -0.612 | 0.612
$\widehat{ATT}_{DR\\_opt\\_ranf}$ | -0.240 | 0.242 | -0.208 | 0.209 | -0.663 | 0.664 | -0.624 | 0.624
$\widehat{ATT}_{DR\\_ebcw\\_ranf}$ | -0.228 | 0.230 | -0.197 | 0.198 | -0.638 | 0.639 | -0.599 | 0.599
$\widehat{ATT}_{DR\\_ps\\_ranf}$ | -0.228 | 0.230 | -0.197 | 0.198 | -0.640 | 0.642 | -0.599 | 0.599
$\widehat{ATT}_{DR\\_RF\\_ranf}$ | -0.330 | 0.332 | -0.284 | 0.285 | -0.968 | 0.969 | -0.897 | 0.897
$\widehat{ATT}_{DR\\_BCART\\_ranf}$ | -0.325 | 0.327 | -0.280 | 0.281 | -0.958 | 0.959 | -0.886 | 0.886
$\widehat{ATT}_{DR\\_XGBoost\\_ranf}$ | -0.318 | 0.320 | -0.274 | 0.275 | -0.871 | 0.872 | -0.817 | 0.817
$\widehat{ATT}_{DR\\_Lasso\\_ranf}$ | -0.271 | 0.275 | -0.235 | 0.238 | -0.727 | 0.733 | -0.688 | 0.691
$\bm{\widehat{ATT}_{DR\\_DLW\\_ranf}}$ | -0.232 | 0.235 | -0.201 | 0.202 | -0.663 | 0.667 | -0.621 | 0.622
Table A3: Performance comparison of our method ($\bm{\widehat{ATT}_{DR\\_DLW\\_bart}}$) with weighting-based methods under the doubly-robust framework where BART is used for the outcome function component. | d = 8 | d = 16
---|---|---
| N = 5,000 | N = 10,000 | N = 5,000 | N = 10,000
$s_{c}$ | Estimator | Bias | RMSE | Bias | RMSE | Bias | RMSE | Bias | RMSE
0.2 | $\widehat{ATT}_{base}$ | -0.743 | 0.745 | -0.749 | 0.750 | -1.364 | 1.366 | -1.370 | 1.370
$\widehat{ATT}_{bart}$ | -0.008 | 0.031 | -0.005 | 0.022 | -0.024 | 0.037 | -0.017 | 0.026
$\widehat{ATT}_{DR\\_ebal\\_bart}$ | -0.003 | 0.030 | 0.000 | 0.022 | -0.006 | 0.030 | -0.002 | 0.020
$\widehat{ATT}_{DR\\_cbps\\_bart}$ | -0.003 | 0.030 | 0.000 | 0.022 | -0.006 | 0.030 | -0.002 | 0.020
$\widehat{ATT}_{DR\\_opt\\_bart}$ | -0.003 | 0.030 | 0.000 | 0.021 | -0.008 | 0.029 | -0.003 | 0.021
$\widehat{ATT}_{DR\\_ebcw\\_bart}$ | -0.003 | 0.030 | 0.000 | 0.022 | -0.006 | 0.030 | -0.002 | 0.020
$\widehat{ATT}_{DR\\_ps\\_bart}$ | -0.003 | 0.030 | 0.000 | 0.022 | -0.006 | 0.030 | -0.002 | 0.020
$\widehat{ATT}_{DR\\_RF\\_bart}$ | -0.005 | 0.031 | -0.003 | 0.022 | -0.019 | 0.034 | -0.012 | 0.024
$\widehat{ATT}_{DR\\_BCART\\_bart}$ | -0.007 | 0.032 | -0.003 | 0.023 | -0.017 | 0.034 | -0.011 | 0.025
$\widehat{ATT}_{DR\\_XGBoost\\_bart}$ | -0.005 | 0.034 | -0.002 | 0.022 | -0.016 | 0.033 | -0.010 | 0.023
$\widehat{ATT}_{DR\\_Lasso\\_bart}$ | -0.004 | 0.030 | -0.001 | 0.021 | -0.009 | 0.030 | -0.004 | 0.020
$\bm{\widehat{ATT}_{DR\\_DLW\\_bart}}$ | -0.012 | 0.056 | 0.001 | 0.023 | -0.008 | 0.034 | -0.002 | 0.021
0.4 | $\widehat{ATT}_{base}$ | -1.281 | 1.282 | -1.288 | 1.288 | -2.136 | 2.137 | -2.142 | 2.143
$\widehat{ATT}_{bart}$ | -0.016 | 0.039 | -0.010 | 0.025 | -0.047 | 0.063 | -0.046 | 0.049
$\widehat{ATT}_{DR\\_ebal\\_bart}$ | -0.004 | 0.039 | 0.000 | 0.024 | -0.007 | 0.050 | -0.014 | 0.036
$\widehat{ATT}_{DR\\_cbps\\_bart}$ | -0.004 | 0.039 | 0.000 | 0.024 | -0.009 | 0.049 | -0.015 | 0.034
$\widehat{ATT}_{DR\\_opt\\_bart}$ | -0.005 | 0.037 | -0.002 | 0.024 | -0.012 | 0.046 | -0.018 | 0.028
$\widehat{ATT}_{DR\\_ebcw\\_bart}$ | -0.004 | 0.039 | 0.000 | 0.024 | -0.007 | 0.050 | -0.014 | 0.036
$\widehat{ATT}_{DR\\_ps\\_bart}$ | -0.004 | 0.039 | 0.000 | 0.024 | -0.006 | 0.051 | -0.014 | 0.035
$\widehat{ATT}_{DR\\_RF\\_bart}$ | -0.012 | 0.038 | -0.007 | 0.025 | -0.039 | 0.058 | -0.039 | 0.043
$\widehat{ATT}_{DR\\_BCART\\_bart}$ | -0.010 | 0.039 | -0.007 | 0.026 | -0.037 | 0.058 | -0.035 | 0.040
$\widehat{ATT}_{DR\\_XGBoost\\_bart}$ | -0.012 | 0.039 | -0.004 | 0.026 | -0.032 | 0.054 | -0.037 | 0.043
$\widehat{ATT}_{DR\\_Lasso\\_bart}$ | -0.007 | 0.038 | -0.002 | 0.023 | -0.015 | 0.050 | -0.023 | 0.033
$\bm{\widehat{ATT}_{DR\\_DLW\\_bart}}$ | 0.001 | 0.044 | 0.001 | 0.024 | -0.005 | 0.067 | -0.018 | 0.037
Table A4: Performance comparison of our methods ($\bm{\widehat{ATT}_{DLW}}$ and $\bm{\widehat{ATT}_{DR\\_DLW\\_bart}}$) with frequently used methods. | d = 8 | d = 16
---|---|---
| N = 5,000 | N = 10,000 | N = 5,000 | N = 10,000
$s_{c}$ | Estimator | Bias | RMSE | Bias | RMSE | Bias | RMSE | Bias | RMSE
0.2 | $\widehat{ATT}_{base}$ | -0.743 | 0.745 | -0.749 | 0.750 | -1.364 | 1.366 | -1.370 | 1.370
$\bm{\widehat{ATT}_{OLS}}$ | 0.000 | 0.029 | 0.001 | 0.020 | 0.000 | 0.026 | -0.001 | 0.021
$\widehat{ATT}_{PSM}$ | -0.005 | 0.043 | 0.001 | 0.029 | -0.006 | 0.055 | -0.005 | 0.039
$\widehat{ATT}_{ARB}$ | -0.001 | 0.030 | 0.001 | 0.020 | -0.005 | 0.026 | 0.000 | 0.020
$\widehat{ATT}_{CF}$ | -0.133 | 0.137 | -0.104 | 0.106 | -0.579 | 0.580 | -0.510 | 0.511
$\widehat{ATT}_{CEVAE}$ | -0.104 | 0.112 | -0.098 | 0.109 | -0.579 | 0.581 | -0.359 | 0.362
$\bm{\widehat{ATT}_{DLW}}$ | -0.043 | 0.149 | -0.014 | 0.058 | -0.067 | 0.113 | -0.049 | 0.098
$\bm{\widehat{ATT}_{DR\\_DLW\\_bart}}$ | -0.012 | 0.056 | 0.001 | 0.023 | -0.008 | 0.034 | -0.002 | 0.021
0.4 | $\widehat{ATT}_{base}$ | -1.281 | 1.282 | -1.288 | 1.288 | -2.136 | 2.137 | -2.142 | 2.143
$\bm{\widehat{ATT}_{OLS}}$ | -0.002 | 0.033 | 0.001 | 0.021 | 0.001 | 0.035 | -0.002 | 0.023
$\widehat{ATT}_{PSM}$ | -0.004 | 0.054 | -0.002 | 0.034 | -0.011 | 0.096 | -0.008 | 0.060
$\widehat{ATT}_{ARB}$ | 0.003 | 0.034 | 0.000 | 0.023 | -0.004 | 0.044 | -0.003 | 0.031
$\widehat{ATT}_{CF}$ | -0.256 | 0.258 | -0.200 | 0.201 | -1.008 | 1.009 | -0.905 | 0.905
$\widehat{ATT}_{CEVAE}$ | -0.203 | 0.208 | -0.146 | 0.157 | -0.962 | 0.965 | -0.676 | 0.678
$\bm{\widehat{ATT}_{DLW}}$ | -0.053 | 0.103 | -0.042 | 0.076 | -0.131 | 0.258 | -0.139 | 0.200
$\bm{\widehat{ATT}_{DR\\_DLW\\_bart}}$ | 0.001 | 0.044 | 0.001 | 0.024 | -0.005 | 0.067 | -0.018 | 0.037

### A2. Results of experiments on synthetic data for Setting 3

From Tables A5 - A8, we find that the comparison among different estimators in
Setting 3 is similar to Setting 2, which further confirms the robustness of
our method when dealing with different model specifications. Please refer to
the main text for further explanations of the comparison results.

Table A5: Performance comparison of our method ($\bm{\widehat{ATT}_{DLW}}$) with weighting-based methods. | d = 8 | d = 16
---|---|---
| N = 5,000 | N = 10,000 | N = 5,000 | N = 10,000
$s_{c}$ | Estimator | Bias | RMSE | Bias | RMSE | Bias | RMSE | Bias | RMSE
0.2 | $\widehat{ATT}_{base}$ | -2.527 | 2.530 | -2.537 | 2.539 | -5.761 | 5.766 | -5.792 | 5.794
$\widehat{ATT}_{ebal}$ | -2.529 | 2.532 | -2.536 | 2.538 | -5.760 | 5.764 | -5.792 | 5.794
$\widehat{ATT}_{cbps}$ | -2.529 | 2.532 | -2.536 | 2.538 | -5.760 | 5.764 | -5.792 | 5.794
$\widehat{ATT}_{opt}$ | -2.529 | 2.532 | -2.536 | 2.538 | -5.762 | 5.766 | -5.793 | 5.795
$\widehat{ATT}_{ebcw}$ | -2.529 | 2.532 | -2.536 | 2.538 | -5.760 | 5.764 | -5.792 | 5.794
$\widehat{ATT}_{ps}$ | -2.530 | 2.533 | -2.537 | 2.538 | -5.762 | 5.766 | -5.793 | 5.794
$\widehat{ATT}_{RF}$ | -1.458 | 1.460 | -1.335 | 1.336 | -4.717 | 4.721 | -4.583 | 4.585
$\widehat{ATT}_{BCART}$ | -1.337 | 1.340 | -1.196 | 1.197 | -4.629 | 4.633 | -4.480 | 4.481
$\widehat{ATT}_{XGBoost}$ | -0.601 | 0.604 | -0.547 | 0.549 | -1.945 | 1.949 | -1.974 | 1.976
$\widehat{ATT}_{Lasso}$ | -2.528 | 2.531 | -2.537 | 2.538 | -5.763 | 5.767 | -5.792 | 5.794
$\bm{\widehat{ATT}_{DLW}}$ | -0.075 | 0.161 | -0.113 | 0.134 | -0.821 | 1.518 | -0.371 | 0.448
0.4 | $\widehat{ATT}_{base}$ | -3.500 | 3.503 | -3.508 | 3.509 | -7.407 | 7.410 | -7.432 | 7.434
$\widehat{ATT}_{ebal}$ | -3.500 | 3.502 | -3.508 | 3.509 | -7.409 | 7.412 | -7.434 | 7.436
$\widehat{ATT}_{cbps}$ | -3.500 | 3.502 | -3.508 | 3.509 | -7.409 | 7.413 | -7.434 | 7.436
$\widehat{ATT}_{opt}$ | -3.501 | 3.503 | -3.508 | 3.509 | -7.410 | 7.413 | -7.434 | 7.436
$\widehat{ATT}_{ebcw}$ | -3.500 | 3.502 | -3.508 | 3.509 | -7.409 | 7.412 | -7.434 | 7.436
$\widehat{ATT}_{ps}$ | -3.501 | 3.504 | -3.508 | 3.509 | -7.410 | 7.413 | -7.434 | 7.436
$\widehat{ATT}_{RF}$ | -2.045 | 2.047 | -1.865 | 1.866 | -6.274 | 6.277 | -6.109 | 6.110
$\widehat{ATT}_{BCART}$ | -1.875 | 1.878 | -1.677 | 1.678 | -6.202 | 6.205 | -6.017 | 6.018
$\widehat{ATT}_{XGboost}$ | -0.756 | 0.759 | -0.728 | 0.730 | -2.958 | 2.960 | -2.868 | 2.869
$\widehat{ATT}_{Lasso}$ | -3.501 | 3.503 | -3.508 | 3.509 | -7.407 | 7.410 | -7.432 | 7.434
$\bm{\widehat{ATT}_{DLW}}$ | -0.122 | 0.151 | -0.072 | 0.105 | -0.162 | 1.133 | -0.197 | 0.612
Table A6: Performance comparison of our method ($\bm{\widehat{ATT}_{DR\\_DLW\\_ranf}}$) with weighting-based methods under the doubly-robust framework where random forest is used for the outcome function component. | d = 8 | d = 16
---|---|---
| N = 5,000 | N = 10,000 | N = 5,000 | N = 10,000
$s_{c}$ | Estimator | Bias | RMSE | Bias | RMSE | Bias | RMSE | Bias | RMSE
0.2 | $\widehat{ATT}_{base}$ | -2.527 | 2.530 | -2.537 | 2.539 | -5.761 | 5.766 | -5.792 | 5.794
$\widehat{ATT}_{ranf}$ | -1.304 | 1.306 | -1.092 | 1.093 | -4.788 | 4.792 | -4.519 | 4.520
$\widehat{ATT}_{DR\\_ebal\\_ranf}$ | -1.304 | 1.306 | -1.098 | 1.099 | -4.705 | 4.709 | -4.450 | 4.451
$\widehat{ATT}_{DR\\_cbps\\_ranf}$ | -1.304 | 1.306 | -1.098 | 1.099 | -4.705 | 4.709 | -4.450 | 4.451
$\widehat{ATT}_{DR\\_opt\\_ranf}$ | -1.304 | 1.306 | -1.099 | 1.099 | -4.707 | 4.710 | -4.450 | 4.451
$\widehat{ATT}_{DR\\_ebcw\\_ranf}$ | -1.304 | 1.306 | -1.098 | 1.099 | -4.705 | 4.709 | -4.450 | 4.451
$\widehat{ATT}_{DR\\_ps\\_ranf}$ | -1.304 | 1.306 | -1.098 | 1.099 | -4.706 | 4.709 | -4.450 | 4.451
$\widehat{ATT}_{DR\\_RF\\_ranf}$ | -1.059 | 1.061 | -0.866 | 0.867 | -4.369 | 4.372 | -4.081 | 4.082
$\widehat{ATT}_{DR\\_BCART\\_ranf}$ | -1.032 | 1.034 | -0.840 | 0.840 | -4.344 | 4.347 | -4.055 | 4.056
$\widehat{ATT}_{DR\\_XGBoost\\_ranf}$ | -0.873 | 0.875 | -0.730 | 0.731 | -3.447 | 3.450 | -3.312 | 3.313
$\widehat{ATT}_{DR\\_Lasso\\_ranf}$ | -1.304 | 1.306 | -1.099 | 1.099 | -4.707 | 4.710 | -4.450 | 4.452
$\bm{\widehat{ATT}_{DR\\_DLW\\_ranf}}$ | -0.726 | 0.730 | -0.638 | 0.639 | -3.063 | 3.096 | -2.778 | 2.780
0.4 | $\widehat{ATT}_{base}$ | -3.500 | 3.503 | -3.508 | 3.509 | -7.407 | 7.410 | -7.432 | 7.434
$\widehat{ATT}_{ranf}$ | -1.865 | 1.866 | -1.574 | 1.574 | -6.303 | 6.305 | -5.973 | 5.975
$\widehat{ATT}_{DR\\_ebal\\_ranf}$ | -1.864 | 1.865 | -1.578 | 1.579 | -6.217 | 6.219 | -5.901 | 5.902
$\widehat{ATT}_{DR\\_cbps\\_ranf}$ | -1.864 | 1.865 | -1.578 | 1.579 | -6.217 | 6.219 | -5.901 | 5.902
$\widehat{ATT}_{DR\\_opt\\_ranf}$ | -1.864 | 1.865 | -1.579 | 1.579 | -6.219 | 6.221 | -5.901 | 5.903
$\widehat{ATT}_{DR\\_ebcw\\_ranf}$ | -1.864 | 1.865 | -1.578 | 1.579 | -6.217 | 6.219 | -5.901 | 5.902
$\widehat{ATT}_{DR\\_ps\\_ranf}$ | -1.864 | 1.865 | -1.578 | 1.579 | -6.218 | 6.220 | -5.901 | 5.902
$\widehat{ATT}_{DR\\_RF\\_ranf}$ | -1.521 | 1.523 | -1.249 | 1.250 | -5.843 | 5.845 | -5.484 | 5.485
$\widehat{ATT}_{DR\\_BCART\\_ranf}$ | -1.482 | 1.483 | -1.211 | 1.212 | -5.823 | 5.825 | -5.460 | 5.461
$\widehat{ATT}_{DR\\_XGBoost\\_ranf}$ | -1.213 | 1.215 | -1.025 | 1.026 | -4.681 | 4.682 | -4.443 | 4.444
$\widehat{ATT}_{DR\\_Lasso\\_ranf}$ | -1.864 | 1.866 | -1.579 | 1.579 | -6.218 | 6.220 | -5.901 | 5.903
$\bm{\widehat{ATT}_{DR\\_DLW\\_ranf}}$ | -1.095 | 1.096 | -0.900 | 0.901 | -3.696 | 3.728 | -3.572 | 3.578
Table A7: Performance comparison of our method ($\bm{\widehat{ATT}_{DR\\_DLW\\_bart}}$) with weighting-based methods under the doubly-robust framework where BART is used for the outcome function component. | d = 8 | d = 16
---|---|---
| N = 5,000 | N = 10,000 | N = 5,000 | N = 10,000
$s_{c}$ | Estimator | Bias | RMSE | Bias | RMSE | Bias | RMSE | Bias | RMSE
0.2 | $\widehat{ATT}_{base}$ | -2.527 | 2.530 | -2.537 | 2.539 | -5.761 | 5.766 | -5.792 | 5.794
$\widehat{ATT}_{bart}$ | -0.037 | 0.065 | -0.028 | 0.040 | -0.820 | 0.830 | -0.473 | 0.478
$\widehat{ATT}_{DR\\_ebal\\_bart}$ | -0.037 | 0.065 | -0.028 | 0.040 | -0.820 | 0.829 | -0.473 | 0.478
$\widehat{ATT}_{DR\\_cbps\\_bart}$ | -0.037 | 0.065 | -0.028 | 0.040 | -0.820 | 0.829 | -0.473 | 0.478
$\widehat{ATT}_{DR\\_opt\\_bart}$ | -0.037 | 0.065 | -0.028 | 0.040 | -0.820 | 0.830 | -0.473 | 0.478
$\widehat{ATT}_{DR\\_ebcw\\_bart}$ | -0.037 | 0.065 | -0.028 | 0.040 | -0.820 | 0.829 | -0.473 | 0.478
$\widehat{ATT}_{DR\\_ps\\_bart}$ | -0.037 | 0.065 | -0.028 | 0.040 | -0.820 | 0.830 | -0.473 | 0.478
$\widehat{ATT}_{DR\\_RF\\_bart}$ | -0.015 | 0.055 | -0.009 | 0.030 | -0.744 | 0.755 | -0.405 | 0.410
$\widehat{ATT}_{DR\\_BCART\\_bart}$ | -0.013 | 0.055 | -0.005 | 0.030 | -0.738 | 0.749 | -0.403 | 0.409
$\widehat{ATT}_{DR\\_XGBoost\\_bart}$ | 0.010 | 0.055 | 0.010 | 0.032 | -0.498 | 0.514 | -0.231 | 0.241
$\widehat{ATT}_{DR\\_Lasso\\_bart}$ | -0.037 | 0.065 | -0.028 | 0.040 | -0.821 | 0.830 | -0.473 | 0.478
$\bm{\widehat{ATT}_{DR\\_DLW\\_bart}}$ | -0.009 | 0.064 | -0.005 | 0.032 | -0.492 | 0.632 | -0.207 | 0.231
0.4 | $\widehat{ATT}_{base}$ | -3.500 | 3.503 | -3.508 | 3.509 | -7.407 | 7.410 | -7.432 | 7.434
$\widehat{ATT}_{bart}$ | -0.066 | 0.086 | -0.033 | 0.051 | -1.192 | 1.201 | -0.651 | 0.656
$\widehat{ATT}_{DR\\_ebal\\_bart}$ | -0.066 | 0.086 | -0.033 | 0.051 | -1.192 | 1.200 | -0.651 | 0.656
$\widehat{ATT}_{DR\\_cbps\\_bart}$ | -0.066 | 0.086 | -0.033 | 0.051 | -1.192 | 1.200 | -0.651 | 0.656
$\widehat{ATT}_{DR\\_opt\\_bart}$ | -0.066 | 0.086 | -0.033 | 0.051 | -1.192 | 1.201 | -0.651 | 0.656
$\widehat{ATT}_{DR\\_ebcw\\_bart}$ | -0.066 | 0.086 | -0.033 | 0.051 | -1.192 | 1.200 | -0.651 | 0.656
$\widehat{ATT}_{DR\\_ps\\_bart}$ | -0.066 | 0.086 | -0.033 | 0.051 | -1.192 | 1.201 | -0.651 | 0.656
$\widehat{ATT}_{DR\\_RF\\_bart}$ | -0.036 | 0.065 | -0.006 | 0.039 | -1.102 | 1.111 | -0.571 | 0.577
$\widehat{ATT}_{DR\\_BCART\\_bart}$ | -0.032 | 0.062 | -0.003 | 0.041 | -1.100 | 1.109 | -0.568 | 0.574
$\widehat{ATT}_{DR\\_XGBoost\\_bart}$ | -0.001 | 0.053 | 0.019 | 0.043 | -0.780 | 0.791 | -0.330 | 0.339
$\widehat{ATT}_{DR\\_Lasso\\_bart}$ | -0.066 | 0.086 | -0.033 | 0.051 | -1.192 | 1.201 | -0.651 | 0.656
$\bm{\widehat{ATT}_{DR\\_DLW\\_bart}}$ | -0.052 | 0.071 | 0.001 | 0.044 | -0.655 | 0.785 | -0.238 | 0.306
Table A8: Performance comparison of our methods ($\bm{\widehat{ATT}_{DLW}}$ and $\bm{\widehat{ATT}_{DR\\_DLW\\_bart}}$) with frequently used methods. | d = 8 | d = 16
---|---|---
| N = 5,000 | N = 10,000 | N = 5,000 | N = 10,000
$s_{c}$ | Estimator | Bias | RMSE | Bias | RMSE | Bias | RMSE | Bias | RMSE
0.2 | $\widehat{ATT}_{base}$ | -2.527 | 2.530 | -2.537 | 2.539 | -5.761 | 5.766 | -5.792 | 5.794
$\widehat{ATT}_{OLS}$ | -2.529 | 2.532 | -2.537 | 2.538 | -5.762 | 5.766 | -5.792 | 5.794
$\widehat{ATT}_{PSM}$ | -2.420 | 2.435 | -2.447 | 2.455 | -5.732 | 5.742 | -5.758 | 5.762
$\widehat{ATT}_{ARB}$ | -2.530 | 2.533 | -2.537 | 2.538 | -5.762 | 5.767 | -5.794 | 5.795
$\widehat{ATT}_{CF}$ | -1.105 | 1.107 | -0.827 | 0.827 | -4.692 | 4.696 | -4.383 | 4.384
$\widehat{ATT}_{CEVAE}$ | -0.407 | 0.422 | -0.222 | 0.255 | -1.970 | 1.983 | -1.935 | 1.950
$\bm{\widehat{ATT}_{DLW}}$ | -0.075 | 0.161 | -0.113 | 0.134 | -0.821 | 1.518 | -0.371 | 0.448
$\bm{\widehat{ATT}_{DR\\_DLW\\_bart}}$ | -0.009 | 0.064 | -0.005 | 0.032 | -0.492 | 0.632 | -0.207 | 0.231
0.4 | $\widehat{ATT}_{base}$ | -3.500 | 3.503 | -3.508 | 3.509 | -7.407 | 7.410 | -7.432 | 7.434
$\widehat{ATT}_{OLS}$ | -3.501 | 3.503 | -3.508 | 3.509 | -7.409 | 7.412 | -7.434 | 7.436
$\widehat{ATT}_{PSM}$ | -3.394 | 3.413 | -3.403 | 3.416 | -7.419 | 7.426 | -7.383 | 7.387
$\widehat{ATT}_{ARB}$ | -3.501 | 3.504 | -3.508 | 3.510 | -7.412 | 7.415 | -7.436 | 7.437
$\widehat{ATT}_{CF}$ | -1.527 | 1.528 | -1.136 | 1.137 | -6.115 | 6.117 | -5.705 | 5.706
$\widehat{ATT}_{CEVAE}$ | -0.576 | 0.589 | -0.268 | 0.284 | -2.853 | 2.866 | -2.809 | 2.824
$\bm{\widehat{ATT}_{DLW}}$ | -0.122 | 0.151 | -0.072 | 0.105 | -0.162 | 1.133 | -0.197 | 0.612
$\bm{\widehat{ATT}_{DR\\_DLW\\_bart}}$ | -0.052 | 0.071 | 0.001 | 0.044 | -0.655 | 0.785 | -0.238 | 0.306

### A3. Model Convergence of experiments on synthetic data

Figures A1 - A3 show the negative log-likelihood (NLL) of covariates in both
control and treatment groups during the learning process. We can see that the
NLL converges within a small number of epochs, under different model settings
with respect to the dimensionality, confounding strength and sample size. The
model convergence empirically implies that the estimated distribution
approaches the underlying true distribution. Blue and red denote training and
validation, respectively.

Figure A1: Convergence of NNL for Setting 1. Figure A2: Convergence of NNL for
Setting 2. Figure A3: Convergence of NNL for Setting 3.

### A4. Hyper-parameter tuning

We use the grid search to tune all involved parameters based on the log-
likelihood on the validation set. The searching pool of the hyper-parameters
and their optimal values are listed in Table A9.

Table A9: Hyper-parameter tuning. Parameter | Searching pool |  Optimal value
---|---|---
The number of transformations K | {2, 4, 6, 8, 10} |  K=6 for synthetic data
K=2 for real data
The number of hidden layers in $NN_{1}$ | {1, 2, 4} |  2
The number of hidden units in $NN_{1}$ | {32, 64, 128} |  64
The number of hidden layers of $NN_{2}$ | {1, 2, 4} |  1
The number of hidden units of $NN_{2}$ | {8, 16} |  8
Batch size | {64, 128, 256} |  128
Learning rate | {0.01, 0.001, 0.0001} |  0.0001