Treatment Effect Heterogeneity and Importance Measures for Multivariate Continuous Treatments

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 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 et al., 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 et al., 2010) and Gaussian processes (Banerjee et al., 2013) do not immediately apply. 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 et al., 2020; Semenova and Chernozhukov, 2021; Fan et al., 2022; Shin and Antonelli, 2023) or with multiple levels (Chang and Roy, 2024). In the binary setting, heterogeneity is typically summarized by the conditional average treatment effect (CATE) function, $\operatorname{\mathbb{E}}\{Y(1)-Y(0)|\bm{X}=\bm{x}\}$ where $Y(w)$ denotes the potential outcome under a binary exposure level $w$. While the CATE is a functional estimand, a number of summaries of this function have been proposed to simplify heterogeneity further. One 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 et al. (2022) and Li et al. (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 both define estimands for continuous, multivariate exposures and for estimating 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 that 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 towards 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 et al., 2022) for interactions between exposures and covariates to allow for heterogeneity without imposing strong parametric assumptions.

Throughout, we assume that we observe $\bm{\mathcal{D}}=\{(Y_{i},\bm{X}_{i},\bm{W}_{i})\}_{i=1}^{n}$ for each individual where $Y_{i}\in\mathbb{R},\bm{X}_{i}\in\mathbb{R}^{p}$, and $\bm{W}_{i}\in\mathbb{R}^{q}$. Also, we adopt the potential outcomes framework: $Y_{i}(\bm{w})$ denotes the potential outcome that would be observed under exposure $\bm{w}$. In order to denote potential outcomes in this manner, and to identify causal effects from the observed data, we make the following assumptions:

1. (i)

   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}(\bm{W}_{i})$.

2. (ii)

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

3. (iii)

   Unconfoundedness: $Y(\bm{w})\mathchoice{\mathrel{\hbox to0.0pt{$\displaystyle\perp$\hss}\mkern 2.0mu{\displaystyle\perp}}}{\mathrel{\hbox to0.0pt{$\textstyle\perp$\hss}\mkern 2.0mu{\textstyle\perp}}}{\mathrel{\hbox to0.0pt{$\scriptstyle\perp$\hss}\mkern 2.0mu{\scriptstyle\perp}}}{\mathrel{\hbox to0.0pt{$\scriptscriptstyle\perp$\hss}\mkern 2.0mu{\scriptscriptstyle\perp}}}\bm{W}|\bm{X}$ for all $\bm{w}$.

The positivity assumption guarantees that each unit has a non-zero probability to be exposed to each treatment level for all possible values of pre-treatment variables, at least in large samples. 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.

Under the assumptions introduced in Section 2, for any fixed $\bm{w}$ we can identify $\operatorname{\mathbb{E}}\{Y(\bm{w})|\bm{X}=\bm{x}\}$ from the observed data by $\operatorname{\mathbb{E}}(Y|\bm{X}=\bm{x},\bm{W}=\bm{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 $\bm{X}$, the main effect of exposures $\bm{W}$, and the interactions of covariates and exposures on the outcome. Therefore, we write our model as follows:

$\displaystyle\operatorname{\mathbb{E}}(Y|\bm{X},\bm{W})=c+f(\bm{X})+g(\bm{W})+h(\bm{X},\bm{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.

The first simplification we make is to set $h(\bm{X},\bm{W})=\sum_{j=1}^{p}h_{j}(X_{j},\bm{W})$ where $X_{j}$ is the $j$-th component of $\bm{X}$. This restricts interactions between covariates and exposures to be limited to interactions between a single covariate and the multivariate exposures. We feel this is a mild assumption that is reasonable in practice, and will improve estimation efficiency in our model. The first hurdle to identification of the individual functions is that $h_{j}(X_{j},\bm{W})$ may capture the main effect of either the covariate or exposures. For example, we could have $h_{j}(X_{j},\bm{W})=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}(X_{j},\bm{W})$ to be of the form $h^{cov}_{j}(X_{j})h^{exp}_{j}(\bm{W})$ where $h^{cov}_{j}$ and $h^{exp}_{j}$ are arbitrary functions restricted to not be constant functions. This allows us to prevent the interaction function $h_{j}(X_{j},\bm{W})$ from simply absorbing the main effects of $X_{j}$ and $\bm{W}$. We describe our strategy for each of these separable functions in the following sections.

Next, we study the rate at which the posterior distributions concentrate in the model described in Section 3. Let $\mathcal{S}(\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 et al. (2022), we make the following assumptions about the data-generating process (Condition A).

1. (A1)

   $Y_{i}\sim\operatorname{Normal}(c^{*}+f^{*}(\bm{X}_{i})+g^{*}(\bm{W}_{i})+h^{*}(\bm{X}_{i},\bm{W}_{i}),1)$

2. (A2)

   The range of $\bm{X}_{i}$ and $\bm{W}_{i}$ are $[0,1]^{p}$ and $[0,1]^{q}$, respectively.

3. (A3)

   $f^{*}\in\mathcal{S}(\alpha_{x},p,d_{x})$, $g^{*}\in\mathcal{S}(\alpha_{w},p,d_{w})$, and $h^{*}\in\mathcal{S}(\alpha_{h},p+q,d_{h})$

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 $i$-th individual’s $j$-th covariate, is replaced with its quantile value among all the $j$-th covariate values, and the same for the exposures. Further, we must make assumptions about the SoftBART prior distribution $\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(\bm{x},\bm{w})$, though we emphasize that Li et al. (2022) 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 $\operatorname{\mathbb{E}}_{0}$ denote the expectation with respect to the true data-generating mechanism, $\Pi_{n}$ denote the posterior distribution, and $\|\mu\|_{n}^{2}=\dfrac{1}{n}\sum_{i=1}^{n}\mu(\bm{X}_{i},\bm{W}_{i})^{2}$ where $\mu(\bm{X}_{i},\bm{W}_{i})=c+f(\bm{X}_{i})+g(\bm{W}_{i})+h(\bm{X}_{i},\bm{W}_{i})$. In Appendix A, we prove the following theorem.

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:

	$\displaystyle\bm{X}_{i}\overset{\text{i.i.d.}}{\sim}N(\bm{0},\textbf{I}_{5}),\quad\bm{W}_{i}\sim N(\vec{\mu}_{i},\vec{\Sigma})\qquad\text{where}$
	$\displaystyle\vec{\mu}_{i}=\left(e^{X_{i1}}/(1+e^{X_{i1}})-0.5,0.1X_{i2}^{2}-0.1,0.3X_{i3},\sin(X_{i2}),0.05X_{i4}^{3}\right)^{T},$
	$\displaystyle\Sigma_{ij}=\begin{cases}1\qquad&i=j\\ 0.3\qquad&i\neq j\end{cases}$

and $\Sigma_{ij}$ denotes the $(i,j)$ element of $\vec{\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

	$\displaystyle Y_{i}=f^{*}(\bm{X}_{i})+g^{*}(\bm{W}_{i})+\sum_{j=1}^{5}h^{*}_{j}(X_{ij},\bm{W}_{i})+\epsilon_{i},\qquad\epsilon_{i}\overset{\text{i.i.d.}}{\sim}N(0,1)$
	$\displaystyle f^{*}(\bm{X}_{i})=X_{i1}+X_{i2}-0.5X_{i3}$
	$\displaystyle g^{*}(\bm{W}_{i})=\text{I}{\left\{W_{i1}>0\right\}}+W_{i1}e^{0.3W_{i3}}+\arctan(W_{i2})+\sin(W_{i2}W_{i3}\pi)+\min(|W_{i3}|,1).$

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

	No interaction:	$\displaystyle h^{*}_{1}(X_{i1},\bm{W})=h^{*}_{2}(X_{i2},\bm{W})=0$
	Moderate:	$\displaystyle h^{*}_{1}(X_{i1},\bm{W})=0.2\arctan(4X_{i1})g^{*}(\bm{W}),\quad h^{*}_{2}(X_{i2},\bm{W})=0.2\cos(X_{i2}\pi)g^{*}(\bm{W})$
	Strong:	$\displaystyle h^{*}_{1}(X_{i1},\bm{W})=0.4\arctan(4X_{i1})g^{*}(\bm{W}),\quad h^{*}_{2}(X_{i2},\bm{W})=0.4\cos(X_{i2}\pi)g^{*}(\bm{W})$

while $h^{*}_{j}(X_{ij},\bm{W})=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_{\bm{w}_{1},\bm{w}_{0}}(\bm{X})\equiv\operatorname{\mathbb{E}}\{Y(\bm{w}_{1})-Y(\bm{w}_{0})|\bm{X}\}$ at 100 randomly chosen locations from the distribution of $\bm{X}$, as well as the variable importance metrics $\psi_{j}$ for all covariates. We choose two fixed exposures levels $\bm{w}_{0}=(-0.5,-0.5,-0.5,-0.5,-0.5)^{T}$ and $\bm{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|\bm{X},\bm{W})=\bm{X}\bm{\beta}+m(\bm{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|\bm{X},\bm{W})=m(\bm{X},\bm{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 (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_{\bm{w}_{1},\bm{w}_{0}}(\bm{X})$. We find that our method is the only method that achieves the nominal 95% 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.

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,...,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 $\phi$ 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, i.e. 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 non-zero, 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.

We now use our proposed approach to gain important epidemiological insights about the effect of PM_2.5 components and ozone on mortality. We have demographic, socioeconomic, and mortality 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 these individuals, 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 the zip-code of each individual from the United States Census Bureau and the Center for Disease Control’s Behavioral Risk Factor Surveillance System. These area-level covariates consist of average body mass index, smoking rates, median household income, median house value, education, population density, and percent owner occupied housing. We also adjust for both temperature and precipitation 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^{\circ}\times 0.01^{\circ}$) 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 1km by 1km 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 is defined as 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 5000 iterations, discarding the first 1000 as a burn-in, and thinning every fourth sample.

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 and older individuals.

While this paper provides strong evidence of heterogeneous treatment effects for environmental mixtures, and provides users with a 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 $\bm{w}_{0}$. While we expect most reasonable choices of $\bm{w}_{0}$ to lead to similar values of the MTE-VIM, results could be sensitive to this choice. We recommend users utilize a choice of $\bm{w}_{0}$ that does not violate the positivity assumption, which is a crucial consideration when analyzing multivariate exposures as noted in Antonelli and Zigler (2024). An additional limitation of our approach is that we made the assumption that the treatment effect heterogeneity was additive in the covariates, which allowed us to write $h(\bm{X},\bm{W})=\sum_{j=1}^{p}h_{j}(X_{j},\bm{W})$. 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.

Lastly, there are many future research directions to consider. One 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 et al., 2021), and these ideas could be extended to heterogeneous treatment effects seen here. Another important 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.