Accounting for shared covariates in semi-parametric Bayesian additive regression trees

TL;DR

This work addresses identifiability and bias challenges in semi-parametric Bayesian additive regression trees when covariates overlap between the linear main-effects predictor and the BART component. It introduces CSP-BART, which permits shared covariates and employs novel double-grow and double-prune tree moves along with a hierarchical prior on the linear coefficients to isolate primary effects while capturing nonparametric interactions. Through simulation studies on Friedman-type data and TIMSS 2019 applications, CSP-BART demonstrates lower bias for targeted main effects and the ability to uncover meaningful interactions, with competitive predictive performance relative to SSP-BART and VCBART. The approach delivers interpretable parameter estimates, handles missing data considerations via pre-screening or flexible extensions, and is implemented in an available R package, enabling application to education data and other complex predictive tasks.

Abstract

We propose some extensions to semi-parametric models based on Bayesian additive regression trees (BART). In the semi-parametric BART paradigm, the response variable is approximated by a linear predictor and a BART model, where the linear component is responsible for estimating the main effects and BART accounts for non-specified interactions and non-linearities. Previous semi-parametric models based on BART have assumed that the set of covariates in the linear predictor and the BART model are mutually exclusive in an attempt to avoid poor coverage properties and reduce bias in the estimates of the parameters in the linear predictor. The main novelty in our approach lies in the way we change the tree-generation moves in BART to deal with this bias and resolve non-identifiability issues between the parametric and non-parametric components, even when they have covariates in common. This allows us to model complex interactions involving the covariates of primary interest, both among themselves and with those in the BART component. Our novel method is developed with a view to analysing data from an international education assessment, where certain predictors of students' achievements in mathematics are of particular interpretational interest. Through additional simulation studies and another application to a well-known benchmark dataset, we also show competitive performance when compared to regression models, alternative formulations of semi-parametric BART, and other tree-based methods. The implementation of the proposed method is available at \url{https://github.com/ebprado/CSP-BART}.

Figures (10)

• Figure 1: An example of a tree generated from BART in $4$ different instances. In principle, BART does not generate only one tree but rather a set of trees which, summed together, are responsible for the final prediction. As indicated in panel (a), observations are pushed to the left child node when the splitting criterion is satisfied. The tree is represented as $\mathcal{T}^{(r)}_{1}$, where $r = \{1,2,3,4\}$ denotes the number of the iteration in which the tree is updated. The splitting rules (covariates and their split-points) are presented inside the rectangles which represent the internal nodes. The predicted values $\mu_{t \ell}$ are shown inside the circles which correspond to the terminal nodes. $\mathcal{T}^{(1)}_{1}$ illustrates the tree at iteration one with two internal nodes and three terminal nodes. From $\mathcal{T}^{(1)}_{1}$ to $\mathcal{T}^{(2)}_{1}$, the grow move is illustrated, as $\mu_{13}$ in $\mathcal{T}^{(1)}_{1}$ is split into $\mu_{13}$ and $\mu_{14}$ in $\mathcal{T}^{(2)}_{1}$ by using $x_{3} \le 2$. In addition, the prune move can be seen when $\mathcal{T}^{(2)}_{1}$ reverts to $\mathcal{T}^{(1)}_{1}$. The change move is shown when comparing $\mathcal{T}^{(2)}_{1}$ and $\mathcal{T}^{(3)}_{1}$, as the splitting rule that defines $\mu_{13}$ and $\mu_{14}$ is changed from $x_{3} \le 2$ to $x_{4} \le 0.75$. Finally, the swap move is illustrated in the comparison of $\mathcal{T}^{(3)}_{1}$ and $\mathcal{T}^{(4)}_{1}$.
• Figure 2: Boxplots of simulation results obtained across $50$ replicate datasets generated according to the Friedman equation, considering $n = 1000$, $p = (10, 50)$, and $\sigma^{2} = (1, 10)$. The y-axis exhibits the bias related to the parameter estimates $\hat{\beta}_{4}$ and $\hat{\beta}_{5}$ for the novel CSP-BART and various competitors. Recall that the GAM has been given the true model structure so its superior performance is expected.
• Figure 3: An illustration of the tree structure used to generate the response via Equation \ref{['tree_based_equation']}. In if-else format this can be written as $\mathcal{T}_{1} \:\vert\:\mathbf{x}_i = f(x_{i1}, x_{i2}, x_{i3})= 4\mathds{1}(x_{i1} \le 0.5) \times \mathds{1}(x_{i2} \le 0.5) -7 \mathds{1}(x_{i1} \le 0.5)\times \mathds{1}(x_{i2} > 0.5) + 3 \mathds{1}(x_{i1} > 0.5)\times \mathds{1}(x_{i3} \le 0.5) - 8 \mathds{1}(x_{i1} > 0.5)\times \mathds{1}(x_{i3} > 0.5)$, where $\mathds{1}(\cdot)$ denotes the indicator function. Note that the tree splits on both primary ($x_{1}$ and $x_{2}$) and non-primary ($x_3$) covariates.
• Figure 4: Boxplots of the simulation results obtained across $50$ replicate datasets generated according to Equation \ref{['tree_based_equation']}, considering $n = 1000$, $p = (10, 50)$, and $\sigma^{2} = (1, 10)$. The y-axis exhibits the bias related to the parameter estimates $\hat{\beta}_{1}$ and $\hat{\beta}_{2}$ for the novel CSP-BART and various competitors. Recall that the GAM has been given the true model structure so its superior performance is expected.
• Figure 5: Boxplots of the simulation results obtained across 50 replicate datasets generated according to Equation \ref{['smooth_function_with_interaction']}, considering $n = 1000$, $p = (10, 50)$, and $\sigma^{2} = (1, 10)$. The y-axis exhibits the bias related to the parameter estimates $\hat{\beta_{1}}$ and $\hat{\beta_{2}}$ for the novel CSP-BART and various competitors. Recall that the GAM has been given the true model structure so its superior performance is expected.