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