Individualised Dose-Response Estimation using Generative Adversarial Nets

1 INTRODUCTION . Most of the methods developed in the causal inference literature focus on learning the effects of binary or categorical treatments ( Bertsimas et al. , 2017 ; Alaa et al. , 2017 ; Alaa & van der Schaar , 2017 ; Athey & Imbens , 2016 ; Wager & Athey , 2018 ; Yoon et al. , 2018 ) . These treatments , though , are often administered at a certain dosage which can take on continuous values ( such as vasopressors ( Döpp-Zemel & Groeneveld , 2013 ) ) . In medicine , using a high dosage of a drug can lead to toxic effects while using a low dosage can result in no effect on the patient outcome ( Wang et al. , 2017 ) . Moreover , the dosage levels used when choosing between multiple treatments for a patient are crucial for the decision ( Rothwell et al. , 2018 ) . While admissible dosage intervals for drugs are often determined from clinical trials ( Cook et al. , 2015 ) , these trials often have a small number of patients and use simplistic mathematical models to assign dosage levels to patients that do not take into account patient heterogeneity ( Ursino et al. , 2017 ) . After drugs are approved through clinical trials , observational data collected about different treatment dosages prescribed to a diverse set of patients offers us the opportunity to learn individualized responses . As the relationships between treatment dosage efficacy , toxicity and patient features become more complex , estimating dose-response from observational data becomes particularly important in order to identify optimal dosages for each patient . Fortunately , there is a wealth of observational data available in the medical domain from electronic health records ( Henry et al. , 2016 ) . Learning from observational data already presents significant challenges in the binary treatment setting . As explained by Spirtes ( 2009 ) , in an observational treatment-effect dataset , only the factual outcome is present ( i.e . the outcome for the treatment that was actually given ) - the counterfactual outcomes are not observed . This problem is exacerbated in the dose-response setting in which the number of counterfactuals is no longer even finite . Moreover , the treatment assignment is non-random and instead is assigned according to the features associated with each sample . Due to the continuous nature of the dosage parameter , adjusting for the bias in the dosage assignments is significantly more complex than for binary ( or even multiple ) treatments . Thus , standard methods for adjusting for treatment selection bias can not be easily extended to handle bias in the dosage parameter . In this paper we address the problem of dose-response estimation from observational data by building on the framework introduced in GANITE ( Yoon et al. , 2018 ) . The key idea is to modify the GAN framework ( Goodfellow et al. , 2014 ) to generate the unobserved counterfactual outcomes from a standard treatment-effect dataset . Already , GANITE presents a significant modification to the original GAN framework - rather than the discriminator discriminating between entirely real or entirely fake samples , the discriminator is attempting to pick out the real component from a vector containing the real ( factual ) outcome from the dataset and the fake ( counterfactual ) outcomes generated by the generator . We also inherit this key difference from a standard GAN , but in addition we must make further modifications to the original GANITE framework in order to address the dosage problem . A naive attempt to extend Yoon et al . ( 2018 ) to the dosage setting might involve trying to define a discriminator that takes as input an entire dose-response curve for each treatment from the generator ( with the outcome for the observed treatment-dosage pair replacing the generated one ) and that tries to determine the factual treatment-dosage pair . This fails for two reasons : ( 1 ) we do not wish to assume prior knowledge of the functional form of the dose-response curves and so will have access to the generated dose-response curves only by evaluating them at given points ( and so `` entire '' dose-response curves can not be passed to the discriminator ) ; ( 2 ) substituting the generator output for the factual treatment-dosage pair with the factual outcome will almost always create a discontinuity in the response curve and thus the factual treatment-dosage pair would be very easy to identify . We overcome these two hurdles by defining a discriminator that , rather than acting on the entire dose-response curves , acts on a finite set of points from each curve , as shown in Fig . 1 . From among the chosen points , the discriminator will then attempt to pick out the factual one . To ensure that the entire dose-response curve is well-estimated , we sample the set of points randomly each time an input would be passed to the discriminator . If we were to fix a set of points in advance to compare for all treatments and samples then only the outcomes associated with these dosage levels would be well estimated . As our discriminator will be taking as input a set of random dosage-outcome pairs , we need to condition its behaviour to be like that of a function on a set . In particular , we draw on ideas from Zaheer et al . ( 2017 ) to ensure that the discriminator acts as a function on sets and its output does not depend on the order in which the elements of the set are given as input . In addition , we model the generator as a multi-task deep network capable of taking dosages as an input ; this gives us the flexibility to learn heterogeneous dose-response curves for the different treatments . We also develop a hierarchical discriminator which breaks down the job of the discriminator into determining the factual treatment and determining the factual dosage using separate networks . We show in the experiments section that this approach significantly improves performance and is more stable than using a single network discriminator . Our contributions in this paper are 3-fold : ( 1 ) we propose DRGAN , a significantly modified GAN framework , capable of dose-response estimation , ( 2 ) we propose novel architectures for each of our networks , ( 3 ) we propose a new semi-synthetic data simulation for use in the dose-response setting . We show , using semi-synthetic experiments , that our model outperforms existing benchmarks . 2 RELATED WORK . Methods for estimating the outcomes of treatments with an exposure dosage parameter that only employ observational data make use of the generalized propensity score ( GPS ) ( Imbens , 2000 ; Imai & Van Dyk , 2004 ; Hirano & Imbens , 2004 ) or build on top of balancing methods for multiple treatments . Schwab et al . ( 2019 ) developed a neural network based method to estimate counterfactuals for multiple treatments and continuous dosages . The proposed Dose Response networks ( DRNets ) in Schwab et al . ( 2019 ) consist of a three level architecture with shared layers for all treatments , multi-task layers for each treatment and additional multi-task layers for dosage sub-intervals . More specifically , for each treatment w , the dosage interval [ aw , bw ] is subdivided into E equally sized sub-intervals and a multi-task head is added for each sub-interval . Their model architecture extends the one in Shalit et al . ( 2017 ) by adding the multi-task heads for the dosage strata . However , the main advantage of using multi-task heads for dosage intervals would be the added flexibility in the model to learn potentially very different functions over different regions of the dosage interval . DRNets does not determine the dosage intervals dynamically and thus much of this flexbility is lost . We demonstrate in our experiments that DRGAN outperforms both GPS and DRNets . For a discussion of works that address treatment-response estimation without a dosage parameter , see Appendix A . Note that for such methods we can not treat the dosage as an additional input due to the bias associated with its assignment . 3 PROBLEM FORMULATION . We consider receiving observations of the form ( xi , tif , y i f ) for i = 1 , ... , N , where , for each i , these are independent realizations of the random variables ( X , Tf , Yf ) . We refer to X as the feature vector lying in some feature space X , containing pre-treatment covariates ( such as age , weight and lab test results ) . The treatment random variable , Tf , is in fact a pair of values Tf = ( Wf , Df ) where Wf ∈ W corresponds to the type of treatment being administered ( e.g . chemotherapy or radiotherapy ) which lies in the discrete space of k treatments , W = { w1 , ... , wk } , andDf corresponds to the dosage of the treatment ( e.g . number of cycles , amount of chemotherapy , intensity of radiotherapy ) , which , for a given treatment w lies in the corresponding treatment ’ s dosage space , Dw , which in the most general case is some continuous space ( e.g . the interval [ 0 , 1 ] ) . We define the set of all treatmentdosage pairs to be T = { ( w , d ) : w ∈ W , d ∈ Dw } . Following Rubin ’ s potential outcome framework ( Rubin , 1984 ) , we assume that for all treatmentdosage pairs , ( w , d ) , there is a potential outcome Y ( w , d ) ∈ Y ( e.g . 1-year survival probability ) . The observed outcome is then defined to be Yf = Y ( Wf , Df ) . We will refer to the unobserved ( potential ) outcomes as counterfactuals . The goal of dose-response estimation is to derive unbiased estimates of the potential outcomes for a given set of input covariates : µ ( t , x ) = E [ Y ( t ) |X = x ] ( 1 ) for each t ∈ T , x ∈ X . We refer to µ ( · ) as the individualised dose-response function . In general , this quantity is not the same as E [ Y |X = x , Tf = t ] ( which can be easily estimated from observational data ) in the presence of selection bias which often presents itself in observational datasets . In order for these two quantities to be equal , we must make the following common assumption . Assumption 1 . ( Unconfoundedness ) The treatment assignment , Tf , and potential outcomes , Y ( w , d ) , are conditionally independent given the covariates X , i.e . { Y ( w , d ) |w ∈ W , d ∈ Dw } ⊥⊥ Tf |X . ( 2 ) This assumption is also commonly referred to as no hidden confounding . In addition , in order to make µ ( · ) identifiable we must also assume that any treatment-dosage pair could be assigned to any given sample . Assumption 2 . ( Overlap ) For each x ∈ X such that p ( x ) > 0 , we have 1 > p ( t|x ) > 0 for each t ∈ T . 4 DOSE-RESPONSE GAN . We propose estimating µ by first training a generator to generate dose-response curves for each sample within the training dataset . The learned generator can then be used to train an inference network using standard supervised methods . We build on the idea presented in Yoon et al . ( 2018 ) , using a modified GAN framework to generate potential outcomes conditional on the observed features , treatment and factual outcome . Several changes must be made to both the generator and discriminator architectures and learning paradigms in order to produce a model capable of handling the dose-response setting .

Summary: The paper proposes a generative adversarial net model for heterogeneous dose-response causal effect estimation. The idea is to generate counterfactual dose-response curves using the generator that can fool the discriminator that tries to distinguish between factual data and counterfactual data. Factual data along with counterfactual data generated by the GAN can then be used for heterogeneous causal effect estimation.