A Deep Variational Approach to Clustering Survival Data
Authors: Laura Manduchi
Ričards Marcinkevičs
Michela C. Massi
Thomas Weikert
Alexander Sauter
Verena Gotta
Timothy Müller
Flavio Vasella
Marian C. Neidert
Marc Pfister
Bram Stieltjes
Julia E Vogt
Venue: ICLR 2022 Poster (2022)
Keywords: survival analysis
clustering
healthcare
variational autoencoders
deep generative models

Abstract: 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.

Published as a conference paper at ICLR 2022
A DEEP VARIATIONAL APPROACH TO CLUSTERING SURVIVAL DATA
Laura Manduchi," Ričards Marcinkeviës.
Michela C. Massi, Thomas Weikert," Alexander Sauter,* Verena Gotta, Timothy Müller, Flavio Vasella, Marian C. Neidert, ? Mare Pfister, Bram Stieltjes & Julia E. Vogt'
'ETH Zürich; Politecnico di Milano; CADS, Human Technopole; Universityy Hospital Basel; University Children's 's Hospital Basel; University of Zürich; St. Gallen Cantonal Hospital
ABSTRACT
In this work, we study the problem of clustering survival data - challenging and far under-explored task. We introduce novel semi-supervised probabilistic ap- proach to cluster survival data by leveraging recent advances in stochastic gradient variational inference. contrast to previous work, our proposed method employs a deep generative model to uncover the underlying distribution of both the ex- planatory variables and censored survival times. We compare our model 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 identifying clus- ters and is competitive at predicting survival times. Relying on novel generative assumptions, the proposed model offers a holistic perspective clustering sur- vival data and holds a promise of discovering subpopulations whose survival is regulated by different generative mechanisms.
INTRODUCTION
Survival analysis (Rodríguez, 2007; D. G. Altman, 2020) has been extensively used in variety of medical applications to infer a relationship between explanatory variables and po- tentially censored survival outcome. The latter indicates the time 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 (Buck- ley & James, 1979). Recently, many machine learning tech- niques 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).
Patient population
Group
Group
Age
Late nouut
Group
Tregneel Taggor
Figure 1: Survival clustering.
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 no guarantee that the discovered subgroups are correlated with patient survival (Bair & Tibshirani, 2004). For this reason, we focus on semi- supervised learning approach to cluster survival data that jointly considers explanatory variables and censored outcome as indicators for patient's '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 more personalised discase management (Fenstermacher et al., 2011). Figure 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 distribu- tions do not need differ between clusters: compare groups 1 and
1 Equal contribution. Correspondence {laura. manduchi, ricardsm)e nf ethz.

---

Published as a conference paper at ICLR 2022
Table 1: Comparison of the proposed model the related approaches: semi-supervised clustering (SSC), profile regression (PR) for survival data, survival cluster analysis (SCA), and deep survival machines (DSM). Here, t denotes survival time, a denotes explanatory variables, corresponds to latent representations, K stands for the number of clusters, and ((; is the likelihood function.
SSC
PR
SCA
DSM
VaDeSC
Predicts f?
Learns #?
X
Maximises & (zz, t)?
X
/
X
x
Scalable?
X
X
/
/
Does not require K?
/
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 high-dimensional, unstruc- tured data (Liverani et al., 2020), or they focus on the discovery of purely outcome-drine (Chapfuwa et al., 2020; Nagpal al., 2021a), that is clusters characterised entirely by survival time. The latter might fail in applications where the survival distribution alone is not sufficiently informa- tive 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).
clusters
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 gener- ative 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 sur- vival 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.
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 the related work on elustering and mixture models for survival data on synthetic and real-world datasets. In particular, we show that VaDeSC outperforms baseline methods in terms of identifying clusters and is comparable in terms of time-to-event predictions. (iii) We apply our 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 patient characteristics, in agreement with previous medical findings.
Our main contributions are as follows: (i)
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 differ- ences 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 mixtures, 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.c. 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 sur- vival 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. DCM does not specify generative model and its loss is derived empirically by combining the VAE loss with the
2

---

Published as a conference paper at ICLR 2022
likelihood of survival times. On the contrary, method is probabilistic and has interpretable generative process from which an ELBO of the joint likelihood can be derived.
The approach by Liverani et al. (2020) is, 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 al. (2010)). In particular, they introduce Dirichlet process mixture model with cluster- specific parameters for the Weibull distribution. However, their method is unable to tackle high- dimensional 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 al., 2019). Table compares our and related methods w.r.t. range of properties. For an overview of other lines of work and detailed comparison see Appendices B and C.
3 METHOD
We present VaDeSC - a novel variational deep survival clustering model. Figure 2 provides summary of our approach: the input vector x is mapped to latent repre- sentation z using VAE with Gaussian mixture prior. The survival density function is given by mixture of Weibull distributions with cluster- specific parameters G. The parameters of the Gaussian mixture and Weibull dis- tributions are then optimised jointly using both the ex- planatory input variables and survival outcomes.
Encoder
KL divergence
Decoder
Preliminaries We consider the following setting: let D = {(0;, 8,, +)) be dataset of N three-tuples, one for each patient. Herein, æ; denotes the explana- tory variables, or features. &; is the censoring indica- tor: 5 0 if the survival time of the i-th patient was censored, and 8;
Cluster
Cluster 2
Output:
Cluster assignment: pfe 1ix.t) pic
Latent Space
otherwise. Finally, is the po- tentially censored survival time. A maximum likelihood approach to survival analysis seeks to model the survival distribution S(t|x) = (T t|x) (Cox, 1972). Two challenges of survival analysis are (i) the censoring of survival times and (iz) complex nonlinear relationship between x and t. When clustering survival data, additionally consider latent cluster assignment variable 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 æ; and c;.
Survival distribution: 1.50 5i1l.
Figure Summary the VaDeSC.
..., K) unobserved
Generative Model Following the problem definition above, we assume that data are generated from a random process consist- ing of the following steps (see Figure 3). First, a cluster as- signment p(c;
[L....,K} is sampled from a categorical distri-
ne. Then, continuous latent em- bedding, z E RJ, is sampled from Gaussian distribution, whose mean and variance depend on the sampled cluster (2|c; {MI,.., px) (E {E...., Σx}) = (z: p.s E,). The ex- planatory variables on 2: x ~ p(æ/2; y). where p(æ|z; 7) = Bernoulli(2; p-) for binary-valued features and N (æ; pyy diag (o²)) for real-valued features. Herein, Mr and o² are produced by f(z; 7) decoder neural network parameterised by 7. Finally, the survival time de- pends on the cluster assignment indicator 8, i.e. t ~~ p (t]2z.c). Similarly to conventional survival analysis, assume non-informative censoring (Rodríguez, 2007).
bution:
generated from distribution conditioned
&
Figure 3: Generative model.
latent vector 2, and censoring
Survival Model Above, p (t|z, refers to the cluster-specific survival model. We follow an approach similar Ranganath et al. (2016) and Liverani et al. (2020); in particular,
3

---

Published as a conference paper at ICLR 2022
sume that given and the uncensored survival time follows the Weibull distribution given by Weibull (softplus (x Be) k). where softplus(z)) = log (1 + exp(z)); k is the shape parameter; and Ge are cluster-speeific survival parameters. Note that we omitted the bias term 30,0 for the sake of brevity. Observe that softplus (z (.)) corresponds to the scale parameter of the Weibull distribution. We assume that the shape parameter k is global: however, an adaptation to cluster-specific parame- ters, as proposed by Liverani al. (2020), is straightforward. The Weibull distribution with scale
and shape & has the probability density function given by f(z; X, * (5) k- exp (- (x) k
for x 2 0. Consequently, adjusting for renht-censoring yields the following distribution:
p(||2, c; B, k) = f(;; X 1:) ((1),, c)'=
k
R-1
exp (-
exp (
(I)
where 3 = [B 3x X = softplus (z B2): and ((t) - f.a f(t: Xe, k) is the survival function. Henceforth, we will use p(1|z, c) as a shorthand notation for p(4|2, c; B,k). In this paper, only consider right-censoring: however, the proposed model can be extended tackle other forms of censoring.
Joint Probability Distribution Assuming the generative process described above, the joint prob- p(æ, 1, 2, c) = =ʃ- =f p(a))t, 2)) (1, e). It is important to note that æ and are independent given so are and c. Hence, we can rewrite the joint probability of the data, also referred to as the likelihood function, given the parameters π μ., E. y. B, k as
ability of x and can be written as p(r, 1== f 1
p(x, t; 1, E, 3,
/
/ p(al2; ()t(1 B, B,k)p(z|c n4, E)p(c T),
(2)
where /a (11,,-.. [/l1,-..kb.. E EE Ek), and B (0 Bk)
Evidence Lower Bound Given the data generating assumptions stated before, the objective is to infer the parameters π, μa, E, 7 and 3 which better explain the covariates and survival outcomes (æ,, /- Since the likelihood function in Equation 2 is intractable, we maximise a lower bound of the log marginal probability of the data:
E(p(: (3)
log p(x, t; nI, μ. E, 7, 2 Ez m.t) log
(2,,t
We approximate the probability of the latent variables and given the observations with varia- tional distribution q(z, c/æ t) = (()()))((()))) f), where the first term the encoder parameterised by neural network. The second term is equal to the true probability p(c|2,t):
p(z, t(c)p(c)
p(t c)p((z))p(())
((l2,,)) p(c|z,t)
(4)
Et p(z, t(c)p(c)
K
K p(tlz.c)l(zlcppll
Thus, the evidence lower bound (ELBO) can be written as
C(x,t) (=(())))(() log p(x|2; 7) + B logp(t|zz,; B,k) DRL. (q (2, () p (2,c; n,, x))
(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 the latent embeddings and cluster assignments q(2, c|ar, t). It can be then one assigned to one cluster. The ELBOO 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.
mixture of survival distributions, each
be computed from the distribution p(c|z,t) 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 p(=|c)p(e)
Missing Survival Time The hard cluster assignments
individual is unknown, using the Bayes' rule we instead compute p(c|z) Σe_ p(zce)p(c)
4

---

Published as a conference paper at ICLR 2022
4 EXPERIMENTAL SETUP
Datasets We evaluate VaDeSC on a range of synthetic, semi-synthetic (SurvMNIST: Pölsterl (2019)), and real-world survival datasets with varying numbers of data points, explanatory vari- ables, 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. (2021)); an observational co- hort 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. (2019); Bakr et al. (2017); Clark et al. (2013); Weikert et al. (2019)). Detailed descrip- tion of the datasets and preprocessing can be found Appendices E and G. For (semi- semii-syythhetic data, we focus on the clustering performance of the considered methods; whereas for real-world data, where the true cluster structure unknown, we compare time-to-event predictions. In addi- tion, we provide an in-depth cluster analysis for the NSCLC (see Section 5.3) and Hemodialysis (see Appendix H.8) datasets.
Table 2: Summary of the datasets. Here, N the total number of data points, D is the number of explanatory variables, K is the number of clusters if known. We report the percentage of censored observations and whether the cluster sizes are balanced if known.
Dataset 60,000
N
D
% censored
Data type
K
Balanced?
Section
Synthetic
1,000
30
Tabular
3
Y
5.1, H.2
SurvMNIST
70,000
28x28
52
Image
S
N
5.1, H.6
SUPPORT
9.105
59
32
Tabular
5.2
FLChain
6,524
7
70
Tabular
11.7
HGG
453
147
25
Tabular
5.2
Hemodialysis
1,493
57
91
Tabular
5.2, H.8
NSCLC
961
64
33
Image
5.3, H.9
Baselines & Ablations We compare our method to several well-established baselines: the semi- supervised 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 neural network techniques, we use the same encoder architectures and numbers of latent dimensions. Although the profile regression approach of Liverani et al. (2020) closely related to ours, 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 full-scale comparison is impossible due to computational limitations. Appendix H.3 contains 'down-scaled' experiment with the profile regression on synthetic data. Additionally, we consider k-means and regularised Cox PH and Weibull AFT models (Simon 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 VAE with auxiliary Weibull survival loss term, which is similar to the deep survival analysis (DSA; Ranganath et al. (2016) combined with k-means; (ii) training a completely unsupervised version without modelling the survival, which is similar to VaDE (Jiang et al., 2017); and (iti) predicting cluster assignments when the survival time unobserved. Appendix G contains further implementation details.
Evaluation We evaluate the clustering performance of models, when possible, in terms of accu- racy (ACC), normalised mutual information (NMI). and adjusted Rand index (ARI). For the time- to-event predictions, we evaluate the ability of methods to rank individuals by their risk using the concordance index (CI; Raykar et al. (2007)). Predicted median survival times are evaluated using the relative absolute error (RAE; Yu et al. (2011)) and calibration slope (CAL), as implemented by Chapfuwa et al. (2020). We report RAE on both non-censored (RAE,)) and censored (RAE.) data points (see Equations 14 and 15 in Appendix F). The relative absolute error quantifies the relative deviation of median predictions from the observed survival times; while the calibration slope in- dicates whether model tends to under- or overestimate risk average. For the (semi-) -eeth- datasets we average all results across independent simulations, i.e. dataset replicates; while for the real-world data we the Monte Carlo cross- cross-validae validation procedure.
5

---

Published as a conference paper at ICLR 2022
5 RESULTS
5.1 CLUSTERING
We first compare clustering performance on the nonlinear (semi-)syntheiis data. Table shows the results averaged across several simulations. In addition to the clustering, we evaluate the con- cordance index to verify that the methods can adequately model time-to-event in these datasets. Training set results are reported in Table 12 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 help identify clusters since the completely unsupervised VaDE achieves significantly worse results. 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 ("). Furthermore, performing k-means clustering in the latent space of VAE with the Weibull survival loss ("VAE + Weibull") clearly does not lead to the identification 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 essential role in inferring clusters.
Interestingly, while SCA and DSM achieve good results on SurvMNIST, both completely fail to iden- tify clusters correctly on synthetic data, for which the generative process is very similar the one assumed by VaDeSC. In the synthetic data, the clusters do not have prominently different survival distributions (see Appendix E. they are rather characterised by different associations between the covariates and survival times whereas the two baseline methods tend to discover clusters with disparate survival distributions. The SSC offers little no gain over the conventional k-means per- formed on the complete feature set. Last, we note that in both datasets VaDeSC has a significantly 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 MurVM data. It appears that SCA and DSM, 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
minor differences in KM curves and learns
Table 3: Test set clustering performance on synthetic and SurvMNIST data. "VAE Weibull" corresponds to an ablation of VaDeSC w/o the Gaussian mixture prior. "w/o t" corresponds to the cluster assignments made by VaDeSC when the survival time not given. Averages and standard deviations are reported across 5 and 10 independent simulations, respectively. Best results are shown in bold, second best in italic.
Dataset
Method
ACC
NMI
ARI
CI
k-means
0.44
0.06
0.05+F0.03
Cox PH
0.77+0.02
SSC
0.45.10.03
0.08-10.04
0.06-1:002
SCA
0.45
0.05+10.05
0.04-10.05
0.82+00.22
Synthetic
DSM
0.37++0.02
0.01+00.00
0.01+-0.00
0.76+10.02
VAE+ + Weibull
0.46=:0.06
0.09+1.00
0.09:00.0
0.71
VaDE
0.74+10.1
0.53+0..1
0.55+10.20
VaDeSC (w/o t)
0.8810.03
0.6010.07
0.6710.07
0.84
VaDeSC (ours)
0.90+0.02
0.66+00.5
0.73+0.05
k-means
0.4910.06
0.31
0.22-10.04
Cox PH
0.74
SSC
0.49
0.3140..4
0.22+0004
SCA
0.56+0.09
0.46+00.06
0.33+0.10
0. 0.79-0.06
SurvMNIST
DSM
0.54+0.11
0.40
0.31+0.14
0. 0.79+00.5
VAE + Weibull
0.49
0.32
0.24
0.76
VaDE
0.47+00007
0.38
0.24-10.08
VaDeSC (w/o t)
0.57+00.9
0.5110.00
0.3740.10
0.80=10.05
VaDeSC (ours)
0.18
0.55=10.11
0.39

---

RA Eme Published as a conference paper at ICLR 2022
Ou2er 2
(a) Ground truth (b) VAE Weibull
SCA
(d) DSM
VaDeSC (ours)
Figure 4: Cluster-specific Kaplan- Meierr (KM) curves (top) and t-SNE visualisation of latent rep- resentations (bottom), coloured according to survival times (yellow and blue correspond to higher and lower survival times, respectively), learnt by different models (b-e) from one replicate of the survMNIST dataset. Panel (a) shows KM curves of the ground truth clusters. Plots were generated using 10,000 data points randomly sampled from the training set; similar results were observed on the test set.
Table Test set ev-event performance on SUPPORT, HGG, and Hemodialysis datasets. Aver- ages and standard deviations are reported across 5 independent train-test splits.
Dataset
Method
CI
RAE
CAL
Cox PH
0.84.10.01
Wcibull AFT
0.84+0001
0.62++0.01
0.13+0011
1.27+00.02
SCA
0.83+:0.02
0.78+t0.13
0.06
1.74+00.5
SUPPORT
DSM
0.87+00011
0.56+00.2
0.13+004
1.43+0.07
VAE + Weibull
0.84-10.01
0.56110.02
0.20-10.02
1.28
VaDeSC (ours)
0.85110.11
0.53
0.23
1.24:10.05
Cox PH
0.74::0.05
Weibull AFT
0.74+:0.5
0.56+10.04
0.14
1.16
SCA
0.63:40.08
0.97+40.05
0.00
2.59
HGG
DSM
0.75+-0.04
0.57+00.05
0.18+00.7
1.09-+0.08
VAE + Weibull
0.75++0.05
0.52
0.12+0.07
1.14+-0.11
VaDeSC (ours)
0.74110.05
0.53
0.13
1.12:0.09
Cox PH
0.83+F0.04
Wcihull AFT
0.83++0.05
0.81+00.33
0.01
4.46+
SCA
0.75110.05
0.86-10.07
0.0210.02
7.23113.22
Hemodialysis
DSM
0.80110.06
0.85
0.02-10.44
8.23
VAE + Weibull
0.77-10.06
0.80-10.06
0.0210001
4.49
VaDeSC (ours)
0.80++0.05
0.78++0.05
0.01
3.74+
representations that clearly reflect the covariate and survival variability, Similar differences can be observed for the synthetic data (see Appendix H.2).
5.2 TIME-TO-EVENT PREDICTION
We now assess time-to-event prediction on clinical data. As these datasets do not have ground truth clustering labels, do not evaluate the clustering performance. However, for Hemodialysis, provide in-depth qualitative assessment of the learnt clusters in Appendix H.8.
Table shows the time-to-event prediction performance on SUPPORT, HGG, and Hemodialysis. Results for FLCCain 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 promising and suggest that VaDeSC remains competitive at time-to- tnee-to-event modelling. offers overall balanced predictions, and not prone to extreme overfitting even when applied to simple clinical datasets that low-dimensional contain few non-censored patients.

---

Published as a conference paper at ICLR 2022
5.3 APPLICATION TO COMPUTED TOMOGRAPHY DATA
We further demonstrate the viability of our model in real-world application using collection of several CT image datasets acquired from NSCLC patients (see Appendix E). The resulting dataset poses two major challenges. The first 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, representation learning model that captures cluster-specific associations between medical images and survival times could be highly beneficial to both explore the sources of variability in this dataset and to understand 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-even- prediction baselines, we fit Cox PH and Weibull AFT models on radiomics features. The latter are extracted using manual pixel-wise tumour segmentation maps, which time-consuming and expensive obtain. Note that neither DSM nor VaDeSC requires 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 slightly better calibration necessary for survival analysis on CT scans. Similarly tabular clinical datasets (see Table 4), DSM and VaDeSC have comparable performance. Therefore, we investigate cluster assignments qualitatively highlight the differences between the two methods.
average. This suggests that laborious manual tumour segmentation is
Table 5: Test set time-to-event prediction performance on NSCLC data. For Cox PH and Weibull AFT, radiomics features were extracted using tumour segmentation maps. Averages and standard deviations are reported across 100 independent train-test splits, stratified by survival time.
Method
CI
RAFm
RAE
CAL
Radiomics + Cox PH
0.60
Radiomics + Weibull AFT
0.60-1.02
0.70
0.45
1.26110.04
DSM
0.59-+0.04
0.72++0.03
0.34
1.24+:0.07
VaDeSC (ours)
0.02
0.71+-0.03
0.35+
1.21+
In Figure 5, we plot the KM curves and corresponding centroid CT images for the clusters discov- ered by DSM and VaDeSC on one train-test split. By performing several independent experiments (see Appendix H0), we observe that both methods discover patient groups with different empirical survival time distributions. However, contrast to DSM, VaDeSC clusters are consistently associ- ated with the tumour location. On the contrary, the centroid CT images of the DSM clusters show
Survival Probability
0.8
Survival Probability
0.8
Cluster
Cluster
Cluster
0.6
0.6
Cluster
0.4
0.2
0.0
0.0
0.0
0.2
0.4
0.6
0.8
1.0
0.0
0.2
0.4
0.6
0.8
1.0
Time
Time
(a) DSM
(b) VaDeSC (ours)
Figure 5: Cluster-specifie Kaplan-Mes Meier curves and corresponding centroid CT images, computed by averaging all samples assigned to each cluster by (a) DSM and (b) VaDeSC on the NSCLC data.
8

---

Published as a conference paper at ICLR 2022
Table 6: Cluster-specific statistics for a few demographic and clinical variables (not used during training) from the NSCLC dataset for DSM and VaDeSC. T. Vol. stands for the tumour volume; M1 denotes spread of cancer distant organs and tissues; and 2 T3 denotes tumour stage of least 3. Kruskal-War -allis Httest p-values are reported at the significance levels of 0.001, 0.01, and 0.05.
DSM
VaDeSC (ours)
Variable
2
p-val.
3
p-val.
Vol., cm
39
51
<
le-3
40
5e-2
Age, yrs
67
68
68
69
0.11
62
69
67
70
< le-3
Female, %
29
30
26
21
0.3
36
19
38
23
S le-3
Smoker, %
84
100
80
89
0.9
67
94
87
100
0.12
M1, %
40
55
16
42
0.4
20
45
44
45
0.2
2 T3,
12
32
0.2
10
35
0.7
no visible difference. This is further demonstrated in Figure 6, 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 clear association with tumour location, clusters spond to the upper section of the lungs. Indeed, multiple studies have suggested higher five-year survival rate in patients with the tumour in the upper lobes (Lee et al., 2018). also interesting to observe the amount of variability within each cluster: every generated scan is characterised by a unique lung shape and tumour size. Since DSM is not a generative model, we instead plot the original samples assigned to each cluster by DSM in Ap- pendix H.9. Finally, in Table 6 we compute cluster- specific statisties for a few important demographic and clinical variables (Etiz et al., 2002; Agarwal al., 2010; Bhatt al., 2014). We observe that VaDeSC tends to discover clusters with more dis- parate characteristics and thus able to stratify patients not only by risk but also by clinical con- ditions. Overall, the results above agree with our findings on fies clusters informed by both the covariates and survival outcomes, leading to stratification very different from previous approaches.
and
corre-
Mean
Generated Samples
Figure 6: CT images generated by (i) sampling latent representations from the Gaussian mix- ture learnt by VaDeSC and (ii) decoding rep- resentations using the decoder network.
data: VaDeSC identi-
6 CONCLUSION
In this paper, we introduced novel deep probabilistic model for clustering survival data. In contrast existing approaches, our method can retrieve clusters driven by both the explanatory variables and survival information and it can be trained efficiently 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 identifies meaningful clusters from a challenging medical imaging dataset. The analysis of these clusters provides interesting insights that could be useful for clinical decision-man decision-making. To conclude, the proposed VaDeSC model offers a holistic perspective on clustering survival data by learning structured representations which reflect the covariates, outcomes, and varying relationships between them.
Limitations & Future Work The proposed model has few limitations. requires fixing a number of mixture components a priori since global parameters are not treated in a fully Bayesian manner, as opposed to the profile 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
our model to facilitate its application in the medical domain and (ii) a fully Bayesian treatment of global parameters. In- depth interpretation of the clusters we found in clinical datasets beyond the scope of this paper, but we plan investigate these clusters together with our clinical collaborators in follow-up work.

---

Published as a conference paper at ICLR 2022
ACKNOWLEDGEMENTS
We thank Alexandru Tifrea, Nicolò Ruggeri, and Kieran Chin-Cheong for valuable discussion and comments and Dr. Silvia Liverani, Paidamoyo Chapfuwa, and Chirag Nagpal for sharing the code. Ričards Mareinkeviis is supported by the SNSF grant #320038 18909.. Laura Manduchi sup- ported by the PHRT SHFN grant #1 11-0000018-000 057: SWISSHEART.
ETHICS STATEMENT
The in-house PET/CT and HGG data were acquired at the University Hospital Basel and Univer- sity Hospital Zürich, respectively, and retrospective, observational studies were approved by local ethies committees (Ethikkommission Nordwest- und Zentralschweiz, no. 2016-01649; Kantonale Ethikkommission Zürich, no. PB-2017-00093).
REPRODUCIBILITY STATEMENT
To ensure the reproducibility of this work several measures were taken. The experimental setup outlined in Section 4. Appendix E details benchmarking datasets and simulation procedures. Appendix defines metrics used for model comparison. Appendix G contains data preprocess- ing, architecture, and hyperparameter tuning details. For the computed tomography data we as- sess the stability of the obtained results in Appendix H.10, Most datasets used for compari- son are either synthetic publicly available. HGG, Hemodialysis, and in-house PET/CT data could not be published due to medical confidentiality. The code is publicly available at https: //github. com/ 160924677777 vadees..
REFERENCES
M. Abadi, A. Agarwal, P. Barham, E. Brevdo, Z. Chen, C. Citro, G. S. Corrado, A. Davis, J. Dean, M. Devin, S. Ghemawat, I. Goodfellow, A. Harp, G. Irving, M. Isard, Y. Jia, R. Jozefow- icz, L. Kaiser, M. Kudlur, J. Levenberg. D. Mané, R. Monga, S. Moore, D. Murray, C. Olah, M. Schuster, J. Shlens, B. Steiner, Sutskever, K. Talwar, P. Tucker, V. Vanhoucke, V. Va- sudevan, F. Viégas, O. Vinyals, P. Warden, M. Wattenherg, M. Wicke, Y. Yu, and X. Zheng. Large-scale machine learning on heterogeneous systems, 2015. URL http: //tensorflow. org/.
TensorFlow:
H. J. W. L. Aerts, E. R. Velazquez, R. Leijenaar, C. Parmar, P. Grossmann, S. Cavalho, J. Bussink, R. Monshouwer, B. Haibe-Kains, D. Rietveld, Hoebers, M. Rietbergen, R. Leemans, A. Dekker, J. Quackenbush, R. Gillies, and P. Lambin. Decoding tumour phenotype by non- invasive imaging using a quantitative radiomics approach. Nature Communications, 5, 2014.
J. W. L. Aerts, E. Rios Velazquez, R. T. H. Leijenaar, C. Parmar, P. Grossmann, S. Carvalho, J. Bussink, R. Monshouwer, B. Kaibe-Kains, D. Rietveld, F. Hoebers, M. M. Rietbergen, C. R. Gillies, and P. Lambin. Data from NSCLC-radiomin genomics, 2015. URL attps://wikk www.eiii..aaneaeaaaaaaaeaeeaaeae
Leemans, A. Dekker, Quackenbush, cancer www.nneriimmiingrcccc.c.m
/ALLL...
H. W. I.. Aerts, L. Wee, Rios Velazquez, valho, J. Bussink, R. Monshouwer, B. Haibe-Kains D. Rietveld, F. Hoebers, M. M. Rietbergen, C. R. Leemans, A. Dekker, Quackenbush, J. Gillies, and Lambin. Data from NSCLC- radiomics, 2019. URL https: petttps:/wwikke /ttww/iiii///////// cancer www.iiggarccaaicmmmmmmmm hwwp://wwww......eeeeeeeeeeeeeee net/x/FgL1..
H. Leijenarr, Parmar, P. Grossmann, Car-
M. Agarwal, G. Brahmanday, G. W. Chmielewski, R. J. Welsh, and K. P. Ravikrishnan. Age, tumor size, type of surgery, and gender predict survival early stage (stage I and II) non-small cell lung cancer after surgical resection. Lung Cancer, 68(3):398 0029984--,,20 2010.
E. Ahlqvist, Storm, A. Käräjämäki, Martinell, M. Dorkhan, A. Carlsson, P. Vikman, R. B. Prasad, D. M. Aly, P. Almgren, Y. Wessman, N. Shaat, P. Spégel, H. Mulder, E. Lindholm, O. Melander, O. Hansson, U. Malmqvist, Á. Lernmark, K. Lahti, T. Forsén, T. Tuomi, A. H. Rosengren, and L. Groop. Novel subgroups of adult-onset diabetes and their association with outcomes: a data-driven cluster analysis six variables. The Lancet Diabetes Endocrinology, 6(5):361-369, 2018.
10

---

Published as a conference paper at ICLR 2022
S. Amaral, W. Hwang, B. Fivush, A. Neu, D. Frankenfield, and S. Furth. Serum albumin level and risk for mortality and hospitalizatton in adolescents on hemodialysis. Clinical journal of the American Society of Nephrology, 3(3):759-767, 2008.
E. Bair and R. Tibshirani. Semi-supervised methods to predict patient survival from gene expression data. PLoS Biology, 204) 2004.
Bakr, O. Gevaert, S. Echegaray, K. Ayers, M. Zhou, M. Shafiq, Zheng, Zhang, A. Le- ung, M. Kadoch, J. Shrager, A. Quon, D. Rubin, S. Plevritis, and S. Napel. Data for NSCLC radiogenomics collection, 2017. URL https: /ttt://wkkkk. cancerimagingarchieee net/
x x/WGGAQQ.
S. Bakr, O. Gevaert, S. Echegaray, K. Ayers, M. Zhou, M. Shafiq, H. Zheng, J. Anthony B., W. Zhang, A. N. C. Leung. Kadoch, C. D. Hoang, J. Shrager, A. Quon, Rubin, S. K. Plevritis, and S. Napel. A radiogenomic dataset of non-small cell lung cancer. Scientific Data,
(1), 2018.
G. A. Bello, T. J. W. Dawes, Duan, C. Biffi, Marvao, S. Howard, S. R. Gibbs, M. R. Wilkins, S. A. Cook, D. Rueckert, and D. P. O'Regan. Deep learning cardiac motion analysis for human survival prediction. Nature machine intelligence, 1:95-104, 2019.
A. Ben-Hur, A. Elisseeff, and I. Guyon. stability based method for discovering structure in clustered data. Pacific Symposium on Biocompuling. Pacific Symposium Biocomputing, pp- 6-17, 2002.
V. Bhatt, R. Batra, Silberstein, R. Loberiza, and A. K. Ganti. Effect of smoking on sur- vival from non-small cell lung cancer: a retrospective Veterans' Affairs Central Cancer Registry (VACCR) cohort analysis. Medical Oncology, 32(1), 2014.
Buckley and I. James. Linear regression with censored data. Biometrika, 66(3):429-4366 1979.
J. C. Buckner. Factors influencing survival in high- grade gliomas. Seminars in Oncology, 30:10-14 14, 2003.
Chapfuwa, C. Tao, C. Li, C. Page, B. Goldstein, L. C. Duke, and R. Henao. Adversarial time- to-event modeling. In Proceedings of the 35th International Conference on Machine Learning, volume 80, pp. 735-744. PMLR, 2018
Chapfuwa, C. Li, N. Mchta, L. Carin, and R. Henao. Survival cluster analysis. In Proceedings of the ACM Conference on Health, Inference, and Learning, pp. 60-68, Association for Computing Machinery, 2020.
Chen and C. Guestrin. XGBoost: A scalable tree boosting system. Proceedings of the 22nd ACM SIGKDD International Conference Knowledge Discovery and Data Mining, pp. 785- 794. ACM, 2016.
K. Clark, B. Vendt, K. Smith, Freymann, J. Kirby, P. Koppel, S. Moore, S. Phillips, D. Maf- fitt, M. Pringle, L. Tarbox, and F. Prior. The cancer imaging archive (TCIA): Maintaining and operating a public information repository. Journal of Digital Imaging, 26(6):1045- 1057, 2013.
S. Collins and Varmus. new initiative on precision medicine. New England Journal of Medicine, 372(9):793-795, 2015.
D. R. Cox. Regression models and life-tables. Journal of the Royal Statistical Society: Series B (Methodological), 3472):177-22,, 1972.
D. G. Altman. Analysis of survival times. In Practical Statistics for Medical Research, chapter 13, pp. 365-395. Chapman and Hall/CRC Texts in Statistical Science Series, 2020.
J. Daugirdas. Second generation logarithmic estimates of single-pool variable volume Kt/V: an analysis of error. Journal of the American Society of Nephrology JASN, ((5)::205-1111 1993.
J. Daugirdas and D. Schneditz. Overestimation of hemodialysis dose depends dialysis efficiency by regional blood flow but not by conventional two pool urea kinetic analysis. ASAIO journal, (3):M719-M77) 1995.
11

---

Published as a conference paper at ICLR 2022
C. Davidson- Pilon, J. Kalderstam, N. Jacobson, S. Reed, B. Kuhn, P. Zivich, M. Williamson, J. K. Abdeali, D. Datta, A. Fiore-Garllnn, A. Parij, D. Wilson, Gabriel, L. Moneda, A. Moncada- Torres, K. Stark. H. Gadgil, Jona, K. Singaravelan, L. Besson, M. Sancho Peña, S. Anton, A. Klintherg, J. Growth, J. Noorbakhsh, M. Begun, R. Kumar, S. Hussey, S. Seabold, and D. Gol- land. lifelines: v0.25.9, 2021. URL hww.h:tttt..cc....t. /ooi org/ /0/./1...00 008155811222e
4505728.
N. Dilokthanakul, P. A. M. Mediano, M. Garnelo, M. C. H. Lee, H. Salimbeni, K. Arulkumaran, and M. Shanahan. Deep unsupervisee clustering with Gaussian mixture variational autoencoders, 2016.
A. Dispenzieri, J. A. Katzmann, R. A. Kyle, D. R. Larson, T. M. Thernea,, C. L. Colby, R. J. Clark, G. P. Mead, S. Kumar, L. J. Melton, and S. V. Rajkumar. Use of nonclonal serum immunoglobulin free light chains to predict overall survival in the general population. Mayo Clinic Proceedings, 87(6): 2012.
D. Etiz, L. B. Marks, S.-M. Zhou, G. C. Bentel, R. Clough, M. L.. Hernando, and P. A. Lind. Influence of tumor volume on survival in patients irradiated for non-small-cell lung cancer. Inter- national Journal of Radiation Oncology, Biology, Physics, 53(4):835-846, 2002.
D. Faraggi and R. Simon. A neural network model for survival data. Statistics in Medicine, 14(1): 73-82, 1995.
V. T. Farewell. The use of mixture models for the analysis of survival data with long-term survivors. Biometrics, 38(4): 1982.
D. A. Fenstermacher, R. M. Wenham, D. E. Rollison, and W. S. Dalton. Implementing personalized medicine cancer center. The Cancer Journal, (7(6))528-538 2011.
R. Foley, P. Parfrey, Harnett, G. M. Kent, D. Murray, and P. Barre. Hypoalbuminemia, cardiac morbidity, and mortality in end-stage renal disease. Journal of the American Society of Nephrol- ogy: JASN, 5:728-36 -6888-6,111 1996.
O. Gevaert, J. Xu, C. D. Hoang, A. N. Leung, Y. Xu, A. Quon, D. L. Rubin, S. Napel, and S. K. Plevritis. Non-small cell lung cancer: Identifying prognostic imaging biomarkers by leveraging public gene expression microarray data methods and preliminary results. Radiology, 264(2): 387-396, 2012.
M. Goldstein, X. Han, A. Puli, A. Perotte, and R. Ranganath. X-CAL: Explicit calibration for survival analysis. Advances in neural information processing systems, 33: 13:11896-18807 2020.
V. Gotta, O. Marsenic, and M. Pfister. Age- and weight- based differences in haemodialysis prescrip- tion and delivery in children, adolescents and young adults. Nephrology Dialysis Transplantation, 33:1649-1660, 2018
V. Gotta, O. Marsenic, and M. Pfister. Understanding urea kinetic factors that enhance personalized hemodialysis prescription in children. ASAIO Journal, 66:115 123, 2019a.
V. Gotta, M. Pfister, and O. Marsenic. Ultrafiltration rates in children on chronic hemodialysis routinely exceed weight based adult limit. Hemodialysis International, 23, 2019b.
V. Gotta, G. Tancev, O. Marsenic, J. E. Vogt, and M. Pfister. Identifying key predictors mortality in young patients chronic haemodialysis machine learning approach. Nephrology, dialy- sis, transplantation: : official publication of the European Dialysis and Transplant Association European Renal Association, 2020.
V. Gotta, O. Marsenic, A. Atkinson, and M. Pfister. Hemodialysis (HD) dose and ultrafiltration rate associated with survival in pediatric and adolescent patients chronic hd-a large observa- tional study with follow-up to young adult age. Pediatric Nephrology, pp. 1-12, 2021.
C. Haarburger, P. Weitz, O. Rippel, and D. Merhof. Image-based survival prediction for lung cancer patients using CNNs. In 2019 IEEE 16th International Symposium on Biomedical Imaging (ISBI 2019). IRFE, 2019.
12

---

Published as a conference paper at ICLR 2022
H. Ishwaran, U. B. Kogalur, E. H. Blackstone, and M. S. Lauer. Random survival forests. Annals of Applied Statistics, 2(3):841-860) 2008.
R. A. Jacobs, M. I. Jordan, S. J. Nowlan, and G. E. Hinton. Adaptive mixtures of local experts. Neural Computation, 3(1):79-87, 1991.
Z. Jiang, Y. Zheng, H. Tan, B. Tang, and H. Zhou. Variational deep embedding: An unsupervised and generative approach to clustering. In Proceedings of the 26th International Joint Conference Artificial Intelligence, pp. 1965-1972. AAAI Press, 2017.
M. J. Johnson, D. K. Duvenaud, A. Wiltschko, R. P. Adams, and S. R. Datta. Composing graphical models with neural networks for structured representations and fast inference. In Advances in Neural Information Processing Systems, volume 29, 2946-2954, Curran Associates, Inc., 2016.
J. L. Katzman, U. Shaham, A. Cloninger, J. Bates, T. Jiang, and Y. Kluger. DeepSurv: Personalized treatment recommender system using a Cox proportional hazards deep neural network. BMC Medical Research Methodology, 18(1), 2018.
D. Kingma and M. Welling. Auto-encoding variational Bayes. In 2nd International Conference on Learning Representations, 2014.
W. A. Knaus, E. Harrell, J. Lynn, Goldman, S. Phillips, A. F. Connors, N. V. Dawson, W. J. Fulkerson, R. M. CalifI. N. Desbiens, P. Layde, R. K. Oye, P. E. Bellamy, R. B. Hakim, and D. P. Wagner. The SUPPORT prognostic model: Objective estimates of survival for seriously ill hospitalized adults. Annals of Internal Medicine, 122(3):191-22) 1995.
H. Kvamme, Ø. Borgan, and I. Scheel. Time-to-event prediction with neural networks and Cox regression. Journal of Machine Learning Research, 20(129):1-30. 2019.
R. A. Kyle, T. M. Therneau, S. V. Rajkumar, D. R. Larson, M. F. Plevak, J. R. Offord, A. Dispenzieri, J. A. Katzmann, and L. J. Melton. Prevalence of monoclonal gammopathy of undetermined significance. New England Journal of Medicine, 354(13);1362-136) 2006.
Lange, V. Roth, M. L. Braun, and M. Buhmann. Stability-based validation of clustering solu- tions. Neural Computation, 16: 199--133,,,, 2004.
Y. LeCun, C. Cortes, and Burges. MNIST handwritten digit database. ATT Labs, 2, 2010. URL http:/ /yannnnn... lecun. com/ extpb////////////////////////t// mt.st..
H. W. Lee, C.-H. Lee, and Y. S. Park. Location of stage 1-I1II non-small cell lung cancer and survival rate: Systematic review and meta-analysic Thoracic Cancer, 912):1614-1622, 2018.
Lehr and Capek. Histogram equalization of CT images. Radiology, 154(1):163-16) 1985.
S. Liverani, D. I. Hastie, L. Azizi, M. Papathomas, and S. Richardson. PReMiuM: An R package for profile regression mixture models using Dirichlet processes. Journal of Statistical Software, 64(7), 2015.
Liverani, L. Leigh, Hudson, and J. E. Byles. Clustering method for censored and collinear survival data. Computational Statistics, 2020.
Lundberg and S.-I. Lee. A unified approach to interpreting model predictions. Advances in Neural Information Processing Systems 30, pp. 4765-4774, Curran Associates, Inc., 2017.
M. Lundberg, Erion, H. Chen, A. DeGrave, M. Prutkin, B. Nair, R. Katz, Himmelfarb, N. Bansal, and S.-L. Lee. From local explanations to global understanding with explainable AI for trees. Nature machine intelligence, 2(1):56-67, 2020.
Y. Luo, T. Tian, J. Shi, J. Zhu, and B. Zhang. Semn-crowdsourredd clustering with deep generative models. Advances in Neural Information Processing Systems, volume 31, pp. 3212-3222. Curran Associates, Inc., 2018.
13

---

Published as a conference paper at ICLR 2022
L. Manduchi, K. Chin-Cheong, H. Michel, S. Wellmann, and J. E. Vogt. Deep conditional gaussian mixture model for constrained clustering, In Advances in Neural Information Processing Systems, volume 35, 2021.
O. Marsenic, M. Anderson, and K. G. Couloures. Relationship between interdialytic weight gain and blood pressure in pediatric patients on chronic hemodialysis. BioMed Research International, 2016.
G. J. McLachlan and D. Peel. Finite mixture models. John Wiley & Sons, 2004.
E. Min, X. Guo, Q. Liu, G. Zhang, J. Cui, and J. Long. A survey of clustering with deep learning: From the perspective of network architecture. IEEE Access, 6:39501-39514, 2018.
Molitor, M. Papathomas, M. Jerrett, and S. Richardson. Bayesian profile regression with an application to the national survey of children's health. Biostatistics, 11(3):484 498, 2010.
S. C. Mouli, B. Ribeiro, and J. Neville. A deep learning approach for survival clustering without end-of-life signals, 2018 URL https: https///pppeevieeeeeeee net/ forum wiwiammiii---ii
2R-
E. Movilli, Gaggia, R. Zubani, C. Camerini, Valerio Vizzardi, G. Parrinello, S. Savoldi, M. Fis- cher, F. Londrino, and G. Cancarini. Association between high ultrafiltration rates and mortality in uraemic patients on regular hacmodialysis. A 5-year prospective observational multicentre study. Nephrology, dialysis, transplantation: Official publication of the European Dialysis and Trans- plant Association European Renal Association, 22(12)354473352 2007.
C. Nagpal, X. R. Li, and A. Dubrawski. Deep survival machines: Fully parametric survival re- gression and representation learning for censored data with competing risks. IEEE Journal of Biomedical and Health Informatics, 2021a
C. Nagpal, S. Yadlowsky, N. Rostamzadeh, and K. Heller. Deep Cox mixtures for survival regres- sion, 2021b. arXiv:2101 100.26661...66....
F. Pedregosa, G. Varoquaux, A. Gramfort, V. Michel, B. Thirion, O. Grisel, M. Blondel, P. Pretten- hofer, R. Weiss, V. Dubourg, J. Vanderplas, A. Passos, D. Cournapeau, M. Brucher, M. Perrot, and E. Duchesnay. Scikit-learn: Machine learning in Python. Journal of Machine Learning Research, 12:2825-2830, 2011.
L. Perez and J. Wang. The effectiveness of data augmentation in image classification using deep learning, 2017, arXiv:1712.04621.
S. Pölsterl. Survival analysis for deep learning, 2019. URL
201 9/07/ survivall. vanianall wwwl-aalllsssssaaaaoaaa for---for-eep- le tearning//
Core Team. R: A Language and Environment for Statistical Computing. R Foundation for Statis- tical Computing, Vienna, Austria, 2020. URL https: /www. R- pr--projee jerrjjectooom org/.
R. Ranganath, A. Perotte, N. Elhadad, and D. Blei. Deep survival analysis. In Proceedings of the 1st Machine Learning for Healthcare Conference, volume 56, pp. 101-114. PMLR, 2016.
V. C. Raykar, H. Steck, B. Krishnapuram, C. Dehing--berijyy and P. Lambin. On ranking in survival analysis: Bounds on the concordance index. In Proceedings of the 20th International Conference on Neural Information Processing Systems, pp. 1209-1216. Curran Associates Inc., 2007.
D. Rezende, S. Mohamed, and D. Wierstra. Stochastic backpropagation and approximate infer- ence in deep generative models. In Proceedings of the 31 International Conference on Machine Learning, volume 32, pp. 1278-1286 PMIR, 2014.
R. Rodríguez. Lecture notes on generalized linear models, 2007. URL https::/data. princeton wwtpwwwwww/weeeccwwwes //tees/////////eeeoe
O. Rosen and M. Tanner. Mixtures of proportional hazards regression models. Statistics in Medicine, 18(9): 119-1131, 1999.
L. S. Shapley. A value for n n-person games. In Contributions to the Theory of Games, volume 2, pp. 307-317. Princeton University Press, 1953.
14

---

Published as a conference paper at ICLR 2022
N. Simon, J. Friedman, T. Hastie, and R. Tibshirani. Regularization paths for Cox's proportional hazards model via coordinate descent. Journal of Statistical Software, 39(5), 2011.
K. Simonyan and A. Zisserman. Very deep convolutional networks for large-scale image recogni- tion. In 3rd International Conference on Learning Representations, 2015.
Tanniou, van der Tweel, S. Teerenstra, and K. C. B. Roes. Subgroup analyses confirmatory clinical trials: time to be specific about their purposes. BMC Medical Research Methodology, 16 (1), 2016.
L. van der Maaten and G. Hinton. Visualizing data using t-SNE. Journal of Machine Learning Research, 9(86):2579-2600 2008.
J. M. van Griethuysen, A. Fedorov, C. Parmar, A. Hosny, N. Aucoin, V. Narayan, R. G. H. Beets- Tan, J.-C. Fillion-Robin. S. Pieper, and H. J. W. L. Aerts. Computational radiomics system to decode the radiographic phenotype. Cancer Research, 77211)e104-1107 2017.
J. Weikert, T. A. Antonoli, Bremerich, B. Stieltjes, Sommer, A. W. Sauter. Evaluation of an Al-powered lung nodule algorithm for detection and 3D segmentation of primary lung tumors. Contrast Media & Molecular Imaging. 2019.
M. Weller, M. van den Bent, M. Preusser, E. Le Rhun, J. C. Tonn, G. Minniti, M. Bendszus, C. Bal- ana, O. Chinot, L. Dirven, P. French, M. E. Hegi. A. S. Jakola, M. Platten, P. Roth, R. Rudå, S. Short, M. Smits, M. J. B. Taphoorn, A. von Deimling, M. Westphal, R. Soffietti, G. Reifen- berger, and W. Wick. EANO guidelines the diagnosis and treatment of diffuse gliomas of adulthood. Nature Reviews Clinical Oncology, 18(3):170-186) 2020.
M. Winkel. Statistical lidelsme-models, 2007. URL http:/ //////wwww httt:/tttst....t.ttttttttt...ss. stats. ox. ac .k/ wi nike //bs3b00_11111 B. pdd. Lecture notes. University of Oxford.
E. Xia, X. Du, Mei, W. Sun, Tong, Z. Kang, J. Sheng. Li, C. Ma, Dong, and S. Outcome-driven clustering of acute coronary syndrome patients using multi-task neural network with attention, 2019. arXiv:1903.00197.
C.-N. Yu, R. Greiner, H.-C. Lin, and V. Baracos. Learning patient-specific cancer survival distri- butions sequence of dependent regressors. In Advances Neural Information Processing Systems, volume 24. Curran Associates, Inc., 2011.
15

---

Published as a conference paper at ICLR 2022
A CENSORING IN SURVIVAL ANALYSIS
Censoring is one of the key features of survival anal- ysis that distinguish 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 ob- servation period. The latter units are referred to as censored, and the observation time for these units is lower bound on the time-to-event. The censoring event deseribed above is commonly known as right censor- ing (Winkel, 2007) and is the only form of censoring considered throughout the current paper. Figure 7 de- picts 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 non-informative, i.e.
Dead
Consored
1000 1500
2000
since Enrolment,
Figure 7: Survival time, in days, measured from the enrolment in the study for a few subjects from an observational cohort of high-grade glioma patients. Here, x corre- sponds to death (the event of interest), and denotes censoring.
of the schemes assume that that the censoring time is independent of unit's sur- vival beyond it (Rodríguez, 2007). This assumption is made by the VaDeSC model as well.
B RELATED WORK
Herein, we provide a detailed overview of the related work. The proposed VaDeSC model (see Section 3) builds on three major topies: (i) probabilistic deep unsupervised clustering, (ii) mixtures of regression models, and (ii) survival analysis.
Deep Unsupervised Clustering Long- 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 nigh-dimensiooal data (Min et al., 2018). Several techniques have been proposed in the literature, featuring wide range of neural network architectures. Among them, the VAE (Kingma & Welling, 2014; Rezende et al., 2014) combines variational Bayesian methods with the flexibility and scalability of neural networks. Gaussian mixture variational autoencoders (GMM- VAE; Dilokthanakul et al. (2016)) and variational deep embedding (VaDE; Jiang et al. (2017)) are variants of the VAE in which the prior 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 the inclusion of side information in the form of pairwise constraints (Manduchi et al., 2021), overcoming the restrictive i.i.d. assumption of the samples, or on a fully Bayesian treatment of global parameters (Luo et al., 2018; Johnson et al., 2016), alleviating the need to fix the number of clusters a priori.
Mixtures of Regression Models Mixtures of regression (MeLachlan & Peel, 2004) model a re- sponse variable by mixture of individual regressions with the help of latent cluster assignments. Such mixtures do not have to be limited to finite number of components; for example, profile re- gression (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 intro- duced by Jacobs al. (1991).
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 fitting 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 fit 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 first machine learning approaches for survival analysis was Faraggi- Simon's '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
16

---

Published as a conference paper at ICLR 2022
developed: random survival forests (Ishwaran et al., 2008), deep survival analysis (Ranganath al., 2016), neural networks based on the Cox PH model (Katzman et al., 2018; Kvamme et al., 2019), adversarial time-t--event modelling (Chapfuwa et al., 2018), and deep survival machines (Nagpal et al., a0) etc. An exhaustive overview of these and other techniques beyond the scope of the current paper.
C COMPARISON WITH RELATED WORK
Generative Assumptions In addition to the differences summarised in Table 1, VaDeSC and the related approaches (Chapfuwa et al., 2020; Nagpal et al., 202 la; Liverani et al., 2020) assume differ- ent data generating processes. Simplified graphical models assumed by these techniques are shown in Figure 8. Both SCA and DSM are not generative w.r.t. explanatory variables On the other hand, both profile
and thus, do
maximise the joint likelihood C (x.t) and are outcome-drive regression and VaDeSC specify model for covariates as well as for survival time I. However, PR does not infer latent representations z and importantly, in its simplest form assumes that x A/c. The latter assumption can limit the discriminative power of mixture with few components. By default, VaDeSC does not make such restrictive assumptions.
(a) SCA
(b) DSM
(c) PR
(d) VaDeSC (ours)
Figure 8: Comparison of the generative assumptions made by the models related to the VaDeSC: (a) survival cluster analysis (SCA; Chapfuwa et al. (2020). (b) deep survival machines (DSM; Nag- pal al. (2021a)), (c) and profile regression for survival data (PR; Liverani et al. (2020)). For the sake of brevity, the censoring indicator 6, model parameters and hyperparameters were omitted. Here, t denotes survival time, æ denotes explanatory variables, corresponds to the latent represen- tation, and stands for the cluster assignment. Shaded nodes correspond to observed variables. The presence of dashed edges depends on the modelling choice.
Deep Cox Mixtures The concurrent work by Nagpal et al. (2021b) on deep Cox mixtures allows leveraging latent representations learnt by a VAE capture nonlinear relationships in the data. The hidden representation is simultaneously used to fit mixture of K Cox models. In contrast to our method, DCM does not specify generative model explicitly. The VAE loss 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 mainly used to enforce Gaussian prior in the latent space, resulting in more compact representations. It can be replaced with regularised AE, or even MLP encoder, as mentioned by the authors (Nagpal et al., 202 Thus, the maximisation L (x,t). rather than of L(t), is not the main emphasis of the DCMs. On the contrary, probabilistic and specifies a clear and interpretable generative process from which an ELBO of the joint likelihood be derived formally. Additionally, VaDeSC enforces more structured representation in the latent space, since uses a Gaussian mixture prior.
the joint likelihood approach
Empirical Contribution Last but not least, the experimental setup of the current paper different from those of Bair & Tibshirani (2004); Chapfuwa et al. (2020); Liverani et al. (2020); 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).
17

---

Published as a conference paper at ICLR 2022
D ELBO TERMS
In this appendix, we provide the complete derivation of the ELBO of VaDeSC using the SGVB estimator. From Equation 5 rewrite the ELBO
C(x, t) = (=(=))))(() log æ; 7) + logp(t|z, (, + logp(zle; E)
(6)
+ logg(c; T) log c(,,))
In the following we investigate each term of the equation above.
The first term is also known as the reconstruction term classical VAE set- ting. Herein, we provide the derivation for p(z|>; 7) = Bernoullixz; pyy, Lee for binary-valued explanatory variables:
Reconstruction
E (o( logp(z|2; 7) B)(((() logp(z/z; 7) log p 00; 7)
Fa log yy, + (1 xa) log(1 - Hlya)
L d=1
where is the number of Monte Carlo samples in the SGVB estimator; 2a f(x(); 7): 20) (<]|==) N(2; Ho:r)): [no, log r3) g(x; 0) and g(: 0) is encoder neural network parameterised by x. refers to the d-th component of æ and D is the number of features. For other data types, the distribution p(æ z; y) 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
(x logp(t|z,,; c; B.
log logp(t/> == = B,k)p(c 112(11
(8)
K log
plc= - 112(9,0)
15
exp -
A
where 25 softplus (2()T T B.))
Clustering The third term in Equation 6 corresponds the clustering loss and is defined
E (l logp(z/c, c; ,4, Σ) Byy
plc 0| log 1ogg(2/c== 14, E)
(9)
p(c= = ,0 log ,1, E).
Prior
The fourth term be approximated
By log,(c; т)
p(c= f|2(0)1,111))))) log p(c v; πτ)
(10)
Variational Finally, the fifth term of the ELBO corresponds to the entropy of the variational dis- tribution q(2, l):
- (ol log log9(2, c/a, t) (e Dlg <(=[=) + log p(cl2, =y log g(z (x) (B(())))((c ((c)(()),))
(11)
18

---

Published as a conference paper at ICLR 2022
where the first term is the entropy of multivariate Gaussian distribution with diagonal covariance matrix. We can then approximate the expression in Equation 11 by
K EE
1 T.
(10g(2T) + 1) +
log OB
plc u/ (1220,11))0000000)11 logglc v/2(9,1) (12)
2
1=
where J is the number of latent space dimensions.
E DATASETS
Below 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 multilayer perceptron and is, therefore, nonlinear. Appendix E. contains 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ölsterl (2019), used for benchmarking nonlinear survival analysis models, egg. by Goldstein et al. (2020). 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-specific scale parameter. Appendix E.2 provides further details on the generation procedure.
SUPPORT To compare models real clinical data, consider the Study Understand Prog- noses and Preferences for Outcomes and Risks of Treatments (SUPPORT; Knaus et al. (1995)) consisting of seriously ill hospitalised adults from five 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 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. FI.Chain is by far the most low- dimensional 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), type of brain tumour (Buckner, 2003). The cohort includes patients with glioblastoma multiforme and anaplastic astroeytoma (Weller et al., 2020) treated surgically the University Hospital Zürich between 03/2008 and 06/2017. The dataset con- sists of demographic variables, treatment assignments, pre- and post-operative volumetric variables, molecular markers, performance scores, histology findings, 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.
Hemodialysis Another clinical dataset we include is an observational cohort of patients from the DaVita Kidney Care (Da Vita Inc., Denver, CO, USA) dialysis centres (Gotta et al., 2018: 201 20;9;;;; 2020; 2021). The cohort is composed of patients who started chronic hemodialysis (HD) in child- hood 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 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), fluid removal in terms of UFR (mL/kg/h) and total ultrafiltration (UF, % per kg dry weigth per session), and the in- terdialytic weight gain (IDWG) (Marsenic et al., 2016). It should be noted that this dataset is rather
19

---

Published as a conference paper at ICLR 2022
challenging. given the high percentage of censored observations (see Table 2). The scientific use of the deidentified standardised electronic medical records was approved by DaVita (DaVita Clinical Research@, Minncapolis, MN); IRB approval was not required since retrospective analysis was per- formed on the deidentified dataset. complete description the variables and measurements used provided by Gotta et al. (2021) in the Methods section.
NSCLC Finally, we perform experiments on a pooled dataset of five observational cohorts con- sisting of patients diagnosed with non-small cell lung cancer (NSCLC). We consider large set of pretreatment CT scans and CT components of positron emission tomography (PET//C scans from the following sources:
An in-house PET/CT dataset consisting of 392 patients treated at the University Hospital Basel (Weikert et al., 2019). It includes manual delineations, clinical and survival data. A retrospective, observational study was approved by the local ethics committee (Ethikkom- mission Nordwest- und Zentralschweiz, no. 2016-01649).
Lungl 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 The Cancer Imaging Archive (TCIA; Clark al. (2013)).
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 profiles. It is is publicly available in TCIA.
NSCLC Radiogenomics dataset acquired from 211 NSCLC patents from the Stanford Uni- versity School of Medicine and Palo Alto Veterans Affairs Healthcare System (Bakr al., 2017; 2018; Gevaert et al., 201 2). The dataset comprises CT and PET/CT scans, segmenta- tion 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 (PETyCT with the maximum tumour area outside the lungs, e.g. in the brain or abdomen, are excluded from the analysis. Thus, the final dataset includes 961 subjects.
E. SYNTHETIC DATA GENERATION
Herein, 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:
1. Let n. 1 for S c K.
2. Sample a ~ (at(TT), for 1 Si S N.
3. Sample Me, unif (-2 1). for 1 cos K and و J.
4. Let E. = Sc, where S. E J is a random symmetric, positive-definite matrix,' for
1 K.
5. Sample z; ~ N (Ha, Ea), for S.
6. Let (z) ) = W2ReLU (AReLU(W,RR)) (WR (w + bp) - bb) + b2, where Wo E RXX Wi E R W2 RDxh and bo R² bi Rh ba RDD are random matrices and vectors."
and h is the number of hidden units.
7. Let x; (x,) for 1 1 N.
8. Sample Be, unif (-10, 10). for S K and 0 05isJ.
9. Sample u, ~ Weibull (softplus (z] B...) k). for N.
'Generated using make spd matrix from scikit-learn (Pedregosa et al., 2011).
"Wo Wi, W. are generated using make oow rank naxrix from scikit-lcam (Pedregosa et al., 201 1). We omitted bias terms for the sake of brevity.
20

---

Published as a conference paper at ICLR 2022
10. Sample 5, Bernoulli1 Peens); for 1 N.
11. Let t, = u, 6, = 1, and sample t, ~ unif(0, u,) otherwise, for 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; is the number of features; &o is the shape parameter of the Weibull distribution (see Equation 1); and Peens denotes the probability of censoring. The key difference between clusters in how latent variables 2 generate the survival time I: each cluster has a different weight vector 33 that determines the scale of the Weibull distribution.
In our experiments, we fix K = 3, N = 60000, J = 16, D = 1000, == 1. and Peens 0.3. We hold out 18000 data points as a test 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.
Cluster
1:3ke Dimeesiin 2
Proteeteyy
20
0.6
-20
-40
Clatter Cluster Cluster
-60
-80
-40
-20
0.0
300
(a) E-SNE plot
(b) Kaplan-Meier plot
Figure 9: Visualisation of one replicate of the synthetic data. t-SNE plot (on the left) is based on the explanatory variables only. Kaplan-Meier curves (on the right) are plotted separately for each cluster. Different colours correspond different clusters. As can be seen, clusters are determined by both differences in (a) covariate and (b) survival distributions.
E.2 SURVIVAL MNIST DATA GENERATION
We generate semi-synthetic SurvMNIST data as described by Pölsterl (2019) (original implementa- tion available at https: / gut..b...... com/ sebp/ swrv/vurviallllea enn-ennn---eee est imator). The gener- ative process can be summarised by the following steps:
1. Assign every digit in the original MNIST dataset (LeCun al., 2010) to one of K clusters (ensure that every cluster contains at least one digit)
2. Sample risk score Ve unif (2, 15), for S c K.
3. Let > + exp(r.), for 1 S c K, where Ty is the specified mean survival time and 1/2 is the baseline hazard.
4. Sample a; unif(0, 1) and let Pom. (uI:N), ), where (-) denotes the a-th quantile.
log a, for 1 <ii < N. Ac.
5. Let 9eens = Pom
6. Sample feems ~ unif (minigigN U4 (.ems).
us < feens, and 8, otherwise, for 5 N. 8. Let f; u; if 8; = 1, and f; = fgems otherwise, for 1 5 S N.
7. Let 6;
Observe that here Peens is only lower bound on the probability of censoring and the fraction of censored observations can much larger than Peens-
21

---

Published as a conference paper at ICLR 2022
Table An example of an assignment of MNIST digits to clusters and risk scores under K 5.
In our experiments, fix K 5 and Peens 0.3. We repeat simulations independently times. The train-test split is defined in the original MNIST data (60,000 10,000). Table 7 shows an example of cluster and risk assignments in of the simulations. The resulting dataset is visualised in Figure 10. Observe that the clus- ters are not well-distinguishable based on covariates alone. Also note that some clusters contain many censored data points, whereas some do not contain any, shown in the Kaplan-Meier plot.
Cluster
Digits
Risk Score
{4.8) 8)
7.62
2
{0.3.6}
13.12
3
(1,9}
5.17
4
(2,7)
11.60
5
{5}
2.03
L-3RE Drmaniioz
25-
Suroeel Frobebilty
0.6
-25
0.4
-50
-75
Clanter
-30
-00
400
-20
20
40
80
80
0.2
0.6
t.S Dimension
Time
(a) t-SNE plot
(b) Kaplan-Meier plot
Figure 10: Visualisation of one replicate of the survMNIST dataset. Generally, SurvMNIST tends to have more disparate survival distributions across clusters than the synthetic data (ef. Figure 9).
F EVALUATION METRICS
Below 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) implemented in the sklearn. met rict module of scikit-lear (Pedregosa et al., 2011). Clustering accuracy is evaluated by finding an optimal map- ping between assignments and true cluster labels using the Hungarian algorithm, as implemented in sk le rn. ut i s. inear ass gnment...
The concordance index (Raykar al., 2007) evaluates how well survival model is able rank individuals in terms of their risk. Given observed survival times t,, predicted risk scores 73+ and censoring indicators d; the concordance index is defined
1, <6. 17
1= E 1,,
CI
(13)
Thus, perfect ranking achieves a CI of 1.0, whereas a random ranking is expected to have a Cl of 0.5. We use the lifelines implementation of the CI (Davidson-Pinn
al., 2021).
Relative absolute eтror (Yu et al., 2011) evaluates predicted survival times in terms of their relative deviation from the observed survival times. particular, given predicted survival times f,, for non- censored data points, RAE is defined as
ΣN (t, - 12) /2, 6z
RAE
(14)
ΣN1 6;
On censored data, similarly to Chapfuwa et al. (2020), we define the metric as follows:
RAE Σ (b. - 4)) / (1 5,) 1
(15)
ΣN (1 - 6z)
22

---

Published as a conference paper at ICLR 2022
Last but not least, similar to Chapfuwa al. (2020) the calibration slope to evaluate the overall calibration of models. It desirable for the predicted survival times to match with the empirical marginal distribution of t. Calibration slope indicates whether model tends to under- or overestimate risk on average, and thus, an ideal calibration slope is 1.0.
G IMPLEMENTATION DETAILS
Herein, we provide implementation details for the experiments described Section 5.
G.1 PREPROCESSING
For all datasets, we re-scale survival times be in [0.001, 1.001]. For SurvMNIST, we re-scale the features to be in [0, 1]. For tabular datasets, re-scale the features to zero-mean and unit- variance. Categorical features are encoded using dummy variables. For Hemodialysis data, the same preprocessing routine is adopted as described by Gotta et al. (2021).
For NSCLC CT images, preprocessing consists of several steps. For each subject, slices within 15 mm from slice with maximal transversal tumour area are averaged to produce single 2D image. 2D images are then cropped around the lungs and normalised. Subsequently, histogram equalisation applied (Lehr & Capek, 1985). The images then downscaled to resolution of 64 x 64 pixels. Figure 11 provides an example raw and preprocessed images from the Lung1 dataset. Finally, during VaDeSC, SCA, and DSM training, augment images by randomly reducing/increasing the brightness, adding Poisson noise, flipping 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 scaling, and the predicted survival outcome. Neither VaDeSC, nor DSM, achieve predic- tive performance comparable to radiomics based features in the absence of augmentation.
Raw
(b) Preprocessed
(c) Brightness
(d) Poisson noise
(e) Flip
(f) Rotate
(g) Blur
(h) Zoom
Stretch
(i) Shift
Figure 11: A full resolution Lungl image averaged across 11 CT slices (a) before and (b) after
cropping, normalisation, and histogram equalisation. Panels (c)-(j) depict augmentations applied to images during neural network training.
For the radiomics-based Cox PH and Weibull AFT models, features extracted from the pre- processed 2D images using the Radi omi csFeatureTxt ractor provided by the PyRadiomics library (van Griethuysen al., 2017). Tumour segmentation maps were given as input during feature extraction, to detect the ROI.
23

---

Published as a conference paper at ICLR 2022
G.2 MODELS, ARCHITECTURES & HYPERPARAMETERS
The following implementations of the models were used:
we implemented our model in TensorFlow 2.4 (Abadi et al., 2015). The code available at https: / /github.. com/ i 66994667 4 /vadesc.
VaDeSC:
k-means: we used the implementation available in the scikit-learn library (Pedregosa et al., 2011).
Cox PH & Weibull AFT: used the implementation available in the lifelines library (Davidson-Pinl Pialdsn--Pilonn- et al., 2021).
Radiomics: radiomic features were extracted using the PyRadiomics library (van Griethuy- sen et al., 2017).
SSC: we re- re-llmpeleeeed the model in Python based on the original R (R Core Team, 2020) code provided by Bair Tibshirani (2004).
SCA: adapted the original code by Chapfuwa al. (2020) available https:/ //
github.com. peteomppaiii papppaiaaa /aa.ddamyy/suuuvvaaoooluttttses Sattrvalllllserrrreeeee
DSM: we adapted the original code by Nagpal et al. (2021a) available at https:
//githubs wtw.hh....ooaaaeeeooomooooeeeeee wwt. aut onl ab www///eeeuuiiivieeeeeeeeeeeeeeee va iachi nes.
Profile 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 di- mensions were kept fixed across all neural-network- 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 = Monte Carlo samples for the SGVB estimator. In most datasets, do not pretrain autoencoder weights. Generally, we observe little variability in performance when adjusting the number of latent dimensions, shape of the Weibull distribution 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, truncate the Dirichlet process at the true number of clusters, when known. Similarly, for the DSM, the mini-batel size, learning rate, and the number of training epochs were tuned. To make the comparison with VaDeSC fair, we fixed the discount parameter for the censoring loss at 1.0, since VaDeSC does not discount censored observations 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 the beginning of the training. To prevent this, the encoder and decoder networks pretrained using an autoencoder loss. In the most of our experiments, we did not observe such behaviour. Hence, 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 fitted Gaussian mixture model to initialise the parameters π, Σ.
usually
24

---

Published as a conference paper at ICLR 2022
Table 8: Encoder and decoder architectures used for synthetic, SurvMNIST, and HGG data. tfk1 stands for tensorf low keras layers. encoded_s dii.d.iiia ze corresponds to the number of la- tent dimensions, J; and inp_shape is the number of input dimesions, D. For survMNIST, the activation of the last decoder layer set to 'sigmoid'.
Encoder
ffkl Dense (500, activatiio='relu')
t ffk1 Dense (500, activation='relu") 2
3
tfk1. Dense (2000, eetivation='ellu)
4
.fkki.Dnsse lencoded.size, activation=None);
sigma t.fk1.Dess Dense (encoded.size, activation=None)
Decoder
1
tfkl Dense (2000, activation='relu")
2 tfkl Dense (500, cetivation='rell))
3
tfk1. Dense (500, setivation='relll relu')
4 tfkl Dense (inp shape)
Table 9: Encoder and decoder architectures used for SUPPORT, FLChain, and Hemodialysis data.
Encoder
Dense (50, activation='relu")
2
Dense (100, activation='relu" (encoded.size, activation=None); sigma fkk Dense (encoded.size, activation=None)
3
mu
Decoder
fkll Dense (100, activation 'rel')
2
ffk1 Dense (50, activation='rell
3
fkkl. Dense (inp shape)
25

---

Published as a conference paper at ICLR 2022
Table 10:
Encoder and decoder architectures used for NSCLC data, based on VGG
(de-Jconvolutional blocks (Simonyan & Zisserman, 2015).
Encoder
VGG convolutional block
1
t.k1. Conv2D (32, 3, activation='rell)
2
tfkl. Conv2D (32, 3, activation='relu"
3
tfkl MaxPooling2D (2, 2))
// VGG convolutional block
4
tfk1 Co.11.Conv2 (64, 3, activation='relu))
5
tfkl ConvDD (64, 3, activation='zeln
6
tfkl Maw.ooliiogoDD ( (2, 2))
7
Lfkl.Flatten()
8
t Dense (encoded_size, activation-None); sigma tfkl.Dense (encoded.size, activation-None,
Decoder
tfkl Dense (10816, activation=Nome
2
Reshape (target.shape- (13, 13, 64))
// VGG deconvolutional block
3
t L.f....tsSammmlnn222((,, 2))
4
tfk1. www.ooonntrrrasssseeeeee (64, activation='relu))
5
tfkl. w - (64, 3, activation='rell. relu')
// VGG deconvolutional block
6
tfkl UpSeuuppmmllllg2((2...... 2))
7
tfkl. Conv2brransoose (32, 3, activation='zelu")
8
tfkl Con (32, 3, oetivationn'relul
9
tfkl. Cow.oDoransoosseeee.eeooeot (1, 3, activation=None)
Table 11: Hyperparameters used for training the VaDeSC across all datasets. Herein, J denotes the number latent dimensions; K the number clusters; and & is the shape parameter of the Weibull distribution. MSE stands for mean squared error; BCE for binary cross entropy.
Mini-batch Size
Learning Rate
Pretrain # #pochs
Dataset
J
K
# epochs
Rec. Loss
Synthetic
16
256
lc-3
1,000
MSE
SurvMNIST
16
5
1.0
256
le-3
300
BCE
SUPPORT
16
4
2.0
256
le-3
300
MSE
FLChain
4
2
0.5
256
1e-3
300
MSE
0
HGG
16
3
2.0
256
1c-3
300
MSE
Hemodialysis
16
3.0
256
MSE
NSCLC
32
4
1.0
64
le-3
1,500
MSE
26

---

Published as a conference paper at ICLR 2022
H FURTHER RESULTS
H.1 TRAINING SET CLUSTERING PERFORMANCE
Table 12: Training clustering performance on synthetic and SurvMNIST data. Averages and standard deviations are reported across 5 and 10 independent simulations, respectively. Best results are shown in bold, second best in italic. The results are similar the test-set performance reported in Table 3.
Dataset
Method
ACC
NMI
ARI
CI
k-means Cox PH
0.44
0.06
0.03
0.78
SSC
0.44+0..2
0.08
0.06+00.02
SCA
0.45:10.09
0.05
0.05
0.831.0.02
Synthetic
DSM
0.39-10.04
0.02+00.11
0.02+10.02
0.76:10.01
VAE Weibull
0.46:10.06
0.09
0.09+00.04
0.72
VaDE
0.73+:021
0.53
0.55+0.21
VaDeSC (w/o f)
0.8810.03
0.60:10.07
0.6710.07
0.801
VaDeSC (ours)
0.9010.02
0.66
0.7410.05
k-means
0.48
0.30
0.21
Cox PH
0.77
SSC
0.48
0.30
0.21
SCA
0.56+0.09
0.47+0007
0.34+0.11
0.83
M
DSM
0.52+0.10
0.38
0.28+
0.82
VAE Weibull
0.48+0006
0.31
0.24+0006
0.80
VaDE
0.47-10.07
0.38
0.24100.07
VaDeSC (w/o f)
0.57100.0
0.52100.10
0.3710.10
0.83
VaDeSC (ours)
0.58.10.10
0.54
0.39
H.2
QUALITATIVE RESULTS: SYNTHETIC DATA
Oaxter
Olster Clusker
Gaxter2
Oaster
Dlubber
(a) Ground truth
(b) VAE+ Weibull
(c) SCA
DSM
(e) VaDeSC (ours)
Figure 12: Cluster-specific Kaplan-Meier (KM) curves (top) and (-SNE visualisation of latent rep- resentations (bottom), coloured according to survival times (yellow and blue correspond to higher and lower survival times, respectively), learnt by different models (b-e) from one replicate of the synthetic dataset. Panel (a) shows KM curves of the ground truth clusters. Plots were generated using 10,000 data points randomly sampled from the training set, similar results were observed on the test set.
H.3 COMPARISON WITH PROFILE REGRESSION
Herein, we provide additional experimental results to compare VaDeSC with the profile 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 profile regression, we reduce the dimensionality the dataset by retaining only the first 100 principal components, which preserve, approximately, 90% of the variance. We fit profile
27

---

Published as a conference paper at ICLR 2022
survival regression model with the cluster-specifie 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 the test set of 1500 data points. VaDeSC, both given survival times at prediction and without them, outperforms profile regression by a margin. Never- theless, PR offers a noticeable improvement over k-means clustering and even more sophisticated SCA and DSM (ef. Table 3). While in some instances, PR manages to identify some cluster struc- ture. in others, it assigns all data points single cluster, hence, a large standard deviation across simulations. It is also interesting to see that reduction in the dataset size has led to significant decrease in the performance of VaDeSC. We attribute this to a highly nonlinear structure of this dataset.
Table 13: Test set clustering performance on subsample of the synthetic dataset. Averages and standard deviations are reported across independent simulations. VaDeSC outperforms the baseline by large margin.
Method
ACC
NMI
ARI
k-means
0.44
0.07=0.04
0.05=00022
PR
0.47+t0.13
0.15=0.17
0.15=0.17
VaDeSC (w/o f)
0.73+:0.88
0.34
0.38-10.13
VaDeSC (ours)
0.77+L0.09
0.41+
0.45+00.16
This experiment has clear shortcomings: due to computational limitations, profile regression model on a transformed version of the dataset with fewer dimensions which might not reflect 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 efficient approach.
only able to fit
H.4 VARYING THE FRACTION OF CENSORED DATA POINTS
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. where vary the percentage of censored data while keeping all other simulation parameters fixed. 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.
335
DSM
2.30
WoluulAkr
0.8
80
20
60 85
(a)
(b)
(c)
(d)
Figure 13: Performance of the VaDeSC the synthetic data (see Appendix E. varying percent- ages of censored data points (10-90%): (a) clustering accuracy, (b) concordance index, (c) relative absolute error (RAE) evaluated on non-censored and (d) censored data. Averages and standard de- viations are reported across 5 independent simulations, evaluation was performed on test data. For reference, we report performance of the unsupervised clustering model VaDE and of DSM and Weibull AFT models for survival analysis.
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, average, VaDeSC is still more accurate and stable than the completely
28

---

Published as a conference paper at ICLR 2022
unsupervised VaDE. For time-to-event predictions, all models behave expected: with increase in the percentage of censored observations, test-set CI decreases, RAEac increases, while RAE, 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 real-world dataset with moderate percentage of censored observations, e.g. SUPPORT.
H.5
VARYING THE NUMBER OF COMPONENTS VADESC
0.850
0.80
0.800
Set
0.70
0.750
0.65
0.725
0.700
0.60
6
Number of Clusters in VaDeSC
Number Clusters
(h)
Figure 14: Average test set (a) CI and (b) RAE synthetic data with K clusters achieved by VaDeSC models with different numbers mixture components. Error bars correspond to standard deviations across 5 independent simulations. The CI plot (left) exhibits pronounced 'elbow" the true number of clusters.
0.9
0.85
0.8
Test Set CI
0.7
0.75
0.6
0.70
0.5
12
10 Number of Clusters VaDeSC
8
12
Number of Clusters VaDeSC
(b)
Figure 15: Average test set (a) and (b) RAE on SurvMNIST with K clusters achieved by VaDeSC models with different numbers of mixture components. Error bars correspond to standard deviations across 10 independent simulations. The CI plot (left) appears to saturate at the true number of clusters, although the 'elbow" is not quite as sharp for the synthetic data (cf. Figure 14).
29

---

Published as a conference paper at ICLR 2022
6.6 RECONSTRUCTIONS & GENERATED SURVMNIST SAMPLES
7
9
2
8
7
9
2
8
Figure 16: Reconstructions of survMNIST digits. Original digits are shown in the top row, their reconstructions by VaDeSC in the bottom
4
4
9
9
7
7
7
7
0
0
C
2
8 2 9 0
6
3
/
4
9
4
9
7
7
7
0
0
0
2
2 6 2
8
3
3 5
3
9 9
4
9 4
4
7
7777
7
0
8
6
5
3 5
7
4
9
7 7
7
0
6
8
(a) Cluster 1.
(b) Cluster
(c) Cluster 3.
(d) Cluster
(c) Cluster 5.
Figure 17: SurvMNIST samples generated by VaDeSC. In this simulation, the true clustering is given
by ((7), {0,9), [2, 4, 6), (8}, (1,3, (.),5)) Cluster appears to align with {0, 9); cluster - with
(7); cluster 4 with {2, 4, 63: cluster 5 with {1,3.5}.
H.7 TIME-TO-EVENT PREDICTION: FLCHAIN
Table 14: Time-to-event test set performance on FLChain. Averages and standard deviations are reported across 5 train-test splits, All of the methods perform comparably w.r.t. CI. SCA achieves better RAE and calibration.
Method
CI
RAE,
RAE,
CAL
Cox PH Weibull AFT
0.80=0.01
0.00
0.72-10.11
0.02-10.00
2.18:10.07
SCA
0.78=10022
0.69
0.05+00.55
1.33:10.24
DSM
0.79-1000/
0.76:10.05
0.02+10.01
2.35
Weibull
0.80=0.01
0.76110001
0.02:10.00
2.55
VaDeSC (ours)
0.00=0001
0.76+10.01
0.02-10.00
2.52+0.08
H.8 QUALITATIVE RESULTS: HEMODIALYSIS DATA
In addition to the results reported in Section 5.2, we have investigated qualitatively the clustering obtained by VaDeSC when applied the Hemodialysis dataset. The true structure of this data unknown, however several studies have suggested stratification of patients according to age and dialysis doses (Gotta et al., 2021). We compute the optimal number of clusters using geir-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 show a higher survival probability over time compared to those in cluster
30

---

Published as a conference paper at ICLR 2022
Cluster Cluster
Survual Sunivall Protabilltyy
0.8
L.SAE Dimersion
20
0-
0.4
-20
-40
0.2
955
Cluster
-60
0.0
-40
20
20
0.0
0.7
0.4
0.6
0.8
1.0
SNE Dimenzion
(a) t-SNE plot
Kaplan Meier plot
Figure 18: Visualisation of (a) the t-SNE decomposition of the Hemodialysis data in the embedded space and (b) the Kaplan-Mr Meier curves for the clusters discovered by VaDeSC. The two clusters have substantially different survival distributions.
We now characterise each subgroup of patients by iden- tifying the most influential covariates in determining the cluster assignment by VaDeSC. In particular, we ap- ply state-of-the-art method based Shapley values (Shapley, 1953), namely the TreeExplainer by Lundberg et al. (2020), to explain XGBoost classifier (Chen & Guestrin, 2016) trained to predict the VaDeSC cluster labels from patient characteristics (see Appendix 1). To mimic the VaDeSC clustering model, we included survival times additional input to the classifier. In Figure 20, we present the most important covariates together with their cluster-wise distributions. It is important note that the survival time was also ranked among the most impor- tant, but we did not include it into Figure 20, Serum al- bumin level, previously identified an important survival indicator patients receiving hemodialysis (Foley et al., 1996; Amaral et al., 2008), emerges as most important predictor. Additionally, dialysis dose in terms of spKt/V and fluid balance in terms of ultrafiltration (UF) rate prove to be crucial in classifying the data into high-risk (cluster 1, )) and low-risk (cluster 2, groups. By analysing the cluster-wise distributions, high-risk cluster is characterised by lower values albumin, dialysis dose, and ultrafiltration,, in agreement with previous studies (Gotta et al., 2020; 2021; Movilli al., 2007).
Cluste
Cluder2
Figure 19:
Cluster-wise 3 parameters of the Weibull distributions discovered on the Hemodialysis dataset.
notice that the
Cluster-Specific Survival Models Ultimately, our model should discover subgroups of patients characterised not only by their risk but, more importantly, by the relationships between the covari- ates and survival outcomes. In other words, the cluster-specific parameters of the Weibull distribu- tion should ideally vary across clusters, highlighting different associations between the covariates. We demonstrate this in Figure 19 by plotting the cluster-specific survival parameters, Go (see Equa- tion I), 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 depend- ing on the considered cluster.
31

---

Published as a conference paper at ICLR 2022
Cucter
maunnalbumin
gHiL
of
solgght
spity, choek,
Likgro.87)
mean (dh
100
220
300
400
500
392
700
800
PatientGerder_
3.
02
24
0,6
6R
1.4
maan row
10
mean LPCR
0.0
1.2
144
girgMay
meannneutroohis
Jymphotytes
Covariates Distribution across Clusters
Mean([Shapley Valuel)
The most important predictors for VaDeSC cluster assignments, according to the average Shapley values (see Appendix I). The survival time was not included into this plot, although it was ranked among the most important features. The two clusters are characterised by disparate distributions of the considered clinical variables, where the high-risk cluster shows lower values of albumin, dialysis dose, and ultrafiltration in agreement with previous studies.
Figure 20:
H.9 QUALITATIVE RESULTS: NSCLC
Mean
Assigned Samples
Mean
Assigned Samples
(a) DSM
(b) VaDeSC (ours)
Figure 21: A random selection CT images assigned by (a) DSM and (b) VaDeSC to each cluster (denoted by colour) with the corresponding centroid images, computed across all samples. The colours correspond to the same clusters as Figure 5, Clusters discovered by VaDeSC are strongly correlated with the tumour location; DSM does not appear discover such association.
32

---

Published as a conference paper at ICLR 2022
Mean
Generated Samples
An extended version of Figure 6. CT images generated by (i) sampling latent represen- tations from the Gaussian mixture learnt by VC and (ii) decoding the representations using the decoder network. The colours correspond to the same clusters as Figure 5.
Figure 22:
H. CLUSTERING STABILITY: 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 replicates. However, unsu- pervised learning algorithms based on deep neural networks have been shown to be highly unstable real-world and complex data (Jiang 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 com- plex due 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 defined a priori, some clusters collapse, Le. 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 's runs resulted in collapsed clusters, compared to 45% for DSM. Thus, VaDeSC is relatively stable w.r.t. the number discovered clusters.
We also test the variability in the survival information encoded in each cluster, also shown by the Kaplan-Meiee curves in Figure 5. For each experiment, we compute the median survival time of the samples assigned by the algorithm deviation across the experiments of the cluster-specific median survival time. To have a meaningful average, we sort clusters by the median survival time, highest lowest. As discussed in Section 5.3, DSM tends to discover clusters with more disparate survival curves than VaDeSC, which is reflected 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.
each cluster. In Table 15 report the average and standard
The average and the standard deviation of the cluster-specific median survival time com- puted across 20 independent experiments. For each experiment we re-train the model on a different train-test split. The survival time has been scaled between and 1. We also report the percentage of experiments that resulted in one or more collapsed clusters.
Table 15:
Collapsed Clusters,
Average Median Survival Time SD
Method
Cluster
Cluster
Cluster
Cluster
DSM 45
0.296 + 0.087
0.240 + 0.083
0.138 + 0.055
0.093 + 0.032
VaDeSC
0.220 0.044
0.157 0.016
0.139 0.016
0.086 + 0.038
Finally, test whether the association between tumour locations and clusters is also consistent across experiments. We first compute the means of the mixture of Gaussians learnt by VaDeSC for each experiment, iee. train-test split. As these vectors live in the latent space of VaDeSC,
33

---

Published as a conference paper at ICLR 2022
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, upper row corresponds to patients with higher overall survival probability. It is evident that in most experiments, the upper row corresponds 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.
sort the clusters according to the median survival times such that the
Figure 23: CT images generated by training the VC model 20 times using different train-test splits. For each trained model, we (i) select the learnt means of the Gaussian mixture and (ii) decode them using the decoder network. Each column corresponds to different seed of the train-test split used to train the model. Each row corresponds to a different cluster. For each seed, we order the rows according the median survival within the cluster, hence the upper rows associated with higher survival time. We observe a clear association with tumour location, consistent across most experiments.
I EXPLAINING VADESC CLUSTER ASSIGNMENTS
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 desirable.
Background: SHAP algorithm for interpreting model predictions SHapley Additive exPlana- tion (SHAP), by Lundberg & Lee (2017), is an algorithm that aims at explaining the prediction of an instance æ by computing the contribution of each feature to the prediction. This contribution estimated in terms of Shapley regression values, 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(=') = do + Σド= Ø5ΣY where g(-) is the explanation model, z' 1) effect of the j-th feature on the output fxx of the original model.
F is the binary feature vector of dimension F, and Φj R. is the
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
S - |S| - 1)! [fsut] (æsug)) - fs(xs))
Σ SCAU
(16)
[F
where J (1,..., F} is the set of all predictor variables, (sS()) and æs are input vectors from æ, composed of the features in the subset C F with and without the j-th feature, respectively, and f()) is the original model fitted on the set of predictor variables defined by the subseript.
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 sistent with the local explanations, since the Shapley values for instance are the atomic unit of the global interpretations. Indeed, the global importance of the j-th feature is computed
N
1
1 N
Σ
(17)
T=1
34

---

Published as a conference paper at ICLR 2022
where N is the total number of observations used to explain the model globally.
Despite the theoretical advantages of SHAP values, their practical use hindered by the complexity of estimating E (x) [ efficiently. TreeExplainer (Lundberg al., 2020) is a variant of SHAP algorithm that computes the classical Shapley values from the game theory specifically for trees and tree ensembles reducing the complexity from exponential polynomial time.
Preprocessing With the goal of explaining VaDeSC cluster assignments in terms of input covari- ates. we preselect features manually, 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. E) in VaDeSC depends on survival time (see Sec- tion 3), we included the survival times as an additional input to the classifier. Note that for this post hoc analysis we used training data only, as were interested in explaining what our model learnt from the raw data to perform clustering.
Explanations To characterise patient subgroups identified by VaDeSC, needs identify the features in the input space that maximise cluster separability in the embedding space. Note that classifier trained to recognise cluster labels from the raw covariates is inherently seeking those dimensions of maximal separability in the input. Therefore, we first fitted a classifier model to extract the explanations in terms of the imput features most relevant for the prediction of cluster labels.
In particular, we trained an eXtreme Gradient Boosting (XGBoostt binary classifier (Chen & Guestrin, 2016), to predict patient labels e {1,2} assigned by VaDeSC. We choose XGBoost since it has high accuracy in the presence of nonlinear relationships, naturally captures feature in- teractions, suffers from low bias, and supports fast exact computation of Shapley values by TreeEx- plainer (Lundberg et al., 2020). XGBoost was trained for 5,000 boosting rounds with a maximum tree depth of 6, learning rate equal to 0.01, subsampling rate of 0.5, and early stopping after 200 rounds without an improvement on the validation set (set to 10% the training data). After training the classifier, determinant features for each cluster label, as in Equation 17. To profile the clusters precisely and avoid the noise induced by weakly assigned observations, we computed SHAP values on sample of cluster prototypes, i.e. observations with m p(c = v/ 22 t) 2 0.75. We retained at most 500 prototypes from each cluster. Due to the cluster imbalance, that resulted in 500 observations from cluster and 212 from cluster 2.
used the TreeFxplainer
compute SHAP values and identify the most
35


Reviews:
[{"review_id": "u2iUSo3JUwS", "reviewer": "Reviewer_VEke", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "The paper is well written and easy to follow. \n\nNonetheless, Figure 4 is hard to read and interpret, it would be interesting to do this analysis on a dataset with a smaller number of clusters. Further description of the clustering evaluation would also be valuable: in figure 4, is the latent representation obtained on the test split only? It seems the authors use a Monte Carlo cross-validation, it would be beneficial to add this point in the experiment section.\nFinally, in Figure 5, were the number of clusters and their shapes consistent over multiple iterations? An experiments measuring this stability would be a good addition.\n\nFinally, it would be necessary to further articulate the technical contributions of the paper, as the proposed model is similar to [1] in which they also propose a variational approach to obtain a better embedding that influences the baseline distributions. The main difference resides in the baseline distributions which are parametric in [1]. The proposed appraoch merges [1] and [2] and extends it to a fully generative model. This should be clearly stated and a comparison with the distribution obtained in [1] would be a valuable addition.\n\n[1] Nagpal, C., Yadlowsky, S., Rostamzadeh, N. and Heller, K., 2021. Deep Cox mixtures for survival regression. \n[2] Nagpal, C., Li, X.R. and Dubrawski, A., 2021. Deep survival machines: Fully parametric survival regression and representation learning for censored data with competing risks. ", "summary_of_the_paper": "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.\nSynthetic and real world datasets demonstrate the competitiveness of the method on both discriminative and clustering performances.", "summary_of_the_review": "An interesting work, well-described and with in-depth experiments. However, the technical contribution has to be described in greater details.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "k3C0bjaPnhc", "reviewer": "Reviewer_6hmP", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "his paper proposes a new latent variable model for clustering survival data. The new method aims to identify clusters that have characteristic generative survival mechanisms or associations between the auxiliary covariates and the survival times, contrary to methods that identify clusters solely based on the survival times or solely based on covariates. The model is formulated as a latent variable model that combines a general latent variable model with neural network parameterized decoder likelihood for covariates as well as Weibull distribution model for the survival function. Model inference is developed suing the auto-encoding variational Bayes. Amortised variational inference combines a neural network parameterized encoder for the covariates as well as a separate encoder block for the cluster assignments that is defined directly by the generative model, which also naturally handles data points where the survival time is missing/censored. Overall, the model specification seems well-motivated, and the description of the model/results admirably clear.\n\nThe proposed method is compared against a large collection of previous methods on both synthetic and real data sets. Performance metrics include clustering assigngment accuracy and survival time prediction accuracy. Overall, experiments with simulated and real data demonstrate that the proposed method can provide more accurate/meaningful cluster assignments than previous methods, while at the same time providing survival predictions that are on par with the existing methods. \n\nThis manuscript proposes an interesting latent mixture survival model that makes a useful contribution to the field. Technical details appear correct.  I do not have any major concerns with the manuscript. \n\nConcerning the experiments with simulated data, the differences in cluster-specific Kaplan-Meyer curves appear perhaps unreasonably large compared to what is typically expected in real applications.  Can you provide additional experiments with modified versions of the data set where differences between the KM curves would be smaller?  Results for such simulations would also reveal how the performance of different methods behave when you make your dataset more difficult. \n", "summary_of_the_paper": "This paper proposes a new latent variable model for clustering survival data. Experiments with simulated and real data demonstrate that the proposed method can identify accurate/meaningful cluster assignments while simultaneously providing survival predictions that are on par with the existing methods.", "summary_of_the_review": "This manuscript proposes an interesting latent mixture survival model that makes a useful contribution to the field.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works."}, "extra_text": {}}, {"review_id": "cWAPWfTAHv", "reviewer": "Reviewer_k5xJ", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "The novelty of the proposed work is modest. As the authors point out, similar clustering-based approaches for survival analysis have already been proposed. Likewise, encoding data types such as images into a latent representation using neural networks is common. To the best of my knowledge, though, combining these approaches is novel.\n\nIn terms of technical soundness, the proposed approach is reasonable. I didn’t verify it in detail, but the ELBO approximation appears to follow typical conventions. The experiments are well designed with good baselines and evaluation methodology. The results seem to be presented fairly, and the authors discuss both the strengths and limitation of their approach.\n\nThe paper is generally well written and easy to follow.\n\nIn terms of reproducibility and contribution to the community, the authors provide technical details in the appendix. They also make their code available. While I did not run the code, it does adhere to a clear style, so I believe that would be a helpful contribution to the community.", "summary_of_the_paper": "In this work, the authors propose a model for clustering survival data which accounts for both survival time as well as covariates, such as patient demographics. The model is similar to previous approaches, though it incorporates a neural network encoder to handle unstructured data types as input. A modest set of experiments suggest that the proposed approach identifies more meaningful clusters than similar approaches. Also, while not the focus of the proposed approach, empirical evaluation shows that the proposed method performs similarly to existing approaches on time-to-event predictions.", "summary_of_the_review": "Overall, the proposed work is an incremental step forward for survival analysis. While the impact may not be huge, I believe the proposed approach is a reasonable, novel view on the problem. The experimental results are not overwhelming, but they show that the proposed approach does succeed at clustering patients. This is key for precision medicine, so other researchers could also build on this formulation and approach.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "XBBB79fXlpU", "reviewer": "Reviewer_Z7tL", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Strengths:\n- The literature review on related clinical survival analysis research is thoroughly summarized with clear statement on the relevance and differences with comparison to the proposed method.\n- Experiment design is comprehensive and rigorous with various important evaluation metrics of clustering analysis and survival predictions, on different types of datasets (synthetic, semi-synthetic and real-world benchmark data)\n- Detailed experiment assumptions, parameters, implementations, as well as results (enough Appendix references and supplementary material to reproduce the proposed work as well as baseline models)\n- Meaningful discussion of method limitations and future work\n\nWeaknesses:\n- My major complaint originated from the potential fair practical utility of the proposed method in clinical applications indicated by Figure 5, the comparison of Kaplan-Meier curves between DSM and VaDeSC among identified clusters. One major purpose of the survival analysis is to guide bedside clinicians and nurses for timely intervention for better survival outcome and the clustering analysis could offer meaningful grouping of risk-stratified patients with intervention priorities. Thus a clear separation of different shapes of KL curves might be more helpful in real-time bedside practice and guidance, than the visual separation of images and allocation of tumors. In this case, DSM seems to perform better for this purpose in indicate \"When to intervene\" on the NSCLC dataset. However, VaDeSC, by its nature to differentiate variants, might be better utilized in diagnosis process for image data. I also wonder whether it's the same case for tabular data.\n- Color visualization of Figure 4: KM curves of assigned clusters are very hard to see and compare. Yellow and blue might be better replaced with e.g. red and green for better contrast. \n- Is there any justification for why FLChain results are in Appendix instead of main text? section 5.2 on page 7. Methods and datasets are not uniform across the paper e.g. table 2, 3 and 4 are of different subsets. If space is limited, is it possible to provide complete table of results comparison of all metrics, all methods and all datasets in the Appendix. \n- The content of the paper is very dense to read due to the great efforts the author(s) have put. Sometimes the results lack immediate clear interpretation in the main text and referring to Appendix seems to be too frequent. It will be better if the main content could be more targeted and focused.   \n\nAdditional question:\n- Would you pls comment how the performance of different models vary by the %censored of the datasets? \n", "summary_of_the_paper": "The paper proposed a novel variational deep survival clustering (VadeSC) method to discover underlying distribution of both the explanatory variables and censored survival times. Its main contributions are jointly modeling variables and censored survival outcomes and comprehensive comparison of proposed methods on state-of-art existing models on both synthetic and real-world datasets.", "summary_of_the_review": "The paper which proposed a novel survival clustering algorithm is well-written with thorough literature review, clear analysis assumptions, comprehensive experiment design and results comparison with state-of-art baselines models and closely related benchmark datasets, which is a good example for dedicated research work. ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works."}, "extra_text": {}}]

Author rebuttal:

Meta-review: Four knowledgeable referees recommend Accept. I also think the paper provides a unique contribution to the field of deep survival models and I, therefore, recommend Accept
Decision: Accept (Poster)