Bayesian Additive Distribution Regression

TL;DR

This paper introduces DistBART, a Bayesian nonparametric method for distribution regression that models the relationship between distributions and responses using Bayesian additive regression trees, with theoretical and empirical validation.

Contribution

It proposes a novel Bayesian additive distribution regression model with theoretical guarantees and scalable inference techniques, connecting tree-based methods to kernel approaches.

Findings

DistBART performs well on synthetic and real data.

Theoretical results show adaptive posterior concentration.

Scalable inference is achieved via random-feature approximation.

Abstract

Distribution regression, where the goal is to predict a scalar response from a distribution-valued predictor, arises naturally in settings where observations are grouped and outcomes depend on group-level characteristics rather than on individual measurements. We introduce DistBART, a Bayesian nonparametric approach to distribution regression that models the regression function as a linear functional with the Riesz representer assigned a Bayesian additive regression trees (BART) prior. We argue that shallow decision tree ensembles encode reasonable inductive biases for tabular data, making them appropriate in settings where the functional depends primarily on low-dimensional marginals of the distributions. We show this both empirically on synthetic and real data and theoretically through an adaptive posterior concentration result. We also establish connections to kernel methods, and use this connection to motivate variants of DistBART that can learn nonlinear functionals. To enable scalability to large datasets, we develop a random-feature approximation that samples trees from the BART prior and reduces inference to sparse Bayesian linear regression, achieving computational efficiency while retaining uncertainty quantification.