Causal Estimation of Exposure Shifts with Neural Networks

The field of causal inference has seen immense progress in the past couple of decades with the development of targeted doubly-robust methods yielding desirable theoretical efficiency guarantees on estimates for various causal effects (Van der Laan et al., 2011; Kennedy, 2016). These advancements have recently been incorporated into the neural network (NN) literature for causal inference via targeted regularization (TR) (Shi et al., 2019; Nie et al., 2021). TR methods produce favorable properties for causal estimation by incorporating a regularization term into a supervised neural network model. However, it remains an open task to develop a NN method that specifically targets the causal effect of a shift in the distribution for a continuous exposure/treatment variable (Muñoz and Van Der Laan, 2012). We call this problem shift-response function (SRF) estimation. Many scientific questions can be formulated as an SRF estimation task (Muñoz and Van Der Laan, 2012). Some notable examples in the literature include estimating the health effects from shifts to the distribution of environmental, socioeconomic, and behavioral variables (e.g., air pollution, income, exercise habits) (Muñoz and Van Der Laan, 2012; Díaz and Hejazi, 2020; Smith et al., 2023).

One of our objectives is to develop a neural network technique that addresses a timely and highly prominent regulatory question. Specifically, the EPA is currently considering whether or not to revise the National Ambient Air Quality Standards (NAAQS), potentially lowering the acceptable annual average $\text{PM}_{2.5}$ concentration from 12 to 11, 10, or 9 $\mathrm{\upmu g/m^{3}}$. We anticipate that the revision to the NAAQS will ultimately result in a shift to the distribution of $\text{PM}_{2.5}$ concentrations. Our goal is to estimate, for the first time, the reduction in deaths that would result from this anticipated shift using causal methods for SRFs.

Contributions. We develop a novel method, called Targeted Regularization for Exposure Shifts with Neural Networks (tresnet), which introduces two necessary and generalizable methodological innovations to the TR literature. First, we use a TR loss targeting SRFs, ensuring that our estimates retain the properties expected from TR methods such as asymptotic efficiency and double robustness (Kennedy, 2016). Given standard regularity conditions, these properties guarantee that the SRF is consistently estimated when either the outcome model or a density-ratio model for the exposure shift is correctly specified, achieving the best possible efficiency rate when both models are consistently estimated. Second, tresnet accommodates non-continuous outcomes belonging to the exponential family of distributions (such as mortality counts) that frequently arise in real-world scenarios, including our motivating application. In addition to its suitability for our application, we evaluate the performance of tresnet in a simulation study tailored for SRF estimation, demonstrating improvements over neural network methods not specifically designed for SRFs.

Our results contribute to the public debate informing the US Environmental Protection Agency (EPA) on the effects of modifying air quality standards. A preview of the results (fully developed in Section 5) is presented in Figure 1. The figure presents the estimated reduction in deaths (%) resulting from various shifts to the distribution of $\text{PM}_{2.5}$ across every ZIP-code in the contiguous US between 2000 and 2016. These shifts limit the maximum concentration of $\text{PM}_{2.5}$ to the cutoff value for every ZIP-code that exceeds the cutoff, and otherwise leaves ZIP-codes that do not exceed the cutoff unchanged. We vary the cutoff in this SRF between 6 $\mathrm{\upmu g/m^{3}}$ and 16 $\mathrm{\upmu g/m^{3}}$ (x-axis). The y-axis represents the % reduction in deaths that corresponds with each cutoff threshold. Notably, a NAAQS threshold of 9 $\mathrm{\upmu g/m^{3}}$, would have had the effect of decreasing elder mortality by 4%. These findings present a data-driven perspective on the potential health benefits of the EPA’s proposal.

We let $(A,Y,{\bm{X}})$ denote a unit from the target population, where $A\in{\mathcal{A}}$ is a continuous exposure variable (also known as the treatment), $Y\in{\mathcal{Y}}$ is the outcome of interest, and ${\bm{X}}\in{\mathcal{X}}$ are covariates. We assume a sample of size $n$ and denote it as $\{{\bm{X}}_{i},A_{i},Y_{i}\}_{i=1}^{n}$. For instance, in the application of interest, which is described in detail in Section 5, $i$ represents a zip code location in a given year, $A_{i}$ represents its annual average concentration of $\text{PM}_{2.5}$ ($\mu g/m^{3}$), $Y_{i}$ is the number of deaths, and ${\bm{X}}_{i}$ includes demographic and socioeconomic variables. For simplicity, we will omit the subscript $i$ unless required for clarity or when describing statistical estimators from samples.

We assume the following general non-parametric structural causal model (SCM)(Pearl, 2009):

where $U_{Y},U_{A},U_{\bm{X}}$ are exogenous variables. The functions $f_{Y},f_{A},f_{\bm{X}}$ are deterministic and unknown. We will avoid measure-theoretic formulations to keep the presentation simple.

The quantities $Y^{a}=f_{Y}({\bm{X}},a,U_{Y})$ are known as potential outcomes. Potential outcomes can be divided into two parts. The factual outcome corresponds with the observed outcome when $a=A$. The counterfactual outcomes are hypothetical results at any point $a\in{\mathcal{A}}$ where $a\neq A$. We will typically refer to $Y^{a}$ as the counterfactual, as SRFs are tooled for predicting the potential outcome from unobserved exposures (i.e. the counterfactuals), unless specifically stated otherwise.

There are two functions that are pivotal in defining and estimating causal effects. First, the outcome function represents the expected (counter)factual outcome conditional on covariates, expressed as

Unless otherwise specified, all expectations are taken from the data distribution in Equation 1. The second key function, the (generalized) propensity score, is the distribution of the exposure conditional on the covariates, denoted as $p(a|{\bm{x}}):=p(A=a|{\bm{X}}={\bm{x}})$. Note that $p$ is also used generically to represent any density function with the given arguments; the propensity score being a special case.

As we have suggested earlier, an estimator of the SRF can be derived from estimators of the outcome and density ratio functions. Using TR, we will obtain an estimator $\hat{{\bm{\psi}}}^{\text{tr}}$ that “converges efficiently” to the true ${\bm{\psi}}$ at a rate in accordance with the prevailing semiparametric efficiency theory surrounding robust causal effect estimation (Kennedy, 2022).

We conducted simulation experiments to validate the design choices of tresnet. The so-called fundamental problem of causal inference is that the counterfactual responses are never observed in real data. Thus, we must rely on widely used semi-synthetic datasets to evaluate the validity of our proposed estimators. First, we consider two datasets introduced by Nie et al. (2021), which are continuous-treatment adaptions to the popular datasets ihdp (Hill, 2011) and news (Newman, 2008). We also apply our methods to (Nie et al., 2021)’s fully simulated data, sim-N. In addition, we consider the fully simulated dataset described in (Bahadori et al., 2022), which features a continuous treatment and has been previously used for calibrating models in air pollution studies. Finally, we also consider three variants of the tcga dataset presented in Bica et al. (2020). The three variants consist of three different dose specifications as treatment assignments and the corresponding dose-response as the outcome.

The datasets described here have been employed without substantial modifications from the original source studies to facilitate fair comparison. Note that for each of these synthetic and semi-synthetic datasets, we have access to the true counterfactual outcomes, which allows us to compute sample SRFs exactly. We will consider the estimation task of 20 equally-spaced percent reduction shifts between 0-50% from the current observed exposures. More specifically, $\tilde{A}=(1-c)A$ for values of $c$ in $0-50\%$.

Evaluation metric and task. Given a semi-synthetic dataset ${\mathcal{D}}$, consider an algorithm that produces an estimator $\hat{{\bm{\psi}}}^{(s)}_{\tilde{p},{\mathcal{D}}}$ of ${{\bm{\psi}}}^{(s)}_{\tilde{p},{\mathcal{D}}}$ given an exposure shift $\tilde{p}$ and random seed $s$. To evaluate the quality of the estimator, we use the mean integrated squared error $\textsc{mise}_{\mathcal{D}}=(n_{\text{seeds}}|\tilde{{\mathcal{P}}}|)^{-1}\sum_{s}\sum_{\tilde{p}}\lvert\hat{{\bm{\psi}}}^{(s)}_{\tilde{p},{\mathcal{D}}}-{\bm{\psi}}^{(s)}_{\tilde{p},{\mathcal{D}}}\rvert^{2}.$ This metric is a natural adaption of an analogous metric commonly used in dose-response curve estimation (Bica et al., 2020).

Experiment 1: Does targeting SRFs specifically improve their estimation (compared to no targeting or standard ERF targeting?) We evaluate two variants of tresnet against alternative SRF estimators, including vcnet (Nie et al., 2021) and drnet (Schwab et al., 2020), two prominent methods used in causal TR estimation for continuous treatments. Since these methods were not designed specifically for SRFs, we expect that their performance for this task can be improved by adapting TR through their baseline architectures.

The first variant of tresnet uses varying-coefficient layers based on splines—see the discussion in Section 3.2 for background. We compare this variant, named tresnet${}_{\textsc{vc}}$, with the following baselines:

1. (1)

   $\textsc{aipw}_{\textsc{vc}}$, the AIPW-type estimator for SRFs discussed in Section 3.1, wherein we fit separate outcome and density ratio models which are then substituted into Equation 5;

2. (2)

   $\textsc{outcome}_{\textsc{vc}}$, which uses the same outcome model as tresnet${}_{\textsc{vc}}$, but without the TR and density ratio heads;

3. (3)

   vcnet, which uses a similar overall architecture as tresnet${}_{\textsc{vc}}$, but with a TR designed for ERFs rather than for SRFs.

The second variant of tresnet uses varying coefficients based on piecewise linear functions instead of splines. This variant, tresnet${}_{\textsc{pl}}$, is compared against the following analogous baselines:

1. (1)

   $\textsc{aipw}_{\textsc{pl}}$, which is the piecewise analogue to $\textsc{aipw}_{\textsc{vc}}$;

2. (2)

   drnet, based on Schwab et al. (2020), functioning as the piecewise linear analogue to $\textsc{outcome}_{\textsc{vc}}$;

3. (3)

   $\textsc{drnet}+\textsc{tr}_{\textsc{erf}}$, which is the analogue to vcnet that can be constructed by adding a density ratio head and TR designed for ERFs rather than SRFs.

Table 1(a) shows the results of this experiment. For both architectures, the tresnet variants achieve the best performance. tresnet${}_{\textsc{vc}}$ is somewhat better than vcnet, which have comparable architectures although tresnet uses a different TR implementation specific for SRF estimation rather than ERF estimation. Likewise, tresnet${}_{\textsc{pl}}$ outperforms $\textsc{drnet}+\textsc{tr}_{\textsc{erf}}$. These moderate but consistent performance improvements suggest the importance of SRF-specific forms of TR. We also see strong advantages against outcome-based predictions and AIPW estimators, suggesting that the TR loss and the shared learning architecture is a boon to performance. These results are compatible with observations from previous work in the TR literature (Nie et al., 2021; Shi et al., 2019).

Experiment 2: Does TR improve estimation when count-valued outcomes are observed? For these experiments, we used the spline-based variant and evaluate whether tresnet with the Poisson-specific TR, explained in Section 3, performs better than the mean-squared error (MSE) loss variant when the true data follows a Poisson distribution. This evaluation is important since our application consists of count data requiring a Poisson model which are widely used to investigate the effects of $\text{PM}_{2.5}$ on health (Wu et al., 2020; Josey et al., 2022). We construct similar semi-synthetic datasets as in Experiment 1, however in this experiment the outcome-generating mechanism samples from a Poisson distribution rather than a Gaussian distribution. The results of this experiment are clear–using the correct exponential family for the outcome model is critical, regardless of whether TR is implemented.

We implemented tresnet for count data to estimate the health benefits caused by shifts to the distribution of $\text{PM}_{2.5}$ that would result from lowering the NAAQS—the regulatory threshold for the annual-average concentration of $\text{PM}_{2.5}$ enforced by the EPA.

Data. The dataset is comprised of Medicare recipients²²2Access to Medicare data is restricted without authorization by the Centers for Medicare & Medicaid Services since it contains sensitive health information. The authors have completed the required IRB and ethical training to handle these datasets. from 2000–2016, involving 68 million unique individuals. The data includes measurements on participant race/ethnicity, sex, age, Medicaid eligibility, and date of death, which are subsequently aggregated to the annual ZIP-code level. The $\text{PM}_{2.5}$ exposure measurements are extracted from a previously fit ensemble prediction model (Di et al., 2019). The confounders include measurements on meteorological information, demographics, and the socioeconomic status of each ZIP-code. Calendar year and census region indicators are also included to account for spatiotemporal trends. To compile our dataset, we replicated the steps and variables outlined by Wu et al. (2020).

Exposure shifts. We consider two types of $\text{PM}_{2.5}$ shifts, cutoff shits and percent reduction shift, each providing different perspectives and insights. The counterfactuals implied from these scenarios are illustrated in figures 2(a) and 2(b), respectively. First, a cutoff shift, parameterized by a threshold $d$, encapsulates scenarios in which every ZIP-code year that exceeded some threshold are truncated to that maximum threshold. Mathematically, the shift is defined by the transformation $\tilde{A}=\min(A,c)$. To be more succinct, the exposure shift defines a counterfactual scenario. For this application, the threshold $c$ is evaluated at equally spaced points starting with 15 $\mathrm{\upmu g/m^{3}}$ moving down to 6 $\mathrm{\upmu g/m^{3}}$. We expect that at 15 $\mathrm{\upmu g/m^{3}}$ there will be little to no reduction in deaths since $>99\%$ of observations fall below that range. We can contemplate the proposed NAAQS levels through $d$ assuming that full compliance to the new regulation holds for incompliant ZIP-codes. The exposure shift should otherwise not affect already compliant ZIP-codes. Second, we onsider percent reduction shifts. This scenario assumes that all ZIP-code years reduce their pollution levels proportionally from their observed value. More precisely, the shift is defined as $\tilde{A}=A(1-c)$. We considered a range of percent reduction shifts between $c\in(0,50)\%$. We can interpret these shifts in terms of the NAAQS by mapping each percent reduction to a compliance percentile. For instance, Figure 4 shows that, under a 30% overall reduction in historical values, approximately 82% would comply with a NAAQS of 9 $\mathrm{\upmu g/m^{3}}$.

Implementation. We implement tresnet using varying-coefficient splines as in Section 4. We select a NN architecture using the out-of-sample prediction error from a 20/80% test-train split to choose the number of hidden layers (1-3 layers) and hidden dimensions (16, 64, 256). We found no evidence of overfitting in the selected models. To account for uncertainty in our estimations, we train the model on 100 bootstrap samples, each with random initializations, thereby obtaining an approximate posterior distribution. Deep learning ensembles have been previously shown to approximate estimation uncertainty well in deep learning tasks (Izmailov et al., 2021).

Results. Figure 1 in the introduction presents the effects of shifting the $\text{PM}_{2.5}$ distribution at various cutoffs on the expected reduction in deaths. The slope is steeper at stricter/lower cutoffs, likely because lower cutoffs affect a larger fraction of the observed population and reduce the overall $\text{PM}_{2.5}$. For instance, figure Figure 1 shows that had no ZIP-code years exceeded 12 $\mathrm{\upmu g/m^{3}}$, the observed death counts would have decreased by around 1%. If the cutoff is lowered to 9 $\mathrm{\upmu g/m^{3}}$, then deaths could have fallen by around 4%. The slope becomes increasingly steeper as the $\text{PM}_{2.5}$ threshold is reduced, suggesting an increasing benefit from lowering the standard concentration level. Another way to interpret this result is to say that there is a greater gain to reducing mortality caused by $\text{PM}_{2.5}$ from lowering the concentration level from 10 to 8 $\mathrm{\upmu g/m^{3}}$ than there is from lowering it from 12 to 10 $\mathrm{\upmu g/m^{3}}$, despite the obvious observation that both contrasts examine a $\text{PM}_{2.5}$ reduction of 2 $\mathrm{\upmu g/m^{3}}$.