Variational Deep Survival Machines: Survival Regression with Censored Outcomes
Qinxin Wang Jiayuan Huang Junhui Li Jiaming Liu


Understanding and Improving CNNs with Complex Structure Tensor: A Biometrics Study
    Josef Bigun, Fernando Alonso-Fernandez
    April 29, 2024
==================================================================================


  Survival regression aims to predict the time when an event of interest will take place, typically a death or a failure. A fully parametric method <cit.> is proposed to estimate the survival function as a mixture of individual parametric distributions in the presence of censoring.
  In this paper, We present a novel method to predict the survival time by better clustering the survival data and combine primitive distributions. We propose two variants of variational auto-encoder (VAE), discrete and continuous, to generate the latent variables for clustering input covariates. The model is trained end to end by jointly optimizing the VAE loss and regression loss.
  Thorough experiments on dataset SUPPORT and FLCHAIN show that our method can effectively improve the clustering result and reach competitive scores with previous methods. We demonstrate the superior result of our model prediction in the long-term. Our code is available at <https://github.com/qinzzz/auton-survival-785>
  







§ INTRODUCTION

Survival Analysis is a machine learning problem for analyzing the expected duration of time until an event occurs, typically the death of a patient in medical treatment. It has been extensively used in many medical applications. Formally, the problem can be described as a regression problem to estimate the conditional survival curve P(T>t|X), which T is the event of interest and X is covariate. Given X, the survival function predicts the probability of an event T that happens after time t. In practice, the presence of censoring, where the outcome measurement is only partially known during the whole experiment, makes the survival analysis more difficult. For example, if the study ends before the event happens, we would not know the exact outcome for this individual's data collection. The censored observations cannot be simply discarded because they contain crucial information <cit.>. <cit.> has also shown the importance of censored instances for deriving the accurate estimation.

Traditional statistical approaches are usually non-parametric or semi-parametric for survival regression. However, non-parametric methods will face curse of dimensional problem. Meanwhile semi-parametric methods depend on strong modelling assumptions which are usually unrealistic. Recently, some fully parametric deep learning approaches have been proposed for survival regression. For example, Deep Survival Machine<cit.> estimates the survival function as a mixture of individual parametric distributions in the presence of censoring and it doesn't need strong assumptions. <cit.> generalizes DSM by assuming within each latent group the proportional hazards assumption holds.

Clustering is also a valuable tool in survival analysis. Patients can usually be clustered into different latent groups and each of these groups would have an individual distribution. <cit.> presents a novel method for clustering survival data — variational
deep survival clustering (VaDeSC) that discovers groups of patients characterised by different generative mechanisms of survival outcome.

In this project, we propose two models named VDSM-cat and VDSM-clus, which are novel approaches combine DSM and VAE to both cluster data and predict survival time. The input data will first pass through VAE's encoder and we will get latent variables. These latent variables will be transformed as weights of DSM's mixed individual distributions. Compared to original DSM, our model achieve better results in the long term, which suggests the embedded VAE could give a better cluster outcome and is helpful for final prediction.

Our main contributions are as follows:

(1) We propose two new methods VDSM-cat and VDSM-clus, which corresponds to categorical distribution and Gaussian Mixture distribution for the latent variable respectively. 

(2) We compare our approaches with DSM on SUPPORT and FLCHAIN dataset and find our methods could achieve superior results in long time prediction.






















§ RELATED WORKS
 


Time-to-event analysis with censoring is an important problem which could be applied into many domains including bio-statistics and health informatics<cit.>, actuarial sciences<cit.> and econometrics<cit.>. It has also been increasingly popular in engineering, including applications to predict maintenance time of equipments. <cit.> proposed a Bayesian hierarchical model to predict the end of life (EoL) and end of discharge (EoD) of Li-ion batteries using load profiles and discharge data. <cit.> used survival regression model to track the degradation of the system and plan optimal maintenance strategy.

Research in survival analysis pays more attention to deep learning approaches in recent years <cit.> <cit.>. <cit.> proposed a fully parametric model named Deephit to predict survival outcomes in competing risks situation. However, their architecture could only predict failure times over a discrete set of fixed size. When survival horizons become long, this approach would require large number of parameters to be learnt, which can be very impractical. Another limitation of this approach is that it can be more sensitive to events at shorter horizons and it will not model long term horizon well. <cit.> propose using black box optimization to adaptively select the best model for a given event horizon.
Deep Survival Machine (DSM) <cit.> is proposed as a fully-parametric approach to estimate time-to-event in the presence of censoring and competing risks. 
In DSM, the survival function is estimated as a mixture of individual parametric survival distributions, where the distributions parameters are learnt by training a Multilayer Perceptron. The strong assumptions of proportional hazards is not required. In this way, it enables the model to learn a rich distributed representation of the input covariates. 
<cit.> proposed a new approach based on learning mixtures of Cox regressions to model individual survival distributions. It involves estimating hazard ratios within latent clusters followed non-parametric estimation of the baseline survival rates.

Semi-supervised clustering for survival data has been first studied in <cit.>. It proposed pre-selecting variables based on univariate Cox regression hazard scores and then performing k-means clustering on the subset of features to discover patient subpopulations. <cit.> uses Cox regression to explore differences across subgroups of diabetic patients discovered by hierarchical clustering. <cit.> propose a deep clustering approach to differentiate between long- and short- term survivors based on a modified Kuiper statistic in the absence of end-of-life signals.<cit.> used a multitask learning approach for the outcome-driven clustering of acute syndrome patients. In <cit.> and <cit.>, DSM and DCM set a mixture of survival models through representations learnt by an encoder neural network.Recently, Variational AutoEndoers (VAE) has also been involved in time-to-event analysis and clustering. <cit.> introduces a varitional time-to-event prediction model, which uses Variational Survival Inference (VSI) to predict the time-to-event distribution without the need to specify a parametric form for the baseline distribution. In this way, they avoid the restrictive assumptions in classical survival analysis models. <cit.> proposed a semi-supervised approach to cluster survival data. It outperforms baseline methods through clustering the latent variables and is comparable in terms of time-to-event predictions.



























§ OUR METHOD



 §.§ Problem Definition

We consider a dataset of right censored observations 𝒟 = {(x_i, δ_i, u_i)}^N_i=1, where x_i are the covariates of an individual i, δ_i is an indicator of whether an event occurred or not and u_i is either the time of event or censoring as indicated by δ_i. Our goal is to learn a survival function S(t|x) = ℙ(T > t| X = x) from the input data. We denote the uncensored subset (δ=1) of data as 𝒟_𝒰 and the censored subset (δ=0) as 𝒟_𝒞. 




 §.§ Deep Survival Machine

In this section, we briefly describe the model proposed by Deep Survival Machine (DSM) <cit.>. Given the above variable x and z, the survival time t can be modeled as 

    p(T|X) = ∑_Z p(T|X,Z) p(Z|X)


The model aims to maximize the probability over the entire dataset. Considering both censored and uncensored data, we model the probability by 

    p(T|X,Z) = ∏_i=1^|D| p(T=t_i|Z)^δ_i p(T>t_i|Z)^1-δ_i

And the loss can be expressed as

    ℒ_DSM = ln p(T|X,Z) = ∑_i=1^|D|δ_i ln p(T=t_i|Z) + ∑_i=1^|D| (1-δ_i) ln p(T>t_i|Z)

We denote it as combining two separate loss for uncensored and censored data:

    ℒ_DSM = ℒ_U + ℒ_C


The log likelihood of the uncensored dataset can be denoted as


    ℒ_U = ln P(𝒟_𝒰;Θ)
           = ln∏_i=1^|D_U|ℙ (T=t_i |X = x_i;Θ) 
       = ∑_i=1^|D_U|ln (∑_k=1^Kℙ (T=t_i |Z) ℙ(Z | X=x_i;w)) 
       = ∑_i=1^|D_U|ln (𝔼_Z ∼ p(z|x_i,w)ℙ (T=t_i |Z)) 
       ≥∑_i=1^|D_U| (𝔼_Z ∼ p(z|x_i,w) [lnℙ (T=t_i |Z)])


For the censored data, we  maximize the log likelihood by computing the probability of event happening after t_i.

    ℒ_C = ln P(𝒟_𝒞;Θ)
           = ln∏_i=1^|D_C|ℙ (T>t_i |X = x_i;Θ) 
       ≥∑_i=1^|D_C| (𝔼_Z ∼ p(z|x_i,w) [lnℙ (T>t_i |Z)])










 §.§ Proposed Model

As in DSM, we choose to model the conditional distribution of ℙ(T|X) as a mixture over K well-defined distributions. 
We refer to them as primitive distributions. To satisfy certain properties of survival time prediction, DSM experimented with two types of distributions, the Weibull and the Log-Normal distribution (see Table <ref>). For each primitive distribution, we need a independent set of parameters {η_k, β_k}. The final individual survival distribution ℙ(T|X) is a weighted average over these K distributions.

In DSM paper, they use a linear transform x^T W to get a z with K classes. Though simple and intuitive, this classification method fails to learn a good weight over K categories in a unsupervised setting.
We propose a variant of DSM<cit.>, Variational Deep Survival Machine (VDSM), by introducing Variational AutoEncoder<cit.> into the stage of sampling z from latent space given input x. Our goal is to learn a better cluster assignment with VAE models.
Figure <ref> shows an overview of the model structure. The input X is a set of features an individual patient. We use a VAE encoder to generate the latent variable Z. The Z constains of K classes, where each cluster has an independent underlying survival distribution. The final predictied distribution is ℙ(T|X) is a weighted average over the K distributions.

In the VAE encoder and decoder module, we experimented with two different VAE for classification and clustering respectively. We refer to them as VDSM-cat and VDSM-clus.










  
Variational Deep Survival Machine with categorical VAE (VDSM-cat)

We introducing Categorical Variational AutoEncoder proposed in <cit.> into the stage of sampling z from latent space given input x, which is illustrated in Figure <ref>.
Consider a latent space with K clusters { 1, ..., K }, z is sampled from a categorical distribution z ∼ p(z) = Cat(K; π).


We know that categorical distribution is discrete and not differentiable. To make the model amenable to gradient based optimization, we use the Gumbel-Max trick<cit.> to generate samples which can be used in backpropagation.The Gumbel-Max trick is an efficient method to generate samples from a categorical distribution:






    z = one_hot(min_i[g_i+logπ_i])


where g_i is the samples drawn from Gumbel(0,1), π_i is the ith-class possibility of the categorical distribution. To cover the gap that argmax is not differentiable, the Gumbel-Softmax<cit.> is proposed as a continuous approximation:


    y_i = exp((log(π_i)+g_i)/ τ)/∑_j=1^k exp((log(π_j)+g_j)/ τ)      for i = 1,...,k.




The loss for this categorical VAE would be the Kullback-Leibler divergence between generated distribution of z with the expected categorical distribution p(z), plus the reconstruction loss of X.

    ℒ_VAE = 𝕂𝕃(p(z)||q(z|x)) + 𝔼_z ∼ q(z|x) [log p(x|z)]




  
Variational Deep Survival Machine with generative clustering (VDSM-clus) 



The assumption behind categorical VAE is that x belongs to one of the K categories. However, the latent variable of VAE assumes a Gaussian mixture as prior. We believe the latent cluster assumption of Gaussian Mixture Model (GMM) can be applied to extend it with a mixture of Gaussian distributions. 
Here we first describe the process of the generative model. A cluster c is sampled from a categorical distribution c ∼ p(c) = Cat(π). Every cluster is a Gaussian distribution with mean μ_c and variance σ_c. Then a continuous latent embedding z is sampled from the Gaussian Distribution. The input covariant x is generated from a distribution conditioned on z. Finally, the survival time t depends on z and c, i.e., t ∼ p(t|z,c).
The main difference here is that z conforms to a normal distribution depends on cluster c.

The generative process above can be defined by the following,


    p(x, z, c)=p(x|z)p(z|c)p(c)

where

    p(c)    = Cat(c, π)
    
        p(z|c)    = 𝒩(z|μ_c, σ_c^2)
    
        p(x|z)    = 𝒩(x|μ_x, σ_x^2)


<cit.> proposed an unsupervised generative clustering framework VaDE that combines VAE and GMM together, as is shown in Figure <ref>. Inspired by their method, we optimized the VAE by maximizing the evidence lower bound (ELBO) using the SGVB estimator and the reparameterization trick.
We borrowed the learning objective derived in the paper, which is defined as:

    ℒ_VAE   = ℒ_ELBO(x) 
        = E_q(z,c|x) [ logp(x,z,c)/q(z,c|x) ] 
       = E_q(z,c|x) [log p(x|z)] - D_KL(q(z,c|x)||p(z|c))

The first term in Equation <ref> is the reconstruction term between encoder input and decoder output. The second term is the Kullback-Leibler divergence between the Mixture-of-Gaussian prior p(z|c) to the posterior q(z,c|x). The training procedure aims to maximize the ELBO with regard to parameters. The cluster assignment is obtained by the approximation of q(c|x):

    q(c|x) = p(c|z) = p(c) p(z|c)/∑_c'=1^K p(c') p(z|c')




Finally, the combined loss of our two proposed models are:

    ℒ   = ℒ_DSM + ℒ_VAE












§ EXPERIMENTAL SETUP




 §.§ Baseline










  
Deep Survival Machine
Deep Survival Machine is a novel approach to estimate time-to-event in the presence of censoring. By leveraging a hierarchical graphical model parameterized by
neural networks, the model learns distributional representations of the input covariates and mitigate existing challenges in survival regression. DSM estimates the conditional survival function
𝕊(·| X) as a mixture of individual parametric survival distributions without a strong assumptions of proportional hazards, and enables learning with time-varying risks. In this way, DSM allows for learning of rich distributed representations of the input covariates, helping knowledge transfer across multiple competing risks.












 


 §.§ Dataset




  
SUPPORT (Study to understand prognoses and preferences for outcomes and risks of
 <cit.>):
This Dataset is composed of 9,105 terminally ill patients on life support. The median of studied patients' survival time was 58 days. A majority with 79% population were recorded as ‘White’, while the rest were coded as ‘Black’, ‘Hispanic’ and ‘Asian'.




  
FLCHAIN (Assay of Serum Free Light Chain): This is a public dataset developed by <cit.> aiming to study the relationship between serum free light chain and mortality. It includes covariates like age, gender, serum creatinine and presence of monoclonal gammapothy. We used a subset of this dataset, where all the individuals with missing covariates are removed. The remaining dataset consists of 6,524 individuals.










 §.§ Experiment Metrics

We evaluate the baseline model DSM, the proposed model VDSM-cat and VDSM-clus with two popular metrics in survival analysis literature. Similar to DSM settings <cit.>, we evaluate these two metrics with different truncation event horizon quantiles of 25%, 50% and 75% to measure how good the model can be at showing the risks at different scale of time.



  
Concordance index C^td It is used to evaluate the proportion of all comparable pairs in which the predictions and outcomes are concordant. More specifically, Concordance intuitively means that two samples were ordered correctly by the model. If the one with a higher estimated risk score has a shorter actual survival time, we can say the two samples are concordant.



  
ROC-AUC In survival analysis, the area under receiver operating characteristic (ROC) curve is used by defining sensitivity and specificity as time-dependent measures. All individuals that experienced an event prior to or at time t are Cumulative cases and dynamic controls are those individuals experienced an event after the time.




 §.§ Hyper-parameter Searching

To achieve the best performance, we use grid search to find the best hyper-parameter settings for the DSM model. As reported in the DSM paper <cit.>, we evaluate the model with the following parameter settings: the number of clusters is chosen from {4,6,8}, the discounting factor is chosen from {0.5,0.75,1}, and the learning rate of optimizer Adam is chosen from {1e-3, 1e-4}. The best performance is achieved with 4 clusters, 0.5 as the discounting factor and 1e-4 as the learning rate.

For our proposed VDSM-cat and VDSM-clus, we research the hyper-parameter settings. VAEs are added to the DSM to get better clustering representation and therefore it may behave differently with different number of clusters. 

Our experiment shows that the VDSM-clus outperforms the VDSM-cat with every cluster number in {4,6,8} according to the Concordance index and ROC-AUC. For SUPPORT dataset, 10 clusters gives the best result and for FLCHAIN dataset, 6 clusters shows the best performance.









§ RESULTS AND DISCUSSION



 §.§ Experiment Results

We run Deep Survival Machine (DSM) <cit.> and our proposed models on SUPPORT and FLCHAIN dataset. 
Our code for DSM is adapted from the open-source framework Auton-Survival <cit.>.
The results are shown in Table <ref>, <ref>, <ref>, and <ref>.




 §.§ Clustering Results

We run multiple experiments on SUPPORT and FLCHAIN dataset to collect results. First, we generate the latent variable after training the model. Since the latent variable represents a distribution over latent clusters, we choose the category with the largest probability as the label and draw a plot to visualize the clustering result with the t-SNE algorithm.

Figure <ref> and Figure <ref> show the clustering result on SUPPORT and FLCHAIN dataset, respectively. 

Compared to DSM, VDSM-clus has a better clustering result, which is not surprising because the reconstruction loss forces the VAE to learn the latent embedding better.

It should be noted that VDSM-cat does not show better clustering result. In categorical VAE, the latent variable consists of latent dimension N and category dimension K. In our implementation, we set K as the number of latent clusters in DSM, and regard N as a multi-nominal distribution of with N independent trails. 

To be specific, the probability of selecting category c is P(c |x, n=1,2,...,N)/∑_c' P(c'|x, n=1,2,...,N). Since we assume the N latent dimensions are independent, we have P(c|x) = P(c|x, n=1,2,...,N) = ∏_i^N P(c|x,n=i).
This may not be the best way to handle the latent dimension in categorical VAE and is probably the reason why the clustering result of VDSM-cat is not as expected.









 §.§ Better Prediction at Long-time Scale

We compare the results between our models and DSM based on the Concordance Index. The Concordance Index is a commonly used metric in survival analysis, and it evaluate the probability of the pair risk of different patient, at the time one event has already occurred to one of the patients, which means, ideally, that patient should be assigned with a higher risk. The bigger this probability is, the more consistent the prediction is with the truth.

It shows the potential of VAE to capture the long-time-dependent features with the latent embedding of inputs. Another possible explanation is that the latent representation clustering is more important to survival prediction in a long time scale. We are still working on this to prove our assumptions. The ROC-AUC also shows similar results.



 §.§ Effect of Clustering on Time-to-event Prediction

We also found that different quality of clustering does not affect the time-to-event prediction much. According Figure <ref> and Figure <ref>, on both SUPPORT and FLCHAIN datasets the VDSM-clus has a better result than the VDSM-cat. However, experiment results show that there is no much difference between these two models.

Manduchi et al.<cit.> also show similar results by comparing different clustering method. Their proposed model is better at clustering but only stay at competitive with other state-of-the-art survival models. Perhaps the clustering is not the best way to represent the input data.




















§ CONCLUSION & FUTURE WORK


We proposed to incorporate  Variational AutoEncoder to the Deep Survival Machines, which is a semi-supervised deep probabilistic model for clustering survival data. By experimentation, we demonstrated that the VAE model leads to a better result in clustering patients to different latent groups than the naive approach in DSM architecture. Furthermore, the grouping information allows us to get more accurate survival time prediction for patients comparing with state-of-art models. Interesting future work directions include trying other variational clustering methods for survival data, e.g. Random Forest, a Multiple Task Variational Autoencoder to better study the latent relationship between features. Furthermore, the model can be modified to be able to predict survival time based on image input. 



§ DIVISION OF WORK

Jiaming Liu -  Jiaming co-researched the survival models, collected references about VAE and other deep learning models for survival regression and took part in the design of VAE-cat and VAE-clus. He pre-processed MNIST to do some early experiments and implemented several visualization tools to help analyze model results. He also wrote Introduction and Related Works part in this report.

Jiayuan Huang - Jiayuan co-researched the survival models used in survival analysis, and designed the method of incorporating categorical into DSM. He implemented the proposed model VDSM-cat and conducted experiments on two datasets to collect results. He also wrote the Result and Discussion part in this report.

Junhui Li - Junhui co-researched the survival models used in survival used in survival analysis. She transferred existing VAE methods to Pytorch framework for experiment.She also researched into potential improvement methods and wrote the Dataset, Conclusion and Future Work part in this report. 


Qinxin Wang - Qinxin co-researched the survival models (DSM and DCM), and worked on the mathematical formulation of our new methods. She surveyed multiple VAEs for clustering, implemented the proposed model VDSM-clus, and conducted experiments. She wrote the Method part and also drew model diagrams in this report. 


























































































































plain



