Title: VaRT: Variational Regression Trees
Venue: NeurIPS 2023
Accepted: yes
Keywords: Probabilistic Machine Learning

Variational Inference

Bayesian Inference

Bayesian Nonparametrics
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

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.

1 Introduction and Related Work

The use of binary trees for classification and regression problems is widespread in the Machine Learning community (see Breiman [2001], Criminisi [2011]). Chipman et al. [1998] proposed the Bayesian CART method, which offers a fully Bayesian treatment of the problem, aiming to learn a distribution over trees as opposed to a single tree. This innovation led to 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 a posterior distribution over the space of trees via variational inference, a compelling MCMC alternative based on an optimization problem as elaborated in Blei et al. [2017]. Our approach introduces a unique stochastic process over the space of trees, serving as a model prior. Notably, despite the infinite parameter specification of this stochastic process, we show that variational inference remains feasible.

This paper is also related to work on soft decision trees (Kontschieder et al. [2015]), where the goal is to learn a stochastic decision tree of a fixed structure by minimizing some sort of loss function (see Section 2.1). In this paper, we adopt a fully Bayesian perspective to this problem and instead learn a distribution over soft decision trees of varying structure. A fully Bayesian approach has the advantage in that it allows us 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 of trees, a variational inference algorithm for approximating the posterior distribution over trees, and a demonstration of the performance of our model on both regression and causal inference tasks. We believe that this work opens up new avenues for future research in the area of Bayesian Nonparametrics and decision trees.

2 Model

2.1 Stochastic Decision Trees

A stochastic tree, denoted as T , is a binary tree where each of its b -1 internal nodes is associated with a soft decision rule. This rule indicates the probability of a 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 i takes the left branch ) = σ ( β T i X ) . Moreover, we introduce a prior probability measure, denoted as Π(Θ , T ) , that encompasses the space of trees T (defined by splits and decision rules) and the associated leaf parameters Θ . The generative process for the data is divided into two parts:

1. Prior Distribution:

Unlike BART or Bayesian CART, our prior distribution is designed differently, incorporating a sampling algorithm that generates a collection of parameters used to fully and uniquely reconstruct a tree. The second part of the model is the likelihood that gives rise to the observed data

2. Likelihood:

The likelihood represents the conditional distribution of y i given x i , T , and Θ . In our case, it is usually modeled as a Normal distribution. Further details regarding the likelihood and the overall model will be discussed in subsequent sections.

2.2 Likelihood

We associate each of the b leafs of a stochastic tree T with a parameter Θ := { θ 1 , ..., θ b } . Every parameter θ i ∈ Θ specifies a density over Y , denoted by f ( y | θ i ) . 1 . Throughout, we assume that the features x ∈ X are fixed, and that the target variables y ∈ Y are sampled from a distribution whose parameters stochastically determined by having x ∈ X traverse down the tree. In other words, if x ∈ X is assigned to leaf i , then y ∈ Y is sampled from f ( y | θ i ) . If a feature x ∈ X is assigned to leaf node i with a certain probability, this implies that the likelihood p ( y | X,T, Θ) is a mixture distribution over the densities f ( y | θ i ) and is given by

2.3 A Bayesian non-parametric process prior over the space of stochastic trees Π(Θ , T )

We define a prior distribution over a stochastic decision trees by specifying a sampling algorithm. The joint distribution over decision trees T and it's corresponding leaf parameters Θ is specified through an infinite number of parameters that have a one-to-one correspondence with characteristics of a stochastic decision tree. The parameters required to draw a sample from the joint distribution over ( T, Θ) -which we denote by Π( T, Θ) are as follows 2

{ γ i } ∞ i =1 ∈ [0 , 1] ∞ specify the probabilities of splitting a node in the tree. In other words, node i in the tree is split with probability γ i .

{ µ i } ∞ i =1 are the parameters that define the distribution over the weights specifying the splitting rules.

{ ζ i } ∞ i =1 are the parameters that define the conditional likelihoods over the observations falling in node i , provided that it is a leaf.

S6=0

S3=1

S=0

S1 ～Bern(1)

S2 ～Bern(2)

S6 ～Bern(6)

S3～Bern()

S7～Bern()

3

3

b

β

β

02,06,07

2

β3

3

0～ExpFam(i)

2

β3

Y

βi～ExpFam(μi)

b

b

6

T

y

x

N

s

β

θ

γ

µ

ζ

∞

y

x

N

Given these parameters, a stochastic tree is sampled as follows:

For each node node i ∈ { 1 , 2 , ... }

(a) Sample s 1 , ..., s k , ... ∼ Bernoulli ( γ i )

(b) Sample β 1 , ..., β k , ... ∼ ExpFam ( µ i )

(c) Sample θ 1 , ..., θ k , ... ∼ ExpFam ( ζ i )

Note that the parameters ( s 1 , ..., β 1 , ..., θ 1 , ... ) suffice to completely specify a binary stochastic tree. We give a concrete example of this generative process in Figure 1 and the graphical model associated to this data-generating process is shown in 3.

2.4 A variational family over the space of trees and tree parameters q (Θ , T )

A variational family is obtained by truncating the Bayesian Nonparametric process described in section 2.3 at a fixed tree depth d . This truncation is depicted as a graphical model in figure 4. A similar pattern is observed in variational inference for Dirichlet Process Mixtures proposed by Blei and Jordan [2006], where the stick-breaking process of the Dirichlet Process mixture is also

s

β

θ

˜

γ

˜

µ

˜

ζ

2

d

-

1

2

d

y

x

N

truncated after a fixed number of steps. It 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. ˜ γ, ˜ µ, ˜ ζ ) 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 to the model parameters: 3

Where λ is a 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-Gumbel approximation is used for the Bernoulli random variables (Jang et al. [2017]). Other approaches for variance-reduction like Rao-Blackwellization are possible but these ideas were not explored in this paper (see Ranganath et al. [2013]). Given a sample from the variational family ( s 1 , ..., s 2 d -1 , β 1 , ..., β 2 d -1 , θ 1 , ..., θ 2 d ) , we evaluate each of the terms in the ELBO as follows:

In order to make use of reparametrization gradients, each of the three terms in the ELBO must be written in terms of the model parameters ( i.e. { γ , µ , ζ , ˜ γ , ˜ µ , ˜ ζ } ) and a reparametrization for each of the samples must be specified. These calculations 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 as a natural choice for the splitting and prediction vectors of the prior and variational families, as follows: 4

Where S i and Ψ i are diagonal covariance matrices. We also assign a normal distribution to each of the leaves in the tree by setting f ( y | θ i ) to a normal distribution with parameters N ( θ T i x, σ 2 i ) . The normal distribution is a natural choice for regression problems and allows for an easy application of

the reparametrization trick. For classification, adding a probit link function to the likelihood should result in a 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 ∇ ϕ E q ϕ ( x ) ( f ( x )) = E q ϕ ( f ( x ) ∇ ϕ log q ϕ ( x )) ≈ 1 N ∑ N i =1 f ( x ( i ) ) ∇ ϕ log q ϕ ( x ( i ) ) . 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 variance-reduction trick through a reparametrization of the samples obtained from the variational family. In particular, we rewrite samples from the variational as a function of an auxiliary noise variable as follows

The reparametrization of the normal distribution of equations 13 and 14 is straightforward and is given below

Reparametrizing the Bernoulli random variables requires use of the Softmax-Gumbel trick of Jang et al. [2017]. The key insight relies on the fact that samples from a Bernoulli distribution with success probability γ may be obtained as follows

Where g 1 , g 2 are independent samples from a Gumbel (0 , 1) distribution. In practice, we approximate equation 20 using the following limitting representation of the onehot argmax function

During training, we set τ to a fixed value and then anneal it towards zero as 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 to rewrite each of the terms in equations 8 to 12 solely in terms of the sample parameters ( s, β, θ ) and the distributional parameters ( γ, ˜ γ, µ, S, ˜ µ, ˜ S, ζ, Ψ , ˜ ζ, ˜ Ψ) . 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 s = ( s 1 , ..., s 2 d -1 ) . If we denote the internal and leaf node indicator functions for node i by ψ i and κ i respectively, we may rewrite each of the terms in the ELBO as follows

Where ψ i and κ i are given by

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 a significant advantage of the Bayesian approach compared to other blackbox machine learning algorithms. Given a dataset D = { ( x i , y i ) } N i =1 , confidence intervals for the predictions are obtained by approximating the predictive posterior distribution as follows:

Equation 31 shows that the predictive posterior is a mixture of the posterior and the likelihood. Thus, to obtain samples from the predictive posterior distribution, we sample a tree with its parameters (Θ , T ) from the posterior and and then use (Θ , T ) to 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 ( T, Θ) as an approximation. For prediction purposes, we estimate the posterior predictive mean using the mixture representation of equation 32.

4 Experiments

This section presents a 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 et al. [2010]) and a Multilayer Perceptron (MLP, LeCun et al. [2015]) is deferred to the Appendix. We conducted experiments on 18 distinct Datasets from the UCI Machine Learning Repository to benchmark our algorithm. In addition, we explored the use of VaRT for estimating of conditional treatment effects in causal inference problems, motivated by the fact that BART is a popular algorithm for such problems.

4.1 Performance on Regression Datasets

The results of our regression experiments are shown in Table 2. We compare the performance of VaRT to three popular Boosting Algorithms: CatBoost, XGBoost, and RandomForest. We compared an ensemble of 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 et 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 a 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 a fair comparison on the off-the-shelf performance of each of the algorithms. In practice, both the regularization and truncation parameters could be chosen through cross-validation. Early stopping by evaluating the RMSE on a training or validation set everytime a fixed number of epochs are completed should also increase performance.

Compared to other tree-based Bayesian algorithms, such as 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

Table A: Comparison of RMSE values for different boosting algorithms and VaRT. The RMSE values correspond to the average RMSE on the test sets of a 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 .

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. This was validated by our experiments on datasets with over 40,000 datapoints, where training took less than 3 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). Each experiment took roughly ∼ 1 -3 minutes per dataset. Code with the random seeds needed to replicate the results of 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 [2011]. 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 Y (1) and Y (0) , respectively. The observed outcome is denoted by Y and is related to the potential outcomes through the formula Y = AY (1) + (1 -A ) Y (0) , where A is a dichotomous variable indicating whether a particular individual in the population recieved treatment. Our goal is to estimate different quantities that quantify the effect of a treatment for various subpopulations.

In section 4.2.1, 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.2, our goal is to estimate the Average Treatment Effect (ATE), which under standard causal inference assumptions can be given by:

While we cannot provide a comprehensive treatment of causal inference in this paper, we refer interested readers to Pearl et al. [2016], M. [2023].

4.2.1 Experiment on Toy Data

We begin by generating two toy synthetic datasets using the following generative process

In the "partially overlapping" setting (Figure 6), the means of the normal distributions used to generate the data are ( µ 1 , µ 2 ) = (40 , 20) , as in Hill [2011]. In contrast, in the "overlapping" setting, we set ( µ 1 , µ 2 ) = (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 Y (1) , while the blue curve (bottom) represents the conditional expectation of the potential outcome under no treatment Y (0) . 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 areas 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 on the fit of the model and helps us avoid overfitting.

EstimationofConditionalAverageTreatmentEffectCurves

EstimationofConditionalAverageTreatmentEffectCurves

140

Y(0) pred VaRT

with80%BayesianConfidenceBands

140

Y(0) pred VaRT

with80%BayesianConfidenceBands

Y(1) pred VaRT

Y(1) pred VaRT

130

Y(0)

130

Y(0)

Y(1)

Y(1)

X(0) density [not to scale]

X(0) density [not to scale]

120

X(1) density [not to scale]

120

X(1) density [not to scale]

110

110

OutcomeY

Outcome Y

100

100

90

06

High overlap area

High overlap area

80

leads to low uncertainty

08

leads to lowuncertainty

70

Low overlap

leads to higl

70

Low overlap area

leads to high uncertainty

09

10

20

30

40

50

60

70

09

10

20

30

40

50

09

70

Covariates X

Covariates X

4.2.2 Experiments on the Infant Health and Development Program (IHDP) Dataset

We use the IHDP dataset to simulate the two response surfaces of Hill [2011] 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

The second response surface is highly non-linear and it is generated as follows

Here, ω is chosen so that the CATE is equal to 4 for every X (hence the ATE is also 4 ). The results of these two experiments are reported in table 1 and suggest that VaRT clearly outperforms BART in the task of Causal Effect Estimation on both response surfaces.

5 Conclusion and Future Work

In this paper, we presented a novel Bayesian non-parametric process over the space of stochastic decision trees and derived a variational approximation to it. We also proposed a variational inference algorithm to approximate the Bayesian posterior distribution, which we evaluated on both regression and causal inference tasks. Our experiments showed that our model's performance is competitive with other state-of-the-art methods in machine learning.

Looking ahead, future research could explore extensions of the non-parametric 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 a promising avenue for future research in Bayesian non-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.

6 Aknowledgements

This research greatly benefited from discussions with Aaron Schein, Nihar Mauskar, and Ansaf Salleb-Aouissi. Their support was invaluable. Additionally, I thank David Blei for his enlightening seminar on Applied Causal Inference, which inspired this paper. Special thanks to Emily and Michael Teti for their unwavering support during the rebuttal process.

This research was partially supported by NIH/NLM (www.nlm.nih.gov) grant R01LM013327.

References

Leo Breiman. Machine Learning , 45(1):5-32, 2001. doi: 10.1023/a:1010933404324. URL https: //doi.org/10.1023/a:1010933404324 .

Antonio Criminisi. Decision forests: A unified framework for classification, regression, density estimation, manifold learning and semi-supervised learning. Foundations and Trends® in Computer Graphics and Vision , 7(2-3):81-227, 2011. doi: 10.1561/0600000035. URL https://doi.org/ 10.1561/0600000035 .

Hugh A. Chipman, Edward I. George, and Robert E. McCulloch. Bayesian cart model search. Journal of the American Statistical Association , 93(443):935-948, 1998. ISSN 01621459. URL http://www.jstor.org/stable/2669832 .

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.1214/09-aoas285. URL https: //doi.org/10.1214%2F09-aoas285 .

David M. Blei, Alp Kucukelbir, and Jon D. McAuliffe. Variational inference: A review for statisticians. Journal of the American Statistical Association , 112(518):859-877, apr 2017. doi: 10.1080/ 01621459.2017.1285773. URL https://doi.org/10.1080%2F01621459.2017.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.1109/iccv.2015.172. URL https://doi.org/10.1109/iccv.2015. 172 .

David M. Blei and Michael I. Jordan. Variational inference for Dirichlet process mixtures. Bayesian Analysis , 1(1):121 - 143, 2006. doi: 10.1214/06-BA104. 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.

Eric Jang, Shixiang Gu, and Ben Poole. Categorical reparameterization with gumbel-softmax, 2017.

Rajesh Ranganath, Sean Gerrish, and David M. Blei. Black box variational inference, 2013.

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-444, May 2015. doi: 10.1038/nature14539. URL https://doi.org/10.1038/nature14539 .

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 DeVito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala. Pytorch: An imperative style, high-performance deep learning library, 2019.

Eli 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.1198/jcgs.2010.08162.

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

Table 2: Root Mean-Squared Error (RMSE) for VaRT, BART, and a Feedforward Neural Network on 18 regression datasets from the UCI Machine Learning Repository. For VaRT we show the 10 th and 90 th 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 et al. [2010]. The Markov Chain of 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.

VaRT Fit for apiecewise mixture of Gaussians

VaRT Fit for a piecewise mixture of Gaussians

posterior predictive median

posterior predictive mean

A(5,62)f(x) = 0.005 *x² +1,x ≤ 2

(5,4²)f(x)=-0.003 *x,2<x≤8

D

(12.42)fx)= -0.025 *x +4,8<× ≤15

{20,4²)f(x) =0.01*x,x >15

-2

-3

posterior predictive median

posterior predictive mean

10

Q

10

20

30

-20

-i0

Q

10

20

OE

X

Peer reviews:
comments: 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.

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.
2, 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.
3, 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.

QUESTIONS:
1, what is function g(\phi,\epsilon) in equation (16)?
2, what is the symbol \beta_{pa_k} in equation (27)?
3, 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.

LIMITATIONS:
When compared with CART, the proposed method seems time-consumming
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-06

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n1. 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.\n\nWEAKNESSES:\nThere 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.\n\nQUESTIONS:\n1. 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?\n\nLIMITATIONS:\nNot much, see Weaknesses and Questions.", "TITLE": "Official Review", "RECOMMENDATION": 8, "REVIEWER_CONFIDENCE": 5, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-07"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\n* 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.\n\nWEAKNESSES:\n* 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.\n\nQUESTIONS:\n1. 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\nLIMITATIONS:\nSome limitations and future work are discussed in the paper. I don't foresee any potential negative societal impact.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-07"}

comments: 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.

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.

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.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-07