P ERFECT M ATCH : A S IMPLE M ETHOD FOR L EARNING R EPRESENTATIONS F OR C OUNTERFACTUAL I NFERENCE W ITH N EURAL N ETWORKS
Anonymous authors
Paper under double-blind review
A BSTRACT
Learning representations for counterfactual inference from observational data is of high practical relevance for many domains, such as healthcare, public policy and economics.
Counterfactual inference enables one to answer ”What if...?”
questions, such as ”What would be the outcome if we gave this patient treatment t 1 ?”.
However, current methods for training neural networks for counterfactual inference on observational data are either overly complex, limited to settings with only two available treatment options, or both.
Here, we present Perfect Match (PM), a method for training neural networks for counterfactual inference that is easy to implement, compatible with any architecture, does not add computational complexity or hyperparameters, and extends to any number of treatments.
PM is based on the idea of augmenting samples within a minibatch with their propensity-matched nearest neighbours.
Our experiments demonstrate that PM outperforms a number of more complex state-of-the-art methods in inferring counterfactual outcomes across several real-world and semi-synthetic datasets.
I NTRODUCTION
Estimating individual treatment effects 1 (ITE) from observational data is an important problem in many domains.
In medicine, for example, we would be interested in using data of people that have been treated in the past to predict what medications would lead to better outcomes for new patients (Shalit et al.
Similarly, in economics, we would for example want to determine how effective job programs would be based on results of past job training programs (LaLonde (1986)).
ITE estimation from observational data is difﬁcult for two reasons: Firstly, we never observe all potential outcomes.
If a patient is given a treatment to treat her symptoms, we never observe what would have happened if the patient was prescribed a potential alternative treatment in the same situation.
Secondly, the assignment of cases to treatments is typically biased such that cases for which a given treatment is more effective are more likely to have received that treatment.
The distribution of samples may therefore differ signiﬁcantly between the treated group and the overall population.
A supervised model na¨ıvely trained to minimise the factual error would overﬁt to the properties of the treated group, and thus not generalise well to the entire population.
To address these problems, we introduce Perfect Match (PM), a simple method for training neural networks for counterfactual inference that extends to any number of treatment options.
PM effectively controls for biased assignment of treatments in observational data by augmenting every sample within a minibatch with its closest matches by propensity score from the other treatment options.
PM is easy to use with existing neural network architectures, simple to implement, and does not add any hyperparameters or computational complexity.
We perform experiments that demonstrate that PM outperforms a number of more complex state-of-the-art methods in inferring counterfactual outcomes on several real-world and semi-synthetic datasets.
Our experiments also show that
PM is more robust to a high level of treatment assignment bias than existing methods.
We believe that our simple and effective method for ITE estimation that extends to any number of treatment options could potentially make counterfactual inference more accessible, particularly for medical applications, where multiple available treatment options are the norm.
Contributions.
This work contains the following contributions:
• We introduce Perfect Match, a simple methodology for learning neural representations for counterfactual inference based on propensity score matching within minibatches.
• We perform extensive experiments on semi-synthetic and real-world data in both the binary and multiple treatment setting.
The experiments show that PM outperforms a number of more complex state-of-the-art methods in inferring counterfactual outcomes.
R ELATED W ORK
There are ﬁve main categories of approaches to learning to estimate ITEs:
Matching-based Methods.
Matching methods estimate the counterfactual outcome of a sample
X with respect to treatment t using the factual outcomes of its nearest neighbours having received t , with respect to a metric space.
k-Nearest-Neighbour (kNN) (Ho et al.
(2007)) operates in the potentially high-dimensional covariate space, and therefore might suffer from the curse of dimensionality.
Propensity Score Matching (PSM) (Rosenbaum & Rubin (1983)) matches on the scalar probability p ( t | X ) of t given the covariates X .
PSM may require under- or oversampling the original data.
Adjusted Regression Methods.
Adjusted regression models apply regression models with both treatment and covariates as input.
The simplest case is Ordinary Least Squares, which can either be used for building one model, with the treatment as an input feature, or multiple separate models, one for each treatment (Kallus (2017)).
More complex regression models, such as Treatment-Agnostic
Representation Networks (TARNET) (Shalit et al.
(2017)) may be used to capture non-linear relationships.
Methods that make use of propensity scores, and a model of the outcomes, such as
Propensity Dropout (PD) (Alaa et al.
(2017)), are referred to as doubly robust (Funk et al.
Tree-based Methods.
Tree-based methods train many weak learners to build expressive ensemble models.
Examples of tree-based methods are Bayesian Additive Regression Trees (BART)
(Chipman et al.
; Chipman & McCulloch (2016)) and Causal Forests (CF) (Wager & Athey (2017)).
Representation-balancing Methods.
Representation-balancing methods seek to learn a highlevel representation for which the covariate distributions are balanced across treatment groups.
Balancing Neural Networks (Johansson et al.
(2016)) attempt to ﬁnd such representations by minimising the discrepancy distance (Mansour et al.
(2009)) between treatment groups.
Counterfactual Regression Networks (CFRNET, Shalit et al. (2017)) use different metrics such as the Wasserstein distance.
Distribution-modeling Methods.
Generative Adversarial Nets for inference of Individualised
Treatment Effects (GANITE)
(Yoon et al.
address ITE estimation using counterfactual and
ITE generators.
GANITE uses a complex architecture with many hyperparameters and submodels that may be difﬁcult to implement and optimise.
Causal Effect Variational Autoencoders (CEVAEs) (Louizos et al.
(2017)) use a latent variable modeling approach that is robust to hidden confounding, but does not address biased treatment assignment.
Causal Multi-task Gaussian Processes (CMGP) (Alaa & van der Schaar (2017)) apply a multi-task Gaussian Process to ITE estimation.
The optimisation of CMGPs involves a matrix inversion of O ( n 3 ) complexity that limits their scalability.
In contrast to existing methods, PM is a simple doubly robust method built on propensity score matching that can be used to train expressive non-linear neural network models for ITE estimation from observational data in settings with any number of treatments (Table 1).
While the underlying idea behind PM is simple and effective, it has, to the best of our knowledge, not yet been explored.
PM
CF
BART BNN CFRNET
[Table]
M ETHODOLOGY
Problem Setting.
We consider a setting in which we are given N observed samples X , where each sample consists of p covariates x i with i ∈ [0 , p ) .
For each sample, the potential outcomes are represented as a vector Y with k entries y j where each entry corresponds to the outcome when applying one treatment t j out of the set of k available treatments T
= { t 0 , ..., t k − 1 } with j ∈ [0 , k ) .
As training data, we receive observed samples X and their factual outcomes y j when applying one treatment t j .
The set of available treatments can contain two or more treatment options.
We refer to the special case of two available treatment options as the binary treatment setting.
Given the observational training data with factual outcomes, we wish to train a predictive model that is able to accurately produce a predicted potential outcomes vector ˆ Y with k entries ˆ y j .
In literature, this setting is known as the Rubin-Neyman potential outcomes framework (Rubin (2005)).
Precision in Estimation of Heterogenous Effect (PEHE).
The primary metric that we optimise for when training models to estimate ITE is the PEHE (Hill (2011)).
In the binary setting, the PEHE measures the ability of a predictive model to estimate the difference in effect between two treatments t 0 and t 1 for samples X .
To compute the PEHE, we measure the mean squared error between the true difference in effect y 1 ( n ) − y 0 ( n ) , drawn from the noiseless underlying outcome distributions µ 1 and µ 0 , and the predicted difference in effect ˆ y 1 ( n ) − ˆ y 0 ( n ) indexed by n over N samples:
[Equation]
When the underlying noiseless distributions µ j are not known, the true difference in effect y 1 ( n ) − y 0 ( n ) can be estimated using the noisy ground truth outcomes y i (Appendix A).
As a secondary metric, we consider the error (cid:15) ATE in estimating the average treatment effect (ATE) (Hill (2011)).
The
ATE measures the average difference in effect across the whole population (Appendix B).
The ATE is not as important as PEHE for models optimised for ITE estimation, but can be a useful indicator of how well an ITE estimator performs at comparing two treatments across the entire population.
We can neither calculate (cid:15) PEHE nor (cid:15) ATE without knowing the outcome generating process.
Multiple Treatments.
Both (cid:15) PEHE and (cid:15) ATE can be trivially extended to multiple treatments by considering the average PEHE and ATE between every possible pair of treatments.
Note that we lose the information about the precision in estimating ITE between speciﬁc pairs of treatments by averaging over all
[Equation]
pairs.
However, one can inspect the pair-wise PEHE to get the whole picture.
[Equation]
Perfect Match (PM).
We consider fully differentiable neural network models optimised via minibatch stochastic gradient descent to predict potential outcomes ˆ Y given a sample X .
To address the treatment selection bias inherent in observational data, we propose the use of a training methodology based on propensity score matching.
A propensity score is the conditional probability p ( t | X ) of a given sample X receiving a speciﬁc treatment t (Rosenbaum & Rubin (1983);
Ho et al.
The propensity score is a balancing score , meaning that if the distribution of propensity scores across treatment groups is the same, then the distribution of their covariates x i is also guaranteed to be the same across treatment groups (Ho et al.
By matching samples with their nearest neighbours by propensity score we can, in the optimal case, balance the covariates x i .
This breaks the dependence of treatment assignment on X if all relevant covariates are observed, and therefore effectively removes any potential treatment assignment bias (Ho et al.
However, in practice, we do not have access to the true treatment assignment probability p ( t | X ) .
We therefore have to estimate the propensity score p ( t | X ) by training a predictive model to predict the likelihood of receiving a treatment t for a given sample X .
Note that propensity scores are not the only balancing score (Ho et al.
For example, X itself is another balancing score.
However, matching on the onedimensional propensity score is preferable because it avoids the curse of dimensionality that would be associated with matching on the potentially high-dimensional X directly.
In PM, we match every sample within a minibatch with its nearest neighbours by propensity score from all other treatments (Algorithm 1, more details in Appendix C), taken from the training set.
Every minibatch the model is trained on therefore contains the same number of samples for each treatment group, and the covariates x i of each treatment group are approximately balanced on average.
The intuition behind matching on the minibatch level, rather than the dataset level (Ho et al.
(2011)), is that it ensures that every gradient step is done in a way that is approximately unbiased by the treatment assignment bias inherent in observational data.
In this sense, PM can be seen as a minibatch sampling strategy (Csiba & Richt´arik (2018)) speciﬁcally designed to improve learning for counterfactual inference. procedure P
X out ← do 4: p ( t | X ) ← E ) 5: for 0 to k do 6: 7: ps
← p Add Model Selection. Besides accounting for the treatment assignment bias, the other major issue in learning for counterfactual inference from observational data is that, given multiple models, it is not trivial to decide which one to select.
The root problem is that we do not have direct access to the true error in estimating counterfactual outcomes, only the error in estimating the observed factual outcomes.
This makes it difﬁcult to perform parameter and hyperparameter optimisation, as we do not know which models are better than others for counterfactual inference on a given dataset.
To rectify this problem, we use a nearest neighbour approximation ˆ (cid:15) NN-PEHE of the ˆ (cid:15) PEHE metric for the binary (Shalit et al.
and multiple treatment settings for model selection.
The ˆ (cid:15) NN-PEHE estimates the treatment effect of a given sample by substituting the true counterfactual outcome with the outcome y j from a respective nearest neighbour NN matched on X using the Euclidean distance.
[Equation]
Analogously to Equations (2) and (3), the ˆ (cid:15) NN-PEHE metric can be extended to the multiple treatment setting by considering the mean ˆ (cid:15) NN-PEHE between all (cid:0) [Equation] (cid:1) possible pairs of treatments (Appendix D). Model Architecture.
No different from other applications, the chosen architecture plays a key role in the performance of neural networks when attempting to learn representations for counterfactual inference.
Shalit et al. (2017) claimed that the na¨ıve approach of appending the treatment index t j may perform poorly if X is high-dimensional, because the inﬂuence of t j on the hidden layers may be lost during training. Shalit et al. (2017) subsequently introduced the TARNET architecture to rectify this issue.
Since the original TARNET was limited to the binary treatment setting, we extended the TARNET architecture to the multiple treatment setting (Figure 1).
We did so by using k head networks, one for each treatment option over a set of shared base layers, each with L layers.
In TARNET, the j th head network is only trained on samples from treatment t j .
The shared layers are trained on all samples.
By using a head network for each treatment, we ensure t j maintains an appropriate degree of inﬂuence on the network output at all points during training.
[Figure]
[Table]
E XPERIMENTS
Our experiments aimed to answer the following questions:
1 What is the comparative performance of PM in inferring counterfactual outcomes in the binary and multiple treatment setting compared to existing state-of-the-art methods?
2 Does model selection by NN-PEHE outperform selection by factual MSE?
3 How does the relative number of matched samples within a minibatch affect performance?
4 How well does PM cope with an increasing treatment assignment bias in the observed data?
How does the presence of hidden confounders inﬂuence the performance of PM? 6 How do the learning dynamics of minibatch matching compare to dataset-level matching?
D ATASETS
We performed experiments on four real-world and semi-synthetic datasets (Table 2) with binary and multiple treatment options in order to gain a better understanding of the empirical properties of PM.
Infant Health and Development Program (IHDP).
The IHDP dataset (Hill (2011)) contains data from a randomised study on the impact of specialist visits on the cognitive development of children, and consists of 747 children with 25 covariates describing properties of the children and their mothers.
Children that did not receive specialist visits were part of a control group.
The outcomes were simulated using the NPCI package from Dorie (2016).
The IHDP dataset is biased because the treatment groups had a biased subset of the treated population removed (Shalit et al.
We used the same simulated outcomes 2 as Shalit et al.
Jobs
The Jobs dataset (LaLonde (1986)) is a blend of data from randomised and observational studies on the effect of professional training programs on unemployment, and is a commonly used benchmark dataset for counterfactual inference in the binary setting.
Each sample consisted of demographic covariates, such as age, gender and previous income, and the treatment was enrolment in a job training program.
All samples that were not enrolled in a training program were considered control cases.
The task was to predict the potential unemployment given the covariates and treatment.
We used the same feature and sample set as described in Dehejia & Wahba (2002); Smith &
Todd (2005); Shalit et al.
A subgroup of the data was randomised and we therefore were able to estimate the ground truth average treatment effect on the treated (ATT) (Shalit et al.
Treatment assignment in Jobs is biased because parts of the data come from an observational study.
lllllll l ll l ll l l ll llllllllll l l l ll l l lll l llll l l l l ll ll l l l l l lll l l ll l lllllllll l l l l l l l lll l l lll l ll llllllllll ll ll l l ll l l l lll l llll l lll l l l lll l lllllllllllllllll l l l l lllll l l lllll l lllllllll l l ll l ll l lll lll l l l lll l l l l l llllll ll l l llllll l ll l l l lllllll l ll l llllllll l lllllllllllllllllllllllll l l l llllll l l ll l l l lllll l l l ll l l l lllllllllllllllllllll l l l ll l l lll l ll l l ll l llllll l l l llllllll l l l l l l ll l lll l l l llllllll l l l l ll l l l lllllllll l l ll l lllll l l l llllll l l l l l ll lllllllll l l lllll l l l lll ll lll l lllllllllllllll l ll l l l lll ll lllllll l l l lll l lllllllll l l ll l l l l l l ll l l lll l l l l l l l l llllllllllll l lll l l lllll
l l l lllll l llll l llll l l lllll l l l llllllll l ll l l l l l ll lllllll l l l llll l ll l l ll l llllllll l llllll l lllllllllllllll l ll l l lllllll l l ll l l l ll lll llllllllllllll l llll l l l l l l l l ll l l llllllllll l l l lll l l lllllll l l ll lll l lllll l l l l l l ll ll l lllllll l l ll l lllllllll l
llllllllllll l ll l ll l l ll l l lllllllllllllllllllllllllllll l l l l lllllll l l l ll l l ll llllll l l l l l l l l l l l llllllllllllllllllll l lllll l llllll l l l l l l l l l l l l l ll l l l l lllll l l l l l l llllll l ll l l l ll l l ll l llllllllllll l ll ll l lllllll l llllllllllllllllllllllllllllllllllll l lll l l l ll l ll lll lll l l ll l lllllll l llll l ll l l lll llllllllllllllllllll l l ll l lllll l ll l l l l l llll l ll l l ll l llll l l l ll llll l lllll ll l l l lllllllll l l lllllllllll l ll llllllllllllll l lllll l lllll l l l l llllll ll l l l l l l lll
{+l l l l l l llll l l l ll l l l lllllllll lll ll llllll l llll l lllll l lllll l l llll l llllllllllllllllllllllllllllllllll l l ll ll lllllllll l l l ll ll ll ll l ll lllll l lllllllll l ll ll lll l l ll l l l l l l lllll l l lll l lllllllllllll l l ll l lllll l l l lll ll lllllllll    
+}
MSE
PE H E r =0.
lllllll l ll l ll l l l l llllllllll l l l ll l l lll l llll l l l l ll ll l l l l l lll l l ll l lllllllll
l l l l l l l ll l l l lll l ll llllllllll ll ll l l ll l l l l ll l l l ll l lll l l l lll l l lll ll lllllllllll l l l l ll l ll l l ll l ll l llllll l ll l l ll l ll l lll lll l l l lll l l l l l llllll l l l l llllll l ll l l l lllllll l ll l l l lll lll l lllll l l l ll l lllll l lll l ll l l l l l llllll l l ll l l l lllll l l l ll l l l lllllllllllllllllllll l l l ll l l lll l ll l l ll l l l ll l l l l l llllllll l l l l l l ll l l l l l l l llllllll l l l l ll l l l ll l lll ll l l l ll l l ll l l l l l l l llll l l l l l ll lllllllll l l lll l l l l l lll l l lll l llllllllllll ll l l ll l l l lll l l lllllll l l l lll l lllllllll l l l l l l l l l l ll l l lll l l l l l l l l lllllllllll l l ll l l l lllll
l l l lllll l llll l llll l l l llll l l l llllllll l ll l l l l l l l lllllll l l l llll l ll l l ll l llllllll l llllll l lllllllllllllll l ll l l lllllll l l l l l l l l l l l l llllllllllllll l llll l l l l l l l l ll l l llllllllll l l l lll l l lllllll l l l l l ll l lllll l l l l l l ll ll l lllllll l l l l l llll l ll l l l llllllllllll l ll l l l l l ll l l ll l llllllllllllllllllllllllll l l l l llllll l l l l ll l l ll ll l lll l l l l l l l l l l l llllllllllllllllllll l lllll l lll l l l l l l l l l l l l l l l l ll l l l l lllll l l l l l l llllll l ll l l l ll l l ll l llllllllllll l ll ll l lllllll l lllllll l llll lll
lll l lllllllllllllllll l lll l l l ll l ll ll l l l l l l ll l lllllll l llll l l l l l lll lllllllll ll l l l llll l l l l l l l lllll l ll l l l l l llll l l l l l ll l llll l l l ll llll l lllll ll l l l lllllllll l l lllllllllll l ll llllllllllllll l lllll l l l l ll l l l l llllll ll l l l l l l lll
{+l l l l l l llll l l l l l l l l lllllllll l ll ll llllll l l lll l lllll l l llll l l llll l ll l l l l ll llllllllllllllllllllllllll l l ll ll lllllllll l l l ll ll ll ll l ll lll l l l l llllllll l ll l l lll l l ll l l l l l l lllll l l lll l lllllllllllll l l l l l lllll l l l l l l l l lllllllll 2 4 6 8 10 12 
+}
NN−PEHE
PE H E r =0.
[Figure]
[Figure]
l l l l l l l
Treatment Assignment Bias k e A T E l l l l l l l l PMCFRNETTARNETPDCF
Treatment Assignment Bias k e PE
H E
l l l l l l l l PMCFRNETTARNETPDCF
News. The News dataset was ﬁrst proposed as a benchmark for counterfactual inference by Johansson et al.
and consists of 5000 randomly sampled news articles from the NY Times corpus 3 .
The News dataset contains data on the opinion of media consumers on news items.
The samples X represent news items consisting of word counts x i ∈ N , the outcome y j ∈ R is the reader’s opinion of the news item, and the k available treatment options represent various devices that could be used for viewing, e.g. smartphone, tablet, desktop, television or others (Johansson et al.
We extended the original dataset speciﬁcation in (Johansson et al.
to enable the simulation of arbitrary numbers of viewing devices.
To model that consumers prefer to read certain media items on speciﬁc viewing devices, we train a topic model on the whole NY Times corpus and deﬁne z ( X ) as the topic distribution of news item X .
We then randomly pick k + 1 centroids in topic space, with k centroids z j per viewing device and one control centroid z c .
We assigned a random Gaussian outcome distribution with mean µ j ∼ N (0 . 45 , 0 . 15) and standard deviation σ j ∼ N (0 . 1 , 0 . 05) to each centroid.
For each sample, we drew ideal potential outcomes from that Gaussian outcome distribution ˜
y
j ∼ N ( µ j , σ j )+ (cid:15) with (cid:15) ∼ N (0 , 0 . 15) .
We then deﬁned the unscaled potential outcomes ¯ y j = ˜
y
j ∗ [ D ( z ( X ) , z j ) + D ( z ( X ) , z c )] as the ideal potential outcomes ˜ y
j weighted by the sum of distances to centroids z j and the control centroid z c using the Euclidean distance as distance D. We assigned the observed treatment t using t | x ∼ Bern ( softmax ( κ ¯ y j )) with a treatment assignment bias coefﬁcient κ , and the true potential outcome y j = C ¯ y j as the unscaled potential outcomes ¯ y j scaled by a coefﬁcient C = 50 .
We used four different variants of this dataset with k = 2 , 4 , 8 , and 16 viewing devices, and κ = 10 , 10 , 10 , and 7 , respectively.
Higher values of κ indicate a higher expected treatment assignment bias depending on ¯ y j .
κ = 0 indicates no assignment bias.
The Cancer Genomic Atlas (TCGA).
The TCGA project collected gene expression data from various types of cancers in 9659 individuals (Weinstein et al.
There were four available clinical treatment options: (1) medication, (2) chemotherapy, (3) surgery, or (4) both surgery and chemotherapy.
We used a synthetic outcome function that simulated the risk of cancer recurrence after receiving either of the treatment options based on the real-world gene expression data.
Before further processing, we standardised the gene expression data using the mean and standard deviations of gene expression at each gene locus for normal tissue in the training set.
To produce the potential outcomes Y , we ﬁrst selected k + 1 gene expression phenotypes as representative phenotypes: One phenotype X j for each treatment option and one as a control phenotype X c .
Analogously to the
News dataset, we assigned a random Gaussian outcome distribution with mean µ j ∼ N (0 . 45 , 0 . 15) and standard deviation σ t ∼ N (0 . 1 , 0 . 05) to each phenotype.
[-l l l l llllll ll l l l l l l lll l l l l l l llll l l l ll l l l lllllllll lll ll llllll l llll l lllll l lllll l l llll l llllllllllllllllllllllllllllllllll l l ll ll lllllllll l l l ll ll ll ll l ll lllll l lllllllll l ll ll lll l l ll l l l l l l lllll l l lll l lllllllllllll l l ll l lllll l l l lll ll lllllllll    
 MSE PE H E r =0.-]For each sample, we drew ideal potential outcomes from that Gaussian outcome distribution ˜
[-l l l l l l l ll l l l lll l ll llllllllll ll ll l l ll l l l l ll l l l ll l lll l l l lll l l lll ll lllllllllll l l l l ll l ll l l ll l ll l llllll l ll l l ll l ll l lll lll l l l lll l l l l l llllll l l l l llllll l ll l l l lllllll l ll l l l lll lll l lllll l l l ll l lllll l lll l ll l l l l l llllll l l ll l l l lllll l l l ll l l l lllllllllllllllllllll l l l ll l l lll l ll l l ll l l l ll l l l l l llllllll l l l l l l ll l l l l l l l llllllll l l l l ll l l l ll l lll ll l l l ll l l ll l l l l l l l llll l l l l l ll lllllllll l l lll l l l l l lll l l lll l llllllllllll ll l l ll l l l lll l l lllllll l l l lll l lllllllll l l l l l l l l l l ll l l lll l l l l l l l l lllllllllll l l ll l l l lllll l l l lllll l llll l llll l l l llll l l l llllllll l ll l l l l l l l lllllll l l l llll l ll l l ll l llllllll l llllll l lllllllllllllll l ll l l lllllll l l l l l l l l l l l l llllllllllllll l llll l l l l l l l l ll l l llllllllll l l l lll l l lllllll l l l l l ll l lllll l l l l l l ll ll l lllllll l l l l l llll l ll l l l llllllllllll l ll l l l l l ll l l ll l llllllllllllllllllllllllll l l l l llllll l l l l ll l l ll ll l lll l l l l l l l l l l l llllllllllllllllllll l lllll l lll l l l l l l l l l l l l l l l l ll l l l l lllll l l l l l l llllll l ll l l l ll l l ll l llllllllllll l ll ll l lllllll l lllllll l llll lll lll l lllllllllllllllll l lll l l l ll l ll ll l l l l l l ll l lllllll l llll l l l l l lll lllllllll ll l l l llll l l l l l l l lllll l ll l l l l l llll l l l l l ll l llll l l l ll llll l lllll ll l l l lllllllll l l lllllllllll l ll llllllllllllll l lllll l l l l ll l l l l llllll ll l l l l l l lll l l l l l l llll l l l l l l l l lllllllll l ll ll llllll l l lll l lllll l l llll l l llll l ll l l l l ll llllllllllllllllllllllllll l l ll ll lllllllll l l l ll ll ll ll l ll lll l l l l llllllll l ll l l lll l l ll l l l l l l lllll l l lll l lllllllllllll l l l l l lllll l l l l l l l l lllllllll 2 4 6 8 10 12 
-]y
j ∼ N ( µ j , σ j ) + (cid:15) with (cid:15) ∼ N (0 , 0 . 15) .
We then deﬁned the unscaled potential outcomes ¯ y j = ˜
y
j ∗
[ D ( X, X t j ) + D ( X, X c )] as the ideal potential outcomes ˜ y
j weighted by the sum of distances to phenotype X t j and the control phenotype X c using the cosine similarity as distance metric D. We assigned the observed treatment t using t | x ∼ Bern ( softmax ( κ ¯ y j )) with a treatment assignment bias coefﬁcient κ = 10 , and the true potential outcome y j = C ¯ y j as the unscaled potential outcomes ¯ y j scaled by a coefﬁcient C = 50 .
All datasets with the exception of IHDP were split into a training (63%), validation (27%) and test set (10% of samples).
For IHDP we used exactly the same splits as previously used by Shalit et al.
We repeated experiments on IHDP, Jobs, News and TCGA 1000, 10, 50, and 5 times, respectively.
We reassigned outcomes and treatments with a new random seed for each repetition. [Table]
4. E XPERIMENTAL S ETUP
Models.
We evaluated PM, ablations, baselines [Table] and all relevant state-of-the-art methods: kNN (Ho et al.
, BART (Chipman et al.
; Chipman & McCulloch (2016)), Random Forests (RF) (Breiman (2001)), CF (Wager & Athey (2017)), GANITE (Yoon et al.
, Balancing
Neural Network (BNN) (Johansson et al.
, TARNET (Shalit et al.
(2017)), Counterfactual
Regression Network using the Wasserstein regulariser (CFRNET Wass ) (Shalit et al.
, CEVAE (Louizos et al.
, CMGP (Alaa & van der Schaar (2017)), and PD (Alaa et al.
We trained a Support Vector Machine (SVM) with probability estimation (Pedregosa et al.
to estimate p ( t | X ) for PM on the training set.
We also evaluated preprocessing the entire training set with PSM using the same matching routine as PM (PSM PM ) and the ”MatchIt” package (PSM MI , Ho et al.
before training a TARNET (Appendix E).
In addition, we trained an ablation of PM where we matched on the covariates X (+ on X ) directly, if X was low-dimensional ( p < 200 ), and on a 50-dimensional representation of X obtained via principal components analysis (PCA), if X was high-dimensional, instead of on the propensity score.
We also evaluated PM with a multi-layer perceptron (+ MLP) that received the treatment index t j as an input instead of using a TARNET.
[Figure]
[Table]
Architectures. To ensure that differences between methods of learning counterfactual representations for neural networks are not due to differences in architecture, we based the neural architectures for TARNET, CFRNET Wass , PD and PM on the same, previously described extension of the TARNET architecture (Shalit et al.
, Details in Appendix F) to the multiple treatment setting.
Hyperparameters.
For the IHDP, Jobs, News and TCGA dataset we respectively used 30, 30, 10, and 5 optimisation runs for each method using randomly selected hyperparameters from predeﬁned ranges (Appendix G).
We selected the best model across the runs based on validation set ˆ (cid:15) NN-PEHE .
Metrics.
We calculated the (cid:15) PEHE and (cid:15) ATE for datasets for which we know the outcome generating process.
For the Jobs dataset, we did not have access to counterfactual outcomes.
We therefore calculated the ATT =
[-| T 1 ∩ E | (cid:80) Nn =
 y 1 ( n ) − | T 
 ∩ E | (cid:80)-][Equation]
Nn =0 y 0 ( n ) where E was the randomised subset of the dataset, and reported the error in estimating ATT (cid:15) ATT = | ATT − | T 1 ∩ E | (cid:80)
Nn [-=
-]=
[Equation] ˆ y 1 ( n ) | , and the policy risk R Pol ( π ) (Shalit et al.
(2017); Yoon et al.
The policy risk is the average loss in value when treating according to the policy π implied by an ITE estimator (Appendix H).
R ESULTS AND D ISCUSSION
Counterfactual Inference.
We evaluated the counterfactual inference performance of the listed models in the binary setting (Table 3) and multiple treatment setting (Table 4).
On IHDP, PM reached the third best performance in terms of √ (cid:15) PEHE after CFRNET and CMGP, and the second best (cid:15) ATE after CMGP.
On Jobs, PM reached the second best R Pol ( π ) after GANITE.
On the binary News-2, PM outperformed all other methods in terms of √ (cid:15) PEHE and (cid:15) ATE .
On the News-4/8/ datasets with more than two treatments, PM consistently outperformed all other methods - in some cases by a large margin - on both metrics with the exception of the News-4 dataset, where PM came second to PD.
The strong performance of PM across a wide range of datasets with varying amounts of treatment options is remarkable considering how simple it is compared to other, highly specialised methods.
Notably, PM consistently outperformed both CFRNET, with the exception of the IHDP dataset, which accounted for covariate imbalances between treatments via regularisation rather than matching, and PSM MI , which accounted for covariate imbalances by preprocessing the entire training set with a matching algorithm (Ho et al.
We also found that matching on the propensity score was, in almost all cases, not signiﬁcantly different from matching on X directly when X was low-dimensional, or a low-dimensional representation of X when X was high-dimensional (+ on
X ).
This indicates that PM is effective with any low-dimensional balancing score. In addition, using PM with the TARNET architecture outperformed the MLP (+ MLP) in almost all cases, with the exception of the low-dimensional IHDP and Jobs.
We conclude that matching on the propensity score or a low-dimensional representation of X and using the TARNET architecture are sensible default conﬁgurations, particularly when X is high-dimensional. [Figure]
Model Selection.
To judge whether NN-PEHE is more suitable for model selection for counterfactual inference than MSE, we compared their respective correlations with the PEHE on IHDP.
We found that NN-PEHE correlates signiﬁcantly better with the PEHE than MSE (Figure 2).
Number of Matches per Minibatch.
To determine the impact of matching fewer than 100% of all samples in a batch, we evaluated PM on News-8 trained with varying percentages of matched samples on the range 0 to 100% in steps of 10% (Figure 3).
We found that including more matches indeed consistently reduces the counterfactual error up to 100% of samples matched.
Interestingly, we found a large improvement over using no matched samples even for relatively small percentages ( < 40%) of matched samples per batch.
This shows that propensity score matching within a batch is indeed highly effective at improving the training of neural networks for counterfactual inference.
Treatment Assignment Bias.
To assess how the predictive performance of the different methods is inﬂuenced by increasing amounts of treatment assignment bias, we evaluated their performances on News-8 while varying the assignment bias coefﬁcient κ on the range of 5 to 20 (Figure 4).
We found that PM handles high amounts of assignment bias better than existing state-of-the-art methods.
Hidden Confounding.
To analyse the sensitivity of the different methods to hidden confounding, we successively retrained the models on a reduced set of covariates x i in increasing steps of 10% ( = 2053 gene loci) in the range of 10% to 90% (Figure 5).
We removed the covariates from the model but kept them for computing assignments and outcomes, i.e. they were hidden confounders.
The results showed that PM was competitive with TARNET and CFRNET in terms of robustness to hidden confounding.
Notably, PM, TARNET and CFRNET were relatively stable up to levels of 60% of hidden confounding, likely because some gene loci are highly correlated. Comparing Minibatch and Dataset Matching.
As outlined previously, it is possible to optimise the counterfactual error using the observed factual samples if (1) we manage to break the dependence of treatment assignment on X , and (2) we observe all relevant variables (Ho et al. If we were successful in balancing the covariates, we would expect that the counterfactual error is implicitly and consistently improved alongside the factual error, despite optimising the neural network for the observed factual loss only. To elucidate to what degree this is the case when using the various evaluated matching-based methods, we evaluated the respective training dynamics of PM, PSM PM and PSM MI (Figure 6). We found that PM better conforms to the desired behavior than PSM PM and
PSM MI .
PSM PM , which used the same matching strategy as PM but on the dataset level, showed a much higher variance than PM.
PSM MI was overﬁtting to the treated group.
Limitations. A limitation of this work is that the theory of matching to balance the covariates only holds in the strongly ignorable setting, i.e. under the assumption that there are no unobserved confounders.
However, our experiments (Figure 5) and related works showed that even the presence of large amounts of hidden confounders may not necessarily decrease the performance of ITE estimators in practice if we observe proxy variables (Montgomery et al.
; Louizos et al.
In future work, we would like to integrate PM with diagnostic tools that address hidden confounding (Kallus & Zhou (2018)), and the ability to learn from multiple related tasks (Schwab et al.
C ONCLUSION
We presented PM, a new and simple method for training neural networks for estimating ITEs from observational data that extends to any number of available treatment options.
In addition, we extended the PEHE metric to settings with more than two treatments and introduced a nearest neighbour approximation of PEHE and mPEHE that can be used for model selection without having access to counterfactual outcomes.
We performed experiments on several real-world and semisynthetic datasets that showed that PM outperforms a number of more complex state-of-the-art methods in inferring counterfactual outcomes.
We also found that the NN-PEHE correlates signiﬁcantly better with real PEHE than MSE, that including more matched samples in each minibatch improves the learning of counterfactual representations, and that PM handles an increasing treatment assignment bias better than existing state-of-the-art methods.
PM may be used for settings with any amount of treatment options, is compatible with any existing neural network architecture, simple to implement, and does not introduce any additional hyperparameters or computational complexity
A P RECISION IN E STIMATION OF H ETEROGENOUS E FFECT (PEHE)
When the underlying noiseless distributions µ j are not given, PEHE is calculated using:
[Equation]
B A VERAGE T REATMENT E FFECT (ATE)
The ATE measures the average difference in effect across the whole population (Hill (2011)).
It can be useful indicator of how well an ITE estimator performs at comparing two treatments across the entire population.
[Equation]
Similar to Equation (A.1), the counterfactual treatment effect can be estimated using the noisy ground truth outcomes if the underlying noiseless distributions µ j are not given:
[Equation]
C D ETAILS P ERFECT M ATCH (PM)
Algorithm 1 outlines the procedure of using PM to augment a batch with propensity score matches for training a neural network with stochastic gradient descent.
We trained a Support Vector Machine (SVM) with probability estimation as propensity score estimator E PS .
To speed up recalling nearest neighbours, we additionally prepared an index per treatment into X train sorted by propensity score.
At augmentation time, we used binary search on this index to ﬁnd the nearest neighbours by propensity score in O(log N t ) where N t is the number of samples X in X train assigned to the treatment group t .
To avoid overﬁtting to speciﬁc edge samples when propensity scores are not distributed evenly in the training set, we chose at random from one of the k =6 closest samples by propensity score.
See Table S1 for a comparison of the effects on predictive performance of choosing varying values of k .
For the ”+ on X ” model we matched, using the Euclidean distance, directly on the covariates X on IHDP and Jobs, and using a low dimensional representation of X , obtained using principal component analysis (PCA) with 50 principal components, on News.
[Table]
D N EAREST N EIGHBOUR A PPROXIMATION OF PEHE FOR M ULTIPLE T REATMENTS (NN- M PEHE)
The ˆ (cid:15) NN-PEHE metric can be extended to the multiple treatment setting by considering the mean ˆ (cid:15) NN-PEHE between all (cid:0) [Equation] possible pairs of treatments:
[Equation]
D. [Equation]
E P SEUDOCODE P ROPENSITY S CORE M ATCHING (PSM)
Algorithm S2 outlines the procedure of preprocessing a training set using PSM MI .
For PSM PM , we used the PM matching algorithm (see Algorithm 1) to ﬁnd matches between samples from the treatment groups.
For PSM MI , we used the MatchIt package (Ho et al.
(2011)) on setting ”nearest” to ﬁnd matches between sets of samples from two treatment groups (function get matched samples ).
Algorithm S2 Preprocessing a training set using Propensity Score Matching (PSM).
After preprocessing, the training set contains an equal number of samples from each treatment group and the covariates x i across all treatment groups are approximately balanced.
Input Output
1: procedure P REPROCESS T RAINING S ET : 2: t base ← index of treatment with smallest number of samples in X train 3: X base ← all samples in X train with treatment t base 4: for i = 0 to k − 1 do 5: if i 6 = t base then 6:
X current ← all samples in X train with treatment i 7: X matched ← get matched samples ( X base , X current ) 8: Add all samples in X matched to X out 9:
Add all samples in X base to X out F TARNET, CFRNET AND PD FOR M ULTIPLE T REATMENTS
In TARNET models, we used ELU nonlinearities between the L hidden layers with M hidden units.
L and M were hyperparameters that we optimised during hyperparameter optimisation (Section G).
We did not use batch normalisation (BN).
To extend CFRNET to the multiple treatment setting, we deﬁned the ﬁrst treatment option as the control treatment and regularised all treatment options to have the same activation distribution in the topmost shared layer (Shalit et al.
For PD, we only used propensity dropout and the propensity estimation network.
We did not make use of the alternating training schedule proposed in (Alaa et al.
The ”+ MLP” model was a simple MLP with L hidden layers of M hidden units that received the treatment option index t j as an additional input along with the covariates X , and output a k -dimensional potential outcome vector ˆ Y .
The MLP used ELU nonlinearities between the L hidden layers, and also did not use BN.
G H YPERPARAMETERS
To ensure a fair comparison, we used a standardised approach to hyperparameter optimisation for those methods for which we did not have previously reported performance numbers.
In particular, each method we trained was optimised over the same amount of hyperparameter optimisation runs.
For the methods that used neural network models (TARNET, CFRNET, PD, PSM PM , PSM MI , including ”+ on X ” and ”+ MLP”, and PM), we chose hyperparameters at random from predeﬁned ranges (Table S2).
For CFRNET, we additionally varied the weight of the imbalance penalty at random between 0.1, 1.0, and 10.0.
All methods that used a neural network model used the TARNET architecture to ensure differences in performance are not due to architectural differences.
To train PM for the Jobs dataset, we used ﬁxed hyperparameters: A batch size of 50, 60 hidden units per layer, and 3 hidden layers.
For optimisation, we used the Adam optimiser with a learning rate of 0 .
001 for a maximum of 100 (News, TCGA) or 400 (IHDP, Jobs) with an early stopping patience of 30 on the factual MSE.
We used the default hyperparameters for BART and CF from the ”bartMachine” (Kapelner & Bleich (2013)) and ”grf” (Athey et al.
R-packages.
For GANITE, we used our own implementation since there was no open source implementation available, and - in addition to the parameters in Table S2 - optimised over the supervised loss weights α and β (Yoon et al.
(2018)) between 0.1, 1, and 10.
For the GANITE generators and discriminators, we used MLP architectures with L hidden layers of M hidden units each.
[Table]
We deﬁne the policy risk R Pol ( π ) as the expected loss in value when treating according to the policy π implied by an ITE estimator (Shalit et al.
[Equation]
where the policy π ( X ) for a sample X implied by an ITE estimator that produces potential outcomes ˆ y is deﬁned as:
[Equation]
I O RIGINS OF P ERFORMANCE N UMBERS IN T ABLE
Where available, performance numbers in Table 3 were taken from the original authors for IHDP and Jobs (Table S3).
For News-2/4/8/16 and TCGA, performance numbers were produced by us.
[Table]