Extended Fiducial Inference for Individual Treatment Effects via Deep Neural Networks

Causal inference is a fundamental problem in many disciplines such as medicine, econometrics, and social science. Formally, let $\{(y_{1},{\boldsymbol{x}}_{1},t_{1}),(y_{2},{\boldsymbol{x}}_{2},t_{2}),\ldots,(y_{n},{\boldsymbol{x}}_{n},t_{n})\}$ denote a set of observations drawn from the following data-generating equations:

where ${\boldsymbol{x}}_{i}\in\mathbb{R}^{d}$ represents a vector of covariates of subject $i$, $t_{i}\in\{0,1\}$ represents the treatment assignment to subject $i$; $c(\cdot)$ represents the expected outcome of subject $i$ if assigned to the control group (with $t_{i}=0)$, and $\tau({\boldsymbol{x}}_{i})$ is the expected treatment effect of subject $i$ if assigned to the treatment group (with $t_{i}=1$); $\sigma>0$ is the standard deviation, and $z_{i}$ represent a standardized random error that is not necessarily Gaussian. Under the potential outcome framework (Rubin, 1974), each individual receives only one assignment of the treatment with $t_{i}=0$ or 1, but not both. The goal of causal inference is to make inference for the average treatment effect (ATE) or individual treatment effect (ITE).

The ATE is defined as

where $\mathcal{X}$ denotes the sample space of ${\boldsymbol{x}}$, and $F({\boldsymbol{x}})$ denotes the cumulative distribution function of ${\boldsymbol{x}}$. To estimate ATE, a variety of methods, including outcome regression, augmented/inverse probability weighting (AIPW/IPW) and matching, have been developed. See Imbens (2004) and Rosenbaum (2002) for overviews.

The ITE is often defined as the conditional average treatment effect (CATE):

see e.g., Shalit et al. (2017) and Lu et al. (2018). Recently, Lei and Candès (2021) proposed to make predictive inference of the ITE by quantifying the uncertainty of

where $Y_{i}(t_{i})$ denotes the potential outcome of subject $i$ with treatment assignment $t_{i}\in\{0,1\}$. Henceforth, we will call $\tilde{\tau}_{i}$ the predictive ITE.

It is known that ATE and ITE are identifiable if the conditions ‘strong ignorability’ and ‘overlapping’ are satisfied. The former means that, after accounting for observed covariates, the treatment assignment is independent of potential outcomes; and the latter ensures that every subject in the study has a positive probability of receiving either assignment, allowing for meaningful comparisons between treatment and control groups. Mathematically, the two conditions can be expressed as:

where $T\in\{0,1\}$ represents the treatment assignment variable, and $\models$ denotes conditional independence. Together, they ensure that the causal effect can be correctly estimated without bias. See e.g. Guan and Yang (2019) for more discussions on this issue.

However, even under these assumptions, accurate inference for ATE and ITE can still be challenging. Specifically, the inference task can be complicated by unknown nonlinear forms of $c({\boldsymbol{x}})$ and $\tau({\boldsymbol{x}})$. To address these issues, some authors have proposed to approximate them using a machine learning model, such as random forest (RF) (Breiman, 2001), Bayesian additive regression trees (BART) (Chipman et al., 2010), and neural networks. Refer to e.g., Foster et al. (2011), Hill (2011), Shalit et al. (2017), Wager and Athey (2018), and Hahn et al. (2020) for the details. Unfortunately, these methods often yield point estimates for the ATE and ITE, while failing to correctly quantifying their uncertainty due to the complexity of the machine learning models. Quite recently, Lei and Candès (2021) proposed to quantify the uncertainty of the predictive ITE using the conformal inference method (Vovk et al., 2005; Shafer and Vovk, 2008). This method provides coverage-guaranteed confidence intervals for the predictive ITE, but the intervals may become overly wide when the machine learning model is not consistently estimated. In short, while machine learning models, particularly neural networks, can effectively model complex, nonlinear functions such as $c(\cdot)$ and $\tau(\cdot)$ for causal inference, performing accurate uncertainty quantification with these models remains a significant challenge. This is because these models typically have a complex functional form and involve a large number of parameters.

In this paper, we propose to conduct causal inference using an extended fiducial inference (EFI) method (Liang et al., 2024), with the goal of addressing the uncertainty quantification issue associated with treatment effect estimation. EFI provides an innovative framework for inferring model uncertainty based solely on observed data, aligning with the goal of fiducial inference (Fisher, 1935; Hannig, 2009). Specifically, it aims to solve the data-generating equations by explicitly imputing the unobserved random errors and approximating the model parameters from the observations and imputed random errors using a neural network; it then infers the uncertainty of the model parameters based on the learned neural network function and the imputed random errors (see Section 2 for a brief review). To make the EFI method feasible for causal effect estimation with accurate uncertainty quantification, we extend the method in two key aspects:

• (i)

  We approximate each of the unknown functions, $c({\boldsymbol{x}})$ and $\tau({\boldsymbol{x}})$, by a deep neural network (DNN) model. The DNN possesses universal approximation capability (Hornik et al., 1989; Hornik, 1991; Kidger and Lyons, 2020), meaning it can approximate any continuous function to an arbitrary degree of accuracy, provided it is sufficiently wide and deep. This property makes the proposed method applicable to a wide range of data-generating processes.

• (ii)

  We theoretically prove that the dimensions (i.e., the number of parameters) of the DNN models used to approximate $c({\boldsymbol{x}})$ and $\tau({\boldsymbol{x}})$ are allowed to increase with the sample size $n$ at a rate of $O(n^{\zeta})$ for some $0<\zeta<1$, while the uncertainty of the DNN models can still be correctly quantified. That is, we are able to correctly quantify the uncertainty of the causal effect although it has to be approximated using large models.

In this paper, we regard a model as ‘large’ if its dimension increases with $n$ at a rate of $1/2\leq\zeta<1$. We note that part (ii) represents a significant theoretical innovation in statistical inference for large models. In the literature on this area, most efforts have focused on linear models, featuring techniques such as desparsified Lasso (Javanmard and Montanari, 2014; van de Geer et al., 2014; Zhang and Zhang, 2014), post-selection inference (Lee et al., 2016), and Markov neighborhood regression (Liang et al., 2022a). For nonlinear models, the research landscape appears to be more scattered. Portnoy (1986, 1988) showed that for independently and identically distributed (i.i.d) random vectors with the dimension $p$ increasing with the sample size $n$, the central limit theorem (CLT) holds if $p=O(n^{\zeta})$ for some $0\leq\zeta<1/2$. It is worth noting that Bayesian methods, despite being sampling-based, do not permit the dimension of the true model to increase with $n$ at a higher rate. For example, even in the case of generalized linear models, to ensure the posterior consistency, the dimension of the true model is only allowed to increase with $n$ at a rate $0\leq\zeta<1/4$ (see Theorem 2 and Remark 2 of Jiang (2007)). Under its current theoretical framework developed by Liang et al. (2024), EFI can only be applied to make inference for the models whose dimension is fixed or increases with $n$ at a very low rate. This paper extends the theoretical framework of EFI further, establishing its applicability for statistical inference of large models.

It is worth noting that a DNN model with size $p=O(n^{\zeta})$, where $\zeta$ is close to (but less than) 1, has been shown to be sufficiently large for approximating many data generation processes. This is supported by the theory established in Sun et al. (2022) and Farrell et al. (2021). In Sun et al. (2022), it is shown that, as $n\to\infty$, a sparse DNN model of this size can provide accurate approximations for multiple classes of functions, such as bounded $\alpha$-Hölder smooth functions (Schmidt-Hieber, 2020), piecewise smooth functions with fixed input dimensions (Petersen and Voigtlaender, 2018), and functions representable by an affine system (Bolcskei et al., 2019). Similar results have also been obtained in Farrell et al. (2021), where it is shown that a multi-layer perceptron (MLP) with this model size and the ReLU activation function can provide an accurate approximation to the functions that lie in a Sobolev ball with certain smoothness. The approximation capability of DNNs of this size has also been empirically validated by Hestness et al. (2017), where a neural scaling law of $p=O(n^{\zeta})$ with $0.5\leq\zeta<1$ was identified through extensive studies across various model architectures in machine translation, language modeling, image processing, and speech recognition.

To highlight the strength of EFI in uncertainty quantification and to facilitate comparison with the conformal inference method, this study focuses on inference for predictive ITEs, although the proposed method can also be extended to ATE and CATE. Our numerical results demonstrate the superiority of the proposed method over the conformal inference method.

The remaining part of this paper is organized as follows. Section 2 provides a brief review of the EFI method. Section 3 extends EFI to statistical inference for large statistical models. Section 4 provides an illustrative example for EFI. Section 5 applies the proposed method to statistical inference for predictive ITEs, with both simulated and real data examples. Section 6 concludes the paper with a brief discussion.

While fiducial inference was widely considered as a big blunder by R.A. Fisher, the goal he initially set —inferring the uncertainty of model parameters on the basis of observations — has been continually pursued by many statisticians, see e.g. Zabell (1992), Hannig (2009), Hannig et al. (2016), Murph et al. (2022), and Martin (2023). To this end, Liang et al. (2024) developed the EFI method based on the fundamental concept of structural inference (Fraser, 1966, 1968). Consider a regression model:

where $Y\in\mathbb{R}$ and ${\boldsymbol{X}}\in\mathbb{R}^{d}$ represent the response and explanatory variables, respectively; ${\boldsymbol{\theta}}\in\mathbb{R}^{p}$ represents the vector of parameters; and $Z\in\mathbb{R}$ represents a scaled random error following a known distribution $\pi_{0}(\cdot)$. For the model (1), the treatment assignment $T$ should be included as a part of ${\boldsymbol{X}}$.

Suppose that a random sample of size $n$ has been collected from the model, denoted by $\{(y_{1},{\boldsymbol{x}}_{1}),(y_{2},{\boldsymbol{x}}_{2}),\ldots,(y_{n},{\boldsymbol{x}}_{n})\}$. In the point of view of structural inference (Fraser, 1966, 1968), they can be expressed in the data generating equations as follow:

This system of equations consists of $n+p$ unknowns, namely, $\{{\boldsymbol{\theta}},z_{1},z_{2},\ldots,z_{n}\}$, while there are only $n$ equations. Therefore, the values of ${\boldsymbol{\theta}}$ cannot be uniquely determined by the data-generating equations, and this lack of uniqueness of unknowns introduces uncertainty in ${\boldsymbol{\theta}}$.

Let ${\boldsymbol{Z}}_{n}=\{z_{1},z_{2},\ldots,z_{n}\}$ denote the unobservable random errors, which are also called latent variables in EFI. Let $G(\cdot)$ denote an inverse function/mapping for the parameter ${\boldsymbol{\theta}}$, i.e.,

It is worth noting that the inverse function is generally non-unique. For example, it can be constructed by solving any $p$ equations in (6) for ${\boldsymbol{\theta}}$. As noted by Liang et al. (2024), this non-uniqueness of inverse function mirrors the flexibility of frequentist methods, where different estimators of ${\boldsymbol{\theta}}$ can be designed for different purposes.

As a general method, Liang et al. (2024) proposed to approximate the inverse function $G(\cdot)$ using a sparse DNN, see Figure 1 for illustration. They also introduced an adaptive stochastic gradient Langevin dynamics (SGLD) algorithm, which facilitates the simultaneous training of the sparse DNN and simulation of the latent variables ${\boldsymbol{z}}$. This is briefly described as follows.

Let $\hat{{\boldsymbol{\theta}}}_{i}:=\hat{g}(y_{i},{\boldsymbol{x}}_{i},z_{i},{\boldsymbol{w}}_{n})$ denote the DNN prediction function parameterized by the weights ${\boldsymbol{w}}_{n}$ in the EFI network, and let

which serves as an estimator of $G(\cdot)$. The EFI network has two output nodes defined, respectively, by

where $\tilde{y}_{i}=f({\boldsymbol{x}}_{i},z_{i},\bar{{\boldsymbol{\theta}}})$, $f(\cdot)$ is as specified in (6), and $d(\cdot)$ is a function that measures the difference between $y_{i}$ and $\tilde{y}_{i}$. For example, for a normal linear/nonlinear regression, it can be defined as

For logistic regression, it is defined as a squared ReLU function, see Liang et al. (2024) for the details. Furthermore, EFI defines an energy function as follows:

for some regularization coefficient $\eta>0$, where first term measures the fitting error of the model as implied by equation (10), and the second term regularizes the variation of $\hat{{\boldsymbol{\theta}}}_{i}$, ensuring that the neural network forms a proper estimator of the inverse function. Given this energy function, we define the likelihood function as

for some constant $\epsilon$ close to 0. As discussed in Liang et al. (2024), the choice of $\eta$ does not have much affect on the performance of EFI as long as $\epsilon$ is sufficiently small.

Subsequently, the posterior of ${\boldsymbol{w}}_{n}$ is given by

where $\pi({\boldsymbol{w}}_{n})$ denotes the prior of ${\boldsymbol{w}}_{n}$; and the predictive distribution of ${\boldsymbol{Z}}_{n}$ is given by

In EFI, ${\boldsymbol{w}}_{n}$ is estimated through maximizing the posterior $\pi_{\epsilon}({\boldsymbol{w}}_{n}|{\boldsymbol{X}}_{n},{\boldsymbol{Y}}_{n})$ given the observations $\{{\boldsymbol{X}}_{n},{\boldsymbol{Y}}_{n}\}$. By the Bayesian version of Fisher’s identity (Song et al., 2020), the gradient equation $\nabla_{{\boldsymbol{w}}_{n}}\log\pi_{\epsilon}({\boldsymbol{w}}_{n}|{\boldsymbol{X}}_{n},{\boldsymbol{Y}}_{n})$ $=0$ can be re-expressed as

which can be solved using an adaptive stochastic gradient MCMC algorithm (Liang et al., 2022b; Deng et al., 2019). The algorithm works by iterating between two steps:

• (a)

  Latent variable sampling: draw ${\boldsymbol{Z}}_{n}^{(k+1)}$ according to a Markov transition kernel that leaves $\pi_{\epsilon}({\boldsymbol{z}}|{\boldsymbol{X}}_{n},{\boldsymbol{Y}}_{n},{\boldsymbol{w}}_{n}^{(k)})$ to be invariant;

• (b)

  Parameter updating: update ${\boldsymbol{w}}_{n}^{(k)}$ toward the maximum of $\log\pi_{\epsilon}({\boldsymbol{w}}_{n}|{\boldsymbol{X}}_{n},{\boldsymbol{Y}}_{n},{\boldsymbol{Z}}_{n})$ using stochastic approximation (Robbins and Monro, 1951), based on the sample ${\boldsymbol{Z}}_{n}^{(k+1)}$.

See Algorithm 1 for the pseudo-code. This algorithm is termed “adaptive” because the transition kernel in the latent variable sampling step changes with the working parameter estimate of ${\boldsymbol{w}}_{n}$. The parameter updating step can be implemented using mini-batch SGD, and the latent variable sampling step can be executed in parallel for each observation $(y_{i},{\boldsymbol{x}}_{i})$. Hence, the algorithm is scalable with respect to large datasets.

Under mild conditions for adaptive stochastic gradient MCMC algorithms (Deng et al., 2019; Liang et al., 2022b), it was shown in Liang et al. (2024) that

where ${\boldsymbol{w}}_{n}^{*}$ denotes a solution to equation (15) and $\stackrel{{\scriptstyle p}}{{\to}}$ denotes convergence in probability, and that

in 2-Wasserstein distance, where $\stackrel{{\scriptstyle d}}{{\rightsquigarrow}}$ denotes weak convergence.

To study the limit of (18) as $\epsilon$ decays to 0, i.e.,

where $p_{n}^{*}({\boldsymbol{z}}|{\boldsymbol{Y}}_{n},{\boldsymbol{X}}_{n},{\boldsymbol{w}}_{n}^{*})$ is referred to as the extended fiducial density (EFD) of ${\boldsymbol{Z}}_{n}$ learned in EFI, it is necessary for ${\boldsymbol{w}}_{n}^{*}$ to be a consistent estimator of ${\boldsymbol{w}}_{*}$, the parameters of the underlying true EFI network. To ensure this consistency, Liang et al. (2024) impose some conditions on the structure of the DNN and the prior distribution $\pi({\boldsymbol{w}}_{n})$. Specifically, they assume that ${\boldsymbol{w}}_{n}$ takes values in a compact space $\mathcal{W}$; $\pi({\boldsymbol{w}}_{n})$ is a truncated mixture Gaussian distribution on $\mathcal{W}$; and the DNN structure satisfies certain constraints given in Sun et al. (2022), e.g., the width of the output layer (i.e., the dimension of ${\boldsymbol{\theta}}$) is fixed or grows very slowly with $n$. They then justify the consistency of ${\boldsymbol{w}}_{n}^{*}$ based on the sparse deep learning theory developed in Sun et al. (2022). The consistency of ${\boldsymbol{w}}_{n}^{*}$ further implies that

serves as a consistent estimator for the inverse function/mapping ${\boldsymbol{\theta}}=G({\boldsymbol{Y}}_{n},{\boldsymbol{X}}_{n},{\boldsymbol{Z}}_{n})$.

By Theorem 3.2 in Liang et al. (2024), for the target model (1), which is a noise-additive model, the EFD of ${\boldsymbol{Z}}_{n}$ is invariant to the choice of the inverse function, provided that $d(\cdot)$ is specified as in (10) in defining the energy function. Further, by Lemma 4.2 in Liang et al. (2024), $p_{n}^{*}({\boldsymbol{z}}|{\boldsymbol{Y}}_{n},{\boldsymbol{X}}_{m},{\boldsymbol{w}}_{n}^{*})$ is given by

where $P_{n}^{*}({\boldsymbol{z}}|{\boldsymbol{X}}_{n},{\boldsymbol{Y}}_{n},{\boldsymbol{w}}_{n}^{*})$ represents the cumulative distribution function (CDF) corresponding to $p_{n}^{*}({\boldsymbol{z}}|{\boldsymbol{X}}_{n},{\boldsymbol{Y}}_{n},{\boldsymbol{w}}_{n}^{*})$; $\mathcal{Z}_{n}=\{{\boldsymbol{z}}:U_{n}({\boldsymbol{Y}}_{n},{\boldsymbol{X}}_{n},{\boldsymbol{Z}}_{n},{\boldsymbol{w}}_{n}^{*})=0\}$ represents the zero-energy set, which forms a manifold in the space $\mathbb{R}^{n}$; and $\nu$ is the sum of intrinsic measures on the $p$-dimensional manifold in $\mathcal{Z}_{n}$. That is, under the consistency of ${\boldsymbol{w}}_{n}^{*}$, $p_{n}^{*}({\boldsymbol{z}}|{\boldsymbol{X}}_{n},{\boldsymbol{Y}}_{n},{\boldsymbol{w}}_{n}^{*})$ is reduced to a truncated density function of $\pi_{0}^{\otimes n}({\boldsymbol{z}})$ on the manifold $\mathcal{Z}_{n}$, while $\mathcal{Z}_{n}$ itself is also invariant to the choice of the inverse function as shown in Lemma 3.1 of Liang et al. (2024). In other words, for the model (1), the EFD of ${\boldsymbol{Z}}_{n}$ is asymptotically invariant to the inverse function we learned given its consistency.

Let $\Theta:=\{{\boldsymbol{\theta}}\in\mathbb{R}^{p}:{\boldsymbol{\theta}}=G^{*}({\boldsymbol{Y}}_{n},{\boldsymbol{X}}_{n},{\boldsymbol{z}}),{\boldsymbol{z}}\in\mathcal{Z}_{n}\}$ denote the parameter space of the target model, which represents the set of all possible values of ${\boldsymbol{\theta}}$ that $G^{*}(\cdot)$ takes when ${\boldsymbol{z}}$ runs over $\mathcal{Z}_{n}$. Then, for any function $b({\boldsymbol{\theta}})$ of interest, its EFD $\mu_{n}^{*}(\cdot|{\boldsymbol{Y}}_{n},{\boldsymbol{X}}_{n})$ associated with $G^{*}(\cdot)$ is given by

where $\mathcal{Z}_{n}(B)=\{{\boldsymbol{z}}\in\mathcal{Z}_{n}:b(G^{*}({\boldsymbol{Y}}_{n},{\boldsymbol{X}}_{n},{\boldsymbol{z}}))\in B\}$. The EFD provides an uncertainty measure for $b({\boldsymbol{\theta}})$. Practically, the EFD of $b({\boldsymbol{\theta}})$ can be constructed based on the samples $\{b(\bar{{\boldsymbol{\theta}}}_{1}),b(\bar{{\boldsymbol{\theta}}}_{2}),\ldots,b(\bar{{\boldsymbol{\theta}}}_{M})\}$, where $\{\bar{{\boldsymbol{\theta}}}_{1},\bar{{\boldsymbol{\theta}}}_{2},\ldots,\bar{{\boldsymbol{\theta}}}_{M}\}$ denotes the fiducial $\bar{{\boldsymbol{\theta}}}$-samples collected at step (iv) of Algorithm 1.

Finally, we note that, as discussed in Liang et al. (2024), the invariance property of $\mathcal{Z}_{n}$ is not crucial to the validity of EFI, although it does enhance the robustness of the inference. Additionally, for a neural network model, its parameters are only unique up to certain loss-invariant transformations, such as reordering hidden neurons within the same hidden layer or simultaneously altering the sign or scale of certain connection weights, see Sun et al. (2022) for discussions. Therefore, in EFI, the consistency of ${\boldsymbol{w}}_{n}^{*}$ refers to its consistency with respect to one of the equivalent solutions to (15), while mathematically ${\boldsymbol{w}}_{n}^{*}$ can still be treated as unique. Refer to Section $\S$1.1 (of the supplement) for more discussions on this issue.

In this section, we first establish the consistence of the inverse function/mapping learned in EFI for large models, and then discuss its application for uncertainty quantification of deep neural networks.

To illustrate how EFI works for statistical inference problems, we consider a linear regression example:

where $T_{i}\in\{0,1\}$ is a binary variable indicating the treatment assignment, $\tau$ is the treatment effect, ${\boldsymbol{x}}_{i}\in\mathbb{R}^{d}$ are confounders/covariates, $z_{i}\sim N(0,1)$ is the standardized random noise, and ${\boldsymbol{\beta}}\in\mathbb{R}^{d}$ and $\sigma\in\mathbb{R}_{+}$ are unknown parameters. For this example, $\tau$ represents the ATE as well as the CATE, due to its independence of the covariates ${\boldsymbol{x}}$. In the simulation study, we set $\tau=1$, $\mu=1$, $d=4$, and ${\boldsymbol{\beta}}=(-1,1,-1,1)^{\top}$; generate ${\boldsymbol{x}}_{i}\sim N(0,I_{d})$; and generate the treatment variable via a logistic regression:

where $\nu=1$ and ${\boldsymbol{\xi}}=(-1,1,-1,1)^{\top}$. We consider three different cases with the sample size $n=250$, 500 and 1000, respectively. For each case, we generate 100 datasets.

Statistical inference for the parameters in the model (22) can be made with EFI under its standard framework. Let ${\boldsymbol{\theta}}=(\tau,\mu,{\boldsymbol{\beta}}^{\top},\log\sigma)^{\top}$ be the parameter vector. EFI approximates the inverse function ${\boldsymbol{\theta}}=g(y,T,{\boldsymbol{x}},z)$ by a DNN, for which $(y,T,{\boldsymbol{x}},z)$ serves as input variables and ${\boldsymbol{\theta}}$ as output variables. The results are summarized in Table 1.

For comparison, a variety of methods, including Unadj (Imbens and Rubin, 2015), inverse probability weighting (IPW) (Rosenbaum, 1987), double-robust (DR) (Robins et al., 1994; Bang and Robins, 2005), and BART (Hill, 2011), have been applied to this example. These methods fall into distinct categories. The Unadj is straightforward, estimating the ATE by calculating the difference between the treatment and control groups, i.e., $\hat{\tau}=\frac{1}{n_{t}}\sum_{i=1}^{n_{t}}Y_{i}(1)-\frac{1}{n_{c}}\sum_{i=1}^{n_{c}}Y_{i}(0)$, where the effect of confounders is not adjusted. Both IPW and DR are widely used ATE estimation methods, which adjust the effect of confounders based on propensity scores. They both are implemented using the R package drgee (Zetterqvist and Sjölander, 2015). The BART employs Bayesian additive regression trees to learn the outcome function, which naturally accommodates heterogeneous treatment effects as well as nonlinearity of the outcome function. It is implemented using the R package bartcause (Dorie and Hill, 2020).

The comparison indicates that EFI performs very well for this standard ATE estimation problem. Specifically, EFI generates confidence intervals of nearly the same length as DR, but with more accurate coverage rates. This is remarkable, as DR has often been considered as the golden standard for ATE estimation and is consistent if either the outcome or propensity score models is correctly specified, and locally efficient if both are correctly specified. Furthermore, EFI produces much shorter confidence intervals compared to Unadj, IPW, and BART, while maintaining more accurate coverage rates.

We attribute the superior performance of EFI on this example to its fidelity in parameter estimation, an attractive property of EFI as discussed in Liang et al. (2024). As implied by (14), EFI essentially estimates ${\boldsymbol{\theta}}$ by maximizing the predictive likelihood function $\pi_{\epsilon}({\boldsymbol{Z}}_{n}|{\boldsymbol{X}}_{n},{\boldsymbol{Y}}_{n},{\boldsymbol{\theta}})\propto\pi_{0}^{\otimes n}({\boldsymbol{Z}}_{n})e^{-U_{n}({\boldsymbol{Y}}_{n},{\boldsymbol{X}}_{n},{\boldsymbol{Z}}_{n},{\boldsymbol{w}})/\epsilon}$, which balances the likelihood of ${\boldsymbol{Z}}_{n}$ and the model fitting errors coded in $U_{n}(\cdot)$. In contrast, the maximum likelihood estimation (MLE) method sets $\hat{{\boldsymbol{\theta}}}_{MLE}=\arg\max_{{\boldsymbol{\theta}}}\pi_{0}^{\otimes n}({\boldsymbol{Z}}_{n})$, where ${\boldsymbol{Z}}_{n}$ is expressed as a function of $({\boldsymbol{Y}}_{n},{\boldsymbol{X}}_{n},{\boldsymbol{\theta}})$. In general, MLE is inclined to be influenced by the outliers and deviations of covariates especially when the sample size is not sufficiently large. It is important to note that the MLE serves as the core for all the IPW, DR and BART methods in estimating the outcome and propensity score models. For this reason, various adjustments for confounding and heterogeneous treatment effects have been developed in the literature.

Compared to the existing causal inference methods, EFI works as a solver for the data-generating equation (as $\epsilon\downarrow 0$), providing a coherent way to address the confounding and heterogeneous treatment effects and resulting in faithful estimates for the model parameters and their uncertainty as well. This example illustrates the performance of EFI in ATE estimation when confounders are present, while the examples in the next section showcase the performance of EFI in dealing with heterogeneous treatment effects via DNN modeling. Extensive comparisons with BART and other nonparametric modeling methods are also presented.

In this example, we omit the estimation of the propensity score model. As discussed in Section 6, the proposed method can be extended by including an additional DNN to approximate the propensity score, enabling the use of inverse probability weighting for ATE estimation. However, the ATE estimation is not the focus of this work.

This section demonstrates how EFI can be used to perform statistical inference of the predictive ITE for the data-generating model (1). Let ${\boldsymbol{\theta}}_{c}$ denote the vector of parameters for modeling the function $c({\boldsymbol{x}})$, let ${\boldsymbol{\theta}}_{\tau}$ denote the vector of parameters for modeling the function $\tau({\boldsymbol{x}})$, and let ${\boldsymbol{\theta}}=\{{\boldsymbol{\theta}}_{c},{\boldsymbol{\theta}}_{\tau},\log(\sigma)\}$ denote the whole set of parameters for the model (1). We model the inverse function ${\boldsymbol{\theta}}=g(y,T,{\boldsymbol{x}},z)$ by a DNN. Also, we can model each of the functions $c({\boldsymbol{x}})$ and $\tau({\boldsymbol{x}})$ by a DNN if their functional forms are unknown. For convenience, we refer to the DNN for modeling $c({\boldsymbol{x}})$ as ‘$c$-network’ and that for modeling $\tau({\boldsymbol{x}})$ as ‘$\tau$-network’, and ${\boldsymbol{\theta}}_{c}$ and ${\boldsymbol{\theta}}_{\tau}$ represent their weights, respectively. As mentioned previously, we can restrict the sizes of the $c$-network and $\tau$-network to the order of $O(n^{\tilde{\zeta}})$ for some $0<\tilde{\zeta}<1$.

Note that in solving the data generating equations (1), the proposed method involves two types of neural networks: one for modeling causal effects and the other for approximating the inverse function ${\boldsymbol{\theta}}=g(y,T,{\boldsymbol{x}},z)$. While we still refer to the proposed method as ‘Double-NN’, it actually involves three DNNs.