Year: 2018

Abstract:
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 t1?”. 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.

heading: 1 INTRODUCTION
text: Estimating individual treatment effects1 (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. (2017)). 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 difficult 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 significantly between the treated group and the overall population. A supervised model naı̈vely trained to minimise the factual error would overfit 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.
1The ITE is sometimes referred to as Conditional Average Treatment Effect (CATE).
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.

heading: 2 RELATED WORK
text: There are five 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. (2011)).
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. (2010); 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 find 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. (2018)) address ITE estimation using counterfactual and ITE generators. GANITE uses a complex architecture with many hyperparameters and submodels that may be difficult 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(n3) 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.

heading: 3 METHODOLOGY
text: Problem Setting. We consider a setting in which we are given N observed samples X , where each sample consists of p covariates xi with i ∈ [0, p). For each sample, the potential outcomes are represented as a vector Y with k entries yj where each entry corresponds to the outcome when applying one treatment tj out of the set of k available treatments T = {t0, ..., tk−1} with j ∈ [0, k). As training data, we receive observed samples X and their factual outcomes yj when applying one treatment tj . 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 t0 and t1 for samples X . To compute the PEHE, we measure the mean squared error between the true difference in effect y1(n) − y0(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:
PEHE = 1
N
N∑
n=0
( Eyj(n)∼µj(n)[y1(n)− y0(n)]− [ŷ1(n)− ŷ0(n)] )2 (1)
When the underlying noiseless distributions µj are not known, the true difference in effect y1(n)− y0(n) can be estimated using the noisy ground truth outcomes yi (Appendix A). As a secondary metric, we consider the error 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 PEHE nor ATE without knowing the outcome generating process.
Multiple Treatments. Both PEHE and 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 specific pairs of treatments by averaging over all ( k 2 ) pairs. However, one can inspect the pair-wise PEHE to get the whole picture.
̂mPEHE = 1( k 2
) k−1∑
i=0
i−1∑
j=0 ̂PEHE,i,j (2) ̂mATE = 1( k 2
) k−1∑
i=0
i−1∑
j=0
̂ATE,i,j (3)
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 sampleX receiving a specific treatment t (Rosenbaum & Rubin (1983); Ho et al. (2007)). 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 xi is also guaranteed to be the same across treatment groups (Ho et al. (2007)). By matching samples with their nearest neighbours by propensity score we can, in the optimal case, balance the covariates xi. 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. (2007)). 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. (2007)). 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-dimensionalX 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 xi 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 reduces the variance during training which in turn leads to better expected performance for counterfactual inference (Appendix D). In this sense, PM can be seen as a minibatch sampling strategy (Csiba & Richtárik (2018)) specifically designed to improve learning for counterfactual inference.
Algorithm 1 Batch Augmentation using Perfect Match (PM). After augmentation, each batch contains an equal number of samples from each treatment group and the covariates xi across all treatment groups are approximately balanced.
Input: Batch of B random samples Xbatch with assigned treatments t, training set Xtrain of N samples, number of treatment options k, propensity score estimator EPS to calculate the probability p(t|X) of treatment assigned given a sample X Output: Batch Xout consisting of B × k matched samples 1: procedure PERFECT MATCH: 2: Xout ← Empty 3: for sample X with treatment t in Xbatch do 4: p(t|X)← EPS(X) 5: for i = 0 to k − 1 do 6: if i 6= t then 7: psi ← p(t|X)i 8: Xmatched ← get closest match to propensity score psi with treatment i from Xtrain 9: Add sample Xmatched to Xout
10: Add X to Xout
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 difficult 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 ̂NN-PEHE of the ̂PEHE metric for the binary (Shalit et al. (2017)) and multiple treatment settings for model selection. The ̂NN-PEHE estimates the treatment effect of a given sample by substituting the true counterfactual outcome with the outcome yj from a respective nearest neighbour NN matched onX using the Euclidean distance.
̂NN-PEHE = 1
N
N∑
n=0
( [y1(NN(n))− y0(NN(n))]− [ŷ1(n)− ŷ0(n)] )2 (4)
Analogously to Equations (2) and (3), the ̂NN-PEHE metric can be extended to the multiple treatment setting by considering the mean ̂NN-PEHE between all ( k 2 ) possible pairs of treatments (Appendix E).
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 tj may perform poorly if X is high-dimensional, because the influence of tj 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 jth head network is only trained on samples from treatment tj . The shared layers are trained on all samples. By using a head network for each treatment, we ensure tj maintains an appropriate degree of influence on the network output at all points during training.
Under review as a conference paper at ICLR 2019

heading: 4 EXPERIMENTS
text: 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? 5 How does the presence of hidden confounders influence the performance of PM? 6 How do the learning dynamics of minibatch matching compare to dataset-level matching?

heading: 4.1 DATASETS
text: 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. (2017)). We used the same simulated outcomes2 as Shalit et al. (2017).
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. (2017). 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. (2017)). Treatment assignment in Jobs is biased because parts of the data come from an observational study.
2Available at: http://www.mit.edu/˜fredrikj/files/IHDP-1000.tar.gz
News. The News dataset was first proposed as a benchmark for counterfactual inference by Johansson et al. (2016) and consists of 5000 randomly sampled news articles from the NY Times corpus3. The News dataset contains data on the opinion of media consumers on news items. The samples X represent news items consisting of word counts xi ∈ N, the outcome yj ∈ 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. (2016)). We extended the original dataset specification in (Johansson et al. (2016)) to enable the simulation of arbitrary numbers of viewing devices. To model that consumers prefer to read certain media items on specific viewing devices, we train a topic model on the whole NY Times corpus and define z(X) as the topic distribution of news item X . We then randomly pick k + 1 centroids in topic space, with k centroids zj per viewing device and one control centroid zc. 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)+ with ∼ N (0, 0.15). We then defined the unscaled potential outcomes ȳj = ỹj ∗ [D(z(X), zj) + D(z(X), zc)] as the ideal potential outcomes ỹj weighted by the sum of distances to centroids zj and the control centroid zc 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 coefficient κ, and the true potential outcome yj = Cȳj as the unscaled potential outcomes ȳj scaled by a coefficient 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. (2013)). 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 first selected k + 1 gene expression phenotypes as representative phenotypes: One phenotype Xj for each treatment option and one as a control phenotype Xc. 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. For each sample, we drew ideal potential outcomes from that Gaussian outcome distribution ỹj ∼ N (µj , σj) + with ∼ N (0, 0.15). We then defined the unscaled potential outcomes ȳj = ỹj ∗ [D(X,Xtj ) + D(X,Xc)] as the ideal potential outcomes ỹj weighted by the sum of distances to phenotype Xtj and the control phenotype Xc 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 coefficient κ = 10, and the true potential outcome yj = Cȳj as the unscaled potential outcomes ȳj scaled by a coefficient 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. (2017). 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.
3https://archive.ics.uci.edu/ml/datasets/bag+of+words

heading: 4.2 EXPERIMENTAL SETUP
text: Models. We evaluated PM, ablations, baselines, and all relevant state-of-the-art methods: kNN (Ho et al. (2007)), BART (Chipman et al. (2010); Chipman & McCulloch (2016)), Random Forests (RF) (Breiman (2001)), CF (Wager & Athey (2017)), GANITE (Yoon et al. (2018)), Balancing Neural Network (BNN) (Johansson et al. (2016)), TARNET (Shalit et al. (2017)), Counterfactual Regression Network using the Wasserstein regulariser (CFRNETWass) (Shalit et al. (2017)), CEVAE (Louizos et al. (2017)), CMGP (Alaa & van der Schaar (2017)), and PD (Alaa et al. (2017)). We trained a Support Vector Machine (SVM) with probability estimation (Pedregosa et al. (2011)) 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 (PSMPM) and the ”MatchIt” package (PSMMI, Ho et al. (2011)) before training a TARNET (Appendix F). 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 tj as an input instead of using a TARNET.

heading: RF † 26.59 ± 3.02 † 18.03 ± 3.18 † 23.77 ± 2.14 † 12.40 ± 2.29 † 26.13 ± 2.48 † 15.91 ± 2.00
text: 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, CFRNETWass, PD and PM on the same, previously described extension of the TARNET architecture (Shalit et al. (2017), Details in Appendix G) 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 predefined ranges (Appendix H). We selected the best model across the runs based on validation set ̂NN-PEHE.
Metrics. We calculated the PEHE and 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 = 1|T1∩E| ∑N n=0 y1(n) − 1|T0∩E| ∑N n=0 y0(n) where E was the randomised subset of the dataset, and reported the error in estimating ATT ATT = |ATT− 1|T1∩E| ∑N n=0 ŷ1(n)|, and the policy risk RPol(π) (Shalit et al. (2017); Yoon et al. (2018)). The policy risk is the average loss in value when treating according to the policy π implied by an ITE estimator (Appendix I).

heading: 5 RESULTS AND DISCUSSION
text: 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 √ PEHE after CFRNET and CMGP, and the second best ATE after CMGP. On Jobs, PM reached the second best RPol(π) after GANITE. On the binary News-2, PM outperformed all other methods in terms of √ PEHE and ATE. On the News-4/8/16 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 PSMMI, which accounted for covariate imbalances by preprocessing the entire training set with a matching algorithm (Ho et al. (2011)). We also found that matching on the propensity score was, in almost all cases, not significantly 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 configurations, particularly when X is high-dimensional.
Under review as a conference paper at ICLR 2019 PM PSMPM PSMMI
ε P EH
E
Desired
Figure 6: Comparison of the learning dynamics during training (normalised training epochs; from start = 0 to end = 100 of training, x-axis) of several matching-based methods on the validation set of News-8. The coloured lines correspond to the mean value of the factual error ( √ MSE; red) and the counterfactual error ( √ ̂PEHE; blue) across 10 randomised hyperparameter configurations (lower is better). The shaded area indicates the standard deviation. All methods used exactly the same model and hyperparameters and only differed in how they addressed treatment assignment imbalance. The leftmost figure shows the desired behavior of the counterfactual and factual error jointly decreasing until convergence. PM exhibits the desired behavior of aligning the optimisation of the factual and counterfactual error. PSMPM shows a weaker degree of alignment and much higher variance. In contrast, PSMMI shows the undesired behavior of the counterfactual error increasing as the factual error decreases, indicating that the models were overfitting to the properties of the treated group.
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 significantly 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 influenced by increasing amounts of treatment assignment bias, we evaluated their performances on News-8 while varying the assignment bias coefficient κ 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 xi 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 onX , and (2) we observe all relevant variables (Ho et al. (2007)). 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, PSMPM and PSMMI (Figure 6). We found that PM better conforms to the desired behavior than PSMPM and PSMMI. PSMPM, which used the same matching strategy as PM but on the dataset level, showed a much higher variance than PM. PSMMI was overfitting 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. (2000); Louizos et al. (2017)). 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. (2018)).

heading: 6 CONCLUSION
text: 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 significantly 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.

heading: A PRECISION IN ESTIMATION OF HETEROGENOUS EFFECT (PEHE)
text: When the underlying noiseless distributions µj are not given, PEHE is calculated using:
̂PEHE = 1
N
N∑
n=0
( [y1(n)− y0(n)]− [ŷ1(n)− ŷ0(n)] )2 (A.1)

heading: B AVERAGE TREATMENT EFFECT (ATE)
text: 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.
ATE = || 1
N
N∑
n=0
Ey(n)∼µ(n)[y(n)]− 1
N
N∑
n=0
ŷ(n)||22 (B.1)
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:
̂ATE = || 1
N
N∑
n=0
y(n)− 1 N
N∑
n=0
ŷ(n)||22 (B.2)

heading: C DETAILS PERFECT MATCH (PM)
text: 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 EPS. To speed up recalling nearest neighbours, we additionally prepared an index per treatment into Xtrain sorted by propensity score. At augmentation time, we used binary search on this index to find the nearest neighbours by propensity score in O(log Nt) where Nt is the number of samples X in Xtrain assigned to the treatment group t. To avoid overfitting to specific 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 S1: Comparison of PM with varying numbers of nearest neighbours k considered for randomised matching on the News-2/4/8/16 datasets (k=1 corresponds to no randomisation; k=6, in cursive, are the results reported in the main body). We report the mean value ± the standard deviation of √ ̂PEHE on the test sets over 50 runs. There was no significant differences between different levels of k and k=6 (significance level α < 0.05). In particular, there was no significant difference between using randomisation and not using randomisation.
News-2 News-4 News-8 News-16 Method √ ̂PEHE √ ̂PEHE √ ̂PEHE √ ̂PEHE
PM (k=1) 16.79 ± 1.16 21.86 ± 2.24 20.70 ± 1.87 20.06 ± 1.43 PM (k=3) 16.82 ± 1.17 21.49 ± 2.43 20.50 ± 1.48 20.28 ± 1.52 PM (k=6) 16.76 ± 1.26 21.58 ± 2.58 20.76 ± 1.86 20.24 ± 1.46 PM (k=9) 16.91 ± 1.18 21.56 ± 2.58 20.75 ± 1.53 20.47 ± 1.94 PM (k=12) 17.11 ± 0.94 21.16 ± 2.20 21.10 ± 2.14 20.28 ± 1.64 PM (k=15) 16.87 ± 1.14 20.94 ± 2.05 20.92 ± 1.80 20.48 ± 1.80

heading: D REDUCED VARIANCE IS LINKED TO IMPROVED PERFORMANCE IN ESTIMATING TREATMENT EFFECTS
text: In the standard supervised setting, the variance added by minibatch SGD leads to slower convergence (Csiba & Richtrik (2018)), but will not typically lead to models that perform worse if they are trained until a convergence criterion has been met, e.g. early stopping. It is therefore perhaps surprising that we observe a difference not just in convergence (Figure 6), but also in counterfactual estimation performance (Tables 3 and 4) when comparing models trained with and without batch matching.
A potential cause of the observed difference in treatment effect estimation performance is that we do not have access to the true counterfactual error in observational data to select the best model encountered during training. We therefore have to resort to using the factual error (e.g. factual MSE) or an estimator based on the factual error (e.g. NN-PEHE) to select the best encountered model. However, when the variance during training is high, chances are that we select a model based on the factual error that at this point during the training happens to be suboptimal in terms of counterfactual error. Through this mechanism, added variance during training for counterfactual inference is directly linked to worse expected performance in treatment effect estimation (Figure S1). We do not observe the same behavior in standard supervised learning tasks because we have direct access to the underlying objective, and are therefore able to select the best encountered model regardless of the variance during optimisation, i.e. models that during training do not perform well momentarily will not be selected because their measured error is high. We confirmed experimentally that gradient steps are more likely to be in opposing directions when not using PM (Figure S2).
ε P EH
E
Supervised Setting Counterfactual Setting Low Variance High Variance Low Variance High Variance
Figure S1: Illustration of the impact of added variance introduced by minibatch SGD during training. In the supervised setting (left), high variance does in general not impact the final performance in terms of expected factual error (red), because we can select (dotted black line) the best encountered model during training using the observed error. In the counterfactual setting (right), we are not able to compute the true counterfactual error (blue) in observational data, and must therefore resort to selecting based on the factual error or a measure derived from the factual error. Even though we encounter models that perform similarly well, we are less likely to select these models under high variance, because we do not have direct access to the counterfactual error and the minimum of the factual error in general does not correspond to the minimum of the counterfactual error, particularly when optimisation of counterfactual and factual error are not well aligned. The dashed blue line indicates the counterfactual error of the model selected based on factual error.
M SE
εPEHEΔ εPEHEΔ εPEHEΔ εPEHEΔ
Δ
M SE Δ
M SE Δ
M SE Δ
PM PSMPM PSMMIDesired
lo w er hi gh er
higherlower higherlower higherlower higherlower
lo w er hi gh er
lo w er hi gh er
lo w er hi gh er
Figure S2: Comparison of the frequencies of individual gradient steps during training of several matching-based methods on the validation set of News-8. The heatmap entries show the frequencies of pairs of differences in factual error (∆MSE; y-axis) and counterfactual error (∆ PEHE; x-axis) after individual gradient steps. All methods used the same model and only differed in how they addressed treatment assignment bias. The leftmost heatmap shows the desired behavior of gradient steps changing the counterfactual and factual error jointly during training. PM most closely resembles the desired behavior of jointly optimising counterfactual and factual error. In contrast, PSMPM and PSMMI have a much higher relative frequency of gradient steps that do not move into the same magnitude or direction (off-diagonal entries) for both factual and counterfactual error.

heading: E NEAREST NEIGHBOUR APPROXIMATION OF PEHE FOR MULTIPLE
text: TREATMENTS (NN-MPEHE)
The ̂NN-PEHE metric can be extended to the multiple treatment setting by considering the mean ̂NN-PEHE between all ( k 2 ) possible pairs of treatments:
̂NN-mPEHE = 1( k 2
) k−1∑
i=0
i−1∑
j=0
̂NN-PEHE,i,j (E.1)

heading: F PSEUDOCODE PROPENSITY SCORE MATCHING (PSM)
text: Algorithm S2 outlines the procedure of preprocessing a training set using PSMMI. For PSMPM, we used the PM matching algorithm (see Algorithm 1) to find matches between samples from the treatment groups. For PSMMI, we used the MatchIt package (Ho et al. (2011)) on setting ”nearest” to find 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 xi across all treatment groups are approximately balanced.
Input: N training samples Xtrain with assigned treatments t, Number of treatment options k, function get matched samples to find nearest matches by propensity score between sets of samples Output: Preprocessed training set Xout consisting of matched samples from each treatment group 1: procedure PREPROCESS TRAINING SET: 2: tbase ← index of treatment with smallest number of samples in Xtrain 3: Xbase ← all samples in Xtrain with treatment tbase 4: for i = 0 to k − 1 do 5: if i 6= tbase then 6: Xcurrent ← all samples in Xtrain with treatment i 7: Xmatched ← get matched samples(Xbase, Xcurrent) 8: Add all samples in Xmatched to Xout 9: Add all samples in Xbase to Xout

heading: G TARNET, CFRNET AND PD FOR MULTIPLE TREATMENTS
text: 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 H). We did not use batch normalisation (BN). To extend CFRNET to the multiple treatment setting, we defined the first 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. (2017)). 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. (2017)). The ”+ MLP” model was a simple MLP with L hidden layers of M hidden units that received the treatment option index tj 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.

heading: H HYPERPARAMETERS
text: 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, PSMPM, PSMMI, including ”+ on X” and ”+ MLP”, and PM), we chose hyperparameters at random from predefined 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 fixed 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. (2016)) 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 S2: Hyperparameter ranges used in the performed experiments.
Hyperparameter IHDP News/TCGA
Batch size B 4, 8, 50, 100 50 Number of units per hidden layer M 50, 100, 200 40, 60, 80 Number of hidden layers L 1, 2, 3 2, 3
I POLICY RISK RPOL(π)
We define the policy risk RPol(π) as the expected loss in value when treating according to the policy π implied by an ITE estimator (Shalit et al. (2017)):
RPol(π) = 1− (E[y1|π(X) = 1] · p(π = 1) + E[y0|π(X) = 0] · p(π = 0)) (I.1) where the policy π(X) for a sampleX implied by an ITE estimator that produces potential outcomes ŷ is defined as:
π(X) = { 1 if ŷ1,X − ŷ0,X > λ (to treat) 0 otherwise (to not treat)
(I.2)

heading: J ORIGINS OF PERFORMANCE NUMBERS IN TABLE 3
text: 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 S3: Origins of the performance results reported in Table 3. Entries show the reference to the original report of the performance results, or ”ours” if we produced the results. n.r. = not reported
Method IHDP Jobs
PM ours ours + on X ours ours + MLP ours ours
kNN (Shalit et al. 2017) (Shalit et al. 2017) PSMPM n.r. ours PSMMI ours ours
RF (Shalit et al. 2017) (Shalit et al. 2017) CF (Shalit et al. 2017) (Shalit et al. 2017) BART (Shalit et al. 2017) (Shalit et al. 2017)
GANITE (Yoon et al. 2018) (Yoon et al. 2018) BNN (Shalit et al. 2017) (Shalit et al. 2017) PD n.r. n.r. TARNET (Shalit et al. 2017) (Shalit et al. 2017) CFRNETWass (Shalit et al. 2017) (Shalit et al. 2017)
CEVAE (Louizos et al. 2017) (Louizos et al. 2017) CMGP (Yoon et al. 2018) (Yoon et al. 2018)

Review summaries:
This paper proposes an augmentation of traditional neural network learning to allow for the inference of causal effects. Specifically, they modify the data sampling procedure of SGD during training to use matched samples that are paired via propensity score matching. Experimental results on a number of dataset show that the proposed methodology is comparable to alternative machine learning based causal inference methods.