Orthogonal Random Forest for Causal Inference

Orthogonal Random Forest for Causal Inference

\twocolumn[ \icmltitle{Orthogonal Random Forest for Causal Inference}

\icmlaffiliation{msr}{Microsoft Research--New England} \icmlaffiliation{umn}{University of Minnesota--Twin Cities}

\icmlcorrespondingauthor{Miruna Oprescu}{[email removed] } \icmlcorrespondingauthor{Vasilis Syrgkanis}{[email removed]} \icmlcorrespondingauthor{Zhiwei Steven Wu}{[email removed]}

\vskip 0.3in ]

\printAffiliationsAndNotice

Introduction

Many problems that arise in causal inference can be formulated in the language of conditional moment models: given a target feature $x$ find a solution $\theta_0(x)$ to a system of conditional moment equations

given access to $n$ i.i.d. samples from the data generating distribution, where $\psi$ is a known score function and $h_0$ is an unknown nuisance function that also needs to be estimated from data. Examples include non-parametric regression, heterogeneous treatment effect estimation, instrumental variable regression, local maximum likelihood estimation and estimation of structural econometric models.\footnote{See e.g. Reiss2007 and examples in Chernozhukov2016locally,Chernozhukov2018plugin} The study of such conditional moment restriction problems has a long history in econometrics (see e.g. Newey1993,Ai2003,Chen2009,Chernozhukov2015).

\swedit{In this general estimation problem, the main} goal is to estimate the target parameter at a rate that is robust to the estimation error of the nuisance component. This allows \swedit{the use of} flexible models to fit the nuisance functions and enables asymptotically valid inference. Almost all prior work on the topic has focused on two settings: i) they either assume the target function $\theta_0(x)$ takes a parametric form and allow for a potentially high-dimensional parametric nuisance function, e.g. Chernozhukov2016locally,Chernozhukov2017,Chernozhukov2018plugin, ii) \swedit{or} take a non-parametric stance at estimating $\theta_0(x)$ but do not allow for high-dimensional \swedit{nuisance functions} WA,ATW17.

We propose Orthogonal Random Forest (ORF), a random forest-based estimation algorithm, which performs non-parametric estimation of the target parameter while permitting more complex nuisance functions with high-dimensional parameterizations. Our estimator is also asymptotically normal and hence allows for the construction of asymptotically valid confidence intervals via plug-in or bootstrap approaches. Our approach combines the notion of Neyman orthogonality of the moment equations with a two-stage random forest based algorithm, which generalizes prior work on Generalized Random Forests ATW17 and the double machine learning (double ML) approach proposed in Chernozhukov2017. \swedit{To support our general algorithm, we also provide a novel nuisance estimation algorithm---Forest Lasso---that effectively recovers high-dimensional nuisance parameters provided they have locally sparse structure.} This result combines techniques from Lasso theory Lasso-Book with concentration inequalities for $U$-statistics Hoeffding.

{As a concrete example and as a main application of our approach, we consider the problem of heterogeneous treatment effect estimation. This problem is at the heart of many decision-making processes, including clinical trial assignment to patients, price adjustments of products, and ad placement by a search engine. In many situations, we would like to take the heterogeneity of the population into account and estimate the heterogeneous treatment effect (HTE)---the effect of a treatment $T$ (e.g. drug treatment, price discount, and ad position), on the outcome $Y$ of interest (e.g. clinical response, demand, and click-through-rate), as a function of observable characteristics $x$ of the treated subject (e.g. individual patient, product, and ad). HTE estimation is a fundamental problem in causal inference from observational data imbens_rubin_2015,WA,ATW17, and is intimately related to many areas of machine learning, including contextual bandits, off-policy evaluation and optimization Swaminathan16,Wang17,NW17, and counterfactual prediction SwaminathanJ15,Hartford. }

{The key challenge in HTE estimation is that the observations are typically collected by a policy that depends on confounders or control variables $W$, which also directly influence the outcome. Performing a direct regression of the outcome $Y$ on the treatment $T$ and features $x$, without controlling for a multitude of other potential confounders, will produce biased estimation. This leads to a regression problem that in the language of conditional moments takes the form:

where $\theta_0(x)$ is the heterogeneous effect of the treatment $T$ (discrete or continuous) on the outcome $Y$ as a function of the features $x$ and $f_0(x,W)$ is an unknown nuisance function that captures the direct effect of the control variables on the outcome. Moreover, unlike active experimentation settings such as contextual bandits, when dealing with observational data, the actual treatment or logging policy $\mathbb{E}\left[ T | x, W \right] = g_0(x, W)$ that could potentially be used to de-bias the estimation of $\theta_0(x)$ is also unknown.

There is a surge of recent work at the interplay of machine learning and causal inference that studies efficient estimation and inference of treatment effects. doubleML propose a two-stage estimation method called double machine learning that first orthogonalizes out the effect of high-dimensional confounding factors using sophisticated machine learning algorithms, including Lasso, deep neural nets and random forests, and then estimates the effect of the lower dimensional treatment variables, by running a low-dimensional linear regression between the residualized treatments and residualized outcomes. They show that \swedit{even} if the estimation error of the first stage is \swedit{not particularly accurate}, the second-stage estimate can still be ${n^{-1/2}}$-asymptotically normal. However, their approach requires a parametric specification of $\theta_0(x)$. In contrast, another line of work that brings machine learning to causal inference provides fully flexible non-parametric HTE estimation based on random forest techniques WA,ATW17,Powers2017. However, these methods heavily rely on low-dimensional assumptions.

Our algorithm ORF, when applied to the HTE problem (see Section (ref)) allows for the non-parametric estimation of $\theta_0(x)$ via forest based approaches while simultaneously allowing for a high-dimensional set of control variables $W$. This estimation problem is of practical importance when a decision maker (DM) wants to optimize a policy that depends only on a small set of variables, e.g. due to data collection or regulatory constraints or due to interpretability of the resulting policy, while at the same time controlling for many potential confounders in the existing data that could lead to biased estimates. Such settings naturally arise in contextual pricing or personalized medicine. In such settings the DM is faced with the problem of estimating a conditional average treatment effect conditional on a small set of variables while controlling for a much larger set. Our estimator provably offers a significant statistical advantage for this task over prior approaches.

In the HTE setting, the ORF algorithm follows the residual-on-residual regression approach analyzed by Chernozhukov2016locally to formulate a locally Neyman orthogonal moment and then applies our orthogonal forest algorithm to this orthogonal moment. Notably, ATW17 also recommend such a residual on residual regression approach in their empirical evaluation, which they refer to as “local centering”, albeit with no theoretical analysis. Our results provide a theoretical foundation of the local centering approach through the lens of Neyman orthogonality. Moreover, our theoretical results give rise to a slightly different overall estimation approach than the one in ATW17: namely we residualize locally around the target estimation point $x$, as opposed to performing an overall residualization step and then calling the Generalized Random Forest algorithm on the residuals. The latter stems from the fact that our results require that the nuisance estimator achieve a good estimation rate only around the target point $x$. Hence, residualizing locally seems more appropriate than running a global nuisance estimation, which would typically minimize a non-local mean squared error. Our experimental findings reinforce this intuition (see e.g. comparison between ORF and the GRF-Res benchmark). Another notable work that combines the residualization idea with flexible heterogeneous effect estimation is that of NW17, who formulate the problem as an appropriate residual-based square loss minimization over an arbitrary hypothesis space for the heterogeneous effect function $\theta(x)$. Formally, they show robustness, with respect to nuisance estimation errors, of the mean squared error (MSE) of the resulting estimate in expectation over the distribution $X$ and for the case where the hypothesis space is a reproducing kernel Hilbert space (RKHS). Our work differs primarily by: i) focusing on sup-norm estimation error at any target point $x$ as opposed to MSE, ii) using forest based estimation as opposed to finding a function in an RKHS, iii) working with the general orthogonal conditional moment problems, and iv) providing asymptotic normality results and hence valid inference.

We provide a comprehensive empirical comparison of ORF with several benchmarks, including three variants of GRF. We show that by setting the parameters according to what our theory suggests, ORF consistently outperforms all of the benchmarks. Moreover, we show that bootstrap based confidence intervals provide good finite sample coverage.

Finally, to motivate the usage of the ORF, we applied our technique to Dominick's dataset, a popular historical dataset of store-level orange juice prices and sales provided by University of Chicago Booth School of Business. The dataset is comprised of a large number of covariates $W$, but economics researchers might only be interested in learning the elasticity of demand as a function of a few variables $x$ such as income or education. We applied our method (see Appendix (ref) for details) to estimate orange juice price elasticity as a function of income, and our results, depicted in Figure (ref), unveil the natural phenomenon that lower income consumers are more price-sensitive.

\vscomment{

Related Work

we need to add an extensive comparison to many relevant literatures.}

Estimation via Local Orthogonal Moments

\vsedit{ We study non-parametric estimation of models defined via conditional moment restrictions, in the presence of nuisance functions. Suppose we have a set of $2n$ observations $Z_1, \ldots , Z_{2n}$ drawn independently from some underlying distribution $\mathcal{D}$ over the observation domain $\mathcal{Z}$. Each observation $Z_i$ contains a feature vector $X_i\in \mathcal{X} \mathrel{\vcenter{\baselineskip0.5ex \lineskiplimit0pt \hbox{\scriptsize.}\hbox{\scriptsize.}}} = [0, 1]^d$. }

\vsedit{ Given a target feature $x\in \mathcal{X}$, our goal is to estimate a parameter vector $\theta_0(x)\in \mathbb{R}^p$ that is defined via a local moment condition, i.e. for all $x\in \mathcal{X}$, $\theta_0(x)$ is the unique solution with respect to $\theta$ of:

where $\psi\colon \mathcal{Z}\times \mathbb{R}^p \times \mathbb{R}^\ell \rightarrow \mathbb{R}^p$ is a score function that maps an observation $Z$, parameter vector $\theta(x)\in \Theta \subset \mathbb{R}^p$, and nuisance vector $h(x, w)$ to a vector-valued score $\psi(z; \theta(x), h(x, w))$ and $h_0 \in H \subseteq \left( \mathbb{R}^d\times \mathbb{R}^L\rightarrow \mathbb{R}^\ell \right)$ is an unknown nuisance function that takes as input $X$ and a subvector $W$ of $Z$, and outputs a nuisance vector in $\mathbb{R}^\ell$. For any feature $x \in \mathcal{X}$, parameter $\theta\in \Theta$, and nuisance function $h \in H$, we define the moment function as:

We assume that the dimensions $p, \ell, d$ are constants, while the dimension $L$ of $W$ can be growing with $n$. }

\vsedit{ We will analyze the following two-stage estimation process.

}

\vsedit{ In practice, our framework permits the use of any method to estimate the nuisance function in the first stage. However, since our description is a bit too abstract let us give a special case, which we will also need to assume for our normality result. Consider the case when the nuisance function $h$ takes the form $h(x, w) = g(w; \nu(x))$, for some known function $g$ but unknown function $\nu: \mathcal{X}\rightarrow \mathbb{R}^{d_{\nu}}$ (with $d_{\nu}$ potentially growing with $n$), i.e. locally around each $x$ the function $h$ is a parametric function of $w$. Moreover, the parameter $\nu_0(x)$ of the true nuisance function $h_0$ is identified as the minimizer of a local loss:

Then we can estimate $\nu_0(x)$ via a locally weighted and penalized empirical loss minimization algorithm. In particular in Section (ref) we will consider the case of local $\ell_1$-penalized estimation that we will refer to as forest lasso and which provides formal guarantees in the case where $\nu_0(x)$ is sparse. }

The key technical condition that allows us to reliably perform the two-stage estimation is the following local orthogonality condition, which can be viewed as a localized version of the Neyman orthogonality condtion Neyman,doubleML around the neighborhood of the target feature $x$. Intuitively, the condition says that the score function $\psi$ is insensitive to local perturbations in the nuisance parameters around their true values.

Orthogonal Random Forest

\vsedit{We describe our main algorithm orthogonal random forest (ORF) for calculating the similarity weights in the second stage of the two stage estimation. In the next section we will see that we will be using this algorithm for the estimation of the nuisance functions, so as to perform a local nuisance estimation.} At a high level, ORF can be viewed as an orthogonalized version of GRF that is more robust to the nuisance estimation error. Similar to GRF, the algorithm runs a tree learner over $B$ random subsamples $S_b$ (without replacement) of size $s< n$, to build $B$ trees such that each tree indexed by $b$ provides a tree-based weight $a_{ib}$ for each observation $Z_i$ in the input sample. \vsedit{Then the ORF weight $a_i$ for each sample $i$ is the average over the tree-weights $a_{ib}$.}

The tree learner starts with a root node that contains the entire $\mathcal{X}$ and recursively grows the tree to split $\mathcal{X}$ into a set of leaves until the number of observations in each leaf is not too small. The set of neighboods defined by the leaves naturally gives a simlarity measure between each observation and the target $x$. Following the same approach of GRF,WA, we maintain the following tree properties in the process of building a tree.

The key modification to GRF's tree learner is our incorporation of orthogonal nuisance estimation in the splitting criterion. \vsdelete{In order to accurately capture the heterogeneity in the target parameter $\theta$ in the presence of high-dimensional nuisance, we perform a two-stage estimation locally at each internal node to decide the next split. } While the splitting criterion does not factor into our theoretical analysis (similar to GRF), we find it to be an effective practical heuristic.

\paragraph{Splitting criterion with orthogonalization.} \vsedit{At each internal node $P$ we perform a two-stage estimation over $(P\cap S^1)$, i.e. the set of examples in $S^1$ that reach node $P$: 1) compute a nuisance estimate $\hat h_P$ using only data $P\cap S^1$ (e.g. by estimating a parameter $\hat \nu_P$ that minimizes $\sum_{i\in (P\cap S^1)} \ell(Z_i; \nu) + \lambda \|\nu\|_1$ and setting $\hat{h}_P(\cdot) = g(\cdot; \hat{\nu}_P)$), and then 2) form estimate $\hat \theta_P$ using $\hat h_P$:\footnote{In our implementation we actually use a cross-fitting approach, where we use half of $P\cap S^1$ to compute a nuisance function to apply to the other half and vice versa.} \[ \textstyle{\hat \theta_P \in\operatorname*{\mathrm{argmin}}_{\theta\in \Theta} \left\|\sum_{i\in (P\cap S^1)} \psi(Z_i; \theta, \hat h_P(W_i)) \right\|} \] We now generate a large random set of candidate axis-aligned splits (satisfying Specification (ref) and we want to find the split into two children $C_1$ and $C_2$ such that if we perform the same two-stage estimation separately at each child, the new estimates $\hat \theta_{C_1}$ and $\hat \theta_{C_2}$ take on very different values, so that the heterogeneity of the two children nodes is maximized. Performing the two-stage estimation of $\hat \theta_{C_1}$ and $\hat \theta_{C_2}$ for all candidate splits is too computationally expensive. Instead, we will approximate these estimates by taking a Newton step from the parent node estimate $\hat \theta_P$: for any child node $C$ given by a candidate split, our proxy estimate is: \[ \textstyle{\tilde \theta_{C} = \hat \theta_P - \frac{1}{|C\cap S^1|}\sum_{i\in C_j\cap S^1}} A_P^{-1} \psi(Z_i; \hat \theta_P, \hat h_P(X_i, W_i))} \] where $A_P = \frac{1}{|P\cap S^1|} \sum_{i\in P\cap S_b^1} \nabla_\theta \psi(Z_i; \hat \theta_P, \hat h_P(X_i, W_i))$.} \vsedit{We select the candidate split that maximizes the following proxy heterogeneity score: for each coordinate $t\in [p]$ let

where $\rho_{t,i} = A_P^{-1} \psi_t(Z_i; \hat \theta_P, \hat h_P(X_i, W_i))$.

We then create a single heterogeneity score per split as a convex combination that puts weight $\eta$ on the mean and $(1-\eta)$ on the maximum score across coordinates. $\eta$ is chosen uniformly at random in $[0,1]$ at each iteration of splitting. Hence, some splits focus on heterogeneity on average, while others focus on creating heterogeneity on individual coordinates. }

\paragraph{ORF weights and estimator.} \vsedit{For each tree indexed $b\in [B]$ based on subsample $S_b$, let $L_b(x)\subseteq \mathcal{X}$ be the leaf that contains the target feature $x$. We assign tree weight and ORF weight to each observation $i$:

}

WA show that under the structural specification of the trees, the tree weights are non-zero only around a small neighborhood of $x$; a property that we will leverage in our analysis.

Convergence and Asymptotic Analysis

\vsedit{The ORF estimate $\hat{\theta}$ is computed by solving the weighted moment condition in (ref), using the ORF weights as described in the previous section. We now provide theoretical guarantees for $\hat{\theta}$ under the following assumption on the moment, score fuction and the data generating process.}

All the results presented in the remainder of the paper will assume these conditions and we omit stating so in each of the theorems. Any extra conditions required for each theorem will be explicitly provided. Note that except for the local orthogonality condition, all of the assumptions are imposing standard boundedness and regularity conditions of the moments.

For asymptotic normality we will restrict our framework to the case of parametric nuisance functions, i.e. $h(X, W) = g(W; \nu(X))$ for some known function $g$ and to a particular type of nuisance estimators that recover the true parameter $\nu_0(x)$. Albeit we note that the parameter $\nu(X)$ can be an arbitrary non-parametric function of $X$ and can also be high-dimensional. We will further assume that the moments also have a smooth co-variance structure in $X$, i.e. if we let $$V=\psi(Z; \theta_0(x), g(W; \nu_0(x)))$$ then $\text{Var}(V \mid X=x')$ is Lipschitz in $x'$ for any $x'\in [0,1]^d$.

Given the result in (ref), we can follow the same approach of Bootstrap of Little Bags by ATW17,SEXTON2009801 to build valid confidence intervals.

Nuisance Estimation: Forest Lasso

Next, we study the nuisance estimation problem in the first stage and provide a general nuisance estimation method that leverages locally sparse parameterization of the nuisance function, permitting low error rates even for high-dimensional problems. \vsedit{Consider the case when the nuisance function $h$ takes the form $h(x, w) = g(w; \nu(x))$ for some known functional form $g$, for some known function $g$ but unknown function $\nu: \mathcal{X}\rightarrow \mathbb{R}^{d_{\nu}}$, with $d_{\nu}$ potentially growing with $n$. Moreover, the parameter $\nu_0(x)$ of the true nuisance function $h_0$ is identified as the minimizer of some local loss, as defined in Equation (ref).}

\vsedit{We consider the following estimation process: given a set of observations $D_1$, we run the same tree learner in Section (ref) over $B$ random subsamples (without replacement) to compute ORF weights $a_i$ for each observation $i$ over $D_1$. Then we apply a local $\ell_1$ penalized $M$-estimation:

}

\vsedit{To provide formal guarantees for this method we will need to make the following assumptions.}

\vsedit{Under Assumption (ref) we show that the local penalized estimator achieves the following parameter recovery guarantee.}

Heterogeneous Treatment Effects

Now we apply ORF to the problem of estimating heterogeneous treatment effects. We will consider the following extension of the partially linear regression (PLR) model due to Robinson. \footnote{The standard PLR model Robinson considers solely the case of constant treatment effects, $Y = \left\langle \theta_0, T \right\rangle + f_0(X, W) + \varepsilon$, and the goal is the estimation of the parameter $\theta_0$.} We have $2n$ i.i.d. observations $D = \{Z_i = (T_i, Y_i, W_i, X_i)\}_{i=1}^{2n}$ such that for each $i$, $T_i$ represents the treatment applied that can be either real-valued (in $\mathbb{R}^p$) or discrete (taking values in $\{0, e_1, \ldots , e_p\}$, where each $e_j$ denotes the standard basis in $\mathbb{R}^p$), $Y_i\in \mathbb{R}$ represents the outcome, $W_i\in [-1, 1]^{d_\nu}$ represents potential confounding variables (controls), and $X_i\in \mathcal{X} = [0, 1]^d$ is the feature vector that captures the heterogeneity. The set of parameters are related via the following equations:

where $\eta, \varepsilon$ are bounded unobserved noises such that $\operatorname*{\mathbb{E}}_{}\left[\varepsilon \mid W, X, T\right] = 0$ and $\operatorname*{\mathbb{E}}_{}\left[\eta \mid X, W, \varepsilon\right] = 0$.\swcomment{added the last $\varepsilon$} In the main equation (ref), $\mu_0\colon \mathbb{R}^d\times \mathbb{R}^{d_\nu}\rightarrow [-1, 1]^p$ represents the treatment effect function. Our goal is to estimate conditional average treatment effect (CATE) $\theta_0(x)$ conditioned on target feature $x$:

The confounding equation (ref) determines the relationship between treatments variable $T$ and the feature $X$ and confounder $W$. To create an orthogonal moment for identifying $\theta_0(x)$, we follow the classical residualization approach similar to doubleML. First, observe that

Let us define the function $q_0(X, W) = \operatorname*{\mathbb{E}}_{}\left[Y\mid X, W\right]$, and consider the residuals

Then we can simplify the equation as $\tilde Y = \mu_0(X, W) \cdot \tilde T + \varepsilon$. \swedit{As long as $\eta$ is independent of $\mu_0(X, W)$ conditioned on $X$ (e.g. $\eta$ is independent of $W$ or $\mu_0(X, W)$ does not depend on $W$), } we also have $$\operatorname*{\mathbb{E}}_{}\left[\mu_0(X, W) \mid X, \eta\right] = \operatorname*{\mathbb{E}}_{}\left[\mu_0(X, W)\mid X\right] = \theta(X).$$ Since $\operatorname*{\mathbb{E}}_{}\left[\varepsilon \mid X, \eta\right] = \mathbb{E}\left[ \mathbb{E}\left[ \varepsilon\mid X, W, T \right]\mid X, \eta \right]d X, \eta}=0$, then $$\operatorname*{\mathbb{E}}_{}\left[\tilde Y \mid X, \tilde T\right] = \operatorname*{\mathbb{E}}_{}\left[\mu_0(X, W) \mid X\right]\cdot \tilde T = \theta(X) \cdot \tilde T.$$ This relationship suggests that we can obtain an estimate of $\theta(x)$ by regressing $\tilde Y$ on $\tilde T$ locally around $X=x$. We can thus define the orthogonalized score function: for any observation $Z = (T, Y, W, x)$, any parameter $\theta\in \mathbb{R}^p$, any estimates $q$ and $g$ for functions $q_0$ and $g_0$, the score $\psi(Z; \theta, h(X, W))$ is:

where $h(X, W) = (q(X, W), g(X, W))$. In the appendix, we show that this moment condition satisfies local orthogonality, and it identifies $\theta_0(x)$ as long as as the noise $\eta$ is independent of $\mu_0(X,W)$ conditioned on $X$ and the expected matrix $\operatorname*{\mathbb{E}}_{}\left[\eta \eta^\intercal \mid X=x\right]$ is invertible. Even though the approach applies generically, to obtain formal guarantees on the nuisance estimates via our Forest Lasso, we will restrict their functional form.

\paragraph{Real-valued treatments.} Suppose $f_0$ and each coordinate $j$ of $g_0$ and $\mu_0$ are given by high-dimensional linear functions: $f_0(X, W) = \langle W, \beta_0(X) \rangle$, $\mu_0^j(X, W) = \langle W, u_0^j(X) \rangle, \quad g_0^j(X, W) = \langle W, \gamma_0^j(X)\rangle$, where $\beta_0(X), \gamma_0^j(X), u_0^j(X)$ are $k$-sparse vectors in $\mathbb{R}^{d_\nu}$. Consequently, $q_0(X, W)$ can be written as a $k^2$-sparse linear function over degree-2 polynomial features $\phi_2(W)$ of $W$. Then as long as $\gamma_0, \beta_0$ and $\mu_0$ are Lipschitz in $X$ and the confounders $W$ satisfy $$\operatorname*{\mathbb{E}}_{}\left[\phi_2(W) \phi_2(W)^\intercal\mid X\right] \succeq \Omega(1) I,$$ then we can use Forest Lasso to estimate both $g_0(x, w)$ and $q_0(x,w)$. Hence, we can apply the ORF algorithm to get estimation error rates and asymptotic normality results for $\hat{\theta}$. (see Appendix (ref) for formal statement).

\paragraph{Discrete treatments.} We now describe how our theory can be applied to discrete treatments. Suppose $f_0$ and each coordinate $j$ of $g_0$ are of the form: $f_0(X, W) = \langle W, \beta_0(X) \rangle$ and $g_0^j(X, W) = \mathcal{L}(\langle W, \gamma_0^j(X)\rangle)$, where $\mathcal{L}(t) = 1/(1 + e^{-t})$ is the logistic function. Note in this case $\eta$ is not independent of $W$ since $$\text{Var}(\eta_j) = g_0^j(X,W)(1 - g_0^j(X, W)).$$ To maintain the conditional independence between $\mu_0(X, W)$ and $\eta$ conditioned on $X$, we focus on the setting where $\mu_0$ is only a function of $X$, i.e. $\mu(X, W) = \theta(X)$ for all $W, X$. In this setting we can estimate $g_0$ by running a forest logistic lasso for each treatment $j$. Then we can estimate $q_0(x, W)$ as follows: For each $t\in \{e_1, \ldots, e_p\}$ estimate the expected counter-factual outcome function: $m_0^t(x, W) = \mu_0^t(x, W) + f_0(x, W)$, by running a forest lasso between $Y$ and $X, W$ only among the subset of samples that received treatment $t$. Similarly, estimate $f_0(x, W)$ by running a forest lasso between $Y$ and $X, W$ only among the subset of samples that received treatment $t=0$.

Then observe that $q_0(x, W)$ can be written as a function of $f_0$, $g_0^t$ and $m_0^t$. Thus we can combine these estimates to get an estimate of $q_0$. Hence, we can obtain a guarantee similar to that of Corollary (ref) (see appendix).

\paragraph{Doubly robust moment for discrete treatments.} In the setting where $\mu$ also depends on $W$ and treatments are discrete, we can formulate an alternative orthogonal moment that identifies the CATE even when $\eta$ is correlated with $\mu(X, W)$. This moment is based on first constructing unbiased estimates of the counterfactual outcome $m_0^t(X, W) = \mu_0^t(X, W) + f_0(X, W)$ for every observation $X, W$ and for any potential treatment $t$, i.e. even for $t\neq T$. The latter is done by invoking the doubly robust formula robins1995semiparametric,cassel1976some,kang2007demystifying:

with the convention that

Then we can identify the parameter $\theta^t(x)$ using the moment: $\textstyle{\mathbb{E}[Y^{(t)}-Y^{(0)} | X=x] = \theta_t(x).}$

One can easily show that this moment satisfies the Neyman orthogonality condition with respect to the nuisance functions $m$ and $g$ (see appendix). In fact this property is essentially implied by the fact that the estimates $Y^{(t)}$ satisfy the double robustness property, since double robustness is a stronger condition than orthogonality. We will again consider $\mu_0^j(X, W) = \langle W, u_0^j(X) \rangle$. Then using similar reasoning as in the previous paragraph, we see that with a combination of forest logistic lasso for $g_0^t$ and forest lasso for $m_0^t$, we can estimate these nuisance functions at a sufficiently fast rate for our ORF estimator (based on this doubly robust moment) to be asymptotically normal, assuming they have locally sparse linear or logistic parameterizations.

Monte Carlo Experiments

We compare the empirical performance of ORF with other methods in the literature (and their variants).\footnote{The source code for running these experiments is available in the git repo \href{https://github.com/Microsoft/EconML/tree/master/prototypes/orthogonal_forests}{Microsoft/EconML}.}

The data generating process we consider is described by the following equations:

Moreover, $x_i$ is drawn from the uniform distribution $U[0, 1]$, $W_i$ is drawn from $\mathcal{N}(0, I_p)$, and the noise terms $\varepsilon_i \sim U[-1, 1], \eta_i \sim U[-1, 1]$. The $k$-sparse vectors $\beta_0, \gamma_0 \in \mathbb{R}^{p}$ have coefficients drawn independently from $U[0, 1]$. The dimension $p = 500$ and we vary the support size $k$ over the range of $\{1, 5, 10, 15, 20, 25, 30\}$. We examine a treatment function $\theta(x)$ that is continuous and piecewise linear (detailed in (ref)). In Appendix (ref) we analyze other forms for $\theta(x)$. \vsdelete{we show that ORF also outperforms other methods with discontinuous $\theta(x)$ that is piecewise constant or piecewise polynomial.}

For each fixed treatment function, we repeat $100$ experiments, each of which consists of generating $5000$ observations from the DGP, drawing the vectors $\beta_0$ and $\gamma_0$, and estimating $\hat{\theta}(x)$ at 100 test points $x$ over a grid in $[0, 1]$. We then calculate the bias, variance and root mean squared error (RMSE) of each estimate $\hat{\theta}(x)$. \iffull Here we report summary statistics of the median and $5-95$ percentiles of these three quantities across test points, so as to evaluate the average performance of each method. \fi We compare two variants of ORF with two variants of GRF ATW17 (see Appendix (ref) for a third variant) and two extensions of double ML methods for heterogeneous treatment effect estimation Chernozhukov2017.

ORF variants. (1) ORF: We implement ORF as described in (ref), setting parameters under the guidance of our theoretical result: subsample size $s \approx (n/\log(p))^{1/(2\tau+1)}$, Lasso regularization $\lambda_\gamma, \lambda_q \approx \sqrt{\log(p) s / n}/20$ (for both tree learner and kernel estimation), number of trees $B=100 \geq n/s$, a max tree depth of $20$, and a minimum leaf size of $r=5$. (2) ORF with LassoCV (ORF-CV): we replaced the Lasso algorithm in ORF's kernel estimation, with a cross-validated Lasso for the selection of the regularization parameter $\lambda_\gamma$ and $\lambda_q$. ORF-CV provides a more systematic optimization over the parameters.

GRF variants. (1) GRF-Res-Lasso: We perform a naive combination of double ML and GRF by first residualizing the treatments and outcomes on both the features $x$ and controls $W$, then running GRF R package by GRF on the residualized treatments $\hat T$, residualized outcomes $\hat Y$, and features $x$. A cross-validated Lasso is used for residualization. (2) GRF-Res-RF: We combine DoubleML and GRF as above, but we use cross-validated Random Forests for calculating residuals $\hat T$ and $\hat Y$.

Double ML with Polynomial Heterogeneity (DML-Poly). An extension of the classic Double ML procedure for heterogeneous treatment effects introduced in Chernozhukov2017. This method accounts for heterogeneity by creating an expanded linear base of composite treatments (cross products between treatments and features). (1) Heterogeneous Double ML using LassoCV for first-stage estimation (HeteroDML-Lasso): In this version, we use Lasso with cross-validation for calculating residuals on $x\cup W$ in the first stage. (2) Heterogeneous Double ML using random forest for first-stage estimation (HeteroDML-RF): A more flexible version that uses random forests to perform residualization on treatments and outcomes. The latter performs better when treatments and outcomes have a non-linear relationship with the joint features of $(x, W)$.

• Generalized Random Forest variants:

• Double ML with Polynomial Heterogeneity (DML-Poly): An extension of the classic Double ML procedure for heterogeneous treatment effects introduced in Chernozhukov2017. This method accounts for heterogeneity by creating an expanded linear base of composite treatments (cross products between treatment and features).

\end{enumerate} \fi

We generated data according to the Monte Carlo process above and set the parameters to $n=5000$ samples, $p=500$ controls, $d=1$ features and support size $k \in \{1, 5, 10, 15, 20, 25, 30\}$ and three types of treatment effect functions. In this section, we present the results for a piecewise linear treatment effect function.

In Figure (ref), we inspect the goodness of fit for the chosen estimation methods across 100 Monte Carlo experiments. We note the limitations of two versions of the GRF-Res estimators, GRF-Res-Lasso and GRF-Res-RF, in capturing the treatment effect function well. The GRF-Res-RF estimations have a consistent bias as the Random Forest residualization cannot capture the dependency on the controls $W$ given their high-dimensionality. The HeteroDML methods are not flexible enough to capture the complexity of the treatment effect function. The best performers are the ORF-CV, ORF, and GRF-Res-Lasso, with the latter estimator having a larger bias and variance.

We analyze these estimators as we increase the support size of $W$. Figures (ref) illustrate the evaluation metrics across different support sizes. The ORF-CV performs very well, with consistent bias and RMSE across support sizes and treatment functions. The bias, variance and RMSE of the ORF grow with support size, but this growth is at a lower rate compared to the alternative estimators. The ORF-CV and ORF algorithms perform better than the GRF-Res methods on all metrics for this example. We observe this pattern for the other choices of support size, sample size and treatment effect function (see Appendix (ref)). In Figure (ref), we provide a snapshot of the bootstrap confidence interval coverage for this example.

Acknowledgements

ZSW is supported in part by a Google Faculty Research Award, a J.P. Morgan Faculty Award, a Mozilla research grant, and a Facebook Research Award. Part of this work was completed while ZSW was at Microsoft Research-New York City.