Generalized Random Forests

1 Regression Random Forests Recap

Random forests extend decision trees through ensembling, and in their classical regression form they are designed to estimate the conditional mean function

\[ m(x) = \mathbb{E}[Y \mid X = x], \]

given training data \(\{(X_i, Y_i)\}_{i=1}^n\). A regression tree partitions the covariate space \(\mathcal{X}\) into \(M\) disjoint terminal nodes \(A_1, \dots, A_M\). Here, \(|A|\) denotes the number of training observations with \(X_i\in A\). For a new input \(x \in A_m\), the tree prediction is the sample average of outcomes in that region,

\[ \hat f_{A_m} = \frac{1}{|A_m|} \sum_{i: X_i \in A_m} Y_i, \]

with trees fitted on randomly sampled subsets of data and features. A random forest then averages a large number of such trees to reduce variance and improve stability. We have shown previously that the forest prediction is also a kernel method. More importantly, the essential component we will look at in this lecture is the splitting rule used in tree construction, which is reduction of variance in a regression problem:

\[ \begin{aligned} &\Var(A) - \left( \frac{|A_L|}{|A|} \Var(A_L) + \frac{|A_R|}{|A|} \Var(A_R) \right)\\ =& \frac{1}{|A|}\sum_{i: X_i \in A} (Y_i - \bar Y_A)^2 - \left( \frac{|A_L|}{|A|} \frac{ \sum_{i: X_i \in A_L} (Y_i - \bar Y_{A_L})^2 }{|A_L|} + \frac{|A_R|}{|A|} \frac{ \sum_{i: X_i \in A_R} (Y_i - \bar Y_{A_R})^2 }{|A_R|} \right) \\ =& \frac{1}{|A|}\left[\sum_{i: X_i \in A} (Y_i - \bar Y_A)^2 - \left( \sum_{i: X_i \in A_L} (Y_i - \bar Y_{A_L})^2 + \sum_{i: X_i \in A_R} (Y_i - \bar Y_{A_R})^2 \right)\right], \end{aligned} \]

which corresponds to the ANOVA decomposition of the total sum of squares. This splitting rule is intuitive and effective for estimating conditional means, but it does not generalize to other parameters of interest. The generalized random forest (GRF) framework proposed by Athey et al. (2019) overcomes this limitation by replacing the variance reduction splitting rule with a more general and theoretically justified criterion based on influence functions.

2 Influence Function

The influence function is a fundamental concept in statistics. It is typically a semiparametric and robust statistics concept, but here we demonstrate it through simple scalar examples that connect with the score function, Fisher information, and M-estimators. The influence function provides a way to understand the robustness and efficiency of an estimator. It measures the sensitivity of an estimator to small changes in the data distribution, which then allows us to analyze the asymptotic behavior of the estimator. Let \(P\) denote the distribution of a random variable \(Y\), then a statistical parameter can often be written as a functional of \(P\). For example,

• Mean: \(T(P) = \mathbb{E}_P[Y] = \int y \, dP(y)\);
• Variance: \(T(P) = \Var_P(Y)\);
• Quantile: \(T(P) = F_P^{-1}(\tau)\).

Suppose we slightly perturb \(P\) by introducing a small amount of point mass at an observation \(y\). This new perturbed distribution can be expressed as

\[ P_\varepsilon = (1 - \varepsilon) P + \varepsilon \, \delta_y, \]

where \(\delta_y\) is the Dirac measure at \(y\) and \(\varepsilon > 0\) is small. The influence function then examines how the statistical parameter \(T(P)\) changes as we vary \(\varepsilon\):

\[ \begin{aligned} \text{IF}(y; T, P) &= \lim_{\varepsilon \to 0} \frac{T(P_\varepsilon) - T(P)}{\varepsilon} \\ &= \left. \frac{d}{d\varepsilon} T\big( \varepsilon \, \delta_y + (1 - \varepsilon) P \big) \right|_{\varepsilon = 0} \end{aligned} \]

which is the Gateaux derivative of the functional \(T\) at \(P\) in the direction of \(\delta_y - P\). You may think of the distribution \(P\) as an infinite vector of probabilities assigned to each possible observation, and the influence function measures the sensitivity of the functional \(T\) to a small change in this vector. This derivative quantifies the infinitesimal effect that adding a small mass at \(y\) has on the parameter estimate. With this understanding, we can construct a first-order functional approximation of the estimator \(T\) at its truth \(T(P)\):

\[ \begin{aligned} T(P_\varepsilon) &= T(P) + \int \text{IF}(u; T, P) d(P_\varepsilon - P)(u) + R(\varepsilon) \\ &= T(P) + \int \text{IF}(u; T, P) d( \varepsilon \delta_y + (1 - \varepsilon) P - P)(u) + R(\varepsilon) \\ &= T(P) + \varepsilon \int \text{IF}(u; T, P) d(\delta_y - P)(u) + R(\varepsilon) \end{aligned} \]

where \(R(\varepsilon)\) is a remainder term that goes to zero faster than \(\varepsilon\) as \(\varepsilon \to 0\). With this understanding, we can analyze the estimator \(T\) on the observed empirical distribution \(P_n\) based on a sample \(Y_1, \dots, Y_n\), i.e.,

\[ P_n = \frac{1}{n} \sum_{i=1}^n \delta_{Y_i}. \]

Noticing that the difference of the empirical distribution \(P_n\) and the true distribution \(P\) can be expressed as

\[ P_n - P = \frac{1}{n} \sum_{i=1}^n (\delta_{Y_i} - P) = \sum_{i=1}^n \varepsilon_i (\delta_{Y_i} - P) \]

with all \(\varepsilon_i = 1/n\). Hence the difference between \(P_n\) and \(P\) can be viewed as a sum of small perturbations, each adding a small mass \(\varepsilon_i\) at the observed data point \(Y_i\) and removing the same amount of mass proportionally to \(P\). Assuming sufficient differentiability of \(T\) and remainder control, we can apply the functional Taylor expansion to approximate \(T(P_n)\):

\[ \begin{aligned} T(P_n) &= T(P) + \int \text{IF}(y; T, P) d(P_n - P)(y) + R(P_n - P) \\ &= T(P) + \int \text{IF}(y; T, P) d P_n(y) - \int \text{IF}(y; T, P) d P(y) + R(P_n - P) \\ &= T(P) + \frac{1}{n} \sum_{i=1}^n \text{IF}(Y_i; T, P) - \int \text{IF}(y; T, P) dP(y) + R(P_n - P) \\ &= T(P) + \frac{1}{n} \sum_{i=1}^n \text{IF}(Y_i; T, P) + R(P_n - P) \end{aligned} \]

with the understanding that the integral of the influence function with respect to the distribution \(P\) is zero¹. Hence this decomposition shows that the error of the estimator \(T(P_n) - T(P)\) can be approximated by the average of the influence functions evaluated at the observed data points, plus a small remainder term assumed to be \(o_p(n^{-1/2})\) under these conditions.

\[ T(P_n) - T(P) \approx \frac{1}{n} \sum_{i=1}^n \text{IF}(Y_i; T, P) \]

Note that the right hand side is the average of a set of i.i.d. terms. Assuming the influence function has finite variance, this allows us to establish asymptotic normality via the central limit theorem. Since these terms are independent and have mean 0, we estimate the variance of their average using the estimated influence values:

\[ \widehat{\Var}(T(P_n)) = \frac{1}{n^2} \sum_{i=1}^n \widehat{\text{IF}}^2(Y_i; T, P). \]

2.1 Example: Mean Estimation

Let’s take a look at an example of the influence function for the mean estimation. For the mean, we have \(\theta = T(P) = \mathbb{E}_P[Y]\). Consider \(P_\varepsilon = (1-\varepsilon)P + \varepsilon \delta_y\), i.e., putting a small point mass at \(y\). Then

\[ T(P_\varepsilon) = (1-\varepsilon)\mathbb{E}_P[Y] + \varepsilon y \]

And we can calculate the influence function as

\[ \begin{aligned} \text{IF}(y;T,P) &= \left.\frac{d}{d\varepsilon}T(P_\varepsilon)\right|_{\varepsilon=0}\\ &= y - \mathbb{E}_P[Y]. \end{aligned} \]

Note that \(\int \text{IF}(y;T,P)\, dP(y) = \mathbb{E}_P[Y] - \mathbb{E}_P[Y] = 0\) confirms our earlier claim. Therefore, the influence function for the mean functional is simply the deviation of the observation \(y\) from the population mean \(\mathbb{E}_P[Y]\). Empirically, to calculate the influence function at each data point \(Y_i\), we can replace the population mean with the sample mean \(\bar{Y}\):

\[ \widehat{\text{IF}}_i = Y_i - \bar{Y}. \]

To obtain the variance, note that the true influence function satisfies \(\Var(\text{IF}(Y_i;T,P)) = \Var(Y_i - \mathbb{E}_P[Y]) = \Var(Y_i)\). We estimate this using the empirical variance:

\[ \widehat{\Var}(\text{IF}(Y_i;T,P)) = \frac{1}{n} \sum_{i=1}^n (Y_i - \bar{Y})^2. \]

This uses the empirical-variance denominator \(n\); the usual unbiased sample variance instead divides by \(n-1\).

3 Generalized Random Forests

The idea of the splitting rule in generalized random forests is to use the influence function at each node to guide the splitting, instead of using variance reduction. This is typically a semiparametric model topic, but here we will only give a brief introduction. The key idea is to consider a more general parameter of interest \(\theta\) defined through a moment condition, rather than just the mean. For example, we use \(\psi\) to denote an estimating function corresponding to the parameter of interest:

\[ \mathbb{E}[\psi(Y; \theta)] = 0, \]

At each internal node, we have a subset of data points \(\{(X_i, Y_i): X_i \in A\}\) corresponding to the region \(A\). By treating the data in this node as coming from i.i.d. samples from the same distribution, we can compute the influence function for each observation in this node with respect to the parameter of interest. We then use these influence function values to perform a variance reduction type of splitting, similar to the classical regression tree. Why does this make sense? Recall that the influence function measures the sensitivity of the estimator to small changes in the data distribution. By splitting the node in a way that maximizes the difference in average influence function values between child nodes, we are effectively trying to put observations that have positive influence on the parameter estimate into one side of the split, and those with negative influence into the other side. This leads to child nodes that are more homogeneous in terms of their contribution to the parameter estimate, which reduces the variance of the estimator within each node.

Let’s consider a regression problem where we observe \(\{x_i, y_i\}_{i=1}^n\). The parameter of interest is the conditional mean \(m(x) = \E( Y \mid X=x)\). However, at each internal node, we will simply treat all observations as i.i.d., and estimate the node parameter \(\theta\). After calculating the influence function for each observation in the node, we will use these influence function values as pseudo-responses to perform a variance reduction type of splitting. For the regression case, since the influence function is the centered \(Y_i\) values within the node, the splitting criterion becomes exactly the same as the classical regression tree splitting rule. But we could look at some other examples:

3.1 Binary Outcome

Suppose \(Y \in \{0, 1\}\), and we are interested in estimating the conditional probability \(\theta(x) = \mathbb{P}(Y=1 \mid X=x)\). The moment condition at a given node, treating the data as i.i.d. within that node, is

\[ \mathbb{E}[Y - \theta \mid X\in A] = 0. \]

Here, let \(P\) denote the distribution of \(Y\) conditional on \(X\in A\). To calculate the influence function, we consider the perturbed distribution \(P_\varepsilon = (1-\varepsilon)P + \varepsilon \delta_y\). Then following our definition of the influence function, we have

\[ T(P_\varepsilon) = (1-\varepsilon)\mathbb{E}_P[Y] + \varepsilon y = (1-\varepsilon)\theta + \varepsilon y \]

and

\[ \text{IF}(y;T,P) = \left.\frac{d}{d\varepsilon}T(P_\varepsilon)\right|_{\varepsilon=0} = y - \theta. \]

Empirically, we can estimate the influence function at each data point \(Y_i\) in the node by replacing \(\theta\) with the sample proportion \(\hat{\theta} = \frac{1}{|A|} \sum_{i: X_i \in A} Y_i\):

\[ \widehat{\text{IF}}_i = Y_i - \hat{\theta}. \]

Hence, we use variance reduction based on the pseudo-responses \(Y_i - \hat{\theta}\). Subtracting the same parent-node mean does not change the within-child variances. For binary outcomes, Gini impurity is twice the empirical variance, so the two criteria give the same split rankings. Entropy can give different rankings.

3.2 Treatment Effect Estimation

Suppose we have a binary treatment \(A \in \{0, 1\}\) and an outcome \(Y\). We are interested in estimating the conditional average treatment effect

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

Under consistency, conditional unconfoundedness, and overlap, this can be identified as

\[ \tau(x) = \mu_1(x) - \mu_0(x), \qquad \mu_a(x) = \mathbb{E}[Y \mid X=x, A=a]. \]

Let

\[ e(X) = \mathbb{P}(A=1 \mid X) \]

be the propensity score, and let

\[ m(X) = \mathbb{E}[Y \mid X] = e(X)\mu_1(X) + (1-e(X))\mu_0(X) \]

be the marginal outcome regression. Notice that \(m(X)\) does not condition on treatment. We can estimate it by regressing \(Y\) on \(X\) while ignoring the treatment label.

Let’s first see how these quantities give us a moment condition for the treatment effect. Since \(\mathbb{E}[A-e(X) \mid X]=0\),

\[ \begin{aligned} \mathbb{E}[(A-e(X))(Y-m(X)) \mid X] &= \mathbb{E}[(A-e(X))Y \mid X] \\ &= e(X)\mu_1(X) - e(X)m(X) \\ &= e(X)(1-e(X))\{\mu_1(X)-\mu_0(X)\} \\ &= e(X)(1-e(X))\tau(X). \end{aligned} \]

Also, because \(A\) is binary,

\[ \mathbb{E}[(A-e(X))^2 \mid X] = \operatorname{Var}(A \mid X) = e(X)(1-e(X)). \]

Combining these two equations gives

\[ \boxed{ \mathbb{E}\left[ (A-e(X))\{Y-m(X)-\tau(X)(A-e(X))\} \mid X \right] = 0. } \]

This has a familiar regression interpretation. After removing the parts of \(Y\) and \(A\) predictable from \(X\), we regress the remaining outcome on the remaining treatment. The local slope is the treatment effect.

To connect this moment condition with our influence-function discussion, consider a fixed node \(\mathcal{L}\). First treat \(e\) and \(m\) as known, and write

\[ U = A-e(X), \qquad V = Y-m(X). \]

Within this node, consider the scalar moment equation

\[ \mathbb{E}[U(V-\tau_{\mathcal{L}} U) \mid X\in \mathcal{L}] = 0. \]

Solving it gives

\[ \tau_{\mathcal{L}} = \frac{\mathbb{E}[UV \mid X\in \mathcal{L}]} {\mathbb{E}[U^2 \mid X\in \mathcal{L}]}. \]

If the treatment effect is constant within the node, this slope equals that treatment effect. When the treatment effect varies within the node, \(\tau_{\mathcal{L}}\) is an average weighted by \(e(X)(1-e(X))\). Thus, interpreting one slope as the treatment effect throughout the node is a local approximation.

Now we can calculate its influence function. Let \(P_{\mathcal{L}}\) denote the distribution within the fixed node, and write

\[ a_{\mathcal{L}} = \mathbb{E}_{P_{\mathcal{L}}}[UV], \qquad q_{\mathcal{L}} = \mathbb{E}_{P_{\mathcal{L}}}[U^2] > 0, \qquad \tau_{\mathcal{L}} = \frac{a_{\mathcal{L}}}{q_{\mathcal{L}}}. \]

Adding a small point mass at an observation with residuals \((u,v)\) changes the slope to

\[ \tau_{\mathcal{L},\varepsilon} = \frac{(1-\varepsilon)a_{\mathcal{L}}+\varepsilon uv} {(1-\varepsilon)q_{\mathcal{L}}+\varepsilon u^2}. \]

Using the quotient rule,

\[ \begin{aligned} \operatorname{IF}_{\mathcal{L}}(u,v) &= \left.\frac{d}{d\varepsilon}\tau_{\mathcal{L},\varepsilon}\right|_{\varepsilon=0} \\ &= \frac{(uv-a_{\mathcal{L}})q_{\mathcal{L}}-a_{\mathcal{L}}(u^2-q_{\mathcal{L}})}{q_{\mathcal{L}}^2} \\ &= \boxed{\frac{u(v-\tau_{\mathcal{L}}u)}{q_{\mathcal{L}}}}. \end{aligned} \]

The numerator is the residual moment for this observation. The denominator measures how much treatment variation is available within the node. An observation has positive influence when its contribution pushes the local slope upward, and negative influence when it pushes the slope downward.

In practice, we estimate the nuisance functions and form residuals

\[ \widetilde A_i = A_i-\widehat e(X_i), \qquad \widetilde Y_i = Y_i-\widehat m(X_i). \]

For this scalar local model, the node estimate and corresponding estimated influence values are

\[ \widehat\tau_{\mathcal{L}} = \frac{\sum_{i:X_i\in \mathcal{L}}\widetilde A_i\widetilde Y_i} {\sum_{i:X_i\in \mathcal{L}}\widetilde A_i^2}, \]

and

\[ \widehat\rho_i = \frac{\widetilde A_i(\widetilde Y_i-\widehat\tau_{\mathcal{L}}\widetilde A_i)} {\frac{1}{n_{\mathcal{L}}}\sum_{j:X_j\in \mathcal{L}}\widetilde A_j^2}, \]

where \(n_{\mathcal{L}}\) is the number of observations in the node. These influence values provide pseudo-responses for finding splits that separate different local treatment effects.

For clarity, this derivation uses a scalar local regression without an intercept. The full causal forest also fits a local intercept. The essential idea remains to learn neighborhoods suited to treatment-effect estimation and then perform a locally weighted regression of residualized outcomes on residualized treatment.

Finally, estimating \(e\) and \(m\) introduces another statistical issue. The conditional residual moment above is orthogonal: at the true treatment effect, small nuisance-function errors have no first-order effect on its conditional expectation. Cross-fitting, or out-of-bag (OOB) nuisance predictions, helps avoid fitting and evaluating the nuisance functions on the same observations. Honesty addresses a different use of the data by separating observations used to choose splits from those used to estimate parameters within leaves.

3.3 A General Numerical Calculation

In general, we may not need to derive the influence function analytically since the estimator we use could be complex. However, realizing that the influence function measures sensitivity to a small perturbation, we can approximate it numerically. For a sufficiently smooth functional \(T\), the leave-one-out approximation is

\[ \widehat{\text{IF}}_i \approx (n-1)\left\{T(P_n)-T(P_n^{(-i)})\right\}. \]

This means that to estimate the influence function at a data point \(Y_i\), we subtract the estimator \(T\) computed without \(Y_i\) (i.e., leave-one-out) from the estimator on the full dataset, and multiply by \(n-1\). This approach can be restricted to the data points within a node when constructing the tree, with \(n\) replaced by the node sample size, for sufficiently smooth parameter functionals. A common nonzero scaling of these differences does not change the variance-reduction split ranking within a node. This numerical approach is particularly useful when the estimator \(T\) is complex and we do not have a closed-form expression for its influence function.

4 Forest Estimation and Variance

Realizing that random forests use a kernel method, or weighted average method, we can express the target \(\theta(x)\) through the local moment condition

\[ \E[\psi(Y; \theta(x)) \mid X=x] = 0 \]

where \(\psi\) is an estimating function corresponding to the parameter of interest. It may be a likelihood score or a negative loss gradient, but a general moment condition need not come from a loss function. For causal estimation, the estimating function also depends on treatment and nuisance functions, whose arguments are suppressed here.

• For squared error loss, we can use \(\psi(Y; \theta) = Y - \theta\), which is proportional to the negative loss gradient. This gives our moment condition for conditional mean estimation.
• For a Bernoulli likelihood, parameterizing by the log odds \(\theta(x)=\log\{p(x)/(1-p(x))\}\) gives the likelihood score \(\psi(Y; \theta) = Y - \text{logit}^{-1}(\theta) = Y - p(x)\). This gives our moment condition for conditional probability estimation.

The locally weighted moment condition can be expressed as solving for \(\widehat\theta(x)\) with

\[ \sum_{i=1}^n w_i(x) \psi(Y_i; \widehat\theta(x)) = 0 \]

where \(w_i(x)\) are the forest weights for each observation, learned from the training data for the target point \(x\). They are properly normalized to sum to 1. In each tree, the weight is shared equally among the observations used for estimation in the terminal node containing \(x\), and these weights are then averaged across trees. Hence the weight reflects how each observation contributes to the estimation of \(\theta(x)\), with observations in the learned neighborhood of \(x\) receiving higher weights.

We can actually derive the asymptotic variance of the forest estimator \(\hat{\theta}(x)\) using the influence function. The key idea is to approximate its estimation error by a weighted average of influence contributions evaluated at each observation. Assuming the estimating function is sufficiently smooth, a Taylor expansion gives

\[ \begin{aligned} \mathbf{0} &= \sum_{i=1}^n w_i(x) \psi(Y_i; \widehat{\theta}(x)) \qquad \text{by definition of} \, \widehat\theta(x) \\ &= \sum_{i=1}^n w_i(x) \psi(Y_i; \theta(x)) + \sum_{i=1}^n w_i(x) \left[ \left. \frac{\partial}{\partial \theta} \psi(Y_i; \theta) \right|_{\theta = \theta(x)} \right] (\widehat{\theta}(x) - \theta(x)) + r_n(x), \end{aligned} \]

where \(r_n(x)\) is the Taylor remainder. Think of this as a local linear approximation to the estimating equation, with \(\widehat{\theta}\) lying close to \(\theta\). Hence, if we define the Jacobian matrix

\[ H(x) = \mathbb{E}\left[ \left. \frac{\partial}{\partial \theta} \psi(Y_i; \theta) \right|_{\theta = \theta(x)} \mid X = x \right], \]

this measures how the expected estimating function changes with the parameter. When \(\psi\) is a likelihood score, it equals the negative Fisher information under the usual regularity conditions. More generally, it is the moment Jacobian, and we estimate it with

\[ \widehat{H}(x) = \sum_{i=1}^n w_i(x) \left. \frac{\partial}{\partial \theta} \psi(Y_i; \theta) \right|_{\theta = \widehat{\theta}(x)}. \]

Then, assuming \(H(x)\) is nonsingular, the weighted Jacobians consistently estimate it, and the remainder is negligible relative to the forest estimator’s standard error, we can re-express \(\widehat{\theta}(x) - \theta(x)\) approximately as

\[ \widehat{\theta}(x) - \theta(x) \approx - \widehat{H}(x)^{-1} \sum_{i=1}^n w_i(x) \psi(Y_i; \theta(x)). \]

The relevant standard-error scale depends on the forest’s localization and subsampling; it need not be \(n^{-1/2}\). For confidence intervals for the true \(\theta(x)\), the local approximation bias must also be negligible relative to this standard error.

Following our previous discussion on estimating the variance of an estimator using the influence function, we take the variance of the entire weighted estimating function. Define

\[ \Gamma_n(x) = \Var\left\{\sum_{i=1}^n w_i(x)\psi(Y_i;\theta(x))\right\}, \]

where the variance includes the randomness of the training data and the resulting forest weights. Given a consistent estimate \(\widehat\Gamma_n(x)\), this leads to the sandwich form:

\[ \widehat{\Var}(\widehat{\theta}(x)) = \widehat{H}(x)^{-1}\widehat\Gamma_n(x)\widehat{H}(x)^{-T}. \]

Here \(H^{-T}\) denotes the transpose of \(H^{-1}\); for a scalar parameter, the expression is \(\widehat\Gamma_n(x)/\widehat H(x)^2\). The estimating function \(\psi\) becomes an influence contribution after multiplication by \(-H(x)^{-1}\); \(\psi\) alone is generally not the influence function. However, in practice, there are three considerations in estimating this variance:

1. The weights \(w_i(x)\) are random and depend on the data, so the weighted estimating functions need not be independent. The simple sum \(\sum_i w_i(x)^2\psi(Y_i;\widehat\theta(x))^2\) from a fixed-weight calculation does not generally estimate \(\Gamma_n(x)\). In grf, the “honest” approach separates the observations used to choose splits from those used to estimate parameters within each tree’s leaves. This addresses one source of data reuse, but observations are still shared across trees, and forest-specific variance estimation is needed.
2. Out-of-bag prediction excludes an observation from the trees used to predict for that observation. It is also useful when estimating nuisance functions for the causal moment. However, it does not ensure independence of all the weighted estimating functions or justify treating the learned weights as fixed.
3. In grf, \(\Gamma_n(x)\) is estimated using grouped trees and a bootstrap-of-little-bags construction (Athey et al. 2019). Comparing variation within and across groups accounts for the forest’s sampling variability. We shall discuss the variance estimation in more detail in the uncertainty lecture.

5 Example: Using the grf Package

We can illustrate CATE estimation using a causal forest in the grf package. By default, the package uses honest trees, dividing each tree’s subsample in half: one half chooses the splits, and the other estimates parameters within the leaves. The package also provides variance estimates for the CATE. In the code, \(W\) denotes the treatment indicator \(A\). Calling predict(cf) without new data produces out-of-bag predictions for the training observations.

  library(grf)
  
  # Simulated data
  n <- 500
  p <- 5
  X <- matrix(rnorm(n * p), n, p)
  W <- rbinom(n, 1, 0.5)  
  tau <- 2 * X[,1] # the true CATE function
  
  Y <- X[,1] * 2 + W * tau + rnorm(n)

  par(mfrow=c(1,2))
  # plot data with two regression lines
  plot(X[,1], Y, col = ifelse(W==1, "deepskyblue", "darkorange"), pch=16, xlab="X1", ylab="Y")
  abline(lm(Y[W==1] ~ X[W==1,1]), col="deepskyblue", lwd=2)
  abline(lm(Y[W==0] ~ X[W==0,1]), col="darkorange", lwd=2) 
  
  # Fit a causal forest
  cf <- causal_forest(X, Y, W,
                      honesty = TRUE,
                      honesty.fraction = 0.5)
  
  # Estimate conditional treatment effects using out-of-bag predictions
  cf.pred <- predict(cf, estimate.variance = TRUE)
  tau_hat <- cf.pred$predictions
  
  # plot the estimated CATE against the true CATE
  plot(X[,1], tau_hat, pch=16, xlab="X1", ylab="Estimated CATE", col="black")
  points(X[,1], tau, col="red", pch=16, cex = 0.5)
  legend("bottomright", legend=c("Estimated CATE", "True CATE"), 
         col=c("black", "red"), pch=16, cex = 1.5)

The following plot shows pointwise 95% confidence intervals for the CATE at 20 observations.

  # get pointwise 95% confidence intervals
  ci <- 1.96 * sqrt(cf.pred$variance.estimates)
  
  # plot the estimated CATE with confidence intervals for the first 20 observations
  # order them by their true CATE values
  tau_hat_ci <- data.frame(tau = tau, 
                           tau_hat = tau_hat,
                           lower = tau_hat - ci,
                           upper = tau_hat + ci)[1:20, ]
  
  tau_hat_ci <- tau_hat_ci[order(tau_hat_ci$tau), ]
  
  plot(tau_hat_ci$tau, pch=19, xlab="Index (ordered by true CATE)", 
       ylab="Estimated CATE with 95% CI", 
       col = "red", ylim = c(min(tau_hat_ci), max(tau_hat_ci)), 
       main = "Estimated CATE with 95% CI vs True CATE")
  points(tau_hat_ci$tau_hat, pch=19, col = "black", cex = 0.7)
  segments(seq_len(nrow(tau_hat_ci)),
           tau_hat_ci$lower, 
           seq_len(nrow(tau_hat_ci)),
           tau_hat_ci$upper, 
           lwd = 1.5)
  legend("topleft", legend=c("True CATE", "Estimated CATE", "95% CI"), 
         col=c("red", "black", "black"), pch=c(19,19,NA), 
         lwd=c(NA,NA,1.5), cex = 1.2)

---

Athey, Susan, Julie Tibshirani, and Stefan Wager. 2019. “Generalized Random Forests.” The Annals of Statistics 47 (2): 1148–78.

---

1. A simple intuitive explanation is that when we move \(P\) towards the direction of itself, i.e., \((1 - \varepsilon) P + \varepsilon P\), the influence function would on average not change the value of our estimator. For mean estimation, this is easier to see. For a more rigorous argument, we realize that \[ \begin{aligned} \int \text{IF}(y; T, P) dP(y) &= \int \frac{d}{d\varepsilon} T((1 - \varepsilon) P + \varepsilon \delta_y) \Big|_{\varepsilon = 0} dP(y) \\ &= \frac{d}{d\varepsilon} \int T((1 - \varepsilon) P + \varepsilon \delta_y) dP(y) \Big|_{\varepsilon = 0} \\ &= \frac{d}{d\varepsilon} \int T( P + \varepsilon (\delta_y - P) ) dP(y) \Big|_{\varepsilon = 0} \qquad (\ast) \end{aligned} \] For the term inside, we are going to linearize it around \(\varepsilon = 0\) if \(T\) is nicely behaved. This means we again break down \(T( P + \varepsilon (\delta_y - P) )\) into \(T(P)\) and linear terms of influence functions: \[ T(P + \varepsilon (\delta_y - P) ) = T(P) + \varepsilon \int \text{IF}(u; T, P) d(\delta_y - P)(u) + R(\varepsilon) \] For the integral, we have \[ \int \text{IF}(u; T, P) d \delta_y(u) = \text{IF}(y; T, P) \] and \[ \int \text{IF}(u; T, P) d P(u) \] being the thing we wanted to prove to be zero. Denote this by \(A\). Then, the integral in \((\ast)\) becomes \[ \begin{aligned} \int T( P + \varepsilon (\delta_y - P) ) dP(y) &= \int \left[ T(P) + \varepsilon (\text{IF}(y; T, P) - A) + R(\varepsilon) \right] dP(y)\\ &= T(P) + \varepsilon (A - A) + o(\varepsilon) \\ &= T(P) + o(\varepsilon) \end{aligned} \] Taking derivative with respect to \(\varepsilon\) and evaluating at \(\varepsilon = 0\) gives us zero, which completes the proof.↩︎