Title: INDIVIDUALISED DOSE-RESPONSE ESTIMATION USING GENERATIVE ADVERSARIAL NETS

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. Patient outcome Dosage DRGAN Counterfactual Generator Dosage DRGAN Discriminator

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

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 [a w , b w ] 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.

3. PROBLEM FORMULATION:

We consider receiving observations of the form (x i , t i f , y i f ) for i = 1, ..., N , where, for each i, these are independent realizations of the random variables (X, T f , Y f ). 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, T f , is in fact a pair of values T f = (W f , D f ) where W f ∈ W corresponds to the type of treatment being administered (e.g. chemotherapy or radiotherapy) which lies in the discrete space of k treatments, W = {w 1 , ..., w k }, and D f 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, D w , 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 ∈ D w }. 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 Y f = Y (W f , D f ). 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:

EQUATION

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, T f = 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, T f , and potential outcomes, Y (w, d), are conditionally independent given the covariates X, i.e.

EQUATION

)

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.

4.1. COUNTERFACTUAL GENERATOR:

Our generator, G : X × T × Y × Z → Y T takes features, x ∈ X , factual outcome, y f ∈ Y, received treatment and dosage, t f = (w f , d f ) ∈ 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, t f , y f , z)(•) : T → Y. We can then write

EQUATION

to denote our generated counterfactual outcome for the treatment-dosage pair t. We will write

Ŷcf (t) = G(X, T f , Y f , 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, L S , to be

EQUATION

where the expectation is taken over X, T f , Y f and Z.

4.2. COUNTERFACTUAL DISCRIMINATOR:

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 n w ∈ Z + be the number of dosage levels we will compare for treatment w ∈ W1 .

For each w ∈ W, let Dw = {D w 1 , ..., D w nw } be a random subset2 of D w of size n w , where for the factual treatment, W f , DW f contains n W f -1 random elements along with D f . We define Ỹw = (D w i , Ỹ w i ) nw i=1 ∈ (D w × Y) nw to be the vector of dosage-outcome pairs for treatment w where

Ỹ w i = Y f if W f = w and D f = D w i Ŷcf (w, D w i ) else (5)

and will write Ỹ = ( Ỹw ) w∈W . We will write d w j , ỹw and ỹ to denote realisations of D w j , Ỹw and ỹ.

Our discriminator, D : X × w∈W (D w × Y) nw → [0, 1] nw , will take the features x ∈ X together with the (random) set of generated outcomes ỹ ∈ 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

EQUATION

where the expectation is taken over X, T f , Ỹ and { Dw : w ∈ W}, D w,d corresponds to the discriminator output for treatment-dosage pair (w, d). The minimax game is then given by min

EQUATION

where λ is used to control the trade-off between L and L S (we set λ = 1 in the experiments).

The task of the discriminator (i.e. picking out the factual dosage from k j=1 n wj treatment-dosage pairs) becomes increasingly difficult as we increase n w or k because the dimension of the discriminator output space, n w , increases. Although we control n w , if we set it too low, then the set ỹ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 n w 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 n w .

First observe that the probability P((W f , D f ) = (w, d)|X, Dw , Ỹ) can be written as

EQUATION

We can therefore break down the discriminator into a hierarchical model by learning one discriminator, D W , that outputs P(W f = w|X, Dw , Ỹ) which we will refer to as the treatment discriminator, and then a discriminator, D w , for each treatment, w ∈ W, that outputs P(D f = d|W f = w, X, Dw , Ỹ) which we will refer to as the dosage discriminator for treatment w.

The treatment discriminator, D W : X × w∈W (D w × Y) nw → [0, 1] k ,

takes the features, x, and generated potential outcomes, ỹ, and outputs a probability for each treatment, w 1 , ..., w k . Writing D w W to denote the output of D W corresponding to treatment w, we define the loss, L W , to be

EQUATION

where, again, the expectation is taken over X, W f , D f , Ỹ and { Dw } w∈W .

Then, for each w ∈ W, D w : X × (D w × Y) nw → [0, 1] nw is a map that takes the features, x, and generated potential outcomes, ỹw , corresponding to treatment w and outputs a probability for each dosage level, d w 1 , ..., d w nw , in a given realisation of Dw . Writing D j w to denote the output of D w corresponding to dosage level D w j , we define the loss of each dosage discriminator to be

EQUATION

where the expectation is taken over X, Dw , Ỹw , W f and D f . The I {W f =w} term ensures that only samples for which the factual treatment is w are used to train dosage discriminator D w (otherwise there would be no factual dosage for that sample).

We define the overall discriminator D H : X × w∈W (D w × Y ) nw → [0, 1] nw by defining its output corresponding to the treatment-dosage pair (w, d w j ) as

EQUATION

)

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) + λL S (G) D * H w,j = D * W w × D * w j D * W = arg min D W L W (D W ; G * ) D * w = arg min Dw L d (D w ; G * ), ∀w ∈ W (12)

Fig. 2 depicts our generator and hierarchical discriminator. Pseudo-code for our algorithm can be found in Appendix C.

4.3. INFERENCE NETWORK:

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.

5. ARCHITECTURE:

In this section, we describe in detail the novel architectures that we adopt to model each of the functions G, D, D W , D w1 , ..., D w k 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 w f , d f , y f or z as inputs.

5.1. GENERATOR ARCHITECTURE:

Shared layer w f < l a t e x i t s h a 1 _ b a s e 6 4 = " ( n u l l ) " > ( n u l l ) < / l a t e x i t > < l a t e x i t s h a 1 _ b a s e 6 4 = " ( n u l l ) " > ( n u l l ) < / l a t e x i t > < l a t e x i t s h a 1 _ b a s e 6 4 = " ( n u l l ) " > ( n u l l ) < / l a t e x i t > < l a t e x i t s h a 1 _ b a s e 6 4 = " ( n u l l ) " > ( n u l l ) < / l a t e x i t > d f < l a t e x i t s h a 1 _ b a s e 6 4 = " ( n u l l ) " > ( n u l l ) < / l a t e x i t > < l a t e x i t s h a 1 _ b a s e 6 4 = " ( n u l l ) " > ( n u l l ) < / l a t e x i t > < l a t e x i t s h a 1 _ b a s e 6 4 = " ( n u l l ) " > ( n u l l ) < / l a t e x i t > < l a t e x i t s h a 1 _ b a s e 6 4 = " ( n u l l ) " > ( n u l l ) < / l a t e x i t > y f < l a t e x i t s h a 1 _ b a s e 6 4 = " ( n u l l ) " > ( n u l l ) < / l a t e x i t > < l a t e x i t s h a 1 _ b a s e 6 4 = " ( n u l l ) " > ( n u l l ) < / l a t e x i t > < l a t e x i t s h a 1 _ b a s e 6 4 = " ( n u l l ) " > ( n u l l ) < / l a t e x i t > < l a t e x i t s h a 1 _ b a s e 6 4 = " ( n u l l ) " > ( n u l l ) < / l a t e x i t > z < l a t e x i t s h a 1 _ b a s e 6 4 = " ( n u l l ) " > ( n u l l ) < / l a t e x i t > < l a t e x i t s h a 1 _ b a s e 6 4 = " ( n u l l ) " > ( n u l l ) < / l a t e x i t > < l a t e x i t s h a 1 _ b a s e 6 4 = " ( n u l l ) " > ( n u l l ) < / l a t e x i t > < l a t e x i t s h a 1 _ b a s e 6 4 = " ( n u l l ) " > ( n u l l ) < / l a t e x i t > d < l a t e x i t s h a 1 _ b a s e 6 4 = " ( n u l l ) " > ( n u l l ) < / l a t e x i t > < l a t e x i t s h a 1 _ b a s e 6 4 = " ( n u l l ) " > ( n u l l ) < / l a t e x i t > < l a t e x i t s h a 1 _ b a s e 6 4 = " ( n u l l ) " > ( n u l l ) < / l a t e x i t > < l a t e x i t s h a 1 _ b a s e 6 4 = " ( n u l l ) " > ( n u l l ) < / l a t e x i t >

x < l a t e x i t s h a 1 _ b a s e 6 4 = " ( n u l l ) " > ( n u l l ) < / l a t e x i t > < l a t e x i t s h a 1 _ b a s e 6 4 = " ( n u l l ) " > ( n u l l ) < / l a t e x i t > < l a t e x i t s h a 1 _ b a s e 6 4 = " ( n u l l ) " > ( n u l l ) < / l a t e x i t > < l a t e x i t s h a 1 _ b a s e 6 4 = " ( n u l l ) " > ( n u l l ) < / l a t e x i t > ŷ(w, d)

< l a t e x i t s h a 1 _ b a s e 6 4 = " ( n u l l ) " > ( n u l l ) < / l a t e x i t > < l a t e x i t s h a 1 _ b a s e 6 4 = " ( n u l l ) " > ( n u l l ) < / l a t e x i t > < l a t e x i t s h a 1 _ b a s e 6 4 = " ( n u l l ) " > ( n u l l ) < / l a t e x i t > < l a t e x i t s h a 1 _ b a s e 6 4 = " ( n u l l ) " > ( n u l l ) < / l a t e x i t > w < l a t e x i t s h a 1 _ b a s e 6 4 = " ( n u l l ) " > ( n u l l ) < / l a t e x i t > < l a t e x i t s h a 1 _ b a s e 6 4 = " ( n u l l ) " > ( n u l l ) < / l a t e x i t > < l a t e x i t s h a 1 _ b a s e 6 4 = " ( n u l l ) " > ( n u l l ) < / l a t e x i t > < l a t e x i t s h a 1 _ b a s e 6 4 = " ( n u l l ) " > ( n u l l ) < / l a t e x i t > Given observations, (x, t f , y f ), a noise vector z, and a target treatment-dosage pair, t = (w, d), we define

EQUATION

)

Each of g, g w1 , ..., g w k are modelled as fully connected networks. Fig. 3 depicts our generator architecture.

5.2. DISCRIMINATOR ARCHITECTURES:

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.

5.2.1. PERMUTATION INVARIANCE AND PERMUTATION EQUIVARIANCE:

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 : U m × V → C is permutation invariant with respect to the space U m if for every u = (u 1 , ..., u m ) ∈ U m , every v ∈ V and every permutation, σ, of {1, ..., m} we have

EQUATION

)

Definition 2. A function f : U m × V → C m is permutation equivariant with respect to the space U m if for every u = (u 1 , ..., u m ) ∈ U m , every v ∈ V and every permutation, σ, of {1, ..., m} we have

EQUATION

where f j (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 f inv (u) = σ(1 b 1 T m (φ(u 1 ), ..., φ(u m ))) , (16) where 1 l is a vector of 1s of dimension l, φ is any function φ : U → R q 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

EQUATION

where I m is the m × m identity matrix, λ and γ are scalar parameters and Θ is a vector of weights.

5.2.2. HIERARCHICAL DISCRIMINATOR ARCHITECTURE:

the form given in Eq. ( 16) that each individually act on the spaces (D w × Y) nw . That is, for each treatment, w ∈ W we define a map h w inv : (D w × Y) nw → H w H by substituting ỹw for u in Eq. ( 16). We then define H H = w∈W H w H and h 1 (ỹ) = (h w1 inv (ỹ w1 ), ..., h w k inv (ỹ w k )). We want each dosage discriminator, D w , to be permutation equivariant with respect to ỹw . To achieve this each D w will consist of two layers of the form given in Eq. ( 17) with the equivariance input, u, to the first layer being ỹ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.

6. EVALUATION:

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.

6.1. EXPERIMENTAL SETUP:

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, v w 1 , v w 2 , v w 3 . For each run of the experiment, these parameters are sampled randomly by first sampling a vector, u w i , from N (0, 1) and then setting v w i = u w i /||u w i || where || • || is the standard Euclidean norm. The shape of the dose-response curve for each treatment, f w (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, d w , for each treatment from a beta distribution, d w |x ∼ Beta(α, β w ). The parameter α ≥ 1 controls the dosage selection bias4 and the parameter β w is set to

β w = α-1 d * 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 w f |x ∼ Categorical(softmax(κf (x, d w )) 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 (w f , d w f ). Unless otherwise specified, we set κ = 2 and α = 2.

We consider 3 different shapes for f w to demonstrate learning heterogeneous dose-response curves. The first curve can be broken down into two terms, a linear (in d) increasing term (v 1 1 ) T x+12(v 1 2 ) T xd and a quadratic (in d) decreasing term -12(v 1 3 ) T xd 2 . 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).

Treatment:



Dose-Response:

Optimal dosage

1 f1(x, d) = C((v 1 1 ) T x + 12(v 1 2 ) T xd -12(v 1 3 ) T xd 2 ) d * 1 = (v 1 2 ) T x 2(v 1 3 ) T x 2 f2(x, d) = C((v 2 1 ) T x + sin(π( v 2 2 T x v 2 3 T x )d)) d * 2 = (v 2 3 ) T x 2(v 2 2 ) T x 3 f3(x, d) = C((v 3 1 ) T x + 12d(d -b) 2

, where b = 0.75

(v 3 2 ) T x (v 3 3 ) T x ) b 3 if b ≥ 0.75 1 if b < 0.75

Table 1 : Dose response curves used to generate semi-synthetic outcomes for patient features x. In the experiments, we set C = 10. v w 1 , v w 2 , v w 3 are the parameters associated with each treatment w.

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.

6.2. SOURCE OF GAIN:

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 (+ L S ), ( 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 n w on TCGA. Similar results for News can be found in Appendix I.1. The performance of the single discriminator causes significant performance drops around n w = 9 across all metrics. As previously noted, this is due to the dimension of the output space (which for n w = 9 is 27) being too large. Conversely, we see that our hierarchical discriminator shows much more stable performance even when n w = 19. We investigate in Appendix I.1 the hyperparameter λ.

6.3. BENCHMARKS COMPARISON:

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. PE √ MISE √ PE DRGAN 1.89 ± 0.05 0.25 ± 0.05 3.71 ± 0.05 3.90 ± 0.05 2.09 ± 0.12 0.32 ± 0.05 DRNet 3.64 ± 0.12 0.67 ± 0.05 4.98 ± 0.12 4.17 ± 0.11 4.45 ± 0.12 1.44 ± 0.05 DRN-W 3.71 ± 0.12 0.63 ± 0.05 5.07 ± 0.12 4.56 ± 0.12 4.47 ± 0.12 1.37 ± 0.05 GPS 4.83 ± 0.01 1.60 ± 0.01 6.97 ± 0.01 24.1 ± 0.05 7.39 ± 0.00 20.2 ± 0.01 MLP-M 3.96 ± 0.12 1.20 ± 0.05 5.17 ± 0.12 5.82 ± 0.16 4.97 ± 0.16 1.59 ± 0.05 MLP 4.31 ± 0.05 0.97 ± 0.05 5.48 ± 0.16 6.45 ± 0.21 5.34 ± 0.16 1.65 ± 0.05

Table 3 : Performance of individualized treatment-dose response estimation on three datasets. Bold indicates the method with the best performance for each dataset.

In Appendix I.4 we compare DRGAN with DRNET and GPS for an increasing number of treatments.

6.4. TREATMENT AND DOSAGE SELECTION BIAS:

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

7. CONCLUSION:

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.

APPENDIX A EXPANDED RELATED WORKS:

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.

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, Dw . The loss is given by

EQUATION

)

where Sample generator noise z j = (z j 1 , ..., z j nz ) from Unif([0, 1] nz ) for j = 1, ..., n mb 6:

Ỹ (w, d) is Y f if T f = (w,

for j = 1, ..., n mb do 7:

for w ∈ W do 

E BENCHMARKS:

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/drnet which uses the causaldrf R 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. 

G DATASET DESCRIPTIONS:



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 

EQUATION

)

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.

I ADDITIONAL RESULTS:



I.1 INVESTIGATING HYPERPARAMETER SENSITIVITY (n w AND λ):

Here we present additional results for our investigation of the hyperparameters n w and λ. Fig. 8 reports each of the 3 performance metrics as we increase the number of dosage samples, n w , 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 n w 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, L S . 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. Table 8 : Performance of DRGAN and the benchmarks when we increase the number of treatments in the dataset to 6 and 9. Bold indicates the method with the best performance for each dataset.

:

In practice we set all nw to be the same. The default setting is 5 in the experiments.

In practice, when Dw = [0, 1], each D w j is sampled independently and uniformly from [0, 1]. Note that for each training iteration, Dw is resampled (see Section 1).

Details of each dataset along with suggested interpretations of the synthetic treatments and outcomes for each dataset can be found in Appendix G

When α = 1, βw = 1 and Beta(α, βw) reduces to the uniform distribution. See Appendix H.

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

Figures and Tables:

Figure 1: Comparison of DRGAN and GANITE highlighting the key difference between the two different problems they address (dose-response estimation vs. standard treatment-response estimation).

Figure2: Overview of our model for the setting with two treatments (w f corresponds to the factual treatment and w cf to the counterfactual treatment). The generator is used to generate an output for each dosage level in each Dw , these outcomes together with the factual outcome, y f , are used to create the set of dosage-outcome pairs, ỹ, which is passed directly to the treatment discriminator. Each dosage discriminator receives only the part of ỹ corresponding to that treatment, i.e. ỹw . These discriminators are combined (Eq. 11) to define D H which is used to give feedback to the generator.

Figure 3: Generator architecture.

Figure 5: Performance of single vs. hierarchical discriminator when increasing the number of dosage samples (n w ) on TCGA dataset.

Figure 6: Performance of the 4 methods on datasets with varying bias levels.

.., d w,j n d ) independently and uniformly from (D w ) n d parameters θ I ← θ I + αg I 12: Output: I

Figure 8: Performance metrics when increasing the number of dosage samples on News dataset.

Figure 9: Performance metrics when increasing λ on TCGA dataset.

Source of gain analysis for our model. Metrics are reported as Mean ± Std.

Table content: <table><tr><td/><td>√</td><td>MISE</td><td>TCGA √ DPE</td><td>√ PE</td><td>√ MISE</td><td>News √ DPE</td><td>√</td><td>PE</td></tr><tr><td>Baseline</td><td colspan="8">4.18 ± 0.32 2.06 ± 0.16 1.93 ± 0.12 6.17 ± 0.27 6.97 ± 0.27 6.20 ± 0.21</td></tr><tr><td>+ LS</td><td colspan="8">3.37 ± 0.11 1.14 ± 0.05 0.84 ± 0.05 4.51 ± 0.16 4.46 ± 0.12 4.40 ± 0.11</td></tr><tr><td>+ Multitask</td><td colspan="8">3.15 ± 0.12 0.85 ± 0.05 0.67 ± 0.05 4.11 ± 0.11 4.33 ± 0.11 4.31 ± 0.11</td></tr><tr><td colspan="9">+ Hierarchical 2.54 ± 0.05 0.36 ± 0.05 0.45 ± 0.05 4.07 ± 0.05 4.24 ± 0.11 4.17 ± 0.12</td></tr><tr><td>+ Inv/Eqv</td><td colspan="8">1.89 ± 0.05 0.31 ± 0.05 0.25 ± 0.05 3.71 ± 0.05 4.14 ± 0.11 3.90 ± 0.05</td></tr></table>

d) or given by the generator if T f = (w, d). The expectation is taken over X, T f , Y f , Z and Dw . Training of the inference network in DRGAN 1:Input: dataset C = {(x i , t i f , y i f ) : i = 1, ...,N }, trained generator G, batch size n mb , number of dosages per treatment n d , dimensionality of noise n z , learning rate α 2: Initialize: θ I , 3: while I has not converged do , (w 1 , d 1 ), y 1 ), ..., (x n mb , (w n mb , d n mb ), y n mb ) from C

Table content: <table><tr><td>D.1 PSEUDO-CODE FOR TRAINING THE INFERENCE NETWORK</td></tr><tr><td>Algorithm 2 4: Sample (x 1 5:</td></tr></table>

Hyperparameters search range for DRNet. *: For the DRNet model using Wasserstein regularization only.

Table content: <table><tr><td>Hyperparameter</td><td>Search range</td></tr><tr><td>Batch size</td><td>32, 64, 128</td></tr><tr><td colspan="2">Number of units per hidden layer 24, 48, 96, 192</td></tr><tr><td>Number of hidden layers</td><td>2, 3</td></tr><tr><td>Dropout percentage</td><td>0.0, 0.2</td></tr><tr><td>Imbalance penalty weight  *</td><td>0.1, 1.0, 10.0</td></tr><tr><td/><td>Fixed</td></tr><tr><td>Number of dosage strata E</td><td>5</td></tr><tr><td>Hyperparameter</td><td>Search range</td></tr><tr><td>Batch size</td><td>64, 128, 256</td></tr><tr><td>Number of units per hidden layer</td><td>32, 64, 128</td></tr><tr><td>Size of invariant and equivariant representations</td><td>16, 32, 64, 128</td></tr><tr><td/><td>Fixed</td></tr><tr><td>Number of hidden layers per multitask head</td><td>2</td></tr><tr><td>Number of dosage samples</td><td>5</td></tr><tr><td>λ</td><td>1</td></tr><tr><td>Optimization</td><td>Adam Moment Optimization</td></tr></table>

Hyperparameters search range for DRGAN.

Table content: <table/>

The datasets are split into 64/16/20% for training, validation and testing respectively. The validation dataset is used for hyperparameter optimization.

Table content: <table><tr><td/><td colspan="3">TCGA News MIMIC</td></tr><tr><td>Number of samples</td><td>9659</td><td>10000</td><td>3000</td></tr><tr><td>Number of features</td><td>4000</td><td>2858</td><td>9</td></tr><tr><td>Number of treatments</td><td>3*</td><td>3</td><td>2</td></tr></table>

Summary description of datasets. *: for our final experiment in Appendix I.4 we increase the number of treatments in TCGA to 6 and 9.

Table content: <table><tr><td>d  *  w</td><td>+2-α</td></tr></table>

Table 7 we report the Dosage Policy Error (DPE) corresponding to Section 6.3 in the main paper. 6.40 ± 0.01 * 1.41 ± 0.12 * MLP-M 0.92 ± 0.05 * 4.94 ± 0.16 * 0.77 ± 0.05 * MLP 1.04 ± 0.05 * 5.18 ± 0.12 * 0.80 ± 0.05 * Performance of individualized treatment-dose response estimation on three datasets. Bold indicates the method with the best performance for each dataset. *: performance improvement is statistically significant. MISE √ DPE √ PE DRGAN 2.37 ± 0.12 0.43 ± 0.05 0.32 ± 0.05 2.79 ± 0.05 0.51 ± 0.05 0.54 ± 0.05 DRNET 4.09 ± 0.16 0.52 ± 0.05 0.71 ± 0.05 4.31 ± 0.12 0.59 ± 0.05 0.74 ± 0.05 GPS 6.62 ± 0.01 2.04 ± 0.01 2.61 ± 0.00 7.58 ± 0.01 3.14 ± 0.01 2.91 ± 0.01

Table content: <table><tr><td>Methods</td><td>TCGA √ DPE</td><td>News √ DPE</td><td>MIMIC √ DPE</td></tr><tr><td>DRGAN</td><td>0.31 ± 0.05</td><td>4.14 ± 0.11</td><td>0.51 ± 0.05</td></tr><tr><td>DRNet</td><td colspan="2">0.51 ± 0.05  *  4.39 ± 0.11  *</td><td>0.52 ± 0.05</td></tr><tr><td>DRN-W</td><td>0.50 ± 0.05  *</td><td>4.21 ± 0.11</td><td>0.53 ± 0.05</td></tr><tr><td>GPS</td><td>1.38 ± 0.01</td><td/><td/></tr></table>

Peer review (iclr2020):
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. To tailor to this specific heterogeneous dose-response curve estimation problem, the paper makes various architectural design modification from a regular GAN, such as using factual predictive performance as regularization, multitask for treatment option and dosage in the generator, hierarchical decomposition of the dose and response prediction in the discriminator, and permutation invariance when predicting the dosage-response curves for various treatment in the discriminator. These design choices are reasonably explained and the empirical gains due to these design choices are well documented (Table 2). A few concerns: 1. It is unclear why using a GAN that seeks to produce samples that are indistinguishable from the factual ones will be the right way to synthesize counterfactual samples. Intuition and further justification need to be provided. Related to this concern, do the authors also implicitly reason under the consistency assumption (the observed/factual samples are the realizations potential outcomes of the corresponding treatment-dosage pairs in the potential outcome framework) ? If so, the authors should also reflect that in their list of assumptions. 2. It is also unclear what kind of theoretical guarantees the proposed method could deliver. Since the setting considered by the paper is to formally estimate the dosage-response effect, such guarantees, while potentially challenging to derive, are nonetheless crucial to better understand the circumstances under which the proposed method is reliable/unreliable. This is especially the case given the elusive use of GAN to generate counterfactual outcomes as mentioned in the previous concern. 3. It will also be interesting to see whether the proposed method can be degenerated to handle simpler cases such as heterogeneous treatment effect estimation of binary treatment. Or alternatively, when there is only one treatment option, where when treated there is a dosage-response curve while when not treated the dosage is zero. Related to this question, is this also the way to parameterize whether a subject receives a treatment or not in the model?