SurvITE: Learning Heterogeneous Treatment Effects from Time-to-Event Data
Authors: Alicia Curth
Changhee Lee
Mihaela van der Schaar
Venue: NeurIPS 2021 Poster (2021)
Keywords: Heterogeneous Treatment Effects
Survival Analysis
Counterfactual Estimation
Causal inference

Abstract: We study the problem of inferring heterogeneous treatment effects from time-to-event data. While both the related problems of (i) estimating treatment effects for binary or continuous outcomes and (ii) predicting survival outcomes have been well studied in the recent machine learning literature, their combination -- albeit of high practical relevance -- has received considerably less attention. With the ultimate goal of reliably estimating the effects of treatments on instantaneous risk and survival probabilities, we focus on the problem of learning (discrete-time) treatment-specific conditional hazard functions. We find that unique challenges arise in this context due to a variety of covariate shift issues that go beyond a mere combination of well-studied confounding and censoring biases. We theoretically analyse their effects by adapting recent generalization bounds from domain adaptation and treatment effect estimation to our setting and discuss implications for model design. We use the resulting insights to propose a novel deep learning method for treatment-specific hazard estimation based on balancing representations. We investigate performance across a range of experimental settings and empirically confirm that our method outperforms baselines by addressing covariate shifts from various sources.

SurvITE: Learning Heterogeneous Treatment Effects from Time-to-Event Data
Alicia Curth University of Cambridge
Changhee Lee" Chung- Ang University changheelee@cau. actauuacc.aa. ac kr
Mihaela van der Schaar University of Cambridge University California, Los Angeles The Alan Turing Institute mv4720cam ac uk
amc2530cam ac.
Abstract
We study the problem of inferring heterogeneous treatment effects from time-to- event data. While both the related problems of (i) estimating treatment effects for binary or continuous outcomes and (ii) predicting survival outcomes have been well studied in the recent machine learning literature, their combination albeit of high practical relevance has received considerably less attention. With the ultimate goal of reliably estimating the effects of treatments on instantaneous risk and survival probabilities, we focus on the problem of learning (discrete-time) treatment- specific conditional hazard functions. We find that unique challenges arise in this context due to a variety covariate shift issues that go beyond mere combination of well- studied confounding and censoring biases. We theoretically analyse their effects by adapting recent generalization bounds from domain adaptation and treatment effect estimation to setting and discuss implications for model design. We use the resulting insights to propose a novel deep learning method for treatment- sppecifcc hazard estimation based on balancing representations. We investigate performance across a range of experimental settings and empirically confirm that our method outperforms baselines by addressing covariate shifts from various sources.
Introduction
The demand for methods evaluating the effect treatments, policies and interventions on individuals is rising as interest moves from estimating population effects to understanding effect heterogeneity in fields ranging from economics to medicine. Motivated by this, the literature proposing machine learning (ML) methods for estimating the effects of treatments on continuous (or binary) end-points has grown rapidly, most prominently using free-based methods [ [I] 12
4 5] Gaussian processes 2
[7 and, in particular, neural networks (NNs) g 9 IO
B3
14
15 In comparison, the ML
literature on heterogeneous treatment effect (HTE) estimation with time-to-event outcomes is rather sparse. This is despite the immense practical relevance of this problem e.g. many clinical studies consider time-to-event outcomes; this could be the time onset or progression of disease, the time to occurrence of adverse event such as stroke heart attack, the time until death of patient.
In part, the scarcity of HTE methods may be due to time-to-event outcomes being inherently more challenging to model, which is attributable to two factors [16]: (i) time-to-even outcomes differ from standard regression targets as the main objects of interest are usually not only expected survival times but the dynamics of the underlying stochastic process, captured by hazard and survival functions, and (ii) the presence censoring. This has led the development of rich literature on survival analysis particularly in (bio)statistics, see e.g. [16 17]. Classically, the effects of treatments in clinical studies with time-to-event outcomes are assessed by examining the coefficient of a treatment
Equal contribution
35th Conference on Neural Information Processing Systems (NcurIPS 2021).

---

indicator in a (remi-)parametrin model, e.g. Cox proportional hazards model [18]. which relies on the often unrealistic assumption that models are correctly specified. Instead, we therefore adopt the nonparametric viewpoint of van der Laan and colleagues [19 20
21 22] who have developed tools to incorporate MI. methods into the estimation of treatment-specific population average parameters. Nonparametrically investigating treatment effect heterogeneity, however, has been studied in much less detail in the survival context. While a number of tree-based methods have been proposed recently |23 24 25 26). NN-based methods lack extensions to the time-tO-event setting despite their successful adoption for estimating the effects of treatments on other outcomes - the only exception being 127]. who directly model event times under different treatments with generative models.
Instead of modeling event times directly like in [27], we consider adapting machine learning methods, with special focus on NNs, for estimation of (discrete-time) treatment- sppecifc hazard functions. We do so because many target parameters of interest in studies with time-to-event outcomes are functions of the underlying temporal dynamics; that is, hazard functions can be used to directly compute (differences in) survival functions, (restricted) mean survival time, and hazard ratios. We begin by exploring and characterising the unique features of the survival treatment effect problem within the context of empirical risk minimization (ERM); to the best of our knowledge, such an investigation is lacking in previous work. In particular, we show that learning treatfent-specifie hazard functions a challenging problem due to the potential presence of multiple sources of covariate shift: (i) -on-manddmized treatment assignment (confounding), (ii) informative censoring and (iii) a form of shift we term event cennnddddeed covariate shift, all of which can impact the quality of hazard function estimates. We then theoretically analyze the effects of said shifts on ERM. and use our insights to propose a new NN-hased model for treatment effect estimation in the survival context.
Contributions (i) We identify and formalize key challenges of heterogeneous treatment effect estimation in time-to-event data within the framework of ERM. In particular, as discussed above, we show that when estimating treatment-specifie hazard functions, multiple sources of covariate shift arise. (ii) We theoretically analyse their effects by adapting recent generalization bounds from domain adaptation and treatment effect estimation to our setting and discuss implications for model design. This analysis provides new insights that are of independent interest also in the context of hazard function estimation in the absence of treatments. (iii) Based on these insights, propose a new model (SurvITE) relying on balanced representations that allows for estimation of rreament-specine target parameters (hazard and survival functions) in the survival context, well as a sister model (SurvIHE), which can be used for individualized hazard estimation in standard survival settings (without treatments). We investigate performance across range of experimental settings and empirically confirm that SurvITE outperforms a range of natural baselines by addressing covariate shifts from various sources.
2 Problem Definition
In this section, we discuss the problem setup of heterogeneous treatment effect estimation from time-to-event data, our target parameters and assumptions. In Appendix A, we present a self- contained introduction to and comparison with heterogeneous treatment effect estimation with standard (binary/continuos) outcomes.
Problem setup. Assume we observe time-to-event dataset D = {(a;, x1, f., 6,) comprising realizations of the tuple X, A) IP for patients. Here, X E X' and A E {0, are random variables for a covariate vector describing patient characteristics and an indicator whether binary treatment was administered at baseline, respectively. Let T E T and C T denote random variables for the time-to-event and the time-to-censoring; here, events are usually adverse, e.g. progression/onset of discase even death, and censoring indicates loss of follow-up for a patient. Then, the observed time-to- toveent outcomes of each patient are described by T min C) and A = 1(T < C"), which indicate the time elapsed until either an event or censoring occurs and whether the event observed not, respectively. Throughout, treat survival time discrete" and the time horizon as finite with pre-defined maximum mmxx, so that the set of possible survival times is T= (1,.... Imax)
'Where necessary, discretization can be performed by transforming continuous valued times into set of contiguous time intervals, ie., T' = r implies I' E [tr, & + öt) where öt implies the temporal resolution.

---

We transform the short data structure outlined above to a so-called long data structure which can be used directly estimate conditional hazard functions using standard machine learning methods 20 We define two counting processes (1) and (c( which track events and censoring, i.e. (t) 1( < 1, A 1) and (t = 1(f < 1 0) for E T: both are zero until either an event or censoring occurs. By convention, we let Nr()) = Nc(() = O. Further, let n Nr (t-1)
0) be the indicator for an event occuring exactly time t; thus, for individual with 7 7 and A 1, (()) for all t * T, and (t)) 1 at the event time t= T. The conditional hazard is the probability that an event occurs at time т given that it does not occur before time T, hence it can be defined as [22
(()) = 1 (t)
M(Ta, a, x) P T, A 1 PP (7) - 1| 1|
(T-1) 0, a,
It is easy to see from I that given data in long format, \(T|a, x) can be estimated for any T by solving standard classification problem with (T) as target variable, considering only the samples at-risk at time T in each treatment arm (individuals for which neither event nor censoring has occurred
until that time point; i.c. the set I(T, def {i E [n] (T Nc(T- 1); = A, a}). Finally, given the hazard, the associated survival function s() = P(T 7|A a, x) can then be computed as S(T(a, x) = -Ilsrr (1 X(tla, x)). The censoring hazard Ac((aa, x) and
survival function Sc z) can be defined analogously.
Target parameters. While the main interest in the standard treatment effect estimation setup with continuous outcomes usually lies in estimating only the (difference between) conditional outcome means under different treatments, there is broader range of target parameters of interest in the time-to-event context, including both treatment-specific target functions and contrasts that represent some form of heterogeneous treatment effect (HTE). We define the treatment-specific (conditional) hazard and survival functions as
X (T|z) = (T = T 2 r, do - a, C 2 T), =x) x)
(T|x) = (T > C2, x
a. a.
EST (1 (1- (t(x))
Here, do(-) denotes [28] 's do-operator which indicates an intervention; in our context, do(AA
a, C,C 2 -) ensures that every individual is assigned treatment and observed at (not censored before) the time-step of interest [20). Below we discuss assumptions that are necessary to identify such interventional quantities from observational datasets in the presence of censoring.
Given (T(zz) and (T(zz) possible HTEs of interes
3 iincluee the difference in treatient-speciffc (T 5th (T(zz, the difference in restricted mean
survival times at time T, i... MITEeuro(T survival time (RMST) up to time i.e. HTE (x) = Ensl (z) (t&l/=))) (tx - 44-1). and hazard ratios. In the following, we will focus on estimation of the treatment specific hazard functions {X*
IeET as this can be used to compute survival functions and causal contrasts.
Assumptions. (1. Identification) To identify interventional quantities from observational data, it is necessary make number of untestable assumptions on the underlying data- generating process (DGP) (28 this generally limits the ability to make causal claims to set- tings where sufficient domain knowledge is available. Here, as 220 21), we assume the data was generated from the fairly general di- rected acyclic graph (DAG) presented in Fig. I As there are no arrows originating in hidden nodes entering treatment or censoring nodes, this graph formalizes (I.a) The 'No Hidden Confounders' Assumption in static treatment effect estimation and (1.b) The "Censoring At
Shift 1:
Shift
Shift 3:
Nr 1)
(t)
X2
x
X₂
Ne(t 1)
(r(t)
Figure I: The assumed underlying DAG, Covariates X can be split into (possibly overlapping) subsets X1, X2 and Xs. determining treatment selection, informative censoring, and event times, respectively.
Note: All parameters of interest to us are heterogeneous (also sometimes referred to as individualized), i.e. a function of the covariates X. while the majority of existing literature in (bio)statisties considers population average parameters that functions of quantities such as P(1 > 7 |do(AA = a)). which average over all X.
3

---

Random' Assumption in survival analysis [0]. treatment on event time implicitly requires that censoring can be 'switched off' i.c. intervened on. This graph implicitly also formalizes (1.c) The Consistency Assumption, i.e. that observed outcomes are 'potential' outcomes under the observed intervention, as each node in a DAG defined as a function of its ancestors and exogenous noise [20]. Under these assumptions, (T (r) = A(r(a, x)).
The latter is necessary here as estimating an effect of
(2. Estimation) To enable nonparametric estimation of (T(2)) for some T, we additionally need to assume that the interventions of interest are observed with non-zero probability; within
different literatures these assumptions are known under the label of 'overlap' or 'positivity' [ 19
In particular, for 0 < 61, C2, < 1 we need that (2.a) < P(A a|X = 2) < 1 i.c. treatment assignment is non-deterministic, and that (2.b) (t) for all < т, iee no individual will be deterministically censored before т. Finally, because (T[E) is probability defined conditional on survival up to time т, we need to assume that (2.c) ((Nr(7-1) = A all assumptions in more detail in Appendix C.
0,4 a.
X x) > 63 > O for it to be well-defined. We formally state and discuss
3 Challenges in Learning Treatment-Specti Hazard Functions using ERM
Preliminaries: ERM under Covariate Shift. Recall that in problems with covariate shift, the training distribution X,Y Y ~ Qo(+) used for ERM and target distribution X, Y (-) are mis- matched: One assumes that the marginals do not match, i.c. QX()) * Q1 (X), while the conditionals If the hypothesis class H. used in ERM does not contain the truth (or in the presence of heavy regularization), this can lead to suboptimal hypothesis choice as arg minge B ()(Y h())] o arg minkeR E E () h())) in general.
remain the same, i.e. QY [X) Q1 (Xl)) 22).
3.1 Sources of Covariate Shift in Learning Treatment-Speeiffe Hazard Functions
We now consider how to learn a treatment-specifie hazard function X (r|zz) from observational data using ERM. As detailed in Section 2 we exploit the long data format by realizing that X (T(zz can be estimated by solving a standard classification problem with (()) as dependent variable and as covariates, using only the samples risk with treatment status a, i.e. I(T,, a), which corresponds to solving the empirical analogue of the problem
(T(x) arg min Ex,(r)))P., () (X)] harEH
(3)
where (T with P (X) (r- Nc(r-1) 0, = a) = P(X > A a). If the loss function l is chosen to be the log-loss, this corresponds to optimizing the likelihood of the hazard.
Pa.r to refer the observational (at-risk) distribution Pa.r (T))
The observational (at-risk) covariate distribution (X), however, is not our target distribution: instead, to obtain reliable hazard estimates for the whole population, we wish to optimize the fit over the population at baseline, i.c. the marginal distribution X ~ P(X) which we will refer to as Po below to emphasize being the baseline at-risk distribution Here, differences between PX( and the population at-risk IP,, (X) can arise due to three distinct sources of covariate shift:
(Shift 1) Confounding/treatmen selection bias: if treatment is not assigned completely random, then PXX]AA = a) * Po(X) and the distribution of characteristics across the treatment arms differs already at baseline, thus (P, * (X(()) in general.
(Shifi 2) Censoring bias: regardless of the presence of confounding, if the censoring hazard is not independent of covariates, i.e. (c(Tla, x) * (c(T)a), then the population at-risk changes over time such that Pa.r, (X) * Po.r, (X) # Po(X) in general. If, in addition, there are differences between the treatment-specife censoring hazards, then the at-risk distribution will also differ across treatment arms at any given time-point, i.c. Fa,r()) + (X) for T > 1 in general.
(Shift 3) Event-induced shifts: Counterintuitively, even in the absence of both confounding and censoring, there will be covariate shift in the at-risk population if the event-hazard depends on covariates, i.e. if X x) & X(r(a) then Pa.. (X) + Pa (X) A (X) in general. Further, if there are heterogenous treatment effects, then P. a,r (X) + IP a (X) for > 1 in general.
"ith slight abuse of notation, will use Po and Pa.r also to refer to densities of continuous æ

---

What makes the survival treatment effect estimation problem unique? While Shift arises also in the standard treatment effect estimation setting, Shift and Shift 3 arise uniquely due to the nature of time-to-event data Thus, estimating treatment effects from time-to-event data is inherently more involved than estimating treatment effects in the standard static setup, as covariate shift at time horizon т > 1 can arise even in randomized control trial (RCT). Thus, in addition to the overall at-risk population changing over time, both treatment effect heterogeneity and treatment-dependent censoring can lead to differences in the composition of the population at-risk in each treatment arm. Further, Shifts I, 2 and 3 can also interact to create more extreme shifts; e.g. if treatment selection is based on the same covariates as the event process (i.e. X2 X3 in Fig. then event-induce shift can amplify the selection effect over time (refer to Appendix E for a synthetic example of this).
3.2 Possible Remedies and Theoretical Analysis
A natural solution to tackle bias in ERM caused by covariate shift is to use importance weighting
330 i.e to reweight the empirical risk by the density ratio of target (X and observed distribution
Far (X). If we wanted to obtain hazard estimator for (T, a), optimized towards the marginal
population, optimal importance weights are given by (x) = (r) = (2))))(().T) with Pr.a =
Pu(z)
Pi.o
P( 2 A a). (x) = P(A = a|X = x> the propensity score, and r" r) = P(T 2 r|A= a, = z) the probability to be at risk, i.e. the probability that neither event nor censoring occurred before time r. These weights are well-defined due to the overlap assumptions detailed in Sec. 2 however, they are in general unknown as they depend on the unknown larget parameters (T|x) through r (2., r). Further, especially for large r, these weights might be very extreme even if known, which can lead to highly unstable results [31] making biased yet stabilized weighting schemes, c.g. truncation. a good alternative. Therefore, we only assume access to some (possibly imperfect) weights Wa,r (z) s.t. Ex~P..- [wa,r (z)] = I, so that we can create a weighted distribution Pa, Wa,- (1) PC (1)) (Note: (z) can be recovered by using Wa, 1.)
Either instead of [ [ or in addition to weighting [T 12 14 32], the literature on learning balanced representations for static treatment effect estimation has focused on finding different remedy for distributional differences between treatment arms: creating representations $ X - R which have similar (weighted) distributions across arms measured by an integral probability metric (IPM), motivated by generalization bounds. As we show below, we can exploit a similar fcature in our context by finding a representation that minimizes the IPM term not between treatment arms, but between covariate distribution at baseline Po and pur,-. The proposition below bounds the target risk
of a hazard estimator ) (x) h (z)) relying on any representation. The proof, which relies on the concept of excess target information loss, proposed recently to analyze aammain-dvvesaane training (33), and the standard IPM arguments made in e.g. (32), stated in Appendix C.
Proposition 1. For fixed a, T and representation + X - R, let PP P and Te denote the target, observational, and weighted observational distribution of the representation T. Define the pointwise losses
(n.e(2;a,7) Ey
A a] ln del Ey (TC) - -AA=aa
(4)
of (hazard) hypothesis h = harr R J0. 1] w.r.t. distributions in covariate and representation space, respectively. Assume there exists constant Co > S.s.L. Co Lo " (0,a.7) 9 for some family of functions G. Then we have that
(ew.eoen..e (X;a,)( & Ex [War ( (X))x, +Co IPMc or Pur.
(h)
(5)
Target Risk
Weighted observational risk
Distance in 4 drspace
Info loss
'Interestingly, changes of the at-risk population over time arise also in standard survival problems (without treatments); yet in the context of prediction these do not matter: as the at-risk population at any time-step is also the population that will be encountered at test-time, this shift in population over time is not problematic, unless is caused by censoring. If, however, our goal is estimation the best target parameter (here: the hazard at a specific point in time 7) over the whole population, this corresponds to a setting where the ideal evaluation is performed on a population different from the observed one and hence requires careful consideration of the consequences of the covariate shifts discussed above.
5

---

where IPMg(IP (P, Q) - supges |f g ((z) - [x)))d and we define the excess target informa- tion loss The (h) analogously to 133] as ni. (h) def Ex~r Keng (X) - Spm.. (xm.,.r(X)) pX)] with E dhr
21.00 (0: a, r) - (n, a. For invertible D 1/4 (h) = 50e ((() = 0.
Unlike the bounds provided in 9 invertible; we consider this feature important as none of the works listed actually enforced invertibility in their proposed algorithms. Given bound ( it is easy to see why non-invertibilty can be useful: for any (possibly non-invertible) representation for which it holds that (T)) (), A, it also (z)
this bound does not rely on representations to be
holds that n (h) = fpº ze (z)) = Spu.* p(z) = O and the causally identifying restrictions continue to hold. A simple representation for which this property holds is selection mechanism that chooses only the causal parents of Y (r) from within X; if X can be partitioned into variables affecting the instantaneous risk (X3 in Fig. and variables affecting only treatment assignment (X1 \ Xx)) and/or censoring mechanism (X2 X 3), then the IPM term can be reduced by representation which drops the latter two sets of variables- or irrelevant variables correlated with any such variables without affecting v (h). As a consequence, event-induced covariate shift can generally not be fully corrected for using non-invertible representations (unless the variables affecting event time are different at every time-step). Further, given perfect importance weights w", both ni (h) and IPM term are zero.
Except for the dependence n (h), this bound differs from the regression -based bound for survival treatment effects stated in [27) (which is identical to the original treatment effect bound in (9) in that we have dependence on in the IPM term, which, among other things, explicitly captures the effect of censoring. Our bound motivates that, instead of finding representations that balance treatment- and control group at baseline (or at each time step) we should find representations that balance F2, towards the baseline distribution Pd for each time step, which motivates our method detailed below. If, instead, we would apply the IPM-term to encourage only the arm-specific at-risk distributions at each time-step to be similar, this would correct only for shifts due to (i) confounding at baseline, (ii) treatment-inducee differences in censoring and (iii) erattenn--nnduced differences in events. It will, however, not allow handling the event- and censoring-induce status. Note that this bound therefore also motivates the use of balanced representations for modeling time-to-event outcomes in the presence of informative censoring even in the standard prediction setting, which is finding that could be of independent interest for the ML survival analysis literature.
shifts that occur regardless of treatment
3.3 From hazards to survival functions
If the ultimate goal is to use the hazard function to estimate survival functions as (r(zz) = His, (1 (1- (1]=)) the best target population to consider during hazard estimation may not be the marginal distribution. Instead, the optimal target population may depend on the metric by which we wish to evaluate the resulting survival function. If we wanted to find the survival function that maximises the complete data likelihood (corresponding to the hypothetical setting in which we intervened set A a and C' 2 7), the target population (at each time step t) would be P( 2 Ł, do(AA = a, C 2 f)) - the population that preserves event-induced shift but removes selection- and censoring-inducee shifts. If, instead, we focused on the MSE of estimating the survival function (as in our experiments), it becomes more difficult to derive an exact target population for estimating the hazards. If we assume access to a perfect estimate of the survival function for the first 1-1 time steps (i.e. S(r-1|x) = 1(x) and focus only on estimating the next hazard, X2 (T|X)) we can write
Ex-e, -((r)/))(((((())
Ex 100e)(1077--1133)--xx)x))) -1
= (3w_=/157--11099((/1111111
and notice that the MSE will implicitly down-weigh individuals with lower survival probability
Defining exact target population for the hazard when the goal is to also estimate the survival function well is thus not straightforward, making exact importance weighting difficult. Additionally, unlike the marginal population, interventional populations which change over time, such P(X 2 t, do(A A = a, C 2 t)), are never observed in practice and hence cannot be used to perform balancing
'Due to the square in the term (r 1[x)², this will be even more extreme than exact up spighiitting of the population [1 2 v, do(A a, C r)).

---

Big ay Rareer
Cr
x
(x))
a,7 (0p)
Ro.1
Agz
RoAAmax
Figure SurvITE architecture.
regularization of representation using empirical IPM-terms. Therefore, we refrain from using importance weighting in our method (which is described in the next section), and resort to using the marginal population for balancing regularization the representation throughout. Intuitively, expect that doing will not over-correct for event-induced
as outlined in the previous section, shifts that are predictive of outcome (and should hence be preserved) as such "over-balancing" would reduce the predictive power of the representation, which would immediately be penalized by the presence of the expected loss component in the bound" Additionally, we expect that using the marginal population for balancing could be useful also for estimating the survival function even the absence of selection- and censoring-induced induced shifts, as it may help to remove the effect variables that appear spuriously correlated with outcome over time.
4 SurvITE: Estimating HTEs from Time-to-Event Data
Based the theoretical analysis above, propose novel deep learning approach to HTE estimation from observed time-to-event data, which we refer to as SurvITE (Individualized Treatment Effect estimator for Survival analysis) The network architecture is illustrated in Figure Note that even in the absence of treatments we can use this architecture for estimation of hazards and survival functions by using only one treatment SurvIHE (Individualized Hazard Estimator for Survival analysis) - is of independent interest in the standard survival setting, it tackles Shifts & 3. Below, describe the empirical loss functions we use to find representation O and hypotheses ha.r-
0. As we show in the experiments, this version of method
Let & X R denote the representation (parameterized by 0g) and harr R [0, 1] the hazard estimator for treatment a and time т (parameterized by Oh...). each implemented as a fully- connected neural network. While the output heads are thus unique each treatment-gne groppp time-step combination, we allow hazard estimators to share information by using one shared representation for all hazard functions. This allows for both borrowing of information across different a, and significantly reduces the number of parameters of the network. Then, given the time-to-event data D, use the following empirical loss functions for the observational risk and the IPM term:
trex
Crisk(Ap tmax 1=
Σ
m, + (a)
fax
Cipm (0s)
Σ
Wass I (0()))
(E(0,1)
where W Wass(:;-) is the finite-sample Wasserstein distance 34 refer Appendix D for further detail. Note that Lipm (9s), which penalizes the discrepancy between the baseline distribution and each at-risk distribution - simultaneously tackles all three sources of shifts. Further, na, = [I(r, a) is the number of samples at-risk in each treatment arm; its presence ensures that each a, 7-combination contributes equally to the loss. Overall, we can find and ha. 's that optimally trade off balance and
In practice, we ensure this by weighting contribution the IPM term by a hyperparameter that is chosen to preserve predictive performance the representation (sce Appendix D).
"The source code for SurvITE is available in https:t: htt///ggtttt.t//////eeeeee com/ http:/chh88//uu1111
7

---

predictive power as suggested by the generalization bound ( by minimizing the following loss: (0s, 0%.) Crieh 0n) + (pmm)
(6)
where 9,. (fh..-)arn (o. and 3 > 0 is a hyper-parameter. The pseudo-code of SurvITE, the details of how to obtain W ass(;;)) and how we set / can be found in Appendix D.
5 Related work
Heterogeneous treatment effect estimation (non-survival) has been studied in great detail in the
recent ML literature. While early work built mainly tree-based methods [ 2 4. many other methods, such as Gaussian processes 6 and GANS (35, have been adapted to estimate
HTEs. Arguably the largest stream of work 8 9
IO III 13 14 [5] built on NNs, due their flexibility and ease of manipulating loss functions, which allows for easy incorporation of balanced representation learning as proposed in 8 9 and motivated also the approach taken in this paper.
Another popular approach has been to consider model-agnostic (or 'meta-learner" [36] strategies,
which provide a 'recipe' for estimating HTEs using any predictive ML method [36 37 38 [5).
Because of their simplicity, the single model (S-learner) which uses the treatment indicator as an additional covariate in otherwise standard model-fitting- - and two model (T-learner) which splits the sample by treatment status and fit two separate models strategies [6]] can be directly applied to the survival setting by relying on standard survival (prediction) method as base-learner.
ML methods for survival prediction continue to multiply; here we focus on the most related class of methods namely on those nonparametrically modeling conditional hazard or survival functions and not on those relying on flexible implementations of the Cox proportional hazards model (e.g.
One popular nonparametric estimator of survival functions is [4] 's random survival forest, which relies on the Nelson-Aalen estimator to nonparametrically estimate the cumulative hazard within tree-leaves. The idea of modeling discrete-time hazards directly using any arbitrary classifier and long data- structures goes back to least 49), with implementations using NN-based methods presented e.g.
39 40 4T) or modeling (log-)time
regression problem (e.g.
42
43
44
45
46
[50 51 52 53 544 models the probability mass function instead of the hazard, and [55] use labels 1[T > the to estimate the survival function directly using multi-task logistic regression. For a more detailed overview of different strategies for estimating survival functions, refer to Appendix B.
Estimating HTES from time-to-event data has been studied in much less detail. tree-based nearest-neighbor estimates to estimate expected differences in survival time directly, and (24] use a BART-based S-Iearner to output expected differences in log-survival time. [56]
23 25 use
performed simulation study using different survival prediction models as base-lcarners for two- model approach to estimating the difference in median survival time. Based on ideas from the semi-paramen parammeaeetric efficiency literature, (26) and [57] propose estimators that target the (restricted) mean survival time directly and consequently do not output estimates of the treatment- pratment-specific hazard or survival functions. We consider the ability to output treatment- -eeccie---pecifi predictions an important feature of model if the goal use model output to give decision support, given that it allows the maaeer to trade-off relative improvement with the baseline risk of patient. Finally, (27) recently proposed generative model for treatment-speche representations to balance only the treatment groups at baseline. This model does not output hazard- or survival functions, but can provide approximations by performing Monte-Card sampling.
event times which relies on balancing
Experiments
Unfortunately, when the goal is estimating (differences of) survival functions (instead of predicting survival), evaluation on real data will not reflect performance w.r.t. the intended baseline population. Therefore, we conduct range of synthetic experiments with known ground truth. We evaluate the effects of different shifts separately by starting with survival estimation without treatments, and then introduce treatments. Finally, we use the real-world dataset Twins (58) which has uncensored survival outcomes for twins (where the treatment is 'being born heavier'), ) and is hence free of Shifts & 2.
Baselines. Throughout, we use Cox regression (Cox), model using separate logistic regression to solve the hazard classification problem at each time-step (L.R-sep), random survival forest (RSF), and deep learning- beasengg-based time-to-event method (54] (DeepHit) natural basclines; when there
8

---

51: Shrtko
S2: 5t.00
53: 5 IC 3,2- (9,10))
54: 5P8g) IC= 3,5 (0.00000000000 (0,30))
0.10
0.4
0.31
1.35
6.15
035
6.11
0.22
0.81
15
040
10
15
0.22
30
32
14
15
20
Cox
LRoep
C5A
SurvIT (CFR-11
SurVITE (CFR-2)
SurvIHE (no IPMI/ SurVITE (no IPM0
SuNVHE SurVITE
NM7 by IC
53: o0ll by C (P- 1011
MMM by i (2
St: N7E by One (9, 10300
0.33
1.00
0.09
0377
0.07
10)
0.55
(8.35)
040
01.21
0.5
040
Timee
Time
Time
Figure 3: RMSE of estimating the survival function (t]x)) (top) and the treatment effect HT (t)r) (bottom) for different time steps across synthetic settings. Averaged across 5 runs.
are treatments, we use them in two-model (T-learner) approach. In settings with treatments, we additionally use the CSA-INFO model of [27] (CSA), where we use its generative capabilities to approximate target quantities via monte-carle sampling. Finally, consider ablations of SurvITE (and SurvIHE); in addition to removing the IPM term (SurvITE (no IPM)). we consider two variants of SurvITE based on I CFRNet balancing term: SurvITE (CFR-1) creates a representation balancing treatment groups at baseline only, and SUrVITE (CFR-2) creates a representation optimiz- ing for balance of treatment groups each time step (i.e. balancing towards Po). We diseuss implementation in Appendix D.
Synthetie Experiments. We consider range of synthetic simulation setups (S1-S4) to highlight and isolate the effects of the different types of covariate shift. As event and censoring processes, we
(0.10(-52 a (1ss 2 0) + 0.5) for S 10
(0.10 -a a (1 2 0) 0.5)) fort > 10 >c((|x) 0101z
with treatment assignment mechanism a ~ Bern o kp)), with a the sigmoid function. Additionally, we assume administrative censoring at / = 30 throughout, i.e., A (30 1, marking e.g. the end of a hypothetical clinical study. Covariates are generated from a 10-dimensional multivariate normal distribution with correlations, i.c. X N(, Σ) where Σ (1 - p)1 + p11 with 0.2. We use 5000 independently generated samples each for training and testing.
In S1, begin with the simplest case treatments and censoring using only (t|x=) to generate events, considering only event-induced shift (Shift 3). In S2, we introduce informative censoring using (t(zrz (Shift 2+3). In S3, use treatments and consider biased treatment assignment (without censoring) (Shift 1+3). In S4, we consider the most difficult case with all three types of shift (Shift 1+2+3). In the latter two settings, vary treatment selection by changing (i) whether the covariate set overlaps with the event-inducing covariates (P=(1, ( or not (P= (9, 10)) and (ii) the selection strength & E 3). We present exploratory plots of these DGPs in Appendix
Fig. (top) shows performance on estimating (t]z) -Ilac (1 (k|x)) for all scenarios and methods, while Fig. 3 (bottom) shows performance on estimating the difference in survival functions (HT' (t) (t) for selection of methods (for readability. full results in Appendix F). In Table [ we further evaluate the estimation of differences in RMST (HTErmec(zz) Results for hazard estimation and additional performance metrics for survival prediction are reported in Appendix F. We observe that SurvITE (/SurvIHE) performs best throughout, and that introduction of the IPM term leads to substantial improvements across all scenarios. In S1 with only event-induced covariate shift and in S3/4 when treatment selection and event- weent-inuucing covariates overlap (P=(1,2)) balancing cannot remove all shift as the shift-inducing covariates are predictive of outcome; however, even here the IPM-term helps it encourages dropping other covariates (which appear imbalanced due to correlations X). While our method was motivated by theory for estimation of hazard functions, thus indeed also leads to gains in survival function estimation. As expected, both Cox and LR-sep

---

Table RMSE on estimation of HTE (x) (mean + --)-CCI for different times for the Synthetic and Twins datasets (Ls are the 25 & 75th and 75 & 95th percentiles of event times, respectively).
Methods
83 (C
10
3, no overlag)
20
S41C
10
3, no overlop)
L
Twins (no censoring)
30
180
Twins (vensoring)
30
180
Cox
0.434-00.03
0.03
0.424+:0.02
1.04700044
2.85
20.33-1-0.50
2:88:+0009
20.60++0.50
RSF
0.328-0002
1.027--0.03
0.312:1:002
0.
3.15:0077
22.424-0.36
A.1840088
22.02:4:0:66
LR-sep
0.:2::0.02
11111:00007
0.418:1:002
1.149-10.04
2.4400.1
20.60:1:0.53
2.4-:0.10
20.06::0.52
Deepliit
0.347=#0.03
0.821
0.361++0.08
0.830=0..15
2.95
20.89++1.91
2.86-+0..9
20.59
CSA
0.01
2.098:10:26
0.406:1:0001
19952:0.12
3.42:0.12
1..20:1:121
4.41:0554
47.79:1:.55
SurvITE (no IPM)
0.275=00004
0.84320.11
0.310:20:05
0,930:00.11
2,8020010
19.80-21.011
2.85:20.22
20.00
SurvITB (CFR-1)
0.269=00.04
+-009
0.341++0002
1001=+0.10
0.68
19.16-+00.7
2.67:
19.10+00885
SureITE (CFR-2)
0:23:=0001
0:691-100008
0.291:20.00
0815:20.15
2.61
18.69:2:0.64
2.69:2:0222
19:20=#1.44
SurvITE
0.225 0.03
0.687 0.008
0.237
0.703::0.06
2.53
18.34
2.63:40.10
14.6610..5
do not perform well as they are misspecified, while the nonparametrie RSF is sufficiently flexible to capture the underlying DGP and usually performs similarly to SurvITE (architecture only), but is outperformed once the IPM term is added. For readability, we did not include DeepHit in Fig. 3 using table F.1 presented in Appendix F, we observe that DeepHit performs worse than the SurvITE architecture without IPM term, indicating that our model architecture alone is better suited for estimation of treatment-spedd specific survival functions (note that [5:]] focused mainly on discriminative (predictive) performance, and not on the estimation of the survival function itself). Therefore, upon addition of the IPM-terms, the performance gap between SurvITE and DeepHit only becomes larger.
A comparison with ablated versions highlights the effect of using the baseline population to define balance; naive balancing across treatment arms (either at baseline - SurvITE(CFR-I) or over time- SurvITE(CFR-22 2)) is not as effective as using the baseline population as target, especially at the later time steps where the effects of tine-varying shifts worsen. While SUrvITE(CFR-2) almost matches the performance of the full SurvITE in S3, it performs considerably worse in S4, indicating that this form of balancing suffers mainly due to its ignorance of censoring. Finally, a comparison with CSA highlights the value of modeling hazard functions directly: we found that Monte-Caid Carlo approximation of the survival function using the generated event times gives very badly calibrated survival curves as event times generated by CSA were concentrated in a very narrow interval, leading to survival estimates of O and elsewhere. Its performance on estimation of RMST was likewise poor; we conjecture that this is due to (i) CSA modeling continuous time, while the outcomes were generated using a coarse discrete time model, and (ii) the significant presence of administrative censoring.
Real data: Twins. Finally, we consider the Twins benchmark dataset, containing survival times (in days, administratively censored at t=365) of 1 |1400 pairs of twins, which is used in (58) 355 to measure HTES of birthweight on infant mortality. We split the data 50/50 for training and testing (by twin pairs), and similar to (35), use a rovariate-based sampling mechanism to select only one twin for training to emulate selection bias. Further, we consider a second setting where we additionally introduce dewariate-dependee censoring. For all discrete-time models, we use a non-uniform discretization to construct classification tasks because most events are concentrated in the first weeks. more detailed description of the data and experimental setup can be found in Appendix E. As the data is real and ground truth probabilities are unknown, HTEyysse (ar) is suited best to evaluate performance on estimating effect heterogeneity. The results presented in Table largely confirm our findings on relative performance in the synthetic experiments; only RSF performs relatively worse on this dataset.
Conclusion
We studied the problem of inferring heterogeneous treatment effects from time-to-oe to-enent data by focusing on the challenges inherent to reatment-specinn hazard estimation. We found that a variety of covariate shifts play role in this context, theoretically analysed their impact, and demonstrated across a range of experiments that our proposed method SurvITE successfully mitigates them.
Limitations. Like all methods for inferring causal effects from observational data, SurvITE relies on a set of strong assumptions which should be evaluated by domain expert prior to deployment in practice. Here. the time-to-event nature of our problem adds additional assumption ('random cen- soring') ) to the standard 'no hidden confounders' assumption in classical treatment effect estimation. If such assumptions are not properly assessed in practice, any causal conclusions may be misleading.
10

---

Acknowledgments and Disclosure of Funding
We thank anonymous reviewers well members of the vanderschaar-lab for many insightful comments and suggestions. AC gratefully acknowledges funding from AstraZeneca. CL was supported through the IITP grant funded by the Korea government(MSIT) (No. 2021-0-01341, AI Graduate School Program, CAU), Additionally, MvdS received funding from the Office of Naval Research (ONR) and the National Science Foundation (NSF, grant number 1722516).
References
[1] Jennifer L Hill. Bayesian nonparametric modeling for causal inference. Journal of Computa- tional and Graphical Statistics, 20(1):217-240, 2011.
[2] Susan Athey and Guido Imbens. Recursive partitioning for heterogeneous causal effects. Proceedings of the National Academy of Sciences, 1 113(27):7353-7360 2016.
[3] Stefan Wager and Susan Athey. Estimation and inference of heterogeneous treatment effects using random forests. Journal of the American Statistical Association, 113(523):1228--2222 2018.
[4] Susan Athey, Julie Tibshirani, Stefan Wager, et al. Generalized random forests. The Annals of Statistics, 47(2): 148-11788 2019.
[5] P Richard Hahn, Jared S Murray, Carlos M Carvalho, et al. Bayesian regression tree models for causal inference: Regularization, confounding, and heterogencous effects (with discussion). Bayesian Analysis, 15(3):955-1055 2020.
[6] Ahmed M Alaa and Mihaela van der Schaar. Bayesian inference of individualized treatment effects using multi-task gaussian processes. Advances Neural Information Processing Systems, 10:3424-3433 2017.
[7] Ahmed Alaa and Mihaela van der Schaar. Limits of estimating heterogeneous treatment effects: Guidelines for practical algorithm design. In International Conference on Machine Learning, pages 129 138, 2018.
[8] Fredrik Johansson, Uri Shalit, and David Sontag. Learning representations for counterfactual inference. In International conference on machine learning, pages 3020-3029. PMLR, 2016.
[9] Uri Shalit, Fredrik D Johansson, and David Sontag. Estimating individual treatment effect: generalization bounds and algorithms. In International Conference on Machine Learning, pages 0076-3085. PMLR, 2017.
[10] Fredrik D Johansson, Nathan Kallus, Uri Shalit, and David Sontag. Learning weighted repre- sentations for generalization across designs. arXiv preprint arXiv:1802.088 1802.08598, 2018.
[11] Claudia Shi, David Blei, and Victor Veitch. Adapting neural networks for the estimation of treatment effects. In Advances in Neural Information Processing Systems, pages 2507-2517, 2019.
[12] Negar Hassanpour and Russell Greiner. Counterfactual regression with importance sampling weights. In IJCAI, pages 5880-5887, 2019.
[13] Negar Hassanpour and Russell Greiner. Learning disentangled representations for counterfactual regression. In International Conference on Learning Representations, 2020.
[14] Serge Assaad, Shuxi Zeng, Chenyang Tao, Shounak Datta, Nikhil Mehta, Ricardo Henao, Fan Li, and Lawrence Carin Duke. Counterfactual representation learning with balancing weights. In International Conference Artificial Intelligence and Statistics, pages 1972-1980. PMLR. 2021.
[15] Alicia Curth and Mihaela van der Schaar. Nonparametric estimation heterogeneous treatment effects: From theory to learning algorithms. Proceedings of the 24th International Confer- ence on Artificial Intelligence and Statistics (AISTATS) 2021 (To Appear); arXiv preprint 21X1v:2101.10943, 2021.

---

[16] Gerhard Tutz, Matthias Schmid, et al. Modeling discrete time-to-event data. Springer. 2016.
[17] John P Klein and Melvin L Moeschberger. Survival analysis: techniques for censored and truncated data, volume 1230. Springer, 2003.
[18] David R. Cox. Regression models and life tables (with discussion). Journal of the Royal Statistical Society. Series B, 34:187-220, 1972.
[19] Mark J Van der Laan and Sherri Rose. Targeted learning: causal inference observational and experimental data. Springer Science & Business Media, 2011.
[20] Ori M Stitelman and Mark van der Laan. Collaborative targeted maximum likelihood for time to event data. The International Journal of Biostatistics, 6(1), 2010.
[21] Ori M Stitelman, C William Wester, Victor De Gruttola, and Mark J van der Laan. Targeted maximum likelihood estimation of effect modification parameters in survival analysis. The international journal of biostatistics,, 7(1), 2011.
[22] Weixin Cai and Mark J. der Laan. One-step targeted maximum likelihood estimation for time-to-event outcomes. Biometric Methodology, 2019.
[23] Sami Tabib and Denis Larocque. Non-parametric individual treatment effect estimation for survival data with random forests. Bioinformatics, 36(2):629-636, 2020.
[24] Nicholas C Henderson, Thomas A Louis, Gary L Rosner, and Ravi Varadhan. Individualized treatment effects with censored data via fully nonparametric bayesian accelerated failure time models. Biostatistics, 21(1):50-68, 2020.
[25] Weijia Zhang, Thuc Duy Le, Lin Liu, Zhi-Hua Zhou, and Jiuyong Li. Mining heterogencous causal effects for personalized cancer treatment. Bioinformatics, 33(15):2372-2378 2017.
[26] Yifan Cui, Michael R Kosorok, Stefan Wager, and Ruoging Zhu. Estimating heteroge- neous treatment effects with right-censored data via causal survival forests. arXiv preprint W:Xiv:200100988) 2020.
[27] Paidamoyo Chapfuwa. Serge Assaad, Shuxi Zeng, Michael J Pencina, Lawrence Carin, and Ricardo Henao. Enabling counterfactual survival analysis with balanced representations. In Proceedings of the Conference Health, Inference, and Learning, pages 133-145, 2021.
[28] Judea Pearl. Causality. Cambridge University Press, 2009.
[29] Wouter M Kouw and Marco Loog. An introduction to domain adaptation and transfer learning. arXiv preprinl arXiv: 1812. 18806, 2018.
[30] Hidetoshi Shimodaira. Improving predictive inference under covariate shift by weighting the log-likelihood function. Journal of statistical planning and inference, 90(2):227-244, 2000.
[31] Corinna Cortes, Yishay Mansour, and Mehryar Mohri. Learning bounds for importance weighting. Nips, volume 10, pages 442-450. Citeseer, 2010.
[32] Fredrik D Johansson, Uri Shalit, Nathan Kallus, and David Sontag. Generalization bounds and representation learning for estimation of potential outcomes and causal effects. arXiv preprint arXiv:2001.07426 2020.
[33] Fredrik D Johansson, David Sontag, and Rajesh Ranganath. Support and invertibility in domain- invariant representations. The 22nd International Conference on Artificial Intelligence and Statistics, pages 527-536. PMLR, 2019.
[34] Marco Cuturi and Arnand Doucet. Fast computation of wasserstein barycenters. In International Conference on Machine Learning, pages 685-693. PMIR, 2014.
[35] Jinsung Yoon, James Jordon, and Mihaela van der Schaar. Ganite: Estimation of individualized treatment effects using generative adversarial nets. In International Conference on Learning Representations, 2018.
12

---

[36] Sören R Künzel, Jasjeet S Sekhon, Peter J Bickel, and Bin Yu. Metalearners for estimating heterogeneous treatment effects using machine learning. Proceedings of the national academy of sciences, 1 116100)4156-416 2019.
[37] Edward H Kennedy, Optimal doubly robust estimation of heterogeneous causal effects, arXiv preprint arXiv:2004. 14497, 2020.
[38] Xinkun Nie and Stefan Wager. Quasi-oracle estimation of heterogeneous treatment effects. Biometrika, 108(2):299-3 219. 2021.
[39] David Faraggi and Richard Simon. A neural network model for survival data. Statistics in medicine, 14(1):73-82, 1995.
[40] Jared Katzman, Uni Shaham, Jonathan Bates, Alexander Cloninger, Tingting Jiang, and Yuval Kluger. Deep survival: A deep cox proportional hazards network. arXiv preprint 1606, 2016.
[41] Margaux Luck, Tristan Sylvain, Héloïse Cardinal, Andrea Lodi, and Yoshua Bengio. Deep learning for patient-specific kidney graft survival analysis. arXiv preprint 100454,, 2017.
[42] Torsten Hothorn, Kurt Hornik, and Achim Zeileis Unbiased recursive partitioning: condi- tional inference framework, Journal of Computational and Graphical Statistics, 15(3):651-674 2006.
[43] R. Ranganath, A. Perotte, N. Elhadad, and Blei. Deep survival analysis. In Proceedings of the 1st Machine Learning for Healthcare Conference (MLHC 2016), 2016.
[44] Paidamoyo Chapfuwa, Chenyang Tao, Chunyuan Li, Courtney Page, Benjamin Goldstein, Lawrence Carin Duke, and Ricardo Henao, Adversarial time-to-event modeling. In International Conference Machine Leurning, pages 735-744. PMLR, 2018.
[45] Jon Arni Steingrimsson, Liqun Diao, and Robert L Strawderman. Censoring unbiased regression trees and ensembles. Journal of the American Statistical Association, 114655))370--33,, 2019.
[46] Jon Arni Steingrimsson and Samantha Morrison. Deep learning for survival outcomes. Statistics in medicine, (9(17)23339-23 2349, 2020.
[47] Anand Avati, Tony Duan, Sharon Zhou, Kenneth Jung, Nigam H Shah, and Andrew Y Ng. Countdown regression: sharp and calibrated survival predictions. In Uncertainty Artificial Intelligence, pages 145-155. PMLR, 2020.
[48] Hemant Ishwaran, Udaya B. Kogalur, Eugene H. Blackstone, and Michael S. Lauer. Random survival forests. The Annals of Applied Statistics, 2(3):841-860, September 2008
[49] Charles C Brown. On the use of indicator variables for studying the time-dependence of parameters in response-time model. Biometrics, pages 863-872, 1975.
[50] Elia Biganzoli, Patrizia Boracchi, Luigi Mariani, and Ettore Marubini. Feed forward neural networks for the analysis of censored survival data: a partial logistic regression approach. Statistics in medicine, 17(10): 169-1 186, 1998.
[51] Michael F Gensheimer and Balasubramanian Narasimhan. A scalable discrete-time survival model for neural networks. PeerJ, 7:e6257, 2019.
[52] Kan Ren, Jiarui Qin, Lei Zheng, Zhengyu Yang, Weinan Zhang, Lin Qiu, and Yong Yu. Deep recurrent survival analysis. In Proceedings of the 33rd AAAI Conference Artificial Intelligence (AAAI 2019), 2019.
[53] Hävard Kvamme and Ømulf Borgan. Continuous and discrete-time survival prediction with neural networks. arXiv preprint 4rXiv:1910.06724, 2019.
[54] Changhee Lee, William R. Zame, Jinsung Yoon, and Mihaela van der Schaar. Deephit: A deep learning approach to survival analysis with competing risks. In Proceedings of the 32th AAAI Conference Artificial Intelligence (AAAI 2018), 2018.
13

---

[55] Chun-Nam Yu, Russell Greiner, Hsiu-Chin Lin, and Vickie Baracos. Learning vatient-speciife cancer survival distributions as a sequence of dependent regressors. Advances in Neural Information Processing Systems, 24: 14:1885-1853 2011.
[56] Liangyuan Hu, Jiayi Ji, and Fan Li. Estimating heterogeneous survival treatment effect in observational data using machine learning. arXiv preprint arXiv:2008.07044, 2020.
[57] Iván Díaz, Oleksandr Savenkov, and Karla Ballman. Targeted learning ensembles for optimal individualized treatment rules with time-to- to-event outcomes. Biometrika, 105(3):723-738, 2018.
[58] Christos Louizos, Uri Shalit, Joris Mooij, David Sontag, Richard Zemel, and Max Welling. Causal effect inference with deep latent-variable models. In Proceedings of the 31st Conference on Neural Information Processing Systems (NeurIPS 2017), 2017.
14


Reviews:
[{"review_id": "GQfsMJC9BBw", "reviewer": "Reviewer_2JWN", "summary": "This paper tackles the problem of individual treatment effect prediction in survival analysis situations, which is the case the outcome variable is the time to event. In this situation, there are two additional bias(covariate shift) problems compared with the standard ITE prediction. They showed the generalization bound for this problem, then applied counter factual regression to minimize it through the IPM term.", "questions": "", "limitations": "", "rating": 7, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "I think the ITE prediction for the survival outcome is an interesting topic. Such a situation occurs not only in clinical studies but also in marketing, retail, and some other business situations.\n\nThe main contribution of this work is defining the generalization bound for this problem then minimize it through the CFR scheme. I think their approach technically sounds.\n\nIt was quite time-consuming for me to understand this paper. Mainly because of the lack of knowledge about the survival analysis. So I appreciate it if you could add some explanation for the problem this paper tries to solve without using the term for the survival analysis.\n\nFor minor issues\n\n- line 156 there are two dots before \"Here,\"\n- line 361 RMST -> RMSE?\n- Is there any reason to use ITE and HTE alternately?", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "They adequately addressed the limitation in their paper.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "6"}, "extra_text": {}}, {"review_id": "903awm2RU96", "reviewer": "Reviewer_ukxC", "summary": "In this paper, the authors present a framework for heterogeneous treatment effect estimation for time-to-event outcomes. In recent times, both areas have been explored independently and this paper proposes an approach for the censored time-to-event setting. SurvITE (individual treatment effect) and SurvHTE (heterogeneous treatment effect). Authors present experimental results for the same comparing against ablations and survival analysis baselines.", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "**originality** - the proposed problem setup brings together significant body of work in both areas (survival analysis and heterogeneous/individual treatment effect estimation) together. The proposed formulation in Eq (6) uses ideas similar to those proposed in 8,9,10 and 32 from the references.\n\n**quality and clarity** I do believe the draft appears to be technically sound and well presented in general. My only concern is with the way the IPM is applied in this paper. While authors do claim that this help obtaining shift-invariant representations, it is not directly clear for me on how penalizing IPM for each (a,$\\tau$) combination w.r.t the baseline is more better than doing it across the treatments explicitly. It might help to improve the writeup here, as I do think this is critical for the paper.  \n\nIn addition, I do think the experimental analysis section does appear incomplete.  I also don't completely agree with authors claim on line 44 that reference 27 is doing plain regression on targets. I think a more comprehensive discussion is needed w.r.t this paper in the experiments also if possible.\n\n**significance** - Making a contribution in this space of HTE for time-to-event survival outcomes is very much needed. \n\n**Post author feedback**\n\nI have gone through the author response and I have changed my score for this paper. I do believe clarifying these concerns will improve the paper quality.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Yes", "review_text": "", "extra_scores": {"ethical_concerns": "No", "needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "2.5 hours"}, "extra_text": {}}, {"review_id": "031_u33OiEE", "reviewer": "Reviewer_4Jiu", "summary": "In this work the authors quickly remind the reader of the pitfalls one encounter in the survival analysis framework, in order to motivate a new loss and corresponding deep model for modelling of the group-specific hazard functions.\nThe main contribution of this paper resides in Proposition 1: an upper bound of the population risk (unknown) in terms of quantities that can be estimated from the data. From this proposition the authors propose to minimize the upper-bound and then provide an example to motivate their intuition.", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "I find this paper interesting and fairly well written (given the constraints of the NeurIPS format). The bound in Proposition 1 is not only interesting from a purely theoretical point of view but also enables practical methods, of which the authors give an example. The bound is similar to previous bounds in the literature but does not rely on invertibility, which makes it much more practical as most models in use today are not invertible and when they are (e.g. Normalizing Flows), are computationally expensive.\n\nThe authors correctly point out that it's possible to solve the same problem by reweighting the ERM loss. It would have been worthwhile to study the link between their proposed method and propensity methods as it is known that reweighting and optimal transport plans (present here through the Wasserstein metric) are equivalent and since bounds on the excess risk in the reweighted ERM framework under censorship are known (see e.g. [1], or [2, 3]).\n\n[1] G. Ausset, S. Clémençon, and F. Portier, “Empirical Risk Minimization under Random Censorship: Theory and Practice,” 2019, Available: http://arxiv.org/abs/1906.01908\n[2] W. Stute, “The Central Limit Theorem Under Random Censorship,” Ann. Statist., vol. 23, no. 2, pp. 422–439, Apr. 1995, doi: 10.1214/aos/1176324528.\n[3] T. A. Gerds, J. Beyersmann, L. Starkopf, S. Frank, J. V. D. Laan, and M. Schumacher, “The Kaplan-Meier Integral in the Presence of Covariates: A Review,” From Statistics to Mathematical Finance, pp. 25–42, 2017, doi: 10.1007/978-3-319-50986-0.\n\n\nI also find the choice of competing methods in the experimental section to be puzzling and in need of better motivation: the authors exclude a wide range of methods (namely the \"deep\" Cox extensions like [4] or [5]) that are known to often be state of the art, for reasons I do not fully understand. This is especially puzzling since they later use the vanilla Cox model as a baseline.\n\n[4] J. Katzman, U. Shaham, J. Bates, A. Cloninger, T. Jiang, and Y. Kluger, “DeepSurv: Personalized Treatment Recommender System Using A Cox Proportional Hazards Deep Neural Network,” BMC Med Res Methodol, vol. 18, no. 1, p. 24, Dec. 2018, doi: 10.1186/s12874-018-0482-1.\n[5] C. Lee, W. Zame, J. Yoon, and M. V. D. Schaar, “DeepHit: A Deep Learning Approach to Survival Analysis With Competing Risks,” 2018.\n\n\nDespite those shortcomings I think, the paper is still overall a good fit for NeurIPS.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Yes, the authors are aware of the potential limitations of their work in terms of fairness. There is, however, not much they can do if they want to obtain bounds.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "4"}, "extra_text": {}}]

Author rebuttal:

Meta-review: The authors formalize heterogeneous treatment effect estimation from time-to-event data under empirical risk minimization. This is a natural extension of the prior work done for continuous and binary outcomes with the added challenges from to the time-to-event data such as censoring bias. The paper being a straightforward follow the recipe type of paper doesn't detract from its quality. The main thing I'd suggest would be complete the references in the related work for ML methods to include things like:

@inproceedings{avati2020countdown,
  title={Countdown regression: sharp and calibrated survival predictions},
  author={Avati, Anand and Duan, Tony and Zhou, Sharon and Jung, Kenneth and Shah, Nigam H and Ng, Andrew Y},
  booktitle={Uncertainty in Artificial Intelligence},
  pages={145--155},
  year={2020},
  organization={PMLR}
}

and

@inproceedings{ranganath2016deep,
  title={Deep survival analysis},
  author={Ranganath, Rajesh and Perotte, Adler and Elhadad, No{\'e}mie and Blei, David},
  booktitle={Machine Learning for Healthcare Conference},
  pages={101--114},
  year={2016},
  organization={PMLR}
}



Decision: Accept (Poster)