Title: GANITE: ESTIMATION OF INDIVIDUALIZED TREAT-MENT EFFECTS USING GENERATIVE ADVERSARIAL NETS

Abstract: Estimating individualized treatment effects (ITE) is a challenging task due to the need for an individual's potential outcomes to be learned from biased data and without having access to the counterfactuals. We propose a novel method for inferring ITE based on the Generative Adversarial Nets (GANs) framework. Our method, termed Generative Adversarial Nets for inference of Individualized Treatment Effects (GANITE), is motivated by the possibility that we can capture the uncertainty in the counterfactual distributions by attempting to learn them using a GAN. We generate proxies of the counterfactual outcomes using a counterfactual generator, G, and then pass these proxies to an ITE generator, I, in order to train it. By modeling both of these using the GAN framework, we are able to infer based on the factual data, while still accounting for the unseen counterfactuals. We test our method on three real-world datasets (with both binary and multiple treatments) and show that GANITE outperforms state-of-the-art methods.

1. INTRODUCTION:

Individualized treatment effects (ITE) estimation using observational data is a fundamental problem that is applicable in a wide variety of domains. For instance, (1) in understanding the heterogeneous effects of drugs (Shalit et al. (2017) ; Alaa & van der Schaar (2017); Alaa et al. ( 2017)); (2) in evaluating the effect of a policy on unemployment rates (LaLonde (1986) ; Smith & Todd (2005) );

(3) in verifying which factor causes a certain disease (Höfler (2005) ) and ( 4) estimating the effects of pollution on the weather (Hannart et al. (2016) ).

As explained in Spirtes (2009) , the problem of ITE estimation differs from the standard supervised learning problem. First, among the potential outcomes, only the factual outcome is actually observed (revealed), counterfactual outcomes are not observed and so the entire vector of potential outcomes can never be obtained. Second, unlike randomized controlled trials (RCT), observational studies are prone to treatment selection bias. For instance, left ventricular assist device (LVAD) treatment is mostly applied to high-risk patients with severe cardiovascular diseases before heart transplantation, the distribution of features among these patients will be significantly different to the distribution among non-LVAD treated patients (Kirklin et al. (2010) ). The sample distribution can vary drastically across different choices of treatments and therefore, if we were to apply a supervised learning framework for each treatment separately, the learned models would not generalize well to the entire population.

Classical works in this domain solved the problem of estimating the average treatment effects from observational data (Dehejia & Wahba (2002b) ; Lunceford & Davidian (2004) ). These works account for the selection bias using propensity scores (the estimated probability of receiving a treatment) to create unbiased estimators of the average treatment effect. Dehejia & Wahba (2002b) used a one-toone matching methodology to pair treated and control patients with similar features while Lunceford & Davidian (2004) used propensity scoring weighing to account for the selection bias. More recent works focus on individualized treatment effects (Chipman et al. (2010) ; Wager & Athey (2017) ; Athey & Imbens (2016) ; Lu et al. (2017) ; Alaa & van der Schaar (2017); Porter et al. (2011) ; Johansson et al. (2016) ; Alaa et al. (2017) ; Louizos et al. (2017) ; Shalit et al. (2017) ). Detailed qualitative comparisons to these works will be discussed in the next subsection and numerical comparisons can be found in Section 5.

In this paper, we propose a novel approach that attempts to not only fit a model to the observed factual data, but also account for the unseen counterfactual outcomes. We view the factual outcome as an observed label and consider the counterfactual outcomes to be missing labels. Missing labels are generated by the well-known Generative Adversarial Nets (GAN) framework (Goodfellow et al. (2014) ). More specifically, the counterfactual generator of GANITE attempts to generate counterfactual outcomes in such a way that when given the combined vector of factual and generated counterfactual outcomes the discriminator of GANITE cannot determine which of the components is the factual outcome. With the complete labels (combined factual and estimated counterfactual outcomes), the ITE estimation function can then be trained for inferring the potential outcomes of the individual based on the feature information in a supervised way. By also modelling this ITE estimation function using a GAN framework, we are able not only to predict the expected outcomes but also provide confidence intervals for the predictions, which is very important in, for example, the medical setting.

Unlike many other state-of-the-art methods, our method naturally extends to -and in fact is defined in the first place for -any number of treatments. We conduct experiments with three real-world observational datasets (with both binary and multiple treatments), and GANITE outperforms stateof-the-art methods.

1.1. RELATED WORKS:

Previous works on ITE estimation can be divided into three categories. In the first, a separate model is learned for each treatment; this approach does not account for selection bias and so each model learned will be biased toward the distribution of that treatment's population. In the second, the treatment is considered a feature, with one model learned for everything, and the mismatch between the entire sample distribution and treated and control distributions is adjusted in order to account for selection bias. 2017) used deep learning approaches to solve the ITE problem under this one model methodology. Inherent to the approach of learning a balanced representation is that the representation must trade off between containing predictive information and reducing biased information. This is because often it will be the case that information that is biased is also highly predictive (in fact in the medical setting this is precisely why it is biased -because the doctors will assign treatments based on predictive features). On the other hand, our framework is not forced to make this information trade-off -the dataset we learn our final ITE estimator on contains the original dataset, and so contains at least as much information as that one. In the experiment section, we show that our proposed framework outperforms Shalit et al. (2017) , particularly when the bias is high. In the third category, Alaa & van der Schaar (2017); Alaa et al. ( 2017) used a multi-task model approach. Alaa et al. (2017) used multi-task neural nets to estimate (1) the selection bias, (2) the controlled outcome and (3) the treated outcome with shared layers across these three tasks. Alaa & van der Schaar (2017) used a Gaussian Process approach in the multi-task model setting. Our work is perhaps most similar to Alaa & van der Schaar (2017) since there, too, they attempted to account for the counterfactuals and were similarly able to provide confidence in their estimates using credible intervals. They were able to access counterfactuals through a posterior distribution which was then accounted for in the learning of their model.

2. PROBLEM FORMULATION: ESTIMATION OF INDIVIDUALIZED TREATMENT EFFECTS:

Let X denote the s-dimensional feature space and Y the set of possible outcomes. Consider a joint distribution, µ, on X × {0, 1} k × Y k where k is the number of possible treatments. Suppose that (X, T, Y) ∼ µ. We call X ∈ X the (s-dimensional) feature vector, T ≡ (T 1 , ..., T k ) ∈ {0, 1} k the treatment vector and Y ≡ (Y 1 , ..., Y T ) ∈ Y k the vector of potential outcomes (or the Individualized Treatment Effects (ITE)). We assume that (with probability 1), there is precisely one non-zero component of T and we denote by η the index of this component. Denote by µ X the marginal distribution of X and by µ Y (x) the conditional distribution of Y given X = x, for x ∈ X (marginalized over T). This setting is known as the Rubin-Neyman causal model (Rubin ( 2005)).

We introduce two1 assumptions about the distribution µ in the Rubin-Neyman causal model. Assumption 1. (Overlap) For all x ∈ X , for all i ∈ {1, ..., k},

0 < P(T i = 1|X = x) < 1.

This assumption ensures that at every point in the feature space, there is a non-zero probability of being given treatment i for every i. Assumption 2. (Unconfoundedness) Conditional on X, the potential outcomes, Y, are indepen-

dent of T, Y |= T|X.

This assumption is also referred to as no unmeasured confounding and requires that all joint influences on Y and T are measured. Note that this assumption means that µ Y (x) no longer needs to be marginalized over T, since, under this assumption, they are independent.

Assume now that we observe samples of (X, T, Y η ) (whose joint distribution we denote by µ f ), so that our dataset, D, is given by D

= (x(n), t(n), y η(n) (n)) N n=1

. Importantly, we only observe the component of the potential outcome vector that corresponds to the assigned treatment, we call this the factual outcome, and refer to unobserved potential outcomes as counterfactual outcomes or just counterfactuals. We denote by y f (n) and y cf (n) the factual outcome and (vector of) counterfactual outcome(s), respectively. From this point forward, we omit the dependence on n for ease of notation.

In this setting, we wish to be able to draw samples from µ Y (x) for any x ∈ X . We measure the performance of the generator, I(x), using two different metrics depending on whether k = 2 (i.e. binary treatments) or k > 2 (i.e. multiple treatments).

For k = 2 we use the expected Precision in Estimation of Heterogeneous Effects, P EHE , introduced in Hill (2011), given by:

EQUATION

)

For k > 2 we use the expected mean squared error:

M SE = E x∼µ X ||E y∼µ Y (x) y -E ŷ∼I(x) [ŷ]|| 2 2 (2)

where

|| • || 2 is the standard 2 -norm in R k .

In order to achieve this goal, we separate the problem into two parts. First, we attempt to generate proxies for the unobserved counterfactual outcomes using a counterfactual generator (G) to create a complete dataset. Then, using this proxy dataset, we learn the ITE generator, I.

3. GANITE: GENERATIVE ADVERSARIAL NETS FOR INFERENCE OF INDIVIDUALIZED TREATMENT EFFECT ESTIMATION:



3.1. OVERVIEW:

The objective of GANITE is to generate potential outcomes for a given feature vector x. However, due to the lack of counterfactual outcomes we are unable to learn the distribution of potential outcomes directly. To account for these counterfactuals, we first attempt to generate samples, ỹcf , using a counterfactual generator, G, from the distribution µ Y cf (x, t, y f ) (the conditional distribution of the counterfactual outcomes, Y cf given that X = x, T = t and Y η = y f ) for each sample in our dataset. We can then combine these proxy counterfactuals with the original dataset to obtain a com-

plete dataset D = {x(n), t(n), ỹ(n)} N n=1

where ỹ is the combination of y f and ỹcf (with ỹη = y f ). The ITE generator, I, can then be optimized using D.

We follow a conditional GAN framework similar to the one set out in Mirza & Osindero (2014) to model the latter of these generators. For the former, we have to use a different discriminator in order to capture the same idea. More specifically, GANITE consists of two blocks: a counterfactual imputation block and an ITE block, each of which consists of a generator and a discriminator. We describe each of these blocks and their components in more detail in the following subsection.

3.2. A DETAILED BREAKDOWN:

Counterfactual generator (G): The counterfactual generator, G, uses the feature vector, x, the treatment vector, t, and the factual outcome, y f , to generate a potential outcome vector, ỹ. We let g be a function g :

EQUATION

The goal now is to find a function, g, such that G(x, t, y f ) ∼ µ Y (x, t, y f ). We write ỹ to denote a sample of G and ȳ to denote the vector obtained by replacing ỹη with y f . Observe that the ηth component of a sample from µ Y (x, t, y f ) will be y f , since we are sampling Y conditional on Y η = y f .

Counterfactual discriminator (D G )::

We introduce a discriminator, D G , which maps pairs (x, ȳ) to vectors in [0, 1] k with the i-th component, written D G (x, ỹ) i , representing the probability that the i-th component of ỹ is the factual outcome, equivalently the probability that η = i. This is in contrast to the standard GAN framework in which the discriminator is given a single sample from one of two distributions and it attempts to determine which distribution it came from. Here the discriminator is given a sample consisting of components from two different distributions and attempts to determine which components came from which distribution.

We train D G to maximize the probability of correctly identifying η. We then train G to maximize the probability of D G incorrectly identifying η (equivalently we try to minimize the probability of a correct identification -this is the adversarial method of learning between G and D G ).

Following the framework in Goodfellow et al. (2014) , we note that this formulation is captured by a minimax problem given by

min G max D G E (x,t,y f )∼µ f E z G ∼U ((-1,1) k ) t T log D G (x, ỹ) + (1 -t) T log(1 -D G (x, ỹ)) (4)

where log is performed element-wise and T denotes the transpose operator.

After training the counterfactual generator, we use it to generate the dataset D and pass this dataset to the ITE block.

ITE generator (I)::

The ITE generator, I, uses only the feature vector, x, to generate a potential outcome vector, ŷ. Similar to our approach with G, let h be a function h

: X × [-1, 1] k → Y k and z I ∼ U((-1, 1) k ).

We define the random variable I(x) as

I(x) = h(x, z I ) (5)

and similarly, the goal is to find a function, h, such that I(x) ∼ µ Y (x). We write ŷ to denote a sample from I(x).

ITE discriminator (D I ): Again, we introduce a discriminator, D I , but this time, since we have access to a complete dataset D, we can use a standard conditional GAN discriminator -it takes a pair (x, y * ) and returns a scalar corresponding to the probability that y * was from the data D (rather than drawn from I). Again, we train the generator and discriminator in an adversarial fashion using the following minimax criteria

EQUATION

where again log is taken element-wise.

4. GANITE: OPTIMIZATION:

In this section, we describe the empirical loss functions that are used to optimize each component of GANITE. The Pseudo-code is summarized in the Appendix.

4.1. COUNTERFACTUAL BLOCK (G, D G )::

Based on equation 4, the empirical objective of the minimax problem for G and D G can be defined by

V CF (x(n), t(n), ȳ(n)) = t(n) T log(D G (x(n), ȳ(n))) + (1 -t(n)) T log(1 -D G (x(n), ȳ(n))).

We also introduce the following 'supervised' loss in order to enforce the restriction that g η should be equal to

y f . L G S (y f (n), ỹη(n) (n)) = (y f (n) -ỹη(n) (n)) 2

More specifically, due to the structure of G, it outputs a full vector of potential outcomes, and so it not only outputs counterfactuals, but also gives a value for the one factual that was used as input. We account for this by using L G S to force the generated factual outcome to be close to the actually observed factual outcome. This is because, as noted above, conditional on observing y f , the component of y corresponding to y f should clearly be equal to y f . With the above two objective functions, G and D G are iteratively optimized with k G minibatches as follows:

min D G - k G n=1 V CF (x(n), t(n), ȳ(n)) min G k G n=1 V CF (x(n), t(n), ȳ(n)) + αL G S (y f (n), ỹη(n) (n))

where α ≥ 0 is a hyper-parameter.

4.2. ITE BLOCK (I, D I )::

After training the counterfactual block (G, D G ), GANITE optimizes the ITE block (I, D I ). Based on equation 6, the empirical objective of the minimax problem for I and D I can be defined by

V IT E (x(n), ȳ(n), ŷ(n)) = log(D I (x(n), ȳ(n))) + log(1 -D I (x(n), ŷ(n))).

Furthermore, in order to optimize the performance with respect to equations 1 and 2, we additionally introduce supervised losses (for the respective cases of k = 2 (binary treatments) and k > 2 (multiple treatments)) that are defined as follows:

(k = 2) : L I S (ȳ(n), ŷ(n)) = ((ȳ 1 (n) -ȳ0 (n)) -(ŷ 1 (n) -ŷ0 (n))) 2 (k > 2) : L I S (ȳ(n), ŷ(n)) = ||ȳ(n) -ŷ(n)|| 2

2 . I and D I are then iteratively optimized with k I minibatches as follows:

min D I - k I n=1 V IT E (x(n), ȳ(n), ŷ(n)) min I k G n=1 V IT E (x(n), ȳ(n), ŷ(n)) + βL I S (ȳ(n), ŷ(n))

where β ≥ 0 is a hyper-parameter.

Empirical justification for the inclusion of all the above losses can be found in Section 5.3 where we explore the effect of training with and without each of the losses, see Table 6 . We demonstrate there that using a combination of both (for both G and I) gives the best performance. In addition to this, by using a GAN loss for I we learn the conditional distribution of the potential outcomes rather than just the expectations (which would be the case if we only used the supervised loss L I S . See Table 1 in Section 5.3). This allows us to capture the uncertainty of the outcomes, which is very important in the medical setting when treatment decisions need to be made by doctors on the basis of these types of estimations.

Due to the lack of ground truth, it is often difficult in causal inference tasks to optimize the hyperparameters. More specifically, we do not have access to the true loss function (either PEHE or MSE) that we are trying to minimize, and so it is not possible to select hyper-parameters that minimise the true loss. One of the advantages of GANITE, however, is that our target loss can be estimated from the generated counterfactuals, unlike other methods such as in Shalit et al. (2017) . Therefore, we can directly optimize the hyper-parameters that minimize this estimated PEHE/MSE over the hyper-parameter space -details of our hyper-parameter optimization and the achieved optimal hyperparameters are illustrated in the Appendix.

5. EXPERIMENTS:



5.1. DATASETS:

Due to the nature of the problem, it is very difficult to evaluate the performance of the algorithm on real-world datasets -we never have access to the ground truth. Previous works, such as Shalit et al. (2017) ; Louizos et al. (2017) , use both semi-synthetic datasets (either the treatments or the potential outcomes are synthesized) and datasets collected from randomized controlled trials (RCT) to evaluate the ITE generator. We use two semi-synthetic datasets, IHDP and Twins, and one realworld dataset, Jobs, to evaluate the performance of GANITE with various state-of-the-art methods. These datasets are the same as the ones used in Shalit et al. (2017) ; Louizos et al. (2017) . Below, we give a detailed explanation of Twins. The details of IHDP and Jobs are well described in Shalit et al. (2017) ; Hill (2011); Dehejia & Wahba (2002a) and the Appendix.

Twins: This dataset is derived from all births in the USA between 1989-1991 (Almond et al. (2005) ). Among these births, we only focus on the twins. We define the treatment t = 1 as being the heavier twin (and t = 0 as being the lighter twin). The outcome is defined as the 1-year mortality. For each twin-pair we obtained 30 features relating to the parents, the pregnancy and the birth: marital status; race; residence; number of previous births; pregnancy risk factors; quality of care during pregnancy; and number of gestation weeks prior to birth. We only chose twins weighing less than 2kg and without missing features (list-wise deletion). This creates a complete dataset (without missing data). The final cohort is 11,400 pairs of twins whose mortality rate for the lighter twin is 17.7%, and for the heavier 16.1%. In this setting, for each twin pair we observed both the case t = 0 (lighter twin) and t = 1 (heavier twin); thus, the ground truth of individualized treatment effect is known in this dataset. In order to simulate an observational study, we selectively observe one of the two twins using the feature information (creating selection bias) as follows: t|x ∼ Bern(Sigmoid(w T x + n)) where w T ∼ U((-0.1, 0.1) 30×1 ) and n ∼ N (0, 0.1).

5.2. PERFORMANCE METRICS AND SETTINGS:

We use four different performance metrics: expected Precision in Estimation of Heterogeneous Effect (PEHE), average treatment effect (ATE) (Hill ( 2011)), policy risk (R pol (π)), and average treatment effect on the treated (ATT) (Shalit et al. (2017) ). In this subsection, we only provide definitions for PEHE and R pol (π). ATE and ATT are explained (and reported) in the Appendix.

If both factual and counterfactual outcomes are generated from a known distribution (so that we are able to compute the expectations of the outcomes, like in the IHDP dataset), and the treatment is binary, the empirical PEHE ( P EHE ) can be defined as follows:

P EHE = 1 N N n=1 E (y1(n),y0(n))∼µ Y (x(n)) [y 1 (n) -y 0 (n)] -[ŷ 1 (n) -ŷ0 (n)] 2

where y 1 (n), y 0 (n) are treated and controlled outcomes drawn from the ground truth (µ Y (x)) and ŷ1 (n), ŷ0 (n) are their estimations.

If both factual and counterfactual outcomes are observed but the underlying distribution is unknown (like in the Twins dataset), ˆ P EHE can be defined as follows:

ˆ P EHE = 1 N N n=1 [y 1 (n) -y 0 (n)] -[ŷ 1 (n) -ŷ0 (n)] 2

If only factual outcomes are available but the testing set comes from a randomized controlled trial (RCT), such as in the Jobs dataset, Policy risk (R pol (π)) can be defined as follows (Shalit et al. (2017) ):

R pol (π) = 1 N N n=1 1 - k i=1 1 |Π i ∩ T i ∩ E| x(n)∈Πi∩Ti∩E y i (n) × |Π i ∩ E| |E| where Π i = {x(n) : i = arg max ŷ}, T i = {x(n) : t i (n) = 1}

, and E is the subset of RCT.

Each dataset is divided 56/24/20% into training/validation/testing sets. Hyper-parameters such as the number of hidden layers α and β are chosen using Random Search (Bergstra & Bengio (2012) ). Details about the hyper-parameters are discussed in the Appendix. We run each algorithm 100 times (except for the IHDP dataset, on which we run each algorithm 1,000 times which is the same setting in Shalit et al. ( 2017)) with new training/validation/testing splits and report the mean and standard deviation of the performances.

5.3. EXPERIMENTAL RESULTS:

In our first simulation we focus on demonstrating the effect that including each of the losses introduced in Section 4 has on the performance of the algorithm. As is demonstrated below, inclusion of all four losses gives the best results.

We generate a synthetic dataset as follows: we draw 10,000 10-dimensional feature vectors x ∼ N (0 10×1 , 0.5×(Σ+Σ T )) where Σ ∼ U((-1, 1) 10×10 ). The treatment assignment is then generated as t|x ∼ Bern(Sigmoid(w T t x + n t )) where w T t ∼ U((-0.1, 0.1) 10×1 ) and n t ∼ N (0, 0.1). The potential outcome vector is then generated as y|x ∼ (w T y x + n y )) where w T y ∼ U((-1, 1) 10×2 ) and n y ∼ N (0 2×1 , 0.1 × I 2×2 ). We use 8,000 instances for training and 2,000 instances for testing. We repeat this 100 times and report the average P EHE on the testing set. Table 1 shows the performance of the GANITE architecture using different combinations of the four losses in Section 4. For each generator component, we can use 3 different combinations of loss:

(1) Supervised loss (S loss) only (this reduces the corresponding component to a standard neural network), (2) GAN loss only, (3) both S loss and GAN loss. The top-left most entry corresponds to simply using a standard neural network to first impute the counterfactuals and then using another standard neural network to learn an ITE estimator from the imputed dataset. As can be seen, this already performs well, but by adding the GAN losses for both the imputation step and the estimation step, a significant gain is shown (15.6%) (the bottom-right entry).

In Fig. 2 we show that GANITE is robust to an increased selection bias. We generate 10,000 10-dimensional treated samples from x 1 ∼ N (µ 1 , 0.5 × (Σ + Σ T )) and controlled samples from x 0 ∼ N (µ 0 , 0.5 × (Σ + Σ T )) where Σ ∼ U((-1, 1) 10×10 ). Fixing µ 0 and varying µ 1 , we generate various datasets with different Kullback-Leibler divergences (KL divergences) of µ 1 with respect to µ 0 . A higher KL divergence indicates a higher selection bias (a larger mismatch) between treated and controlled distributions. As seen in Fig. 2 , GANITE robustly outperforms state-of-the-art methods such as Shalit et al. (2017) ; Alaa & van der Schaar (2017) across the entire range of tested divergences.

Binary treatments: In this section, we evaluate GANITE for estimating individualized treatment effects for binary treatments. We use three datasets and report the P EHE both in-sample and outof-sample (for ATE and ATT see the appendix). We compare GANITE with least squares regression using treatment as a feature (OLS/LR 1 ), separate least squares regressions for each treatment (OLS/LR 2 ), balancing linear regression (BLR) (Johansson et al. (2016) As can be seen in Table 2 , GANITE achieves significant performance gains on the Twins and Jobs datasets in comparison with state-of-the-art methods (both in-sample and out-of-sample2 ). GANITE achieves a much higher gain for individualized treatment effect estimations (such as ˆ P EHE and R pol (π)) than average treatment effect estimations (such as ˆ AT E and AT T ). On IHDP, GANITE is competitive with BART and BNN but is outperformed by TARNET, CFR W ASS and CMGP. We believe this is due to the fact that GANITE has a large number of parameters to be optimized and IHDP is a relatively small dataset (747 samples). This belief is backed up by our significant gains over these methods in both Twins and Jobs, where the number of samples is much larger, (11400 and 3212 samples, respectively). Multiple treatments: GANITE is naturally defined for estimating multiple treatment effects. In this subsection, we further preprocess the Twins data to create a dataset containing multiple treatments.

Methods:

The multiple treatments are determined as follows: (1) t = 1: lower weight, female sex, (2) t = 2: lower weight, male sex, (3) t = 3: higher weight, female sex, (4) t = 4: higher weight, male sex. Therefore, we have 4 possible treatments for each sample.

We use the mean-squared error:

MSE y = 1 N × |T i | N i=1 t∈Ti y t (x i ) -ŷt (x i ) 2

as the performance metric to evaluate the multiple treatment effects (other metrics, such as PEHE, do not have a natural extension to the multiple treatments setting). For comparison with state-of-the-art methods, we naively extend BLR, C Forest, BNN, TARNET, CFR W ASS , and CMGP for multiple treatments: one of the four treatments is selected as a control treatment and then the remaining three create three separate binary ITE estimation problems (all against the same chosen control treatment).

As can be seen in Table 3 compared with Table 2 GANITE significantly outperforms other state-ofthe-art methods such as TARNET and CFR W ASS (17.7% and 19.1% gains in terms of out sample MSE, respectively). This is because, GANITE is designed for multiple treatments; the model is jointly trained for all treatments. On the other hand, other methods are designed for binary treatments and only naively extend to multiple treatments by training pairs of the available treatments.

5.4. DISCUSSION:

The experimental results provide various intuitions of the GANITE framework for ITE estimation. First, GANITE can be easily extended to any number of treatments and performs well in this multiple treatment setting. As can be seen in Section 4, however, a different loss for binary treatment and multiple treatments must be used because PEHE is only defined for the binary treatment settingthe MSE is not a natural generalisation of the PEHE and we believe exploration of possible loss functions in this setting would be an interesting future work.

A further extension to this problem would be to consider a setting in which a patient may receive several treatments (rather than just one). While this work can handle this problem naively (by treating each combination of treatments as a separate 'treatment') we believe this would also be an interesting problem to explore in a future work.

6. CONCLUSION:

In this paper we introduced a novel method for dealing with the ITE estimation problem. We have shown empirically that our method is more robust to large selection biases and performs better on standard benchmark datasets than other state-of-the-art methods. Our method also achieves significant performance gains over state-of-the-art when estimating ITE for multiple treatments because it is able to jointly estimate the representations across the multiple treatments. In this subsection, we use two different performance metrics for average treatment effect (ATE) estimation: average treatment effect (ATE) (Hill ( 2011)), and average treatment effect on the treated (ATT) (Shalit et al. (2017) ).

If both factual and counterfactual outcomes are generated from a known distribution (and so we are able to compute the expected value, such as in IHDP), the error of ATE ( AT E ) is defined as:

AT E = || 1 N N n=1 E y(n)∼µ Y (x(n)) [y(n)] - 1 N n i=1 ŷ(n)|| 2 2

where ŷ is the estimated potential outcome.

If both factual and counterfactual outcomes are observed but the underlying distribution is unknown (like in Twins), ˆ AT E is defined as:

ˆ AT E = || 1 N N n=1 y(n) - 1 N N n=1 ŷ(n)|| 2 2

If only factual outcomes are available (such as in Jobs), treatment is binary, and the testing set comes from a randomized controlled trial (RCT), the true average treatment effect on the treated (ATT) is defined as follows (Shalit et al. (2017) ):

AT T = 1 |T 1 ∩ E| xi∈T1∩E Y 1 (x i ) - 1 |T 0 ∩ E| xi∈C∩E Y 0 (x i ) AT T = |AT T - 1 |T 1 ∩ E| xi∈T1∩E Ŷ1 (x i ) -Ŷ0 (x i )|

where T 1 is the subset corresponding to treated samples, T 0 is the subset corresponding to controlled samples, and E is the subset corresponding to the randomized controlled trials.

Table. 5 shows the performance of the various algorithms with respect to these metrics.

As can be seen in Table 5 , the GANITE achieves competitive performances for Average Treatment Effect (ATE) estimation but not the best model to estimate ATE (except the Jobs dataset). However, we do not believe that it is an important metric for distinguishing models where the task is predicting treatment effects on an individual level. The problem we address with GANITE is to estimate the ITE. We used the ATE performance as a sanity check for our method -and believe it passes the sanity check, being competitive with most other methods.

HYPER-PARAMETER OPTIMIZATION:

We optimize our hyper-parameters in GANITE by estimating the PEHE (MSE in the case of multiple treatments) on the dataset generated by G and minimizing this with respect to the hyper-parameters. The table below indicates specifics of this process, including the values we search over. and what we want to compute is the conditional probability P(y|x). Therefore, zero out the impact of t and y f can be done as follows.

P(y|x) = P(y|x, t, y f )P (y f , t|x)dtdy f

The P(y|x, t, y f ) is learned by G and P (y f , t|x) can be easily learned using supervised learning framework (all the labels (y f , t) are available) with drop-out approach (to approximate the integral as the sample mean of multiple samples). We use multi-layer perceptron (MLP) with multiple outputs to learn the function P (y f , t|x). We called this as zero-out GANITE. We compare the performance of zero-out GANITE to original GANITE in Table 9 . As can be seen in Table 9 , the performance of original GANITE is marginally better than zero-out GANITE in three different datasets. We do not believe that zero-out GANITE would be any simpler than our proposed structure -in both cases we would still have 2 learning stages ( 

:

All works cited in the Related Works section make this assumption.

Note that both "in-sample" and "out-sample" experiments evaluate the performance of I networks.

Figures and Tables:

For instance, Chipman et al. (2010); Wager & Athey (2017); Athey & Imbens (2016); Lu et al. (2017) used tree-based models, Porter et al. (2011) used doubly-robust methods, Dehejia & Wahba (2002b); Lunceford & Davidian (2004), k-nearest neighbor (kNN) Crump et al. (2008) used propensity and matching based methods, and Johansson et al. (2016); Shalit et al. (

Figure 2: Performance comparison between GANITE and state-of-the-art methods as the selection bias is varied (Kullback-Leibler divergence of treated with respect to controlled distributions)

Figure 3: (a) Underlying distribution of potential outcomes (Y), (b) Underlying distribution of treatment assignments P(T |X), (c) Training data (factual outcomes) sampled from distributions explained in (a) and (b), (d) Potential outcomes sampled from trained ITE generator (I).

Performance with various combinations of GANITE (G: Counterfactual generator, I: ITE generator, S loss: Supervised loss). Details of each cell are illustrated in the Appendix as Figures.

Table content: <table><tr><td>G</td></tr></table>

Performance of ITE estimation with three real-world datasets. Bold indicates the method with the best performance for each dataset.

Table content: <table><tr><td/><td/><td/><td/><td colspan="2">Datasets (Mean ± Std)</td><td/></tr><tr><td>Methods</td><td>IHDP (</td><td>√</td><td>P EHE )</td><td>Twins ( √</td><td>ˆ P EHE )</td><td colspan="2">Jobs (R pol (π))</td></tr><tr><td/><td colspan="3">In-sample Out-sample</td><td>In-sample</td><td>Out-sample</td><td colspan="2">In-sample Out-sample</td></tr><tr><td>GANITE</td><td>1.9 ± .4</td><td/><td>2.4 ± .4</td><td>.289 ± .005</td><td>.297 ± .016</td><td>.13 ± .01</td><td>.14 ± .01</td></tr><tr><td>OLS/LR 1 OLS/LR 2 BLR</td><td>5.8 ± .3  *  2.4 ± .1 5.8 ± .3  *</td><td/><td>5.8 ± .3  *  2.5 ± .1 5.8 ± .3  *</td><td colspan="3">.319 ± .001  *  .320 ± .002 .320 ± .003  *  .21 ± .00  *  .318 ± .007 .22 ± .00  *  .312 ± .003  *  .323 ± .018 .22 ± .01  *</td><td>.23 ± .02  *  .24 ± .01  *  .25 ± .02  *</td></tr><tr><td>k-NN</td><td>2.1 ± .1</td><td/><td>4.1 ± .2  *</td><td colspan="2">.333 ± .001  *  .345 ± .007  *</td><td>.02 ± .00</td><td>.26 ± .02  *</td></tr><tr><td>BART</td><td>2.1 ± .1</td><td/><td>2.3 ± .1</td><td>.347 ± .009  *</td><td colspan="2">.338 ± .016 .23 ± .00  *</td><td>.25 ± .02  *</td></tr><tr><td>R Forest</td><td>4.2 ± .2  *</td><td/><td>6.6 ± .3  *</td><td>.306 ± .002  *</td><td colspan="2">.321 ± .005 .23 ± .01  *</td><td>.28 ± .02  *</td></tr><tr><td>C Forest</td><td>3.8 ± .2  *</td><td/><td>3.8 ± .2  *</td><td>.366 ± .003  *</td><td colspan="2">.316 ± .011 .19 ± .00  *</td><td>.20 ± .02  *</td></tr><tr><td>BNN</td><td>2.2 ± .1</td><td/><td>2.1 ± .1</td><td>.325 ± .003  *</td><td colspan="2">.321 ± .018 .20 ± .01  *</td><td>.24 ± .02  *</td></tr><tr><td>TARNET</td><td>.88 ± .02</td><td/><td>.95 ± .02</td><td>.317 ± .005  *</td><td colspan="2">.315 ± .003 .17 ± .01  *</td><td>.21 ± .01  *</td></tr><tr><td colspan="2">CFR W ASS .71 ± .02 CMGP .65 ± .44</td><td/><td>.76 ± .02 .77 ± .11</td><td>.315 ± .007  *  .320 ± .002</td><td colspan="2">.313 ± .008 .17 ± .01  *</td><td>.21 ± .01  *</td></tr></table>

Metric: MSE y Performance of multiple treatment effects estimations using Twins data. Bold indicates the method with the best performance for each dataset.

Table content: <table><tr><td/><td>In Sample</td><td>Gain (%)</td><td>Out Sample</td><td>Gain (%)</td></tr><tr><td>GANITE</td><td>.0427 ± .0161</td><td>(-)</td><td>.0723 ± .0183</td><td>(-)</td></tr><tr><td>OLS/LR 1</td><td>.0855 ± .0096</td><td>50.1%</td><td>.0871 ± .0142</td><td>17.0%</td></tr><tr><td>OLS/LR 2 BLR</td><td>.0857 ± .0099 .0996 ± .0081  *</td><td>50.2% 57.1%</td><td>.0883 ± .0147 .1017 ± .0127</td><td>18.1% 28.9%</td></tr><tr><td>KNN</td><td>.0930 ± .0101  *</td><td>54.1%</td><td>.1008 ± .0236</td><td>28.3%</td></tr><tr><td>BART</td><td>.1097 ± .0084  *</td><td>61.1%</td><td>.1037 ± .0283</td><td>30.3%</td></tr><tr><td>R Forest</td><td>.0442 ± .0069</td><td>3.4%</td><td>.0927 ± .0138</td><td>22.0%</td></tr><tr><td>C Forest</td><td>.1607 ± .0014  *</td><td>73.4%</td><td>.1665 ± .0035  *</td><td>56.6%</td></tr><tr><td>BNN</td><td>.0602 ± .0102</td><td>29.1%</td><td>.1031 ± .0145</td><td>29.9%</td></tr><tr><td>TARNET</td><td>.0854 ± .0091</td><td>50.0%</td><td>.0879 ± .0030</td><td>17.7%</td></tr><tr><td colspan="2">CFR W ASS .0896 ± .0036  *  CMGP .0844 ± .0073  *</td><td>52.3% 49.4%</td><td>.0894 ± .0057 .0793 ± .0191</td><td>19.1% 8.3%</td></tr></table>

Performance of average treatment effect estimation. Bold represents the best performance.

Table content: <table><tr><td/><td/><td/><td colspan="2">Datasets (Mean ± Std)</td><td/><td/></tr><tr><td/><td colspan="2">IHDP ( AT E )</td><td colspan="2">Twins (ˆ AT E )</td><td colspan="2">Jobs ( AT T )</td></tr><tr><td/><td colspan="2">In-sample Out-sample</td><td>In-sample</td><td>Out-sample</td><td colspan="2">In-sample Out-sample</td></tr><tr><td>GANITE</td><td>.43 ± .05</td><td>.49 ± .05</td><td>.0058 ± .0017</td><td>.0089± 0.0075</td><td>.01 ± .01</td><td>.06 ± .03</td></tr><tr><td>OLS/LR 1</td><td>.73 ± .04  *</td><td>.94 ± .06  *</td><td>.0038 ± .0025</td><td>.0069 ± .0056</td><td>.01 ± .00</td><td>.08 ± .04</td></tr><tr><td>OLS/LR 2 BLR</td><td>.14 ± .01 .72 ± .04  *</td><td>.31 ± .02 .93 ± .05  *</td><td colspan="2">.0039 ± .0025 .0057 ± .0036 .0334 ± .0092  *  .0070 ± .0059</td><td>.01 ± .01 .01 ± .01</td><td>.08 ± .03 .08 ± .03</td></tr><tr><td>k-NN</td><td>.14 ± .01</td><td>.90 ± .05  *</td><td>.0028 ± .0021</td><td colspan="2">.0051 ± .0039 .21 ± .01  *</td><td>.13 ± .05</td></tr><tr><td>BART</td><td>.23 ± .01</td><td>.34 ± .02</td><td colspan="2">.1206 ± .0236  *  .1265 ± .0234  *</td><td>.02 ± .00</td><td>.08 ± .03</td></tr><tr><td>R Forest</td><td>.73 ± .05  *</td><td>.96 ± .06  *</td><td>.0049 ± .0034</td><td>.0080 ± .0051</td><td>.03 ± .01</td><td>.09 ± .04</td></tr><tr><td>C Forest</td><td>.18 ± .01</td><td>.40 ± .03</td><td colspan="2">.0286 ± .0035  *  .0335 ± .0083  *</td><td>.03 ± .01</td><td>.07 ± .03</td></tr><tr><td>BNN</td><td>.37 ± .03</td><td>.42 ± .03</td><td>.0056 ± .0032</td><td>.0203 ± .0071</td><td>.04 ± .01</td><td>.09 ± .04</td></tr><tr><td>TARNET</td><td>.26 ± .01</td><td>.28 ± .01</td><td>.0108 ± .0017  *</td><td>.0151 ± .0018</td><td>.05 ± .02</td><td>.11 ± .04</td></tr><tr><td>CFR W ASS</td><td>.25 ± .01</td><td>.27 ± .01</td><td colspan="2">.0112 ± .0016  *  .0284 ± .0032  *</td><td>.04 ± .01</td><td>.09 ± .03</td></tr><tr><td>CMGP</td><td>.11 ± .10</td><td>.13 ± .12</td><td>.0124 ± .0051</td><td>.0143 ± .0116</td><td>.06 ± .06</td><td>.09 ± .07</td></tr></table>

Hyper-parameters of GANITE

Table content: <table><tr><td>Blocks</td><td>Sets of Hyper-parameters</td></tr><tr><td>Initialization</td><td>Xavier Initialization for Weight matrix, Zero initialization for bias vector.</td></tr><tr><td>Optimization</td><td>Adam Moment Optimization</td></tr><tr><td>Batch size (k G , k I )</td><td>{32, 64, 128, 256}</td></tr><tr><td>Depth of layers</td><td>{1, 3, 5, 7, 9}</td></tr><tr><td colspan="2">Hidden state dimension {s, int(s/2), int(s/3), int(s/4), int(s/5)}</td></tr><tr><td>α, β</td><td>{0, 0.1, 0.5, 1, 2, 5, 10}</td></tr></table>

As explained in Section 4, I tries to learn the potential distribution that consists of factual outcomes and counterfactual outcomes that are generated by G. To evaluate how well I is able to learn from G, we compare the performance of G and I of the GANITE framework in Table8in terms of ITE estimation. As can be seen in Table8, the in-sample ITE estimation performance of GANITE (I) is competitive with GANITE (G). It experimentally verifies that GANITE (I) learns well from the outputs of GANITE (G).

Table content: <table><tr><td/><td>: Optimal Hyper-parameters of GANITE</td></tr><tr><td>Dataset</td><td>Optimal Hyper-parameters</td></tr><tr><td colspan="2">IHDP k ADDITIONAL EXPERIMENTS</td></tr><tr><td colspan="2">PERFORMANCE GAP BETWEEN G AND I</td></tr></table>

Table 9 experimentally verifies this). Performance comparison of ITE estimation between GANITE and zero-out GANITE with three real-world datasets.

Table content: <table><tr><td/><td/><td/><td/><td colspan="2">Datasets (Mean ± Std)</td><td/></tr><tr><td>Methods</td><td>IHDP (</td><td>√</td><td>P EHE )</td><td>Twins ( √</td><td>ˆ P EHE )</td><td colspan="2">Jobs (R pol (π))</td></tr><tr><td/><td colspan="3">In-sample Out-sample</td><td>In-sample</td><td colspan="3">Out-sample In-sample Out-sample</td></tr><tr><td>GANITE</td><td>1.9 ± .4</td><td/><td>2.4 ± .4</td><td colspan="3">.289 ± .005 .297 ± .016 .13 ± .01</td><td>.14 ± .01</td></tr><tr><td>zero-out GANITE</td><td>2.2 ± .6</td><td/><td>2.6 ± .7</td><td colspan="3">.297 ± .008 .308 ± .022 .14 ± .02</td><td>.17 ± .01</td></tr></table>

Peer review (iclr2018):
 Summary: This paper proposes to estimate the individual treatment effects (ITE) through training two separate conditional generative adversarial networks (GANs). First, a counterfactual GAN is trained to estimate the conditional distribution of the potential outcome vector, which consists of factual outcome and all other counterfactual outcomes, given 1) the feature vector, 2) the treatment variable, and 3) the factual outcome. After training the counterfactual GAN, the complete dataset containing the observed potential outcome vector can be generated by sampling from its generator. Second, a ITE GAN is trained to estimate the conditional distribution of the potential outcome vector given only the feature vector. In this way, for any test sample, its potential outcomes can be estimated using the generator of the trained ITE GAN. Given its pontential outcome vector, its ITE can be estimated as well. Experimental results on the synthetic data shows the proposed approach, called GANITE, is more robust to the existence of selection bias, which is defined as the mismatch between the treated and controlled distributions, compared to its competing alternatives. Experiments on three real world datasets show GANITE achieves the best performance on two datasets, including Twins and Jobs. It does not perform very well on the IHDP dataset. The authors also run experiments on the Twins dataset to show the proposed approach can estimate the multiple treatment effects with better performance. Comments 1) This paper is well written. The background and related works are well organized. 2) To the best of my knowledge, this is the first work that applies GAN to ITE estimation. 3) Experiments on the synthetic data and the real-world data demonstrate the advantage of the proposed approach. 4) The authors directly present the formulation without providing sufficient motivations. Could the authors provide more details or intuitions on why GAN would improve the performance of ITE estimation compared to approaches that learn representations to minimize the distance between the distributions of different treatment groups, such as CFR_WASS? 5) As is pointed out by the authors, the proposed approach does not perform well when the dataset is small, such as the IHDP data. However, in practice, a lot of real-world datasets might have small sample size, such as the LaLonde dataset. Did the authors plan to extend the model to handle those small-sized data sets without completely changing the model. 6) When training the ITE GAN, the objective is to learn the conditional distribution of the potential outcome vector given the feature vector. Did the authors try the option of replacing ITE GAN with multi-task regression? Will the performance become worse using multi-task regression? I think this comparison would be a sanity check on the utility of using GAN instead of regression models for ITE estimation.