SurvITE: Learning Heterogeneous Treatment Effects from Time-to-Event Data

1 Introduction . The demand for methods evaluating the effect of 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 tree-based methods [ 1 , 2 , 3 , 4 , 5 ] , Gaussian processes [ 6 , 7 ] , and , in particular , neural networks ( NNs ) [ 8 , 9 , 10 , 11 , 12 , 13 , 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 to onset or progression of disease , the time to occurrence of an adverse event such as a stroke or heart attack , or the time until death of a 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-event 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 of censoring . This has led to the development of a 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 ( NeurIPS 2021 ) . indicator in a ( semi- ) parametric 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 ML 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 [ 27 ] , 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-specific 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 treatment-specific hazard functions is a challenging problem due to the potential presence of multiple sources of covariate shift : ( i ) non-randomized treatment assignment ( confounding ) , ( ii ) informative censoring and ( iii ) a form of shift we term event-induced 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-based 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-specific 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 , we propose a new model ( SurvITE ) relying on balanced representations that allows for estimation of treatment-specific target parameters ( hazard and survival functions ) in the survival context , as well as a sister model ( SurvIHE ) , which can be used for individualized hazard estimation in standard survival settings ( without treatments ) . We investigate performance across a 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 selfcontained introduction to and comparison with heterogeneous treatment effect estimation with standard ( binary/continuous ) outcomes . Problem setup . Assume we observe a time-to-event dataset D = { ( ai , xi , τ̃i , δi ) } ni=1 comprising realizations of the tuple ( A , X , T̃ , ∆ ) ∼ P for n patients . Here , X ∈ X and A ∈ { 0 , 1 } are random variables for a covariate vector describing patient characteristics and an indicator whether a binary treatment was administered at baseline , respectively . Let T ∈ 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 disease or even death , and censoring indicates loss of follow-up for a patient . Then , the observed time-to-event outcomes of each patient are described by T̃ = min ( T , C ) and ∆ = 1 ( T ≤ C ) , which indicate the time elapsed until either an event or censoring occurs and whether the event was observed or not , respectively . Throughout , we treat survival time as discrete2 and the time horizon as finite with pre-defined maximum tmax , so that the set of possible survival times is T = { 1 , · · · , tmax } . 2Where necessary , discretization can be performed by transforming continuous-valued times into a set of contiguous time intervals , i.e. , T = τ implies T ∈ [ tτ , tτ + δ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 to directly estimate conditional hazard functions using standard machine learning methods [ 20 ] . We define two counting processes NT ( t ) and NC ( t ) which track events and censoring , i.e . NT ( t ) = 1 ( T̃ ≤ t , ∆ = 1 ) and NC ( t ) = 1 ( T̃ ≤ t , ∆ = 0 ) for t ∈ T ; both are zero until either an event or censoring occurs . By convention , we let NT ( 0 ) = NC ( 0 ) = 0 . Further , let Y ( t ) = 1 ( NT ( t ) = 1 ∩NT ( t−1 ) = 0 ) be the indicator for an event occuring exactly at time t ; thus , for an individual with T̃ = τ and ∆ = 1 , Y ( t ) = 0 for all t 6= τ , and Y ( t ) = 1 at the event time t = τ . The conditional hazard is the probability that an event occurs at time τ given that it does not occur before time τ , hence it can be defined as [ 22 ] λ ( τ |a , x ) = P ( T̃ = τ , ∆ = 1|T̃ ≥ τ , A = a , X = x ) = P ( Y ( τ ) = 1|NT ( τ−1 ) =C ( τ−1 ) = 0 , A = a , X = x ) ( 1 ) It is easy to see from ( 1 ) that given data in long format , λ ( τ |a , x ) can be estimated for any τ by solving a standard classification problem with Y ( τ ) as target variable , considering only the samples at-risk at time τ in each treatment arm ( individuals for which neither event nor censoring has occurred until that time point ; i.e . the set I ( τ , a ) def= { i ∈ [ n ] : NT ( τ−1 ) i = NC ( τ−1 ) i = 0 ∩ Ai = a } ) . Finally , given the hazard , the associated survival function S ( τ |a , x ) = P ( T > τ |A = a , X = x ) can then be computed as S ( τ |a , x ) = ∏ t≤τ ( 1− λ ( t|a , x ) ) . The censoring hazard λC ( t|a , x ) and survival function SC ( t|a , x ) 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 a 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 λa ( τ |x ) = P ( T = τ |T ≥ τ , do ( A = a , C ≥ τ ) , X = x ) Sa ( τ |x ) = P ( T > τ |do ( A = a , C ≥ τ ) , X = x ) = ∏ t≤τ ( 1− λa ( t|x ) ) ( 2 ) Here , do ( · ) denotes [ 28 ] ’ s do-operator which indicates an intervention ; in our context , do ( A = a , C ≥ τ ) ensures that every individual is assigned treatment a and is 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 λa ( τ |x ) and Sa ( τ |x ) , possible HTEs of interest3 include the difference in treatment-specific survival times at time τ , i.e . HTEsurv ( τ |x ) = S1 ( τ |x ) − S0 ( τ |x ) , the difference in restricted mean survival time ( RMST ) up to time L , i.e . HTErmst ( x ) = ∑ tk≤L ( S1 ( tk|x ) −S0 ( tk|x ) ) · ( tk− tk−1 ) , and hazard ratios . In the following , we will focus on estimation of the treatment specific hazard functions { λa ( t|x ) } a∈ { 0,1 } , t∈T 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 to make a number of untestable assumptions on the underlying datagenerating process ( DGP ) [ 28 ] – this generally limits the ability to make causal claims to settings where sufficient domain knowledge is available . Here , as [ 20 , 21 ] , we assume the data was generated from the fairly general directed acyclic graph ( DAG ) presented in Fig . 1 . As there are no arrows originating in hidden nodes entering treatment or censoring nodes , this graph formalizes ( 1.a ) The ‘ No Hidden Confounders ’ Assumption in static treatment effect estimation and ( 1.b ) The ‘ Censoring At 3Note : 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 ) statistics considers population average parameters that are functions of quantities such as P ( T > τ |do ( A = a ) ) , which average over all X . Random ’ Assumption in survival analysis [ 20 ] . The latter is necessary here as estimating an effect of treatment on event time implicitly requires that censoring can be ‘ switched off ’ – i.e . 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 is defined as a function of its ancestors and exogenous noise [ 20 ] . Under these assumptions , λa ( τ |x ) = λ ( τ |a , x ) . ( 2 . Estimation ) To enable nonparametric estimation of λa ( τ |x ) 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 ’ [ 9 , 19 ] . In particular , for 0 < 1 , 2 , < 1 we need that ( 2.a ) 1 < P ( A = a|X = x ) < 1 − 1 , i.e . treatment assignment is non-deterministic , and that ( 2.b ) P ( NC ( t ) = 0|A = a , X = x ) > 2 for all t < τ , i.e . no individual will be deterministically censored before τ . Finally , because λa ( τ |x ) is a probability defined conditional on survival up to time τ , we need to assume that ( 2.c ) P ( NT ( τ−1 ) = 0|A = a , X = x ) > 3 > 0 for it to be well-defined . We formally state and discuss all assumptions in more detail in Appendix C .

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