Causal Estimation of Exposure Shifts with Neural Networks
A fundamental task in causal inference is estimating the effect of
distribution shift in the treatment variable. We refer to this problem as
shift-response function (SRF) estimation. Existing neural network methods for
causal inference lack theoretical guarantees and practical implementations for
SRF estimation. In this paper, we introduce Targeted Regularization for
Exposure Shifts with Neural Networks (TRESNET), a method to estimate SRFs with
robustness and efficiency guarantees. Our contributions are twofold. First, we
propose a targeted regularization loss for neural networks with theoretical
properties that ensure double robustness and asymptotic efficiency specific to
SRF estimation. Second, we extend targeted regularization to support loss
functions from the exponential family to accommodate non-continuous outcome
distributions (e.g., discrete counts). We conduct benchmark experiments
demonstrating TRESNET's broad applicability and competitiveness. We then apply
our method to a key policy question in public health to estimate the causal
effect of revising the US National Ambient Air Quality Standards (NAAQS) for PM
2.5 from 12 ${\mu}g/m^3$ to 9 ${\mu}g/m^3$. This change has been recently
proposed by the US Environmental Protection Agency (EPA). Our goal is to
estimate the reduction in deaths that would result from this anticipated
revision using data consisting of 68 million individuals across the U.S.
Mauricio Tec, Kevin Josey, Oladimeji Mudele, Francesca Dominici
# Causal Shift-Response Functions with Neural Networks:

###### Abstract

Policymakers are required to evaluate the health benefits of reducing the National Ambient Air Quality Standards (NAAQS; i.e., the safety standards) for fine particulate matter (PM\({}_{2.5}\)) before implementing new policies. We formulate this objective as a _shift-response function_ (SRF) and develop methods to analyze the problem using methods for causal inference, specifically under the _stochastic interventions_ framework. SRFs model the average change in an outcome of interest resulting from a hypothetical shift in the observed exposure distribution. We propose a new broadly applicable doubly-robust method to learn SRFs using targeted regularization with neural networks. We evaluate our proposed method under various benchmarks specific for marginal estimates as a function of continuous exposure. Finally, we implement our estimator in the motivating application that considers the potential reduction in deaths from lowering the NAAQS from the current level of \(12~{}\upmu\mathrm{g}/\mathrm{m}^{3}\) to levels that are recently proposed by the Environmental Protection Agency in the US (10, 9, and 8 \(\upmu\mathrm{g}/\mathrm{m}^{3}\)).

Machine Learning, ICML, ICML

## 1 Introduction

On January 6, 2023, the US Environmental Protection Agency (EPA) asked for public comment on a proposal to lower the National Ambient Air Quality Standard (NAAQS) for annual PM\({}_{2.5}\) concentrations from 12 \(\upmu\mathrm{g}/\mathrm{m}^{3}\) to 9 or 10 \(\upmu\mathrm{g}/\mathrm{m}^{3}\), considering lowering to as low as 8 \(\upmu\mathrm{g}/\mathrm{m}^{3}\), or as high as 11 \(\upmu\mathrm{g}/\mathrm{m}^{3}\). The EPA estimated that lower standards could result in as much as \(\$43\) billion in net health benefits in 2032 (The New York Times, 2023). The final decision on the new NAAQS will be made later this year. The EPA's inquiry into lowering PM\({}_{2.5}\) concentrations remains contentious when considering all possible economic, social, and health consequences. We will focus on the potential health benefits gained from decreasing PM\({}_{2.5}\) concentrations and lowering the enforced limits. We formulate our analysis by introducing a general-purpose framework for neural networks that learns a _shift-response function_ (SRF). An SRF characterizes the causal effect for an outcome of interest (e.g. mortality) that can be attributed to a shift in the distribution of an exposure variable (e.g. a reduction of PM\({}_{2.5}\)). SRFs can be formally defined using _stochastic interventions_ a causal inference framework that is useful for evaluating changes to a continuous exposure like PM\({}_{2.5}\)(Munoz & Van Der Laan, 2012).

To study the effects of lowering the NAAQS in the US, this paper proposes **T**argeted **R**egularization for **E**xposure **S**hifts with **Neural **N**etworks (tresnet), a novel method for estimating SRFs. tresnet adapts previous work on targeted regularization for estimating discrete and average-dose response curves with neural networks, incorporating ideas from targeted maximum likelihood estimation. tresnet uses a neural network with distinct prediction heads for the outcome and an importance-sampled density ratio of the observed and shifted exposure distributions while satisfying an estimating equation that ensures various desirable properties. A key property of our estimator is double-robustness. A doubly-robust estimator of the SRF yields unbiased estimates when either the outcome model or a density ratio

Figure 1: Estimated reduction to early deaths among 68.5 million elderly adults enrolled in Medicare in the US for different NAAQS thresholds using tresnet. The curve shows the effect of implementing a cutoff exposure shift to reduce all observations above the NAAQS to the proposed upper limit.
of the observed and shifted exposures are correctly specified. In addition, the doubly-robust SRF estimator achieves the semiparametric efficiency bound when both models are correct (Munoz and Van Der Laan, 2012; Kennedy, 2016).

Currently, there is limited evidence that can directly inform the new NAAQS rule, even when numerous studies report adverse health outcomes attributable to high levels of PM\({}_{2.5}\)(Wu et al., 2020). Only a few peer reviewed articles have taken a causal perspective into non-linear effect estimation for changing PM\({}_{2.5}\) concentrations. Among these reports, none directly analyze policy interventions by supposing the changes manifest as an exposure shift. Instead, most analyses into the associations between PM\({}_{2.5}\) and adverse health outcomes focus on estimating exposure-response functions (ERF). While ERFs are useful for a variety of purposes (e.g. identifying potential thresholds), analyzing exposure shifts more concisely evaluates the potential costs or benefits for proposed interventions, as this paper will highlight concerning the EPA's desire to implement new ambient air standards. While trenet is motivated by this application, it can generalize to many settings with different types of exposure shifts, particularly within the domains of environmental health and epidemiology. A preview of the results is shown in Figure 1, showing a steep decrease in deaths with a lower standard. Details are provided in Section 7.

ContributionsThis paper contributes to causal machine-learning methods and the public health domain, summarized as follows: (1) We use state-of-the-art neural network methods to take an important step toward answering an open public health issue in the US. Our results clearly demonstrate the health benefits of lowering the acceptable air quality standards. Further, they show evidence of increasing marginal benefits as the threshold is further reduced. (2) To our knowledge, we provide the first causal inference method using neural networks for continuous stochastic interventions or count data. The proposed methods are highly competitive and often outperform others in our simulation studies with synthetic/semi-synthetic datasets.

## 2 Background on Stochastic Interventions and Estimands of Interest

NotationThe exposure1 and outcome of interest are denoted as \(A_{i}\in\mathcal{A}\) and \(Y_{i}\in\mathcal{Y}\), respectively. We suppose that \(A_{i}\) is continuous. Each unit \(i\in\mathcal{I}\) is described by a set of covariates \(\mathbf{X}_{i}\in\mathcal{X}\). We denote \(\mathbb{P}^{X}\) as the marginal distribution of \(\mathbf{X}_{i}\). The observed exposure density function is denoted as \(\pi_{0}(a\mid\mathbf{x})\triangleq\Pr(A_{i}=a\mid\mathbf{X}_{i}=\mathbf{x})\). This quantity is widely known as the _generalized propensity score_, which has broad utility for causal inference (Hirano and Imbens, 2005; Imai and van Dyk, 2004). Throughout we assume a sample \((\mathbf{X}_{i},A_{i},Y_{i})_{i=1}^{n}\) of size \(n\).

Footnote 1: The terms exposure, intervention, and treatment are used interchangeably throughout this manuscript.

### Potential outcomes and assumptions

We adopt the potential outcomes framework, also known as the Rubin causal model, to describe the assumptions necessary to identify the target SRF parameter. Let \(Y_{i}^{a}\) denote the _potential outcome_ of unit \(i\) under a hypothetical exposure \(a\in\mathcal{A}\)(Imbens and Rubin, 2015). Potential outcomes require the stable-unit treatment value assumption (SUTVA), having two conditions. The first condition is _consistency_, which entails \(Y_{i}^{A_{i}}=Y_{i}\), meaning that the potential outcome of the factual treatment corresponds with the observed outcome. The second condition is _no interference_, which states that \(Y^{A_{i}}\) does not depend on the treatment of other units.

**Assumption 2.1** (Sutva).: (a) \(Y_{i}^{A_{i}}=Y_{i}\) (consistency); (b) \(Y_{i}^{\mathbf{A}}=Y_{i}^{A_{i}}\) where \(\mathbf{A}=(A_{i})_{i=1}^{n}\) is the vector of all treatment assignments (no interference).

SUTVA is widely assumed in most causal inference applications, particularly in previously proposed methods for exploring the effects of air pollution on human health (Dominici and Zigler, 2017). We also assume that, conditionally to all the measured confounders included in \(\mathbf{X}_{i}\), there are no unmeasured confounders (variables correlated with both the treatment and outcome); or more formally:

**Assumption 2.2** (unconfoundedness).: \(\mathbf{X}_{i}\) consists of pre-treatment covariates only (it is not caused by the treatment or outcome) and \(A_{i}\perp(Y_{i}^{a})_{a\in\mathcal{A}}\mid\mathbf{X}_{i}\), where \(\perp\) denotes conditional independence.

### Stochastic interventions and estimand

The _stochastic intervention_ (SI) framework (Hubbard and Van der Laan, 2008) encompasses methods for estimating an expected potential outcome value when replacing the observed exposure distribution, \(\pi_{0}\), with a counterfactual exposure described by \(\pi_{d}\). The intervention is stochastic since \(\pi_{d}\) is a distribution rather than a constant. In this work, we assume that \(\pi_{d}\) is defined implicitly by a transformation of \(A_{i}\), denoted by \(A_{i}^{d}\). We let \(\mathbf{O}_{i}^{d}=(\mathbf{X}_{i},A_{i},A_{i}^{d},Y_{i})\) denote the augmented data and \(\mathbb{P}_{d}\) denote its distribution. Table 1 shows common examples of relevant exposure shifts, although a specific exposure shift should ideally be directed by the application itself and by input from domain experts. For instance, one way to approach our motivating application is to have \(d\) represent the new NAAQS threshold (e.g., 12, 10, 9) using a cutoff SI. Alternatively, we could use an expert-defined hypothetical exposure distribution that uses ZIP code level information not captured by \(\mathbf{X}_{i}\). We identify various possible unit-level transforms in the table. In general, such \(\pi_{d}\) cannot be computed as a simple closed-form expression in terms \(\pi_{0}\).
The causal estimand of interest corresponding with an SRF can be defined as follows. First, define the expected potential outcome as:

\[\begin{split} q(a,\mathbf{x})&\triangleq\mathbb{E}[Y_{i}^{a }\mid\mathbf{X}_{i}=\mathbf{x},A_{i}=a]\\ \mu_{d}&\triangleq\mathbb{E}_{\mathbb{P}^{d}}\left[q( \mathbf{X}_{i},A_{i}^{d})\right].\end{split} \tag{1}\]

The _shift-response function_ (SRF) follows the mapping \(\pi_{d}\mapsto\mu(\pi_{d})\). It is useful to contrast SRFs with ERFs. An ERF represents the expected potential outcome at the same fixed treatment level averaged all units. An ERF can be seen as a limiting case of an SRF where \(\pi_{d}\) is a point mass distribution (though this estimator is no longer stochastic).

Figure 2 illustrates the above concepts in an hypothetical air pollution example. Figure 2a shows the cutoff SI (inspired by our application), and its corresponding SRF is plotted in Figure 2b. The SRF is highly informative of the effect of the cutoff. It shows that reducing the cutoff from 12 to 8 provides a sharp decline in deaths, but flattens out after 6 \(\upmu\mathrm{g}/\mathrm{m}^{3}\). For additional demonstration, Figure 2c contains an alternative example of SI for the same data by reducing every observed value by 10%. Lastly, Figure 2d shows a comparison with the exposure-response function (ERF). Since every point in the ERF implies the same hypothetical exposure for all units, it is less informative of the causal effect of a shift, particularly under non-linear effects. The estimand associated with the ERF has the form \(\xi(a)\triangleq\mathbb{E}_{\mathbf{X}_{i}\sim\mathbb{P}^{X}}[q(\mathbf{X}_{i},a)]\), yielding a curve \(a\mapsto\xi(a)\).

Observe that the expectation defining \(q\) in Equation (1) requires the true causal model of the conditional distribution of the potential outcomes \(Y_{i}^{a}\). However, this distribution is unknown and cannot be observed from the data in general. However, the following well-known result establishes conditions of its identifiability from the observed data (Imbens and Rubin, 2015):

**Proposition 2.3**.: _Suppose Assumptions 2.1 and 2.2 hold. Then \(q(a,\mathbf{x})=\mathbb{E}[Y_{i}\mid\mathbf{X}_{i}=\mathbf{x},A_{i}=a]\). The right-hand side expression does not involve the potential outcomes._

Importance sampling/inverse weightingDefine the density ratio \(w_{d}(a|\mathbf{x}):=\pi_{d}(a|\mathbf{x})/\pi_{0}(a|\mathbf{x})\), which is a measure of the domain shift. Given an estimator of \(\hat{w}_{d}\), the simplest kind of estimator of \(\mu_{d}\) is based on importance sampling (Munoz and Van Der Laan, 2012). It is the generalization of the well-known inverse probability weight (IPW) estimator:

\[\hat{\mu}_{d}^{\text{pw}}=n^{-1}{\sum}_{i=1}^{n}\hat{w}_{d}(A_{i}|\mathbf{X}_{i}) Y_{i}. \tag{2}\]

When \(w_{d}(A_{i}|\mathbf{X}_{i})\) is correctly estimated, then \(\hat{\mu}_{d}^{\text{pw}}\) is unbiased. However, this estimator is known to suffer from high variance (Munoz and Van Der Laan, 2012). This motivates the discussion on double robustness in the next section.

### Doubly-robust inference of SIs

Doubly-robust estimation requires estimating two underlying quantities - the importance sampling ratio and the regression function (Munoz and Van Der Laan, 2012). A doubly-robust stochastic intervention effect estimator is consistent if either model of the respective nuisance parameter is consistently estimated (Kennedy, 2016). This double-robust property also achieves local efficiency when both models are consistently estimated, empowering thoughtful machine-learning methods. We formulate this section for the case of a single \(d\). Indeed, in many applications of interest, including the one motivating this paper, only a few finite values of \(d\) are relevant (for example, relevant NAAQS are 8, 9, 10, 11, 12). Results for a single value will generally extend to a small finite number of values (Nie et al., 2021).

Let \(\hat{w}_{d}\), \(\hat{q}\), and \(\hat{\mu}_{d}\) be estimators of \(w_{d}\), \(q\), and \(\mu_{d}\), respec

\begin{table}
\begin{tabular}{l|l|l} \hline shift & \(A_{i}^{d}\) & \(\pi_{d}\) \\ \hline Cutoff & \(\min(A_{i},d)\) & \(\Pr(A_{i}<d)\pi_{0}(a|\mathbf{x})+\Pr(A_{i}\geq d)\delta_{d}(a)\) \\ \hline Additive & \(A_{i}-d\) & \(\pi_{0}(a+d|\mathbf{x})\) \\ \hline Percent & \(A_{i}(1-d)\) & \((1-d)\pi_{0}(a/(1-d)|\mathbf{x})\) \\ \hline Bijection & \(d(A_{i},\mathbf{x})\) & \(\pi_{0}(d^{-1}(a,\mathbf{x})|\mathbf{x})\frac{|\hat{w}_{d}|}{\hat{w}_{d}}d(a,\mathbf{x})|\) \\ \hline Unit-level & \(d(A_{i},\mathbf{x},i)\) & closed-form not available in general \\ \hline \end{tabular}
\end{table}
Table 1: Common examples of exposure shifts. \(\delta_{d}(a)\) denotes a point mass distribution at \(a=d\).

Figure 2: **(a)** Example stochastic intervention (SI) setting an upper cutoff at \(d=9\). **(b)** is the resulting shift-response function (SRF) for varying cutoffs \(d\). **(c)** Another example of SI where the shift is a 10% reduction. **(d)** The exposure-response function (ERF) for comparison with panel (a); a point in the curve implies the same hypothetical exposure for all units, being less informative of the causal effect of a shift.
tively. In short, doubly robust estimation for an SRF is achieved by specifying the estimating equation

\[\mathbb{E}_{\mathbb{P}^{d}}[\Psi_{\pi_{d}}(\mathbf{O}_{i}^{d},\hat{w}_{d},\hat{q}_{d},\hat{\mu}_{d})]=0. \tag{3}\]

Here, \(\Psi_{\pi_{d}}\) is the _efficient-influence function_ (EIF) of \(\mu_{d}\)(Kennedy, 2016). For additional background, Hines et al. (2022) provides an excellent introduction to semi-parametric inference and the relevance of EIFs in causal analyses. The following two results are taken from Munoz and Van Der Laan (2012) and characterize the estimating equation for stochastic interventions. We refer to this article for proofs of the statements in this section.

**Proposition 2.4** (Eif).: _The efficient-influence function (EIF) of \(\mu_{d}\) is_

\[\begin{split}&\Psi_{\pi_{d}}(\mathbf{O}^{d},w_{d},q,\mu_{d})\triangleq\\ & w_{d}(A_{i}|\mathbf{X}_{i})(Y_{i}-q(A_{i},\mathbf{X}_{i}))+q(A_{i}^{ \delta},\mathbf{X}_{i})-\mu_{d}.\end{split} \tag{4}\]

Solving the estimating equation yields an unbiased estimator of \(\mu_{d}\). To see this, notice that if \(\mathbb{E}_{\mathbb{P}^{d}}[\Psi_{\pi_{d}}(\mathbf{O}_{i}^{d},\hat{w}_{d},\hat{q} _{d},\mu_{d})]=0\) holds (\(\mu_{d}\) here is the true parameter), then we can rearrange the terms so that

\[\mu_{d}=\mathbb{E}[\hat{q}(A_{i}^{\delta},\mathbf{X}_{i})]+\mathbb{E}[\hat{w}_{d} (A_{i}|\mathbf{X}_{i})(Y_{i}-\hat{q}(A_{i},\mathbf{X}_{i}))]. \tag{5}\]

The following result shows that the estimating equation is solved when either the density ratio or outcome model is correctly specified and that the contraction rate under finite sample sizes is controlled by the composite contraction rate for these two estimators.

**Proposition 2.5** (Doubly-robust estimation).: _Suppose Assumptions 2.1 and 2.2 hold, and that \(w_{d}(a|\mathbf{x})<\infty\) for all \(a\in\mathcal{A}\) and \(\mathbf{x}\in\mathcal{X}\). Then \(\mathbb{E}_{\mathbb{P}^{d}}[\Psi_{\pi_{d}}(\mathbf{O}_{i}^{d},\hat{w}_{d},\hat{q}, \mu_{d})]=0\) if either \(\hat{q}=q\) or \(\hat{w}_{d}=w_{d}\)._

Augmented IPWThe most straightforward way to construct a doubly-robust estimator so that the estimating equation holds is by plugging in the unknown quantities in Equation (5) with corresponding estimates. This estimator is known as _augmented inverse probability weighting_ (AIPW Munoz and Van Der Laan (2012):

\[\hat{\mu}_{d}^{\text{AIPW}}=n^{-1}{\sum_{i=1}^{n}}\hat{w}_{d}(A_{i}|\mathbf{X}_{i })(Y_{i}-\hat{q}(\mathbf{X}_{i},A_{i}))+\hat{q}(A_{i}^{d},\mathbf{X}_{i}). \tag{6}\]

It has been argued by Shi et al. (2019); Nie et al. (2021) that AIPW-type estimators may be unstable in small finite samples despite being asymptotically efficient due to the dependence on the bias term. This motivates the introduction of targeted regularization.

## 3 Tresnet: Targeted Regularization for Exposure Shifts with Neural Networks

OverviewWe now explain the key idea behind tresnet. We provide a summary and details in the next section. Equation (5) shows that the estimating equation can be interpreted as a bias correction in the plugin estimator due to the domain shift induced by the new exposure distribution. Targeted regularization was proposed to mitigate the unreliability of the bias term in the context of estimating the average treatment effect of a binary intervention (Shi et al., 2019; Nie et al., 2021). The strategy behind targeted regularization consists of learning an outcome model highly predictive of the observed data, then learning a small perturbation so that the bias term vanishes. In this way, the asymptotic and double robustness properties are preserved while achieving strong performance in finite samples (Shi et al., 2019).

### Targeted regularization

Targeted regularization (TR) begins with two models, \(q_{\theta}\) and \(w_{d,\theta}\), parameterized by neural networks and introduces a perturbation parameter \(\varepsilon_{d}\) to be optimized along with these models. The TR loss is given by

\[\begin{split}\mathcal{L}_{\text{TR}}(q_{\theta},w_{d,\theta}, \varepsilon_{d})&\triangleq\frac{1}{2n}\sum_{i=1}^{n}w_{d,\theta }\left(Y_{i}-\tilde{q}_{d,\theta}(\mathbf{X}_{i},A_{i})\right)^{2}\\ \tilde{q}_{d,\theta}(\mathbf{X}_{i},A_{i})&\triangleq q _{\theta}(\mathbf{X}_{i},A_{i})+\varepsilon_{d}\end{split} \tag{7}\]

The key observation is that \(\frac{\partial}{\partial\varepsilon_{d}}\mathcal{L}_{\text{tr}}(q_{\theta},w_{ d,\theta},\varepsilon_{d})=0\) if and only if the bias term vanishes, i.e., \(\frac{1}{n}\sum_{i=1}^{n}w_{d,\theta}(Y_{i}-\tilde{q}_{d,\theta},i)=0\). Upon successful optimization, we can use the perturbed estimator

\[\tilde{\mu}_{d,\theta}=\frac{1}{n}\sum_{i=1}^{n}\tilde{q}_{d,\theta},i, \tag{8}\]

which satisfies the finite-sample estimating equation \(\frac{1}{n}\sum_{i=1}^{n}\Psi_{\mathbb{P}^{d}}(\mathbf{O}_{i}^{d},w_{d,\theta}, \tilde{q}_{\theta,d},\tilde{\mu}_{d})=0\).

When learning an SRF for a finite set of values \(\mathcal{D}\), we can use the integrated loss. \(\mathcal{L}_{\text{TR}}^{\mathcal{D}}(q_{\theta},w_{d,\theta},\varepsilon_{ \mathcal{D}})=\frac{1}{|\mathcal{D}|}\sum_{d\in\mathcal{D}}\mathcal{L}_{\text{ TR}}^{d}(q_{\theta},w_{d,\theta},\varepsilon_{d})\) with \(\varepsilon_{\mathcal{D}}=(\varepsilon_{d})(d\in\mathcal{D})\). When \(\mathcal{D}\) is infinite or very large, additional analytical considerations will be required in general (Nie et al., 2021). Our focus here is a small finite set of values as required in our application of interest and many realistic settings.

The targeted loss is learned along with the empirical risk losses of \(q_{\theta}\) and \(w_{d,\theta}\):

\[\begin{split}\mathcal{L}(q_{\theta},w_{d,\theta},\varepsilon_{d})= \\ \mathcal{L}_{Y}(q_{\theta})+\mathcal{L}_{A}(w_{d,\theta})+\beta \mathcal{L}_{\text{TR}}(q_{\theta},w_{d,\theta}.\varepsilon_{d})\end{split} \tag{9}\]

In what follows, we explain the empirical risk losses in more detail. Recall that the idea is to calibrate \(q_{\theta}\) to the observed data so that \(\varepsilon_{d}\) is only a small perturbation. For generic applications, we use the quadratic loss
\(\mathcal{L}_{Y}(q_{\theta})=\frac{1}{2n}\sum_{i=1}^{n}(Y_{i}-q_{\theta}(\mathbf{X}_{i},A_{i}))^{2}\). However, we will use a more appropriate loss for count data in our motivating application. Since specifying \(\mathcal{L}_{A}(w_{\theta,d})\) will require considering different cases, we explain in the next section.

### Density ratio learning

We discuss three strategies for density ratio learning. The first one, and the most obvious one, is to leverage that the shifted exposure can be written as a transformation \(\pi_{d}=F(\pi_{0})\) (see Table 1). Accordingly, we can let \(\pi_{\theta}\) be parameterized by a neural network mean to approximate \(\pi_{0}\) and use the average negative loglikelihood loss \(\mathcal{L}_{W}^{\text{\sc cps}}(\pi_{\theta})=-\frac{1}{n}\sum_{i=1}^{n}\log( \pi_{\theta}(A_{i}|\mathbf{X}_{i}))\), and then set \(w_{d,\theta}=F(\pi_{\theta})/\pi_{\theta}\). Here gps stands for the generalized propensity score.

However, it has been noted that direct density ratio estimation may be easier than learning a density (Sugiyama et al., 2012). Furthermore, there may be cases when \(A_{i}^{d}\) is defined in such a way that \(F\) is not available. To develop a method applicable to these cases, we depart from the following known result.

**Lemma 3.1**.: _Let \(Z_{i}\sim_{\text{\sc iid}}\operatorname{Bernoulli}(1/2)\), \(A_{i}^{1}\sim_{\text{\sc iid}}\pi_{1}\) and \(A_{i}^{2}\sim_{\text{\sc iid}}\pi_{2}\), and define the auxiliary variable \(\bar{A}_{i}\triangleq Z_{i}A_{i}^{1}+(1-Z_{i})A_{i}^{2}\). Then \(\Pr(Z_{i}=1|\bar{A}_{i})=\sigma(\log(\pi_{1}(\bar{A}_{i})/\pi_{2}(\bar{A}_{i}))\). where \(\sigma(l)=(1+e^{-l})\) is the sigmoid function._

See Sugiyama et al. (2012) for the proof of this fact. According to the lemma, we can learn the log-density ratio by learning a classifier between samples \(A_{i}^{d}\) and \(A_{i}\). For convenience, denote \(p_{d,\theta}(a|\mathbf{x})=\sigma(\log(w_{d,\theta}(a|\mathbf{x})))\). Then we can use the classification loss

\[\mathcal{L}_{A}(w_{d,\theta})=\] \[\quad-\frac{1}{2n}\sum_{i=1}^{n}\left(\log(p_{d,\theta}(A_{i}^{d} |\mathbf{X}_{i}))+\log(1-p_{d,\theta}(A_{i}|\mathbf{X}_{i})\right).\]

Even when a transformation of \(\pi_{0}\) is known. It may be possible to entirely learn the density ratio with the above classification loss and the parameterization \(w_{d,\theta}=F(\pi_{\theta})/\pi_{\theta}\). We call this the hybrid-classifier (hybrid-class) strategy. We found it to be very competitive with the loglikelihood of loss in our experiments. We suspect it is less prone to overfitting, particularly when combined with label smoothing. Finally, the full-classifier (full-class) case is when we have a direct model for \(w_{d,\theta}\) because it can't be expressed in terms of \(\pi_{0}\). In the next section we specify a suitable architecture for this case.

### Neural Network Architecture

When \(\pi_{d}\) can be written in terms of \(\pi_{0}\), our choice of neural network architecture is identical to vcnet(Nie et al., 2021). However, we remark that the loss functions used to train the network differ from the ones used in vcnet, since the latter concerns the estimation of ERFs. We will provide a brief overview of this architecture for completeness, and explain the changes necessary for the case when \(\pi_{d}\) cannot be written in terms of \(\pi_{0}\). Nie et al. (2021) also provides a visual representation of the architecture.

**vcnet backbone** The initial part of the network learns a representation of the covariates \(\mathbf{Z}_{i}=f_{\theta}(\mathbf{X}_{i})\) using a feed-forward network. Then, the outcome model takes the representation \(\mathbf{Z}_{i}\) as an input and passes it through a varying-coefficient feed-forward network, where the weights and biases are functions of the exposure. This effect is achieved by parameterizing the weights and biases as learnable linear combinations of a spline basis.

For the generalized propensity score, Nie et al. (2021) propose to discretize the exposure domain in \(k\) subintervals and use a single linear layer taking \(\mathbf{Z}_{i}\) as input and having \(k\) outputs, where the \(k\)-th output represents the conditional probability of the \(k\)-th subinterval. They focus on the case when \(\mathcal{A}=(0,1)\), but any bounded interval could be considered. To reduce the effect of discretization, they propose to use interpolation, although we refer to the original vcnet paper for details. As the authors remark, alternatives like Mixture density networks (Bishop, 1994), or normalizing flows (Rezende and Mohamed, 2015) could be considered.

**Ratio model** Motivated by this architecture, we propose the following strategy for the case when \(w_{d,\theta}\) needs to be estimated directly. First, if only one value of \(d\) is of interest, we can use a feed-forward network with \(k\) outputs, each giving the value of \(\log w_{d,\theta}\) in a subinterval of \(\mathcal{A}\). Instead, when multiple values of \(d\) are considered, we can use a varying-coefficient feed-forward network parameterized by \(d\). This will ensure the ability to learn for various values of \(d\) without significantly increasing the number of parameters.

**Perturbation parameters** We parameterize each \(\epsilon_{d}\) as an independent free parameter, in contrast to vcnet, which uses a spline parameterization. We found that independent parameters work better when considering only a small number of possible values for \(d\) as in our applications.

## 4 Targeted Regularization For Count Data

The targeted loss needs to be modified to be useful with count data. Count data are very common within public health and epidemiology studies, for instance, to represent the number of deaths, infections, or hospitalizations. A common modeling choice for regression with count data outcome is to assume Poisson distribution of the count variable. For causal inference, this would imply that \(Y_{i}^{a}\sim\operatorname{Poisson}(q(\mathbf{X}_{i},a))\). Its most important property is that the error is proportional to the mean,
\(\sqrt{\operatorname{Var}[Y_{i}^{a}]}=q(\mathbf{X}_{i},a)\). This property is often referred to as the _multiplicative errors_ property. When the true data has this property, regressing with the quadratic loss (which assumes homoscedastic errors) will result in a misspecified model that will be very sensitive to large values of \(Y_{i}\). In this section, we demonstrate how to adapt our proposed targeted regularization to deal with multiplicative errors, focusing on the Poisson model.

The first step is to adapt the outcome model for Poisson regression. For this purpose, we replace the mean-squared error (Gaussian) loss with the Poisson (negative log-likelihood) loss. Up to a constant, it can be written as

\[\mathcal{L}_{Y}^{\text{Poi}}(q_{\theta}) \triangleq\frac{1}{n}\sum_{i=1}\left(q_{\theta}(\mathbf{X}_{i},A_{i} )-Y_{i}\log(q_{\theta}(\mathbf{X}_{i},A_{i}))\right) \tag{10}\] \[q_{\theta}(\mathbf{X}_{i},A_{i}) \triangleq e^{f_{\theta}(\mathbf{X}_{i},A_{i})}.\]

Here \(f_{\theta}\) is the neural network model, and we are using the canonical exp/log-link function. We now express the targeted loss for Poisson regression. The two key differences are the use of multiplicative perturbation and the re-weighted negative loglikelihood loss:

\[\mathcal{L}_{\text{TR}}^{\text{Poi}}(q_{\theta},w_{d,\theta}, \varepsilon_{d})\triangleq \tag{11}\] \[\qquad\frac{1}{n}\sum_{i=1}w_{d,\theta}\left(q_{\theta}(\mathbf{X}_{ i},A_{i})-Y_{i}\log(\tilde{q}_{d,\theta}(\mathbf{X}_{i},A_{i}))\right)\] \[\tilde{q}_{d,\theta}(\mathbf{X}_{i},A_{i})\triangleq e^{\varepsilon_{ d}}q_{\theta}(\mathbf{X}_{i},A_{i}).\]

It holds that \(\frac{\partial}{\partial\varepsilon_{d}}\mathcal{L}_{\text{TR}}^{\text{Poi}} (q_{\theta},w_{d,\theta},\varepsilon_{d})=0\) if and only if the bias term in the estimating equation vanishes, i.e., \(\frac{1}{n}\sum_{i=1}^{n}w_{d,\theta}(Y_{i}-\tilde{q}_{d,\theta,i})=0\).

Finally, we remark that count data regression often requires an _offset_ term \(N_{i}\). For example, if unit \(i\) represents a geographic area with a population (the offset) \(N_{i}\), then by defining \(q_{\theta}(\mathbf{X}_{i},A_{i})\triangleq N_{i}e^{f_{\theta}(\mathbf{X}_{i},A_{i})}\) we can interpret the network output as the (logarithmic) unit rate of occurrence. In our application, the offset is the population enrolled in Medicare for a given county for a given year.

## 5 Related Work

Neural networks for estimating causal effects with continuous treatmentstresnet contributes to the growing neural network literature on causal inference with continuous treatments. This literature is mainly divided among works aiming to estimate individualized effects (e.g., (Bica et al., 2020; Yoon et al., 2018)) and those estimating ERFs, also known as average dose-response functions (ADRFs). tresenet is more closely related to the latter group. In particular, we implement tresenet on top of the architecture proposed by Nie et al. (2021) and Schwab et al. (2020).

Doubly-robust estimation for stochastic interventionsEIFs have been widely studied for constructing asymptotically efficient estimators (Bickel et al., 1993; Kennedy, 2016; Van der Laan et al., 2011). See Hines et al. (2022) for an accessible review. Targeted regularization is a special case. It was originally proposed for estimation with binary treatments by Shi et al. (2019) and extended to ERF estimation by Nie et al. (2021). We have adopted this framework for the case of SIs with continuous exposure. Notably, Duong et al. (2021) use network methods for SIs for binary treatments. To our knowledge, no previous work has explored the estimation of SIs with neural networks. Our derivation of doubly robust inference for SIs closely follows Munoz and Van Der Laan (2012).

## 6 Synthetic Data Experiments

We conducted a series of experiments to evaluate the performance of tresenet2. We use semi-synthetic datasets since evaluating causal inference requires access to the true counterfactuals, which are rarely available in real data. The semi-synthetic datasets are ihdf Hill (2011) and news Newman (2008) which are widely used as causal inference benchmarks. The continuous treatment-adapted versions of these datasets which we use have been proposed by Nie et al. (2021). We also apply the same fully synthetic dataset (sim) proposed in that study. We explain the main experiment here and refer to the Appendix for additional ones. These experiments include an ablation on the density ratio strategies, showing that all of them are fairly comparable (surprising to us); and a variant for the Poisson regression, showing it outperforms the Gaussian case for count data. For the tresenet experiments presented here, we use the hybrid-class loss introduced in Section 3.2.

MetricsFor each dataset, we use the counterfactuals to generate the true SRF corresponding to a percent reduction exposure shift. We use 20 uniformly spaced reductions \(\{d_{k}\}_{k=1}^{20}\subset(0,0.5)\). We evaluate the performance using the average bias bias and root-mean-squared error mse along the SRF. Letting \(\hat{\mu}_{d}\) be the estimate of \(\mu_{d}\), these quantities are defined as bias\(\triangleq 20^{-1}{\sum_{k=1}^{20}}(\hat{\mu}_{d_{k}}-\mu_{d_{k}})\) and \(\text{rMSE}\triangleq(20^{-1}{\sum_{k=1}^{20}}(\hat{\mu}_{d_{k}}-\mu_{d_{k}} )^{2})^{1/2}\). We perform multiple simulations from 100 random seeds and obtain the average and standard deviation of these metrics. Each simulation generates a different train-test split, and we only report the metric values on the test set. We provide additional details in the appendix.

BaselinesWe evaluate ipw and aipw (Equations (2) and (6)). For the latter, we use the same architecture but without targeted regularization. We call this baseline ensnet-aipw. In addition, we compare with the outcome models
(no regularization) of vcnet Nie et al. (2021) and drnet Schwab et al. (2020). We compare these baselines against tresnet, implemented with the hybrid density ratio estimator and the tresnet backbone for the outcome. The hyper-parameter \(\beta\) was fixed at \(0.1\) without any tuning. Two variants of tresnet are also considered, one in which the estimator is fed as an input to Equation (6) (tresnet-aipw), and one in which we replace the outcome model backbone based on vcnet with drnet.

**Results**  Table 2 shows the results. We observe that tresnet consistently outperforms ipw, aipw and vcnet in all settings. Contrary to our expectations, vanilla drnet was surprisingly strong on the hldp dataset. Combining it with tresnet improved the bias but affected the rmse. Overall, we suspect that these benchmarks are fairly easy for the outcome model and validation in more complex scenarios is required.

## 7 The Health Benefits of Stricter Air Pollution Regulation in the US

This section uses tresnet to consider the EPA petition for public comment regarding their proposal to lower the National Ambient Air Quality Standard (NAQS) for annual PM\({}_{2.5}\) pollution from their current value of 12 \(\upmu\mathrm{g}/\mathrm{m}^{3}\). The new proposes standards suggest values of 9-10 \(\upmu\mathrm{g}/\mathrm{m}^{3}\), but could consider going as low as 8 \(\upmu\mathrm{g}/\mathrm{m}^{3}\).

**Overview** To reflect the EPA guidelines, we consider two types of exposure shifts from Table 1: cutoff intervention and percent reductions, each providing different perspectives and valuable information. An explanation and visual intuition about these exposure shifts have been provided in Table 1 and Figure 2. For each case, we estimate the corresponding SRF using tresnet for Poisson regression (Section 4). The results are shown in Figures 1 and 4. The exposure shifts can be described as follows:

1. A cutoff shift \(A_{i}^{d}=\min(A_{i},d)\) where \(d\) goes from 15 \(\upmu\mathrm{g}/\mathrm{m}^{3}\) down to 6 \(\upmu\mathrm{g}/\mathrm{m}^{3}\). We expect that at 15 \(\upmu\mathrm{g}/\mathrm{m}^{3}\) there will be no reduction in deaths since more than 99% of observations fall below that range. Analyzing this exposure shift would shed light on a counterfactual world where pollution exposure stays as they are for all levels below the standard, but only high values are shrunk.
2. An exposure shift \(A_{i}^{d}=A_{i}(1-d)\) implementing an overall percent reduction. We take \(d\) in the \(0-50\%\) range in the simulations. We can interpret these shifts in terms of the NAAQS by mapping each percent reduction to percentile and quantile metrics.

**Data** The application dataset comes from a previous study assessing the effect of PM\({}_{2.5}\) on elder mortality by Wu et al. (2020). We provide an overview here and refer to their article for details. The dataset includes Medicare data from the 2000-2016 cohort3. This cohort comprises 68,503,979 individuals and 27,106,639 early deaths. Measurements include age, sex, race/ethnicity, residential ZIP code, Medicaid eligibility, and date of death. The records are aggregated at the annual ZIP code level, resulting in 573,370,257 ZIP code years. The PM\({}_{2.5}\) exposure measurements are derived from the ensemble prediction model of Di et al. (2019). The ensemble was trained on daily PM\({}_{2.5}\) concentrations measured at 2,156 US EPA monitoring sites using predictor variables derived from satellite data, land-use data, and atmospheric measurements. It is then averaged at the annual ZIP code level. Finally, these sources are combined with demographic, socio-economic, and environmental and weather confounding factors (Tec et al., 2023) derived from the US Census Bureau, gridMET (Abatzoglou, 2013), and risk factors from the US Centers for Disease Control and Prevention.

Footnote 3: Access to Medicare data is available from The Center of Medicare and Medicaid Services to protect sensitive information under a restricted usage agreement from. The authors have completed the required IRB and ethical training to manipulate these data.

**Model and architecture** We conduct a Poisson regression for count data and apply targeted regularization for Poisson regression (Equation (11)). Each unit in the regression dataset is a ZIP code year. The year is included as an additional covariate. We use the same neural network architecture and hyper-parameters as in the simulation study. For the cutoff shift, we use the classifier density ratio loss. We preferred it because it provides a smoothing effect over the mass point in the shifted density (see Table 1). For the percent reduction, we use the parameterized density ratio loss. We fit the model using the 80-20% train-test split, taking the model that performed best in the reweighted validation error (ref). We repeat the procedure for 100 runs with different

\begin{table}
\begin{tabular}{c c|c c c c|c c c} \hline \hline Dataset & Metric & ipw & center-aipw & vcnet & denet & **tresnet** & **tresnet-aipw** & **tresnet-bnet** \\ \hline \multirow{3}{*}{hldp} & bias & \(-0.38\pm 0.08\) & \(-0.30\pm 0.24\) & \(-0.30\pm 0.24\) & \(-0.22\pm 0.10\) & \(-0.23\pm 0.26\) & \(-0.23\pm 0.26\) & \(-0.20\pm 0.1\) \\  & MSE & \(1.02\pm 0.07\) & \(0.46\pm 0.25\) & \(0.47\pm 0.25\) & \(\mathbf{0.34\pm 0.06}\) & \(0.40\pm 0.24\) & \(0.40\pm 0.24\) & \(0.38\pm 0.09\) \\ \hline \multirow{3}{*}{news} & bias & \(-0.02\pm 0.03\) & \(\mathbf{0.00\pm 0.03}\) & \(\mathbf{0.00\pm 0.03}\) & \(\mathbf{0.00\pm 0.03}\) & \(\mathbf{0.00\pm 0.03}\) & \(0.00\pm 0.03\) & \(0.02\pm 0.03\) \\  & MSE & \(\mathbf{0.03\pm 0.02}\) & \(\mathbf{0.03\pm 0.02}\) & \(\mathbf{0.03\pm 0.01}\) & \(\mathbf{0.03\pm 0.01}\) & \(\mathbf{0.03\pm 0.02}\) & \(\mathbf{0.03\pm 0.02}\) & \(\mathbf{0.03\pm 0.01}\) \\ \hline \multirow{3}{*}{sim} & bias & \(0.10\pm 0.06\) & \(0.06\pm 0.06\) & \(0.06\pm 0.06\) & \(0.07\pm 0.06\) & \(\mathbf{0.04\pm 0.07}\) & \(\mathbf{0.04\pm 0.07}\) & \(0.07\pm 0.06\) \\  & MSE & \(0.12\pm 0.05\) & \(0.08\pm 0.04\) & \(0.08\pm 0.04\) & \(0.10\pm 0.04\) & \(\mathbf{0.07\pm 0.04}\) & \(\mathbf{0.07\pm 0.05}\) & \(0.10\pm 0.04\) \\ \hline \hline \end{tabular}
\end{table}
Table 2: Comparison in benchmarks for continuous treatments. The task is to estimate an SRF for a percent reduction up to 50%.
seeds and data splits. We found no evidence of overfitting with the selected models. We approximately quantify uncertainty by taking the IQR (25% and 75% quantiles) over the predictions in the test set of the multiple re-runs. While it is not a formal estimation procedure, ensembling neural networks gives reasonable results (Izmailov et al., 2021).

ResultsFigure 1 presents the effects of cutoff intervention on deaths and it shows that there is a decrease in deaths as the NAAQS cutoff is reduced. The slope is steeper as the cutoff is reduced likely because lower cutoffs affects a larger fraction of the observed population. For instance, the figure shows that by enforcing 12 \(\upmu\mathrm{g}/\mathrm{m}^{3}\) current observed death counts would decrease by around 1%. In addition, if the cutoff is lowered to 9 \(\upmu\mathrm{g}/\mathrm{m}^{3}\), then deaths can fall up to around 4%. The slope in increasingly steeper, being strongly suggestive of increasing marginal benefits of lowering the standard. For instance, this means that there is as much or more gain from lowering the standard of 12 to 10 \(\upmu\mathrm{g}/\mathrm{m}^{3}\), as there is from lowering it from 10 to 8 \(\upmu\mathrm{g}/\mathrm{m}^{3}\).

The results for the percent reduction exposure shift are presented in Figure 4. A striking result is that decrease in deaths is linear with respect to the percent decrease in pollution. The model is not enforcing this result and, in fact, the simulation study show that ability to fit highly non-linear SRFs. As such, the SRF shows an approximate 0.5% decrease in deaths resulting from a 10% decrease in \(\mathrm{PM}_{2.5}\). If we take as reference the 75% currently at 12 \(\upmu\mathrm{g}/\mathrm{m}^{3}\) (Figure 4), we can identify interesting points along the SRF. Lowering to 10, 9, and 8 \(\upmu\mathrm{g}/\mathrm{m}^{3}\) requires a 15, 25, and 35 percent reduction in \(\mathrm{PM}_{2.5}\), respectively. These reductions correspond in turn to 0.7, 1.1, and 1.5 percent decrease in mortality.

Furthermore, Figure 3 shows that under a 20% reduction starting from 9 \(\upmu\mathrm{g}/\mathrm{m}^{3}\), about 30% of the ZIP code years would still exceed the threshold. It also shows that 25% percent of factual observations were above the current standard of 12 \(\upmu\mathrm{g}/\mathrm{m}^{3}\). In addition, the quantile view presented in that same figure shows that we will require a 32% overall reduction to bring down the 75% quantile from 12 to 8 \(\upmu\mathrm{g}/\mathrm{m}^{3}\). Comparing the results of the two exposure shifts suggests that further gains are obtained by enforcing the standard as a hard limit rather than a homogeneous reduction allowing a few locations to deviate.

## 8 Discussion and Ethics Statement

We have taken an important step toward answering the open public health issue concerning the health benefits of lowering the safety air quality standards in the US. To answer this question, we have provided the first causal inference method using neural networks for continuous stochastic interventions. Moreover, we extended this method to deal with count data as necessitated by our application, and perhaps many applications in the public health and epidemiology domains.

There are several key limitations worth mentioning. First, we only assessed uncertainty using an approximate cross-validating procedure. Thus, we advise against interpreting the statistical significance of the results based on the current confidence bounds, particularly for actionable policy considerations. Future work could explore combining fresnet with Bayesian methods. Secondly, our application considers hypothetical exposure shifts (representative of the EPA rule) but does not address the compliance question of estimating the most likely exposure shift produced by the new rule. Follow-up work will focus on this part of the application, which can also xuse tools found in the causal inference literature. Third, the annual average \(\mathrm{PM}_{2.5}\) assessments evaluated at the ZIP-code level are predictions rather than the true observed values, making our analysis subject to attenuation bias caused by measurement error. However, several experiments on measurement error with clustered air pollution exposures have demonstrated that the attenuation pulls the causal effect towards a null result, implying that the benefits of reducing the NAAQS may be greater than what we report (Josey et al., 2022; Wei et al., 2022). Finally, a potential disadvantage (and advantage) of the SI framework is that it places additional burden onto the analyst designing the exposure shift. Thus, SRF estimates have a propensity for misuse if the assumptions behind the exposure shift considered are not clearly and carefully stated.

Figure 4: Estimated SRF of the total deaths (%) for different overall percent reductions.

Figure 3: _(Left)_ Fraction (%) of observed units remaining above \(\mathrm{PM}_{2.5}\) limit as a function of reduction (%) considering different NAAQS (current NAAQS is at 12 \(\upmu\mathrm{g}/\mathrm{m}^{3}\)). _(Right)_ Quantiles of the annual \(\mathrm{PM}_{2.5}\) after the percent reduction.
## Acknowledgements

The work was supported by the National Institutes of Health (R01AG066793, R01ES030616, R01ES034373), the Alfred P. Sloan Foundation (G-2020-13946), and the Ren Che Foundation supporting the climate-smart public health project in the Golden Lab. Computations ran on the FASRC Cannon cluster supported by the FAS Division of Science Research Computing Group at Harvard University.
