Title: TOPIC MODELS WITH SURVIVAL SUPERVISION: ARCHETYPAL ANALYSIS AND NEURAL APPROACHES

Abstract: We introduce two approaches to topic modeling supervised by survival analysis. Both approaches predict time-to-event outcomes while simultaneously learning topics over features that help prediction. The high-level idea is to represent each data point as a distribution over topics using some underlying topic model. Then each data point's distribution over topics is fed as input to a survival model. The topic and survival models are jointly learned. The two approaches we propose differ in the generality of topic models they can learn. The first approach finds topics via archetypal analysis, a nonnegative matrix factorization method that optimizes over a wide class of topic models encompassing latent Dirichlet allocation (LDA), correlated topic models, and topic models based on the "anchor word" assumption; the resulting survival-supervised variant solves an alternating minimization problem. Our second approach builds on recent work that approximates LDA in a neural net framework. We add a survival loss layer to this neural net to form an approximation to survival-supervised LDA. Both of our approaches can be combined with a variety of survival models. We demonstrate our approach on two survival datasets, showing that survival-supervised topic models can achieve competitive time-to-event prediction accuracy while outputting clinically interpretable topics.

1. INTRODUCTION:

Predicting time-to-event outcomes arises in a variety of applications. For example, in healthcare, we may be interested in predicting how much time a patient has to live. In criminology, we may be interested in predicting when a convicted criminal might reoffend. In e-commerce and on streaming platforms, companies with subscription services like Amazon and Netflix may be interested in predicting when users might cancel their subscriptions. In many such applications, we can now collect an enormous number of measurements per person/subject. However, how all of these measurements relate is typically unknown. In this paper, we aim to address the twin objectives of learning how measurements relate in the form of a topic model, and learning how topics can assist in predicting time-to-event outcomes via a survival analysis model.

For ease of exposition, we phrase the problem we consider in the classical survival analysis context of predicting time until death. We assume that we have access to a training dataset of n subjects. For each subject, we know how many times each of d "words" appears, where the dictionary of words is pre-specified. As an example, in a clinical context, one word might correspond to "low blood pressure reading"; for a given subject, we can count how many such readings the subject has had recorded in the past. We denote X i,u to be the number of times word u ∈ {1, . . . , d} appears for subject i ∈ {1, . . . , n}. Viewing X as an n-by-d matrix, the i-th row of X can be thought of as the feature vector for the i-th subject. As for the training label for the i-th subject, we have two recordings: event indicator δ i ∈ {0, 1} specifies whether the i-th subject died, and observed time Y i ∈ R + is the i-th subject's "survival time" (time until death) if δ i = 1 or the "censoring time" if δ i = 0. The censoring time gives a lower bound on the survival time for the i-th subject. For example, when we stop collecting data, some subjects will still be alive, so we know they live at least as long as when we stopped collecting training data.

Our goal is to discover topics for the d words that help predict survival times of unseen test subjects. Note that an unsupervised topic model like latent Dirichlet allocation (LDA) (Blei et al., 2013) would not use any of the training labels (the event indicators δ i 's and observed times Y i 's), learning topics using only the word counts matrix X. Meanwhile, in survival analysis, a standard approach would involve learning a survival model using all the patients' feature vectors and labels but the model would not learn thematic structure in the different features, e.g., topics. Jointly learning both a topic model and a survival model was first done by Dawson & Kendziorski (2012) , who combined LDA with a Cox proportional hazards model (Cox, 1972) . Using LDA with r topics, Dawson and Kendziorski represent the i-th subject as a probability vector W i ∈ [0, 1] r specifying the subject's membership in each of the r topics; then W i 's are treated as the input covariates to the Cox model. Dawson and Kendziorski called this joint model SURVLDA and derived a variational EM algorithm to estimate its parameters.

In this paper, we build on SURVLDA by proposing two new survival-supervised topic modeling approaches, both of which allow for either the topic or the survival model to be replaced. Our contributions are as follows:

• (Section 3) We show how to take a discriminative approach to jointly learning topic and survival models, where topics are estimated via archetypal analysis (Cutler & Breiman, 1994; Javadi & Montanari, 2019) . Archetypal analysis represents each subject as a convex combination of "archetypes", which are optimized to be diverse yet still be close to the convex hull of the subjects' feature vectors. Applied to topic modeling, the archetypes are the topics, with each archetype specifying a particular topic's word distribution. Archetypal analysis does not assume a parametric model and can learn a wide class of topic models. We describe how to combine archetypal analysis with any survival analysis model for which we can take a specific partial derivative.

• (Section 4) We approximate Dawson and Kendziorski's SURVLDA model in a neural net framework, which allows for different choices of topic and survival models to be combined. This approach requires that the topic and survival models already have neural net approximations or formulations. For example, LDA and some variants of it can already be approximated using variational autoencoders (Srivastava & Sutton, 2017; Card et al., 2018) . In particular, Card et al. (2018) show how to approximate supervised LDA (McAuliffe & Blei, 2008) in a neural net framework that they call SCHOLAR; they specifically consider classification as the supervised task although they mention that their framework could be used to predict other real-valued outputs. We specifically combine their approach with that of Katzman et al. (2018) to handle survival supervision.

• (Section 5) We apply our two proposed approaches to two survival analysis datasets (predicting how long pancreatitis patients stay in an intensive care unit, and time until death for breast cancer subjects), comparing against a number of classical and recently developed deep survival analysis baselines. Survival-supervised topic models have time-to-event prediction accuracy that is competitive with top-performing existing baselines while producing clinically interpretable topics.

2. BACKGROUND:

We begin with some background on archetypal analysis, topic modeling, and survival analysis.

Along the way, we introduce notation that recurs throughout the paper. As a reminder, we assume that we have access to training data

(X 1 , Y 1 , δ 1 ), (X 2 , Y 2 , δ 2 ), . . . , (X n , Y n , δ n ),

where the i-th training subject has feature vector X i ∈ R d , observed time Y i ∈ R + , and event indicator δ i ∈ {0, 1}. Throughout this paper, we generally take X i,u (for u ∈ {1, 2, . . . , d}) to be the number of times word d appears, for some user-specified dictionary of d words. We let X i,u denote the fraction of times a word appears for a specific subject, meaning that

X i,u = X i,u d v=1 X i,v

.

Note that X is an n-by-d matrix, and we use X i to denote the i-th row of X. We use this indexing notation for other matrices as well.

2.1. ARCHETYPAL ANALYSIS AND TOPIC MODELING:

Archetypal analysis (Cutler & Breiman, 1994; Javadi & Montanari, 2019 ) posits that each training vector X i can be well-approximated by a convex combination of r different unknown "archetypes"

H 1 , H 2 , . . . , H r ∈ R d : X i ≈ r g=1 W i,g H g (2.1)

for some weights W i,1 , . . . , W i,r ∈ [0, 1] that sum to 1, i.e., the vector W i := (W i,1 , . . . , W i,r ) resides in the probability simplex ∆ r := {w ∈ [0, 1] r : r g=1 w g = 1}. By stacking the archetypes H 1 , . . . , H r as rows to form the matrix H, equation (2.1) can be expressed as X ≈ W H. Archetypal analysis aims to estimate W and H given X.

If the archetypes H 1 , . . . , H r are constrained to be in the probability simplex ∆ d , then we get a topic model, and each archetype corresponds to a word distribution. For example, if rows of W are generated i.i.d. from a Dirichlet distribution, and rows of H are generated i.i.d. from another Dirichlet distribution, then we get LDA (Blei et al., 2013) . As a slight modification of this setup, if the rows of W are instead generated from a logistic normal distribution that allows correlation between topics, we get the correlated topic model (Lafferty & Blei, 2006) . For an example that is not generative, if the archetypes are on a probability simplex, and for each archetype g ∈ {1, . . . , r}, there exists a word w that only appears in archetype g (i.e., H g,w > 0 and H h,w = 0 for all h = g), then we have a topic model satisfying the separability or "anchor word" assumption (Donoho & Stodden, 2004; Arora et al., 2012a; b; 2013) . Archetypal analysis optimizes over matrices W and H that include all of the aforementioned topic models above as special cases. In fact, archetypal analysis does not require that archetypes be on a probability simplex or that they be nonnegative; the input matrix X need not consist of word frequencies and could be positive or negative real-valued measurements. Crucially, the error in approximation (2.1) should be small; precise details including identifiability and degeneracy issues can be found in Section 3 of Javadi & Montanari (2019) .

To estimate weights W and archetypes H, Javadi and Montanari proposed the following approach. First, for a point u ∈ R d and a matrix V ∈ R m×d , we define the distance from u to the convex hull of the rows of V as D(u, V ) := min w∈∆ m u -V w 2 . The vector w ∈ ∆ m that achieves the minimum consists of the convex combination weights that best combine rows of V to approximate the point u. Then Javadi and Montanari (approximately) minimize the nonconvex loss

L arch (W , H; λ) := n i=1 X i -H W i 2 2 ♠ +λ r g=1 D 2 (H g , X) ♥ (2.2)

subject to the constraint that W i ∈ ∆ r for i = 1, . . . , n; constant λ ≥ 0 is a user-specified regularization parameter. Minimizing term ♠ (error of approximating input data X i 's as convex combination of archetypes) encourages the archetypes to be far apart and have a convex hull that contains the input data. However, this term does not prevent the archetypes from taking on extreme values; for example, if the archetypes already have a convex hull that contains the X i 's (so that ♠ = 0), we can move the archetypes even farther apart and still have their convex hull contain the X i 's (so we still have ♠ = 0). We prevent this behavior by minimizing term ♥, which encourages each archetype to be close to the convex hull of the input data.

To learn a topic model, we enforce that the archetypes correspond to distributions over words by requiring each row of H to be in probability simplex ∆ d . The resulting optimization problem is ( W , H) ∈ argmin W ∈R n×r , H∈R r×d s.t. Wi∈∆ r for all i, Hg∈∆ d for all g L arch (W , H; λ).

(2.3)

A local minimum can be found by alternating between minimizing W with H fixed, and vice versa.

2.2. SURVIVAL ANALYSIS:

Archetypal analysis and topic models are unsupervised methods. To predict time-to-event outcomes, we turn toward survival analysis models. Suppose we take the i-th subject's feature vector to be W i ∈ R r instead of X i . As this notation suggests, when we combine topic and survival models, W i corresponds to the i-th subject's archetype/topic combination weights; this strategy for combining topic and survival models was also done by Dawson & Kendziorski (2012) , which in turn is based on the supervised LDA formulation (McAuliffe & Blei, 2008) . We treat the training data as (W 1 , Y 1 , δ 1 ), . . . , (W n , Y n , δ n ), disregarding the X i and X i values from earlier.

We aim to reason about the survival time of an unseen test feature vector W 0 ∈ R r . Specifically, let random variable T 0 denote the survival time corresponding to feature vector W 0 (treated as a random variable). Then our goal is to produce an estimate S of the (conditional) survival function S(t|w) := P(T 0 > t | W 0 = w) for t ≥ 0 and w ∈ R r .

Importantly, for a given feature vector w ∈ R d , note that S( 

i ≤ C i ), then Y i = T i and δ i = 1; otherwise Y i = C i and δ i = 0.

This setup is referred to as random censoring. Details can be found in a survival analysis textbook (e.g., Kalbfleisch & Prentice, 2002 ).

As a concrete example of how survival function S can be computed via minimizing a loss function, we next present the classical Cox proportional hazards model (Cox, 1972 ). As we discuss shortly, this is just one example of a survival model that can be combined with our proposed survivalsupervised archetyptal analysis or neural topic modeling approaches.

Example 1 (Cox proportional hazards) Recall that survival function S is 1 minus the CDF of the distribution of survival time T 0 given feature vector W 0 = w. We denote the CDF of this distribution as F (t|w) and assume it has a probability density function f (t|w) = ∂ ∂t F (t|w). Then the Cox model constrains S through the so-called hazard function h of S, given by

h(t|w) := - ∂ ∂t log S(t|w) = -∂ ∂t [1 -F (t|w)] S(t|w) = f (t|w) S(t|w) , (2.4)

which is the instantaneous rate of death at time t divided by the probability of surviving up to time t, all conditioned on the feature vector being w. Specifically, the Cox model assumes that hazard function h factors as h(t|w) = h 0 (t)e β w , where the two parameters are the baseline hazard function h 0 : R + → R + , and the vector of regression coefficients β ∈ R r . Under random censoring (and actually more general censoring models), we can estimate β without knowing h 0 via maximizing a profile likelihood, which is equivalent to minimizing the loss

L Cox (β|W ) := n i=1 δ i -β W i +log n j=1 s.t. Yj ≥Yi exp(β W j ) .

(2.5)

Given an estimate of β, we can deterministically compute a nonparametric estimate for baseline hazard function h 0 ; this estimation procedure is standard and can be found in Section 7.8 of Cox & Oakes (1984) . Once we have estimates h 0 and β for h 0 and β, then for any test feature vector w ∈ R r , we can estimate this test feature vector's corresponding hazard function via

h(t|w) = h 0 (t)e β w .

Using the first equality of equation (2.4), note that S(t|w) = exp -t 0 h(τ |w)dτ . We can plug estimate h for h into this equation to get an estimate of S:

S(t|w) = exp - t 0 h(τ |w)dτ .

(2.6)

In practice, the integral is computed via a summation.

Other survival losses are possible aside from L Cox (β|W ). As a second example, we provide the survival loss function for the Weibull accelerated failure time (AFT) model in Appendix B. The critical requirement of our proposed methods to follow is that the survival loss used is differentiable with respect to W . For example, the elastic-net-regularized Cox proportional hazards model by Park & Hastie (2007) also satisfies this condition.

3. SURVIVAL-SUPERVISED ARCHETYPAL ANALYSIS:

Survival supervision can readily be incorporated into the archetypal analysis optimization problem (2.3) by adding a survival loss L surv (W , θ) to the objective function, where θ here is the collection of all parameters specifying the survival model. For example, we could have L surv (W , θ) = L Cox (β|W ) as defined in equation (2.5) with parameters θ = β.1 Specifically, letting Θ denote the set of possible values that parameter θ can take on, and η > 0 denote a user-specified importance weight of the survival loss, we now instead solve

( W , H) ∈ argmin W ∈R n×r , H∈R r×d , θ∈Θ s.t. Wi∈∆ r for all i, Hg∈∆ d for all g L arch (W , H; λ) + ηL surv (W , θ), (3.1)

where L arch is given in equation (2.2). Javadi & Montanari (2019) solve the unsupervised archetypal analysis optimization problem (2.3) using the Proximal Alternating Linearized Minimization (PALM) algorithm by Bolte et al. (2014) . We augment this algorithm to handle survival supervision, resulting in an algorithm we call SURVIVAL-ARCHETYPES. We first state what SURVIVAL-ARCHETYPES is before explaining our algorithmic modifications to the unsupervised variant.

3.1. THE SURVIVAL-ARCHETYPES ALGORITHM:

In what follows, we let Π U (V ) denote the version of V where each of its rows has been projected onto the set U . Formally, the i-th row of Π U (V ) is given by

[Π U (V )] i = min u∈U u -V i 2 .

For example, if V ∈ R n×r consists of nonnegative entries where each row's sum is strictly greater than 0, then

[Π ∆ r (V )] i = V i / r j=1 V i,j

. Next, we denote the convex hull of the rows of a matrix

V ∈ R m×d by conv(V ) := m i=1 w i V i : w ∈ ∆ m .

Then the SURVIVAL-ARCHETYPES algorithm repeats the following steps until convergence:

1. Update archetypes: with step size parameter γ 1 = 2 W W F , where • F denotes the Frobenius norm, set

H ← H - 1 γ 1 W (W H -X), H ← Π ∆ d H - λ λ + γ 1 ( H -Π conv(X) ( H)) .

2. Update convex combination weights: with step size parameter γ 2 found using backtracking line search (Parikh & Boyd, 2014, Section 4. 3), set

W ← Π ∆ r W - 1 γ 2 (W H -X)H + η ∂L surv (W , θ) ∂W .

3. Update survival model: θ ← argmin θ∈Θ L surv (W , θ). This step amounts to fitting the survival model with rows of W treated as the feature vectors and can just use the model's existing fitting code as a black box.

Initialization. Following Javadi & Montanari (2019) , we use the successive projections algorithm by Araújo et al. (2001) to initialize archetypes H. We can initialize each row of W by setting

W i ← argmin w∈∆ r X i -H w 2 .

Lastly, how the survival model parameters θ is initialized depends on the survival model. For example, for the Cox proportional hazards model, we can initialize β to be the all zeros vector. For the Weibull AFT model, we can initialize µ and β to be all zeros, and σ to be 1.

A key technical requirement of SURVIVAL-ARCHETYPES is that we need to be able to compute the gradient ∂Lsurv(W ,θ)

∂W:

.

As illustrative examples, we show what this gradient is equal to for the Cox and Weibull AFT models in Appendix C.

3.2. RELATING TO THE UNSUPERVISED ARCHETYPAL ANALYSIS PALM ALGORITHM:

The original PALM algorithm for unsupervised archetypal analysis can be recovered by setting η = 0 and removing step 3. Moreover, the step size in step 2 need not be found using backtracking line search. In particular, when η = 0, step 2 takes a proximal gradient step, where the gradient is

(W H -X)H ,

which has Lipschitz modulus 2 HH F ; hence, we can set step size parameter γ 2 = 2 HH F (Bolte et al., 2014, Remark 7(ii) ). When η > 0, the Lipschitz modulus can vary by the survival model used and in general does not have a closed-form expression, so we use a line search. Lastly, if we are not constraining the archetypes to correspond to word distributions, then the projection onto ∆ d in step 1 can be removed.

4. NEURAL SURVIVAL-SUPERVISED TOPIC MODELS:

Our proposed approach to a neural survival-supervised topic modeling builds on the SCHOLAR framework by Card et al. (2018) . Card et al. do not explicitly consider survival analysis in their setup although they mention that predicting different kinds of real-valued outputs can be incorporated by using different label networks. We use their same setup and have the final label network perform survival analysis via the same approach as Katzman et al. (2018) ; note that Katzman et al. specifically consider the Cox proportional hazards model but their neural net approach works with some other survival models as well such as the Weibull AFT model. We first give an overview of SCHOLAR and then explain how to implement the final survival analysis label network.

For ease of exposition, we present the SCHOLAR framework without what Card et al. refer to as "covariates" (auxiliary information known about subjects in addition to the word count matrix X). The SCHOLAR framework specifies a generative model for the data, including how each individual word in each subject is generated. In particular, recall that X i,u denotes the number of times the word u ∈ {1, 2, . . . , d} appears for the i-th subject. Let n i denote the number of words for the i-th subject, i.e., n i = d u=1 X i,u . We now define the random variable ψ i, ∈ {1, 2, . . . , d} to be what the -th word for the i-th subject is (for i = 1, 2, . . . , n and = 1, 2, . . . , n i ). Then the generative process for SCHOLAR with r topics is as follows, stated for the i-th subject:

1. Generate the i-th subject's topic distribution: (a) Sample W i from a logistic normal distribution with mean vector µ ∈ R r and covariance matrix Σ ∈ R r×r . (b) Set the topic weights vector for the i-th subject to be W i = softmax( W i ).

2. Generate the i-th subject's words:

(a) Set word parameter η i = f word (W i ), where f word is a generator network. (b) For word = 1, 2, . . . , n i : Sample ψ i, ∼ Multinomial(softmax(η i )).

3. Generate the i-th subject's output label: Sample Y i from a distribution parameterized by label network f label (W i ).

There are a wide variety of choices for the parameters µ, Σ, f word , and f label . For example, to approximate supervised LDA (McAuliffe & Blei, 2008) where the topic distributions are sampled from a symmetric Dirichlet distribution with parameter α > 0 and the output label is continuous and has unit variance, we set µ to be the all zeros vector, Σ = diag((r -1)/(αr)), f word (w) = w H where H ∈ R r×d has a Dirichlet prior per row, f label (w) = w φ for a parameter vector φ ∈ R r , and set Y i to be generated from a Gaussian with mean f label (w) and variance 1. Card et al. (2018) also explain how to approximate the correlated topic model by Lafferty & Blei (2006) . To estimate the model parameters, Card et al. use a sampling-based variational autoencoder framework (Kingma & Welling, 2014; Rezende et al., 2014) .

Survival supervision. To incorporate survival analysis, we follow the same approach as Katzman et al. (2018) and change step 3 of the generative process above to be deterministic and instead output the variable

Ξ i = f label (W i ) := β W i for parameter vector β ∈ R r .

In particular, we do not actually model how observed times Y i 's are generated; modeling Ξ i 's is sufficient. Then we can minimize the Cox proportional hazards loss:

L Cox (β|W ) = n i=1 δ i -Ξ i + log n j=1 s.t. Yj ≥Yi exp(Ξ i ) ,

where

z i = log Yi-µ-Ξi σ

. Regularization on β can easily be added (e.g., lasso, elastic net). Other losses are also possible. The Weibull AFT loss given in Appendix B uses the same label network as the Cox example above, namely f label (W i ) = β W i . For both the Cox and Weibull AFT examples, the label network could instead be a multilayer perceptron or a more complex neural net rather than a simple inner product. We refer to SCHOLAR with a survival loss as SURVIVAL-SCHOLAR.

5. EXPERIMENTS:

We apply SURVIVAL-ARCHETYPES and SURVIVAL-SCHOLAR to two survival analysis datasets focusing on two diseases: pancreatitis and breast cancer. For pancreatitis, we use the MIMIC III electronic health records dataset (Johnson et al., 2016) , looking only at the pancreatitis patients admitted to the intensive care unit (ICU) and who did not die while in the ICU; this amounted to 371 patients where we extracted 2557 features (preprocessing details are in Appendix D.1). We predict how long each patient will stay in the ICU. For breast cancer, we use the METABRIC dataset (Curtis et al., 2012) , which consists of 1981 patients. We use the same 79 one-hot encoded features as Lee et al. (2018) to predict time until death per subject. Some features are continuous and need to be discretized for use with our topic models (resulting in 100 total features; see Appendix D.2 for details). For both datasets, we randomly divide the dataset into a 75%-25% train-test split.

We benchmark our approaches against a total of 10 baselines: 7 classical methods (lasso-regularized Cox proportional hazards with and without PCA preprocessing, Weibull AFT with and without PCA, k-nearest neighbor survival analysis (Beran, 1981; Lowsky et al., 2013) with and without PCA, and random survival forests (Ishwaran et al., 2008) ), 2 deep learning methods (DEEPSURV (Katzman et al., 2018) and DEEPHIT (Lee et al., 2018) ), and Dawson and Kendziorski's SURVLDA (Dawson & Kendziorski, 2012) . For lasso-regularized Cox, our hyperparameter sweep does include an approximation to the standard unregularized Cox model. For simplicity, for our archetypal analysis and neural approaches, we use the standard Cox model as the survival model. For all methods, if the method does not already have a hyperparameter selection procedure, we use 5-fold cross-validation on the training data to select hyperparameters prior to training on the complete training data using the best parameters found; hyperparameter search grids are in Appendix E. For the pancreatitis dataset, due to the number of subjects being small, we use repeated 5-fold cross-validation with 10 repeats. Repeated k-fold cross validation has been found to be useful in such small dataset regimes (Braga-Neto & Dougherty, 2004) . For both cross-validation and for evaluating test set accuracy, we use the standard survival analysis metric of concordance index (Harrell Jr et al., 1982) , which is the fraction of pairs of validation/test subjects correctly ordered by the prediction algorithm in terms of which subject has a longer survival time (amongst pairs that can be ordered). and SURVIVAL-SCHOLAR all do nearly as well. Also, note that despite SURVIVAL-SCHOLAR being a neural approximation of SURVLDA, the two methods' accuracies are different; this phenomenon has also been reported for SCHOLAR and the various topic models it approximates (Card et al., 2018) . Overall, there is no single best survival estimator. The three survival-supervised topic models also jointly estimate topics, and per topic, tells us whether presence of that topic leads to greater or lower probability of survival. As we discuss next, despite SURVIVAL-SCHOLAR having only a concordance index of 0.59 on the pancreatitis dataset, it still manages to produce clinically interpretable topics predictive of whether pancreatitis patients will stay longer in the ICU.

Test set concordances are reported in:

We now give a brief summary of learned topics. Note that for both datasets, the vast majority of words we used require clinical expertise to interpret. For ease of exposition, we defer examples of actual topics learned to Appendix F, where per topic, we list its top 20 most probable words along with the topic's Cox β coefficient-a higher coefficient corresponds to predicting a shorter ICU length of stay in the pancreatitis dataset and a shorter time until death in the breast cancer dataset.

A topic with β coefficient 0 gets ignored for prediction.

Pancreatitis. SURVIVAL-ARCHETYPES identified one archetype with a nonzero Cox β coefficient (4.5) corresponding to a healthy group with lower-risk interventions (e.g., smaller-bore IV, normal MCV and HCT, top words do not have data elements related to severe illness). All other archetypes have β coefficient 0 and correspond to sicker patient characteristics (e.g., atypical lab tests and toxicology panels). SURVIVAL-SCHOLAR separated clinical events into 3 meaningful topics: one for laboratory tests, one for patient presentation characteristics, and one for procedures, precautions, monitoring, and vitals (this last topic has the smallest β coefficient associated with longer ICU length of stay). SURVLDA also produced interpretable topics such as critical illness, normal health state, and acid base disorders and liver involvement.

Breast cancer. SURVIVAL-SCHOLAR found topics that distinguish elderly, advanced cancers (β coefficient 0.71) from ones with early and younger hormone positive characteristics (β coefficients -0.74 and -0.79). SURVLDA also produced topics with identifiable characteristics; however more topics were found (7 topics) with two overlapping topics indicative of elderly stage 2 breast cancer, and three other overlapping topics (all indicative of hormone positive, cellular, and proliferative features). SURVIVAL-ARCHETYPES has noticeably lower prediction accuracy on this dataset, which is reflected in the topics it learns: the two topics with nonzero β coefficients have opposite β coefficient signs yet have mostly the same top words, suggesting too much topic overlap.

6. CONCLUSIONS:

Many methodological advances have been made in survival analysis especially with the help of deep learning. The advances have largely focused on prediction accuracy and less on interpreting timeto-event predictions in the application domains of interest. This interpretation can be challenging when the number of features is large and how features relate is not obvious. In this paper, we show that survival-supervised topic modeling can address this challenge: the topics learned reveal feature co-occurrences and have relative weights indicating their impact on predicting longer or shorter survival times. These topics can be used by practitioners to check if the models agree with existing domain knowledge and to help with model debugging. These survival-supervised topic models are flexible and can be used with a variety of topic and survival models.

B THE SURVIVAL LOSS FUNCTION OF THE WEIBULL ACCELERATED FAILURE TIME (AFT) MODEL:

As a second example of a survival model that can be combined with either of our proposed survivalsupervised topic modeling approaches, we present the Weibull AFT model.

Example 2 (Weibull AFT) The Weibull AFT model assumes each subject's (possibly unobserved) survival time T i satisfies log T i = µ + β W + σε i , where µ ∈ R, β ∈ R r , and σ > 0 are model parameters, and noise variable ε has probability density function f ε (z) = exp(z -exp(z)) for z ∈ R. Under random censoring, maximum likelihood estimation for µ, β, and σ amounts to minimizing the loss

L AFT (β, µ, σ|W ) := - n i=1 δ i log f ε (z i ) -δ i log σ + (1 -δ i ) log S ε (z i ) , (B.1)

where

z i = log Y i -µ -β W i σ , S ε (t) = ∞ t f ε (u)du = e -e t .

Note that the AFT model is actually a proportional hazards model with hazard function

h(t|w) = exp(-µ/σ)t 1/σ-1 σ exp(β w).

Hence, after minimizing the loss function L AFT (β, µ, σ|W ) to estimate β, µ, and σ, we can plug in their estimates into the hazard formula above and compute an estimate for survival function S using equation (2.6).

To use the Weibull AFT survival model with our survival-supervised neural topic modeling approach in Section 4, we would minimize loss (B.1), where we redefine

z i = log Y i -µ -Ξ i σ .

As a reminder, Ξ i here is the output of the label network for subject i. In the simplest case, we would set Ξ i = β W i . However, we could replace this inner product with a neural net.

C COX AND WEIBULL AFT SURVIVAL LOSS GRADIENTS:

For archetypal analysis, Cox and Weibull AFT survival models could both be readily used by using the gradients given below.

Example 3 (Cox survival loss gradient) When L surv (W , θ) = L Cox (β|W ) as given in equation (2.5) with θ = β, we have

∂L surv (W , θ) ∂W ,g = n i=1 s.t. Yi≤Y δ i exp(β W ) n j=1 s.t. Yi≤Yj exp(β W j ) -δ β g .

Example 4 (Weibull AFT survival loss gradient) When L surv (W , θ) = L AFT (β, µ, σ|W ) as given in equation (B.1) with θ = (β, µ, σ), we have

∂L surv (W , θ) ∂W ,g = - 1 σ β g exp Y -β W -µ σ -δ . D DATASET PREPROCESSING D.1 PANCREATITIS

The pancreatitis dataset we evaluated on is created from MIMIC III, a critical care health records database containing 52 thousand individuals and their hospital encounters involving admission to the ICU at Beth Israel Deaconess center between 2001 and 2012 (Johnson et al., 2016)) . Experiments were conducted using a subset of MIMIC III version 1.4 dataset consisting of patients having pancreatitis requiring admission to the ICU. Patients were included in the study if they have an ICU admission with a primary billing code of pancreatitis, resulting in a cohort of 371 individuals. For patients who are admitted to the ICU multiple times, we only consider their first visit to the ICU. This subset of the data has no right-censoring.

Features extracted include demographics, medications, billing codes, procedures, laboratory measurements, events recorded into charts, and vitals. Features were extracted from the relational database into a 4-column format for patient id, time, event, and event value. To prevent erroneous merging of different events into a single event, and to provide more informative events, event strings are concatenations of the event descriptor prefixed with the table from which they are derived and additional relevant information such as measurement type, measurement units, etc. Because events recorded in charts are sometimes automated and sometimes manually entered, a physician-developed mapping and lower-casing all fields were used to resolve duplicate entries. As we aim to predict the patient length of stay in ICU, we extract clinical events from the subjects' electronic health records strictly before ICU admission. Our definitions of clinical events mean that a subject can have multiple instances of one event; for example, one patient might have multiple results for a particular lab test on file.

Single-occurrence categorical events (e.g., gender) are one-hot encoded, with one category removed as the reference category. Multiple-occurrence categorical events (e.g., urine color) are encoded by counting each categorys occurrences in a single subjects records. Numeric clinical events are treated as categorical by mapping observed values to equally spaced ranges by quantile (5 bins of roughly equal number of subjects per bin). Missing records are not imputed as missing certain events can have clinical significance. Therefore, for features with incomplete records, an additional feature is added solely to indicate whether missingness is observed for each subject; this approach to handling missing data is motivated by the work of Lipton et al. (2016) . After preprocessing, the total number of features used for prediction is 2557.

D.2 BREAST CANCER:

The Molecular Taxonomy of Breast Cancer International Consortium (METABRIC) dataset contains genetic and clinical information for 1981 subjects with primary breast tumors Curtis et al. (2012) .

We use the preprocessed version by Lee et al. (2018) , who one-hot encoded 21 clinical features to obtain 79 features.2 Of the 1981 subjects, 55% (1093) are right-censored. After one-hot encoding the categorical features, a total of 79 features are used for prediction in experiments that do not involve topic modeling. As the survival-supervised topic models we use do not handle continuous features, we discretize each continuous feature into 5 equal-sized bins based on feature value quantiles, the same way we did for the pancreatitis dataset. This results in a total of 100 features.

E HYPERPARAMETER SEARCH GRIDS:

We perform grid search for all models except for DEEPSURV and DEEPHIT, which use their own hyperparameter selection. We specify the search grid per method below.

Lasso-regularized Cox: lasso regularization parameter: {0.0001, 0.025, 0.05, 0.075, 0. {4, 8, 16, 32, 64, 128, 256, 512}; number of PCA components: {3, 5, 7, 10, 25} Random survival forest: number of trees: {16, 32, 64, 128, 256, 512} SURVLDA: number of topics r: {3, 5, 7, 10}; Dirichlet α: {0.1, 10} SURVIVAL-ARCHETYPES: number of archetypes r: {3, 5, 7, 10}; archetypal analysis regularization parameter λ: {10, 100, 1000}; survival loss weight η: {10 4 , 10 5 , 10 6 , 10 7 } SURVIVAL-SCHOLAR: number of topics:

{3, 5, 7, 10}; survival loss weight η: {10 1 , 10 2 , 10 3 , 10 4 , 10 5 , 10 

:

For the Weibull AFT loss given in Appendix B that has parameters β ∈ R r , µ ∈ R, and σ > 0, we would set Lsurv(W , θ) = L AFT (β, µ, σ|W ) as defined in equation (B.1) with parameters θ = (β, µ, σ).

https://github.com/chl8856/DeepHit

Figures and Tables:

EXAMPLES OF TOPICS LEARNEDNote that for each survival-supervised topic model learned, what matters in terms of how topics predict survival time is the relative values of the Cox β coefficients for that specific model. A β coefficient of 0 corresponds to a topic that gets ignored by prediction. Below, we present the topics with nonzero β coefficient learned using all training data (with hyperparameters chosen via crossvalidation) for SURVIVAL-ARCHETYPES and SURVIVAL-SCHOLAR. Per topic, we indicate the β coefficient and list the top 20 most probable words for that topic. F.1 PANCREATITIS SURVIVAL-ARCHETYPES finds 10 topics but only one that has a nonzero Cox β coefficient: Topic 1 -Cox β coefficient 4.54308798 Word Probability within topic 20 gauge insertion date:::1 0.00259127 lab:blood:chemistry:sodium:::˜q4(139.0-142.0) count 0.00154494 lab:blood:chemistry:urea nitrogen:::˜q1(2.0-12.0) count 0.00153877 lab:blood:hematology:mcv:::˜q4(90.0-94.0) count 0.00152838 lab:blood:hematology:hematocrit:::˜q5(36.5-51.8) count 0.00147893 lab:blood:chemistry:gamma glutamyltransferase:::˜missing flag 0.00147567 lab:blood:chemistry:ntprobnp:::˜missing flag 0.00147308 lab:blood:chemistry:immunoglobulin g:::˜missing flag 0.00147301 chart:cvp alarm low:mmhg:::˜missing flag 0.00147280 chart:cvp alarm high:mmhg:::˜missing flag 0.00147280 chart:alt:iu/l:::˜missing flag 0.00147145 chart:ast:iu/l:::˜missing flag 0.00147145 lab:blood:hematology:nucleated red cells:::˜missing flag 0.00147084 lab:blood:hematology:granulocyte count:::˜missing flag 0.00147082 lab:blood:chemistry:protein electrophoresis:::˜missing flag 0.00147058 lab:blood:chemistry:cholesterol, ldl, measured:::˜missing flag 0.00146990 chart:troponin-t:ng/ml:::˜missing flag 0.00146951 prescribed:sodium phosphate via iv:mmol:::˜missing flag 0.00146929 chart:tank b psi.:na:::˜missing flag 0.00146928 chart:tank a psi.:na:::˜missing flag 0.00146928 SURVIVAL-SCHOLAR finds three topics, all with nonzero Cox β coefficients: Topic 1 -Cox β coefficient 0.30225980 Word Probability within topic lab:blood:chemistry:glucose:::˜q2(95.0-112.0) count 0.28912985 lab:blood:chemistry:anion gap:::˜q3(13.0-15.0) count 0.28761890 lab:blood:chemistry:urea nitrogen:::˜q1(2.0-12.0) count 0.28706142 lab:blood:hematology:rdw:::˜q4(15.8-17.4) count 0.28468764 lab:blood:hematology:hematocrit:::˜q3(30.4-33.3) count 0.28418380 lab:blood:hematology:platelet count:::˜q2(114.0-184.0) count 0.28227943 lab:blood:hematology:mcv:::˜q3(87.0-90.0) count 0.28181386 lab:blood:chemistry:sodium:::˜q3(137.0-139.0) count 0.28117300 lab:blood:chemistry:glucose:::˜q3(112.0-134.0) count 0.28062224 lab:blood:chemistry:potassium:::˜q4(4.3-4.8) count 0.27537790 lab:blood:chemistry:potassium:::˜q3(4.0-4.

On the pancreatitis dataset, SURVLDA followed by SURVIVAL-ARCHETYPES outperform all the other methods, and the two deep learning baselines (DEEPSURV and DEEPHIT) perform worse than standard Cox proportional hazards as well as many of the other classical baselines. Meanwhile, on the breast cancer dataset, DEEPSURV achieves the best performance although Weibull AFT, k-nearest neighbors with PCA preprocessing, DEEPHIT,

Table content: <table><tr><td>Model</td><td colspan="2">Dataset Pancreatitis Breast Cancer</td></tr><tr><td>Lasso Cox</td><td>0.56</td><td>0.65</td></tr><tr><td>Lasso Cox + PCA</td><td>0.53</td><td>0.64</td></tr><tr><td>Weibull AFT</td><td>0.50</td><td>0.66</td></tr><tr><td>Weibull AFT + PCA</td><td>0.60</td><td>0.65</td></tr><tr><td>k-nearest neighbors</td><td>0.59</td><td>0.58</td></tr><tr><td>k-nearest neighbors + PCA</td><td>0.53</td><td>0.66</td></tr><tr><td>Random Survival Forest</td><td>0.56</td><td>0.60</td></tr><tr><td>DEEPSURV</td><td>0.55</td><td>0.67</td></tr><tr><td>DEEPHIT</td><td>0.53</td><td>0.66</td></tr><tr><td>SURVLDA</td><td>0.64</td><td>0.64</td></tr><tr><td>SURVIVAL-ARCHETYPES</td><td>0.63</td><td>0.63</td></tr><tr><td>SURVIVAL-SCHOLAR</td><td>0.59</td><td>0.66</td></tr></table>

Test set concordance indices for various methods on the pancreatitis and metabric datasets. Per dataset, we use bold for the highest and italics for the second highest concordance indices.

Table content: <table/>

Peer review (iclr2020):
This paper explores the problem of simultaneous learning for topic modeling and survival prediction. Their contributions are: (1) integrating two different topic modeling approaches with survival models for joint learning and (2) showing results on two medical datasets with some brief analysis of what topics are recovered. This is an interesting task to be using with topic modeling. I appreciate that the approach could be used in medical systems where interpretability of models is very important. I think the high-level idea of the algorithms is fine, but the main drawback is that the experimental section needs to be expanded on, perhaps with larger datasets and more analysis. Unless more experiments are provided, I might lean towards rejection. I think it*s ok that some of the empirical results are inconclusive, but it should lead to more quantitative error analysis. I definitely think the smaller size of the datasets could be one factor in the negative results, which is why I think the authors should try experimenting with a larger dataset. More specific comments and suggestions: - Each dataset is small (371 and 1981 data points respectively). This may be too small to effectively train some of these neural models, which I believe may be one reason why they underperform, especially on the pancreatitis data. In order for the experiments to be more conclusive, maybe you can switch to a different dataset. - Related to the size of the data, I suspect that many of the differences in the table are statistically insignificant. Can the authors please specify which of the results are significantly better than the others? - While I liked the examples in Section F of the appendix, they are difficult to compare between models, and not all of the models are included. More quantitative analysis is needed to really understand how the models’ learned topics differ in terms of coherency, cohesiveness, interpretability, etc. Providing really detailed quantitative comparisons might strengthen the analysis section of the paper. - I liked the preciseness of the descriptions of different algorithms in the background section. They were all explained nicely and easy to understand. However, the background section is a bit too detailed at times. Authors may consider condensing a bit more.