Abstract
Causal mediation analysis provides insights into the mechanisms through which treatments affect outcomes. While mediation effects often vary across individuals, most existing methods focus solely on population-average effects, overlooking individual-level heterogeneity. To address this limitation, we propose a Bayesian regression tree ensemble method that flexibly models nonlinear relationships and captures treatment-by-mediator interactions in the mediation process. Using hierarchical posterior sampling, our approach provides credible intervals with nominal coverage rates for inferring heterogeneous mediation effects. Additionally, we leverage regression tree summaries to identify subgroups with distinct mediation effects and employ SHapley Additive exPlanation values to highlight key moderators and their influence on the mediation process. Comprehensive simulations demonstrate the method’s accuracy in estimating and inferring heterogeneous mediation effects. Finally, we apply our method to investigate the heterogeneous mediation role of Alzheimer’s disease pathology burden in the effect of apolipoprotein E genotype on late-life cognition.
Keywords: Bayesian tree ensembles, causal mediation analysis, heterogeneous effects, moderation mechanisms, nonlinear interactions
1. Introduction
Understanding complex causal pathways is essential in scientific research, especially in answering questions about how, for whom, and under what conditions a treatment affects an outcome. Mediation analyses address the question of how by uncovering the pathways through which a treatment effect (TE) is generated. Under the classical mediation framework (eg, Baron and Kenny, 1986), a treatment influences a mediator, which subsequently impacts the outcome. The total TE can be decomposed into an indirect effect transmitted through the mediator and a direct effect that operates directly or through other unspecified pathways. Across disciplines, empirical researchers use mediation analysis to uncover the mechanisms through which a treatment affects an outcome, thereby informing the design of targeted interventions and policies.
Notably, the direction and magnitude of the mediation effects are often heterogeneous across different subgroups and settings, known as moderators. Therefore, it is important to further address the questions of for whom and under what conditions a treatment works, beyond the question of how, through a heterogeneous (moderated) mediation analysis (eg, Edwards and Lambert, 2007). Figure 1A illustrates the heterogeneous mediation mechanism. The treatment or exposure impacts the outcome through the mediator, while moderators may change the direction or strength of one or more paths that transmit the treatment’s effect on the outcome.
Figure 1.
Heterogeneous mediation mechanism and a causal structure of AD pathology in APOE genotype-cognition association. (A) Heterogeneous mediation mechanism diagram. (B) Causal structure illustrating how AD pathology mediates the relationship between APOE genotype and late-life cognition.
However, most existing causal mediation methods focus primarily on population-average effects, potentially obscuring important variations in causal mediation mechanisms across different individual and contextual characteristics. Investigating these variations can help researchers improve and tailor interventions or policies to specific individuals and contexts. Qin and Wang (2024) addressed this gap by incorporating moderators into the regression-based causal mediation analysis. They developed the general definition, identification, estimation, and sensitivity analysis for causal moderated mediation (MM) effects under the potential outcomes framework (Rubin, 1978). Unlike the traditional MM analysis (eg, Preacher et al., 2007), Qin and Wang (2024) considered the treatment-by-mediator interaction when defining indirect and direct effects, accounted for observed confounders of the relationships among the treatment, mediator, and outcome, and evaluated the sensitivity of analysis results to unmeasured confounders. However, their method relies on parametric assumptions about the relationships among the variables. The causal effect estimates may be biased if the models are misspecified. Additionally, their method requires prior knowledge to select appropriate moderators, which can be challenging when such knowledge is limited or when the choice of moderators is not straightforward.
To address these limitations and enable more robust and flexible modeling, it is desirable to develop a machine learning method for heterogeneous causal mediation analysis. Among the machine learning methods used in causal inference, Bayesian Additive Regression Trees (BART) (Chipman et al., 2010) have consistently ranked among the top-performing methods in the Atlantic Causal Inference Conference Data Challenge over several years (Dorie et al., 2019). Hahn et al. (2020) and Carnegie et al. (2019) explored BART’s capabilities for examining TE heterogeneity and identifying subgroups with different levels of TE. Compared to other methods, BART-based methods are more likely to detect subgroup signals without overfitting (Thal and Finucane, 2023) and perform substantially better in real data (Chipman et al., 2006). Moreover, BART’s ability to provide credible intervals offers a unique advantage for inferring causal effects.
Linero and Zhang (2025) extended the application of BART to causal mediation analysis, enabling robust estimation, and inference of the mediation effects in the presence of nonlinear relationships. While they focused on estimating population-average mediation effects, Ting and Linero (2025) further considered estimating heterogeneous causal mediation effects with BART. They introduced a linear varying-coefficient model based on BART to identify and regularize heterogeneous causal mediation effects. The method allows for flexible investigations of the heterogeneity of causal mediation effects. However, it has several limitations. First, their model assumes no interaction between the treatment and mediator, which can lead to bias when such interactions are present and influential. Second, the method relies on linear varying-coefficient models for the outcome, which may not adequately capture potential nonlinear relationships between the mediator and outcome. Third, they assume small and mostly homogeneous mediation effects, which could make the method conservative in estimating heterogeneity in practical applications.
To overcome these limitations, we develop a nonparametric BART model for heterogeneous causal mediation analysis, termed Heterogeneous Mediation with BART (HMBART). HMBART allows for the treatment-by-mediator interaction, accommodates nonlinear relationships, and can capture large and heterogeneous mediation effects. It extends the approach proposed by Linero and Zhang (2025) to address the more challenging problem of estimating individual mediation effects. Through various simulation settings, we demonstrate that HMBART outperforms existing methods in reducing bias, achieving the nominal coverage rate (CR), and maintaining adequate statistical power.
We organize this article as follows. In Section 2, we introduce our method for conducting heterogeneous causal mediation analysis with BART. In Section 3, we present a series of simulations that demonstrate the superior performance of our method. Section 4 illustrates the application of HMBART in a real-data analysis, in which we assess the heterogeneous mediating roles of Alzheimer’s disease (AD) pathology burden underlying the relationship between the apolipoprotein E (APOE) genotype and late-life cognitive function. Finally, Section 5 discusses the advantages and limitations of our proposed method.
2. Method
2.1. Definition of the mediation effects
We illustrate the general definitions of heterogeneous causal mediation effects for a binary treatment, which can be easily extended to accommodate a discrete or continuous treatment. We denote the treatment, mediator, and outcome as 
and 
, respectively. The vector 
represents pretreatment covariates that may function as moderators, confounders, or both.
Under the potential outcomes framework, each individual 
is associated with two potential mediators, 
and 
, corresponding to the two treatment conditions. Similarly, there are two potential outcomes, 
and 
. In practice, only the potential outcomes corresponding to the assigned treatment 
are observed, which we denote as 
and 
for 
. To account for the influence of treatment on the mediator, which in turn affects the outcome, one may alternatively view the potential outcome as a function of both the treatment and the potential mediator under the same treatment condition, that is, 
.
These definitions of the potential mediators and potential outcomes rely on the stable unit treatment value assumption (SUTVA) (Rubin, 1990), which assumes that (1) there is only one version of each treatment condition and (2) an individual’s potential mediators do not depend on other individuals’ treatment status, and an individual’s potential outcomes do not depend on other individuals’ treatment status or mediator values.
The total TE for each individual 
can be defined as a contrast of the potential outcomes between the two treatment conditions: 
By introducing an additional potential outcome, 
, which denotes individual 
’s potential outcome if they were assigned to the treatment condition 
, yet the mediator takes the value that they would have taken under the counterfactual control condition 
, we decompose the total TE into the natural indirect effect (NIE) and the natural direct effect (NDE) (Pearl, 2022; Robins and Greenland, 1992):
 |
which represents the effect of the treatment on individual 
’s outcome under the treatment condition transmitted solely through the treatment-induced change in the mediator, and
 |
which represents the effect of the treatment on individual 
’s outcome if the mediator were held constant at the level that would have been realized under the control condition.
In addition to the NIE/NDE decomposition that 
, the total TE can be alternatively decomposed into the pure indirect effect (PIE), 
, and the total direct effect (TDE), 
(Robins and Greenland, 1992). The two decompositions are different in the presence of the treatment-by-mediator interaction. Nguyen et al. (2021) provided guidelines for selecting the most suitable decomposition based on specific research objectives. When the goal is to detect an indirect effect and estimate its magnitude, the NIE/NDE decomposition is preferred. If the interest is in assessing the direct effect beyond the indirect effect, the PIE/TDE decomposition is recommended. For simplicity, we focus on NIE and NDE in this article, while the estimation of PIE and TDE follows the same procedure.
The primary goal of this study is to evaluate heterogeneity in causal mediation effects across individuals. Therefore, we define the heterogeneous mediation effects as the mediation effects varying by individual characteristics, captured by the conditional expectation of each effect given the individual characteristics 
,
 |
(1) |
 |
(2) |
2.2. Identification of the mediation effects
We observe 
for the treated group and 
for the control group, whereas 
remains counterfactual and unobserved for all. Hence, to identify the conditional expectation of causal mediation effects as defined in (1) and (2), we need to relate unobserved quantities to observed data. The identification of the effects depends on the sequential ignorability assumption (Imai et al., 2010).
Assumption 1:
Strong ignorability of treatment. Given the pretreatment covariates
, the treatment is assumed to be independent of all potential mediators and potential outcomes,
where
and
can be equal or unequal and
for
, known as the positivity assumption. This assumes no omitted pretreatment confounders of the treatment-mediator or treatment-outcome relationships within levels of
.
Assumption 2:
Strong ignorability of the mediator. Given the pretreatment covariates
, the potential mediators are assumed to be independent of the potential outcomes within or across the treatment conditions,
where
for
. This assumption indicates that there are no omitted pretreatment confounders (those preceding the treatment) and no post-treatment confounders (those influenced by the treatment) in the mediator-outcome relationship, controlling for pretreatment covariate levels.
Under the sequential ignorability assumption, the conditional expectation of the potential outcomes given the individual characteristics can be identified as
 |
(3) |
where 
and 
can be equal or unequal. Correspondingly, we can identify the conditional causal mediation effects as defined in (1) and (2).
2.3. BART models
Based on the identification results in Equation 3, we estimate the conditional expectation of each potential outcome and, correspondingly, the conditional causal mediation effects defined in (1) and (2), by fitting models for the treatment, mediator, and outcome. We fit the models through BART, which allows us to flexibly capture nonlinear interactions, including the treatment-by-mediator interaction, and to provide inference for the effects.
For a 
-dimensional vector of predictors 
and a continuous response 
, the BART model is represented as 
, where 
is the sum of regression trees capturing the relationship between the predictors and the response, and 
. To obtain the sum-of-trees model 
, BART uses Bayesian backfitting Markov chain Monte Carlo (MCMC) (Chipman et al., 2010), which iteratively constructs and fits residuals; thus, each tree is updated by fitting the partial residuals after removing the contributions of all other trees.
We sequentially fit three BART models for a binary treatment 
, a continuous mediator 
, and a continuous outcome 
. For individual 
, the models are specified as follows:
(1) Treatment model: we model the probability of treatment 
as a function of the covariates 
through a probit BART model 
, where 
is the standard normal cumulative distribution function.
(2) Mediator model: we model the mediator 
as a function of the treatment 
and covariates 
: 
(3) Outcome model: we model the outcome 
as a function of the treatment 
, mediator 
, and covariates 
: 
In the BART fitting process, prior specifications of the model parameters play an important role. Using naive regularization priors in the presence of confounding can inadvertently lead models to attribute causal effects more to the treatment and mediator than to the confounders (Hahn et al., 2020; Linero, 2024). As a result, the discovery of truly influential confounders can be distorted by regularization, which is known as regularization-induced confounding (RIC). An effective way of mitigating RIC is through clever covariate adjustments (Hahn et al., 2020; Linero and Zhang, 2025; Ting and Linero, 2025). This involves adjusting for the propensity score of 
in the mediator and outcome models and adjusting for the prediction of 
in the outcome model, known as clever covariates. The clever covariates serve as summary measures of the confounders in the treatment-mediator, treatment-outcome, and mediator-outcome relationships, encouraging the models to incorporate confounders and thus reducing RIC.
Correspondingly, the refined mediator and outcome models are
 |
(4) |
 |
(5) |
where 
and 
are augmented covariate vectors incorporating the clever covariates, with the propensity score of 
as 
, and the predictions of 
under both treatment conditions as 
and 
. Both 
and 
are included to account for the difference in the sets of confounders of the mediator-outcome relationship across treatment conditions.
2.4. Bayesian estimation
BART employs an ensemble of Bayesian regression trees to predict response variables. Within this Bayesian framework, the sum-of-trees model requires prior specifications both for regularization and to prevent overfitting (Hill et al., 2020). A detailed prior specification for BART models is provided in Web Appendix A. We denote the sets of priors for models of 
, 
, and 
as 
, 
, and 
, respectively. To optimize the models and enhance robustness, we tune these prior specifications via cross-validation, as demonstrated by Chipman et al. (2006) and Dorie et al. (2019).
With specified priors, we can fit the BART models via the Bayesian backfitting MCMC procedure. In each iteration, a posterior sample of tree parameters is generated, based on which the posterior predictions of 
, 
, and 
(including the propensity score 
and clever covariates 
and 
) are generated. Figure 2 shows the detailed fitting process. We implement BART using the R package bartMachine (Kapelner and Bleich, 2016), which provides fast and reliable model fitting, prediction, and diagnostic checking.
Figure 2.
HMBART modeling process. The gray circles of 
(pretreatment covariates), 
(treatment), 
(mediator), and 
(outcome) represent the observed data, while the white circles indicate model-derived estimates. The models for 
, 
, and 
, fitted through BART, are shown within rounded rectangles, with arrows denoting the corresponding priors feeding into them. The augmented covariates are represented as 
and 
, where 
accounts for the mediator under different treatment conditions.
2.5. Inference of the mediation effects
We apply the Monte Carlo 
-computation approach to estimate mediation effects in nonlinear settings. This method computes the necessary counterfactuals by predicting potential outcomes under a given treatment level, while setting the mediator to its predicted value under the same or a different treatment level. To implement this 
-computation framework within our nonparametric Bayesian context, we develop a two-stage hierarchical posterior sampling strategy. This strategy constitutes the core of our inference: we first draw posterior samples from the mediator model and then from the outcome model. To explicitly account for uncertainty arising from the BART predictions, we incorporate both the predicted expectations and associated errors. The errors from the mediator model contribute to the prediction of the outcomes and are thus non-negligible. Therefore, when sampling the potential mediators, we add a random draw from the mediator model’s error distribution to ensure our estimates reflect the full variability in the data.
The step-by-step procedure is as follows. First, we fit the treatment model 
with all the covariates as predictors and the treatment as the target. Then we estimate propensity scores 
for the treatment and construct augmented covariates 
by including the propensity scores in the original covariates 
. This helps balance the distribution of covariates across treatment groups. Then, we fit the mediator model 
with both the treatment 
and the augmented covariates 
as predictors. The predictions of the mediator for individual 
under both treatment levels, 
and 
, are included in constructing augmented covariates 
. They serve as clever covariates to improve the precision of the outcome model. Finally, we fit the outcome model 
.
Given the fitted models, the two-stage hierarchical posterior sampling is formulated as follows. In the first stage, we draw 
posterior samples of the potential mediators for each individual 
under both treatment conditions from 
. The 
th posterior sample representing the potential value of the mediator for individual 
is denoted as 
. Note that 
used in the clever covariates is the average of 
across 
posterior samples. Specifically, we draw 
from their marginal distributions, where 
are random draws from 
representing the mediator model’s error terms shared across treatment levels to reduce Monte Carlo variability (Rene et al., 2023), and 
is the BART-estimated variance associated with the 
th posterior sample. With 
, the predicted expectations and associated errors of mediators can be incorporated into the prediction of the outcomes.
In the second stage, we draw 
posterior samples of the predicted potential outcomes for each of the 
posterior samples of the predicted potential mediators. The 
th posterior sample of the potential outcomes obtained conditional on the 
th posterior sample of the potential mediators is denoted as 
. Specifically, we compute 
, 
, and 
.
This creates a nested structure that propagates the uncertainty from the mediator model to the final outcome. Eventually, we have 
posterior samples of each potential outcome for each individual. Following Equations 1 and 2, each posterior sample of the heterogeneous mediation effects is calculated as 
and 
. The point estimates of the heterogeneous mediation effects are then obtained as an average of all posterior samples of the effects:
 |
The 
credible intervals of the individual effects 
and 
are then obtained by taking the 
and 
quantiles over 
posterior samples of 
and 
, respectively.
In addition, to mitigate high estimation uncertainty in small samples, we use empirical Bayes shrinkage to borrow strength across similar individuals and stabilize noisy estimates. Empirical Bayes shrinkage is applied to the individual estimates to produce final estimates and corresponding credible intervals (see Web Appendix B). Algorithm 1 details the procedure for fitting sequential models, predicting potential outcomes, and inferring the corresponding heterogeneous mediation effects. After estimating heterogeneous mediation effects and their credible intervals, we further examine how covariates contribute to this heterogeneity. Web Appendix C outlines the procedure for identifying influential covariates and subgroups.
3. Simulations
We evaluate the performance of our proposed HMBART method by comparing it with two alternative approaches for heterogeneous causal mediation analysis: regression-based causal MM (Qin and Wang, 2024) and varying-coefficient Bayesian causal mediation forest (VCBCMF) (see Web Appendix D for detailed model comparisons).
We generate datasets with 250, 500, and 1000 independent observations. Each covariate in 
is independently drawn from either a uniform or normal distribution, as detailed for each case below. The treatment 
, mediator 
, and outcome 
are respectively generated from 
, 
, and 
. For each sample size, we examine four cases to generate linear or nonlinear functions, 
, 
, and 
, leading to a total of 
scenarios. The detailed simulation settings are provided in Web Appendix E.
Simulations are conducted at the individual level, with evaluation metrics computed for each individual and then averaged across the population over all replications. We evaluate the accuracy of the estimated heterogeneous mediation effects and frequentist properties of their associated credible intervals. For each simulation scenario, we run 200 replications and compare the performance of the three methods based on various evaluation criteria, including the difference between the corresponding average mediation effect and the true effect (bias), root mean squared error (RMSE), proportion of 95% credible intervals containing the true value (CR), and the proportion of 95% credible intervals not containing 0 (power). In addition, we report the variance estimation. For MM, the reported variance is the mean of the estimated mediation-effect variances across all individuals. For VCBCMF and HMBART, the variance is calculated as the variance across all posterior samples of mediation effects for each individual, and then averaged across individuals.
The simulation results are shown in Table 1. For reference, we computed 95% confidence intervals for the observed coverage probabilities using the binomial distribution, under the null hypothesis that the true coverage probability equals 0.95. The number of trials in the binomial model was taken as the product of the sample size and the number of replications. For sample sizes of 250, 500, and 1000, these confidence intervals are 
, 
, and 
, respectively. We define methods achieving the nominal CR as having CRs falling above these lower confidence limits, indicating no inflated type I error rate. We define optimal coverage as the CR closest to the nominal level while controlling the type I error rate. The following results are highlighted in bold: the lowest bias and RMSE, the smallest variance paired with nominal coverage, the optimal coverage, and the highest power among methods maintaining that nominal coverage.
Table 1.
Comparison of model performance for NDE and NIE estimates across different cases with sample sizes equal to 250, 500, and 1000. Case 1 has linear functions, and Cases 2–4 have nonlinear functions.
| NDE | NIE | |||||||||||
|---|---|---|---|---|---|---|---|---|---|---|---|---|
| Case | Sample Size | Model | Bias | RMSE | Variance | CR | Power | Bias | RMSE | Variance | CR | Power |
| (1) | 250 | MM | −0.016 | 0.271 | 0.105 | 0.990 | 0.783 | 0.005 | 0.459 | 0.243 | 0.968 | 0.559 |
| VCBCMF | 0.022 | 0.420 | 0.196 | 0.947 | 0.590 | 0.474 | 0.977 | 0.088 | 0.495 | 0.145 | ||
| HMBART | −0.083 | 0.453 | 0.888 | 0.998 | 0.248 | 0.106 | 0.752 | 0.808 | 0.973 | 0.082 | ||
| 500 | MM | 0.018 | 0.190 | 0.048 | 0.961 | 0.847 | −0.037 | 0.307 | 0.117 | 0.945 | 0.680 | |
| VCBCMF | 0.047 | 0.342 | 0.179 | 0.967 | 0.585 | 0.463 | 0.889 | 0.088 | 0.511 | 0.230 | ||
| HMBART | −0.099 | 0.367 | 0.815 | 0.997 | 0.275 | 0.078 | 0.629 | 0.658 | 0.973 | 0.147 | ||
| 1000 | MM | 0.015 | 0.133 | 0.023 | 0.915 | 0.895 | −0.077 | 0.206 | 0.058 | 0.999 | 0.743 | |
| VCBCMF | 0.035 | 0.272 | 0.163 | 0.981 | 0.595 | 0.462 | 0.816 | 0.086 | 0.538 | 0.280 | ||
| HMBART | −0.060 | 0.287 | 0.637 | 0.997 | 0.406 | 0.036 | 0.464 | 0.486 | 0.984 | 0.262 | ||
| (2) | 250 | MM | −0.036 | 0.430 | 0.092 | 0.830 | 0.676 | −0.075 | 0.395 | 0.085 | 0.856 | 0.566 |
| VCBCMF | 0.066 | 0.545 | 0.086 | 0.748 | 0.460 | 0.225 | 0.453 | 0.032 | 0.691 | 0.121 | ||
| HMBART | −0.066 | 0.469 | 0.214 | 0.948 | 0.387 | 0.025 | 0.402 | 0.289 | 0.978 | 0.042 | ||
| 500 | MM | −0.041 | 0.350 | 0.044 | 0.750 | 0.772 | −0.099 | 0.283 | 0.039 | 0.831 | 0.691 | |
| VCBCMF | 0.133 | 0.441 | 0.079 | 0.827 | 0.440 | 0.206 | 0.394 | 0.034 | 0.744 | 0.226 | ||
| HMBART | 0.005 | 0.392 | 0.132 | 0.936 | 0.476 | −0.012 | 0.287 | 0.248 | 0.986 | 0.122 | ||
| 1000 | MM | −0.078 | 0.288 | 0.020 | 0.723 | 0.848 | −0.074 | 0.242 | 0.019 | 0.819 | 0.740 | |
| VCBCMF | 0.144 | 0.374 | 0.070 | 0.868 | 0.459 | 0.206 | 0.356 | 0.033 | 0.778 | 0.266 | ||
| HMBART | −0.001 | 0.322 | 0.084 | 0.951 | 0.580 | −0.017 | 0.214 | 0.182 | 0.993 | 0.176 | ||
| (3) | 250 | MM | 0.010 | 0.646 | 0.482 | 0.965 | 0.400 | 0.054 | 0.822 | 0.270 | 0.761 | 0.315 |
| VCBCMF | 0.475 | 0.670 | 0.070 | 0.526 | 0.507 | 0.278 | 0.642 | 0.031 | 0.452 | 0.907 | ||
| HMBART | −0.026 | 0.399 | 0.210 | 0.967 | 0.646 | 0.138 | 0.576 | 0.249 | 0.935 | 0.013 | ||
| 500 | MM | 0.008 | 0.469 | 0.215 | 0.951 | 0.581 | 0.050 | 0.727 | 0.118 | 0.604 | 0.550 | |
| VCBCMF | 0.325 | 0.471 | 0.065 | 0.669 | 0.668 | 0.280 | 0.558 | 0.029 | 0.563 | 0.813 | ||
| HMBART | −0.019 | 0.346 | 0.155 | 0.966 | 0.689 | 0.101 | 0.480 | 0.230 | 0.967 | 0.086 | ||
| 1000 | MM | 0.030 | 0.297 | 0.102 | 0.983 | 0.705 | 0.030 | 0.649 | 0.055 | 0.453 | 0.789 | |
| VCBCMF | 0.234 | 0.341 | 0.055 | 0.787 | 0.703 | 0.278 | 0.477 | 0.025 | 0.617 | 0.704 | ||
| HMBART | 0.029 | 0.233 | 0.090 | 0.965 | 0.745 | 0.055 | 0.389 | 0.173 | 0.983 | 0.209 | ||
| (4) | 250 | MM | 0.122 | 0.880 | 0.613 | 0.919 | 0.724 | −0.175 | 1.281 | 0.684 | 0.765 | 0.380 |
| VCBCMF | 0.345 | 0.819 | 0.283 | 0.776 | 0.935 | 0.592 | 1.066 | 0.108 | 0.586 | 0.218 | ||
| HMBART | −0.096 | 0.714 | 1.346 | 0.995 | 0.681 | 0.215 | 0.928 | 1.345 | 0.989 | 0.036 | ||
| 500 | MM | 0.081 | 0.669 | 0.274 | 0.874 | 0.849 | −0.107 | 1.095 | 0.290 | 0.625 | 0.661 | |
| VCBCMF | 0.263 | 0.618 | 0.247 | 0.854 | 0.888 | 0.544 | 0.997 | 0.109 | 0.567 | 0.318 | ||
| HMBART | −0.033 | 0.554 | 1.466 | 0.999 | 0.662 | 0.129 | 0.824 | 1.270 | 0.991 | 0.080 | ||
| 1000 | MM | 0.094 | 0.601 | 0.132 | 0.750 | 0.913 | −0.083 | 1.006 | 0.139 | 0.483 | 0.876 | |
| VCBCMF | 0.181 | 0.497 | 0.215 | 0.896 | 0.876 | 0.523 | 0.944 | 0.109 | 0.572 | 0.368 | ||
| HMBART | −0.029 | 0.449 | 1.493 | 1.000 | 0.620 | 0.049 | 0.671 | 1.141 | 0.992 | 0.135 | ||
The method with the smallest bias and RMSE in each case is bold. The smallest variance with nominal CR is bold. The CR for methods achieving optimal coverage is bold. The highest-power method among those achieving the nominal CR is bold.
The results demonstrate that HMBART achieves nominal CRs and optimal inferential performance for heterogeneous mediation-effect inference, outperforming existing methods, particularly in nonlinear scenarios. In nonlinear cases, HMBART consistently yields lower bias and RMSE than MM and VCBCMF, emphasizing its capability to handle complex nonlinear mediation mechanisms with heterogeneous effects. As the sample size increases, HMBART exhibits clear improvements over competing methods, with reduced bias and RMSE, improved CRs, and greater power. This indicates consistent improvement in performance as sample size increases.
In scenarios with large, heterogeneous mediation effects, such as Cases 3 and 4, HMBART exhibits higher variance than other methods. However, this larger variance reflects a more conservative assessment of uncertainty, which ensures that the type I error rate is not inflated. In contrast, MM and VCBCMF, with lower variances, struggle to achieve nominal coverage, suggesting their estimates may be overly shrunk.
Overall, these findings highlight the flexibility and effectiveness of HMBART, especially in scenarios involving nonlinear relationships and larger datasets. In contrast, the performance of VCBCMF is inferior because it overlooks treatment-by-mediator interactions, and its linear varying-coefficient framework lacks the flexibility needed to capture complex relationships. In addition, VCBCMF may overshrink the effects when the true mediation effects are large and heterogeneous. HMBART not only offers reduced bias and ensures valid inference but also demonstrates superior flexibility in modeling nonlinear relationships, making it a strong candidate for complex heterogeneous causal mediation analyses.
To assess efficiency, average runtimes were measured for Cases 3 and 4 (
, 10 CPU cores). Using 2000 burn-in and 500 sampling iterations, MM, VCBCMF, and HMBART averaged 6.50, 4.87, and 6.05 min, respectively.
4. Real data application
In this study, we examine the relationship between the APOE genotype and late-life cognition, mediated by AD pathology, using data from the Religious Orders Study and the Memory and Aging Project (ROSMAP). Data preprocessing procedures are described in Web Appendix F. The exposure variable is an indicator for 
allele carrier status. The mediator is global AD pathology burden, defined as the average of neuritic plaques, diffuse plaques, and neurofibrillary tangles, assessed microscopically across five brain regions: mid-frontal cortex, mid-temporal cortex, inferior parietal cortex, entorhinal cortex, and hippocampus. The outcome is the last recorded cognitive test score. Covariates include age at the cognitive visit, gender, race, and years of education. The analytic sample consists of 1112 participants without the 
allele and 317 participants with the allele.
The APOE 
4 allele is a well-established genetic risk factor for cognitive decline and AD-related dementia (Farrer et al., 1997). Nichols et al. (2023) analyzed participants from the ROSMAP dataset with linear mediation models. They found that APOE 
’s harmful impact on cognition operated through both AD and non-AD neurodegenerative pathologies, with the strongest mediation among participants with dementia. In this application, we aim to assess both the direct effect of the APOE genotype on late-life cognition and the indirect effects mediated through AD pathology. Figure 1B illustrates the mediation mechanism, with age serving as a moderator example.
Table 2 compares NDE and NIE estimation and inference across the three methods MM, VCBCMF, and HMBART, reporting the estimated average mediation effects, 95% intervals constructed with 2.5% and 97.5% quantiles of all estimated individual mediation effects, and the proportion of significant individual effects whose credible/confidence intervals do not contain 0.
Table 2.
The effect of APOE genotype 
allele on late-life cognition mediated by AD pathology in the ROSMAP dataset.
| NDE | NIE | |||||
|---|---|---|---|---|---|---|
| Method | Mean effect | 95% distributional interval | Significant (%) | Mean effect | 95% distributional interval | Significant (%) |
| MM | −0.143 | [−0.352, 0.131] | 22.5 | −0.376 | [−0.615, −0.194] | 95.9 |
| VCBCMF | −0.049 | [−0.077, −0.029] | 0 | −0.145 | [−0.265, −0.080] | 65.4 |
| HMBART | −0.216 | [−0.287, −0.142] | 48.3 | −0.286 | [−0.345, −0.215] | 17.9 |
We present the estimated mean mediation effects and 95% distributional interval of individual effects constructed from the 2.5% and 97.5% quantiles of heterogeneous effects, and calculate the percentage of individuals with significant effects.
The mean estimates of NDE and NIE differ across methods, with MM and HMBART providing estimates of larger magnitude than VCBCMF. The relatively narrow 95% intervals of VCBCMF indicate that it tends to shrink the effects toward homogeneity. According to HMBART’s results, 48.3% of individuals exhibit a significant NDE, and 17.9% exhibit a significant NIE. It is well-documented that the APOE 
4 promotes the accumulation of amyloid plaques and neurofibrillary tangles, which subsequently result in cognitive decline (Holtzman et al., 2012). All models indicate that the indirect effects through AD pathology are larger in magnitude than the direct effects, while VCBCMF cannot identify any significant direct effects.
Figure 3A displays SHapley Additive exPlanation (SHAP) variable importance (Lundberg and Lee, 2017) summary plots for the NDE and NIE, ranking moderators by their importance in explaining mediation effects. Positive SHAP values indicate that the corresponding observations increase the predicted values. For example, males tend to have less negative NDE and NIE than females, which suggests males are less affected by the APOE 
4 allele’s direct and indirect impact on cognition compared to females. Age emerges as the most influential moderator. Furthermore, Figure 3B presents decision tree visualizations for the NDE and NIE, illustrating heterogeneity across subgroups defined by key moderators. The trees show how age, race, gender, and years of education interact to influence these effects.
Figure 3.
Identifying important moderators of the effect of APOE on cognition mediated by global AD pathology burden. (A) The SHAP feature importance of NDE and NIE. Moderators are ranked by importance scores (indicated by the adjacent values) in moderating mediation effects. Dots represent their contribution to NDE or NIE, with color indicating the moderator value. (B) Decision trees to detect distinct subgroups for NDE and NIE. Each branch condition defines a subgroup based on specific moderators. The values in the nodes represent the estimated effects for each subgroup, while the percentages indicate the proportion of individuals in that subgroup.
Figure 4 shows a clear trend of how age modifies NDE and NIE. In Figure 4A, HMBART credible intervals are shown in yellow for significant effects and gray for nonsignificant effects. In Figure 4B, we fit a generalized additive model (GAM) (Hastie, 2017) with NDE and NIE estimates as response variables and splines of age as predictors, showing smoothed curves that reveal the trend of the HMBART estimates varying with age. The most significant effects are between the ages of 70 and 90. The significant direct effects tend to appear at earlier ages than the significant indirect effects. For older individuals, their cognitive decline is attributed more to the AD pathology-mediated indirect effects, rather than the direct effect of the APOE genotype. For individuals aged 90 or older, the reduced effects may be attributed to survival bias. Furthermore, this phenomenon may be explained by cognitive resilience, in which certain individuals maintain stable cognitive function despite significant AD pathology. In individuals younger than 85, we observe a worsening trend in the NIE, indicating that aging amplifies the indirect effect of APOE 
4, leading to cognitive decline.
Figure 4.
Individual NDE and NIE change with age. (A) Age trend of NDE and NIE. Individual estimates are shown in dots with orange or black. Credible intervals are shown in yellow or gray, with orange dots and yellow shades signifying statistically significant estimates. (B) GAM fitted trends of NDE and NIE. Each dot represents an individual estimate, and the curve traces the smoothed trend of NDE/NIE changes with age.
5. Discussion
We introduce a novel framework for heterogeneous causal mediation analysis that can flexibly capture complex nonlinear relationships and the treatment-by-mediator interaction in mediation processes. Compared with existing methods, our approach achieves more accurate estimation with lower bias and RMSE for nonlinear data structures. Our method is designed to assess heterogeneity in mediation effects, providing credible intervals with nominal CRs that competing methods often fail to achieve. We implement HMBART as a user-friendly R package. To the best of our knowledge, there is currently a lack of comprehensive software packages specifically tailored for nonlinear heterogeneous causal mediation analysis. The package also includes interpretive tools, such as SHAP variable importance, dependence plots, and regression tree summary plots, to assess the influence of moderators on the magnitude, direction, and significance of a mediation effect.
Future research can extend our approach in several directions. First, sensitivity analysis tools can be developed to assess robustness against unmeasured pretreatment confounding. Second, the framework can be expanded to handle continuous or multi-valued treatments, mediators, and outcomes. Third, more rigorous methods for assessing the importance of moderators and identifying subgroups can be explored. Finally, we focused on one single mediator and required moderators to be pretreatment covariates. Future work will consider a more comprehensive mediation framework involving multiple mediators that may interact with one another. Overall, HMBART is a useful tool that provides low-bias point estimates and valid inference for heterogeneous mediation effects.
Supplementary Material
Web Appendices, Tables, and Figures referenced in Sections 2–4, along with the code for implementing the proposed method and conducting simulations, are available with this paper at the Biometrics website on Oxford Academic.
Contributor Information
Chen Liu, Department of Biostatistics and Health Data Science, University of Pittsburgh, Pittsburgh, PA 15261, United States.
Xu Qin, Department of Biostatistics and Health Data Science, University of Pittsburgh, Pittsburgh, PA 15261, United States; Department of Health and Human Development, University of Pittsburgh, Pittsburgh, PA 15260, United States.
Victor B Talisa, Department of Biostatistics and Health Data Science, University of Pittsburgh, Pittsburgh, PA 15261, United States; Department of Critical Care Medicine, University of Pittsburgh, Pittsburgh, PA 15213, United States.
Jiebiao Wang, Department of Biostatistics and Health Data Science, University of Pittsburgh, Pittsburgh, PA 15261, United States; Clinical and Translational Science Institute, University of Pittsburgh, Pittsburgh, PA 15261, United States.
Funding
Research reported in this publication was supported by the National Institutes of Health (NIH) under Award Numbers R21AG087057 and R01AG080590, National Science Foundation Faculty Early Career Development Program (CAREER) Award (2337612), and a U.S. Department of Education Institute of Education Sciences Grant (R305D200031). This research was supported in part by the University of Pittsburgh Center for Research Computing and Data, RRID:SCR_022735, through the resources provided. Specifically, this work used the HTC cluster, which is supported by NIH award number S10OD028483.
Conflicts of interest
None declared.
Data availability
The materials, data, and code that support the findings in this paper are available in the supplementary materials associated with this article. The R package implementing the proposed HMBART method is also available on GitHub at https://github.com/Lewis-ChenLiu/HMBART. The dataset used in this paper is from the Religious Orders Study and the Memory and Aging Project and available upon request through https://www.radc.rush.edu/.
References
- Baron R. M., Kenny D. A. (1986). The moderator–mediator variable distinction in social psychological research: Conceptual, strategic, and statistical considerations. Journal of Personality and Social Psychology, 51, 1173. [DOI] [PubMed] [Google Scholar]
- Carnegie N., Dorie V., Hill J. L. (2019). Examining treatment effect heterogeneity using BART. Observational Studies, 5, 52–70. [Google Scholar]
- Chipman H., George E., McCulloch R. (2006). Bayesian ensemble learning. Advances in Neural Information Processing Systems, 19, 265–272. [Google Scholar]
- Chipman H. A., George E. I., McCulloch R. E. (2010). BART: Bayesian additive regression trees. The Annals of Applied Statistics, 4, 266–298. [Google Scholar]
- Dorie V., Hill J., Shalit U., Scott M., Cervone D. (2019). Automated versus do-it-yourself methods for causal inference: Lessons learned from a data analysis competition. Statistical Science, 34, 43–68. [Google Scholar]
- Edwards J. R., Lambert L. S. (2007). Methods for integrating moderation and mediation: a general analytical framework using moderated path analysis. Psychological Methods, 12, 1. [DOI] [PubMed] [Google Scholar]
- Farrer L. A., , Cupples L. A., Haines J. L., Hyman B., Kukull W. A., Mayeux R. et al. (1997). Effects of age, sex, and ethnicity on the association between apolipoprotein E genotype and Alzheimer disease: a meta-analysis. JAMA, 278, 1349–1356. [PubMed] [Google Scholar]
- Hahn P. R., Murray J. S., Carvalho C. M. (2020). Bayesian regression tree models for causal inference: Regularization, confounding, and heterogeneous effects (with discussion). Bayesian Analysis, 15, 965–1056. [Google Scholar]
- Hastie T. J. (2017). Generalized additive models. In: Statistical Models in S, Boca Raton, FL, USA: Routledge. pp. 249–307. [Google Scholar]
- Hill J., Linero A., Murray J. (2020). Bayesian additive regression trees: A review and look forward. Annual Review of Statistics and Its Application, 7, 251–278. [Google Scholar]
- Holtzman D. M., Herz J., Bu G. (2012). Apolipoprotein E and apolipoprotein E receptors: normal biology and roles in Alzheimer disease. Cold Spring Harbor Perspectives in Medicine, 2, a006312. [DOI] [PMC free article] [PubMed] [Google Scholar]
- Imai K., Keele L., Tingley D. (2010). A general approach to causal mediation analysis. Psychological Methods, 15, 309. [DOI] [PubMed] [Google Scholar]
- Kapelner A., Bleich J. (2016). bartMachine: Machine learning with Bayesian additive regression trees. Journal of Statistical Software, 70, 1–40. [Google Scholar]
- Linero A. R. (2024). In nonparametric and high-dimensional models, Bayesian ignorability is an informative prior. Journal of the American Statistical Association, 119, 2785–2798. [Google Scholar]
- Linero A. R., Zhang Q. (2025). Mediation analysis using Bayesian tree ensembles. Psychological Methods, 30, 60–82. [DOI] [PubMed] [Google Scholar]
- Lundberg S. M., Lee S.-I. (2017). A unified approach to interpreting model predictions. In: Proceedings of the 31st International Conference on Neural Information Processing Systems, NIPS’17, 4768–4777., Red Hook, NY, USA: Curran Associates Inc. [Google Scholar]
- Nguyen T. Q., Schmid I., Stuart E. A. (2021). Clarifying causal mediation analysis for the applied researcher: Defining effects based on what we want to learn. Psychological Methods, 26, 255. [DOI] [PMC free article] [PubMed] [Google Scholar]
- Nichols E., Brickman A. M., Casaletto K. B., Dams-O’Connor K., George K. M., Kumar R. G. et al. (2023). AD and non-AD mediators of the pathway between the APOE genotype and cognition. Alzheimer’s and Dementia, 19, 2508–2519. [DOI] [PMC free article] [PubMed] [Google Scholar]
- Pearl J. (2022). Direct and indirect effects. In: Probabilistic and Causal Inference: The Works of Judea Pearl, New York, NY, USA: ACM. pp. 373–392. [Google Scholar]
- Preacher K. J., Rucker D. D., Hayes A. F. (2007). Addressing moderated mediation hypotheses: Theory, methods, and prescriptions. Multivariate Behavioral Research, 42, 185–227. [DOI] [PubMed] [Google Scholar]
- Qin X., Wang L. (2024). Causal moderated mediation analysis: Methods and software. Behavior Research Methods, 56, 1314–1334. [DOI] [PubMed] [Google Scholar]
- Rene L., Linero A. R., Slate E. (2023). Causal mediation and sensitivity analysis for mixed-scale data. Statistical Methods in Medical Research, 32, 1249–1266. [DOI] [PMC free article] [PubMed] [Google Scholar]
- Robins J. M., Greenland S. (1992). Identifiability and exchangeability for direct and indirect effects. Epidemiology, 3, 143–155. [DOI] [PubMed] [Google Scholar]
- Rubin D. B. (1978). Bayesian inference for causal effects: The role of randomization. The Annals of Statistics, 6, 34–58. [Google Scholar]
- Rubin D. B. (1990). Formal mode of statistical inference for causal effects. Journal of Statistical Planning and Inference, 25, 279–292. [Google Scholar]
- Thal D. R., Finucane M. M. (2023). Causal methods madness: Lessons learned from the 2022 ACIC competition to estimate health policy impacts. Observational Studies, 9, 3–27. [Google Scholar]
- Ting A., Linero A. R. (2025). Estimating heterogeneous causal mediation effects with Bayesian decision tree ensembles. Journal of the American Statistical Association, 120, 1400–1413. [Google Scholar]
Associated Data
This section collects any data citations, data availability statements, or supplementary materials included in this article.
Supplementary Materials
Web Appendices, Tables, and Figures referenced in Sections 2–4, along with the code for implementing the proposed method and conducting simulations, are available with this paper at the Biometrics website on Oxford Academic.
Data Availability Statement
The materials, data, and code that support the findings in this paper are available in the supplementary materials associated with this article. The R package implementing the proposed HMBART method is also available on GitHub at https://github.com/Lewis-ChenLiu/HMBART. The dataset used in this paper is from the Religious Orders Study and the Memory and Aging Project and available upon request through https://www.radc.rush.edu/.