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 introduces Dose Response Generative Adversarial Network (DRGAN) that is aimed at generating entire dose-response curve from observational data with single dose treatments. This work is an extension of GANITE (Yoon et al., 2018) for the case of real-valued treatments (i.e., dosage). The proposed model consists of 3 blocks: (1) a generator, (2) a discriminator, and (3) an inference block. In this paper, GANITE’s generator and discriminator architectures are modified to be able to handle real-valued treatments.