# 2D score based estimation of heterogeneous treatment effects

Steven Siwei Ye
Department of Statistics
University of California, Los Angeles
Los Angeles, CA 90095
<EMAIL_ADDRESS>
&Yanzhen Chen
Department of ISOM
Hong Kong University of Science and Technology
Clear Water Bay, Kowloon Hong Kong
<EMAIL_ADDRESS>
&Oscar Hernan Madrid Padilla
Department of Statistics
University of California, Los Angeles
Los Angeles, CA 90095
<EMAIL_ADDRESS>

###### Abstract

In the study of causal inference, statisticians show growing interest in
estimating and analyzing heterogeneity in causal effects in observational
studies. However, there usually exists a trade-off between accuracy and
interpretability when developing a desirable estimator for treatment effects.
To make efforts to address the issue, we propose a score-based framework for
estimating the Conditional Average Treatment Effect (CATE) function in this
paper. The framework integrates two components: (i) leverage the joint use of
propensity and prognostic scores in a matching algorithm to obtain a proxy of
the heterogeneous treatment effects for each observation, (ii) utilize non-
parametric regression trees to construct an estimator for the CATE function
conditioning on the two scores. The method naturally stratifies treatment
effects into subgroups over a 2d grid whose axis are the propensity and
prognostic scores. We conduct benchmark experiments on multiple simulated data
and demonstrate clear advantages of the proposed estimator over state of the
art methods. We also evaluate empirical performance in real-life settings,
using two observational data from a clinical trial and a complex social
survey, and interpret policy implications following the numerical results. The
R code for implementing the method introduced in the paper is publicly
available on one of the author’s Github page
(https://github.com/stevenysw/causal_pp).

_K_ eywords observational data $\cdot$ subgroup treatment effects $\cdot$
regression trees $\cdot$ matching

##  1 Introduction

The questions that motivate many scientific studies in disciplines such as
economics, epidemiology, medicine, and political science, are not
associational but causal in nature. Understanding causality, which often
refers to the study of causal inference, is an emerging area of statistics.
Many researchers are interested in inferring average treatment effects, which
provide a good sense of whether treatment is likely to deliver more benefit
than the control among a whole community. However, the same treatment may
affect different individuals very differently. Therefore, a substantial amount
of works focus on analyzing heterogeneity in treatment effects, of which the
term refers to variation in the effects of treatment across individuals. This
variation may provide theoretical insights, revealing how the effect of
interventions depends on participants’ characteristics or how varying features
of a treatment alters the effect of an intervention.

In this paper, we follow the binary outcome framework for causal inference
[56, 67], where each unit is assigned into either the treatment or the control
group. Each unit has an observed outcome variable with a set of covariates. In
randomized experiments and observational studies, it is desirable to replicate
a sample as closely as possible by obtaining subjects from the treatment and
control groups with similar covariate distributions when estimating causal
effects. However, it is almost impossible to match observations exactly the
same in both treatment and control groups in observational studies. To address
this problem, it is usually preferred to define prespecified subgroups under
certain conditions and estimate the treatment effects varying among subgroups.
Accordingly, the conditional average treatment effect [31, CATE;] function is
designed to capture heterogeneity of a treatment effect across subpopulations.
In most cases, the function is conditioned on some component(s) of the
covariates or a single statistic, like propensity score [66] and prognostic
score [34]. Propensity scores are the probabilities of receiving the treatment
of interest; prognostic scores model the potential outcome under a control
group assignment. To understand treatment effect heterogeneity in terms of
propensity and prognostic scores, we assume that equal or similar treatment
effects are observed along some intervals of the two scores.

We target at constructing an accurate and interpretable estimator for
treatment effects that conditions on both propensity and prognostic scores and
assumes a piecewise constant structure in treatment effects. We take a step
further from score-based matching algorithms and propose a data-driven
approach that integrates the joint use of propensity and prognostic scores in
a matching algorithm and a partition over the entire population via a non-
parametric regression tree. In the first step, we estimate propensity scores
and prognostic scores for each observed unit in the data. Secondly, we perform
a $K$-nearest-neighbor matching of units of the treatment and control groups
based on the two estimated scores and forth construct a proxy of individual
treatment effects for all units. The last step involves growing a binary tree
regressed on the two estimated scores.

The complementary nature of propensity and prognostic score methods supports
that conditioning on both the propensity and prognostic scores has the
potential to reduce bias and improve the precision of treatment effect
estimates, and it is affirmed in the simulation studies by [50] and [3]. We
also demonstrate such advantage for our proposed estimator across almost all
scenarios examined in the simulation experiments.

Besides high precision in estimation, our proposed estimator demonstrates its
superiority over state-of-arts methods with a few attractive properties as
follows:

  * •

The estimator is computationally efficient. Propensity and prognostic scores
can be easily estimated through simple regression techniques. Our matching
algorithm based on the two scores largely reduces dimensionality compared to
full matching on the complete covariates. Moreover, growing a single
regression tree saves much time than other tree-based estimation methods, such
as BART [32] and random forests [81, 6].

  * •

Many previous works in subgroup analysis, such as [4] and [1], set
stratification on the sample with a fixed number of subgroups before
estimating treatment effects. These approaches require a pre-determination on
the number of subgroups contained in the data, and they inevitably introduce
arbitrariness into the causal inference. In comparison, our proposed method
simultaneously identifies the underlying subgroups in observations through
binary split according to propensity and prognostic scores and provides a
consequential estimation of treatment effects on each subgroups.

  * •

Although random forests based methods [81, 6] achieve great performance in
minimizing bias in estimating treatment effects, these ensemble methods are
often referred to as "black boxes". It is hard to capture the underlying
reason why the collective decision with the high number of operations involved
is made in their estimation process. On the contrary, our proposed method
carries more clear interpretations by providing a 2d summary of treatment
effects. As a result, given the covariates of an observation, one can easily
deduce the positiveness and magnitude of its treatment effect according to its
probability of treatment receipt and potential outcome following the structure
of the regression tree.

We review relevant literature on matching algorithms and estimation of
heterogeneous treatment effects in Section 2. In Section 3, we provide the
theoretical framework and preliminaries for the causal inference model. We
propose our method for estimation and prediction in Section 4. Section 5 lists
the results of numerical experiments on multiple simulated data sets and two
real-world data sets, following with the comparison with state-of-the-art
methods in existing literature and the discussion on policy implications under
different realistic scenarios.

##  2 Relevant literature

Statistical analysis of causality can be dated back to [56]. Causal inference
can be viewed as an identification problem [47], for which statisticians are
dedicated to learn the true causality behind the data. In reality, however, we
do not have enough information to determine the true value due to a limited
number of observations for analysis. This problem is also summarized as a
"missing data" problem [23], which stems from the fundamental problem of
causal inference [40], that is, for each unit at most one of the potential
outcomes is observed. Importantly, the causal effect identification problem,
especially for estimating treatment effects, can only be resolved through
assumptions. Several key theoretical frameworks have been proposed over the
past decades. The potential outcomes framework by [67], often referred to as
the Rubin Causal Model (RCM) [40], is a common model of causality in
statistics at the moment. [21] develops a decision theoretic approach to
causality that rejects counterfactuals. [59, 60] advocates for a model of
causality based on non-parametric structural equations and path diagrams.

### Matching

To tackle the "missing data" problem when estimating treatment effects in
randomized experiments in practice, matching serves as a very powerful tool.
The main goal of matching is to find matched groups with similar or balanced
observed covariate distributions [74]. The exact $K$-nearest-neighbor matching
[67] is one of the most common, and easiest to implement and understand
methods; and ratio matching [72, 69, 54], which finds multiple good matches
for each treated individual, performs well when there is a large number of
control individuals. [64], [30], [86], and [87] developed various optimal
matching algorithms to minimize the total sum of distances between treated
units and matched controls in a global sense. [2] studied the consistency of
covariate matching estimators under large sample assumptions. Instead of
greedy matching on the entire covariates, propensity score matching (PSM) by
[70] is an alternative algorithm that does not guarantee optimal balance among
covariates and reduces dimension sufficiently. [44] improved propensity score
matching with regression adjustment. The additional matching on prognostic
factors in propensity score matching was first considered by [69]. Later, [50]
demonstrated the superiority of the joint use of propensity and prognostic
scores in matching over single-score based matching in low-dimensional
settings through extensive simulation studies. [3] extended the method to fit
to high-dimensional settings and derived asymptotic results for the so-called
doubly robust matching estimators.

### Subclassification

To understand heterogeneity of treatment effects in the data,
subclassification, first used in [17], is another important research problem.
The key idea is to form subgroups over the entire population based on
characteristics that are either immutable or observed before randomization.
[66, 65] and [52] examined how creating a fixed number of subclasses according
to propensity scores removes the bias in the estimated treatment effects, and
[83] developed a similar methodology in settings with more than two treatment
levels. Full matching [63, 33, 75] is a more sophisticated form of
subclassification that selects the number of subclasses automatically by
creating a series of matched sets. [71] presented three measures derived using
the theory of order statistics to claim heterogeneity of treatment effect
across subgroups. [76] pioneered in exploiting standard regression tree
methods [14] in subgroup treatment effect analysis. Further, [5] derived a
recursive partition of the population according to treatment effect
heterogeneity. [37] was the first work to advocate for the use of Bayesian
additive regression tree models [16, BART;] for estimating heterogeneous
treatment effects, followed by a significant number of research papers
focusing on the seminal methodology, including [29], [38], and [32]. [1]
introduced endogenous stratification to estimate subgroup effects for a fixed
number of subgroups based on certain quantiles of the prognostic score. More
recently, [57] combined the fused lasso estimator with score matching methods
to lead to a data-adaptive subgroup effects estimator.

### Machine Learning for Causal Inference

For the goal of analyzing treatment effect heterogeneity, supervised machine
learning methods play an important role. One of the more common ways for
accurate estimation with experimental and observational data is to apply
regression [45] or tree-based methods [43]. From a Bayesian perspective, [36]
provided a principled way of adding priors to regression models, and [78]
developed Bayesian non-parametric approaches for both linear regression and
tree models. The recent breakthrough work by [81] proposed the causal forest
estimator arising from random forests from [13]. More recently, [6] took a
step forward and enhanced the previous estimator based on generalized random
forests. [42] adapted an estimator from the Support Vector Machine (SVM)
classifier with hinge loss [82]. [11] studied treatment effect estimators with
lasso regularization [80] when the number of covariates is large, and [48]
applied group lasso for simultaneous covariate selection and robust estimation
of causal effects. In the meantime, a series of papers including [61], [49]
and [77], focused on developing meta-learners for heterogeneous treatment
effects that can take advantage of various machine learning algorithms and
data structures.

### Applied Work

On the application side, the estimation of heterogeneous treatment effects is
particularly an intriguing topic in causal inference with broad applications
in scientific research. [28] estimated heterogeneous treatment effects in
randomized experiments in the context of political science. [22] explored the
use of propensity score matching for nonexperimental causal studies with
application in economics. [20] investigated heterogeneous treatment effects to
provide the evidence base for precision medicine and patient-centred care.
[84] proposed the Survival Causal Tree (SCT) method to discover patient
subgroups with heterogeneous treatment effects from censored observational
data. [62] examined three classes of approaches to identify heterogeneity of
treatment effect within a randomized clinical trial, and [79] rendered a
subgroup treatment estimate for drug trials.

##  3 Preliminaries

Before we introduce our method, we need to provide some mathematical
background for treatment effect estimation. We follow Rubin’s framework on
causal inference [67], and assume a superpopulation or distribution
$\mathcal{P}$ from which a realization of $n$ independent random variables is
given as the training data. That is, we are given
$\\{(Y_{i}(0),Y_{i}(1),X_{i},Z_{i})\\}_{i=1}^{n}$ independent copies of
$(Y(1),Y(0),X,Z)$, where $X_{i}\in\mathbb{R}^{d}$ is a $d$-dimensional
covariate or feature vector, $Z_{i}\in\\{0,1\\}$ is the treatment-assignment
indicator, $Y_{i}(0)\in\mathbb{R}$ is the potential outcome of unit $i$ when
$i$ is assigned to the control group, and $Y_{i}(1)$ is the potential outcome
when $i$ is assigned to the treatment group.

One important and commonly used measure of causality in a binary treatment
model is the average treatment effect [44, ATE;], that is, the mean outcome
difference between the treatment and control groups. Formally, we write the
ATE as

$\text{ATE}:=\mathbb{E}[Y(1)-Y(0)].$

With the $n$ units in the study, we further define the individual treatment
effect (ITE) of unit $i$ denoted by $D_{i}$ as

$D_{i}:=Y_{i}(1)-Y_{i}(0).$

Then, an unbiased estimate of the ATE is the sample average treatment effect

$\bar{Y}(1)-\bar{Y}(0)=\frac{1}{n}\sum_{i=1}^{n}D_{i}.$

However, we cannot observe $D_{i}$ for any unit because a unit is either in
the treatment group or in the control group, but not in both.

To analyze heterogeneous treatment effects, it is natural to divide the data
into subgroups (e.g., by gender, or by race), and investigate if the average
treatment effects are different across subgroups. Therefore, instead of
estimating the ATE or the ITE directly, statisticians seek to estimate the
conditional average treatment effect (CATE), defined by

$\tau(x):=\mathbb{E}[Y(1)-Y(0)\ |\ X=x].$ (1)

The CATE can be viewed as an ATE in a subpopulation defined by $\\{X=x\\}$,
i.e. the ATE conditioned on membership in the subgroup.

We also recall the propensity score [66], denoted by $e(X)$, and defined as

$e(X)=\mathbb{P}(Z=1\ |\ X).$

Thus, $e(X)$ is the the probability of receiving treatment for a unit with
covariate $X$. In addition, we consider prognostic scores, denoted by $p(X)$,
for potential outcomes and we use the conventional definition as the predicted
outcome under the control condition:

$p(X)=E[Y(0)\ |\ X].$

We restrict our attention to the case of no effect modification so that there
is a single prognostic score $p(X)$, satisfying the following condition [34,
Proposition 1 in]:

$Y(0)\perp\\!\\!\\!\perp X\ |\ p(X).$

We are interested in constructing a 2d summary of treatment effects based on
propensity and prognostic scores. Instead of conditioning on the entire
covariates or a subset of it in the CATE function, we express our estimand,
named as scored-based subgroup CATE, by conditioning on the two scores:

$\tau(x):=\mathbb{E}[Y(1)-Y(0)\ |\ e=e(x),p=p(x)].$ (2)

For interpretability, we assume that treatment effects are piecewise constant
over a 2d grid of propensity and prognostic scores. Specifically, there exists
a partition of intervals $\\{I_{1}^{e},...,I_{s}^{e}\\}$ of $[0,1]$ and
another partition of intervals $\\{I_{1}^{p},...,I_{t}^{p}\\}$ of $\mathbb{R}$
such that for any $i\in\\{1,...,s\\}$ and $j\in\\{1,...,t\\}$, we have

$\tau(x)\equiv C_{i,j}\ \ \text{ for }x\text{ s.t. }e(x)\in I^{e}_{i},p(x)\in
I^{p}_{j},$

where $C_{i,j}\in\mathbb{R}$ is a constant.

Moreover, our estimation of treatment effects relies on the following
assumptions:

######  Assumption 1.

Throughout the paper, we maintain the Stable Unit Treatment Value Assumption
[45, SUTVA;], which consists of two components: no interference and no hidden
variations of treatment. Mathematically, for unit $i=1,...,n$ with outcome
$Y_{i}$ and treatment indicator $Z_{i}$, it holds that

$Y_{i}(Z_{1},Z_{2},...,Z_{n})=Y_{i}(Z_{i}).$

Thus, the SUTVA requires that the potential outcomes of one unit should be
unaffected by the particular assignment of treatments to the other units.
Furthermore, for each unit, there are no different forms or versions of each
treatment level, which lead to different potential outcomes.

######  Assumption 2.

The assumption of probabilistic assignment holds. This requires the assignment
mechanism to imply a non-zero probability for each treatment value, for every
unit. For the given covariates $X$ and treatment-assignment indicator $Z$, we
must have

$0<\mathbb{P}(Z=1\ |\ X)<1,$

almost surely.

This condition, regarding the joint distribution of treatments and covariates,
is also known as overlap in some literature (See Assumption 2.2 in [44],
Imbens04 and [19], DAmour21), and it is necessary for estimating treatment
effects everywhere in the defined covariate space. Note that
$\mathbb{P}(Z_{i}=1\ |\ X_{i})$ is the propensity score. In other words,
Assumption 2 requires that the propensity score, for all values of the
treatment and all combinations of values of the confounders, be strictly
between 0 and 1.

######  Assumption 3.

We make the assumption that

$(Y(0),Y(1))\perp\\!\\!\\!\perp Z\ |\ e(X),p(X)$

holds.

This assumption is inspired by the usual unconfoundedness assumption:

$(Y(0),Y(1))\perp\\!\\!\\!\perp Z\ |\ X$ (3)

Combining Assumption 2 and that in Equation (3), the conditions are typically
referred as strong ignorability defined in [66]. Strong ignorability states
which outcomes are observed or missing is independent of the missing data
conditional on the observed data. It allows statisticians to address the
challenge that the "ground truth" for the causal effect is not observed for
any individual unit. We rewrite the conventional assumption by replacing the
vector of covariates $x$ with the joint of propensity score $e(x)$ and $p(x)$
to accord with our estimation target.

Provided that Assumptions 1-3 hold, it follows that

$\mathbb{E}[Y(z)\ |\ e=e(x),p=p(x)]=\mathbb{E}[Y\ |\ e=e(x),p=p(x),Z=z],$

and thus our estimand (2) is equivalent to

$\tau(x)=\mathbb{E}[Y\ |\ e=e(x),p=p(x),Z=1]-\mathbb{E}[Y\ |\
e=e(x),p=p(x),Z=0].$ (4)

Thus, in this paper we focus on estimating (4), which is equivalent to (2) if
the assumptions above hold, but might be different if Assumption 3 is
violated.

##  4 Methodology

We now formally introduce our proposal of a three-step method for estimating
heterogeneous treatment effects and the estimation rule for a given new
observation. We assume a sample of size $n$ with covariate $X$, treatment
indicator $Z$, and outcome variable $Y$, where the notations inherit from the
previous section. Generally, we consider a low-dimensional set-up, where the
sample size $n$ is larger than the covariate dimension $d$. An extension of
our proposed method to the high-dimensional case is discussed in this section
as well.

### Step 1

We first estimate propensity and prognostic scores for all observations in the
sample. For propensity score, we apply a logistic regression on the entire
covariate $X$ and the treatment indicator $Z$ by solving the optimization
problem

$\hat{\alpha}=\underset{\alpha\in\mathbb{R}^{d}}{\arg\min}\
-\sum_{i=1}^{n}\Big{[}Z_{i}\log\sigma(X_{i}^{\top}\alpha)+(1-Z_{i})\log(1-\sigma(X_{i}^{\top}\alpha))\Big{]},$
(5)

where $\sigma(x)=\frac{1}{1+\exp(-x)}$ is the logistic function. With the
coefficient vector $\hat{\alpha}$, we compute the estimated propensity scores
$\hat{e}$ by

$\hat{e}_{i}=\sigma\left(X_{i}^{\top}\hat{\alpha}\right).$

For prognostic score, we restrict to the controlled group, and regress the
outcome variable $Y$ on the covariate $X$ through ordinary least squares: we
solve

$\hat{\theta}=\underset{\theta\in\mathbb{R}^{d}}{\arg\min}\
\sum_{i:Z_{i}=0}(Y_{i}-X_{i}^{\top}\theta)^{2},$ (6)

and we estimate prognostic scores as

$\hat{p}_{i}=X_{i}^{\top}\hat{\theta}.$

### Step 2

Next, we perform a nearest-neighbor matching based on the two estimated scores
from the previous step. We adapt the notation from [2], and use the standard
Euclidean norm as the distance metric in the matching algorithm. Formally, for
the units $i$ and $j$ with estimated propensity scores
$\hat{e}_{i},\hat{e}_{j}$ and propensity scores $\hat{p}_{i},\hat{p}_{j}$, we
define the score-based Euclidean distance between $i$ and $j$ by

$d(i,j)=\sqrt{(\hat{e}_{i}-\hat{e}_{j})^{2}+(\hat{p}_{i}-\hat{p}_{j})^{2}}.$

Let $j_{k}(i)$ be the index $j\in\\{1,2,...,n\\}$ that solves $Z_{j}=1-Z_{i}$
and

$\sum_{l:Z_{l}=1-Z_{i}}\textbf{1}\left\\{d(l,i)\leq d(j,i)\right\\}=k,$

where $\textbf{1}\\{\cdot\\}$ is the indicator function. This is the index of
the unit that is the $k$th closest to unit $i$ in terms of the distance
between two scores, among the units with the treatment opposite to that of
unit $i$. We can now construct the $K$-nearest-neighbor set for unit $i$ by
the set of indices for the first $K$ matches for unit $i$,

$\mathcal{J}_{K}(i)=\\{j_{1}(i),...,j_{K}(i)\\}.$

We then compute

$\tilde{Y}_{i}=(2Z_{i}-1)\left(Y_{i}-\frac{1}{K}\sum_{j\in\mathcal{J}_{K}(i)}Y_{j}\right).$
(7)

Intuitively, the construction of $\tilde{Y}$ gives a proxy of the individual
treatment effect (ITE) on each unit. We find $K$ matches for each unit in the
opposite treatment group based on the similarity of their propensity and
prognostics scores, and the mean of the $K$ matches is used to estimate the
unobserved potential outcome for each unit.

### Step 3

The last step involves denoising of the point estimates of the individual
treatment effects $\tilde{Y}$ obtained from Step 2. The goal is to partition
all units into subgroups such that the estimated treatment effects would be
constant over some 2d intervals of propensity and prognostic scores (see the
left of Figure 1).

Figure 1: Left:  a hypothetical partition over the 2d space of propensity and
prognostic scores with the true values of piecewise constant treatment
effects; Right: a sample regression tree $T$ constructed in Step 3.

To perform such stratification, we grow a regression tree on $\tilde{Y}$,
denoted by $T$, and the regressors are the estimated propensity scores
$\hat{e}$ and the estimated prognostic scores $\hat{p}$ from Step 1. We follow
the very general rule of binary recursive partitioning to build the tree $T$:
allocate the data into the first two branches, using every possible binary
split on every covariate; select the split that minimizes Gini impurity,
continue the optimal splits over each branch along the covariate’s values
until the minimum node size is reached. To avoid overfitting, we set the
minimum node size as 20 in our model. Choosing other criteria such as
information gain instead of Gini impurity is another option for splitting
criteria. A 10-fold cross validation is also performed at meantime to prune
the large tree $T$ for deciding the value of cost complexity. Cost complexity
is the minimum improvement in the model needed at each node. The pruning rule
follows that if one split does not improve the overall error of the model by
the chosen cost complexity, then that split is decreed to be not worth
pursuing. [35, See more details in Section 9.2 of]

The final tree $T$ (see the right plot of Figure 1) contains a few terminal
nodes, and these are the predicted treatment effects for all units in the
data. The values exactly represent a piecewise constant stratification over
the 2d space of propensity and prognostic scores.

### Estimation on a new unit

After we obtain the regression tree model $T$ in Step 3, we can now estimate
the value of the individual treatment effect corresponding to a new unit with
covariate $x_{\text{new}}$.

We first compute the estimated propensity and prognostic scores for the new
observation by

$\hat{e}_{\text{new}}=\sigma(x_{\text{new}}^{\top}\hat{\alpha}),\
\hat{p}_{\text{new}}=x_{\text{new}}^{\top}\hat{\theta},$

where $\hat{\alpha}$ and $\hat{\theta}$ are the solutions to Equations (2) and
(3) respectively. Then with the estimated propensity score
$\hat{e}_{\text{new}}$ and prognostic score $\hat{p}_{\text{new}}$, we can get
an estimate of the treatment effect for this unit following the binary
predictive rules in the tree $T$.

### High-Dimensional Estimator

In a high-dimensional setting where the covariate dimension $d$ is much larger
than the sample size $n$, we can estimate the propensity and prognostic scores
by adding a lasso ($l1$-based) penalty [80] instead. This strategy was first
proposed and named as "doubly robust matching estimators" (DRME) by [3]. The
corresponding optimization problems for the two scores can be written as

$\displaystyle\hat{\alpha}$
$\displaystyle=\underset{\alpha\in\mathbb{R}^{d}}{\arg\min}\
-\sum_{i=1}^{n}\Big{[}Y_{i}\log\sigma(X_{i}^{\top}\alpha)+(1-Y_{i})\log(1-\sigma(X_{i}^{\top}\alpha))\Big{]}+\lambda_{1}\sum_{j=1}^{d}|\alpha_{j}|,$
$\displaystyle\hat{\theta}$
$\displaystyle=\underset{\theta\in\mathbb{R}^{d}}{\arg\min}\
\sum_{i:Z_{i}=0}(Y_{i}-X_{i}^{\top}\theta)^{2}+\lambda_{2}\sum_{j=1}^{d}|\theta_{j}|.$

The selection of the tuning parameter $\lambda_{1},\lambda_{2}$ can be
determined by any information criteria (AIC, BIC, and etc.). In practice, we
use $10$-fold cross-validation (CV) to select the value of $\lambda$. Then we
perform the $K$-nearest-neighbor matching based on propensity and prognostic
scores to get the estimates of individual treatment effects using Equation
(7).

We extend the above estimator with our proposal of applying a regression tree
on the estimated propensity and prognostic scores. The procedure for
estimation of subgroup heterogeneous treatment effects and estimation on a new
unit remains the same as Step 3 for low-dimensional set-ups.

######  Remark.

The choice of the number of nearest neighbors is a challenging problem.
Updating distance metrics for every observation is computationally expensive,
and choosing a value that is too small leads to a higher influence of noise on
estimation. With regards to the application of nearest neighbor matching in
causal inference, [2] derived large sample properties of matching estimators
of average treatment effects with a fixed number of nearest neighbors, but the
authors did not provide any details on how to select the exact number of
neighbors. Conventional settings on the number of nearest neighbors in current
literature is to set $K=1$ [74, 8, one-to-one matching;]. However, [55]
suggested that in observational studies, substantially greater bias reduction
is possible through matching with a variable number of controls rather than
exact pair matching.

In Appendix A, we conduct a simulation study following one of the generative
models from Section 5 to show how sensitive estimation accuracy is to the
number of nearest neighbors selected and setting $K$ to a large number other
than 1 is more ‘sensible’ to reduce estimation bias. Although it is usually
difficult to select a perfect value of $K$ in practice, simply setting
$K\approx\log(n)$ as suggested by [15] leads to reasonable results for a data
sample of size $n$. Throughout all our experiments in the next section,
setting $K$ to the integer closest to the value $\log(n)$ provides estimates
with high accuracy and does not require too much computational cost.

### Computational Complexity

Our method is composed of three steps as introduced above. We first need to
implement a logistic regression for estimating propensity score for a sample
with size $n$ and ambient dimension $d$. This computation has a complexity of
$O(nd)$. The estimation of prognostic scores requires a complexity of
$O(nd^{2})$ when $n>d$ and of $O(d^{3})$ for high-dimensional settings [27].
The complexity of a $K$-nearest-neighbor matching based on the two estimated
scores in the second step is of $O(Kn)$ [53], and the selection of
$K\approx\log(n)$ leads to a complexity of $O(n\log n)$. In the third step, we
grow a regression tree based on two estimated scores, and it requires a
computational complexity of $O(n\log n)$.

Overall, the eventual computational complexity of our method depends on the
comparison between the order of $d^{2}$ and $\log(n)$. For the settings where
the sample size $n$ is greater than the ambient dimension $d$, our method
attains a computational complexity of $O(nd^{2})$ if the order of $d^{2}$ is
greater than that of $\log(n)$. Otherwise, the complexity becomes $O(n\log
n)$. For high-dimensional settings, we have $d>n$, and the order of $d^{3}$ is
greater than that of $n\log(n)$. Hence, the resulting computational complexity
of our method is $O(d^{3})$.

##  5 Experiments

In this section, we will examine the performance of our proposed estimator
(PP) in a variety of simulated and real data sets. The baseline estimators we
compete against are leave-one-out endogenous stratification [1, ST;], causal
forest [81, CF;], single-score matching including propensity-score matching
(PSM) and prognostic-score matching. Note that in the original research by
[1], the authors restricted their attention to randomized experiments, because
this is the setting where endogenous stratification is typically used.
However, they mentioned the possibility of applying the method on
observational studies. We take this into consideration, and make their method
as one of our competitors.

We implement our methods in R, using the packages "FNN" [10] for $K$-nearest-
neighbor matching and "rpart" for growing a non-parametric regression tree.
Throughout, we set the number of nearest neighbors, $K$, to be the closest
integer to $\log(n)$, where $n$ is the sample size. Regression tree pruning is
set as the default in the package. For causal forest, we directly use the R
package "grf" developed by [6], following with a default selection of the
minimum leaf size $k=1$ and the number of trees $B=2000$. Software that
replicate all the simulations is available on the authors’ Github page.

We evaluate the performance of each method according to two aspects, accuracy
and uncertainty quantification. The results for single-score matching
algorithms are not reported in this paper because of very poor performance
throughout all scenarios.

###  5.1 Simulated Data

We first examine on the following simulated data sets under six different data
generation mechanisms. We get insights from the simulation study in [50] for
the models considered in Scenarios 1-4. The propensity score and outcome
(prognosis) models in Scenarios 1 and 4 are characterized by additivity and
linearity (main effects only), but with different piecewise constant
structures in the true treatment effects over a 2d grid of the two scores. We
add non-additivity and non-linear terms to both propensity and prognosis
models in Scenarios 2 and 3. In other words, both propensity and prognostic
scores are expected to be misspecified in these two models if we apply
generalized linear models directly in estimation. Scenario 5 comes from [1],
with a constant treatment effect over all observations. Scenario 6 is
considered in [81] (see Equation 27 there), in which the propensity model
follows a continuous distribution instead of a linear structure. A high-
dimensional setting ($d>>n$) is examined in Scenario 7, where the generative
model inherits from Scenario 1.

We first introduce some notations used in the experiments: the sample size
$n$, the ambient dimension $d$, as well as the following functions:

$\displaystyle\text{true treatment effect:
}\tau^{*}(X)=\mathbb{E}\left[Y(1)-Y(0)\big{|}X\right],$
$\displaystyle\text{treatment propensity: }e(x)=\mathbb{P}(Z=1|X=x),$
$\displaystyle\text{treatment logit:
}\text{logit}(x)=\log\left(\frac{e(x)}{1-e(x)}\right).$

Throughout all the models we consider, we mantain the unconfoundedness
assumption discussed in Section 3, generate the covariate $X$ following a
certain distribution, and entail homoscedastic Gaussian noise $\epsilon$.

We evaluate the accuracy of an estimator $\tau(X)$ by the mean-squared error
for estimating $\tau^{*}(X)$ at a random example $X$, defined by

$\text{MSE}(\hat{\tau}(X)):=\frac{1}{n}\sum_{i=1}^{n}\left[\hat{\tau}_{i}(X)-\tau_{i}^{*}(X)\right]^{2}.$

We record the averaged MSE over 100 Monte Carlo trials for each scenario. In
terms of uncertainty quantification, we measure the coverage probability of
$\tau(X)$ with a target coverage rate of 0.95. For endogenous stratification
and our proposed method, we use non-parametric bootstrap to construct the
empirical quantiles for each unit. The details on the implementation of non-
parametric bootstrap methods are presented in Appendix B. For causal forest,
we construct 95$\%$ confidence intervals by estimating the standard errors of
estimation using the "grf" package.

Scenario 1. With $d\in\\{2,10,50\\}$, $n\in\\{1000,5000\\}$, for $i=1,...,n$,
we generate the data as follows:

$\displaystyle Y_{i}=p(X_{i})+Z_{i}\cdot\tau_{i}^{*}+\epsilon_{i},$
$\displaystyle\tau_{i}^{*}=\textbf{1}_{\\{e(X_{i})<0.6,p(X_{i})<0\\}},$
$\displaystyle\text{logit}(X_{i})=X_{i}^{\top}\beta^{e},$ $\displaystyle
p(X_{i})=X_{i}^{\top}\beta^{p},$ $\displaystyle
X_{i}\overset{i.i.d.}{\sim}\mathcal{U}[0,1]^{d},$
$\displaystyle\epsilon_{i}\overset{i.i.d.}{\sim}\mathcal{N}(0,1),$

where

$\beta^{e},\beta^{p}\overset{i.i.d.}{\sim}\mathcal{U}[-1,1]^{d}.$

Scenario 2. We now add some interaction terms to the propensity and prognostic
models in Scenario 1, while keeping the set-ups of the covariate $X$, the
response $Y$, the true treatment effect $\tau^{*}$ and the error term
$\epsilon$ unchanged. We set $d=10$ and $n=3000$ in this case.

$\displaystyle\text{logit}(X_{i})$
$\displaystyle=X_{i}^{\top}\beta^{e}+0.5X_{i1}X_{i3}+0.7X_{i2}X_{i4}+0.5X_{i3}X_{i5}$
$\displaystyle\,\,\,\ +0.7X_{i4}X_{i6}+0.5X_{i5}X_{i7}+0.5X_{i1}X_{i6}$
$\displaystyle\,\,\,\ +0.7X_{i2}X_{i3}+0.5X_{i3}X_{i4}+0.5X_{i4}X_{i5}$
$\displaystyle\,\,\,\ +0.5X_{i5}X_{i6}$ $\displaystyle p(X_{i})$
$\displaystyle=X_{i}^{\top}\beta^{p}+0.5X_{i1}X_{i3}+0.7X_{i2}X_{i4}+0.5X_{i3}X_{i8},$
$\displaystyle\,\,\,\ +0.7X_{i4}X_{i9}+0.5X_{i8}X_{i10}+0.5X_{i1}X_{i9}$
$\displaystyle\,\,\,\ +0.7X_{i2}X_{i3}+0.5X_{i3}X_{i4}+0.5X_{i4}X_{i8}$
$\displaystyle\,\,\,\ +0.5X_{i8}X_{i9}.$

Scenario 3. Similar to Scenario 2, we add some nonlinear terms to the model in
Scenario 1, with $d=10$ and $n=3000$, as follows:

$\displaystyle\text{logit}(X_{i})$
$\displaystyle=X_{i}^{\top}\beta^{e}+X_{i2}^{2}+X_{i4}^{2}-X_{i7}^{2},$
$\displaystyle p(X_{i})$
$\displaystyle=X_{i}^{\top}\beta^{p}+X_{i2}^{2}+X_{i4}^{2}-X_{i10}^{2}.$

Scenario 4. In this case, we define the true treatment effect with a more
complicated piecewise constant structure over the 2d grid, under the same
model used in Scenario 1, with $d=10$ and $n=3000$:

$\displaystyle\tau_{i}^{*}$ $\displaystyle=\begin{cases}0&\text{if
}e(X_{i})\leq 0.6,p(X_{i})\leq 0,\\\ 1&\text{if }e(X_{i})\leq
0.6,p(X_{i})>0\text{ or }e(X_{i})>0.6,p(X_{i})\leq 0,\\\ 2&\text{if
}e(X_{i})>0.6,p(X_{i})>0.\end{cases}$

Scenario 5. Setting $d=10$ and $n=4000$, the data is generated as:

$\displaystyle Y_{i}$ $\displaystyle=1+\beta^{\top}X_{i}+\epsilon_{i},$
$\displaystyle X_{i}$
$\displaystyle\overset{i.i.d.}{\sim}\mathcal{N}(0,\textbf{I}_{d\times d}),$
$\displaystyle\epsilon_{i}$
$\displaystyle\overset{i.i.d.}{\sim}\mathcal{N}(0,100-d),$

where $\beta=(1,...,1)^{\top}\in\mathbb{R}^{d}$. Moreover, the treatment
indicators for the simulations are such that $\sum_{i}Z_{i}=\lceil n/2\rceil$.
By construction, the vector of treatment effects satisfies $\tau^{*}=0$.

Scenario 6. The data satisfies

$\displaystyle Y_{i}$
$\displaystyle=2X_{i}^{\top}\textbf{e}_{1}-1+\epsilon_{i},$ $\displaystyle
Z_{i}$ $\displaystyle\sim\text{Binom}(1,e(X_{i})),$ $\displaystyle X_{i}$
$\displaystyle\overset{i.i.d.}{\sim}\mathcal{U}[0,1]^{d},$ $\displaystyle
e(X_{i})$
$\displaystyle=\frac{1}{4}[1+\beta_{2,4}(X_{i}^{\top}\textbf{e}_{1})],$
$\displaystyle\epsilon_{i}$
$\displaystyle\overset{i.i.d.}{\sim}\mathcal{N}(0,1),$

where $\textbf{e}_{1}=(1,0,...,0)$. We compare the performance of different
methods under two settings: $d=2,n=1000$ and $d=10,n=3000$. Note that in this
data model, $\tau_{i}^{*}=0$ for all $i\in\\{1,...,n\\}.$

Scenario 7. In the last case, we study the performance of different estimators
on a high-dimensional data. The data model follows

$\displaystyle X_{i}\overset{i.i.d.}{\sim}\mathcal{U}[0,1]^{d},$
$\displaystyle Y_{i}=p(X_{i})+Z_{i}\cdot\tau_{i}^{*}+\epsilon_{i},$
$\displaystyle\tau_{i}^{*}=\textbf{1}_{\\{e(X_{i})<0.6,p(X_{i})<0\\}},$
$\displaystyle\text{logit}(X_{i})=0.4X_{i1}+0.9X_{i2}-0.4X_{i3}-0.7X_{i4}-0.3X_{i5}+0.6X_{i6},$
$\displaystyle
p(X_{i})=0.9X_{i1}-0.9X_{i2}+0.2X_{i3}-0.2X_{i4}+0.9X_{i5}-0.9X_{i6},$
$\displaystyle\epsilon_{i}\overset{i.i.d.}{\sim}\mathcal{N}(0,1).$

Figure 2: Comparison of mean squared errors over 100 Monte Carlo simulations
under different generative models for our proposed method (PP), endogenous
stratification (ST) and causal forest (CF).

We select $n=3000$ and $d=5000$ for examination.

The boxplots that depict the distribution of MSEs obtained under all scenarios
are presented in Figure 2. We can see that for Scenario 1, our proposed
estimator achieves better accuracy when the sample size $n$ is large, and it
is the best among the three estimators in these cases. The good performance of
our method under a large-size setting is consistent when we assume a more
complex partition on the defined 2d grid. In addition, variation in accuracy,
measured by the difference between the upper quartiles and the lower quartiles
(also referred as the interquartile ranges) in each boxplot becomes smaller,
accompanied with the increase in $d$ and $n$. In Scenarios 2 and 3, we
introduce non-additivity and non-linear terms into the data model. Although
linear assumptions are violated for both propensity and prognostic models, our
method performs better compared to the other two methods regarding accuracy
and variability. For a potential outcome model with randomized assignment of
treatment and constant treatment effects, as in Scenario 5, our method still
has the best accuracy compared to the benchmarks, even though large noise is
added to the true signal. Only in Scenario 6 where we assume a continuous
distribution on the propensity model, causal forests outperform our estimator
in terms of variation. In a high-dimensional setting such as Scenario 7, we
consider modified methods with lasso-regularized regressions for both our
methodology and endogenous stratification, and our method maintains its
superiority as in the low-dimensional set-ups.

In summary, our proposed method achieves a comparably good accuracy, with the
smallest variance across 100 Monte Carlo simulations in most cases.

We now take a careful look at the visualization comparison between the true
treatment effect and the predictions obtained from our method for Scenarios 1
and 4. We confine both the true signal and the predictive model in a 2d grid
scaled on the true propensity and prognostic scores, as shown in Figure 3. It
is not surprising that our proposed estimators provide a descent recovery of
the piecewise constant partition in the true treatment effects over the 2d
grid, with only a small difference in the magnitude of treatment effects.

(a) $\tau^{*}_{0}$ of Scenario 1

(b) Prediction on Scenario 1

(c) $\tau^{*}_{0}$ of Scenario 4

(d) Prediction on Scenario 4

Figure 3: One instance comparison between the true treatment effects and the
estimates from propensity-prognostic score based method.

With regard to uncertainty quantification, we examine coverage rates with a
target confidence level of 0.95 for each method under different scenarios, and
the corresponding results are recorded in Table 1. It is quite clear that our
proposed method achieves nominal coverage over the other two methods in almost
all scenarios. Considering the small variation in accuracy as shown in the
boxplots above for most scenarios, our method is the most robust one among the
three candidates.

Table 1: Reported coverage rate with a target confidence level of 0.95. Scenario | $n$ | $d$ | PP | ST | CF | $n$ | $d$ | PP | ST | CF
---|---|---|---|---|---|---|---|---|---|---
1 | 1000 | 2 | 0.981 | 0.424 | 0.719 | 5000 | 10 | 0.978 | 0.791 | 0.317
2 | 1000 | 10 | 0.988 | 0.354 | 0.140 | 3000 | 10 | 0.944 | 0.126 | 0.270
3 | 1000 | 10 | 0.999 | 0.915 | 0.449 | 3000 | 10 | 0.976 | 0.907 | 0.293
4 | 1000 | 10 | 0.991 | 0.639 | 0.606 | 3000 | 10 | 0.985 | 0.522 | 0.528
5 | 1000 | 10 | 1.000 | 1.000 | 0.997 | 4000 | 10 | 1.000 | 1.000 | 0.980
6 | 1000 | 2 | 0.998 | 1.000 | 0.928 | 3000 | 10 | 1.000 | 1.000 | 0.992
7 | 1000 | 2000 | 0.816 | 0.751 | 0.726 | 3000 | 5000 | 0.745 | 0.578 | 0.441

###  5.2 Real Data Analysis

To illustrate the behavior of our estimator, we apply our method on the two
real-world data sets, one from a clinical study and the other from a complex
social survey. Propensity score based methods are frequently used for
confounding adjustment in observational studies, where baseline
characteristics can affect the outcome of policy interventions. Therefore, the
results from our method are expected to provide meaningful implications for
these real data sets. However, due to the complicated sampling nature of
complex survey, we will take extra care on dealing with cluster sampling
weight in order to apply our score-based method.

####  5.2.1 Right Heart Catheterization Analysis

While randomized control trials (RCT) are widely encouraged as the ideal
methodology for causal inference in clinical and medical research, the lack of
randomized data due to high costs and potential high risks leads to the
studies based on observational data. In this section, we are interested in
examining the association between the use of right heart catheterization (RHC)
during the first 24 hours of care in the intensive care unit (ICU) and the
short-term survival conditions of the patients. Right Heart Catheterization
(RHC) is a procedure for directly measuring how well the heart is pumping
blood to the lungs. RHC is often applied to critically ill patients for
directing immediate and subsequent treatment. However, RHC imposes a small
risk of causing serious complications when administering the procedure.
Therefore, the use of RHC is controversial among practitioners, and scientists
want to statistically validate the causal effects of RHC treatments. The
causal study using observational data can be dated back to [18], where the
authors implemented propensity score matching and concluded that RHC treatment
lead to lower survival than not performing the treatment. Later, [39] proposed
a more efficient propensity-score based method and the recent study by [51]
using a modified propensity score model suggested RHC significantly affected
mortality rate in a short-term period.

A dataset for analysis was first used in [18], and it is suitable for the
purpose of applying our method because of its its extremely well-balanced
distribution of confounders across levels of the treatment [73]. The treatment
variable $Z$ in the data indicates whether or not a patient received a RHC
within 24 hours of admission. The binary outcome $Y$ is defined based on
whether a patient died at any time up to 180 days since admission. The
original data consisted of 5735 participants with 73 covariates. We preprocess
the full data in the way suggested in [39] and [51], by removing all
observations that contain null values in covariates, dropping the singular
covariate in the reduced data, and encoding categorical variables into dummy
variables. The resulted data contains 2707 observations and 72 covariates,
with 1103 in the treated group ($Z=1$) and 1604 in the control group ($Z=0$).
Among the 72 observed covariates, there are 21 continuous, 25 binary, and 26
dummy variables transformed from the original 6 categorical variables.

Figure 4: Prediction model for Right Heart Catheterization (RHC) data.

The result of the prediction model from our proposed method is reported in
Figure 4 below. We observe that the sign of estimated treatment effects varies
depending on the value of the propensity score and prognostic score. This
particular pattern implies that RHC procedures indeed offer both benefits and
risks in affecting patients’ short-term survival conditions. Specifically, we
are interested in the occurrence of large positive treatment effects (increase
in chance of death) from the estimation. An estimated treatment effect of 0.16
is observed on the group of patients with propensity scores less than 0.62 and
prognostic scores less than 0.74, and this group accounts for 55$\%$ of the
entire sample. Under the scenario of RHC data, a smaller propensity score
means that the patient is less likely to receive RHC procedures after
admitting to the ICU, and it is related to the availability of RHC procedures
at the hospital to which the patient is admitted. A smaller prognostic score
tells that the patient has lower underlying chance of death. One possible
explanation for this significant positive treatment on this certain group is
that drastic change in treatment procedures that were applied to patients who
do not actually need the aggressive style of care largely undermine patients’
health conditions after admission and increase the mortality rate. Another
large positive treatment effect is found on the group with propensity scores
greater than 0.81 and prognostic scores less than 0.48. This would be
consistent with the findings of [12], where the authors found that the
hospitals with the higher than predicted use of RHC had higher than expected
than expected mortality rates. In summary, our findings generally agree with
the results and explanations in [18] and they offer some insights for
practitioners to decide whether they should apply RHC procedures to patients.

####  5.2.2 National Medical Expenditure Survey

For the next experiment, we analyze a complex social survey data. In many
complex surveys, data are not usually well-balanced due to potential biased
sampling procedure. To incorporate score-based methods with complex survey
data requires an appropriate estimation on propensity and prognostic scores.
[25] suggested that combining a propensity score method and survey weighting
is necessary to achieve unbiased treatment effect estimates that are
generalizable to the original survey target population. [7] conducted
numerical experiments and showed that greater balance in measured baseline
covariates and decreased bias is observed when natural retained weights are
used in propensity score matching. Therefore, we include sampling weight as an
baseline covariate when estimating propensity and prognostic scores in our
analysis.

In this study, we aim to answer the research question: how one’s smoking habit
affects his or her medical expenditures over lifetime, and we use the same
data set as in [46], which is originally extracted from the 1987 National
Medical Expenditure Survey (NMES). The NMES included detailed information
about frequency and duration of smoking with a large nationally representative
data base of nearly 30,000 adults, and that 1987 medical costs are verified by
multiple interviews and additional data from clinicians and hospitals. A large
amount of literature focus on applying various statistical methods to analyze
the causal effects of smoking on medical expenditures using the NMES data. In
the original study by [46], the authors first estimated the effects of smoking
on certain diseases and then examined how much those diseases increased
medical costs. In contrast, [68], [41], and [85] proposed to directly estimate
the effects of smoking on medical expenditures using propensity-score based
matching and subclassification. [32] applied Bayesian regression tree models
to assess heterogeneous treatment effects.

For our analysis, we explore the effects of extensive exposure to cigarettes
on medical expenditures, and we use pack-year as a measurement of cigarette
measurement, the same as in [41] and [32]. Pack-year is a clinical
quantification of cigarette smoking used to measure a person’s exposure to
tobacco, defined by

$\text{pack-year}=\frac{\text{number of cigarettes per
day}}{20}\times\text{number of years smoked}.$

Following that, we determine the treatment indicator $Z$ by the question
whether the observation has a heavy lifetime smoking habit, which we define to
be greater than 17 pack-years, the equivalent of 17 years of pack-a-day
smoking.

The subject-level covariates $X$ in our analysis include age at the times of
the survey (between 19 and 94), age when the individual started smoking,
gender (male, female), race (white, black, other), marriage status (married,
widowed, divorced, separated, never married), education level (college
graduate, some college, high school graduate, other), census region
(Northeast, Midwest, South, West), poverty status (poor, near poor, low
income, middle income, high income), seat belt usage (rarely, sometimes,
always/almost always), and sample weight. We select the natural logarithm of
annual medical expenditures as the outcome variable $Y$ to maintain the
assumption of heteroscedasticity in random errors. We preprocess the raw data
set by omitting any observations with missing values in the covariates and
excluding those who had zero medical expenditure. The resulting restricted
data set contains 7903 individuals, with 4014 in the treated group ($Z=1$) and
3889 in the controlled group ($Z=0$).

Figure 5: Prediction model for National Medical Expenditure Survey (NMES)
data.

The prediction model obtained from our method, as shown in Figure 5, is simple
and easy to interpret. We derive a positive treatment effect across the entire
sample, and the effect becomes significant when the predicted potential
outcome is relatively low (less than 5.8). These results indicate that more
reliance on smoking will deteriorate one’s health condition, especially for
those who currently do not have a large amount of medical expenditure.
Moreover, we observe a significant positive treatment effect of 1.4, in other
word, a substantial increase in medical expenditure for the subgroup with
propensity score less than 0.18. It is intuitive to assume that a smaller
possibility of engaging in excessive tobacco exposure is associated with
healthier living styles. This phenomenon is another evidence that individuals
who are more likely to stay healthy may suffer more from excessive exposure to
tobacco products. In all, these results support policymakers and social
activists who advocate for nationwide smoking ban.

##  6 Conclusions

Our method is different from existing methods on estimating heterogeneous
treatment effects in a way that we incorporate both matching algorithms and
non-parametric regression trees in estimation, and the final estimate can be
regarded as a 2d summary on treatment effects. Moreover, our method exercises
a simultaneous stratification across the entire population into subgroups with
the same treatment effects. Subgroup treatment effect analysis is an important
but challenging research topic in observational studies, and our method can be
served as an efficient tool to reach a reasonable partition.

Our numerical experiments on various simulated and real-life data lay out
empirical evidence of the superiority of our estimator over state-of-the-art
methods in both accuracy and interpretability. We also discovered that our
method is powerful in investigating subpopulations with significant treatment
effects. Identifying representative subpopulations that receive extreme
results after treatment is a paramount task in many practical contexts.
Through empirical experiments on two real-world data sets from observational
studies, our method demonstrates its ability in identifying these significant
effects.

Although our method shows its outstanding performance in estimating treatments
effects under the piecewise constant structure assumption, it remains
meaningful and requires further study to develop more accurate recovery of
such structure. For example, a potential shortcoming of using conventional
regression trees for subclassification is that the binary partition over the
true signals is not necessarily unique. Using some variants of CART, like
optimal trees [9] and dyadic regression trees [24], would be more appropriate
for estimation under additional assumptions. Applying other non-parametric
regression techniques, such as $K$-nearest-neighbor fused lasso [58], is
another direction if we assume a more complicated piecewise constant structure
in treatment effects other than a rectangular partition on 2d data. It is also
worth improving the estimation of propensity and prognostic scores using
similar non-parametric based methods if a piecewise constant assumption hold
for the two scores as well.

##  Appendix A Study on the choice of the number of nearest neighbors

In this section, we examine how the number of nearest neighbors in the
matching algorithm affects the estimation accuracy. Recall in Step 2 of our
proposed method, we implement a $K$-nearest-neighbor algorithm based on the
two estimated scores for a sample of size $n$. The computational complexity of
this $K$-NN algorithm is of $O(Kn)$. Although a larger $K$ typically leads to
a higher estimation accuracy, more computational costs become the
corresponding side-effect. Therefore, a smart choice of $K$ is essential to
balance the trade-off between accuracy and computational expense.

We follow the same generative model in Scenario 1 from Section 5 and compute
the averaged mean squared error over 100 Monte Carlo simulations for
$K=1,...,50$ with a fixed sample size $n=5000$. The results in Figure 6 show
that the averaged MSE continuously decreases as the number of nearest
neighbors $K$ selected in the matching algorithm grows. However, the speed of
improvement in accuracy gradually slows down when $K$ exceeds 10, which is
close to $\log(5000)$. This suggests that an empirical choice of
$K\approx\log(n)$ is sufficient to produce a reasonable estimate on the target
parameter and this choice is more ’sensible’ than the conventional setting of
$K=1$.

Figure 6: The plot of averaged MSE against the number of nearest neighbors.

##  Appendix B Non-parametric Bootstrap in Simulation Studies

In Section 5, we use non-parametric bootstrap to construct confidence
intervals for endogenous stratification and our proposed method. We use these
bootstrap samples to compute coverage rates with respect to a target level of
$95\%$ as a measurement of uncertainty. The bootstrap method, introduced by
[26], is a simple but powerful tool to obtain a robust non-parametric estimate
of the confidence intervals through sampling from the empirical distribution
function of the observed data. In this appendix, we introduce the details on
how we implement non-parametric bootstrap for the purpose of computing
coverage rates in the simulation experiments.

For each scenario in Section 5, we start with generating a sample following
the defined data generation model with a sample size $n$. Next, we create 1000
random resamples with replacement from this single set of data, also with the
sample size $n$. We then apply both methods on these simulation repetitions,
and obtain a series of estimations on each unit in the original set. Following
these estimations, we calculate the corresponding $2.5\%$ and $97.5\%$
quantiles for all units in the original sample. Coverage rates of a $95\%$
confidence level are thus the frequencies of the original units falling inside
the intervals between the two quantiles computed in the previous step.

## References

  * [1] A. Abadie, M.. Chingos and M.. West  “Endogenous stratification in randomized experiments theory”  In _The Review of Economics and Statistics_ C.4, 2018, pp. 567–580
  * [2] A. Abadie and G.. Imbens  “Large sample properties of matching estimators for average treatment effects”  In _Econometrica_ 74.1, 2006, pp. 235–267
  * [3] J. Antonelli, M. Cefalu, N. Palmer and D. Agniel  “Doubly robust matching estimators for high dimensional confounding adjustment”  In _Biometrics_ 74.4, 2018, pp. 1171–1179
  * [4] S.. Assmann, S.. Pocock, L.. Enos and L.. Kasten  “Subgroup analysis and other (mis)uses of baseline data in clinical trials”  In _The Lancet_ 355.9209, 2000, pp. 1171–1179
  * [5] S. Athey and G.. Imbens  “Recursive partitioning for heterogeneous causal effects”  In _Proceedings of the National Academy of Sciences_ 113.27, 2016, pp. 7353–7360
  * [6] S. Athey, J. Tibshirani and S. Wager  “Generalized random forests”  In _The Annals of Statistics_ 47.2, 2019, pp. 1148–1178
  * [7] P.. Austin, N. Jembere and M. Chiu  “Propensity score matching and complex surveys”  In _Statistical Methods in Medical Research_ 27.4, 2018, pp. 1240–1257
  * [8] P.. Austin and T. Schuster  “The performance of different propensity score methods for estimating absolute effects of treatments on survival outcomes: a simulation study”  In _Statistical Methods in Medical Research_ 25.5, 2016, pp. 2214–2237
  * [9] D. Bertsimas and J. Dunn  “Optimal classification trees”  In _Machine Learning_ 106.7, 2017, pp. 1039–1082
  * [10] A. Beygelzimer et al.  “FNN: fast nearest neighbor search algorithms and applications”, 2013
  * [11] A. Bloniarz et al.  “Lasso adjustments of treatment effect estimates in randomized experiments”  In _Proceedings of the National Academy of Sciences_ 113.27, 2016, pp. 7383–7390
  * [12] M.. Blumberg and G.. Binns  “Swan-Ganz catheter use and mortality of myocardial infarction patients”  In _Health Care Financing Review_ 15.4, 1994, pp. 91–103
  * [13] L. Breiman  “Random forests”  In _Machine Learning_ 45, 2001, pp. 5–32
  * [14] L. Breiman, J. Friedman, C.. Stone and R.. Olshen  “Classification and regression trees”  Taylor $\&$ Francis, 1984
  * [15] M.. Brito, E.. Chávez, A.. Quiroz and J.. Yukich  “Connectivity of the mutual k-nearest-neighbor graph in clustering and outlier detection”  In _Statistics $\&$ Probability Letters_ 35.1, 1997, pp. 33–42
  * [16] H.. Chipman, E.. George and R.. McCulloch  “BART: Bayesian additive regression trees”  In _The Annals of Applied Statistics_ 4.1, 2010, pp. 266–298
  * [17] W.. Cochran  “The effectiveness of adjustment by subclassification in removing bias in observational studies”  In _Biometrics_ 24.2, 1968, pp. 295–313
  * [18] Jr Connors et al.  “The effectiveness of right heart catheterization in the initial care of critically ill patients”  In _Journal of the American Medical Association_ 276.11, 1996, pp. 889–897
  * [19] A. D’Amour et al.  “Overlap in observational studies with high-dimensional covariates”  In _Journal of Econometrics_ 221, 2021, pp. 644–654
  * [20] I.. Dahabreh, R. Hayward and D.. Kent  “Using group data to treat individuals: understanding heterogeneous treatment effects in the age of precision medicine and patient-centred evidence”  In _International Journal of Epidemiology_ 45.6, 2016, pp. 2184–2193
  * [21] A.. Dawid  “Causal inference without counterfactuals”  In _Journal of the American Statistical Association_ 95.450, 2000, pp. 407–424
  * [22] R.. Dehejia and S. Wahba  “Propensity score matching methods for nonexperimental causal studies”  In _Review of Economics and Statistics_ 84.1, 2002, pp. 151–161
  * [23] P. Ding and F. Li  “Causal inference: a missing data perspective”  In _Statistical Science_ 33.2, 2018, pp. 214–237
  * [24] D.. Donoho  “CART and best-ortho-basis: a connection”  In _The Annals of Statistics_ 25.5, 1997, pp. 1870–1911
  * [25] E.. DuGoff, M. Schuler and E.. Stuart  “Generalizing observational study results: applying propensity score methods to complex surveys”  In _Health Service Research_ 49.1, 2014, pp. 284–303
  * [26] B. Efron  “Bootstrap methods: another look at the Jackknife”  In _The Annals of Statistics_ 7.1, 1979, pp. 1–26
  * [27] B. Efron, T. Hastie, I. Johnstone and R. Tibshirani  “Least angle regression”  In _The Annals of Statistics_ 32.2, 2004, pp. 407–499
  * [28] B. Gaines and J. Kuklinski  “Estimation of heterogeneous treatment effects related to self-selection”  In _American Journal of Political Science_ 55.3, 2011, pp. 724–736
  * [29] D.. Green and H.. Kern  “Modeling heterogeneous treatment effects in survey experiments with Bayesian additive regression trees”  In _The Public Opinion Quarterly_ 76.3, 2012, pp. 491–511
  * [30] X.. Gu and P.. Rosenbaum  “Comparison of multivariate matching methods: structures, distances, and algorithms”  In _Journal of Computational and Graphical Statistics_ 2.4, 1993, pp. 405–420
  * [31] J. Hahn  “On the role of the propensity score in efficient semiparametric estimation of average treatment effects”  In _Econometrica_ 66.2, 1998, pp. 315–331
  * [32] P.. Hahn, J.. Murray and C.. Carvalho  “Bayesian regression tree models for causal inference: regularization, confounding, and heterogeneous effects”  In _Bayesian Analysis_ 15.3, 2020, pp. 965–1056
  * [33] B.. Hansen  “Full matching in an observational study of coaching for the SAT”  In _Journal of the American Statistical Association_ 99.467, 2004, pp. 609–618
  * [34] B.. Hansen  “The prognostic analogue of the propensity score”  In _Biometrika_ 95.2, 2008, pp. 481–488
  * [35] T. Hastie, R. Tibshirani and J. Friedman  “The elements of statistical learning: data mining, inference, and prediction”  New York: Springer, 2001
  * [36] J.. Heckman, H.. Lopes and R. Piatek  “Treatment effects: a Bayesian perspective”  In _Econometric Reviews_ 33.1-4, 2014, pp. 36–67
  * [37] J.. Hill  “Bayesian nonparametric modeling for causal inference”  In _Journal of Computational and Graphical Statistics_ 20.1, 2011, pp. 217–240
  * [38] J.. Hill and Y. Su  “Assessing lack of common support in causal inference using Bayesian nonparametrics: implications for evaluating the effect of breastfeeding on children’s cognitive outcomes”  In _The Annals of Applied Statistics_ 7.3, 2013, pp. 1386–1420
  * [39] K. Hirano and G.. Imbens  “Estimation of causal effects using propensity Score weighting: an application to data on Right Heart Catheterization”  In _Health Services $\&$ Outcomes Research Methodology_ 2.3, 2001, pp. 259–278
  * [40] P.. Holland  “Statistics and causal Inference”  In _Journal of the American Statistical Association_ 81.396, 1986, pp. 945–960
  * [41] K. Imai and D.. Dyk  “Causal inference with general treatment regimes”  In _Journal of the American Statistical Association_ 99.467, 2004, pp. 854–866
  * [42] K. Imai and M. Ratkovic  “Estimating treatment effect heterogeneity in randomized program evaluation”  In _The Annals of Applied Statistics_ 7.1, 2013, pp. 443–470
  * [43] K. Imai and A. Strauss  “Estimation of heterogeneous treatment effects from randomized experiments, with application to the optimal planning of the get-out-the-vote campaign”  In _Political Analysis_ 19, 2011, pp. 1–19
  * [44] G.. Imbens  “Nonparametric estimation of average treatment effects under exogeneity: a review”  In _Review of Economics and Statistics_ 86, 2004, pp. 4–29
  * [45] G.. Imbens and D.. Rubin  “Causal inference for statistics, social, and biomedical sciences: an introduction”  Cambridge University Press, 2015
  * [46] E. Johnson, F. Dominici, M. Griswold and S.. Zeger  “Disease cases and their medical costs attributable to smoking: an analysis of the national medical expenditure survey”  In _Journal of Econometrics_ 112.1, 2003, pp. 135–151
  * [47] L. Keele  “The statistics of causal inference: a view from political methodology”  In _Political Analysis_ 23.3, 2015, pp. 313–335
  * [48] B. Koch, D.. Vock and J. Wolfson  “Covariate selection with group lasso and doubly robust estimation of causal effects”  In _Biometrics_ 74.1, 2018, pp. 8–17
  * [49] S.. Künzel, J.. Sekhon, P.. Bickel and B. Yu  “Meta-learners for estimating heterogeneous treatment effects using machine learning”  In _Proceedings of the National Academy of Sciences_ 116.10, 2019, pp. 4156–4165
  * [50] F.. Leacy and E.. Stuart  “On the joint use of propensity and prognostic scores in estimation of the average treatment effect on the treated: a simulation study”  In _Statistics in Medicine_ 33.20, 2014, pp. 3488–3508
  * [51] W.. Loh and S. Vansteelandt  “Confounder selection strategies targeting stable treatment effect estimators”  In _Statistics in Medicine_ 40.3, 2021, pp. 607–630
  * [52] J.. Lunceford and M. Davidian  “Stratification and weighting via the propensity score in estimation of causal treatment effects: a comparative study”  In _Statistics in Medicine_ 23.19, 2004, pp. 2937–2960
  * [53] U. Luxburg  “A tutorial on spectral clustering”  In _Statistics and Computing_ 17.4, 2007, pp. 395–416
  * [54] K. Ming and P.. Rosenbaum  “A note on optimal matching with variable controls using the assignment algorithm”  In _Journal of Computational and Graphical Statistics_ 10.3, 2001, pp. 455–463
  * [55] K. Ming and P.. Rosenbaum  “Substantial gains in bias reduction from matching with a variable number of controls”  In _Biometrics_ 56.1, 2000, pp. 118–124
  * [56] J. Neyman  “On the application of probability theory to agricultural experiments”  In _The Annals of Agricultural Sciences_ 10, 1923, pp. 1–51
  * [57] O… Padilla, P. Ding, Y. Chen and G. Ruiz  “A causal fused lasso for interpretable heterogeneous treatment effects estimation”, 2021  arXiv:2110.00901
  * [58] O… Padilla, J. Sharpnack, Y. Chen and D.. Witten  “Adaptive nonparametric regression with the K-nearest neighbour fused lasso”  In _Biometrika_ 107.2, 2020, pp. 293–310
  * [59] J. Pearl  “Causal diagrams for empirical research”  In _Biometrika_ 82.4, 1995, pp. 669–710
  * [60] J. Pearl  “Causality: models, reasoning, and inference”  Cambridge University Press, 2009
  * [61] M. Qian and S.. Murphy  “Performance guarantees for individualized treatment rules”  In _The Annals of Statistics_ 39.2, 2011, pp. 1180–1210
  * [62] A. Rekkas et al.  “Predictive approaches to heterogeneous treatment effects: a scoping review”  In _BMC Medical Research Methodology_ 20.1, 2020, pp. 264
  * [63] P.. Rosenbaum  “A characterization of optimal designs for observational studies”  In _Journal of the Royal Statistical Society, Series B (Statistical Methodology)_ 53.3, 1991, pp. 597–610
  * [64] P.. Rosenbaum  “Optimal matching for observational studies”  In _Journal of the American Statistical Association_ 84.408, 1989, pp. 1024–1032
  * [65] P.. Rosenbaum and D.. Rubin  “Constructing a control group using multivariate matched sampling methods that incorporate the propensity score”  In _The American Statistician_ 39.1, 1985, pp. 33–38
  * [66] P.. Rosenbaum and D.. Rubin  “The central role of the propensity score in observational studies for causal effects”  In _Biometrika_ 70.1, 1983, pp. 41–55
  * [67] D.. Rubin  “Estimating causal effects of treatments in randomized and nonrandomized studies”  In _Journal of Educational Psychology_ 66.5, 1974, pp. 688–701
  * [68] D.. Rubin  “Using propensity scores to help design observational studies: application to the tobacco litigation”  In _Health Services and Outcomes Research Methodology_ 2, 2001, pp. 169–188
  * [69] D.. Rubin and N. Thomas  “Combining propensity score matching with additional adjustments for prognostic covariates”  In _Journal of the American Statistical Association_ 95.450, 2000, pp. 573–585
  * [70] D.. Rubin and N. Thomas  “Matching using estimated propensity scores: relating theory to practice”  In _Biometrics_ 52.1, 1996, pp. 249–264
  * [71] I.. Schou and I.. Marschner  “Methods for exploring treatment effect heterogeneity in subgroup analysis: an application to global clinical trials”  In _Pharmaceutical Statistics_ 14.1, 2015, pp. 44–55
  * [72] H. Smith  “Matching with multiple controls to estimate treatment effects in observational studies”  In _Sociological Methodology_ 27.1, 1997, pp. 325–353
  * [73] M.. Smith et al.  “Introduction to computational causal inference using reproducible Stata, R, and Python code: A tutorial”  In _Statistics in Medicine_ , 2021
  * [74] E.. Stuart  “Matching methods for causal inference: a review and a look forward”  In _Statistical science: a review journal of the Institute of Mathematical Statistics_ 25.1, 2010, pp. 1–21
  * [75] E.. Stuart and K.. Green  “Using full matching to estimate causal effects in non-experimental studies: examining the relationship between adolescent marijuana use and adult outcomes”  In _Developmental Psychology_ 44.2, 2008, pp. 395–406
  * [76] X. Su et al.  “Subgroup analysis via recursive partitioning”  In _Journal of Machine Learning Research_ 10, 2009, pp. 141–158
  * [77] V. Syrgkanis et al.  “Machine learning estimation of heterogeneous treatment effects with instruments”  In _Advances in Neural Information Processing Systems_ , 2019, pp. 15167–15176
  * [78] M. Taddy, M. Gardner, L. Chen and D. Draper  “A nonparametric Bayesian analysis of heterogenous treatment effects in digital experimentation”  In _Journal of Business $\&$ Economic Statistics_ 34.4, 2016, pp. 661–672
  * [79] J. Tanniou, I. Tweel, S. Teerenstra and K… Roes  “Estimates of subgroup treatment effects in overall nonsignificant trials: to what extent should we believe in them?”  In _Pharmaceutical Statistics_ 16.4, 2017, pp. 280–295
  * [80] R. Tibshirani  “Regression shrinkage and selection via the lasso”  In _Journal of the Royal Statistical Society, Series B (Statistical Methodology)_ 58.1, 1996, pp. 267–288
  * [81] S. Wager and S. Athey  “Estimation and inference of heterogeneous treatment effects using random forests”  In _Journal of the American Statistical Association_ 113.523, 2018, pp. 1228–1242
  * [82] G. Wahba  “Soft and hard classification by reproducing kernel Hilbert space methods”  In _Proceedings of the National Academy of Sciences_ 99.26, 2002, pp. 16524–16530
  * [83] S. Yang et al.  “Propensity score matching and subclassification in observational studies with multi-level treatments”  In _Biometrics_ 72.4, 2016, pp. 1055–1065
  * [84] W. Zhang et al.  “Mining heterogeneous causal effects for personalized cancer treatment”  In _Bioinformatics_ 33.15, 2017, pp. 2372–2378
  * [85] D.. Zhao and K. Imai  “Propensity score-based methods for causal inference in observational studies with non-binary treatments”  In _Statistical Methods in Medical Research_ 29.3, 2020, pp. 709–727
  * [86] J.. Zubizarreta  “Using mixed integer programming for matching in an observational study of kidney failure after surgery”  In _Journal of the American Statistical Association_ 107.500, 2012, pp. 1360–1371
  * [87] J.. Zubizarreta and L. Keele  “Optimal multilevel matching in clustered observational studies: a case study of the effectiveness of private schools under a large-scale voucher system”  In _Journal of the American Statistical Association_ 112.518, 2017, pp. 547–560