SlowGuess/NeurIPS_2023_OCR huggingface dataset:
We investigate the performance of Variational 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:
$$
\mathbb {E} (Y (1) - Y (0) | X = x) = \mathbb {E} (Y | A = 1, X = x) - \mathbb {E} (Y | A = 0, X = x). \tag {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:
$$
\mathbb {E} (Y (1) - Y (0)) = \mathbb {E} (\mathbb {E} (Y | A = 1, X) - \mathbb {E} (Y | A = 0, X)). \tag {34}
$$
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
$$
A _ {i} \sim \text {B e r n o u l l i} (0. 5) \tag {35}
$$
$$
X _ {i} \mid A _ {i} = a _ {i} \sim a _ {i} \mathcal {N} \left(\mu_ {1}, 1 0 ^ {2}\right) + \left(1 - a _ {i}\right) \mathcal {N} \left(\mu_ {2}, 1 0 ^ {2}\right) \tag {36}
$$
$$
Y _ {i} \mid A _ {i} = a _ {i}, X _ {i} = x _ {i} \sim a _ {i} \mathcal {N} \left(9 0 + e ^ {0. 0 6 x _ {i}}, 1\right) + (1 - a _ {i}) \mathcal {N} \left(7 2 + 3 \sqrt {x _ {i}}, 1\right) \tag {37}
$$
In the "partially overlapping" setting (Figure 6), the means of the normal distributions used to generate the data are  $(\mu_1, \mu_2) = (40, 20)$ , as in Hill [2011]. In contrast, in the "overlapping" setting, we set  $(\mu_1, \mu_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.
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
$$
\pi \sim \operatorname {C a t} ((0. 5, 0. 2, 0. 1 5, 0. 1, 0. 0 5)) \tag {38}
$$
$$
\beta_ {i} \sim \pi \cdot (0, 1, 2, 3, 4) \quad i \in \{1, \dots , 2 5 \} \tag {39}
$$
$$
Y (0) \sim \mathcal {N} \left(X ^ {T} / \beta , 1\right) \tag {40}
$$
$$
Y (1) \sim \mathcal {N} \left(X ^ {T} \beta + 4, 1\right) \tag {41}
$$
Response Surface	VaRT	BART	Ground Truth
Linear	3.86 ± 0.12	3.6 ± 0.6	4.00
Non-Linear	3.80 ± 0.19	2.9 ± 0.8	4.00
Table 1: Estimation of the ATE for the 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  $\lambda = 1$ . The algorithm whose mean values were closer to the Ground Truth are shown in bold.
The second response surface is highly non-linear and it is generated as follows
$$
\pi \sim \operatorname {C a t} ((0. 6, 0. 1, 0. 1, 0. 1, 0. 1)) \tag {42}
$$
$$
\beta_ {i} \sim \pi \cdot (0., 1., 2., 3., 4) \quad i \in \{1, \dots , 2 5 \} \tag {43}
$$
$$
Y (0) \sim \mathcal {N} \left(\exp \left(X ^ {T} \beta + 1 / 2\right), 1\right) \tag {44}
$$
$$
Y (1) \sim \mathcal {N} \left(X ^ {T} \beta - \omega , 1\right) \tag {45}
$$
Here,  $\omega$  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.
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.
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.
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.0751±0.0031	0.0708±0.0031	0.085±0.005	0.0710±0.0020*	0.088±0.013
concreteslump(N=103, d=7)	0.049±0.023	0.034±0.017	0.069±0.031	0.149±0.011	0.16±0.10
forest(N=517, d=12)	0.205±0.005	0.194±0.005	0.219±0.008	0.1945±0.0026*	0.1977±0.0026
solar(N=1066, d=10)	0.114±0.006	0.108±0.004	0.124±0.013	0.1038±0.0004	0.1044±0.0009
stock(N=536, d=11)	0.0351±0.0014	0.0336±0.0009	0.040±0.005	0.051±0.004	0.060±0.017
yacht(N=308, d=6)	0.059±0.007	0.053±0.007	0.069±0.012	0.0684±0.0018	0.080±0.004
airfoil(N=1503, d=5)	0.0828±0.0026	0.0776±0.0020	0.0894±0.0029	0.087±0.005	0.092±0.008
autos(N=159, d=25)	0.107±0.007	0.0948±0.0031	0.121±0.012	0.075±0.006	0.093±0.004
breastcancer(N=194, d=33)	0.244±0.017	0.220±0.014	0.273±0.023	0.219±0.008	0.205±0.010
concrete(N=1030, d=8)	0.0837±0.0013	0.0786±0.0020	0.0922±0.0025	0.0768±0.0012*	0.085±0.005
fertility(N=100, d=9)	0.197±0.023	0.180±0.018	0.214±0.027	0.160±0.005	0.158±0.013
housing(N=506, d=13)	0.074±0.005	0.068±0.004	0.082±0.006	0.069±0.006*	0.077±0.006
machine(N=209, d=7)	0.093±0.007	0.083±0.009	0.104±0.009	0.0695±0.0013	0.14±0.04
pendulum(N=630, d=9)	0.060±0.006	0.054±0.006	0.073±0.009	0.0697±0.0016	0.072±0.006
servo(N=167, d=4)	0.105±0.013	0.090±0.012	0.127±0.015	0.0827±0.0034	0.110±0.006
skillcraft(N=3338, d=19)	0.0692±0.0006	0.0676±0.0007	0.0738±0.0031	0.0699±0.0007	0.0692±0.0025
sml(N=4137, d=26)	0.0356±0.0029	0.0342±0.0022	0.041±0.004	0.0416±0.0009	0.0384±0.0033
wine(N=1599, d=11)	0.0746±0.0022	0.0713±0.0020	0.0806±0.0034	0.0947±0.0031	0.0751±0.0015
TOTAL	8	-	-	8	2
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.
Figure A: 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. 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.