A Deep Variational Approach to Clustering Survival Data

1 INTRODUCTION Survival analysis ( Rodrı́guez , 2007 ; D. G. Altman , 2020 ) has been extensively used in a variety of medical applications to infer a relationship between explanatory variables and a potentially censored survival outcome . The latter indicates the time to a certain event , such as death or cancer recurrence , and is censored when its value is only partially known , e.g . due to withdrawal from the study ( see Appendix A ) . Classical approaches include the Cox proportional hazards ( PH ; Cox ( 1972 ) ) and accelerated failure time ( AFT ) models ( Buckley & James , 1979 ) . Recently , many machine learning techniques have been proposed to learn nonlinear relationships from unstructured data ( Faraggi & Simon , 1995 ; Ranganath et al. , 2016 ; Katzman et al. , 2018 ; Kvamme et al. , 2019 ) . Clustering , on the other hand , serves as a valuable tool in data-driven discovery and subtyping of diseases . Yet a fully unsupervised clustering algorithm does not use , by definition , the survival outcomes to identify clusters . Therefore , there is no guarantee that the discovered subgroups are correlated with patient survival ( Bair & Tibshirani , 2004 ) . For this reason , we focus on a semi-supervised learning approach to cluster survival data that jointly considers explanatory variables and censored outcome as indicators for patient ’ s state . This problem is particularly relevant for precision medicine ( Collins & Varmus , 2015 ) . The identification of such patient subpopulations could , for example , facilitate a better understanding of a disease and a more personalised disease management ( Fenstermacher et al. , 2011 ) . Figure 1 schematically depicts this clustering problem : here , the overall patient population consists of three groups characterised by different associations between the covariates and survival , resulting in disparate clinical conditions . The survival distributions do not need to differ between clusters : compare groups 1 ( ) and 3 ( ) . Clustering of survival data , however , remains an under-explored problem . Only few methods have been proposed in this context and they either have limited scalability in high-dimensional , unstructured data ( Liverani et al. , 2020 ) , or they focus on the discovery of purely outcome-driven clusters ( Chapfuwa et al. , 2020 ; Nagpal et al. , 2021a ) , that is clusters characterised entirely by survival time . The latter might fail in applications where the survival distribution alone is not sufficiently informative to stratify the population ( see Figure 1 ) . For instance , groups of patients characterised by similar survival outcomes might respond very differently to the same treatment ( Tanniou et al. , 2016 ) . To address the issues above , we present 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 . It extends previous variational approaches for unsupervised deep clustering ( Dilokthanakul et al. , 2016 ; Jiang et al. , 2017 ) by incorporating cluster-specific survival models in the generative process . Instead of only focusing on survival , our approach models the heterogeneity in the relationships between the covariates and survival outcome . Our main contributions are as follows : ( i ) We propose a novel , deep probabilistic approach to survival cluster analysis that jointly models the distribution of explanatory variables and censored survival outcomes . ( ii ) We comprehensively compare the clustering and time-to-event prediction performance of VaDeSC to the related work on clustering and mixture models for survival data on a range of synthetic and real-world datasets . In particular , we show that VaDeSC outperforms considered baseline methods in terms of identifying clusters and is comparable in terms of time-to-event predictions . ( iii ) We apply the proposed model to computed tomography imaging data acquired from non-small cell lung cancer patients and assess obtained clustering qualitatively . We demonstrate that VaDeSC discovers clusters associated with well-known clinical patient characteristics , in agreement with previous medical findings . 2 RELATED WORK . Clustering of survival data has been first explored by Bair & Tibshirani ( 2004 ) ( semi-supervised clustering ; SSC ) . The authors propose pre-selecting variables based on univariate Cox regression hazard scores and subsequently performing k-means clustering on the subset of features to discover patient subpopulations . More recently , Ahlqvist et al . ( 2018 ) use Cox regression to explore differences across subgroups of diabetic patients discovered by k-means and hierarchical clustering . In the spirit of the early work by Farewell ( 1982 ) on mixtures of Cox regression models , Mouli et al . ( 2018 ) 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 . Xia et al . ( 2019 ) adopt a multitask learning approach for the outcome-driven clustering of acute coronary syndrome patients . Chapfuwa et al . ( 2020 ) propose a survival cluster analysis ( SCA ) based on a truncated Dirichlet process and neural networks for the encoder and time-to-event prediction model . Somewhat similar techniques have been explored by Nagpal et al . ( 2021a ) who introduce finite Weibull , named deep survival machines ( DSM ) . DSM fits a mixture of survival regression models on the representations learnt by an encoder neural network . From the modelling perspective , the above approaches focus on outcome-driven clustering , i.e . they recover clusters entirely characterised by different survival distributions . On the contrary , we aim to model cluster-specific associations between covariates and survival times to discover clusters characterised not only by disparate risk but also by different survival generative mechanisms ( see Figure 1 ) . In the concurrent work , Nagpal et al . ( 2021b ) introduce deep Cox mixtures ( DCM ) jointly fitting a VAE and a mixture of Cox regressions . As opposed to our approach , DCM does not specify a generative model and its loss is derived empirically combining the VAE loss with the likelihood of survival times . As a result , DCM does not formally maximise the joint likelihood of the data . On the contrary , we focus on a probabilistic method with an interpretable generative process from which an ELBO of the joint likelihood can be derived formally . The approach by Liverani et al . ( 2020 ) is , on the other hand , the most closely related to ours . The authors propose a clustering method for collinear survival data based on the profile regression ( PR ; Molitor et al . ( 2010 ) ) . In particular , they introduce a Dirichlet process mixture model with clusterspecific parameters for the Weibull distribution . However , their method is unable to tackle highdimensional unstructured data , since none of its components are parameterised by neural networks . This prevents its usage on real-world complex datasets , such as medical imaging ( Haarburger et al. , 2019 ; Bello et al. , 2019 ) . Table 1 compares our and related methods w.r.t . a range of properties . For an overview of other lines of work and a detailed comparison see Appendices B and C. 3 METHOD In this section , we present VaDeSC — a novel variational deep survival clustering model . Figure 2 provides a schematic summary of the proposed approach : the input vector x is mapped to a latent representation z using a VAE regularised by a Gaussian mixture prior . The survival density function is given by a mixture of Weibull distributions with cluster-specific parameters β . The parameters of the Gaussian mixture and the Weibull distributions are then optimised jointly using both the explanatory input variables and survival outcomes . Preliminaries We consider the following setting : let D = { ( xi , δi , ti ) } Ni=1 be a dataset consisting ofN threetuples , one for each patient . Herein , xi denotes the explanatory variables , or features . δi is the censoring indicator : δi = 0 if the survival time of the i-th patient was censored , and δi = 1 otherwise . Finally , ti is the potentially censored survival time . A maximum likelihood approach to survival analysis seeks to model the survival distribution S ( t|x ) = P ( T > t|x ) ( Cox , 1972 ) . Two key challenges of survival analysis are ( i ) the censoring of observed survival times and ( ii ) a potentially complex nonlinear relationship between x and t. When clustering survival data ( Liverani et al. , 2020 ) , we additionally consider a latent cluster assignment variable ci ∈ { 1 , ... , K } unobserved at training time . Here , K is the total number of clusters . The problem then is twofold : ( i ) to infer unobserved cluster assignments and ( ii ) model the survival distribution given xi and ci . c z π µ Σ β γ xtδ k N K time . Figure 3 summarises the random process graphically . Survival Model Above , p ( t|z , c ) refers to the cluster-specific survival model . We follow an approach similar to Ranganath et al . ( 2016 ) and Liverani et al . ( 2020 ) ; in particular , we as- sume that given z and c , the uncensored survival time follows the Weibull distribution given by Weibull ( softplus ( z > βc ) , k ) , where softplus ( x ) = log ( 1 + exp ( x ) ) ; k is the shape parameter ; and βc are cluster-specific survival parameters . Note that we omitted the bias term βc,0 for the sake of brevity . Observe that softplus ( z > βc ) corresponds to the scale parameter of the Weibull distribution . We assume that the shape parameter k is global ; however , an adaptation to cluster-specific parameters , as proposed by Liverani et al . ( 2020 ) , is straightforward . The Weibull distribution with scale λ and shape k has the probability density function given by f ( x ; λ , k ) = kλ ( x λ ) k−1 exp ( − ( x λ ) k ) , for x ≥ 0 . Consequently , adjusting for right-censoring yields the following distribution : p ( t|z , c ; β , k ) = f ( t ; λzc , k ) δS ( t|z , c ) 1−δ = [ k λzc ( t λzc ) k−1 exp ( − ( t λzc ) k ) ] δ [ exp ( − ( t λzc ) k ) ] 1−δ , ( 1 ) where β = { β1 , . . . , βK } ; λzc = softplus ( z > βc ) ; and S ( t|z , c ) = ∫∞ t=t f ( t ; λzc , k ) is the survival function . Henceforth , we will use p ( t|z , c ) as a shorthand notation for p ( t|z , c ; β , k ) . In this paper , we only consider right-censoring ; however , the proposed model can be extended to tackle other forms of censoring . Joint Probability Distribution Assuming the generative process described above , the joint probability of x and t can be written as p ( x , t ) = ∫ z ∑K c=1 p ( x , t , z , c ) = ∫ z ∑K c=1 p ( x|t , z , c ) p ( t , z , c ) . It is important to note that x and t are independent given z , so are x and c. Hence , we can rewrite the joint probability of the data , also referred to as the likelihood function , given the parameters π , µ , Σ , γ , β , k as p ( x , t ; π , µ , Σ , γ , β , k ) = ∫ z K∑ c=1 p ( x|z ; γ ) p ( t|z , c ; β , k ) p ( z|c ; µ , Σ ) p ( c ; π ) , ( 2 ) where µ = { µ1 , ... , µK } , Σ = { Σ1 , ... , ΣK } , and β = { β1 , ... , βK } . Evidence Lower Bound Given the data generating assumptions stated before , the objective is to infer the parameters π , µ , Σ , γ , and β which better explain the covariates and survival outcomes { xi , ti } Ni=1 . Since the likelihood function in Equation 2 is intractable , we maximise a lower bound of the log marginal probability of the data : log p ( x , t ; π , µ , Σ , γ , β , k ) ≥ Eq ( z , c|x , t ) log [ p ( x|z ; γ ) p ( t|z , c ; β , k ) p ( z|c ; µ , Σ ) p ( c ; π ) q ( z , c|x , t ) ] . ( 3 ) We approximate the probability of the latent variables z and c given the observations with a variational distribution q ( z , c|x , t ) = q ( z|x ) q ( c|z , t ) , where the first term is the encoder parameterised by a neural network . The second term is equal to the true probability p ( c|z , t ) : q ( c|z , t ) = p ( c|z , t ) = p ( z , t|c ) p ( c ) ∑K c=1 p ( z , t|c ) p ( c ) = p ( t|z , c ) p ( z|c ) p ( c ) ∑K c=1 p ( t|z , c ) p ( z|c ) p ( c ) . ( 4 ) Thus , the evidence lower bound ( ELBO ) can be written as L ( x , t ) = Eq ( z|x ) p ( c|z , t ) log p ( x|z ; γ ) + Eq ( z|x ) p ( c|z , t ) log p ( t|z , c ; β , k ) −DKL ( q ( z , c|x , t ) ‖ p ( z , c ; µ , Σ , π ) ) . ( 5 ) Of particular interest is the second term which encourages the model to maximise the probability of observing the given survival outcome t under the variational distribution of the latent embeddings and cluster assignments q ( z , c|x , t ) . It can be then seen as a mixture of survival distributions , each one assigned to one cluster . The ELBO can be approximated using the stochastic gradient variational Bayes ( SGVB ) estimator ( Kingma & Welling , 2014 ) to be maximised efficiently using stochastic gradient descent . For the complete derivation we refer to Appendix D. Missing Survival Time The hard cluster assignments can be computed from the distribution p ( c|z , t ) of Equation 4 . However , the survival times may not be observable at test-time ; whereas our derivation of the distribution p ( c|z , t ) depends on p ( t|z , c ) . Therefore , when the survival time of an individual is unknown , using the Bayes ’ rule we instead compute p ( c|z ) = p ( z|c ) p ( c ) ∑K c=1 p ( z|c ) p ( c ) .

Summary: This work tackles the problem of clustering in the context of survival data using a generative model. A variational autoencoder is used for modelling the data, while the latent representation is leveraged to model the survival outcome conditionned on the assigned cluster following a Weibull distribution. This approach allows to leverage both the survival outcome and the covariates for clustering. Moreover, the latent state associated to each cluster is then used for interpretability of the observed cluster. Synthetic and real world datasets demonstrate the competitiveness of the method on both discriminative and clustering performances.