Causal Forests / Heterogeneous Treatment Effects

Motivating Example: Who Benefits from the Drug?#

A hospital conducts a randomized trial of a new drug for heart disease. The trial shows that, on average, the drug reduces the risk of a heart attack by 5 percentage points. But the hospital administrator asks: which patients benefit most?

Perhaps the drug works well for patients with high cholesterol but provides little benefit for those with low cholesterol. Perhaps it is more effective for older patients. Perhaps there is a complex interaction between age, cholesterol, blood pressure, and diabetes status.

Traditional subgroup analysis requires you to pre-specify: "Let me check patients above vs. below age 65." But with 50 potential effect modifiers, this approach encounters the multiple testing problem directly. And if the analysis is exploratory rather than pre-specified, it is likely to produce "significant" heterogeneity even when there is none.

Causal forests address this challenge by using the random forest algorithm, modified to target treatment effect estimation rather than outcome prediction (Athey & Imbens, 2016) (Wager & Athey, 2018). They discover which covariates drive heterogeneity without requiring you to pre-specify the subgroups.

From Prediction to Causal Estimation#

A standard random forest predicts the average outcome within each leaf, much like OLS predicts conditional means. A causal forest instead estimates the treatment effect within each leaf, enabling estimation of treatment effect heterogeneity. The forest's output is a Conditional Average Treatment Effect (CATE):

This formula tells you: for a person with characteristics $x$, what is the expected effect of treatment?

How Causal Trees Differ from Regular Trees#

Regular trees split on variables that best predict $Y$. Causal trees split on variables that best separate treatment effects. At each node, the split is chosen to maximize the variance of treatment effects across the two child nodes. If splitting on "age > 65" produces groups with 10% and 1% treatment effects, that separation is a good split. If splitting on "gender" produces two groups with 5% effects each, it is uninformative.

Honest Estimation#

A key innovation is honest estimation. In a standard random forest, the same data determine tree structure and estimate predictions within leaves. This dual use creates overfitting.

In an honest causal forest, data are split:

• Structure sample: Determines where to split
• Estimation sample: Estimates treatment effects within each leaf

This separation ensures CATE estimates are not biased by the tree-growing process.

Common Confusions#

"Is variable importance from a causal forest causal?" No. Variable importance tells you which variables best predict heterogeneity in treatment effects. It does not tell you that those variables cause the heterogeneity. Zip code might have high importance because it proxies for diet, exercise, and healthcare access.

"Do I need an experiment?" Not necessarily, but unconfoundedness is required. Causal forests work with both experimental and observational data, as long as treatment is independent of potential outcomes conditional on $X$. They are most credible with experimental data.

"How is this approach different from subgroup regressions?" Subgroup analysis requires pre-specified subgroups. Causal forests discover relevant subgroups from the data and handle complex interactions that would be difficult to specify manually. The trade-off is interpretability. For estimating an average treatment effect with ML-assisted confounding control, see double/debiased machine learning.

"Can I use causal forests for policy targeting?" Yes. If you estimate $\hat{\tau}(x)$ for each person, you can target treatment to those with the largest predicted benefits. This application is called optimal treatment assignment or personalized policy learning.

The Target Estimand#

Identifying Assumptions#

1. Unconfoundedness: $\{Y(0), Y(1)\} \perp\!\!\!\perp D \mid X$
2. Overlap: $0 < P(D = 1 \mid X = x) < 1$ for all $x$
3. SUTVA: No interference between units

The Causal Forest Estimator#

Wager and Athey (2018) show that:

where the weights $\alpha_i(x)$ are determined by the forest structure. Under regularity conditions:

• Consistency: $\hat{\tau}(x) \to \tau(x)$
• Asymptotic normality: $(\hat{\tau}(x) - \tau(x)) / \hat{\sigma}(x) \to N(0, 1)$, enabling pointwise confidence intervals

Imagine a scatterplot of patients with age on one axis and cholesterol on another. Each point is colored by treatment effect. A causal forest draws a heat map over this space, estimating the treatment effect at every point and discovering regions where effects cluster together.

Explore how targeting treatment to high-CATE subgroups improves the average effect per treated individual:

1# Requires: grf
2# grf: Generalized Random Forests (Athey, Tibshirani, Wager)
3library(grf)
4
5# --- Step 1: Fit the causal forest ---
6# causal_forest() grows an ensemble of honest causal trees
7# honesty = TRUE: separate subsamples for tree structure and leaf estimation
8# This separation ensures valid confidence intervals for CATEs
9cf <- causal_forest(
10X = as.matrix(df[, covariate_cols]),
11Y = df$outcome,
12W = df$treatment,
13num.trees = 2000,  # more trees = more stable estimates (diminishing returns past ~2000)
14honesty = TRUE,    # required for valid inference; do not turn off
15seed = 42
16)
17
18# --- Step 2: Estimate individual-level CATEs ---
19# predict() returns out-of-bag CATE estimates for each observation
20# estimate.variance = TRUE enables pointwise confidence intervals
21cate <- predict(cf, estimate.variance = TRUE)
22df$tau_hat <- cate$predictions          # estimated treatment effect for each unit
23df$tau_se <- sqrt(cate$variance.estimates)  # SE via infinitesimal jackknife
24
25# --- Step 3: Estimate the average treatment effect (ATE) ---
26# target.sample = "all" averages CATEs across the full sample
27ate <- average_treatment_effect(cf, target.sample = "all")
28cat("ATE:", ate[1], "(SE:", ate[2], ")\n")
29# Interpretation: the population-average causal effect of treatment
30
31# --- Step 4: Variable importance ---
32# Measures which covariates the forest uses most for splitting
33# NOTE: this is predictive importance for heterogeneity, NOT causal moderation
34varimp <- variable_importance(cf)
35varimp_vec <- setNames(as.vector(varimp), covariate_cols)
36sort(varimp_vec, decreasing = TRUE)[1:10]
37
38# --- Step 5: Calibration test for heterogeneity ---
39# Tests whether estimated CATEs predict actual treatment effect variation
40# A significant "differential forest prediction" suggests genuine heterogeneity
41test_calibration(cf)
42
43# --- Step 6: Best linear projection ---
44# Projects CATEs onto specific covariates for interpretable coefficients
45# Use this (not variable importance) to test specific moderation hypotheses
46blp <- best_linear_projection(cf,
47A = as.matrix(df[, c("age", "cholesterol")]))
48print(blp)

1. Calibration test. test_calibration() in grf tests whether estimated CATEs predict actual treatment effect variation. A significant "differential forest prediction" coefficient suggests genuine heterogeneity.

2. CATE distribution. Plot $\hat{\tau}(x)$. Tight concentration around the ATE suggests homogeneous effects. Wide spread suggests substantial heterogeneity.

3. Variable importance. Report which covariates the forest uses most for splitting. Remember: predictive importance, not causal.

4. Best linear projection. Project CATEs onto key covariates for interpretable coefficients.

5. Overlap check. Verify propensity scores are bounded away from 0 and 1.

6. Out-of-bag predictions. The grf package uses OOB predictions by default for honest estimation.

Interpreting Your Results#

Significant heterogeneity detected: Report the CATE distribution, top variables by importance, and best linear projection. Discuss who benefits most and policy implications.

No significant heterogeneity: This null finding is informative. The ATE is a good summary for everyone.

Variable importance caveat: "Age has the highest variable importance" means age best predicts which people have larger or smaller effects. It does NOT mean age causes the difference. It is important to caveat this distinction.

Calibration test p-value: 0.002
Variable importance: age (0.15), education (0.12), prior_earnings (0.11)

• OLS with interactions: When the dimensions of heterogeneity are known a priori and can be specified parametrically — simpler and more interpretable when theory provides clear subgroup hypotheses.
• Pre-specified subgroup analysis: When the research question concerns a small number of pre-specified groups (e.g., by gender, age bracket) rather than a continuous heterogeneity surface.
• DML: When the goal is an average treatment effect with high-dimensional confounders, rather than treatment effect heterogeneity.
• Bayesian Additive Regression Trees (BART): An alternative ML-based approach to heterogeneous treatment effects with built-in uncertainty quantification and a different regularization philosophy.