PaperlayerView exact revision on arXiv

Source-linked AI summary

Splitting the Difference: Interpretable Causal Forests for Treatment Effect Heterogeneity and Bias

Nicolas Alexander Ihlo, Merle Behr

arXiv:2609.16971v1stat.MLcs.LG

TL;DR

Individual treatment-effect prediction needs interpretation as well as accuracy, particularly when observational confounding can be mistaken for genuine heterogeneity. IntCF changes random-forest split selection to distinguish bias correction from effect heterogeneity, and is reported to improve interpretability while remaining competitive in prediction accuracy.

  • Problem

    Individual treatment-effect estimates need to be interpretable as well as accurate, while the unobserved counterfactual makes the treatment effect unavailable as a direct regression target.

    Source: §1 · p. 2, §2.1 · p. 6

  • Method

    IntCF retains tree-based forest predictions but combines heterogeneity and bias criteria at split selection, with validation of their contributions at each node.

    Source: §3.3 · p. 13, §6.1 · p. 29

  • Results

    IntCF exceeds substantially more complex competitors in interpretability while remaining competitive with them in prediction accuracy.

    Source: §6.2 · p. 29

  • Takeaways & Limitations

    The fitted tree structure can directly support interpretation of whether splits reflect bias correction or treatment-effect heterogeneity.

    Source: §6.1 · p. 29

    The paper does not establish IntCF estimator consistency, and a typical proof assumption is violated by its discontinuous change-point model.

    Source: §6.3 · p. 30

Abstract

from arXiv · show

In various fields, such as medicine and marketing, accurately predicting individual treatment effects holds significant promise. However, achieving reliable predictions alone is often insufficient for making informed decisions; it is equally important to understand why the treatment effect is higher for some individuals than for others. To address this two-fold challenge of prediction and interpretation, we introduce an algorithm based on decision trees and random forests for estimating individual treatment effects. Our algorithm is simple: it operates exactly like a standard random forest, but with a different splitting criterion, and requires no additional workarounds such as double machine learning or orthogonalization as used in Generalized random forests. It handles observational studies with varying treatment propensities without requiring separate estimation of the full propensity function. This is achieved by combining two splitting criteria---one targeting heterogeneity in the treatment effect, the other targeting bias correction for the average treatment effect---which together improve split point selection and automatically distinguish confounders from features responsible for heterogeneity. As a result, interpretation follows directly from the fitted tree structure itself, that is, from which features the trees split on and with which split statistics, without requiring separate post-hoc analysis. For the theoretical analysis of this algorithm, we consider a change point model with step functions for potential outcomes and treatment propensity and provide insights into the theoretical underpinnings of our approach. Simulation studies show that our simple algorithm achieves comparable, and often better, prediction accuracy than existing methods, while substantially improving interpretability.

1. Introduction

The paper targets interpretable individual treatment-effect prediction, especially in observational settings where features can signal either confounding or genuine effect heterogeneity. IntCF modifies random-forest split selection to distinguish these roles while retaining a single-ensemble structure.

1. Introduction

  • Individual treatment-effect estimates need interpretation as well as accuracy, and observational confounding can make a feature appear relevant because it affects treatment propensity rather than effect heterogeneity.The paper assumes confounders are observed and rules out hidden confounding.Source: §1 · p. 2

1.1. Previous Tree-based Methods for Causal Inference

  • IntCF changes causal-forest split selection to account explicitly for confounders, aiming to distinguish bias-driven splits from genuine treatment-effect heterogeneity while preserving a single ensemble.Other tree-based approaches can sacrifice the direct, interpretable output of one ensemble through separate outcome forests or orthogonalization.Source: §1.1 · pp. 3–4

1.2. Confounding versus Heterogeneity: an Illustrative Example

  • In the example, IntCF splits on X1 to correct bias, while causal forest splits on X2 and mistakes a confounding artifact for heterogeneity.IntCF yields unbiased estimates in both child nodes; the causal-forest split leaves confounding from X1 unaddressed.Source: Figure 1 · p. 5, §1.2 · p. 5

1.3. Outline of Paper

  • The paper moves from the setting and algorithm to change-point theory, then evaluates the method in simulations, a semi-synthetic benchmark, and a real-data example.Source: §1.3 · p. 6

2. Preliminaries

The preliminaries define individual treatment-effect estimation in a potential-outcomes setting and review how random forests recursively partition features to make predictions.

2.1. Data and Potential Outcomes Model

  • The target is the conditional average treatment effect τ(x), interpreted as the individual treatment effect available for covariates x, although each unit reveals only its outcome under the assigned treatment.Because the other potential outcome is unobserved, the individual treatment effect cannot be directly observed for standard regression.Source: §2.1 · p. 6

2.2. Decision Trees and RFs

  • A random forest averages predictions from trees that recursively split the feature space using thresholds on individual features.For regression, tree predictions are usually combined by taking their average.Source: §2.2 · p. 7
  • At each node, candidate splits form two child subsets; the split with the greatest impurity decrease is chosen, and the process repeats until a stopping criterion is reached.The parent sample set is partitioned into disjoint left and right child sets.Source: §2.2 · p. 7

3. Interpretable Causal Tree and Forest Algorithm

This section describes IntCF’s causal-tree and forest procedure, which separates treatment-effect heterogeneity from bias correction in split selection and uses held-out data to validate split contributions and produce honest predictions. It also covers storing the split statistic v, growing forest trees, and calculating and scaling validation-based feature importance.

3.1. Correcting Heterogeneity Splitting Rule for Confounding

  • Within-node unconfoundedness remains a standing assumption, while departures from the weighted parent-child estimate relationship can indicate split-specific confounding beyond finite-sample noise.The proposed criterion targets the latter violation; it does not assess or correct unobserved confounding remaining within a node.Source: §3.1 · pp. 9–10
  • IntCF selects splits by maximizing either heterogeneity or bias improvement, with child-node propensity differences signaling confounding that the bias criterion can correct.The heterogeneity criterion captures variation in predicted treatment effects; the bias criterion measures the split’s correction of violations in the sample-size-weighted relationship between parent and child estimates.Source: §3.1 · p. 10

3.2. Honest Validation of Splits

  • Held-out samples validate split attribution at every node, because recomputing the training criteria out of bag alone does not remove confounding that later splits correct.The procedure extends honest prediction to inner nodes and uses the fully grown tree or forest to address residual confounding in validation.Source: §3.2 · pp. 11–12
  • Validation compares forward node estimates with backward predictions to assess heterogeneity and bias, whose criteria sum to the decrease in mean squared error from final predictions.Backward predictions aggregate final validation predictions from leaves upward; their heterogeneity criterion measures child-node differences, while bias validation measures the decrease in squared bias.Source: §3.2 · pp. 12–13

3.3. Summary of Algorithm

  • IntCT grows trees using the combined splitting criterion on training data, then validates node contributions and uses validation-based forward estimates for leaf predictions.The forest IntCF combines multiple IntCT trees.Source: §3.3 · p. 13, §3.3.1 · pp. 13–14

3.3.1. Interpretable Causal Tree (IntCT)

3.3.2. Interpretable Causal Forest (IntCF)

  • Each IntCF tree uses an independently assigned training/validation split, considers every sample in one role, and tests a random subset of features at each split.The per-tree split plays a role analogous to bootstrap and out-of-bag sampling in standard random forests.Source: §3.3.2 · p. 14

3.3.3. Feature Importance

  • Feature importance sums validation-criterion contributions across nodes splitting on a feature, using heterogeneity by default, with bias or both criteria available as alternatives.Source: §3.3.3 · p. 14
  • Because bias-validation contributions can be negative, importance scores that include them may be negative; scores are scaled so their nonnegative values sum to 1.This differs from the customary scaling of importance scores to sum to 1.Source: §3.3.3 · pp. 14–15

4. Population-level Analysis of the Splitting Criteria under a Change Point Model

The section analyzes population behavior under a single-covariate change-point model to explain how the proposed criteria separate confounding bias from genuine treatment-effect heterogeneity.

4.1. Population-level Analysis

  • The analysis studies population criteria rather than finite-sample guarantees, deriving their limits for fixed splits under simplifying assumptions.The data-independent split assumption is violated by the algorithm; the authors leave the technical uniform-convergence extension to data-selected regions for future work.Source: §4 · p. 15, §4.1 · pp. 15–17

4. Population-level Analysis of the Splitting Criteria under a Change Point Model

  • The section’s theoretical aim is to explain how the combined criterion distinguishes confounding-driven bias from genuine treatment-effect heterogeneity.The analysis uses a stylized model with step functions for treatment effect, propensity, and main outcome.Source: §4 · p. 15

4.2. Marginalization Effects

  • Because tree splits marginalize over other features, analyzing a candidate split along one variable reduces to a univariate change-point model.The model allows the treatment effect, propensity, and main outcome to change at the same threshold.Source: §4.2 · p. 17

4.3. Bias–Heterogeneity Decomposition of the Change Point Model

  • Theorem 5 decomposes the model into a constant-treatment-effect component and an unbiased-difference-in-means component, isolating bias correction from genuine heterogeneity.The decomposition is not unique, but it separates the two sources of improvement in a split.Source: §4.3 · pp. 18–19

4.4. Identifying Heterogeneity and Bias with Signed Splitting Criteria

  • At the true change point, the heterogeneity criterion has a local extremum when treatment effects vary, while the bias criterion is maximized there when difference-in-means is biased.The results hold for the specified change-point model under the positivity assumption.Source: §4.4 · pp. 20–22

5. Simulations and Application

The section evaluates IntCF’s prediction and interpretation trade-off on benchmarks, synthetic data with known effects, and a real-world NHEFS application.

5.1. Prediction Accuracy on Real-data Benchmarks

  • IntCF performs best in 2 benchmark settings, second best in 3, and third of 4 in the remaining 3, making it competitive with tree-based alternatives.It is mostly comparable with causal forest, with some improvement on IHDP and ACIC 2016; causal forest gives better ATE estimates on Lalonde datasets.Source: Table 2 · p. 23, §5.1 · p. 23

5.2. Interpretation Accuracy on Synthetic Data

  • IntCF achieves the best PEHE results in most synthetic models, especially interaction and mixed; Forest DML and GRF perform best on the change point model.IntCF beats causal forest in every case except the interaction model with s = 1.Source: Figure 4 · p. 26, §5.2 · pp. 26–27
  • IntCF’s feature importance reaches ROC-AUC almost 1 across the considered synthetic models and complexity settings, while most alternatives score lower.The statistic tests whether importance distinguishes features in the treatment-effect formula from those outside it.Source: §5.2 · pp. 26–27, Figure 5 · p. 27

5.3. Application on NHEFS Data set

  • In the 1566-smoker NHEFS analysis, IntCF attributes splits to bias correction or heterogeneity without requiring covariates to be labeled beforehand.The analysis uses the same nine baseline covariates as Hernan and Robins [2020].Source: §5.3 · p. 28
  • In NHEFS, IntCF estimates a 3.2 kg weight gain after quitting smoking; wt71 has highest heterogeneity importance, while age is the most relevant confounder.The estimated gain is slightly below the 3.4 kg reported by Hernan and Robins [2020].Source: §5.3 · pp. 28–29, Figure 6 · p. 28

6. Discussion

IntCF modifies only the random forest splitting criterion, combining heterogeneity and bias criteria so tree splits can be interpreted as heterogeneity or confounding correction. The paper reports competitive prediction accuracy and improved recovery of heterogeneity features, while identifying theoretical and analytical questions for future work.

6.1. Summary

  • In simulations, IntCF achieved competitive prediction accuracy and substantially outperformed established tree-based methods in recovering heterogeneity features, especially with confounders.The paper’s theoretical analysis supports the criterion’s attribution of splits to bias correction or genuine heterogeneity under a change-point model.Source: §6.1 · p. 29

6.2. Simplicity as a Feature, not a Limitation

  • IntCF changes only the splitting criterion, preserving direct tree-average predictions while distinguishing heterogeneity from bias correction at splits.Unlike generalized random forest and CausalForestDML, its prediction mechanism does not require an additional orthogonalization step; causal forest retains the simple structure but does not make this split-level distinction.Source: §6.2 · p. 29

6.3. Limitations and Future Work

  • The theoretical analysis is limited to the population level and a change-point model, and marginalization effects were studied with only one confounder at a time.The paper did not establish consistency of the IntCF estimator; it identifies extending the analysis to finite samples and multiple confounders as future work.Source: §6.3 · p. 30

A.1. Proofs of Section 3

The appendix proofs invoke stated assumptions and algebraic identities to establish their lemmas. The supplied passages identify some proof steps but do not state the lemmas’ full conclusions.

A.1. Proofs of Section 3

  • The Lemma 1 computation uses the parent-to-children partition identity and Assumption (12) in successive steps.The passage identifies the identities used, but does not give the lemma’s conclusion.Source: §A.1 · p. 34
  • The proof of Lemma 2 applies Assumption (13) and direct computation.The supplied passage does not state the resulting identity.Source: §A.1 · p. 35
  • The proof of Lemma 3 uses direct computation, including an expression involving child-node treatment-effect estimates and their parent estimate.The displayed fragment gives a relation involving the left and right child estimates and the parent estimate.Source: §A.1 · p. 35

A.2. Proof of Theorem 5

The proof of Theorem 5 analyzes a change-point model, taking the covariate to be uniform on [0, 1] without loss of generality. It decomposes the outcome and treatment effect into bias and heterogeneity components and includes a setting in which the difference-in-means estimator is unbiased.

A.2. Proof of Theorem 5

  • The proof assumes a uniform covariate on [0, 1], noting that an affine rescaling gives an equivalent model on another interval.The rescaling preserves treatment-effect estimates and signed splitting criteria.Source: §A.2 · p. 37
  • Theorem 5’s proof decomposes the outcome function and treatment effect into bias and heterogeneity components.The passage states that μ(x) = μB(x) + μH(x) and τ(x) = τB(x) + τH(x).Source: §A.2 · p. 37
  • Under the specified parameter condition, the difference-in-means estimator is unbiased in expectation for the treatment effect.The passage explicitly states that its expectation equals the treatment effect.Source: §A.2 · p. 38

A.3. Proof of Theorem 6

The proof of Theorem 6 studies signed splitting criteria in a change-point model, assuming a uniform covariate and restricting the change point to the unit interval. Its results identify an extremum of the bias criterion and a local extremum of the heterogeneity criterion at the change point.

A.3. Proof of Theorem 6

  • The proof assumes X is uniform on [0, 1] and takes the change-point parameter γ to lie in (0, 1).The proof says parameter values outside that interval can be represented by an equivalent model inside it with the same expected predictions and signed splitting criteria.Source: §A.3 · p. 38
  • When the difference-in-means estimator is biased in the bias model, the bias splitting criterion has a global extremum at the change point γ.The proof rules out local extrema on either side of γ before identifying the global extremum there.Source: §A.3 · pp. 41–42
  • For a non-constant treatment effect in the heterogeneity model, the heterogeneity splitting criterion has a local extremum at γ.In the specified model, the difference-in-means estimate tends to the correct average treatment effect.Source: §A.3 · pp. 42–43

Related papers

Loading 2609.16971v1…