Extracted main text — title through conclusion, appendix excluded. This is what our citation measures are computed over, published so the extraction can be checked by eye.
69,973 characters · 14 sections · 76 citation commands
Quasi-Oracle Estimation of Heterogeneous Treatment Effects
Keywords: Boosting; Causal inference; Empirical risk minimization; Kernel regression; Penalized regression.
The problem of heterogeneous treatment effect estimation in observational studies arises in a wide variety application areas athey2017beyond, ranging from personalized medicine obermeyer2016predicting to offline evaluation of bandits \citep*{dudik2011doubly}, and is also a key component of several proposals for learning decision rules athey2017efficient,hirano2009asymptotics. 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 imai2013estimating, recursive partitioning \citep*{athey2015machine,su2009subgroup}, BART \citep*{hahn2020bayesian,hill2011bayesian}, random forests \citep*{wager2017estimation}, boosting powers2018some, neural networks \citep*{shalit2017estimating}, etc., as well as combinations thereof \citep*{kunzel2019metalearners}; see dorie2019automated for a recent survey and comparisons.
However, although this line of work has led to many promising methods, the literature has not yet settled on a comprehensive answer as to how machine learning methods should be adapted for treatment effect estimation in observational studies. The process of developing causal variants of machine learning methods is in practice a labor intensive process, effectively requiring the involvement of specialized researchers. Moreover, with some exceptions, the above methods are mostly justified via numerical experiments, and come with no formal convergence guarantees or error bounds proving that the methods in fact succeed in isolating causal effects better than a simple non-parametric regression-based approach would.
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 error bounds to oracle methods that know everything about the data-generating distribution except the treatment effects. Conceptually, our approach fits into a research program---outlined by van2003cross and later developed by \citet*{chernozhukov2017double}, luedtke2016super and references therein---whereby we pair ideas on doubly robust estimation with oracle inequalities and cross-validation to develop loss functions that can be used for principled statistical estimation using generic machine learning tools.
We formalize our problem in terms of the potential outcomes framework neyman1923applications,rubin1974estimating. The analyst has access to $n$ independent and identically distributed examples $(X_i, \, Y_i, \, W_i)$, $i = 1, \, ..., \, n$, where $X_i \in \xx$ denotes per-person features, $Y_i \in \RR$ is the observed outcome, and $W_i \in \cb{0, \, 1}$ is the treatment assignment. We posit the existence of potential outcomes $\cb{Y_i(0), \, Y_i(1)}$ corresponding to the outcome we would have observed given the treatment assignment $W_i = 0$ or $1$ respectively, such that $Y_i = Y_i(W_i)$, and seek to estimate the conditional average treatment effect (CATE) function $\tau^*(x) = E\{Y(1) - Y(0) \cond X = x\}.$ In order to identify $\tau^*(x)$, we assume unconfoundedness, i.e., the treatment assignment is randomized once we control for the features $X_i$ rosenbaum1983central.
We write the treatment propensity as \smash{$e^*(x) = \PP{W = 1 \cond X = x}$} and the conditional response surfaces as \smash{$\mu^*_{(w)}(x) = E\{Y(w) \cond X = x\}$} for $w\in \{0,1\}$; throughout this paper, we use $*$-superscripts to denote unknown population quantities. Then, under unconfoundedness,
Given this setup, it is helpful to re-write the CATE function $\tau^*(x)$ in terms of the conditional mean outcome \smash{$m^*(x) = \EE{Y \cond X = x} = \mu_{(0)}^*(X_i) + e^*(X_i) \tau^*(X_i)$} as follows, with the shorthand $\varepsilon_i := \varepsilon_i(W_i)$,
This decomposition was originally used by robinson1988root to estimate parametric components in partially linear models, and has received considerable attention in recent years. \citet*{athey2018generalized} rely on it to grow a causal forest that is robust to confounding, robins2004optimal builds on it in developing $G$-estimation for sequential trials, and chernozhukov2017double present it as a leading example on how machine learning methods can be put to good use in estimating nuisance components for semiparametric inference. All these results, however, consider estimating parametric models for $\tau(\cdot)$ or, in the case of athey2018generalized, local parametric modeling.
The goal of this paper is to study how we can use the Robinson's transfomation (ref) for flexible treatment effect estimation that builds on modern machine learning approaches such as boosting or deep learning. Our main result is that we can use this representation to construct a loss function that captures heterogeneous treatment effects, and that we can then accurately estimate treatment effects---both in terms of empirical performance and asymptotic guarantees---by finding regularized minimizers of this loss function.
As motivation for our approach, note that (ref) can equivalently be expressed as robins2004optimal
and so an oracle who knew both the functions $m^*(x)$ and $e^*(x)$ a priori could estimate the heterogeneous treatment effect function $\tau^*(\cdot)$ by empirical loss minimization,
where the term $\Lambda_n\p{\tau(\cdot)}$ is interpreted as a regularizer on the complexity of the $\tau(\cdot)$ function. 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 in practice 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 (ref) is not feasible.
Given these preliminaries, we here study the following class of two-step estimators using cross-fitting chernozhukov2017double, schick1986asymptotically motivated by the above oracle procedure:
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 the $R$-learner in recognition of the work of robinson1988root and to emphasize the role of residualization. We will also refer to the squared loss \smash{$\hL_n\{\tau(\cdot)\}$} as the $R$-loss.
This paper makes the following contributions. First, we implement variants of our method based on penalized regression, kernel ridge regression, and boosting. In each case, we find that the $R$-learner exhibits promising performance relative to existing proposals. Second, we prove that---in the case of penalized kernel regression---error bounds for the feasible estimator for \smash{$\htau(\cdot)$} asymptotically match the best available bounds for the oracle method \smash{$\ttau(\cdot)$}. The main point here is that, heuristically, the rate of convergence of \smash{$\htau(\cdot)$} depends only on the functional complexity of $\tau^*(\cdot)$, and not on the functional complexity of $m^*(\cdot)$ and $e^*(\cdot)$. More formally, provided we estimate $m^*(\cdot)$ and $e^*(\cdot)$ at $o(n^{-1/4})$ rates in root-mean squared error, we show that we can achieve considerably faster rates of convergence for \smash{$\htau(\cdot)$}---and these rates only depend on the complexity of $\tau^*(\cdot)$. We note that the oracle version (ref) of our loss function is a member of a class of loss functions for heterogeneous treatment effect estimation considered in luedtke2016super, and that results in that paper immediately imply large-sample consistency of the minimizer of this oracle loss. Our contribution is the result on rates---specifically, that the estimation error in nuisance components does not affect our excess loss bounds for \smash{$\htau(\cdot)$}.
The $R$-learning approach has several practical advantages over existing, more ad hoc proposals. Any good heterogeneous treatment effect estimator needs to achieve two goals: First, it needs to eliminate spurious effects by controlling for correlations between $e^*(X)$ and $m^*(X)$, and then it needs to accurately express $\tau^*(\cdot)$. Most existing machine learning approaches to treatment effect estimation seek to provide an algorithm that accomplishes both tasks at once powers2018some,shalit2017estimating,wager2017estimation. In contrast, the $R$-learner cleanly separates these two tasks: We eliminate spurious correlations via the structure of the loss function \smash{$\hL_n$}, while we can induce a representation for \smash{$\htau(\cdot)$} by choosing the method by which we optimize (ref).
This separation of tasks allows for considerable algorithmic flexibility: Optimizing (ref) is an empirical minimization problem, and so can be efficiently solved via off-the-shelf software such as glmnet for high-dimensional regression \citep*{friedman2010regularization}, XGboost for boosting chen2016xgboost, or TensorFlow for deep learning tensorflow. Furthermore, we can tune any of these methods by cross validating on the loss \smash{$\hL_n$}, which avoids the use of more sophisticated model-assisted cross-validation procedures as developed in athey2015machine or powers2018some. Relatedly, the machine learning method used to optimize (ref) only needs to find a generalizable minimizer of \smash{$\hL_n$} rather than to also control for spurious correlations, and thus we can confidently use black-box methods without auditing their internal state to check that they properly control for confounding. Instead, we only need to verify that they in fact find good minimizers of \smash{$\hL_n$} on holdout data.
Under unconfoundedness (Assumption (ref)), the CATE function can be written as $\tau^*(x) = \mu_{(1)}^*(x) - \mu_{(0)}^*(x)$, with $\mu_{(w)}^*(x) = \EE{Y \cond X = x, \, W = w}$. As a consequence of this representation, it may be tempting to first estimate \smash{$\hmu_{(w)}(x)$} on the treated and control samples separately, and then set \smash{$\htau(x) = \hmu_{(1)}(x) - \hmu_{(0)}(x)$}. This approach, however, is often not robust: Because \smash{$\hmu_{(1)}(x)$} and \smash{$\hmu_{(0)}(x)$} are not trained together, their difference may be unstable. As an example, consider fitting the lasso tibshirani1996regression to estimate $\hat{\mu}_{(1)}(x)$ and $\hat{\mu}_{(0)}(x)$ in the following high-dimensional linear model, \smash{$Y_i(w) = X_i^\top \beta^*_{(w)} + \varepsilon_i(w)$} with \smash{$X_i, \beta^*_{(w)} \in \mathbb{R}^d$}, and \smash{$\EE{\varepsilon_i(w) \cond X_i} = 0$}. A naive approach would fit two separate lassos to the treated and control samples,
and then use it to deduce a treatment effect function, \smash{$\htau(x) = x^\top (\hbeta_{(1)} - \hbeta_{(0)})$}. However, the fact that both \smash{$\hbeta_{(0)}$} and \smash{$\hbeta_{(1)}$} are regularized towards 0 separately may inadvertently regularize the treatment effect estimate \smash{$\hbeta_{(1)} - \hbeta_{(0)}$} away from 0, even when $\tau^*(x) = 0$ everywhere. This problem is especially acute when the treated and control samples are of different sizes; see \citet*{kunzel2019metalearners} for some striking examples.
The recent literature on heterogeneous treatment effect estimation has proposed several ideas on how to avoid such regularization bias. Some recent papers have proposed structural changes to various machine learning methods aimed at focusing on accurate estimation of $\tau(\cdot)$ athey2015machine,hahn2020bayesian,imai2013estimating,powers2018some,shalit2017estimating, su2009subgroup,wager2017estimation. For example, with the lasso, imai2013estimating advocate replacing (ref) with a single lasso as follows,
where then \smash{$\htau(x) = x^\top \hdelta$}. This approach always correctly regularizes towards a sparse $\delta$-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 \citet*{shalit2017estimating} for a neural network architecture designed to highlight treatment effect heterogeneity without being affected by confounders.
Here, instead of trying to modify the algorithms underlying different machine learning tools to improve their performance as treatment effect estimators, we focus on modifying the loss function used to training generic machine learning methods. In doing so, we build on the research program developed in van2003cross, van2006targeted and \citet*{van2007super}, and later fleshed out for the context of individualized treatment rules by luedtke2016optimal,luedtke2016statistical,luedtke2016super. In an early technical report, van2003cross discuss choosing the best among a potentially growing set of generic statistical rules by cross-validating on a doubly robust objective. In the case without nuisance components, an $\varepsilon$-net version of this procedure was shown to have good asymptotic properties \citep*{van2003cross,van2006cross}. Meanwhile, luedtke2016super discuss a class of valid objectives for learning either individualized treatment rules or heterogeneous treatment effects---the oracle version (ref) of our loss function fits within this class---and discuss properties of model averaging and cross-validation with these objectives. Our contributions with respect to this line of work include using the $R$-loss for treatment effect estimation via generic machine learning and developing strong excess loss bounds \smash{$\htau(\cdot)$} that hold for a computationally tractable and widely used approach to non-parametric estimation, namely penalized regression over a reproducing kernel Hilbert space.
Another closely related trend in the literature has focused on meta-learning approaches that are not closely tied to any specific machine learning method. \citet*{kunzel2019metalearners} proposed two approaches to heterogeneous treatment effect estimation via generic machine learning methods. One, called the $X$-learner, first estimates $\hmu_{(w)}(x)$ via appropriate non-parametric regression methods. Then, on the treated observations, it defines pseudo-effects \smash{$D_i = Y_i - \hmu_{(0)}^{(-i)}(X_i)$}, and uses them to fit \smash{$\htau_{(1)}(X_i)$} via a non-parametric regression. Another estimator \smash{$\htau_{(0)}(X_i)$} is obtained analogously, and the two treatment effect estimators are aggregated as
Another method, called the $U$-learner, starts by noticing that
and then fitting $U_i$ on $X_i$ using any off-the-shelf method. Relatedly, athey2015machine and tian2014simple develop methods for heterogeneous treatment effect estimation based on weighting the outcomes or the covariates with the propensity score; for example, we can estimate $\tau^*(\cdot)$ by regressing $Y_i \{W_i - e^*(X_i)\} / \{e^*(X_i)(1 - e^*(X_i))\}$ on $X_i$. In our experiments, we compare our method at length to those of kunzel2019metalearners. Again, relative to this line of work, our main contribution is our method, the $R$-learner, which provides meaningful improvements over baselines in a variety of settings, and our analysis, which provides a quasi-oracle error bound for the conditional average treatment effect function, i.e., where the error of \smash{$\htau$} may decay faster than that of \smash{$\he$} or \smash{$\hatm$}.
The closest result to us in this line of work is from \citet*{zhao2017selective}, who combine Robinson's transformation with the lasso to provide valid post-selection inference on effect modification in the high-dimensional linear model. To our knowledge, our paper is the first to use Robinson's transformation to motivate a loss function that is used in a general machine learning context.
Our formal results draw from the literature on semiparametric efficiency and constructions of orthogonal moments including robinson1988root and, more broadly, \citet*{belloni2017program}, \citet*{bickel1993efficient}, \citet*{chernozhukov2017double}, newey1994asymptotic, robins2004optimal, robins1, \citet*{robins2017minimax}, tsiatis2007semiparametric, van2011targeted, etc., that aim at $\sqrt{n}$-rate estimation of a target parameter in the presence of nuisance components that cannot be estimated at a $\sqrt{n}$ rate. Algorithmically, our approach has a close connection to targeted maximum likelihood estimation \citep*{scharfstein1999adjusting,van2006targeted}, which starts by estimating nuisance components non-parametrically, and then uses these first stage estimates to define a likelihood function that is optimized in a second step. We also note that using held-out prediction for nuisance components, also known as cross-fitting, is an increasingly popular approach for making machine learning methods usable in classical semiparametrics athey2017efficient,chernozhukov2017double,schick1986asymptotically,van2011targeted,wager2016high.
The main difference between this literature and our results is that existing results typically focus on estimating a single (or low-dimensional) target parameter, whereas we seek to estimate an object $\tau^*(\cdot)$ that may also be quite complicated itself. Another research direction that also uses ideas from semiparametrics to estimate complex objects is centered on estimating optimal treatment allocation rules \citep*{athey2017efficient,dudik2011doubly,laber2015tree,luedtke2016super,zhang2012estimating}. This problem is closely related to, but subtly different from the problem of estimating $\tau^*(\cdot)$ under squared-error loss; see kitagawa2018should.
Finally, we note that all results presented here assume a sampling model where observations are drawn at random from a population, and we define our target estimand $\tau(\cdot)$ in terms of moments of that population. \citet*{ding2019decomposing} consider heterogeneous treatment effect estimation in a strict randomization inference setting, where we the features and potential outcomes \smash{$\cb{X_i, \, Y_i(0), \, Y_i(1)}_{i = 1}^n$} are taken as fixed and only the treatment $W_i$ is random imbens_rubin_2015; they then show how to estimate the projection of the realized treatment heterogeneity $Y_i(1) - Y_i(0)$ onto the linear span of the $X_i$. It would be interesting to consider whether it is possible to derive useful results on non-parametric (regularized) heterogeneous treatment effect estimation under randomization inference.
To see how the $R$-learner works in practice, we consider an example motivated by \citet*{arceneaux2006comparing}, who studied the effect of paid get-out-the-vote calls on voter turnout. A common difficulty in comparing the accuracy of heterogeneous treatment effect estimators on real data is that we do not have access to the ground truth. From this perspective, a major advantage of this application is that arceneaux2006comparing found no effect of get-out-the-vote calls on voter turnout, which suggests that the underlying effect is close to nonexistant. We then spike the original dataset with a synthetic treatment effect $\tau^*(\cdot)$ such as to make the task of estimating heterogeneous treatment effects non-trivial. In other words, both the baseline signal and propensity scores are from real data; however, $\tau^*(\cdot)$ is chosen by us, and so we can check whether different methods in fact succeed in recovering it.
The design of arceneaux2006comparing was randomized separately by state and competitiveness of the election, and accounting for varying treatment propensities is necessary for obtaining correct causal effects: A naive analysis ignoring variable treatment propensities estimates the average effect of a single get-out-the-vote call on turnout as 4%, whereas an appropriate analysis finds with high confidence that any treatment effect must be smaller than 1% in absolute value. Although the randomization probabilities were known to the experimenters, we here hide them from our algorithm, and require it to learn a model \smash{$\he(\cdot)$} for the treatment propensities. We also note that, in the original data, not all voters assigned to be contacted could in fact answer the phone call, meaning that all effects should be interpreted as intent to treat effects. We focus on $d = 11$ covariates (including state, county, age, gender, etc.). Both the outcome $Y$ and the treatment $W$ are binary. The full sample has $1,895,468$ observations, of which $59,264$ were assigned treatment. For our analysis, we focused on a subset of $148,160$ samples containing all the treated units and a random subset of the controls; thus, 2/5 of our analysis sample was treated. We further divided this sample into a training set of size 100,000, a test set of size 25,000, and a holdout set with the rest.
As discussed above, for the purpose of this evaluation, we assume that the treatment effect in the original data is 0, and spike in a synthetic treatment effect $\tau^*(X_i) = - \text{VOTE00}_i / (2 + 100 / \text{AGE}_i)$, where $\text{VOTE00}_i$ indicates whether the $i$-th unit voted in the year 2000, and $\text{AGE}_i$ is their age. Because the outcomes are binary, we add in the synthetic treatment effect by strategically flipping some outcome labels. Denote the original unflipped outcomes as $Y_i^*$. To add in a treatment effect $\tau^*(\cdot)$, we first draw Bernoulli random variables $R_i$ with probability $\abs{\tau^*(X_i)}$. Then, if $R_i = 0$, we set $Y_i(0) = Y_i(1) = Y_i^*$, whereas if $R_i = 1$, we set $\{Y_i(0), \, Y_i(1)\}$ to $(0, \, 1)$ or $(1, \, 0)$ depending on whether $\tau^*(X_i) > 0$ or $\tau^*(X_i) < 0$ respectively. Finally, we set $Y_i = Y_i(W_i)$. As is typical in causal inference applications, the treatment heterogeneity here is quite subtle, with $\textrm{var}\{\tau^*(X)\} = 0.016$, and so a large sample size is needed in order to reject a null hypothesis of no treatment heterogeneity.
To use the $R$-learner, we first estimated \smash{$\he(\cdot)$} and \smash{$\hatm(\cdot)$} to form the $R$-loss function in (ref). To do so, we fit models for the nuisance components via both boosting and the lasso with tuning parameters selected via cross-validation. Then, we chose the model that minimized cross-validated error. This criterion lead us to pick boosting for both \smash{$\he(\cdot)$} and \smash{$\hatm(\cdot)$}. Another option would have been to combine predictions from the lasso and boosting models, as advocated by \citet*{van2007super}.
Next, we optimized the $R$-loss function. We again tried methods based on both the lasso and boosting. This time, the lasso achieved a slightly lower training set cross-validated $R$-loss than boosting, namely 0.1816 versus 0.1818. Because treatment effects are so weak and so there is potential to overfit even in cross-validation, we also examined $R$-loss on the holdout set. The lasso again came out ahead, and the improvement in $R$-loss is stable, 0.1781 versus 0.1783. We thus chose the lasso-based \smash{$\htau(\cdot)$} fit as our final model for $\tau^*(\cdot)$. As an aside, we note that although the improvement in $R$-loss is stable, the loss itself is somewhat different between the training and holdout samples. This appears to be due to the term \smash{$n^{-1} \sum_i \{Y_i - \mu^*_{(W_i)}(X_i)\}^2$} induced by irreducible outcome noise. This term is large and noisy in absolute terms; however, it gets canceled out when comparing the accuracy of two models. This phenomenon plays a key role in understanding the behavior of model selection via cross-validation wager2020cross,yang2007consistency.
\sloppy{ Given the constructed CATE function $\tau^*(\cdot)$ in our semi-synthetic data generative distribution, we can evaluate the oracle test set mean-squared error, \smash{$1/n_{test} \sum_{\cb{i \in test}} \{\htau(X_i) - \tau^*(X_i)\}^2$}. Here, it is clear that the lasso did substantially better than boosting, achieving a mean-squared error of $0.47 \times 10^{-3}$ versus $1.23 \times 10^{-3}$. The right panel of Figure (ref) compares \smash{$\htau(\cdot)$} estimates from minimizing the $R$-loss using the lasso and boosting respectively. The lasso is somewhat biased, but boosting is noisy, and the bias-variance trade-off favors the lasso in this case. With a larger sample size, we'd expect boosting to achieve lower mean-squared error.}
We also compared our approach to both the single lasso approach (ref), and a popular non-parametric approach to heterogeneous treatment effect estimation via BART hill2011bayesian, with the estimated propensity score added in as a feature following the recommendation of \citet*{hahn2020bayesian}. The single lasso got an oracle test set error of $0.61 \times 10^{-3}$, whereas BART got $4.05 \times 10^{-3}$. It thus appears that, in this example, there is value in using a non-parametric method for estimating \smash{$\he(\cdot)$} and \smash{$\hatm(\cdot)$}, but then using the simpler lasso for \smash{$\htau(\cdot)$}. In contrast, the single lasso approach uses linear modeling everywhere (thus leading to potential model misspecification and confounding), whereas BART uses non-parametric modeling everywhere, which can make it difficult to obtain a stable $\tau(\cdot)$ fit. Section (ref) has a more comprehensive simulation evaluation of the $R$-learner relative to several baselines, including the meta-learners of \citet*{kunzel2019metalearners}.
In the previous section, we considered an example application where we were willing to carefully consider the estimation strategies used in each step of the $R$-learner. In other cases, however, a practitioner may prefer to use some off-the-shelf treatment effect estimators as the starting point for their analysis. Here, we discuss how to use the $R$-learning approach to build a consensus treatment effect estimate via a variant of stacking \citep*{breiman1996stacked,luedtke2016super,van2007super,wolpert1992stacked}.
Suppose we start with $k = 1, \, ..., \, K$ different treatment effect estimators \smash{$\htau_k$}, and that we have access to out-of-fold estimates \smash{$\htau_k^{(-i)}(X_i)$} on our training set. Suppose, moreover, that we have trusted out-of-fold estimates $\he^{(-i)}(X_i)$ and $\hatm^{(-i)}(X_i)$ for the propensity score and main effect respectively. Then, we propose building a consensus estimate \smash{$\htau(\cdot)$} by taking the best positive linear combination of the \smash{$\htau_k(\cdot)$} according to the $R$-loss:
For flexibility, we also allow the stacking step (ref) to freely adjust a constant treatment effect term $c$, and we add an intercept $b$ that can be used to absorb any potential bias of \smash{$\hatm$}.
We test this approach on the following data-generation distributions. In both cases, we drew $n = 10,000$ i.i.d. samples from a randomized study design, $X_i \sim \nn\p{0, \,I_{d \times d}}$
for different choices of $\tau^*(\cdot)$ and $\sigma$, and with $d = 10$. We consider both a smooth treatment effect function $\tau^*(X_i) = 1/(1 + e^{X_{i1} - X_{i2}})$, and a discontinuous $\tau^*(X_i) = \mathbbm{1}\{(X_{i1} > 0)\}/(1 + e^{- X_{i2}})$. Given this data-generating process, we tried estimating $\tau(\cdot)$ via BART \citep*{hahn2020bayesian,hill2011bayesian}, causal forests \citep*{athey2018generalized,wager2017estimation}, and a stacked combination of the two using (ref). We assume that the experimenter knows that the data was randomized, and used \smash{$\he(x) = 0.5$} in any place a propensity score was needed. For stacking, we estimated \smash{$\hatm(\cdot)$} using a random forest.
Results are shown in Figure (ref). In the example with a smooth $\tau^*(\cdot)$, BART slightly out-performs causal forests, while stacking does better than either on its own until the noise level $\sigma$ gets very large---in which case none of the methods do much better than a constant treatment effect estimator. Meanwhile, the setting with the discontinuous $\tau^*(\cdot)$ appears to be particularly favorable to causal forests, at least for lower noise levels. Here, stacking is able to automatically match the performance of the more accurate base learner.
As discussed in the introduction, the high-level goal of our formal analysis is to establish error bounds for $R$-learning that only depend on the complexity of $\tau^*(\cdot)$, and that match the error bounds we could achieve if we knew $m^*(\cdot)$ and $e^*(\cdot)$ a-priori. In order to do so, we focus on a variant of the $R$-learner 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 bartlett2006empirical,caponnetto2007optimal,cucker2002mathematical,mendelson2010regularization, steinwart2008support, and thus provides an ideal case study for examining the asymptotic behavior of the $R$-learner.
We study \smash{$\Norm{\cdot}_\hh$}-penalized kernel regression, where $\hh$ is a reproducing kernel Hilbert space (RKHS) with a continuous, positive semi-definite kernel function $\mathcal{K}$. Let $\pp$ be a non-negative measure over the compact metric space $\xx \subset \mathbb{R}^d$, and let $\mathcal{K}$ be a kernel with respect to $\pp$. Let \smash{$T_\mathcal{K}: L_2(\pp)\to L_2(\pp)$} be defined as \smash{$T_\mathcal{K}(f)(\cdot) = E\{\mathcal{K}(\cdot, \, X) f(X)\}$}. By Mercer's theorem cucker2002mathematical, there is an orthonormal basis of eigenfunctions $(\psi_j)_{j=1}^\infty$ of $T_\mathcal{K}$ with corresponding eigenvalues $(\sigma_j)_{j=1}^\infty$ such that $\mathcal{K}(x,y) = \sum_{j=1}^\infty \sigma_j \psi_j(x) \psi_j(y)$. Consider the function \smash{$\phi: \xx \to l_2$} defined by \smash{$\phi(x) = (\sqrt{\sigma_j}\psi_j(x))_{j=1}^\infty$}. Following mendelson2010regularization, we define the RKHS $\hh$ to be the image of $l_2$: For every $t\in l_2$, define the corresponding element in $\hh$ by \smash{$f_t(x) = \langle \phi(x), t \rangle$} with the induced inner product \smash{$\langle f_s, f_t \rangle_\hh = \langle t, s \rangle$}.
To interpret the assumption above, note that we do not assume that $\tau^*(\cdot)$ has a finite $\hh$-norm; rather, we only assume that we can make it have a finite $\hh$-norm after a sufficient amount of smoothing. More concretely, with \smash{$\alpha = 0$}, \smash{$T_\mathcal{K}^\alpha$} would be the identity operator, and so this assumption would be equivalent to the strongest possible assumption that \smash{$\Norm{\tau^*(\cdot)}_\hh < \infty$} itself. Then, as \smash{$\alpha$} grows, this assumption gets progressively weaker, and at $\alpha = 1/2$ it would devolve to simply asking that \smash{$\tau^*(\cdot)$} belong to the space \smash{$L_2(\mathcal{P})$} of square-integrable functions.
We study oracle penalized regressions $\ttau(\cdot)$ that minimize the following objective,
as well as feasible analogues obtained by cross-fitting chernozhukov2017double,schick1986asymptotically:
Adding the upper bound \smash{$\Norm{\tau}_\infty \leq 2M$} (or, in fact, any finite upper bound on $\tau$) enables us to rule out some pathological behaviors.
We seek to characterize the accuracy of our estimator \smash{$\htau(\cdot)$} by bounding its regret \smash{$R(\htau)$},
Recall that, by the expansion (ref), we have \smash{$E\{Y_i - m^*(X_i) \cond X_i, \, W_i\} = \tau^*(X_i)\{W_i - e^*(X_i)\}$}, implying that $$ L(\tau) = E[\textrm{var}\{Y_i - m^*(X_i) \cond X_i, \, W_i\}] + E[ \{\tau(X_i)-\tau^*(X_i)\}^2\p{W_i-e^*(X_i)}^2], $$ and \smash{$R(\tau) = E[ \{\tau(X_i)-\tau^*(X_i)\}^2\{W_i-e^*(X_i)\}^2]$}. Thus if we have overlap, i.e., there is an $\eta > 0$ such that $\eta < e^*(x) < 1 - \eta$ for all $x \in \xx$, then
meaning that regret bounds translate into squared-error loss bounds for $\tau(\cdot)$, and vice-versa. We note that when the overlap parameter $\eta$ gets close to 0, the coupling (ref) gets fairly loose.
The sharpest regret bounds for the oracle learner (ref) under Assumptions (ref) and (ref) are due to mendelson2010regularization (see also \citet*{steinwart2009optimal}), and scale as
where the \smash{$\too_P$}-notation hides logarithmic factors. In the case $\alpha = 0$ where $\tau^*$ is within the RKHS used for penalization, we recover the more familiar $n^{-1/(1+p)}$ rate established by caponnetto2007optimal. Again, our goal is to establish excess loss bounds for our feasible estimator \smash{$\htau$} that match the bound (ref) available to the oracle that knows $m^*(\cdot)$ and $e^*(\cdot)$ a-priori.
In order to do so, we first need to briefly review the proof techniques underlying (ref). The argument of mendelson2010regularization relies on the following quasi-isomorphic coordinate projection lemma of bartlett2008fast. To state this result, write
for the radius-$c$ ball of $\hh$ capped by $2M$, let \smash{$\tau^*_c = \argmin\cb{L(\tau) : \tau \in \hh_c}$} denote the best approximation to $\tau^*$ within $\hh_c$, and define $c$-regret \smash{$R(\tau; \, c) = L(\tau) - L(\tau^*_c)$} over $\tau \in \hh_c$. We also define the estimated and oracle $c$-regret functions $\hR_n$ and $\tR_n$ written in terms of the estimated and oracle losses $\hL_n$ and $\tL_n$:
\smash{$\hat{R}_n(\tau; \, c)$} is not observable as it depends on $\tau^*_c$; however, this does not hinder us from establishing high-probability bounds for it. The lemma below is adapted from bartlett2008fast.
In other words, the above lemma reduces the problem of deriving regret bounds to establishing quasi-isomorphisms as in (ref), and any with-high-probability quasi-isomorphism guarantee yields a with-high-probability regret bound. In particular, we can use this approach to prove the regret bound (ref) for the oracle learner as follows. We first need a with-high-probability quasi-isomorphism of the following form,
mendelson2010regularization provide such a bound for $\rho_n(c)$ that scales as
Lemma (ref) then immediately implies that penalized regression over $\hh_C$ with the oracle loss function \smash{$\tL(\cdot)$} and regularizer $\kappa_1\rho_n(c)$ satisfies the bound below with high probability:
Furthermore, following Corollary 2.7 in mendelson2010regularization, for any $1 \leq c \leq C$, we also have
Note that mendelson2010regularization considers the case where $C = \infty$; here, instead, we only take $C$ to be large enough for our argument; see the proof for details. Finally, noting the scaling of $\rho_n(c)$ in (ref) and the approximation error bound
established by smale2003estimating under the setting of Assumption (ref), we achieve a practical regret bound by choosing $c = c_n$ to optimize the right-hand side of (ref). The specific rate in (ref) arises by setting \smash{$c_n = n^{{\alpha}/(p + (1 - 2\alpha))}$}.
For our purposes, the upshot is that if we can match the strength of the quasi-isomorphism bounds (ref) with our feasible loss function, i.e., get an analogous bound in terms of \smash{$\hR_n$} as opposed to \smash{$\tR_n$}, then we can also match the rate of any regret bounds proved using the above argument. The proof of the following result relies several concentration results, including Talagrand's inequality and generic chaining talagrand2006generic, and makes heavy use of cross-fitting style arguments chernozhukov2017double,schick1986asymptotically,van2011targeted.
This result implies that we can turn any quasi-isomorphism for the oracle learner (ref) with error $\rho_n(c)$ into a quasi-isomoprhism bound for \smash{$\hR(\tau)$} with error inflated by the right hand side of (ref). Thus, given any regret bound for the oracle learner built using Lemma (ref), 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.
In other words, we have found that with penalized kernel regression, the $R$-learner can match the best available performance guarantees available for the oracle learner (ref) that knows everything about the data generating distribution except the true treatment effect function---both the feasible and the oracle learner satisfy
As we approach the semiparametric case, i.e., $\alpha, \, p \rightarrow 0$, we recover the well-known result from the semiparametric inference literature that, in order to get $1/\sqrt{n}$-consistent inference for a single target parameter, we need 4-th root consistent nuisance parameter estimates; see chernozhukov2017double for a review and references. We also note that after we disseminated a first draft of our paper, several authors have established further quasi-oracle type results for the $R$-learner and related methods; see in particular foster2019orthogonal and kennedy2020optimal.
We emphasize that our quasi-oracle result depends on a local robustness property of the $R$-loss function, and does not hold for general meta-learners; for example, it does not hold for the $X$-learner of \citet*{kunzel2019metalearners}. To see this, we argue by contradiction: We show that it is possible to make $o(n^{-1/4})$-changes to the nuisance components \smash{$\hmu_{(w)}(x)$} used by the $X$-learner that induce changes in the $X$-learner's \smash{$\htau(\cdot)$} estimates that dominate the error scale in (ref). Thus, there must be some choices of $o(n^{-1/4})$-consistent \smash{$\hmu_{(w)}(x)$} with which the $X$-learner does not converge at the rate (ref). The contradiction arises as follows: Pick $\xi > 0$ such that $0.25 + \xi < (1 - 2\alpha)/[2\{p + (1 - 2\alpha)\}]$, and modify the nuisance components used to form the $X$-learner in (ref) such that \smash{$\hmu_{(0)}(x) \leftarrow \hmu_{(0)}(x) - c / n^{0.25 + \xi}$} and \smash{$\hmu_{(1)}(x) \leftarrow \hmu_{(1)}(x) + c / n^{0.25 + \xi}$}. Recall that the $X$-learner fits \smash{$\hat{\tau}_{(1)}(\cdot)$} by minimizing ${n_1^{-1}} \sum_{W_i = 1} \{Y_i - \hat{\mu}_{(0)}^{(-i)}(X_i) - \tau_{(1)}(X_i)\}^2$, and fits \smash{$\hat{\tau}_{(0)}(\cdot)$} by solving an analogous problem on the controlled units. Combining the $\hat{\tau}_{(w)}$ estimates from these two loss functions, we see by inspection that its final estimate of the treatment effect is also shifted by \smash{$\htau(x) \leftarrow \htau(x) + c / n^{0.25 + \xi}$}. The perturbations \smash{$c / n^{0.25 + \xi}$} are vanishingly small on the $n^{-1/4}$ scale, and so would not affect conditions analogous to those of Theorem (ref); yet they have a big enough effect on $\htau(x)$ to break any convergence results on the scale of (ref). We note that kunzel2019metalearners do have some quasi-oracle type results; however, they only focus on the case where the number of control units $\abs{\cb{W_i = 0}}$ grows much faster than the number of treated units $\abs{\cb{W_i = 1}}$. 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, \smash{$\mu^*_{(0)}(x) = \EE{Y_i(0) \cond X_i = x}$}. Intriguingly, in this special case, we have \smash{$m^*(x) \approx \mu^*_{(0)}(x)$} and $e^*(x) \approx 0$, and so the $R$-learner as in (ref) is roughly equivalent to the $X$-learner procedure (ref). Thus, at least qualitatively, we can interpret the result of kunzel2019metalearners as a special case of our result in the case where the number of controls dominates the number of treated units, or vice-versa.
Our approach to heterogeneous treatment effect estimation via learning objectives can be implemented using any method that is framed as a loss minimization problem, such as boosting, decision trees, etc. In this section, we focus on simulation experiments using the $R$-learner, a direct implementation of (ref) based on lasso, kernel ridge regression, and boosting. We follow the nomenclature of kunzel2019metalearners and consider the following methods for heterogeneous treatment effect estimation as baselines. The $S$-learner fits a single model for \smash{$f(x, \, w) = \EE{Y \cond X = x, \, W = w}$}, and then estimates \smash{$\htau(x) = \hf(x, 1) - \hf(x, \, 0)$}; the $T$-learner fits the functions \smash{$\mu^*_{(w)}(x) = \EE{Y \cond X = x, \, W = w}$} separately for $w\in\{0,1\}$, and then estimates \smash{$\htau(x) = \hmu_{(1)}(x) - \hmu_{(0)}(x)$}; the $X$-learner and $U$-learner are as described in Section (ref). In addition, for the boosting-based experiments, we consider the causal boosting algorithm (denoted by $CB$ in Section (ref)) proposed by powers2018some.
Finally, for the lasso-based experiments, we consider an additional variant of our method, the $RS$-learner, that combines the spirit of $R$- and $S$-learners by adding an additional term in the loss function and then separately penalizes the main and treatment effect terms as in imai2013estimating. Specifically, we use \smash{$\htau(x) = x^\top \hdelta$}, where \smash{$\hb$} and \smash{$\hdelta$} minimize
Heuristically, one may hope that the $RS$-learner may be more robust, as it has an additional term to eliminate confounders.
In all simulations, we generate data as follows: for different choices of $X$-distribution $P_d$ indexed by dimension $d$, noise level $\sigma$, propensity function $e^*(\cdot)$, baseline main effect $b^*(\cdot)$, and treatment effect function $\tau^*(\cdot)$:
We consider the following specific simulation designs. Setup A has difficult nuisance components and an easy treatment effect function. We use the scaled friedman1991multivariate function for the baseline main effect \smash{$b^*(X_i) = \sin(\pi X_{i1} X_{i2})+ 2(X_{i3}-0.5)^2+ X_{i4}+0.5 X_{i5}$}, along with \smash{$X_i \sim \textrm{Unif}(0,1)^d$}, \smash{$e^*(X_i)=\textrm{trim}_{0.1} \{\sin(\pi X_{i1} X_{i2})\}$} and $\tau^*(X_i) = (X_{i1}+X_{i2})/2$, where $\textrm{trim}_\eta(x) = \max\{\eta, \, \min(x, 1-\eta)\}$. Setup B employs a randomized trial. Here, $e^*(x) = 1/2$ for all $x \in \RR^d$, so it is possible to be accurate without explicitly controlling for confounding. We take \smash{$X_i \sim \mathcal{N}(0,I_{d\times d})$}, \smash{$\tau^*(X_i)=X_{i1} + \log(1+e^{X_{i2}})$}, and \smash{$b^*(X_i) = \max\{X_{i1}+X_{i2}, X_{i3},0\} + \max\{X_{i4}+X_{i5},0\}$}. Setup C has an easy propensity score and a difficult baseline. In this setup, there is strong confounding, but the propensity score is much easier to estimate than the baseline: \smash{$X_i \sim \mathcal{N}(0,I_{d\times d})$}, \smash{$e^*(X_i) = {1}/(1+e^{X_{i2}+X_{i3}})$}, \smash{$b^*(X_i) = 2\log(1+e^{X_{i1}+X_{i2}+X_{i3}})$}, and the treatment effect is constant, $\tau^*(X_i)= 1$. Setup D has unrelated treatment and control arms, with data generated as \smash{$X_i \sim \mathcal{N}(0,I_{d\times d})$, $e^*(X_i) = 1/(1+e^{-X_{i1}}+e^{-X_{i2}})$}, $\smash{\tau^*(X_i)= \max\{X_{i1}+X_{i2}+X_{i3},0\}} - \smash{\max\{X_{i4}+X_{i5},0\}}$, and \smash{$b^*(X_i) = \p{\max\{X_{i1}+X_{i2}+X_{i3},0\} + \max\{X_{i4}+X_{i5},0\}}/2$}. Here, $\mu^*_{(0)}(X)$ and $\mu^*_{(1)}(X)$ are uncorrelated, and so there is no upside to learning them jointly.
In this section, we compare $S$-, $T$-, $X$-, $U$-, and our $R$- and $RS$-learners implemented via the lasso on simulated designs. For the $S$-learner, we follow imai2013estimating using (ref), while for the $T$-learner, we use (ref). For the $X$-, $R$-, and $RS$-learners, we use $L_1$-penalized logistic regression to estimate propensity $\hat{e}$, and the lasso for all other regression estimates.
For all estimators, we run the lasso on the pairwise interactions of a natural spline basis expansion with 7 degrees of freedom on $X_i$. We generate $n$ data points as the training set and generate a separate test set also with $n$ data points, and the reported mean-squared error is on the test set. The penalty parameter is chosen by 10-fold cross validation. For the $R$- and $RS$-learners, we use 10-fold cross-fitting on $\hat{e}$ and $\hat{m}$ in (ref). All methods are implemented via glmnet \citep*{friedman2010regularization}. We note that the $U$-learner suffers from high variance and instablity due to dividing by the propensity estimates. Therefore, we set a cutoff for the propensity estimate to be at level 0.05. We have also found empirically $U$-learner achieves much lower estimation error if we choose to use the largest regularization parameter that achieves 1 standard error away from the minimum in the cross validation step. Therefore, the $U$-learner uses lambda.1se as its cross validation parameter; other learners use lambda.min in glmnet.
In Figure (ref), we compare the performance of our 6 considered methods to an oracle that runs the lasso on (ref), for different values of sample size $n$, dimension $d$, and noise level $\sigma$. As is clear from these 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. The raw numbers depicted in Figure (ref) is available in Appendix (ref).
In Setups $A$ and $C$, where there is complicated confounding that needs to be overcome before we can estimate a simple treatment effect function $\tau^*(\cdot)$, the $R$- and $RS$-learners stand out. All methods do reasonably well in the randomized trial (Setup $B$) where it was not necessary to adjust for confounding, and the $X$-, $S$-, and $R$-learners do best. Finally, having completely disjoint functions for the treated and control arms is unusual in practice. However, we consider this possibility in Setup $D$, where there is no reason to model \smash{$\mu^*_{(0)}(x)$} and \smash{$\mu^*_{(1)}(x)$} jointly, and find that the $T$-learner---which in fact models them separately---performs well.
Overall, the $R$- and $RS$-learner consistently achieve good performance and, in most simulation specifications, essentially match the performance of the oracle (ref) in terms of the mean-squared error. The $U$-learner suffers from high loss due to its instability.
We move on to compare $S$-, $T$-, $X$-, $U$-, and $R$- learners implemented via kernel ridge regression with a Gaussian kernel. We use a variant of the KRLS package krls, available at \url{https://github.com/lukesonnet/KRLS}, that allows for weighted regression. For fitting the objective in each subroutine in all methods, we run a 5-fold cross validation to search through the width of the Gaussian kernel and the ridge regularization parameter both from a grid of $10^k$s with $k$ ranging in $[-3,\,-2.5,\, -2,\,\ldots,2, \, 2.5,\, 3]$. We experiment on the same set of setups and parameter variations, including variations on sample size $n$, dimension $d$, and noise level $\sigma$ as in Section (ref), and include all numbers depicted in Figure (ref) in Appendix (ref). In Figure (ref), we again observe that the $R$-learner does particularly well in Setups $A$ and $C$, where the treatment effect functions are relatively simple and the treatment propensity is not constant.
Finally, we compare $S$-, $T$-, $X$-, $U$-, and $R$- learners implemented via gradient boosting, as well as the causal boosting ($CB$) algorithm. We use the causalLearning R package for $CB$, while all other methods are implemented via XGboost chen2016xgboost. For fitting the objective in each subroutine in all methods, we draw a random set of 10 combinations of hyperparmaeters from the following grid: subsample$=[0.5,\, 0.75,\, 1]$, colsample_bytree$=[0.6,\, 0.8,\, 1]$, eta$=[5\text{e-}3,\, 1\text{e-}2,\, 1.5\text{e-}2,\, 2.5\text{e-}2,\, 5\text{e-}2,\, 8\text{e-}2,\, 1\text{e-}1,\, 2\text{e-}1]$, max_depth$=[3,\cdots, 20]$, \texttt{gamma}=\texttt{Uniform}$(0,\, 0.2)$, \texttt{min_child_weight}$=[1, \cdots, 20]$, \texttt{max_delta_step}$=[1, \cdots, 10]$, and cross validate on the number of boosted trees for each combination with an early stopping of 10 iterations. We experiment on the same set of setups and parameter variations including variations on sample size $n$, dimension $d$, and noise level $\sigma$ as in Section (ref), and include all numbers depicted in Figure (ref) in Appendix (ref).
In Figure (ref), we observe again that $R$-learner stands out in Setup $A$ and $C$, with all methods performing reasonably well in the randomized control setting of Setup $B$; in Setup $D$, the $T$-learner performs best since the the treated and control arms are generated from unrelated functions.
Before we conclude this section, we note that in both sets of the experiments, for simplicity of illustration, we have used lasso, kernel ridge regression, and boosting respectively to learn $\hat{m}(\cdot)$ and $\hat{e}(\cdot)$. In practice, we recommend cross validating on a variety of black-box learners, e.g. lasso, random forests, neural networks, etc. that are tuned for prediction accuracy to learn these two pilot quantities. All simulation results above can be replicated using the publicly available rlearner package for R CRAN, available at \url{https://github.com/xnie/rlearner}.
A natural generalization of our setup arises when, in some applications, we need to work with multiple treatment options. For example, in medicine, we may want to compare a control condition to multiple different experimental treatments. If there are $k$ different treatments along with a control arm, we can encode $W \in \cb{0, \, 1}^k$, and note that a multivariate version of Robinson's transformation suggests the following estimator,
where the angle brackets indicate an inner product, $e(x) = \EE{W \cond X = x} \in \RR^k$ is a vector, and $\tau_l(x)$ measures the conditional average treatment effect of the $l$-th treatment arm at $X_i = x$, for $l = 1, \, ..., \, k$. When implementing variants of this approach in practice, different choices of $\Lambda_n\{\tau(\cdot)\}$ may be needed to reflect relationships between the treatment effects of different arms, e.g. whether there is a natural ordering of treatment arms, or if there are some arms that we believe a priori to have similar effects.
It would also be interesting to consider extensions of the $R$-learner to cases where the treatment assignment $W_i$ is not unconfounded, and we need to rely on an instrument to identify causal effects. chernozhukov2017double discusses how Robinson's approach to the partially linear model generalizes naturally to this case, and \citet*{athey2018generalized} adapt their causal forest to work with instruments. The underlying estimating equations, however, cannot be interpreted as loss functions as easily as (ref), especially in the case where instruments may be weak, and so we leave this extension of the $R$-learner to future work.
We are grateful for enlightening conversations with Susan Athey, Emma Brunskill, John Duchi, Tatsunori Hashimoto, Guido Imbens, S\"oren K\"unzel, Percy Liang, Whitney Newey, Mark van der Laan, Alejandro Schuler, Robert Tibshirani, Bin Yu and Yuchen Zhang as well as for helpful comments and feedback from seminar participants at several universities and workshops and from the referees. This research was partially supported by a grant from the National Science Foundation, a Facebook Faculty Award, and a 2018 Stanford Human-Centered AI seed grant. The first author was awarded a Thomas R. Ten Have Award based on this work at the 2018 Atlantic Causal Inference Conference.
{0.2pt plus 0.3ex}