Year: 2019

Abstract:
The problem of estimating treatment responses from observational data is by now a well-studied one. Less well studied, though, is the problem of treatment response estimation when the treatments are accompanied by a continuous dosage parameter. In this paper, we tackle this lesser studied problem by building on a modification of the generative adversarial networks (GANs) framework that has already demonstrated effectiveness in the former problem. Our model, DRGAN, is flexible, capable of handling multiple treatments each accompanied by a dosage parameter. The key idea is to use a significantly modified GAN model to generate entire dose-response curves for each sample in the training data which will then allow us to use standard supervised methods to learn an inference model capable of estimating these curves for a new sample. Our model consists of 3 blocks: (1) a generator, (2) a discriminator, (3) an inference block. In order to address the challenge presented by the introduction of dosages, we propose novel architectures for both our generator and discriminator. We model the generator as a multitask deep neural network. In order to address the increased complexity of the treatment space (because of the addition of dosages), we develop a hierarchical discriminator consisting of several networks: (a) a treatment discriminator, (b) a dosage discriminator for each treatment. In the experiments section, we introduce a new semi-synthetic data simulation for use in the dose-response setting and demonstrate improvements over the existing benchmark models.

heading: 1 INTRODUCTION
text: 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 cannot 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 cannot 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.

heading: 2 RELATED WORK
text: 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 cannot treat the dosage as an additional input due to the bias associated with its assignment.

heading: 3 PROBLEM FORMULATION
text: 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 .

heading: 4 DOSE-RESPONSE GAN
text: 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.

heading: 4.1 COUNTERFACTUAL GENERATOR
text: Our generator, G : X ×T ×Y ×Z → YT takes features, x ∈ X , factual outcome, yf ∈ Y , received treatment and dosage, tf = (wf , df ) ∈ T , and some noise, z ∈ Z (typically multivariate uniform or Gaussian), as inputs. The output will be a dose-response curve for each treatment (as shown in Fig. 1), so that the output is a function from T to Y , i.e. G(x, tf , yf , z)(·) : T → Y . We can then write
ŷcf (t) = G(x, tf , yf , z)(t) (3)
to denote our generated counterfactual outcome for the treatment-dosage pair t. We will write Ŷcf (t) = G(X, Tf , Yf ,Z)(t) (i.e. the random variable induced by G).
While the job of the counterfactual generator is to generate outcomes for the treatment-dosage pairs which were not observed, Yoon et al. (2018) demonstrated that the performance of the counterfactual generator is improved by adding a supervised loss term that regularises its output for the factual treatment (in our case treatment-dosage pair). We define the supervised loss, LS , to be
LS(G) = E [ (Yf −G(X, Tf , Yf ,Z)(Tf ))2 ] , (4)
where the expectation is taken over X, Tf , Yf and Z.

heading: 4.2 COUNTERFACTUAL DISCRIMINATOR
text: As noted in Section 1, our discriminator will act on a random set of points from each of the generated dose-response curves. Similar to Yoon et al. (2018), we define a discriminator, D, that will attempt to pick out the factual treatment-dosage pair from among the (random set of) generated ones.
Formally, let nw ∈ Z+ be the number of dosage levels we will compare for treatment w ∈ W1. For each w ∈ W , let D̃w = {Dw1 , ..., Dwnw} be a random subset2 of Dw of size nw, where for the factual treatment, Wf , D̃Wf contains nWf − 1 random elements along with Df . We define Ỹw = (D w i , Ỹ w i ) nw i=1 ∈ (Dw×Y)nw to be the vector of dosage-outcome pairs for treatment w where
Ỹ wi = { Yf if Wf = w and Df = Dwi Ŷcf (w,D w i ) else
(5)
and will write Ỹ = (Ỹw)w∈W . We will write dwj , ỹw and ỹ to denote realisations of D w j , Ỹw and ỹ. Our discriminator, D : X × ∏w∈W(Dw × Y)nw → [0, 1]∑nw , will take the features x ∈ X together with the (random) set of generated outcomes ỹ ∈ Y ∑ nw , and output a probability for each treatment-dosage pair indicating the discriminator’s belief that that pair is the factual one.
As in the standard GAN framework, we define a minimax game by defining the value function to be L(D,G) = E [ ∑ w∈W ∑ d∈D̃w I{Tf=(w,d)} logD w,d(X, Ỹ)+I{Tf 6=(w,d)} log(1−Dw,d(X, Ỹ)) ] , (6)
where the expectation is taken over X, Tf , Ỹ and {D̃w : w ∈ W}, Dw,d corresponds to the discriminator output for treatment-dosage pair (w, d).
1In practice we set all nw to be the same. The default setting is 5 in the experiments. 2In practice, when Dw = [0, 1], each Dwj is sampled independently and uniformly from [0, 1]. Note that for
each training iteration, D̃w is resampled (see Section 1).
The minimax game is then given by min G max D L(D,G) + λLS(G) , (7) where λ is used to control the trade-off between L and LS (we set λ = 1 in the experiments). The task of the discriminator (i.e. picking out the factual dosage from ∑k j=1 nwj treatment-dosage pairs) becomes increasingly difficult as we increase nw or k because the dimension of the discriminator output space, ∑ nw, increases. Although we control nw, if we set it too low, then the set ỹw may not well-represent the dose-response curve, particularly if the dose-response curve is complex. In practice we found that even for moderate settings of nw and only 2 treatments, modelling the discriminator as a single function resulted in poor performance. In order to overcome this problem, we introduce a novel hierarchical discriminator which involves a treatment discriminator with output dimension k and several dosage discriminators, one for each treatment, with output dimensions nw.
First observe that the probability P((Wf , Df ) = (w, d)|X, D̃w, Ỹ) can be written as P(Wf = w|X, D̃w, Ỹ)× P(Df = d|Wf = w,X, D̃w, Ỹ) . (8)
We can therefore break down the discriminator into a hierarchical model by learning one discriminator, DW , that outputs P(Wf = w|X, D̃w, Ỹ) which we will refer to as the treatment discriminator, and then a discriminator, Dw, for each treatment, w ∈ W , that outputs P(Df = d|Wf = w,X, D̃w, Ỹ) which we will refer to as the dosage discriminator for treatment w.
The treatment discriminator, DW : X× ∏ w∈W(Dw × Y)nw → [0, 1]k, takes the features, x, and generated potential outcomes, ỹ, and outputs a probability for each treatment, w1, ..., wk. Writing DwW to denote the output of DW corresponding to treatment w, we define the loss, LW , to be
LW(DW ;G) = −E [ ∑ w∈W I{Wf=w} logD w W(X, Ỹ) + I{Wf 6=w} log(1−DwW(X, Ỹ)) ] , (9)
where, again, the expectation is taken over X,Wf , Df , Ỹ and {D̃w}w∈W . Then, for each w ∈ W , Dw : X × (Dw × Y)nw → [0, 1]nw is a map that takes the features, x, and generated potential outcomes, ỹw, corresponding to treatment w and outputs a probability for each dosage level, dw1 , ..., d w nw , in a given realisation of D̃w. Writing Djw to denote the output of Dw corresponding to dosage level Dwj , we define the loss of each dosage discriminator to be
Ld(Dw;G) = −E [ I{Wf=w} nw∑ j=1 I{Df=Dwj } logD j w(X, Ỹw)+I{Df 6=Dwj } log(1−D j w(X, Ỹw)) ] ,
(10)
where the expectation is taken over X, D̃w, Ỹw,Wf and Df . The I{Wf=w} term ensures that only samples for which the factual treatment is w are used to train dosage discriminator Dw (otherwise there would be no factual dosage for that sample).
We define the overall discriminator DH : X × ∏ w∈W(Dw × Y )nw → [0, 1] ∑ nw by defining its output corresponding to the treatment-dosage pair (w, dwj ) as
Dw,jH (x, ỹ) = D w W(x, ỹ)×Djw(x, ỹw) . (11)
Instead of the minimax game in Eq. 7, the generator and discriminator are trained according to the minimax game defined by seeking G∗, D∗H that solve
G∗ = arg min G L(D∗H ;G) + λLS(G) D∗Hw,j = D∗Ww ×D∗wj D∗W = arg min DW LW(DW ;G∗) D∗w = arg min Dw Ld(Dw;G∗),∀w ∈ W (12)
Fig. 2 depicts our generator and hierarchical discriminator. Pseudo-code for our algorithm can be found in Appendix C.

heading: 4.3 INFERENCE NETWORK
text: Once we have learned the counterfactual generator, we can use it only to access (generated) doseresponse curves for all samples in the dataset. To generate dose-response curves for a new sample we use the counterfactual generator along with the original data to train an inference network, I : X × T → Y . Details of the loss and pseudo-code can be found in Appendix D.

heading: 5 ARCHITECTURE
text: In this section, we describe in detail the novel architectures that we adopt to model each of the functions G,D,DW ,Dw1 , ...,Dwk which draws from the ideas in Zaheer et al. (2017). The inference network, I, has the same architecture as the generator, but does not receive wf , df , yf or z as inputs.
5.1 GENERATOR ARCHITECTURE
We adopt a multi-task deep learning model for G by defining a function g : X × T × Y × Z → H for some latent space H (typically Rl for some l) and then for each treatment w ∈ W we introduce a multitask "head", gw : H×Dw → Y taking inputs from H and a dosage, d, to produce an outcome ŷ(w, d) ∈ Y . Given observations, (x, tf , yf ), a noise vector z, and a target treatment-dosage pair, t = (w, d), we define
G(x, tf , yf , z)(t) = gw(g(x, tf , yf , z), d) . (13)
Each of g, gw1 , ..., gwk are modelled as fully connected networks. Fig. 3 depicts our generator architecture.

heading: 5.2 DISCRIMINATOR ARCHITECTURES
text: As noted in Section 1, our discriminators need to act as functions of sets (of randomly selected dosageoutcome pairs). While we could require that our discriminators try to learn this during training, by enforcing them to be functions of sets through their architecture, we reduce the complexity of learning the discriminators (they no longer need to "rule out" functions which are not functions of sets). This results in better performing discriminators, which in turn improves the performance of the generator.
In practice, the treatment discriminator receives all of the sets (i.e. one set for each treatment) of dosage-outcome pairs and outputs a probability for each treatment (i.e. there is one output corresponding to each set). In order to define such a function, we treat each input set as a vector but require that the outputs be invariant to (i.e. should not depend on) the ordering of the set as a vector.
Each dosage discriminator receives the set corresponding to a given treatment and is tasked with outputting a probability for each element in the set. In order to define such a function, we consider the input and output as vectors but then require that if we permute the elements of the input vector, the output should be permuted in the same way. We formalise the required notions - permutation invariance and permutation equivariance (Zaheer et al., 2017) - in the following subsection.

heading: 5.2.1 PERMUTATION INVARIANCE AND PERMUTATION EQUIVARIANCE
text: The notions of what it means for a function to be permutation invariant and permutation equivariant with respect to (a subset of) its inputs are given below in definitions 1 and 2, respectively. Let U , V , C be some spaces. Let m ∈ Z+. Definition 1. A function f : Um × V → C is permutation invariant with respect to the space Um if for every u = (u1, ..., um) ∈ Um, every v ∈ V and every permutation, σ, of {1, ...,m} we have
f(u1, ..., um, v) = f(uσ(1), ..., uσ(m), v) . (14)
Definition 2. A function f : Um×V → Cm is permutation equivariant with respect to the space Um if for every u = (u1, ..., um) ∈ Um, every v ∈ V and every permutation, σ, of {1, ...,m} we have
f(uσ(1), ..., uσ(m), v) = (fσ(1)(u, v), ..., fσ(m)(u, v)) , (15)
where fj(u, v) is the jth element of f(u, v).
To build up functions that are permutation invariant and permutation equivariant we make the following observations: (1) the composition of any function with a permutation invariant function is permutation invariant, (2) the composition of two permutation equivariant functions is permutation equivariant.
Zaheer et al. (2017) provide several possible building blocks to use to construct invariant and equivariant deep networks. The basic building block we will use for invariant functions will be a layer of the form
finv(u) = σ(1b1 T m(φ(u1), ..., φ(um))) , (16)
where 1l is a vector of 1s of dimension l, φ is any function φ : U → Rq for some q (in this paper we use a standard fully connected layer) and σ is some non-linearity.
The basic building block for equivariant functions is defined in terms of an equivariance input, u, and an auxiliary input, v, by
fequi(u,v) = σ(λImu + γ(1m1 T m)u + (1mΘ T )v) , (17)
where Im is the m×m identity matrix, λ and γ are scalar parameters and Θ is a vector of weights.
5.2.2 HIERARCHICAL DISCRIMINATOR ARCHITECTURE
In the case of the hierarchical discriminator, we want the treatment discriminator, DW , to be permutation invariant with respect to ỹw for each treatment, w ∈ W . To achieve this we define a function h1 : ∏ w∈W(Dw × Y)nw → HH and require that this function be permutation invariant with respect to each of the spaces (Dw × Y)nw . We then concatenate the output of h1 with the features x and pass these through a fully connected network h2 : X ×HH → [0, 1]k so that
DW(x, ỹ) = h2(x, h1(ỹ)). (18)
To construct h1, we concatenate the outputs of several invariant layers of
the form given in Eq. (16) that each individually act on the spaces (Dw × Y)nw . That is, for each treatment, w ∈ W we define a map hwinv : (Dw × Y)nw → HwH by substituting ỹw for u in Eq. (16). We then defineHH = ∏ w∈W HwH and h1(ỹ) = (hw1inv(ỹw1), ..., hwkinv(ỹwk)).
We want each dosage discriminator, Dw, to be permutation equivariant with respect to ỹw. To achieve this each Dw will consist of two layers of the form given in Eq. (17) with the equivariance input, u, to the first layer being ỹw and to the second layer being the output of the first layer and the auxiliary input, v, to the first layer being the features, x, and then no auxiliary input to the second layer.
Diagrams depicting the architectures of the treatment discriminator and dosage discriminators can be found in Fig. 4(a) and Fig. 4(b) respectively.

heading: 6 EVALUATION
text: The nature of the treatment-effects estimation problem in even the binary treatments setting does not allow for meaningful evaluation on real-world datasets. While there are well-established benchmark synthetic models for use in the binary (or multiple) case, no such models exist for the dosage setting. We propose our own semi-synthetic data simulation to evaluate our model against several benchmarks.

heading: 6.1 EXPERIMENTAL SETUP
text: Semi-synthetic data generation: We simulate data as follows. We obtain features, x, from a real dataset (in this paper we use TCGA (Weinstein et al., 2013), News (Johansson et al., 2016; Schwab et al., 2019)) and MIMIC III (Johnson et al., 2016))3. We consider 3 treatments each accompanied by a dosage. Each treatment, w ∈ W , is associated with a set of parameters, vw1 ,vw2 , vw3 . For each run of the experiment, these parameters are sampled randomly by first sampling a vector, uwi , from N (0,1) and then setting vwi = uwi /||uwi || where || · || is the standard Euclidean norm. The shape of the dose-response curve for each treatment, fw(x, d), is given in Table 1, along with a closed-form expression for the optimal dosage. We add ∼ N (0, 0.2) noise to the outcomes. We assign factual treatment-dosage pairs to each sample by first sampling a dosage, dw, for each treatment from a beta distribution, dw|x ∼ Beta(α, βw). The parameter α ≥ 1 controls the dosage selection bias4 and the parameter βw is set to βw = α−1d∗w + 2− α, with d ∗ w being the optimal dosage for each treatment5. This setting of βw ensures that the mode of the Beta distribution is the optimal dosage. Once we have sampled a dosage for each treatment, we assign a treatment according to wf |x ∼ Categorical(softmax(κf(x, dw)) where a higher κ will result in a stronger selection bias, and κ = 0 results in the treatments being assigned completely randomly. The factual treatment-dosage pair is then given by (wf , dwf ). Unless otherwise specified, we set κ = 2 and α = 2.
We consider 3 different shapes for fw to demonstrate learning heterogeneous dose-response curves. The first curve can be broken down into two terms, a linear (in d) increasing term (v11) Tx+12(v12) Txd and a quadratic (in d) decreasing term −12(v13)Txd2. This first term could, for example, represent the improved efficacy of higher dosages of chemotherapy in reducing the size of a tumour, while the quadratic term could represent the increasing toxicity of chemotherapy as the dosage increases. This type of trade-off presents itself in many other settings where there are both costs and rewards.
For metrics, we use Mean Integrated Square Error (MISE), Dosage Policy Error (DPE) and Policy Error (PE) (Silva, 2016; Schwab et al., 2019). Details can be found in Appendix F.
Benchmarks: We compare against two benchmarks: Generalized Propensity Score (GPS) (Imbens, 2000) and Dose Reponse Networks (DRNet) (Schwab et al., 2019). For DRNets, we compare against both the standard model architecture described by Schwab et al. (2019) as well as with Wasserstein regularization (DRN-W).
3Details of each dataset along with suggested interpretations of the synthetic treatments and outcomes for each dataset can be found in Appendix G
4When α = 1, βw = 1 and Beta(α, βw) reduces to the uniform distribution. See Appendix H. 5If the optimal dosage is 0, we sample dw from 1− Beta(α, βw) where βw is set as though d∗w = 1. This
results in the dosage being sampled symmetrically for d∗w = 0 and d∗w = 1.
As a baseline for comparison, we also use a standard multilayer perceptron (MLP) that takes as input the patient features, the treatment and dosage and estimates the patient outcome and a multitask variant (MLP-M) that has a designated head for each treatment. See Appendix E for details of the benchmark models and their hyperparameter optimisation.

heading: 6.2 SOURCE OF GAIN
text: Before comparing against the benchmarks, we investigate how each component of our model affects performance. We start with a baseline model in which both the generator and discriminator consist of a single fully connected network. One at a time, we add in the following components (cumulatively until we reach our full model): (1) the supervised loss in Eq. 4 (+ LS), (2) multitask heads in the generator (+ Multitask), (3) hierarchical discriminator (+ Hierarchical) and (4) invariance/equivariance layers in the treatment and dosage discriminators (+Inv/Eqv). We report the results in Table 2 for TCGA and News for all 3 error metrics (MISE, DPE and PE), computed over 30 runs.
We see that the addition of each component results in a performance improvement for our model, with the final row (which corresponds to our full model) demonstrating the best performance across both datasets and for all metrics.
To further demonstrate the advantages of our hierarchical discriminator, in Fig. 5 we investigate how our hierarchical discriminator compares with a single network discriminator (all other components are included in both models, see Appendix B for details of the single discriminator) when we vary the hyperparameter nw on TCGA. Similar results for News can be found in Appendix I.1.
The performance of the single discriminator causes significant performance drops around nw = 9 across all metrics. As previously noted, this is due to the dimension of the output space (which for nw = 9 is 27) being too large. Conversely, we see that our hierarchical discriminator shows much more stable performance even when nw = 19. We investigate in Appendix I.1 the hyperparameter λ.

heading: 6.3 BENCHMARKS COMPARISON
text: We now compare DRGAN against the benchmarks on our 3 semi-synthetic datasets. For Mimic, due to the low number of samples available, we use only two treatments - 2 and 3. We report √ MISE and√
PE in Table 3, with results for √
DPE given in Appendix I.3. We see that DRGAN demonstrates a statistically significant improvement over every benchmark across all 3 datasets, confirming that DRGAN is able to learn response-curves on top of very different underlying patient features.
In Appendix I.4 we compare DRGAN with DRNET and GPS for an increasing number of treatments.

heading: 6.4 TREATMENT AND DOSAGE SELECTION BIAS
text: In this section, we assess the robustness of each method to varying treatment and dosage bias. We report results for √ MISE on TCGA here. For the other metrics see Appendix I.2. Fig. 6(a) shows the performance of the 4 methods for κ between 0 (no bias) and 10 (strong bias). Fig. 6(b) shows the performance for α between 1 (no bias) and 8 (strong bias). We see that our model shows consistent performance, significantly outperforming the benchmark methods across the entire ranges of κ and α.

heading: 7 CONCLUSION
text: In this paper we proposed a novel framework for estimating dose-response curves from observational data. Our method modified the GAN framework, introducing a novel hierarchical discriminator for use in the dose-response setting. We also proposed novel architectures for the networks involved in our model and introduced a new semi-synthetic data simulation for use as a benchmark in this setting. On this data we demonstrated significant improvements over the benchmarks.

heading: A EXPANDED RELATED WORKS
text: Most methods for performing causal inference in the static setting focus on the scenario with two or multiple treatment options and no dosage parameter. The approaches taken by such methods to estimate the treatment effects involve either building a separate regression model for each treatment (Stoehlmacher et al., 2004; Qian & Murphy, 2011; Bertsimas et al., 2017) or using the treatment as a feature and adjusting for the imbalance between the different treatment populations. The former does not generalise to the dosage setting due to the now infinite number of possible treatments available. In the latter case, methods for handling the selection bias involve propensity weighting (Crump et al., 2008; Alaa et al., 2017; Shi et al., 2019), building sub-populations using tree based methods (Chipman et al., 2010; Athey & Imbens, 2016; Wager & Athey, 2018; Kallus, 2017) or building balancing representations between patients receiving the different treatments (Johansson et al., 2016; Shalit et al., 2017; Li & Fu, 2017; Yao et al., 2018). An additional approach involves modelling the data distribution of the factual and counterfactual outcomes (Alaa & van der Schaar, 2017; Yoon et al., 2018).
Silva (2016) leverages observational and interventional data to estimate the effects of discrete dosages for a single treatment. In particular, Silva (2016) uses observational data to construct a non-stationary covariance function and develop a hierarchical Gaussian process prior to build a distribution over the dose response curve. Then, controlled interventions are employed to learn a non-parametric affine transform to reshape this distribution. The setting in Silva (2016) differs significantly from ours as we do not assume access to any interventional data.

heading: B SINGLE DISCRIMINATOR MODEL
text: In the paper we developed a hierarchical discriminator and demonstrated that it performs significantly better than the single discriminator setup that we now describe in this section.
B.1 SINGLE DISCRIMINATOR
In the single model, we will aim to learn a single discriminator, D, that outputs P((Wf , Df ) = (w, d)|X, D̃w, Ỹ) for each w ∈ W and d ∈ D̃w. We will write Dw,d(·) to denote the output of D that corresponds to the treatment-dosage pair (w, d). We define the loss, LD, to be
LD(D;G) = −E [ ∑ w∈W ∑ d∈D̃w I{Tf=(w,d)} logD w,d(X, Ỹ) + I{Tf 6=(w,d)} log(1−Dw,d(X, Ỹ)) ] (19) where the expectation is taken over X, {D̃w}w∈W , Ỹ,Wf and Df and we note that the dependence on G is through Ỹ. Our single discriminator will be trained to minimise this loss directly. The generator GAN-loss, LG, is then defined by
LG(G) = −LD(D∗;G) (20) where D∗ is the optimal discriminator given by minimising LD. The generator will be trained to minimise LG + λLS .
B.2 SINGLE DISCRIMINATOR ARCHITECTURE
In the case of the single discriminator, we want the output of D corresponding to each treatment w ∈ W , i.e. (Dw,1, ...,Dw,nw), to be permutation equivariant with respect to ỹw and permutation invariant with respect to each ỹv for v ∈ W\{w}. To achieve this, we first define a function f : ∏ w∈W(Dw×Y)nw → HS and require that this function be permutation invariant with respect to each of the spaces (Dw × Y)nw . For each treatment, w ∈ W , we introduce a multitask head, fw : X ×HS × (Dw ×Y)nw → [0, 1]nw , and require that each of these functions be permutation equivariant with respect to their corresponding input space (Dw × Y)nw but they can depend on the features, x ∈ X , and invariant latent representation coming from f arbitrarily. Writing f jw to denote the jth output of fw, the output of the discriminator given input features, x, and generated outcomes, ỹ, is defined by
Dw,j(x, ỹ) = f iw(x, f(ỹ), ỹw). (21)
To construct the function f , we concatenate the outputs of several invariant layers of the form given in Eq. (16) that each individually act on the spaces (Dw × Y)nw . That is, for each treatment, w ∈ W we define a map fwinv : (Dw × Y)nw → HwS by substituting ỹw for u in Eq. (16). We then define HS = ∏ w∈W HwS and f(ỹ) = (fw1inv(ỹw1), ..., f wk inv(ỹwk)).
Each fw will consist of two layers of the form given in Eq. (17) with the equivariance input, u, to first layer being ỹw and to the second layer being the output of the first layer and the auxiliary input, v, to the first layer being the concatenation of the features and invariant representation, i.e. (x, f(ỹ)) and then no auxiliary input to the second layer.
A diagram depicting the architecture of the single discriminator model can be found in Fig. 7.

heading: C COUNTERFACTUAL GENERATOR PSEUDO-CODE
text: Algorithm 1 Training of the generator in DRGAN 1: Input: dataset C = {(xi, tif , yif ) : i = 1, ..., N}, batch size nmb, number of dosages per
treatment nd, number of discriminator updates per iteration nD, number of generator updates per iteration nG, dimensionality of noise nz , learning rate α
2: Initialize: θG, θW , {θw}w∈W 3: while G has not converged do
Discriminator updates 4: for i = 1, ..., nD do 5: Sample (x1, (w1, d1), y1), ..., (xnmb , (wnmb , dnmb), ynmb) from C 6: Sample generator noise zj = (z j 1, ..., z j nz ) from Unif([0, 1]
nz ) for j = 1, ..., nmb 7: for w ∈ W do 8: for j = 1, ..., nmb do 9: Sample D̃jw = (d w,j 1 , ..., d w,j nd
) independently and uniformly from (Dw)nd 10: Set ỹjw according to Eq. 5 11: Calculate gradient of dosage discriminator loss
gw ← ∇θw −  ∑ {j:wj=w} nd∑ k=1 I{dj=dw,jk } logDw(xj , ỹ j w) + I{dj 6=dw,jk } log(1−Dw(xj , ỹ j w))  12: Update dosage discriminator parameters θw ← θw + αgw 13: Set ỹj = (ỹjw)w∈W 14: Calculate gradient of treatment discriminator loss
gW ← ∇θW − nmb∑ j=1 ∑ w∈W I{wj=w} logDW(xj , ỹj) + I{wj 6=w} log(1−DW(xj , ỹj))  15: Update treatment discriminator parameters θW ← θW + αgW
Generator updates 16: for i = 1, ..., nG do 17: Sample (x1, (w1, d1), y1), ..., (xnmb , (wnmb , dnmb), ynmb) from C 18: Sample generator noise zj = (z j 1, ..., z j nz ) from Unif([0, 1]
nz ) for j = 1, ..., nmb 19: Sample (D̃jw)w∈W from Πw∈W(Dw)nd for j = 1, ..., nmb 20: Set ỹ according to Eq. 5 21: Calculate gradient of generator loss
gG ← ∇θG [ nmb∑ j=1 ∑ w∈W nd∑ l=1 I{wj=w,dj=dw,jl } log(D w W(xj , ỹj)w ×Dlw(xj , ỹjw)l)
+I{wj 6=w,dj 6=dw,jl } log(1− (D w W(xj , ỹj)×Dlw(xj , ỹjw))) ] 22: Update generator parameters θG ← θG + αgG 23: Output: G
D INFERENCE NETWORK
To generate dose-response curves for new samples, we learn an inference network, I : X× T → Y . This inference network is trained using the original dataset and the learned counterfactual generator. As with the training of the generator and discriminator, we train using a random set of dosages, D̃w. The loss is given by
LI(I) = E [ ∑ w∈W ∑ d∈D̃w (Ỹ (w, d)− I(X, (w, d)))2 ] , (22)
where Ỹ (w, d) is Yf if Tf = (w, d) or given by the generator if Tf 6= (w, d). The expectation is taken over X, Tf , Yf ,Z and D̃w.
D.1 PSEUDO-CODE FOR TRAINING THE INFERENCE NETWORK
Algorithm 2 Training of the inference network in DRGAN 1: Input: dataset C = {(xi, tif , yif ) : i = 1, ..., N}, trained generator G, batch size nmb, number
of dosages per treatment nd, dimensionality of noise nz , learning rate α 2: Initialize: θI , 3: while I has not converged do 4: Sample (x1, (w1, d1), y1), ..., (xnmb , (wnmb , dnmb), ynmb) from C 5: Sample generator noise zj = (z j 1, ..., z j nz ) from Unif([0, 1]
nz ) for j = 1, ..., nmb 6: for j = 1, ..., nmb do 7: for w ∈ W do 8: Sample D̃jw = (d w,j 1 , ..., d w,j nd
) independently and uniformly from (Dw)nd 9: Set ỹjw according to Eq. 5 5
10: Calculate gradient of inference network loss
gI ← ∇θI [ nmb∑ j=1 ∑ w∈W nd∑ l=1 (ỹjw)l − I(xj , (w, dw,jl ))2 ]
11: Update inference network parameters θI ← θI + αgI 12: Output: I

heading: E BENCHMARKS
text: We use the publicly available GitHub implementation of DRNet provided by Schwab et al. (2019): https://github.com/d909b/drnet. Moreover, we also used a GPS implementation similar to the one from https://github.com/d909b/drnetwhich uses the causaldrfR package (Galagate, 2016). More spcifically, the GPS implementation uses a normal treatment model, a linear treatment formula and a 2-nd degree polynomial for the outcome. Moreover, for the TCGA and News datasets, we performed PCA and only used the 50 principal components as input to the GPS model to reduce computational complexity.
Hyperparameter optimization: The validation split of the dataset is used for hyperparameter optimization. For the DRNet benchmarks we use the same hyperparameter optimization proposed by Schwab et al. (2019) with the hyperparameter search ranges described in Table 4. For DRGAN, we use the hyperparameter optimization method proposed in GANITE (Yoon et al., 2018), where we use the complete dataset from the counterfactual generator to evaluate the MISE on the inference network. We perform a random search (Bergstra & Bengio, 2012) for hyperparameter optimization over the search ranges in Table 5. For a fair comparison, for the MLP-M model we used the same architecture used in the inference network of DRGAN. Similarly, for the MLP model we use the same architecture as for the MLP-M, but without the multitask heads.

heading: F METRICS
text: The Mean Integrated Square Error (MISE) measures how well the models estimates the patient outcome across the entire dosage space:
MISE = 1
N
1
k ∑ w∈W N∑ i=1 ∫ Dw ( yi(w, u)− ŷi(w, u) )2 du . (23)
In addition to this, we also compute the mean dosage policy error (DPE) (Schwab et al., 2019) to assess the ability of the model to estimate the optimal dosage point for every treatment for each individual:
DPE = 1
N
1
k ∑ w∈W N∑ i=1 ( yi(w, d∗w)− yi(w, d̂∗w) )2 , (24)
where d∗w is the true optimal dosage and d̂ ∗ w is the optimal dosage identified by the model. The optimal dosage points for a model are computed using SciPy’s implementation of Sequential Least SQuares Programming.
Finally, we compute the mean policy error (PE) (Schwab et al., 2019) which compares the outcome of the true optimal treatment-dosage pair to the outcome of the optimal treatment-dosage pair as selected by the model:
PE = 1
N N∑ i=1 ( yi(w∗, d∗w)− yi(ŵ∗, d̂∗w) )2 , (25)
where w∗ is the true optimal treatment and ŵ∗ is the optimal treatment identified by the model. The optimal treatment-dosage pair for a model is selected by first computing the optimal dosage for each treatment and then selecting the treatment with the best outcome for its optimal dosage.
Each of these metrics are computed on a held out test-set.

heading: G DATASET DESCRIPTIONS
text: TCGA: The TCGA dataset consists of gene expression measurements for cancer patients (Weinstein et al., 2013). There are 9659 samples for which we used the measurements from the 4000 most variable genes. The gene expression data was log-normalized and each feature was scaled in the [0, 1] interval. Moreover, for each patient, the features x were scaled to have norm 1. We give meaning to our treatments and dosages by considering the treatment as being chemotherapy/radiotherapy/immunotherapy and their corresponding dosages. The outcome can be thought of as the risk of cancer recurrence (Schwab et al., 2019).
News: The News dataset consists of word counts for news items. We extracted 10000 samples each with 2858 features. As in (Johansson et al., 2016; Schwab et al., 2019), we give meaning to our treatments and dosages by considering the treatment as being the viewing device (e.g. phone, tablet etc.) used to read the article and the dosage as being the amount of time spent reading it. The outcome can be thought of as user satisfaction.
MIMIC III: The Medical Information Mart for Intensive Care (MIMIC III) (Johnson et al., 2016) database consists of observational data from patients in the ICU. We extracted 3000 patients that receive antibiotics treatment and we used as features 9 clinical covariates measured during the day of ICU admission. Again, the features were scaled in the [0, 1] interval. In this setting, we can considered as treatments the different antibiotics and their corresponding dosages.
For a summary description of the datasets, see table 6. The datasets are split into 64/16/20% for training, validation and testing respectively. The validation dataset is used for hyperparameter optimization.

heading: H DOSAGE BIAS
text: In order to create dosage-assignment bias in our dataset, we assign dosages according to dw|x ∼ Beta(α, βw). The selection bias is controlled by the parameter α ≥ 1. When we set βw = α−1d∗w +2−α (which ensures that the mode of our distribution is d∗w), we can write the variance of dw in terms of α and d∗w as follows
Var(dw) = α2−α d∗w + 2α− α2
(α−1d∗w + 2) 2(α−1d∗w + 3)
≈ cα 2
dα3 . (26)
We see that the variance of our Beta distribution therefore decreases with α, resulting in the sampled dosages being closer to the optimal dosage, thus resulting in higher dosage-selection bias. In addition we note that the Beta(1, 1) distribution is in fact the uniform distribution, corresponding to the dosages being sampled independently of the patient features, resulting in no selection bias when α = 1.

heading: I ADDITIONAL RESULTS
text: I.1 INVESTIGATING HYPERPARAMETER SENSITIVITY (nw AND λ)
Here we present additional results for our investigation of the hyperparameters nw and λ. Fig. 8 reports each of the 3 performance metrics as we increase the number of dosage samples, nw, used to train the discriminators on the News dataset. As with the TCGA results in the main paper we see that the single discriminator suffers a significant performance decrease when nw is set too high.
Fig. 9 and Fig. 10 report each of the performance metrics when we increase λ, the hyperparameter that trades off between the GAN loss, L, and the supervised loss, LS . The results shown are on the TCGA and News datasets, respectively. As we would expect, for λ = 0, the performance is significantly worse than for λ = 1, since this corresponds to no supervised loss (this result is also backed up by our source of gain experiments). In addition, as we increase λ we see performance slowly degrades, this is because the model starts to behave more and more like MLP-M since the supervised loss becomes dominant.
I.2 ADDITIONAL RESULTS ON SELECTION BIAS
In Fig. 11 we report the DPE and PE for our treatment and dosage bias experiment from Section 6.4 of the main paper.
I.3 DOSAGE POLICY ERROR FOR BENCHMARK COMPARISON
In Table 7 we report the Dosage Policy Error (DPE) corresponding to Section 6.3 in the main paper.
I.4 VARYING THE NUMBER OF TREATMENTS
In our final experiment, we increase the number of treatments by defining 3 or 6 additional treatments. The parameters vw1 ,v w 2 ,v w 3 are defined in exactly the same way as for 3 treatments. The outcome shapes for treatments 4 and 7 are the same as for treatment 1, similarly for 5, 8 and 2 and for 6, 9 and 3. In Table 8 we report MISE, DPE and PE on the TCGA dataset with 6 treatments (TCGA-6) and with 9 treatments (TCGA-9). Note that we use 3 dosage samples for training DRGAN in this experiment.

Review summaries:
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.

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.