Bayesian Additive Distribution Regression
| Authors | Antonio R. Linero et al. |
| Year | 2026 |
| Field | AI / ML |
| arXiv | 2603.06462 |
| Download | |
| Categories | stat.ME, stat.ML |
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.
Engineering Breakdown
Plain English
This paper introduces DistBART, a method for predicting scalar outcomes from distribution-valued predictors—situations where you have grouped observations and need to predict outcomes based on the characteristics of entire distributions rather than individual data points. The authors model the regression function as a linear functional and assign it a Bayesian additive regression trees (BART) prior, leveraging the inductive bias that shallow tree ensembles work well for tabular data where predictions depend mainly on low-dimensional marginals. The approach is validated both empirically on synthetic and real datasets and theoretically, demonstrating that this framework effectively captures the relationship between distribution inputs and scalar outputs.
Core Technical Contribution
The core novelty is combining distribution regression with BART through a Riesz representation framework, where the regression function is expressed as a linear functional with its representer governed by a BART prior. This is non-trivial because standard BART operates on fixed-dimensional vector inputs, while distributions are infinite-dimensional objects; the Riesz representation elegantly reduces this to a finite-dimensional problem amenable to BART's tree-based structure. The key insight is that shallow tree ensembles encode useful inductive biases for tabular data, making them appropriate when the functional depends primarily on low-dimensional marginals rather than complex distributional features. Prior work either treated distributions naively or used complex nonparametric methods without leveraging the specific structure of tree-based models.
How It Works
The input is a set of grouped observations where each group has an associated distribution (e.g., a distribution over individual measurements within that group) and a scalar outcome. The method first represents the regression function f(·) as a linear functional via Riesz representation, meaning f(P) = ∫ r(x) dP(x) for some representer function r. Rather than directly estimating r from infinite-dimensional function space, DistBART places a BART prior on r, decomposing it as a sum of decision trees: r(x) = Σ T_m(x), where each tree T_m is shallow (typically depth 2-4). During inference, the algorithm updates the tree structure and leaf values using Bayesian methods, effectively learning which low-dimensional marginals matter most for prediction by selecting splitting variables in the trees. The output is a posterior distribution over possible representers r, enabling both point predictions and uncertainty quantification.
Production Impact
This approach directly addresses a common but underserved problem: predicting outcomes for grouped or aggregated data without access to individual-level features. In practice, this applies to scenarios like predicting customer lifetime value from purchase history distributions, forecasting treatment effects from patient population characteristics, or estimating loan default risk from credit score distributions within cohorts. Engineers adopting DistBART would replace naive approaches (e.g., computing hand-crafted summary statistics like mean/variance and feeding them to standard regression) with a principled method that automatically learns which distributional properties matter. The main trade-off is computational cost—BART requires iterative Bayesian inference and MCMC sampling, making it slower than point-estimate methods, though modern implementations are reasonably fast for typical tabular problem sizes (thousands to millions of samples). Integration into production requires careful handling of how to encode distributions (empirical samples, kernel estimates, or fitted parametric forms) and managing posterior sampling for real-time predictions.
Limitations and When Not to Use This
The paper assumes distributions can be represented and compared meaningfully, which breaks down if you lack sufficient samples from each distribution or if the distributions are high-dimensional with complex dependence structures where low-dimensional marginals are uninformative. The shallow tree assumption—that predictions depend primarily on marginal distributions rather than complex interactions or tail dependencies—may not hold for applications requiring sensitivity to distributional shape beyond first/second moments. The method inherits BART's limitations around high-dimensional tabular data and does not naturally handle very large feature spaces without careful feature engineering. Additionally, the paper's theoretical analysis appears incomplete in the abstract (it trails off mid-sentence), and real-world guidance on hyperparameter selection (tree depth, number of trees) and when to prefer this over alternatives remains unclear.
Research Context
This work builds on decades of BART research (Chipman et al., 2010) and recent advances in distribution regression, extending these ideas to a Bayesian nonparametric setting. It positions itself as an improvement over both naive approaches (summarizing distributions into fixed statistics) and complex alternatives like kernel methods or deep neural networks that may overfit on small-to-medium datasets. The paper likely benchmarks against standard regression baselines and other distribution regression methods on established datasets, though specific benchmark comparisons aren't detailed in the abstract. This opens research directions in combining other tree-based methods (random forests, XGBoost) with functional data analysis, extending to multivariate or distributional outputs, and developing scalable approximate inference for large-scale applications.
:::tip Subscribe Get weekly breakdowns of papers like this in AI Letters - the newsletter for engineers building production AI systems. :::