Learning Survival Distribution with Implicit Survival Function

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, $t_{x}$ denotes the true survival time, $x=(x_{1},x_{2},\dots,x_{p})^{T}$ denotes covariates of samples, $w=(w_{1},w_{2},\dots,w_{p})^{T}$ denotes parameters of the linear regression, and $h_{0}(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óyland (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.

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=\{t^{p}_{0},t^{p}_{1},\cdots t^{p}_{k-1}\}$. DeepHit Lee et al. (2018) uses a fully-connected network to directly predict occurrence probability $\hat{p}(t^{p}_{i}|x)$ defined as:

where $t^{p}_{i}$ is a time point in the discrete time space $t^{p}_{i}\in 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.

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\geq 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 $\left\{\beta_{k},\eta_{k}\right\}_{k=1}^{K}$ and their weights $\left\{\alpha_{k}\right\}_{k=1}^{K}$ are estimated through MLP. Thus, the final individual survival distribution $\hat{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.

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\in\mathbb{R}^{d}$ using a Multilayer Perceptron (MLP) denoted by encoder $E(\cdot)$:

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

where $H(\cdot)$ 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.

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 $t_{i}$ as:

Therefore, the estimated survival rate $\hat{S}(t_{i}|x)$ is defined as:

The estimated occurrence probability $\hat{p}(t|x)$ is approximated through:

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

For numerical stability, we manually set $\hat{S}(0|x)=1$ and $\hat{S}(t_{max}|x)=0$, where $t_{max}$ is ensured to be larger than any possible survival time in the dataset.

Analytical solutions for integration in Eq. 4.2 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,t_{max})$ is split into $K$ intervals $\{(t_{i}^{p},t_{i+1}^{p}]\}^{K-1}_{i=0}$ with time points $T=\{t_{i}^{p}\}^{K}_{i=0}$, where $t_{0}^{p}=0$ and $t_{k}^{p}=t_{max}$. In this paper, we set $t_{i+1}^{p}=t_{i}^{p}+\epsilon$ for convenience.

Let $g(t,x)$ denote $\ln(1-\hat{h}(t|x))$. Therefore, the integration in Eq. 4.2 for $t_{i}^{p}\in T$ is calculated using Simpson Formula as:

Thus, the event rate (CDF) is estimated as $\hat{W}(t_{i}^{p}|x)=1-\hat{S}(t_{i}^{p}|x)$.

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 $\hat{p}(t|x)$, though forms of loss functions for these two types of samples are the same.

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.

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.

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, 5-fold cross validation is applied following DeepHit Lee et al. (2018).

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 $t_{x_{i}}$ denotes the true survival time of $x_{i}$.

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

ISF is implemented with $PyTorch$. Number of hidden units of $E(\cdot)$ defined in Eq. 10 and $H(\cdot)$ defined in Eq. 4.1 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 $\epsilon$ will be discussed in the ablation study.

The reproduction of DeepHit and DRSA is based on the official code of DRSA¹¹1https://github.com/rk2900/drsa. And the reproduction of DSM refers to the official package $auto\_survival$²²2https://autonlab.github.io/auton-survival/models/dsm.

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 $t^{p}_{i+1}=t^{p}_{i}+1$, $\epsilon$ in Eq. 4.2 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.

For further understanding of ISF, we conduct experiments on ISF with variation of $\epsilon$ in Eq. 4.2 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\propto 1/\epsilon$.

In this paper, We use a hyperparameter $\epsilon$ 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 $\epsilon$ 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 $\epsilon$ leads to high computational complexity as discussed in Section 4.5, the generation ability of ISF enables models trained with relatively high $\epsilon$ to generate acceptable results of survival prediction.

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 $t^{p}_{i}$ instead of observed time $t^{o}_{x}$. If ISF maximizes $\hat{p}(t^{o}_{x}|x)$ or $\hat{S}(t^{o}_{x}|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^{2}K)$ sample-time pairs with numerical integration in $K$ intervals for optimization based on $t^{o}_{x}$. 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 $\epsilon$ 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 $\epsilon$. 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.

In real-world applications, right-censoring is most common in datasets, which indicates that the true survival time is larger than the observed time $t^{x}>t^{o}_{x}$. 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.