Adaptive Transformer Modelling of Density Function for Nonparametric Survival Analysis

We integrate the variation of last-observation-carried-forward (LOCF) method to handle missing data. It duplicates the value of the last observation to replace the following missing values until $\Delta_{\tau}=3$, and the ones that are still missing after LOCF are imputed by mean/mode of all previous points for continuous/binary covariates until $T_{max}$, making sure no missing data for all time points $T_{j}$ [18]. Next, we extract latent representations of time-invariant features $\boldsymbol{x}_{n}$ and time-varying features $\boldsymbol{x}_{v}$ by static (-s) and tabular-data-based dynamic (-d1) extraction modules separately. These modules are constructed using MLPs to deal with numerical tabular data. The representation of $\boldsymbol{x}_{n}$ is replicated $T_{max}+1$ times by encompassing $T_{0}$ as well. For static modelling, these are subsequently transmitted to encoder. For dynamic modelling, these representations are concatenated with the representation of $\boldsymbol{x}_{v}$ at their respective time points before the transmission. Meanwhile, a convolutional neural network (CNN) variation (-d2) of the image-based dynamic extraction module can be used to address the image-like input formats. As shown in Fig. 1(a), the extracted feature can be shared across a predefined time window $T_{w}$ in light of the data sparsity.

The core of the encoder module is a Transformer, which treats each patient as a ‘sentence’ and the embedded features as ‘words’ of the sentence. For an input sample, the number of words correspond to the duration $t\in\{0,1,...,T_{max}\}$, where we predefined $T_{0}=0$ and $T_{max}$ is a hyper-parameter selected based on the longest temporal data of a dataset. The previous concatenated representations pass through a MLP followed by layer normalization to get the embedded features. Following the conventional approach of a Transformer, we utilize the sine-cosine positional embedding [34] as temporal embedding in this work, and add it onto the set of embedded features, whose length is set as the embedding dimension $d_{m}$ of the Transformer. The Transformer encoder then processes embedded features and produces $T_{max}+1$ outputs, each with shape $1\times d_{m}$. It is worth noting that the self-attention layers in the encoder is modified to prevent positions from attending to subsequent positions. Specifically, it prohibits each position from attending to subsequent positions, and the attention scores for all illegal connections are masked out by assigning them with $-\infty$ [34]. Next, all-time point outputs are fed into an exclusive 2-layer MLP. The first layer is followed by rectified linear unit (ReLU) and layer normalization, with shape of $d_{m}\times\frac{d_{m}}{2}$. The second layer, with shape of $d_{m}\times 1$, is followed by a softmax layer to produce the individual estimated PDF. Further, the estimated survival function $\hat{S}(t\mid(\boldsymbol{x}_{n},\boldsymbol{x}_{v}))$ can be calculated based on Eq. 2, which ensures its monotonicity is preserved. Moreover, in discrete survival analysis, the mean lifetime $\mu$ can be approximated by the sum of the survival probabilities up to $T_{max}$. We could get the estimated mean lifetime by further involving Eq. 2 as

where the employment of the approximately equal symbol in the equation is attributed to the presence of time point $T_{0}=0$. The variance of distribution is computed as

The Study to Understand Prognoses Preferences Outcomes and Risks of Treatment (SUPPORT) [14] is a large static survival dataset of seriously ill hospitalized adults. The Molecular Taxonomy of Breast Cancer International Consortium (METABRIC) [3] is a static breast cancer dataset aiming to distinguish its subtypes based on the molecular characteristics. Their pre-processing strategies follow DeepSurv [13].

We also generate a static synthetic dataset (SYNTH-s) of the style of that in [17] but without competing risks. The dataset contains $N=15,100$ examples drawn from the stochastic process

where $\boldsymbol{x}_{n}^{i}$ is a vector of 4-dimensional variables and $\boldsymbol{\gamma}_{n}=\boldsymbol{10}$. We randomly select $50\%$ patients to be right-censored with random censoring time uniformly drawn from $[0,T^{i}]$. More details are listed in Tab. 1.

On the basis of SYNTH-s, we further generate dynamic synthetic dataset (SYNTH-d) by adding additional dynamic variables following Weibull distribution, and introducing temporal noise disturbances to make them variable over time as

where $\boldsymbol{x}_{v}^{i}$ is a $20\times T_{max}$ dynamic variable matrix for all time points, $a$ is the shape parameter, $\beta$ is the scale parameter, $\boldsymbol{\gamma}_{v_{1}}=\boldsymbol{\gamma}_{v_{2}}=\boldsymbol{5}$, and $max(\cdot)$ and $min(\cdot)$ are operations on the temporal dimension. Besides, $v_{1}$ and $v_{2}$ are two randomly selected subsets that satisfy $v_{1}\cap v_{2}=\varnothing$, $v_{1}\cup v_{2}=v$. $T^{i}$ is then resampled, and the method of introducing censoring cases remains the same as before.

Moreover, MSReactor [22] dataset is a quantifiable, objective collection on cognition via longitudinal computerized test for Multiple Sclerosis (MS), integrating with other 8 static covariates. In each test, patients are instructed to respond as quickly as possible to onscreen stimuli, and their reaction time is recorded in millisecond (ms). The test includes 3 different tasks for testing their psychomotor function, attention and working memory. Each patient undergoes the test a number of times after the diagnosis and prior to the occurrence of the event/censoring (with at least one-month interval between every two adjacent tests). The survival event is characterized by EDSS progression through the six-month disability worsening confirmation rule [10]. Numerous research investigations have indicated that utilizing reaction data could potentially offer a more responsive approach for detecting subclinical cognitive impairment in comparison to current cognitive assessment methods [6, 27, 37].

The longitudinal reaction test will be considered as time-varying covariates. However, due to certain redundancies present in MSReactor, evident through pronounced inter-column associations, the characteristics of adjacent columns exhibit robust correlations, deviating from the conventional tabular data extraction where each column represents highly streamlined information. Without the application of specialized data preprocessing techniques and innovative model architectures, the existing survival models may encounter difficulties in extracting meaningful latent patterns from this data. Therefore, we transform the longitudinal tabular data part into a composite “reaction tensor” after monthly imputation, and utilize the certain module to deal with it. Specifically, each patient has a unique 3-dimensional reaction tensor¹¹1More details are illustrated in Appendix A., as shown in Fig. 2. Its Z-axis is corresponding to the 3 different tasks. X-axis is response dimension with fixed length of 30, corresponding to the 30 times the patient needs to finish three tasks separately in each test. Y-axis corresponds to the times patient has undergone tests per month with fixed length from the start time $T_{0}$ to the end time $T_{max}$. The reaction tensor is divided into several smaller tensors along Y-axis by $T_{w}$ as in Fig. 1(a).

C-index [33] is able to estimate ranking ability by comparing relative risks across all pairs in the test set as

where $\mathbb{I}(\star)$ is an indicator function, and $\delta_{i}=0$ if $T^{i}$ is uncensored and 1 otherwise.

MAE-Uncensored (MAE-U) can compensate for the inability of C-index to measure the mean absolute value of the estimated risk score. It is computed as

MAE-Hinge (MAE-H) is a one-sided MAE for only censoring cases, opposite with MAE-U for uncensoring only. It considers only if the predicted time $\hat{\mu}$ is earlier than the censored time $T$ as follow

The area under ROC curve for survival analysis involves treating survival issue as binary classification across various quantiles of event times and defining the sensitivity and specificity as time-dependent measures [16]. The cumulative AUC measures model’s capability of discriminating individuals who fail by a specified $t$ ($T_{j}\leq t$) from subjects who fail after this time ($T_{j}>t$). We compute the mAUC by integrating the cumulative AUC over all time range $(T_{j},T_{j}+1)$.

In terms of C-index, our UniSurv-s secures the first position on SYNTH-s and the second position on both SUPPORT and METABRIC. It also reaches the best mAUC on METABRIC and SYNTH-s and the second best on SUPPORT. DeepSurv shows comparable ranking performance on two real-world datasets. This illustrates that parametric model still hold a slight advantage over non-parametric model, rely on its robust probability distribution assumptions. Meanwhile, the performances of the other four models vary, creating a competitive landscape. This makes it difficult to definitively judge their performance under single ranking metrics.

Meanwhile, our UniSurv performs well in MAE-U and exhibits notably superior performance in the realm of MAE-H, with statistical significance compared to other models. Only DeepSurv in SUPPORT is comparable with ours in both two MAEs. Conversely, the performance of TDSM, while excelling in MAE-U, lags notably behind in MAE-H. This is because the loss design of TDSM leads to overfitting on uncensored data throughout the learning process, failing to capture the fact that most censored samples have longer survival times. Further, the inadequacy of TDSM’s predictions for censoring is also evident by Fig. 3, in which we represent the difference between true censoring time and estimated mean lifetime with red lines for some censoring cases. We show the METABRIC results from TDSM, UniSurv-s and the second-best MAE-U model DeepHit here. The more and longer red lines, the model have less sensitivity of censoring prediction. It can be observed that UniSurv has the capability to provide accurate predictions for the majority of censoring cases. This outcome can be attributed to the incorporation of the MAE-margin concept within the Margin-Mean loss $\mathcal{L}_{mm}$ in Eq. 6, as it leverages prior knowledge from the training dataset to effectively “enforce” predicted survival time to exceed the censoring time. On the other hand, DeepHit exhibits significant inefficiency in forecasting longer censoring times. Similar to TDSM, this is also due to the absence of certain constraints within its loss designs beyond the censoring time, which may give rise to a systemic bias in predicting censoring cases.

As depicted in Tab. 3, UniSurv-d1 demonstrates superior performance over two other models for longitudinal datasets, as evidenced by higher values in C-index, mAUC and lower values in two MAEs. However, the performance of these three methods is generally suboptimal, as their C-index values remain below $0.6$. This occurrence likely arises from the fact that the temporal data in MSReactor diverges from conventional survival tabular data, instead representing a reaction testing approach applied to MS patients. Traditional models struggle to effectively extract meaningful insights from this intricate and redundant information. Notably, when we preprocess the computerized test data into ”reaction tensor” and employ CNN to extract latent features, the performance of UniSurv-d2 surpasses the others with statistically significant improvements. However, this ”tensor” method has not demonstrated effectiveness for SYNTH-d, primarily due to the isotropic distribution of each variable $x_{v}$ during data generation, resulting in their mutual independence and lack of correlation.

As shown in Fig. 4, all five histograms depict the distribution of survival times skewed towards the early segment of the time horizon, while censoring times tend to cluster in the latter half, especially in SUPPORT, SYNTH-s, MSReactor and SYNTH-d. This leads to survival models facing difficulty in maintaining predictive accuracy over time, as evidenced by the time-dependent AUC (TD-AUC). For example, all the performances of UniSurv-s, DeepHit and DSM, or their dynamic variants (UniSurv-d2, DDH, RDSM) exhibit a consistent decline in TD-AUC as time progresses. However, UniSurv still outperforms others, especially on two dynamic datasets. For METABRIC, due to its relatively low censoring rate and evenly distributed censoring cases, all three models maintain their TD-AUC quite well, with some even showing an upward trend, particularly UniSurv. It affirms that Transformer encoder based on Margin-Mean-Variance loss learning can effectively alleviate the challenges posed by survival datasets characterized by long-tail distributions.

Fig. 6(b) shows the effect of $T_{w}$. We can observe that when $T_{w}=8$, the model can achieve the highest C-index and lowest MAE-H, which is associated with the progression rate of MS. However, during the same period, MAE-U demonstrates its poorest performance. It is also apparent that the fluctuations in MAE-U and MAE-H exhibit a contrasting pattern. This disparity can be attributed to the disparate distributions of censoring and uncensoring within the MSReactor as in Fig. 4. Meanwhile, this underscores that there exists potential for enhancing the robustness of UniSurv.

As the number of losses increases, finding the optimal weight combination indeed becomes challenging, but grid search can take care of this. The four losses do not need to be standardized to a similar magnitude. The unique characteristics of different datasets can lead to distinct optimal weights for losses. We assessed the sensitivities of $\lambda_{m}$ and $\lambda_{v}$ particularly on MSReactor in Fig. 6(c) and Fig. 6(d), and some selected PDFs are shown in Fig. 6(g). The model exhibits robustness when small variations occur in $\lambda_{m}$ or $\lambda_{v}$, as the performance near their optimal values does not exhibit significant degradation. In some cases, two MAEs even perform better. This phenomenon is attributed to the opposing fluctuation trends exhibited by the MAEs, indicating a trade-off made by UniSurv during training. Notably, the C-index appears to be more sensitive to changes in $\lambda_{v}$ compared to variations in $\lambda_{m}$. As the variations in both weights increase, deviations in the PDF gradually emerge, with its peak drifting further away from the actual event time and assuming irregular shapes.

To emphasize UniSurv’s reliability for higher dimensionality datasets and its robustness to data noise, we have expanded the number of features for the existing synthetic datasets SYNTH-s and SYNTH-d without altering event or censoring settings, resulting in new SYNTH-sk and SYNTH-dk datasets, where k denotes the dimension of $\boldsymbol{x}_{n}^{i}$ in Eq. 11 and Eq. 12 is increased from $4$ to $4\cdot 5^{k}$, and the dimension of $\boldsymbol{x}_{v}^{i}$ in Eq. 12 is increased from $20$ to $20\cdot 5^{k}$. Additionally, we introduced noise $\boldsymbol{\epsilon}^{i}\sim\epsilon_{0}\cdot\mathcal{N}(0,\mathbf{I})$ to $\boldsymbol{x}_{n}^{i}$ and $\boldsymbol{x}_{v}^{i}$ separately in all datasets. As the results shown in Tab. 4, UniSurv performs well on high-dimensional datasets and exhibits robustness to small levels of noise interference.