📊 statistics

Forest Kernel Balancing Weights: Outcome-Guided Features for Causal Inference

This paper proposes Forest Kernel Balancing, a causal inference method that leverages the implicit kernels of tree-based models like random forests and BART to select outcome-predictive features for covariate balancing, thereby achieving superior computational and statistical performance compared to standard kernel methods that ignore outcome information.

Original authors: Andy A. Shen, Eli Ben-Michael, Avi Feller, Luke Keele, Jared Murray

Published 2026-08-19

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Andy A. Shen, Eli Ben-Michael, Avi Feller, Luke Keele, Jared Murray

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

In the world of scientific research, understanding cause and effect is often a game of comparison. To know if a new medicine works, or if a specific training program improves skills, researchers ideally run a randomized experiment. In such a study, participants are assigned by chance to either receive the treatment or not, creating two groups that are virtually identical in every way except for the intervention. This randomness ensures that any difference in the outcome is truly due to the treatment. However, in many real-world situations, from studying the impact of education policies to understanding the long-term effects of conflict, researchers cannot assign treatments by chance. They must rely on observational data, where people choose their own paths or are assigned by circumstances beyond the researcher's control.

When groups are not formed by chance, they often differ in important ways before the study even begins. One group might be older, wealthier, or more educated than the other. These pre-existing differences, known as confounders, can distort the results, making it look like a treatment caused an outcome when it was actually just a reflection of the group's background. To fix this, statisticians use a technique called balancing. They assign weights to individuals in the study, essentially giving more importance to some people and less to others, so that the treated and untreated groups look statistically similar on paper. The challenge lies in deciding exactly which characteristics to balance. If researchers only compare simple, straight-line differences, they might miss complex, hidden relationships that drive the results. If they try to balance too many complicated features, the math can become unstable and unreliable.

A team of researchers has proposed a new way to solve this puzzle by borrowing a tool from machine learning. They suggest using the hidden patterns found in decision trees to guide the balancing process. In their approach, they first train a computer model on the data to predict the outcome. This model, built from many small decision trees, learns which combinations of factors are most important for predicting what happens to a person. The researchers then use this learned structure to define a map of similarity between individuals. Instead of just asking if two people have the same income or age, the model asks if they would end up in the same final category if they were sorted through the decision trees. This method, which the authors call forest kernel balancing, allows the data itself to reveal which complex relationships matter most for the outcome.

The researchers tested this idea through a series of computer simulations and real-world case studies. In their simulations, they created artificial data where the relationship between the treatment and the outcome was complex and non-linear, meaning it did not follow a simple straight line. They compared their new method against traditional approaches that rely on fixed mathematical formulas to measure similarity. The results showed that the forest kernel method was significantly better at reducing error and bias. It managed to find the correct answer more often than the standard methods, particularly when the data contained tricky, non-linear patterns that the old formulas missed. The study suggests that by letting the outcome guide the search for important features, researchers can create more accurate comparisons without needing to guess in advance which complex interactions to look for.

To see if this worked in the real world, the team applied their method to two distinct scenarios. The first involved a study of a mathematics tutoring program. In this case, researchers had data from a group of students who were randomly assigned to receive math or vocabulary training, providing a perfect "gold standard" answer to compare against. They then tried to replicate this result using only the observational data from students who had chosen their own training paths. When the researchers used standard methods to balance the groups, especially after transforming the data to make the relationships more complex, their estimates drifted away from the true answer. However, when they used the forest kernel method, their estimates stayed much closer to the known truth, successfully recovering the experimental result.

The second application looked at the long-term effects of child soldiering in Uganda. Researchers wanted to understand how being abducted by a rebel group affected a young person's years of schooling. Because abduction was not random, the researchers had to carefully balance the characteristics of those who were abducted against those who were not. Using the new method, they found that the impact of abduction on education was slightly larger than what traditional methods suggested. More importantly, the new method proved to be more stable. When the researchers changed the number of complex features they included in the analysis, the traditional method produced wildly different answers, while the forest kernel method remained consistent. This stability suggests that the new approach is better at distinguishing between features that truly matter for the outcome and those that are just random noise.

The core insight of this work is that the best way to compare two groups is not to rely on a fixed rulebook of what to compare, but to let the data teach the researcher what is important. By using the structure of machine learning models to define similarity, the researchers created a tool that adapts to the specific complexities of the problem at hand. This does not mean the method is perfect or that it solves every problem in observational research, but it offers a powerful new way to handle the messy, non-linear realities of human behavior. The findings indicate that when researchers are willing to use outcome information to guide their search for balance, they can uncover more reliable truths about cause and effect, even when they cannot run a controlled experiment.

Technical Summary: Forest Kernel Balancing Weights

Problem Statement
In observational causal inference, achieving covariate balance between treated and control groups is essential for unbiased estimation of treatment effects (e.g., the Average Treatment Effect on the Treated, or ATT). While weighting methods like Inverse Propensity Score Weighting (IPW) and balancing weights aim to re-weight the control group to match the treated group, a persistent challenge is selecting the correct set of features to balance. Traditional approaches often rely on raw covariates or pre-specified transformations (e.g., polynomials), which may fail to capture complex nonlinearities and interactions critical for predicting outcomes and addressing confounding.

Kernel balancing offers a data-driven alternative by balancing a low-dimensional summary of a kernel matrix, which implicitly represents an infinite set of features. However, existing kernel balancing methods are "design-based": they construct kernels using only baseline covariates, ignoring outcome information. Consequently, these kernels may capture complex covariate relationships that are irrelevant to the outcome, failing to adequately address confounding mechanisms that drive the treatment effect.

Methodology: Forest Kernel Balancing
The authors propose Forest Kernel Balancing, a method that leverages outcome information to learn the kernel matrix used for balancing. The core insight relies on the fact that tree-based machine learning models—specifically Random Forests (RF) and Bayesian Additive Regression Trees (BART)—implicitly define a kernel based on the co-occurrence of observations in the same terminal leaf node.

The proposed procedure involves the following steps:

Outcome-Guided Kernel Construction: Instead of using a fixed design-based kernel (e.g., Gaussian or polynomial), the method trains a tree ensemble (RF or BART) on a pilot sample (consisting of control units for ATT estimation) to predict the outcome. The resulting model defines a similarity metric (kernel) where two units are similar if they frequently share leaf nodes across the ensemble. This ensures the kernel captures feature representations that are predictive of the outcome.
Sample Splitting (Cross-Fitting): To preserve the design-based principle of observational studies (where outcomes should not influence the design phase) and prevent overfitting, the data is split. The pilot sample trains the model; the analysis sample (remaining controls and all treated units) is used to compute the kernel matrix, solve for weights, and estimate the treatment effect. This process is repeated via cross-fitting to reduce variance.
Spectral Approximation: Given that the forest kernel induces a high-dimensional feature space, the method employs spectral truncation. It computes the eigendecomposition of the kernel matrix and balances the top $r$ principal components (scaled by singular values) rather than the full feature set.
Hybrid Balancing: The authors advocate for balancing a combination of the forest-derived kernel features and the raw covariates. To ensure neither dominates the objective function, they propose a "block normalization" strategy where raw covariates and kernel principal components are scaled to have equal aggregate variance.
Weight Estimation: The balancing weights are estimated by minimizing a convex objective function that balances the weighted difference in the feature space (imbalance) against the dispersion of the weights (variance), similar to standard linear balancing weights but applied to the forest kernel features.

Key Contributions

Outcome-Guided Feature Learning: The paper introduces a framework where the kernel used for balancing is learned from outcome data via tree-based models, ensuring that the features balanced are those most relevant to predicting the outcome and addressing confounding.
Robustness to Transformations: Unlike Gaussian or polynomial kernels, which are sensitive to marginal transformations of covariates (e.g., log-transforms), forest kernels are largely invariant to such transformations because they rely on the relative ordering of splits in the tree structure.
Integration of Machine Learning and Causal Design: The work bridges the gap between machine learning (specifically the implicit kernels of RF and BART) and causal inference design, demonstrating how "honest" sample splitting can allow the use of outcome data for feature selection without violating causal identification assumptions.

Results

Simulation Studies: Extensive simulations demonstrate that forest kernel balancing reduces both bias and Root Mean Squared Error (RMSE) compared to:
- Balancing on raw covariates alone.
- Balancing on design-based kernels (Gaussian) alone.
- Balancing on a combination of raw covariates and design-based kernels.
  The improvements are most pronounced when the number of principal components is moderate and when the data involves complex nonlinearities. The method remains robust across varying covariate dimensions ( $q \in \{10, 50, 100\}$ ).
Empirical Applications:
- Mathematics Tutoring (Within-Study Comparison): In a study where a logarithmic transformation of covariates was applied to distort the linear basis, forest kernel balancing successfully recovered the experimental benchmark (the true treatment effect). In contrast, methods using raw covariates or design-based Gaussian kernels failed to recover the benchmark after the transformation, highlighting the forest kernels' invariance and ability to capture outcome-relevant structure.
- Child Soldiering in Uganda: In an observational study of the effects of forced youth military service, forest kernel balancing produced estimates that were more stable across varying numbers of principal components compared to design-based kernels. The forest-based estimates suggested a larger negative effect on educational attainment, and the stability of these estimates across hyperparameter choices suggested they were capturing consistent confounding mechanisms.

Significance and Claims
The paper claims that forest kernel balancing offers a meaningful computational and statistical improvement over standard kernel methods by incorporating outcome information into the feature selection process without compromising the validity of causal inference. By leveraging the implicit kernels of RF and BART, the method automatically identifies nonlinearities and interactions that are pertinent to the outcome, thereby better addressing confounding. The authors position this work as a middle ground between pure outcome modeling and pure weighting, suggesting it as a promising direction for future research in doubly robust estimation and sensitivity analysis for unobserved confounding. The authors explicitly state that their approach improves upon existing methods in scenarios where covariate-outcome relationships are complex and not well-captured by linear or standard design-based kernels.

Technical Summary: Forest Kernel Balancing Weights

More like this