#  flexBART: Flexible Bayesian regression trees with categorical predictors

Sameer K. Deshpande Dept. of Statistics, University of Wisconsin–Madison.
<EMAIL_ADDRESS>
Most implementations of Bayesian additive regression trees (BART) one-hot
encode categorical predictors, replacing each one with several binary
indicators, one for every level or category. Regression trees built with these
indicators partition the discrete set of categorical levels by repeatedly
removing one level at a time. Unfortunately, the vast majority of partitions
cannot be built with this strategy, severely limiting BART’s ability to
partially pool data across groups of levels. Motivated by analyses of baseball
data and neighborhood-level crime dynamics, we overcame this limitation by re-
implementing BART with regression trees that can assign multiple levels to
both branches of a decision tree node. To model spatial data aggregated into
small regions, we further proposed a new decision rule prior that creates
spatially contiguous regions by deleting a random edge from a random spanning
tree of a suitably defined network. Our re-implementation, which is available
in the flexBART package, often yields improved out-of-sample predictive
performance and scales better to larger datasets than existing implementations
of BART.

##  1 Introduction

###  1.1 Motivation

Bayesian additive regression trees (BART; Chipman et al., , 2010) is a
Bayesian sum-of-trees model that approximates an unknown function $f$ using a
large ensemble of regression trees (i.e. piecewise constant step functions).
When $f$ is non-linear or involves complicated high-order interactions between
several predictors, correctly specifying $f$’s functional form with a
parametric model is extremely challenging, if not impossible. With BART,
however, users often obtain accurate predictions of function evaluations along
with reasonably well-calibrated uncertainty intervals without pre-specifying
the functional form of $f$. Better still, despite BART’s dependence on several
hyperparameters, users tend to obtain excellent results using default values
suggested by Chipman et al.,  (2010). BART’s ease-of-use and generally
excellent performance have made it a popular “off-the-shelf” modeling tool.

Despite BART’s popularity and success, the way most implementations handle
categorical predictors leaves much to be desired. Briefly, these
implementations one-hot encode categorical predictors, replacing each with
several binary indicators (one for each level or category). We demonstrate
that one-hot encoding severely limits BART’s ability to “borrow strength”
across different levels of a categorical variable. We re-implemented BART in a
way that permits much more flexible modeling with categorical predictors. Our
implementation is built using a new C++ backend that enables writing further
BART extensions. To motivate our development, consider the following problems.

Pitch framing in baseball. In baseball, when the batter decides not to swing
at a pitch, the umpire classifies the pitch as a ball or a strike. In
principle, ball/strike decisions are deterministic: if any part of the
baseball passes through the strike zone — an imaginary rectangle over home
plate that extends from the middle of the batter’s chest to their kneecaps —
it must be called a strike. In practice, observed ball/strike decisions are
incredibly noisy, varying considerably within and between umpires. Many
professional teams have heavily invested in catchers adept at pitch framing,
or systematically influencing umpires’ ball/strike decisions. To quantify the
effect of pitch framing, Deshpande and Wyner,  (2017) estimated each catcher’s
effect on called strike probabilities using a hierarchical logistic regression
model that “borrowed strength” across umpires. Although their model allowed
the effect of pitch location to vary across umpires, it assumed constant
player effects. As a result, they could not probe the possibility that some
(but not all) catchers could influence some (but not all) umpires on pitches
thrown to some (but not all) locations.

A potentially more realistic model would incorporate interactions between
umpires, pitch location, and players. Extending Deshpande and Wyner,  (2017)’s
model to do so, however, introduces a huge number of additional parameters.
For instance, nearly 100,000 parameters are needed to capture potential
interactions between the roughly 1000 batters and 100 umpires. The difficulty
in adding player, umpire, and location interactions to Deshpande and Wyner,
(2017)’s parametric model makes a nonparametric approach rather compelling.

Crime in Philadelphia. Balocchi and colleagues have recently used Bayesian
hierarchical models to study the spatiotemporal variation of crime at the
census-tract level in the city of Philadelphia (Balocchi and Jensen, , 2019;
Balocchi et al., , 2021, 2022). These works find that although models that
“share strength spatially” often outperform models that do not, failing to
account for spatial discontinuities can produced highly biased crime
forecasts. Spatial discontinuities can occur along visible geographic
landmarks like highways, parks, or rivers. But they can also coincide with
less visible differences in neighborhood-level demographic and socioeconomic
measures that have resulted from discriminatory policies like redlining.

Figure 1 illustrates the challenges faced when modeling crime in Philadelphia.
Figure 1(a) shows the spatial variability in the mean monthly crime density
(defined as crime counts per square mile; see Section 4.2 for details) in each
of Philadelphia’s 384 census tracts averaged across the 16 year period from
January 2006 to December 2021. Although some spatially adjacent tracts have
very similar crime dynamics, other tracts display markedly different patterns
than their neighbors (e.g. the starred tract in Figure 1(a) and Figure 1(b)).

(a)

(b)

Figure 1: (a) Average crime densities within each of the 384 census tracts,
with red corresponding to higher crime density and blue corresponding to lower
crime density. (b) Time series of the circled census tracts in (a) displaying
sharp spatial discontinuities (top) and spatial similarity (bottom). The
starred tract in (a) corresponds to the lowest curve in the top panel of (b).

Our exploratory analysis suggests that an accurate model of the spatiotemporal
variability in crime density must account for (i) non-linearities in time and
(ii) interactions between time and the census tract. Unfortunately, correctly
pre-specifying such non-linear interactions in a parametric model is
challenging, motivating a more flexible, nonparametric approach. Our ideal
model would exploit spatial similarities when they are present but also adapt
to spatial discontinuities suggested by the data. In other words, we want to
model spatial variation in a data-adaptive fashion, rather than imposing
strong spatial smoothness through, say, rigid priors on spatial random
effects.

###  1.2 Our contributions

In both motivating problems, we must estimate a function that may (i) be
highly non-linear; (ii) depend on the complex interaction and interplay
between categorical and continuous covariates; and (iii) have a functional
form that is difficult to pre-specify correctly in a parametric model. From
this perspective, BART appears ideally suited for both problems.

As we argue conceptually in Section 3.1 and demonstrate empirically in Section
4, however, existing implementations of BART can be improved with a more
thoughtful treatment of the categorical predictors. Briefly, in the baseball
application, by one-hot encoding player and umpire identities, existing
implementations of BART are unable to pool data across many small groups of
players and umpires. Further, in the crime modeling application, these
implementations cannot adapt to any pockets of spatial smoothness.

To overcome these challenges, we re-implemented BART in a way that avoids one-
hot encoding and yields a richer range of partitions of discrete space. Our
implementation, which is available in the flexBART package, is built using a
new C++ class structure that facilitates extending BART to accommodate more
complicated decision rules. The remainder of the paper is structured as
follows. We review the basic BART model in Section 2 before revealing how one-
hot encoding categorical variables severely restricts the representational
flexibility of regression tree ensembles in Section 3.1. We describe our
implementation in Section 3.2. In Section 3.3, we illustrate the extensibility
afforded by our implementation by describing a conceptually simple yet
powerful approach to fitting BART models over networks. We illustrate the
improvements in predictive quality and runtime offered by our implementation
in Section 4. We conclude in Section 5 with a discussion of several
methodological and computational extensions of our modeling framework.

##  2 A brief review of BART

In each motivating application, we want to estimate a function without
imposing strong structural assumptions about its form. BART facilitates such
estimation by approximating the function with a sum of binary regression
trees. For simplicity, we describe BART only in the homoscedastic setting with
independent Gaussian errors.

Suppose that we observe $n$ pairs
$(\bm{x}_{1},y_{1}),\ldots,(\bm{x}_{n},y_{n})$ of predictors
$\bm{x}\in\mathcal{X}$ and scalar outcomes $y\in\mathbb{R}$ from the model
$y\sim\mathcal{N}\left(f(\bm{x}),\sigma^{2}\right).$ We assume throughout that
we have $p_{\text{cont}}$ continuous predictors and $p_{\text{cat}}$
categorical predictors, for a total of $p=p_{\text{cont}}+p_{\text{cat}}$
predictors. Without loss of generality, we assume that the continuous
predictors have been re-scaled to lie in the interval $[0,1].$ For
$j=1,\ldots,p_{\text{cat}},$ we assume that the $j$-th categorical predictor
takes values in a discrete set $\mathcal{K}_{j}.$ We concatenate our
predictors into a single vector
$\bm{x}=(\bm{x}^{\top}_{\text{cont}},\bm{x}^{\top}_{\text{cat}})^{\top}$ whose
first $p_{\text{cont}}$ entries record the values of the continuous predictors
and whose last $p_{\text{cat}}$ entries record the values of the categorical
predictors. This way, $\mathcal{X}$ is a product space:
$\mathcal{X}=[0,1]^{p_{\text{cont}}}\times\mathcal{K}_{1}\times\cdots\times\mathcal{K}_{p_{\text{cat}}}.$

Formally, a regression tree over $\mathcal{X}$ is a triplet
$(T,\mathcal{D},\bm{\mu})$ consisting of (i) a finite, rooted binary tree $T$
containing several terminal or leaf nodes and non-terminal nodes; (ii) a
collection of decision rules $\mathcal{D},$ one for each non-terminal node in
$T;$ and (iii) a collection $\bm{\mu}$ of scalars or jumps, one for each leaf
node in $T.$ Every non-terminal node is connected to two children nodes, a
left child and a right child. Decision rules take the form
$\\{X_{j}\in\mathcal{C}\\}$ where $\mathcal{C}$ is a subset of the available
values of the $j$-th covariate $X_{j}$ at the associated non-terminal node of
$T.$ When $X_{j}$ is continuous, $\mathcal{C}$ is a half-open interval
$[0,c)\subset[0,1]$ and when $X_{j}$ is categorical, $\mathcal{C}$ is a
proper, non-empty subset of $\mathcal{K}_{j-p_{\text{cont}}}.$

Given a decision tree $(T,\mathcal{D})$ and an $\bm{x}\in\mathcal{X},$ we can
trace a path down the tree by following the decision rules. Starting from the
root node, when the path reaches a node with decision rule
$\\{X_{j}\in\mathcal{C}\\},$ it proceeds to the left child if
$x_{j}\in\mathcal{C}$ and to the right child otherwise. We restrict our
attention to decision trees that partition $\mathcal{X}$ in the sense that (i)
every leaf node is reached by the decision-following path of at least one
$\bm{x}\in\mathcal{X}$ and (ii) the decision-following path of every
$\bm{x}\in\mathcal{X}$ reaches a single, unique leaf. By associating each leaf
of a partitioning tree with its own scalar, the regression tree
$(T,\mathcal{D},\bm{\mu})$ represents a piecewise constant function over
$\mathcal{X}$ (see Figure 2(a) and 2(b)).

(a)

(b)

(c)

Figure 2: (a) A regression tree defined over $\mathcal{X}=[0,1]^{2}.$ (b)
Representation of the tree in (a) as a piecewise constant function. (c) A
regression tree defined over
$\mathcal{X}=[0,1]\times\\{c_{1},c_{2},\ldots,c_{6}\\}$ that can assign
multiple categorical levels to each branch. flexBART is built using trees like
the one in (c).

In principle, we can always find a piecewise constant function that
approximates a function $f$ arbitrarily well. When $f$ involves non-
linearities or complicated interactions, the approximating function typically
corresponds to a very deep regression tree with many leaf nodes. Rather than
attempt to learn a single, deep tree approximating $f,$ BART instead focuses
on learning an ensemble of much simpler trees that sum to an accurate
approximation of $f.$

###  2.1 Drawing regression trees from the prior

Chipman et al.,  (2010) introduced a regression tree prior that, building on
an earlier proposal by Chipman et al.,  (1998), strongly regularizes towards
relatively shallow trees with only a few leaves. We describe the regression
tree prior implicitly with a three-step procedure for simulating prior draws.
In the first step, we draw $T$ by simulating a branching process. Then,
conditionally on $T,$ we sequentially draw random decision rules at each non-
terminal node in a way that ensures $(T,\mathcal{D})$ partitions
$\mathcal{X}.$ Finally, we draw the $L(T)$ jumps in $\bm{\mu}.$

Drawing $T$. We grow $T$ recursively starting from a single root node, which
is considered to be a terminal node of depth zero. Each time a new terminal
node is created at depth $d$, we grow the tree from that node by creating two
child nodes with probability $0.95(1+d)^{-2}.$ Because the growth probability
decays quadratically with tree depth, the branching process almost surely
terminates with finitely many leaves.

Drawing $\mathcal{D}|T.$ Drawing independent decision rules can produce trees
that do not partition $\mathcal{X}.$ We instead draw decision rules
sequentially, only sampling the rule at a non-terminal node after all its
ancestors in $T$ have an associated rule. We draw the rule at given non-
terminal node by first picking $j,$ the index of the decision variable,
uniformly at random. Then, we set $\mathcal{C}$ to be a random subset of
$\mathcal{A},$ the set of all available values of $X_{j}$ at the current tree
node. We can compute $\mathcal{A}$ by taking intersections of the
$\mathcal{C}$’s (or their complements) at the node’s ancestors in the tree.
When $X_{j}$ is continuous, $\mathcal{A}$ is just an interval and we set
$\mathcal{C}=[0,c)$ where $c$ is drawn uniformly from $\mathcal{A}.$ When
$X_{j}$ is categorical, $\mathcal{A}$ is a discrete set and we simply assign
each element of $\mathcal{A}$ to $\mathcal{C}$ with probability $1/2.$

Drawing $\bm{\mu}|T,\mathcal{D}.$ We draw the scalar jump for each leaf
independently from a $\mathcal{N}\left(\mu_{0},\tau^{2}/M\right)$
distribution. In this way, the marginal prior of $f(\bm{x})$ is
$\mathcal{N}\left(\mu_{0},\tau^{2}\right)$ for every $\bm{x}\in\mathcal{X}.$
Chipman et al.,  (2010) recommended setting $\mu_{0}$ and $\tau$ so that this
marginal distribution places substantial probability over the range of the
observed $y_{i}$’s.

We complete the BART prior with a conditionally conjugate inverse gamma prior,
$\sigma^{2}\sim\text{Inv.\ Gamma}(3/2,3\lambda/2),$ where $\lambda$ is set to
ensure that the prior places 90% probability on the event that $\sigma$ is
less than the standard deviation of the $y_{i}$’s. Although these data-
dependent prior settings are not formally coherent, they work exceptionally
well in practice.

######  Remark 1.

When there are no categorical predictors (i.e. when $\mathcal{X}=[0,1]^{p}$),
the BART prior induces a prior over partitions of $[0,1]^{p}$ into axis-
parallel, rectangular boxes. Every such partition corresponds to at least one
decision tree that receives positive prior probability. So the induced prior
over partitions of $[0,1]^{p}$ into rectangular boxes has full support.
Further, the BART prior places positive probability around every piecewise
constant functions supported over axis-parallel boxes. These full support
properties imbue BART with tremendous representation flexibility.

###  2.2 Posterior computation

Chipman et al.,  (2010) used a Gibbs sampler to simulate posterior samples of
the regression tree ensemble. In each iteration, every regression tree is
updated conditionally fixing the remaining $M-1$ trees and $\sigma.$ Then,
$\sigma$ is updated conditionally fixing all regression trees. Two steps are
used to update an individual regression tree $(T,\mathcal{D},\bm{\mu}).$
First, given the data, $\sigma,$ and all other trees, a new decision tree
$(T,\mathcal{D})$ is drawn from its marginal posterior. Then, a new collection
of jumps $\bm{\mu}$ is drawn conditionally on the just-drawn
$(T,\mathcal{D}).$

The new decision tree $(T,\mathcal{D})$ is drawn with a Metropolis-Hasting
step. While there are several proposal distributions encouraging both small
and large modifications to the decision tree (see, e.g., Pratola, , 2016), we
focus on the simplest that randomly prunes or grows the decision tree. Pruning
$(T,\mathcal{D})$ involves (i) deleting two leafs sharing a common parent and
their incident edges in $T$; (ii) deleting the decision rule associated with
the parent of the deleted leafs; and (iii) leaving the rest of the tree and
decision rules unchanged. Growing $(T,\mathcal{D})$ in contrast, involves (i)
selecting an existing leaf and connecting two new children to it; (ii) drawing
a new decision rule for the selected node; and (iii) leaving the rest of the
tree and decision rules unchanged. When growing a tree, we draw the new
decision rule from the prior to simplify the acceptance probability
calculation; see Appendix B.

In our simple regression setting with i.i.d. errors, the jumps in $\bm{\mu}$
are conditionally independent and normally distributed. When the number of
observations in leaf is large, the conditional posterior distribution of
corresponding jump is sharply concentrated around the average of a particular
partial residual (see Equation B7 in Appendix B for a precise formula). In
this way, each tree performs adaptive partial pooling by partitioning
observations into several cells and then averaging within each cell.

######  Remark 2.

The individual trees in the BART ensembles are not identified. To see this,
note that permuting the order or labels of the trees leaves their sum, and
hence the likelihood function, unchanged. More subtly, we can represent the
same step function as a single tree or sum-of-trees in multiple ways (see
Figure B4 in Appendix B). Consequently, the BART posterior over tree space is
characterized by several facially different modes that nevertheless correspond
to exactly the same piecewise constant functions. On this view, the Gibbs
sampler’s slow mixing over tree space is analogous to the label switching
phenomenon observed in Bayesian mixture models. Practically, despite BART’s
failure to mix over tree space, we often obtain accurate point estimates and
reasonably well-calibrated uncertainty intervals for evaluations of
$f(\bm{x}_{i})$ after running many short chains of the Gibbs sampler. We defer
further discussion of this point to Appendix B.

##  3 Introducing flexBART

Most BART implementations one-hot encode categorical predictors, replacing
each with a collection of binary indicators, one for every level. That is,
rather than work with the original predictor space $\mathcal{X},$ which is a
product of Euclidean and discrete spaces, they instead work with an extended
predictor space $\tilde{X}=[0,1]^{\tilde{p}}$ where
$\tilde{p}=p_{\text{cont}}+\lvert\mathcal{K}_{1}\rvert+\cdots+\lvert\mathcal{K}_{p_{\text{cat}}}\rvert.$
At first glance, one-hot encoding appears innocuous since no information about
the predictor is lost. Further, one-hot encoding is de rigueur within
statistics and machine learning, with many popular textbooks recommending the
practice, especially for (generalized) linear modeling. In light of Remark 1,
it would appear that by replacing categorical predictors with continuous ones,
these BART implementations inherit desirable full support properties.

###  3.1 A one-hot encoded hole in the BART prior

The situation is, regrettably, not quite so simple. Although the prior places
positive probability on every partition of the extended space
$\tilde{\mathcal{X}}$ into axis-parallel boxes, it places zero probability on
a staggeringly overwhelming majority of partitions of the original space
$\mathcal{X}.$

To illustrate such lack-of-support, consider a setting with only one
categorical predictor $X$ taking values in $\\{c_{1},\ldots,c_{K}\\}.$ One-hot
encoding $X$ introduces $K$ binary indicators, $X_{1},\ldots,X_{K},$ where
$X_{k}=\mathbbm{1}\left(X=c_{k}\right).$ Suppose that
$(T,\mathcal{D},\bm{\mu})$ is a regression tree with at least two leaf nodes
defined over the extended space $[0,1]^{K}.$ Since $(T,\mathcal{D})$
partitions the entire space $[0,1]^{K}$, it automatically partitions the set
of standard basis vectors $\\{\bm{e}_{1},\ldots,\bm{e}_{K}\\}$ where the
$k$-th entry of $\bm{e}_{k}$ is equal to one and all other entries are zero.
By identifying $\bm{e}_{k}$ with the level $c_{k}$ of $X,$ we conclude that
$(T,\mathcal{D})$ also partitions $X$’s levels.

In fact, $(T,\mathcal{D})$ partitions $X$’s levels into several singleton sets
containing one level and at most one set containing multiple levels. To
understand why, consider the $\bm{e}_{k}$’s decision-following paths. When
they encounter the decision rule $\\{X_{j}\in[0,c)\\},$ only $\bm{e}_{j}$’s
path can proceed to the right while all others proceed to the left. Thus, the
decision tree partitions $\\{c_{1},\ldots,c_{K}\\}$ by repeatedly removing one
element at a time. Such a strategy can only produce $2^{K}-K$ distinct
partitions of $\\{c_{1},\ldots,c_{K}\\}.$ When $K>2,$ this number is much
smaller than the $K$-th Bell number, which counts the number of partitions of
$K$ objects. For instance, after one-hot encoding, the BART prior places
positive probability on just 27 of the 52 possible partitions when $K=5$ and
on less than 1% of partitions when $K=10$ (1,014 out of 115,975). The scale of
this lack-of-support phenomenon is magnified in the presence of multiple
predictors. By one-hot encoding categorical predictors, most BART
implementations cannot partially pool data across small groups of levels.

Figure 3(a) illustrates typical partitions of the Philadelphia census tracts
supported by the BART prior after one-hot encoding. The corresponding
regression trees pool together data from most tracts into one big cluster
while separately modeling a few singleton tracts. Although such modeling may
be desirable when all but a few census tracts display similar crime dynamics,
it is wholly inadequate in light of Figure 1. Indeed, a cursory look at that
figure suggests forming several small clusters containing multiple tracts with
similar dynamics would be more effective. Unfortunately, such clusters cannot
be formed by the BART prior after one-hot encoding.

(a)

(b)

(c)

Figure 3: Examples of partitions of the Philadelphia census tracts induced by
regression trees drawn from the BART prior after one-hot encoding (a), without
one-hot encoding (b), and using the network-based approach described in
Section 3.3. We colored the census tracts according to their corresponding
leaf in the regression tree.

###  3.2 A more flexible and extensible implementation

To overcome the limitations of one-hot encoding, we developed flexBART, a new
implementation of BART that uses regression trees that can assign multiple
levels to each branch to encourages richer partial pooling across categorical
levels. Figure 2(c) shows an example of such a tree defined over the space
$\mathcal{X}=[0,1]\times\left\\{c_{1},\ldots,c_{6}\right\\}.$ In writing
flexBART, we largely followed the example set by the authors of BART, writing
flexBART’s core sampling functions in C++ and interfacing those functions with
R through Rcpp (Eddelbuettel and François, , 2011). Although our overall
design is similar to theirs, there are several key differences. Briefly, we
created a new, highly extensible class to represent decision rules and
eliminated many redundant calculations performed by BART, which led to
substantial run-time improvements. flexBART is available at
https://github.com/skdeshpande91/flexBART.

As described in Pratola et al.,  (2014), BART represents regression trees as a
collection of nodes, which are instances of a custom class containing only six
data members: (i) pointers to separate class instances for the parent and
children of the node; (ii) an integer for the decision variable index $j$;
(iii) an integer recording the index of the cutpoint $c$111BART pre-computes a
finite grid of cutpoints for each predictor instead of drawing them from
$[0,1]$.; and (iv) a double for the jump. Armed with this lightweight class,
BART’s authors wrote a recursive functions to that trace the decision-
following path of each $\bm{x}\in\mathcal{X}.$ Although the class and this
function facilitate extremely fast and memory-efficient MCMC simulations, they
are insufficient for representing regression trees like the one in Figure
2(c).

We consequently re-wrote the regression tree node class and replaced the two
integers parametrizing the decision rule with an instance of a new class
representing decision rules. The data members of our new decision rule class
include (i) an integer for the decision variable index $j$; (ii) a Boolean
recording whether $X_{j}$ is a continuous or categorical predictor; (iii) a
double recording the value of the cutpoint $c$ for continuous $X_{j}$’s; and
(iv) two sets storing the levels assigned to the left and right child of the
node when $X_{j}$ is categorical. To accommodate our new classes, we rewrote
the function that determines the leaf node corresponding to a give $\bm{x}.$
Briefly, given an instance of the rule class and $\bm{x},$ that function
checks whether the decision variable is continuous or categorical. If the
variable is continuous, the function checks whether the corresponding entry in
$\bm{x}$ is less than the cutpoint. Otherwise, for a categorical splitting
variable, the function checks whether the relevant entry in $\bm{x}$ is
assigned to the left child of the node. We also introduced a separate function
(draw_rule) to draw new decision rules from the prior for use in the
Metropolis-Hastings proposal. After drawing the decision variable index $j$,
the function checks whether $X_{j}$ is continuous or categorical and
determines the set of available values $\mathcal{A}.$ If $X_{j}$ is
continuous, the function draws a cutpoint $c$ uniformly from the interval
$\mathcal{A}.$ If $X_{j}$ is categorical, the function draws a random subset
of the discrete set $\mathcal{A}$ and assigns it to the left child. See
Appendix C for technical details about the implementation.

To summarize, whereas BART hard-coded information about the decision rule into
its regression tree node class, we abstracted this information into a separate
class. Since we now store the levels assigned to the left and right child in
the case of a categorical decision rule, this abstraction increases the memory
footprint of our implementation. However, we have observed that this overhead
is often negligible in practice, even with many categorical levels. Further,
as we sketch in Section 5, such abstraction makes extending BART to use more
complicated decision rules relatively straightforward.

While implementing our sampler, we noticed that BART performs many redundant
calculations in each MCMC iteration. Specifically, when BART grows a tree from
or prunes a tree back to a particular node, it loops over all observations to
determine which observations’ decision-following paths reach that node.
Because at most two leaf nodes are changed during any tree update, the mapping
from observations to leaf nodes does not change much iteration-to-iteration.
In other words, before updating a given regression tree, BART re-derives
information that it had computed when updating the tree in the previous
iteration.s In flexBART, we eliminated these redundancies by saving and
manipulating a separate data structure that records which observations reach
which leafs. By maintaining these structure, we can eliminate unnecessary
loops over the entire dataset, yielding substantial speed-ups at scale. See
Appendix C2 for more details.

###  3.3 Splitting a network-structured categorical predictor

Figure 3(b) shows typical partitions of the Philadelphia census tracts induced
by the BART prior when we randomly assign each categorical level to the left
or right child instead of one-hot encoding. The corresponding regression trees
promote partial pooling by forming several large clusters, each containing
multiple tracts. While these new partitions are a marked improvement over
those in Figure 3(a), they are still not wholly satisfactory for our crime
data. To wit, while Figure 1 suggests the presence of spatially contiguous
clusters of tracts with similar crime dynamics, all clusters in Figure 3(b)
contain tracts that are geographically far apart. That this is so is entirely
expected: in assigning tracts independently to the children of a decision
node, the BART prior described in Section 2.1 does not leverage any spatial
adjacency information! To induce partitions with spatially contiguous
clusters, we must change the decision rule prior. Practically, using the
backend of flexBART, the main work involves updating the function draw_rule.

To describe this modification, suppose we have a categorical predictor $X$
whose levels display a known similarity structure. Suppose further that we
would like our regression trees to partition the levels so that similar levels
clustered together more often than not. To achieve this, we form a network
$\mathcal{G}$ where each vertex corresponds to a distinct level of $X$ and an
edge is drawn between vertices corresponding to similar levels. Our goal, is
to divide $\mathcal{G}$ recursively into connected components. In the context
of our crime data, $\mathcal{G}$’s vertices correspond to census tract and
edges represent spatial adjacency.

It suffices to describe a random process to partition a connected subgraph of
$\mathcal{G}$ into two connected components. Specifically, suppose that in
generating a regression tree from the prior that we have selected our network-
structured categorical predictor $X$ as the decision variable. Further,
suppose that the set of available levels $\mathcal{A},$ induces a connected
subgraph $\mathcal{G}[\mathcal{A}]$ of $\mathcal{G}.$ Unless
$\mathcal{G}[\mathcal{A}]$ is a complete graph, independently assigning levels
in $\mathcal{A}$ to $\mathcal{C}$ with a fixed probability tends not to
partition $\mathcal{G}[\mathcal{A}]$ into two connected components. We instead
follow the three-step procedure illustrated in Figure 4. First, we draw a
spanning tree of $\mathcal{G}[\mathcal{A}]$ using Wilson’s algorithm (Wilson,
, 1996), which involves running a loop-erased random walk on the vertices of
$\mathcal{G}[\mathcal{A}].$ Then, we delete a uniformly selected edge from the
spanning tree, thereby partitioning the spanning tree into two connected
components. Finally, we set $\mathcal{C}$ to contain the vertices of one of
these components.

(a)

(b)

(c)

(d)

Figure 4: Cartoon illustration of our graph partitioning process. (a) Induced
subgraph of levels available at a particular non-terminal node in the decision
tree. (b) Uniformly drawn random spanning tree. (c) Deleting a uniformly
selected edge (the edge $(2,8)$ in this case) from the spanning tree
disconnects the tree and partitions the vertices into two subsets. (d) The
subgraphs induced by the two vertex sets produced in (c) remain connected.

Embedding this process into the decision rule prior induces a prior over
partitions of $\mathcal{G}$ into connected components, which turns out to have
full support. That is, every partition of the network $\mathcal{G}$ into
connected components can be represented with a decision tree receiving
positive prior probability. In our motivating crime example, extending the
BART prior in this way enables us to consider all spatially contiguous
partitions of the census tracts.

######  Remark 3.

Forming spatially contiguous clusters by deleting spanning tree edges is an
emerging idea in the spatial clustering literature (see, e.g., Teixeira et
al., , 2015, 2019; Li and Sang, , 2019; Luo et al., 2021b, ). In the context
of BART, Luo et al., 2021a and Luo et al.,  (2022) introduced a decision rule
prior that is very similar to ours. Whereas we were motivated by categorical
predictors displaying network structure, their procedure recursively
partitions a continuous manifold. Briefly, they first triangulate the manifold
and then form a graph capturing the adjacency relationships. Then they
recursively remove edges from spanning trees of the graph to partition the
manifold into contiguous regions.

##  4 Empirical results

We performed two experiments, one for each motivating dataset, comparing the
predictive performance, runtime, and memory utilization of flexBART and BART.
In Section 4.1, we report the results of the experiment with the baseball data
that highlights some of the benefits of directly partitioning categorical
levels and showcases flexBART’s ability to scale to large datasets. Then, in
Section 4.2, we report the results of a full leave-one-out analysis of the
Philadelphia crime data that demonstrates the benefits of our partially
pooling data flexibly across vertices in a network.

We ran all experiments on a shared high-throughput computing cluster. All
results are based on simulating a single Markov chain for 2,000 iterations.
Although we would simulate more chains in actual applied practice, a single
chain sufficed for our comparative study. Code to reproduce our experiments is
available as Supplementary Materials.

In both experiments, we created training-testing splits where some categorical
levels appeared only in the testing dataset. When BART one-hot encodes
categorical predictors, it internally creates a design matrix with columns for
each observed continuous predictor and each categorical level. In this design
matrix, categorical levels not present in the training dataset correspond to
constant columns. By default, BART drops these constant columns before
starting the MCMC simulation. In some sense, by dropping constant columns,
BART loses information about the original predictor, replacing it with a
modified version that can take on only a subset of values. In our baseball
application, because each training dataset contains a different number of
unique players and umpires, BART effectively uses a different set of
predictors on each fold. While this behavior raises interesting conceptual
questions — namely, is it appropriate to average across cross-validation folds
if different predictors are used on each — we would not anticipate dropping or
retaining constant columns to make much practical difference. This is largely
because the source code for the Metropolis-Hastings step in BART is hardcoded
to reject transitions to trees with fewer than five observations in any leaf.
The restriction effectively prevents the sampler from accepting trees that
split on constant predictors. Nevertheless, we fit two BART models using the
BART package in our experiments, one that drops categorical levels unseen
during training (BART-default) and one that does not drop those unobserved
levels (BART-alt)222BART-alt was run by setting the argument rm.const = FALSE
in BART::wbart and BART::pbart..

The occurrence of categorical levels in testing but not training data poses no
real difficulties to flexBART. In our implementation, users pass a list
enumerating all possible levels of all categorical predictors (i.e. the sets
$\mathcal{K}_{j}$ from Section 2) alongside the training data. With such
information, even if a level is not observed in the training data, flexBART
can still meaningfully assign that level to the left or right child of a
decision node. Consequently, flexBART tends not to return the exact same
prediction for all levels unseen during training.

###  4.1 Pitch framing

We scraped pitch-by-pitch data from each Major League Baseball season between
2013 and 2019 using the baseballr package (Petti and Gilani, , 2022). Our
dataset includes the ball/strike decision; the horizontal and vertical
coordinates of each pitch as it crosses the front edge of home plate; and the
identities of the batter, catcher, pitcher, and umpire involved in a pitch.
Each season, there were about 350,000 called pitches and around 1000 batters,
100 catchers, 800 pitchers, and 90 umpires. Letting $y$ be a binary indicator
of ball ($y=0$) or strike ($y=1$) and concatenating pitch location and the
identities of the players and umpire into the vector $\bm{x},$ we modeled
$\mathbb{P}(y=1)=\Phi(f(\bm{x}))$ where $\Phi$ is the standard normal
cumulative distribution function and used BART to approximate $f.$ For each
season, we generated 10 training-testing splits and computed the posterior
mean called strike probability for each pitch in every training and testing
dataset. We assessed the predictive performances of each implementation using
misclassification rate.

The average in-sample misclassification rates were 8.5% (flexBART), 9.6%
(BART-default), and 9.9% (BART-alt), while the average out-of-sample rates
were 8.6% (flexBART), 9.7% (BART-default), and 62.3% (BART-alt). BART-alt’s
apparent degradation in out-of-sample predictive performance likely is an
artifact BART’s internal design matrix construction.

Comparing flexBART to BART-default, it initially appears that the former’s
more flexible partial pooling across levels only slightly improved predictive
accuracy. To better contextualize the one percentage point improvement in
misclassification error, we fit another model with flexBART to predict called
strike probability using only pitch location. Interestingly, the location-only
model achieved an 8.9% misclassification rate both in- and out-of-sample,
placing it squarely between flexBART and BART-default. The similarity in
predictions across models suggest that, after accounting for location, player
and umpire identities have a small effect on called strike probability.
Nevertheless, across every fold, flexibly “borrowing strength” across players
and umpires with flexBART conferred a slight predictive advantage over the
location-only model. Simultaneously, one-hot encoding these identities with
BART yielded somewhat worse predictions than the location-only model. We
observed qualitatively similar results with other error metrics like mean
squared error and log-loss; see Figure A1 in Appendix A.

Although flexBART did not dramatically improve the predictive accuracy for
these data, our experiment revealed other important practical advantages.
Recall from Section 3.2 that BART repeatedly re-calculates which observations
reach leafs in every tree. Removing these redundancies led to a nearly six
times speed-up on these data: flexBART taking roughly 72 minutes to draw 2,000
posterior samples while the BART implementations needed over nine hours.
flexBART was also much more memory-efficient. Because it one-hot encoded the
player and umpire identifies, BART internally created design matrices with
approximately 350,000 rows and 2,000 columns. All together, the BART
implementations utilized about 25GB of memory per fold. In contrast, no run of
flexBART required more than 3GB of memory, the vast majority of which was
dedicated to storing the matrix of posterior samples.

###  4.2 Philadelphia crime data

We obtained our crime data from opendataphilly.org, where the Philadelphia
Police Department releases the type, time, and location (longitude and
latitude) of every reported crime in the city. Following Balocchi et al.,
(2022), in each tract, we computed the monthly crime density, defined as
number of crimes per square mile, and applied an inverse hyperbolic sine
transformation to counteract the considerable skewness. Formally, let
$y_{v,t}$ be the transformed crime density in census tract $v$ at time $t,$
with $t=1$ corresponding to January 2006 and $t=192$ corresponding to December
2021. We model $y_{v,t}\sim\mathcal{N}\left(f(t,v),\sigma^{2}\right),$ where
the function $f$ depends on a ordered predictor (the time index $t$) and an
unordered categorical predictor (the census tract label $v$), which displays a
natural network adjacency structure. In our experiments, we modified the
decision rule prior to split only at the observed time indices333By default,
flexBART draws cutpoints uniformly but users can override this behavior by
specifying a grid of cutpoints for each continuous variable..

We performed a full leave-one-out experiment in which we trained our BART
models on all data from all but one census tracts and attempted to predict the
entire trajectory of transformed crime densities for the held-out tract. For
each cross-validation fold, we did not delete the held-out vertex and its
incident edges from the network encoding tract adjacency. In a sense, our
experiment probes the ability of each BART implementation to interpolate at an
existing network vertex rather than extrapolate to entirely new vertices. In
addition to BART-default and BART-alt, we fit two models using flexBART, one
which does not account for the network structure (flexBART) and one that does
(networkBART).

Each implementation makes predictions about the held-out tract using a
weighted average of predictions from the other tracts. Whereas networkBART is
designed to weight nearby tracts more heavily, flexBART tends to average over
an essentially random collection of tracts. BART-default and BART-alt tend to
weigh every other tract about equally.

Each implementation fit the training data extremely well: the average in-
sample root mean square errors (RMSE) were 1.22 (flexBART and networkBART) and
1.29 (BART-default and BART-alt). Perhaps unsurprisingly, networkBART produced
more accurate out-of-sample predictions, achieving an average out-of-sample
RMSE of 3.00, which is smaller than 5.69 (flexBART), 3.73 (BART-default), and
6.17 (BART-alt). Like in our baseball example, we suspect the large gap in
performance between BART-default and BART-alt derives from the way BART
constructs its internal design matrices.

More interestingly, networkBART was not uniformly superior to the other
implementations across every cross-validation fold. In fact, networkBART
outperformed BART-default in only 231 of the 384 folds. We noticed that BART-
default generally outperformed networkBART when the held-out tract’s crime
dynamics happened to be closer to the overall city-wide average than to the
average of the neighboring tracts. As a concrete example, consider the starred
tract in Figure 1(a), which displayed consistently lower crime density than
all its neighbors. When we held out this tract, networkBART’s predictions,
being primarily driven by the tract’s neighbors, over-estimated the actual
crime densities. Because this tract’s crime densities happened to be closer to
the global average than the average in a local neighborhood, BART-default
returned much more accurate predictions. We also noticed that although
flexBART pooled data across multiple tracts, it often demonstrated worse
predictive performance than the other implementations. Essentially, in
borrowing strength across tracts in a totally random fashion, flexBART was
unable to adapt to and exploit the spatial smoothness present in Figure 1(a).

For each cross-validation fold, the training dataset contained over 73,000
total observations. Because it repeatedly simulated random walks on networks,
networkBART is predictably slower than flexBART, drawing 2,000 samples in 36
minutes compared to flexBART’s 16 minutes. Nevertheless, networkBART is still
faster than both BART implementations, which required over seven hours to draw
the same number of samples, further underscoring the importance of eliminating
redundant calculations.

##  5 Discussion

Existing implementations of BART one-hot encode categorical predictors,
operationalizing them as several binary indicators. By adopting such a
representation, they effectively specify a prior that places zero probability
on the vast majority of partitions of the levels, limiting their ability to
partially pool data across categories. We overcome this limitation with a new
implementation of BART using a more general class of regression trees. In our
trees, multiple levels of a categorical variable can be assigned to both the
left and right child of decision nodes. BART ensembles built using our new
trees can borrow strength much more flexibly across categorical levels.
Motivated by a spatial analysis of crime in Philadelphia, we introduced a
stochastic process that, when incorporated into our decision rule prior,
enables our regression trees to form spatially contiguous partitions of small
geographic regions.

In that stochastic process, we deleted a uniformly selected edge from a
uniformly sampled spanning tree of a suitably constructed network. This was
done to ensure that set $\mathcal{C}$ and its complement
$\mathcal{A}\setminus\mathcal{C}$ induce connected subgraphs. We have found
that the process can create very small clusters containing only a few vertices
(see, e.g., the middle and bottom panels of Figure 3(c)). To encourage more
balanced partitions, we could instead delete an edge with probability
proportional to the size of the smallest cluster that results when the edge is
deleted. We plan to explore other network partitioning processes that still
ensure that $\mathcal{C}$ induces a connected subgraph. For instance, we could
form $\mathcal{C}$ by randomly selecting an element of $\mathcal{A}$ and then
including all vertices in the network that are within a randomly chosen graph
distance from the selected element.

While we focused on data aggregated within census tracts, we can seamlessly
model variation at multiple spatial resolutions using BART. Specifically, we
could introduce additional categorical predictors that encode membership in
other, not necessarily nested geographic areas, like counties or police
precincts. We could further form adjacency networks for each of these
predictors to encourage regression trees to form spatially contiguous
partitions at each resolution. We plan to explore such multi-resolution models
more fully in future work.

More substantively, our argument against one-hot encoding categorical
predictors when fitting BART models is based on a particular lack-of-support
phenomenon. It is important to note that the cause of this phenomenon is not
the one-hot encoding per se. Rather, the problem is the combination of one-hot
encoding and the use of axis-aligned decision rules based on one predictor.
One can, in fact, overcome the lack-of-support problem by utilizing decision
rules based on linear combinations of the binary indicators created by one-hot
encoding. After all, for a categorical predictor $X$ and discrete set
$\mathcal{C},$ the indicator $\mathbbm{1}\left(X\in\mathcal{C}\right)$ is just
the sum of binary indicators $\mathbbm{1}\left(X=c\right)$ for each
$c\in\mathcal{C}.$

More generally, one can imagine running BART using oblique regression trees
built with decision rules of the form
$\\{\phi^{\top}\bm{x}_{\text{cont}}\in(-\infty,c)\\}$ where
$\bm{x}_{\text{cont}}$ is the vector of continuous predictors (see Cattaneo et
al.,  (2022) and references therein for a recent, comprehensive review).
Although we have not yet implemented oblique trees, our abstract decision rule
class renders such extension relatively easy. It essentially suffices to
implement a random procedure to draw the pair $(\phi,c)$ in such a way that
the resulting tree leafs correspond to non-empty subsets of $\mathcal{X}.$
That is, the main challenge is ensuring every internal node’s decision rule is
compatible with the constraints implied by the rules of its ancestors in the
tree. Whether BART ensembles built using oblique trees offer any practical or
theoretical advantages over the more conventional axis-aligned trees is a
decidedly open question. Nevertheless, we anticipate that the source code
underlying flexBART can enable empirical investigations of this question.

## Acknowledgements

I would like to thank Prof. Cecilia Balocchi (Edinburgh) for several helpful
discussions about flexible spatial modeling of areal crime data that motivated
this work. I am deeply grateful to my colleagues Profs. Keith Levin and Jun
Zhu at Wisconsin for many productive conversations about networks, spanning
trees, and spatial clustering. Thanks are also due to Susan Glenn, Ajinkya
Kokandakar, and Kehui Yao, who provided immensely helpful feedback on early
versions of the flexBART package. Finally, I am indebted to Prof. Ed George
for providing constructive feedback on an early draft of this manuscript.

This research was performed using the computing resources and assistance of
the UW–Madison Center For High Throughput Computing (CHTC) in the Department
of Computer Sciences. The CHTC is supported by UW–Madison, the Advanced
Computing Initiative, the Wisconsin Alumni Research Foundation, the Wisconsin
Institutes for Discovery, and the National Science Foundation, and is an
active member of the OSG Consortium, which is supported by the National
Science Foundation and the U.S. Department of Energy’s Office of Science.

Support was also provided by the University of Wisconsin–Madison, Office of
the Vice Chancellor for Research and Graduate Education with funding from the
Wisconsin Alumni Research Foundation.

## References

  * Balocchi et al., (2022) Balocchi, C., Deshpande, S. K., George, E. I., and Jensen, S. T. (2022).  Crime in Philadelphia: Bayesian clustering with particle optimization.  Journal of the American Statistical Association.
  * Balocchi et al., (2021) Balocchi, C., George, E. I., and Jensen, S. T. (2021).  Clustering areal units at multiple levels of resolution to model crime incidence in Philadelphia.  arXiv:2112.02059.
  * Balocchi and Jensen, (2019) Balocchi, C. and Jensen, S. T. (2019).  Spatial modeling of trends in crime over time in Philadelphia.  Annals of Applied Statistics, 13(4):2235–2259.
  * Cattaneo et al., (2022) Cattaneo, M. D., Chandak, R., and Klusowski, J. M. (2022).  Convergence rates of oblique regression trees for flexible function libraries.  arXiv preprint 2210.14429v1.
  * Chipman et al., (1998) Chipman, H. A., George, E. I., and McCulloch, R. E. (1998).  Bayesian CART model search.  Journal of the American Statistical Association, 93(443):935–948.
  * Chipman et al., (2010) Chipman, H. A., George, E. I., and McCulloch, R. E. (2010).  BART: Bayesian additive regression trees.  Annals of Applied Statistics, 4(1):266–298.
  * Deshpande and Wyner, (2017) Deshpande, S. K. and Wyner, A. J. (2017).  A hierarchical Bayesian model of pitch framing.  Journal of Quantitative Analysis in Sports, 13(3):95–112.
  * Eddelbuettel and François, (2011) Eddelbuettel, D. and François, R. (2011).  Rcpp: Seamless R and C++ integration.  Journal of Statistical Software, 40(8):1–18.
  * Kim and Ročková, (2023) Kim, J. and Ročková, V. (2023).  On mixing rates for Bayesian CART.  arXiv preprint 2306.00126.
  * Li and Sang, (2019) Li, F. and Sang, H. (2019).  Spatial homogeneity pursuit of regression coefficients for large datasets.  Journal of the American Statistical Association, 114(527):1050–1062.
  * (11) Luo, Z. T., Sang, H., and Mallick, B. (2021a).  Bast: Bayesian additive regression spanning trees for complex constrained domain.  In Advances in Neural Information Processing Systems.
  * (12) Luo, Z. T., Sang, H., and Mallick, B. (2021b).  A Bayesian contiguous partitioning method for learning clustered latent variables.  Journal of Machine Learning Research, 22(37):1–52.
  * Luo et al., (2022) Luo, Z. T., Sang, H., and Mallick, B. (2022).  BAMDT: Bayesian additive semi-multivariate decision trees for nonparametric regression.  In Proceedings of the 39th International Conference on Machine Learning.
  * Petti and Gilani, (2022) Petti, B. and Gilani, S. (2022).  baseballr: Acquiring and analyzing baseball data.
  * Pratola, (2016) Pratola, M. T. (2016).  Efficient Metropolis-Hastings proposal mechanisms for Bayesian regression tree models.  Bayesian Analysis, 11(3):885–911.
  * Pratola et al., (2014) Pratola, M. T., Chipman, H. A., Gattiker, J. R., Higdon, D. M., McCulloch, R., and Rust, W. N. (2014).  Parallel Bayesian additive regression trees.  Journal of Computational and Graphical Statistics, 23(3):830–852.
  * Ronen et al., (2022) Ronen, O., Saarinen, T., Tan, Y. S., Duncan, J., and Yu, B. (2022).  A mixing time lower bound for a simplified version of BART.  arXiv preprint 2210.09352v1.
  * Teixeira et al., (2015) Teixeira, L. V., Assunção, R. M., and Loschi, R. H. (2015).  A generative spatial clustering model for random data through spanning trees.  In 2015 IEEE International Conference on Data Mining.
  * Teixeira et al., (2019) Teixeira, L. V., Assunção, R. M., and Loschi, R. H. (2019).  Bayesian space-time partitioning by sampling and pruning spanning trees.  Journal of Machine Learning Research, 20(85):1–35.
  * Wilson, (1996) Wilson, D. B. (1996).  Generating random spanning trees more quickly than the cover time.  In Proceedings of the twenty-eighth annual ACM symposium on Theory of Computing.

Supplementary Materials

##  Appendix A Additional experiments

###  A.1 Called strike probabilities

In Section 4.1 of the main text, we reported the in- and out-of-sample
misclassification rates averaged over all 70 training and testing folds of our
baseball data. Figure A1 shows boxplots comparing three different error
metrics — misclassification rate, mean square error, and log-loss — computed
on each fold. Each error metric was computed using the posterior mean estimate
of the called strike probability. We included the errors for the model that
estimates called strike probability using pitch location only but omitted the
out-of-sample results for BART-alt as those tended to be substantially larger
than those of other methods. We found that across all error metrics, flexBART
achieved smaller error than BART-default and the location-only model on every
training and testing fold.

Figure A1: The misclassification rate, mean square error, and log-loss of
BART-default, BART-alt, flexBART, and the location-only BART model on every
training (top row) and testing (bottom row) fold. Out-of-sample errors for
BART-alt were consistently much larger than those of other methods and were
omitted from the figure.

###  A.2 Philadelphia crime data

In Section 4.2 of the main text, we reported the results of a leave-one-out
analysis of the Philadelphia crime data. In that experiment, we repeatedly
trained a BART model using all data from $n-1=383$ census tracts and attempted
to predict all the observations in the held-out tract. We found that
networkBART, a version of the BART prior that recursively partitioned the
network of census tracts into spatially contiguous clusters, performed the
best. Essentially networkBART made predictions at the held-out tract by
averaging predictions made in neighboring tracts. Given the spatial smoothness
evident in Figure 1a of the main text, our leave-one-out experiment arguably
favors networkBART over BART-based implementations, which cannot exploit the
spatial structure.

For an arguably fairer comparison, we performed a second cross-validation
experiment in which we repeatedly trained BART models using 90% of all
observed data and evaluated the predictions on the remaining 10%. We held out
data uniformly at random, ensuring that every census tract appeared in every
training dataset, and allowing us to probe how partially pooling data across a
network helps predict at vertices already seen during training.

Figure A2 compares the in- and out-of-sample RMSEs of each method across 100
different training-testing folds. In this experiment, BART-default and BART-
alt performed virtually identically, which is not entirely unexpected. To
elaborate, because every tract appeared in every training dataset, BART-
default and BART-alt constructed identical internal design matrices for the
training and testing data, which were subsequently passed to the same sampler.

(a)

(b)

Figure A2: In-sample (a) and out-of-sample (b) RMSE across 100 training-
testing splits of the crime data.

Similar to what we observed in the leave-one-out experiment from Section 4.2
of the main text, our second experiment reveals that pooling data across
multiple census tracts can improve predictions. However, unlike that
experiment, Figure 2(b) reveals that networkBART and flexBART achieved
extremely similar out-of-sample RMSEs. That is, while leveraging the adjacency
structure helped in our leave-one-out experiment, in our second experiment,
how we partially pooled data did not make a material difference. While these
finds appear to contradict one another, they actually align with our basic
intuition about Bayesian shrinkage. Namely, when partially pooling across many
groups, the effects of shrinkage are more apparent for groups with very little
data. Recall that no data was observed in the held-out tract during training
in our leave-one-out experiment. As a result, the effects of our network-
informed shrinkage are quite evident there. By contrast, in our second
experiment, because we observed quite a lot of data from each tract during
training, we would not expect to see much evidence of shrinkage.

##  Appendix B Gibbs sampler derivation

To describe the Gibbs sampler used to simulate posterior draws of the tree
ensemble, we require some additional notation. Recall that the decision tree
$(T,\mathcal{D})$ partitions $\mathcal{X}$ into disjoint, non-empty regions,
one for each leaf in the tree $T.$ For each $\bm{x}\in\mathcal{X},$ let
$\ell(\bm{x})$ be the leaf corresponding to the region containing the point
$\bm{x}.$ For brevity, we will say that the leaf $\ell(\bm{x})$ contains the
point $\bm{x}.$ Further, by associating leaf $\ell$ with a scalar jump
$\mu_{\ell},$ we can identify the regression tree $(T,\mathcal{D},\bm{\mu})$
with a piecewise constant function. Formally, we introduce the evaluation
function $g:\mathcal{X}\rightarrow\mathbb{R},$ which is defined as

$g(\bm{x};T,\mathcal{D},\bm{\mu})=\mu_{\ell(\bm{x})},$ (B1)

that returns the jump associated to the leaf containing the point $\bm{x}.$

Given observations $(\bm{x}_{1},y_{1}),\ldots,(\bm{x}_{n},y_{n})$ and a
decision tree $(T,\mathcal{D}),$ let $I(\ell)$ be the set of indices for those
observations contained in leaf $\ell.$ That is,
$I(\ell)=\\{i:\ell(\bm{x}_{i})=\ell\\}.$ For compactness, we have suppressed
the dependence of $\ell(\bm{x})$ and $I(\ell)$ on $(T,\mathcal{D})$ in our
notation. The collection of $I(\ell)$’s determines a mapping from the set of
observations to the set of leafs in $T.$ This mapping plays a critical role in
the conditional regression tree updates (see Equations (B6) and (B7)). When
$n$ is large, computing the $I(\ell)$’s is the main computational bottleneck
of each individual tree update. As we describe later in Appendix C.2, avoiding
redundant re-computations of the $I(\ell)$’s yields considerable speed-up.

Recall that our model is $y\sim\mathcal{N}\left(f(\bm{x}),\sigma^{2}\right)$
and that we wish to estimate $f$ using a regression tree ensemble. To this
end, we formally introduce a collection
$\mathcal{E}=\\{(T_{m},\mathcal{D}_{m},\bm{\mu}_{m}):m=1,\ldots,M\\}$ of $M$
regression trees such that

$f(\bm{x})=\sum_{m=1}^{M}{g(\bm{x};T_{m},\mathcal{D}_{m},\bm{\mu}_{m})}.$

Let $\bm{f}=(f(\bm{x}_{1}),\ldots,f(\bm{x}_{n}))^{\top}$ be the vector of
regression function evaluations at the observed inputs. We specify independent
and identical priors on each $(T_{m},\mathcal{D}_{m},\bm{\mu}_{m})$ and use a
Gibbs sampler to simulate posterior draws of $(\mathcal{E},\sigma)|\bm{y}.$ At
a high-level, the sampler alternates between updating $\mathcal{E}$ keeping
$\sigma$ fixed and updating $\sigma$ keeping $\mathcal{E}$ fixed. When
sampling $\mathcal{E},$ the sampler sequentially updates each regression tree
conditionally fixing the other $M-1$ trees and $\sigma.$

###  B.1 A Bayesian backfitting algorithm

In what follows, we will describe the update for the $m$-th tree in the
ensemble, conditionally given all other trees, the collection of which we
denote as $\mathcal{E}^{(-m)}.$ The conditional posterior distribution of the
$m$-th tree has density

$p(T,\mathcal{D},\bm{\mu}|\bm{y},\sigma,\mathcal{E}^{(-m)})\propto
p(\bm{y}|T,\mathcal{D},\bm{\mu},\sigma,\mathcal{E}^{(-m)})p(T,\mathcal{D},\bm{\mu}).$

The sampler draws a tree from this distribution in two steps. First, it draws
a new decision tree $(T_{m},\mathcal{D}_{m})$ from the marginal posterior
distribution $(T,\mathcal{D})|\left[\bm{y},\sigma,\mathcal{E}^{(-m)}\right]$
with density

$p(T,\mathcal{D}|\bm{y},\sigma,\mathcal{E}^{(-m)})\propto\int{p(\bm{y}|T,\mathcal{D},\bm{\mu},\sigma,\mathcal{E}^{(-m)})p(T,\mathcal{D},\bm{\mu})d\bm{\mu}}.$

Then, we draw new jumps from the conditional posterior of
$\bm{\mu}|\left[T_{m},\mathcal{D}_{m},\bm{y},\sigma,\mathcal{E}^{(-m)}\right].$

The conditional likelihood of the $m$-th tree plays a crucial role in both
updates. Given $\mathcal{E}^{(-m)}$ and $\sigma,$ we have

$p(\bm{y}|T,\mathcal{D},\bm{\mu},\sigma,\mathcal{E}^{(-m)})=(2\pi\sigma^{2})^{-n/2}\exp\left\\{-\frac{1}{2\sigma^{2}}\sum_{i=1}^{n}{\left[r_{i}-g(\bm{x}_{i};T,\mathcal{D},\bm{\mu})\right]^{2}}\right\\},$
(B2)

where

$r_{i}=y_{i}-\sum_{m^{\prime}\neq
m}{g(\bm{x}_{i};T_{m^{\prime}},\mathcal{D}_{m^{\prime}},\bm{\mu}_{m^{\prime}})}$
(B3)

is the $i$-th partial residual. Let $\bm{r}=(r_{1},\ldots,r_{n})^{\top}$
denote the vector of partial residuals.

We can rewrite the sum over all observations in Equation (B2) as the following
double sum

$\sum_{i=1}^{n}{\left[r_{i}-g(\bm{x}_{i};T,\mathcal{D},\bm{\mu})\right]}=\sum_{\ell}{\sum_{i\in
I(\ell)}{\left[r_{i}-\mu_{\ell}\right]}},$ (B4)

where the outer sum is over the leafs of $T$ and the inner sum is over the
observations contained in the leaf. From this decomposition, we conclude that
the conditional likelihood of the $m$-th regression tree factorizes over the
leafs in $T:$

$p(\bm{y}|T,\mathcal{D},\bm{\mu},\sigma,\mathcal{E}^{(-m)})=(2\pi\sigma^{2})^{-n/2}\times\prod_{\ell}{\exp\left\\{-\frac{1}{2\sigma^{2}}\sum_{i\in
I(\ell)}{\left[r_{i}-\mu_{\ell}\right]^{2}}\right\\}}.$ (B5)

Recall from Section 2.1 of the main text that we specified independent
$\mathcal{N}\left(\mu_{0},\sigma^{2}\right)$ priors for each jump in
$\bm{\mu}.$ We conclude, then, that the entire joint distribution of
$(\bm{y},T,\mathcal{D},\bm{\mu})$ factorizes over the leaf nodes:

$\displaystyle\begin{split}p(\bm{y},T,\mathcal{D},\bm{\mu}|\sigma,\mathcal{E}^{(-m)})&=(2\pi\sigma^{2})^{-n/2}\times\prod_{\ell}{\left\\{-\frac{1}{2\sigma^{2}}\sum_{i\in
I(\ell)}{\left[r_{i}-\mu_{\ell}\right]^{2}}\right\\}}\\\
~{}&~{}\times\prod_{\ell}{\left[\left(2\pi\tau^{2}\right)^{-1/2}\exp\left\\{-\frac{(\mu_{\ell}-\mu_{0})^{2}}{2\tau^{2}}\right\\}\right]}.\end{split}$
(B6)

The expression in Equation (B6) immediately implies that the jumps in
$\bm{\mu}$ are conditionally independent given the decision tree
$(T,\mathcal{D}).$ In fact, we conclude that

$\mu_{\ell}|T,\mathcal{D},\bm{y},\sigma,\mathcal{E}^{(-m)}\sim\mathcal{N}\left(P_{\ell}^{-1}\theta_{\ell},P_{\ell}^{-1}\right)$

where

$\displaystyle P_{\ell}$ $\displaystyle=\sigma^{-2}\lvert
I(\ell)\rvert+\tau^{-2}$ $\displaystyle\theta_{\ell}$
$\displaystyle=\sigma^{-2}\sum_{i\in I(\ell)}{r_{i}}+\tau^{-2}\mu_{0}.$ (B7)

Thus, conditional on the decision tree $(T,\mathcal{D}),$ we must compute (i)
the number of observations contained in the leaf and (ii) the sum of the
partial residuals in each leaf in order to update the jumps.

###  B.2 The grow/prune transition kernel

Suppose that $(T_{m},\mathcal{D}_{m})=(T,\mathcal{D})$ at the start of a the
decision tree update. We draw a new value of the $m$-th decision tree with a
Metropolis-Hastings step. Specifically, we first propose a new tree
$(T^{\star},\mathcal{D}^{\star})$ by randomly perturbing $(T,\mathcal{D}).$
Then we accept the proposal and set
$(T_{m},\mathcal{D}_{m})=(T^{\star},\mathcal{D}^{\star})$ with probability

$\alpha((T,\mathcal{D})\rightarrow(T^{\star},\mathcal{D}^{\star}))=\min\left\\{1,\frac{p(T^{\star},\mathcal{D}^{\star}|\bm{y},\sigma,\mathcal{E}^{(-m)})q(T,\mathcal{D}|T^{\star},\mathcal{D}^{\star})}{p(T,\mathcal{D}|\bm{y},\sigma,\mathcal{E}^{(-m)})q(T^{\star},\mathcal{D}^{\star}|T,\mathcal{D})}\right\\},$
(B8)

where $q(\cdot|\cdot)$ is the transition density of the to-be-specific kernel.
If we reject the proposal, we leave the decision tree unchanged (i.e. we set
$(T_{m},\mathcal{D}_{m})=(T,\mathcal{D})).$

Before describing the transition kernel, observe that the ratio of posterior
densities in Equation (B8) decomposes as a product of a prior density ratio
and a marginal likelihood ratio

$\frac{p(T^{\star},\mathcal{D}^{\star}|\bm{y},\sigma,\mathcal{E}^{(-m)})}{p(T,\mathcal{D}|\bm{y},\sigma,\mathcal{E}^{(-m)})}=\frac{p(T^{\star},\mathcal{D}^{\star})}{p(T,\mathcal{D})}\times\frac{p(\bm{y}|T^{\star},\mathcal{D}^{\star},\sigma,\mathcal{E}^{(-m)})}{p(\bm{y}|T,\mathcal{D},\sigma,\mathcal{E}^{(-m)})}.$

In our simple regression setting, we obtain a closed-form expression for the
marginal decision tree likelihood by integrating out $\bm{\mu}$ from Equation
(B6):

$p(\bm{y}|T,\mathcal{D},\sigma,\mathcal{E}^{(-m)})=\prod_{\ell}{\left[\tau^{-1}\times
P(\ell)^{-1}\times\exp\left\\{-\frac{\Theta(\ell)^{2}}{2P(\ell)}\right\\}\right]},$
(B9)

where $P(\ell)$ and $\theta(\ell)$ are defined in Equation (B7).

We implemented a simple kernel that randomly grows or prunes a tree (see
Figure B3 for a cartoon illustration). Notice that we can reverse a grow
transition with a prune transition and vice versa. Conditional on growing the
tree, we select a leaf node uniformly at random, attach two new children to
the selected leaf, and draw a corresponding decision rule from the prior.
Conditional on pruning the tree, we select a so-called “no grandchild” node
uniformly at random, remove its two children from the tree, and delete its
decision rule.

Growing and pruning change at most two leaves in the tree, leaving the rest
unchanged. In Figure B3, for instance, the current tree is otherwise identical
to the pruned tree except at node $\texttt{6}.$ Similarly, the current tree
differs from the grown tree only at node $\texttt{7}.$ The local nature of the
transition kernel introduces considerable cancellation in the MH acceptance
probability. We derive these acceptance probabilities for grow moves and prune
moves below.

Figure B3: Cartoon illustration of the grow and prune proposals used to draw
new decision trees in our sampler. Observe that at most two leaf nodes are
changed in both proposals. To implement the grow/prune kernel, one must be
able to sample decision rules from the prior.

Grow moves. Suppose first that we grew $(T,\mathcal{D})$ from a leaf nx to
obtain the proposal $(T^{\star},\mathcal{D}^{\star}).$ In the proposed tree,
let $\\{X_{j(\texttt{nx})}\in\mathcal{C}(\texttt{nx})\\}$ be the random
decision rule associated with $\texttt{nx}.$ To obtain
$(T^{\star},\mathcal{D}^{\star})$ from $(T,\mathcal{D})$ from our kernel, we
must (i) decide to grow, (ii) select nx uniformly at random; and (iii) draw
the pair $(j(\texttt{nx}),\mathcal{C}(\texttt{nx}))$ from the decision rule
prior, conditional on all the decision rules at nx’s ancestors in $T.$ The
transition density for proposing $(T^{\star},\mathcal{D}^{\star})$ grown from
$(T,\mathcal{D})$ is therefore

$q(T^{\star},\mathcal{D}^{\star}|T,\mathcal{D})=q(\text{grow}|T)\times\texttt{n\\_leaf}(T)^{-1}\times
p(j(\texttt{nx}),\mathcal{C}(\texttt{nx})|T,\mathcal{D})$

where $q(\text{grow}|T)$ is the probability of attempting a grow move;
$\texttt{n\\_leaf}(T)$ counts the number of leaf nodes in $T;$ and
$p(j(\texttt{nx}),\mathcal{C}(\texttt{nx})|T,\mathcal{D})$ is the prior
density for the new decision. To compute the acceptance probability, we need
the density of the reverse transition, which is

$q(T,\mathcal{D}|T^{\star},\mathcal{D}^{\star})=q(\text{prune}|T^{\star})\times\texttt{n\\_nog}(T^{\star})^{-1},$

where $\texttt{n\\_nog}(T^{\star})$ counts the number of non-terminal nodes in
$T^{\star}$ whose children are leaf nodes (i.e. the nodes with “no
grandchildren” in $T^{\star}$).

If $(T^{\star},\mathcal{D}^{\star})$ is obtained by growing $(T,\mathcal{D})$
from $\texttt{nx},$ the ratio of prior densities turns out to be

$\frac{p(T^{\star})p(\mathcal{D}^{\star}|T^{\star})}{p(T)p(\mathcal{D}|T)}=\frac{p(j(\texttt{{nx}}),\mathcal{C}(\texttt{nx}))\times\left[0.95(1+d(\texttt{nx}))^{-2}\right]\times\left[1-0.95(2+d(\texttt{nx}))^{-2}\right]^{2}}{\left[1-0.95(1+d(\texttt{nx}))^{-2}\right]},$

where $d(\texttt{nx})$ is the depth of the node nx in tree $T.$ Hence for grow
moves, we compute

$\displaystyle\begin{split}\frac{p(T^{\star},\mathcal{D}^{\star})q(T,\mathcal{D}|T^{\star},\mathcal{D}^{\star})}{p(T,\mathcal{D})q(T^{\star},\mathcal{D}^{\star}|T,\mathcal{D})}&=\frac{q(\text{prune}|T^{\star})/\texttt{n\\_nog}(T^{\star})}{q(\text{grow}|T)/\texttt{n\\_leaf}(T)}\\\
~{}&\times\frac{\left[0.95(1+d(\texttt{nx}))^{-2}\right]\times\left[1-0.95(2+d(\texttt{nx}))^{-2}\right]^{2}}{\left[1-0.95(1+d(\texttt{nx}))^{-2}\right]}.\end{split}$
(B10)

Turning to the ratio of marginal decision tree likelihoods. Notice that there
is substantial overlap in the set of leafs in $T$ and $T^{\star}.$ The only
differences are (i) nx is a leaf in $T$ but not $T^{\star}$ and (ii)
$T^{\star}$ contains two leaf nodes (the children of nx in $T^{\star},$ which
we call nxl and nxr) that are not themselves contained in $T.$ As a result, we
have

$\displaystyle\begin{split}\frac{p(\bm{y}|T^{\star},\mathcal{D}^{\star},\sigma,\mathcal{E}^{(-m)})}{p(\bm{y}|T,\mathcal{D},\sigma,\mathcal{E}^{(-m)})}&=\tau^{-1}\times\left(\frac{P(\texttt{nxr})P(\texttt{nxl})}{P(\texttt{nx})}\right)^{-\frac{1}{2}}\\\
~{}&\times\exp\left\\{\frac{\Theta(\texttt{nxl})^{2}}{2P(\texttt{nxl})}+\frac{\Theta(\texttt{nxr})^{2}}{2P(\texttt{nxr})}-\frac{\Theta(\texttt{nx})^{2}}{2P(\texttt{nxr})}\right\\}\end{split}$
(B11)

Prune moves. Now suppose that we pruned $(T,\mathcal{D})$ by removing the
children nxl and nxr of a no grandchild node of nx in $T$ to obtain
$(T^{\star},\mathcal{D}^{\star}).$ Using essentially the same arguments for
grow moves, we compute

$\displaystyle\begin{split}\frac{p(T^{\star},\mathcal{D}^{\star})q(T,\mathcal{D}|T^{\star},\mathcal{D}^{\star})}{p(T,\mathcal{D})q(T^{\star},\mathcal{D}^{\star}|T,\mathcal{D})}&=\frac{q(\text{grow}|T^{\star})/\texttt{n\\_leaf}(T^{\star})}{q(\text{prune}|T)/\texttt{n\\_nog}(T)}\\\
~{}&\times\frac{\left[1-0.95(1+d(\texttt{nx}))^{-2}\right]}{\left[0.95(1+d(\texttt{nx}))^{-2}\right]\times\left[1-0.95(2+d(\texttt{nx}))^{-2}\right]^{2}}.\end{split}$
(B12)

and

$\displaystyle\begin{split}\frac{p(\bm{y}|T^{\star},\mathcal{D}^{\star},\sigma,\mathcal{E}^{(-m)})}{p(\bm{y}|T,\mathcal{D},\sigma,\mathcal{E}^{(-m)})}&=\tau\times\left(\frac{P(\texttt{nx})}{P(\texttt{nxr})P(\texttt{nxl})}\right)^{-\frac{1}{2}}\\\
&~{}\times\exp\left\\{\frac{\Theta(\texttt{nx})^{2}}{2P(\texttt{nx})}-\frac{\Theta(\texttt{nxl})^{2}}{2P(\texttt{nxl})}-\frac{\Theta(\texttt{nxr})^{2}}{2P(\texttt{nxr})}\right\\}.\end{split}$
(B13)

Observe that to compute the MH acceptance ratio, we do not need to evaluate
$P(\ell)$ and $\Theta(\ell)$ for every leaf in $T$ or $T^{\star}.$ To wit, we
need only compute these quantities for exactly three nodes in $T^{\star}$
(resp. $T$) for grow (resp. prune) moves. Crucially, this means that not every
observation contributes to the MH acceptance ratio. We exploit this fact in
our implementation to avoid looping over all $n$ observations in our tree
updates, which unnecessarily slows down the sampler. In our implementation, we
compute the log of the MH acceptance ratio calculations in the functions
grow_tree and prune_tree.

###  B.3 Mixing & identifiability

Remark 2 of the main text points out the slow mixing of our Gibbs sampler. The
root cause of the slow mixing over tree space is fact that the same function
can often be written as a sum-of-trees, or indeed as a single tree, in several
ways. Figure 4(a) shows two representations of a step function defined over
$[0,1]^{2}$ that assumes four different values.

(a)

(b)

(c)

Figure B4: (a) Piecewise constant functions can often be represented as a
single regression tree in multiple ways. (b & c) They can also be represented
as sums of regression trees in multiple ways.

Now consider a simplified BART model in which our ensemble contains only one
tree. Insofar as the two trees in Figure 4(a) have the same structure and
induce the same partition of the observations, they will have the same
posterior densities. Further suppose that the posterior has modes at these two
trees. To navigate between these modes using only grow and prune moves, we
must fully prune one tree back to the root node and then grow the other tree.
The intermediate states visited by the sampler will have rather low posterior
density. In other words, navigating between two modes requires traversing a
valley of low probability, which is extremely difficult using only local
moves. Recently, Ronen et al.,  (2022) and Kim and Ročková,  (2023) have
formalized this intuition, providing lower bounds of the mixing time for a
single-tree model that are exponential in the number of observations.

The situation is even more fraught in the sum-of-trees setting. For starters,
the posterior landscape display a huge degree of symmetry. Indeed, given a
particular tree ensemble $\mathcal{E},$ we can obtain a notionally different
ensemble by permuting the labels of the trees. Beyond these symmetric
representations, we can obtain equivalent fits to the data with very different
tree ensembles. For instance, in the context of Figure 4(b), we can define an
two-tree ensemble in which the tree splitting on $X_{1}$ (resp. $X_{2})$ is
the first (resp. second) tree and another in which it is the second (resp.
first) tree. And as Figure 4(c) illustrates, we can represent the single trees
in Figure 4(a) as a sum of similar tree and a stump, which consists only of
the root node. It is, in our opinion, essentially hopeless to expect our local
grow/prune transition kernel to navigate efficiently between such
representations.

The non-identifiability actually introduces more challenges beyond slow
mixing. It is exceedingly common to determine variable importance by counting
the number of decision rules based on each predictor and to “identify”
interactions by tracking the variables appearing in decision rules. The fact
that $\mathcal{E}$ is not identified complicates the use of these heuristics.
Looking at the tree on the left of Figure 4(a), we might conclude that (i)
$X_{2}$ is more important than $X_{1}$ because it appears in two decision
rules and that (ii) the function captures an interaction between $X_{1}$ and
$X_{2}.$ Looking at the tree on the right, on the other hand, we would
conclude that $X_{1}$ is more important than $X_{2}.$ The sum of trees
representation in Figure 4(b) paints yet another picture of the same function
as a sum of two main effects, one for $X_{1}$ and one for $X_{2}.$

##  Appendix C Implementation details

We now provide further details about our implementation. First, we outline
some of the important class and data structures that we introduced to
represent regression trees, decision rules, and the sets $I(\ell).$ Then, we
describe how carefully tracking and updating our representations of the
$I(\ell)$’s facilitated considerable speed-up relative to BART.

###  C.1 Class & data structures

The tree class. Regression tree nodes in our implementation are instances of a
custom C++ class, which we named $\texttt{tree}.$ The class data members
include (i) pointers to the parent, left child, and right child of the node;
(ii) an instance of a decision rule class (described below); and (iii) a
double for the jump $\mu.$ Most of the class’s function members are
convenience functions for getting and setting individual data members. For
instance, the function tree::set_mu assigns the value of the jump $\mu$
associated with the node while the function tree::get_nid returns a canonical
label for each node. These labels are computed recursively: starting from the
root node, which is labelled 1, the left (resp. right) child of a node with
label nx is $2\texttt{nx}$ (resp. $2\texttt{nx}+1$). The class also contains
functions for computing summary statistics about the tree including the
numbers of leaf nodes (tree::get_nbots) and “no grandchild” nodes, whose
children are leaf nodes (tree::get_nnog).

Beyond these basic functions, the class contains some more substantive
functions. There are two functions (tree::birth and tree::death) that perform
the necessary book-keping associated with our grow/prune transition kernel.
For instance, when we accept a grow transition that adds two children to the
node nx, tree::birth creates two new class instances, one for each child of
nx; sets the parent pointers in each of these instances; and updates the left
and right child pointers in nx’s tree instance. Similarly, when we accept a
prune transition that removes two nodes, tree::prune deletes the class
instances for these nodes; and sets the child pointers of the parent to the
null pointer. The class additionally contains functions that read in a
predictor vector $\bm{x}$444Technically, the function takes two pointers as
input, one for the continuous predictors and one for the categorical
predictors and return a pointer to the leaf node $\ell(\texttt{nx})$
containing nx (tree::get_bn) and the value $g(\bm{x};T,\mathcal{D},\bm{\mu})$
(tree::evaluate). Both functions trace the full decision-following path from
the root to the leaf $\ell(\bm{x}).$ Although we do not need to call these in
our main MCMC loop, we do utilize them when drawing posterior predictive
samples of regression function evaluations $f(\bm{x}^{\star})$ at new
predictors $\bm{x}^{\star}$. Finally, the tree class contains two functions
used to compute the set $\mathcal{A}$ of available predictor values at each
regression tree node, one for continuous predictors and one for categorical
predictors. Both functions work by recursing up the tree (by following parent
pointers) and checking whether any ancestors of a given node also split on the
same variable. These functions are called when drawing a new decision rule for
grow transitions.

In our implementation, we represent the regression tree ensemble $\mathcal{E}$
with a vector of tree instances.

The rule_t class. Unlike BART’s source code, which hardcoded decision rule
information into the tree class, we created separate class called rule_t
encoding the decision rules. Our class contains several data members including
(i) the decision variable index $j$; (ii) a Boolean encoding whether that
splitting variable is continuous or categorical; (iii) a cutpoint $c$; and
(iv) two sets, l_vals and r_vals, that record the categorical levels assigned
to the left and right child. Note that the cutpoint is set and accessed only
when $X_{j}$ is continuous and the two sets are populated only when $X_{j}$ is
categorical. Otherwise, they remain at default values ($c$’s default value is
0 and l_vals and r_vals default to std::set<int>’s of length zero).

When we propose grow moves in our sampler, we draw a new decision rule using
the function draw_rule. That function first selects the splitting variable and
then checks whether it is continuous or categorical. If it is continuous, the
function calls the function tree::get_rg_cont to determine the interval of
available values and then draws a cutpoint uniformly from that interval. If
the splitting variable is categorical, it uses tree::get_rg_cat to determine
the set of available levels. Then, draw_rule checks whether the levels are
network-structured. If they are not, it repeatedly sweeps over the levels,
assigning each to l_vals with probability 0.5, until both l_vals and r_vals
contain at least one element. If the levels are network-structured, draw_rule
forms the induced subgraph $\mathcal{G}[\mathcal{A}];$ samples a uniform
spanning tree with Wilson’s algorithm; deletes an edge from the spanning tree;
and assigns the vertices of the two connected components to l_vals and r_vals.

The suff_stat_map structure. Recall from Equation (B6) that the joint density
of the data & $m$-th regression tree, conditional on all other regression
trees and $\sigma,$ factorizes over the leaf nodes. Consequently, all
calculation required to update an individual regression tree involve
determining which observations are contained in which leafs of a tree. In a
sense, the map from observations to leaf node is a sufficient statistic,
conditional on $\sigma$ and $\mathcal{E}^{(-m)}.$ To keep track of these maps,
we created a separate structure, which we called suff_stat_map. The structure
is implemented as a std::map whose keys are integers corresponding to the
canonical labels of the leaf nodes and whose values are integer vectors of
observation indices. In our implementation, we keep track of a vector of
suff_stat_map’s, one for each regression tree. During our MCMC simulation, we
update both the regression trees and the sets $I(\ell)$ whenever we accept a
grow or prune transition. We wrote two function, compute_suff_stat_grow and
compute_suff_stat_prune that compute the $I(\ell)$’s for the proposal tree
$T^{\star}.$

###  C.2 Avoiding redundant computations

In each iteration of our MCMC simulation, we must compute the vector of
partial residuals $\bm{r}$ before updating the $m$-th tree. Perhaps the
simplest approach, based solely on Equation (B3), is to compute the fit of
each tree at each observed predictor vector $\bm{x}.$ That is, we would loop
over the trees in $\mathcal{E}^{(-m)}$ and all $n$ observations to compute
each
$g(\bm{x}_{i};T_{m^{\prime}},\mathcal{D}_{m^{\prime}},\bm{\mu}_{m^{\prime}})$
using the function tree::evaluate. Since the trees tend to be relative shallow
(thanks to the regularization prior), individual calls of tree::evaluate are
quite fast. Nevertheless, computing $M\times(M-1)\times n$ tree fits from
scratch in every MCMC iteration gets expensive, especially when $n$ is large.
Instead, following BART’s example, we keep track of two quantities,
allfit_train, which is a running estimate of $\bm{f},$ and residual, which
stores the full residual $\bm{y}-\bm{f}.$ At the beginning of the $m$-th
regression tree update, we temporarily remove the fit of the $m$-th tree from
allfit_train and add the tree fit to residual. We then pass pointers to these
arrays to other functions that perform the MH update. Finally, after updating
the regression tree, we compute the fit of the $m$-th tree, restore it to
allfit_train, and remove it from residual. In this way, when updating the
$m$-th regression tree, we need only evaluate the fit of that tree at every
observation.

Although carefully updating running estimates of $\bm{f}$ and $\bm{r}$ avoids
looping over all trees and all observations in every tree update, it still
requires a loop over the observations to compute
$g(\bm{x}_{i};T_{m},\mathcal{D}_{m},\bm{\mu}_{m}).$ When $n$ is large, even
this single loop can be prohibitively expensive. It turns out that we can
avoid calling tree::evaluate altogether by using our suff_stat_map structures.
Specifically, we can update the entries in allfit_train and residual with an
outer loop over each leaf and an inner loop over observations contained in the
leaf. Within the outer loop, we use the class function tree::get_mu to look up
the value of the jump $\mu_{\ell}$ associated with the current leaf. Then,
within the same outer loop, we loop over the observations contained in the
leaf and add or subtract (resp. add) $\mu_{\ell}$ to the entries of
allfit_train (resp. residual). We perform a similar loop after the tree has
been updated, except this time we subtract (resp. add) the jumps to the
entries of allfit_train (resp. residual). While these loops introduce some
computational overhead due to the non-contiguous access to allfit_train and
residual, they importantly avoid having to compute each
$g(\bm{x}_{i};T,\mathcal{D},\bm{\mu})$ anew each iteration. The speedup of
flexBART relative to BART largely stems from this avoidance.