A D EEP V ARIATIONAL A PPROACH TO C LUSTERING S URVIVAL D ATA
Anonymous authors
Paper under double-blind review
A BSTRACT
In this work, we study the problem of clustering survival data — a challenging and so far under-explored task.
We introduce a novel semi-supervised probabilistic approach to cluster survival data by leveraging recent advances in stochastic gradient variational inference.
In contrast to previous work, our proposed method employs a deep generative model to uncover the underlying distribution of both the explanatory variables and censored survival times.
We compare our model to the related work on clustering and mixture models for survival data in comprehensive experiments on a wide range of synthetic, semi-synthetic, and real-world datasets, including medical imaging data.
Our method performs better at identifying clusters and is competitive at predicting survival times.
Relying on novel generative assumptions, the proposed model offers a holistic perspective on clustering survival data and holds a promise of discovering subpopulations whose survival is regulated by different generative mechanisms.
I NTRODUCTION
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 deﬁnition, 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 identiﬁcation 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 ( (cid:4) ) and 3 ( (cid:4) ).
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.
[Table]
The latter might fail in applications where the survival distribution alone is not sufﬁciently 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) [Table]
To address the issues above
[Table]
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-speciﬁc 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 ﬁndings.
R ELATED W ORK
Clustering of survival data has been ﬁrst 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 modiﬁed 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 ﬁnite Weibull, named deep survival machines (DSM).
DSM ﬁts 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-speciﬁc 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 ﬁtting 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 proﬁle regression (PR;
Molitor et al. (2010)).
In particular, they introduce a Dirichlet process mixture model with clusterspeciﬁc 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.
M ETHOD
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-speciﬁc 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 = { ( x i , δ i , t i ) }
Ni =1 be a dataset consisting of N threetuples, one for each patient.
Herein, x i denotes the explanatory variables, or features.
δ i is the censoring indicator: δ i =
if the survival time of the i -th patient was censored, and δ i = 1 otherwise.
Finally, t i 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 c i ∈ { 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 x i and c i .
[Figure]
Generative Model Following the problem deﬁnition above, we assume that data are generated from a random process consisting of the following steps.
First, a cluster assignment c ∈ { 1 , . . .
, K } is sampled from a categorical distribution: c ∼ p ( c ; π ) =
π c .
Then, a continuous latent embedding, z ∈ R J , is sampled from a Gaussian distribution, whose mean and variance depend on the sampled cluster c : z ∼ p ( z | c ; { µ 1 , ..., µ K } , { Σ 1 , ..., Σ K } ) =
N ( z ; µ c , Σ c ) .
The explanatory variables x are generated from a distribution conditioned on z : x ∼ p ( x | z ; γ ) , where p ( x | z ; γ ) =
Bernoulli ( x ; µ γ ) for binary-valued features and N (cid:0) x ; µ γ , diag (cid:0) σ 2 γ (cid:1)(cid:1) for real-valued features.
Herein, µ γ and σ 2 γ are produced by f ( z ; γ ) – a decoder neural network parameterised by γ .
Finally, the survival time t depends on the cluster assignment c , latent vector z , and censoring indicator δ , i.e. t ∼ p ( t | z , c ) .
Similarly to conventional survival analysis, we assume non-informative censoring (Rodr´ıguez, 2007), i.e. that the censoring time is independent of the uncensored survival time.
Figure 3 summarises the random process graphically.
Survival Model Above, p ( t | z , c ) refers to the cluster-speciﬁc survival model.
We follow an approach similar to Ranganath et al.
(2016) and Liverani et al.
; in particular, we as- sume that given z and c , the uncensored survival time follows the Weibull distribution given by
Weibull (cid:0) softplus (cid:0) z
(cid:62) β c (cid:1) , k (cid:1) , where softplus ( x ) = log (1 + exp( x )) ; k is the shape parameter; and β c are cluster-speciﬁc survival parameters.
Note that we omitted the bias term β c, 0 for the sake of brevity.
Observe that softplus (cid:0) z (cid:62) β c (cid:1) corresponds to the scale parameter of the Weibull distribution.
We assume that the shape parameter k is global; however, an adaptation to cluster-speciﬁc 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λ (cid:0) xλ (cid:1) k − 1 exp (cid:16) − (cid:0) xλ (cid:1) k (cid:17) , for x ≥ 0 .
Consequently, adjusting for right-censoring yields the following distribution:
[Equation]
where β = { β 1 , . .
. , β K } ; λ z c = softplus (cid:0) z (cid:62) β c (cid:1) ; and S ( t | z , c ) = (cid:82)
∞ t = t f
( t ; λ z c , 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
[Equation]
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
[Equation]
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 { x i , t i } Ni =1 .
Since the likelihood function in Equation 2 is intractable, we maximise a lower bound of the log marginal probability of the data:
[Equation]
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 ﬁrst term is the encoder parameterised by a neural network.
The second term is equal to the true probability p ( c | z , t ) :
[Equation]
Thus, the evidence lower bound (ELBO) can be written as
[Equation]
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 efﬁciently 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
[Equation]
E XPERIMENTAL S ETUP
Datasets We evaluate VaDeSC on a range of synthetic, semi-synthetic (survMNIST; P¨olsterl (2019)), and real-world survival datasets with varying numbers of data points, explanatory variables, and fractions of censored observations (see Table 2).
In particular, real-world clinical datasets include two benchmarks common in the survival analysis literature, namely SUPPORT (Knaus et al., 1995) and FLChain (Kyle et al., 2006; Dispenzieri et al., 2012); an observational cohort of pediatric patients undergoing chronic hemodialysis (Hemodialysis; Gotta et al.
; an observational cohort of high-grade glioma patients (HGG); and an aggregation of several computed tomography (CT) image datasets acquired from patients diagnosed with non-small cell lung cancer (NSCLC;
Aerts et al.
; Bakr et al.
; Clark et al.
Detailed description of the datasets and preprocessing can be found in Appendices E and G. For (semi-)synthetic data, we focus on the clustering performance of the considered methods; whereas for real-world data, where the true cluster structure is unknown, we compare time-to-event predictions.
In addition, we provide an in-depth cluster analysis for the NSCLC (see Section 5.3) and
Hemodialysis (see Appendix H.8) datasets.
[Table]
Baselines & Ablations We compare our method to several well-established baselines: the semisupervised clustering (SSC; Bair & Tibshirani (2004)), survival cluster analysis (Chapfuwa et al., 2020), and deep survival machines (Nagpal et al., 2021a).
For the sake of fair comparison, in
SCA we truncate the Dirichlet process at the true number of clusters, if known.
For all neuralnetwork-based techniques, we use the same encoder architectures and numbers of latent dimensions.
Although the proﬁle regression approach of Liverani et al.
(2020) is closely related to ours, it is not scalable to large unstructured datasets, such as survMNIST and NSCLC, since it relies on MCMC methods for Bayesian inference and is not parameterised by neural networks.
Therefore, a fullscale comparison is impossible due to computational limitations.
Appendix H.3 contains a ‘downscaled’ experiment with the proﬁle regression on synthetic data.
Additionally, we consider k -means and regularised Cox PH and Weibull AFT models (Simon et al., 2011) as na ¨ ıve baselines.
For the
VaDeSC, we perform several ablations: ( i ) removing the Gaussian mixture prior and performing post hoc k -means clustering on latent representations learnt by a VAE with an auxiliary Weibull survival loss term, which is similar to the deep survival analysis (DSA; Ranganath et al.
combined with k -means; ( ii ) training a completely unsupervised version without modelling the survival, which is similar to VaDE (Jiang et al., 2017); and ( iii ) predicting cluster assignments when the survival time is unobserved.
Appendix G contains further implementation details.
Evaluation We evaluate the clustering performance of models, when possible, in terms of accuracy (ACC), normalised mutual information (NMI), and adjusted Rand index (ARI).
For the timeto-event predictions, we evaluate the ability of methods to rank individuals by their risk using the concordance index (CI; Raykar et al.
Predicted median survival times are evaluated using the relative absolute error (RAE; Yu et al.
and calibration slope (CAL), as implemented by
Chapfuwa et al.
We report RAE on both non-censored (RAE nc ) and censored (RAE c ) data points (see Equations 14 and 15 in Appendix F).
The relative absolute error quantiﬁes relative deviation of median predictions from the observed survival times; while the calibration slope indicates whether a model tends to under- or overestimate risk on average.
For the (semi-)synthetic datasets we average all results across independent simulations , i.e. dataset replicates; while for the real-world data we use the Monte Carlo cross-validation procedure.
R ESULTS 5. C LUSTERING
We ﬁrst compare clustering performance on the nonlinear (semi-)synthetic data.
Table 3 shows the results averaged across several simulations.
In addition to the clustering, we evaluate the concordance index to verify that the methods can adequately model time-to-event in these datasets.
Training set results are reported in Table 12 in Appendix H.
As can be seen, in both problems, VaDeSC outperforms other models in terms of clustering and achieves performance comparable to SCA and DSM w.r.t.
the CI.
Including survival times appears to be helpful for identifying clusters, since the completely unsupervised VaDE achieves signiﬁcantly worse results.
It is also assuring that the model is able to predict clusters fairly well even when the survival time is not given at test time (“w/o t ”).
Furthermore, performing k -means clustering in the latent space of a VAE with the Weibull survival loss (“VAE + Weibull”) clearly does not lead to the identiﬁcation of the correct clusters.
In both datasets, k -means on VAE representations yields almost no improvement over k -means on raw features.
This suggests that the Gaussian mixture structure incorporated in the generative process of VaDeSC plays an essential role in inferring clusters.
Interestingly, while SCA and DSM achieve good results on survMNIST, both completely fail to identify clusters correctly on synthetic data, for which the generative process is very similar to the one assumed by VaDeSC.
In the synthetic data, the clusters do not have prominently different survival distributions (see Appendix E.1); they are rather characterised by different association between the covariates and survival times — whereas the two baseline methods tend to discover clusters with disparate survival distributions.
The SSC offers little to no gain over the conventional k -means performed on the complete feature set.
Last, we note that in both datasets VaDeSC has a signiﬁcantly better CI than the Cox PH model likely due to its ability to capture nonlinear relationships between the covariates and outcome.
Figure 4 provides a closer inspection of the clustering and latent representations on survMNIST data.
It appears that SCA and DSM, as expected, fail to discover clusters with similar Kaplan–
Meier (KM) curves and have a latent space that is driven purely by survival time.
While the VAE +
Weibull model learns representations driven by both the explanatory variables (digits in the images) and survival time, the post hoc k -means clustering fails at identifying the true clusters.
By contrast,
VaDeSC is capable of discovering clusters with even minor differences in KM curves and learns
[Table]
[Figure]
[Table]
representations that clearly reﬂect the covariate and survival variability.
Similar differences can be observed for the synthetic data (see Appendix H.2).
5. T IME - TO - EVENT P REDICTION
We now assess time-to-event prediction on clinical data.
As these datasets do not have ground truth clustering labels, we do not evaluate the clustering performance.
However, for Hemodialysis, we provide an in-depth qualitative assessment of the learnt clusters in Appendix H.8.
Table 4 shows the time-to-event prediction performance on SUPPORT, HGG
[Table]
and Hemodialysis.
Results for FLChain are reported in Appendix H.7 and yield similar conclusions.
Surprisingly, SCA often has a considerable variance w.r.t.
the calibration slope, sometimes yielding badly calibrated predictions.
Note, that for SCA and DSM, the results differ from those reported in the original papers likely due to a different choice of encoder architectures and numbers of latent dimensions.
In general, these results are promising and suggest that VaDeSC remains competitive at time-to-event modelling, offers overall balanced predictions, and is not prone to extreme overﬁtting even when applied to simple clinical datasets that are low-dimensional or contain few non-censored patients.
5. A PPLICATION TO C OMPUTED T OMOGRAPHY D ATA
We further demonstrate the viability of our model in a real-world application using a collection of several CT image datasets acquired from NSCLC patients (see Appendix E).
The resulting dataset poses two major challenges.
The ﬁrst one is the high variability among samples, mainly due to disparate lung tumour locations.
In addition, several demographic characteristics, such as patient’s age, sex, and weight, might affect both the CT scans and survival outcome.
Thus, a representation learning model that captures cluster-speciﬁc associations between medical images and survival times could be highly beneﬁcial to both explore the sources of variability in this dataset and to understand the different generative mechanisms of the survival outcome.
The second challenge is the modest dataset size ( N = 961 ).
To this end, we leverage image augmentation during neural network training (Perez & Wang, 2017) to mitigate spurious associations between survival outcomes and clinically irrelevant CT scan characteristics (see Appendix G).
We compare our method to the DSM, omitting the SCA, which seems to yield results similar to the latter.
As well-established time-to-event prediction baselines, we ﬁt Cox PH and Weibull AFT models on radiomics features .
The latter are extracted using manual pixel-wise tumour segmentation maps, which are time-consuming and expensive to obtain.
Note that neither DSM, nor VaDeSC require segmentation as an input.
Table 5 shows the time-to-event prediction results.
Interestingly,
DSM and VaDeSC achieve performance comparable to the radiomics-based models, even yielding a slightly better calibration on average.
This suggests that laborious manual tumour segmentation is not necessary for survival analysis on CT scans.
Similarly to tabular clinical datasets (see Table 4),
DSM and VaDeSC have comparable performance.
Therefore, we investigate cluster assignments qualitatively to highlight differences between the two methods.
[Table]
In Figure 5, we plot the KM curves and corresponding centroid CT images for the clusters discovered by DSM and VaDeSC on one train-test split.
By performing several independent experiments (see Appendix H.10), we observe that both methods discover patient groups with different empirical survival time distributions.
However, in contrast to DSM, VaDeSC clusters are consistently associated with the tumour location.
On the contrary, the centroid CT images of the DSM clusters show [Figure] [Table] no visible difference.
This is further demonstrated in Figure
[Table]
where we generated several CT scans for each cluster by sampling from the multivariate Gaussian distribution in the latent space of the VaDeSC.
We observe a clear association with tumour location, as clusters 1 ( (cid:4) ) and 3 ( (cid:4) ) correspond to the upper section of the lungs.
Indeed, multiple studies have suggested a higher ﬁve-year survival rate in patients with the tumour in the upper lobes (Lee et al., 2018).
It is also interesting using the decoder network. to observe the amount of variability within each cluster: every generated scan is characterised by a unique lung shape, and tumours size.
Since DSM is not a generative model, we instead plot the original samples assigned to each cluster by DSM in Appendix H.
Finally, in Table we compute clusterspeciﬁc statistics for a few important demographic and clinical variables (Etiz et al., 2002; Agarwal et al., 2010; Bhatt et al., 2014). We observe that
VaDeSC tends
to discover
clusters
with more disparate characteristics and is thus able to stratify patients not only by risk,
but also
by clinical conditions.
Overall, the results above agree with our ﬁndings on (semi synthetic data: VaDeSC identiﬁes clusters informed
by both
the
covariates
and
survival outcomes
, leading to a
stratiﬁcation
very
different from previous approaches.
C ONCLUSION
In this paper, we introduced a novel deep probabilistic model for clustering survival data.
In contrast to existing approaches, our method can retrieve clusters driven by both the explanatory variables and survival information and it can be trained efﬁciently in the framework of stochastic gradient variational inference.
Empirically, we showed that our model offers an improvement in clustering performance compared to the related work while staying competitive at time-to-event modelling.
We also demonstrated that our method identiﬁes meaningful clusters from a challenging medical imaging dataset.
The analysis of these clusters provides interesting insights that could be useful for clinical decision-making.
To conclude, the proposed VaDeSC model offers a holistic perspective on clustering survival data by learning structured representations which reﬂect the covariates, outcomes, and varying relationships between them.
Limitations & Future Work
The proposed model has a few limitations.
It requires ﬁxing a number of mixture components a priori , since global parameters are not treated in a fully Bayesian manner, as opposed to the proﬁle regression.
Although the obtained cluster assignments can be explained post hoc , the relationship between the raw features and survival outcomes remains unclear.
Thus, further directions of work include ( i ) improving the interpretability of our model to facilitate its application in the medical domain and ( ii ) a fully Bayesian treatment of global parameters.
Indepth interpretation of the clusters we found in clinical datasets is beyond the scope of this paper, but we plan to investigate these clusters together with our clinical collaborators in follow-up work. R EPRODUCIBILITY S TATEMENT
To ensure the reproducibility of this work several measures were taken.
The experimental setup is outlined in Section 4.
Appendix E details benchmarking datasets and simulation procedures. Appendix F deﬁnes metrics used for model comparison. Appendix G contains data preprocessing, architecture and hyperparameter tuning details. For the computed tomography data we assess the stability of the obtained results in Appendix H.10.
The code was made publicly available.
1 Most datasets used for comparison are either synthetic or publicly available.
HGG, Hemodialysis and in-house PET/CT data could not be published due to medical conﬁdentiality
A C ENSORING IN S URVIVAL A NALYSIS
Censoring is one of the key features of survival analysis that distinguish it from the classical regression modelling (Rodr´ıguez, 2007).
For some units in the study the event of interest, e.g. death, is observed, whereas for others it does not occur during the observation period.
The latter units are referred to as censored , and the observation time for these units is a lower bound on the time-to-event.
The censoring event described above is commonly known as right censoring (Winkel, 2007) and is the only form of censoring considered throughout the current paper.
Figure 7 depicts an example of the survival data: 30 patients from an observational cohort, 13 of whom were censored.
Although there exist several types of censoring, most of the schemes assume that it is non-informative , i.e. that the censoring time is independent of unit’s survival beyond it (Rodr ´ ıguez, 2007).
This assumption is made by the VaDeSC model as well.
B R ELATED W ORK
Herein, we provide a detailed overview of the related work.
The proposed VaDeSC model (see
Section 3) builds on three major topics: ( i ) probabilistic deep unsupervised clustering, ( ii ) mixtures of regression models, and ( iii ) survival analysis.
Deep Unsupervised Clustering Long-established clustering algorithms, such as k-means and
Gaussian Mixture Models, have been recently combined with deep neural networks to learn better representations of high-dimensional data (Min et al., 2018).
Several techniques have been proposed in the literature, featuring a wide range of neural network architectures.
Among them, the VAE (Kingma & Welling, 2014; Rezende et al., 2014) combines variational Bayesian methods with the ﬂexibility and scalability of neural networks.
Gaussian mixture variational autoencoders (GMMVAE; Dilokthanakul et al.
and variational deep embedding (VaDE; Jiang et al.
(2017)) are variants of the VAE in which the prior is a Gaussian mixture distribution.
With this assumption, the data is clustered in the latent space of the VAE and the resulting inference model can be directly optimised within the framework of stochastic gradient variational Bayes.
Further extensions focus on a fully Bayesian treatment of global parameters with a Dirichlet prior over the mixing coefﬁcients of the Gaussian mixture model (GMM)
(Luo et al., 2018; Johnson et al., 2016), alleviating the need to ﬁx the number of clusters a priori .
Mixtures of Regression Models Mixtures of regression (McLachlan & Peel, 2004) model a response variable by a mixture of individual regressions with the help of latent cluster assignments.
Such mixtures do not have to be limited to a ﬁnite number of components; for example, proﬁle regression (Molitor et al., 2010) leverages the Dirichlet mixture model for cluster assignment.
Along similar lines, within the machine learning community, mixture of experts models (MoE) were introduced by Jacobs et al.
Mixture models have been successfully leveraged for survival analysis as well (Farewell, 1982;
Rosen & Tanner, 1999; Liverani et al., 2020; Chapfuwa et al., 2020; Nagpal et al., 2021a;b).
For instance, Farewell (1982) considered ﬁtting separate models for two subpopulations: short-term and long-term survivors.
Rosen & Tanner (1999) extended the classical Cox PH model with the MoE.
Recent neural network approaches typically ﬁt mixtures of survival models on rich representations produced by an encoding neural network (Chapfuwa et al., 2020; Nagpal et al., 2021a;b).
Nonlinear Survival Analysis One of the ﬁrst machine learning approaches for survival analysis was Faraggi–Simon’s network (Faraggi & Simon, 1995), which was an extension of the classical
Cox PH model (Cox, 1972).
Since then, an abundance of machine learning methods have been developed: random survival forests (Ishwaran et al., 2008), deep survival analysis (Ranganath et al., , neural networks based on the Cox PH model (Katzman et al., 2018; Kvamme et al., 2019), adversarial time-to-event modelling (Chapfuwa et al., 2018), and deep survival machines (Nagpal et al., 2021a) etc.
An exhaustive overview of these and other techniques is beyond the scope of the current paper.
C C OMPARISON WITH R ELATED W ORK
Generative Assumptions
In addition to the differences summarised in Table 1, VaDeSC and the related approaches (Chapfuwa et al., 2020; Nagpal et al., 2021a; Liverani et al., 2020) assume different data generating processes.
Simpliﬁed graphical models assumed by these techniques are shown in Figure 8.
Both SCA and DSM are not generative w.r.t. explanatory variables x , and thus, do not maximise the joint likelihood L ( x , t ) and are outcome-driven. On the other hand, both proﬁle regression and VaDeSC specify a model for covariates x , as well as for survival time t . However, PR does not infer latent representations z and importantly, in its simplest form assumes that x ⊥⊥ t | c .
The latter assumption can limit the discriminative power of a mixture with few components.
By default, VaDeSC does not make such restrictive assumptions.
[Figure]
Deep Cox Mixtures The concurrent work by Nagpal et al.
(2021b) on deep Cox mixtures allows leveraging latent representations learnt by a VAE to capture nonlinear relationships in the data.
The hidden representation is simultaneously used to ﬁt a mixture of K Cox models.
In contrast to our method, DCM does not specify a generative model explicitly.
The VAE loss is added as a weighted term to the main loss derived from the Cox models and is optimised jointly using a Monte Carlo
EM algorithm.
The VAE is mainly used to enforce a Gaussian prior in the latent space, resulting in more compact representations.
It can be replaced with a regularised AE, or even an MLP encoder, as mentioned by the authors (Nagpal et al., 2021b).
Thus, the maximisation of the joint likelihood
L ( x , t ) , rather than of L ( t ) , is not the main emphasis of the DCMs.
On the contrary, our approach is probabilistic and speciﬁes a clear and interpretable generative process from which an ELBO of the joint likelihood can be derived formally.
Additionally, VaDeSC enforces a more structured representation in the latent space, since it uses a Gaussian mixture prior.
Empirical Contribution Last but not least, the experimental setup of the current paper is different from those of Bair & Tibshirani (2004);
Chapfuwa et al.
(2020); Liverani et al.
; Nagpal et al.
(2021a;b).
Previous work on mixture modelling and clustering for survival data has focused on applications to tabular datasets.
Survival analysis on image datasets (Haarburger et al., 2019; Bello et al., 2019) poses a challenge pertinent to many application areas, e.g. medical imaging.
An important empirical contribution of this work is to apply the proposed technique and demonstrate its utility on computed tomography data (see Section 5.3).
In this appendix, we provide the complete derivation of the ELBO of VaDeSC using the SGVB estimator.
From Equation 5 we can rewrite the ELBO as
[Equation]
In the following we investigate each term of the equation above.
Reconstruction The ﬁrst term is also known as the reconstruction term in a classical VAE setting.
Herein, we provide the derivation for p ( x | z ; γ ) =
Bernoulli ( x ; µ γ ) , i.e. for binary-valued explanatory variables:
[Equation]
where L is the number of Monte Carlo samples in the SGVB estimator; µ ( l ) γ = f
[Equation] ; γ ) ; z ( [Equation]
) ∼ q ( z | x ) = N ( z ; µ θ , σ 2 θ ) ; [ µ θ , log σ 2 θ ]
= g ( x ; θ ) and g ( · ; θ ) is an encoder neural network parameterised by θ ; x d refers to the d -th component of x and D is the number of features.
For other data types, the distribution p ( x | z ; γ ) in Equation 7 needs to be chosen accordingly.
Survival The second term of the ELBO includes the survival time.
Following the survival model described by Equation 1, the SGVB estimator for the term is given by
[Equation]
where λ z ( l ) ν
[Equation]
Clustering The third term in Equation 6 corresponds to the clustering loss and is deﬁned as
[Equation]
Prior The fourth term can be approximated as
[Equation]
Variational Finally, the ﬁfth term of the ELBO corresponds to the entropy of the variational distribution q ( z , c | x , t ) :
[Equation]
where the ﬁrst term is the entropy of a multivariate Gaussian distribution with a diagonal covariance matrix.
We can then approximate the expression in Equation 11 by
[Equation]
where J is the number of latent space dimensions.
E D ATASETS
Below we provide detailed summaries of the survival datasets used in the experiments (see Table 2).
Synthetic As the simplest benchmark problem, we simulate synthetic survival data based on the generative process reminiscent of the VaDeSC generative model (see Figure 3).
In this dataset, both covariates and survival times are generated based on latent variables which are distributed as a Gaussian mixture.
The dependence between latent and explanatory variables is given by a multilayer perceptron and is, therefore, nonlinear.
Appendix E.1 contains a detailed description of these simulations and a visualisation of one replicate (see Figure 9).
survMNIST Another dataset we consider is the semi-synthetic survival MNIST (survMNIST) introduced by P¨olsterl (2019), used for benchmarking nonlinear survival analysis models, e.g. by
Goldstein et al.
In survMNIST, explanatory variables are given by images of handwritten digits (LeCun et al., 2010).
Digits are assigned to clusters and synthetic survival times are generated using the exponential distribution with a cluster-speciﬁc scale parameter.
Appendix E.2 provides further details on the generation procedure.
SUPPORT To compare models on real clinical data, we consider the Study to Understand Prognoses and Preferences for Outcomes and Risks of Treatments (SUPPORT; Knaus et al.
consisting of seriously ill hospitalised adults from ﬁve tertiary care academic centres in the United
States.
It includes demographic, laboratory, and scoring data acquired from patients diagnosed with cancer, chronic obstructive pulmonary disease, congestive heart failure, cirrhosis, acute renal failure, multiple organ system failure, and sepsis.
FLChain A second clinical dataset was acquired in the course of the study of the relationship between serum free light chain (FLChain) and mortality (Kyle et al., 2006; Dispenzieri et al., 2012) conducted in Olmsted County, Minnesota in the United States.
The dataset includes demographic and laboratory variables alongside with the recruitment year.
FLChain is by far the most lowdimensional among the considered datasets (see Table 2).
Both SUPPORT and FLChain data are publicly available.
HGG Additionally, we perform experiments on an in-house dataset consisting of 453 patients diagnosed with high-grade glioma (HGG), a type of brain tumour (Buckner, 2003).
The cohort includes patients with glioblastoma multiforme and anaplastic astrocytoma (Weller et al., 2020).
The dataset consists of demographic variables, treatment assignments, pre- and post-operative volumetric variables, molecular markers, performance scores, histology ﬁndings, and information on tumour location.
This dataset has by far the least amount of patients (see Table 2) and therefore, could be a good benchmark for the ‘low data’ regime.
Another clinical dataset we include is an observational cohort of patients from the
DaVita Kidney Care (DaVita Inc., Denver, CO, USA) dialysis centres (Gotta et al., 2018; 2019a;b; 2020; 2021).
The cohort is composed of patients who started chronic hemodialysis (HD) in childhood and have then received thrice-weekly HD between 05/2004 and 03/2016, with a maximum follow-up until < 30 years of age.
The dataset consists of demographic factors, such as the age from the start of dialysis, gender, as well as etiology of kidney disease and comorbidities, dialysis dose in terms of spKt/V, eKt/V (Daugirdas, 1993; Daugirdas & Schneditz, 1995), ﬂuid removal in terms of UFR (mL/kg/h) and total ultraﬁltration (UF, % per kg dry weigth per session), and the interdialytic weight gain (IDWG) (Marsenic et al., 2016).
It should be noted that this dataset is rather challenging, given the high percentage of censored observations (see Table 2).
The scientiﬁc use of the deidentiﬁed standardised electronic medical records was approved by DaVita (DaVita Clinical
Research®, Minneapolis, MN); IRB approval was not required since retrospective analysis was performed on the deidentiﬁed dataset.
A complete description of the variables and measurements used is provided by Gotta et al.
(2021) in the Methods section.
NSCLC Finally, we perform experiments on a pooled dataset of ﬁve observational cohorts consisting of patients diagnosed with non-small cell lung cancer (NSCLC).
We consider a large set of pretreatment CT scans and CT components of positron emission tomography (PET)/CT scans from the following sources:
• An in-house PET/CT dataset consisting of 392 patients treated at a university hospital.
It includes manual delineations, clinical and survival data.
A retrospective, observational study was approved by the local ethics committee.
• Lung1 dataset consisting of 422 NSCLC patients treated at Maastro Clinic, the Netherlands (Aerts et al., 2014; 2019).
For these patients, CT scans, manual delineations, clinical, and survival data are available.
Lung1 is publicly available in The Cancer Imaging Archive (TCIA; Clark et al.
• Lung3 dataset consisting of 89 NSCLC patients treated at Maastro Clinic, the Netherlands (Aerts et al., 2014; 2015).
Lung3 includes pretreatment CT scans, tumour delineations, and gene expression proﬁles.
It is is publicly available in TCIA.
• NSCLC
Radiogenomics dataset acquired from 211 NSCLC patients from the Stanford University School of Medicine and Palo Alto Veterans Affairs Healthcare System (Bakr et al., 2017; 2018; Gevaert et al., 2012).
The dataset comprises CT and PET/CT scans, segmentation maps, gene mutation and RNA sequencing data, clinical data, and survival outcomes.
It is is publicly available in TCIA.
Only subjects with a tumour segmentation map and transversal CT or PET/CT scan available are retained.
Patients who have the slice of (PET)/CT with the maximum tumour area outside the lungs, e.g. in the brain or abdomen, are excluded from the analysis.
Thus, the ﬁnal dataset includessubjects.
E. S YNTHETIC D ATA G ENERATION
Herein, we provide details on the generation of synthetic survival data used in our experiments (see
Section 4).
The procedure for simulating these data is very similar to the generative process assumed by the proposed VaDeSC model ( cf .
Figure 3) and can be summarised as follows:
Let π c = 1 K , for 1 ≤ c ≤ K . 2.
Sample c i ∼ Cat ( π ) , for 1 ≤ i ≤ N . 3.
Sample µ c,j ∼ unif (cid:0) − 12 , 12 (cid:1) , for 1 ≤ c ≤ K and 1 ≤ j ≤ J . 4.
Let Σ c
=
I J S c , where
S c ∈ R
J × J is a random symmetric, positive-deﬁnite matrix, 2 for 1 ≤ c ≤ K . 5.
Sample z i ∼ N ( µ c i , Σ c i ) , for 1 ≤ i ≤ N . 6.
Let g dec ( z )
=
W 2 ReLU ( W 1 ReLU ( W 0 z + b 0 )
+ b 1 ) + b 2 , where W 0 ∈ R
h × J , W 1 ∈ R h × h , W 2 ∈ R D × h and b 0 ∈ R h , b 1 ∈ R h , b 2 ∈ R D are random matrices and vectors,and h is the number of hidden units.
Let x i =
g dec ( z i ) , for 1 ≤ i ≤ N . 8.
Sample β c,j ∼ unif ( − 10 , 10) , for 1 ≤ c ≤ K and 0 ≤ j ≤ J . 9.
Sample u i ∼ Weibull (cid:0) softplus (cid:0) z
(cid:62) i
β c i (cid:1) , k (cid:1) , 4 for 1 ≤ i ≤ N . 10.
Sample δ
i ∼ Bernoulli (1 − p cens ) , for 1 ≤ i ≤ N . 11.
Let t i
= u
i if δ i
= 1 , and sample t i ∼ unif (0 , u i ) otherwise, for 1 ≤ i ≤ N .
Here, similar to Section 3, K is the number of clusters; N is the number of data points; J is the number of latent variables; D is the number of features; k is the shape parameter of the Weibull distribution (see Equation 1); and p cens denotes the probability of censoring.
The key difference between clusters is in how latent variables z generate the survival time t : each cluster has a different weight vector β that determines the scale of the Weibull distribution.
Generated In our experiments, we ﬁx K = 3 , N = 60000 , J = 16 , D = 1000 , k = 1 , and p cens = 0 .
We hold out 18000 data points as a test set and repeat simulations 5 times independently.
A visualisation of one simulation is shown in Figure 9.
As can be seen in the t -SNE (van der Maaten & Hinton, 2008) plot, three clusters are not as clearly separable based on covariates alone; on the other hand, they have visibly distinct survival curves as shown in the KM plot.
[Figure]
E. S URVIVAL MNIST D ATA G ENERATION
We generate semi-synthetic survMNIST data as described by P ¨ olsterl (2019) (original implementation available at https://github.com/sebp/survival-cnn-estimator ).
The generative process can be summarised by the following steps:
Assign every digit in the original MNIST dataset (LeCun et al., 2010) to one of K clusters (ensure that every cluster contains at least one digit) 2.
Sample risk score r c ∼ unif (cid:0) 12 , 15 (cid:1) , for 1 ≤ c ≤ K . 3.
Let λ c = 1 T 0 exp( r c ) , for 1 ≤ c ≤ K , where T 0 is the speciﬁed mean survival time and 1 Tis the baseline hazard.
Sample a i ∼ unif (0 , 1) and let u i = − log a
i λ ci , for 1 ≤ i ≤ N . 5.
Let q cens =
q 1 − p cens ( u 1: N ) , where q α ( · ) denotes the α -th quantile.
Sample t cens ∼ unif (min 1 ≤ i ≤ N u
i , q cens ) .
Let δ i
= 1 if u i ≤ t cens , and δ i =
otherwise, for 1 ≤ i ≤ N . 8.
Let t i
= u
i if δ i
= 1 , and t i =
t cens otherwise, for 1 ≤ i ≤ N .
Observe that here p cens is only a lower bound on the probability of censoring and the fraction of censored observations can be much larger than p cens .
[Table]
[Table] In our experiments, we ﬁx K = 5 and p cens = 0 .
We repeat simulations independently 10 times.
The train-test split is deﬁned as in the original MNIST data (60,000 vs.
Table 7 shows an example of cluster and risk score assignments in one of the simulations.
The resulting dataset is visualised in Figure 10.
Observe that the clusters are not well-distinguishable based on covariates alone.
Also note that some clusters contain many censored data points, whereas some do not contain any, as shown in the
Kaplan–Meier plot.
[Figure]
to have more disparate survival distributions across clusters than the synthetic data ( cf
F E VALUATION M ETRICS
Below we provide details about the evaluation metrics used to compare clustering and time-to-event prediction performance of the models.
For clustering, we employ the adjusted Rand index (ARI) and normalised mutual information (NMI) as implemented in the sklearn.metrics module of scikit-learn (Pedregosa et al., 2011).
Clustering accuracy is evaluated by ﬁnding an optimal mapping between assignments and true cluster labels using the Hungarian algorithm, as implemented in sklearn.utils.linear assignment .
The concordance index (Raykar et al., 2007) evaluates how well a survival model is able to rank individuals in terms of their risk.
Given observed survival times t i , predicted risk scores η i , and censoring indicators δ i , the concordance index is deﬁned as
[Equation]
[Equation] Thus, a perfect ranking achieves a CI of 1.0, whereas a random ranking is expected to have a CI of 0.5.
We use the lifelines implementation of the CI (Davidson-Pilon et al., 2021).
Relative absolute error (Yu et al., 2011) evaluates predicted survival times in terms of their relative deviation from the observed survival times.
In particular, given predicted survival times ˆ t
i , for noncensored data points, RAE is deﬁned as [Equation]
[Equation]
On censored data, similarly to Chapfuwa et al.
(2020), we deﬁne the metric as follows:
[Equation]
Last but not least, similar to Chapfuwa et al.
we use the calibration slope to evaluate the overall calibration of models.
It is desirable for the predicted survival times to match with the empirical marginal distribution of t .
Calibration slope indicates whether a model tends to under- or overestimate risk on average, and thus, an ideal calibration slope is 1 . 0 .
G I MPLEMENTATION D ETAILS
Herein, we provide implementation details for the experiments described in Section 5.
G. P REPROCESSING
For all datasets, we re-scale survival times to be in [0 . 001 , 1 . 001] .
For survMNIST, we re-scale the features to be in [0 , 1] .
For tabular datasets, we re-scale the features to zero-mean and unitvariance.
Categorical features are encoded using dummy variables.
For Hemodialysis data, the same preprocessing routine is adopted as described by Gotta et al.
For NSCLC CT images, preprocessing consists of several steps.
For each subject, slices withinmm from a slice with a maximal transversal tumour area are averaged to produce a single 2D image.
2D images are then cropped around the lungs and normalised.
Subsequently, histogram equalisation is applied (Lehr & Capek, 1985).
The images are then downscaled to a resolution of 64 × 64 pixels.
Figure 11 provides an example of raw and preprocessed images from the Lung1 dataset.
Finally, during VaDeSC, SCA, and DSM training, we augment images by randomly reducing/increasing the brightness, adding Poisson noise, ﬂipping horizontally, blurring, zooming, stretching, and shifting both vertically and horizontally (see Figure 11).
During initial experimentation, we observed that image augmentation was crucial for decorrelating clinically irrelevant scan characteristics, such as rotation or scaling, and the predicted survival outcome.
Neither VaDeSC, nor DSM, achieve predictive performance comparable to radiomics-based features in the absence of augmentation.
[Figure]
images during neural network training.
For the radiomics-based Cox PH and Weibull AFT models, features were extracted from the preprocessed 2D images using the RadiomicsFeatureExtractor provided by the PyRadiomics library (van Griethuysen et al., 2017).
Tumour segmentation maps were given as input during feature extraction, to detect the ROI.
The following implementations of the models were used:
• VaDeSC : we implemented our model in TensorFlow 2.4 (Abadi et al., 2015).
Currently, the code is available at https://anonymous.4open.science/r/ 8d1dd2d7-5475-485d-b203-57a60f5ca114 .
• k -means : we used the implementation available in the scikit-learn library (Pedregosa et al., 2011).
• Cox PH & Weibull AFT : we used the implementation available in the lifelines library (Davidson-Pilon et al., 2021).
• Radiomics : radiomic features were extracted using the PyRadiomics library (van Griethuysen et al., 2017).
• SSC : we re-implemented the model in Python based on the original R (R Core Team, 2020) code provided by Bair & Tibshirani (2004).
• SCA : we adapted the original code by Chapfuwa et al.
(2020) available at https:// github.com/paidamoyo/survival_cluster_analysis .
• DSM : we adapted the original code by Nagpal et al.
(2021a) available at https: //github.com/autonlab/DeepSurvivalMachines .
• Proﬁle regression : we used the original implementation available in the R package PReMiuM (Liverani et al., 2015).
Throughout our experiments, we use several different encoder and decoder architectures described in Tables 8, 9, and 10.
As mentioned before, the encoder architectures and numbers of latent dimensions were kept ﬁxed across all neural-network-based techniques, namely VaDeSC, SCA, and
DSM, for the sake of fair comparison.
Hyperparameters VaDeSC hyperparameters across all datasets are reported in Table 11.
We use L = 1 Monte Carlo samples for the SGVB estimator.
In most datasets, we do not pretrain autoencoder weights.
Generally, we observe little variability in performance when adjusting the number of latent dimensions, shape of the Weibull distribution or mini-batch size.
For the SCA model, we tuned the concentration parameter of the Dirichlet process, mini-batch size, learning rate, and the number of training epochs.
For the sake of fair comparison, we truncate the Dirichlet process at the true number of clusters, when known.
Similarly, for the DSM, the mini-batch size, learning rate, and the number of training epochs were tuned.
To make the comparison with VaDeSC fair, we ﬁxed the discount parameter for the censoring loss at 1 . 0 , since VaDeSC does not discount censored observations in the likelihood, and used the same number of mixture components as for the VaDeSC.
Pretraining VAE-based models (Dilokthanakul et al., 2016; Jiang et al., 2017; Kingma & Welling, 2014) tend to converge to undesirable local minima or saddle points, due to high reconstruction error at the beginning of the training.
To prevent this, the encoder and decoder networks are usually pretrained using an autoencoder loss.
In the most of our experiments, we did not observe such behaviour.
Hence, we did not pretrain the model, except for the Hemodialysis dataset.
The latter showed better performance if pretrained for only 1 epoch.
After the encoder-decoder network has been pretrained, we projected the training samples into the latent space and ﬁtted a Gaussian mixture model to initialise the parameters π , µ , Σ .
[Table] [Table]
[Table]
Herein, J denotes the number of latent dimensions; K is the number of clusters; and k is the shape parameter of the
Weibull distribution.
MSE stands for mean squared error; BCE – for binary cross entropy.
[Table]
H F URTHER R ESULTS H. T RAINING S ET C LUSTERING P ERFORMANCE
[Table]
H. Q UALITATIVE R ESULTS : S YNTHETIC D ATA
[Figure]
H. C OMPARISON WITH P ROFILE R EGRESSION
Herein, we provide additional experimental results to compare VaDeSC with the proﬁle regression (PR) for survival data (Liverani et al., 2020).
The comparison is performed on a subsample of the original synthetic dataset (see Appendix E.1) consisting of N = 5000 data points.
Prior to performing proﬁle regression, we reduce the dimensionality of the dataset by retaining only the ﬁrst 100 principal components, which preserve, approximately, 90% of the variance.
We ﬁt a proﬁle survival regression model with the cluster-speciﬁc shape parameter of the Weibull distribution by running 10000 iterations of the MCMC in the burn-in period and 2000 iterations after the burn-in.
To facilitate fair comparison, we set the initial number of clusters to the ground truth K = 3 .
For
VaDeSC, the same hyperparameters are used as reported in Table 11.
Table 13 reports clustering performance on the test set of 1500 data points.
VaDeSC, both given survival times at prediction and without them, outperforms proﬁle regression by a margin.
Nevertheless, PR offers a noticeable improvement over k -means clustering and even more sophisticated
SCA and DSM ( cf .
Table 3).
While in some instances, PR manages to identify some cluster structure, in others, it assigns all data points to a single cluster, hence, a large standard deviation across simulations.
It is also interesting to see that a reduction in the dataset size has led to a signiﬁcant decrease in the performance of VaDeSC.
We attribute this to a highly nonlinear structure of this dataset.
[Table]
This experiment has clear shortcomings: due to computational limitations, we are only able to ﬁt a proﬁle regression model on a transformed version of the dataset with fewer dimensions which might not reﬂect the full capacity of the PR model.
Nevertheless, we believe that poor scalability of the
MCMC, both in the number of covariates and in the number of training data points, warrant a more computationally efﬁcient approach.
H. V ARYING THE F RACTION OF C ENSORED D ATA P OINTS
Herein we investigate the behaviour of the proposed approach under different fractions of censored observations.
We perform a controlled experiment on the synthetic dataset (see Appendix E.1) where we vary the percentage of censored data while keeping all other simulation parameters ﬁxed.
We evaluate the proposed algorithm, VaDeSC, and several baselines (VaDE, DSM, and Weibull AFT) under 10-90% of censored observations and compare their clustering and time-to-event predictive performance.
[Figure]
Figure 13 shows the results of the experiment. VaDeSC stably achieves clustering accuracy of, approximately, 90% for 10, 30, and 50% of censored observations. However, for 70 and 90%, its accuracy drops considerably, and variance across simulations increases. Nevertheless, even at 90% of censored data points, on average, VaDeSC is still more accurate and stable than the completely unsupervised VaDE.
For time-to-event predictions, all models behave as expected: with an increase in the percentage of censored observations, test-set CI decreases, RAE nc increases, while RAE c decreases. Overall, VaDeSC, DSM, and Weibull AFT appear to behave and scale very similarly w.r.t. changes in censoring. In the future, it would be interesting to perform a similar experiment on a real-world dataset with a moderate percentage of censored observations, e.g. SUPPORT.
H. V ARYING THE N UMBER OF C OMPONENTS IN V A D E SC
[Figure]
[Figure] [Figure] [Figure] [Figure]
[Figure]
[Figure]
T IME - TO - EVENT P REDICTION : FLC HAIN
[Table]
[Table] [Figure] [Table] : H EMODIALYSIS D ATA
In addition to the results reported in Section 5.2, we have investigated qualitatively the clustering obtained by VaDeSC when applied to the Hemodialysis dataset.
The true structure of this data is unknown, however several studies have suggested a stratiﬁcation of patients according to age and dialysis doses (Gotta et al., 2021).
We compute the optimal number of clusters using grid-search by measuring the time-to-event prediction performance.
With two clusters our model achieved the best results.
In Figure 18(a), we plot the t -SNE decomposition of the latent embeddings, z (see
Figure 2), obtained by VaDeSC.
Observe that the patients clustered together tend to be close in the latent space, following the mixture of Gaussian prior.
We also plot Kaplan–Meier curves for each patient subgroup according to the predicted cluster in Figure 18(b).
Interestingly, the two resulting curves differ substantially from each other.
The patients in cluster 2 ( (cid:4) ) show a higher survival probability over time compared to those in cluster 1 ( (cid:4) ).
[Figure]
We now characterise each subgroup of patients by identifying the most inﬂuential covariates in determining the cluster assignment by VaDeSC.
In particular, we apply a state-of-the-art method based on Shapley values (Shapley, 1953), namely the TreeExplainer by Lundberg et al.
(2020), to explain an XGBoost classiﬁer (Chen &
Guestrin, 2016) trained to predict the VaDeSC cluster labels from patient characteristics (see Appendix I).
To mimic the VaDeSC clustering model, we included survival times as an additional input to the classiﬁer.
In Figure 20, we present the most important covariates together with their cluster-wise distributions.
It is important to note that the survival time was also ranked among the most important, but we did not include it into Figure 20.
Serum albumin level, previously identiﬁed as an important survival indicator for patients receiving hemodialysis (Foley et al., 1996; Amaral et al., 2008), emerges as the most important predictor.
Additionally, dialysis dose in terms of spKt/V and ﬂuid balance in terms of ultraﬁltration (UF) rate prove to be crucial in classifying the data into high-risk (cluster 1 , (cid:4) ) and low-risk (cluster 2 , (cid:4) ) groups.
By analysing the cluster-wise distributions, we notice that the high-risk cluster is characterised by lower values of albumin, dialysis dose, and ultraﬁltration, in agreement with previous studies (Gotta et al., 2020; 2021; Movilli et al., 2007).
Cluster-Speciﬁc Survival Models Ultimately, our model should discover subgroups of patients characterised not only by their risk but, more importantly, by the relationships between the covariates and survival outcomes.
In other words, the cluster-speciﬁc parameters of the Weibull distribution should ideally vary across clusters, highlighting different associations between the covariates.
We demonstrate this in Figure 19 by plotting the cluster-speciﬁc survival parameters, β c (see Equation 1), trained using the Hemodialysis data.
We observe a clear difference between the two clusters, with several latent dimensions described by both positive and negative survival parameters depending on the considered cluster. H. Q UALITATIVE R ESULTS : NSCLC
[Figure]
H. C LUSTERING S TABILITY : NSCLC
The notion of stability has been widely explored as a validation technique for clustering algorithms (Ben-Hur et al., 2002; Lange et al., 2004).
The intuition is that a ‘ good ’ algorithm should discover consistent clustering structures across repeated random simulations or replicates.
However, unsupervised learning algorithms based on deep neural networks have been shown to be highly unstable on real-world and complex data (Jiang et al., 2017).
Herein we test the stability of VaDeSC and compare it to a competitive baseline, DSM, on the NSCLC dataset.
This dataset is extremely complex due to both the high-dimensionality and the limited number of samples.
We run both methods on 20 independent train-test splits (with random seeds) and we observe the differences in the learnt clustering structures.
First of all, note that the number of discovered clusters varies across experiments.
Even though the number of mixture components in survival distribution is deﬁned a priori , some clusters collapse, i.e. they do not contain any samples from both the training and test set.
We measure this source of instability by computing the percentage of experiments that resulted in at least one collapsed cluster in Table 15.
We observe that only 5% of VaDeSC’s runs resulted in collapsed clusters, compared to 45% for DSM.
Thus, VaDeSC is relatively stable w.r.t.
the number of discovered clusters.
We also test the variability in the survival information encoded in each cluster, also shown by the
Kaplan–Meier curves in Figure 5.
For each experiment, we compute the median survival time of the samples assigned by the algorithm to each cluster.
In Table 15 we report the average and standard deviation across the experiments of the cluster-speciﬁc median survival time.
To have a meaningful average, we sort clusters by the median survival time, highest to lowest.
As discussed in Section 5.3,
DSM tends to discover clusters with more disparate survival curves than VaDeSC, which is reﬂected by the averages reported.
Here, we instead focus on the standard deviation.
Overall, VaDeSC shows lower standard deviations than DSM.
VaDeSC appears to be stable w.r.t.
the survival information captured by each cluster.
[Table]
Finally, we test whether the association between tumour locations and clusters is also consistent across experiments.
We ﬁrst compute the means of the mixture of Gaussians learnt by VaDeSC for each experiment, i.e. train-test split.
As these vectors live in the latent space of VaDeSC, we decode them, using the decoder network, to project them into the input space of the CT images.
Resulting images are shown in Figure 23, where each column corresponds to a different seed of the train-test split.
As before, we sort the clusters according to the median survival times such that the upper row corresponds to patients with higher overall survival probability.
It is evident that in most experiments, the upper row corresponds to the upper section of the lungs, whereas the bottom row — to the lower section of the lungs.
Thus, VaDeSC discovers clusters consistently associated with tumour location.
I E XPLAINING V A D E SC C LUSTER A SSIGNMENTS
In this appendix, we detail the post hoc analysis performed to enrich our case study on the real-world
Hemodialysis dataset (see Appendix H.8) by explaining VaDeSC cluster assignments in terms of input covariates.
Similar analysis can be performed for any survival dataset where an explanation of
VaDeSC clusters in terms of salient clinical concepts is desirable.
Background: SHAP algorithm for interpreting model predictions SHapley Additive exPlanation (SHAP), by Lundberg & Lee (2017), is an algorithm that aims at explaining the prediction of an instance x by computing the contribution of each feature to the prediction.
This contribution is estimated in terms of Shapley regression values, an additive feature attribution method inspired by the coalitional game theory (Shapley, 1953).
Additive feature attribution methods imply that the explanation model is a linear function of binary variables: g ( z (cid:48) )
=
φ 0 + (cid:80) Fj =
φ j z (cid:48)
j where g ( · ) is the explanation model, z (cid:48) ∈ { 0 , 1 } F is the binary feature vector of dimension F , and φ j ∈ R is the effect of the j -th feature on the output f ( x ) of the original model.
Shapley regression values estimate this effect by exploiting ideas drawn from the coalitional game theory.
In particular, for the j -th variable the Shapley regression values are computed as
[Equation]
16) where F = { 1 , ..., F } is the set of all predictor variables, x S∪{ j } and x S are input vectors from x , composed of the features in the subset S ⊆ F with and without the j -th feature, respectively, and f ( · ) is the original model ﬁtted on the set of predictor variables deﬁned by the subscript.
The feature importances estimated by SHAP are the Shapley values of a conditional expectation function of the original model (Lundberg et al., 2020).
With SHAP, global interpretations are consistent with the local explanations, since the Shapley values for instance i are the atomic unit of the global interpretations.
Indeed, the global importance of the j -th feature is computed as
[Equation]
where N is the total number of observations used to explain the model globally.
Despite the theoretical advantages of SHAP values, their practical use is hindered by the complexity of estimating E [ f ( x ) | x S ] efﬁciently.
TreeExplainer (Lundberg et al., 2020) is a variant of SHAP algorithm that computes the classical Shapley values from the game theory speciﬁcally for trees and tree ensembles reducing the complexity from exponential to polynomial time.
Preprocessing With the goal of explaining VaDeSC cluster assignments in terms of input covariates, we preselect features manually, to make our explanations robust and relevant from a clinical perspective.
Guided by the medical expertise, we exclude highly correlated or redundant features.
While the VaDeSC can handle the entire raw dataset meaningfully, the presence of highly correlated features might affect their relevance.
Since p ( c | z , t ) in VaDeSC depends on survival time (see Section 3), we included the survival times as an additional input to the classiﬁer.
Note that for this post hoc analysis we used training data only, as we were interested in explaining what our model learnt from the raw data to perform clustering.
Explanations To characterise patient subgroups identiﬁed by VaDeSC, one needs to identify the features in the input space that maximise cluster separability in the embedding space.
Note that a classiﬁer trained to recognise cluster labels from the raw covariates is inherently seeking those dimensions of maximal separability in the input.
Therefore, we ﬁrst ﬁtted a classiﬁer model to extract the explanations in terms of the input features most relevant for the prediction of cluster labels.
In particular, we trained an eXtreme Gradient Boosting (XGBoost) binary classiﬁer (Chen &
Guestrin, 2016), to predict patient labels c ∈ { 1 , 2 } assigned by VaDeSC.
We choose XGBoost since it has high accuracy in the presence of nonlinear relationships, naturally captures feature interactions, suffers from low bias, and supports fast exact computation of Shapley values by TreeExplainer (Lundberg et al., 2020).
XGBoost was trained for 5,000 boosting rounds with a maximum tree depth of 6, a learning rate equal to 0.01, a subsampling rate of 0.5, and early stopping after 200 rounds without an improvement on the validation set (set to 10% of the training data).
After training the classiﬁer, we used the TreeExplainer to compute SHAP values and identify the most determinant features for each cluster label, as in Equation 17.
To proﬁle the clusters precisely and avoid the noise induced by weakly assigned observations, we computed SHAP values on a sample of cluster prototypes , i.e. observations with max ν ∈{ 1 , 2 } p ( c = ν | z , t ) ≥ 0 .
We retained at most 500 prototypes from each cluster.
Due to the cluster imbalance, that resulted in 500 observations from cluster 1 and 212 – from cluster 2.