Balance-Guided Forest for Causal Effect Estimation
Abstract
Covariate balancing is widely used for estimating causal effects in observational studies. A central challenge is determining which covariate features should be balanced. Forest-kernel balancing provides a data-driven way to learn informative balancing features. Existing forest-kernel balancing methods first learn a tree-based representation and then construct balancing weights. However, this procedure creates a mismatch, as limited tree capacity may be spent on features that are already well represented. We propose Balance-Guided Forest (BGF), which uses the current balancing weights to guide subsequent tree construction. BGF grows trees using a novel splitting score that combines residual imbalance with an outcome contrast, favoring splits that separate outcome-relevant regions where imbalance remains. This couples representation learning with covariate balancing rather than treating balancing as a downstream step. We establish an oracle-type bound that links the iterative construction to weighting error and causal bias. Simulations exhibit improvements in causal estimation relative to existing forest-kernel approaches. Real-data applications illustrate the practical behavior of BGF across observational studies.
Related
- Graph-Split Bayesian Causal Forest for Spatial Heterogeneous Treatment Effect Estimation
- Feature Space Selection and Heterogeneous Effect Estimation for Blood-Brain Barrier Permeability: A Random Forest to the Generalized Random Forest Pipeline
- Network Effect Decomposition for Causal Effect Estimation under Network Interference
- Proximal Balancing for Causal Effect Estimation under Unmeasured Confounding
- Identification and Estimation of Causal Estimands with Missing Not at Random Data
BibTeX
@misc{yi2026balanceguided,
title = {{Balance-Guided Forest for Causal Effect Estimation}},
author = {Feiyang Yi and Jelena Bradic},
howpublished = {arXiv},
year = {2026},
month = nov,
doi = {10.48550/arxiv.2611.05723},
eprint = {2611.05723},
archivePrefix = {arXiv},
url = {https://arxiv.org/abs/2611.05723},
}
RIS
TY - UNPB
TI - Balance-Guided Forest for Causal Effect Estimation
AU - Feiyang Yi
AU - Jelena Bradic
PY - 2026
DA - 2026/11/03
PB - arXiv
DO - 10.48550/arxiv.2611.05723
UR - https://arxiv.org/abs/2611.05723
AB - Covariate balancing is widely used for estimating causal effects in observational studies. A central challenge is determining which covariate features should be balanced. Forest-kernel balancing provides a data-driven way to learn informative balancing features. Existing forest-kernel balancing methods first learn a tree-based representation and then construct balancing weights. However, this procedure creates a mismatch, as limited tree capacity may be spent on features that are already well represented. We propose Balance-Guided Forest (BGF), which uses the current balancing weights to guide subsequent tree construction. BGF grows trees using a novel splitting score that combines residual imbalance with an outcome contrast, favoring splits that separate outcome-relevant regions where imbalance remains. This couples representation learning with covariate balancing rather than treating balancing as a downstream step. We establish an oracle-type bound that links the iterative construction to weighting error and causal bias. Simulations exhibit improvements in causal estimation relative to existing forest-kernel approaches. Real-data applications illustrate the practical behavior of BGF across observational studies.
ER -