An Improved Neural Network Model for Treatment Effect Estimation

An improved neural network model for treatment effect estimation Niki Kiriakidou?[0000−0003−1729−4124] and Christos Diou[0000−0002−2461−1928] Department of Informatics and Telematics, Harokopio University, Greece {kiriakidou,cdiou}@hua.gr Abstract. Nowadays, in many scientific and industrial fields there is an increasing need for estimating treatment effects and answering causal questions. The key for addressing these problems is the wealth of observational data and the processes for leveraging this data. In this work, we propose a new model for predicting the potential outcomes and the propensity score, which is based on a neural network architecture. The proposed model exploits the covariates as well as the outcomes of neighboring instances in training data. Numerical experiments illustrate that the proposed model reports better treatment effect estimation performance compared to state-of-the-art models. Keywords: Causal inference ·Dragonnet ·treatment effect ·potential outcomes ·propensity score. 1 Introduction For decades, causal inference has been a crucial research topic in many scientific fields, such as healthcare [4], education [7] and economics [20]. Causal inference aims at answering questions regarding the effect of interventions, (e.g., a new drug, a new educational method or a new pricing policy) to the target outcome variables (e.g., health, learning or financial indicators, respectively). The inference of causal effects is a challenging problem and the most effectual way to infer causality is through randomized controlled trials (RCTs). In many cases, however, it is expensive, time-consuming, unethical or even impossible to conduct an RCT. Nowadays, the abundance of observational data presents an opportunity for accurate estimation of causal effects, however, observational data contain recorded information about samples, such as actions and outcomes along with appropriate context, but there is way to directly influence the mechanism that caused the action. Furthermore, in observational data may exist confounding variables, which affect both treatment and outcome. If these are not adjusted, they could lead to incorrect and misleading results. In this work, a neural network model is proposed for treatment effect estimation through the prediction of the conditional outcomes and the propensity score. The model extends the state-of-the-art Dragonnet architecture [18] to exploit the ?Corresponding author arXiv:2205.11106v1 [stat.ML] 23 May 2022 2 N. Kiriakidou and C. Diou covariates along with information from the outcomes of the instances contained in the training data. The rationale behind the proposed approach is to enrich the inputs of the model with the average outcomes of the nearest neighbors from the control and treatment group along with the covariates, in order to reduce bias and increase the prediction accuracy. To estimate treatment effects, the proposed method first trains a model for the prediction of conditional outcomes and the propensity score and then the trained model is used by a downstream estimator. Our experiments illustrate that the proposed approach maintains state of the art performance for the estimation of average treatment effect (ATE), while it leads to significant improvement in estimating the individual treatment effect (ITE). The remainder of this paper is organised as follows: Section 2 presents a review of neural network based models for the estimation of treatment effects. Section 3 presents a comprehensive description of the proposed modified model and its architecture. Section 4 provides information about the data. Section 5 presents a detailed experimental analysis, focusing on the evaluation of the proposed model. Section 6 summarizes the main findings and conclusions of this research, and some interesting directions for future work. 2 Related Work During the last decade, a lot of research has been conducted towards more accurate and reliable estimation of treatment effects. Most of this research is based on the use of neural networks, exploiting the predictive power of these machine learning models. Johansson et al. [10] proposed a new algorithmic framework for counterfactual inference. More specifically, they formulated the causal inference problem as a domain adaptation problem and developed a new class of representation algorithms for the calculation of treatment effects. They highlighted that learning representations, which enforce similarity between control and treated groups, is able to lead to better estimations of causal effects. They compared a variant of the proposed algorithm based on a neural network approach, named Balancing Neural Network (BNN), against traditional models, which reported the best overall performance. Shalit et al. [17] proposed a new theoretical analysis and a new framework, named Counterfactual Regression (CFR) for predicting individual treatment effects. The proposed framework aims on learning a balanced representation using a prediction model, so that the distributions of control and treated group look similar. To measure the distances between two distributions they utilized the integral probability metrics: Maximum Mean Discrepancy (MMD) [5] and Wasserstein distance (Wass) [21]. Additionaly, the major contribution of their work is the introduction of a generalization-bound for the estimation of individual treatment effect, where every individual is only identified by its features. In their experiments, they compared the performance of two proposed models, CFR (MMD) and CFR (Wass), which use MMD and Wass distances, re- An improved neural network model for treatment effect estimation 3 spectively, against state-of-the-art models. Furthermore, they included a variant without balance regularization, named Treatment Agnostic Representation Network (TARNet). Based on their experimental analysis, they stated that all proposed models presented the best performance in terms of estimating treatment effects. Another approach for estimating individual treatment effect was proposed by Yoon et al. [22], which is based on Generative Adversarial Nets (GANs). The rationale behind the proposed approach is to simulate the uncertainty in the counterfactual distributions by considering learning them using a GAN model. Along this line, they developed a novel model, named Generative Adversarial Nets for inference of Individalized Treatment Effects (GANITE), which was able to provide confidence intervals for its predictions. Their numerical experiments revealed that the proposed method exhibited promising performance. Louizos et al. [13] highlighted the significance of handling confounders for inferring treatment effects from observational data. More specifically, they stated that there is a strong possibility of existing uncertain and noisy “proxy variables”, in case there is no access to all confounders. To address the previous difficulties they proposed a new model, called CEVAE, based on variational autoencoders. A considerable advantage of their approach is that the data generating process as well as the structure of the hidden confounders requires substantially weaker assumptions. Finally, the authors presented that CEVAE exhibited more robust behaviour against hidden confounders in the case of noisy proxies. Shi et al. [18] proposed a novel neural network model for estimating treatment effects from observational data. The proposed model, named Dragonnet, focuses on improving the estimations through the sufficiency of the propensity score. Additionally, the authors proposed targeted regularization, which constitutes a procedure to induce bias based on non-parametric estimation theory and aims to further improve the estimation of treatment effect. Finally, the authors provided experimental evidence about the superiority of Dragonnet against BNN, CEVAE, GANITE, TARNet, CFR (MMD) and CFR (Wass) using two benchmark datasets. In this work, we propose a neural network model for predicting the potential outcomes and the propensity score. The proposed model architecture is a modification of Dragonnet’s architecture. The major difference between the proposed model and Dragonnet is that the former’s inputs contain information from the covariates as well as from the outputs of control and treated group.Our numerical experiments provide empirical and statistical evidence about the efficacy and efficiency of our approach. 3 Modified Dragonnet model In this section, we present the proposed model for the estimation of treatment effects. The rationale behind our approach is to enrich the training data with information from the outcomes, which can be exploited by the proposed model in order to obtain more accurate predictions. 4 N. Kiriakidou and C. Diou 3.1 Calculation of average outcome vectors We limit our discussion to the case of binary treatments. Let Xdenote the d-dimensional space of covariates and consider a joint distribution Πon X × {0,1} × Y. Suppose that (X, T, Y )∼Π, are random variables with domains X,{0,1}and Y, corresponding to the covariates, treatment and outcome for a single sample, respectively. Let also Y0denote the outcome for a sample when T= 0 and Y1stand for the outcome of a sample when T= 1. Given a dataset (xi, ti, yi),i= 1,2, . . . , n where xi∈ X ,ti∈ {0,1}and yi∈ Y our goal is to estimate the average treatment effect ψ=E[Y|X, T = 1] −E[Y|X, T = 0] (1) For each observed sample in the dataset, either ti= 0 (Y0is factual) or ti= 1 (Y0is counterfactual) and yi=tiY1+ (1 −ti)Y0, based on the framework of Neyman-Rubin [16]. The main idea of our model is to reduce bias in treatment effect estimation, by utilizing the average outcomes of knearest neighbors y(0) iof the control group and y(1) iof the treatment group for each available sample i. Algorithm 1 presents a pseudocode for the calculation of y(0) iand y(1) i. The algorithm takes as inputs the design matrix X, whose rows correspond to the covariate vectors of samples, the binary vector of treatment values t, the outcome vector yin the dataset, as well as the number of nearest neighbors k. Initially, y(0) and y(1) are initialized to 0. (Step 1). Next, for every instance xiwe calculate the average outcomes for control and treated group (Steps 2-7). More specifically, we calculate the k-nearest neighbors of xiin X, contained in the control group (i.e T= 0) and append their corresponding indices in the index set S0(Step 4). Then, we calculate the average of the outcomes of these neighbors, y(0) i=1 kPj∈S0yj(Step 5) Similarly, we calculate the average outcome of the k-nearest neighbors of xi, contained in the treatment group (i.e T= 1)(Step 6-7) Algorithm 1 Inputs: X: design matrix t: vector of treatment values t y: vector of outcome values y k: number of nearest neighbors Output: y(0): vector with average of k-nearest outcomes from control group for each sample y(1): vector with average of k-nearest outcomes from treatment group for each sample An improved neural network model for treatment effect estimation 5 Step 1: Set y(0) =0and y(1) =0 Step 2: for xi, i = 1,2, . . . , n do Step 3: xi=X[i, :] Step 4: Calculate the index set S0containing the indices of the k-nearest neighbors of xiwith T= 0 Step 5: y(0) i=1 kX j∈S0 yj Step 6: Calculate the index set S1containing the indices of the k-nearest neighbors of xiwith T= 1 Step 7: y(1) i=1 kX j∈S1 yj Based on the presented iterative process the average outcome from control and treated group is obtained and stored in y(0) and y(1), respectively. Notice that these will be used by the proposed model for the prediction of the conditional outcomes Q(t, x) = E(Y|X=x, T =t)and the propensity score g(x) = P(T= 1|X=x). 3.2 Modified Dragonnet architecture The proposed model consists of a modification of the state-of-the-art Dragonnet model [18]. The model takes as inputs the design matrix Xand the average outcomes from control and treated group, y(0) and y(1), respectively, while its three-headed architecture produces the predictions of propensity score ˆg(·)and conditional outcomes ˆ Q(0,·,·,·)and ˆ Q(1,·,·,·). Figure 1 presents a high-level architecture of the proposed modified Dragonnet model. Initially, a number of dense layers are utilized in order to produce a representation layer Z(X)∈Rp. Next, the output of Z(X)is concatenated with y(0) and the combined information is further processed by dense layers for the prediction of the outcome ˆ Q(0,·,·,·). Similarly, the output of Z(X)is concatenated with y(1) and through a number of dense layers the model provides the outcome ˆ Q(1,·,·,·,·). Additionally, the shared representation Z(X)is used for predicting ˆg(·), through the use of a simple linear map followed by a sigmoid activation function. The model is trained by minimizing the following loss function ˆ θ= arg min θ ˆ R(θ;X,y(0),y(1))(2) where θis the parameter vector and ˆ Ris defined by ˆ R(θ;X,y(0),y(1)) = 1 nX ih(Qnn(ti,xi, y(0) i, y(1) i;θ)−yi)2+αf(gnn(xi;θ), ti)i (3) 6 N. Kiriakidou and C. Diou where Qnn(ti,xi, y(0) i, y(1) i;θ)and gnn(xi;θ)are the output heads, fis the cross entropy function and α > 0is a hyperparameter used for weighting the two loss components. Fig. 1. Modified Dragonnet architecture Additionally, in order to increase the performance of the proposed model we utilized targeted regularization [18], which constitutes a modification to the loss function (2), by introducing a regularization term and an extra parameter. More specifically, the modified Dragonnet model is trained by minimizing the following loss ˆ θ, ˆ= arg min θ, "ˆ R(θ;X,y(0),y(1)) + β1 nX i γ(yi, ti,xi, y(0) i, y(1) i;θ, )#(4) where β,  are positive parameters, ˆ R(θ;X,y(0),y(1))is defined by Eq. (3) and the regularization term γ(yi, ti,xi, y(0) i, y(1) i;θ, )is defined by γ(yi, ti,xi, y(0) i, y(1) i;θ, )=(yi−˜ Q(ti,xi, y(0) i, y(1) i;θ, ))2 ˜ Q(ti,xi, y(0) i, y(1) i;θ, ) = Qnn(ti,xi, y(0) i, y(1) i;θ) + ti gnn(xi;θ)−1−ti 1−gnn(xi;θ) The rationale behind the loss function (4) is based on non-parametric estimation theory and consists on improving the model’s estimation of treatment effects. Additionally, under conditions, the following estimator of ψ ˆ ψtreg =1 n n X i=1 [ˆ Qtreg(1,xi, y(0) i, y(1) i)−ˆ Qtreg(0,xi, y(0) i, y(1) i)] where ˆ Qtreg =˜ Q(·,·,·,·;ˆ θ, ˆ), has the following properties [11] : 1. ˆ ψwill fast converge to ψeven in case ˆ Qand ˆgconverge slowly to Qand g. 2. asymptotically ˆ ψhas the lowest variance from any other considered estimator of ψ. An improved neural network model for treatment effect estimation 7 4 Data Considering that real-world data for causal inference are rarely available, we scarcely have access to the ground truth causal effects. Therefore, to overcome this problem we rely on semi-synthetic data for the empirical evaluation of causal estimation procedures. We used the semi-synthetic IHDP dataset introduced by Hill [8]. This dataset was constructed from the Infant Health and Development Program and the outcome and treatment assignment are fully known. It comprises 25 features regarding childs and mothers and 747 units, in which 139 belong to the treatment group and the rest 608 belong to the control group. In order to have comparable results, we used 1000 realizations from the NPCI package [2] similar to Shi et al. [18]. 5 Experimental results In this section, we evaluate the prediction performance of the proposed modified Dragonnet model against the state-of-the-art Dragonnet model. It is worth mentioning, that we selected to compare the proposed model against Dragonnet, since it outperforms all other state-of-the-art models. The performance of each model was measured using the metrics absolute error in ATE [17] |AT E|and expected Precision in Estimation of Heterogeneous Effect [8] P EHE , which are respectively defined by: |AT E|= 1 n n X i=1 [Q(1,xi)−Q(0,xi)] −ˆ ψtreg and P EHE =1 n n X i=1 h(Q(1,xi)−Q(0,xi)) −(ˆ Qtreg(1,xi, y(0) i, y(1) i)−ˆ Qtreg(0,xi, y(0) i, y(1) i))i2 It is worth noticing, that |AT E|and P EHE metrics are used to compare the evaluated models as estimators and predictors, respectively and have been also used in [10,13,17,22]. In our experiments, the state-of-the art model Dragonnet was used with its default optimized parameter settings [18], while the proposed model followed a similar architecture and hyper-parameter selection with Dragonnet. More specifically, we utilize three dense layers (of 200 neurons with Exponential Linear Unit (ELU) activation function) in order to produce a representation layer Z(X). Next, the output of Z(X)is concatenated with y(0) and the combined information is further processed by two dense layers (of 100 neurons each with ELU activation function and kernel regularizer of 10−2) for the prediction of the outcome of the control group. A similar approach was used for providing the outcome of the treated group. The hyperperameters were set as k= 10,α= 1 and β= 1 8 N. Kiriakidou and C. Diou and 20% of the training data were utilized for validation as in Dragonnet. Both evaluated models were trained using stochastic gradient descent with momentum [15]. The performance of the proposed modified Dragonnet utilizing three different distance metrics i.e Euclidean, Manhattan and Chebychev. These distances constitute the most widely used in the literature [14,19]. It is worth mentioning that these distances belong to the class of Minkowski distances, which is defined by kx−ykp= d X i |xi−yi|p! 1 p where x, y ∈Rdand p∈N∗. In case, p= 1,p= 2 and p=∞the Minkowski distance is reducted to the Manhattan, Euclidean and Chebychev metric, respectively. The detailed experimental results for each model and realization of IHDP can be found in https://github.com/kiriakidou/Modified_Dragonnet. The implementation code was written in Python 3.7 using Keras library [6] and run on a PC (3.2GHz Quad-Core processor, 16GB RAM) using Windows operating system. Given into consideration that a small number of simulations tend to dominate the benchmarking process, the cumulative total for a performance metric over all simulations seem to be too uninformative and misleading. For this reason, we used Dolan and Moré’s [1] performance profiles, which removes the influence of such simulations on the benchmarking process and provides us information such as probability of success, efficiency and robustness in compact form. In more detail, each profile plots the fraction Pof simulations for which any given model is within a factor τof the best model. Additionally, in order to examine and reject the hypothesis that both models perform equally and provide statistical evidence about the superiority of the proposed model, we utilize the methodology presented in [12]. More specifically, we apply the non-parametric Friedman Aligned-Ranks (FAR) test [9] in order to rank the models and the post-hoc Finner test [3] for examining the existence of significant differences. Next, we evaluate the performance of: –“Dragonnet”, which stands for Dragonnet model of Shalit et al. [18]. –“Modified Dragonnet (Euclidean)”, which stands for the proposed model using Euclidean distance for the calculation of the average of the outcomes of nearest instances. –“Modified Dragonnet (Manhattan)”, which stands for the proposed model using Manhattan distance for the calculation of the average of the outcomes of nearest instances. –“Modified Dragonnet (Chebychev)”, which stands for the proposed model using Chebychev distance for the calculation of the average of the outcomes of nearest instances. An improved neural network model for treatment effect estimation 9 Figure 2 presents the performance profiles of the three versions of the proposed model and the Dragonnet, based on |AT E|metric. Obviously, all compared models reported similar performance. More specifically, Modified Dragonnet (Euclidean) solves 30% of the simulations with the lowest error ATE, while both Dragonnet and Modified Dragonnet (Chebychev) solve 28%. Additionally, Modified Dragonnet (Manhattan) reported the worst performance solving 25% of simulations. Fig. 2. Performance profiles of all evaluated models based on |AT E | Figure 3 presents the performance profiles of the three versions of the proposed model and the Dragonnet, based on P EHE metric. The proposed model considerably outpeformed the state-of-the-art Dragonnet with any used distance metric, in terms of P EHE. All versions of Modified Dragonnet solve 34% of the simulations with the best (lowest) error PEHE, while Dragonnet solves only 8% of the simulations. Table 1 presents the statistical comparison between the three versions of the proposed model and the Dragonnet based on |AT E|metric. Clearly, Modified Dragonnet (Euclidean) reported the best performance, slightly outperforming all compared models. Additionally, it was the only version of the proposed model, which reported better FAR ranking than the state-of-the-art model Dragonnet. However, the interpretation of Finner post-hoc test illustrated that there are not considerable differences, which results that all models performed equally well. Table 2 presents that the proposed model considerably outpeformed the Dragonnet in terms of P EHE with every utilized distance metric, which is statistically confirmed by FAR and Finner tests. Modified Dragonnet (Euclidean) reported the best performance since it exhibited top ranking. However, Finner post-hoc test reveals that all versions of the model perform equally well and there are no significant statistical differences in their performances.