Learning Survival Distribution with Implicit Survival Function

1 Introduction

Survival analysis is a typical statistical task for tracking occurrence of the event of interest through modeling relationship between covariates and event occurrence. In some medical situations [Courtiol et al., 2019; Zadeh Shirazi et al., 2020], researchers model the death probability of some diseases using survival analysis to explore effects of prognostic factors. However, some samples lose tracking (censored) during observation. For example, some patients are still alive at the end of observation, whose survival times are unavailable. Such censored samples are valuable for analysis of favorable prognosis. Therefore, censorship is one key problem in survival analysis as well as survival distribution modeling.

The most widely-used survival analysis model Cox proportional hazard method [Cox, 1992] predicts a hazard rate, which assumes that the relationship between covariates and hazard is time-invariant. For optimization, Cox model and its extensions [Tibshirani, 1997; Li et al., 2016; Katzman et al., 2018; Zhu et al., 2016] maximize the ranking accuracy of comparable pairs including comparison between uncensored samples and censored samples.

∗Corresponding author: Weimin Tan and Bo Yan. This work is supported by NSFC (Grant No.: U2001209, 61902076) and Natural Science Foundation of Shanghai (21ZR1406600). Our code is available at https://github.com/Bcai0797/ISF.

Lately, some works introduce deep neural networks to survival analysis. DeepSurv [Katzman et al., 2018] and DeepConvSurv [Zhu et al., 2016] simply replace the linear regression in the Cox model with neural networks for non-linear representations. These methods maintain the strong assumption of hazards' time-invariance in Cox model, leading to weak generalization of networks in real-world applications.

To avoid strong assumption on survival distribution, researchers try to estimate a distribution in a discrete time space instead of predicting a time-invariant risk. DeepHit [Lee et al., 2018] is proposed to learn occurrence probabilities at preset time points directly without assumptions about underlying stochastic process. Deep Recurrent Survival Analysis

(DRSA) [Ren et al., 2019] builds a recurrent network to capture the sequential patterns of the feature over time in survival analysis. Therefore, both DeepHit and DRSA learn a discrete survival distribution. Compared to the cross-entropy loss, the log-likelihood loss obtains better prediction for DeepHit and DRSA [Zadeh and Schmid, 2021]. On the basis of predicted occurrence probabilities in the discrete time space, the loglikelihood is naturally estimated in DeepHit and DRSA for both censored and uncensored samples.

Differing from discrete distribution estimation in DeepHit and DRSA, DSM [Nagpal et al., 2021] estimates the average mixture of parametric distributions. In implementation, DSM employs Weibull and Log-Normal distributions for analytical solutions of the cumulative distribution functions (CDF) and support limited in the space of positive reals. Therefore, DSM includes censored samples during optimization through CDF estimation. However, DSM also introduces assumptions on survival distribution through parametric distribution selection.

In this paper, we propose Implicit Survival Function (ISF) based on Implicit Neural Representation which is widelyused in 2D and 3D image representation [Mildenhall et al., 2020; Chen et al., 2020]. As shown in Figure 1(a), ISF estimates a conditional hazard rate with the given sample and time. To capture time patterns, we embed the input time through Positional Encoding [Vaswani et al., 2017]. The aggregated vector of encoded sample feature and time embedding is fed to a the regression module for conditional hazard rate estimation without strong assumptions on survival distribution. As shown in Figure 1(b), we employ numerical integration with predicted conditional hazard rates for survival distribution prediction.

For optimization, we maximize likelihood of both censored and uncensored samples on the basis of approximated CDF of survival in a discrete time space. And experimental results prove that ISF is robust to the hyperparameter setting of the discrete time space.

To summarize, the contributions of this paper can be listed as:

• The proposed Implicit Survival Function (ISF) directly models the conditional hazard rate without strong assumptions on survival distribution, and captures the effect of time through Positional Encoding.
• To estimate survival distribution with ISF, numerical integration is used to approximate the cumulative distribution function (CDF). Therefore, ISF can handle censorship common in survival analysis through maximum likelihood estimation based on approximated CDF.
• Though survival distribution estimation of ISF is based on a discrete time space, ISF has capability to represent a continuous survival distribution through Implicit Neural Representation. And experimental results show that ISF is robust to the setting of the discrete time space.
• To demonstrate performance of the proposed model compared with the state-of-the-art methods, experiments are built on several real-world datasets. Experimental results show that ISF outperforms the state-of-the-art methods.

2 Formulation

Survival analysis models aim at modeling the probabilistic density function (PDF) of tracked event defined as:

where t denotes time, and tx denotes the true survival time.

Thus, the survival rate that the tracked event occurs after time ti is defined as:

Similarly, the event rate function of time ti is defined as the cumulative distribution function (CDF):

The conditional hazard rate h(t|x) is defined as:

3 RelatedWork

In this section, we describe several related approaches. The previous methods are divided into three parts based on their target of estimation: proportional hazard rate, discrete survival distribution and distribution mixture.

3.1 Proportional Hazard Rate

The Cox proportional hazard method proposed in [Cox, 1992] is a widely-used method in survival analysis tasks. Cox model assumes that the hazard rate of occurrence of a certain event is constant with time and the log of hazard rate can be represented by a linear function. Thus, the basic form of Cox model is:

where t denotes time, tx denotes the true survival time, x = (x1, x2, . . . , xp)T denotes covariates of samples, w = (w1,w2, . . . ,wp)T denotes parameters of the linear regression, and h0(t) denotes a fixed time-dependent baseline hazard function. Parameters w can be estimated by minimizing the negative log partial likelihood.

However, the time-invariance assumption of hazard in Cox model weakens its generalization. Other methods make different assumptions about the survival function such as Exponential distribution [Lee and Wang, 2003], Weibull distribution [Ranganath et al., 2016], Wiener process [Doksum and H´oyland, 1992] and Markov Chain [Longini et al., 1989]. These methods with strong assumptions about the underlying stochastic processes fix the form of survival functions, which suffers from generalization problem in real-world situations.

The outstanding capability of deep learning in non-linear regression achieve researchers' high attention. Therefore, many approaches introduce deep learning to survival analysis. DeepSurv [Katzman et al., 2018] replaces the linear regression of Cox model with a deep neural network for non-linear representation, but maintains the basic assumption of Cox model. Some works [Zhu et al., 2016; Li et al., 2019] extend DeepSurv with a deep convolutional neural network for unstructured data such as images.

3.2 Discrete Probability Distribution

To avoid strong assumptions about the survival time distribution, previous methods model the survival analysis problem in a discrete space with K time points T = {tp0 , tp1 , · · · tpk-1 }. DeepHit [Lee et al., 2018] uses a fully-connected network to directly predict occurrence probability ˆ p(tpi |x) defined as:

where tpi is a time point in the discrete time space tpi ∈ T.

DRSA [Ren et al., 2019] employs standard LSTM units [Hochreiter and Schmidhuber, 1997] to capture sequential patterns of features over time, and predicts a conditional hazard rate defined as:

Hence, DRSA defines occurrence probability of event as:

Although both DeepHit and DRSA predicts directly predict survival distribution without strong assumption, they only estimate probabilities at discrete time points.

3.3 DistributionMixture

Discrete probability distribution estimation methods only estimate a fixed number of probabilities, which limits their applications. To generate a continuous probability distribution, DSM [Nagpal et al., 2021] learns a mixture of K well-defined parametric distributions. Assuming that all survival times follows t ≥ 0, DSM selects distributions which only have support in the space of positive reals. And for gradient based optimization, CDF of selected distributions require analytical solutions. In implementation, DSM employs Weibull and Log-Normal distributions, namely primitive distributions.

During inference, parameters of K primitive distributions {βk, ηk}Kk=1 and their weights {αk}Kk=1 are estimated through MLP. Thus, the final individual survival distribution ˆ p(t|x) is defined as the weighted average of K primitive distributions:

However, DSM introduces assumptions of survival distributions since primitive distribution selection is taken as a hyperparameter.

4 Methodology

To model the survival distribution, we propose Implicit Survival Function (ISF) to estimate conditional hazard rate with positional encoding of time. In this section, we will demonstrate details of ISF as illustrated in Figure 2.

4.1 Implicit Survival Function

The proposed ISF aims at predicting h(t|x) defined in Eq. 4. For a given sample x, ISF first generates a feature vector z ∈ Rd using a Multilayer Perceptron (MLP) denoted by encoder E(·):

To capture the effect of time, Positional Encoding (PE) of time t is added to the feature vector z. Then, our hazard rate regression ˆh(t|x) is defined as:

where H(·) is implemented with a MLP.

Positional Encoding maps time t to a embedding of d dimensions using pre-defined sinusoidal functions [Vaswani et al., 2017]:

The sinusoidal function based Positional Encoding provides shift-invariant representations, and let MLP learn high frequency functions [Tancik et al., 2020]. Therefore, ISF employs Positional Encoding defined in Eq.12 for embedding of time in survival analysis.

4.2 Survival Distribution Estimation

For survival distribution estimation with ISF, we first estimate survival rate S(t|x) defined in Eq. 2, and then approximate occurrence probability p(t|x) defined in Eq. 1 through difference of survival rate.

From Eqs. 2 and 4, we can derive the log survival rate at time ti as:

Therefore, the estimated survival rate ˆ S(ti|x) is defined as:

The estimated occurrence probability ˆ p(t|x) is approximated through:ˆ

where ϵ is a hyperparameter. The setting of ϵ depends on the precision of annotations in the dataset. Corresponding discussion is included in Section 5.5

For numerical stability, we manually set ˆ S(0|x) = 1 and ˆ S(tmax|x) = 0, where tmax is ensured to be larger than any possible survival time in the dataset.

4.3 Numerical Integration

Analytical solutions for integration in Eq. 14 is unavailable for ISF. To overcome such problem, we use numerical integration to approximate CDF in a discrete time space.

The duration of survival time [0, tmax) is split into K intervals {(tpi , tpi+1 ]}K-1 i=0 with time points T = {tpi }Ki=0 , where tp0 = 0 and tpk = tmax. In this paper, we set tpi+1 = tpi + ϵ for convenience.

Let g(t, x) denote ln(1-ˆh(t|x)). Therefore, the integration in Eq. 14 for tpi ∈ T is calculated using Simpson Formula as:

Thus, the event rate (CDF) is estimated as ˆW(tpi |x) = 1ˆ S(tpi |x).

4.4 Loss Function

Like existing approaches [Lee et al., 2018; Ren et al., 2019; Nagpal et al., 2021], we construct loss functions on the basis of maximum likelihood estimation. Although ISF provides a conditional hazard rate in the continuous time space, the optimization is performed in the discrete time space for CDF approximation. In this section, for easily understanding, we describe the proposed loss function separately for censored and uncensored samples in the view of predicting ˆ p(t|x), though forms of loss functions for these two types of samples are the same.

Censored Samples

For a censored sample, the true survival time tx is unknown but the latest observation time tox is available, which indicates tx > tox . Thus, the loss function is expected to maximize ˆ S(tox |x). For simplification, we maximize ˆ S(tpi |x) where tox ∈ (tpi , tpi+1 ].

Therefore, the loss function for censored samples is defined as:

where the latest observation time tox ∈ (tpi , tpi+1 ].

Uncensored Samples

Given an uncensored sample (x, tox ), the observation time tox is equal to the true survival time tx. Thus, we maximize ˆ p(tpi |x) where the true survival time tox ∈ (tpi-1 , tpi ]:

Unified Loss

According to Lcs in Eq. 17 and Lucs in Eq. 18, loss for both uncensored and censored samples can be represented as sum of ˆ p(tpi |x) in the discrete time space. For unification, we first define an indicator vector Y x ∈ RK in the discrete time space including K + 1 time points as:

Thus, the proposed loss function can be unified as:

The unified loss function L(·) handles both censored and uncensored samples. We use indicator vector Y x to control likelihood calculation. Hence, the proposed loss function is suitable for any type of censorship.

4.5 Computational Complexity

As discussed in Sections 4.2, 4.3 and 4.4, estimation and optimization of ISF is performed in a discrete time space with K time intervals. For N samples, ISF predicts O(NK) occurrence probabilities for survival distribution estimation. However, such process can be accelerated in the parallel computation situation because of independent positional encoding of time points.

4.6 Difference fromExistingMethods

In this section, we compare the proposed model ISF with deep-learning models DeepHit, DRSA and DSM whose survival distribution estimation is close to that of ISF. We illustrate brief frameworks of these models and ISF in Figure 3.

ISF vs DeepHit

As shown in Figure 3(a), DeepHit directly regresses occurrence probabilities at preset time points through MLP. Therefore, the number of parameters dependents on the number of time points in the discrete time space.

Since ISF takes positional encoding of time as input, the number of parameters in ISF is independent to the amount of time points. Therefore, ISF has better expansibility for time space variation.

ISF vs DRSA

According to Eqs. 7 and 11, the goal of both ISF and DRSA is conditional hazard rate estimation. With estimated hazard rate, occurrence probability can be easily derived as shown in Eqs. 8 and 15.

The main difference between ISF and DRSA is the method of capturing time effect. As shown in Figure 3(b), DRSA applies RNN to learn sequential patterns in a discrete time space and serially processes preset time points, while ISF uses positional encoding to exploit time information in the real field through parallel computation.

ISF vs DSM

DSM models continuous survival distribution with mixture of parametric distributions as shown in Figure 3(c). Instead of explicit distribution representation in Eq. 9, ISF learns a function H(·) taking time as input defined in Eq. 11 to directly estimate conditional hazard rate. Therefore, the implicit representation of survival distribution in ISF avoids strong assumptions on survival distribution.

With decrease of ϵ in Eq. 15, precision of occurrence probability approximation increase, and thus ISF can be regarded as approximation of a continuous survival distribution. Distribution mixture in DSM directly models a continuous survival distribution, but distribution selection is a hyperparameter with strong assumptions about the stochastic process.

5 Experiments

In this section, we compare the proposed method ISF with the state-of-the-art deep-learning survival distribution estimation methods including DeepHit, DRSA and DSM. DeepHit predicts the occurrence probability ˆ p(t|x) directly with a fullyconnected neural network [Lee et al., 2018]. DRSA estimates a conditional hazard rate ˆh(t|x) with LSTM units to capture sequential patterns [Ren et al., 2019]. Both DeepHit and DRSA perform survival analysis in the discrete time space, while DSM estimates a continuous survival distribution through the mixture of parametric distributions [Nagpal et al., 2021]. Besides, we also compare ISF with Cox [Cox, 1992], its deep-learning extension DeepSurv [Katzman et al., 2018] and random forest based survival analysis method RSF [Ishwaran et al., 2008].

5.1 Datasets

To demonstrate the performance of the proposed method, experiments are conducted on several public real-world dataset:

• CLINIC tracks patients' clinic status [Knaus et al., 1995]. The tracked event is the biological death. Survival analysis in CLINIC is to estimate death probability with physiologic variables.
• MUSIC is a user lifetime analysis containing about 1000 users with entire listening history [Jing and Smola, 2017]. The tracked event is the user visit to the music service. The goal of survival analysis is to predict the time elapsed from the last visit of one user to the next visit.
• METABRIC dataset contains gene expression profiles and clinical features of the breast cancer from 1,981 patients [Curtis et al., 2012]. Following the experimental setting of DeepHit, 21 clinical features are used during evaluation [Lee et al., 2018].

The statistics of three datasets is shown in Table 1. The training and testing split of CLINIC and MUSIC follows the setting of DRSA [Ren et al., 2019]. For METABRIC, 5fold cross validation is applied following DeepHit [Lee et al., 2018].

5.2 Metric

Concordance Index (C-index, CI) is a widely-used evaluation metric in survival analysis for measuring the probability of accurate pair-wise order of comparable samples' event time. However, the ordinary CI [Harrell et al., 1982] for proportional hazard models assumes the predicted value is time-invariant [Cox, 1992; Tibshirani, 1997; Katzman et al., 2018], while distribution estimation based methods predict a time-dependent distribution of survival. Thus, following DeepHit and DSM, we perform time-dependent concordance index [Antolini et al., 2005], which is defined as:

where txi denotes the true survival time of xi.

5.3 Implementation Details

For fair comparison, the discrete time space in experiments is set as {(0, 1], (1, 2], . . . , (K - 1,K]} following setting of DeepHit and DRSA. According to the maximum time shown in Table 1, tmax is set as 400, and K = tmax.

ISF is implemented with PyTorch. Number of hidden units of E(·) defined in Eq. 10 and H(·) defined in Eq. 11 are corresponding set as {256, 512, 256} and {256, 256, 1} for all experiments.

During training, we perform Adam optimizer. Models of the best CI is selected with variation in hyperparameters of learning rate {10-3 , 10-4 , 10-5 }, weight of decay {10-3 , 10-4 , 10-5 } and batch size {8, 16, 32, 64, 128, 256}. The influence of ϵ will be discussed in the ablation study.

The reproduction of DeepHit and DRSA is based on the official code of DRSA 1 . And the reproduction of DSM refers to the official package auto survival2 .

5.4 Performance Comparison

To evaluate performance of ISF, we conduct experiments in three public datasets CLINIC, MUSIC and METABRIC compared with several existing methods. Since compared discrete time space methods DeepHit and DRSA set time points as tpi+1 = tpi + 1, ϵ in Eq. 15 which controls precision of ISF is set as 1 during training and evaluation for fair comparison.

As shown in Table 2, ISF achieve the best CI in three datasets which censoring rates are 0.132, 0.351 and 0.552. Therefore, ISF is robust to censoring rate. Besides, the large number of samples in MUSIC dataset contributes to performance improvement of ISF, while ISF has relatively low improvement in METABRIC containing fewer samples.

5.5 Ablation Study

For further understanding of ISF, we conduct experiments on ISF with variation of ϵ in Eq. 15 which controls precision to study the effect of precision. As discussed in Section 4.5, ISF predicts O(NK) occurrence probabilities for N samples with K time intervals where K ∝ 1/ϵ.

1 https://github.com/rk2900/drsa

2 https://autonlab.github.io/auton-survival/models/dsm

Table 3: CI performance comparison with variation of ϵ during training in CLINIC. During inference, ϵ of all models is fixed to 1 for fair comparison and accurate evaluation.

Training Precision

Since survival time annotations in CLINIC are saved as integer, the ideal ϵ for CLINIC is ϵ = 1. Therefore, we evaluate CI of ISF on CLINIC with variation of ϵ during training in this section. For fair comparison and accurate evaluation, ϵ in inference in this section is fixed to ϵInference = 1.

As defined in Eq. 15, ϵ determines precision of ISF. In CLINIC dataset, estimation precision of ISF is higher than annotation precision when ϵTrain < 1 during training. On the contrary, if ϵTrain > 1, annotation precision is higher than estimation precision. In such case, ISF predicts occurrence probabilities at unseen time points.

In Table 3, results of ϵTrain from 0.1 to 10. For ϵTrain ∈ [0.1, 1), ISF achieves close CI since estimation precision of these models is higher than annotation precision. For ϵTrain ∈ {2, 5}, the performance is also close to that of ISF with ϵTrain = 1, which indicates that ISF is capable of extrapolating in a certain range of time and robust to ϵTrain variation. In the extreme case of ϵTrain = 10, CI of ISF

significantly decreases since the maximum survival time in CLINIC is 82.

Inference Precision

In this section, we study generalization ability of ISF with variation of ϵInference during evaluation. Based on ISF trained with ϵTrain = 1, we adjust ϵInference from 0.1 to 1 during inference, and evaluate corresponding CI performance in three public datasets. In ϵInference < 1 experiments, ISF predicts conditional hazard rates at time points unseen in training. Hence, results of CI demonstrate generalization ability of ISF.

As shown in Table 4, ISF performance has little decrease when ϵInference < ϵTrain . Hence, ISF has high generalization for occurrence probability prediction at time points beyond the preset discrete time space, which proves that ISF manages to capture patterns of time through representations from sinusoidal positional encoding.

6 Discussion

In this section, we discuss some features of ISF in details.

6.1 Estimation Precision

In this paper, We use a hyperparameter ϵ to control the sampling density of the discrete time space, which has impact on the estimation precision of ISF. Experimental results of the ablation study in Section 5.5 show that ISF with varied ϵ achieves close CI performance in a certain range, even if the estimation precision is lower than annotation precision.

ISF captures time patterns through positional encoding as defined in Eq. 12. Representation based on sinusoids is shift-variation and enables MLP learn high frequency functions [Tancik et al., 2020]. Therefore, ISF manages to extrapolate occurrence probabilities unseen during training.

Although low ϵ leads to high computational complexity as discussed in Section 4.5, the generation ability of ISF enables models trained with relatively high ϵ to generate acceptable results of survival prediction.

6.2 Discrete Time Space

ISF estimates conditional hazard rates in a discrete uniform time space for optimization and inference. For N samples with K time intervals, ISF processes O(NK) pairs of sample and time during training and inference. In this section, we discuss the necessity of uniform time sampling.

In Section 4.4, we maximize occurrence probabilities at time points tpi instead of observed time tox . If ISF maximizes

ˆ p(tox |x) or ˆ S(tox |x) during optimization, the number and distribution of processed sample-time pairs depends on the training set. In the extreme case that the training set contains N samples with highly discrete survival time, ISF processes O(N 2K) sample-time pairs with numerical integration in K intervals for optimization based on tox . And the distribution of these sample-time pairs relies on the distribution of observed time, which perhaps introduces prior of the survival time distribution in the training set. Though ISF based on the discrete time space replaces the observed time with preset time points, the optimization process is based on adjustable uniform sampling of time. And the adjustment of the discrete time space is independent to the model architecture of ISF.

The ablation study of ϵ also proves that the preset discrete uniform time space based optimization and inference provides enough accuracy for survival analysis. Moreover, the estimation precision of ISF can be easily changed without model architecture modification through variation of hyperparameter ϵ. Hence, occurrence probabilities prediction in a discrete time space through ISF like previous works [Lee et al., 2018; Ren et al., 2019] is reasonable and robust.

6.3 Unified Loss Function

In real-world applications, right-censoring is most common in datasets, which indicates that the true survival time is larger than the observed time tx > tox . Therefore, existing discrete or continuous distribution prediction methods only considers right-censoring in loss functions [Lee et al., 2018; Ren et al., 2019; Nagpal et al., 2021].

Instead of establishing two distinct loss functions for censored and uncensored samples, the proposed loss function uses indicator vector Y defined in Eq. 19 for likelihood calculation. Therefore, a unified loss function defined in Eq. 20 is proposed for both censored and uncensored samples and is easy to be extended for any type of censoring.

7 Conclusion

In this paper, we propose Implicit Survival Function (ISF) for conditional hazard rate estimation in survival analysis. ISF employs sinusoidal positional encoding to capture time patterns. Two MLP are used to encode input covariates and regress conditional hazard rates. For survival distribution estimation, ISF performs numerical integration to approximate CDF for survival rate prediction.

Compared with existing methods, ISF estimates survival distribution without strong assumptions about survival distribution and models a continuous distribution through Implicit Neural Representation. Therefore, ISF models based on different settings of the discrete time space share a common architecture of the network. Moreover, ISF has robustness to estimation precision controlled by the discrete time space whether the estimation precision is higher than the annotation precision or not. Experimental results show that ISF outperforms the state-of-the-art survival analysis models on Concordance Index performance in three public datasets with varied censoring rates.