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

Abstract:
We study the problem of inferring heterogeneous treatment effects from time-toevent 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 wellstudied 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.

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

heading: 2 Problem Definition
text: 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.

heading: 3 Challenges in Learning Treatment-Specific Hazard Functions using ERM
text: Preliminaries: ERM under Covariate Shift. Recall that in problems with covariate shift, the training distribution X,Y ∼ Q0(·) used for ERM and target distribution X,Y ∼ Q1(·) are mismatched: One assumes that the marginals do not match, i.e. Q0(X) 6= Q1(X), while the conditionals remain the same, i.e. Q0(Y |X) = Q1(Y |X) [29]. If the hypothesis classH used in ERM does not contain the truth (or in the presence of heavy regularization), this can lead to suboptimal hypothesis choice as arg minh∈H EX,Y∼Q1(·)[`(Y, h(X))] 6= arg minh∈H EX,Y∼Q0(·)[`(Y, h(X))] in general.

heading: 3.1 Sources of Covariate Shift in Learning Treatment-Specific Hazard Functions
text: We now consider how to learn a treatment-specific hazard function λa(τ |x) from observational data using ERM. As detailed in Section 2, we exploit the long data format by realizing that λa(τ |x) can be estimated by solving a standard classification problem with Y (τ) as dependent variable and X as covariates, using only the samples at risk with treatment status a, i.e. I(τ, a), which corresponds to solving the empirical analogue of the problem
λ̂a(τ |x) ∈ arg min ha,τ∈H EX,Y (τ)∼Pa,τ (·)[`(Y (τ), ha,τ (X)] (3)
where we use Pa,τ to refer to the observational (at-risk) distribution Pa,τ (X,Y (τ)) = λaT (τ |X)Pa,τ (X) with Pa,τ (X) = P(X|NT (τ−1) = NC(τ−1) = 0, A = a) = P(X|T̃ ≥ τ,A = a). If the loss function ` is chosen to be the log-loss, this corresponds to optimizing the likelihood of the hazard.
The observational (at-risk) covariate distribution Pa,τ (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.e. the marginal distribution X ∼ P(X) which we will refer to as P0(X) below to emphasize it being the baseline at-risk distribution4. Here, differences between P0(X) and the population at-risk Pa,τ (X) can arise due to three distinct sources of covariate shift:
• (Shift 1) Confounding/treatment selection bias: if treatment is not assigned completely at random, then P(X|A = a) 6= P0(X) and the distribution of characteristics across the treatment arms differs already at baseline, thus Pa,τ (X) 6= P0(X) in general.
• (Shift 2) Censoring bias: regardless of the presence of confounding, if the censoring hazard is not independent of covariates, i.e. λC(τ |a, x) 6= λC(τ |a), then the population at-risk changes over time such that Pa,τ1(X) 6= Pa,τ2(X) 6= P0(X) in general. If, in addition, there are differences between the treatment-specific censoring hazards, then the at-risk distribution will also differ across treatment arms at any given time-point, i.e. Pa,τ (X) 6= P1−a,τ (X) for τ > 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 λ(τ |a, x) 6= λ(τ |a) then Pa,τ1(X) 6= Pa,τ2(X) 6= P0(X) in general. Further, if there are heterogenous treatment effects, then Pa,τ (X) 6= P1−a,τ (X) for τ > 1 in general.
4With slight abuse of notation, we will use P0 and Pa,τ also to refer to densities of continuous x
What makes the survival treatment effect estimation problem unique? While Shift 1 arises also in the standard treatment effect estimation setting, Shift 2 and Shift 3 arise uniquely due to the nature of time-to-event data5. 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 a 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 1, 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. X1 = X3 in Fig. 1) then event-induced shift can amplify the selection effect over time (refer to Appendix E for a synthetic example of this).

heading: 3.2 Possible Remedies and Theoretical Analysis
text: A natural solution to tackle bias in ERM caused by covariate shift is to use importance weighting [30]; i.e. to reweight the empirical risk by the density ratio of target P0(X) and observed distribution Pa,τ (X). If we wanted to obtain a hazard estimator for (τ, a), optimized towards the marginal population, optimal importance weights are given by w∗a,τ (x) = P0(x) Pa,τ (x) = pτ,a ea(x)ra(x,τ) with pτ,a = P(T̃ ≥ τ,A = a), ea(x) = P(A = a|X = x) the propensity score, and ra(x, τ) = P(T̃ ≥ τ |A = a,X = x) the probability to be at risk, i.e. the probability that neither event nor censoring occurred before time τ . 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 target parameters λa(τ |x) through ra(x, τ). Further, especially for large τ , these weights might be very extreme even if known, which can lead to highly unstable results [31] – making biased yet stabilized weighting schemes, e.g. truncation, a good alternative. Therefore, we only assume access to some (possibly imperfect) weights wa,τ (x) s.t. EX∼Pa,τ [wa,τ (x)] = 1, so that we can create a weighted distribution Pwa,τ = wa,τ (x)Paτ (x). (Note: Paτ (x) can be recovered by using wa,τ (x) = 1.) Either instead of [8, 9] or in addition to weighting [10, 12, 14, 32], the literature on learning balanced representations for static treatment effect estimation has focused on finding a different remedy for distributional differences between treatment arms: creating representations Φ : X → R which have similar (weighted) distributions across arms as measured by an integral probability metric (IPM), motivated by generalization bounds. As we show below, we can exploit a similar feature in our context by finding a representation that minimizes the IPM term not between treatment arms, but between covariate distribution at baseline P0 and Pwa,τ . The proposition below bounds the target risk of a hazard estimator λ̂a(τ |x) = h(Φ(x)) relying on any representation. The proof, which relies on the concept of excess target information loss, proposed recently to analyze domain-adversarial training [33], and the standard IPM arguments made in e.g. [32], is stated in Appendix C.
Proposition 1. For fixed a, τ and representation Φ : X → R, let PΦ0 , PΦa,τ and Pw,Φa,τ denote the target, observational, and weighted observational distribution of the representation Φ. Define the pointwise losses
`h,Q(x; a, τ) def = EY (τ)|x,a∼Q[`(Y (τ), h(Φ(X)))|X = x,A = a]
`h,QΦ(φ; a, τ) def = EY (τ)|φ,a∼QΦ [`(Y (τ), h(Φ))|Φ = φ,A = a]
(4)
of (hazard) hypothesis h ≡ ha,τ : R → [0, 1] w.r.t. distributions in covariate and representation space, respectively. Assume there exists a constant CΦ > 0 s.t. CΦ−1`h,Pw,Φa,τ (φ, a, τ) ∈ G for some family of functions G. Then we have that
EX∼P0 [`h,P(X; a, τ)]︸ ︷︷ ︸ Target Risk ≤ EX∼Pa,τ [wa,τ (X)`h,P(X; a, τ)]︸ ︷︷ ︸ Weighted observational risk +CΦ IPMG(PΦ0 ,Pw,Φa,τ )︸ ︷︷ ︸ Distance in Φ-space + ηlΦ(h)︸ ︷︷ ︸ Info loss (5)
5Interestingly, 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 it is caused by censoring. If, however, our goal is estimation of the best target parameter (here: the hazard at a specific point in time τ ) 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.
where IPMG(P,Q) = supg∈G ∣∣∫ g(x)(P(x)−Q(x))dx∣∣ and we define the excess target informa-
tion loss η`Φ(h) analogously to [33] as η ` Φ(h) def = EX∼P[ξPΦ0 ,P(X) − ξPw,Φa,τ ,P(X)] with ξQΦ,Q(x) def = `h,QΦ(φ; a, τ)− `h,Q(x; a, τ). For invertible Φ, η`Φ(h) = ξQΦ,Q(x) = 0.
Unlike the bounds provided in [9, 10, 32, 14, 27], this bound does not rely on representations to be invertible; we consider this feature important as none of the works listed actually enforced invertibility in their proposed algorithms. Given bound (5), it is easy to see why non-invertibilty can be useful: for any (possibly non-invertible) representation for which it holds that Y (τ) |=X|Φ(X), A, it also holds that η`Φ(h) = ξPΦ,P(x) = ξPw,Φa,τ ,P(x) = 0 and the causally identifying restrictions continue to hold. A simple representation for which this property holds is a selection mechanism that chooses only the causal parents of Y (τ) from within X; if X can be partitioned into variables affecting the instantaneous risk (X3 in Fig. 1), and variables affecting only treatment assignment (X1 \X3) and/or censoring mechanism (X2 \X3), then the IPM term can be reduced by a representation which drops the latter two sets of variables – or irrelevant variables correlated with any such variables – without affecting η`Φ(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 η`Φ(h) and IPM term are zero.
Except for the dependence on η`Φ(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 treatmentand control group at baseline (or at each time step) we should find representations that balance PΦa,τ towards the baseline distribution PΦ0 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-induced differences in censoring and (iii) treatment-induced differences in events. It will, however, not allow handling the event- and censoring-induced shifts that occur regardless of treatment 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 a finding that could be of independent interest for the ML survival analysis literature.

heading: 3.3 From hazards to survival functions
text: If the ultimate goal is to use the hazard function to estimate survival functions as Ŝa(τ |x) =∏ t≤τ ( 1− λ̂a(t|x) ) , 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 to set A = a and C ≥ τ ), the target population (at each time step t) would be P(X|T ≥ t, do(A = a,C ≥ t)) – the population that preserves event-induced shift but removes selection- and censoring-induced 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 time steps (i.e. Ŝ(τ−1|x) = S(τ−1|x)) and focus only on estimating the next hazard, λa(τ |X), we can write
EX∼P0 [(Sa(τ |X)−Ŝa(τ |X))2] = EX∼P0 [(Sa(τ−1|X)(1−λa(τ |X))−Ŝa(τ−1|X)(1−λ̂a(τ |X)))2] = EX∼P0 [Sa(τ−1|X)2(λ̂a(τ |X)−λa(τ |X))2]
and notice that the MSE will implicitly down-weigh individuals with lower survival probability6.
Defining an 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 as P(X|T ≥ t, do(A = a,C ≥ t)), are never observed in practice and hence cannot be used to perform balancing
6Due to the square in the term Sa(τ − 1|X)2, this will be even more extreme than exact up-weighting of the population P(X|T ≥ τ, do(A = a,C ≥ τ)).
regularization of a 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 of the representation throughout. Intuitively, as outlined in the previous section, we expect that doing so will not over-correct for event-induced 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 bound7. Additionally, we expect that using the marginal population for balancing could be useful also for estimating the survival function even in the absence of selection- and censoring-induced shifts, as it may help to remove the effect of variables that appear spuriously correlated with outcome over time.

heading: 4 SurvITE: Estimating HTEs from Time-to-Event Data
text: Based on the theoretical analysis above, we propose a 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).8 The network architecture is illustrated in Figure 2. 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 a = 0. As we show in the experiments, this version of our method – SurvIHE (Individualized Hazard Estimator for Survival analysis) – is of independent interest in the standard survival setting, as it tackles Shifts 2 & 3. Below, we describe the empirical loss functions we use to find representation Φ and hypotheses ha,τ .
Let Φ : X → R denote the representation (parameterized by θφ) and ha,τ : R → [0, 1] the hazard estimator for treatment a and time τ (parameterized by θha,τ ), each implemented as a fullyconnected neural network. While the output heads are thus unique to each treatment-group 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, we use the following empirical loss functions for the observational risk and the IPM term:
Lrisk(θφ, θh) = 1
tmax tmax∑ t=1 ∑ i:τ̃i≥t n−11,tai` ( yi(t), h1,t(Φ(xi)) ) + n−10,t (1−ai)` ( yi(t), h0,t(Φ(xi)) ) ,
Lipm(θφ) = ∑
a∈{0,1} tmax∑ t=1 Wass ( {Φ(xi)}ni=1, {Φ(xi)}i:τ̃i≥t,ai=a ) ,
whereWass(·, ·) is the finite-sample Wasserstein distance [34]; refer to Appendix D for further detail. Note that Lipm(θφ), which penalizes the discrepancy between the baseline distribution and each at-risk distribution PΦa,τ , simultaneously tackles all three sources of shifts. Further, na,t = |I(τ, a)| is the number of samples at-risk in each treatment arm; its presence ensures that each a, τ -combination contributes equally to the loss. Overall, we can find Φ and ha,τ ’s that optimally trade off balance and
7In practice, we ensure this by weighting the contribution of the IPM term by a hyperparameter that is chosen to preserve predictive performance of the representation (see Appendix D).
8The source code for SurvITE is available in https://github.com/chl8856/survITE.
predictive power as suggested by the generalization bound (5) by minimizing the following loss:
Ltarget(θφ, θh) = Lrisk(θφ, θh) + βLipm(θφ) (6) where θh = {θha,τ }a∈{0,1},τ∈T , and β > 0 is a hyper-parameter. The pseudo-code of SurvITE, the details of how to obtain Wass(·, ·) and how we set β can be found in Appendix D.

heading: 5 Related work
text: Heterogeneous treatment effect estimation (non-survival) has been studied in great detail in the recent ML literature. While early work built mainly on tree-based methods [1, 2, 3, 4], many other methods, such as Gaussian processes [6, 7] and GANS [35], have been adapted to estimate HTEs. Arguably the largest stream of work [8, 9, 10, 11, 12, 13, 14, 15] built on NNs, due to 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, 15]. 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 [36], can be directly applied to the survival setting by relying on a 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. [39, 40, 41]) or modeling (log-)time as a regression problem (e.g. [42, 43, 44, 45, 46, 47]). One popular nonparametric estimator of survival functions is [48]’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 datastructures goes back to at least [49], with implementations using NN-based methods presented in e.g. [50, 51, 52, 53]. [54] models the probability mass function instead of the hazard, and [55] use labels 1{T > t}t∈T 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. [23, 25] use tree-based nearest-neighbor estimates to estimate expected differences in survival time directly, and [24] use a BART-based S-learner to output expected differences in log-survival time. [56] performed a simulation study using different survival prediction models as base-learners for a twomodel approach to estimating the difference in median survival time. Based on ideas from the semi-parametric efficiency literature, [26] and [57] propose estimators that target the (restricted) mean survival time directly and consequently do not output estimates of the treatment-specific hazard or survival functions. We consider the ability to output treatment-specific predictions an important feature of a model if the goal is to use model output to give decision support, given that it allows the decision-maker to trade-off relative improvement with the baseline risk of a patient. Finally, [27] recently proposed a generative model for treatment-specific event times which relies on balancing representations to balance only the treatment groups at baseline. This model does not output hazardor survival functions, but can provide approximations by performing Monte-Carlo sampling.

heading: 6 Experiments
text: 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 a 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 1 & 2.
Baselines. Throughout, we use Cox regression (Cox), a model using a separate logistic regression to solve the hazard classification problem at each time-step (LR-sep), random survival forest (RSF), and a deep learning-based time-to-event method [54] (DeepHit) as natural baselines; when there
are treatments, we use them in a 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-carlo sampling. Finally, we consider ablations of SurvITE (and SurvIHE); in addition to removing the IPM term (SurvITE (no IPM)), we consider two variants of SurvITE based on [9]’s CFRNet balancing term: SurvITE (CFR-1) creates a representation balancing treatment groups at baseline only, and SurvITE (CFR-2) creates a representation optimizing for balance of treatment groups at each time step (i.e. no balancing towards P0). We discuss implementation in Appendix D.
Synthetic Experiments. We consider a 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 use
λa(t|x) = {
0.1σ(−5x21 − a · (1{x3 ≥ 0}+ 0.5)) for t ≤ 10 0.1σ(10x2 − a · (1{x3 ≥ 0}+ 0.5))) for t > 10 , λC(t|x) = 0.01σ(10x24)
with treatment assignment mechanism a ∼ Bern(ξ · σ( ∑ p∈P xp)), with σ the sigmoid function. Additionally, we assume administrative censoring at t = 30 throughout, i.e., λC(30|x) = 1, marking e.g. the end of a hypothetical clinical study. Covariates are generated from a 10-dimensional multivariate normal distribution with correlations, i.e. X ∼ N (0,Σ) where Σ = (1− ρ)I + ρ11> with ρ = 0.2. We use 5000 independently generated samples each for training and testing.
In S1, we begin with the simplest case – no treatments and no censoring – using only λ0(t|x) to generate events, considering only event-induced shift (Shift 3). In S2, we introduce informative censoring using λC(t|x) (Shift 2+3). In S3, we 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, we vary treatment selection by changing (i) whether the covariate set overlaps with the event-inducing covariates (P={1, 2}) or not (P={9, 10}) and (ii) the selection strength ξ ∈ {1, 3}. We present exploratory plots of these DGPs in Appendix E.
Fig. 3 (top) shows performance on estimating S0(t|x) = ∏ k≤t ( 1− λ0(k|x) ) for all scenarios and methods, while Fig. 3 (bottom) shows performance on estimating the difference in survival functions (HTEsurv(t|x)) for a selection of methods (for readability, full results in Appendix F). In Table 1, we further evaluate the estimation of differences in RMST (HTErmst(x)). 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-inducing 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 as it encourages dropping other covariates (which appear imbalanced due to correlations in X). While our method was motivated by theory for estimation of hazard functions, it thus indeed also leads to gains in survival function estimation. As expected, both Cox and LR-sep
do not perform well as they are misspecified, while the nonparametric 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-specific survival functions (note that [54] 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-1), or over time – SurvITE(CFR-2)) is not as effective as using the baseline population as a target, especially at the later time steps where the effects of time-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-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 0 and 1 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 11400 pairs of twins, which is used in [58, 35] 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 covariate-based sampling mechanism to select only one twin for training to emulate selection bias. Further, we consider a second setting where we additionally introduce covariate-dependent 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. A 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, HTErmst(x) is suited best to evaluate performance on estimating effect heterogeneity. The results presented in Table 1 largely confirm our findings on relative performance in the synthetic experiments; only RSF performs relatively worse on this dataset.

heading: 7 Conclusion
text: We studied the problem of inferring heterogeneous treatment effects from time-to-event data by focusing on the challenges inherent to treatment-specific hazard estimation. We found that a variety of covariate shifts play a 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 a domain expert prior to deployment in practice. Here, the time-to-event nature of our problem adds an additional assumption (‘random censoring’) 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.
Acknowledgments and Disclosure of Funding
We thank anonymous reviewers as well as 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).

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

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.

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.