Dirichlet-Tree Allocation Priors for Bayesian Causal Forests
Abstract
Bayesian Causal Forests (BCF) estimate heterogeneous treatment effects with a separate ensemble of treatment trees. The prior on the leaf values of the treatment trees determines whether sampling variability is attributed to the treatment effect. Independent shrinkage priors exert no control over how a tree allocates variability across its leaves, and their total prior variance grows as leaves are added. We propose the Dirichlet-Tree Bayesian Causal Forest, which jointly regularizes the leaf values of each treatment tree through a Dirichlet-Tree prior on their variance shares, generated by stick-breaking along the tree topology. We characterize the spike and tails of the marginal leaf prior, show that the prior adaptively bounds the number of regions of effect heterogeneity, and establish posterior consistency for the forest. In simulations, the model outperforms BCF when heterogeneity is confined to a few regions and reduces spurious heterogeneity elsewhere, and, unlike independent shrinkage priors, its accuracy does not deteriorate as the trees grow. We also evaluate the model on real data, including a pancreatic cancer survival study.
Related
- Bayesian Neural Networks via Dirichlet Process Priors on the Data-Generating Distribution for Randomized Value Functions in Deep Reinforcement Learning
- Factorial Multivariate Bayesian Causal Forests: heterogeneous main and interaction effects of multiple treatments on correlated outcomes
- Bayesian Causal Forests for Zero-Inflated Continuous Outcomes
- Proximal R2 Priors for Structured Variance Allocation
- Sieve-free ReLU Bayesian priors with near-minimax posterior contraction
BibTeX
@misc{aguilar2026dirichlettree,
title = {{Dirichlet-Tree Allocation Priors for Bayesian Causal Forests}},
author = {Javier Enrique Aguilar and Daniel Klippert and Yusufhan Balci and Alexander Marx},
howpublished = {arXiv},
year = {2026},
month = oct,
doi = {10.48550/arxiv.2610.26728},
eprint = {2610.26728},
archivePrefix = {arXiv},
url = {https://arxiv.org/abs/2610.26728},
note = {CC BY 4.0},
}
RIS
TY - UNPB
TI - Dirichlet-Tree Allocation Priors for Bayesian Causal Forests
AU - Javier Enrique Aguilar
AU - Daniel Klippert
AU - Yusufhan Balci
AU - Alexander Marx
PY - 2026
DA - 2026/10/19
PB - arXiv
DO - 10.48550/arxiv.2610.26728
UR - https://arxiv.org/abs/2610.26728
AB - Bayesian Causal Forests (BCF) estimate heterogeneous treatment effects with a separate ensemble of treatment trees. The prior on the leaf values of the treatment trees determines whether sampling variability is attributed to the treatment effect. Independent shrinkage priors exert no control over how a tree allocates variability across its leaves, and their total prior variance grows as leaves are added. We propose the Dirichlet-Tree Bayesian Causal Forest, which jointly regularizes the leaf values of each treatment tree through a Dirichlet-Tree prior on their variance shares, generated by stick-breaking along the tree topology. We characterize the spike and tails of the marginal leaf prior, show that the prior adaptively bounds the number of regions of effect heterogeneity, and establish posterior consistency for the forest. In simulations, the model outperforms BCF when heterogeneity is confined to a few regions and reduces spurious heterogeneity elsewhere, and, unlike independent shrinkage priors, its accuracy does not deteriorate as the trees grow. We also evaluate the model on real data, including a pancreatic cancer survival study.
ER -