Learning Objectives for Treatment Effect Estimation

arXiv:1712.04912v1 [stat.ML] 13 Dec 2017

Xinkun Nie
xinkun@stanford.edu

Stefan Wager
swager@stanford.edu

Draft version December 2017
Abstract
We develop a general class of two-step algorithms for heterogeneous treatment effect
estimation in observational studies. We first estimate marginal effects and treatment
propensities to form an objective function that isolates the heterogeneous treatment
effects, and then optimize the learned objective. This approach has several advantages over existing methods. From a practical perspective, our method is very flexible
and easy to use: In both steps, we can use any method of our choice, e.g., penalized
regression, a deep net, or boosting; moreover, these methods can be fine-tuned by
cross-validating on the learned objective. Meanwhile, in the case of penalized kernel
regression, we show that our method has a quasi-oracle property, whereby even if our
pilot estimates for marginal effects and treatment propensities are not particularly accurate, we achieve the same regret bounds as an oracle who has a-priori knowledge of
these nuisance components. We implement variants of our method based on both penalized regression and convolutional neural networks, and find promising performance
relative to existing baselines.

1

Introduction

The problem of heterogeneous treatment effect estimation in observational studies arises in
a wide variety application areas [Athey, 2017], ranging from personalized medicine [Obermeyer and Emanuel, 2016] to offline evaluation of bandits [Dudı́k et al., 2011], and is also
a key component of several proposals for learning decision rules [Athey and Wager, 2017,
Hirano and Porter, 2009]. There has been considerable interest in developing flexible and
performant methods for heterogeneous treatment effect estimation. Some notable recent
advances include proposals based on the lasso [Imai and Ratkovic, 2013], recursive partitioning [Athey and Imbens, 2016, Su et al., 2009], BART [Hill, 2011], random forests [Wager
and Athey, 2017], boosting [Powers et al., 2017], deep neural networks [Shalit et al., 2017],
etc.; see Dorie et al. [2017] for a recent survey and comparisons.
However, although this line of work has led to many promising methods, the literature
does not yet provide a comprehensive answer as to how machine learning methods should
be adapted to a treatment effect estimation context. First of all, the process of developing
“causal” variants of machine learning methods is still a fairly labor intensive process—
effectively requiring the involvement of specialized researchers—and there is no concrete
guidance on how to turn a good generic predictor into a good treatment effect estimator.1
1 As an exception to this claim, Künzel et al. [2017] recently proposed a different meta-learner approach
with an aim similar to ours. We discuss their method at length when presenting details of our proposal.

1

Second, with some exceptions, the above methods are mostly justified via numerical experiments, and come with no formal convergence guarantees or regret bounds proving that the
methods in fact succeed in isolating causal effects.
In this paper, we discuss a new approach to estimating heterogeneous treatment effects
that addresses both of these concerns. Our framework allows for fully automatic specification of heterogeneous treatment effect estimators in terms of arbitrary loss minimization
procedures. Moreover, we show how the resulting methods can achieve comparable regret
bounds to oracle methods that know everything about the data-generating distribution except the treatment effects.
We formalize our problem in terms of the potential outcomes framework [Neyman, 1923,
Rubin, 1974]. The analyst has access to n independent and identically distributed examples
(Xi , Yi , Wi ), i = 1, ..., n, where Xi ∈ X denotes per-person features, Yi ∈ R is the observed
outcome, and Wi ∈ {0, 1} is the treatment assignment. We posit the existence of potential outcomes {Yi (0), Yi (1)} corresponding to the outcome we would have observed given
treatment assignment Wi = 0 or 1 respectively, such that Yi = Yi (W

 i ), and seek to estimate the conditional average treatment effect function τ ∗ (x) = E Y (1) − Y (0) X = x .
In order to identify τ ∗ (x), we assume unconfoundedness as below, i.e., that the treatment assignment is as good as random once we control
 for the features Xi [Rosenbaum
∗
X
=
x
for the treatment propensity, and
and Rubin, 1983].
We
write
e
(x)
=
P
W
=
1


µ∗(w) (x) = E Y (w) X = x for the conditional response surfaces. Throughout this paper,
we use ∗-superscripts to denote unknown population quantities.
Assumption 1. The treatment assignment Wi is unconfounded, {Yi (0), Yi (1)} ⊥
⊥ Wi Xi .
Given this setup, we can use a few lines of algebra to check that the conditional
average

treatment effect can be expressed as follows in terms of m∗ (x) = E Y X = x :


Yi − m∗ (Xi ) = (Wi − e∗ (Xi )) τ ∗ (Xi ) + εi , εi = Yi − E Yi Xi , Wi .
(1)
This decomposition was originally used by Robinson [1988] to estimate parametric components in partially linear models, and has regularly been used in both statistics and econometrics ever since; see, e.g., Chernozhukov et al. [2017] and Zhao et al. [2017] for recent
examples. With the representation (1), an oracle who knew both the functions m∗ (x) and
e∗ (x) a priori could estimate the heterogeneous treatment effect function τ ∗ (·) by simple
loss minimization:


!2
n

1 X
(Yi − m∗ (Xi )) − (Wi − e∗ (Xi )) τ (Xi ) + Λn (τ (·)) ,
(2)
τ̃ (·) = argminτ

n
i=1

where the term Λn (τ (·)) is interpreted as a regularizer on the complexity of the τ (·) function.
In practice, this regularization could be explicit as in penalized regression, or implicit, e.g.,
as provided by a carefully designed deep neural network.
The difficulty, however, is that we never know the weighted main effect function m∗ (x)
and usually don’t know the treatment propensities e∗ (x) either, and so the estimator (2) is
not feasible. In this paper, we study the following class of two-step estimators motivated by
the oracle procedure (2):
1. Fit m̂(x) and ê(x) via appropriate methods tuned for optimal predictive accuracy,
then
2

2. Estimate treatment effects via a plug-in version of (2), where ê(−i) (Xi ), etc., denote
held-out predictions, i.e., predictions made without using the i-th training example,
( n

1 X
Yi − m̂(−i) (Xi )
τ̂ (·) = argminτ
n i=1
(3)
)


2
(−i)
− Wi − ê
(Xi ) τ (Xi ) + Λn (τ (·)) .
In other words, the first step learns an approximation for the oracle objective, and the
second step optimizes it. We refer to this approach as R-learning in recognition of the work
of Robinson [1988], and also to emphasize the role of residualization.
This approach has several advantages over existing, more ad hoc proposals. The steps
1 and 2 capture the conceptual aspects of heterogeneous treatment effect estimation in a
general way, and so implementing new variants of these methods reduces to routine work.
Moreover, the key distinctive step using (3) is simply a loss minimization problem, and so
can be efficiently solved via off-the-shelf software such as TensorFlow [Abadi et al., 2016] or
PyTorch [Paszke et al., 2017]. Finally, it is straightforward to tune the regularizer Λn (τ ) by
simply cross validating the squared-error loss in (3), which avoids the use of more sophisticated model-assisted cross-validation procedures as developed in Athey and Imbens [2016]
or Powers et al. [2017].
This paper makes the following contributions. First, we implement variants of our
method based on penalized regression and convolutional neural networks, and in each case,
our method improves upon existing state-of-the-art. Second, in the case of penalized kernel
regression, we prove regret bounds for the feasible estimator for τ̂ (·) that, asymptotically,
match the best available bounds for the oracle method τ̃ (·). Crucially, this quasi-oracle behavior holds even if the first-step predictors ê(x) and m̂(x) are up to an order of magnitude
less accurate than τ̂ (·) itself. Both results highlight the promise of R-learning as a general
purpose framework for estimating heterogeneous treatment effects in observational studies.

2

Related Work

Most existing methods for estimating heterogeneous treatment effects [e.g., Athey and Imbens, 2016, Imai and Ratkovic, 2013, Powers et al., 2017, Shalit et al., 2017, Su et al., 2009,
Wager and Athey, 2017] are based on structural changes to machine learning methods that
effectively regularize the resulting treatment effect estimates. This principle is simplest to
∗
discuss in the context of the high-dimensional
linear model, Yi (w) = Xi> β(w)
+ εi (w) with

∗
d
Xi , β(w) ∈ R , and E εi (w) Xi = 0, where treatment effects are to be fit via the lasso [Tibshirani, 1996]. A naive approach would fit two separate lassos to the treated and control
samples,


2
 X 

β̂(w) = argminβ(w)
Yi − Xi> β(w) + λ(w) β(w) 1 ,
(4)


{i:Wi =w}

and then use it to deduce a treatment effect function, τ̂ (x) = x(β̂(1) − β̂(0) ). Here, however,
the fact that both β̂(0) and β̂(1) are regularized towards 0 separately may inadvertently regularize the treatment effect estimate β̂(1) − β̂(0) away from 0, even when τ (x) = 0 everywhere.
This problem is especially acute when the treated and control samples are of different sizes.
3

In order to avoid such regularization bias, Imai and Ratkovic [2013] advocate instead using
τ̂ (x) = x> δ̂ with
( n 
)
2


X
>
>
b̂, δ̂ = argminb, δ
Yi − Xi b + (Wi − 0.5)Xi δ + λb kbk1 + λδ kδk1 ,
(5)
i=1

and note that this approach always correctly regularizes towards a sparse δ-vector for treatment heterogeneity. The other approaches cited above present variants and improvements
of similar ideas in the context of more sophisticated machine learning methods; see, for example, Figure 1 of Shalit et al. [2017] for a neural network architecture designed to highlight
treatment effect heterogeneity without being affected by confounders.
In contrast, we do not rely on opening up machine learning methods to stabilize the way
they model treatment effect components; rather our approach based on learning objectives
reduces the problem of heterogeneous treatment effect estimation to several rounds of loss
minimization on standard architectures. Recently, Künzel, Sekhon, Bickel, and Yu [2017]
proposed another approach to heterogeneous treatment effect estimation via generic machine
learning methods, called the X-learner. They proceed as follows (see Künzel et al. [2017]
for details):
1. Estimate µ̂(w) (x) via appropriate non-parametric regression methods, then
(−i)

2. On the treated observations, define pseudo-effects Di = Yi − µ̂(0) (Xi ), and use them
to fit τ̂(1) (Xi ) via a non-parametric regression, and on the control observations, define
(−i)
Di = µ̂(1) (Xi ) − Yi , and use them to fit τ̂(0) (Xi ),
3. Report τ̂ (x) = (1 − ê(x))τ̂(1) (x) + ê(x)τ̂ (0).
Just like with our approach, variants of this method can also be implemented via off-theshelf software. We compare our approach to that of Künzel et al. [2017] in our experiments.
Künzel et al. [2017] also provide some theoretical guarantees for X-learner; however, these
guarantees are weaker than the ones we derive for our method. In particular, the rates they
obtain for the X-learner are limited by the better of the rates achieved for µ̂(0) (x) or µ̂(1) (x),
whereas we prove results (3) that allow an order-of-magnitude improvement over the rate
at which we can learn the nuisance components.

3

Experiments

Our approach to heterogeneous treatment effect estimation via learning objectives can be
implemented using any method that framed as a loss minimization problem, such as boosting, decision trees, etc. In our experiments, we focus on variants of (3) implemented via
penalized regression and convolutional neural networks.
In each case, we compare our R-learner to the following 3 baselines (we follow the nomenclature
of Künzel

 et al. [2017]): the S-learner fits a single model for
f (x, w) = E Y X = x, W = w , and then estimate τ̂ (x) = fˆ(x, 1) − fˆ(x, 0); the T -learner
fits the functions µ∗(w) (x) = E Y X = x, W = w separately for w ∈ {0, 1}, and then
estimate τ̂ (x) = µ̂(1) (x) − µ̂(0) (x); and the X-learner is as described in Section 2.

4

3.1

Lasso Example

We start by comparing all four methods implemented via the lasso on simulated designs.
For S-learning, we follow Imai and Ratkovic [2013] using (5), while for T -learning, we use
(4). For X-learning and R-learning, we use L1 -penalized logistic regression to estimate
propensity ê, and the lasso for all other regression estimates. In all simulations, we generate
data as follows:
Xi ∼ Pd , Wi | Xi ∼ Bernoulli(e∗ (Xi )), εi | Xi ∼ N (0, 1),
Yi = b∗ (Xi ) + (Wi − 0.5)τ ∗ (Xi ) + σεi ,

(6)

for different choices X-distribution Pd indexed by dimension d, noise level σ, propensity
function e∗ (·), baseline main effect b∗ (·), and treatment effect function τ ∗ (·). We consider
the following two specific setups, each of which depends on an overlap parameter η, with
the second setup depend on a sparsity level k.
1. We use the Friedman function [Friedman, 1991] for the baseline main effect b∗ (Xi ) =
10 sin(πXi1 Xi2 ) + 20(Xi3 − 0.5)2 + 10Xi4 + 5Xi5 , with Xi ∼ Unif(0, 1)d , e∗ (Xi ) =
sin(πXi1 Xi2 ), τ ∗ (Xi ) = (Xi1 + Xi2 )/2.
√
Pk
∗
∗
2. We use Xi ∼ N (0, Id×d ), θi =
j=1 Xij / k, b (Xi ) = max{0, θi }, e (Xi ) =
2
∗
trimη (0.5 · (1 + sign(θi )θi )), τ (Xi ) = sin(2πXi1 ).
We run the lasso on a natural spline basis expansion with 7 degrees of freedom on Xi for
all estimators. We split the n data points into a training set and a test set with ratio 4:1,
and the reported the mean squared error (MSE) is on the test set. All lasso regularization
penalties are chosen from 10-fold cross validation. As discussed earlier, for R-learning, we
use cross-fitting on ê and m̂ in (3). Simulation results are averaged across 200 runs. All
methods are implemented via glmnet [Friedman et al., 2010].
In Figure 1, we compare the performance of our 4 considered methods to an oracle that
runs the lasso on (2). As is clear from this illustrations, the considered simulation settings
differ vastly in difficulty, both in terms of the accuracy of the oracle, and in terms of the
ability of feasible methods to approach the oracle (note that all MSE numbers are on a
log-scale).
Encouragingly, we see that R-learning outperforms its baselines across a wide variety of
settings. Furthermore, in the easier Setup 2, the R-learner essentially matches the performance of the oracle in terms of the MSE (whereas other methods do not). Meanwhile, in the
more difficult Setup 1, the R-learner eventually falls short of the oracle, but still performs
respectably. In terms of the other methods, S gets second place, and X improves over T .
We provide the numbers behind Figure 1 in Table 2, Table 3, and Table 4 in Section 6 at
the end of the paper.

3.2

MNIST Rotation Task via Convolutional Neural Networks

We next compare all four methods on a task built on top of the classic MNIST dataset of
LeCun [1998]. In this task, each image is rotated by a random amount Yi that depends on a
treatment Wi . Our goal is to estimate the effect of the treatment Wi on the rotation Yi . For
motivation, one could imagine that we are comparing two different processes for generating
the images, and want to understand whether one process rotates the images differently from
the other.
5

S
T
X
R

●

●
●●

−5

−4

−3

−2

−1

0

1

●
●●
●●
●
●
●
●
●●
●
●●

●

−1.0

1
●
●

●
● ●

●
●
●
●
●● ●●
●
● ●● ●
●
●● ● ●
●
●
●
●

●
●

●
●

●

●
●

●

−0.5

●

0.0

log−MSE

2

log−MSE

3

●

0.5

4

S
T
X
R

●●
●
●●●●●
●

−1.2

●
●●● ●
●

−1.1

●
●
●●●● ●●
● ●
●● ●

−1.0

●●
●
● ●
●●
● ● ●●●
●● ●●●

−0.9

−0.8

oracle log−MSE

oracle log−MSE

Setup 1

Setup 2

−0.7

−0.6

Figure 1: Performance of lasso-based S-, T -, X-, and R-learners, relative to a lasso-based
oracle learner (2), across simulation setups described in Section 3.1. All mean squared error
(MSE) numbers are aggregated over 200 runs.
Specifically, we generate data as follows. Let Xi denote the raw image, and let C(Xi ) ∈
{0, 1, ..., 9} denote the digit corresponding to the image. Then, we obtain Wi and Yi via:
Wi ∼ Bernoulli(e∗ (Xi )),

Yi = b∗ (Xi ) + Wi τ ∗ (Xi ) + εi ,

iid

εi ∼ N (0, 4),

e∗ (Xi ) = 0.2,

b∗ (Xi ) = 40 + 2Ci

if Ci ∈ [0, 4],

∗

∗

if Ci ∈ [5, 9],

e (Xi ) = 0.8,

b (Xi ) = −40 + Ci
∗

τ (Xi ) = 10 + Ci
∗

τ (Xi ) = −10 − Ci

(7)

if Ci is even,
if Ci is odd.

The analyst then has access to the value of the angle Yi , the treatment assignment Wi , and
a version of the image Xi rotated by an angle Yi in degrees (as illustrated in Figure 2). The
task is to estimate the heterogeneous treatment effect τ ∗ (Xi ).
We consider variants of S-, T -, X-, and R-learning, all implemented using convolutional
neural networks. In S-learning, we have one network that predicts the outcome Y while
jointly learning a 5-dimensional embedding for the treatment W . In T -learning, we learn two
separate networks for the treated and control samples respectively. In X-learning, we learn
two separate networks to estimate µ∗(1) (·) and µ∗(0) (·), two separate networks to estimate
∗
∗
τ(1)
(·) and τ(0)
(·), and another separate network to estimate e∗ (·). In R-learning, note
that because we use held-out predictions of ê(·) and m̂(·), we learn a separate network for
estimating e∗ (·) and m∗ (·) for each of the held-out set. We also have a network to estimate
τ ∗ (·), with loss using the predictions from the ê and m̂ networks. All networks have two
convolutional layers with kernel size 5, and a dropout layer. All ê networks are trained with
the cross entropy loss, and all other networks with the mean squared error loss. We use the
standard MNIST training and test splits, and evaluate all methods on the test set.
We implemented all methods in PyTorch [Paszke et al., 2017] training with the Adam update rule [Kingma and Ba, 2014] over 20 epochs. We tuned the same set of hyper-parameters
6

Figure 2: Examples of the input dataset, as generated by (7). Each image Xi is rotated by
the number of degrees specified by Yi . The task is to estimate the heterogeneous treatment
effect τ ∗ (Xi ) for each image.

Mean
Std. Dev.

S
781.7
48.9

T
100.4
13.3

X
36.7
3.3

R
19.6
2.1

Table 1: The mean and standard deviation of the mean squared error (MSE) in estimating
τ ∗ (·) on the MNIST rotation task. Results reported over 10 runs.
for all methods using a simple grid search of learning rate from [1e−4, 3e−4, 1e−3], and batch
size from [32, 64, 128], and fixed all other hyper-parameters by their default values. In Table
1, we reported the mean and standard deviation of the mean squared error (MSE) over 10
runs, with all methods using the same set of 10 random seeds.
Encouragingly, we again find R-learning to perform well, with a typical error of ∼ 4.5
degrees. Here, X-learning also performs respectably, with a typical error of ∼ 6 degrees,
whereas S- and T -learning provide essentially useless estimates of τ ∗ (·). To understand the
difference in performance of the S-learners here and with the lasso, recall that with the lasso
we used a carefully designed S-learner following Imai and Ratkovic [2013], whereas here we
took more of an out-of-the-box approach. S-learning may thus be particularly sensitive to
architecture design.

7

4

Asymptotics of Heterogeneous Treatment Effect Estimation via Penalized Kernel Regression

In order to illustrate formal aspects of R-learning, we study a variant of (3) based on
penalized kernel regression. The problem of regularized kernel learning covers a broad class
of methods that have been thoroughly studied in the statistical learning literature [see, e.g.,
Bartlett and Mendelson, 2006, Caponnetto and De Vito, 2007, Cucker and Smale, 2002,
Steinwart and Christmann, 2008, Mendelson and Neeman, 2010], and thus provides an ideal
case study for highlighting the promise of R-learning. Throughout this section, we will seek
to derive performance guarantees for R-learning that match the best available guarantees
for the infeasible oracle learner (2). More formally, we conduct our analysis in the following
setting.
Assumption 2. We study k·kH -penalized kernel regression, where H is a reproducing kernel
Hilbert space (RKHS) with a continuous, positive semi-definite kernel function K. Let P
be a non-negative measure over the compact metric space X ⊂ Rd , and that K is a kernel
with respect to P such that K(x,Px) ≤ 1 for all x ∈ X , that the eigenvalues σj of K belong
∞
to lp for some 0 < p < 1, i.e., j=1 σjp < ∞, and that the eigenfunctions ψj (·) of K are
uniformly bounded, i.e., supj kψj k∞ ≤ A < ∞, with kψj kH = 1. Finally, we assume that
the outcomes Yi are almost surely bounded, |Yi | ≤ M .
We do not assume that τ ∗ (·) has a finite H-norm; rather, we only assume that we can
make it have a finite H-norm after a sufficient amount of smoothing. Specifically, writing
α
TK for the linear mapping TK (f )(·) = E [K(·, X)f (X)], we assume that kTK
(τ ∗ (·))kH < ∞
α
for some 0 < α < 1/2. Here, with α = 0, TK would be the identity operator, and so this
assumption would be equivalent to the strongest possible assumption that kτ ∗ (·)kH < ∞
itself. Then, as α grows, this assumption gets progressively weaker, and at α = 1/2 it would
devolve to simply asking that τ ∗ (·) belong to the space L2 (P) of square-integrable functions,
1/2
because kTK (f )kH = kf kL2 (P) for all f ∈ L2 (P).


Assumption 3. The oracle treatment effect function τ ∗ (x) = E Yi (1) − Yi (0) Xi = x
α
(τ ∗ (·))kH < ∞ for some 0 < α < 1/2.
satisfies kTK
Given this setup, we study oracle penalized regression rules of the following form,
( n 
)
2
1X
∗
∗
(Yi − m (Xi )) − (Wi − e (Xi )) τ (Xi ) + Λn (kτ kH ) , (8)
τ̃ (·) = argminτ ∈H
n i=1
as well as feasible analogues obtained by cross-fitting [Chernozhukov et al., 2017, Schick,
1986]:
( n

1 X
Yi − m̂(−q(i)) (Xi )
τ̂ (·) = argminτ ∈H
n i=1
(9)
)


2
(−q(i))
− Wi − ê
(Xi ) τ (Xi ) + Λn (kτ kH ) ,
where q is a mapping from the i = 1, ..., n sample indices to Q evenly sized data folds,
such that ê(−q(i)) (x) and m̂(−q(i)) (x) are each trained without considering observations in
8

the q(i)-th data fold. (Typically we set Q to 5 or 10.) We seek to characterize the accuracy
of our estimator τ̂ (·) by bounding its regret R(τ̂ ),
"
2 #
R(τ ) = L(τ ) − L(τ ∗ ), L(τ ) = E
(Yi − m∗ (Xi )) − τ (Xi ) (Wi − e∗ (Xi ))
.
(10)
Note that if we have overlap, i.e., there is an η > 0 such that η < e∗ (x) < 1 − η
for all x ∈ X , then following from R(τ ) = E[(Wi − e∗ (Xi ))2 (τ (Xi ) − τ ∗ (Xi ))2 ], we have
R(τ ) ≤ E[(τ (Xi ) − τ ∗ (Xi ))2 ] ≤ η −2 R(τ ), meaning that regret bounds directly translate into
squared-error loss bounds for τ (·), and vice-versa.
The sharpest available regret bounds for (8) given Assumptions 2 and 3 are due to
Mendelson and Neeman [2010] (see also Steinwart et al. [2009]), and scale as


1−2α
eP n− p+(1−2α) ,
(11)
R (τ̃ ) = O
e
where the O-notation
hides log factors. In the case α = 0 where τ ∗ is within the RKHS used
for penalization, we recover the more familiar n−1/(1+p) rate established by Caponnetto and
De Vito [2007]. Again, our goal is to establish regret bounds for our feasible estimator τ̂
that match the best available regret bounds for τ̃ .

4.1

Fast Rates and Isomorphic Coordinate Projections

In order to establish regret bounds for τ̂ , we first need to briefly review the proof techniques
underlying (11). The argument of Mendelson and Neeman [2010] relies on the following
quasi-isomorphic coordinate projection lemma of Bartlett [2008]. To state this result, write
Hc = {τ : kτ kH ≤ c} for the radius-c ball of H, let τc∗ = argmin {L(τ ) : τ ∈ Hc } denote the
best approximation to τ ∗ within Hc , and define c-regret R(τ ; c) = L(τ ) − L(τc∗ ) over τ ∈ Hc .
We also define estimated and oracle regret functions
bn (τ ; c) = L
b n (τ ) − L
b n (τ ∗ ),
R
c
where

en (τ ; c) = L
e n (τ ) − L
e n (τ ∗ ),
R
c

2
n 
X
e n (τ ) = 1
L
(Yi − m∗ (Xi ) − τ (Xi )(Wi − e∗ (Xi )
n i=1

(12)

(13)

is the oracle loss function on the n samples used for empirical minimization, and
n 
2
X
b n (τ ) = 1
(Yi − m̂(−q(i)) (Xi ) − τ (Xi )(Wi − ê(−q(i)) (Xi )
L
n i=1

(14)

is the feasible cross-fitted loss function. Note that R̂n (τ ; c) is not actually observable in
practice as it depends on τc∗ ; however, this does not hinder us from establishing highprobability bounds for it.
Lemma 1 (Bartlett [2008]). Let Ľn (τ ) be any loss function, and Řn (τ ; c) = Ľn (τ ) − Ľn (τc∗ )
be the associated regret. Suppose that, for every c ≥ 1, the following inequality holds:
1
Řn (τ ; c) − ρn (c) ≤ R(τ ; c) ≤ 2Řn (τ ; c) + ρn (c) for all τ ∈ Hc .
2

9

(15)

Then, there exist constants κ1 and κ2 such that any solution to the empirical minimization
problem with regularizer Λn (c) ≥ ρn (c),

(16)
τ̌ ∈ argminτ ∈H Ľ(τ ) + κ1 Λn (kτ kH ) ,
also satisfies the following risk bound
L (τ̌ ) ≤ inf {L(τ ) + κ2 Λn (kτ kH )} .

(17)

τ ∈H

In other words, the above lemma reduces the problem of deriving regret bounds to establishing quasi-isomorphisms as in (15) (and any with-high-probability quasi-isomorphism
guarantee yields a with-high-probability regret bound). In order to prove a regret bound
(11) following the approach of Mendelson and Neeman [2010], we first establish a quasiisomorphism in terms of the oracle empirical regret,
1e
en (τ ; c) + ρn (c),
Rn (τ ; c) − ρn (c) ≤ R(τ ; c) ≤ 2R
2

(18)

and then get a regret bound in terms of the regularized risk as in (17). Mendelson and
Neeman [2010] then conclude their argument by noting that, for any c > 0,
inf {L(τ ) + κ2 ρn (kτ kH )} ≤ L (τ ∗ ) + (L (τc∗ ) − L (τ ∗ )) + κ2 ρn (c),

τ ∈H

(19)

and that by optimizing this bound (up to polylogarithm factors) and setting
cn = nα/(p+(1−2α)) , we recover the regret bound (11) for the oracle learner. Their argument for choosing the optimal value of cn to use in (19) leverages the approximation error
bound
2α−1
1/α
2
α
kτc∗ − τ ∗ kL2 (P) ≤ c α kTK
(τ ∗ (·))kH ,
(20)
established by Smale and Zhou [2003] given Assumption 3.
For our purposes, the upshot is that if we can match the strength of the quasiisomorphism bounds (18) with our feasible loss function, then we can also match the rate
of any regret bounds proved using Lemma 1. We do so via the following result, the proof
of which relies heavily on the cross-fitting construction used in (9), as advocated by Chernozhukov et al. [2017] and Schick [1986].
Lemma 2. Suppose that both pilot regression adjustment ê(x) and m̂(x) converge at rate
h
i
h
i

2
2
E (m̂(X) − m∗ (X)) , E (ê(X) − e∗ (X)) = O a2n
(21)
for some sequence an → 0. Suppose, moreover, that we have overlap, i.e., η < e∗ (x) < 1 − η
for some η > 0, and that Assumptions 2 and 3 hold. Then, for any ε > 0, there is a constant
U (ε) such that the regret functions induced by the oracle learner (8) and the feasible learner
(9) are coupled as
bn (τ ; c) − R
en (τ ; c)
R
1−p p

≤ U (ε) a2n R(τ ; c)

+

cp+

1−p
2

R(τ ; c) 2
√
cp +
n

(2α−1)(1−p)
2α

√

n

+

c

c

+

2
p(3−p)
1 2α−1 (1−p)
p+1 + 2
α
1+p

n1/(1+p)
10

c

p(3−p)
p+1

R(τ ; c)
n1/(1+p)

!
,

(1−p)2
1+p

(22)

simultaneously for all c ≥ 1 and τ ∈ Hc , with probability at least 1 − ε, for large enough
values of n.
Thus, given any quasi-isomorphism for the oracle learner (18) with error ρn (c), we can
b ). Thus, given any regret bound for
use (22) to derive a quasi-isomoprhism bound for R(τ
the oracle learner built using Lemma 1, we can also get an analogous regret bound for the
feasible learner provided we regularize just a little bit more. The following result makes this
formal.
Theorem 3. Given the conditions of Lemma 1, suppose that we obtain τ̂ (·) via a penalized kernel regression variant of R-learning (9), with a properly chosen penalty of the form
Λn (kτ̂ kH ) specified in the proof. Suppose, moreover, that the nuisance parameter estimates
satisfy (21) with a rate an bounded as
(1 − α)(1 + p) 1 − 2α
.
(23)
4
p + (1 − 2α)
eP (n−(1−2α)/(p+(1−2α)) ).
Then τ̂ (·) satisfies the same regret bound (11) as τ̃ (·), i.e., R (τ̂ ) = O
an = O n−κ

4.2



with κ >

Discussion

The above result shows that, in the context of penalized kernel regression, R-learning can
match the best available performance guarantees available for the oracle learner (8) that
knows everything about the data generating distribution except the true treatment effect
function. In other words, the best available regret bounds for each method scale as
1−2α

eP (rn2 ), with rn = n− 21 p+(1−2α) .
R(τ̂ ), R(τ̃ ) = O

(24)

The key condition in Theorem 3 is the bound (23) on how fast the nuisance components
need to be estimated; which, in terms of the notation from (24), can be written as
(1−α)(1+p)
2

an  rn

.

(25)

We note that (1 − α)(1 + p)/2 < 1 for any 0 < α < 1/2 and 0 < p < 1; for example, when
α = 1/3, and p = 1/2, we have rn = n−1/4 , and αn  n−1/8 . This bound allows an to
decay slower than rn (i.e., we can have rn  an ). This substantiates our earlier claim about
R-learning, whereby we can achieve quasi-oracle behavior even if m̂(·) and ê(·) are an order
of magnitude less accurate than our target error rate for τ̂ (·).2
Künzel et al. [2017] also present a quasi-oracle result for their X-learning procedure,
but their result is considerably less general than ours. They focus on the case where the
number of control units |{Wi = 0}| grows much faster than the number of treated units
|{Wi = 1}|; and, in this case, they show that the X-learner performs as well
 as an oracle

who already knew the mean response function for the controls, µ∗(0) (x) = E Yi (0) Xi = x .
In this context, it is interesting to note that under the asymptotics of Künzel et al. [2017]
with heavy treatment-control imbalance, we have m∗ (x) ≈ µ∗0 (x) and e∗ (x) ≈ 0, and so
R-learning as in (9) is roughly equivalent to the X-learning procedure presented in Section
2. Thus, at a high level, we can interpret the result of Künzel et al. [2017] as a special
case of our result in the case where the number of controls dominates the number of treated
units.
2 As we approach the semiparametric case, i.e., α, p → 0, we recover the well-known result from the
√
semiparametric inference literature that, in order to get 1/ n-consistent inference for a single target parameter, then need 4-th root consistent nuisance parameter estimates; see Robinson [1988], Chernozhukov
et al. [2017], and references therein.

11

5

Proofs

Before beginning our proof, we present an inequality that we will use frequently. Under
Assumption P
2, if η < e∗ (x) < 1 − η for some η > 0, then there is a constant B depending
∞
on A, p and j=1 σjp such that, for all τ with kτ kH ≤ c,
p

1−p

kτ − τc∗ k∞ ≤ B kτ − τc∗ kH kτ − τc∗ kL2 (P) ≤ 2p η −(1−p) B cp R(τ ; c)

1−p
2

.

(26)

The first inequality in (26) follows directly from Lemma 5.1 of Mendelson and Neeman
[2010]. Meanwhile, thanks to overlap, we know that
h
i
2
2
2
R(τ ; c) = EP (Wi − e∗ (x)) (τ (Xi ) − τc∗ (Xi )) ≥ η 2 kτ (Xi ) − τc∗ (Xi )kL2 (P) ,
(27)
and finally kτ − τc∗ kH ≤ 2c by the triangle inequality so (26) in fact holds.

5.1

Proof of Lemma 2

b ) as follows:
We start by decomposing the feasible loss function L(τ
n



2
X 
b )= 1
Yi − m̂(−q(i)) (Xi ) − τ (Xi ) Wi − ê(−q(i)) (Xi )
L(τ
n i=1
=

n


1 X
(Yi − m∗ (Xi )) + m∗ (Xi ) − m̂(−q(i)) (Xi )
n i=1

− τ (Xi ) (Wi − e∗ (Xi )) − τ (Xi )(e∗ (Xi ) − ê(−q(i)) (Xi ))

2

n

=

1X
2
((Yi − m∗ (Xi )) − τ (Xi ) (Wi − e∗ (Xi )))
n i=1
+

n

2
1 X  ∗
m (Xi ) − m̂(−q(i)) (Xi ) − τ (Xi )(e∗ (Xi ) − ê(−q(i)) (Xi ))
n i=1

+

n


1X
2 (Yi − m∗ (Xi )) m∗ (Xi ) − m̂(−q(i)) (Xi )
n i=1

n


1X
2 (Yi − m∗ (Xi )) e∗ (Xi ) − ê(−q(i)) (Xi ) τ (Xi )
−
n i=1
n

−



1X
2 (Wi − e∗ (Xi )) m∗ (Xi ) − m̂(−q(i)) (Xi ) τ (Xi )
n i=1

+



1X
2 (Wi − e∗ (Xi )) e∗ (Xi ) − ê(−q(i)) (Xi ) τ (Xi )2 .
n i=1

n

Furthermore, we can verify that some terms cancel out when we restrict attention to our
b ; c) − R(τ
e ; c) = L(τ
b ) − L(τ
b c∗ ) − L(τ
e ) + L(τ
e c∗ ); in particular, note
main object of interest R(τ

12

e ):
that the first summand above is exactly L(τ
b ; c) − R(τ
e ; c)
R(τ
n

=



−2 X  ∗
m (Xi ) − m̂(−q(i)) (Xi ) e∗ (Xi ) − ê(−q(i)) (Xi ) (τ (Xi ) − τc∗ (Xi ))
n i=1
+

n
2

1 X ∗
e (Xi ) − ê(−q(i)) (Xi )
τ (Xi )2 − τc∗ (Xi )2
n i=1

−

n


1X
2 (Yi − m∗ (Xi )) e∗ (Xi ) − ê(−q(i)) (Xi ) (τ (Xi ) − τc∗ (Xi ))
n i=1

−



1X
2 (Wi − e∗ (Xi )) m∗ (Xi ) − m̂(−q(i)) (Xi ) (τ (Xi ) − τc∗ (Xi ))
n i=1

n

n




1X
+
2 (Wi − e∗ (Xi )) e∗ (Xi ) − ê(−q(i)) (Xi ) τ (Xi )2 − τc∗ (Xi )2 .
n i=1
Call these 5 summands Ac1 (τ ), Ac2 (τ ), B1c (τ ), B2c (τ ) and B3c (τ ) respectively. We now proceed
to bound them, each on their own.
Starting with Ac1 (τ ), a simple application of the Cauchy-Schwarz inequality implies that
v
v
u n
u n
X
u1 X
u

2
1
2
c
∗
(−q(i))
t
|A1 (τ )| ≤ 2
m (Xi ) − m̂
(Xi ) t
e∗ (Xi ) − ê(−q(i)) (Xi ) kτ − τc∗ k∞ .
n i=1
n i=1
This inequality is deterministic, and so trivially holds simultaneously across all τ ∈ Hc .
Now, the two square-root terms denote the mean squared errors of the m- and e-models
respectively, and decay at rate OP (an ) by Assumption 3 and a direct application of Markov’s
inequality. Thus, applying (26) to bound the infinity-norm discrepancy between τ and τc∗ ,
we find that simultaneously for all c ≥ 1,


1−p
sup {|Ac1 (τ )| : τ ∈ Hc } = OP cp R(τ ; c) 2 a2n .
(28)
We can bound Ac2 (τ ) via a similar argument. By another application of Cauchy-Schwarz,
n
2
1 X ∗
e (Xi ) − ê(−q(i)) (Xi ) τc∗ (Xi )
|Ac2 (τ )| ≤ 2 kτ − τc∗ k∞
n i=1
n
2
2 X ∗
e (Xi ) − ê(−q(i)) (Xi ) τ ∗ (Xi )
n i=1
v
v
u n
u n
X
u
u1 X

1
4
2
∗
∗
(−q(i))
t
+ 2 kτ − τc k∞
e (Xi ) − ê
(Xi ) t
(τ ∗ (Xi ) − τ ∗ (Xi )) .
n i=1
n i=1 c

≤ kτ − τc∗ k∞



Now, recall that τ ∗ (x) = E Yi (1) − Yi (0) Xi = x and that outcomes are bounded by M .
Thus, kτ ∗ k∞ ≤ 2M , and so the first term above can be bounded just like Ac1 (τ ). Then, for
the second term, we note that
v
u n
n
2
u1 X
4

1 X ∗
t
e∗ (Xi ) − ê(−q(i)) (Xi ) ≤
e (Xi ) − ê(−q(i)) (Xi ) = OP a2n ,
n i=1
n i=1

13

while (20) lets us bound the term due to approximation error via Markov’s inequality. The
upshot is that the second term is eventually negligible relative to the first one, and so Ac2 (τ )
can also be uniformly bounded as in (28).
We now move to bounding B1c (τ ). We start by bounding this quantity over sets indexed
by c and δ such that kτ − τc∗ kL2 (P) ≤ δ, i.e., we bound
n
o
B1c, δ = sup |B1 (τ )| : τ ∈ Hc , kτ − τc∗ kL2 (P) ≤ δ .
To bound this term, it is P
helpful to decompose this error in terms of our cross-fitting strucQ
ture, noting that B1c, δ ≤ q=1 B1c, δ (q) for

( P
∗
∗
(−q(i))
(Xi ) (τ (Xi ) − τc∗ (Xi ))
{i:q(i)=q} 2 (Yi − m (Xi )) e (Xi ) − ê
c, δ
:
B1 (q) = sup
|{i : q(i) = q}|
)
τ ∈ Hc , kτ − τc∗ kL2 (P) ≤ δ .
Note that because ê(−q(i)) (·) is computed without considering any observations in the set
{i : q(i) = q}, we can effectively treat this term as deterministic when bounding B1c, δ (q)
(moreover, this error is bounded by Assumption 3). Furthermore, noting that the L2 (P)entropy number of Hc scales as (c/ε)2p and that kτ − τc∗ k∞ ≤ cp δ 1−p on this set by (26), we
can use a standard arguments following Bartlett and Mendelson [2002] with a Rademacher
complexity bound as in equation (2.4) of Koltchinskii [2006] to check that

s   
p(3−p)
(1−p)2
p (1−p)
p+1 δ 1+p
1
c
δ
c

B1c, δ (q) {Xi , Wi , Yi : q(i) = q} = log
O √
+
(29)
u
n
n1/(1+p)
with probability at least 1 − u, for any u > 0. Because the right-hand side does not depend
on {Xi , Wi , Yi : q(i) = q}, this bound also holds unconditionally, and finally B1c, δ is also
bounded at this rate.
1 2α−1
Next, thanks to (20) there is a constant D such that kτc∗ − τc∗0 kL2 (P) ≤ Dc 2 α . Thus,
0
for any c < c ,
n
o
sup B1c (τ ) : τ ∈ Hc , kτ − τc∗ kL2 (P) ≤ δ
n
o
1 2α−1
≤ sup B1c (τ ) : τ ∈ Hc0 , kτ − τc∗0 kL2 (P) ≤ δ + Dc 2 α
.
Applying (29) at c = 2kc and δ = 2−kδ , for kc , kδ = 0, 1, 2, ..., and setting the probability
threshold u ∝ 2−kc −kδ , we can us a union bound to check that

 
2 
(1−p)

 (1−p)
p(3−p)
1+p
1 2α−1
r   cp δ + c 21 2α−1
α
c p+1 δ + c 2 α


c 

 , (30)
√
B1c (τ ) = OP 
log
+



1/(1+p)
δ
n
n
simultaneously for all τ , c ≥ 1 and δ ≤ 4M 2 ; then, by concavity,


 
2

p(3−p)
(1−p)2
1 2α−1 (1−p)
1 (2α−1)(1−p)
p+1
1+p
2
α
1+p
p
(1−p)
c
δ
+
c
2
α
+c

c δ
eP 
.
√
B1c (τ ) = O
+


1/(1+p)
n
n

14

(31)

Similar arguments apply to B2c (τ ) and B3c (τ ). Finally, recalling that R(τ ; c) is within a
2
constant factor of kτ − τc∗ kL2 (P) given overlap, we obtain our desired result.

5.2

Proof of Theorem 3

As discussed earlier, the arguments of Mendelson and Neeman [2010] can be used to get
regret bounds for the oracle learner. In order to extend their results, we first review their
result analysis briefly. Their results imply the following facts. For any ε > 0, there is a
constant U (ε) such that
ρn (c) = U (ε) 1 + log(n) + log log c + e

1





p

(c + 1) log(n)
√
n

2
 1+p

(32)

satisfies, for large enough n with probability at least 1 − ε, the condition
en (τ ; c) − ρn (c) ≤ R(τ ; c) ≤ 1.99R
en (τ ; c) + ρn (c).
0.51R

(33)

Thus, thanks to Lemma 1 we know that


α
L(τc∗n ) − L(τ ∗ ) + ρn (cn ) with cn = n p+(1−2α) ,

R(τ̃ ) ≤ OP

and then pairing (20) with the form of ρn (c) in (32), we conclude that



1−2α
e n− p+(1−2α) .
R(τ̃ ) .P max L(τc∗n ) − L(τ ∗ ), ρn (cn ) = O

(34)

(35)

Our present goal is to extend this argument to get a bound for R(τ̂ ).
Towards this end, we first note that, by concavity,
1−p
1 − p − 1+p
1 + p 1−p
1 − p − 1+p
γn 2 (R(τ ; c) − γn ) =
γn 2 +
γn 2 R(τ ; c),
≤ γn 2 +
2
2
2

 1 (1−p)2
2
2
1 (1−p)
−1
1 (1 − p)2 12 (1−p)
1 (1 − p)2
2
1+p
R(τ ; c),
R(τ, c)
≤ 1−
γn 1+p
γn2 1+p +
2 1+p
2 1+p

R(τ, c)

1−p
2

for any choice of γn . Applying this result with
4κ

γn = log(n)8/(1+p) n− 1+p c2n ,
we√can re-express the main result from Lemma 2 as follows (recall that given (23), a2n 
1/ n). For any ε > 0, there is a (different) constant U (ε) such that, with probability at
least 1 − ε
bn (τ ; c) − R
en (τ ; c) ≤ 0.01R(τ ; c)
R
1−p
2

+ U (ε) γn

a2n cp

+

cp+

(2α−1)(1−p)
2α

√

n

+

c

2
p(3−p)
1 2α−1 (1−p)
p+1 + 2
α
1+p

n1/(1+p)

+

c

p(3−p)
p+1

2
1 (1−p)
1+p

γn2

n1/(1+p)

simultaneously for all c ≤ log2 (n)cn and τ ∈ Hc , for large enough values of n.

15

! (36)

We now show that the second term on the right-hand side of (36) is smaller than ρn (c).
Thanks to (23), there is a constant ε > 0 such that





 p 
1−p
2(1 − p)
α (1 − p) + αp
c
2 p
2
e
γn an c = O exp log(n) −κ 2 +
+
1+p
p + (1 − 2α)
cpn



 p 
log(n)
4 (1 − α)(1 + p)
c
= O exp
−
+α
p + (1 − 2α)
1+p
4
cpn


ρ(cn ) cp
=O
= o (ρn (cn ))
p
log2 (n) cn
for all c ≤ cn log2 (n). Similarly,
p+

cn

(2α−1)(1−p)
2α

1

n− 2 = n

2αp+(2α−1)(1−p)−p−1+2α
2(p+(1−2α))

−1+2α

= n p+(1−2α) ,

and so, for all c ≤ cn log2 (n),
cp+

(2α−1)(1−p)
2α

1

n− 2 = o (ρn (cn )) .

Finally, for the third term,
p(3−p)

cn p+1

+ 21

2α−1 (1−p)2
α
1+p

1

n− 1+p = n

2αp(3−p)+(2α−1)(1−p)2 −2p−2+4α
2(p+1)(p+(1−2α))

−1+2α

(2α−1)p−(1−p)2

= n p+(1−2α) n 2(p+1)(p+(1−2α)) ,

(2α−1)/α)

), and so the same argument applies
while the last term is smaller because γn = o(cn
to both.
The upshot is that, combining (33) with (36), we find that for large enough n with
probability at least 1 − ε,
bn (τ ; c) − 2ρn (c) ≤ R(τ ; c) ≤ 2R
en (τ ; c) + 2ρn (c),
0.5R

(37)

for all c ≤ cn log2 (n). Then applying the same argument as above, we us Lemma 1 to check
that the estimator defined as
n
o
b ) + 2κ1 ρn (kτ k ) : kτ̂ k ≤ log2 (n)cn
τ̂ ∈ argminτ ∈H L(τ
(38)
H
H
has regret bounded on the order of
R(τ̂ ) .P



L(τc∗n ) − L(τ ∗ ) + ρn (cn ) . ρn (cn ),

(39)

which establishes the desired decay rate on regret. Finally, note that ρn (cn log2 (n))  ρ(cn ),
and so the constraint kτ̂ kH ≤ log2 (n)cn will be inactive with high probability for large n.
Thus, we also obtain the same with-high-probability regret bounds without this constraint.

6

Appendix

For completeness, we include the mean squared error (MSE) numbers behind Figure 1 for
the lasso simulations in Section 3.1.

16

n
500
500
500
500
500
500
500
500
500
500
500
500
500
500
500
500
500
500
500
500
1000
1000
1000
1000
1000
1000
1000
1000
1000
1000
1000
1000
1000
1000
1000
1000
1000
1000
1000
1000

p
6
6
6
6
6
6
6
6
6
6
12
12
12
12
12
12
12
12
12
12
6
6
6
6
6
6
6
6
6
6
12
12
12
12
12
12
12
12
12
12

σ
0.1
0.1
1.0
1.0
2.0
2.0
5.0
5.0
10.0
10.0
0.1
0.1
1.0
1.0
2.0
2.0
5.0
5.0
10.0
10.0
0.1
0.1
1.0
1.0
2.0
2.0
5.0
5.0
10.0
10.0
0.1
0.1
1.0
1.0
2.0
2.0
5.0
5.0
10.0
10.0

η
0.1
0.2
0.1
0.2
0.1
0.2
0.1
0.2
0.1
0.2
0.1
0.2
0.1
0.2
0.1
0.2
0.1
0.2
0.1
0.2
0.1
0.2
0.1
0.2
0.1
0.2
0.1
0.2
0.1
0.2
0.1
0.2
0.1
0.2
0.1
0.2
0.1
0.2
0.1
0.2

S
2.42
2.46
2.25
2.34
1.99
2.24
2.83
2.30
6.36
5.50
2.32
2.20
2.18
2.13
1.90
2.16
3.39
3.39
7.75
8.30
2.53
2.51
2.36
2.39
2.24
2.28
1.90
1.93
3.97
3.28
2.41
2.31
2.33
2.29
2.15
2.17
2.19
2.25
4.54
4.58

T
17.91
18.21
18.89
19.05
19.81
20.37
28.11
28.30
32.69
31.90
19.16
19.08
20.22
20.30
23.26
24.05
27.64
27.56
30.93
31.16
17.48
17.39
17.36
17.48
18.18
18.27
23.53
23.54
30.15
30.69
17.83
17.71
18.10
17.93
19.62
19.43
24.74
25.43
28.74
28.40

X
17.50
17.81
18.22
18.33
18.13
18.60
21.28
20.80
23.50
22.86
17.95
17.91
18.15
18.24
19.38
19.71
20.92
21.08
23.37
23.88
17.25
17.16
17.06
17.15
17.56
17.66
19.29
19.34
21.84
22.58
17.29
17.20
17.26
17.19
17.98
17.67
19.32
19.72
21.59
21.09

R
1.56
1.53
1.53
1.60
1.57
1.71
2.14
2.04
4.01
3.68
1.88
1.78
1.90
1.79
1.70
1.94
3.34
2.87
6.06
5.44
1.41
1.41
1.38
1.37
1.44
1.38
1.49
1.45
2.33
2.41
1.61
1.50
1.58
1.54
1.49
1.54
1.66
1.80
2.65
3.45

oracle
0.01
0.01
0.17
0.18
0.32
0.30
0.87
0.90
3.43
3.43
0.02
0.02
0.19
0.18
0.34
0.36
1.06
1.09
4.06
3.09
0.00
0.00
0.12
0.11
0.24
0.24
0.63
0.45
1.21
1.57
0.01
0.01
0.12
0.11
0.25
0.24
0.66
0.59
1.48
1.97

Table 2: Mean squared error (MSE) from Setup 1. Results are averaged across 200 runs.
Note that in certain settings the oracle appears to have 0 MSE due to rounding.

17

n
500
500
500
500
500
500
500
500
500
500
500
500
500
500
500
500
500
500
500
500
500
500
500
500
500
500
500
500
500
500
500
500

p
6
6
6
6
6
6
6
6
6
6
6
6
6
6
6
6
12
12
12
12
12
12
12
12
12
12
12
12
12
12
12
12

k
2
2
2
2
3
3
3
3
4
4
4
4
5
5
5
5
2
2
2
2
3
3
3
3
4
4
4
4
5
5
5
5

σ
0.1
0.1
1.0
1.0
0.1
0.1
1.0
1.0
0.1
0.1
1.0
1.0
0.1
0.1
1.0
1.0
0.1
0.1
1.0
1.0
0.1
0.1
1.0
1.0
0.1
0.1
1.0
1.0
0.1
0.1
1.0
1.0

η
0.1
0.2
0.1
0.2
0.1
0.2
0.1
0.2
0.1
0.2
0.1
0.2
0.1
0.2
0.1
0.2
0.1
0.2
0.1
0.2
0.1
0.2
0.1
0.2
0.1
0.2
0.1
0.2
0.1
0.2
0.1
0.2

S
0.65
0.49
0.65
0.59
0.62
0.46
0.69
0.60
0.62
0.46
0.68
0.58
0.63
0.48
0.73
0.61
0.59
0.46
0.66
0.57
0.56
0.46
0.67
0.57
0.61
0.48
0.67
0.61
0.59
0.49
0.64
0.57

T
1.39
0.98
1.23
0.95
1.23
0.86
1.20
0.95
1.17
0.81
1.17
0.96
1.15
0.82
1.18
0.95
1.37
0.94
1.15
0.94
1.19
0.83
1.10
0.90
1.12
0.79
1.10
0.94
1.10
0.80
1.06
0.93

X
1.33
0.92
1.03
0.80
1.15
0.80
0.97
0.77
1.07
0.75
0.94
0.79
1.05
0.76
0.96
0.77
1.29
0.86
0.92
0.77
1.08
0.76
0.88
0.74
0.99
0.72
0.89
0.76
0.97
0.72
0.85
0.75

R
0.53
0.41
0.57
0.56
0.44
0.39
0.57
0.55
0.44
0.40
0.57
0.57
0.42
0.41
0.57
0.55
0.46
0.42
0.57
0.55
0.43
0.43
0.58
0.54
0.46
0.42
0.57
0.56
0.45
0.43
0.55
0.55

oracle
0.47
0.39
0.54
0.53
0.43
0.37
0.55
0.55
0.43
0.37
0.55
0.55
0.40
0.38
0.55
0.53
0.43
0.38
0.55
0.54
0.40
0.37
0.55
0.54
0.41
0.38
0.54
0.55
0.38
0.39
0.56
0.53

Table 3: Mean squared error (MSE) from Setup 2, with n = 500. Results are averaged
across 200 runs.

18

n
1000
1000
1000
1000
1000
1000
1000
1000
1000
1000
1000
1000
1000
1000
1000
1000
1000
1000
1000
1000
1000
1000
1000
1000
1000
1000
1000
1000
1000
1000
1000
1000

p
6
6
6
6
6
6
6
6
6
6
6
6
6
6
6
6
12
12
12
12
12
12
12
12
12
12
12
12
12
12
12
12

k
2
2
2
2
3
3
3
3
4
4
4
4
5
5
5
5
2
2
2
2
3
3
3
3
4
4
4
4
5
5
5
5

σ
0.1
0.1
1.0
1.0
0.1
0.1
1.0
1.0
0.1
0.1
1.0
1.0
0.1
0.1
1.0
1.0
0.1
0.1
1.0
1.0
0.1
0.1
1.0
1.0
0.1
0.1
1.0
1.0
0.1
0.1
1.0
1.0

η
0.1
0.2
0.1
0.2
0.1
0.2
0.1
0.2
0.1
0.2
0.1
0.2
0.1
0.2
0.1
0.2
0.1
0.2
0.1
0.2
0.1
0.2
0.1
0.2
0.1
0.2
0.1
0.2
0.1
0.2
0.1
0.2

S
0.50
0.37
0.62
0.49
0.51
0.39
0.60
0.49
0.50
0.39
0.61
0.51
0.50
0.38
0.62
0.52
0.46
0.37
0.59
0.51
0.48
0.38
0.61
0.51
0.50
0.39
0.60
0.52
0.50
0.38
0.60
0.51

T
1.46
0.95
1.32
0.94
1.24
0.81
1.20
0.90
1.13
0.74
1.14
0.87
1.02
0.68
1.17
0.86
1.41
0.96
1.19
0.92
1.21
0.80
1.11
0.87
1.11
0.77
1.09
0.86
1.07
0.71
1.08
0.84

X
1.42
0.92
1.21
0.84
1.19
0.77
1.01
0.77
1.07
0.71
0.96
0.75
0.97
0.66
0.96
0.74
1.36
0.90
1.03
0.77
1.13
0.76
0.90
0.73
1.02
0.72
0.88
0.72
0.98
0.67
0.87
0.71

R
0.39
0.33
0.48
0.46
0.37
0.32
0.46
0.44
0.35
0.32
0.47
0.45
0.33
0.32
0.47
0.46
0.34
0.33
0.50
0.47
0.35
0.33
0.50
0.47
0.36
0.33
0.50
0.48
0.34
0.32
0.50
0.47

oracle
0.38
0.32
0.49
0.44
0.36
0.31
0.46
0.44
0.34
0.32
0.47
0.45
0.34
0.30
0.47
0.45
0.35
0.32
0.49
0.46
0.34
0.31
0.49
0.47
0.34
0.32
0.48
0.47
0.32
0.31
0.49
0.47

Table 4: Mean squared error (MSE) from Setup 2, with n = 1000. Results are averaged
across 200 runs.

19

References
Martı́n Abadi, Paul Barham, Jianmin Chen, Zhifeng Chen, Andy Davis, Jeffrey Dean,
Matthieu Devin, Sanjay Ghemawat, Geoffrey Irving, Michael Isard, Manjunath Kudlur,
Josh Levenberg, Rajat Monga, Sherry Moore, Derek G. Murray, Benoit Steiner, Paul
Tucker, Vijay Vasudevan, Pete Warden, Martin Wicke, Yuan Yu, and Xiaoqiang Zheng.
TensorFlow: A system for large-scale machine learning. In 12th USENIX Symposium on
Operating Systems Design and Implementation (OSDI 16), pages 265–283, 2016.
Susan Athey. Beyond prediction: Using big data for policy problems. Science, 355(6324):
483–485, 2017.
Susan Athey and Guido Imbens. Recursive partitioning for heterogeneous causal effects.
Proceedings of the National Academy of Sciences, 113(27):7353–7360, 2016.
Susan Athey and Stefan Wager. Efficient policy learning. arXiv preprint arXiv:1702.02896,
2017.
Peter L Bartlett. Fast rates for estimation error and oracle inequalities for model selection.
Econometric Theory, 24(2):545–552, 2008.
Peter L Bartlett and Shahar Mendelson. Rademacher and Gaussian complexities: Risk
bounds and structural results. Journal of Machine Learning Research, 3:463–482, 2002.
Peter L Bartlett and Shahar Mendelson. Empirical minimization. Probability Theory and
Related Fields, 135(3):311–334, 2006.
Andrea Caponnetto and Ernesto De Vito. Optimal rates for the regularized least-squares
algorithm. Foundations of Computational Mathematics, 7(3):331–368, 2007.
Victor Chernozhukov, Denis Chetverikov, Mert Demirer, Esther Duflo, Christian Hansen,
Whitney Newey, and James Robins. Double/debiased machine learning for treatment and
structural parameters. The Econometrics Journal, 2017.
Felipe Cucker and Steve Smale. On the mathematical foundations of learning. Bulletin of
the American Mathematical Society, 39(1):1–49, 2002.
Vincent Dorie, Jennifer Hill, Uri Shalit, Marc Scott, and Dan Cervone. Automated versus doit-yourself methods for causal inference: Lessons learned from a data analysis competition.
arXiv preprint arXiv:1707.02641, 2017.
Miroslav Dudı́k, John Langford, and Lihong Li. Doubly robust policy evaluation and learning. In Proceedings of the 28th International Conference on Machine Learning, pages
1097–1104, 2011.
Jerome Friedman, Trevor Hastie, and Rob Tibshirani. Regularization paths for generalized
linear models via coordinate descent. Journal of Statistical Software, 33(1):1, 2010.
Jerome H Friedman. Multivariate adaptive regression splines. The annals of statistics, pages
1–67, 1991.
Jennifer L Hill. Bayesian nonparametric modeling for causal inference. Journal of Computational and Graphical Statistics, 20(1), 2011.
20

Keisuke Hirano and Jack R Porter. Asymptotics for statistical treatment rules. Econometrica, 77(5):1683–1701, 2009.
Kosuke Imai and Marc Ratkovic. Estimating treatment effect heterogeneity in randomized
program evaluation. The Annals of Applied Statistics, 7(1):443–470, 2013.
Diederik Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv
preprint arXiv:1412.6980, 2014.
Vladimir Koltchinskii. Local Rademacher complexities and oracle inequalities in risk minimization. The Annals of Statistics, 34(6):2593–2656, 2006.
Sören Künzel, Jasjeet Sekhon, Peter Bickel, and Bin Yu. Meta-learners for estimating
heterogeneous treatment effects using machine learning. arXiv preprint arXiv:1706.03461,
2017.
Yann LeCun. The mnist database of handwritten digits. 1998.
Shahar Mendelson and Joseph Neeman. Regularization in kernel learning. The Annals of
Statistics, 38(1):526–565, 2010.
Jersey Neyman. Sur les applications de la théorie des probabilités aux experiences agricoles:
Essai des principes. Roczniki Nauk Rolniczych, 10:1–51, 1923.
Ziad Obermeyer and Ezekiel J Emanuel. Predicting the futurebig data, machine learning,
and clinical medicine. The New England journal of medicine, 375(13):1216, 2016.
Adam Paszke, Sam Gross, Soumith Chintala, and Gregory Chanan. PyTorch: Tensors and
dynamic neural networks in Python with strong GPU acceleration, 2017. URL https:
//github.com/pytorch/pytorch.
Scott Powers, Junyang Qian, Kenneth Jung, Alejandro Schuler, Nigam H Shah, Trevor
Hastie, and Robert Tibshirani. Some methods for heterogeneous treatment effect estimation in high-dimensions. arXiv preprint arXiv:1707.00102, 2017.
Peter M Robinson. Root-n-consistent semiparametric regression. Econometrica: Journal of
the Econometric Society, pages 931–954, 1988.
Paul R Rosenbaum and Donald B Rubin. The central role of the propensity score in
observational studies for causal effects. Biometrika, 70(1):41–55, 1983.
Donald B Rubin. Estimating causal effects of treatments in randomized and nonrandomized
studies. Journal of Educational Psychology, 66(5):688, 1974.
Anton Schick. On asymptotically efficient estimation in semiparametric models. The Annals
of Statistics, pages 1139–1151, 1986.
Uri Shalit, Fredrik D. Johansson, and David Sontag. Estimating individual treatment effect:
generalization bounds and algorithms. In Proceedings of the 34th International Conference
on Machine Learning, pages 3076–3085, 2017.
Steve Smale and Ding-Xuan Zhou. Estimating the approximation error in learning theory.
Analysis and Applications, 1(01):17–41, 2003.

21

Ingo Steinwart and Andreas Christmann. Support Vector Machines. Springer Science &
Business Media, 2008.
Ingo Steinwart, Don R Hush, and Clint Scovel. Optimal rates for regularized least squares
regression. In Conference on Learning Theory, 2009.
Xiaogang Su, Chih-Ling Tsai, Hansheng Wang, David M Nickerson, and Bogong Li. Subgroup analysis via recursive partitioning. The Journal of Machine Learning Research, 10:
141–158, 2009.
Robert Tibshirani. Regression shrinkage and selection via the lasso. Journal of the Royal
Statistical Society: Series B (Statistical Methodology), pages 267–288, 1996.
Stefan Wager and Susan Athey. Estimation and inference of heterogeneous treatment effects
using random forests. Journal of the American Statistical Association, (just-accepted),
2017.
Qingyuan Zhao, Dylan S Small, and Ashkan Ertefaie. Selective inference for effect modification via the lasso. arXiv preprint arXiv:1705.08020, 2017.

22

