Topic Models with Survival Supervision: Archetypal Analysis and Neural Approaches

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 Xi , 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 Yi ∈ 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 Yi ’ 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 Wi ∈ [ 0 , 1 ] r specifying the subject ’ s membership in each of the r topics ; then Wi ’ 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 ( X1 , Y1 , δ1 ) , ( X2 , Y2 , δ2 ) , . . . , ( Xn , Yn , δn ) , where the i-th training subject has feature vector Xi ∈ Rd , observed time Yi ∈ R+ , and event indicator δi ∈ { 0 , 1 } . Throughout this paper , we generally take Xi , 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 Xi , u denote the fraction of times a word appears for a specific subject , meaning that Xi , u = Xi , u∑d v=1Xi , v . Note that X is an n-by-d matrix , and we use Xi 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 Xi can be well-approximated by a convex combination of r different unknown “ archetypes ” H1 , H2 , . . . , Hr ∈ Rd : Xi ≈ r∑ g=1 Wi , gHg ( 2.1 ) for some weights Wi,1 , . . . , Wi , r ∈ [ 0 , 1 ] that sum to 1 , i.e. , the vector Wi : = ( Wi,1 , . . . , Wi , r ) resides in the probability simplex ∆r : = { w ∈ [ 0 , 1 ] r : ∑r g=1 wg = 1 } . By stacking the archetypes H1 , . . . , Hr as rows to form the matrixH , equation ( 2.1 ) can be expressed asX ≈WH . Archetypal analysis aims to estimate W and H given X . If the archetypes H1 , . . . , Hr 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. , Hg , w > 0 and Hh , w = 0 for all h 6= 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 ∈ Rd and a matrix V ∈ Rm×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 Larch ( W , H ; λ ) : = n∑ i=1 ‖Xi −H > Wi‖22︸ ︷︷ ︸ ♠ +λ r∑ g=1 D2 ( Hg , X ) ︸ ︷︷ ︸ ♥ ( 2.2 ) subject to the constraint thatWi ∈ ∆r for i = 1 , . . . , n ; constant λ ≥ 0 is a user-specified regularization parameter . Minimizing term ♠ ( error of approximating input data Xi ’ 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 Xi ’ s ( so that ♠ = 0 ) , we can move the archetypes even farther apart and still have their convex hull contain the Xi ’ 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∈Rn×r , H∈Rr×d s.t.Wi∈∆r for all i , Hg∈∆d for all g Larch ( W , H ; λ ) . ( 2.3 ) A local minimum can be found by alternating between minimizing W with H fixed , and vice versa .

Summary: The paper addresses the problem of survival analysis (predicting time until, e.g., death) using topic modeling. The point of introducing topic modeling here is to gain better insight into what helps predict survival times of unseen test subjects. This contrasts with much of the earlier work in survival analysis, which is mainly or only concerned with getting the best prediction without concern with the “thematic structure” of features. Such a thematic structure could be useful to clinicians, but analysis of the learned topic structure by clinical experts is apparently left for future work. Empirical results on pancreatitis and metabric datasets show that the methods of the paper are comparable to several established baselines.