Title: PERFECT MATCH: A SIMPLE METHOD FOR LEARN-ING REPRESENTATIONS FOR COUNTERFACTUAL IN-FERENCE WITH NEURAL NETWORKS

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 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. 1 The ITE is sometimes referred to as Conditional Average Treatment Effect (CATE).

1. INTRODUCTION:

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. (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.

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(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.

Table 1 : Comparison of several representative state-of-the-art methods for counterfactual inference showing whether they were designed for > 2 treatments and their complexity of implementation.

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 Ŷ with k entries ŷ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 ŷ1 (n) -ŷ0 (n) indexed by n over N samples:

PEHE = 1 N N n=0 E yj (n)∼µj (n) [y 1 (n) -y 0 (n)] -[ŷ 1 (n) -ŷ0 (n)] 2 (1)

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 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.

EQUATION

Perfect Match (PM). We consider fully differentiable neural network models optimised via minibatch stochastic gradient descent to predict potential outcomes Ŷ 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 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 x i 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 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. (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-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 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 & Richtá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 x i across all treatment groups are approximately balanced.

Input: Batch of B random samples X batch with assigned treatments t, training set X train of N samples, number of treatment options k, propensity score estimator E PS to calculate the probability p(t|X) of treatment assigned given a sample X Output: Batch X out consisting of B × k matched samples 1: procedure PERFECT MATCH:

2:

X out ← Empty 3:

for sample X with treatment t in X batch do 4:

p(t|X) ← E PS (X)

5:

for i = 0 to k -1 do 6: if i = t then 7: ps i ← p(t|X) i 8:

X matched ← get closest match to propensity score ps i with treatment i from X train 9:

Add sample X matched to X out 10:

Add X to X out 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 y j from a respective nearest neighbour NN matched on X using the Euclidean distance.

EQUATION

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 t j may perform poorly if X is high-dimensional, because the influence 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 jth 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 influence on the network output at all points during training.

X t { L shared layers { L head layers } k heads ŷ 1 ŷ k-1 ŷ k t=0 t=k-2 t=k-1

Figure 1 : The TARNET architecture with k heads for the multiple treatment setting. 4

Under review as a conference paper at ICLR 2019 

4. EXPERIMENTS:

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?

4.1. DATASETS:

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 outcomes 2 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.

2 Available at: http://www.mit.edu/ ˜fredrikj/files/IHDP-1000.tar.gz q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q 0 100 200 300 400 0 20 40 60 80 MSE PEHE ρ=0.34 q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q qq q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q qq q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q qq q q q q q q q qq q q q q q q q q q q q q q q q q qq q q q q qq q q q q q q q q q q q q q q q q q q q q q q q q q q q q qq q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q qq q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q qq q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q qq q q q q q q q q q q q q q qq q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q qq q q q q q q q q q q q q q q q q qq q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q qq q q q q q q q q q q qq q q q q q q q q q q q q q q q q q q q q q q q q qq q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q qq q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q qq q q q q q q q q q q q q qq q q q q q q q q q q qq q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q qq q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q qq q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q qq q q q q q q q q q q qq q q q q q q q q q q q q q q q q q q q q q q q qq q q qq q q q q q q q qq q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q qq q q q q q q q q q q q q q q q q q q q q q q q qq q q q q q q q q q q qq q q q q q q q q q q q q q q q qq q q q q q q q q q q q q q q q q q q q qq q q q q q q qq q q q q q q q q q q q q q q q q q q q q q q q q q q qq q q q q q q q q q q q q q q q q q q q qq q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q 2 4 6 ε PEHE q q q q q q q q q q q Figure 3: Change in error (y-axes) in terms of precision in estimation of heterogenous effect (PEHE) and average treatment effect (ATE) when increasing the percentage of matches in each minibatch (x-axis). Symbols correspond to the mean value of ˆ ATE (red) and √ ˆ PEHE (blue) on the test set of News-8 across 50 repeated runs with new outcomes (lower is better). Performance improves with more matches added.

Under review as a conference paper at ICLR 2019 Treatment Assignment Bias κ ε ATE q q q q q q q q PM CFRNET TARNET PD CF 6 8 10 12 14 16 18 20 18 20 22 24 26 28 30 Treatment Assignment Bias κ ε PEHE q q q q q q q q PM CFRNET TARNET PD CF 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 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. (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 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 ỹj ∼ N (µ j , σ j )+ with ∼ N (0, 0.15). We then defined the unscaled potential outcomes ȳj = ỹj * [D(z(X), z j ) + D(z(X), z c )] as the ideal potential outcomes ỹ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(κ ȳj )) with a treatment assignment bias coefficient κ, and the true potential outcome y j = C ȳj as the unscaled potential outcomes ȳ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 ȳ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 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. For each sample, we drew ideal potential outcomes from that Gaussian outcome distribution ỹj ∼ N (µ j , σ j ) + with ∼ N (0, 0.15).

We then defined the unscaled potential outcomes ȳj = ỹj * [D(X, X tj ) + D(X, X c )] as the ideal potential outcomes ỹj weighted by the sum of distances to phenotype X tj 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(κ ȳj )) with a treatment assignment bias coefficient κ = 10, and the true potential outcome y j = C ȳj as the unscaled potential outcomes ȳ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. 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 (PSM PM ) and the "MatchIt" package (PSM MI , 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 t j as an input instead of using a TARNET. ε ATE q q q q q q q q q q PM CFRNET TARNET PD 20 40 60 80 10 12 14 16 18 20 22 Percentage of Hidden Confounding [%] ε PEHE q q q q q q q q q q PM CFRNET TARNET PD ˆ PEHE (right) across 5 repeated runs with new outcomes (lower is better). The shaded area indicates the standard deviation. PM, TARNET and CFRNET are relatively stable up to 60% of hidden confounding. The methods degrade strongly in counterfactual prediction performance at levels of hidden confounding higher than 60%.

Under review as a conference paper at ICLR 2019 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 TAR-NET 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 = 2018)). The policy risk is the average loss in value when treating according to the policy π implied by an ITE estimator (Appendix I).

5. RESULTS AND DISCUSSION:

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 R Pol (π) 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 PSM MI , 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. 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. PSM PM shows a weaker degree of alignment and much higher variance. In contrast, PSM MI 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 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. (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, 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 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) ).

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.

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. 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 ). 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. 

E NEAREST NEIGHBOUR APPROXIMATION OF PEHE FOR MULTIPLE 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: t base ← index of treatment with smallest number of samples in X train 3:

ˆ NN-mPEHE = 1 k 2 k-1 i=0 i-1 j=0 ˆ NN-PEHE,i,j (E.

X base ← all samples in X train with treatment t base 4:

for i = 0 to k -1 do 5:

if i = 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 13

Under review as a conference paper at ICLR 2019 G TARNET, CFRNET AND PD FOR MULTIPLE TREATMENTS

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 t j as an additional input along with the covariates X, and output a k-dimensional potential outcome vector Ŷ . The MLP used ELU nonlinearities between the L hidden layers, and also did not use BN.

H HYPERPARAMETERS:

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 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 "bart-Machine" (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. where the policy π(X) for a sample X implied by an ITE estimator that produces potential outcomes ŷ is defined as: π(X) = 1 if ŷ1,X -ŷ0,X > λ (to treat) 0 otherwise (to not treat) (I.2)

J ORIGINS OF PERFORMANCE NUMBERS IN TABLE 3:

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. 

:

https://archive.ics.uci.edu/ml/datasets/bag+of+words

Figures and Tables:

Figure 2: Correlation analysis of the real PEHE (y-axis) with the mean squared error (MSE; left) and the nearest neighbour approximation of the precision in estimation of heterogenous effect (NN-PEHE; right) across over 20'000 model evaluations on the validation set of IHDP. Scatterplots show a subsample of 1'400 data points. ρ indicates the Pearson correlation. NN-PEHE correlates significantly better with PEHE than MSE.

Figure4: Comparison of several state-of-the-art methods for counterfactual inference on the test set of the News-8 dataset when varying the treatment assignment imbalance κ (x-axis), i.e. how much the treatment assignment is biased towards more effective treatments. Symbols correspond to the mean value of ˆ ATE (left) and √ ˆ PEHE (right) across 50 repeated runs with new outcomes (lower is better). The shaded area indicates the standard deviation. PM handles an increasing treatment assignment bias κ better than existing state-of-the-art methods in terms of both ˆ ATE and √ ˆ PEHE .

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 (CFRNET Wass ) (Shalit et al. (2017)), CEVAE (Louizos et al. (2017)), CMGP (Alaa & van der Schaar (2017)), and PD (Alaa et al. (

Figure5: Comparison of several state-of-the-art methods for counterfactual inference on the test set of the TCGA dataset when varying the percentage of hidden confounding (x-axis), i.e. what percentage of the features that were relevant for treatment assignment are visible to the model. Symbols correspond to the mean value of ˆ ATE (left) and √ ˆ PEHE (right) across 5 repeated runs with new outcomes (lower is better). The shaded area indicates the standard deviation. PM, TARNET and CFRNET are relatively stable up to 60% of hidden confounding. The methods degrade strongly in counterfactual prediction performance at levels of hidden confounding higher than 60%.

(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 ŷ1 (n)|, and the policy risk R Pol (π) (Shalit et al. (2017); Yoon et al. (

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 the mean value of the factual error ( √ MSE; red) and the

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.12

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, PSM PM and PSM MI 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.

R POL (π)We define 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. (2017)):R Pol (π) = 1 -(E[y 1 |π(X) = 1] • p(π = 1) + E[y 0 |π(X) = 0] • p(π = 0)) (I.1)

Comparisons of the datasets used in our experiments. We evaluate on three semi-synthetic datasets with varying numbers of treatments and samples, and a real-world clinical trial dataset.

Table content: <table><tr><td>Dataset</td><td colspan="5">Type # Samples # Features # Treatments Counterfactuals</td></tr><tr><td>IHDP Jobs News TCGA</td><td>semi-synthetic real-world semi-synthetic semi-synthetic</td><td>747 3212 5000 9659</td><td>25 7 2870 20531</td><td>2 2 2/4/8/16 4</td><td>available not available available available</td></tr></table>

Comparison of methods for counterfactual inference with two available treatment options on IHDP, Jobs and News-2. We report the mean value ± the standard deviation of √ PEHE , ATE , R Pol and ATT on the test sets over a number of repeated runs. Where available we list the numbers reported by the original authors. n.r. = not reported. † = significantly different from PM (α < 0.05)

Table content: <table><tr><td>Method</td><td>IHDP (1000 repeats) √ PEHE</td><td>ATE</td><td>Jobs (10 repeats) RPol(π)</td><td>ATT</td><td>News-2 (50 repeats) √ ˆ PEHE</td><td>ˆ ATE</td></tr><tr><td colspan="7">PM + on X + MLP kNN PSMPM PSMMI RF CF BART GANITE BNN PD TARNET CFRNETWass  † 0.76 ± 0.02  † 0.27 ± 0.01 0.21 ± 0.0  † 0.09 ± 0.03 16.93 ± 1.12  † 4.54 ± 1.48 0.84 ± 0.61 0.24 ± 0.20 0.18 ± 0.10 0.16 ± 0.07 16.76 ± 1.26 3.99 ± 1.01 0.81 ± 0.57 0.24 ± 0.20 0.23 ± 0.11 0.19 ± 0.09 17.06 ± 1.22 4.14 ± 1.27 0.83 ± 0.57 0.23 ± 0.20 0.19 ± 0.12 0.16 ± 0.07  † 18.38 ± 1.46  † 5.90 ± 2.07  † 4.1 ± 0.2  † 0.79 ± 0.05  † 0.26 ± 0.0 0.13 ± 0.05  † 18.14 ± 1.64  † 7.83 ± 2.55 n.r. n.r. 0.18 ± 0.12 0.18 ± 0.09  † 17.49 ± 1.49  † 5.02 ± 2.34  † 2.70 ± 3.85  † 0.49 ± 0.81 0.20 ± 0.09  † 0.42 ± 0.22  † 17.40 ± 1.30  † 4.89 ± 2.39  † 6.6 ± 0.3  † 0.96 ± 0.06  † 0.28 ± 0.0  † 0.09 ± 0.04  † 17.39 ± 1.24  † 5.50 ± 1.20  † 3.8 ± 0.2  † 0.40 ± 0.03 0.20 ± 0.0  † 0.07 ± 0.03  † 18.36 ± 1.73  † 8.48 ± 2.46  † 2.3 ± 0.1  † 0.34 ± 0.02  † 0.25 ± 0.0  † 0.08 ± 0.03  † 18.53 ± 2.02  † 5.40 ± 1.53  † 2.4 ± 0.4  † 0.49 ± 0.05 0.14 ± 0.01  † 0.06 ± 0.03  † 18.28 ± 1.66  † 4.65 ± 2.12 n.r. n.r.  † 2.1 ± 0.1  † 0.42 ± 0.03 0.24 ± 0.0  † 0.09 ± 0.04 n.r. n.r. n.r. n.r.  † 17.52 ± 1.62 4.69 ± 3.17  † 0.95 ± 0.02  † 0.28 ± 0.01 0.21 ± 0.0 0.11 ± 0.04 17.17 ± 1.25  † 4.58 ± 1.29 CEVAE  † 2.7 ± 0.1  † 0.46 ± 0.02  † 0.26 ± 0.0  † 0.03 ± 0.01 n.r. n.r. CMGP  † 0.77 ± 0.11  † 0.13 ± 0.12 0.24 ± 0.05  † 0.09 ± 0.07 n.r. n.r.</td></tr></table>

Comparison of methods for counterfactual inference with more than two available treatment options on the News-4, News-8 and News-16. Numbers are mean value ± the standard deviation of √ ˆ mPEHE and ˆ mATE over 50 repeated runs. † = significantly different from PM (α < 0.05) 25.05 ± 2.80 † 14.92 ± 3.14 † 24.88 ± 1.98 † 14.00 ± 2.45 † 27.05 ± 2.47 † 16.90 ± 2.01 kNN † 27.92 ± 2.44 † 19.40 ± 3.12 † 26.20 ± 2.18 † 15.11 ± 2.34 † 27.64 ± 2.40 † 17.27 ± 2.10 PSMPM † 22.74 ± 2.58 † 11.62 ± 2.69 † 22.16 ± 1.79 † 8.81 ± 1.96 † 23.57 ± 2.48 † 11.04 ± 2.50 PSMMI † 37.26 ± 2.28 † 30.19 ± 2.47 † 30.50 ± 1.70 † 22.09 ± 1.98 † 28.17 ± 2.02 † 18.81 ± 1.74 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 CF † 27.28 ± 2.59 † 19.04 ± 3.15 † 25.76 ± 2.05 † 15.12 ± 2.20 † 26.85 ± 2.29 † 17.03 ± 1.99 BART † 26.41 ± 3.10 † 17.14 ± 3.51 † 25.78 ± 2.66 † 14.80 ± 2.56 † 27.45 ± 2.84 † 17.50 ± 2.49 GANITE † 24.50 ± 2.27 † ± 2.69 † 23.58 ± 2.48 † 11.20 ± 2.84 † 25.12 ± 3.53 † 13.20 ± 3.28 PD 20.88 ± 3.24 † 8.47 ± 4.51 21.19 ± 2.29 7.29 ± 2.97 † 22.28 ± 2.25 † 10.65 ± 2.22 TARNET † 23.40 ± 2.20 † 13.63 ± 2.18 † 22.39 ± 2.32 † 9.38 ± 1.92 † 21.19 ± 2.01 † 8.30 ± 1.66 CFRNETWass † 22.65 ± 1.97 † 12.96 ± 1.69 † 21.64 ± 1.82 † 8.79 ± 1.68 † 20.87 ± 1.46 † 8.05 ± 1.40

Table content: <table><tr><td>Method</td><td/><td>News-4 (κ = 10) √ ˆ mPEHE</td><td>ˆ mATE</td><td colspan="2">News-8 (κ = 10) √ ˆ mPEHE</td><td>ˆ mATE</td><td>News-16 (κ = 7) √ ˆ mPEHE</td><td>ˆ mATE</td></tr><tr><td>PM + on X + MLP</td><td/><td colspan="3">21.58 ± 2.58 10.04 ± 2.71 20.76 ± 1.86 21.41 ± 1.75  † 8.91 ± 2.00 20.90 ± 2.07</td><td colspan="3">6.51 ± 1.66 20.24 ± 1.46 6.36 ± 1.65 20.67 ± 1.42</td><td>5.76 ± 1.33 5.73 ± 1.16</td></tr></table>

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 find matches between samples from the treatment groups. For PSM MI , 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). 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: N training samples X train 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 X out consisting of matched samples from each treatment group 1: procedure PREPROCESS TRAINING SET:

Table content: <table><tr><td>1)</td></tr><tr><td>F PSEUDOCODE PROPENSITY SCORE MATCHING (PSM)</td></tr><tr><td>Algorithm</td></tr></table>

r. = not reported

Table content: <table><tr><td>Method</td><td>IHDP</td><td>Jobs</td></tr><tr><td>PM + on X + MLP</td><td>ours ours ours</td><td>ours ours ours</td></tr><tr><td>kNN PSMPM PSMMI</td><td>(Shalit et al. 2017) n.r. ours</td><td>(Shalit et al. 2017) ours ours</td></tr><tr><td>RF CF BART</td><td>(Shalit et al. 2017) (Shalit et al. 2017) (Shalit et al. 2017)</td><td>(Shalit et al. 2017) (Shalit et al. 2017) (Shalit et al. 2017)</td></tr><tr><td>GANITE BNN PD TARNET CFRNETWass</td><td>(Yoon et al. 2018) (Shalit et al. 2017) n.r. (Shalit et al. 2017) (Shalit et al. 2017)</td><td>(Yoon et al. 2018) (Shalit et al. 2017) n.r. (Shalit et al. 2017) (Shalit et al. 2017)</td></tr><tr><td>CEVAE CMGP</td><td colspan="2">(Louizos et al. 2017) (Louizos et al. 2017) (Yoon et al. 2018) (Yoon et al. 2018)</td></tr></table>

Peer review (iclr2019):
 Summary: This paper proposed to extend TARNET (Shalit et al. 2017), a representation learning approach for counterfactual inference, in the following ways. First, to extend TARNET to multiple treatment setting, k head networks (instead of 2) were constructed following the shared MLP layers, where each head network modeled the outcome of one treatment. This extension seemed quite straightforward. Second, during training, for every sample in a minibatch, find its nearest neighbors from all other treatments and add them to the minibatch. The distance was measured by the propensity score, which was defined the probability of a sample being assigned to a treatment group and could be learned by a classification model (such as support vector machine used in this work). Therefore, 1) the augmented minibatch would contain the same number of samples for each treatment group; 2) different treatment group were balanced. Third, a model selection strategy was proposed by estimating the PEHE using nearest neighbor search. Comments: This paper is well motivated. The key challenges in counterfactual inference is how to adjust for the bias in treatment assignment and the associated discrepancies in the distribution of different treatment groups. The main idea of this paper, i.e., augmenting the minibatch through propensity score matching for each sample, is well explained in Section 3. However, it could be better if the introduction of model architecture (in Appendix F) was presented in the method section. Did the author need to train (k choose 2) SVMs to compute the propensity scores for samples from k treatment groups? When comparing different approaches, as were shown in Table 3, 4 and Figure 3,4, did the author run any statistical test, such as t-test, to confirm the difference between those distributions were significant? The standard deviations of those errors seemed quite large so the difference could be non-significant. Could the author provide more explanations on why the proposed approach, i.e., minibatch augmentation using propensity score matching, can outperform the TARNET? In TARNET, each sample it only used to update the head network corresponding to the sample*s treatment assignment, why would balancing samples in the minibatch can improve the estimation of treatment effect?