GANITE: E STIMATION OF I NDIVIDUALIZED T REATMENT E FFECTS USING G ENERATIVE A DVERSARIAL
N
ETS
Anonymous authors
Paper under double-blind review
A BSTRACT
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.
I NTRODUCTION
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.
; (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 ¨ oﬂer (2005)) and (4) estimating the effects of pollution on the weather (Hannart et al.
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 signiﬁcantly different to the distribution among non-LVAD treated patients (Kirklin et al.
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.
; Wager & Athey (2017);
Athey & Imbens (2016); Lu et al.
(2017); Alaa & van der Schaar (2017); Porter et al.
; Johansson et al.
; Alaa et al.
; Louizos et al.
; Shalit et al.
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 ﬁt 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.
More speciﬁcally, 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 conﬁdence 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 deﬁned in the ﬁrst 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. R ELATED W ORKS
Previous works on ITE estimation can be divided into three categories.
In the ﬁrst, 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.
For instance, Chipman et al.
; 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.
; Shalit et al.
(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 ﬁnal 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 conﬁdence 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.
P ROBLEM F ORMULATION : E STIMATION OF I NDIVIDUALIZED T REATMENT E FFECTS
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 two 1 assumptions about the distribution µ in the Rubin-Neyman causal model.
Assumption 1.
(Overlap)
For all x ∈ X , for all i ∈ { 1 , ..., k } ,
[Equation]
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 independent of T ,
[Equation]
This assumption is also referred to as no unmeasured confounding and requires that all joint inﬂuences 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 ))
Nn =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, (cid:15) PEHE , introduced in Hill (2011), given by:
[Equation]
For k > 2 we use the expected mean squared error:
[Equation]
where || · || 2 is the standard (cid:96) 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 .
GANITE: G ENERATIVE A DVERSARIAL N ETS FOR INFERENCE OF I NDIVIDUALIZED T REATMENT E FFECT ESTIMATION 3. O VERVIEW
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 ﬁrst attempt to generate samples, ˜ y 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 complete dataset ˜ D
= { x ( n ) , t ( n ) , ˜ y ( n ) } Nn =1 where ˜ y is the combination of y f and ˜ y cf (with ˜ y
η
= y f ).
The ITE generator, I , can then be optimized using ˜ D .
[Figure]
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 speciﬁcally, 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. A D ETAILED B REAKDOWN 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, ˜ y .
We let g be a function g : X × { 0 , 1 } k × Y × [ − 1 , 1] k − 1 → Y k and z G ∼ U (( − 1 , 1) k − 1 ) .
We then deﬁne the random variable G ( x , t , y f ) as
[Equation]
The goal now is to ﬁnd a function, g , such that G ( x , t , y f ) ∼ µ Y ( x , t , y f ) .
We write ˜ y to denote a sample of G and ¯ y to denote the vector obtained by replacing ˜ y
η 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 , ¯ y ) to vectors in [0 , 1] k with the i -th component, written D G ( x , ˜ y ) i , representing the probability that the i -th component of ˜ y 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 identiﬁcation - 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
[Equation]
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, ˆ y .
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 deﬁne the random variable
I ( x ) as
I ( x ) = h ( x , z I ) (5) and similarly, the goal is to ﬁnd a function, h , such that I ( x ) ∼ µ Y ( x ) .
We write ˆ y 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.
GANITE: O PTIMIZATION
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. C OUNTERFACTUAL BLOCK ( G , D G ):
Based on equation 4, the empirical objective of the minimax problem for G and D G can be deﬁned by
[Equation]
We also introduce the following ‘supervised’ loss in order to enforce the restriction that g η should be equal to y f .
[Equation]
More speciﬁcally, 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 GS 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:
[Equation]
where α ≥ 0 is a hyper-parameter.
4. 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 deﬁned by
[Equation]
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 >(multiple treatments)) that are deﬁned as follows:
[Equation]
I and D
I are then iteratively optimized with k I minibatches as follows:
[Equation]
where β ≥ 0 is a hyper-parameter.
Empirical justiﬁcation 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 IS .
See Tablein 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 difﬁcult in causal inference tasks to optimize the hyperparameters.
More speciﬁcally, 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.
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.
E XPERIMENTS 5. D ATASETS
Due to the nature of the problem, it is very difﬁcult 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.
; 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.
; Louizos et al.
Below, we give a detailed explanation of Twins.
The details of IHDP and Jobs are well described in Shalit et al.
; 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.
Among these births, we only focus on the twins.
We deﬁne the treatment t = 1 as being the heavier twin (and t = 0 as being the lighter twin).
The outcome is deﬁned 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 ﬁnal 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 =
(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) .
P ERFORMANCE M ETRICS AND S ETTINGS
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.
In this subsection, we only provide deﬁnitions 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 ( (cid:15) PEHE ) can be deﬁned as follows:
[Equation]
where y 1 ( n ) , y 0 ( n ) are treated and controlled outcomes drawn from the ground truth ( µ Y ( x ) ) and ˆ y 1 ( n ) , ˆ y 0 ( n ) are their estimations.
If both factual and counterfactual outcomes are observed but the underlying distribution is unknown (like in the Twins dataset), ˆ (cid:15) PEHE can be deﬁned as follows:
[Equation]
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 deﬁned as follows (Shalit et al.
[Equation]
where Π i = { x ( n ) : i = arg max ˆ y } , 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.
with new training/validation/testing splits and report the mean and standard deviation of the performances.
5. E XPERIMENTAL R ESULTS
In our ﬁrst 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 Tt x + n t )) where w Tt ∼ 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 Ty x + n y )) where w Ty ∼ 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 (cid:15) PEHE 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 ﬁrst impute the counterfactuals and then using another
[Table]
[Figure]
In Fig.
2 we show that GANITE is robust to an increased selection bias.
We generate 10, 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 (cid:15) PEHE 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)), k-nearest neighbor (k-NN) (Crump et al.
, Bayesian additive regression trees (BART) (Chipman et al.
(2010)), random forests (RForest) (Breiman (2001)), causal forests (C Forest) (Wager & Athey (2017)), balancing neural network (BNN) (Johansson et al.
(2016)), treatment-agnostic representation network (TARNET)
(Shalit et al.
(2017)), counterfactual regression with Wasserstein distance (CFR WASS )
(Shalit et al.
(2017)), and multi-task gaussian process (CMGP) (Alaa & van der
Schaar (2017)).
We evaluate both in-sample and out-of-sample performance in Table 2.
As can be seen in Table 2, GANITE achieves signiﬁcant performance gains on the Twins and Jobs datasets in comparison with state-of-the-art methods (both in-sample and out-of-sample 2 ).
GANITE [-
 [Reviewer 2 - Question 5]-]achieves a much higher gain for individualized treatment effect estimations (such as ˆ (cid:15) PEHE and
R pol ( π ) ) than average treatment effect estimations (such as ˆ (cid:15) ATE and (cid:15) ATT ). On IHDP, GANITE
[Table]
is competitive with BART and BNN but is outperformed by TARNET, CFR WASS 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 signiﬁcant gains over these methods in both Twins and Jobs, where the number of samples is much larger, ( and 3212 samples, respectively). Methods
[Table]
Multiple treatments: GANITE is naturally deﬁned for estimating multiple treatment effects.
In this subsection, we further preprocess the Twins data to create a dataset containing multiple treatments.
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:
[Equation]
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 WASS , 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 signiﬁcantly outperforms other state-ofthe-art methods such as TARNET and CFR WASS (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. D ISCUSSION
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 deﬁned 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.
C ONCLUSION
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 signiﬁcant 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.
A PPENDIX P SEUDO - CODE OF GANITE
Algorithm while convergence of training loss of G and D G do (1) Counterfactual block optimization
Use k C minibatches, iteratively optimize G , D G by stochastic gradient descent (SGD) min D G −
[Equation]
while convergence of training loss of I and D I do (2) ITE block optimization
Use k I minibatches, update I , D I by SGD min
D
I −
[Equation]
( n
[Equation]
( n ))
min I
[Equation]
G (cid:88) n
[Equation]
D ETAILED DESCRIPTION OF THE DATASETS IHDP
Hill (2011) provided a dataset for ITE estimation with the Infant Health and Development Program (IHDP).
The dataset consists of 747 children ( t = 1 : 139, t = 0 608) with 25 features.
We generated potential outcomes from setting A in the NPCI package Dorie (2016).
J OBS
Jobs data studied in LaLonde (1986) is composed of randomized data based on the National Supported Work program and non-randomized data from observational studies.
We use a (random) subset of the randomized data to evaluate the algorithms based on R pol ( π ) and (cid:15) ATT .
The dataset consists of 722 randomized samples ( t = 1 : 297, t = 0 : 425) and 2490 non-randomized samples ( t = 1 : 0, t = 0 : 2490), all with 7 features.
S UMMARY OF THE DATASETS
[Table] [Table]
P ERFORMANCE METRICS AND THE RESULTS OF AVERAGE TREATMENT EFFECT ESTIMATION
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.
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 ( (cid:15) ATE ) is deﬁned as:
[Equation]
where ˆ y is the estimated potential outcome.
If both factual and counterfactual outcomes are observed but the underlying distribution is unknown (like in Twins), ˆ (cid:15) ATE is deﬁned as:
[Equation]
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 deﬁned as follows (Shalit et al. [Equation]
[Equation]
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.
H YPER - 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 speciﬁcs of this process, including the values we search over.
[Table]
For the hyper-parameter optimization of the benchmarks [Table] we follow the hyper-parameter optimzation code published in the github with their main code.
For instance, the hyper-parameters of CFR WASS are optimized using cfr_param_search.py ﬁle which is published in https: //github.com/clinicalml/cfrnet O PTIMAL H YPER - PARAMETERS FOR EACH DATASET
[Table]
of layers: 7, h dim : 8, α : 1 , β :
A DDITIONAL E XPERIMENTS P ERFORMANCE GAP BETWEEN G AND I
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 Table 8 in terms of ITE estimation.
As can be seen in Table 8, the in-sample ITE estimation performance of GANITE (I) is competitive with GANITE (G).
It experimentally veriﬁes that GANITE (I) learns well from the outputs of GANITE (G).
Z ERO OUT THE CONTRIBUTION OF FACTUAL OUTCOME AND TREATMENT
In this section, we use the trained G to compute ITE only with x .
The learned function G needs x , t , and y f as the inputs.
Therefore, in order to compute ITE only with x , we should zero out the contribution of t and y f in the G function. G tries to learn the conditional probability P ( y | x, y f , t )
[Table]
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.
[Equation]
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 (Table 9 experimentally veriﬁes this) [Table]
Methods
[Table]
GANITE
[Table]
[Table] [Equation] [Table] [Equation]
[Figure]
ou
[Figure] [Figure]
[Figure]
Underlying potential outcome distributions
[Figure]
|
[Figure]
Underlying
[Figure]
T
[Figure]
Feature space (X)0.
[Figure]
a c t u r e ou t c o m e s ( Y f ) 0.
Given samples with factual outcomes
T=0T=
Feature space (X)
[Figure] [Figure]
ene
[Figure]
ed ou t
[Figure]
m e
[Figure]
I 0.
[Figure]
Generated potential outcomes by I
T=0T=
[Figure]
(c) (d)
A DDITIONAL FIGURES FOR T ABLE
[Figure]
Supervisedloss ( ℒ
[Figure]
[Table] [Table]
[Table]
(d)
[Table]
: S loss only , I: GAN loss only
[Table]
( 𝑰 ) 𝒙 𝒛 𝑰 𝒚 (cid:3549)
ITE Discriminator ( 𝑫 𝑰 )
GAN loss ( 𝑉 (cid:3010)(cid:3021)(cid:3006) ) 𝒙 𝒚(cid:3365) 𝒚(cid:3365) 𝒚(cid:3549)
[Table]
𝓓(cid:3561)
= 𝒙,𝑦 (cid:3033) ,𝒚(cid:3365) 𝒄𝒇 = 𝒙,𝒚(cid:3365)
[Table]
(e) G: GAN loss only , I: GAN loss only
[Table]
( 𝑰 ) 𝒙 𝒛 𝑰 𝒚(cid:3549)
[Table]
( 𝑫
[Table] [Table]
loss ( 𝑉 (cid:3010)(cid:3021)(cid:3006) ) 𝒙 𝒚(cid:3365) 𝒚(cid:3365) 𝒚(cid:3549)
[Table]
𝓓(cid:3561)
= 𝒙,𝑦 (cid:3033) ,𝒚(cid:3365) 𝒄𝒇 = 𝒙,𝒚(cid:3365)
[Table]
[Table] [Table]
Discriminator ( 𝑫
[Table] [Table]
[Table] [Table]
ITE Generator ( 𝑰 ) 𝒙 𝒛 𝑰 𝒚 (cid:3549)
ITE Discriminator ( 𝑫 𝑰 )
GAN loss ( 𝑉 (cid:3010)(cid:3021)(cid:3006) )
Supervisedloss ( ℒ (cid:3020)(cid:3010) )
𝒙 𝒚(cid:3365) 𝒚(cid:3365) 𝒚(cid:3549)
ITE block 𝓓(cid:3561)
= 𝒙,𝑦 (cid:3033) ,𝒚(cid:3365) 𝒄𝒇 = 𝒙,𝒚(cid:3365)
[Table]
Generator ( 𝑮 ) 𝒛 𝑮 𝑦 (cid:3033) 𝒕 𝒙 𝑦(cid:3556) (cid:3033) Supervisedloss ( ℒ (cid:3020)(cid:3008) ) 𝒚(cid:3557) 𝒄𝒇 𝑦 (cid:3033)
[Table]
(g) G: S loss only , I: S and GAN loss
[Table] [Equation]