Estimation and Inference of Heterogeneous Treatment Effects using Random Forests

Stefan Wager; Susan Athey

Source-linked AI summary

Estimation and Inference of Heterogeneous Treatment Effects using Random Forests

Stefan Wager, Susan Athey

arXiv:1510.04342v4Oct 14, 2015stat.MEmath.STstat.ML

TL;DR

Researchers need methods that can estimate heterogeneous treatment effects and support valid inference. This paper develops causal forests and shows they provide asymptotically valid inference while outperforming k-nearest-neighbor matching in simulations.

Problem
Researchers need to infer heterogeneous treatment effects, but historically datasets were often too small to explore such heterogeneity beyond sample subdivision.
Method
The paper develops causal forests using honest, subsampled trees to estimate treatment effects with tractable asymptotic theory and valid inference.
Results
Causal forests dominate k-nearest-neighbor matching in bias and variance across simulations, maintaining mean-squared error of 0.02 as d grows from 2 to 30.
Takeaways & Limitations
Causal forests combine adaptive feature selection with asymptotically normal, unbiased point estimates and valid confidence intervals for heterogeneous treatment effects.
Takeaways & Limitations
The theory provides only pointwise confidence intervals, while boundary bias and challenging settings such as small samples or many covariates remain open problems.

Abstract

from arXiv · show

Many scientific and engineering challenges -- ranging from personalized medicine to customized marketing recommendations -- require an understanding of treatment effect heterogeneity. In this paper, we develop a non-parametric causal forest for estimating heterogeneous treatment effects that extends Breiman's widely used random forest algorithm. In the potential outcomes framework with unconfoundedness, we show that causal forests are pointwise consistent for the true treatment effect, and have an asymptotically Gaussian and centered sampling distribution. We also discuss a practical method for constructing asymptotic confidence intervals for the true treatment effect that are centered at the causal forest estimates. Our theoretical results rely on a generic Gaussian theory for a large family of random forest algorithms. To our knowledge, this is the first set of results that allows any type of random forest, including classification and regression forests, to be used for provably valid statistical inference. In experiments, we find causal forests to be substantially more powerful than classical methods based on nearest-neighbor matching, especially in the presence of irrelevant covariates.

1 Introduction

The paper develops causal forests to estimate heterogeneous treatment effects with tractable asymptotic theory and valid inference, addressing the limitations of classical methods and prediction-oriented random forests. It combines honest, subsampled trees with theoretical guarantees and finds advantages over nearest-neighbor matching as covariate dimensionality increases.

Motivation: Inference is especially important because searching repeatedly for extreme subgroups can produce spurious heterogeneity, while causal effects lack an observable test-set ground truth.These concerns motivate prespecified subgroup protocols and a central role for statistical theory in evaluating estimation noise.
Motivation: Classical nonparametric methods perform well with few covariates but quickly break down as covariate dimensionality increases, motivating machine-learning approaches.The introduction frames heterogeneous-effect estimation as relevant to medicine, marketing, public policy, and randomized technology experiments.
Contribution: Causal forests estimate treatment effects at tree leaves, providing a forest-based method designed for tractable asymptotic theory and valid statistical inference.The method follows Athey and Imbens [2016] by building forests from causal trees.
Theory: The paper develops consistency and asymptotic normality results for honest, subsampled regression forests, with the infinitesimal jackknife consistently estimating asymptotic variance.Honest trees use separate subsamples for tree construction and leaf prediction; the framework supports confidence intervals for forest estimates.
Empirical results: Causal forests dominate classical k-nearest-neighbor matching in bias and variance across simulations, with the advantage increasing as the number of covariates grows.The paper also examines coverage rates for confidence intervals for heterogeneous treatment effects.
Novelty: The main technical contribution is a theory enabling statistical inference from random-forest predictions, extending rigorous analysis beyond particular variants and black-box heuristics.The authors describe this as the first set of conditions allowing random-forest predictions to support such inference.

2 Causal Forests

Causal forests estimate heterogeneous treatment effects by averaging tree-based within-leaf treatment-effect estimates under unconfoundedness. With appropriately scaled subsamples and honest trees, they are consistent, asymptotically Gaussian and unbiased, with consistently estimable variance.

Causal-effect setup: Under unconfoundedness, causal forests target τ(x) despite observing only one potential outcome per subject, using nearby observations as approximate randomized comparisons.Unconfoundedness makes local methods such as nearest-neighbor matching consistent for τ(x) under continuity assumptions.
Causal trees and forests: Trees adaptively define neighborhoods through leaves, and causal trees estimate treatment effects within leaves whose observations are intended to resemble randomized experiments.Forests average the estimates from B causal trees, while adaptive leaves can narrow where signal changes quickly and widen elsewhere.
Consistency: Causal forests are pointwise consistent for the true treatment effect τ(x) under conditions on the conditional mean functions.The paper establishes consistency for causal forests built from causal trees.
Inference: With appropriately scaling subsample size s, causal-forest predictions are asymptotically Gaussian and unbiased, enabling confidence intervals centered at the forest estimates.The sampling-distribution result uses potential nearest neighbors and classical U-statistic and projection arguments.
Inference: The infinitesimal jackknife variance estimate is consistent for Var[ˆτ(x)], including a finite-sample correction for subsampling without replacement.The correction is n(n−1)/(n−s)^2, and the estimated-to-true variance ratio converges in probability to 1.
Honest causal forests: The theory requires honest trees, which separate split selection from within-leaf estimation; double-sample and propensity-tree procedures satisfy this condition.Subsample re-randomization lets observations participate in both roles across trees, while propensity trees use treatment assignments for splitting and can reduce bias from variation in e(x).

3 Asymptotic Theory for Random Forests

This section develops asymptotic theory showing that suitably honest random-forest predictions are asymptotically Gaussian, with consistently estimable variance and valid confidence intervals. It establishes conditions for these results through subsampling, incrementality, and bias control, while identifying limitations for voting trees.

Random-forest construction: The framework averages trees over all size-s subsamples and auxiliary randomness, while treating Monte Carlo effects as negligible when the number of replicates B is sufficiently large.Finite-B effects are studied by Mentch and Hooker [2016], and Wager et al. [2014] recommend B on the order of n.
Regularity conditions: The results require honest trees and impose regularity conditions including Lipschitz continuity, density assumptions, and restrictions on subsample-size growth.Under these conditions, the theory derives incrementality and bias bounds for honest regular and double-sample trees.
Asymptotic normality and variance estimation: Random-forest predictions are asymptotically Gaussian, and the infinitesimal jackknife consistently estimates their asymptotic variance.The theory therefore supports confidence intervals for the target conditional mean.
Limitations and corrections: For classification forests with k > 1, the proof applies to trees averaging observations within leaves but not to voting trees, whose extension remains open.The finite-sample correction n(n −1)/(n −s)^2 is appropriate for subsampling without replacement and improves practical performance.
Contributions: Theorem 1 establishes conditions under which random-forest predictions are asymptotically unbiased and normal, providing the first such result according to the authors.Earlier work motivated the variance estimator through classical connections but did not formally justify it; this analysis establishes consistency for the random-forest prediction variance.
Incrementality and subsampling: Subsampling converts weakly incremental trees into 1-incremental ensembles, enabling classical asymptotic theory and consistency of the infinitesimal jackknife.The argument first establishes incrementality lower bounds for regression trees and then applies subsampling.

4 Inferring Heterogeneous Treatment Effects

This section extends random-forest asymptotic theory to causal forests for estimating conditional average treatment effects at a prespecified test point. Under unconfoundedness, overlap, and regularity conditions, the resulting estimates are consistent and asymptotically Gaussian and centered.

4 Inferring Heterogeneous Treatment Effects: Causal forests estimate τ(x) by averaging causal-tree estimates trained on subsamples at a prespecified test point x.The training data contain features, outcomes, and binary treatment assignments.
4 Inferring Heterogeneous Treatment Effects: Honest causal trees exclude responses from split selection, while regular trees require sufficient observations from both treatment groups in each leaf.Trees may use treatment assignments when choosing splits, and regularity is defined relative to the test point x.
4 Inferring Heterogeneous Treatment Effects: Under unconfoundedness, overlap, potential-outcome regularity, and specified honest, regular, symmetric random-split conditions, causal-forest predictions are consistent and asymptotically Gaussian and centered.The result applies to forests with α ≤ 0.2 and subsample sizes scaling as in (14).
4 Inferring Heterogeneous Treatment Effects: Regular causal trees cannot generally be constructed simultaneously for every test point, so the proof considers one prespecified point x at a time.The section illustrates this limitation with a one-dimensional treatment assignment determined by whether X_i ≥ 0.

5 Simulation Experiments

Across simulations, causal forests substantially outperform k-NN matching for heterogeneous treatment-effect estimation, maintaining low mean-squared error as dimension increases while revealing coverage limitations caused by bias. The experiments also assess propensity adjustment, adaptation to treatment-effect heterogeneity, sampling variance, Gaussianity, and finite-sample confidence-interval validity.

Study design: The simulations test whether forest-based methods provide asymptotically valid confidence intervals while improving over non-adaptive k-NN in finite samples under unconfoundedness.Performance is evaluated using expected mean-squared error for τ(X) and expected coverage at a target rate of 0.95, with X uniformly distributed and homoscedastic noise.
First setup: propensity adjustment: Causal forests maintain mean-squared error of 0.02 as d increases from 2 to 30, while 10-NN and 100-NN perform an order of magnitude worse.The comparison is from the first setup, where treatment effects are fixed and propensity-treatment-outcome interactions create bias; 100-NN is already bias-dominated at d = 2.
Inference diagnostics: Figure 1 diagnostics show causal-forest sampling variance decreasing with n, infinitesimal-jackknife variance estimates becoming relatively accurate, and predictions approaching Gaussianity.The diagnostics aggregate results over 1,000 randomly sampled test points and evaluate standardized predictions across forests trained on many datasets.
Second setup: treatment-effect heterogeneity: In the heterogeneous-effect setup, causal forests similarly improve over k-NN when d > 2, while unexpectedly improving as the ambient dimension grows from small values.The setting uses randomized treatment with m(x) = 0 and e(x) = 0.5, and treatment-effect heterogeneity supported on the first two features; when d = 2, performance is comparable to 50-NN.
Third setup: boundary and peak effects: Causal forests achieve an order-of-magnitude mean-squared-error improvement over k-NN for d > 2 in the third setting, but their confidence-interval coverage deteriorates.The coverage problem appears to arise from increasing bias: forest intervals are not centered, especially as dimension grows, although the qualitative treatment-effect surface is captured.
Limitations and future work: The study’s main limitation is manually selecting propensity or double-sample forests for each setting, motivating automatic splitting rules and principled selection of s.The authors also identify bias reduction and improved coverage as immediate challenges, including through splitting rules that focus on coordinates with the greatest signal.

6 Discussion

The paper proposes non-parametric causal forests for heterogeneous treatment-effect estimation that combine data-driven feature selection with asymptotically normal, unbiased estimates and valid confidence intervals. It achieves principled inference through honest trees and random-forest subsampling, while leaving global functional inference and boundary-bias correction open.

Contributions: Causal forests enable data-driven feature selection while retaining asymptotically normal, unbiased point estimates and valid confidence intervals.The estimator can be viewed as an adaptive nearest-neighbor method in which the data determine which feature dimensions matter for selecting neighbors.
Inference: Pairing honest trees with random-forest subsampling provides a simple, principled way to address selection bias in adaptive methods used for statistical inference.The discussion identifies selection bias as a central challenge for adaptive inference and situates the approach alongside recent advances by Berk et al. [2013], Chernozhukov et al. [2015], and Taylor and Tibshirani [2015].
Limitations and future work: The current theory provides only pointwise confidence intervals for τ(x), leaving global functional estimation as an open direction.Extending the theory to global functional estimation is identified as a promising avenue for future work.
Limitations and future work: Nearest-neighbor non-parametric estimators can suffer boundary bias, motivating systematic trimming at feature-space boundaries and possible bias correction.The discussion identifies boundary treatment and bias correction as unresolved methodological challenges.

A Additional Simulations

Additional simulations show that causal forests’ advantage over k-NN depends on signal sparsity, while their confidence intervals achieve near-nominal coverage under appropriate variance conditions. As signal density increases, forest gains diminish and can reverse.

Bias-Variance Trade-off: Causal forests face a bias-variance trade-off: bias dominates when s is small relative to n, whereas variance dominates when s is larger relative to n.Confidence intervals achieve close-to-nominal coverage when mean-squared error matches the average variance estimate σ̂2(X).
Signal Density: When signal is spread across many features, forests have less opportunity to improve over baseline methods and do not improve much over them.Earlier experiments benefited from strong, low-dimensional structure that forests could exploit relative to k-NN matching.
Signal Density: Forests vastly outperform k-NN in mean-squared error when signal dimensions are sparse, but lose this advantage when dimensions are comparable and perform worse for dense d = q = 6.When the signal is dense, all surveyed methods have poor coverage except 10-NN, whose intervals are simply very wide.

B Is Honesty Necessary for Consistency?

Honesty appears necessary for pointwise centered asymptotic Gaussian inference: in a corner-point experiment, honest forests remain unbiased while adaptive forests develop increasing bias despite existing average-risk consistency results.

Implication: Pointwise centered asymptotic Gaussianity appears to require honesty, even though most practical random forests are adaptive rather than honest.The framework’s honesty assumption therefore marks its main departure from mainstream random-forest implementations.
Simulation evidence: Honest forests remain unbiased and their mean-squared error decreases with sample size, whereas adaptive forests show bias exceeding sampling variation that increases with n.The comparison uses Figure 3’s corner prediction point, 500 trees, leaf-size k = 1, and 40 replications.
Mechanism: Adaptive CART forests become pointwise biased in corners of x-space because they use outcome information to isolate rare outliers when placing splits.Honest trees avoid this problem by selecting splits on a separate sample that does not reveal where the outliers occur.
Relation to consistency: Growing the minimum leaf size can avert the corner-bias phenomenon, but the resulting bias bounds decay more slowly than random-forest sampling variance and do not support centered confidence intervals.Thus, existing consistency results do not establish the centered inference required here.
Relation to consistency: Adaptive forests can remain risk-consistent at average test points because most large-sample test points lie away from corners, although honest forests retain lower test-set error in this setting.This explains why average prediction consistency does not eliminate the pointwise bias relevant for inference.

B.1 Adaptive versus Honest Predictive Error

Honest and adaptive forests can achieve comparable best predictive error, but honest forests are much less sensitive to the minimum leaf size, performing well even at the default k = 1.

B.1 Adaptive versus Honest Predictive Error: The comparison assesses whether honesty sacrifices mean-squared error at a random test point, a concern motivated by prior findings that honesty can hurt regression-forest performance on some UCI datasets [2014].
B.1 Adaptive versus Honest Predictive Error: With d = 8 and the usual adaptive-forest setting k = 1, adaptive forests perform worse than 50 nearest neighbors, though increasing minimum leaf size improves them.
B.1 Adaptive versus Honest Predictive Error: At their best, adaptive and honest forests have comparable root-mean-squared error, but adaptive forests require careful tuning of the minimum leaf size.The adaptive optimum is k = 64, while honest forests perform well across a wide range of k values, including the default k = 1.

C Proofs

This appendix defines asymptotic-scaling notation and proves Theorem 1 by showing that the estimator’s bias is negligible relative to its standard deviation, enabling Slutsky’s lemma.

Notation: The appendix defines notation for asymptotic equivalence, upper and lower bounds, big-O, big-Omega, and little-o rates.These definitions specify the limiting-ratio meanings used throughout the proofs.
Proof of Theorem 1: Theorem 1 follows by proving |E[ˆµ_n(x)] − µ(x)| / σ_n(x) → 0, after which Slutsky’s lemma completes the argument.The proof combines the preceding theorem statements to replace E[ˆµ_n(x)] with µ(x) in the target result.

C.1 Bounding the Bias of Regression Trees

The proof bounds regression-tree bias by controlling leaf diameters. Regularity, concentration, honesty, and Lipschitz continuity together yield the same bounded-bias conclusion for the forest average.

Lemma 2: Regularity forces sufficiently many splits along each coordinate, and a Chernoff bound converts this into a probabilistic lower bound on coordinate-wise split counts.The argument uses the number of splits leading to L(x) and the random-split-tree property.
Lemma 2: With high probability, leaf diameters contract approximately at the regularity rate because node sample proportions track their Lebesgue measures, except for O(log log n) terminal ancestors.The concentration result holds simultaneously for nearly all parent nodes above L(x), with failure probability controlled asymptotically.
Theorem 3: Under honesty and Lipschitz continuity, bounding the average leaf diameter suffices to bound the conditional-mean bias.Lipschitz continuity contributes the conditional-mean constant, while honesty provides the first observation in the theorem proof.
Theorem 3: The resulting bias bound is O(1), and because forests average trees, the same bound holds for the forest estimator ˆµ(x).The proof combines the diameter bound with the Lipschitz assumption before extending the result from individual trees to their average.

C.2 Bounding the Incrementality of Regression Trees

The section bounds the incrementality of regression trees by controlling k-PNN probabilities and then applies these bounds to establish variance rates for regular and double-sample trees. The proof handles uniform and bounded feature densities, using truncation, Lipschitz moments, honesty, and coupling to control negligible terms.

Proof of Lemma 4: The resulting k-PNN bounds extend from uniform features to any feature density bounded by a constant C on [0, 1]^d.The bounded-density case follows by comparing probabilities with the uniform-density argument.
Proof of Theorem 5: For regular trees, Lemma 4 and Theorem 3 yield the required upper variance control as fixed k becomes negligible relative to shrinking leaf size.The proof bounds k Var[T(x; Z)] using the number of observations in the leaf containing x, then concludes the target convergence.
Proof of Corollary 6: For double-sample trees, applying Theorem 5 with the effective sample size ⌊s/2⌋ gives the desired variance rate, while the complementary term is negligibly small.The proof uses Hoeffding’s standard results together with Lemma 2 and the argument from Theorem 3.

C.3 Properties of Subsampled Incremental Base Learners

This section uses the Efron–Stein ANOVA decomposition to analyze subsampled random forests through their Hájek projections. Under honesty, moment, and subsample-scaling conditions, it establishes asymptotic normality and shows that the infinitesimal jackknife consistently estimates the projection variance.

ANOVA framework: The Efron–Stein ANOVA decomposition provides the abstract framework for analyzing symmetric base learners and random-forest estimators through mean-zero, uncorrelated components.The decomposition also identifies the Hájek projection as the first-order component used in subsequent variance and asymptotic arguments.
Random-forest representation: A subsampled random forest can be rewritten using the ANOVA components of its individual trees, with subsampling multiplicities determining the corresponding projection terms.The derivation relies on how often each training point and each pair of points appears across the possible subsamples.
Asymptotic normality: Under the stated scaling and regularity conditions, the Hájek projection is asymptotically normal, and the forest estimator inherits the required limiting behavior via Slutsky’s lemma.The proof uses Lyapunov’s central limit theorem, honesty-based conditional independence, uniform (2+δ)-moment bounds, and subsample scaling; a possible double-sample tree contributes a factor of 2.

C.4 Extension to Causal Forests

The causal-forest extension reuses the regression-forest bias and incrementality arguments under unconfoundedness, honesty, regularity, and ε-overlap. The resulting bounds retain the same structure but incur worse constants depending on ε.

Extension to Causal Forests: Under ε-overlap, causal-forest bias and incrementality follow through analogues of the regression-forest arguments, with constants worsened by a factor depending on ε.The proof replaces the effective sample size s with εs and reproduces the subsampling-based argument from Section 3.3.2.
Bias: With Lipschitz conditional means, honest trees, bounded potential-outcome means, ε-overlap, and α ≤ 0.2, the causal-forest bias at x is bounded by the stated analogue of Theorem 3.The overlap condition ensures the minimum treated-or-control count satisfies smin/s ≳ ε, enabling the εs replacement in the bias argument.
Incrementality: For honest k-regular causal trees, Lipschitz conditional variances, positive variance, and ε-overlap imply ν(s)-incrementality at x with the stated ε-dependent bound.The proof adapts the minority-class and nearest-neighbor argument within each leaf, then invokes the machinery underlying Lemma 4 and Theorem 5.