VaRT: Variational Regression Trees
Authors: Sebastian Salazar
Venue: NeurIPS 2023 spotlight (2023)
Keywords: Probabilistic Machine Learning
Variational Inference
Bayesian Inference
Bayesian Nonparametrics
TL;DR: We introduce a novel Non-parametric Bayesian model over the space of decision trees and derive a variational approximation to the posterior distribution.

Abstract: Decision trees are a well-established tool in machine learning for classification and regression tasks. In this paper, we introduce a novel non-parametric Bayesian model that uses variational inference to approximate a posterior distribution over the space of stochastic decision trees. We evaluate the model's performance on 18 datasets and demonstrate its competitiveness with other state-of-the-art methods in regression tasks. We also explore its application to causal inference problems. We provide a fully vectorized implementation of our algorithm in PyTorch.

VaRT: Variational Regression Trees
Sebastian Salazar Escobedo
Department of Computer Science
Columbia University New York, NY 10027
aebastian.salazaareee........aae columbia. edu
Abstract
Decision trees are well-established tool in machine learning for classification and regression tasks. In this paper, we introduce novel non-parametric Bayesian model that uses variational inference to approximate posterior distribution over the space of stochastic decision trees. We evaluate the model's performance on 18 datasets and demonstrate its competitiveness with other state-of-the-an methods in regression tasks. We also explore its application to causal inference problems. We provide fully vectorized implementation of our algorithm in PyTorch.
Introduction and Related Work
The use of binary trees for classification and regression problems is widespread in the Machine Learning community (see Breiman [200T] Criminisi 2011]). Chipman et [1998] proposed the Bayesian CART method, which offers fully Bayesian treatment of the problem, aiming to learn distribution over trees as opposed
single tree. This innovation led the development of more sophisticated algorithms, such as Bayesian Additive Regression Trees (BART) (see Chipman et al. (2010), which has stood the test of time and still shows competitive performance to this day. Both BART and Bayesian CART approximate the posterior distribution over trees through the use of Markov Chain Monte Carlo (MCMC) methods.
In the period following the inception of the Bayesian CART methodology, the field has witnessed significant advancements in approximate Bayesian Inference. This paper explores approximating posterior distribution over the space of trees via variational inference, compelling MCMC alternative based on an optimization problem as elaborated in [llii et al. [2017]. Our approach introduces unique stochastic process over the space of trees, serving as model prior. Notably, despite the infinite parameter specification of this stochastic process, we show that variational inference remains feasible.
This paper also related to work soft decision trees (Kontschieder et al. 20155) where the goal is to Iearn a stochastic decision tree of a fixed structure by minimizing some sort of loss function (see Section 2.Ip. In this paper, we adopt fully Bayesian perspective this problem and instead learn distribution over soft decision trees of varying structure. A fully Bayesian approach has the advantage in that allows to naturally quantify the uncertainty of the predictions made by our model.
Overall, our contributions in this paper include a novel stochastic process over the space trees, trees, and
variational inference algorithm for approximating the posterior distribution demonstration of the performance of our model on both regression and causal inference tasks. We believe that this work opens up avenues for future research in the area of Bayesian Non- parametrics and decision trees.
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

2 Model
2.1 Stochastic Decision Trees
A stochastic tree, denoted as T. is binary tree where each of its b - 1 internal nodes is associated with soft decision rule. This rule indicates the probability of particular data point taking the left branch of the tree. In this paper, we focus on decision rules of the form P(datapoint x at node takes the left branch) (3 X). Moreover, we introduce a prior proba- bility measure, denoted as II( T), that encompasses the space of trees T (defined by splits and decision rules) and the associated leaf parameters O The generative process for the data is divided into two parts:
1. Prior Distribution:
(7,0)~I((T,5
Unlike BART or Bayesian CART, our prior distribution designed differently, incorporating sampling algorithm that generates collection of parameters used to fully and uniquely reconstruct tree. The second part of the model the likelihood that gives rise to the observed data
2. Likelihood:
(yi X, T,0) ~ p(ui T)
(2)
The likelihood represents the conditional distribution of y/ given V4. T, and θ In case, it is usually modeled as a Normal distribution. Further details regarding the likelihood and the overall model will be diseussed in subsequent sections.
2.2 Likelihood
We associate each of the leafs of a stochastic tree with a parameter + = (01,-..,0p 0z}- Every parameter 0, θ specifies density over Vi denoted by f (,) Throughout, features x E X are fixed, and that the target variables y E y are sampled from a distribution whose parameters stochastically determined by having E X' traverse down the tree. In other words, if x E X' is assigned to leaf i, then E Y is sampled from f (t)) If a feature x E is assigned to leaf node with certain probability, this implies that the likelihood p(y) X T, O) is mixture distribution over the densities of (1]6;) and is given by
assume that the
plylT,
Σ P(x takes path leading to node (y 0+).
2.3 A Bayesian non-parametric process prior over the space of stochastic trees II(),
We define prior distribution over stochastic decision trees by specifying a sampling algorithm. The joint distribution over decision trees T and it's corresponding leaf parameters is specified through infinite number of parameters that have one-to-one correspondence with characteristics of a stochastic decision tree. The parameters required to draw a sample from the joint distribution over () which we denote by II,T, +) are as follows
1. fr E [o, 1 specify the probabilities of splitting a node in the tree. In other words, node in the tree is split with probability 7,
2. {u: are the parameters that define the distribution over the weights specifying the splitting rules.
3. (4) are the parameters that define the conditional likelihoods over the observations falling in node i, provided that it is leaf.
'The densities are relative to appropriate dominating measure v. In practice, v will usually be the Lebesgue the counting measure.
IIn this paper. we adopt the following convention for numbering the nodes of a binary tree: (a) We start by labeling the root node as 1 (b) recursively label the left children of node j with 2j and the right children of node j with 2j + L.

---

5. 0
5
5, 0
$1 ~ Bern(y))
Berm(y2) Bern(Ya)
Bern Bern(y))
3
P
02. Bg. Bg
B..Bs
Pa
3
2
fa
3
0, ExpFam
02
ExpFam((4
By
Figure I: Taking sample from the generative process
T
B
B
2d-1
Figure Graphical Model as- sociated to the data generating process. The parameters B, and are used to uniquely specify stochastic decision tree and its parameters.
Figure 2: Graphical Model sociated to the data generating process.
Figure 4: Variational approxi- mation to the data generating process.
Given these parameters, stochastic tree sampled as follows:
1. For each node node i E (..).
(a) Sample $1. .... Sk, ~ B Bernoulli>xx)
(b) Sample 31. .... Bk
ExpFam(4:)
(c) Sample 01. ..., O ExpFam (C:)
Note that the parameters (Si,..., $1. ...... suffice to completely specify binary stochastic tree. We give a concrete example of this generative process in Figure and the graphical model associated to this data-generating process is shown in
2.4 A variational family over the space of trees and tree parameters g(e,T)
variational family obtained by truncating Bayesian Nonparametric process described in section 2.3 at fixed tree depth d. This truncation is depicted graphical model in figure 4 similar pattern is observed in variational inference for Dirichlet Process Mixtures proposed by Blei and Jordan 22066), where the stickbbraking process of the Dirichlet Process mixture is also

---

truncated after fixed number of steps. is important to note that only the variational distribution is truncated, not the underlying model itself. Furthermore, we would like to emphasize that the choice of exponential family in the variational distribution aligns with those chosen for the prior distribution. To distinguish them from the prior parameters, we label the variational parameters with a tilde (i.e. Y, F4 () on top.
2.5
Maximizing the ELBO
We obtain an approximation of the posterior distribution by maximizing a weighted version of the evidence lower bound (ELBO) with respect the model parameters
ELBO(q) Ey (log p( (Y|, (KL((((,T)]ll((...)
(4)
(L) = Ey (log ((
(5)
+ AEEyy (log ((,,T)
(6)
XBy (log ()).
(7)
Where is hyperparameter controlling the amount of regularization. The gradients of the variational objective are approximated stochastically with Monte Carlo samples from the variational family. To
obtain low-variance estimates of the gradient estimates, we make use of the reparametrization trick
(see Kingma and Welling [2022), Ruiz et al. (2016), where a Softmax- Gumbeel approximation is
used for the Bernoulli random variables (Jang et al. (0077). Other approaches for variance-reduction are possible but these ideas were not explored in this paper (see Ranganath
like Bao-Blackwellizatti
ct al. 2.13). Given sample from the variational family (81...., 824-1, B1, B2a-1 1, 01, .,..., 024)
evaluate each of the terms in the ELBO as follows:
log (yy(( O, log
Σ Veeeaa
P(x takes a path leading node i),((wl0.)
(8)
log (gg,(TT))
Σ
log (1x)
Σ
log (expFam (0::aam(B::1) 1:))
iEInternal Nodes
Internal Nodes
Σ iELeaves
log (1 75)
Σ
log (expFam(O;;Gi))
(10)
LeLeaves
log g T)
Σ
log (Fx)
Σ
log (expFam(B::
(11)
&EInternal Nodes
(Elnternal Nodes
Σ
log (1 -
Σ
log (ex)Fam(O;;
(12)
Le not in last level
VELeaves
In order to make use of reparametrization gradients, each of the three terms in the ELBO must be written in terms the model parameters (i.e. (7, 14, Ç. 7, p. () ) and a reparametrization for each of
the samples must be specified. These caleulations are deferred to the next section.
3 Variational Regression Trees
In this paper, we explore the use of Variational Regression Trees. We propose using the normal distribution natural choice for the splitting and prediction vectors of the prior and variational families, as follows!
B, ~ N( St))
(13)
0, (Gi, W.))
(14)
Where S; and Vi are diagonal covariance matrices. We also assign normal distribution to each of the leaves in the tree by setting f (0.) to a normal distribution with parameters N/00" x, of). The normal distribution is natural choice for regression problems and allows for an easy application of
Whenever À 1.0, this is equivalent minimizing the KL divergence between the variational family and the posterior distribution over the space of trees. "This shows the parameters of the prior, the corresponding parameters for the variational family are denoted with the same letters but with tilde top.

---

the reparametrization trick. For classification, adding probit link function to the likelihood should result tractable variational inference algorithm. We believe including a logit link on the outcome y is also possible, but would require a Poly-Gamma augmentation (see Polson et al. (2013), ). We leave all of these modifications to future work (see section 5 for an overview on potential future work).
3.1 Learning the Variational Family
The parameters of the variational family are learned by minimizing Monte Carlo approximations of the ELBO with gradient descent. The naive Monte Carlo Estimators for this types of problem are given by V (e) (f())) = E, (f())V & log go (x)) Ri N Ei" f(x()) V a log qo (ax()). Although asymptotically consistent, this estimator exhibits extremely high variance, rendering it impractical for the purposes of stochastic optimization. Instead, we make use of a well-known duriance-reduction trick through a reparametrization of the samples obtained from the variational family. In particular, we rewrite samples from the variational function of an auxiliary noise variable as follows
e~ p(e)
(15)
x x= (0, e). The reparametrization of the normal distribution of equations 13 and 14 is straightforward and is given below
(16)
Grules; Gleaves Miid NO, Ta)
(17)
B = 5 S. Crutes + μ
(18)
8, 4 1/2 Cleaves + Ç.
(19)
Reparametrizing the Bernoulli random variables requires use of the Softmax- Hatmax--ummel trick of Jang et al. (017). The key insight relies on the fact that samples from Bernoulli distribution with success probability y may be obtained as follows
= one hot (arg max (9) + log(n)) 92 + +log(1 - -7))
(20)
(1,0)" v
(21)
Where g1.02 92 are independent samples from a Gumbel 1) distribution. In practice, we approximate equation 20 using the following limitting representation of the onehot argmax function
(91 + +og(?) 92 + log(1 -)))
lim softmax -
(22)
T
During training, we set T to a fixed value and then anneal it towards zero the training progresses. These reparametrizations allow us to use a single sample from the variational family to approximate the ELBO. To tie everything together and make use of the reparametrization trick, we need rewrite each of the terms in equations 8 to 12 solely in terms of the sample parameters (s, B, f) and the
distributional parameters (y, 7, u. S, p. S.C. V. č, W This is done by replacing the sums over the leaves and internal nodes with indicator functions, which may be written as products over the node splitting indicator variables functions for node i by 1/ and K; respectively, we may rewrite each of the terms in the ELBO as follows
(s1,..., 824 1). If we denote the internal and leaf node indicator
logp(0,77 T)
the (log (nx) + 1og (3,;
(23)
+
Ki (log (1 -x) (4))++00 log V.))
(24)
log (g, 7)
Σ the (log (fx) + log.N p4; Si))
(25)
26
+
Σ
Nc log (1 - -x) + Σ
Ky log.NNN((0 No4; V.,)
(26)
1=1
2h-1
[log éj
log ( X) log
Σ
(11))
BT (p))))))) f(y10,)
(27)
A=
5

---

Where th and are given by
Log
4/ Si
(p)
(28)
(1 8t)
I log
(1), 2 fl...., 1}
R2 Jog
(29)
Spa (i), E fb, 2b 20-11 -
3.2 Posterior Predictive Distribution and Bayesian Confidence Intervals
The fully Bayesian treatment of the problem enables us to obtain natural confidence intervals on our predictions, which is significant advantage of the Bayesian approach compared other black- box machine learning algorithms. Given a dataset D = { Mr Jiz confidence intervals for the predictions are obtained by approximating the predictive posterior distribution as follows:
x'))) /
p(y .,, 22,0DD. d') +)
(30)
p(T, (TT, [T, 0 (T,
(31)
/ ((T () [T. (T,,))
(32)
Equation 31 shows that the predictive posterior is mixture of the posterior and the likelihood. Thus, to obtain samples from the predictive posterior distribution, we sample a tree with parameters (O,T) from the posterior and and then use (O, 7) obtain a sample from the likelihood. Moreover, as we do not have access to the true posterior distribution, we use the learned variational family q(", +) approximation. For prediction purposes, estimate the posterior predictive mean using the mixture representation of equation 32
Experiments
This section presents comparison of the performance of Variational Regression Trees with three other widely used boosting algorithms: CatBoost, XGBoost, and RandomForest. A comparison with Bayesian Additive Regression Trees (BART, Chipman al. (2100) and Multilayer Perceptron (MCP, et al. (2013)) is deferred to the Appendix. We conducted experiments on 18 distinct Datasets from the UCI Machine Learning Repository benchmark our algorithm. In addition, explored the use of VaRT for estimating of conditional treatment effects in causal inference problems, motivated by the fact that BART is popular algorithm for such problems.
4.1 Performance Regression Datasets
The results of our regression experiments are shown Table We compare the performance of VaRT to three popular Boosting Algorithms: CatBoost, XGBoost, and RandomForest. We compared an ensemble fifty trees for each algorithm to single VaRT trees of depths 3,5,7, and 10. A comparison with BART and Multilayer Perceptrons is given in the Appendix.
In these experiments, VaRT was trained using gradient descent paired with a Clipped Adam optimizer in PyTorch (Paszke et al. [2019), Bingham al. (2018)). A single VaRT tree for depths 3, 5, 7, and 10 were trained for each dataset. We report the average RMSE values for VaRT across all depths to provide thorough representation of VaRT's performance across various configurations (Table A). The regularization parameters for all runs was set to 10-3 and no hyperparameter tuning was performed to make fair comparison the off-the-shelf performance of each of the algorithms. In practice, both the regularization and truncation parameters could be chosen through cross- valssdvaii.tion.. Early stopping by evaluating the RMSE training validation set everytime fixed number of epochs are completed should also increase performance.
Compared to other tree-based Bayesian algorithms, such BART, VaRT offers distinct advantages that enhance its practicality, Notably, VaRT's framework is inherently parallelizable, a feature that sets it apart from MCMC-based methods like BART which are inherently sequential and lack this

---

Dataset
CatBoost
RandomForest
XGBoost
VaRT3
VaRTs
VaRT7
VaRTIO
autompg
0.121
0.115
0.113
0.116
0.116
0.116
0.109
concreteslump
0.209
0.237
0,252
0.009
0.009
0.028
0.047
forest
0.396
0,403
0,403
0.392
0.394
0.387
0.387
solar
0.086
0.086
0.086
0.114
0.092
0.094
0.093
stock
0.074
0.067
0.072
0.065
0.068
0.065
0.059
yacht
0.036
0.026
0.018
0.041
0.035
0.035
0.046
airfoil
0.073
0.062
0.061
0.151
0.133
0.131
0.131
autos
0.092
0.072
0.076
0.121
0.146
0.134
0.140
breastcancer
0.403
0.362
0.398
0.390
0.476
0.380
0.360
concrete
0.095
0.093
0.081
0.149
0.141
0.158
0.148
fertility
0,402
0.375
0.358
0,481
0.454
0.456
0.387
housing
0.103
0.118
0.120
0.101
0.097
0.098
0.101
machine
0.089
0.112
0.097
0.110
0.123
0.124
0.138
pendulum
0.061
0.068
0.074
0.077
0.064
0.056
0.058
servo
0.066
0.062
0.050
0.096
0.077
0.078
0.109
skillcraft 0.118
0.116
0.120
0.114
0.114
0.114
0.114
sml wine
0.028
0.018
0.020
0.063
0.056
0.049
0.051
0.108
0.103
0.096
0.122
0.119
0.117
0.121
Table A: Comparison of RMSE values for different boosting algorithms and VaRT. The RMSE values correspond the average RMSE on the test of random 90/10% train-test split of the data over ten runs. Bold values indicate the lowest RMSE. Each of the ensemble methods were trained using fifty trees. VaRT's performance is shown for trees of depth 3, 5, 7, and 10.
VaRT
CatBoost
RandomForest
XGBoost
tamielectric (45781 datapoints)
0.508
0.503
0.545
0.504
kin40k (40000 datapoints)
0.099
0.100
0.097
0.110
Table B: RMSE values for different boosting algorithms and VaRT on datasets containing more than 40, 000 datapoints. The RMSE values correspond to the average RMSE on the test sets of a random 90/10% train-test split of the data over six runs. Bold values indicate the lowest RMSE, Each of the ensemble methods were trained using fifty trees, and variational family of depth 5 was used for VaRT.
parallelization potential. This parallelizability not only accelerates the model training process but also empowers efficient exploration of large datasets, making VaRT particularly well-suited for modern computation demands. Furthermore, VaRT leverages gradient-basee optimization techniques, making it compatible with widely used automatic differentiation engines such PyTorch. This was validated by our experiments on datasets with over 40,000 datapoints, where training took less than minutes.
We conducted all experiments on an ASUS Zephyrus G14 laptop with an RTX 2060 Max-Q GPU (6 GB VRAM), a 4900HS AMD CPU, 40 GB of RAM (although RAM usage was kept below 6 GB). Fach experiment took roughly 1-3 minutes per dataset. Code with the random seeds needed to replicate the results this section are provided with the supplementary material.
4.2 Causal Inference case-study
We investigate the performance of Variationial Regression Trees on several causal inference problems using synthetic and semi-synthetic data. Our experiments are similar to those conducted by Hill [201]). To introduce causal inference, we consider the potential outcomes that would have been observed for an individual under treatment and no treatment. We denote these quantities as (()) and (0),, respectively. The observed outcome denoted by and is related to the potential outcomes through the formula Y = AY(1) + (1 (O)( where A is a dichotomous variable indicating whether a particular individual the population recieved treatment. Our goal is to estimate different quantities that quantify the effect of treatment for various subpopulations.

---

we focus on the conditional average treatment effect (CATE), which can be identified through the following formula under standard causal inference assumptions, including overlap. conditional exchangeability. and consistency:
In section 4.2.1
(1Y(1) - Y x) = B(YIA= = x)) A 0, X =x).
(33)
In section 4.2.2 our goal is to estimate the Average Treatment Effect (ATE), which under standard causal inference assumptions can be given by:
((((1) (()) ( x, B(YIA=
(34)
While cannot provide comprehensive treatment of causal inference in this paper, refer interested readers to Pearl et al. 22016] M. 2023).
4.2.1 Experiment on Toy Data
We begin by generating two toy synthetic datasets using the following generative process
A Bernouli(0.5)
(35)
A a, (,,(,, 103) (1 a;) 102)
(36)
A aj X; = H; a,,\ e 1)+
(37)
In the "partially overlapping" setting (Figure 6). the means of the normal distributions used to generate the data are (H1, (22 = (40, 20), as in Hill 2011]. In contrast, in the "overlapping" setting. we set (H1, P2) (30, 40). The figures illustrating these hypothetical scenarios are Figures 5 and 6 In these figures, the red curve (top) represents the conditional expectation of the potential outcome under treatment (1), while the blue curve (bottom) represents conditional expectation of potential outcome under no treatment (()). To showcase the posterior predictive mean of VaRT for each potential outcome, we use dashed lines alongside 80% Bayesian confidence bands. These bands grow wider in areas where the density of the covariates is low and shrink in where the density of the covariates is high. This type of qualitative behavior may be attributed to the Bayesian prior, which acts as a type of regularization the fit of the model and helps us avoid overfitting.
Estimation Conditional Average Treatment Effect Curves with Bayesian Confidence
Average Treatment with 00 Dayesian Conf Eards
... (811 and Helf
givd
230
300
K:01 senvry (est cree
320
21D cenurty (no scree
3.0
210
310
meis liv unuerieeyy
cuertec
Hidd for uncottiity
Loe ouekee
even no aro mndi uncatality
35
Figure
Figure
4.2.2 Experiments on the Infant Health and Development Program (IHDP) Dataset
We the IHDP dataset simulate the two response surfaces of Hill [011] and then compare the performance VaRT relative to BART when estimating the average causal effect. The first response is linear and it is generated as follows
C0., 0.2, 0.15, 0.1, 0.05))
(38)
8, - 20,1, i (1,..., 25}
(39)
Y(0) ~ "IX"
(40)
Y(1) ~ ~X() +44,1)
(41)

---

Response Surface
VaRT
BART
Ground Truth
Linear Non-Linear
3.86
0.12
3.6 = 0.6
4.00
3.80
0.19
2.9 +0.8
4.00
Table Estimation of the ATE for response surfaces described by equations 38 through 45 The results report the mean ACE along with the standard deviation as the error across 5 runs. All of the VaRT trees were trained using a hyperparameter of The algorithm whose mean values were closer to the Grount Truth are shown in bold.
The second response surface is highly non-linear and is generated as follows
Cat((0.6, 0.1,1,0,,,,0111110,,,,,,
(42)
8; (0,.1, E (1,...; 25}
(43)
Y(0) ~ Xe(xx(XX* 1) (()) ~ ~N(XX 3
(44)
(45)
Here, w chosen so that the CATE is equal to 4 every (hence the ATE also 4). The results of these two experiments are reported in table and suggest that VaRT clearly outperforms BART in the task of Causal Effect Estimation on both response surfaces.
Conclusion and Future Work
In this paper, presented novel Bayesian non-parametra process over the space of stochastic decision trees and derived a variational approximation to We also proposed a variational inference algorithm approximate the Bayesian posterior distribution, which evaluated both regression and causal inference tasks. Our experiments showed that our model's performance is competitive with other state-of-the-an -frt methods in machine learning.
Looking ahead, future research could explore extensions of the non con-paraett Bayesian prior over trees to model forests, similar to what we observe in BART. Additionally, other potential lines of work could include modifying the prior to introduce dependencies among the tree parameters or extending the model to classification.
Overall, the findings of this research provide promising avenue for future research in Bayesian con-parametrics and decision trees. We hope that this work will inspire other researchers to build upon our model and improve its versatility and performance in various machine learning applications.

---

Aknowledgements
This research greatly benefited from discussions with Aaron Schein, Nihar Mauskar, and Ansaf Salleb-Aouissi. Their support was invaluable. Additionally, thank David Blei for his enlightening seminar on Applied Causal Inference, which inspired this paper. Special thanks Emily and Michael Teti for their unwavering support during the rebuttal process.
This research was partially supported by NIH/NLM (www.nlm.nhl nih niv. grant RO1LM01332
10

---

References
Lco Breiman. Machine Learning, 45 5-31) 2001. doi: //doi.org/10/ /055//00.02///111109334400333330
URL https
Antonio Criminisi. Decision forests: A unified framework for classification, regression, density estimation, manifold learning and semi-supervised learning. Foundations and Tremdss in Computer Graphics and Vision, 7(2-3)81-227, 2011. 10.. 1561/0600000035
doi: 10.1561/06000000055. URL httpps///doo org/
Hugh A. Chipman, Edward I. George, and Robert E. McCulloch. Bayesian cart model search. Journal of the American Statistical Association, 93(443):935-945 1998. ISSN 01621459. URL.
http:/ w/www/ jstor. org/ Settor.org//ttille
Hugh A. Chipman, Edward I. George, and Robert E. McCulloch, BART: Bayesian additive regression trees. The Annals of Applied Statistics, 4(1), mar 2010. doi: 10.1214009-auu2855. URL https
//doi. org. 10. 121472209 20852855
David M. Blei, Alp Kucukelbir, and Jon D. McAuliffe. Variational inference: A review for statisticians.
Journal of the American Statistical Association, 1 1126118) apr 2017. doi: 10. 1080/
URL httppp/://b....og///.......///// wttp://ww/111000000/000000016660
1285773
Peter Kontschieder, Madalina Fiterau, Antonio Criminisi, and Samuel Rota Bulo. Deep neural decision forests. In 2015 IEEE International Conference on Computer Vision (ICCV). IEEE,
December 2015.
doi: 10. 1105/.cc URL https://doi. merii.org/000 1109///c..////.....o/o 2015
172
David M. Blei and Michael I. Jordan. Variational inference for Dirichlet process mixtures, Bayesian Analysis, 1(1):121 143, 2006. doi: 1..1214006-BA044 URL https://doi. org/10. 1214/
06- BA104
Diederik P Kingma and Max Welling. Auto-encoding variational bayes, 2022.
Francisco J. R. Ruiz, Michalis K. Titsias, and David M. Blei. The generalized reparameterization gradient, 2016.
2017. Rajesh Ranganath, Sean Gerrish, and David M. Blei. Black box variational inference, 2013.
Eric Jang. Shixiang Gu, and Ben Poole. Categorical reparameterization with
Nicholas G. Polson, James G. Scott, and Jesse Windle. Bayesian inference for logistic models using polya-gamma latent variables, 2013.
Yann LeCun, Yoshua Bengio, and Geoffrey Hinton. Deep learning. Nature, 521(7553):436 May 2015. doi: 10. 1038/nature 14539, URL httpp:////// org/10..///..
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Köpf, Edward Yang. Zach Vito,o, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang. Junjie Bai, and Soumith Chintala. Pytorch: An imperative style, high-performanee deep learning library, 2019.
Bingham, Jonathan P. Chen, Martin Jankowiak, Fritz Obermeyer, Neeraj Pradhan, Theofanis Karaletsos, Rohit Singh, Paul Szerlip, Paul Horsfall, and Noah D. Goodman. Pyro: Deep universal probabilistic programming, 2018
Jennifer Hill. Bayesian nonparametric modeling for causal inference. Journal of Computational and Graphical Statistics, 20:217-240 03 2011. doi: 10. 10.1198/jgs.2010.0886
Judea Pearl, M Maria Glymour, and Nicholas P. Jewell. Causal inference in statistics: A primer. 2016.
MIGUEL A.. ROBINS JAMES M., HERNAN. Causal inference: What if. CRC PRESS, 2023.

---

A
Comparison with BART and MLP
VaRT
BART
MLP
Median
top 10%
bottom 10%
50 trees
(100,50) neurons
autompg (N=392, d=7)
0.0051
0.0008:000031
0.085+0.0005
0.0710=0.0020
0.088:100013
concreteslump
0.049=0.023
0.034::0.017
0.069:0.0031
0.149=0.011
0.16:10.10
(N=103, d=7)
forest
0.205
0.194
0.219:10.008
0.9777:0000026
(N=517, d=12)
solar
0.114++0.006
0.000400.004
0.124+-0.013
0.1038 + 0.0004
0. 0.0044
(N=1066, d=10)
stock
0.0051
0.0336+000000
0.040=0.0055
0..004
0.060
(N=566, d=11)
yacht
0.059-100007
0.05340.007
0.06910.012
0.068410.0018
0.080:10.004
(N=308, d=6)
airfoil
0.0026
0.0776+000000
0.0894#0.0029
0.087
0.092+00.008
(N=1503, d=5)
autos
0.007
0.0000
0.221
0.075 0.006
0.093++0.004
(N=159, d=25)
breastcancer
0.244
0.220+000014
0.203++0.02
0.219+
0.205+
(N=194, d=33)
concrete
0.0837=0.0013
0.078610.0000
0.0922:10.0025
0.0768=0.001-
0.085:10.005
(N=1030, d=8)
fertility
0.197
0. 08100:00.018
0.214
0.1601100005
0.158:10.013
(N=100, d=9)
housing
0.074=00005
0.068:00004
0.082+000006
0.069=0.006
0.077=00006
(N=506, d=13)
machine
0.003
0.083+00.009
0.104+0.009
0.0695+0.0013
0.14++0.04
(N=209, d=7)
pendulum
0.000
0.004
0.073+00009
0.0697+0.0016
0.072++0.006
(N=630, d=9)
servo
0.1051:0.013
0.09010.012
0.127
0.082710.0034
0.110
(N=167, d=4)
skillcraft
0.0092
0.06761000007
0.0738+0.0031
0.06991000007
0.0692+1.0025
(N=3338, d=19) sml
0.035610.0029
0.0342+0.0022
0.04110.004
0.041610.0009
0.038410.0033
(N=4137, d=26)
wine
0.0746+0.0022
0.0713+0.0020
0.0806+0.0034
0.0947+0.0031
+000015
(N=1599, d=11) TOTAL
8
8
Table 2: Root Mean-Squared Error (RMSE) for VaRT, BART, and Feedforward Neural Network on 18 regression datasets from the UCI Machine Learning Repository. For VaRT we show the 10th and 90th Bayesian quantiles along with the standard deviation reported as the error across 5 runs. For the purposes of comparison, we use the median RMSE of VaRT to determine whether it outperformed the other algorithms. The RMSE of the best performing algorithm is shown in bold.
*: Indicates that the RMSE among the top 10% draws of VaRT's posterior predictive distribution beats (or overlaps) with the RMSE of the corresponding algorithm.
BART was trained using fifty trees and the off-the-shelf prior parameters of Chipman al. 22110). The Markov Chain BART was given a burn-in period of 200 iterations (i.e. the default choice). Finally, the Feed-forward Neural Network (MLP) had two hidden layers with 100 and 50 neurons respectively, and was trained using an Adam optimizer.
12

---

WIT fit for precense
pmaccis soden
100 -8001 2erat
-1
petarte prodtis Wall
Figure Experiments on toy data generated from a piecewise continuous function. The data was generated from uniform mixture of Gaussians centered on the midpoint of each interval. Analysis was conducted on 100 (left) and 1000 (right) datapoints. Qualitatively it seems like the posterior concentrates around the "true" clusters as we collect more data.
13


Reviews:
[{"review_id": "BN6v0kP8KU", "reviewer": "Reviewer_AUTX", "summary": "The authors propose to use variational inference to train regression trees. The innovation lies in using variational inference as the optimization process rather than Markov Chain Monte Carlo. The authors demonstrate this method on 18 ML UCI problems and some causal inference and toy problems. ", "questions": "The authors mention depth 5-7; what happens at smaller depths? A smaller tree would greatly aid interpretability.", "limitations": "there is no limitations section, although the authors discuss limitations.", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The paper is clearly written and specific about its contribution. The figures are generally good, although they could use better captions (especially Figure 1). ", "weaknesses": "It's hard for me to judge the novelty and potential impact of this work. There are so many decision tree boosting algorithms out there that are blazingly fast (e.g. XGBoost, Catboost), whereas this method takes a few minutes to train. There could be an advantage if the resultant model was more interpretable than others, but that isn't really demonstrated. It is not clear to this reviewer whether BART is really a go-to, off the shelf ML algorithm like XGBoost has become, but the authors focus their side by side comparisons on this algorithm (which makes sense since it is perhaps the closest philosophically). In other words, the comparison would be stronger if a candidate boosted tree algorithm was included in the comparisons. MLPs are known to underperform most tree ensembles on UCI. ", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely appreciate the time and effort you've dedicated to reviewing our manuscript. Your thorough evaluation offers valuable insights that will undoubtedly enhance the quality and impact of this work.\n\nRegarding the concern on training time, we acknowledge that there are blazingly fast boosting algorithms available in the field. Your observation regarding the underperformance of MLPs on tabular datasets is also well-founded. To address these concerns and establish the feasibility and performance of our method across a diverse range of scenarios, we conducted a series of additional experiments. These experiments involved comparing the performance of our approach to well-established boosting algorithms, namely XGBoost, CatBoost, and Random Forest. We are excited to share that the performance of VaRT remained competitive and, in some cases, even outperformed these boosting algorithms across a diverse set of situations. We have included the results of these new experiments in Table A of the attached PDF for your reference and want to thank you for pointing out this gap in our evaluation. \n\nTo showcase the scalability of our method, we also ran experiments on two new datasets containing more than 40,000 datapoints. Despite the increased data volume, the training time remained within the order of one to three minutes (see attached pdf on \"global\" rebuttal). This notable efficiency is attributed to the parallelizability of our approach, which offers a clear computational advantage over other Bayesian algorithms reliant on Markov Chain Monte Carlo.\n\nWe hope that these additional experiments and insights provide a more comprehensive view of the potential impact of our work. Thank you once again for your valuable feedback and insights. Your suggestions have been instrumental in shaping our approach to further align with the expectations and demands of the field. We remain open to answering any other questions.\n\nBest,\n\nThe Authors. "}, {"review_id": "lGHm5B0crx", "reviewer": "Reviewer_MwKc", "summary": "The paper develops a new generating process for soft decision trees. The proposed process is then adopted as a prior model in a Bayesian nonparametric regression setting, and a novel variational inference algorithm is developed using the truncated version of the tree generating process as the variational distribution. Experiments suggest that the proposed method has comparable performance with BART in real data sets and causal inference applications.", "questions": "1. I think it would be beneficial to provide more comparison between the proposed VaRT and conventional Bayesian decision tree models like BART and Bayesian CART, such as:\n\n    1.1 One well-known advantage of BART is that it can perform variable selection to some extend. Does the proposed model have similar capability?\n\n    1.2 What will be the computation benefits of the proposed method compared with the MCMC approach for conventional Bayesian decision tree models?\n\n    1.3 In what scenario will the proposed VaRT outperform/underperform conventional Bayesian decision tree models? \n\n2. Questions/comments on the experiments:\n\n    2.1 Details of the experiments should be included. For example, is the RMSE in Table 1 computed on the test split? If so, how do you obtain the training/test splits from the data sets?\n\n    2.2 In Table 1, I think it will be more fair to also compare with the 0.9 and 0.1 posterior percentiles of BART.\n\n    2.3 Very limited discussion on the experiment results, especially for the UCI datasets.\n\n3. Minor issues:\n\n    3.1 In (3), is the denominator equal to 1? If so, the denominator can be trivially omitted.\n\n    3.2 In Line 74, it is more rigorous to write \"$s_k \\sim \\text{Bernoulli}(\\gamma_k)$ independently for $k = 1, 2, \\ldots$\".\n\n    3.3 Citation in Line 191 does not seem correct.\n", "limitations": "Some limitations and future work are discussed in the paper. I don't foresee any potential negative societal impact.", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "* A good addition to the literature of Bayesian soft decision tree models.\n* The soft decision trees generating process and the variational inference algorithm is technically sound and novel.\n* Competitive performance with BART.", "weaknesses": "* Some important baselines are missing. The proposed method is not compared with other soft decision tree based models such as soft BART (Linero and Yang, 2018; Linero, 2022) and the frequentist version of soft decision trees.\n* Experiments are only based on some default hyperparameters for both VaRT and the baselines; hyperparameter sensitivity is not investigated.\n* The paper can benefit from careful proofreading. For instance, a few notations are used without any definition. Formatting and presentation of the paper can also be improved.\n\nReference:\n\n[1] Linero, A.R. and Yang, Y., 2018. Bayesian regression tree ensembles that adapt to smoothness and sparsity. Journal of the Royal Statistical Society Series B: Statistical Methodology, 80(5), pp.1087-1110.\n\n[2] Linero A.R. 2022. SoftBart: Soft Bayesian Additive Regression Trees. arXiv preprint arXiv:2210.16375.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We would like to express our gratitude for taking the time to review our manuscript. Your insightful comments and suggestions have been instrumental in guiding our work towards refinement, and we appreciate your thorough evaluation.\n\nRegarding your recommendation to evaluate our method against these models and other soft decision tree-based approaches, we want to clarify our approach. In the interest of providing a thorough evaluation within our current constraints, we chose to focus on established boosting methods like XGBoost, Random Forest, and CatBoost when conducting new experiments. The comparisons we've conducted in this rebuttal with widely recognized boosting methods offer insights into the strengths and performance of our approach. We believe these comparisons contribute valuably to our research and the broader field of Bayesian nonparametric regression. \n\nWe fully understand your concern about the comparability of our method and acknowledge the challenges associated with implementing additional benchmarks and locating or implementing specific Python versions of certain models in such a short time-frame. Despite these challenges, we remain committed to addressing your valuable feedback in a meaningful way. In our revised manuscript, we will incorporate a detailed discussion of soft BART models and other relevant soft decision tree-based models. This will underscore the significance of these models in the context of Bayesian soft decision trees and contribute to a comprehensive understanding of the research landscape.\n\nAddressing your specific questions, we do believe our model possess the capacity for variable selection. We believe this could be achieved by incorporating global shrinkage priors into the splitting and prediction rules of VaRT. This clever observation was brought to our attention by reviewer Ya4w. We provide further clarification on this issue in the \"global\" rebuttal section but we remain open to answer any lingering questions.\n\nReviewer 2ow5 suggested that we ran experiments on larger datasets. We took their advice and extended our evaluation to include datasets comprising more than 40,000 data points. Notably, our method demonstrated seamless scalability to handle these larger datasets, owing to its inherent parallelizability. The computational efficiency of our algorithm was evident, with fitting times spanning a range of one to three minutes. It's worth noting the significant contrast with Bayesian MCMC methods like BART, which are constrained by their sequential nature and lack support for parallelization. While we acknowledge the existence of rapid frequentist boosting algorithms such as XGBoost and CatBoost, we recognize the potential for further exploration in this direction. We are excited about the prospect of future advancements in this area and acknowledge that the initial implementation of our algorithm might not outpace these established algorithms in terms of speed.  \n\nMoreover, we hope to have addressed your concerns on the experiments in the \"global\" rebuttal section. In particular, we have included a thorough discussion on our original and new experiments. Please feel free to reach out with any questions and are open to any feedback you may have. We intend to incorporate these discussions into our paper. \n\nLastly, we appreciate your keen attention to detail, which led to the identification of the denominator's equivalence to 1 in equation (3). We have diligently revised our manuscript to reflect this clarification. Your suggestions in points 3.2 and 3.3 have also contributed to improving the clarity of our paper, and we have successfully incorporated these changes.\n\nOnce again, we extend our heartfelt thanks for your invaluable feedback. Your expertise and insights have been instrumental in shaping the trajectory of our work, and we are dedicated to refining our manuscript in accordance with your suggestions.\n\nRespectfully,\n\nThe Authors. \n\n\n"}, {"review_id": "gSEXV9SogH", "reviewer": "Reviewer_Ya4w", "summary": "This paper introduces a non-parametric Bayesian model taking the form of a stochastic decision tree and proposes to approximate the posterior distribution with variational inference. The prior involves three parts, a prior on the tree structure, a prior on the probabilistic splitting criteria, and a prior on the conditional distribution at each leaf. To conduct variational inference, reparameterization tricks of normal and binary variables are used. Prediction can be made based on the posterior predictive distribution and inference can be made with Bayesian credible intervals. Numerical experiments are supportive, including a comparison with benchmarks on prediction and application to toy datasets as well as real datasets.", "questions": "1. Since the data generating model could be viewed as a mixture model, the posterior distribution of the tree also gives a posterior distribution over the clusters of data. Does this posterior distribution over the clusters concentrate around the true clusters when the true data generating model is a mixture model?\n\n2. As mentioned in the paper, the variational regression tree can be generalized to variational classification tree with a binary response variable. How does it generalize to classification of response variables with more than two levels? On a side note, for doing variational inference instead of Gibbs sampling, why would Polya-Gamma augmentation still be needed?\n\n3. The soft splitting criterion is based on a logistic regression of the covariates, so the splitted regions is no longer rectangular. While this may be easier for implementation, how does it compare to the case where the soft splitting criterion are all restricted to be parallel to the axis (a hierarchical prior on beta could do this)? Also, for the obtained posterior samples, does the splitting critria at different nodes tend to be more parallel or more orthogonal?", "limitations": "Not much, see Weaknesses and Questions.", "rating": 8, "confidence": 5, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "1. The paper is clear and well written.\n2. The proposed method is strong in originality.\n3. Numerical experiments are supportive.\n4. It addresses the long-lasting problem that Bayesian tree methods (Bayesian CART or BART) are infeasible for larger datasets and their MCMC chains may never really converge. The combination of stochastic decision trees with variational inference suggests a promising direction in handling this problem.", "weaknesses": "There doesn't appear to be any major weaknesses to me. The following are some minor ones that could be improved:\n1. For Section 4.1, since the task is purely predictive, it would be better to compare with aggregated decision trees or random forest methods as well. Or alternatively, if the comparison is restricted purely to Bayesian models, a Bayesian neural network should be used with parameters inferred either with variational inference or HMC sampler.\n2. It would be nice to compare the posterior distribution approximated by variational inference with the true posterior distribution. For instance, with low dimensional covariates and response, a data augmented Gibbs sampler (with binary latent variables at each splitting point) along with reversible jumps could also be designed and used to sample from the true posterior distribution.\n3. There seems to be too much parameters introduced with gamma, mu, zeta. Some structural formulations such as a cumulative shrinkage prior with just a few prior parameters might be better.\n\nOf course, all of these would require lots of additional effort, so as a first paper in this direction the lack of them is not a problem.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely appreciate your exceptional review of our manuscript. Your insightful feedback has significantly impacted our research, and we are grateful for your valuable insights. Your evaluation not only validates our efforts but also provides us with invaluable perspectives that will undoubtedly shape the trajectory of our paper and future work, regardless of the review outcome.\n\nYour observation regarding the comparison with other frequentist methods is well-taken. To address this concern, we conducted additional experiments comparing the performance of VaRT against established ensemble methods like Random Forests, XGBoost, and CatBoost. Furthermore, we expanded our experiments to include two datasets, each containing over 40,000 points, to address scalability concerns raised by other reviewers. We are excited to report that our method exhibited seamless scalability and comparable performance to other boosting algorithms (see table A). As you rightfully pointed out, we believe methods based on Variational Methods offer a distinctive advantage due to their inherent parallelizability. This is in stark contrast to other Bayesian methods based on Markov Chain Monte Carlo. \n\nYour suggestion regarding exploring a data augmented Gibbs sampler for posterior distribution comparison is insightful. We are currently investigating similar analyses in the literature to gain a better understanding of this approach. If you have any specific references you could suggest, it would greatly aid our exploration.\n\nWe were also particularly intrigued by your mention of cumulative shrinkage priors. Building on this, we plan to explore the application of cumulative shrinkage priors and their implications for our model in future work. We would greatly appreciate your suggestions on relevant references in this regard! (were you thinking of hierarchical priors, like hourshoe priors, on the splitting/leaf nodes to induce sparsity?) \n\nAddressing your specific questions, we conducted experiments on toy data generated from a piecewise continuous function. The data was generated from a uniform mixture of Gaussians centered on the midpoint of each interval. Our results suggest that VaRT correctly \"routes\" each point to the appropriate node with high probability, as evidenced by our analysis on 100 and 1000 datapoints. This suggests that the posterior concentrates around the \"true\" clusters.  \n\nRegarding the use of Pólya-Gamma augmentation, while it may not be strictly necessary, we have encountered cases where Pólya-Gamma random variables have been integrated into Variational Inference schemes to get closed-form updates of the natural gradients. Notably, the paper \"Efficient Gaussian Process Classification Using Pólya-Gamma Data Augmentation\" by Wenzel et al. discusses a similar application. However, we now recognize that their use is not strictly necessary and will explicitly note this in our paper and acknowledge your insight.\n\nFor generalizing our method to more than two response variables, we propose incorporating a matrix β ∈ ℝ^(c × d), attaching a categorical distribution to each leaf nodes as follows $\\text{Categorical}(\\text{SoftMax}(βX))$ and employ discrete reparameterization techniques to fit variational family.\n\nIn response to your question about hierarchical priors and orthogonal splitting criteria, we believe hierarchical priors, such as a horseshoe priors on the variances of β coefficients, could indeed induce sparsity and enhance model interpretability as it would bias the model towards axis-alignted splits. We believe this is an exciting area for future work! Currently our method does not have any particular bias for this and most of the splits end up being non-axis-aligned. \n\nThank you once again for your detailed review! Your feedback has undoubtedly enriched our understanding of our work. Regardless of the review outcome, we intend to acknowledge your valuable contributions in the acknowledgments section of our paper. We'll remain available to answer any other questions you may have during the discussion period. \n\nSincerely,\n\nThe Authors"}, {"review_id": "xuPBLH78Ub", "reviewer": "Reviewer_2ow5", "summary": "The authors propose a non-parametric Bayesian model for Regression trees. Variational inference is used approximate the posterior distribution. Extensive experiments are conducted to show the superiors of the method. ", "questions": "1, what is function g(\\phi,\\epsilon) in equation (16)?\n2, what is the symbol \\beta_{pa_k} in equation (27)?\n3, Are the RMSEs reported in table 1 of VaRT obtained using a single Tree? Please provides some discussions to explain why VaRT with a single Tree outperforms the BART with 50 Trees. It is interesting to provide the prediction results of  CART and random forest.\n", "limitations": "When compared with CART, the proposed method seems time-consumming", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "The authors proposed a non-parametric method to learn regression trees. The paper is well-written and easy to follow. The experimental results on 18 datasets are provided. ", "weaknesses": "1, The basic idea is too simple. The authors proposed several prior distributions for the variables in tree. Then the parameters in the prior distributions are optimized using SGD with some reparametrization tricks. The priors as well as the reparametrization tricks have been proposed already. So the contribution of this paper is limit.\n2, No introduction is provided regarding the 18 datasets. No discussion is provided for the prediction results. The highlighted RMSEs in table 1 are not always the lowest.\n3, The datasets used in 4.1 is too small. It is interesting to show the performance on large scale dataset. The computational complexity of the proposed method seems to be high. \n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Firstly, we extend our gratitude for your assessment of our paper. Your feedback has been instrumental in shaping our work and we value the insights you've shared. We acknowledge your observation about the apparent simplicity of our approach, which centers around non-parametric Bayesian models for regression trees. While the foundational concept of using tree-based models for regression might seem straightforward, we wish to emphasize that the novelty lies in the intricate interplay of the proposed prior distributions, the design of the variational family, and the utilization of these elements within the broader context of Variational Inference to approximate the posterior distribution. We believe that this unique amalgamation contributes to the field by providing a methodologically distinct approach to tackling the challenges associated with Bayesian regression tree modeling.\n\nWe firmly believe that this innovation enriches the field by providing an unparalleled framework for addressing the intricacies associated with regression tree modeling in the context of Bayesian models. As of now, we are not aware of any prior research that encompasses the distinctive amalgamation of techniques featured in our method. Given your expertise in the field, we would greatly appreciate any guidance you can provide in identifying references that might shed light on parallel developments.\n\nIn your review, you noted the application of a reparametrization trick that has been introduced elsewhere in the literature. We acknowledge this point, and we do not claim originality in this aspect. In fact, we explicitly mention its prior introduction in our manuscript (lines 94-96). Our emphasis was on demonstrating how this established technique fits seamlessly into our broader framework.\n\nWe also acknowledge your concerns about the choice of datasets and the absence of discussions regarding prediction results in our paper. We hope to have addressed some of these concerns in the \"global\" rebuttal and remain open to answering any questions you may have. Regarding the prediction results, our intention was to provide transparent evaluation. The notation in Table 1, indicated by \"∗,\" aimed to highlight instances where VaRT's posterior predictive distribution excelled or overlapped with other algorithms. This approach is aligned with our commitment to unbiased reporting.\n\nIn response to your suggestion regarding the dataset sizes employed in our experiments, we took your advice into consideration and extended our evaluation to include datasets comprising more than 40,000 data points. Notably, our method demonstrated seamless scalability to handle these larger datasets, owing to its inherent parallelizability. The computational efficiency of our algorithm was evident, with fitting times spanning a range of one to three minutes. It's worth noting the significant contrast with Bayesian MCMC methods like BART, which are constrained by their sequential nature and lack support for parallelization. While we acknowledge the existence of rapid frequentist boosting algorithms such as XGBoost and CatBoost, we recognize the potential for further exploration in this direction. We are excited about the prospect of future advancements in this area and acknowledge that the initial implementation of our algorithm might not outpace these established algorithms in terms of speed.\n\nRegarding your specific questions, the function $g(\\phi,\\epsilon)$ is the function utilized in the reparametrization trick. While we adopted this notation from (Kingman and Welling, 22'), we regrettably omitted its explicit definition in our manuscript. We appreciate your attention to this detail and will rectify this oversight by including a dedicated line that provides a clear definition of this function within our paper.\n\nFurthermore, we understand that the notation $\\beta_{pa_{k}(i)}$ might have caused confusion. To clarify, this notation represents the $\\beta$ coefficient associated with the $k^{th}$ parent of node $i$ in the tree. We recognize the potential for ambiguity and will address this concern by adding a clarification in our paper to ensure proper understanding.\n\nRegarding the observation that a single tree occasionally outperforms an ensemble of trees, we recognize that this might appear unexpected. The distinction lies in our approach's capacity to endow individual trees with the ability to capture intricate data relationships. This capability is rooted in our algorithm's inherent adaptability, facilitated by more complex splitting and prediction rules. Moreover, as we learn a distribution over the space of trees, each posterior sample can yield a different prediction, resulting in an ensemble-like effect. On average, this allows our model to encapsulate inherently intricate functions. This nuanced aspect likely contributes to instances where our singular VaRT posterior distribution over trees showcases superior performance compared to other boosting methods reliant on ensembles of trees. We acknowledge the importance of accentuating this distinctive attribute within our paper to provide a comprehensive context for the observed outcomes.\n\nWe want to highlight that we really value your critical insights and are eager to address any concerns you may have to ensure the accuracy and quality of our work. We believe that your perspective will greatly contribute to enhancing our manuscript, and we hope that this clarifying response provides a more accurate portrayal of our contributions and intentions. Thank you once again for your thorough review.\n\nRespectfully,\n\nThe Authors. \n\n\n\n\n\n\n\n"}]

Author rebuttal:
We extend our heartfelt gratitude to the reviewers for their insightful evaluations of our paper. The feedback we received has played a pivotal role in the evolution of our work, and we have approached our revisions with utmost seriousness and dedication.

One area that emerged as a significant opportunity for improvement was the lack of experimental comparison with frequentist tree-based boosting methods, a concern raised by reviewers AUTX, Ya4w, and MwKc. As emphasized by reviewer AUTX, our work draws philosophical inspiration from BART, which initially led us to benchmark our algorithm against it. However, we now recognize the necessity of providing a more inclusive comparison. In light of this, we have expanded our experiments to encompass three well-regarded boosting methods: Random Forests, XGBoost, and CatBoost. We are excited to share that even with the introduction of these additional benchmarks, our algorithm remains competitive with these state-of-the-art boosting methods. The results of these experiments are showcased in Table A of the attached pdf. 

In the new experiments of Table A, we compared an ensemble of fifty trees of BART, Random Forests, XGBoost, and CatBoost to a single tree of our algorithm at depths of 3, 5, 7, and 10. Because of our commitment to transparency and rigorous evaluation, we have decided to report the results across all four runs of VaRT, as opposed to just reporting the best one. By adopting this approach, we aim to provide a thorough representation of VaRT's performance across various configurations, showcasing its consistency and versatility in diverse scenarios. 

Our proposed method showcases notable strengths in comparison to established boosting methods as evidenced by the RMSE values presented in Table A. Across a diverse range of datasets, VaRT consistently demonstrates competitive performance, often outperforming or closely aligning with well-regarded boosting techniques such as CatBoost, RandomForest, and XGBoost. The versatility and effectiveness of VaRT's single-tree approach is evident from its ability to yield compelling results on various datasets, highlighting its potential to provide accurate predictions in different problem domains. 

While our experiments were comprehensive, a couple of reviewers rightfully poinited out a lack of depth in our result discussions. To provide context, we employed a 90%/10% train-test split on 18 standard regression datasets from the UCI Machine Learning Repository. For fair comparison, features were standardized within (0,1) and (-1,1) ranges for Table 1 and Table A, respectively. The reported RMSE values in both tables are averaged from five test set runs, accompanied by their standard deviations. 

Closer examination of Tables 1 and A reveals noticeable performance challenges for VaRT arise in datasets 'airfoil,' 'autos,' and 'sml.' This prompts us to scrutinize dataset attributes influencing VaRT's performance in these datasets. As we analyzed the experimental results, we observed variations in VaRT's performance across different datasets. While the majority of the datasets demonstrated strong performance, we noted that the 'autos' and 'airfoil' datasets posed unique challenges. Moreover it is worth pointing out that the 'sml' dataset constitutes a time-series task, which introduces complexities that may not align optimally with VaRT's current framework. These challenges have highlighted specific areas that require further investigation and improvement. 

In comparison to other tree-based Bayesian algorithms, VaRT exhibits distinct advantages that enhance its practical utility. Notably, VaRT's framework is inherently parallelizable, a feature that sets it apart from MCMC-based methods like BART which are inherently sequential and lack this parallelization potential. This parallelizability not only accelerates the model training process but also empowers efficient exploration of large datasets, making VaRT particularly well-suited for modern computation demands. Furthermore, VaRT leverages gradient-based optimization techniques, making it compatible with widely used automatic differentiation engines such as PyTorch. We corroborated this by running new experiments on datasets with 40k+ datapoints where training took < 3 mins (see Table B on the attached pdf).

Reviewer Ya4w gave us with a couple of phenomenal suggestions that have greatly enriched our perspective on our work. One suggestion that particularly piqued our interest was the incorporation of cumulative shrinkage priors into the splitting and regression parameters. We believe this has the potential to significantly enhance the performance and interpretability of VaRT. As these priors are known for inducing sparsity in the posterior distribution, we believe that their addition could also address a question raised by reviewer MwKc on variable selection, as sparse solutions should lead to a model that uses only a subset of the features. We believe this to be a very promising and exciting avenue for future work.  

Moreover, reviewer Ya4w's insight into viewing our model as a mixture model has been enlightening. We conducted experiments using simulated data from a mixture of Gaussians and found that the posterior distribution of our model indeed concentrates around the "true" clusters of the underlying data-generating model (see attached pdf). We hope our interpretation of this insight aligns with the reviewer's Ya4w intention, and we remain open to addressing any questions.

In conclusion, we extend our heartfelt appreciation to the reviewers for their invaluable insights and feedback that have significantly contributed to refining our paper. Your comments have been instrumental in shaping the trajectory of our research and enhancing its quality. In this regard, we will provide a detailed response to each of these points in the "individual" rebuttal section. We sincerely thank you for your time, expertise, and thoughtful guidance.


Meta-review: This paper is in the same family of methods as BART, a very famous Bayesian cousin of the famous CART decision tree learning algorithms. Reasoning and working over the space of probability distributions over trees is not easy. The authors propose a variational inference based method to efficiently sample the posterior distribution, propose an efficient parallel implementation, and provide solid experimental comparisons. My main request is that, for the camera-ready version, the authors should include the additional experimental material included in their rebuttal.
Decision: Accept (spotlight)