Pseudo value-based Deep Neural Networks for Multi-state Survival Analysis

Multi-state survival analysis (MSA) is the problem of analyzing time-to-event data using multi-state models (MSM). Multi-state models (Hougaard, 1999) are models of a continuous-time stochastic process that capture the movement of subjects among a finite number of healthy and/or disease states. Thus, multi-state modeling can provide insights into the disease progression by providing a detailed view of disease or recovery trajectory in patients. This helps to predict the probability of future events after a given history and thus, can improve the clinician’s decision-making ability for survival analysis.

Figure 1 shows an example of a multi-state model for breast cancer progression. Here, a patient who is disease-free or had surgery can transition to locoregional relapse or distant relapse before reaching the death state. Multi-state survival analysis deals with the estimations of the multi-state quantities such as (a) state occupation probability (SOP) (the probability that a subject will be in a state $k$ at time $t$,), (b) the transition probability (the probability of transition to a state $k$ at time $t$ from another state $j$ at time $s$), and (c) dynamic SOP - which is the state occupation probabilities at some future time point $t$, given the event history (such as clinical information) is available up to a given time point $s$. A variety of statistical and machine learning approaches have been developed over the years including non-parametric Aalen-Johansen (AJ) estimators (Aalen and Johansen, 1978), Cox-based semi-parametric methods (De Wreede et al., 2010), parametric multi-state methods (Jackson, 2016), and neural network-based methods -SurvNODE- (Groha et al., 2020) to estimate these multi-state quantities. Often, these approaches make strong assumptions such as linearity, proportional hazard, and Markov assumptions for each state or transition, which rarely hold in practice (Meira-Machado et al., 2009). Moreover, many of these methods do not provide subject-specific predictions and do not handle the censoring well. Furthermore, these existing MSA methods cannot obtain subject-specific predictions of transition probability for non-Markov data because finding consistent estimators for non-Markov data has been understudied in the literature. Thus, new approaches that overcome these issues for multi-state survival analysis are in great demand.

We first describe the derivation of pseudo values for multi-state quantities before discussing our proposed pseudo value-based deep neural networks.

Multi-state Survival Quantities: A multi-state process is a continuous time stochastic process $\{X(t),t\in\mathcal{T}\}$, taking values in the (discrete-state) finite state space $S=\{1,2,...,K\}$, where $\mathcal{T}=[0,\tau],\tau<\infty$. MSA deals with the estimation of the multi-state quantities (Spitoni et al., 2018): Transition probability is the probability of transition to a state k at time t from another state j at time s, defined as ${\mathcal{P}}_{jk}(s,t)=P(X(t)=k|X(s)=j,H(s)),\forall{j,k\in S}$ State Occupation Probability (SOP) is the probability that a subject will be in state k at time t, defined as, $\pi_{k}(t)=P(X(t)=k);k\in S$; and Dynamic SOP is the SOP at some future time point t, given the event history is available up to a given time point s, defined as $\pi_{k}(t|s)=P(X(t)=k|H(s))$.

Pseudo values for multi-state quantities: Multi-state survival datasets are subject to censoring, i.e., incomplete information about the stochastic process (for example, event or transition information missing due to loss to follow-up). Therefore, direct modeling of the event time or status with respect to covariates is challenging for censored observations. Inspired by the recent works (Zhao and Feng, 2020; Rahman et al., 2021), we propose to use pseudo values as a substitute for the estimation of subject-specific multi-state model quantities in the presence of censoring. Thus, we estimate the pseudo values for a multi-state quantity of interest as $\hat{y_{i}}(t^{*})=n\hat{y}(t^{*})-(n-1)\hat{y}^{-i}(t^{*})$, where $\hat{y}(t^{*})$ is an estimate of a consistent estimator, based on a $n$ samples, $\hat{y}^{-i}(t^{*})$ is an estimate from the same estimator based on leave-one-out $(n-1)$ samples, obtained by omitting the $i^{th}$ subject.

Consistent pseudo value derivation via Markov assumption testing: Pseudo values for MSA can be derived from an unbiased and consistent estimator such as the AJ estimator. Theoretical analysis of the consistency of the AJ and LMAJ estimators can be found in (Putter and Spitoni, 2018). However, the AJ estimator is inconsistent for non-Markov data and can result in large estimation errors (Titman and Putter, 2020). Thus, recently, researchers have proposed Landmark AJ (LMAJ) (Putter and Spitoni, 2018) as a consistent and robust estimator for estimating the pseudo values for non-Markov data. However, the AJ estimator is known to be more efficient than LMAJ (Titman and Putter, 2020) (when the Markov assumptions are valid), and in any practical scenario, the appropriateness of the Markov assumptions for a specific dataset remains unknown in advance, thus, making it infeasible to use just one estimator for pseudo value estimation. To address this important challenge, we introduce and describe a pseudo value derivation algorithm shown in Algorithm 1 to efficiently derive pseudo values by selecting consistent estimators by testing the underlying Markovian assumptions. Our algorithm takes as input the multi-state survival data and selects a consistent estimator to obtain pseudo values by testing the Markovianity assumptions in the dataset by using statistical significance tests such as Commenges-Andersen (CA) test (Commenges and Andersen, 1995) or log-rank statistic-based tests (Titman and Putter, 2020). CA and log-tank tests use a test statistic (usually $\chi^{2}$ statistic) and its corresponding p-values to identify the violation of Markov assumptions in the data. Note that $\epsilon$ is chosen based on the minimum size of the population in a landmark state. We fix $\epsilon=1$ in our experiments.

Proposed Model: We propose msPseudo - a first-of-its-kind, pseudo value-based deep learning model for multi-state survival analysis. msPseudo is a simple feedforward deep neural network which performs regression analysis to predict the multi-state quantities, such as state occupation probability (SOP), dynamic SOP, and transition probability (TP), using pseudo values as the response variables, given the covariates. msPseudo captures the complex non-linear hidden relationship between the patient’s characteristics, i.e., the baseline covariates and the multi-state model quantities. For an input $n\times p$ matrix of $p$ baseline covariates with $n$ individuals input, msPseudo returns the predictions of a multi-state quantity (SOP, dynamic SOP, or TP). For a multi-state dataset with $K$ states, predicted output of SOP and dynamic SOP for a subject at a prespecified vector of $M$ time points $\mathbf{t}=\{\tau_{1},\tau_{2},...,\tau_{M}\}$ is a $K\times M$ matrix. For the TP prediction task, the output is a $Q\times M$ matrix, where $Q$ is the number of transitions. We used the mean squared error (the mean squared difference between pseudo values (ground truth) and the predicted multi-state quantity) as the loss function ($L$) to train our msPseudo model.

We conducted experiments on both simulated and real-world datasets to answer the following questions: (a) How do our proposed models compare against the existing MSA approaches for predicting multi-state quantities? (b) How well do our proposed models perform under a variety of censoring settings compared to other models?

Simulation datasets: We generated the following four simulation datasets (two Markov and two non-Markov datasets): (1) Time homogeneous Linear Markov Data; (2) Time homogeneous Nonlinear Markov Data; (3) Linear Reversible Non-Markov Data; (4) Nonlinear Reversible Non-Markov Data - to obtain different datasets with varying Markov and linearity assumptions. For each dataset, we simulated 5000 examples with multiple transitions. The Markov datasets have three states and allow only forward transitions. The Non-Markov multi-state datasets consist of four states: states 1–3 are intermediate and interconnected states, and state 4 is an absorbing state, and they allow reverse transitions (Hoff et al., 2019).

Real-world datasets: We used the following publicly available datasets for our experiments: (1) METABRIC (Rueda et al., 2019) dataset contains 1975 breast cancer patient data with multiple transitions and 20 covariates collected over a 360-month study. This multi-state dataset has four states: Surgery, Locoregional Relapse, Distance Relapse, and Death. (2) EBMT (de Wreede et al., 2011) dataset contains 2279 transplantation patient data collected between 1985 and 1998. In this dataset, an alive patient in remission without recovery or adverse event can move to three possible distinct intermediate states, i.e., recovery, adverse event, and co-occurrence of recovery and adverse event, until one of the two absorbing states (death and relapse) is observed.

Censoring settings: We investigate the impact of the higher rate of censoring (75%) on MSA model performances under two censoring settings: incremental censoring (adding censored observations to a fixed number of uncensored observations) and induced censoring (inducing censored observations by flipping the label of transition status of the uncensored observations) settings (Rahman et al., 2021).

Prediction tasks: Given the covariates, we perform regression for estimating the multi-state quantities such as SOP, dynamic SOP, and TP. We compare the performances of the following multi-state models for these prediction tasks: Non-parametric models: AJ estimator (AJ) (Aalen et al., 2008), LMAJ estimator (LMAJ) (Putter and Spitoni, 2018); Parametric models: Weibull parametric model (msWeibull) (Jackson, 2016), linear Pseudo value model (LinearPseudo) (Andersen and Klein, 2007); Semi-parametric model: Multi-state Cox proportional hazard model (msCox) (De Wreede et al., 2010); Deep learning multi-state model: SurvNODE (Groha et al., 2020); Our proposed model: msPseudo

Evaluations: We evaluate the models in terms of integrated Brier score (iBS) (Spitoni et al., 2018) and integrated AUC (iAUC) (Fawcett, 2006). We perform 5 runs of 5-fold cross-validation and report the average of these evaluation metrics. We train our models using Adam optimizer (Kingma and Ba, 2014) to 10000 epochs with an early stopping criterion. Hyperparameter tuning (over batch size, learning rate, drop-out, number of layers, etc.) is performed to choose the best performing deep learning models. A sigmoid activation function is used in the output layer to obtain the multi-state quantities from the predicted pseudo values.

Simulated data: Table 1 shows that our msPseudo performs significantly better than msCox and SurvNODE for SOP prediction in Non-linear Non-Markov dataset in terms of iAUC and iBS metrics. This shows that our models work well with non-Markov assumptions and can capture non-linearity in the data. Our msPseudo outperforms other models in terms of iAUC and is comparable in terms of iBS for the Non-linear Markov dataset. Table 4 shows that our model msPseudo performs significantly better than msCox and SurvNODE in SOP and Dynamic SOP prediction task on reversible Linear Non-Markov data. In the TP prediction task, msPseudo achieves similar or better results on 7 out of 9 transitions compared to the other models. We also graphically show the time-dependent Brier Score comparison in Figure 2 for dynamic SOP prediction on Linear Non-Markov data. This figure demonstrates that msPseudo achieves $\sim$10% improvement over other multi-state models.

Real-world data: The predictive performances for the real-world clinical data: METABRIC and EBMT are shown in Table 2 & 3. Our model, msPseudo, outperforms all other models on METABRIC data for both SOP and Dynamic SOP prediction tasks. msPseudo also obtains lowest iBS for SOP prediction task on EBMT dataset while msCox performs similar or marginally better for Dynamic SOP predictions. In table 3, we observe that msPseudo gives overall better iBS performance, i.e., better prediction on both METABRIC and EBMT datasets. In some cases, our models give comparable performance to the msCox model due to the absence of covariates interaction effect, negligible violation of the proportional hazards and Markov assumptions. However, when averaged over all states and transitions (shown as the column Avg in all tables), we find that our proposed models outperform msCox and other MSA methods.

Various Censoring Settings: Table 5 shows the iAUC results of different multi-state models on incremental and induced censoring settings for SOP prediction in the time-homogeneous nonlinear Markov dataset. From this table, we see that our msPseudo performs significantly better (more efficient in handling censoring using pseudo values) than other models in both incremental and induced censoring settings with a high censoring rate (75% censoring).

Multi-state survival analysis (MSA) is an important yet under-studied problem in time-to-event literature. Finding consistent estimators for non-Markov data is still an open problem in this field. In this paper, we proposed a first-of-its-kind novel pseudo value-based deep learning model, msPseudo, for estimating multi-state survival quantities in the presence of censoring without making any assumption on the underlying multi-state processes. We show that pseudo values can be a replacement for the estimation of multi-state quantities when derived from a consistent estimator. Through empirical experiments on simulated and real datasets, we demonstrated that our proposed models outperform other multi-state survival models under various censoring settings and for both Markov and non-Markov datasets. We believe this work lays the foundation for future investigations on the use of deep models for MSA, including explaining survival predictions and state-specific transition probabilities in real-world datasets.

This work is supported by grant IIS–1948399 from the US National Science Foundation (NSF).