Before diving into examples, we clarify a few concepts that appear
throughout the package interface.
Shrinkage priors on the step heights
In a BART ensemble each tree contributes a step height (leaf
parameter) to the overall prediction. Classical BART assigns these step
heights a fixed-variance Gaussian prior, which regularises all leaves
equally. In high-dimensional settings, stronger and more adaptive
regularisation is desirable. ShrinkageTrees implements
two shrinkage priors that are placed directly on the step heights via a
scale mixture of normals:
\[
h_\ell \mid \lambda_\ell, \tau, \omega \sim
\mathcal{N}(0,\; \omega\, \lambda_\ell^2\, \tau^2).
\]
Here \(\tau\) is a
global shrinkage parameter shared across all leaves,
\(\lambda_\ell\) is a
local scale specific to leaf \(\ell\), and \(\omega\) is a fixed scaling constant. The
two currently implemented instantiations are:
- Horseshoe (
prior_type = "horseshoe").
Both \(\lambda_\ell\) and \(\tau\) receive independent half-Cauchy
priors. The heavy tails of the half-Cauchy allow individual leaves to
escape shrinkage when the data support a strong effect, while the global
parameter \(\tau\) pulls the bulk of
the estimates toward zero. This is the default prior in
HorseTrees() and CausalHorseForest().
- Half-Cauchy
(
prior_type = "half-cauchy"). Only the local scales \(\lambda_\ell\) receive a half-Cauchy prior;
there is no global shrinkage parameter. This provides per-leaf
adaptivity without the additional pooling across the ensemble.
A forest-wide variant of the horseshoe
(prior_type = "horseshoe_fw") shares a single global \(\tau\) across all trees in the forest
rather than one per tree. These priors can be selected in
ShrinkageTrees() and CausalShrinkageForest()
via the prior_type argument, and can be combined with
DART’s Dirichlet splitting prior for simultaneous structural and
parametric regularisation.
Hyperparameter selection
The two most important hyperparameters are local_hp and
global_hp. These control the horseshoe prior on the step
heights (leaf parameters):
\[
\mu_{jl} \mid \lambda_{jl}, \tau_j \sim \mathcal{N}(0, \lambda_{jl}^2
\tau_j^2),
\qquad \lambda_{jl} \sim \text{C}^+(0, \texttt{local\_hp}),
\qquad \tau_j \sim \text{C}^+(0, \texttt{global\_hp}),
\]
where \(\text{C}^+\) denotes the
half-Cauchy distribution. Smaller values produce stronger shrinkage
toward zero; larger values allow more variation.
HorseTrees() and
CausalHorseForest() provide a convenience
parameter k that sets both scales automatically:
local_hp = global_hp = k / sqrt(number_of_trees). The
default k = 0.1 works well in many settings and is a good
starting point.
ShrinkageTrees() and
CausalShrinkageForest() expose
local_hp and global_hp directly (no
k). A common rule of thumb is:
local_hp = k / sqrt(number_of_trees) with
k in [0.05, 0.5].global_hp = local_hp (symmetric) or a larger value if
you want less overall shrinkage.
The survival functions (SurvivalBART,
SurvivalDART) use k to calibrate the standard
BART leaf prior:
local_hp = range(log(y)) / (2 * k * sqrt(number_of_trees)).
The default k = 2 follows Chipman et al. (2010).
The store_posterior_sample flag
When store_posterior_sample = TRUE, the fitted object
stores the full \(N_\text{post} \times
n\) matrix of posterior draws for predictions. This is needed
for:
predict() on new data (it re-runs the sampler
internally, so posterior samples are always produced);plot(fit, type = "ate") and
plot(fit, type = "cate"), which require the full posterior
distribution;plot(fit, type = "survival") — full posterior credible
bands over both \(\mu_i\) and \(\sigma\) (without samples, only
sigma-uncertainty bands are available);- any custom posterior analysis (e.g. computing posterior credible
intervals for individual predictions).
When FALSE, only posterior means and \(\sigma\) draws are stored, saving memory.
print() and summary() work in both cases, but
predict() will not be available.
Treatment coding (treatment_coding)
All causal model functions — CausalHorseForest(),
CausalShrinkageForest(), SurvivalBCF(), and
SurvivalShrinkageBCF() — decompose the outcome as \[
Y_i = \mu(\mathbf{x}_i) + b_i \cdot \tau(\mathbf{x}_i) + \varepsilon_i,
\] where \(b_i\) is a scalar
that depends on the treatment assignment \(Z_i\). The treatment_coding
argument controls how \(b_i\) is
defined. Four options are available:
"centered" (default). \(b_i = Z_i - 1/2\), so that \(b_i \in \{-1/2,\; 1/2\}\). This is the
original BCF parameterisation.
"binary". \(b_i = Z_i\), so that \(b_i \in \{0,\; 1\}\). Standard binary
coding; the treatment forest captures the full effect of treatment on
the treated.
"adaptive". \(b_i = Z_i - \hat{e}(\mathbf{x}_i)\), where
\(\hat{e}(\mathbf{x}_i)\) is the
estimated propensity score. This follows Hahn, Murray & Carvalho
(2020) and is the coding used in the bcf R package. When
using this option, a propensity vector must be
supplied.
"invariant". Parameter-expanded
(invariant) treatment coding. The coding parameters \(b_0\) and \(b_1\) are assigned \(N(0,\; 1/2)\) priors and estimated within
the Gibbs sampler via conjugate normal updates: \[
Y_i = \mu(\mathbf{x}_i) + b_{Z_i} \cdot \tilde{\tau}(\mathbf{x}_i) +
\varepsilon_i,
\qquad b_0,\; b_1 \sim N(0,\; 1/2).
\] The treatment effect is \(\tau(\mathbf{x}_i) = (b_1 - b_0) \cdot
\tilde{\tau}(\mathbf{x}_i)\), and the posterior draws of \(b_0\) and \(b_1\) are returned in the fitted object.
This parameterisation is invariant to the coding of the treatment
indicator (Hahn et al., 2020, Section 5.2).
The examples below illustrate each option on a simple
continuous-outcome causal model.
set.seed(50)
n_tc <- 60; p_tc <- 5
X_tc <- matrix(rnorm(n_tc * p_tc), n_tc, p_tc)
W_tc <- rbinom(n_tc, 1, 0.5)
tau_tc <- 1.5 * (X_tc[, 1] > 0)
y_tc <- X_tc[, 1] + W_tc * tau_tc + rnorm(n_tc, sd = 0.5)
# Centered (default)
fit_tc_cen <- CausalHorseForest(
y = y_tc,
X_train_control = X_tc, X_train_treat = X_tc,
treatment_indicator_train = W_tc,
treatment_coding = "centered",
number_of_trees = 5, N_post = 50, N_burn = 25,
store_posterior_sample = TRUE, verbose = FALSE
)
cat("Centered — ATE:",
round(mean(fit_tc_cen$train_predictions_treat), 3), "\n")
#> Centered — ATE: 0.238
# Binary
fit_tc_bin <- CausalHorseForest(
y = y_tc,
X_train_control = X_tc, X_train_treat = X_tc,
treatment_indicator_train = W_tc,
treatment_coding = "binary",
number_of_trees = 5, N_post = 50, N_burn = 25,
store_posterior_sample = TRUE, verbose = FALSE
)
cat("Binary — ATE:",
round(mean(fit_tc_bin$train_predictions_treat), 3), "\n")
#> Binary — ATE: 0.521
# Adaptive (requires propensity scores)
ps_tc <- pnorm(0.3 * X_tc[, 1]) # simple propensity model for illustration
fit_tc_ada <- CausalHorseForest(
y = y_tc,
X_train_control = X_tc, X_train_treat = X_tc,
treatment_indicator_train = W_tc,
treatment_coding = "adaptive",
propensity = ps_tc,
number_of_trees = 5, N_post = 50, N_burn = 25,
store_posterior_sample = TRUE, verbose = FALSE
)
cat("Adaptive — ATE:",
round(mean(fit_tc_ada$train_predictions_treat), 3), "\n")
#> Adaptive — ATE: 0.154
# Invariant (parameter-expanded)
fit_tc_inv <- CausalHorseForest(
y = y_tc,
X_train_control = X_tc, X_train_treat = X_tc,
treatment_indicator_train = W_tc,
treatment_coding = "invariant",
number_of_trees = 5, N_post = 50, N_burn = 25,
store_posterior_sample = TRUE, verbose = FALSE
)
cat("Invariant — ATE:",
round(mean(fit_tc_inv$train_predictions_treat), 3), "\n")
#> Invariant — ATE: 0.232
# Posterior draws of b0 and b1 are stored in the fitted object
cat("b0 posterior mean:", round(mean(fit_tc_inv$b0), 3), "\n")
#> b0 posterior mean: 0.068
cat("b1 posterior mean:", round(mean(fit_tc_inv$b1), 3), "\n")
#> b1 posterior mean: 0.153
The survival functions inherit treatment_coding support.
For example,
SurvivalBCF(..., treatment_coding = "invariant") works out
of the box.
