SSL-SurvFormer: A Self-Supervised Learning and Continuously Monotonic Transformer Network for Missing Values in Survival Analysis

Abstract

Survival analysis is a crucial statistical technique used to estimate the anticipated duration until a specific event occurs. However, current methods often involve discretizing the time scale and struggle with managing absent features within the data. This becomes especially pertinent since events can transpire at any given point, rendering event analysis a continuous concern. Additionally, the presence of missing attributes within tabular data is widespread. By leveraging recent developments of Transformer and Self-Supervised Learning (SSL), we introduce SSL-SurvFormer. This entails a continuously monotonic Transformer network, empowered by SSL pre-training, that is designed to address the challenges presented by continuous events and absent features in survival prediction. Our proposed continuously monotonic Transformer model facilitates accurate estimation of survival probabilities, thereby bypassing the need for temporal discretization. Additionally, our SSL pre-training strategy incorporates data transformation to adeptly manage missing information. The SSL pre-training encompasses two tasks: mask prediction, which identifies positions of absent features, and reconstruction, which endeavors to recover absent elements based on observed ones. Our empirical evaluations conducted across a variety of datasets, including FLCHAIN, METABRIC, and SUPPORT, consistently highlight the superior performance of SSL-SurvFormer in comparison to existing methods. Additionally, SSL-SurvFormer demonstrates effectiveness in handling missing values, a critical aspect often encountered in real-world datasets.

1. Introduction

Survival analysis or time-to-event analysis is a fundamental area of statistics concerned with estimating the time until an event occurs, such as a death in biological organisms [1] and failure in mechanical systems [2]. In this paper, we focus on survival regression, which is a type of analysis that relates the time until an event occurs to one or more covariates that may be associated with that duration of time. Censoring is a common issue in survival analysis where events are not always fully observed, often due to limited observation time windows or missing data. We specifically consider right-censored data, where the event is known to have occurred after a certain time but not precisely when. Examples of applications that employ right-censored data include medical treatment recovery time [3], business prediction [4], and churn prediction [5].

Traditional approaches to survival analysis, such as the Kaplan–Meier [6] estimator, the Cox proportional hazard model [7], and the accelerated failure time method [8], were popular before the advent of machine learning. In recent years, machine learning methods, including support vector machines [9,10], random survival forests (RSFs) [11], and gradient boosting [12], have been proposed for survival analysis. Deep learning models like SuMo-Net [13], DeepHit [14], and DeepSurv [15] have advanced survival analysis through neural networks. However, these approaches face significant limitations. Most rely on basic MLP architectures that struggle to capture complex interactions between features. Additionally, many methods (including DeepHit [14], SurvTRACE [16], and TransDSA [17]) discretize continuous survival time into intervals, resulting in information loss.

Transformer, a highly effective sequence modeling architecture, has demonstrated state-of-the-art (SOTA) performance in various applications, including natural language processing (NLP) [18,19,20], speech recognition [21,22], image understanding [23,24,25], and video analysis [26,27]. Recently, Transformer-based methods for tabular data and survival analysis have emerged, including TabTrans [28], FT-Transformer [29], SurvTRACE [16], and TransDSA [17]. According to the published results, these approaches demonstrate superior performance compared to traditional MLP-based methods. Despite these advancements, these models still face the fundamental discretization problem. This issue occurs when continuous survival times must be converted into fixed intervals, introducing several significant limitations: arbitrary interval selection, a trade-off between temporal resolution and data sparsity, and the constraint of predictions to predefined timepoints rather than the full continuous time spectrum.

Another issue in the survival analysis problem is missing features, some of instances in our dataset may have absent features due to various reasons during data collection, processing, storage, or transmission. It is often desirable to know the most likely values of missing data before performing downstream tasks. Existing deep-learning-based methods in survival analysis are not able to handle missing features in survival data. Consequently, they resort to imputation methods like mean/mode imputation and KNN imputation.

In this paper, we address three critical limitations in current survival analysis approaches: (1) discretization of continuous time leading to information loss, (2) architectural constraints of basic MLPs that fail to capture complex feature interactions, and (3) inability to handle missing data without separate imputation steps. We introduce SSL-SurvFormer, which combines a continuously monotonic Transformer architecture with self-supervised learning (SSL) pre-training to overcome these challenges. Our approach enforces positivity constraints on time-related weights, mathematically guaranteeing the monotonicity properties of survival functions without approximation errors. This novel constraint enables direct estimation of the continuous survival function without discretizing time intervals. The model operates in two phases: SSL pre-training with data augmentation followed by fine-tuning of our continuously monotonic Transformer. This integrated approach eliminates the need for separate imputation steps, as missing data handling is directly incorporated into the model. Comprehensive experiments on the FLCHAIN [30], METABRIC [31], and SUPPORT [32] datasets demonstrate that SSL-SurvFormer significantly outperforms existing methods in both prediction accuracy and missing data handling capability. As shown in

Table 1. Comparison between our proposed SSL-SurvFormer and the existing survival analysis methods. MLP stands for Multilayer Perceptron.

Method	Time Continuity	Deep Learning	Architecture	Handle Missing Features
CoxPH [7]	✓	✗	–	✗
DeepSurv [15]	✓	✓	MLP	✗
DeepHit [14]	✗	✓	MLP	✗
SuMo-Net [13]	✓	✓	MLP	✗
SurvTRACE [16]	✗	✓	Transformer	✗
SSL-SurvFormer	✓	✓	Transformer	✓

✓indicates that the method supports the corresponding feature; ✗indicates that the method does not supports the corresponding feature.

, our approach maintains superior performance across datasets with varying levels of missing values, highlighting its robustness and practical utility for real-world survival analysis applications.

2. Related Work

Traditional Survival Prediction: One of the most common survival models in statistical and medical research is the Kaplan–Meier estimator [6], which benefits from learning relatively flexible survival curves. However, it has the drawback of not taking into account the covariates of the patients. The Cox proportional hazard model [7] is capable of incorporating patients’ covariates, but assumes that the hazard rate is constant and that the log of the hazard rate is a linear function of the covariates.

MLP-based Survival Prediction: With the advancement of neural networks, researchers have explored the application of deep learning networks in the field of survival analysis. The Multilayer Perceptron (MLP) [33], a foundational neural network architecture consisting of an input layer, hidden layers with non-linear activation functions, and an output layer, has been widely used for its ability to model complex patterns. MLPs were among the first deep learning approaches applied to survival analysis. By replacing the linear function in Cox model with a MLP, ref. [15] proposed DeepSurv, which is also optimized the model’s partial log-likelihood. However, this approach relies on the same proportional hazards assumption of Cox’s method, which difficult to verify or need to check very carefully every time in medical application. To overcome this issue, ref. [14] introduced DeepHit, which divides the timescale into discrete intervals, then the model is trained by optimizing a mix of the discrete likelihood and a rank-based score. The task becomes similar to a classification in which each outcome is a binary variable reflecting if the patient survives within a specific time interval. As a non-parametric model, this approach offers better discriminative performances when the underlying survival distribution is unknown. However, it suffers from its discretization i.e., determining the breakpoints for time intervals is non-trivial, since too many intervals may lead to unstable model estimation while too few intervals may cause information loss. Furthermore, since the model is optimized for a fixed set of output timepoints, it may not be adaptable for making inferences on other timepoints. There are some attempts to surpass the limitations of timescale discretization. Refs. [34,35] use Ordinary Differential Equations (ODEs) to parameterize the cumulative hazard function, while ref. [36] uses ODEs to parameterize the cumulative distribution function. Closer to our work, ref. [13] described another neural network that avoids the computational burden of ODE, which is a simple MLP survival model capable of directly modeling survival data continuously using a monotonic restriction on the time-dependent weight. While NeuralODEs represent a powerful modeling paradigm for survival analysis, they suffer from significant computational limitations in practical applications. As demonstrated by [37,38], NeuralODEs exhibit computational complexity that grows polynomially with the number of parameters, primarily due to their reliance on numerical ODE solvers that require multiple function evaluations per forward pass. In contrast, our monotonically positive component maintains linear scaling characteristics similar to standard neural networks, offering substantially improved computational efficiency for larger models. Furthermore, by enforcing strict positivity constraints on time-related weights, our approach provides mathematical guarantees of monotonicity for survival functions without introducing the approximation errors inherent in numerical ODE solutions. This mathematical guarantee eliminates potential instabilities and ensures that our model’s outputs consistently satisfy the fundamental monotonicity property required for valid survival function estimation, regardless of parameter initialization or training dynamics. Besides, the deep learning-based models mentioned earlier primarily rely on MLP networks, making them incapable of directly addressing missing values in the dataset.

Transformer-based Survival Prediction: Transformer was first introduced in [39] for language translation and obtained state-of-the-art (SOTA) results in many other language processing tasks. Recently, many models [40,41,42] have successfully applied the Transformer concept to computer vision and tabular data and achieved promising performance. The core idea behind transformer architecture [39] is the self-attention mechanism to capture long-range relationships. It has obtained state-of-the-art in many Natural Language Processing (NLP) tasks. Moreover, Transformers are well suited for parallelization, facilitating training on large datasets. In survival analysis, TransDSA [17] and SurvTRACE [16] utilize Transformer instead of an MLP to better learn survival distribution, model long-range coherence, and obtain higher performance. TransDSA [17] proposed using a Transformer network along with time embeddings, but they treated numerical and categorical data using a single network. On the other hand, SurvTRACE [16] also employed the Transformer architecture but had separate embedding modules for these two types of features, resulting in improved performance. However, they have shared the same discrete time approach with DeepHit [14], dividing the time scale into intervals. This approach posed a common challenge in selecting the appropriate interval. Additionally, neither of these methods offered a solution for handling missing data, requiring the imputation of missing values before training the models.

Handling missing features in survival data: There are several approaches to handling missing data, broadly categorized into deletion and imputation methods. In deletion, entries with missing values are discarded, and this can be performed through pairwise or list-wise deletion [43]. However, as demonstrated in [44], deletion has some weaknesses, as it can introduce bias. On the other hand, imputation involves replacing missing values with predicted values, as outlined in [43]. In simple imputation, missing data can be handled using various methods, including the mode, mean, or median of the available values [45]. Mean imputation is commonly used for numerical data, while mode imputation is typically applied to categorical data. In mode imputation, the mode of an attribute is used to impute the missing value for the corresponding attribute.

A more advanced approach is K-nearest neighbor (KNN) imputation, which use Euclidean distance to impute the missing features using the closest samples. In the era of deep learning (DL), generative models have been utilized to address missing values. For instance, in [46], a GAN-based model is employed to generate synthetic values for the missing data, leveraging the information from the observed values. Similarly, [47] suggested using a generative model, specifically the Variational Autoencoder (VAE) model, to perform the imputation of missing values. Generative approaches for handling missing values typically involve a two-step process. In the first step, a dedicated model is trained for imputation to fill in the missing data; in the second step, a downstream model is trained for the primary task. In this work, we concentrate exclusively on survival analysis tasks involving missing features by leveraging SSL.

3. Proposed Method

3.1. Problem Setup and Notations

Our proposed SSL-SurvFormer aims to model survival prediction in continuous time, thus, we follow the problem setup defined in SuMo-Net [13] as follows. Given a dataset $\mathcal{D}={\{x_i,s_i,e_i\}}_{i=1}^Z$, where Z is the number instances in the dataset $\mathcal{D}$ and i indicates the ith instance (i.e, patient), each instance in the dataset associates with a set of covariates $x_i\in {\mathbb{R}}^N$ for $N\ge 1$, the time when the events occur or the sample is censored $s_i\in {\mathbb{R}}_{\ge 0}$, and event $e_i\in 0,1$ censoring or occurrence at time $s_i$, where $e_i=0$ means censoring. Let T be a random variable representing time to the event; our work only considers the case in which event time T can be observed or censored (right-censored). Our model objective is to estimate the survival function: $S\left(t\right|x)$, which is defined as a (conditional) survival distribution $S\left(t\right|x)=1-H(t\left|x\right)=P(T>t|x)$. In the formula, $H\left(t\right|x)=P(T<t\left|x\right)$ is the cumulative distribution function (CDF) when assuming event time is a continuous random variable. Moreover, the probability density function (PDF) of T is $h\left(t\right|x)={\displaystyle \frac{d}{dt}}H\left(t\right|x)$.

Our goal is summarized as follows: If the survival function of T is denoted by $S\left(t\right|x)$, we aim to predict an estimated distribution $\widehat{S}\left(t\right|x)$ using dataset $\mathcal{D}={\{x_i,s_i,e_i\}}_{i=1}^Z$, such as estimated distribution $\widehat{S}$ as close as truth distribution S.

We first present a continuously monotonic Transformer for survival analysis, named SurvFormer. We then introduce SSL and integrate into SurvFormer to form SSL-SurvFormer.

3.2. SurvFormer: Network Architecture

The overall network architecture of our proposed continuously monotonic Transformer, termed SurvFormer, is given in Figure 1. SurvFormer consists of three components including Feature Embedding to extract the feature representation of an instance, Transformer to model the contextual coherence between attributes, and a Monotonically Positive to restrict the monotonicity constraint. Details of each component are as follows:

3.2.1. Feature Embedding

For a general instance defined by $(x,s,e)$, its data x can be separated into two types: categorical data $x^c\in {\mathbb{R}}^{N^c}$ and numerical data $x^n\in {\mathbb{R}}^{N^n}$ such as $x=[x^c,x^n]\in {\mathbb{R}}^N$ and $N=N^c+N^n$. We transform the categorical data into numbers and normalize the numerical data and the instance’s time s before feeding the model. We separately process numerical data and categorical data as follows.

Numerical data: The numerical data $x^n$ is computed by Equation (1).

$$\begin{array}{c}I=1^D\\ X^n=x^n·I^{\top }\\ F^n=X^n\odot W^n+B^n\end{array}$$

(1)

where $x^n\in {\mathbb{R}}^{N^n}$ is a vector, $I^D$ is the unit vector with D dimensions, $X^n\in {\mathbb{R}}^{N^n\times D}$ is a matrix obtained by replicating the vector $x^n$ D times, ⊙ represents element-wise between $X^n$ and $W^n$, $W^n\in {\mathbb{R}}^{N^n\times D}$ is the embedding matrix for numerical features, and $B^n\in {\mathbb{R}}^{N^n\times D}$ is bias. The output numerical feature representation $F^n$ will be in dimensional space of ${\mathbb{R}}^{N_n\times D}$.

Categorical data: The categorical data, which is commonly represented by one-hot encoding, is embedded as a lookup table through an embedding matrix $W^c\in {\mathbb{R}}^{C\times D}$, where C is the total number of distinct values in all categorical features. For instance, if a patient record contains two categorical variables—occupation (3 unique values) and gender (2 unique values)—then C = 5. The categorical embedding feature of $x^c$ is computed in Equation (2).

$$F^c=O^c·W^c+B^c$$

(2)

where $O^c\in {\mathbb{R}}^{N^c\times C}$ is the one-hot representation of $x^c$ and $B^c\in {\mathbb{R}}^{N^c\times D}$ is bias.

Finally, we stack both the numerical embedding feature $F_n$ and categorical embedding feature $F_c$ to obtain the final feature representation $F=[F_d^c,F_d^n]\in {\mathbb{R}}^{N\times D}$, where $N=N^c+N^n$. The flowchart of the Feature Embedding component is demonstrated in Figure 2.

3.2.2. Transformer

We inherit the merits of PreNorm-variant Transformer architecture defined in [29,48], where the model takes in a sequence of feature embeddings and outputs contextual representations of the same dimension. Following [18], we first append the embeddings of the classification token [CLS] to feature representation F to form a new feature representation $F_0$. The embedding $F_0$ then goes through a Transformer with $\mathfrak{L}$ blocks and we finally obtain the contextual representation $F_{\mathfrak{L}}$ as shown in Figure 1. The detailed network of a Transformer block is given in Figure 3 and the Transformer is described in Algorithm 1.

Algorithm 1 Transformer

Input: Feature representation at first layer $F_0\in {\mathbb{R}}^{(N+1)\times D}$ Parameter: number of Transformer blocks $\mathfrak{L}$ Output: Feature embedding $F_{\mathfrak{L}}\in {\mathbb{R}}^{(N+1)\times D}$   1: function FFN(z)   2:     return $\mathrm{Linear}\left(\mathrm{Dropout}\right(\mathrm{ReGLU}\left(\mathrm{Linear}\right(z\left)\right)\left)\right)$   3: end function   4: function RPN($\mathrm{Module},z$)   5:     return $z+\mathrm{Dropout}\left(\mathrm{Module}\right(\mathrm{Norm}\left(z\right)\left)\right)$   6: end function   7: function Block(z)   8:     return $\mathrm{R}\mathrm{P}\mathrm{N}(\mathrm{F}\mathrm{F}\mathrm{N},\mathrm{R}\mathrm{P}\mathrm{N}(\mathrm{MHSA},z\left)\right)$   9: end function 10: $i\leftarrow 0$ 11: while $i<\mathfrak{L}$ do 12:     $F_i\leftarrow$ Block(Fi−1) 13:     $i\leftarrow i+1$ 14: end while 15: return $F_{\mathfrak{L}}$

In Algorithm 1, FFN and RPN stand for feed-forward network and residual pre-norm module, Norm is LayerNorm [49], and MHSA is multihead self-attention, respectively. We use ReGLU (Regularized Gated Linear Unit) as an activation function due to its reported superiority over the commonly used GELU activation [50,51].

3.2.3. Monotonically Positive

This component aims to restrict the monotonicity constraint and output the estimated survival function. Let $S\left(t\right|x)$ denote the survival function, which is defined as in Equation (3)

$$S\left(t\right|x)=1-H\left(t\right|x)=1-{\int }_0^{t}h\left(t\right|x)dt$$

(3)

Survival functions follow these same three characteristics:

• As t increases, $S\left(t\right)$ should decrease.
• $S\left(0\right)=1$. That is at the start of the study, no one has the event, thus the probability of surviving at $t=0$ is 1.
• $S(\infty )=0$. If the study were to go to $S(\infty )$, then everyone will eventually experience the event and hence the survival probability must be 0.

Inspired by [13,52], our monotonically positive component enforces the network’s weights to be positive to ensure the monotonicity of the survival function. To ensure that the output of the neural network does not decrease as time t increases, we enforce non-negativity constraints on the weights applied to time t.

Our monotonically positive component is defined as an MLP with K layers, where tanh is an activation function of layers $k=0,1,\dots ,K-1$ and sigmoid is an activation function of the last layer. The monotonically positive component takes the last feature representation of $F_{\mathfrak{L}}$, i.e., $f_{CLS}$ and time t as input of the first layer $k=0$. At the first layer $k=0$, the weights on time input t are required to be positive, while the weights on the embeddings $f_{CLS}$ have no constraint. As other layers $k>0$, all the weights are forced to be positive. We use the log space or the square function as in [13] to guarantee the positivity of weights and the existence of a derivative. In this network, sigmoid $\sigma$ is used in the last layer to output a probability-like value, which also is $H\left(t\right|x)$. Finally, the survival probability at time t is equal to the difference of $H\left(t\right|x)$ and 1, i.e., $S\left(t\right|x)=1-H(t\left|x\right)$. Moreover, another output is the PDF of t, which is obtained by computing the derivative of $H\left(t\right|x)$, i.e., $h\left(t\right|x)={\displaystyle \frac{d}{dt}}H\left(t\right|x)$, and is used to compute the loss function in Equation (4). Our monotonically positive network computed the $H\left(t\right|x)$ function through feed-forward. After that, we used PyTorch’s autograd [53] to backpropagate $H\left(t\right|x)$ to obtain $h\left(t\right|x)$ with respect to t. Thus, when the computed gradient of the loss is mentioned, we used autograd again to backpropagate $h\left(t\right|x)$ with respect to the model’s weights. Our implementation follows [13]. The entire flowchart of the monotonically positive component is given in Figure 4.

3.2.4. Loss Function

Our model aims to optimize the negative log-likelihood loss function for censored data, which is often used to train the survival model. As defined in [13], the right-censored log-likelihood score of a set of observations ${\left\{(x_i,s_i,e_i)\right\}}_{i=0}^Z$ is calculated as in Equation (4):

$${\mathcal{L}}_{llh}=-\sum _{i=0}^Z\left[(1-e_i)logS\left(s_i\right|x_i)+e_{i}logh\left(s_i\right|x_i)\right]$$

(4)

Furthermore, we propose to add another loss component ${\mathcal{L}}_{order}$, which will penalize the model when its outputs are not in the true order. In this loss ${\mathcal{L}}_{order}$, we compare the estimated survival function between two instance $x_i$ and $x_j$ observed at the same time $s_i$.

$$\begin{array}{c}{A}_{i,j}\triangleq \mathbb{1}(e_i=1,s_i<s_j)\\ \eta (a,b)=exp\left({\displaystyle \frac{(a-b)}{\sigma }}\right)\\ {\mathcal{L}}_{order}={\displaystyle \sum _{i\ne j}}{A}_{i,j}·\eta (S\left(s_i\right|x_i),S\left(s_i\right|x_j))\end{array}$$

(5)

where A is the indicator of pair $(i,j)$, $S\left(s_i\right|x_i)$ and $S\left(s_i\right|x_j)$ are survival functions of two instances at time $s_i$, and $\eta (a,b)$ represents the function that determines the magnitude of penalization the model incurs when its output is in an incorrect order. This magnitude can be adjusted by manipulating the hyperparameter $\sigma$: a smaller $\sigma$ results in a greater degree of penalization for the model. The optimal value of $\sigma$ is discussed in Appendix C,

Table A4. Hyperparameters of our proposed SurvFormer. Other hyperparameters of Transformer are used similarly to the default configuration in [12].

Hyperparameter	Values
Feature embedding size (D)	$\{16,32,64,128\}$
No. attention heads	$\{2,3,4,8\}$
No. layers Positive	$\{1,2,3\}$
No. nodes per positive layer	$\{16,32,64\}$
Learning rate	$\{5\times {10}^{-3},{10}^{-3},5\times {10}^{-4},{10}^{-4}\}$
Weigh decay	$\{{10}^{-3},{10}^{-4},{10}^{-5}\}$
Batch size	$\{16,32,64,128\}$
Loss $\beta$	$\{0.5,0.6,0.7,0.8,0.9,1\}$
Loss $\sigma$	$\{0.1,0.3\}$

. As a result, our final loss function can be described as in Equation (6):

$$\mathcal{L}={\mathcal{L}}_{llh}+\beta {\mathcal{L}}_{order}$$

(6)

where $\beta$ is the hyperparameter are chosen to trade off between log-likelihood loss and order loss. The optimal value of $\beta$ is discussed in Appendix C, Table A4. Because ${\mathcal{L}}_{order}$ is computed on the difference of two instances’ survival probability at a point of time, we need to choose a set of timepoints for evaluation. In our work, we adapt the discretization approach to divide time duration into equidistant bins in order not to calculate survival function at all the unique timepoints in the dataset. This approach maintains performance but is less time-consuming. In particular, we choose three timepoints at quantiles of 25%, 50%, and 75%. We only use discretization during the training phase as a regularization, so it does not affect flexibility in the inference process.

3.3. SSL-SurvFormer: SurvFormer with SSL

To achieve good performance and avoid random phenomena such as sampling order or weight initialization, most existing DL-based approaches use pre-trained models trained on large-scale datasets such as ImageNet [54] or Kinetics-400 [55]. However, such pre-trained models are not available in tabular data. Self-supervised learning (SSL) [56,57,58] was proposed to overcome this drawback and it has been successfully applied to various domains. A typical workflow to apply SSL is to train the network by learning with a pretext task in the pre-training stage and then fine-tuning the pre-trained network on a target downstream task.

In our research, we extend the use of a self-supervised method to develop a robust representation of survival data, even when dealing with missing values. Survival datasets commonly consist of two parts: fully recorded data and data with some missing features. Our approach involves initially pre-training the model on the complete data subset and subsequently fine-tuning it on the entire dataset. Given that complete data are frequently limited in proportion, we suggest employing an augmentation method to generate additional data samples for the pre-training phase.

3.3.1. Data Augmentation

Data augmentation plays a crucial role in enhancing model robustness when dealing with limited or incomplete datasets. In the context of survival analysis with missing values, we propose a specialized augmentation technique tailored to tabular data structures.

Our approach to data augmentation involves a technique called random feature discarding. This method systematically simulates missing data patterns to help the model learn robust representations even when features are absent. For a given data sample $x_i\in {\mathbb{R}}^N$, we randomly produce a mask vector $m\in {\{0,1\}}^N$, where the dimension matches that of $x_i$ and the value 1 indicates the features that will be discarded. The process is describe in Figure 5.

When dealing with numerical data, we rescale them to fit within the range $[1, 2]$ and represent missing values as 0. This normalization ensures consistent handling across different numerical scales. In the case of categorical data, we treat missing values as a distinct category [M], which allows the model to explicitly recognize and learn patterns associated with missing categorical information.

By employing this method, we are able to generate a considerable amount of data from the original complete dataset for our pre-training objective without necessitating labeled information. This augmentation strategy effectively expands the training distribution, enabling the model to generalize better to various missing data scenarios encountered during inference.

3.3.2. Pre-Training

In the pretext task, we aim to pre-train a feature extraction component corresponding to the two models of Feature Embedding and Transformer. This pre-training module is designed to provide a pre-trained Feature Embedding model and pre-trained Transformer model, which are then transferred to the downstream task of survival prediction. To ensure both Feature Embedding and Transformer are robust against missing values, the proposed pre-training takes a complete feature $x\in {\mathbb{R}}^N$ as its input and outputs the feature $f_{CLS}$, which is capable of being robust against missing values. To achieve this goal, we initiate the process by introducing an augmentation technique, as discussed in Section 3.3.1, to the original complete dataset x. This augmentation process generates synthetic data with missing values, denoted as $\tilde{x}$. This augmented dataset $\tilde{x}$ includes categorical features ${\tilde{x}}^c$, numerical features ${\tilde{x}}^n$, and a corresponding missing mask m. Subsequently, we employ the Feature Embedding and Transformer components to extract feature representations from ${\tilde{x}}^c$ and ${\tilde{x}}^n$, resulting in the feature representation $f_{CLS}$ for the instance $\tilde{x}$. Within this pre-training module, both the Feature Embedding and Transformer models are trained on two pretext tasks: mask prediction and feature reconstruction. For each of these pretext tasks, we apply a Fully Connected Layer (FCN) to the feature $f_{CLS}$.

In the mask prediction task, we denote $\widehat{m}$ as the predicted mask. The loss ${\mathcal{L}}_m$ between the actual missing mask m and the predicted missing mask $\widehat{m}$ is computed using Binary Cross-Entropy (BCE) as follows: ${\mathcal{L}}_m=\mathrm{BCE}(m,\widehat{m})$.

In the feature reconstruction task, we expect both the Feature Embedding and Transformer models to be robust in handling missing values in the corrupted categorical vector, meaning they should be capable of reconstructing $x^c$ given ${\tilde{x}}^c$. The categorical reconstruction loss is defined as ${\mathcal{L}}_c=\mathrm{MBCE}(x^c,{\widehat{x}}^c)$, where MBCE stands for Multi-Label Binary Cross-Entropy Loss. Additionally, by applying the Feature Embedding and Transformer models to the corrupted numerical vector ${\overline{x}}^n$, we aim to reconstruct the vector $x^n$. Hence, the numerical reconstruction loss is defined as ${\mathcal{L}}_n=\mathrm{MSE}(x^n,{\widehat{x}}^n)$, where ${\widehat{x}}^n$ represents the reconstructed numerical vector.

Consequently, the pre-training procedure is trained using the following loss, which combines the mask prediction loss ${\mathcal{L}}_m$, categorical reconstruction loss ${\mathcal{L}}_c$, and numerical reconstruction loss ${\mathcal{L}}_n$.

$${\mathcal{L}}_{pretext}=\alpha {\mathcal{L}}_m+(1-\alpha )({\mathcal{L}}_n+{\mathcal{L}}_c)$$

(7)

where $\alpha$ is a hyperparameter discussed in Appendix C,

Table A5. Hyperparameters for pre-training SSL.

Hyperparameter	Values
Corrupt ratio	$\{0.4,0.5,0.6\}$
Learning rate	$\{5\times {10}^{-3},{10}^{-3},5\times {10}^{-4},{10}^{-4}\}$
Weigh decay	$\{{10}^{-3},{10}^{-4},{10}^{-5}\}$
Batch size	$\{16,32,64,128\}$
Loss $\alpha$	$\{0,0.1,0.2,0.3,0.4,0.5,0.6,0.7,0.8,0.9,1\}$

.

4. Experiments

4.1. Datasets

We evaluate the performance of our model as well as the baseline methods on datasets that include missing values. Specifically, we employ the METABRIC (https://www.cbioportal.org/study/summary?id=brca_metabric (accessed on 17 December 2023)) dataset [31], which consists of 1904 patients and encompasses 30 clinical attributes and 4 genetic attributes. Additionally, we make use of the SUPPORT (https://www.vumc.org/biostatistics/vanderbilt-department-biostatistics (accessed on 17 December 2023)) dataset [32], which comprises 9105 patients and includes 24 clinical features. We also compare our model’s performance with others on a complete dataset. Specifically, we choose another widely-used dataset, FLCHAIN [30], which includes an assay of serum free light chain for 6524 subjects, with 8 covariates. We obtain the dataset from the pycox library (https://github.com/havakv/pycox (accessed on 17 December 2023)). Our model is trained on 5-fold cross-validation, thus, each dataset is divided into 64/16/20 for training/validation/testing. Comprehensive details regarding the percentage of missing values and the maximum number of missing features are available in

Table 2. Datasets summarization. “No. F” is number of features. “Max.No.M.F.” is maximum number of missing features in dataset. “M.R” is missing ratio in dataset.

Datasets	No. F	Max.No.M.F.	M.R (%)
METABRIC	30	12	35.6%
SUPPORT	24	16	96.6%
FLCHAIN	8	0	0%

. Further information on the feature selection process can be found in the Appendix B.

4.2. Evaluation Metrics

In the realm of survival analysis, the concordance index (C-index) [59] and Brier score [60] are important metrics to evaluate the performance of the survival prediction model.

The C-index [59] measures the proportion of pairs of patients for which the predicted survival time agrees with the actual survival time. Specifically, it considers all pairs of patients in which at least one patient has died. For each pair, if the predicted survival time is larger for the patient who lived longer, the prediction is considered consistent with the outcome. The C-index can handle right-censored patients and provides an overall measure of the model’s predictive accuracy, ranging from 0 to 1, with larger values indicating better accuracy.

The Brier score [60] considers both discrimination and calibration. It is decomposed into two primary components: reliability and resolution. Reliability measures the closeness of predicted probabilities to the true probabilities, while resolution measures how much the conditional probabilities differ from the predictive average. The Brier score is a scalar value ranging from 0 to 1, with smaller values indicating superior prediction performance.

However, Harrell’s C-index [59], while widely adopted as the standard measure in survival analysis, has a fundamental limitation: it assumes proportional hazards, where a predictor’s effect on survival remains constant over time. This assumption often fails in clinical settings where treatment effects or risk factors may have varying impacts during different phases of disease progression. The time-dependent C-index [61] that we employ addresses this limitation by evaluating predictive accuracy at specific timepoints rather than assuming uniform effects throughout the follow-up period. This approach is particularly valuable for our study, where the impact of missing data on prediction accuracy may vary across different time horizons. We use C-index to indicate a time-dependent concordance index for now. We benchmark the proposed SSL-SurvFormer with two kinds evaluation based on the C-index and Brier score as follows:

Overall evaluation: The overall metrics include Integrated C-index (IC-index), Integrated Brier Score (IBS), and Integrated Negative Binomial Log-likelihood (INBLL) [62] at all available times of each dataset. The IC-index and IBS provide an overall calculation of the model performance on the C-index and Brier score at all available times, respectively. INBLL provides the measurement of both discrimination and calibration of the estimates of a method under every timepoint. A higher IC-index aand a lower IBS and INBLL are better. We employ integrated metrics (IC-index, IBS, and INBLL) to evaluate performance across the entire time spectrum rather than at isolated timepoints. This provides a more comprehensive evaluation of model performance with missing data, capturing both immediate and long-term predictive accuracy. Integrated metrics are especially appropriate in our context because the effects of missing values may manifest differently across the survival curve, and a comprehensive assessment over all available timepoints yields a more reliable comparison between methods than evaluations at arbitrary timepoints. Those metrics are used to evaluate continuous-time methods.

Quantile evaluation: The quantile metrics include the C-index [61] and Brier score [60] at the dataset-specific 0.25, 0.5, and 0.75 quantiles of the uncensored population event times. Those metrics are used to evaluate discrete-time methods.

The mathematical formulas of the metrics are reported in the Appendix A.

4.3. Implementation Details

Based on two kinds of evaluation metrics, we have implemented the model in two scenarios as follows:

Overall evaluation: We followed [13,15] and evaluated our SSL-SurvFormer models over the entire observed duration. Specifically, we used the models to inference at 1000 timepoints that are evenly distributed across the timescale. Overall metrics are computed on all these timepoints. In this evaluation, we are comparing our model with three other existing effective models: the Cox Proportional Hazard (CoxPH) model [7], DeepSurv [15], and SuMo-Net [13]. The CoxPH [7] model estimates survival outcomes by creating a linear combination of covariates as expressed by the formula: $h\left(x\right)={\beta }^{\top }x$. DeepSurv [15], on the other hand, is an extension of CoxPH that employs a neural network instead of a linear model for its calculations. Meanwhile, SuMo-Net [13] utilizes an MLP network with a positive constraint on the model’s weights to estimate the cumulative distribution function $H\left(x\right)$, which is described in Section 3.2.3. This approach closely resembles our own method. To compare our SSL-SurvFormer with the existing methods for handling missing data, we used two common approaches as baselines: mean and mode (MM) imputation and KNN imputation, as described in Section 2. Specifically, when dealing with categorical data, we applied mode imputation, which involved filling in missing values with the most frequently occurring value in the dataset. For numerical data, we opted for mean imputation, replacing the missing values with the average value. The implementations of MM and KNN followed those in [63].

Quantile evaluation: We followed [14,16] and set the last layer to a vector in five-dimensional space, which reflects the output at the 0, 0.25, 0.5, 0.75, and 1 quantiles. In this case, we compare the performance of our proposed methods with two well-known discrete-time models, i.e., DeepHit [14] and SurvTRACE [16], which are directly optimized on these discrete points. The performance is reported at 0.25, 0.5, and 0.75. As with the overall evaluation, we applied both MM and KNN imputation techniques to address missing data.

Regarding hyperparameters, we have followed the widely used settings [13,28,29] and conducted hyperparameter search spaces to find the optimal configuration. The details of search spaces are included in the Appendix C. We implemented our method in Pytorch. To conduct a fair comparison, we used the AdamW optimizer [64] with early stopping set at 10 iterations for our methods and other existing work. All of our experiments were trained, evaluated on a GPU Tesla T4 16 GB, and conducted under the same conditions in terms of package versions. All the requirements of software libraries and frameworks will be made publicly available along with the code upon publication of the paper.

4.4. Performance Comparison

Our primary objective is to compare our method with other deep learning and Transformer-based approaches that have demonstrated superiority over conventional techniques (such as Cox regression [7] and Kaplan–Meier estimation [6]) in previous studies. We classify survival prediction methods into two groups based on discrete time and continuous time in Table 1. Consequently, we compare our proposed SSL-SurvFormer with the existing SOTA in each group as follows:

Overall Comparison with continuous-time methods: 

Table 3. Overall performance comparison on two missing datasets: METABRIC (MET.) [31] and SUPPORT (SUP.) [32] between our proposed SSL-SurvFormer and the continuous-time methods DeepSurv [15] and Sumo-Net [13]. Results are reported as mean (standard deviation) over the 5-fold cross validation. On each dataset, the best is shown in bold. MM and KNN stand for mean and mode and K-nearest neighbor, respectively.

Dataset	Method	Impute	IC-Index ↑	INBLL ↓	IBS ↓
MET.	CoxPH	MM	0.773 (0.010)	0.441 (0.021)	0.143 (0.008)
		KNN	0.776 (0.009)	0.436 (0.028)	0.140 (0.006)
	DeepSurv	MM	0.778 (0.015)	0.433 (0.022)	0.139 (0.013)
		KNN	0.780 (0.012)	0.426 (0.025)	0.135 (0.010)
	SuMo-Net	MM	0.781 (0.012)	0.425 (0.033)	0.135 (0.011)
		KNN	0.784 (0.011)	0.418 (0.030)	0.131 (0.013)
	SSL-SurvFormer		0.796 (0.014)	0.406 (0.021)	0.129 (0.009)
SUP.	CoxPH	MM	0.658 (0.008)	0.566 (0.021)	0.195 (0.003)
		KNN	0.662 (0.006)	0.563 (0.022)	0.191 (0.004)
	DeepSurv	MM	0.668 (0.009)	0.563 (0.022)	0.188 (0.006)
		KNN	0.668 (0.010)	0.559 (0.019)	0.186 (0.009)
	SuMo-Net	MM	0.673 (0.010)	0.545 (0.012)	0.182 (0.007)
		KNN	0.679 (0.009)	0.540 (0.018)	0.180 (0.007)
	SSL-SurvFormer		0.708 (0.007)	0.528 (0.011)	0.175 (0.004)

↑ indicates higher values are better; ↓ indicates lower values are better.

reports the overall performance comparison on two missing-value datasets between our proposed SSL-SurvFormer and continuous-time methods: Cox Regression [7], DeepSurv [15], and Sumo-Net [13]. We conducted this comparison using both imputation approaches, namely mean and mode (MM) imputation and KNN imputation, for the baselines, as mentioned in Section 4.3. The findings indicate that the KNN imputation method exhibits superior performance when compared to the MM approach. Nevertheless, it is evident from the results that our SSL-SurvFormer model outperforms the other models, regardless of the imputation method used, highlighting its effectiveness in survival analysis, particularly when confronted with datasets containing missing data.

In particular, our SSL-SurvFormer demonstrates a noteworthy advantage when compared to other baseline models. When we compare our model to SuMo-Net [13] and DeepSurv [15], utilizing both KNN imputation and MM imputation on the METABRIC [31] dataset, our model consistently achieves a higher IC-index, with gaps of 0.012, 0.015, 0.016, and 0.018, respectively. Furthermore, our model significantly improves the INBLL [62], with gaps of 0.012, 0.019, 0.020, and 0.027 in comparison to SuMo-Net [13] and DeepSurv [15] using these two imputation methods.

Similarly, when comparing our model to SuMo-Net [13] and DeepSurv [15], utilizing both MM imputation and KNN imputation on the SUPPORT [32] dataset, our SSL-SurvFormer demonstrates a higher IC-index, with gaps of 0.029, 0.035, 0.040, and 0.040, respectively. Additionally, our model improves the INBLL by 0.012, 0.017, 0.031, and 0.035 when compared to SuMo-Net [13] and DeepSurv [15] using these two imputation methods.

In other comparisons, our model consistently surpasses the other baseline models with a competitive margin. This serves as strong evidence of its superiority in addressing survival analysis problems.

Quantile Comparison with discrete-time approaches: 

Table 4. Quantile performance comparison between SSL-SurvFormer and discrete-time methods on two missing datasets. Results are reported as mean (standard deviation) over the 5-fold cross validation. On each dataset, the best is shown in bold. CI and BS stand for C-index and Brier score, respectively. MM and KNN stand for mean and mode and K-nearest neighbor, respectively.

Method			DeepHit		SurvTRACE		SSL-SurvFormer
Impute			MM	KNN	MM	KNN
METABRIC	CI ↑	$q_{0.25}$	0.815 (0.028)	0.815 (0.031)	0.816 (0.022)	0.817 (0.028)	0.826 (0.016)
		$q_{0.5}$	0.799 (0.022)	0.802 (0.029)	0.806 (0.026)	0.808 (0.021)	0.816 (0.020)
		$q_{0.75}$	0.784 (0.029)	0.785 (0.031)	0.785 (0.024)	0.788 (0.026)	0.795 (0.020)
	BS ↓	$q_{0.25}$	0.112 (0.007)	0.110 (0.009)	0.101 (0.005)	0.099 (0.005)	0.095 (0.008)
		$q_{0.5}$	0.162 (0.009)	0.160 (0.012)	0.153 (0.008)	0.148 (0.009)	0.137 (0.013)
		$q_{0.75}$	0.174 (0.012)	0.173 (0.013)	0.164 (0.011)	0.159 (0.010)	0.152 (0.017)
SUPPORT	CI ↑	$q_{0.25}$	0.769 (0.013)	0.771 (0.015)	0.773 (0.020)	0.773 (0.019)	0.785 (0.014)
		$q_{0.5}$	0.717 (0.013)	0.719 (0.012)	0.718 (0.009)	0.721 (0.011)	0.738 (0.007)
		$q_{0.75}$	0.680 (0.012)	0.684 (0.015)	0.682 (0.014)	0.684 (0.012)	0.696 (0.009)
	BS ↓	$q_{0.25}$	0.143 (0.012)	0.143 (0.009)	0.128 (0.006)	0.124 (0.005)	0.118 (0.004)
		$q_{0.5}$	0.204 (0.010)	0.201 (0.006)	0.197 (0.009)	0.195 (0.007)	0.179 (0.002)
		$q_{0.75}$	0.227 (0.006)	0.224 (0.005)	0.214 (0.010)	0.213 (0.007)	0.204 (0.006)

↑ indicates higher values are better; ↓ indicates lower values are better.

reports the quantile performance comparison at 0.25, 0.5, and 0.75 between our SSL-SurvFormer and existing SOTA discrete-time approaches, namely DeepHit [14] and SurvTrace [16]. It is worth noting that the KNN imputation method consistently outperforms the MM approach, aligning with the findings in the overall comparison. However, what remains consistent across all scenarios is that our SSL-SurvFormer model consistently outperforms the two baseline models—DeepHit [14] and SurvTrace—across all quantiles and with both imputation strategies. As an example, when we compare our proposed model to the second-best performer, SurvTRACE [16] with KNN imputation, on the METABRIC [31] dataset, our model demonstrates a significant performance advantage in terms of the C-index across three quantiles. The performance gaps are notable, measuring 0.009, 0.008, and 0.007, respectively. Likewise, when we evaluate our model against the second-best performer, SurvTRACE [16] with KNN imputation, on the SUPPORT [32] dataset, our model outperforms it with substantial performance differences, measuring 0.012, 0.017, and 0.012 across three quantiles. Furthermore, SSL-Survformer outperforms SurvTRACE [16] with KNN imputation by achieving lower Brier scores across three quantiles, with differences of 0.004, 0.011, and 0.007 on the METABRIC [31] dataset and 0.006, 0.016, and 0.009 on the SUPPORT [32] dataset.

Comparision on Complete Data 

Table 5. Overall performance comparison on the no-missing value dataset FLCHAIN [30] between our proposed SSL-SurvFormer and continuous-time methods: Cox Regression [7], DeepSurv [15], and Sumo-Net [13]. Results are reported as mean (standard deviation) over the 5-fold cross validation. The best is shown in bold.

Method	IC-Index ↑	INBLL ↓	IBS ↓
CoxPH	0.782 (0.010)	0.356 (0.016)	0.110 (0.004)
DeepSurv	0.788 (0.008)	0.352 (0.013)	0.107 (0.003)
SuMo-Net	0.788 (0.011)	0.343 (0.017)	0.105 (0.006)
SSL-SurvFormer	0.795 (0.007)	0.333 (0.014)	0.101 (0.005)

↑ indicates higher values are better; ↓ indicates lower values are better.

and

Table 6. Quantile performance comparison between SurvFormer, SSL-SurvFormer, and discrete-time methods on the no-missing value dataset FLCHAIN [30]. Results are reported as mean (standard deviation) over the 5-fold cross validation. The best is shown in bold. CI and BS stand for C-index and Brier score, respectively.

		DeepHit	SurvTRACE	SSL-SurvFormer
CI ↑	$q_{0.25}$	0.798 (0.026)	0.800 (0.030)	0.802 (0.026)
	$q_{0.5}$	0.793 (0.010)	0.796 (0.009)	0.800 (0.010)
	$q_{0.75}$	0.790 (0.008)	0.791 (0.012)	0.794 (0.011)
BS ↓	$q_{0.25}$	0.083 (0.016)	0.059 (0.005)	0.059 (0.006)
	$q_{0.5}$	0.136 (0.014)	0.100 (0.004)	0.099 (0.004)
	$q_{0.75}$	0.210 (0.027)	0.126 (0.009)	0.125 (0.008)

↑ indicates higher values are better; ↓ indicates lower values are better.

display our performance relative to other baseline models on the fully complete FLCHAIN [30] dataset. The findings demonstrate that SSL-SurvFormer excels not only on datasets with missing features but also surpasses other methods on datasets without any missing data. Specifically, our model outperforms the second-best SuMo-Net [13] across all overall metrics, with a noticeable gap of 0.007 on the IC-index, 0.01 on INBLL, and 0.004 on IBS. Moreover, when it comes to quantile metrics, our model achieves competitive results when compared to SurvTRACE [16] and surpasses DeepHit [14] with differences of 0.002, 0.007, and 0.004 on three quantiles, respectively, for the C-Index. Additionally, our model outperforms DeepHit [14] with differences of 0.024, 0.037, and 0.085 on three quantiles for the Brier score.

4.5. Results Analysis

To evaluate the effectiveness of our proposed SSL-SurvFormer, we compare its performance on different kind of data, namely uncensored data and censored data. From Section 4.4, we can see that SSL-SurvFormer surpasses other methods on datasets that contain missing values; therefore, we continue investigating more deeply how our model and the others perform when the missing ratio gradually increases.

Comparison on uncensored data and censored data. To conduct this experiment, we use SUPPORT [32] data as an example. Based on whether or not an event occurs, we separate 1821 samples in the testset into two subsets of 1242 uncensored samples (event occurs) and 579 censored samples (event does not occur). Figure 6 shows the comparison of average survival curves between SSL-SurvFormer and other continuous-time methods on both uncensored data (a) and censored data (b). The figure shows that when an event happens, our method can better detect the occurrence of the event than the others, i.e., our predicted survival probability drops at the time of the event and maintains low values. In contrast, DeepSurv [15] and SuMo-Net [13] have limitations in identifying the event signal and their survival probabilities are still high at the event time. On the censored data, i.e., when an event does not happen, SSL-SurvFormer succeeds in maintaining a high survival probability before the censored timepoint.

Comparison on different missing value ratios. We evaluated our model’s performance using the FLCHAIN [30] dataset, which contains eight features per sample. We systematically removed zero, two, three, four, five, or six random features, corresponding to missing value ratios of 0, 0.25, 0.375, 0.5, 0.625, and 0.75, respectively. Figure 7 demonstrates that SSL-SurvFormer consistently outperforms baseline models across all missing value ratios, with particularly significant performance gaps at ratios of 0.25, 0.375, and 0.5. These results indicate the robustness of SSL-SurvFormer in handling missing data scenarios.

4.6. Ablation Study

We analyze the contribution of SSL in SSL-SurvFormer by comparing the SSL-SurvFormer with SurvFormer by removing SSL pre-training and investigating SSL pre-training by considering the impact of each loss term defined in Equation (7).

Efficiency of SSL in SSL-SurvFormer. We analyze the effectiveness of our proposed SSL pre-training in SSL-SurvFormer by first removing the SSL from SSL-SurvFormer (i.e., it becomes SurvFormer) and then investigating the contribution of individual components in SSL.

Table 7. The impact of each component in SSL pre-training on the FLCHAIN [30] dataset (fully complete dataset). The baseline is our SurvFormer model without pre-training with SSL.

	Baseline	SSL
		${\mathcal{L}}_{\mathbf{m}}$	${\mathcal{L}}_{\mathbf{c}}+{\mathcal{L}}_{\mathbf{n}}$	${\mathcal{L}}_{\mathit{pretext}}$
IC-index ↑	0.793 (0.006)	0.793 (0.007)	0.793 (0.008)	0.795 (0.007)
IBS ↓	0.340 (0.012)	0.336 (0.014)	0.334 (0.016)	0.333 (0.014)
INBLL ↓	0.104 (0.004)	0.104 (0.004)	0.102 (0.006)	0.101 (0.005)

↑ indicates higher values are better; ↓ indicates lower values are better.

compares the performance of SurvFormer and SSL-SurvFormer on the FLCHAIN [30] dataset, which is a dataset with no missing values. Moreover, the SSL pre-train in SSL-SurvFormer is trained by ${\mathcal{L}}_{pretext}$, as defined in Equation (7). In particular, it is trained on augmented data through mask prediction ${\mathcal{L}}_m$ and reconstruction via ${\mathcal{L}}_c+{\mathcal{L}}_n$. We demonstrate the impact of each loss term in Table 7 and

Table 8. The impact of each component in SSL pre-training on the SUPPORT [32] dataset (missing dataset). The baseline is our SurvFormer model without pre-training with SSL. When remove SSL, we use KNN imputation for missing values.

	Baseline	SSL
	KNN Impute	${\mathcal{L}}_{\mathbf{m}}$	${\mathcal{L}}_{\mathbf{c}}+{\mathcal{L}}_{\mathbf{n}}$	${\mathcal{L}}_{\mathit{pretext}}$
IC-index ↑	0.688 (0.009)	0.695 (0.009)	0.700 (0.005)	0.708 (0.007)
IBS ↓	0.181 (0.006)	0.178 (0.008)	0.176 (0.009)	0.175 (0.004)
INBLL ↓	0.537 (0.015)	0.531 (0.014)	0.532 (0.010)	0.528 (0.011)

↑ indicates higher values are better; ↓ indicates lower values are better.

. We also compare the performance of model with each loss on the dataset with missing values, i.e., the SUPPORT [32] dataset, where we pre-train our model with limited complete data (with augmentation technique) and finetune it with all the available data (missing values included). This ablation study indicates that while pre-training with either corrupted mask prediction ${\mathcal{L}}_m$ or reconstruction ${\mathcal{L}}_c+{\mathcal{L}}_n$ can improve performance, the combination of both losses in ${\mathcal{L}}_{pretext}$ yields the best results. On the FLCHAIN [30] dataset, which is entirely complete, SSL leads to a performance improvement of 0.002, 0.007, and 0.003 in terms of the IC-index, IBS, and INBLL, respectively. Furthermore, when dealing with the missing data in the SUPPORT [32] dataset, pre-training with the combined losses in ${\mathcal{L}}_{pretext}$ results in a performance boost of 0.02, 0.006, and 0.009 for the IC-index, IBS, and INBLL, respectively.

5. Conclusions

This paper introduces SSL-SurvFormer, a novel continuously monotonic Transformer-based model designed for robust survival prediction. Our model incorporates Transformer architecture and a monotonically positive neural network, enabling efficient survival prediction without the need to discretize time intervals. Additionally, we leverage self-supervised learning (SSL) techniques, coupled with our proposed data augmentation, to effectively handle missing values in the data. Our extensive experiments and comprehensive ablation study, conducted on three diverse datasets (METABRIC [31], SUPPORT [32], and FLCHAIN [30]) encompassing scenarios with and without missing features, demonstrate the superior performance of SSL-SurvFormer. It outperforms current state-of-the-art methods, both discrete-time and continuous-time methods, in terms of quantile metrics and overall performance metrics.

In our future research endeavors, we intend to explore various self- and semi-supervised learning techniques. This exploration aims to enhance feature learning, potentially further improving our model’s robustness and performance. Additionally, we plan to investigate the development of loss functions that maintain monotonic constraints to expand its applicability to a broader range of survival prediction tasks.