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.
78,397 characters · 19 sections · 11 citation commands
Local Linear Forests
\def\spacingset#1{ {#1}} \spacingset{1}
{\it Keywords:} asymptotic normality; heterogeneous treatment effect; smoothing and nonparametric regression
\spacingset{1.5}
Random forests \citep*{breiman2001} are a popular method for non-parametric regression that have proven effective across many application areas cutler2007random,diaz2006gene,svetnik2003random. A major weakness of random forests, however, is their inability to exploit smoothness in the regression surface they are estimating. As an example, consider the following setup: We simulate $X_1, \dots, X_n$ independently from the uniform distribution on $[0,1]^{20}$, with responses
and our goal is to estimate $\mu({x_0}) = \mathbb{E}[Y \,|\, X = {x_0}]$. The left panel of Figure (ref) shows a set of predictions on this data from a random forest. The forest is unable to exploit strong local trends and, as a result, fits the target function using qualitatively the wrong shape: The prediction surface resembles a step function as opposed to a smooth curve.
In order to address this weakness, we take the perspective of random forests as an adaptive kernel method. This interpretation follows work by \citet*{athey2016}, \citet*{hothorn}, and \citet*{meinshausen2006quantile}, and complements the traditional view of forests as an ensemble method (i.e., an average of predictions made by individual trees). Specifically, random forest predictions can be written as
where the weights $\alpha_i({x_0})$, which are defined in the upcoming display (ref), encode the weight given by the forest to the $i$-th training example when predicting at ${x_0}$. Now, as is well-known in the literature on non-parametric regression, if we want to fit smooth signals without some form of neighborhood averaging (e.g., kernel regression, $k$-NN, or matching for causal inference), it is helpful to use a local regression adjustment to correct for potential misalignment between a test point and its neighborhood \citep*{abadie2011bias,cleveland1988locally,GVK19282144X,heckman1998matching, Load:1999,10.2307/3532868,Stone1977,tibshirani1987local}. These types of adjustments are particularly important near boundaries, where neighborhoods are asymmetric by necessity. With many covariates, these adjustments are also important away from boundaries given that local neighborhoods are often unbalanced due to sampling variation.
The goal of this paper is to improve the accuracy of forests on smooth signals using regression adjustments, potentially in many dimensions. By using the local regression adjustment, it is possible to adjust for asymmetries and imbalances in the set of nearby points used for prediction, ensuring that the weighted average of the feature vector of neighboring points is approximately equal to the target feature vector, and that predictions are centered. The improvement to forests from the regression adjustment is most likely to be large in cases where some features have strong effects with moderate curvature, so that regression adjustments are both effective and important.
In their simplest form, local linear forests take the forest weights $\alpha_i({x_0})$, and use them for local regression:
Here \smash{$\hat{\mu}({x_0})$} estimates the conditional mean function $\mu({x_0})$, and \smash{$\theta({x_0})$} corrects for the local trend in $X_i- x$. The ridge penalty $\lambda ||\theta({x_0})||_2^2$ prevents overfitting to the local trend, and plays a key role both in simulation experiments and asymptotic convergence results. Then, as discussed in Section (ref), we can improve the performance of local linear forests by modifying the tree-splitting procedure used to get the weights $\alpha_i({x_0})$, and making it account for the fact that we will use local regression to estimate $\mu({x_0})$. As a first encouraging result, in the motivating example from Figure (ref), local linear forests have improved upon the fit of standard forests.
These improvements extend to many other types of forests, such as quantile regression forests \citep*{meinshausen2006quantile} or, more broadly, generalized random forests \citep*{athey2016}. An extension of primary interest is to causal forests as proposed by \citet*{athey2016}, which we discuss in Section (ref) in detail; other cases are analogous.
Our main formal result is a Central Limit Theorem for the predictions $\hat{\mu}({x_0})$ from a local linear forest at a given test point $x$, specifying the asymptotic convergence rate and its dependence on subsampling and smoothness of $\mu({x_0})$. This allows us to build pointwise Gaussian confidence intervals, giving practitioners applicable uncertainty quantification. Observe that in Figure (ref), the bias of regression forest predictions affects not only the accuracy prediction curve but also the coverage corresponding confidence intervals, which are not centered on the true function. Local linear forests, in addressing this issue, improve over regression forests in both predictive performance and confidence interval coverage. Strikingly, our local linear forest confidence intervals simultaneously achieve better coverage and are shorter than those built using regression forests.
A simple form of (ref), without regularization or modified tree-splitting procedures, was also considered in a recent paper by \citet*{bloniarz2016supervised}. However, they only report modest performance improvements over basic regression forests; for example, on the “Friedman function” they report roughly a 5% reduction in mean-squared error. In contrast, we find fairly large, systematic improvements from local linear forests; see, e.g., Figure (ref) for corresponding results on the same Friedman function. It thus appears that our algorithmic modifications via regularization and optimized splitting play a qualitatively important role in getting local linear forests to work well. These empirical findings are also mirrored in our theory. For example, in order to prove rates of convergence for local linear forests that can exploit smoothness of $\mu(\cdot)$ and improve over corresponding rates available for regression forests, we need an appropriate amount of regularization in (ref).
Finally, one can also motivate local linear forests from the starting point of local linear regression. Despite working well in low dimensions, classical approaches to local linear regression are not applicable to even moderately high-dimensional problems. (This is a well-known problem. The popular core R function loess \citep*{r-core} allows only 1-4 predictors, while locfit \citep*{locfit} crashes on the simulation from (ref) with $d\ge 7$.) In contrast, random forests are adept at fitting high-dimensional signals, both in terms of their stability and computational efficiency. From this perspective, random forests can be seen as an effective way of producing weights to use in local linear regression. In other words, local linear forests aim to combine the adaptivity of random forests and the ability of local linear regression to capture smoothness.
An implementation of local linear forests, compliant with the assumptions detailed in Section (ref), is available in the R package grf package-grf,r-core.
To illustrate the promise of local linear forests, we consider the problem of predicting the logarithm of wages as a function of covariates including years of education, age, race, and gender; this function plays an important role in the study of labor markets \citep*{RePEc:nbr:nberwo:9732}. This problem has a mix of continuous and categorical variables, to which tree-based methods are well suited. However, wages tend to have a fairly strong and smooth association with age and education, and we might expect local regression adjustments to help with this. In addition, the covariate space is large relative to the size of publicly available administrative data, and there are moderate to strong correlations among many of the covariates, making it challenging to obtain accurate predictions in some regions of the covariate space.
We consider data from the current population survey (CPS), available from the Minnesota Population Center cps. These data describe the wages of 114,291 individuals in 2018 (excluding records that do not contain wage data). To evaluate how model performance varies with sample size, we divide the data into a large test set, which is used to evaluate accuracy overall and in specific regions of the covariate space, and training sets of varying sizes.
We compare local linear forests with ordinary least squares, lasso with interaction terms, gradient boosting, Bayesian additive regression trees, and random forest. For the lasso, random forests, local linear forests and boosting, we chose tuning parameters via cross-validation; in particular, for local linear forests, we tuned on leaf size and $\lambda$. Moreover, for our method, we did use a local linear correction for all variables; rather, we only used non-zero $\theta$-coefficient in (ref) for continuous predictors that had non-zero coefficients in a pilot lasso regression Tibshirani94regressionshrinkage (in general, we have found screening of variables used for a local linear correction in (ref) to benefit both the accuracy and computational performance of our approach). Table (ref) compares predictive performance across several methods, showing that local linear forests can provide a predictive benefit over competing methods.
One motivation for studying wages is to compare wages across different types of workers, which requires accurate predictions even for types of workers who are less frequently observed. We thus evaluate predictive performance in several sparse regions of the covariate space, showing in Table (ref) that local linear forests fit well in these regions. To further explore this idea, Figure (ref) shows plots of observed log wages by predictions from ordinary least squares, lasso, random forests, and local linear forests, on individuals reporting a family size over 6, who amount for $3.3\%$ of the observed population. Cubic spline fits for each method are included to help evaluate calibration on this relatively sparse region of the dataset. Paired t-tests on the sets of squared errors for OLS ($t= 9.96$), lasso ($t=8.59$), boosting ($t = 2.92$), BART ($t = 3.27$), and random forests ($t=2.90$) give evidence for the improvements of local linear forests.
Random forests were first introduced by \citet*{breiman2001}, building on the work of \citet*{breiman1984classification} on recursive partitioning (CART), \citet*{Breiman:1996:BP:231986.231989} on bagging, and \citet*{Amit:1997:SQR:263023.263042} on randomized trees. \citet*{buhlmann2002analyzing} show how bagging makes forests smoother than single trees, while \citet*{Biau:2012:ARF:2503308.2343682} and \citet*{scornet2015} establishes asymptotic risk consistency of random forests under specific assumptions. More sophisticated tree-based ensembles motivated by random forests have been proposed by \citet*{Basu201711236}, who iteratively grow feature-weighted tree ensembles that perform especially well for discovering interactions, \citet*{zhou2018boulevard}, who consider a hybrid between random forests and boosting, and \citet*{zhu2015reinforcement}, who do deeper search during splitting to mitigate the greediness of CART. \citet*{2017arXiv170709461L} propose a Bayesian regression tree ensemble tailored to learning smooth, sparse signals and prove posterior minimaxity under certain conditions, highlighting the promise of tree-based methods that can adapt to smoothness.
The idea of considering random forests as an adaptive kernel method has been proposed by several papers. \citet*{hothorn} suggest using weights from survival trees and gives compelling simulation results, albeit to our knowledge no theoretical guarantees. \citet*{meinshausen2006quantile} proposes this technique for quantile regression forests and gives asymptotic consistency of the resulting predictions. \citet*{athey2016} leverage this idea to present generalized random forests as a method for solving heterogeneous estimating equations. They derive an asymptotic distribution and confidence intervals for the resulting predictions. Local linear forests build on this literature; the difference being that we use the kernel-based perspective on forests to exploit smoothness of $\mu(\cdot)$ rather than to target more complicated estimands (such as a quantile).
Early versions of confidence intervals for random forests, backed by heuristic arguments and empirical evidence, were proposed by \citet*{sexton2009standard} and \citet*{Wager:2014:CIR:2627435.2638587}. \citet*{Mentch:2016:QUR:2946645.2946671} then established asymptotic normality of random forests where each tree depends on a small subsample of training examples (so that there may be asymptotic bias), while \citet*{wager2017} provided a characterization of forests that allows for larger subsamples, deriving both asymptotic normality and valid confidence intervals. The confidence intervals proposed here are motivated by the algorithm of \citet*{sexton2009standard}, and build on the random forest delta method developed by \citet*{athey2016}, taking advantage of improved subsampling rates for improved coverage.
As mentioned in the introduction, a predecessor to this work is a paper by \citet*{bloniarz2016supervised}, who consider local linear regression with supervised weighting functions, including ones produced by a forest. The main differences between our method and that of \citet*{bloniarz2016supervised} is that they do not adapt the tree-splitting procedure to account for the local linear correction, and do not consider algorithmic features---such as ridge penalization---that appear to be needed to achieve good performance both in theory and in practice. Additionally, our method is flexible to forests targeting any heterogeneous estimating equation, and in particular to causal forests. On the formal side, \citet*{bloniarz2016supervised} prove consistency of their method; however, they do not establish rates of convergence and thus, unlike in our Theorem (ref), they cannot use smoothness of $\mu(\cdot)$ to provide theoretical guarantees on improved convergence properties of the forest. They also do not provide a central limit theorem or confidence intervals.
More broadly, there is an extensive body of work on model-based trees that explores different combinations of local regression and trees. \citet*{Torgo:1997:FMR:645526.657280} and \citet*{Gama:2004:FT:990375.990395} study functional models for tree leaves, fitting models instead of local averages at each node. \citet*{Karalic:1992:ELR:145448.146775} suggests fitting a local linear regression in each leaf, and \citet*{Torgo:1997:FMR:645526.657280} highlights the performance of kernel methods in general for MOB tree methods. 10.1007/978-3-642-23783-6_29 propose oblique random forests that learn split directions using the results from ridge regression, similar to our work developing splitting rules for local linear forests but more in the spirit of linear discriminant analysis (LDA). Case-specific random forests, introduced by \citet*{doi:10.1080/10618600.2014.983641}, use local information to upweight training samples not at the prediction step, but during the bootstrap to generate datasets for each tree. \citet*{doi:10.1198/106186008X319331}, and later \citet*{doi:10.1080/00949655.2012.658804}, propose not only prediction, but recursive partitioning via fitting a separate model in each leaf, similar to the residual splitting strategy of local linear forests. Local linear forests complement this literature; they differ, however, in treating forests as a kernel method. The leaf nodes in a local linear forest serve to provide neighbor information, and not local predictions.
Our work is motivated by the literature on local linear regression and maximum likelihood estimation \citep*{abadie2011bias,cleveland1988locally,GVK19282144X,heckman1998matching, Load:1999,10.2307/3532868,Stone1977,tibshirani1987local}. \citet*{Stone1977} introduces local linear regression and gives asymptotic consistency properties. \citet*{Clev1979} expands on this by introducing robust locally weighted regression, and \citet*{fan1992} give a variable bandwidth version. \citet*{cleveland1988locally} explore further uses of locally weighted regression. Local linear regression has been particularly well-studied for longitudinal data, as in \citet*{li2010} and \citet*{yao2005}. \citet*{10.2307/2959068} use local polynomials to estimate the value of a function at the boundary of its domain. \citet*{abadie2011bias} show how incorporating a local linear correction improves nearest neighbor matching procedures.
Local linear forests use a random forest to generate weights that can then serve as a kernel for local linear regression. Suppose we have training data $\smash{(X_1, Y_1), \dots, (X_n, Y_n)}$ with $\smash{Y_i = \mu(X_i) + \epsilon_i}$. Consider using a random forest to estimate the conditional mean function $\mu({x_0}) = \mathbb{E}[Y\mid X={x_0}]$ at a fixed test point ${x_0}$. Traditionally, random forests are viewed as an ensemble method, where tree predictions are averaged to obtain the final estimate. Specifically, for each tree $T_b$ in a forest of $B$ trees, we find the leaf $L_b({x_0})$ with predicted response $\hat{\mu}_b({x_0})$, which is simply the average response of all training data points assigned to $L_b({x_0})$. We then predict the average $\smash{\hat{\mu}({x_0}) = (1/B) \sum_{b=1}^B \hat{\mu}_b({x_0})}$.
An alternate angle, advocated by \citet*{hothorn}, \citet*{meinshausen2006quantile}, and \citet*{athey2016}, entails viewing random forests as adaptive weight generators. Equivalently write $\hat{\mu}({x_0})$ as
where the forest weight $\alpha_i({x_0})$ is
Notice that by construction, for each $i$, $0\le \alpha_i({x_0})\le 1$. Moreover, given that in at least one tree there exists a nonempty cell containing $x_0$, $\sum_{i=1}^n \alpha_i({x_0}) = 1$; otherwise all weights are equal to zero. \citet*{athey2016} use this perspective to harness random forests for solving weighted estimating equations, and give asymptotic guarantees on the resulting predictions.
Local linear forests solve the locally weighted least squares problem (ref) with weights (ref). Equation ((ref)) has a closed-form solution, given below, following the closed-form solutions for ridge regression and classical local linear regression. Throughout this paper, we let $A$ be the diagonal matrix with $A_{i,i} = \alpha_i({x_0})$, and let $J$ denote the $d+1\times d+1$ diagonal matrix with $J_{1,1} = 0$ and $J_{i+1,i+1} = 1$, so as to not penalize the intercept. We define $\Delta$, the centered regression matrix with intercept, as $\Delta_{i,1} = 1$ and $\Delta_{i,j+1} = x_{i,j} - x_{0,j}$. Then the local linear forest estimator can be explicitly written as
Define $\gamma_i = e_i \left(\Delta^T A \Delta + \lambda J \right)^{-1} \Delta^T$, where $e_i$ is a vector of zeroes with 1 in the $i$-th column. Qualitatively, we can think of local linear regression as a weighting estimator, with $\gamma_i \alpha_i(x_0)$ a modulated weighting function whose ${x_0}$-moments are better aligned with the test point ${x_0}$: \smash{$\hat{\mu}({x_0}) = \sum_{i = 1}^n \gamma_i \alpha_i({x_0}) Y_i$} with \smash{$\sum_{i = 1}^n \gamma_i \alpha_i({x_0}) = 1$} and \smash{$\sum_{i = 1}^n \gamma_i \alpha_i({x_0}) X_i\approx {x_0}$}, where the last relation would be exact without a ridge penalty (i.e., with $\lambda = 0$).
With the perspective of generating a kernel for local linear regression in mind, we move to discuss the appropriate splitting rule for local linear forests.
Random forests traditionally use Classification and Regression Trees (CART) from \citet*{breiman1984classification} splits, which proceed as follows. We consider a parent node $P$ with $n_P$ observations $(x_{1}, Y_{1}), \dots, (x_{n_P}, Y_{n_P})$. For each candidate pair of child nodes $C_1, C_2$, we take the mean value of $Y$ inside each child, $\bar{Y}_1$ and $\bar{Y}_2$. Then we choose $C_1, C_2$ to minimize the sum of squared errors
Knowing that we will use the forest weights to perform a local regression, we neither need nor want to use the forest to model strong, smooth signals; the final regression step can model them. Instead, in the parent node $P$, we run a ridge regression to predict $Y_{i}$ from $X_{i}$:
for intercepts $\hat{\alpha}_P$ and $\hat{\beta}_P = (x_P^T x_P + \lambda J)^{-1} x_P^T Y_P.$ We then run a standard CART split on the residuals $\smash{Y_{i} - \hat{Y}_{i}}$, modeling local effects in the forest and regressing global effects back in at prediction. Observe that, much like the CART splitting rule, an appropriate software package can enforce that a forest using this splitting rule splits on every variable and gives balanced splits; hence this splitting rule may be used to grow honest and regular trees (Section (ref)).
To explore the effects of CART and residual splitting rules, we consider this simulation first introduced by \citet*{friedman1991}. Generate $X_1, \dots, X_n$ independently and identically distributed $U[0,1]^5$ and model $Y_i$ from
for $\epsilon \sim N(0,\sigma^2)$. This model has become a popular study for evaluating nonparametric regression methods; see for example \citet*{chipman2010} and \citet*{TaddyBF}. It is a natural setup to test how well an algorithm handles interactions $ \sin(\pi X_{i1} X_{i2})$, its ability to pick up a quadratic signal $20(X_{i3} - 0.5)^2$, and how it simultaneously models strong linear signals $10 X_{i4} + 5 X_{i5} $.
Figure (ref) displays the split frequencies from an honest random forest (left) using standard CART splits, and a local linear forest (right). The x-axis is indexed by variable, here 1 through 5, and the y-axis gives tree depth for the first 4 levels of tree splits. Tiles are darkened according to how often trees in the forest split on that variable; a darker tile denotes more splits at that tree depth. CART splits very frequently on variable 4, which contributes the strongest linear signal, especially at the top of the tree but consistently throughout levels. Local linear forests rarely split on either of the strong linear signals, instead spending splits on the three that are more difficult to model.
Consider the following experiment, which highlights the benefit of the proposed splitting rule. We generate $X_1, \dots, X_n$ independently and uniformly over $[0,1]^{d}$. We hold a cubic signal $20(X_{i1} - 0.5)^3$ constant across simulations, and on each run increase the dimension and add another linear signal. Formally, we let $\xi_j = \mathbbm{1}\{j \le d\}$ and generate responses
For example, at simulation $3$ we have $\xi_1, \xi_2, \xi_3 = 1$ and hence we model $y_i = 20(X_{i1} - 0.5)^3 + 10 X_{i2} +10 X_{i3}$. Root Mean Square Error is displayed in Figure (ref).
In low dimension and with few linear signals, all three methods are comparable. However, they begin to differ quickly. Random forests are not designed for models with so many global linear signals, and hence their Root Mean Square Error increases dramatically with $d$. Moreover, as we add more linear effects, the gap between the two candidate splitting rules grows; heuristically, it becomes more important not to waste splits, and the residual splitting rule gives greater improvements. At a certain point, however, the gap between splitting rules stays constant. Once the forests simply cannot fit a more complex linear function with a fixed amount of data, the marginal benefits of the residual splitting rule level out. We show this to emphasize the contexts in which this splitting rule meaningfully affects the results.
Unless noted otherwise, all random forests used in this paper are grown using a type of sub-sample splitting called “honesty”, used by \citet*{wager2017} to derive the asymptotic distribution of random forest prediction. As outlined in Procedure 1 of \citet*{wager2017}, each tree in an honest forest is grown using two non-overlapping subsamples of the training data, denoted $\mathcal{I}_b$ and $\mathcal{J}_b$. We first choose a tree structure $T_b$ using only the data in $\mathcal{J}_b$, and write ${x_0} \leftrightarrow_{b} x'$ as the boolean indicator for whether the points ${x_0} $ and $ x'$ fall into the same leaf of $T_b$. Then, in a second step, we define the set of neighbors of ${x_0}$ as $L_b({x_0} ) = \{i \in \mathcal{I}_b : {x_0} \leftrightarrow_{b} X_i\}$; this neighborhood function is what we then use to define the forest weights in (ref). We do not use the observed outcomes $y$ from sample $\mathcal{I}_b$ to select split points; but, to ensure that each node has a certain fraction of observations from its parent, we may use the covariates from $\mathcal{I}_b$. This modification allows us to grow honest forests that comply with the upcoming assumption (ref), which says that trees are symmetric in permutations of training data index, split on every variable with nonzero probability, and balance parent observations in each child node. In this way, our theory is consistent and matches the implementation available online.
This type of subsample-splitting lets us control for potential overfitting when growing the tree $T_b$, because the samples $\mathcal{J}_b$ which are in the neighborhood $L_b({x_0})$ were held out when growing $T_b$. Despite considerable interest in the literature, there are no available consistency results for random forests with fully grown trees that do not use honesty. \citet*{Biau:2012:ARF:2503308.2343682} uses a different type of sample splitting, wherein for each tree the data is split into two sets ($\mathcal{D}_n$ and $\mathcal{D}_n'$). $\mathcal{D}_n'$ is used to evaluate the CART criterion at each node during tree growth, and $\mathcal{D}_n$ is used to split. \citet*{biau2008consistency} and \citet*{wager2015uniform} rely on large leaves, while the results of \citet*{scornet2015} on fully grown trees rely on an unchecked high-level assumption. All of these choices come at a cost; forests grown to smaller leaves can model meaningful signal while averaging out erroneous splits. We build honest forests by default.
Empirically, honesty can improve or worsen predictions. In particular, with small samples sizes and strong signals, honesty may limit the expressive power of forests and thus hurt predictive performance; conversely, with large sample sizes and weak signals, honesty may stabilize forests and thus improve performance (see Appendix B of wager2017 for a discussion). In any case, local linear corrections can help mitigate the loss of expressive power due to honesty, and so we may expect that requiring honesty is less onerous with local linear forests than with regression forests.
We recommend selecting ridge penalties by cross-validation, which can be done automatically in the R package grf. It is often reasonable to choose different values of $\lambda$ for forest training and for local linear prediction. During forest growth, equation ((ref)) gives ridge regression predictions $x_i^T \hat{{\bf \beta}}_P$ in each parent leaf. As trees are grown on subsamples of data, over-regularization at this step is a danger even in large leaves. Consequently, small values of $\lambda$ are advisable for penalization on regressions during forest training. Furthermore, as we move to small leaves, computing meaningful regression coefficients becomes more difficult; the ridge regression can begin to mask signal instead of uncovering it. A heuristic that performs well in practice is to store the regression estimates $\hat{{\bf \beta}}_P$ on parent leaves $P$. When the child leaf size shrinks below a cutoff, we use $\hat{{\bf \beta}}_P$ from the parent node to calculate ridge residual pseudo-outcomes, instead of estimating them from unstable regression coefficients on the small child dataset. In practice, this helps to avoid the pitfalls of over-regularizing and of regressing on a very small dataset when growing the forest. At the final regression prediction step ((ref)), however, a larger ridge penalty can control the variance and better accommodate noisy data.
With increasingly high-dimensional data, feature selection before prediction can significantly reduce error and decrease computation time. Often, a dataset will contain only a small number of features with strong global signals. In other cases, a researcher will know in advance which variables are strong predictors that should be included in the linear smoothing, or of special interest. In these cases, it is reasonable to run the regression prediction step on this smaller subset of predictors expected to contribute overarching trends. Such covariates, if they are not already known, can be chosen by a stepwise regression or lasso, or any other technique for automatic feature selection. Last, it is worth noting that these tuning suggestions are pragmatic in nature; the theoretical guarantees provided in Section (ref) are for local linear forests trained without these heuristics.
For conciseness, the majority of this paper focuses on local linear forests for non-parametric regression; however, a similar local linear correction can also be applied to quantile regression forests \citep*{meinshausen2006quantile}, causal forests \citep*{wager2017} or, more broadly, to any instance of generalized random forests \citep*{athey2016}. To highlight this potential, we detail the method and discuss an example for heterogeneous treatment effect estimation using local linear causal forests.
As in \citet*{athey2016}, we frame our discussion in terms of the Neyman-Rubin causal model \citep*{imbens2015causal}. Suppose we have data $(X_i, Y_i, W_i)$, where $X_i$ are covariates, $Y_i \in \mathbb{R}$ is the response, and $W_i \in \{0, \, 1\}$ is the treatment. In order to define the causal effect of the treatment $W_i$, we posit potential outcomes for individual $i$, $Y_i(0)$ and $Y_i(1)$, corresponding to the response the subject would have experienced in the control and treated conditions respectively; we then observe $Y_i = W_iY_i(1) + (1 - W_i)Y_i(0)$. We seek to estimate the conditional average treatment effect (CATE) of $W$, namely $\tau(x) = \mathbb{E}[Y(1) - Y(0) \mid X=x]$. Throughout this paper we assume uncounfoundedness \citep*{doi:10.1093/biomet/70.1.41},
and overlap, $\mathbb{P}[W_i = w] > 0$ for all $i$ and for each $w \in \{0,\, 1\}$. \citet*{wager2017} proposed an extension of random forests for estimating CATEs, and \citet*{athey2016} improved on the method by making it locally robust to confounding using the transformation of \citet*{robinson1988root}. Here, we propose a local linear correction to the method of \citet*{athey2016}, the orthogonalized causal forest, to strengthen its performance when $\tau(\cdot)$ is smooth.
Local linear causal forests start as orthogonalized causal forests do, by estimating the nuisance components
using a local linear forest. We then estimate the conditional average treatment effect via
where the $^{(-i)}$-superscript denotes leave-one-out predictions from the nuisance models. If nuisance estimates are accurate, the intercept $\hat{a}$ should be 0; however, we leave it in the optimization for robustness. We cross-validate local linear causal forests to select $\lambda_\tau$ and $\lambda_a$ by minimizing the $R$-learning criterion recommended by \citet*{nie2017learning}:
Observe that, analogous to the regression case, for very large values of $\lambda_a$ and $\lambda_{\tau}$, we will recover estimates from a causal forest. From the perspective, we can see local linear causal forests as an adapted R-learner method; note that \citet*{nie2017learning} and \citet*{kennedy2020optimal} give quasi-oracle and double-robustness properties of R-learner estimates.
To illustrate the value of the local linear causal forests, we consider a popular dataset from the General Social Survey (GSS) that explores how word choice reveals public opinions about welfare welfare. Individuals filling out the survey from 1986 to 2010 answered whether they believe the government spends too much, too little, or the right amount on the social safety net. GSS randomly assigned the wording of this question, such that the social safety net was either described as “welfare” or “assistance to the poor". This change had a well-documented effect on responses due to the negative perception many Americans have about welfare; moreover, there is evidence of heterogeneity in the CATE surface \citep*{doi:10.1093/poq/nfs036}.
Here, we write $W_i = 1$ if the $i$-th sample received the “welfare” treatment, and define $Y_i = 1$ if the $i$-th response was that the government spends too much on the social safety net. Thus, a positive treatment effect $\tau(x)$ indicates that, conditionally on $X_i= x$, using the phrase “welfare” as opposed to “assistance to the poor” increases the likelihood that the $i$-th subject says the government spends too much on the social safety net. We base our analysis on $d = 12$ covariates, including income, political views, age, and number of children. The full dataset has $N = 28,646$ observations; here, to make the problem interesting, and in particular relevant for practitioners who often have more limited survey data, we test our method on smaller subsamples of the data. Figure (ref) shows boxplots of CATE predictions by category of political views and income, comparing local linear causal forests and causal forests, and indicating possible heterogeneity with an approximately linear pattern.
In order to compare the performance of both methods, we use the transformed outcome metric of \citet*{athey2015machine}. Noting that $\mathbb{E}[(2W_i - 1) Y_i \mid X_i] = \tau(X_i)$, they suggest examining the following test set error criterion
If we can estimate $S_0$ and subtract it out, then (ref) gives an unbiased estimate of the mean-squared error of $\hat{\tau}(\cdot)$. Here, we estimate $S_0$ via out-of-bag estimation on the full dataset with $N = 28,646$, assuming that a local linear forest with such a large sample size has negligible error.
Table (ref) has error estimates for both types of forests using (ref), and verifies that using the local linear correction improves empirical performance across different subsample sizes. Practically, we can view this change as enabling us to get good predictions on less data, a powerful improvement in cases like survey sampling where data can be expensive and difficult to attain. Section (ref) contains a more detailed simulation study of local linear causal forests, comparing them with a wider array of baseline methods.
Returning to the regression case, before we delve into the main result and its proof, we briefly discuss why the asymptotic behavior of local linear forests cannot be directly derived from the existing results of \citet*{athey2016}. This is due to a key difference in the dependence structure of the forest. In the regression case, a random forest prediction at ${x_0}$ is \smash{$\hat{\mu}_{\text{rf}}({x_0}) = \sum_{i=1}^n \alpha_i({x_0}) Y_i$}, where, due to honesty, $Y_i$ is independent of $\alpha_i({x_0})$ given $X_i$. This conditional independence plays a key role in the argument of \citet*{wager2017}. Analogously to \smash{$\hat{\mu}_{\text{rf}}({x_0})$}, we can write the local linear forest prediction as a weighted sum,
where we use notation $\Delta, \, A, \, J$ from (ref). At a first glance, $\hat{\mu}({x_0})$ indeed looks like the output of a regression forest trained on observations $\rho_i$. However, the dependence structure of this object is different. In a random forest, we make $Y_i$ and $\alpha_i({x_0})$ independent by conditioning on $X_i$. For a local linear forest, however, conditioning on $X_i$ will not guarantee that $\rho_i$ and $\alpha_i({x_0})$ are independent, thus breaking a key component in the argument of \citet*{wager2017}.
We now give a Central Limit Theorem for local linear forest predictions, beginning by stating assumptions on the forest following those made in \citet*{wager2017}.
Finally, in our proof, we use the following high-level assumption on the distribution of samples as weighted by the random forest kernel. Let $d_\alpha = \sum_{i = 1}^n \alpha_i(x_0) (X_i - x_0)$ denote the difference between $x_0$ and the $\alpha_i$-weighted average of $X_i$, and let $S_\alpha = \sum_{i = 1}^n \alpha_i (X_i - x_0)^{\otimes 2}$ be the associated quadratic form. From Jensen's inequality, we immediately see that $d_\alpha' S_\alpha^{-1} d_\alpha \leq 1$, with equality only in the degenerate case where $X_i$ has no variation in the direction of $d_\alpha$. The following assumption rules out such degenerate distributions of $X_i$ within leaves, and requires that the $X_i$ enough variation along $d_\alpha$ to separate $d_\alpha' S_\alpha^{-1} d_\alpha$ from 1. We believe this to be a reasonable assumption that should be satisfied by any reasonable tree-growing algorithm; and it would be of interest to derive it from first principles in future work.
Subsampling plays a central role in our asymptotic theory, allowing us to prove asymptotic normality by building on the work of \citet*{efron1981jackknife}. Moreover, subsampling is what we use to tune the bias-variance trade-off of the forest: Forests whose trees are grown on small subsamples have lower bias but higher variance (and vice-versa).
In order to establish asymptotic unbiasedness of forests, \citet*{wager2017} require a subsample size of at least $n^{\beta}$, with
This convergence rate of a traditional honest random forest, however, does not improve when $\mu({x_0})$ is smooth. Here, we show that by using a local regression adjustment and assuming smoothness of $\mu(\cdot)$, we can grow trees on smaller subsamples of size ((ref)) without sacrificing asymptotic variance control. This allows us to decrease the bias (and improve the accuracy) of our estimates.
Our main result establishes asymptotic normality of local linear forest predictions, and gives this improved subsampling rate. The condition $\omega \le 0.2$ allows us to leverage theory from \citet*{wager2017} when controlling part of $\hat{\mu}(x_0)$. We prove this result in Appendix B.
The main draw of this results is that the best attainable rate of convergence $\beta_{\min}$ is improved compared to the rate $\beta_{\text{rf}}$ from (ref) obtained in \citet*{wager2017}. The reason we were able to obtain such an improvement in the rate of convergence is that we have assumed and consequently exploited smoothness via local linear regression.
We note that the condition (ref) on the subsampling rate enforces undersmoothing, i.e., that the error of \smash{$\hat{\mu}_n({x_0})$} will be dominated by variance. Undersmoothing implies that our estimator is asymptotically unbiased, and facilitates construction of confidence intervals. One limitation of this strategy is that it is in general difficult to tune an algorithm for undersmoothing; in particular, tuning via cross-validation does not guarantee undersmoothing. Developing methods for random forest inference that do not rely on undersmoothing following, e.g., \citet*{Hall_2013}, would be of considerable interest; however, this falls beyond the scope of the present paper.
This section complements our main result, as the Central Limit Theorem becomes far more useful when we have valid standard error estimates. Following \citet*{athey2016}, we use the random forest delta method to develop pointwise confidence intervals for local linear forest predictions.
The random forest delta method starts from a solution to a local estimating equation with random forest weights $\alpha_i({x_0})$:
\citet*{athey2016} then propose estimating the error of these estimates as
where $V({x_0}) = \nabla_{(\mu,\theta)} \mathbb{E}[\psi({x_0},Y;\mu,\theta)\mid {x_0}={x_0}]$ is the slope of the expected estimating equation at the optimum, and \smash{$\widehat{H}_n({x_0})$} is an estimate of
The upshot is that \smash{$H_n({x_0})$} measures the variance of an (infeasible) regression forest with response depending on the score function $\psi$ at the optimal parameter values, and that we can in fact estimate \smash{$H_n({x_0})$} using tools originally developed for variance estimation with regression forests. Meanwhile, $V({x_0})$ can be estimated directly using standard methods.
With local linear forests, $(\hat{\mu}, \hat{\theta})$ solve (ref) with score function
where we again use notation defined in (ref) and (ref). First, we note that we have access to a simple and explicit estimator for $V({x_0})$: Twice differentiating (ref) with respect to the parameters $(\mu, \, \theta)$ gives
which we can directly read off of the forest. In this paper, we are only interested in confidence intervals for $\mu({x_0})$, i.e., the first coordinate of $(\mu, \, \theta)$, and to estimate its variance we only need access to the entry in the upper-left corner of (ref), which we call \smash{$\hat{\sigma}_n^2$}. Given our setting, we then note that we can re-express the relevant part of (ref) as follows, in terms of $\zeta' = e_1'M_\lambda^{-1}$:
where \smash{$\widehat{\operatorname{Var}}[]$} refers to an estimate of the variance of the infeasible regression forest defined between the brackets.
Next, we follow \citet*{athey2016}, and proceed using the bootstrap of little bags construction of \citet*{sexton2009standard} to estimate the variance of this infeasible regression forest. At a high level this method is a computationally efficient half-sampling estimator. For any half sample $\mathcal{H}$, let $\Psi_{\mathcal{H}}$ be the average of the empirical scores $\Gamma_i$ averaged over trees that only use data from the half-sample $\mathcal{H}$:
where \smash{$\mathcal{S}_{\mathcal{H}}$} is the set of trees that only use data from the half-sample $\mathcal{H}$, and $L_b({x_0})$ contains neighbors of ${x_0}$ in the $b$-th tree (throughout, we assume that the subsample used to grow each tree has less than $n/2$ samples). Then, a standard half-sampling estimator would simply use \citep*{efron1982jackknife}
Now, carrying out the full computation in (ref) is impractical, and naive Monte Carlo approximations suffer from bias. However, as discussed in \citet*{athey2016} and \citet*{sexton2009standard}, bias-corrected randomized algorithms are available and perform well. Here, we do not discuss these Monte Carlo bias corrections, and instead refer to Section 4.1 of \citet*{athey2016} for details. Simulation results on empirical confidence interval performance are given in Section (ref).
In this section, we compare local linear forests, random forests, BART \citep*{chipman2010}, and gradient boosting \citep*{friedman2001}. We also include a lasso-random forest baseline for local linear forests: on half of the training data, run a lasso \citep*{Tibshirani94regressionshrinkage} regression; on the second half, use a random forest to model the corresponding residuals. Like local linear forests, this method combines regression and forests, making it a natural comparison; it is similar in spirit to the tree-augmented Cox model of \citet*{10.1093/biostatistics/kxi024}, who combine pruned CART trees with proportional hazards regression.
Random forests are trained using the R package grf package-grf, and are cross-validated via the default parameter tuning in grf, which selects values for mtry, minimum leaf size, sample fraction, and two parameters (alpha and imbalance penalty) that control split balance. Local linear forests are tuned equivalently with additional cross-validation for regularization parameters. Variables for the regression at prediction are selected via the lasso. Because existing theoretical results for random forests rely on honesty, all random forests are built with the honest construction. All lasso models are implemented via glmnet \citep*{ref-glmnet-package} and cross-validated with their automatic cross-validation feature. Local linear regression is not included in these comparisons, since the implementations loess and locfit both fail for $d>6$ on this simulation; in Appendix A, we compare local linear regression with this set of methods on lower dimensional linear models. Unless otherwise specified, all reported errors are Root Mean Square Error on 1000 test points averaged over 50 simulation runs.
Gradient boosted trees are implemented by the R package XGBoost xgboost_package. BART for treatment effect estimation is implemented following \citet*{hill2011bayesian}. As is standard, we use the BART package bart_package without any additional tuning. The motivation for not tuning is that if we want to interpret the BART posterior in a Bayesian sense (as is often done), then cross-validating on the prior is hard to justify; and in fact most existing papers do not cross-validate BART.
The first design we study is Friedman's example from equation ((ref)). Figure (ref) shows errors at $n=1000$ fixed, with dimension $d$ varying from 10 to 50. There are two plots shown, to highlight the differences between error variance $\sigma=5$ and $\sigma=20$. Appendix A reports a grid of errors for dimensions 10, 30, and 50, with $n = 1000$ and $5000$, and $\sigma$ taking values of 5 and 20. The second design we consider is given in Section (ref), as in equation ((ref)). Again we test on a grid, letting dimension $d$ take values in 5 and 50, $n$ either 1000 or 5000, and $\sigma$ at 0.1, 1, and 2. Errors are reported in Appendix A.
The third simulation is designed to test how local linear forests perform in a more adversarial setting, where we expect random forests to outperform. We simulate $X_1, \dots, X_n$ i.i.d. $U[0,1]^d$ and model responses as
Here we test dimension $d = 5, 20$ and values of $n = 500, 2000, 10000$. For this simulation, we compare only honest random forests and local linear forests, in order to compare confidence intervals; we compute out of bag Root Mean Square Error and average confidence interval coverage and length. Results are reported in Table (ref).
Figure (ref) shows Root Mean Square Error from equation (ref) at $\sigma=5$ (left) and $\sigma=20$ (right). In Section (ref), we showed that local linear forests and standard regression forests split on very different variables when generating weights. Our intuition is that these are splits we have saved; we model the strong linear effects at the end with the local regression, and use the forest splits to capture more nuanced local relationships for the weights. Local linear forests consistently perform well as we vary the parameters, lending this credibility. The lasso-random forest baseline lines up closely with local linear forests in the high noise case, separating more for low noise. BART and random forests form the next tier of methods on the low noise case; in the high noise case, honest random forest are clustered with local linear forests. BART and boosting separate in the higher noise case, suffering compared to the other methods. Appendix A shows the fuller Root Mean Square Error comparison from Friedman's model.
We move to the second simulation setup, equation (ref), meant to evaluate how methods perform in cases with a strong linear trend in the mean. Tree-based methods will be prone to bias on this setup, as the forests cannot always split on $X_1$, and because the signal is global and smooth. Full error results on the range of competing methods are given in Appendix A. Local linear forests do quite well here; they detect the strong linear signal in the tail, as we saw in Figure (ref), and model it successfully throughout the range of the feature space. Gradient boosted trees perform very well in the low-noise case, but their performance sharply declines when we increase $\sigma$.
We also examine the behavior of our confidence intervals in each of the given simulation setups, shown here in Table (ref). We give average coverage of $95\%$ confidence intervals from 50 repetitions of random forests and local linear forests on draws from the simulation setups in equations (ref), (ref), and (ref), as well as average confidence interval length and Root Mean Square Error. On equation (ref), local linear forest confidence intervals are consistently shorter and closer to $95\%$ coverage, with correspondingly lower mean squared error. Here, both random forests and local linear forests achieve fairly low Root Mean Square Error and coverage at or above $88\%$. For the setup in equation (ref), on the other hand, neither method achieves higher than $75\%$ coverage, and the local linear forest confidence intervals are longer than the random forest confidence intervals. This is an encouraging result, indicating that local linear forests confidence intervals are more adaptable to the context of the problem; we would hope for long confidence intervals when detection is difficult. Moreover, the poor coverage we see sometimes across both methods is likely because the confidence intervals are built on asymptotic results, which may not apply in some relatively low $n$ settings.
We include the approximate step function in equation (ref) to highlight a favorable example for random forests. Local linear forests see equivalent or better coverage on this setup, although at the cost of longer confidence intervals in low dimensions. Especially on small training datasets, local linear forests also improve on random forest predictions in Root Mean Square Error.
The majority of these settings are well-suited to local linear forest success; one can think of several examples where the method is likely to under-perform. With a small number of covariates, the method is similar to local linear regression, possibly worse if the forest has overfit to the data. If the local linear correction does not accurately model any underlying smoothness, cross-validation will select large values for $\lambda$, but given speed and inaccurate assumptions, random forests would be preferred in this setting.
In Section (ref), we introduced a real-data example where the local linear extension of causal forests naturally applies. Evaluating errors empirically, however, is difficult, so we supplement that with a simulation also used by \citet*{wager2017} in evaluating causal forests and \citet*{xlearner}, used to evaluate their meta-learner called the X-learner. Here we let $X \sim U([0,1]^d)$. We fix the propensity $e(x) = 0.5$ and $\mu(x)=0$, and generate a causal effect $\tau$ from each
We will assume unconfoundedness \citep*{doi:10.1093/biomet/70.1.41}; therefore, because we hold propensity fixed, this is a randomized controlled trial.
We compare local linear forests, causal forests, and X-BART, which is the X-learner using BART as a base-learner. Causal forests as implemented by grf are tuned via the automatic self-tuning feature. As in the prediction simulation studies, we do not cross-validate X-BART because the authors recommend X-BART specifically for when a user does not want to carefully tune. We acknowledge that this may hinder its performance. Local linear causal forests are tuned via cross-validation. On these simulations, consider relatively small sample sizes ranging from $n=200$ to $n=1200$ with dimension $d=20$. The goal of this simulation is to evaluate how effectively we can learn a smooth heterogeneous treatment effect in the presence of many noise covariates. \citet*{wager2017} emphasize equation (ref) as a simulation that demonstrates how forests can suffer on the boundary of a feature space, because there is a spike near $x=0$. Root Mean Square Error over 100 repetitions is reported in Table (ref), demonstrating that local linear forests give a significant improvement over causal forests. Both of these setups are reasonable tests for how a method can learn heterogeneity, and demonstrate potential for meaningful improvement with thoughtful variable selection and robustness to smooth heterogeneous signals.
In this paper, we proposed local linear forests as a modification of random forests equipped to model smooth signals and fix boundary bias issues. We presented asymptotic theory showing that, if we can assume smoother signals, we can get better rates of convergence as compared to generalized random forests. We showed on the welfare dataset that local linear forests can model smooth heterogeneous causal effects, and illustrated when and why they outperform competing methods. We also gave confidence intervals from the delta method for the regression case, and demonstrated their effectiveness in simulations.
The regression adjustments in local linear forests prove especially useful when some covariates have strong global effects with moderate curvature. Furthermore, the adjustment provides centered predictions, adjusting for errors due to an asymmetric set of neighbors. It may be that there is a useful polynomial basis corresponding to every situation where local linear forests performed well, but finding such a model would likely require hand-tuning the functional form for competitive performance, and is not automatically suited to a mix of smooth and non-smooth signals. For a departure from frequentist techniques, BART and Gaussian processes are both hierarchical Bayesian methods; BART can be viewed as a form of Gaussian process with a flexible prior, making BART the preferred baseline.
There remains room for meaningful future work on this topic. In some applications, we may be interested in estimating the slope parameter $\theta({x_0})$, rather than merely accounting for it to improve the precision of $\mu({x_0})$. While local linear forests may be an appropriate method for doing so, we have not yet explored this topic and think it could be of significant interest. Extending our theoretical results beyond pointwise convergence would enable finding uniform confidence bands and be of considerable theoretical and practical interest. We have also not considered the theoretical or empirical improvements that could arise from assuming higher order smoothness in the functions we are estimating; searching for additional optimality results in this setting could be another interesting research question.