Perfect Match: A Simple Method for Learning Representations For Counterfactual Inference With Neural Networks

1 INTRODUCTION . 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 . 2 RELATED WORK . 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 . 3 METHODOLOGY . 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

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