TREATMENT EFFECT HETEROGENEITY AND IMPORTANCE MEASURES FOR MULTIVARIATE CONTINUOUS TREATMENTS

Abstract

Estimating the joint effect of a multivariate, continuous exposure is crucial, particularly in environmental health where interest lies in simultaneously evaluating the impact of multiple environmental pollutants on health. We develop novel methodology that addresses two key issues for estimation of treatment effects of multivariate, continuous exposures. We use nonparametric Bayesian methodology that is flexible to ensure our approach can capture a wide range of data generating processes. Additionally, we allow the effect of the exposures to be heterogeneous with respect to covariates. Treatment effect heterogeneity has not been well explored in the causal inference literature for multivariate, continuous exposures, and, therefore, we introduce novel estimands that summarize the nature and extent of the heterogeneity and propose estimation procedures for new estimands related to treatment effect heterogeneity. We provide theoretical support for the proposed models in the form of posterior contraction rates and show that it works well in simulated examples both with and without heterogeneity. Our approach is motivated by a study of the health effects of simultaneous exposure to the components of PM_2.5, where we find that the negative health effects of exposure to environmental pollutants are exacerbated by low socioeconomic status, race and age.

1. Introduction.

An important scientific question is understanding the effect of a multivariate treatment on an outcome, particularly in environmental health where individuals are exposed simultaneously to multiple pollutants and it is of interest to understand the joint impact of each of these pollutants on public health. Further, the effects of each of these pollutants might be heterogeneous with respect to characteristics of the individuals exposed, and it is important to account for and understand this treatment effect heterogeneity. Ignoring this heterogeneity can lead to misspecified outcome models, which lead to biased estimates of the average effect of the treatments on the outcome, and can also mask substantial effects of the treatment on specific subgroups of the population. In this paper we address this problem by developing methodology for estimating heterogeneous treatment effects of continuous, multivariate treatments.

In the past decade, there has been an increasing interest in analyzing the effects of multivariate exposures in environmental health where the exposures are referred to as environmental mixtures (Dominici et al. (2010), Agier et al. (2016), Stafoggia et al. (2017), Gibson, Goldsmith and Kioumourtzoglou (2019), Shin et al. (2023)). Analyzing environmental mixtures is a challenging problem for a number of reasons. For one, the scientific goal is not simply predicting the outcome, but rather understanding the effects of each of the individual exposures, and whether these exposures interact with each other. Therefore, off-the-shelf modern regression strategies such as tree-based models (Breiman (2001), Chipman, George and McCulloch (2010)) and Gaussian processes (Banerjee, Dunson and Tokdar (2013)) do not immediately apply. Consequently, there has been significant interest in developing methods to adapt these models for causal inference (Hill (2011), Hahn, Murray and Carvalho (2020), Ray and van der Vaart (2020)). A number of studies in the environmental statistics literature have tailored these methods for multivariate exposures. A common theme among these methods is the use of nonparametric Bayesian models. Gaussian processes are adapted for this purpose in the popular Bayesian kernel machine regression (BKMR, Bobb et al. (2015)) or in related work that explicitly identifies interactions between exposures (Ferrari and Dunson (2020)). Related Bayesian approaches use basis function expansions to identify important exposures or interactions among exposures for environmental mixtures (Antonelli et al. (2020), Wei et al. (2020), Samanta and Antonelli (2022)). Other, related approaches have been developed all with the goal of estimating the health effects of environmental mixtures and identifying exposures that affect the outcome (Herring (2010), Carrico et al. (2015), Narisetty et al. (2019), Boss et al. (2021), Ferrari and Dunson (2021)).

While the aforementioned approaches are useful for multivariate, continuous exposures, they all implicitly assume that the effect of the environmental exposures is the same across the entire population. This assumption does not hold in the context of ambient air pollution, where effects have been shown to vary by characteristics such as age, race and location (Wang et al. (2020), Bargagli-Stoffi et al. (2020)). Treatment effect heterogeneity has seen an explosion of interest in the causal inference literature, though it is typically restricted to a single treatment, either binary (Athey and Imbens (2016), Wager and Athey (2018), Hahn, Murray and Carvalho (2020), Semenova and Chernozhukov (2021), Fan et al. (2022), Shin and Antonelli (2023)) or with multiple levels (Chang and Roy (2025)). In the binary setting, heterogeneity is typically summarized by the conditional average treatment effect (CATE) function, $\mathbb{E}\left\{Y\right(1)-Y(0)\mid \mathit{X}=\mathit{x}\}$, where $Y\left(w\right)$ denotes the potential outcome under a binary exposure level $w$ and $\mathit{X}$ denotes covariates. While the CATE is a functional estimand, a number of summaries of this function have been proposed to simplify heterogeneity further. Henderson et al. (2020) and Chen et al. (2024) examined the posterior probabilities of differential treatment effect, which measures how much the posterior distributions of each CATE, evaluated at the observed covariates, differ from the sample mean of the CATE. Another straightforward univariate summary is the variance of the CATE (VTE, Levy et al. (2021)), which describes the overall degree of heterogeneity. Related quantities of interest to our study are treatment effect variable importance measures (TE-VIM), which have been recently proposed for univariate treatments in Hines, Diaz-Ordaz and Vansteelandt (2022) and Li, Hubbard and van der Laan (2023) to learn which covariates are key factors driving treatment effect heterogeneity. TE-VIM relates regression-based variable importance measures (Zhang and Janson (2020), Williamson et al. (2021)) and the VTE of Levy et al. (2021) within a causal inference framework. Nearly all of this work on treatment effect heterogeneity focuses on binary treatments, and there has been little to no work on the heterogeneous effects of multivariate, continuous exposures.

There exists a substantial gap in the literature on how to define estimands for continuous, multivariate exposures in the presence of treatment effect heterogeneity and on how to estimate these quantities using observed data and flexible modeling approaches that make as few modeling assumptions as possible. In this paper we aim to fill this gap by first proposing new estimands in this setting, which are interpretable and help describe the complex effect of multiple treatments, and how this effect varies by observed covariates. We extend treatment effect variable importance measures to this more complex setting and describe how to identify and estimate these estimands from observed data. These measures provide researchers with insight into causal effect modifiers for multivariate, continuous exposures and can provide policy-makers with detailed information about how treatments impact outcomes. Additionally, we develop novel estimation strategies using nonparametric Bayesian methodology. By decomposing outcome regression surfaces into distinct, identifiable functions, we are able to explicitly shrink treatment effects toward homogeneity, while still allowing for heterogeneity when it exists. We also make limited modeling assumptions, as we use novel extensions of Bayesian additive regression trees (BART) such as the SoftBART prior distribution (Linero and Yang (2018)) and targeted smoothing BART (Starling et al. (2020), Li, Linero and Murray (2023)) for interactions between exposures and covariates to allow for heterogeneity without imposing strong parametric assumptions.

In addition to methodological development, our nationwide study of the United States Medicare population provides novel epidemiological insights into the health effects of simultaneous exposure to multiple air pollutants. We find that increased exposure to ambient air pollution leads to increases in mortality rates and that this effect varies by race, age and socioeconomic status. This is the first study that we are aware of to investigate heterogeneity of air pollution mixture effects and provides public health researchers with more detailed evidence about the negative health effects of air pollution.

2. Estimands of interest.

Throughout, we assume that we observe $𝓓={\left\{\left(Y_i,{\mathit{X}}_i,{\mathit{W}}_i\right)\right\}}_{i=1}^n$ for each individual where $Y_i\in 𝒴\subset \mathbb{R}$ is the outcome of interest, ${\mathit{X}}_i\in 𝒳\subset {\mathbb{R}}^p$ is a vector of pretreatment covariates, and ${\mathit{W}}_i\in 𝒲\subset {\mathbb{R}}^q$ is a vector of continuous exposures. Also, we adopt the potential outcomes framework: $Y_i(\mathit{w})$ denotes the potential outcome that would be observed under exposure $\mathit{w}$. In order to denote potential outcomes in this manner and to identify causal effects from the observed data, we make the following assumptions.

Assumption 1.

1. SUTVA (Stable Unit Treatment Value Assumption, Rubin (1980)): Treatments of one unit do not affect the potential outcomes of other units (no interference), and there are not different versions of treatments such that $Y_i=Y_i\left({\mathit{W}}_i\right)$.

2. Positivity: $0<f_{W\mid X}(\mathit{w}\mid \mathit{X}=\mathit{x})$ for all $\mathit{x}$ and $\mathit{w}$, where $f_{W\mid X}$ is the conditional density of the treatments given covariates.

3. Unconfoundedness: $Y(\mathit{w})⫫\mathit{W}\mid \mathit{X}$ for all $\mathit{w}$.

The positivity assumption guarantees that each unit has a nonzero probability to be exposed to each treatment level for all possible values of pretreatment variables, at least in large samples. While in principle this assumption is empirically verifiable, there is not a well-established approach to doing so for multivariate, continuous treatments. Though not the focus of our work here, in Appendix D.4 we outline two distinct approaches to assess the plausibility of the multivariate positivity assumption and to examine how robust our results are to this critical assumption. First, we examine the marginal distribution of each exposure given the covariates and check whether each observation satisfies a univariate positivity criterion for each of the exposures. Second, we consider a trimmed causal estimand, which restricts the population being targeted to those observations with higher values of the conditional density of the treatments given the covariates when evaluated at the treatment values of interest. Unconfoundedness requires that the treatment W is independent of the potential outcomes, and it implies that there exist no unmeasured variables that confound the treatment-outcome relationship.

Because exposures and covariates are both multivariate, there are many causal estimands one could define. In the following sections, we first define estimands targeting both average treatment effects and treatment effect heterogeneity as well as variable importance measures summarizing this heterogeneity. While our focus in this paper remains on multivariate continuous exposures, we note that the estimands we propose can be applied to a broader class of exposures with multiple levels, regardless of their dimensionality. In such cases, estimation complexity is significantly reduced due to the lower dimensionality of univariate exposures or the fewer possible interventions for discrete exposures.

2.1. Marginal and heterogeneous effects estimands.

One potential estimand to look at in this scenario is $\mathbb{E}[Y(\mathit{w})\mid \mathit{X}=\mathit{x}]$; however, there are infinitely many values that $\mathit{w}$ can take, and, therefore, it is difficult to interpret this estimand. One option is to look at values of $\mathit{w}$, where all but one of the exposures is fixed, and then visualize this estimand as a function of one exposure. A simpler setting is one in which there are two exposure levels ${\mathit{w}}_1$ and ${\mathit{w}}_0$ under which we want to compare the outcomes. For example, one may be interested in examining what would happen if an intervention was applied that lowered the level of pollution by a specific amount. In this case, one can use the pollution levels without the intervention as ${\mathit{w}}_0$ and those with the intervention as ${\mathit{w}}_1$. Then the effect of the intervention can be described as a function of covariates by

$${\tau }_{{\mathit{w}}_1,{\mathit{w}}_0}\left(\mathit{x}\right)\equiv \mathbb{E}\left\{Y\left({\mathit{w}}_1\right)-Y\left({\mathit{w}}_0\right)\mid \mathit{X}=\mathit{x}\right\},$$

and the extent of treatment effect heterogeneity is dictated by how strongly this depends on X. If interest is in marginal treatment effects, this quantity can be averaged over the covariate distribution to obtain a marginal effect

$${\overline{\tau }}_{{\mathit{w}}_1,{\mathit{w}}_0}=\mathbb{E}\left[{\tau }_{{\mathit{w}}_1,{\mathit{w}}_0}(\mathit{X})\right].$$

This problem effectively reduces to the binary treatment setting where the conditional average treatment effect (CATE) is originally defined, as there are only two exposure levels of interest. Because of this, many existing approaches to summarizing effect heterogeneity from the existing literature can be used. The variance of this function, $\mathrm{V}\mathrm{a}\mathrm{r}\left[{\tau }_{{\mathit{w}}_1,{\mathit{w}}_0}(\mathit{X})\right]$, or the treatment effect variable importance measures of Hines, Diaz-Ordaz and Vansteelandt (2022) can be incorporated analogously. For this reason and in the following section, we focus on the more complex setting where we have more than two exposure levels of interest, yet we want to study heterogeneity and examine the impact of each covariate in the degree of heterogeneity.

2.2. Multivariate treatment effect variable importance measures.

In many scenarios we have more than two exposure levels of interest, yet we still want to define relevant quantities that provide information about treatment effect heterogeneity and which covariates are driving this heterogeneity. For this section we assume that we have a reference level of the exposure denoted by ${\mathit{w}}_0$, which is chosen a priori. While the following estimands are well defined for any choice of ${\mathit{w}}_0$ as long as the positivity assumption holds, we recommend selecting ${\mathit{w}}_0$ within a reasonable range, such as the mean of the exposures or a value between their first and third quartiles, to reduce the extent of extrapolation required. In Appendix D.2, we evaluate how robust the results of Section 6 are to the choice of ${\mathit{w}}_0$ and find that results are stable across different choices.

With ${\mathit{w}}_0$ fixed, we can define the following quantity:

$${\tau }_{{\mathit{w}}_0}(\mathit{x},\mathit{w})\equiv \mathbb{E}\left[Y(\mathit{w})-Y\left({\mathit{w}}_0\right)\mid \mathit{X}=\mathit{x}\right],$$

which describes how the potential outcome surface varies by both $\mathit{x}$ and $\mathit{w}$. This function is difficult to interpret, as it is a function of two multivariate arguments. We propose multivariate treatment effect variable importance measures (MTE-VIM) to summarize how much the heterogeneity of ${\tau }_{{\mathit{w}}_0}(\mathit{x},\mathit{w})$ is driven by each particular covariate.

Before defining the variable importance metric, we must first define the overall amount of heterogeneity of the treatment effect, which we define as

$$\varphi ={\mathbb{E}}_{\mathit{W}}\left[{\mathrm{V}\mathrm{a}\mathrm{r}}_{\mathit{X}}\left\{{\tau }_{{\mathit{w}}_0}(\mathit{X},\mathit{W})\right\}\right].$$

This looks at the variability of the treatment effects as a function of $\mathit{X}$ for a fixed exposure level, but this variability may differ for different values of exposures $\mathit{W}$, and, therefore, we average this variability across the range of exposures. In practice, the expectation and variance are taken with respect to the empirical distributions of exposures and covariates, respectively. Note that instead of the marginal variance of $\mathit{X}$, we could have alternatively used the conditional variance of $\mathit{X}$ given $\mathit{W}$. We believe, however, that the marginal variance is a more appropriate choice for measuring the amount of treatment effect heterogeneity, as it reflects variation of treatment effects that exist regardless of which units receive each level of exposure. Further, we will see in Section 3.3 that the conditional variance would complicate our estimation of this quantity as it would require a model for the distribution of $\mathit{X}$ given $\mathit{W}$, which is difficult to specify given the dimension of both the covariates and exposures. Next, we define

$${\varphi }_j={\mathbb{E}}_{\mathit{W}}\left({\mathrm{V}\mathrm{a}\mathrm{r}}_{\mathit{X}-j}\left[{\mathbb{E}}_{X_j\mid \mathit{X}-j}\left\{{\tau }_{{\mathit{w}}_0}(\mathit{X},\mathit{W})\right\}\right]\right),$$

where ${\mathit{X}}_{-j}$ denotes the $p-1$ remaining covariates without the jth covariate $X_j$. This is similar to $\varphi$, but we first take the conditional expectation of ${\tau }_{{\mathit{w}}_0}(\mathit{X},\mathit{W})$ with respect to the jth covariate, given the remaining covariates, so that the term inside the variance is a function of only ${\mathit{X}}_{-j}$ and $\mathit{W}$. We know by the law of total variance that $\varphi -{\varphi }_j\in [0,\varphi ]$, and this quantity measures the amount heterogeneity of ${\tau }_{{\mathit{w}}_0}(\mathit{X},\mathit{W})$ that cannot be explained without $X_j$.

Assuming some degree of heterogeneity, that is, $\varphi >0$, we propose the following estimand, which we call MTE-VIM, to describe the importance of each covariate to the overall heterogeneity of the causal effect:

$${\psi }_j=1-{\varphi }_j/\varphi \in [0,1].$$

This can be interpreted as the proportion of the treatment effect heterogeneity of ${\tau }_{{\mathit{w}}_0}(\mathit{X},\mathit{W})$ not explained by ${\mathit{X}}_{-j}$. Under the proposed model in Section 3, which implies that the treatment effect heterogeneity is additive in covariates, the sum of all ${\psi }_j$ equals one when covariates are mutually independent. We note that one should carefully interpret ${\psi }_j$ in the presence of highly correlated covariates. If two covariates are highly correlated, then they will both have small values of ${\psi }_j$, regardless of whether they modify the treatment effect. This issue is not unique to our estimand, as it is present in other commonly used variable importance measures that focus on prediction instead of treatment effect heterogeneity (Verdinelli and Wasserman (2024)). We don’t view this as a problem, as these estimands are simply descriptive measures aiming to assign relative importance to each covariate in the heterogeneity of the treatment effect. If one is concerned about this issue, then they can used a grouped version of the MTE-VIM, given by ${\psi }_s$, where $s$ is a subset of $\{1,2,\dots ,p\}$. A detailed description of this can be found in Hines, Diaz-Ordaz and Vansteelandt (2022), though the general idea is to find the proportion of the treatment effect heterogeneity that can’t be explained by the remaining covariates not in $s$. One can use a priori knowledge to select groups of covariates or use the sample correlation matrix of the covariates to identify groups of correlated variables.

Lastly, while our paper focuses mostly on treatment effect heterogeneity and the role that covariates play in modifying the treatment effect, similar ideas could be used to identify which treatments have the largest impact on the outcome. Specifically, one could define an estimand based on ${\mathbb{E}}_{\mathit{X}}\left[{\mathrm{V}\mathrm{a}\mathrm{r}}_{\mathit{W}}\left\{{\tau }_{{\mathit{w}}_0}(\mathit{X},\mathit{W})\right\}\right]$ and ${\mathbb{E}}_{\mathit{X}}\left({\mathrm{V}\mathrm{a}\mathrm{r}}_{{\mathit{W}}_{-k}}\left[{\mathbb{E}}_{W_k\mid {\mathit{W}}_{-k}}\left\{{\tau }_{{\mathit{w}}_0}(\mathit{X},\mathit{W})\right\}\right]\right)$ to identify the proportion of variability in the treatment effect that can not be explained without treatment $k$, which would provide insights about which treatments are driving the treatment effect.

3. Estimation issues.

Under the assumptions introduced in Section 2, for any fixed $\mathit{w}$ we can identify $\mathbb{E}\{Y(\mathit{w})\mid \mathit{X}=\mathit{x}\}$ from the observed data by $\mathbb{E}(Y\mid \mathit{X}=\mathit{x},\mathit{W}=\mathit{w})$. Other identification strategies, such as those involving propensity scores (Rosenbaum and Rubin (1983)) or combinations of outcome models and propensity scores (Bang and Robins (2005)), are common in causal inference. However, these do not apply here due to the difficulty of estimating the density of the multivariate treatment, given the covariates, and because these identification strategies are not well studied for multivariate, continuous treatments. Due to this identification result, all estimands, including the MTE-VIM, are identifiable, given the conditional outcome regression, and, therefore, we focus our estimation on this quantity.

We separate the conditional outcome regression surface into three parts: the main effect of covariates $\mathit{X}$, the main effect of exposures $\mathit{W}$ and the interactions of covariates and exposures on the outcome. Therefore, we write our model as follows:

$$\mathbb{E}(Y\mid \mathit{X},\mathit{W})=c+f(\mathit{X})+g(\mathit{W})+h(\mathit{X},\mathit{W}).$$ (1)

Note that this model is overparameterized, as it currently stands because the $f(\cdot )$ and $g(\cdot )$ functions can be absorbed into the $h(\cdot ,\cdot )$ function. For this reason these functions are not individually identifiable, and only their sum is identified. We write the model in this way so that we can explicitly shrink each of these components separately, and we discuss in this section how to structure the model so that each of these functions is individually identifiable and, therefore, amenable to shrinkage and regularization.

Let ${𝒳}_j$ and ${𝒳}_{-j}$ denote the support of the jth covariate and the remaining covariates, respectively. To simplify the structure of interactions between covariates and exposures, we make the following assumption.

Assumption 2.

For any ${\mathit{w}}_1,{\mathit{w}}_0\in 𝒲,x_j^{\ast },x_j\in {𝒳}_j$ and ${\mathit{x}}_{-j}^{\ast },{\mathit{x}}_{-j}\in {𝒳}_{-j}$,

$$\begin{gathered} \mathbb{E}[Y({\mathit{w}}_1)-Y({\mathit{w}}_0)\mid X_j=x_j^{\ast },{\mathit{X}}_{-j}={\mathit{x}}_{-j}]-\mathbb{E}[Y({\mathit{w}}_1)-Y({\mathit{w}}_0)\mid X_j=x_j,{\mathit{X}}_{-j}={\mathit{x}}_{-j}] \\ =\mathbb{E}[Y({\mathit{w}}_1)-Y({\mathit{w}}_0)\mid X_j=x_j^{\ast },{\mathit{X}}_{-j}={\mathit{x}}_{-j}^{\ast }]-\mathbb{E}[Y({\mathit{w}}_1)-Y({\mathit{w}}_0)\mid X_j=x_j,{\mathit{X}}_{-j}={\mathit{x}}_{-j}^{\ast }] \end{gathered}$$

for $j=1,\dots ,p$.

This assumption ensures that the treatment effect heterogeneity attributable to the jth covariate does not depend on the other covariates. In our model this assumption implies that $h(\mathit{X},\mathit{W})={\sum }_{j=1}^p{h}_j\left(X_j,\mathit{W}\right)$, which restricts interactions between covariates and exposures to be limited to interactions between a single covariate and the multivariate exposures. Given that most environmental mixture studies assume no treatment effect heterogeneity, which is a much stronger assumption, and that previous studies for a single exposure have suggested that the treatment effect of ${\mathrm{P}\mathrm{M}}_{2.5}$ on mortality rate is primarily modified by one covariate at a time (Bargagli-Stoffi et al. (2020), Lee, Small and Dominici (2021)), we feel this is a mild and reasonable assumption in the environmental health context. In simulation studies in Appendix C.2, even when this assumption is violated, we find that our proposed model provides a reasonable estimate compared to existing methods. Nonetheless, if strong, higher order covariate interactions with multivariate exposure effects are expected, we recommend using a more complex model, such as a single BART model with both $\mathit{X}$ and $\mathit{W}$ as input variables that places no additivity restrictions on the model.

The first hurdle to identification of the individual functions is that $h_j\left(X_j,\mathit{W}\right)$ may capture the main effect of either the covariate or exposures. For example, we could have $h_j\left(X_j,\mathit{W}\right)=X_j+W_1+X_j{W}_1$, which captures both main effects and interaction terms, when ideally it would only capture $X_j{W}_1$. Effectively, we want this function to only be nonzero if there is truly an interaction between the exposures and $X_j$. For this reason we restrict $h_j\left(X_j,\mathit{W}\right)$ to be of the form $h_j^{\text{cov}}\left(X_j\right)h_j^{\text{exp}}(\mathit{W})$, where $h_j^{\text{cov}}$ is a nonconstant function of covariate $j$ and $h_j^{\text{exp}}$ is a nonconstant function of the exposures. This allows us to prevent the interaction function $h_j\left(X_j,\mathit{W}\right)$ from simply absorbing the main effects of $X_j$ and $\mathit{W}$. We describe our strategy for each of these separable functions in the following sections.

3.1. Identification through shifting.

The formulation for $h(\mathit{X},\mathit{W})$ in the previous section helped ensure that the interaction functions do not solely contain main effects of the exposures or covariates; however, the individual functions are still only identifiable up to constant shifts. For example, without further restriction, shifting $f(\mathit{X})$ upward by $\delta$ and $g(\mathit{W})$ and $h(\mathit{X},\mathit{W})$ downward by $\delta /2$ leads to the same likelihood. This is problematic, as we would like our interaction functions to be nonzero only when there is heterogeneity of the treatment effect. If these functions are not identifiable, it becomes more difficult to shrink them toward zero when there is no heterogeneity in the model.

One common way to address this is to put moment restrictions on the functions such that ${\mathbb{E}}_{\mathit{X}}\{f(\mathit{X})\}={\mathbb{E}}_{\mathit{W}}\{g(\mathit{W})\}={\mathbb{E}}_{\mathit{X},\mathit{W}}\left\{h\right(\mathit{X},\mathit{W}\left)\right\}=0$ with the additional restriction that ${\mathbb{E}}_{\mathit{X}}\left\{h\right(\mathit{X},\mathit{W}\left)\right\}={\mathbb{E}}_{\mathit{W}}\left\{h\right(\mathit{X},\mathit{W}\left)\right\}=0$. This approach is difficult to implement, however, as enforcing these conditions requires the conditional density of covariates, given exposures, and vice versa, which is difficult to estimate given the dimension of the exposures and covariates. We use an alternative restriction that also leads to identifiability of individual functions but is straightforward to implement. Specifically, we enforce that $f\{\mathbb{E}(\mathit{X})\}=g\{\mathbb{E}(\mathit{W})\}=0$ and $h\{\mathbb{E}(\mathit{X}),\mathit{w}\}=h\{\mathit{x},\mathbb{E}(\mathit{W})\}=0$ for all $\mathit{x}$ and $\mathit{w}$. Under this restriction the sum of the constant term and $f(\mathit{X})$ represents the conditional expectation of the potential outcome as a function of covariates when exposures are set to their mean. The term $g(\mathit{W})$ represents the shifted average exposure-response curve at the mean of covariates, with the shift ensuring that it equals zero at the mean of exposure. Finally, $h(\mathit{X},\mathit{W})$ captures the remaining treatment effect heterogeneity. We show in Appendix B that this leads to identifiability of the individual functions. In practice, this restriction is achieved by centering our estimated functions at each iteration of an MCMC algorithm. If we let all functions with a 0-subscript denote unrestricted functions that do not enforce the aforementioned constraints, then these can be written as

$$\begin{gathered} \mathbb{E}(Y\mid \mathit{X},\mathit{W}) \\ =c_0+f_0(\mathit{X})+g_0(\mathit{W})+\sum _{j=1}^p{h}_{j0}\left(X_j,\mathit{W}\right) \\ =\left[c_0+f_0(\mathbb{E}(\mathit{X}))+g_0(\mathbb{E}(\mathit{W}))+\sum _{j=1}^p{h}_{j0}\left(\mathbb{E}\left(X_j\right),\mathbb{E}(\mathit{W})\right)\right] \\ +\left[f_0(\mathit{X})-f_0(\mathbb{E}(\mathit{X}))+\sum _{j=1}^p\left\{h_{j0}\left(X_j,\mathbb{E}(\mathit{W})\right)-h_{j0}\left(\mathbb{E}\left(X_j\right),\mathbb{E}(\mathit{W})\right)\right\}\right] \\ +\left[g_0(\mathit{W})-g_0(\mathbb{E}(\mathit{W}))+\sum _{j=1}^p\left\{h_{j0}\left(\mathbb{E}\left(X_j\right),\mathit{W}\right)-h_{j0}\left(\mathbb{E}\left(X_j\right),\mathbb{E}(\mathit{W})\right)\right\}\right] \\ +\left[\sum _{j=1}^p\left\{h_{j0}\left(X_j,\mathit{W}\right)-h_{j0}\left(X_j,\mathbb{E}(\mathit{W})\right)-h_{j0}\left(\mathbb{E}\left(X_j\right),\mathit{W}\right)+h_{j0}\left(\mathbb{E}\left(X_j\right),\mathbb{E}(\mathit{W})\right)\right\}\right] \\ =c+f(\mathit{X})+g(\mathit{W})+\sum _{j=1}^p{h}_j\left(X_j,\mathit{W}\right). \end{gathered}$$

At each iteration of the MCMC, the functions are shifted to satisfy this condition, and, therefore, our individual functions are identifiable, and the interaction functions should only be nonzero when there is heterogeneity of the treatment effect. Because of this, we can apply shrinkage priors to the $h_j\left(X_j,\mathit{W}\right)$ functions, which should improve estimation when heterogeneity is not present. Note that the alternative restriction that we impose is baseline-free in the sense that one can replace the chosen baseline, $\mathbb{E}(\mathit{X})$ and $\mathbb{E}(\mathit{W})$, with any values of the covariates and exposures.

3.2. SoftBART and targeted smoothing.

In this section we detail the specific models we use for each of the $f(\mathit{X}),g(\mathit{W})$ and $h(\mathit{X},\mathit{W})$ functions. For the main effect functions, we use the SoftBART prior of Linero and Yang (2018), while we use BART with targeted smoothing (tsBART, Starling et al. (2020), Li, Linero and Murray (2023)) for estimation of the interaction functions. Before detailing each of these BART extensions, we first briefly introduce the original BART model. BART was introduced by Chipman, George and McCulloch (2010) and is a fully Bayesian ensemble-of-trees model that has seen increasing usage in causal inference due to the success of the original BART and its variants (Dorie et al. (2019)). Specifically, it assumes that

$$\begin{gathered} Z=f(\mathit{v})+ϵ,ϵ\stackrel{\text{i.i.d.}}{\sim }N\left(0,{\sigma }^2\right), \\ f(\mathit{v})=\sum _{m=1}^M\mathrm{T}\mathrm{r}\mathrm{e}\mathrm{e}\left(\mathit{v},{𝒯}_m,{ℳ}_m\right), \end{gathered}$$

where $\mathrm{T}\mathrm{r}\mathrm{e}\mathrm{e}\left(\mathit{v},{𝒯}_j,{ℳ}_j\right)$ represents a regression tree, with tree structure ${𝒯}_j$ inducing a step function in $\mathit{v}$, and leaf parameters ${ℳ}_m=\left\{{\mu }_{m1},\dots ,{\mu }_{m{b}_m}\right\}$ for prediction. The standard prior distribution sets ${\mu }_{mk}\sim N\left(0,{\sigma }_{\mu }^2/M\right)$ so that the overall prior variance is $\mathrm{V}\mathrm{a}\mathrm{r}\left(f\right(\mathit{v}\left)\right)={\sigma }_{\mu }^2$. Although we do not detail the hyperparameters of BART in this paper, we note that the default setting of Chipman, George and McCulloch (2010) encourages each tree to be shallow and shrinks leaf parameters toward zero so that Tree(·) can be considered as a weak learner. We refer interested readers to Linero (2017) and Hill, Linero and Murray (2020) for recent reviews of BART.

Although the original BART approach is successful, due to the piecewise constant nature of tree ensembles, it suffers when the underlying truth is smooth, even if it is relatively simple. To overcome this shortcoming, Linero and Yang (2018) propose a smooth modification of BART (SoftBART or SBART) by allowing $\mathit{v}$ to follow a probabilistic path, that is, randomizing the decision rule at each split. Because we expect the effects of environmental exposures or other continuous treatments on the outcome to be smooth, we use SoftBART for estimation of both $f(\mathit{X})$ and $g(\mathit{W})$.

Another smooth variant of BART is tsBART (Starling et al. (2020), Li, Linero and Murray (2023)), which is originally motivated by time-to-event data and density estimation, though we use it for estimation of the interaction functions $h_j\left(X_j,\mathit{W}\right)$. Unlike SoftBART which smooths the regression function over all variables, tsBart smooths over a single targeted variable, say $u$. After centering at $\gamma (u)$, a baseline function of $u$, the model can be written as $f(u,\mathit{v})=\gamma (u)+{\sum }_{j=1}^m\mathrm{t}\mathrm{s}\mathrm{T}\mathrm{r}\mathrm{e}\mathrm{e}\left(u,\mathit{v},{𝒯}_m,{ℳ}_m\right)$ where each terminal node is associated with ${ℳ}_m=\left\{{\mu }_{m1}(u),\dots ,{\mu }_{m{b}_m}(u)\right\}$. Here each element ${\mu }_{mb}(u)$ follows a Gaussian process in $u$ with mean zero and covariance function $\mathrm{\Sigma }\left(u,u^{\prime }\right)$ so that $f(u,\mathit{v})$, the sum of $m$ Gaussian processes, is continuous in $u$. Following Li, Linero and Murray (2023), we approximate the model by $f(u,\mathit{v})=\gamma (u)+{\sum }_{m=1}^M{ℬ}_m(u)\mathrm{T}\mathrm{r}\mathrm{e}\mathrm{e}\left(\mathit{v},{𝒯}_m,{ℳ}_m\right)$, where ${ℬ}_m(u)=\sqrt{2}\mathrm{c}\mathrm{o}\mathrm{s}\left({\omega }_{m}u+b_m\right)$ is a random basis function where ${\omega }_m\stackrel{\text{i.i.d.}}{\sim }N(0,{\rho }^{-2})$ and $b_m\stackrel{\text{i.i.d.}}{\sim }U(0,2\pi )$ so that ${\sum }_{m=1}^M{ℬ}_m(u)\mathrm{t}\mathrm{r}\mathrm{e}\mathrm{e}\left(\mathit{v},{𝒯}_m,{ℳ}_m\right)$ weakly converges to a Gaussian process $\mathrm{G}\mathrm{P}(0,\mathrm{\Sigma }(\cdot ,\cdot \left)\right)$ as $M\to \mathrm{\infty }$ where $\mathrm{\Sigma }\left(u,u^{\prime }\right)={\sigma }_{\mu }^2\mathrm{e}\mathrm{x}\mathrm{p}\left\{-\left(u-u^{\prime }\right)/\left(2{\rho }^2\right)\right\}$.

By fitting tsBART for each interaction function $h_j\left(X_j,\mathit{W}\right)$ in our model by setting $u=X_j$ and $\mathit{v}=\mathit{W}$, we have that each fitted function is necessarily the sum of products of a continuous function of $X_j$ and a flexible tree of $\mathit{W}$ so that it does not include a function solely of $X_j$ or $\mathit{W}$. The default prior specifications from Linero and Yang (2018) and Li, Linero and Murray (2023) are used throughout our implementation. However, to enhance stability, we regularize the interaction function more aggressively by capping the variance of each terminal node, ${\sigma }_{\mu }$, at $3.5/2\sqrt{M}$, which is its default initial specification.

3.3. Estimation of MTE-VIM.

We now detail the estimation strategy for the proposed variable importance measures. Recall that MTE-VIM for the jth covariate, ${\psi }_j$, is composed of the total heterogeneity $\varphi ={\mathbb{E}}_{\mathit{W}}\left[{\mathrm{V}\mathrm{a}\mathrm{r}}_{\mathit{X}}\left\{{\tau }_{{\mathit{w}}_0}(\mathit{X},\mathit{W})\right\}\right]$ and the total heterogeneity accounted for by ${\mathit{X}}_{-j}$, given by ${\varphi }_j={\mathbb{E}}_{\mathit{W}}\left({\mathrm{V}\mathrm{a}\mathrm{r}}_{{\mathit{X}}_{-j}}\left[{\mathbb{E}}_{X_j\mid {\mathit{X}}_{-j}}\left\{{\tau }_{{\mathit{w}}_0}(\mathit{X},\mathit{W})\right\}\right]\right)$. As before, under the assumptions of Section 2, we can write ${\tau }_{{\mathit{w}}_0}(\mathit{x},\mathit{w})=\mathbb{E}(Y\mid \mathit{X}=\mathit{x},\mathit{W}=\mathit{w})-\mathbb{E}(Y\mid \mathit{X}=\mathit{x},\mathit{W}={\mathit{w}}_0$ ), and, therefore, the posterior distribution of ${\tau }_{{\mathit{w}}_0}(\mathit{x},\mathit{w})$ can be obtained once we have the posterior distribution of the outcome regression model. Obtaining $\varphi$ is relatively straightforward, as we use empirical distributions of both the exposures and covariates to approximate moments with respect to $\mathit{W}$ or $\mathit{X}$. We can construct an $n\times n$ matrix for which the ($k,l$)th element is ${\tau }_{{\mathit{w}}_0}\left({\mathit{X}}_l,{\mathit{W}}_k\right)$, where ${\mathit{X}}_l$ and ${\mathit{W}}_k$ represent the observed covariates of the lth observation and the exposures of the kth observation, respectively. Then the sample variance of the kth row of the matrix would estimate ${\mathrm{V}\mathrm{a}\mathrm{r}}_{\mathit{X}}\left\{{\tau }_{{\mathit{w}}_0}\left(\mathit{X},{\mathit{W}}_k\right)\right\}$. Taking the sample mean of these $n$ column sample variances estimates ${\mathbb{E}}_{\mathit{W}}\left[{\mathrm{V}\mathrm{a}\mathrm{r}}_{\mathit{X}}\left\{{\tau }_{{\mathit{w}}_0}(\mathit{X},\mathit{W})\right\}\right]$. This estimation process can be illustrated as follows:

$$\begin{gathered} \left(\begin{array}{cccc}{\tau }_{{\mathit{w}}_0}\left({\mathit{X}}_1,{\mathit{W}}_1\right)& {\tau }_{{\mathit{w}}_0}\left({\mathit{X}}_2,{\mathit{W}}_1\right)& \dots & {\tau }_{{\mathit{w}}_0}\left({\mathit{X}}_n,{\mathit{W}}_1\right)\\ {\tau }_{{\mathit{w}}_0}\left({\mathit{X}}_1,{\mathit{W}}_2\right)& \ddots & & {\tau }_{{\mathit{w}}_0}\left({\mathit{X}}_n,{\mathit{W}}_2\right)\\ ⋮& & \ddots & ⋮\\ {\tau }_{{\mathit{w}}_0}\left({\mathit{X}}_1,{\mathit{W}}_n\right)& {\tau }_{{\mathit{w}}_0}\left({\mathit{X}}_2,{\mathit{W}}_n\right)& \dots & {\tau }_{{\mathit{w}}_0}\left({\mathit{X}}_n,{\mathit{W}}_n\right)\end{array}\right)⟶\left(\begin{array}{c}{\mathrm{V}\mathrm{a}\mathrm{r}}_{\mathit{X}}\left[{\tau }_{{\mathit{w}}_0}\left(\mathit{X},{\mathit{W}}_1\right)\right]\\ {\mathrm{V}\mathrm{a}\mathrm{r}}_{\mathit{X}}\left[{\tau }_{{\mathit{w}}_0}\left(\mathit{X},{\mathit{W}}_2\right)\right]\\ ⋮\\ {\mathrm{V}\mathrm{a}\mathrm{r}}_{\mathit{X}}\left[{\tau }_{{\mathit{w}}_0}\left(\mathit{X},{\mathit{W}}_n\right)\right]\end{array}\right) \\ \downarrow \\ {\mathbb{E}}_{\mathit{W}}\left[{\mathrm{V}\mathrm{a}\mathrm{r}}_{\mathit{X}}\left\{{\tau }_{{\mathit{w}}_0}(\mathit{X},\mathit{W})\right\}\right]. \end{gathered}$$

The same strategy can be used for ${\varphi }_j$ replacing ${\tau }_{{\mathit{w}}_0}\left({\mathit{X}}_l,{\mathit{W}}_k\right)$ with ${\mathbb{E}}_{X_j\mid {\mathit{X}}_{-j}}\left\{{\tau }_{{\mathit{w}}_0}\left({\mathit{X}}_l,{\mathit{W}}_k\right)\right\}$ everywhere, so all that is left is to describe how to obtain ${\mathbb{E}}_{X_j\mid {\mathit{X}}_{-j}}\left\{{\tau }_{{\mathit{w}}_0}\left({\mathit{X}}_l,{\mathit{W}}_k\right)\right\}$. For ease of exposition, we rewrite ${\tau }_{{\mathit{w}}_0}\left({\mathit{X}}_l,{\mathit{W}}_k\right)$ as ${\tau }_{{\mathit{w}}_0}\left(X_{lj},{\mathit{X}}_{l(-j)},{\mathit{W}}_k\right)$, where $X_{lj}$ denotes the jth covariate of the lth individual and ${\mathit{X}}_{l(-j)}$ denotes the remaining covariates of the lth individual. A nonparametric estimate of this mean can be defined as

$${\hat{\mathbb{E}}}_{X_j\mid {\mathit{X}}_{-j}}\left\{{\tau }_{{\mathit{w}}_0}(\mathit{X},\mathit{W})\right\}=\frac{{\sum }_{i=1}^n{\tau }_{{\mathit{w}}_0}\left(X_{ij},{\mathit{X}}_{-j},\mathit{W}\right)K\left({\mathit{X}}_{i(-j)}-{\mathit{X}}_{-j}\right)}{{\sum }_{i=1}^{n}K\left({\mathit{X}}_{i(-j)}-{\mathit{X}}_{-j}\right)},$$

where $K(\cdot )$ is an appropriately defined kernel. We do not provide specific details regarding the choice of the kernel or the bandwidth parameter, as standard approaches to selection of these parameters in nonparametric regression could be applied. The sample correlation matrix of the covariates could help inform this decision, as a larger bandwidth would be appropriate if $X_j$ is assumed to be (approximately) independent of the other covariates. In this case, kernel smoothing simplifies to the sample average of ${\tau }_{{\mathit{w}}_0}\left(X_{ij},{\mathit{X}}_{-j},\mathit{W}\right)$, which we do in the simulation studies in Section 5 where covariates are mutually independent. For a comprehensive overview of kernel smoothing techniques, see Hastie et al. (2009), Ghosh (2018). While the kernel approach is applicable if the dimension of $\mathit{X}$ is small, it won’t work well as the number of covariates grows. In this more difficult setting, we could make parametric assumptions and estimate the mean of $X_j$, given ${\mathit{X}}_{-j}$, using a regression model. For instance, if $X_j$ is continuous, we could assume that it follows a normal distribution, and we can estimate the conditional mean and variance using linear regression. This would provide us an estimate of the conditional density $f_{X_j\mid {\mathit{X}}_{-j}}$ from which we can calculate ${\mathbb{E}}_{X_j\mid {\mathit{X}}_{-j}}\left\{{\tau }_{{\mathit{w}}_0}(\mathit{X},\mathit{W})\right\}$. We use this regression-based approach in the data analysis in Section 6. Since the kernel smoothing and parametric regression approaches discussed above pertain specifically to the calculation of the MTE-VIM, which is a purely descriptive measure of the treatment effect function, they are separated from the MCMC sampling. Also, we believe that the MTE-VIM, which aims to provide some degree of interpretability of the multidimensional treatment effect function, remains useful, even under mild misspecification of the covariate distribution. In principle, one could employ flexible Bayesian nonparametric models for estimating the covariate distribution and incorporate this procedure into the MCMC sampling algorithm to improve the accuracy of MTE-VIM estimation. We do not pursue that approach here, as it would significantly increase computation time, especially for large datasets such as those used in our analysis.

Note that both of these require the construction of an $n\times n$ matrix at each MCMC sample, which is computationally intensive. This is particularly problematic for the analysis of Medicare data in Section 6, where our sample consists of all 38,702 zip codes in the United States. As an approximation to this calculation, one can use a blocking scheme where we construct several submatrices of the data when calculating the variable importance metrics. Specifically, we split the overall sample into $K$ groups of approximately equal size. Letting the sample size in group $k$ be $n_k$, we can estimate the variable importance metric using the subset of the data in group $k$, which only requires the calculation of an $n_k\times n_k$ matrix. We can do this for each group separately and then average results across groups for a final estimate of ${\psi }_j$ for $j=1,\dots ,p$. We see in Appendix C.1 that this blocking scheme performs equally well as using the full sample but is significantly faster on large data sets. Therefore, we use our method with the blocking scheme throughout the paper. We also note that the blocking scheme is entirely independent of the fitting of the conditional outcome models, and, therefore, it should not affect the near-minimax contraction rates for estimating conditional outcomes, which we provide in Section 4.

4. Posterior contraction rates.

Next, we study the rate at which the posterior distributions concentrate in the model described in Section 3. Let $𝒮(\alpha ,p,d)$ denote a set of $\alpha$-Hölder continuous functions on $[0,1{]}^p$ that are constant in all but $d$ of the coordinates. Following Linero and Yang (2018) and Li, Linero and Murray (2023), we make the following assumptions about the data-generating process (Condition A):

• (A1)$Y_i\sim \mathrm{N}\mathrm{o}\mathrm{r}\mathrm{m}\mathrm{a}\mathrm{l}\left(c^{\ast }+f^{\ast }\left({\mathit{X}}_i\right)+g^{\ast }\left({\mathit{W}}_i\right)+h^{\ast }\left({\mathit{X}}_i,{\mathit{W}}_i\right),1\right)$.
• (A2)The range of ${\mathit{X}}_i$ and ${\mathit{W}}_i$ are $[0,1{]}^p$ and $[0,1{]}^q$, respectively.
• (A3)$f^{\ast }\in 𝒮\left({\alpha }_x,p,d_x\right),g^{\ast }\in 𝒮\left({\alpha }_w,p,d_w\right)$ and $h^{\ast }\in 𝒮\left({\alpha }_h,p+q,d_h\right)$.

We assume the error variance is ${\sigma }^2=1$ for simplicity, but it is straightforward to incorporate unknown ${\sigma }^2$ as well. Achieving (A2) is straightforward by applying quantile normalization, where $X_{ij}$, the ith individual’s jth covariate, is replaced with its quantile value among all the jth covariate values and the same for the exposures. Further, we must make assumptions about the SoftBART prior distribution $\mathrm{\Pi }$ used for each regression function $f,g$ and $h$. For brevity, we leave these specific details to Appendix A, and we refer to these assumptions on the prior distribution as Condition P. We have simplified the setup by not incorporating targeted smoothing into $h(\mathit{x},\mathit{w})$, though we emphasize that Li, Linero and Murray (2023) show that for certain choices of basis function, it is not difficult to prove analogous results for targeted smoothing models. We remove targeted smoothing from consideration so that we can use the same Condition P for all model components rather than requiring a separate set of conditions for $h(\cdot ,\cdot )$.

Let ${\mathbb{E}}_0$ denote the expectation with respect to the true data-generating mechanism, ${\mathrm{\Pi }}_n$ denote the posterior distribution, ${\Vert \mu \Vert }_n^2=\frac{1}{n}{\sum }_{i=1}^n\mu {\left({\mathit{X}}_i,{\mathit{W}}_i\right)}^2$, where $\mu \left({\mathit{X}}_i,{\mathit{W}}_i\right)=c+f\left({\mathit{X}}_i\right)+g\left({\mathit{W}}_i\right)+h\left({\mathit{X}}_i,{\mathit{W}}_i\right)$, and ${\mu }^{\ast }$ denote the ground truth of $\mu$. In Appendix A, we prove the following theorem.

Theorem 1.

If Conditions A and P hold,

$${\mathbb{E}}_0\left\{{\mathrm{\Pi }}_n\left({‖\mu -{\mu }^{\ast }‖}_n>M{ϵ}_n\right)\right\}\to 0\mathit{as\; n}\to \mathrm{\infty }$$

for ${ϵ}_n=\text{max}\left({ϵ}_{nf},{ϵ}_{ng},{ϵ}_{nh}\right)$ where ${ϵ}_{nf}=n^{-{\alpha }_x/\left(2{\alpha }_x+d_x\right)}\text{log}{(n)}^{t_f},{ϵ}_{ng}=n^{-{\alpha }_w/\left(2{\alpha }_w+d_w\right)}\text{log}{(n)}^{t_g}$ and ${ϵ}_{nh}=n^{-{\alpha }_h/\left(2{\alpha }_h+d_h\right)}\text{log}{(n)}^{t_h}$, where $t_f={\alpha }_x\left(d_x+1\right)/\left(2{\alpha }_x+d_x\right)$, $t_g={\alpha }_w\left(d_w+\right.1)/\left(2{\alpha }_w+d_w\right)$ and $t_h={\alpha }_h\left(d_h+1\right)/\left(2{\alpha }_h+d_h\right)$.

The rates ${ϵ}_{nf},{ϵ}_{ng},{ϵ}_{nh}$ represent the minimax estimation rates within the respective function classes for ( $f,g,h$ ), up to a logarithmic term, and, therefore, the rate ${ϵ}_n$ is a near-minimax optimal rate for estimation of ${\mu }^{\ast }$. This result shows that our model is able to capture any true outcome regression function as long as it is sufficiently smooth, which is reasonable for the study of environmental pollutants. Further note that, while we have derived posterior contraction rates for $\mu$, these imply the same posterior contraction rates for causal effects, such as ${\tau }_{{\mathit{w}}_0}(\mathit{x},\mathit{w})$, and for variable importance metrics ${\psi }_j$.

5. Simulation studies.

Here we assess the performance of our proposed approach and compare it to existing approaches that do not allow for heterogeneity of the multivariate treatment effect. For the simulation studies, we draw covariates and exposures from the following multivariate normal distributions:

$$\begin{gathered} {\mathit{X}}_i\stackrel{\text{i.i.d.}}{\sim }N\left(0,{\mathbf{I}}_5\right),{\mathit{W}}_i\sim N\left({\mathit{\mu }}_i,\mathbf{\Sigma }\right),\text{where} \\ {\mathit{\mu }}_i={\left(e^{X_{i1}}/\left(1+e^{X_{i1}}\right)-0.5,0.1X_{i2}^2-0.1,0.3X_{i3},\mathrm{s}\mathrm{i}\mathrm{n}\left(X_{i2}\right),0.05X_{i4}^3\right)}^T, \\ {\mathrm{\Sigma }}_{ij}=\left\{\begin{array}{ll}1,& i=j,\\ 0.3,& i\ne j\end{array}\right. \end{gathered}$$

and ${\mathrm{\Sigma }}_{ij}$ denotes the $(i,j)$ element of $\mathbf{\Sigma }$. The covariates are associated with the exposures, and the exposures are correlated with each other, which are both common in observational studies in environmental health. Next, we generate the outcome by

$$\begin{gathered} Y_i=f^{\ast }\left({\mathit{X}}_i\right)+g^{\ast }\left({\mathit{W}}_i\right)+\sum _{j=1}^5{h}_j^{\ast }\left(X_{ij},{\mathit{W}}_i\right)+{ϵ}_i,{ϵ}_i\stackrel{\text{i.i.d.}}{\sim }N(0,1) \\ {f}^{\ast }\left({\mathit{X}}_i\right)=X_{i1}+X_{i2}-0.5X_{i3} \\ {g}^{\ast }\left({\mathit{W}}_i\right)=\mathrm{I}\left\{W_{i1}>0\right\}+W_{i1}{e}^{0.3W_{i3}}+\mathrm{a}\mathrm{r}\mathrm{c}\mathrm{t}\mathrm{a}\mathrm{n}\left(W_{i2}\right)+\mathrm{s}\mathrm{i}\mathrm{n}\left(W_{i2}{W}_{i3}\pi \right)+\mathrm{m}\mathrm{i}\mathrm{n}\left(\left|W_{i3}\right|,1\right). \end{gathered}$$

Further, we consider three scenarios of interactions between covariates and exposures:

$$\mathbf{No\; interaction:}{h}_1^{\ast }\left(X_{i1},\mathit{W}\right)=h_2^{\ast }\left(X_{i2},\mathit{W}\right)=0;$$

$$\begin{gathered} \mathbf{Moderate:}{h}_1^{\ast }\left(X_{i1},\mathit{W}\right)=0.2\text{arctan}\left(4X_{i1}\right)g^{\ast }(\mathit{W}), \\ {h}_2^{\ast }\left(X_{i2},\mathit{W}\right)=0.2\text{cos}\left(X_{i2}\pi \right)g^{\ast }(\mathit{W}); \end{gathered}$$

$$\begin{gathered} \mathbf{Strong:}{h}_1^{\ast }\left(X_{i1},\mathit{W}\right)=0.4\text{arctan}\left(4X_{i1}\right)g^{\ast }(\mathit{W}), \\ {h}_2^{\ast }\left(X_{i2},\mathit{W}\right)=0.4\text{cos}\left(X_{i2}\pi \right)g^{\ast }(\mathit{W}), \end{gathered}$$

while $h_j^{\ast }\left(X_{ij},\mathit{W}\right)=0$ for $j=3,4,5$ under all scenarios. The standard deviations of the heterogeneous treatment effect functions are approximately 0.2 and 0.4 times that of the sum of the two main effects, respectively, which is reasonable, as we expect that interaction effects to be smaller in magnitude than main effects in general. We run the simulation for 300 different data replicates for each scenario and average results across all simulated data sets. We set the sample size to be $n=2000$. We aim to estimate ${\tau }_{{\mathit{w}}_1,{\mathit{w}}_0}(\mathit{X})\equiv \mathbb{E}\left\{Y\left({\mathit{w}}_1\right)-Y\left({\mathit{w}}_0\right)\mid \mathit{X}\right\}$ at 100 randomly chosen locations from the distribution of $\mathit{X}$ as well as the variable importance metrics ${\psi }_j$ for all covariates. We choose two fixed exposures levels ${\mathit{w}}_0={(-0.5,-0.5,-0.5,-0.5,-0.5)}^T$ and ${\mathit{w}}_1={(-0.5,-0.5,-0.5,-0.5,-0.5)}^T$, which roughly correspond to the first and the third quartiles of each exposure, respectively. We consider four different approaches in the simulation. The first is the Bayesian kernel machine regression approach (BKMR) commonly used in the environmental mixtures literature, which fits a model of the form $E(Y\mid \mathit{X},\mathit{W})=\mathit{X}\mathit{\beta }+m(\mathit{W})$. A Gaussian process prior is placed on $m(\cdot )$, which incorporates spike-and-slab prior distributions to remove unnecessary exposures. The second is a standard BART prior that fits a model of the form $E(Y\mid \mathit{X},\mathit{W})=m(\mathit{X},\mathit{W})$ and places a BART prior on $m(\cdot )$ (BART). The third uses a smooth variant of BART for $m(\cdot )$ (SoftBART). Lastly, we use our proposed approach with the blocking scheme (SepBART), which separates the regression surface into separate functions and then uses either SoftBART or tsBART prior distributions for each function separately. Note that BKMR specifically assumes that the effect of covariates is linear, which is the case in our simulation studies, while it does not allow for treatment effect heterogeneity, which makes it misspecified in scenarios with interactions between exposures and covariates.

Figure 1 shows the simulation results for estimating ${\tau }_{{\mathit{w}}_1,{\mathit{w}}_0}(\mathit{X})$. We find that our method is the only method that achieves the nominal $95\mathrm{\%}$ coverage for all interaction scenarios and produces the smallest root mean squared error (RMSE) when there is treatment effect heterogeneity, regardless of its strength. While it is expected that BKMR does not perform well in the presence of interactions between covariates and exposures, as it assumes no treatment effect heterogeneity, it is surprising that SepBART outperforms BKMR, even when there is no heterogeneity and the main effect of covariates is linear, which is the situation that BKMR is designed for. The no-interaction scenario is the least favorable setting for SepBART, as it forces an interaction structure on the model. However, even for this setting, SepBART shows comparable estimation performance to the best performing model in terms of RMSE, which indicates that our model is able to shrink the $h(\cdot )$ functions to zero when they are not required. It is also notable that SepBART improves on the original BART and SoftBART that do not separate the regression function into distinct components, which suggests that we benefit from the proposed model formulation and separation of the effects into different functions by providing balance between flexibility and efficiency.

Fig. 1.

Simulation results for estimating CATE when ${\mathit{w}}_0$ and ${\mathit{w}}_1$ are given. The dashed line in the second row denotes nominal 95% coverage.

We now assess the performance for estimating MTE-VIM introduced in Section 2. As BKMR does not allow for effect heterogeneity, we focus on our approach only here. We calculate ${\hat{\psi }}_j$ for $j=1,2,\dots ,5$ for each data set by taking the posterior mean of estimates and plot the distribution of ${\hat{\psi }}_j$ in Figure 2. The true values of ${\psi }_1$ and ${\psi }_2$ are 0.72 and 0.28, respectively, and the others are zero. We see that, for both interaction scenarios where the true total heterogeneity $\varphi$ is 0.26 and 1.14, respectively (moderate and strong), the posterior mean of each estimate is concentrated near the true value. Along with estimating MTE-VIM, we can conduct a hypothesis test for the difference between two MTE-VIMs, that is, testing $H_0:{\psi }_j={\psi }_k$. This can be done by constructing the posterior distribution of ${\psi }_j-{\psi }_k$ and seeing whether or not the interval contains zero. We calculate how often the null hypothesis is rejected with level $\alpha =0.05$ over 300 data set replicates, and plot the empirical rejection rate of each testing pair in Figure 3. Since the true ${\psi }_1$ is far away from zero, it shows high rejection rates for comparing ${\psi }_1$ and the others (the first column of Figure 3). For ${\psi }_2$, which is nonzero but not large as ${\psi }_1$, it produces slightly weaker power when interactions are moderate but detects almost all differences when interactions are strong. For the last three columns where the null is true, we see rejection rates are controlled under the desired level $\alpha =0.05$, regardless of the scenario, though inference is somewhat conservative.

Fig. 2.

Violin plots of 300 posterior means of the variable importance measures, ${\psi }_j$ for $j=1,2,\dots ,5$ under each of the two scenarios with heterogeneity. The black dots represent the true variable importance measures: 0.72 for ${\psi }_1,0.28$ for ${\psi }_2$ and 0 for the others.

Fig. 3.

Rejection rates for each test where $H_0:{\psi }_i={\psi }_j$ with the level $\alpha =0.05$.

6. Health effects of air pollution mixtures.

We now use our proposed approach to gain important epidemiological insights about the effect of PM_2.5 components and ozone on mortality. We collect information on all Medicare beneficiaries in the United States above the age of 65 for the years 2000–2016. We observe a number of individual-level covariates for the Medicare beneficiaries, such as their age, race, sex and whether they have dual eligibility to Medicaid, which is a proxy for low socioeconomic status. We also observe a number of area-level covariates unique to each zip code, sourced from the United States Census Bureau and the CDC’s Behavioral Risk Factor Surveillance System, regardless of individuals’ age. These area-level covariates consist of average body mass index, smoking rates, median household income, education, population density and percent owner occupied housing. We also adjust for both temperature and humidity variables that are available from the National Climatic Data Center.

We obtain zip code level environmental exposure data from two distinct sources. We obtain estimates of total PM_2.5, ammonium, nitrates and sulfate levels on a (0.01° × 0.01°) monthly grid from the Atmospheric Composition Analysis Group (Van Donkelaar et al. (2019)). We obtain estimates of elemental carbon, organic carbon and ozone on a 1 km by 1 km daily grid from the Socioeconomic Data and Applications Center (Di et al. (2019, 2021), Requia et al. (2021)). We do not have exact residential addresses of individuals in Medicare and only know their residential zip code, and, therefore, all exposures are aggregated to the yearly level at each zip-code. All individual and geographic-level covariates are also aggregated up to the zip-code level by taking their averages or proportions within each zip code so that the outcome of interest is the annual mortality rate in each zip code, which we define as 100 times the number of deaths in that zip code for a particular year divided by the number of person-years in that zip code.

We run our proposed approach as described in Section 3, targeting both marginal and heterogeneous treatment effects to evaluate the extent to which ambient air pollution affects mortality and whether this effect varies by characteristics of zip codes. We run our model for each year separately using the prior year exposures as the pollutants of interest and, therefore, will present results for all years between 2000 and 2016. We run our MCMC algorithm for 10,000 iterations, discarding the first 4000 as a burn-in and thinning every twelfth sample. A detailed discussion on the convergence of MCMC sampling can be found in Appendix D.1. Overall, we find that convergence diagnostics are very good for the average treatment effect across all years studied. Convergence is slightly worse for MTE-VIM values, though this is likely caused by multimodality of the posterior distribution, rather than an issue of MCMC sampling, and is still within an acceptable range.

6.1. Marginal effects of exposures.

First, we investigate the average treatment effect of an increase in environmental mixtures, denoted as $\mathbb{E}\left\{Y\left({\mathit{w}}_1\right)-Y\left({\mathit{w}}_0\right)\right\}$, where ${\mathit{w}}_0$ and ${\mathit{w}}_1$ represent vectors of all exposures corresponding to the first and third quartiles for each exposure, respectively. Hence, a positive ATE would indicate a harmful effect on mortality due to an increase in environmental exposures. Figure 4 illustrates estimates of the ATE for each year. We consistently estimate a positive ATE, implying that increasing levels of all environmental exposures increases the mortality rate across the contiguous United States, which suggests a detrimental impact of environmental pollutants on human health.

Fig. 4.

Posterior means and corresponding 95% credible intervals for the average treatment effects on mortality rates when every pollutant simultaneously increases from its yearly first quartile to the third quartile. The dashed line represents no average treatment effect.

As previously mentioned, the positivity assumption is a critical consideration when analyzing multivariate exposures, as highlighted in Antonelli and Zigler (2024). To address this, we assess the positivity assumption and examine a trimmed estimator that is more robust to positivity violations, with detailed analyses presented in Appendix D.4. We evaluated the positivity assumption for each exposure and found that the majority of observations satisfy a univariate positivity criterion for each exposure simultaneously. Additionally, we examined a trimmed ATE, which targets the ATE for observations with the highest probability of being exposed to either ${\mathit{w}}_0$ and ${\mathit{w}}_1$. The results are very similar between the trimmed ATE and the original ATE, which suggests that our choice of ${\mathit{w}}_1$ and ${\mathit{w}}_0$ does not strongly violate the positivity assumption.

6.2. Heterogeneity of the effect of PM_2.5 components.

While the average treatment effect provides insights into the health effects of air pollution, of even more interest is whether this effect varies across the population, as it is crucially important to understand which communities are at most risk to the detrimental effects of air pollution. Toward this goal, we first investigate the proposed MTE-VIM for each covariate to understand which characteristics drive the heterogeneity of the causal effect. The estimated total heterogeneity, $\varphi$, averaged over all of the study years is 0.016, which implies that the standard deviation of ${\tau }_{{\mathit{w}}_0}(\mathit{W},\mathit{X})$ with respect to $\mathit{X}$ given $\mathit{W}=\mathit{w}$ is, on average, 0.13. Considering that we estimated an average ATE of 0.46 across study years in Figure 4, this suggests there is a nonnegligible amount of treatment effect heterogeneity. Figure 5 shows the posterior mean of the MTE-VIM corresponding to each covariate for each study year.

Fig. 5.

Heatmap of posterior means of the MTE-VIM by year.

We observe some degree of variation in the MTE-VIM for certain covariates over the study years. This is likely due to the inherent difficulty in estimating these parameters in situations with highly correlated exposures and covariates. For this reason there will be more variability in the estimates of treatment effect heterogeneity and the corresponding MTE-VIMs compared with the marginal effect estimates seen in Section 6.1, which leads to more variable results across years. Despite this variability, there is a clear trend that both race and dual eligibility to Medicaid modify the treatment effect the most, followed by age. Figure 6 shows the value of yearly variable importance metrics in Figure 5 when averaged over all study years, and we see that race, dual eligibility to Medicaid and age achieved the largest variable importance measures, with overall means of 0.24, 0.22 and 0.1, respectively. This indicates that, on average, 24% of the treatment effect heterogeneity can only be explained by race, even after controlling for other potential effect modifiers. In situations such as this with high variability around estimates of the importance of each covariate on treatment effect heterogeneity, we can also examine grouped MTE-VIMs. In Appendix D.3 we placed covariates into one of three groups and found consistent results across years, showing that socioeconomic factors (which include race and dual eligibility to Medicaid) had the largest grouped MTE-VIM values.

Fig. 6.

Estimates of the MTE-VIM averaged over the years 2000–2016.

In addition to identifying the characteristics of zip codes that modify the treatment effect the most, it is important to understand the nature and direction of this heterogeneity to better understand which groups are most susceptible to air pollution mixtures. To do this, we examine the conditional average treatment effect, using the same reference exposure levels as in the previous section, denoted by ${\mathit{w}}_1$ and ${\mathit{w}}_0$. Specifically, we plot $\mathbb{E}\left\{Y\left({\mathit{w}}_1\right)-Y\left({\mathit{w}}_0\right)\mid \mathit{X}={\tilde{\mathit{x}}}_{\left(j\right)}\right\}$ as a function of $x_j$, where ${\tilde{\mathit{x}}}_{\left(j\right)}=\left({\overline{x}}_1,\dots ,x_j,\dots ,{\overline{x}}_p\right)$. This sets all covariates to their sample mean and only varies covariate $j$ so that we can examine whether the conditional average treatment effect increases or decreases with covariate $j$. Within the framework of our model and our identification restrictions, $\mathbb{E}\left\{Y\left({\mathit{w}}_1\right)-Y\left({\mathit{w}}_0\right)\mid \mathit{X}={\tilde{\mathit{x}}}_{\left(j\right)}\right\}$ simplifies to the difference in the interaction function, $h_j\left({\mathit{w}}_1,x_j\right)-h_j\left({\mathit{w}}_0,x_j\right)$, which is zero when $x_j$ equals the mean covariate level, ${\overline{x}}_j$. Therefore, when this difference is positive, it implies that individuals with the given covariate level are more susceptible to increases in the environmental mixture compared to those with an average covariate level.

Figure 7 shows the conditional average treatment effect curves, described above, for the covariates with large MTE-VIMs. Across most years, we observed that the treatment effects of air pollution tend to decrease as the proportion of White populations in a zip code increases, while harmful effects are more pronounced in areas with smaller White populations. This trend suggests a potentially greater impact of air pollution on racial minorities. We found that the treatment effect increases significantly with higher rates of dual eligibility, indicating that the detrimental effects of air pollution are more severe in areas with lower socioeconomic status. We also find that areas with older individuals are more adversely affected by increases in pollution, though this effect is not as strong as that of dual eligibility to Medicaid. Overall, our study shows that results from the previous literature (Simoni et al. (2015), Bargagli-Stoffi et al. (2020)) that focused solely on univariate PM_2.5 exposure extend to the more complex setting of multivariate air pollution mixtures. We find a harmful effect of air pollution on mortality and that this effect is heterogeneous, with larger effects in areas with low socioeconomic status, lower proportions of white individuals and older individuals.

Fig. 7.

Estimates of conditional average treatment effects on mortality rates of increasing all pollutants from the first quartiles to the third quartiles as a function of each covariate, $\mathbb{E}\left\{Y\left({\mathit{w}}_1\right)-Y\left({\mathit{w}}_0\right)\mid \mathit{X}={\tilde{\mathit{x}}}_{\left(j\right)}\right\}$, where ${\tilde{\mathit{x}}}_{\left(j\right)}=\left({\overline{x}}_1,\dots ,x_j,\dots ,{\overline{x}}_p\right)$. Each thin gray curve represents the posterior mean for one year, and the thick blue curve represents the average of posterior means over the study years. The red dashed line represents no treatment effect heterogeneity.

7. Discussion.

In this paper we proposed a novel approach to analyze and summarize complex, heterogeneous effects of multivariate continuous exposures. We developed new estimands for this setting, including a treatment effect variable importance measure, which is tailored to multivariate exposure scenarios and provides interpretable quantities that simplify the heterogeneity and provide important epidemiological insights about the effects of air pollution mixtures. To facilitate identification of our estimands and to allow different strengths of regularization for each component of the model, we proposed a separation of the outcome model regression function and integrated a targeted smoothing method into our model to obtain smooth estimates of the degree of heterogeneity by each covariate. Our theoretical results, coupled with simulation studies, provide strong empirical and analytical support for the efficiency and validity of the proposed model. We used the model to gain new insights into the causal impact of multivariate air pollution mixtures on mortality rates in the Medicare population. Our model estimates a detrimental impact of air pollution, consistent with previous literature, and shows that this effect is more pronounced in zip codes with lower socioeconomic status, fewer white individuals and older individuals.

While this paper provides strong evidence of heterogeneous treatment effects for environmental mixtures, and provides users with novel methodology for the multivariate, continuous treatment setting, there are certain limitations with our approach. For one, our definition of the MTE-VIM requires the user to choose a reference exposure level ${\mathit{w}}_0$. While we expect most reasonable choices of ${\mathit{w}}_0$ to lead to similar values of the MTE-VIM, as we have seen in Appendix D.2, results could be sensitive to this choice in certain cases. To avoid violations of the positivity assumption and reliance on model extrapolation, we recommend users select ${\mathit{w}}_0$ from a region of the exposure space with sufficient support in the data.

An additional limitation of our approach is that we made the assumption that treatment effect heterogeneity was additive in the covariates, which allowed us to write $h(\mathit{X},\mathit{W})={\sum }_{j=1}^p{h}_j\left(X_j,\mathit{W}\right)$. We find this to be a reasonable assumption, in most cases, and one that facilitates estimation in a difficult setting, but more complex forms of heterogeneity may not be captured well by this additive structure. For cases where treatment effect heterogeneity is expected to be nonadditive in covariates, one could completely remove the assumption by fitting $h(\mathit{X},\mathit{W})$ with all covariates and exposures at once or relax the assumption to allow for multiplicative interactions between two covariates and exposures. However, this would require adjustments to the proposed identification strategy and the interpretation of each functional component, and it is unclear how much this approach would impact estimation efficiency, particularly in cases with many covariates, as in our real data analysis.

Lastly, there are many future research directions to consider. For one, our focus has been on outcomes measured on a continuous scale, but a natural and important extension would be to allow for binary or count outcomes. The extension to binary outcomes with a probit or logistic link is straightforward using latent variables within our MCMC algorithm (Albert and Chib (1993), Polson, Scott and Windle (2013)). While less trivial, BART has recently been extended to count outcomes (Murray (2021)), and future work could combine our proposed methodology within BART models for count outcomes. One other important direction would be to couple these estimation strategies with sensitivity analysis to unmeasured confounding, a common concern in observational studies such as this one. The recent literature has developed advancements in sensitivity analysis in multiple exposure settings for estimating average treatment effects (Zheng, D’Amour and Franks (2022)), and these ideas could be extended to heterogeneous treatment effects seen here. Another interesting direction would be to combine the approach developed in this manuscript with methodology for optimal policy estimation, in order to provide environmental regulators increased information about the best practices for reducing pollution in the future in order to obtain the largest public health benefit from a proposed reduction in air pollution levels.

Supplementary Material

Proofs and additional results (DOI: 10.1214/25-AOAS2060SUPPA; .pdf). Shin et al. (2025) contains the proofs of the posterior contraction rates and identifiability of the proposed model, and additional simulation and analysis results.

R package (DOI: 10.1214/25-AOAS2060SUPPB; .zip). The SepBART R package is available to implement the proposed approach. The most up-to-date version is available at https://github.com/hshin111/SepBART.