Random Forests for Treatment Probability Estimation

Why Use a Random Forest for Treatment Probabilities?

Inverse probability weighting requires a model for the probability of receiving each treatment, conditional on the measured pre-treatment history. A logistic model can be effective, but its usual linear predictor may not capture thresholds, interactions, or other nonlinear patterns unless they are specified. A random forest offers a flexible alternative: it partitions the covariate space into regions and estimates treatment probabilities from the treatment frequencies in those regions.

Flexibility is not a guarantee of good probabilities. The treatment model must represent the relevant history, its predictions must be suitable for weighting, and the observed data must contain support for the treatment choices being weighted. A forest can adapt to complex patterns while still producing unstable probabilities in sparse regions.

From Forest Leaves to Class Probabilities

A classification tree divides the covariate space into terminal leaves. Within a leaf, a natural estimate of the probability of treatment 1 is the fraction of observations in that leaf that received treatment 1. A random forest builds many trees, typically using bootstrap samples and a random subset of candidate predictors at each split. To obtain a probability estimate, the tree-level class probabilities are averaged across trees.

For a training observation \(i\), write \(A_i\) for its treatment and \(H_i\) for its history. In tree \(b\), let \(L_b(H)\) be the terminal leaf containing prediction history \(H\), and let \(S_b\) be the tree's bootstrap sample, counting repeated sampled observations with their multiplicities. A simple tree-level probability is

The forest probability is \(\hat e(H)=B^{-1}\sum_{b=1}^B\hat e_b(H)\), where \(B\) is the number of trees. Implementations can use different weighting or tree-building details, but the essential object for inverse weighting is a number between 0 and 1 that estimates the conditional class probability. A reported hard class such as “treated” or “untreated” is not enough.

Proof. Each terminal-leaf fraction has a numerator between zero and its denominator, so \(\hat e_b(H)\in[0,1]\). The average of \(B\) numbers in \([0,1]\) is also in \([0,1]\). In the weighted form, \(A_i\in\{0,1\}\) implies \(0\leq\alpha_i(H)A_i\leq\alpha_i(H)\). Summing gives \(0\leq\sum_i\alpha_i(H)A_i\leq\sum_i\alpha_i(H)=1\). \(\square\)

This bound is a mathematical property of the probability calculation, not evidence that the estimated probabilities are accurate. In particular, a forest's class vote fraction—how many trees classify \(H\) as treated—may differ from the average of the trees' treatment probabilities. Use class-probability predictions, not a thresholded class or an unexamined vote output.

Turning Predictions into Inverse Weights

For each person, the denominator must be the estimated probability of the treatment that person actually received. Given \(\hat e_i=\hat e(H_i)\), define the estimated probability of the observed treatment as

When \(A_i=1\), this gives \(\hat p_i=\hat e_i\); when \(A_i=0\), it gives \(\hat p_i=1-\hat e_i\). A common error is to use \(\hat e_i\) in the denominator for everyone. That is correct only for treated observations.

With treatment decisions at multiple times, a fitted probability is needed for each decision given the history available at that time. The denominator for a person's observed treatment sequence is the product of the corresponding observed-treatment probabilities. For example, if two observed treated decisions have fitted probabilities \(0.80\) and \(0.25\), their product is \(0.80(0.25)=0.20\), and the inverse product is \(1/0.20=5\). A small probability at even one decision can make the overall weight large.

Identification and the Role of the Probability Model

The causal justification for inverse weighting does not depend on using a forest. It depends on assumptions about treatment assignment, outcomes, and support, and on estimating the relevant probabilities adequately. For a single decision, let \(Y(a)\) be the potential outcome under treatment \(a\), and suppose \(H\) is measured before treatment.

Proof. Write \(e_a(H)=P(A=a\mid H)\). By iterated expectation,

On the event \(A=a\), consistency gives \(Y=Y(a)\), so the numerator is \(E[I(A=a)Y(a)\mid H]\). Conditional exchangeability makes \(Y(a)\) independent of \(I(A=a)\) given \(H\), yielding

Positivity permits division by \(e_a(H)>0\). Thus the conditional expectation of the weighted observed outcome is \(E[Y(a)\mid H]\); taking expectations gives \(E[Y(a)]\). The integrability assumption ensures the target mean exists, and under the stated identities the absolute expectation of the weighted term is finite. \(\square\)

In practice, the true \(e_a(H)\) is unknown and is replaced by a fitted probability. The identity therefore explains the goal of treatment-probability estimation; it does not guarantee that a particular finite-sample forest estimate recovers the target. A random forest is an estimation strategy, not a substitute for exchangeability, positivity, correct timing of covariates, or adequate probability estimation.

Diagnostics, Comparisons, and Failure Modes

Inspect both the probability estimates and the weights they generate. The following checks help reveal problems, although none proves that the treatment model is correct.

• Check the observed-treatment probabilities. Tabulate or summarize \(\hat p_i\), separately by observed treatment. Values near zero imply large weights. Look at the largest weights and at how many observations contribute to the tail.
• Check treatment support across histories. If a history region contains almost exclusively one treatment, the other class probability is weakly supported. A forest cannot create information about treatment choices absent from the data.
• Compare predicted probabilities with observed treatment frequencies in groups. Large discrepancies can indicate poor probability prediction. Such a check is a diagnostic, not proof of causal validity; strong prediction metrics alone do not establish that the probabilities are suitable for weighting.
• Distinguish training from out-of-bag predictions. Predictions for observations used to construct a tree can be overly optimistic. Out-of-bag predictions exclude an observation from the bootstrap sample for the trees used in its prediction and can help assess predictive performance. They do not, by themselves, establish causal identification or resolve all dependence and inference concerns.
• Review the feature set and timing. Treatment predictors must be available before the decision. Including information measured after treatment can contaminate the assignment model and distort the weights.

Compared with logistic regression, a forest can represent nonlinearities and interactions without prescribing their exact form, but its behavior depends on tree depth, minimum leaf size, the number of candidate predictors, and other choices. A parametric model can be more stable when its form is reasonable and data are limited. Ridge or lasso logistic models impose a different form of regularization, while smoothing sparse treatment-history cells pools information across selected cells. None is universally superior: compare models using probability diagnostics and the resulting weight distribution, not classification accuracy alone.

A forest also does not automatically solve sparse-data or positivity problems. Very small leaves can yield highly variable treatment fractions; large leaves can conceal important differences in treatment assignment. Increasing the number of trees can reduce Monte Carlo variation in the forest average, but it cannot repair inadequate overlap or a systematically unsuitable model. Truncating large weights may reduce variance, but it changes the estimating procedure and can introduce bias; it should not be treated as evidence that the original support problem has disappeared.

Common Mistake