⬇

}

`\minted@fvset}`

# Bayesian Additive Regression Trees for Probabilistic programming

Miriana Quiroga
IMASL-CONICET
&Pablo G Garay
IMASL-CONICET
&Juan Martin Loyola
IMASL-CONICET
Juan M. Alonso
IMASL-CONICET
&Osvaldo A Martin
IMASL-CONICET-UNSL
<EMAIL_ADDRESS>
Av. Italia 1556, San Luis. Argentina.

###### Abstract

Bayesian additive regression trees (BART) is a non-parametric method to
approximate functions. It is a black-box method based on the sum of many trees
where priors are used to regularize inference, mainly by restricting trees’
learning capacity so that no individual tree is able to explain the data, but
rather the sum of trees. We discuss BART in the context of probabilistic
programming languages (PPL), i.e., we present BART as a primitive that can be
used as a component of a probabilistic model rather than as a standalone
model. Specifically, we introduce the Python library PyMC-BART, which works by
extending PyMC, a library for probabilistic programming. We showcase a few
examples of models that can be built using PyMC-BART, discuss recommendations
for the selection of hyperparameters, and finally, we close with limitations
of our implementation and future directions for improvement.

_K_ eywords Bayesian inference $\cdot$ non-parametrics $\cdot$ PyMC $\cdot$
Python $\cdot$ binary trees $\cdot$ ensemble method

##  1 Introduction

Bayesian Additive Regression Trees introduced by Chipman et al. [2010], have
been demonstrated to be competitive compared with alternatives such as
Gaussian processes, random forests, or neural networks. The main reason is
that BART needs minimal user input and tuning while maintaining a good
performance [Chipman et al., 2010, Rockova and Saha, 2018, Hill et al., 2020,
Sparapani et al., 2021]. Additionally, and similar to other probabilistic
methods like Gaussian processes, BART provides uncertainty quantification via
probable intervals.

BART has been applied to solve numerous applications in recent years including
estimation of causal effects [Leonti et al., 2010, Hill, 2011, Hu et al.,
2022, Steele and Schwartz, 2022, Chen et al., 2022], species distribution
modelling [Carlson, 2020], estimating indoor radon concentrations [Kropat et
al., 2015], modeling of asteroid diameters [de Souza et al., 2021], genomics
[Li et al., 2022], just to name a few.

Variable selection has also been an area of interest by BART researchers,
achieving reasonable results [Bleich et al., 2014, Linero, 2018]. The main
approach is based on counting how many times a given covariable is
incorporated into the trees relative to the other covariables in the same
model. In order to improve the performance of this very simple variable
selection procedure, Linero [2018] introduced a sparsity-inducing Dirichlet
hyperprior on the splitting proportions of the regression tree prior.

In the literature, it is common to find specific BART models associated with
specific samplers and thus specific implementations [Chipman et al., 2010,
Pratola et al., 2020, Lamprinakou et al., 2020, Orlandi et al., 2021]. In
other words, general methods to sample from BART posteriors are not common and
instead, conjugate priors are used to define such BART models. Likewise,
packages like BART R [Sparapani et al., 2021], Xbart [He et al., 2019] or
bartpy [Coltman, 2022] have been designed as standalone packages restricted to
predefined models and not as tools to build arbitrary probabilistic models
with BART components.

In this work, we present PyMC-BART an implementation of BART that removes the
conflation of inference and modeling. As far as we know, this is the first
package introducing BART as a primitive for a probabilistic programming
language, and we are only aware of two previous proposals related to
generalizing BART. Namely, Tan and Roy [2019] presented the general BART
framework unifying BART extensions that were previously presented as separated
models and Linero [2022] introduced a reversible jump Markov chain Monte Carlo
algorithm to bypass the need for conjugacy. Abstracting away the
implementation of inference from modeling, as we do with PyMC-BART, has been
key to the success of applied Bayesian modeling in recent years, and the main
reason for the popularity of PPLs like PyMC [Abril-Pla et al., 2023].

As previously mentioned, PyMC-BART works by extending PyMC’s functionality.
From the user’s perspective, the main new object introduced is a BART random
variable that can use together with the built-in PyMC random variables to
build arbitrary probabilistic models in PyMC. We also provided a new inference
method111or step methods, using the PyMC nomeclature called PGBART, which
works exclusively for BART variables. For most users, there is no need to
directly interact with PGBART. Instead, it is designed in such a way that PyMC
will automatically use it whenever a BART variable occurs as part of a model.
Thus, other than using a BART variable in a model, all other modeling aspects
remain the same from a user perspective. Including model building,
conditioning on observed data, fitting the model, computing posterior
predictive samples, etc. Even more, due to the tight integration of PyMC with
ArviZ, tasks such as assessing the fit, comparing to other models, assessing
convergence, and other exploratory analysis of Bayesian models’ tasks [Kumar
et al., 2019, Martin et al., 2021] are immediately accessible to users of
PyMC-BART without the need to learn a new set of tools or syntax. Finally, we
provide a set of tools specifically designed to work with the posterior
distribution of trees, including tools to aid in the interpretation of the
fitted function, to help users perform variable selection, and also to help
diagnose the MCMC samples.

For the rest of this article, we will focus primarily on the practical aspects
of BART. We start by giving a very brief overview of the BART model in Section
2, then we continue with Section 3 describing the PyMC-BART’s API, including
which hyperparameters are available to users. In Section 4 we demonstrate
basic usage through examples. And in Section 5 we discuss how to choose the
number of trees for BART, arguably the most important hyperparameter PyMC-BART
users can change. Finally, we conclude with Section 6 discussing some
limitations of the current state of our implementation and the future of PyMC-
BART. While the target audience of this manuscript is practitioners interested
in adding BART to their Bayesian toolkit, we also provide details of the
sampler we use to approximate the posterior over the trees in Appendix A,
which we hope will help others interested in contributing to the base code.

PyMC-BART is available from the Python Package Index at
https://pypi.org/project/pymc-bart and it can also be installed using conda.
The package documentation, including installation instructions and examples of
how to use BART to conduct different statistical analyses, can be found at
https://www.pymc.io/projects/bart.

The version of PyMC-BART used for this article is 0.5.1 and the version of
PyMC is 5.1.2 [Abril-Pla et al., 2023]. All analyses are supported by
extensive documentation in the form of interactive Jupyter notebooks [Kluyver
et al., 2016] available in the paper repository on GitHub
https://github.com/Grupo-de-modelado-
probabilista/BART/tree/master/experiments, enabling readers to re-run, modify,
and otherwise experiment with the models described here on their own machines.
This repository also includes instructions on how to set up an environment
with all the dependencies used when writing this manuscript.

##  2 The BART model

A BART model can be represented as:

$Y\sim\phi\left(\sum_{i=1}^{m}G_{i}(\boldsymbol{X};\mathcal{T}_{i},\mathcal{M}_{i}),\theta\right)$
(1)

where $\boldsymbol{X}$ are the covariates and $Y$ the response variable. The
sum is done over $m$ trees. Each $G_{i}$ is a binary tree with structure,
$\mathcal{T}_{i}$ i.e., $\mathcal{T}_{i}$ specifies the set of interior nodes
(also known as splitting nodes) and the associated splitting rules, and the
set of terminal nodes (also known as leaf nodes).
$\mathcal{M}_{i}=\\{\mu_{1,i},\mu_{2,i},\cdots,\mu_{b,i}\\}$ represents the
values at the terminal nodes. For an example of a single tree used for
regression, see Figure 1. $\phi$ represents an arbitrary probability
distribution and $\theta$ parameters of $\phi$ that are not modeled as a sum
of trees, like the standard deviation for a Normal likelihood.

Figure 1: Example of a regression tree (left) and the induced partition of the
data space (right). The interior nodes are represented with rectangles. Inside
them, we can find the splitting variable $X$ and the splitting values $c_{1}$
and $c_{2}$. The terminal nodes are represented using rounded rectangles,
inside them, we find the leaf values in blue. In this example, we show a
single tree fitting the data. For BART we use a sum of trees.

The model is completed by specifying priors for $\mathcal{T}$ and
$\mathcal{M}$. For $\mathcal{T}$, independent priors are set for the depth of
the trees, the splitting variables, and the splitting values. Details for such
priors can be found in Appendix A. The overall effect of the BART priors is to
prevent overfitting by making trees shallow, making leaf node values small on
the scale of the data, and regularizing statistical interactions.

We can think of BART as priors over step functions, i.e., priors over
piecewise constant functions. Furthermore, in the limit of the number of trees
$m\to\infty$, BART converges to a nowhere-differentiable Gaussian Process
[Linero and Yang, 2018]. One may object that BART cannot directly model smooth
functions, which are arguably the most common cases for most datasets. Still,
BART is useful in practice, as judged by all its applications. Figure 2 shows
an example of BART fitted to data generated from 3 simple functions, a line, a
sine, and a step function. In all the examples, the sample size is 200. We can
see the effect of m on the result, the HDI (blue band) becomes narrower, and
the mean predicted function better fit the curves, in Section 5 we provide
some guidance on how to select m.

Figure 2: Examples of a BART model fitted to data (orange dots) from 3
different functions (black dashed lines) with m=10 (left column), m=50 (center
column) or m=200 (right column). The mean predicted function and its HDI are
shown in blue and light blue, respectively.

BART is part of the ensemble models family. Members of this family are based
on modeling a function and making predictions, from a combination of simple
models instead of from a single complex one. In the case of BART, the simple
models are the binary trees. Generally, ensemble models have desirable
properties, like being less prone to over-fit [Zhou, 2012, Kuhn and Johnson,
2013]. A drawback of ensemble methods, as well as other black-box methods, is
that a meaningful and direct interpretation of the fitted model is difficult
or even impossible; instead, models are assessed via evaluations at given
values of the covariates. Later, we present examples of partial dependence
plots [Friedman, 2001], which can be efficiently computed with PyMC-BART and
can aid users in the interpretation of the results from BART models.

##  3 API

PyMC-BART seamlessly integrates with PyMC and ArviZ; for example, priors for
$\theta$ in Equation 1, i.e., non-BART related parameters, can be arbitrarily
set by the user using the standard PyMC syntax. The same goes for fitting the
model: users need only to call the pm.sample() function, as they would
normally do for regular models.

Additionally, PyMC uses ArviZ’s InferenceData object [Kumar et al., 2019,
Martin et al., 2022] to store posterior samples, prior/posterior predictive
samples, sampler statistics, etc. (see Figure 5). InferenceData is a rich data
structure based on Xarray [Hoyer and Hamman, 2017]. For PyMC models with a
BART variable from PyMC-BART, the InferenceData object will also store the
variable inclusion record. This data can then be processed by helper functions
provided by PyMC-BART, to obtain variable importance plots, as we show in
Section 4. Users can also use ArviZ for convergence diagnostics including
$\hat{R}$, effective sample size, trace plots, and rank plots [Vehtari et al.,
2021a, Martin et al., 2021].

In the text, we assume the PyMC package has been imported with the pmb alias.

We now discuss the 3 main elements provided by PyMC-BART :

  * •

A BART random variable. This behaves similarly to default PyMC random
variables and can be used together with them, allowing users to build
arbitrary probabilistic models.

  * •

The PGBART step method. A method to sample from a BART random variable. The
user does not need to interact with this sampler, as PyMC will automatically
assign it to a BART variable if such a variable is part of the model.

  * •

A collection of helper functions to aid the user interpret the results.

###  3.1 BART random variables

For those familiar with PyMC syntax, defining BART models is straightforward.
A pmb.BART random variable behaves like other PyMC variables, with a few
caveats. It requires two mandatory arguments: X, a 2D NumPy array or a pandas
DataFrame representing the covariate matrix, and Y, a 1D NumPy array or a
pandas DataFrame representing the response vector. Y is needed because PyMC-
BART uses it to define the starting point when sampling from the BART
posterior and the variance of the leaf nodes.

Another difference is that a pmb.BART variable does not accept other PyMC
random variables as arguments. The main reason is that, as in other
implementations, the priors for the BART variables are not directly set by the
users, but instead hyperparameters are used to adjust them indirectly222While
this seems to work in practice, extensions could be considered in the future,
like using a prior over $m$, instead of a fixed number.. The hyperparameters
that can be changed by the user are:

  * •

The number of trees m. This is a positive integer that defaults to 50. For
some datasets, values as low as 20 could provide a good approximation; for
others, values as high as 200 may be needed. The value of m can be defined
using cross-validation. In our experiments, we found that Pareto smoothed
importance sampling leave-one-out cross-validation (PSIS-LOO-CV) [Vehtari et
al., 2017, 2021b], can be used to find reasonable values of m (see Section 5).

  * •

The value of alpha and beta, controlling the node depth. These parameters are
the same as those originally proposed by Chipman et al. [2010].

  * •

The prior over the splitting variables. This is uniform over the covariates X.
Users can pass an array, of the same length as the total number of covariates,
if they have prior information about the relative importance of the variables.

  * •

The response. How the leaf node values are computed. Available options are
constant (a single value at each leaf node), linear (a linear regression) or
mix (both options selected at random). Defaults to constant. Options linear
and mix are still experimental.

  * •

split_rules. Allows using different split rules for different columns. The
default is ContinuousSplitRule. The two other options are OneHotSplitRule and
SubsetSplitRule [Deshpande, 2023], both meant for categorical variables.

  * •

shape. Specify the shape of the response vector. Defaults to the length of
$Y$. It can be used to model more than one BART random variable per model. By
default, the tree structure is shared for all BART random variables but the
leaf values are independent. This can be changed by separate_trees
hyperparameter.

  * •

separate_trees. Default to False. When training multiple trees (by setting a
shape parameter), the default behavior is to learn a joint tree structure and
only have different leaf values for each. This flag forces a fully separate
tree structure to be trained instead. This is unnecessary in many cases and is
considerably slower, multiplying run-time roughly by the number of dimensions.

###  3.2 PGBART

PyMC is capable of automatically assigning different sampling algorithms to
different parameters in the same model. Thus, PyMC will use the pmb.PGBART
sampler for BART variables, and use one of the built-in PyMC samplers for
$\theta$. If $\theta$ is a continuous parameter, then PyMC will automatically
choose the pm.NUTS sampler [Hoffman and Gelman, 2014]. pmb.PGBART is a sampler
we have specifically developed for BART variables, see Appendix A for details.
These are the hyperparameters related to the sampler:

  * •

num_particles. The number of particles used to sample a new tree. Defaults to
10. In cases where the $\hat{R}$ values [Vehtari et al., 2021a] are too high,
increasing the number of particles can help.

  * •

batch. Number of trees out of the m trees fitted per step. Defaults to "auto",
which is 10% of m during and after tuning. Users can provide a tuple, with the
first element being the batch size during tuning and the second the batch size
after tuning. Increasing batch can help to reduce $\hat{R}$ but in our
experience, it is better to increase the number of particles.

###  3.3 Helper functions

PyMC-BART currently offers 4 utility functions.

  * •

pmb.plot_pdp: A function to generate partial dependence plots [Friedman,
2001]. Its main inputs are a BART random variable (once a model has been
fitted), and the covariate matrix.

  * •

pmb.plot_ice: A function to generate individual conditional expectation plots
[Goldstein et al., 2013]. Its main inputs are a BART random variable (once a
model has been fitted), and the covariate matrix.

  * •

pmb.plot_variable_importance: A function to estimate the variable importance.
Its main inputs are an InferenceData object containing the
"variable_inclusion" record in sample_stats group, a BART random variable
(once a model has been fitted), and the covariate matrix.

  * •

pmb.plot_convergence: A function that plots the empirical cumulative
distribution for the effective sample size and $\hat{R}$ values for the BART
random variables. Its main inputs are InferenceData and the name of the BART
variable.

pmb.plot_variable_importance returns two subplots (see Figure 9). The top one
shows the values of the relative variable importance, estimated as the count
of each variable across all sampled trees and then normalized so the sum of
all variables’ importance is 1. This is a standard plot in the BART
literature. The bottom subpanel shows the square of the Pearson correlation
coefficient333also known as the coefficient of determination between the full
model with all covariates and models with fewer covariates. It can be used to
find the minimal model capable of making predictions that are close to the
full model. To generate this plot, we make two important approximations in
order to reduce the computational cost:

  * •

We do not evaluate all possible combinations of variables, we simply add one
variable at a time, following their relative importance.

  * •

We do not refit the model for 1 to n components; instead, we approximate the
effect of removing variables by traversing trees from the posterior
distribution computed for the full-model and pruning branches without the
variable of interest, see Figure 3. This is similar to the procedure to
compute the partial dependence plots, with the difference that for the plots
we excluded all but one variable, and for the variable importance we start by
excluding all but the most important one, then all but the two most important
ones, until we include all variables.

Figure 3: An example of a pruned tree during the computation of a restricted
model. The tree on the left represents a tree obtained for the full model. The
tree on the right represents the tree that will be used to compute a
restricted model that does not include the variable $X_{0}$.

##  4 Examples

In this section, we use a few examples to show how to build PyMC model with
BART random variables. We are assuming the reader is already familiar with the
PyMC syntax. If that’s not the case we recommend readers spend some time
reading PyMC’s official documentation https://www.pymc.io.

###  4.1 Bikes

For our first example, we will use a dataset from the University of California
Irvine’s Machine Learning Repository
https://archive.ics.uci.edu/ml/datasets/bike+sharing+dataset. As the response
variable, we use the number of bikes rented per hour. And for the covariates
we use the hour of the day, the temperature, the humidity, and the speed of
the wind.

The proposed model is:

⬇

with pm.Model() as model_bikes:

α = pm.Exponential("α", 0.1)

μ = pmb.BART("μ", X, np.log(Y), m=50)

y = pm.NegativeBinomial("y", mu=np.exp(μ), alpha=α, observed=Y)

idata_bikes = pm.sample(tune=2000, draws=2000)

Figure 4: PyMC model for the bikes example. We have followed the import
conventions import pymc as pm and import pymc_experimental as pmx.

Where we have followed the import conventions import pymc as pm and import
pymc-bart as pmb. As was mentioned, the posterior samples are stored in
ArviZ’s InferenceData object [Kumar et al., 2019, Martin et al., 2022]. Figure
5 shows an HTML representation of this object.

Figure 5: Inference data obtained from model_bikes.

Figure 6 shows a visual diagnostics for $\alpha$ the non-BART random variable
in model_bikes such plots are a common, and useful, way to visualize sampling
convergence. On the other hand, this kind of plot can be less useful to
diagnose BART random variables. For that reason PyMC-BART offers the function
pmb.plot_convergence (see Figure 7).

Figure 6: Visual diagnostics for model_bikes generated with the command
az.plot_trace(idata_bikes, var_names=["$\alpha$"], kind="rank_bars"). On the
left, kernel density estimates computed from the posterior distribution of
$\alpha$, and on the right a rank-plot.

From Figure 7 we can see, on the left, that all values of the effective sample
size $\mu$ are above the recommended threshold (dashed vertical line) value of
400 (100 per chain). And on the right, we see that most, but not all, of the
$\mu$ values, have $\hat{R}$ below the recommended threshold, taking a few
more samples could potentially fix this issue. The value of the threshold for
$\hat{R}$ is computed automatically using a multiple comparison adjustment.

Figure 7: Visual diagnostics for model_bikes generated with the command
pmb.plot_convergence(idata_bikes, var_names=["$\mu$"]. On the left is the
empirical cumulative distribution function for the effective samples size and
on the right for $\hat{R}$.

Besides extending PyMC with BART random variables and the PGBART sampler, we
also offer a few helper functions, one of which can be used to compute partial
dependence plots [Friedman, 2001]. For instance, the command
pmb.plot_pdp(idata_bikes, X=X, Y=Y, grid=(2, 2)) generates Figure 8. This
allows the user to analyze the partial contribution of each variable.

Figure 8: Partial contribution of each variable to the number of rented bikes.

From Figure 8, we can conclude that:

  1. 1.

The marginal contribution of the variable hour varies in a more complex way
than for the other variables. Starting from a minimum between 0 and 3 hours
approximately, it increases to a first peak at around 8; then it decreases
slightly to a new, higher peak, at around 17, and finally, it decreases. This
pattern can be interpreted as a relationship between working hours and the
need to rent bikes.

  2. 2.

While the value of temperature increases, the number of rented bikes also
increases, but at some point it stabilizes. This can be explained by saying
that at higher temperatures people are more motivated to go out, but there
comes a point where this is no longer the case.

  3. 3.

humidity seems to contribute in a very slightly negative way, which could be
interpreted as meaning that high humidity does not contribute to people
wanting to ride a bike. Nevertheless, the relationship (if any) seems to be
tenuous.

  4. 4.

windspeed shows practically no contribution to the motivation to ride a bike.

Now we move our focus to the analysis of variable importance. From the top
panel of Figure 9 generated using the code
pmb.plot_variable_importance(idata_bikes, $\mu$, X, samples=100)

From the top panel, we can see that the variables hour and temperature are the
most important covariates and that the other two are less important. Notice
this is in line with the partial dependence plots from Figure 8. This kind of
plot is useful to see the relative importance of a variable but is not very
useful if we want to select a subset of the variables, as we do not have a
clear criterion to separate the variables with “high” importance from those
with “low” importance.

In order to provide a variable selection procedure from the computation of
variable importance, we introduce a new plot. We can see an instance in the
bottom panel of Figure 9. On the x-axis we have the number of variables and on
the y-axis the square of the Pearson correlation coefficient between the
predictions made by the full-model (all variables included) and the restricted
models, i.e., those with only a subset of the variables in the full-model. The
variables are included following the relative variable importance order, as
shown in the top panel. Thus, in the bikes example, the first variable
included is hour, then hour and temperature, followed by hour, temperature and
humidity and finally the four variables hour, temperature, humidity, and
windspeed, i.e., the full model. Hence, from Figure 9 we can see that even a
model with a single component, hour can generate predictions that are very
close to the full model. Moreover, the predictions from the model with the two
components hour and  temperature is, on average, indistinguishable from the
full model. The error bars represent the 94 % highest density interval (HDI)
of the posterior predictive distribution, that is, of the predictions of the
model.

Figure 9: Variable importance results for the Bikes dataset. Top, normalized
variable importance for each variable. Bottom, the square of the Pearson
correlation coefficient between the predictions made by the full model (all 4
variables are included, gray dash line) and the restricted models (1, 2, 3 or
4 variables included, blue line)

Finally, we close this example by showing that the computation of variable
importance is robust with respect to the number of trees m, as can be seen
from Figure 10. We notice this is contrary to the original proposal by Chipman
et al. [2010] where they use a low number of trees (m=20 or 25) for variable
importance and a higher one for inference. With PyMC-BART is possible to use
the same value of m for both tasks. Hence, our suggestion is to calculate the
importance of variables using the same number of trees as employed for
inference, typically ranging from 50 to 200 trees. It’s important to ensure
the adequacy of inference for our specific goals and to verify the absence of
convergence issues before proceeding.

Figure 10: Variable importance for the bike’s model. Hour and temperature are
the most important covariates. The number of trees m has little effect on the
computed values of the variable importance. The black dashed line represents
the uniform importance variable $\left(\frac{1}{4}\right)$.

###  4.2 Friedman function

The second example uses the Friedman function [Chipman et al., 2010], which
consists of generating data for the random variables
$\boldsymbol{X}=(X_{0},X_{1},...,X_{p})$ where $X_{0},X_{1},...,X_{p}$ iid
$\sim\mathcal{U}(0,1)$ and $Y=f(\boldsymbol{X})+\epsilon=10\sin(\pi
X_{0}X_{1})+20(X_{2}-0.5)^{2}+10X_{3}+5X_{4}+\epsilon$, where $\epsilon\sim
N(0,1)$.

We can see that $Y$ only depends on the first five covariates
$\boldsymbol{X}_{0:4}$. Thus, the rest of the covariates
$\boldsymbol{X}_{5:p}$ are completely irrelevant. This fact, plus the
nonlinearities and interactions, make finding $f(x)$ challenging for standard
parametric methods, and thus a good test for BART models.

To fit the data generated from the Friedman function, we use the model:

⬇

with pm.Model() as model_friedman:

σ = pm.HalfNormal(’σ’, 1)

μ = pmb.BART(’μ’, X, Y, m=200)

y = pm.Normal(’y’, μ, σ, observed=Y)

idata_friedman = pm.sample()

Figure 11: PyMC model for the Friedman example. This model is essentially the
same as that of Code Block 4.

We fit this model four times, changing at each instance the number of features
(columns) in variable X, $p\in\\{5,10,100,1000\\}$. That is, we evaluate BART
for an increasing number of irrelevant features.

Figure 12: Inference of Friedman function with dimension
$p\in\\{5,10,100,1000\\}$.

Figure 12 shows the correlation between predicted and observed data with
different numbers of covariables $p\in\\{5,10,100,1000\\}$, the true values
are in the x-axis while the y-axis contains the in-sample predictions (top
panel) and out-of-sample (bottom panel). The error bars represent the 90% HDI.
The more closely the predictions are to the true function, the closer they
will be to the black line at 45∘. For the in-sample predictions, we can see a
very good agreement between predicted and observed data even when the number
of irrelevant features is much larger than the relevant ones. As expected,
out-of-sample predictions are worse than in-sample ones. When the number of
irrelevant features is relatively high PyMC-BART predictions are not-robust
enough to the number of irrelevant features as previously observed Linero
[2018] with a sparsifying prior. Thus, our recommendation for when the number
of irrelevant covariates is very large compared to the relevant ones, like
when p=100 or p=1000 in this example, is to first check variable importance
and if the non-relevant variables represent a very large fraction of the total
number of covariates, then do a second run keeping only the most relevant
covariates. If for some reason the user doesn’t want to remove variables from
the model then increasing the number of tuning steps and the number of
particles can help to improve the out-of-sample fit, at a higher computational
cost.

Figure 13 shows a partial dependence plot for the BART model model_friedman,
with different number of covariables $p\in\\{5,10,100,1000\\}$. We can see
that the variables $\boldsymbol{X}_{5:10}$ have almost zero contribution to
the response variable $Y$444With the rest of the variables
$\boldsymbol{X}_{10:p}$, not shown here but following the same flat pattern.,
while the first 5 variables $\boldsymbol{X}_{0:4}$ have a larger effect on
$Y$. We can also see that, as we increase $p$, the trends are quite robust,
although the response becomes flatter. This effect is more clear for $X_{2}$,
and especially for $p=1000$.

Figure 13: Partial dependence plots with a number of predictors
$p\in\\{5,10,100,1000\\}$.

Next, we compute variable importance, using $p=10$ and model_friedman. We
observe in Figure 14 that the first five variables $\boldsymbol{X}_{0:4}$ are
the more important ones. This is expected from the construction of $Y$, as we
already mentioned that the variables $\boldsymbol{X}_{5:p}$ are unrelated to
the response variable $Y$. This is in agreement with Figure 13. Additionally,
from Figure 14 we can see that variable importance is virtually insensitive to
the values of m when m >= 50. With fewer trees, 10 or 20, the values of the
importance variables have more dispersion, but even in that situation, the
first 5 variables are the most important ones. Again, we can see a qualitative
agreement with Figure 13.

Our results are similar to that of Chipman et al. [2010], with the important
difference that, in our case, the variables $\boldsymbol{X}_{5:p}$ have even
less importance (almost nil), which is a better result.

Figure 14: Relatively variable importance of $\boldsymbol{X}_{0:9}$ for number
of trees m $\in\\{10,20,50,100,200\\}$. The values are normalized so that the
total variable importance sums up to 1. The black dashed line represents the
uniform importance variable $\left(\frac{1}{10}\right)$.

In Figure 15 we can observe that, as the number of irrelevant features
increases, the relative importance of the first five covariates decreases, as
expected, because the total importance of 1 has to be distributed among more
covariates. But we can also see that the relative importance is robust with
respect to an increase in the number of non-relevant covariates, because on
average the variable importance for the non-relevant covariates is smaller
than for the relevant ones.

Figure 15: Relative variable importance for $p\in\\{5,10,100,1000\\}$. The
values are normalized so that the total variable importance sums up to 1. The
mean of the importance of the variables $\boldsymbol{X}_{5:p}$ is represented
with dots.

To evaluate the impact of the hyperparameters alpha and beta, which controls
the prior for the depth of the trees, we use model_friedman for $p=10$,
iterating through m, alpha and beta, for
$\texttt{m}\in\\{10,20,50,100,200\\}$, alpha $\in\\{0.1,0.45,0.95\\}$ and beta
$\in\\{1,2,10\\}$. In Figure 16 we can see that the effect of alpha parameter
and beta parameter is overall small in comparison with the effect of $m$. We
repeat these calculations for models_bikes, models_coal, and the space
influenza toy model with similar outcomes. Because deeper trees are needed to
represent higher-order interactions, we created two modified versions of
Friedman functions by adding, to the first term, one and two more
interactions. Again, the effect of changing alpha is small. All results are
available at https://github.com/Grupo-de-modelado-
probabilista/BART/tree/master/experiments.

After this experiment and following Chipman et al. [2010], we decided to set
alpha=0.95 and beta=2 as the default values.

Figure 16: Box plots for the difference between the true Friedman function and
the predictions for $m\in\\{10,20,50,100,200\\}$ in the panels from left to
right. And $\alpha\in\\{0.1,0.45,0.95\\}$ in the top, middle and bottom row
respectively. In each panel, $\beta\in\\{1,2,10\\}$ is plotted in blue,
orange, and green, respectively

###  4.3 Cox processes

We now show how to use BART for a 1D discretized non-homogeneous Poisson
process. We use the classical coal mining disaster dataset. The same example,
but using Gaussian Processes, can be found in Martin [2018] and the
documentation of GPstuff package [Vanhatalo et al., 2013].

The data consist of timestamps for when the disaster occurred. To be able to
fit this data using a regression model, we bin the data as shown next

⬇

years = int(coal.max() - coal.min())

bins = years // 4

hist, x_edges = np.histogram(coal, bins=bins)

x_centers = x_edges[:-1] + (x_edges[1] - x_edges[0]) / 2

x_data = x_centers[:, None]

y_data = hist

Figure 17: Preprocessing of the coal dataset.

where the values of X correspond to the date of the disasters and Y are the
number of accidents for that date. Because this is a simple function, we use
m=20.

The BART model for this example is:

⬇

with pm.Model() as model_coal:

μ_ = pmb.BART("μ_", X=x_data, Y=np.log(y_data), m=20)

μ = pm.Deterministic("μ", np.exp(μ_))

y_pred = pm.Poisson("y_pred", mu=μ, observed=y_data)

idata_coal = pm.sample()

Figure 18: PyMC model for the coal mining dataset. Compared with the model in
Code Block 4, the main differences are: the use of a Poisson likelihood and
exponential inverse link function np.exp(), and a smaller number of trees
m=20.

The main differences between this model to the ones seen so far are the use of
a Poisson likelihood, the exponential inverse link function np.exp(), and a
smaller number of trees m=20.

In Figure 19, the blue line represents the mean accident rate, and the dark
and light blue bands represent the HDI 50% and 94% respectively. A notable
decrease in accidents can be observed between the years 1880 and 1900.

Figure 19: Intensity estimation for the coal mining disaster dataset.

###  4.4 Heteroscedasticity

So far we have seen examples of PyMC-BART used to model the mean function, but
we can also use it to model a non-constant variance. To exemplify such a
scenario, we are going to make use of the marketing dataset [Kassambara,
2019]. We have the budget spent on YouTube advertisement vs the effect on
sales. We decided to model the mean as a linear model, with a square root
transformation, and let BART be in charge of the standard deviation. The model
is:

⬇

with pm.Model() as model_marketing:

α = pm.HalfNormal("α", 50)

β = pm.HalfNormal("β", 5)

μ = pm.Deterministic("μ", np.sqrt(α + β * X[:, 0]))

σ_ = pmb.BART("σ_", X, np.log(Y), m=50)

σ = pm.Deterministic("σ", np.exp(σ_))

y = pm.Normal("y", μ, σ, observed=Y)

idata = pm.sample()

Figure 20: PyMC model for the marketing example. Notice, how BART is now used
to model the standard deviation of the Normal response.

In Figure 21 we see the mean function as a black line, the 94% HDI of the mean
as a darker blue band, and the 94% HDI of the standard deviation as a lighter
blue band.

Figure 21: Non-constant variance estimation for the marketing dataset.

Alternatively, we could go fully non-parametric with this example and use BART
to model, both the mean and variance.

⬇

with pm.Model() as model_marketing_full:

w = pmb.BART("w", X, Y, m=200, size=2)

y = pm.Normal("y", w[0], np.exp(w[1]), observed=Y)

idata_marketing_full = pm.sample()

Figure 22: Fully non-parametric PyMC model for the marketing example. Notice,
how BART is defined with a size = 2 argument and used to model both the mean
and standard deviation of the Normal response. Alternative to size we can pass
the shape argument, in this case shape(2, 200), with the second dimension
being the sample size.

Notice, how BART is defined with a size = 2 argument and used to model both
the mean and standard deviation of the Normal response. Alternative to size we
can pass the shape argument, in this case it should be shape(2, 200), with the
second dimension being the sample size.

##  5 Defining the number of trees

Finding a good value for the number of trees remains important when using BART
in practice. Intuitively, as this number controls the flexibility of the BART
function, it should be large enough so that the sum of trees is adequate to
explain the data, but not so large that the function becomes too flexible.
Another practical reason to avoid overshooting m is the computational cost of
BART, both in terms of time and memory, which increases with m.

First, Chipman et al. [2010] and others later, have reported that usually the
number of trees should be between 20 and 200. For implementations of BART
without sparsifying prior over the splitting variables, authors recommend a
lower value of m when computing variable importance than when doing inference.
For PyMC-BART we observed that the computation of variable importance is
robust with respect to the value of m (see for example Figures 10 and 14) and
in general we recommend using the same value of m for both inference and
variable importance assessment.

One way to tune m is to perform K-fold cross-validation, as recommended by
Chipman et al. [2010]. Another option is to approximate cross-validation by
using Pareto-smoothed importance sampling leave-one-out cross-validation
(PSIS-LOO-CV) Vehtari et al. [2017, 2021b]. The main advantage of PSIS-LOO-CV
is that we only need to fit the model once for each value of m, instead of
K-times as in K-fold cross-validation. It has been reported that PSIS-LOO-CV
can lead to overfitting Linero and Yang [2018] when used to select m and,
hence, we decided to evaluate both PSIS-LOO-CV and 5-fold cross-validation
with PyMC-BART. We used these two methods because PSIS-LOO-CV is a fast method
to compare models, with strong empirical and theoretical support, and 5-fold
cross-validation has been recommended by Chipman et al. [2010]. Notice that it
is not our intention to perform a one-to-one comparison, instead, we are
interested in providing practical recommendations for users to pick a value of
m.

For five models and datasets, the 3 simple functions from Figure 2, the models
model_bikes, the model_friedman, model_coal, and the space influenza model (a
toy-model binary classification task taken from Martin [2018], Martin et al.
[2021]) we computed PSIS-LOO-CV as implemented in ArviZ and 5-fold cross-
validation. A common pattern we observed when using PSIS-LOO-CV is that, as we
increase the value of m from 10 to 200, in successive steps, the values of the
estimated expected log pointwise predictive density (ELPD) keep increasing or,
at best, stop increasing. We observed the same pattern for 5-fold CV, i.e.,
the root-mean-square deviation (RMSD) decreases with m or, at best,
stabilizes. All the results are available at https://github.com/Grupo-de-
modelado-probabilista/BART/tree/master/experiments.

Regarding these results we want to highlight two facts, the first one is that
we observed that even for low values of m PyMC-BART is able to capture the
mean function approximately well and the effect of increasing m is to refine
the fit with a decrease of the residuals. Second, we want to remind the reader
that as m increases, the values of the leaf nodes are shrunk towards zero and,
thus, this should help mitigate overfitting.

Thus, based on our experiments, and as a rule of thumb, we recommend using
values of m around 50 or less during model building and exploration and
switching to values of m around 200 to obtain the final result. For very
simple models and/or small datasets, like the coal mining example in this
paper, low values of m around 50, could be enough. Nevertheless, probably this
is not the case for most common scenarios involving smooth functions and
datasets around 100 data points or larger.

If values of m around 200 become too costly in terms of time and/or memory, so
it becomes important to find smaller values for m, we recommend using PSIS-
LOO-CV. Let’s say we want to find a suitable value of m for the model in the
first row of Figure 2. We decide to fit the model 3 times with values of m 10,
50, and 100 and hence we will get 3 InferenceData objects, let’s call them
idata_sfl_10, idata_sfl_50, idata_sfl_200. To compute LOO and obtain a plot
like the one in Figure 24 we can do as follows:

⬇

cmp = az.compare({"m=10":idata_sfl_10,

"m=50":idata_sfl_50,

"m=200":idata_sfl_100})

\paraz.plot_compare(cmp,

plot_ic_diff=False, insample_dev=False, legend=False)

Figure 23: Example of model comparison using ArviZ to define m.

From Figure 24, we can see that the value of the ELPD follows the order
$200>50>10$, indicating that we should pick m=200. But it is important to
remark that differences equal to or smaller than 4 are considered irrelevant
and that, additionally, the uncertainty of the estimated ELPD should also be
considered. We can see that there is considerable overlap for m=50 and m=200.
Thus, picking a value of m=50, should save time and memory without a practical
impact on the fitting.

Figure 24: Model comparison using PSIS-LOO-CV for the model and data shown in
the first row of Figure 2.

##  6 Discussion and Conclusion

PyMC-BART extends PyMC, effectively combining the flexibility of BART models
with the power of a modern state-of-the-art probabilistic programming
framework. Through a series of empirical evaluations, we have given an
overview of how PyMC-BART can be used for inference of non-linear regression
models and to assess variable importance or perform variable selection.

We hope that our contribution will help more users to adopt BART models as
part of their toolbox, and we invite researchers and developers to further
improve PyMC-BART, as well as to work on similar implementations in other
probabilistic programming languages. Some aspects we believe could be improved
in the future are: Reducing memory footprint by modifying our tree
implementation (or using a Cython or Rust implementation), increasing the
efficiency of the sampler (for example we might try to use a better proposal
as described in Pratola [2016], or rotating splitting rules to explore a wider
sample space and improve mixing Maia et al. [2022]), or add new visualizations
like those in Inglis et al. [2022]. Additionally, we could increase efficiency
by modifying the initialization and tuning routine of NUTS conditional to a
BART variable in the model. Or explore other distributions for the leaf nodes
such as Gaussian Processes Maia et al. [2022], Wang et al. [2022], which can
reduce the necessary number of trees and improve extrapolation.

##  7 Acknowledgments

This research was supported by the National Agency of Scientific and
Technological Promotion ANPCyT, Grant PICT-2018-02212 (O.A.M.) and National
Scientific and Technical Research Council CONICET, Grant PIP-0087 to O.A.M. We
thank NumFOCUS, a nonprofit 501(c)(3) public charity, for the operational and
financial support of PyMC and PyMC-BART. We also want to thank Marcin
Elantkowski for reporting a bug in an early version of our code, and also for
creating a Rust implementation of PyMC-BART https://github.com/elanmart/rust-
pgbart.

##  Appendix A Sampling from BART

We use a sampler inspired by the Particle Gibbs method introduced by
Lakshminarayanan et al. [2015], but with some modifications in order to be
able to define a generalized version of BART, see Algorithm 1. We implemented
this sampler to work with PyMC, but it is important to note that it can be
implemented in other probabilistic programming languages as well. One
important consideration is that PGBART is designed to sample trees, and we
relied on PyMC’s compound step feature to sample non-BART variables. That is,
PyMC will automatically assign the PGBART sampler to a pmb.BART random
variable, and if other random variables are present in the model it will
assign other samplers to those variables. This makes our implementation of
BART flexible enough to easily accommodate a large family of BART models, and
even a combination of BART and other models, like, for example, linear
regression.

In Algorithm 1 uppercase letters represent arrays, that is $G$ is an array of
$m$ trees, $G_{i}$ represents the $i$-th tree in the array, and $G_{i}^{\mu}$
is the output of the $i$-th tree in the array. $P_{j}^{l}$ represents the log-
likelihood of the $j$-th particle tree. The first three lines are run once to
initialize the algorithm, after that the main loop body (lines 5 to 20)
constitutes one step of the sampling algorithm for a BART variable, which can
be interleaved with one step of another sampling algorithm, like NUTS, for
other variables. To reduce the computational cost at each step, we update a
subset of the $m$ trees $m_{b}\leq m$, by default $m_{b}$ is 10% of $m$. To
simplify the description of the algorithm, we have defined a few functions,
that we now describe.

Algorithm 1 The PGBART sampler.

1:$X,Y,m,n$ $\triangleright$ n is the number of particles

2:$G\leftarrow$ initialize_trees($m$)

3:$\mu$ $\leftarrow\sum_{i=1}^{m}G_{i}^{\mu}$ $\triangleright$ compute the sum
of trees

4:$\varepsilon=f(Y)$ $\triangleright$ compute the variance of the leaf nodes
from $Y$

5:$m_{b}\leq m$

6:for i:=1 to $m_{b}$ do

7: $\mu_{-i}$ $\leftarrow\mu-G_{i}^{\mu}$ $\triangleright$ the sum of trees
without the tree to replace

8: P $\leftarrow$ initialize_particles($n$)

9: while trees grow do

10: for j:=2 to $n$ do

11: $P_{j}\leftarrow$ grow_tree($P_{j},\mu,\varepsilon,X)$ $\triangleright$
attempt to grow a particle-tree

12: $P_{j}^{l}\leftarrow\log p(Y\mid\mu_{-i}+P_{j}^{\mu},\theta)$
$\triangleright$ compute conditional log-likelihood

13: end for

14: $\bar{W}\leftarrow$ normalize_weights($P_{2:N}$)

15: $P_{2:N}\leftarrow$resample($P_{2:N},\bar{W}$)

16: end while

17: $\bar{W}\leftarrow\frac{P_{j}^{l}}{\sum_{j}^{N}P_{j}^{l}}$
$\triangleright$ use the normalized log-likelihoods as weights

18: $G_{i}\leftarrow$ sample($P,\bar{W}$) $\triangleright$ sample a new tree
from all the particles

19: $\mu\leftarrow\mu_{-i}$ \+ $G_{i}$

20: update variable_inclusion

21:end for

initialize_trees

: We set all trees to the mean of $Y$ divided by m. That is, we set the sum of
trees, $\mu$, at the mean of $Y$.

initialize_particles

: In order to propose a new $G_{i}$ tree, we generate $N$ particle-trees. One
of these particles is just the tree we want to replace, $G_{i}$. This ensures
that there is a non-zero probability of keeping the current tree, instead of
accepting a new one. The rest of the particles are grown starting from scratch
(see grow_tree function below) instead of being perturbations of an existing
tree. This helps to explore larger regions.

normalize_weights

: We normalize the weights, so they are between 0 and 1, and they sum up to 1;
this is $\bar{W}$. We also compute $w$, which is the sum of unnormalized
weights, divided by the number of particles. After resampling, all particles
will have the same $w$ weight.

resample

: Based on, $\bar{W}$ we use a systematic resampling, of all but the first two
particles. This will remove particles with low probability and retain those
with higher probability. The number of particles is kept constant, meaning
some particles may be repeated.

sample

: Based on, $\bar{W}$ we select a single particle and its associated tree.

variable_inclusion

: During tuning, this function updates alpha_vec, which is the vector with the
proportions used to sample the splitting variables. After tuning, alpha_vec is
fixed, and we instead update the variable_inclusion vector, which is then
normalized and returned as the estimation of the variable importance. In both
cases, the update is done by counting the splitting variables in the returned
tree.

grow_tree

function attempts to grow a tree based on the following criteria:

node depth

: The probability that a node at depth $d=(0,1,2,...)$ is non-terminal is
given by $\alpha(1+d)^{-\beta}$ with $\alpha\in(0,1)$ and
$\beta\in[0,\infty)$. This prior was proposed and studied by Chipman et al.
[2010]. The default value is $\alpha=0.95$ and $\beta=2$.

splitting variable

: We compute the distribution over the splitting variables from the data. We
begin with a flat categorical distribution, $\alpha_{\text{vec}}$, i.e., all
covariates have the same chance of being used as splitting variables. During
the tuning phase, we continuously update $\alpha_{\text{vec}}$ based on
counting the splitting variables in the accepted trees. After the tuning
phase, we fix this distribution and use it to sample from the splitting
variables.

splitting values

: Uniform over the observed values.

leaf values

: We use $\mathcal{N}(\mu_{\text{pred}},{\varepsilon^{2}})$, where
$\mu_{\text{pred}}$ is computed as the mean of the current sum of trees
divided by the number of trees m. $\varepsilon$ is initially computed from
$Y$, being $\varepsilon=\frac{3}{\sqrt{m}}$ for binomial data and
$\varepsilon=\frac{Y_{\text{std}}}{\sqrt{m}}$ for data other than binomial.
Because the prior variance on the leaf node parameters can be model dependent
during the tuning phase the running variance of the predictions from each
fitted tree is computed and used as a proposal. After tuning the variance
remains fixed.

## References

  * Abril-Pla et al. [2023] O. Abril-Pla, V. Andreani, C. Carroll, L. Dong, C. J. Fonnesbeck, M. Kochurov, R. Kumar, J. Lao, C. C. Luhmann, O. A. Martin, M. Osthege, R. Vieira, T. Wiecki, and R. Zinkov.  Pymc: A modern and comprehensive probabilistic programming framework in python.  _PeerJ Computer Science_ , 9:e1516, 2023.  doi: 10.7717/peerj-cs.1516.
  * Bleich et al. [2014] J. Bleich, A. Kapelner, E. I. George, and S. T. Jensen.  Variable selection for bart: An application to gene regulation.  _The Annals of Applied Statistics_ , 8(3), Sep 2014.  ISSN 1932-6157.  doi: 10.1214/14-aoas755.  URL http://dx.doi.org/10.1214/14-AOAS755.
  * Carlson [2020] C. J. Carlson.  embarcadero: Species distribution modelling with bayesian additive regression trees in r.  _Methods in Ecology and Evolution_ , 11(7):850–858, 2020.  doi: https://doi.org/10.1111/2041-210X.13389.  URL https://besjournals.onlinelibrary.wiley.com/doi/abs/10.1111/2041-210X.13389.
  * Chen et al. [2022] X. Chen, M. O. Harhay, G. Tong, and F. Li.  A bayesian machine learning approach for estimating heterogeneous survivor causal effects: Applications to a critical care trial, 2022.  URL https://arxiv.org/abs/2204.06657.
  * Chipman et al. [2010] H. A. Chipman, E. I. George, and R. E. McCulloch.  BART: Bayesian additive regression trees.  _The Annals of Applied Statistics_ , 4(1):266–298, Mar. 2010.  ISSN 1932-6157.  doi: 10.1214/09-AOAS285.  URL http://projecteuclid.org/euclid.aoas/1273584455.
  * Coltman [2022] J. Coltman.  Bartpy documentation, 2022.  URL https://jakecoltman.github.io/bartpy/.
  * de Souza et al. [2021] R. S. de Souza, A. Krone-Martins, V. Carruba, R. de Cassia Domingos, E. E. O. Ishida, S. Alijbaae, M. H. Espinoza, and W. Barletta.  Probabilistic modeling of asteroid diameters from gaia DR2 errors.  _Research Notes of the AAS_ , 5(8):199, aug 2021\.  doi: 10.3847/2515-5172/ac205e.  URL https://doi.org/10.3847/2515-5172/ac205e.
  * Deshpande [2023] S. K. Deshpande.  flexbart: Flexible bayesian regression trees with categorical predictors, 2023.
  * Friedman [2001] J. H. Friedman.  Greedy function approximation: A gradient boosting machine.  _The Annals of Statistics_ , 29(5):1189 – 1232, 2001.  doi: 10.1214/aos/1013203451.  URL https://doi.org/10.1214/aos/1013203451.
  * Goldstein et al. [2013] A. Goldstein, A. Kapelner, J. Bleich, and E. Pitkin.  Peeking inside the black box: Visualizing statistical learning with plots of individual conditional expectation.  _Journal of Computational and Graphical Statistics_ , 24:44 – 65, 2013.
  * He et al. [2019] J. He, S. Yalov, and P. R. Hahn.  Xbart: Accelerated bayesian additive regression trees.  In K. Chaudhuri and M. Sugiyama, editors, _Proceedings of the Twenty-Second International Conference on Artificial Intelligence and Statistics_ , volume 89 of _Proceedings of Machine Learning Research_ , pages 1130–1138. PMLR, 16–18 Apr 2019.  URL https://proceedings.mlr.press/v89/he19a.html.
  * Hill et al. [2020] J. Hill, A. Linero, and J. Murray.  Bayesian additive regression trees: A review and look forward.  _Annual Review of Statistics and Its Application_ , 7(1):251–278, 2020.  doi: 10.1146/annurev-statistics-031219-041110.  URL https://doi.org/10.1146/annurev-statistics-031219-041110.
  * Hill [2011] J. L. Hill.  Bayesian nonparametric modeling for causal inference.  _Journal of Computational and Graphical Statistics_ , 20(1):217–240, 2011.
  * Hoffman and Gelman [2014] M. D. Hoffman and A. Gelman.  The No-U-Turn Sampler: Adaptively Setting Path Lengths in Hamiltonian Monte Carlo.  _Journal of Machine Learning Research_ , 15(1):1593–1623, 2014.
  * Hoyer and Hamman [2017] S. Hoyer and J. Hamman.  Xarray: N-D Labeled Arrays and Datasets in Python.  _Journal of Open Research Software_ , 5(1), Apr. 2017.  ISSN 2049-9647.  doi: 10.5334/jors.148.
  * Hu et al. [2022] L. Hu, J. Ji, R. D. Ennis, and J. W. Hogan.  A flexible approach for causal inference with multiple treatments and clustered survival outcomes, 2022.  URL https://arxiv.org/abs/2202.08318.
  * Inglis et al. [2022] A. Inglis, A. Parnell, and C. Hurley.  Visualizations for bayesian additive regression trees, 2022.  URL https://arxiv.org/abs/2208.08966.
  * Kassambara [2019] A. Kassambara.  datarium: Data Bank for Statistical Analysis and Visualization, May 2019.  URL https://CRAN.R-project.org/package=datarium.
  * Kluyver et al. [2016] T. Kluyver, B. Ragan-Kelley, F. Pérez, B. Granger, M. Bussonnier, J. Frederic, K. Kelley, J. Hamrick, J. Grout, S. Corlay, P. Ivanov, D. Avila, S. Abdalla, and C. Willing.  Jupyter notebooks – a publishing format for reproducible computational workflows.  In F. Loizides and B. Schmidt, editors, _Positioning and Power in Academic Publishing: Players, Agents and Agendas_ , pages 87 – 90. IOS Press, 2016.
  * Kropat et al. [2015] G. Kropat, F. Bochud, M. Jaboyedoff, J.-P. Laedermann, C. Murith, M. Palacios, and S. Baechler.  Improved predictive mapping of indoor radon concentrations using ensemble regression trees based on automatic clustering of geological units.  _Journal of Environmental Radioactivity_ , 147:51–62, 2015\.
  * Kuhn and Johnson [2013] M. Kuhn and K. Johnson.  _Applied Predictive Modeling_.  Springer-Verlang, New York, 1st ed. 2013, corr. 2nd printing 2018 edition edition, May 2013.  ISBN 978-1-4614-6848-6.
  * Kumar et al. [2019] R. Kumar, C. Carroll, A. Hartikainen, and O. A. Martin.  Arviz a unified library for exploratory analysis of Bayesian models in python.  _Journal of Open Source Software_ , 4(33):1143, 2019.  ISSN 2475-9066.  doi: 10.21105/joss.01143.
  * Lakshminarayanan et al. [2015] B. Lakshminarayanan, D. M. Roy, and Y. W. Teh.  Particle gibbs for bayesian additive regression trees, 2015.  URL https://arxiv.org/abs/1502.04622.
  * Lamprinakou et al. [2020] S. Lamprinakou, E. McCoy, M. Barahona, A. Gandy, S. Flaxman, and S. Filippi.  Bart-based inference for poisson processes, 2020.  URL https://arxiv.org/abs/2005.07927.
  * Leonti et al. [2010] M. Leonti, S. Cabras, C. S. Weckerle, M. N. Solinas, and L. Casu.  The causal dependence of present plant knowledge on herbals—contemporary medicinal plant use in campania (italy) compared to matthioli (1568).  _Journal of Ethnopharmacology_ , 130(2):379–391, 2010.
  * Li et al. [2022] Z. Li, S. Liu, W. Conaty, Q.-H. Zhu, P. Moncuquet, W. Stiller, and I. Wilson.  Genomic prediction of cotton fibre quality and yield traits using Bayesian regression methods.  _Heredity_ , pages 1–10, May 2022.  ISSN 1365-2540.  doi: 10.1038/s41437-022-00537-x.  URL https://www.nature.com/articles/s41437-022-00537-x.  Publisher: Nature Publishing Group.
  * Linero [2018] A. R. Linero.  Bayesian regression trees for high-dimensional prediction and variable selection.  _Journal of the American Statistical Association_ , 113(522):626–636, 2018.
  * Linero [2022] A. R. Linero.  Generalized bayesian additive regression trees models: Beyond conditional conjugacy, 2022.  URL https://arxiv.org/abs/2202.09924.
  * Linero and Yang [2018] A. R. Linero and Y. Yang.  Bayesian regression tree ensembles that adapt to smoothness and sparsity.  _Journal of the Royal Statistical Society B_ , 80(5):1087–1110, 2018.  doi: https://doi.org/10.1111/rssb.12293.  URL https://rss.onlinelibrary.wiley.com/doi/abs/10.1111/rssb.12293.
  * Maia et al. [2022] M. Maia, K. Murphy, and A. C. Parnell.  Gp-bart: a novel bayesian additive regression trees approach using gaussian processes, 2022.  URL https://arxiv.org/abs/2204.02112.
  * Martin [2018] O. Martin.  _Bayesian Analysis with Python: Introduction to Statistical Modeling and Probabilistic Programming Using PyMC3 and ArviZ, 2nd Edition_.  Packt Publishing, 2 edition edition, Dec. 2018.
  * Martin et al. [2022] O. Martin, A. Hartikainen, C. Carroll, and O. Abril-Pla.  Arviz, May 2022.  URL https://doi.org/10.5281/zenodo.6547007.
  * Martin et al. [2021] O. A. Martin, R. Kumar, and J. Lao.  _Bayesian Modeling and Computation in Python_.  Chapman and Hall/CRC, Boca Raton, 1st edition, 2021.  ISBN 978-0-3678-9436-8.
  * Orlandi et al. [2021] V. Orlandi, J. Murray, A. Linero, and A. Volfovsky.  Density regression with bayesian additive regression trees, 2021.  URL https://arxiv.org/abs/2112.12259.
  * Pratola [2016] M. T. Pratola.  Efficient Metropolis–Hastings Proposal Mechanisms for Bayesian Regression Tree Models.  _Bayesian Analysis_ , 11(3):885 – 911, 2016.  doi: 10.1214/16-BA999.  URL https://doi.org/10.1214/16-BA999.
  * Pratola et al. [2020] M. T. Pratola, H. A. Chipman, E. I. George, and R. E. McCulloch.  Heteroscedastic BART via Multiplicative Regression Trees.  _Journal of Computational and Graphical Statistics_ , 29(2):405–417, Apr. 2020.  ISSN 1061-8600.  doi: 10.1080/10618600.2019.1677243.  URL https://doi.org/10.1080/10618600.2019.1677243.  Publisher: Taylor & Francis _eprint: https://doi.org/10.1080/10618600.2019.1677243.
  * Rockova and Saha [2018] V. Rockova and E. Saha.  On theory for bart, 2018.  URL https://arxiv.org/abs/1810.00787.
  * Sparapani et al. [2021] R. Sparapani, C. Spanbauer, and R. McCulloch.  Nonparametric machine learning and efficient computation with bayesian additive regression trees: The bart r package.  _Journal of Statistical Software_ , 97(1):1–66, 2021.  doi: 10.18637/jss.v097.i01.  URL https://www.jstatsoft.org/index.php/jss/article/view/v097i01.
  * Steele and Schwartz [2022] K. M. Steele and M. H. Schwartz.  Causal effects of motor control on gait kinematics after orthopedic surgery in cerebral palsy: a machine-learning approach.  _medRxiv_ , 2022.  doi: 10.1101/2022.01.04.21268561.  URL https://www.medrxiv.org/content/early/2022/01/05/2022.01.04.21268561.
  * Tan and Roy [2019] Y. V. Tan and J. Roy.  Bayesian additive regression trees and the general bart model.  _Statistics in Medicine_ , 38(25):5048–5069, 2019\.  doi: https://doi.org/10.1002/sim.8347.  URL https://onlinelibrary.wiley.com/doi/abs/10.1002/sim.8347.
  * Vanhatalo et al. [2013] J. Vanhatalo, J. Riihimäki, J. Hartikainen, P. Jylänki, V. Tolvanen, and A. Vehtari.  GPstuff: Bayesian modeling with Gaussian processes.  _J. Mach. Learn. Res._ , 2013.
  * Vehtari et al. [2017] A. Vehtari, A. Gelman, and J. Gabry.  Practical Bayesian model evaluation using leave-one-out cross-validation and WAIC.  _Statistics and Computing_ , 27(5):1413–1432, Sept. 2017.  ISSN 1573-1375.  doi: 10.1007/s11222-016-9696-4.  URL https://doi.org/10.1007/s11222-016-9696-4.
  * Vehtari et al. [2021a] A. Vehtari, A. Gelman, D. Simpson, B. Carpenter, and P.-C. Bürkner.  Rank-Normalization, Folding, and Localization: An Improved $\widehat{R}$ for Assessing Convergence of MCMC (with Discussion).  _Bayesian Analysis_ , 16(2):667 – 718, 2021a.  doi: 10.1214/20-BA1221.  URL https://doi.org/10.1214/20-BA1221.
  * Vehtari et al. [2021b] A. Vehtari, D. Simpson, A. Gelman, Y. Yao, and J. Gabry.  Pareto Smoothed Importance Sampling.  _arXiv:1507.02646 [stat]_ , Feb. 2021b.  URL http://arxiv.org/abs/1507.02646.  arXiv: 1507.02646 version: 7.
  * Wang et al. [2022] M. Wang, J. He, and P. R. Hahn.  Local gaussian process extrapolation for bart models with applications to causal inference, 2022.  URL https://arxiv.org/abs/2204.10963.
  * Zhou [2012] Z.-H. Zhou.  _Ensemble Methods: Foundations and Algorithms_.  CRC press, Boca Raton, FL, 1 edition edition, June 2012.  ISBN 978-1-4398-3003-1.


$\displaystyle\leq\frac{\|\mathbf{X}\theta_{2}-\mathbf{X}\theta_{1}\|}{\min(\|\mathbf{X}\theta^{*}-\mathbf{X}\theta_{1}\|,\|\mathbf{X}\theta^{*}-\mathbf{X}\theta_{2}\|)}$
$\displaystyle\leq\frac{\sqrt{n\lambda_{max}}\|\theta_{2}-\theta_{1}\|}{\sqrt{n\lambda_{min}}}$
$\displaystyle=O(\|\theta_{2}-\theta_{1}\|)$

Thus, to get an $\epsilon$-net for $\mathcal{U}$, it suffices to construct an
$O(\epsilon)$-net for $\theta^{*}+B(0,1)$, where $B(0,1)$ is the unit sphere
in ${\mathbb{R}}^{d}$. However, it is known that for this set, we can
construct an $\epsilon$-net with $O((1/\epsilon)^{d})$ elements. From our
previous analysis, we should choose $\epsilon=O(\gamma)$. Thus, by applying a
union bound over the points in the net and then using the Lipschitzness, we
get that with probability at least $1-(1/\gamma)^{d}\delta^{\prime}$,

$\inf_{\lambda\in{\mathbb{R}},\|\theta\|=O(1)}\left(\|\lambda
A\|_{F}+\left\|\frac{\mathbf{X}\theta^{*}-\mathbf{X}\theta}{\|\mathbf{X}\theta^{*}-\mathbf{X}\theta\|}-\lambda
A\tanh\left(\frac{\beta^{*}}{\lambda}\frac{\mathbf{X}\theta^{*}-\mathbf{X}\theta}{\|\mathbf{X}\theta^{*}-\mathbf{X}\theta\|}+\mathbf{X}\theta^{*}\right)\right\|\right)=\Omega(\gamma)$

We would like to choose $\delta^{\prime}$ so that the probability of error is
at most $\delta$. This gives us

$\displaystyle(1/\gamma)^{d}\delta^{\prime}<\delta\implies
d\log(\gamma)+\log(1/\delta^{\prime})>\log(1/\delta)$

If $\gamma=1$, this inequality reduces to $\delta^{\prime}<\delta$. Otherwise,
since $\|A\|_{F}\geq 1$, it suffices to satisfy the inequality

$\log(1/\delta^{\prime})-\frac{d}{2}\log(\log(1/\delta^{\prime})+\log
n)>\log(1/\delta)$

We notice that the preceding inequality is satisfied if we set

$\log(1/\delta^{\prime})=\log(1/\delta)+d\log(\log(1/\delta)+\log n)$

Hence, we conclude that with probability at least $1-\delta$, the infimum is
at least

$\displaystyle\min\left(\frac{\|A\|_{F}}{\sqrt{\log(1/\delta)+\log
n+d\log(\log(1/\delta)+\log n)}},1\right)$
$\displaystyle\quad\quad\quad=\Omega\left(\min\left(\frac{\|A\|_{F}}{\sqrt{\log(1/\delta)+\log
n+d\log\log n+d\log(\log(1/\delta))}},1\right)\right)$

Now, all that remains is to compute the numerator of the rate. Notice that we
are in the event where the eigenvalues of $\mathbb{X}^{\top}\mathbf{X}/n$ are
bounded by positive constants. Hence, suppose $M$ is an upper bound for the
eigenvalues. We have

$\displaystyle d(f_{\theta_{1}},f_{\theta_{2}})$
$\displaystyle=\sqrt{\frac{\|\mathbf{X}\theta\|^{2}}{n}}$ $\displaystyle\leq
M\|\theta_{2}-\theta_{1}\|$

This means that it is enough to find an $\epsilon$-net for the space of
$\theta$’s with respect to the $l_{2}$-norm. This computation was also done in
Theorem 11 and the value we found was

$\sqrt{d\log n+\log(1/\delta)}$

if we want error probability $\delta$.

Hence, we get that for any $\delta>\mathrm{e}^{-n}$, with probability
$1-\delta$ we have:

$\|\mathbf{X}\hat{\theta}-\mathbf{X}\theta^{*}\|\lesssim\sqrt{d\log
n+\log(1/\delta)}\max\left(\frac{\sqrt{\log(1/\delta)+\log n+d\log\log
n+d\log(\log(1/\delta))}}{\|A\|_{F}},1\right)$

Now, we need to consider the constant $C(M)$ that is hidden with $\lesssim$.
The result of Theorem 6 states that $C(M)$ depends single-exponentially on the
bound $M$. In our case, the parameters are bounded. However,
$\mathbf{X}\theta$ is a random quantity, hence it’s maximum entry is not
necessarily bounded. We address this issue now. We immediately notice that
each coordinate of $\mathbf{X}\theta$ is a gaussian with $0$ mean and variance
at most $M$. Also, the coordinate values are all independent, since the rows
of $\mathbf{X}$ are all independent. Hence, the maximum coordinate is
distributed as the maximum of $n$ independent gaussian variables, so it’s
expectation is at most $M\sqrt{2\log n}$. Also, we know that the maximum is
also well concentrated, which means that with probability at least $1-\delta$,

$\|\mathbf{X}\theta\|_{\infty}\leq M\sqrt{2\log n}+\sqrt{\log(1/\delta)}$

Now, if $\delta>1/\text{poly}(n)$, we have that

$\|\mathbf{X}\theta\|_{\infty}=O(\sqrt{\log n})$

This implies that the constant in the rate is at most
$O\left(\mathrm{e}^{\sqrt{\log n}}\right)$, which is smaller than any
polynomial of $n$. The result now follows by recalling that we normalized $A$
by $\|A\|_{2}$ in the beginning.

∎

######  Remark 1.

Since $\|A\|_{F}\geq\|A\|_{2}$, Corollary 12 always guarantees a
$\tilde{O}(d/\sqrt{n})$ consistency rate, regardless of the norms of the
matrix $A$. We have ignored the $\exp(\sqrt{\log n})$ part of the rate, since
it is smaller than any polynomial of $n$.

###  D.2 Sparse Logistic Regression

Another interesting application of the general Theorem 6 involves the case of
Sparse Logistic Regression. In this case, we have the familiar Setting 1 of
logistic regression, with an unknown vector $\theta\in{\mathbb{R}}^{d}$ and a
fixed matrix $\mathbf{X}\in{\mathbb{R}}^{n\times d}$. The extra assumption
comes in the form of sparsity of the vector $\theta$. We will assume, as is
standard in the literature, that $\theta$ is $l_{1}$-sparse. This also enables
us to run MPLE efficiently. Clearly, we expect the rate to depend on $s$ in
that case. The quantity that determines the rate is the metric entropy for
sparse subsets of ${\mathbb{R}}^{d}$. For convenience, we restate it to
include a more detailed description of the assumptions.

######  Theorem 13.

Let $\sigma$ be sampled according to setting 1. Additionally, assume that
$\|\theta^{*}\|_{1}\leq s$. We also assume the restricted eigenvalue property
that is stated in Setting 1, which means that for all $\|\theta\|_{1}\leq s$
we have

$\|\mathbf{X}\theta\|\geq\kappa\sqrt{n}\|\theta\|$

where $\kappa$ is a constant. Let $\hat{\theta},\hat{\beta}$ be the estimates
we get for MPLE constrained in the sparse region. Then, w.pr. $\geq 1-\delta$,

$\|\hat{\theta}-\theta^{*}\|_{2}^{2}+|\hat{\beta}-\beta^{*}|^{2}\lesssim\frac{(n^{2}s\log(d))^{1/3}+\log(1/\delta)}{\|A\|_{F}^{2}}.$

###### Proof.

Notice that since $\hat{\theta},\theta^{*}$ are sparse, by the restricted
eigenvalue property we get

$\displaystyle\frac{\sum_{i=1}^{n}(f_{\hat{\theta}}-f_{\theta^{*}})^{2}}{n}$
$\displaystyle=\frac{\|\mathbf{X}(\hat{\theta}-\theta^{*})\|^{2}}{n}$
$\displaystyle\geq\kappa^{2}\|\hat{\theta}-\theta^{*}\|$

Thus, we can immediately apply Theorem 6 to obtain with probability $1-\delta$

$\|\hat{\theta}-\theta^{*}\|_{2}^{2}+|\hat{\beta}-\beta^{*}|^{2}\lesssim\frac{\inf_{\epsilon\geq
0}\left(\log\frac{n}{\delta}+\epsilon n+\log
N(\mathcal{F},X,\epsilon)\right)}{\|A\|_{F}^{2}}$

All that remains now is to bound the covering number
$N(\mathcal{F},X,\epsilon)$. Since we are in Setting 1, the norms $\|x_{i}\|$
are uniformly bounded. This means that, similarly to Theorems 11, 12, we have
that

$d(f_{\theta_{1}},f_{\theta_{2}})\leq M\|\theta_{1}-\theta_{2}\|$

Thus, is suffices to compute the covering number of the space of $\theta$’s
with respect to the $l_{2}$-norm. The space of $\theta$’s is
$\mathcal{U}=\\{\theta\in{\mathbb{R}}^{d}:\|\theta\|_{1}\leq s\\}$. Thus, we
need the covering numbers for this set. A result from [RWY11] states that

$\log N(\mathcal{U},\|\cdot\|_{2},\epsilon)\leq C\frac{s\log d}{\epsilon^{2}}$

Thus, to find the optimal $\epsilon$, we should minimize the expression

$\epsilon n+\frac{s\log d}{\epsilon^{2}}.$

This is clearly minimized when the two terms are equal, which means

$\epsilon n=\frac{s\log d}{\epsilon^{2}}\implies\epsilon=\left(\frac{s\log
d}{n}\right)^{1/3}$

Hence, the optimal value for the numerator is

$\log(1/\delta)+(n^{2}s\log(d))^{1/3}$

The result follows. ∎

###  D.3 Neural Network Regression

Another application of Theorem 6 involves the case where the regression
function $f_{\theta}$ is not linear, but a neural network. Since neural
networks are often overparametrized, we cannot hope to recover the parameter
$\theta$ is this case. However, we can still apply Theorem 6, which will yield
guarantees for estimating the output of the network. We assume we are in
Setting 2, which means that all the parameters of the network are bounded.
Again, the crucial quantity is the metric entropy for the image of these
neural networks, for which we have sufficiently strong guarantees from the
literature.

See 5

###### Proof.

To begin with, we can directly apply Theorem 6 to obtain with probability
$\geq 1-\delta$:

$\frac{\sum_{i=1}^{n}(f_{\hat{\theta}}-f_{\theta^{*}})^{2}}{n}+|\hat{\beta}-\beta^{*}|^{2}\lesssim\frac{\inf_{\epsilon\geq
0}\left(\log\frac{n}{\delta}+\epsilon n+\log
N(\mathcal{F},X,\epsilon)\right)}{\|A\|_{F}^{2}}$

Hence, we only have to bound the covering numbers of the set $\mathcal{F}$. By
the way we have defined distance $d$, it is essentially the $l_{2}$-norm of
the difference in the outputs of two networks, divided by $\sqrt{n}$. Hence,
we define $\mathcal{H}_{X}$ to be the set of outputs of neural networks
satisfying the bounds of Setting 2, when the input is the matrix $X$, composed
of features $x_{1},\ldots,x_{n}$. Then, it suffices to find an
$\epsilon\sqrt{n}$-net for $\mathcal{H}_{X}$ with respect to the $l_{2}$-norm.
Such a bound is given in Theorem 3.3 of [BFT17], which gives us (using the
fact that we have width bounded by $d$)

$\log N(\mathcal{F},X,\epsilon)\leq\frac{\log
d\sum_{i=1}^{n}\|x_{i}\|^{2}}{(\sqrt{n}\epsilon)^{2}}R^{2}=\frac{nK^{2}R^{2}\log
d}{n\epsilon^{2}}$

Clearly, to find the $\epsilon$ that minimizes the numerator, we set

$n\epsilon=\frac{K^{2}R^{2}\log
d}{\epsilon^{2}}\implies\epsilon=\left(\frac{K^{2}R^{2}\log
d}{n}\right)^{1/3}$

Using this value of $\epsilon$, the numerator becomes

$\log(1/\delta)+(n^{2}K^{2}R^{2}\log d)^{1/3},$

which concludes the proof. ∎

###  D.4 The Curie-Weiss model

The power of obtaining separate rates for estimation of parameters
$\theta,\beta$ can be easily showcased when considering the Curie Weiss model
with external field. In this model, the external field has a fixed direction
$h\in{\mathbb{R}}^{n}$. Without loss of generality, we can assume that $A$ is
the all $1/n$’s matrix, since the error introduced by this modification is
$O(1/n)$. The probability distribution becomes

${\mathbb{P}}[\sigma=y]\propto\exp\left(\theta^{*}\sum_{i=1}^{n}y_{i}h_{i}+\frac{\beta^{*}}{n}\sum_{i,j}y_{i}y_{j}\right)$
(45)

We will identify cases where it is possible to estimate $\theta$ or $\beta$.
For simplicity, we will assume that the coordinates of $h$ lie in
$\\{-1,1\\}$. The result of [DDP19] does not give us any rate, since
$\|A\|_{F}=1$. However, as we will see, the property that determines the rate
has to do with how much $h$ is close to being an eigenvector of $A$.

######  Corollary 1.

Let $\sigma$ be sampled according to (45) with $|\beta^{*}|,|\theta^{*}|\leq
M$. Suppose the external field $h$ has a fraction $\alpha$ of coordinates
equal to $1$ and a fraction $1-\alpha$ equal to $-1$. Let
$\hat{\theta},\hat{\beta}$ be the estimates when we run MPLE on $[-M,M]^{2}$.
Then, with probability at least $1-\delta$

$|\hat{\theta}-\theta^{*}|\leq C\frac{\sqrt{\log
n+\log(1/\delta)}}{\sqrt{4n(1-(2\alpha-1)^{2})}}$

###### Proof.

Let $U,V\setminus U$ be the subsets of vertices that have $h_{i}=1$ and
$h_{i}=-1$ respectively. Let’s start with acquiring a rate for $\theta^{*}$.
Since each parameter is one dimensional and bounded, according to Theorem REF,
with probability $1-\delta$

$|\hat{\theta}-\theta^{*}|\|h\|\leq C\frac{\sqrt{\log
n+\log(1/\delta)}}{\inf_{\lambda\in{\mathbb{R}}}\left(|\lambda|+\left\|\frac{h}{\|h\|}-\lambda
A\tanh\left(\frac{\beta^{*}}{\lambda}\frac{h}{\|h\|}+\theta^{*}h\right)\right\|\right)}$

which implies that

$|\hat{\theta}-\theta^{*}|\leq C\frac{\sqrt{\log
n+\log(1/\delta)}}{\inf_{\lambda\in{\mathbb{R}}}\left(|\lambda|+\left\|h-\lambda
A\tanh\left(\frac{\beta^{*}}{\lambda}h+\theta^{*}h\right)\right\|\right)}$

We have that

$\inf_{\lambda\in{\mathbb{R}}}\left(|\lambda|+\left\|h-\lambda
A\tanh\left(\frac{\beta^{*}}{\lambda}h+\theta^{*}h\right)\right\|\right)\geq\inf_{\lambda\in{\mathbb{R}}}\left\|h-\lambda
A\tanh\left(\frac{\beta^{*}}{\lambda}h+\theta^{*}h\right)\right\|$

The $\tanh$ vector has a special structure. Specifically, coordinates
corresponding to nodes in $U$ have a common value $\mu_{1}(\lambda)$ and nodes
in $V\setminus A$ have a common value $\mu_{2}(\lambda)$. Hence, if we denote
by $\mathbf{1}$ the all ones vector, we have

$\displaystyle\inf_{\lambda\in{\mathbb{R}}}\left\|h-\lambda
A\tanh\left(\frac{\beta^{*}}{\lambda}h+\theta^{*}h\right)\right\|$
$\displaystyle=\inf_{\lambda\in{\mathbb{R}}}\left(\left\|h-\lambda\left(\frac{\alpha
n\mu_{1}(\lambda)+(1-\alpha)n\mu_{2}(\lambda)}{n}\mathbf{1}\right)\right\|_{2}\right)$
$\displaystyle\geq\inf_{\lambda\in{\mathbb{R}}}\left(\|\lambda\mathbf{1}+h\|_{2}\right)$

The latter quantity is the distance of $h$ from the subspace spanned by
$\mathbf{1}$. Hence, the optimal $\lambda^{*}$ corresponds to the projection
of $h$ onto this subspace. By standard calculations, we have that

$\lambda^{*}=-\frac{\langle
h,\mathbf{1}\rangle}{\|\mathbf{1}\|^{2}}=-\frac{2\alpha n-n}{n}=1-2\alpha$

and

$\inf_{\lambda\in{\mathbb{R}}}\left(\|\lambda\mathbf{1}+h\|_{2}\right)=\sqrt{\|h\|_{2}^{2}-\frac{\langle
h,\mathbf{1}\rangle^{2}}{\|\mathbf{1}\|^{2}}}=\sqrt{4n(\alpha-\alpha^{2})}=\sqrt{4n(1-(2\alpha-1)^{2})}$

Therefore, we get that $\hat{\theta}$ is a $\sqrt{\log
n}/\sqrt{4n(1-(2\alpha-1)^{2})S}$-consistent estimator. ∎

## References

  * [Ada+19] Rados Adamczak, Micha Kotowski, Bart Polaczyk and Micha Strzelecki  “A note on concentration for polynomials in the Ising model”  In _Electronic Journal of Probability_ 24 The Institute of Mathematical Statisticsthe Bernoulli Society, 2019
  * [BDF09] Yann Bramoullé, Habiba Djebbari and Bernard Fortin  “Identification of peer effects through social networks”  In _Journal of econometrics_ 150.1 Elsevier, 2009, pp. 41–55
  * [Ber+09] Patrizia Berti, Irene Crimaldi, Luca Pratelli and Pietro Rigo  “Rate of convergence of predictive distributions for dependent data”  In _Bernoulli_ 15.4 Bernoulli Society for Mathematical StatisticsProbability, 2009, pp. 1351–1367
  * [Bes74] Julian Besag  “Spatial interaction and the statistical analysis of lattice systems”  In _Journal of the Royal Statistical Society: Series B (Methodological)_ 36.2 Wiley Online Library, 1974, pp. 192–225
  * [BFT17] Peter Bartlett, Dylan J Foster and Matus Telgarsky  “Spectrally-normalized margin bounds for neural networks”  In _arXiv preprint arXiv:1706.08498_ , 2017
  * [BLM00] Marianne Bertrand, Erzo FP Luttmer and Sendhil Mullainathan  “Network effects and welfare cultures”  In _The Quarterly Journal of Economics_ 115.3 MIT Press, 2000, pp. 1019–1055
  * [BLM13] Stéphane Boucheron, Gábor Lugosi and Pascal Massart  “Concentration inequalities: A nonasymptotic theory of independence”  Oxford university press, 2013
  * [BM18] Bhaswar B Bhattacharya and Sumit Mukherjee  “Inference in Ising models”  In _Bernoulli_ 24.1 Bernoulli Society for Mathematical StatisticsProbability, 2018, pp. 493–525
  * [BN18] Guy Bresler and Dheeraj Nagaraj  “Optimal single sample tests for structured versus unstructured network data”  In _arXiv preprint arXiv:1802.06186_ , 2018
  * [CF13] Nicholas A Christakis and James H Fowler  “Social contagion theory: examining dynamic social networks and human behavior”  In _Statistics in medicine_ 32.4 Wiley Online Library, 2013, pp. 556–577
  * [Cha05] Sourav Chatterjee  “Concentration inequalities with exchangeable pairs”, 2005
  * [Cha07] Sourav Chatterjee  “Estimation in spin glasses: A first step”  In _The Annals of Statistics_ 35.5 Institute of Mathematical Statistics, 2007, pp. 1931–1946
  * [Che+20] Ming Chen, Zhewei Wei, Zengfeng Huang, Bolin Ding and Yaliang Li  “Simple and deep graph convolutional networks”  In _International Conference on Machine Learning_ , 2020, pp. 1725–1735  PMLR
  * [CVV19] Justin Y. Chen, Gregory Valiant and Paul Valiant  “How bad is worst-case data if you know where it comes from?”  In _arXiv_ abs/1911.03605, 2019
  * [Dag+19] Yuval Dagan, Constantinos Daskalakis, Nishanth Dikkala and Siddhartha Jayanti  “Learning from Weakly Dependent Data under Dobrushin’s Condition”  In _Conference on Learning Theory_ , 2019, pp. 914–928
  * [Dag+20] Yuval Dagan, Constantinos Daskalakis, Nishanth Dikkala and Anthimos Vardis Kandiros  “Learning ising models from one or several samples”  In _arXiv preprint arXiv:2004.09370_ , 2020
  * [DDK17] Constantinos Daskalakis, Nishanth Dikkala and Gautam Kamath  “Concentration of multilinear functions of the Ising model with applications to network data”  In _Advances in Neural Information Processing Systems_ , 2017, pp. 12–23
  * [DDP19] Constantinos Daskalakis, Nishanth Dikkala and Ioannis Panageas  “Regression from dependent observations”  In _Proceedings of the 51st Annual ACM SIGACT Symposium on Theory of Computing_ , 2019, pp. 881–889
  * [DMR11] Constantinos Daskalakis, Elchanan Mossel and Sébastien Roch  “Evolutionary Trees and the Ising Model on the Bethe Lattice: A Proof of Steel’s Conjecture”  In _Probability Theory and Related Fields_ 149.1 Springer, 2011, pp. 149–189
  * [DS03] Esther Duflo and Emmanuel Saez  “The role of information and social interactions in retirement plan decisions: Evidence from a randomized experiment”  In _The Quarterly journal of economics_ 118.3 MIT Press, 2003, pp. 815–842
  * [Ell93] Glenn Ellison  “Learning, Local Interaction, and Coordination”  In _Econometrica_ 61.5 JSTOR, 1993, pp. 1047–1071
  * [Fel04] Joseph Felsenstein  “Inferring Phylogenies”  Sinauer Associates Sunderland, 2004
  * [Fen+20] Wenzheng Feng, Jie Zhang, Yuxiao Dong, Yu Han, Huanbo Luan, Qian Xu, Qiang Yang, Evgeny Kharlamov and Jie Tang  “Graph Random Neural Networks for Semi-Supervised Learning on Graphs”  In _Advances in Neural Information Processing Systems_ 33, 2020
  * [Gam03] David Gamarnik  “Extension of the PAC framework to finite and countable Markov chains”  In _IEEE Transactions on Information Theory_ 49.1 IEEE, 2003, pp. 338–345
  * [GG86] Stuart Geman and Christine Graffigne  “Markov Random Field Image Models and their Applications to Computer Vision”  In _Proceedings of the International Congress of Mathematicians_ American Mathematical Society, 1986, pp. 1496–1517
  * [GM18] Promit Ghosal and Sumit Mukherjee  “Joint estimation of parameters in Ising model”  In _arXiv preprint arXiv:1801.06570_ , 2018
  * [GSS96] Edward L Glaeser, Bruce Sacerdote and Jose A Scheinkman  “Crime and social interactions”  In _The Quarterly Journal of Economics_ 111.2 MIT Press, 1996, pp. 507–548
  * [KM15] Vitaly Kuznetsov and Mehryar Mohri  “Learning theory and algorithms for forecasting non-stationary time series”  In _Advances in neural information processing systems_ , 2015, pp. 541–549
  * [KW16] Thomas N Kipf and Max Welling  “Semi-supervised classification with graph convolutional networks”  In _arXiv preprint arXiv:1609.02907_ , 2016
  * [LHG16] Ben London, Bert Huang and Lise Getoor  “Stability and generalization in structured prediction”  In _The Journal of Machine Learning Research_ 17.1 JMLR. org, 2016, pp. 7808–7859
  * [Lon+13] Ben London, Bert Huang, Ben Taskar and Lise Getoor  “Collective stability in structured prediction: Generalization from one example”  In _International Conference on Machine Learning_ , 2013, pp. 828–836
  * [Man93] Charles F Manski  “Identification of endogenous social effects: The reflection problem”  In _The review of economic studies_ 60.3 Wiley-Blackwell, 1993, pp. 531–542
  * [MR09] Mehryar Mohri and Afshin Rostamizadeh  “Rademacher complexity bounds for non-iid processes”  In _Advances in Neural Information Processing Systems_ , 2009, pp. 1097–1104
  * [MR10] Mehryar Mohri and Afshin Rostamizadeh  “Stability bounds for stationary $\varphi$-mixing and $\beta$-mixing processes”  In _Journal of Machine Learning Research_ 11.Feb, 2010, pp. 789–814
  * [MS17] Daniel J. McDonald and Cosma Rohilla Shalizi  “Rademacher complexity of stationary sequences”  In _arXiv preprint arXiv:1106.0730_ , 2017
  * [Pei+20] Hongbin Pei, Bingzhe Wei, Kevin Chen-Chuan Chang, Yu Lei and Bo Yang  “Geom-gcn: Geometric graph convolutional networks”  In _arXiv preprint arXiv:2002.05287_ , 2020
  * [Pes10] Vladimir Pestov  “Predictive PAC learnability: A paradigm for learning from exchangeable input data”  In _Granular Computing (GrC), 2010 IEEE International Conference on_ , 2010, pp. 387–391  IEEE
  * [RWY11] Garvesh Raskutti, Martin J Wainwright and Bin Yu  “Minimax rates of estimation for high-dimensional linear regression over l_q -balls”  In _IEEE transactions on information theory_ 57.10 IEEE, 2011, pp. 6976–6994
  * [Sac01] Bruce Sacerdote  “Peer effects with random assignment: Results for Dartmouth roommates”  In _The Quarterly journal of economics_ 116.2 MIT Press, 2001, pp. 681–704
  * [SK13] Cosma Shalizi and Aryeh Kontorovich  “Predictive PAC Learning and Process Decompositions”  In _Advances in Neural Information Processing Systems_ , 2013
  * [SS14] Allan Sly and Nike Sun  “Counting in two-spin models on d-regular graphs”  In _The Annals of Probability_ 42.6 Institute of Mathematical Statistics, 2014, pp. 2383–2416
  * [TNP08] Justin G Trogdon, James Nonnemaker and Joanne Pais  “Peer effects in adolescent overweight”  In _Journal of health economics_ 27.5 Elsevier, 2008, pp. 1388–1399
  * [Ver10] Roman Vershynin  “Introduction to the non-asymptotic analysis of random matrices”  In _arXiv preprint arXiv:1011.3027_ , 2010
  * [Ver18] Roman Vershynin  “High-dimensional probability: An introduction with applications in data science”  Cambridge university press, 2018
  * [YCS16] Zhilin Yang, William Cohen and Ruslan Salakhudinov  “Revisiting semi-supervised learning with graph embeddings”  In _International conference on machine learning_ , 2016, pp. 40–48  PMLR
  * [Yu94] Bin Yu  “Rates of convergence for empirical processes of stationary mixing sequences”  In _The Annals of Probability_ JSTOR, 1994, pp. 94–116