# Co-data Learning for Bayesian Additive Regression Trees

Jeroen M. Goedhart111Corresponding author, E-mail address:
<EMAIL_ADDRESS>a, Thomas Klauscha, Jurriaan Janssenb, Mark A.
van de Wiela

###### Abstract

Medical prediction applications often need to deal with small sample sizes
compared to the number of covariates. Such data pose problems for prediction
and variable selection, especially when the covariate-response relationship is
complicated. To address these challenges, we propose to incorporate co-data,
i.e. external information on the covariates, into Bayesian additive regression
trees (BART), a sum-of-trees prediction model that utilizes priors on the tree
parameters to prevent overfitting. To incorporate co-data, an empirical Bayes
(EB) framework is developed that estimates, assisted by a co-data model, prior
covariate weights in the BART model. The proposed method can handle multiple
types of co-data simultaneously. Furthermore, the proposed EB framework
enables the estimation of the other hyperparameters of BART as well, rendering
an appealing alternative to cross-validation. We show that the method finds
relevant covariates and that it improves prediction compared to default BART
in simulations. If the covariate-response relationship is nonlinear, the
method benefits from the flexibility of BART to outperform regression-based
co-data learners. Finally, the use of co-data enhances prediction in an
application to diffuse large B-cell lymphoma prognosis based on clinical
covariates, gene mutations, DNA translocations, and DNA copy number data.

Keywords: Bayesian additive regression trees; Empirical Bayes; Co-data; High-
dimensional data; Omics; Prediction

aDepartment of Epidemiology and Data Science, Amsterdam Public Health Research
Institute, Amsterdam University Medical Centers Location AMC, Meibergdreef 9,
the Netherlands

bDepartment of Pathology, Cancer Center Amsterdam, Amsterdam University
Medical Centers Location VUMC, De Boelelaan 1117, the Netherlands

##  1 Introduction

Modern prediction models that deal with many covariates often make use of an
additional structure in the covariate-response relationship to improve
predictions and variable selection. A widely-used structure is sparsity,
employed in lasso regression (Tibshirani [1996]) and the horseshoe (Carvalho
et al. [2010]). Another method to embed a structure is by employing an a
priori weighting scheme for covariates. For instance, in the grouped lasso
(Yuan and Lin [2005]), such scheme corresponds to the amount of regularization
each group of covariates obtains, and for tree-based methods, the scheme
corresponds to the prior probabilities that covariates are selected in the
splitting rules.

Known structures can be embedded when prior knowledge on the application at
hand is available. For example, when few covariates are a priori expected to
have an effect, a sparsity structure may be assumed. Alternatively, one may
learn the structure using external data sources. External data may be discrete
such as a grouping of the covariates and, in the context of Omics, a known
gene signature, or this data may be continuous such as estimated p-values and
effect size estimates of the covariates from a previous related study. We
consider scenarios where such external data, from now on termed co-data
(Neuenschwander et al. [2010]), is available.

Several linear models learn a structure in the covariate-response relationship
from co-data. Such models often utilize a Bayesian framework to naturally
incorporate the co-data through the priors, either by employing empirical or
full Bayes. Most of these learners only handle discrete co-data, often in the
form of an a priori grouping of the covariates. Examples include GRridge (van
de Wiel et al. [2016]), which employs empirical Bayes (EB), and Graper (Velten
and Huber [2019]), which is full Bayes. However, recently, linear models that
allow multiple sources of co-data, both continuous and discrete, were
developed: Ecpc (van Nee et al. [2021]), which is EB-based, and Fwen (Tay et
al. [2023]), which differs from the aforementioned methods by utilizing non-
Bayesian estimation to learn covariate weights.

For more flexible predictions models such as tree-based models, fewer co-data
methods are available, whereas these models may greatly benefit from an
additional structure to reduce the search in parameter space. A notable
contribution along these lines is CoRF, which learns covariate weights from
both discrete and continuous co-data to enrich a random forest model (te Beest
et al. [2017]). Here, we develop an EB-based method to incorporate discrete
and continuous co-data into Bayesian additive regression trees (BART) (Chipman
et al. [2010]).

BART is a sum-of-trees prediction model, embedded in a Bayesian framework,
which may estimate continuous and binary (Chipman et al. [2010]), survival
(Sparapani et al. [2016]), and count responses (Murray [2021]). The main
advantage of a sum-of-trees model is that both additive and interaction
effects are captured and that these effects are learned nonparametrically from
the primary data. To reduce the flexibility of the sum-of-trees, BART utilizes
regularization priors on the tree parameters to prevent overfitting by
favoring shallow trees and small response estimates. BART distinguishes itself
from random forest by mainly modeling low-order interaction effects, whereas a
random forest by default grows deeper trees and hence incorporates high-order
interaction effects. Therefore, BART is more interpretable and it lends itself
better for variable selection compared to a random forest.

BART has already been extended to embed additional structures in the
covariate-response relationship. DART, proposed by Linero [2018], incorporates
a sparse structure, while OG-BART incorporates a discrete, possibly
overlapping, grouping structure with an additional sparsity layer (Du and
Linero [2019]). Both methods employ full Bayes to estimate the structure. A
disadvantage is that those methods are specifically designed for sparse
settings. However, high-dimensional applications may often be rather dense, as
argued for many omics settings (Boyle et al. [2017]).

We present EB-coBART, a BART model that learns from both the primary data and
co-data to embed a structure on the covariates in the form of a weighting
scheme. Bleich et al. [2014] showed that a weighting scheme for the covariates
based on prior knowledge may lead to better predictions compared to default
BART. However, a fairly crude method was proposed by putting twice as much
weight on covariates that were known to be important. It is in general
nontrivial how to determine those weights. Furthermore, it is difficult to
judge the informativeness of the prior knowledge for the data set at hand. EB-
coBART has a more objective and data-driven approach by learning the weighting
scheme by Empirical Bayes combined with a parsimonious co-data model. The
parsimony protects against overfitting and ensures stable estimated weights
for the covariates.

Our method has several assets compared to aforementioned approaches to co-data
learning. First, EB-coBART handles both discrete and continuous co-data.
Second, the weighting scheme of the covariates is learned by both the primary
and the co-data, which protects better against overfitting compared to
approaches that learn the weights only from the primary data. Third, it does
not require sparsity and hence many (small) effects may be modeled. Fourth,
co-data is naturally incorporated using Empirical Bayes (EB), in contrast to
CoRF which relies on a more ad hoc incorporation of the co-data (te Beest et
al. [2017]). Furthermore, EB is easily implemented in existing BART software,
whereas full Bayes approaches require rebuilding the algorithm from scratch.

In addition, our EB framework may also be used to estimate the other
hyperparameters of BART, thereby avoiding cross-validation. While we focus on
estimating co-data driven hyperparameters, we also provide EB estimators for
the other hyperparameters and showcase them in an application.

The remainder of this article is organized as follows. We start by reviewing
BART in Section 2. In Section 3, we first derive a general empirical Bayes
framework that may estimate any hyperparameter of BART. Then, we describe how
the co-data model is included to this framework to estimate the covariate
weights. Next, we present simulations in Section 4 and we illustrate our
method in an application to lymphoma patients in Section 5. We conclude with a
summary and a discussion in Section 6

##  2 Summary of BART

Bayesian additive regression trees (BART), proposed by Chipman et al. [2010],
is a Bayesian sum-of-trees model. The main idea is to sum several weak
learners, i.e. the individual trees, to form a powerful prediction model.
Individual trees are discouraged to have a large effect on the sum by
employing regularization priors that favor shallow trees and small terminal
node values. Because our proposed method relies on several aspects of the
original BART model, we present a detailed summary. We focus on regression,
but BART and our proposed method also generalize to classification using the
data augmentation method of Albert and Chib [1993].

###  2.1 A sum of trees model with regularization priors

Suppose we have a data set
$D=\left\\{\left(y_{i},\boldsymbol{x}_{i}\right)\right\\}_{i=1}^{N}$
consisting of $N$ observations of a normally distributed response $y_{i}$ and
a $p$-dimensional vector of covariates
$\boldsymbol{x}_{i}=\left(x_{1},...,x_{p}\right)$ with $x_{j}$ representing
covariate $j.$ In prediction, we then model $y_{i}$ by
$y_{i}=f\left(\boldsymbol{x}_{i}\right)+\epsilon_{i},\quad\epsilon_{i}\sim\mathcal{N}\left(0,\sigma^{2}\right),$
and we aim to estimate the function $f\left(\boldsymbol{x}_{i}\right).$

BART approximates $f\left(\boldsymbol{x}_{i}\right)$ by a sum-of-trees model
$G\left(\boldsymbol{x}_{i};\boldsymbol{\mathcal{T}},\boldsymbol{\mathcal{M}}\right)$
with input $\boldsymbol{x}_{i}$ and parameters $\boldsymbol{\mathcal{T}}$ and
$\boldsymbol{\mathcal{M}}$:

$f\left(\boldsymbol{x}_{i}\right)\approx
G\left(\boldsymbol{x}_{i};\mathcal{\boldsymbol{\mathcal{T}}},\boldsymbol{\mathcal{M}}\right)\equiv\sum_{t=1}^{K}g_{t}\left(\boldsymbol{x}_{i};\mathcal{T}_{t},\mathcal{M}_{t}\right).$
(1)

Here, $g_{t}$ denotes the $t$th regression tree, having tree structure
parameter $\mathcal{T}_{t}$ and terminal node parameter $\mathcal{M}_{t},$ and
$K$ denotes the number of trees. Parameters $\boldsymbol{\mathcal{T}}$ and
$\boldsymbol{\mathcal{M}}$ collect all tree structures $\mathcal{T}_{t}$ and
all terminal node parameters $\mathcal{M}_{t},$ respectively.

In BART, the sum-of-trees model is embedded in a Bayesian framework with a
Gaussian likelihood

$\pi\left(\boldsymbol{y}\mid\boldsymbol{X},\mathcal{\boldsymbol{\mathcal{T}}},\boldsymbol{\mathcal{M}},\sigma^{2}\right)=\prod_{i=1}^{N}\mathcal{N}\left(y_{i};G\left(\boldsymbol{x}_{i};\mathcal{\boldsymbol{\mathcal{T}}},\boldsymbol{\mathcal{M}}\right),\sigma^{2}\right),$
(2)

with $\boldsymbol{y}=\left(y_{1},...,y_{N}\right)^{T}$ and
$\boldsymbol{X}\boldsymbol{=}\left(\boldsymbol{x}_{1}^{T},...,\boldsymbol{x}_{N}^{T}\right)^{T},$
and a regularization prior distribution
$\pi\left(\boldsymbol{\mathcal{T}},\mathcal{\boldsymbol{M}}\right)$. For the
error variance $\sigma^{2},$ the standard prior
$\pi\left(\sigma^{2}\right)=\mathcal{IG}\left(\sigma^{2};\frac{\nu}{2},\frac{\nu\lambda}{2}\right),$
with hyperparameters $\nu$ and $\lambda,$ is employed.

###  2.2 Prior specification of the tree parameters

The prior $\pi\left(\boldsymbol{\mathcal{T}},\mathcal{\boldsymbol{M}}\right)$
on the tree parameters is simplified by assuming that each tree $t$ is
independently and identically distributed with prior
$\pi\left(\mathcal{T}_{t},\mathcal{M}_{t}\right)=\pi\left(\mathcal{M}_{t}\mid\mathcal{T}_{t}\right)\pi\left(\mathcal{T}_{t}\right):$

$\pi\left(\boldsymbol{\mathcal{T}},\boldsymbol{\mathcal{M}}\right)=\prod_{t=1}^{K}\pi\left(\mathcal{M}_{t}\mid\mathcal{T}_{t}\right)\pi\left(\mathcal{T}_{t}\right),$
(3)

with tree structure prior $\pi\left(\mathcal{T}_{t}\right)$ and terminal node
prior $\pi\left(\mathcal{M}_{t}\mid\mathcal{T}_{t}\right).$

Priors $\pi\left(\mathcal{T}_{t}\right)$ and
$\pi\left(\mathcal{M}_{t}\mid\mathcal{T}_{t}\right)$ are set such that each
tree is expected to only contribute slightly to the overall fit
$f\left(\boldsymbol{x}\right)$. This regularization is achieved by favoring
shallow trees with small terminal node values. Before specifying these priors,
we review the, for us, relevant parameters of a tree.

####  2.2.1 Tree parameterization

The tree structure $\mathcal{T}_{t}$ of tree $t$ is parameterized by a set of
nodes which are internal or terminal

$\mathcal{T}_{t}=\left(\xi_{1t}(x_{j},a_{j},d),\ldots,\xi_{Z_{t}t}(x_{j},a_{j},d),\omega_{1t}(d),\ldots,\omega_{L_{t}t}(d)\right),$
(4)

with $z=1,\ldots,Z_{t}$ indexing the internal nodes $\xi_{zt},$ and
$l=1,\ldots,L_{t}$ indexing the terminal nodes $\omega_{lt}.$ An internal node
$\xi$ is parameterized by a binary splitting rule $\left\\{x_{j}\leq
a_{j}\right\\},$ with $x_{j}$ the chosen splitting variable and $a_{j}$ a
splitting value within the range of covariate $j.$ Both the internal nodes and
the terminal nodes have a depth parameter $d,$ with the root node having
$d=0.$ Index variables $z$ and $l$ move from top to bottom and from left to
right. We omitted the topological structure of the tree.

Terminal node $\omega_{lt}$ also has a parameter $\mu_{lt}$, which represents
the response estimate for the given node. All $L_{t}$ estimates are collected
in the terminal node parameter
$\mathcal{M}_{t}=\left(\mu_{1t},\ldots,\mu_{L_{t}t}\right)$ for tree $t.$

####  2.2.2 Tree structure prior

The tree structure prior $\mathcal{T}_{t}$ is chosen such that nodes have a
probability of $\alpha\left(1+d\right)^{-\beta}$ to be internal and
$1-\alpha\left(1+d\right)^{-\beta}$ to be terminal, with hyperparameters
$\alpha\in\left(0,1\right)$ and $\beta>0.$ Hyperparameters $\alpha$ and
$\beta$ are usually set at $\alpha=0.95$ and $\beta=2.$ Hence, shallow trees
are a priori favored because it becomes more likely that a node is terminal
for a larger depth $d$.

For internal nodes, splitting variables need to be assigned. Splitting
variables $x_{j}$ are chosen from a categorical prior with prespecified
probabilities $\boldsymbol{S}$ for each $j$:
$\boldsymbol{S}=\left(s_{1},\ldots,s_{p}\right).$ Default BART sets equal
covariate weights: $s_{j}=1/p.$ In our proposed method, $\boldsymbol{S}$ will
be estimated using Empirical Bayes and co-data.

The above specification leads to the following prior on tree structure
$\mathcal{T}_{t}:$

$\pi_{\alpha,\beta,\boldsymbol{S}}\left(\mathcal{T}_{t}\right)\propto\left[\prod_{z=1}^{Z_{t}}\textrm{Categorical}\left(x_{jzt};\boldsymbol{S}\right)\right]\left[\prod_{z=1}^{Z_{t}}\alpha\left(1+d_{zt}\right)^{-\beta}\right]\left[\prod_{l=1}^{L_{t}}1-\alpha\left(1+d_{lt}\right)^{-\beta}\right],$
(5)

with $Z_{t}$ the number of internal nodes of tree $t$, $x_{jzt}$ covariate $j$
occurring in the splitting rule of the $z$th internal node of tree $t,$
$L_{t}$ the number of terminal nodes of tree $t$, and $d_{zt}$ and $d_{lt}$
the depths of the $z$th internal node and the $l$th terminal node of tree $t$,
respectively. The subscript denotes the dependence of the prior on
hyperparameters $\alpha,$ $\beta,$ and $\boldsymbol{S}.$ We suppressed the
dependence of the prior on the splitting values $a_{jz}$ because these are not
relevant for the purpose of this paper.

####  2.2.3 Terminal node prior

The $L_{t}$ terminal node parameters $\mu_{lt},$ collected in
$\mathcal{M}_{t},$ of tree $t$ are a priori assumed to be independent and
identically distributed with a centered normal distribution:

$\pi_{k}\left(\mathcal{M}_{t}\mid\mathcal{T}_{t}\right)=\prod_{l=1}^{L_{t}}\mathcal{N}\left(\mu_{lt};0,\sigma_{\mu}^{2}\right),\qquad\sigma_{\mu}=\frac{0.5}{k}\sqrt{K},$
(6)

with hyperparameter $k$ typically $1\leq k\leq 3,$ This prior favors small
terminal node values and hence reduces the effect of a single tree on the
overall fit. This reduction is more prominent for larger $k$ or number of
trees $K.$

###  2.3 Draws from the posterior

By employing the regularizing priors, given by (3), (5), and (6), and the
likelihood, i.e. (2), a posterior of the model parameters
$\left(\mathcal{\boldsymbol{\mathcal{T}}},\boldsymbol{\mathcal{M}},\sigma^{2}\right)$
may be deduced. Because of the large discrete tree space, this posterior is
analytically intractable, and hence a Gibbs sampler is employed. This sampling
algorithm is described elsewhere (Tan and Roy [2019], Kapelner and Bleich
[2016]). In Supplementary Section 1, we provide a short summary of this
algorithm using our notation.

The Gibbs samples represent, after a suitable burn-in period, posterior
samples of the sum-of-trees
$\sum_{t=1}^{K}g_{t}\left(\boldsymbol{x}_{i};\mathcal{T}_{t},\mathcal{M}_{t}\right)$,
which then allow inference on the response $Y.$ We assess convergence of the
MCMC chain by estimating the Gelman-Rubin diagnostic (Gelman and Rubin [1992])
for the error variance $\sigma^{2}$ for continuous responses and for
$\sum_{t=1}^{K}g_{t}\left(\boldsymbol{x}_{i};\mathcal{T}_{t},\mathcal{M}_{t}\right)$
for binary responses, as proposed by Sparapani et al. [2021].

##  3 EB-coBART

It may be beneficial for BART to upweight certain (groups of) covariates that
are expected to be informative for the response $Y.$ In BART, this is
naturally done by modifying hyperparameter
$\boldsymbol{S}=\left(s_{1},\ldots,s_{p}\right)$ of the categorical splitting
variable prior, which normally defaults to a discrete uniform: $s_{j}=1/p$.
Our main contribution is to impose a weighting structure in $\boldsymbol{S}$
in a data-driven manner. To do so, we combine empirical Bayes (EB) with a co-
data model.

Empirical Bayes allows adaptively learning the structure in $\boldsymbol{S}$
with a central role for the primary data
$D=\left\\{\left(y_{i},\boldsymbol{x}_{i}\right)\right\\}_{i=1}^{N}.$ However,
estimating $\boldsymbol{S}$ by pure EB requires estimating a $p$-dimensional
hyperparameter, which likely leads to overfitting. Hence, we model the EB
estimate of $\boldsymbol{S}$ by a co-data model, which renders a substantial
reduction in dimension of the estimated hyperparameter. Additionally, the co-
data model ensures a natural incorporation of external information on the
covariates into BART.

We start by deriving EB for BART with the primary goal to estimate
$\boldsymbol{S}$. In addition, we show EB estimates of the other
hyperparameters $\left(\alpha,\beta,k,\nu,\lambda\right)$, which are normally
estimated by cross-validation. Next, we describe our co-data model that guides
the EB-estimate of $\boldsymbol{S}$. This completes our proposed methodology,
which we call EB-coBART. We end this section with some remarks on choices for
the other hyperparameters.

###  3.1 Empirical Bayes for BART

We collect all hyperparameters of BART in the vector
$\boldsymbol{\rho}=\left(\alpha,\beta,k,\nu,\lambda,\boldsymbol{S}\right),$
which we then estimate by $\hat{\boldsymbol{\rho}}$ using empirical Bayes
(EB). To do so, we maximize the marginal likelihood of BART w.r.t.
$\boldsymbol{\rho}$:

$\hat{\boldsymbol{\rho}}=\arg\max_{\boldsymbol{\rho}}\;m_{\boldsymbol{\rho}}\left(\boldsymbol{y}\mid\boldsymbol{X}\right),$
(7)

with $m_{\boldsymbol{\rho}}\left(\boldsymbol{y}\mid\boldsymbol{X}\right)$ the
marginal likelihood of BART and subscript $\boldsymbol{\rho}$ denoting
dependence on the hyperparameters.

Boatman et al. [2021] showed that, under a reparameterization of the terminal
node prior, the marginal likelihood requires a summation over the full tree
space. This summation is exponentially large and prevents direct maximization
of the marginal likelihood. We therefore rely on an approximation algorithm.

####  3.1.1 Monte Carlo EM algorithm

To approximate (7), we employ the Monte Carlo EM algorithm derived by Casella
[2001]. In essence, this algorithm utilizes Gibbs samples to approximate the
expectation step of the EM algorithm applied to the marginal likelihood. For
BART, this method amounts to the following iterative algorithm:

$\displaystyle\boldsymbol{\hat{\rho}}^{(q+1)}$
$\displaystyle\approx\underset{\boldsymbol{\rho}}{\arg\max}\frac{1}{n_{mc}}\sum_{m=1}^{n_{mc}}\log\left[\pi\left(\boldsymbol{y}\mid\boldsymbol{X},\boldsymbol{\mathcal{T}}_{m}^{(q)},\boldsymbol{\mathcal{M}}_{m}^{(q)},\sigma_{m}^{2(q)}\right)\pi_{\boldsymbol{\rho}}\left(\boldsymbol{\mathcal{T}}_{m}^{(q)},\boldsymbol{\mathcal{M}}_{m}^{(q)},\sigma_{m}^{2(q)}\right)\right]$
$\displaystyle=\underset{\boldsymbol{\rho}}{\arg\max}\sum_{m=1}^{n_{mc}}\log\left[\pi_{\boldsymbol{\rho}}\left(\boldsymbol{\mathcal{T}}_{m}^{(q)},\boldsymbol{\mathcal{M}}_{m}^{(q)},\sigma_{m}^{2(q)}\right)\right],$
(8)

with $m$ indexing the Gibbs samples, $n_{mc}$ the total number of Gibbs
samples, $q$ the iteration index, and
$\left(\boldsymbol{\mathcal{T}}_{m}^{(q)},\boldsymbol{\mathcal{M}}_{m}^{(q)},\sigma_{m}^{2(q)}\right)$
the $m$th posterior sample of the model parameters of BART with the posterior
evaluated at $\boldsymbol{\rho}=\hat{\boldsymbol{\rho}}^{(q)}.$ Thus,
$\pi_{\boldsymbol{\rho}}\left(\boldsymbol{\mathcal{T}}_{m}^{(q)},\boldsymbol{\mathcal{M}}_{m}^{(q)},\sigma_{m}^{2(q)}\right)=\pi_{\alpha,\beta,\boldsymbol{S},k}\left(\boldsymbol{\mathcal{T}}_{m}^{(q)},\boldsymbol{\mathcal{M}}_{m}^{(q)}\right)\pi_{\nu,\lambda}\left(\sigma_{m}^{2(q)}\mid\boldsymbol{\mathcal{T}}_{m}^{(q)},\boldsymbol{\mathcal{M}}_{m}^{(q)}\right)$
evaluates the prior probabilities of these posterior samples as a function of
the hyperparameters $\boldsymbol{\rho}.$ The prior probabilities are evaluated
using (3), (5), (6), and
$\pi_{\nu,\lambda}\left(\sigma_{m}^{2(q)}\mid\boldsymbol{\mathcal{T}}_{m}^{(q)},\boldsymbol{\mathcal{M}}_{m}^{(q)}\right)=\mathcal{IG}\left(\sigma_{m}^{2(q)};\frac{\nu}{2},\frac{\nu\lambda}{2}\right).$
A derivation of (8) is given in Supplementary Section $2$ (eq. $2.5$).

Solving (8) then renders the iterative estimates of the hyperparameters
$\left(\alpha,\beta,k,\nu,\lambda,\boldsymbol{S}\right):$

$\displaystyle\left(\hat{\alpha}^{(q+1)},\hat{\beta}^{(q+1)}\right)$
$\displaystyle=\underset{\alpha,\beta}{\arg\max}\sum_{m=1}^{n_{mc}}\sum_{t=1}^{K}\left[\sum_{z=1}^{Z_{tm}^{(q)}}\log\left(\alpha\left(1+d_{ztm}^{(q)}\right)^{-\beta}\right)+\sum_{l=1}^{L_{tm}^{(q)}}\log\left(1-\alpha\left(1+d_{ltm}^{(q)}\right)^{-\beta}\right)\right],$
(9) $\displaystyle\hat{k}^{(q+1)}$
$\displaystyle=\frac{2\sum_{m=1}^{n_{mc}}\sum_{t=1}^{K}L_{tm}^{(q)}}{\sqrt{\sum_{m=1}^{n_{mc}}\sum_{t=1}^{K}\sum_{l=1}^{L_{tm}^{(q)}}\left(\mu_{ltm}^{(q)}\right)^{2}}\sqrt{K}},$
(10) $\displaystyle\left(\hat{\nu}^{(q+1)},\hat{\lambda}^{(q+1)}\right)$
$\displaystyle=\underset{\nu,\lambda}{\arg\max}\sum_{m=1}^{n_{mc}}\log\mathcal{IG}\left(\sigma_{m}^{2(q)};\frac{\nu}{2},\frac{\nu\lambda}{2}\right),$
(11) $\displaystyle\boldsymbol{\hat{S}}_{1}^{(q+1)}$
$\displaystyle=\left(b_{1}^{(q)}/B^{(q)},\ldots,b_{p}^{(q)}/B^{(q)}\right),$
(12)

with $b_{j}^{(q)}$ the sampled number of splitting rules with covariate $j$ at
iteration $q$, and $B^{(q)}$ the total number of sampled splitting rules. At
iteration $q,$ samples are taken from the posterior of BART using
hyperparameter $\hat{\boldsymbol{\rho}}^{(q)}.$ Derivations of the EB-
estimates are found in Supplementary Section $2$.

Equation (12) shows feature-specific EB estimates
$\boldsymbol{\hat{S}}_{1}^{(q+1)}$ of hyperparameter $\boldsymbol{S},$ which
likely leads to overfitting because $p-1$ hyperparameter estimates are
required. To address this issue, we propose to model the feature-specific
estimates using co-data, rendering a new estimator
$\boldsymbol{\hat{S}}_{2}^{(q+1)}$ of $\boldsymbol{S}.$

###  3.2 Co-data model

Suppose we have $\kappa$ co-data variables of which we have complete
measurements for each covariate $j$. We represent a grouping co-data variable
having $G$ groups by dummy coding indicating which group covariate $j$ belongs
to. Covariates with missing co-data observations may be defined as a separate
group to account for potential information in the missingness (van Nee et al.
[2021]).

We collect the co-data measurements of each $j$ in the vector
$\boldsymbol{c}_{j}\in\mathbb{R^{\kappa}}$ and define the $p\times\kappa$ co-
data model matrix
$\boldsymbol{\mathcal{C}}=\left(\boldsymbol{c}_{1}^{T},\ldots,\boldsymbol{c}_{p}^{T}\right)^{T}.$
We then employ $\boldsymbol{\mathcal{C}}$ for modeling the feature-specific EB
estimates $\boldsymbol{\hat{S}}_{1}^{(q+1)},$ i.e. (12).

To link $\boldsymbol{\hat{S}}_{1}^{(q+1)}$ to co-data
$\boldsymbol{\mathcal{C}},$ we model the counts $b_{j}^{(q)}.$ We start by
noting that each of the in total $B^{(q)}$ splitting rules may be regarded as
a Bernoulli trial for any covariate $j$: covariate $j$ occurs or does not
occur in the given splitting rule. Therefore, count $b_{j}^{(q)}$ is the
outcome of a binomial trial:
$b_{j}^{(q)}\sim\textrm{Bin}\left(w_{j}^{(q)},B^{(q)}\right),$ with unknown
probability of success $w_{j}^{(q)}$ for covariate $j$ at iteration $q.$ We
then model $w_{j}^{(q)}$ using $\boldsymbol{\mathcal{C}}$ by employing a
logistic regression model:

$w_{j}^{(q)}=\textrm{expit}\left(\boldsymbol{c}_{j}^{T}\boldsymbol{\eta}^{(q)}\right),\qquad\textrm{for}\quad
j=1,.\ldots,p,$ (13)

with $\textrm{expit}\left(x\right)=e^{x}/\left(e^{x}+1\right),$ and
$\boldsymbol{{\eta}}^{(q)}\in\mathbb{R^{\kappa}}$ a regression parameter
vector.

We estimate $w_{j}^{(q)}$ by $\hat{w}_{j}^{(q)},$ which will be the co-data
moderated EB-estimates $\boldsymbol{\hat{S}}_{2}^{(q+1)}$ of hyperparameter
$\boldsymbol{S}.$ To determine $\hat{w}_{j}^{(q)},$ we first estimate
$\boldsymbol{\eta}^{(q)}:$

$\hat{\boldsymbol{\eta}}^{(q)}=\arg\max_{\boldsymbol{\eta}^{(q)}}\;\sum_{j=1}^{p}\log\left[\textrm{Bin}\left(\textrm{expit}\left(\boldsymbol{c}_{j}^{T}\boldsymbol{\eta}^{(q)}\right),B^{(q)}\right)\right],$
(14)

i.e. maximum likelihood maximization for logistic regression with covariates
$j$ serving as samples. Equation (14) ignores two dependencies between samples
$j$. First, only one covariate can be used for a given splitting rule. Hence,
success in a given Bernoulli trial for variable $j$ determines failure for the
other variables. The second dependency originates from the dependence between
the covariates, which induces dependencies between the splitting rules.
Because we are only interested in the point estimates
$\hat{\boldsymbol{\eta}}^{(q)}$ and because the estimator (14) is consistent,
we ignore these dependencies.

Estimate $\hat{\boldsymbol{\eta}}^{(q)}$ then determines the estimates
$\hat{w}_{j}^{(q)}=\textrm{expit}\left(\boldsymbol{c}_{j}^{T}\hat{\boldsymbol{\eta}}^{(q)}\right),$
which we collect in the co-data guided empirical Bayes estimator
$\boldsymbol{\hat{S}}_{2}^{(q+1)}$ of $\boldsymbol{S}$ at iteration $q+1$:

$\boldsymbol{\hat{S}}_{2}^{(q+1)}=\left(\hat{w}_{1}^{(q)},\ldots,\hat{w}_{p}^{(q)}\right).$
(15)

Note that $\hat{\boldsymbol{\eta}}^{(q)}$ fully determines (15) given the co-
data $\boldsymbol{\mathcal{C}}.$ Hence, we only estimate $\kappa$
hyperparameters instead of $p-1$ (which is the case for (12)) to determine
$\boldsymbol{\hat{S}}_{2}^{(q+1)}$. We therefore opt for a parsimonious model,
i.e. (13), with $\kappa\ll p$ and only linear effects.

Equation (13) ensures that the co-data and the primary data should be in
agreement. If, for example, the feature specific EB estimates of (12) show a
strong preference for certain covariates, while this preference is not present
in the co-data, the final up/down weighting of the covariates will be
diminished. Vice versa, if the co-data shows a strong preference for a certain
covariate, e.g. by a small p-value, but BART does not recognize this
covariate, the covariate will not be substantially upweighted.

###  3.3 Convergence of EB-coBART

Equation (15) provides iterative updates of hyperparameter $\boldsymbol{S}$
and hence a convergence criterion is required. Typically, either the marginal
likelihood or the estimated hyperparameters, in this case
$\hat{\boldsymbol{\eta},}$ are tracked until they stabilize within a given
tolerance level (Casella [2001]). However, for BART, we found that both
options are not feasible.

First, tracking the marginal likelihood is difficult because its estimation
within reasonable computational time is nontrivial. Known sampling methods
like the harmonic mean estimator (Newton and Raftery [1994]) or averaging the
likelihood from a sample of the prior do not work because too many samples are
required to cover the tree space. Boatman et al. [2021] estimated the marginal
likelihood using prior sampling, but only two covariates were considered. For
larger $p,$ a scenario we are interested in, the tree space becomes too large
for this method to work. Another option for marginal likelihood estimation was
proposed by Chib [1995]. This option is also unfeasible, because it requires
an analytic expression of the full conditional of the tree structure parameter
$\mathcal{\boldsymbol{\mathcal{T}}}$ of BART, which is not available. Second,
we empirically found that tracking hyperparameter $\boldsymbol{\eta}$ until
convergence fails because of overfitting (Supplementary Section 3).

We therefore rely on the widely applicable information criterion (WAIC)
(Watanabe [2013]):

$\textrm{WAIC}=-2\sum_{i=1}^{N}\log\left(E\left[\pi\left(y_{i}|\boldsymbol{\mathcal{T}},\boldsymbol{\mathcal{M}},\sigma^{2},\boldsymbol{x_{i}}\right)\right]\right)+2\sum_{i=1}^{N}\textrm{Var}\left[\log\left(\pi\left(y_{i}|\boldsymbol{\mathcal{T}},\boldsymbol{\mathcal{M}},\sigma^{2},\boldsymbol{x_{i}}\right)\right)\right],$
(16)

with
$\pi\left(y_{i}|\boldsymbol{\mathcal{T}},\boldsymbol{\mathcal{M}},\sigma^{2},\boldsymbol{x_{i}}\right)$
the predictive density, and the expectation and variance taken w.r.t. the
posterior of BART. We then halt the iterative hyperparameter updates
$\boldsymbol{\hat{S}}_{2}^{(q+1)}$ when the WAIC is at minimum.

The WAIC as convergence criterion has some desirable properties. First, the
WAIC is easily computed for nonparametric methods such as BART, contrary to
other information criteria. Second, the WAIC is asymptotically equivalent to
leave-one-out cross-validation of the likelihood. Third, computationally, the
WAIC is an efficient stopping criterion because it may be estimated from the
already available posterior samples of BART (Murray [2021]).

As an alternative to the WAIC, the computationally more demanding cross-
validation (CV) may be used. In Supplementary Section 4, we illustrate in a
simulation that CV as stopping criterion leads to similar predictive
performance and variable selection results compared to WAIC.

###  3.4 Choices for the other hyperparameters

Our derived EB scheme for BART ((9), (10), (11), and (12)) may also be used to
estimate the other hyperparameters $\left(\alpha,\beta,k,\nu,\lambda\right)$
of BART. Using EB for hyperparameter estimation may be an appealing
alternative to cross-validation, which relies on a subjective grid of
hyperparameter choices.

Estimation of multiple (correlated) hyperparameters, however, is intrinsically
difficult. For example, for elastic net, a substantially less flexible
prediction model compared to BART, it was proven that joint estimation of two
hyperparameters already causes identifiability issues (van Nee et al. [2022]).
Also, automatic hyperparameter tuning does not directly account for model
interpretability.

Therefore, our default choice is to fix the tree structure hyperparameters
$\alpha$ and $\beta$ at two flexibility levels: rigid ($\alpha=0.1,$
$\beta=4$), favoring shallow trees and thus interpretability, and flexible
($\alpha=0.95,$ $\beta=2$), favoring less shallow trees. We stress that the
flexible model corresponds to the default settings of BART (Chipman et al.
[2010]). Hyperparameter $k$ is fixed at $k=1$ for the rigid model and $k=2$
for the flexible model, acknowledging that shallow trees have more
observations in the terminal nodes and thus require less shrinkage of the
terminal node parameters. For the error variance hyperparameters
$\left(\nu,\lambda\right),$ we employ the default settings proposed by Chipman
et al. [2010].

However, for high-dimensional data, the predictions of BART may improve by
hyperparameter fine-tuning (Chipman et al. [2010]). In the high-dimensional
application (Section 5), we illustrate that hyperparamater estimation of
$\alpha$ and $k$ using the derived EB-estimator in (9) and (10) renders a
slight improvement in predictive performance compared to fixing $\alpha$ and
$k.$

##  4 Simulations

We compare EB-coBART to BART when informative co-data is available in a high-
dimensional data setting. To do so, we consider two functions $f$ that specify
a covariate-response relationship: a sparse and nonlinear function $f_{SN}$
(subsection 4.1), and a dense and linear function $f_{LD}$ (subsection 4.2).
For $f_{SN}$, we employ discrete co-data in the form of a grouping structure,
while for $f_{LD}$, we employ continuous co-data in the form of a noisy
version of the true linear effect sizes. This set-up ensures that we evaluate
the influence on BART of different types of co-data. Note that OG-BART (Du and
Linero [2019]) can handle discrete co-data as well, but because its
implementation is lacking in the public domain, we do not perform a comparison
with this method. subsection 4.3 deals with uninformative grouping co-data to
evaluate whether EB-coBART recognizes such co-data by not upweighting certain
groups.

For both functions, we consider multiple simulation settings which are
specified in the subsections. In each setting and for each function, we
simulate $N_{sim}=500$ data sets. For each data set, we first fit BART, i.e.
BART with equal covariate weights $s_{j}=1/p.$ Then, we iteratively update
$s_{j}$ according to estimator (15) until the WAIC, (16), is at minimum. BART
fitted with these estimated covariate weights then corresponds to our method
EB-coBART. To fit BART, we employ the R package dbarts (Dorie [2023]), and to
estimate the WAIC, we employ the R package loo (Vehtari et al. [2017]).

We consider a rigid tree model ($\alpha=0.1,$ $\beta=4,$ $k=1$), and a
flexible tree model ($\alpha=0.95,$ $\beta=2,$ $k=2$), as explained in
subsection 3.4. Thus, for each data set, we have four final BART fits: rigid
BART, flexible BART, which is the default, rigid EB-coBART, and flexible EB-
coBART.

We set the other hyperparameters settings of all BART models as follows. We
fix the error variance hyperparameters: $\nu=10$ and $\lambda$ such that the
$75$% quantile of the prior equals
$2/3\hat{\textrm{Var}}\left(\boldsymbol{y}\right),$ with
$\hat{\textrm{Var}}\left(\boldsymbol{y}\right)$ the estimated variance of the
simulated response $\boldsymbol{y}.$ We also fix the number of trees to $K=50$
to balance prediction and variable selection (Chipman et al. [2010], Bleich et
al. [2014]).

For each simulated data set, we estimate the predictive performance on a large
($N_{test}=500$) independent test set for both BART and EB-coBART. We quantify
the performance by the prediction mean square error (PMSE), i.e.
$\textrm{PMSE}=N^{-1}\sum_{i=1}^{N}\left(y_{i}-\hat{y}_{i}\right)^{2}$ and
$\hat{y}_{i}$ denoting the prediction for sample $i.$ Predictions
$\hat{y}_{i}$ are obtained by averaging posterior samples of the sum-of-trees
(Chipman et al. [2010]). We collect variable importance results by monitoring
the co-data-moderated EB estimates $\hat{w}_{j}^{(q)},$ (15), which relate
directly to how often the covariates occur in the tree ensemble.

###  4.1 Sparse and nonlinear setting

Response $y_{i}$ is generated by
$y_{i}=f_{SN}\left(\boldsymbol{x}_{i}\right)+\epsilon_{i},$ with noise
$\epsilon_{i}\sim\mathcal{N}\left(0,1\right),$ for $i=1,\ldots,N,$ and with

$f_{SN}\left(\boldsymbol{x}_{i}\right)=10\sin\left(\pi
x_{i1}x_{i2}\right)+10x_{i3}+20\left(x_{i101}-0.5\right)^{2}+10x_{i102},$ (17)

and $x_{ij}\stackrel{{\scriptstyle
i.i.d.}}{{\sim}}\textrm{Unif}\left(0,1\right),$ for $j=1,\dots,p,$ and
$p=500.$ Thus, covariates $\left\\{1,2,3,101,102\right\\}$ are predictive for
the response and the remaining $495$ covariates are noise.

Co-data are defined as a grouping structure with $G=\\{5,20\\}$ groups. We set
equal-sized groups of size $100$ for $G=5$ and size $25$ for $G=20.$ We then
assign covariates $1,2,\ldots,p/G$ to group $1$, covariates
$p/G+1,\ldots,2p/G$ to group $2$, et cetera. Group $G$ consists of covariates
$j=\left[\left(G-1\right)p/G+1\right],\ldots,p$. This distribution of
covariates among the groups ensures that predictive covariates $\\{1,2,3\\}$
are always in the same group (Group $1$ for $G=5$ and $G=20$) and that
predictive covariates $\\{101,102\\}$ are always in the same group (Group $2$
for $G=5$ and Group $5$ for $G=20$). For both $G$’s, we consider two samples
sizes ($N=100$ and $N=200$).

Our results demonstrate that EB-coBART upweights the predictive groups ($1$
and $2$ for $G=5;$ $1$ and $5$ for $G=20$) and downweights the non-predictive
groups in all simulation settings (Fig. 1). This upweighting effect is
stronger for the rigid tree models because shallow trees include fewer noisy,
non-relevant covariates in this sparse setting. Increasing the sample size
from $N=100$ to $N=200$ reduces the variability in the group-specific
estimates across the data sets as expected.

(a)

(b)

Figure 1: Boxplots of the co-data moderated EB estimates of the covariate
weights of EB-coBART across the $500$ simulated data sets for different
simulation settings. Figure (a) shows results for $G=5$ and Figure (b) for
$G=20,$ for which we only depict the first $10$ groups for visualization. For
each group, left (blue) boxplot corresponds to the rigid tree setting
($\alpha=0.1,$ $\beta=4,$ $k=1$) and right (green) boxplot to the flexible one
($\alpha=0.95,$ $\beta=2,$ $k=2$). Outliers are not shown. The horizontal
dotted lines correspond to equal group weights ($0.2$ for $G=5;$ $0.05$ for
$G=20$).

Table 1 depicts the average PMSE for the four BART models, and three
competitors: random forest, and two other co-data learners that can handle
both discrete and continuous co-data: Ecpc (van Nee et al. [2021]) and CoRF
(te Beest et al. [2017]). For Ecpc, we include a posthoc variable selection
procedure based on the elastic net penalty (van de Wiel et al. [2016], van Nee
et al. [2021]), using the optimal number of covariates $p_{\textrm{sel}}=5.$
Performances were worse with $p_{\textrm{sel}}=\\{10,20\\}.$

Table 1: Average PMSE across data sets for several simulation settings for BART and EB-coBART in the rigid tree setting ($\alpha=0.1,$ $\beta=4,$ $k=1$) and the flexible tree setting ($\alpha=0.95,$ $\beta=2,$ $k=2$). Also included are competitors: random forest, and co-data learners: CoRF and Ecpc with and without posthoc variable selection. | $N=100,$ | $N=100,$ | $N=200,$ | $N=200,$
---|---|---|---|---
| $G=5$ | $G=20$ | $G=5$ | $G=20$
Flexible BART | $11.4$ | $11.4$ | $4.58$ | $4.58$
Flexible EB-coBART | $10.1$ | $7.63$ | $\boldsymbol{4.23}$ | $\boldsymbol{3.23}$
Rigid BART | $9.30$ | $9.30$ | $4.94$ | $4.94$
Rigid EB-coBART | $\boldsymbol{8.81}$ | $\boldsymbol{7.47}$ | $4.66$ | $4.27$
Ecpc (no variable selection) | $24.6$ | $16.9$ | $17.9$ | $11.6$
Ecpc (variable selection) | $11.7$ | $8.43$ | $8.67$ | $7.46$
Random forest | $20.7$ | $20.7$ | $15.4$ | $15.4$
CoRF | $19.2$ | $15.0$ | $14.3$ | $11.6$

For all simulation settings, flexible or rigid EB-coBART performs best (Table
1). Furthermore, both flexible and rigid EB-coBART have a substantially lower
average PMSE compared to the competitors. For most settings, rigid and
flexible BART also outperform the competitors. The exception is for $N=100,$
$G=20$ where Ecpc with variable selection has a lower average PMSE than
flexible BART. Rigid EB-coBART has a lower PMSE than flexible EB-coBART for
$N=100$ and flexible EB-coBART has a lower PMSE for $N=200.$ In this sparse
simulation setting, more regularization on the tree depth is required for
smaller sample sizes.

Fig. 2 compares the relative PMSE of EB-coBART to that of BART demonstrating
that EB-coBART has a smaller PMSE than BART, i.e.
$\textrm{PMSE}_{\textrm{EBcoBART}}/\textrm{PMSE}_{\textrm{BART}}<1,$ for most
data sets in all considered simulation settings and for both tree flexibility
settings. Only for the rigid tree models in the $G=5$ setting, BART has a
smaller PMSE than EB-coBART for a relatively large percentage of simulated
data sets ($19$% for $N=100;$ $22$% for $N=200$).

Figure 2: Boxplot of the ratio
$\textrm{PMSE}_{\textrm{EBcoBART}}/\textrm{PMSE}_{\textrm{BART}}$ across the
$500$ simulated data sets for both the rigid tree models (blue, left) and the
flexible tree models (green, right). The four panels correspond to different
simulation settings.

Increasing the number of groups from $G=5$ to $G=20$ improves the relative
performance of EB-coBART compared to BART, because the co-data has become
relatively more informative. A larger sample size ($N=200$) renders a smaller
relative difference in PMSE, i.e.
$\textrm{PMSE}_{\textrm{EBcoBART}}/\textrm{PMSE}_{\textrm{BART}}$ increases,
compared to a smaller sample sizes ($N=100$). This finding is expected because
prior information is more relevant for smaller sample sizes.

A comparison between the two tree flexibility settings reveals that the
flexible tree models benefit more from the co-data than the rigid tree models
(Fig. 2). The flexible model uses more nodes and can therefore relate better
to the co-data information. Moreover, the co-data helps to prune away non-
informative covariates, which are less frequent for the rigid model.

###  4.2 Dense and linear setting

We briefly discuss results for a dense, linear regression simulation setting
with continuous co-data. Details are found in the Supplement (Section 5). For
variable importance, results are similar as for the sparse nonlinear
simulation: EB-coBART upweights predictive covariates and downweights
nonpredictive covariates. This up/down-weighting effect is stronger for rigid
BART model as this model favors variable selection compared to the flexible
model (Figure S5).

For most data sets, EB-coBART has a lower PMSE than BART. This effect is
stronger for the flexible models. The flexible model has a lower average PMSE
compared to the rigid model for all simulation settings, as expected. For
dense covariate-response relationships, many splitting rules are required,
which are more easily captured by a flexible model compared to a rigid model
(Figure S6).

The linear models Ecpc and ridge regression outperform all BART models because
the data generating mechanism is linear. Ecpc benefits from the co-data, as
evidenced by a substantial decrease in PMSE compared to ridge regression.
Among the nonlinear models, flexible EB-coBART has the lowest average PMSE,
and specifically a lower PMSE than random forest and CoRF (Table S1).

###  4.3 Uninformative co-data

Details for a setting with uninformative grouping co-data are found in
Supplementary Section 6 (Figures S7 and S8). Briefly, EB-coBART and BART
perform similar in this co-data setting. In the rigid tree setting, EB-coBART
does not upweight the groups on average, but there are fluctuations in the
estimated group weights. These fluctuations also induce fluctuations in the
test performance ratio, although on average both BART models have an almost
equal PMSE.

In the flexible tree setting, EB-coBART and BART are practically identical.
EB-coBART does not upweight the groups on average and fluctuations in the
estimated group weights across the data sets are hardly present. The PMSE of
EB-coBART and BART are close to equal for all data sets.

##  5 Application

##### Data

We apply our method to the prognosis of Diffuse large B-cell lymphoma (DLBCL)
patients. DLBCL is a common non-Hodgkin lymphoma for which accurate prognosis
is difficult because of the clinical and biological heterogeneity of the
patients. One well-accepted prognostic clinical covariate is the international
prognostic index (IPI), which scores patients based on their age, the stage of
the tumor, lactose dehydrogenase levels, a mobility measure, and the number of
extranodal sites (Shipp et al. [1986]). However, the predictive power of IPI
is still limited. Therefore, a major branch of DLBCL research focuses on
finding new omics-based markers for DLBCL prognosis.

To this end, we fit EB-coBART to a cohort of $101$ uniformly treated DLBCL
patients for whom we aim to predict two year progression-free survival
(yes/no, $18\%$ no). We treat the outcome as binary because two year is a
clinically well-accepted cut-off and censoring was absent within this time
period. We also have, for each patient, complete measurements of a total of
$p=140$ covariates divided in four types: $67$ DNA copy number variations
(CNV), $69$ point mutations of genes, $3$ translocations, and $1$ clinical
covariate: IPI.

We provide EB-coBART with the following two co-data sources. First, because we
have different types of covariates which may have different scales in relation
to the response, we group covariates by type (CNV, mutation, translocation, or
clinical). IPI is a group on its own because of its known prognostic
importance. Second, we provide continuous co-data by estimating Benjamini-
Hochberg corrected p-values ($-\textrm{logit}$ scale) of all covariates in
association with the binary response from a large ($N=430$) previously
published cohort with DLBCL patients having received a slightly different
treatment compared to the patients in the training cohort. In this way, EB-
coBART makes effective use of available data, while still acknowledging that
the effect of prognostic factors may depend on the type of treatment.

##### Analysis

We fit EB-coBART models to the training cohort. We consider EB-coBART
initialized in the flexible tree setting ($\alpha=0.95,$ $\beta=2,$ $k=2$),
which are the default settings proposed by Chipman et al. [2010], and EB-
coBART initialized in the rigid tree setting ($\alpha=0.95,$ $\beta=2,$
$k=2$). Additionally, we consider, for both tree settings, two hyperparameter
estimation strategies. The first strategy (EB-coBART 1) only estimates the
covariate weight hyperparameter $\boldsymbol{S}$ of BART by estimator (15),
similar to Section 4, and the second strategy (EB-coBART 2) simultaneously
estimates $\boldsymbol{S}$ by (15) and hyperparameters ($\alpha,$ $k$)
according to (9) and (10), with ($\alpha,$ $k$) initialized in the flexible or
the rigid tree setting. Thus, we have a total of four EB-coBART models:
flexible EB-coBART 1, flexible EB-coBART 2, rigid EB-coBART 1, and rigid EB-
coBART 2. Hyperparameter $\beta$ is fixed to $\beta=2$ for flexible models and
$\beta=4$ for the rigid models.

Because BART uses a probit link for binary responses (Albert and Chib [1993]),
we have
$p(y_{i}=1\mid\boldsymbol{x}_{i})=\Phi\left(f(\boldsymbol{x}_{i})\right),$
with $\Phi$ the standard normal cdf and $f\left(\boldsymbol{x}_{i}\right)$ the
latent response, which will be modeled by a sum-of-trees. Therefore, the error
variance equals $\sigma^{2}=1$ and hyperparameters $\left(\nu,\lambda\right)$
for $\pi\left(\sigma^{2}\right)$ are not required. Additionally, probit BART
has terminal node value prior
$\mu_{lt}\sim\mathcal{N}\left(0,3/(k\sqrt{K})\right)$ instead of (6). We fit
BART models using $10$ chains each consisting of $24000$ samples of which
$12000$ are burn-in.

We evaluate predictive performance by the area under the curve (AUC),
estimated using the R package pROC (Robin et al. [2011]), and the average
Brier score:
$\textrm{Brier}=N^{-1}\sum_{i=1}^{N}\left(\hat{y}_{i}-y_{i}\right)^{2}$.
First, we estimate the performance internally on the training cohort using
repeated ($3\times$) $10$-fold cross-validation (CV). Second, we fit the
models on the full training cohort and estimate the performance on an external
test cohort consisting of $83$ patients with the same treatment. In addition,
we investigate the predictive performance of EB-coBART 2 as a function of the
sample size by estimating a learning curve (Goedhart et al. [2023]).

We compare EB-coBART with flexible and rigid BART, i.e. BART models having
equal covariate weights ($s_{j}=1/p$), and with cv-BART, which estimates
hyperparameters $\alpha,$ $k,$ and the number of trees $K$ using $5$-fold CV,
and fixes $\beta=2$ and $s_{j}=1/p.$ We consider the CV grid
$\alpha=\\{0.1,\,0.5,\,0.95\\},$ $k=\\{1,\,2,\,3\\}$, and $K=\\{50,\,150\\}.$
We also fit BART using only IPI as covariate (IPI-BART).

Next to the BART models, we include a comparison with random forest and its
co-data extension CoRF (both fitted using $2000$ trees with the R package
randomForestSRC (Ishwaran and Kogalur [2023]), and to ridge regression and its
co-data extension Ecpc (van Nee et al. [2021, 2023]). For Ecpc, we also
considered a posthoc variable selection with the optimal number of covariates
$p_{\textrm{sel}}=\\{2,5,10,50,80\\}$, but this did not render improved
performances.

##### Results

Predictive performance results are shown in Table 2. To facilitate comparison,
and because the flexible tree setting ($\alpha=0.95,$ $\beta=2,$ $k=2$) is the
default (Chipman et al. [2010]), we show results for this setting here.
Results for the rigid tree setting ($\alpha=0.1,$ $\beta=4,$ $k=1$)
(Supplementary Table S2) are summarized.

Table 2: Predictive performance estimates, based on repeated ($3\times$) $10$-fold cross-validation and an external test cohort, of several prediction models. | Cross-Validation | Test Set Cohort
---|---|---
| AUC | Brier score | AUC | Brier score
flexible BART | 0.68 | 0.173 | 0.557 | 0.162
cv-BART | 0.68 | 0.168 | 0.557 | 0.162
flexible EB-coBART 1 | 0.72 | 0.155 | 0.697 | 0.154
flexible EB-coBART 2 | 0.71 | 0.156 | 0.714 | 0.153
IPI-BART | 0.70 | 0.153 | 0.669 | 0.154
Random Forest | 0.69 | 0.183 | 0.63 | 0.158
CoRF | 0.71 | 0.157 | 0.65 | 0.159
Ecpc | 0.73 | 0.159 | 0.705 | 0.154
Ridge | 0.72 | 0.161 | 0.69 | 0.154

For the cross-validated performances, the prediction models are competitive.
EB-coBART has a marginally improved performance compared to BART and cv-BART
in both tree flexibility settings, while IPI-BART performs similarly to EB-
coBART.

Compared to the cross-validated performance, the differences between the
models are larger when evaluated on the test set. EB-coBART 2 has the best
performance, and slightly better than EB-coBART 1. Because EB-coBART 2
estimates $\alpha=0.61$, it renders a sparser tree model compared to EB-coBART
1, which uses $\alpha=0.95.$ The differences with BART and cv-BART, which
yields $k=2,$ $\alpha=0.5,$ and $K=150$ as hyperparameter estimates, is large.
These differences are significant using DeLong’s paired test (DeLong et al.
[1988]) for the difference in AUC ($p_{\varDelta AUC}=0.013$, $p_{\varDelta
AUC}=0.024,$ respectively) and the Wilcoxon signed rank test for the
difference in brier score ($p_{\varDelta Brier}=0.034$, $p_{\varDelta
Brier}=0.0074,$ respectively). IPI-BART also substantially outperforms BART
and cv-BART.

EB-coBART has a substantially, but not significantly, better test performance
than random forest and CoRF, while Ecpc and ridge regression are competitive.
Ecpc and CoRF benefit only marginally from the co-data, indicated by the
comparative performance of their corresponding base learners.

Rigid EB-coBART has a comparable performance to flexible EB-coBART
(Supplementary Table S2), whereas rigid BART (without co-data) performs better
than the standard, i.e. flexible BART. This suggests a sparse covariate-
response relationship for this application, which, unlike for the standard
BART, is picked up by flexible EB-coBART as well due to the informative co-
data.

The learning curves for flexible EB-coBART 2 (triangles) and flexible BART
(dots) are depicted in LABEL:fig:LearningCurves. Learning curves are created
as follows. For ten different subsample sizes $n$, we sample multiple training
sets without replacement from the training cohort and define corresponding
test sets as left-out samples. We then fit flexible EB-coBART 2 (updating
$s_{j},$ $\alpha,$ $k$) and flexible BART to the training sets and estimate
the AUC on the complementary test sets. We aggregate AUC estimates per
subsample size by the average and plot them to obtain the learning curve.

For small subsample size $n$, BART has a larger AUC than EB-coBART. Moving
from $n=28$ to $n=36,$ the AUC of EB-coBART shows a phase transition (Donoho
and Tanner [2005]) by jumping to a much larger AUC value. This transition is
also present in the estimated weight for IPI, which makes a large jump at this
transition subsample size (Supplementary Figure S9). The AUC of BART remains
relatively constant across the sample size trajectory and only shows a clear
increase in AUC at the end. This indicates that the inclusion of co-data in
the BART algorithm is most beneficial for medium sample sizes, as the sample
size is then large enough to pick up the co-data signal, while still too small
to be picked up without the co-data guidance.

EB-based estimates of hyperparameter $\alpha$ increase for larger sample sizes
(Supplementary Figure S10a), because less tree structure shrinkage is required
when more information is available. Hyperparameter $k$ stayes relatively
stable across different subsample sizes (Supplementary Figure S10b).

(a)

(b)

(c)

(d)

Figure 3: (a) Learning curves for EB-coBART 2 (triangles) and BART (dots). (b)
Estimated WAIC (dots) for $18$ iterations with the minimum indicated by the
dashed vertical line at iteration $9.$ (c) Estimated cumulative covariate
weights for the four types of covariates as a function of external p-values on
the $-\textrm{logit}$ scale for EB-coBART 2. Types of covariates are indicated
by a square (copy number variation), a triangle (mutation), a diamond
(translocation), and a circle (IPI). (d) Partial dependence plots of EB-coBART
2 (triangles) and BART (dots) showing the marginal effect of IPI on the
predictions. On the y-axis, we show the latent response values, i.e.
$Z$-values of the standard normal cdf, because BART models use a probit link
for binary responses. We show the average $\pm$ the standard deviation across
the Gibbs samples of the latent response.

Because flexible EB-coBART 2 (updating $\boldsymbol{S},$ $\alpha,$ $k$)
performs best on the test set, we show this fit in Fig. 3b and c, and discuss
differences with EB-coBART 1 and rigid EB-coBART. LABEL:fig:WAIC shows
estimates of the WAIC across $18$ iterations, and LABEL:fig:GroupWeights shows
co-data moderated EB estimates of the covariate weights $\hat{w}_{j}^{(q)}$,
(15), aggregated by covariate type, at minimum WAIC.

The estimated WAIC of EB-coBART 2 has a clear minimum at iteration $9$
(LABEL:fig:WAIC), which is similar to EB-coBART 1 and the rigid models. At
minimum WAIC, EB-coBART 2 shows a strong preference for IPI, which receives
$92.5$% of the total weight (LABEL:fig:GroupWeights). The remaining weight is
spread out thinly among the omics covariates, with the sum of the weights per
type equal to $6.00$% for the copy number variations, $1.01$% for the
mutations, and $0.49$% for the translocations. IPI receives about $97$% of the
weight when rigid EB-coBART is considered. Devoting much weight to IPI
indicates that EB-coBART automatically finds a relevant signal, because BART
fitted with just IPI (IPI-BART) has a substantially better performance than
BART (Table 2).

Because IPI appears to be dominant for this application, we show the marginal
effect of IPI on the predictions of EB-coBART 2 and BART by the partial
dependence function (Friedman [2001]) (LABEL:fig:PartialDependence).
Marginalized IPI shows a clear effect on the two year progression free
survival (PFS) predictions of EB-coBART, where larger IPI values decrease the
probability of two year PFS. Furthermore, the effect is slightly nonlinear.
The signal in IPI is absent for BART.

##  6 Discussion

We developed EB-coBART, a method that incorporates co-data into BART by
estimating prior covariate weights using empirical Bayes and a co-data model.
This method rendered improved predictions, depending on the informativeness of
the co-data, compared to standard BART in simulations and in an application to
lymphoma patients. Furthermore, this application illustrated that EB-coBART
performs better than CoRF (te Beest et al. [2017]) and competitively to Ecpc
(van Nee et al. [2021]), two state-of-the-art co-data learners. The
comparative performance of EB-coBART and regression-based Ecpc depends on how
well the true model can be approximated by a linear one.

We considered EB-coBART having a flexible tree setting, which is the default
proposed by Chipman et al. [2010], and a rigid tree setting. The rigid model
is sparser because it uses less nodes and it therefore lends itself better for
variable selection and model interpretability compared to the flexible model.
For prediction, we therefore recommend the rigid model when the performance is
comparable, as was the case in the application.

The application illustrated that the specification of co-data may be rather
straightforward, as a simple grouping by covariate type already rendered
informative co-data for BART. Furthermore, the combination of empirical Bayes
and a parsimonious co-data model protects against overfitting when
uninformative co-data is incorporated. This parsimony may be slightly relaxed
by considering nonlinear co-data models such as general additive models
(Hastie and Tibshirani [1986]). Success of such strategy depends on the number
of covariates, as these are effectively the number of samples in the co-data
model. Furthermore, additional penalization in this model may be required to
avoid overfitting (van Nee et al. [2021]).

EB-coBART has to refit BART at each iteration and it is therefore
computationally intensive. However, because we empirically found that $4-12$
iterations is typical, EB-coBART requires less computational time than cross-
validated BART. The computational time of BART, and thus also of EB-coBART,
does not scale well to very high-dimensional settings ($p\approx 10^{4}$). It
may therefore be interesting to apply our method to BART-BMA (Hernández et al.
[2018]), a computationally efficient BART-based method. However, BART-BMA
employs greedy search instead of having a prior on the covariate weights,
which makes the implementation of empirical Bayes not straightforward.

We end with suggestions for future research. First, some co-data presents
itself in a hierarchical structure, such as gene ontology trees. Such
structures may possibly be accommodated by EB-coBART in the co-data model,
e.g. by borrowing ideas from the hierarchical lasso (Yan and Bien [2017]).
Second, for tree-based methods, it is generally difficult to deal with types
of covariates that have a different scale and priority in relation to the
response. Grouping the covariates by type provided a simple solution, but
other, more sophisticated, solutions may improve the tree-based model. It may
be good to favor certain types, e.g. clinical covariates, at the trunk of the
trees, as such covariates may be easier to interpret and may have proven their
use more extensively. Thirdly, we currently did not account for the depth at
which covariates occur when fitting the co-data model; we only accounted for
their occurrence. While the latter may suffice for settings with fairly large
$p$, it may be too granular for low $p$ settings. Depths may be used as
inverse weights in the co-data model to give more weight to covariates close
to the tree’s trunk.

##  7 Data availability and software

For the DLBCL application, treatment and covariate names were anonymized, as
the original data has not been published yet. Anonymized data includes the
training cohort, the test cohort, and the co-data matrix. These data and R
code (version version 4.3.0) to reproduce results presented in Section 4 and
Section 5 are available via https://github.com/JeroenGoedhart/EB_coBART_paper.

##  8 Supplementary Material

Supplementary material is available

## Acknowledgments

The authors thank Bauke Ylstra and Daphne de Jong for providing and discussing
the data presented in Section 5. This research was funded by Hanarth Fonds.

Conflict of Interest: None declared.

## References

  * Albert and Chib [1993] J. H. Albert and S. Chib.  Bayesian analysis of binary and polychotomous response data.  _J Am Stat Assoc_ , 88(422):669–679, 1993.  doi: 10.1080/01621459.1993.10476321.
  * Bleich et al. [2014] J. Bleich, A. Kapelner, E. I. George, and S. T. Jensen.  Variable selection for bart: an application to gene regulation.  _Ann Appl Stat_ , pages 1750–1781, 2014.  doi: 10.1214/14-aoas755.
  * Boatman et al. [2021] J. A. Boatman, D. M. Vock, and J. S. Koopmeiners.  Borrowing from supplemental sources to estimate causal effects from a primary data source.  _Clinical Trials_ , 40(24):5115–5130, 2021.  doi: https://doi.org/10.1002/sim.9114.
  * Boyle et al. [2017] E. A. Boyle, Y. I. Li, and J. K. Pritchard.  An expanded view of complex traits: From polygenic to omnigenic.  _Cell_ , 169(7):1177–1186, 2017.  doi: 10.1016/j.cell.2017.05.038.
  * Carvalho et al. [2010] C. M. Carvalho, N. G. Polson, and J. G. Scott.  The horseshoe estimator for sparse signals.  _Biometrika_ , 97(2):465–480, 2010.  ISSN 0006-3444.  doi: 10.1093/biomet/asq017.
  * Casella [2001] G. Casella.  Empirical Bayes Gibbs sampling.  _Biostatistics_ , 2(4):485–500, 2001.  doi: 10.1093/biostatistics/2.4.485.
  * Chib [1995] S. Chib.  Marginal likelihood from the gibbs output.  _J Am Stat Assoc_ , 90(432):1313–1321, 1995.  doi: 10.1080/01621459.1995.10476635.
  * Chipman et al. [2010] H. A. Chipman, E. I. George, and R. E. McCulloch.  Bart: Bayesian additive regression trees.  _Ann Appl Stat_ , 4(1):266–298, 2010.  doi: 10.1214/09-AOAS285.  URL https://doi.org/10.1214/09-AOAS285.
  * DeLong et al. [1988] E. R. DeLong, D. M. DeLong, and D. L. Clarke-Pearson.  Comparing the areas under two or more correlated receiver operating characteristic curves: a nonparametric approach.  _Biometrics_ , 44(3):837–845, 1988.  doi: 10.2307/2531595.
  * Donoho and Tanner [2005] D. L. Donoho and J. Tanner.  Sparse nonnegative solution of underdetermined linear equations by linear programming.  _Proc Natl Acad Sci USA_ , 102(27):9446–9451, 2005.  doi: 10.1073/pnas.0502269102.
  * Dorie [2023] V. Dorie.  dbarts: Discrete bayesian additive regression trees sampler, 2023.  URL https://CRAN.R-project.org/package=dbarts.  R package version 0.9-23.
  * Du and Linero [2019] J. Du and A. R. Linero.  Incorporating grouping information into bayesian decision tree ensembles.  _Proc Mach Learn Res_ , 97:1686–1695, 2019.  URL https://proceedings.mlr.press/v97/du19d.html.
  * Friedman [2001] J. H. Friedman.  Greedy function approximation: a gradient boosting machine.  _Ann Statist_ , 29(5):1189–1232, 2001.  doi: 10.1214/aos/1013203451.
  * Gelman and Rubin [1992] A. Gelman and D.B. Rubin.  Inference from iterative simulation using multiple sequences.  _Stat Sci_ , 7(4):457–472, 1992.  doi: 10.1214/ss/1177011136.
  * Goedhart et al. [2023] J. M. Goedhart, T. Klausch, and M. A. van de Wiel.  Estimation of predictive performance in high-dimensional data settings using learning curves.  _Comput Stat Data Anal_ , 180:107622, 2023.  doi: https://doi.org/10.1016/j.csda.2022.107622.
  * Hastie and Tibshirani [1986] T. Hastie and R. Tibshirani.  Generalized additive models.  _Stat Sci_ , 1(3):297 – 310, 1986.  doi: 10.1214/ss/1177013604.  URL https://doi.org/10.1214/ss/1177013604.
  * Hernández et al. [2018] B. Hernández, A. E. Raftery, S. R. Pennington, and A. C. Parnell.  Bayesian additive regression trees using bayesian model averaging.  _Stat and Comput_ , 28(4):869–890, 2018.  doi: 10.1007/s11222-017-9767-1.  URL https://doi.org/10.1007/s11222-017-9767-1.
  * Ishwaran and Kogalur [2023] H. Ishwaran and U.B. Kogalur.  _Fast Unified Random Forests for Survival, Regression, and Classification (RF-SRC)_ , 2023.  URL https://cran.r-project.org/package=randomForestSRC.  R package version 3.2.2.
  * Kapelner and Bleich [2016] A. Kapelner and J. Bleich.  bartmachine: Machine learning with bayesian additive regression trees.  _J Stat Softw_ , 70(4):1–40, 2016.  doi: 10.18637/jss.v070.i04.
  * Linero [2018] A.R. Linero.  Bayesian regression trees for high-dimensional prediction and variable selection.  _J Am Stat Assoc_ , 113(522):626–636, 2018.  doi: 10.1080/01621459.2016.1264957.  URL https://doi.org/10.1080/01621459.2016.1264957.
  * Murray [2021] J. S. Murray.  Log-linear bayesian additive regression trees for multinomial logistic and count regression models.  _J Am Stat Assoc_ , 116(534):756–769, 2021.  doi: 10.1080/01621459.2020.1813587.  URL https://doi.org/10.1080/01621459.2020.1813587.
  * Neuenschwander et al. [2010] B. Neuenschwander, G. Capkun-Niggli, M. Branson, and D. J. Spiegelhalter.  Summarizing historical information on controls in clinical trials.  _Clin Trials_ , 7(1):5–18, 2010.  doi: 10.1177/1740774509356002.
  * Newton and Raftery [1994] M. A. Newton and A. E. Raftery.  Approximate bayesian inference with the weighted likelihood bootstrap.  _J Roy Stat Soc B Met_ , 56(1):3–48, 1994.  URL http://www.jstor.org/stable/2346025.
  * Robin et al. [2011] X. Robin, N. Turck, A. Hainard, N. Tiberti, F. Lisacek, J. C. Sanchez, and M. Müller.  proc: an open-source package for r and s+ to analyze and compare roc curves.  _BMC Bioinform_ , 12(1):77, 2011.  doi: 10.1186/1471-2105-12-77.  URL https://cran.r-project.org/web/packages/pROC/index.html.
  * Shipp et al. [1986] M.A. Shipp, D. P. Harrington, M. M. Klatt, M. S. Jochelson, et al.  Identification of major prognostic subgroups of patients with large-cell lymphoma treated with m-bacod or m-bacod.  _Ann Intern Med_ , 104(6):757–765, 1986.  doi: 10.7326/0003-4819-104-6-757.  URL https://www.acpjournals.org/doi/abs/10.7326/0003-4819-104-6-757.
  * 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.  _J Stat Softw_ , 97(1):1–66, 2021.  doi: 10.18637/jss.v097.i01.
  * Sparapani et al. [2016] R. A. Sparapani, B. R. Logan, R. E. McCulloch, and P. W. Laud.  Nonparametric survival analysis using bayesian additive regression trees (bart).  _Stat Med_ , 35(16):2741–2753, 2016.  doi: 10.1002/sim.6893.  URL https://onlinelibrary.wiley.com/doi/abs/10.1002/sim.6893.
  * Tan and Roy [2019] Y. V. Tan and J. Roy.  Bayesian additive regression trees and the general bart model.  _Stad Med_ , 38(25):5048–5069, 2019.  doi: 10.1002/sim.8347.
  * Tay et al. [2023] J. Tay, N. Aghaeepour, T. Hastie, and R. Tibshirani.  Feature-weighted elastic net: using ”features of features” for better prediction.  _Stat Sinica_ , 33:259–279, 2023.  doi: 10.5705/ss.202020.0226.
  * te Beest et al. [2017] D. E. te Beest, S. W. Mes, S. M. Wilting, R. H. Brakenhoff, and M. A. van de Wiel.  Improved high-dimensional prediction with random forests by the use of co-data.  _BMC Bioinform_ , 18(1):584, 2017.  doi: 10.1186/s12859-017-1993-1.
  * Tibshirani [1996] R. Tibshirani.  Regression shrinkage and selection via the lasso.  _J Roy Stat Soc B Met_ , 58(1):267–288, 1996\.  doi: https://doi.org/10.1111/j.2517-6161.1996.tb02080.x.
  * van de Wiel et al. [2016] M. A. van de Wiel, T. G. Lien, W. Verlaat, W. N. van Wieringen, and S. M. Wilting.  Better prediction by use of co-data: adaptive group-regularized ridge regression.  _Stat Med_ , 35(3):368–381, 2016.  doi: https://doi.org/10.1002/sim.6732.
  * van Nee et al. [2021] M. M. van Nee, L. F. A. Wessels, and M. A. van de Wiel.  Flexible co-data learning for high-dimensional prediction.  _Stat Med_ , 40(26):5910–5925, 2021.  doi: https://doi.org/10.1002/sim.9162.
  * van Nee et al. [2022] M. M. van Nee, T. van de Brug, and M. A. van de Wiel.  Fast marginal likelihood estimation of penalties for group-adaptive elastic net.  _J Comput Graph Stat_ , 0(0):1–11, 2022.  doi: 10.1080/10618600.2022.2128809.
  * van Nee et al. [2023] M.M. van Nee, L.F.A. Wessels, and M.A. van de Wiel.  ecpc an r-package for generic co-data models for high-dimensional prediction.  _BMC Bioinform_ , 24(172), 2023.  doi: https://doi.org/10.1186/s12859-023-05289-x.
  * Vehtari et al. [2017] A. Vehtari, A. Gelman, and J. Gabry.  Practical bayesian model evaluation using leave-one-out cross-validation and waic.  _Stat Comput_ , 27:1413–1432, 2017.  doi: 10.1007/s11222-016-9696-4.
  * Velten and Huber [2019] Britta Velten and Wolfgang Huber.  Adaptive penalization in high-dimensional regression and classification with external covariates using variational Bayes.  _Biostatistics_ , 22(2):348–364, 2019.  doi: 10.1093/biostatistics/kxz034.
  * Watanabe [2013] S. Watanabe.  A widely applicable bayesian information criterion.  _J Mach Learn Res_ , 14(1):867–897, 2013.
  * Yan and Bien [2017] X. Yan and J. Bien.  Hierarchical sparse modeling: A choice of two group lasso formulations.  _Stat Sci_ , 32(4):531 – 560, 2017.  doi: 10.1214/17-STS622.  URL https://doi.org/10.1214/17-STS622.
  * Yuan and Lin [2005] M. Yuan and Y. Lin.  Model Selection and Estimation in Regression with Grouped Variables.  _J Roy Stat Soc B Met_ , 68(1):49–67, 2005.  doi: 10.1111/j.1467-9868.2005.00532.x.