Neural Topic Models with Survival Supervision: Jointly Predicting Time-to-Event Outcomes and Learning How Clinical Features Relate

Abstract

In time-to-event prediction problems, a standard approach to estimating an interpretable model is to use Cox proportional hazards, where features are selected based on lasso regularization or stepwise regression. However, these Cox-based models do not learn how different features relate. As an alternative, we present an interpretable neural network approach to jointly learn a survival model to predict time-to-event outcomes while simultaneously learning how features relate in terms of a topic model. In particular, we model each subject as a distribution over “topics”, which are learned from clinical features as to help predict a time-to-event outcome. From a technical standpoint, we extend existing neural topic modeling approaches to also minimize a survival analysis loss function. We study the effectiveness of this approach on seven healthcare datasets on predicting time until death as well as hospital ICU length of stay, where we find that neural survival-supervised topic models achieves competitive accuracy with existing approaches while yielding interpretable clinical “topics” that explain feature relationships.

A Interpreting Topic Heatmaps

In this appendix, we explain how to interpret our topic heatmaps (Fig. 1 and additional plots in our code repository). For many topic models including LDA, a topic is represented as a distribution over d vocabulary words. scholar [5] (and also our survival-supervised version survScholar) reparameterizes these topic distributions; borrowing from SAGE [10], scholar represents a topic as a deviation from a background log-frequency vector. This vector accommodates common words that have similar frequencies across data points. When we visualize a topic, we take this modeling approach into account and only choose to highlight features that have positive log-deviations from the background. Given a topic, having positive log-deviation is analogous to having higher conditional probabilities in the classic topic modeling case but explicitly is relative to background word frequencies (rather than being raw topic word probabilities).

To fill in the details, in step 2(a) of survScholar’s generative process (stated in Sect. 3), each word is drawn from the conditional distribution $\text {softmax}(\gamma + w^T B)$, where $\gamma \in \mathbb {R}^d$ is the background log-frequency vector, $w\in \mathbb {R}^k$ contains a sample’s topic membership weights, and $B\in \mathbb {R}^{k\times d}$ encodes (per topic) every vocabulary word’s log-deviation from the word’s background. This is a reparameterization of how LDA is encoded, which has each word drawn from the conditional distribution $\text {softmax}(w^T H)$ for $H\in \mathbb {R}^{k\times d}$. In particular, note that $H_g = \gamma + B_g$ for every topic $g\in \{1,2,\dots ,k\}$. The background log-frequency vector $\gamma$ is learned during neural net training. Note that SAGE [10] further encourages sparsity in B by adding $\ell _1$ regularization on B.

We found ranking words within a topic by their raw probabilities ($A_g$ in Eq. (2.1)) to be less interpretable than ranking words based on their deviations from their background frequencies ($B_g$) precisely because commonly occurring background words make interpretation difficult. In fact, when Dawson and Kendziorski [8] introduced survLDA, they used an ad hoc pre-processing step to identify background words to exclude from analysis altogether. We avoid this pre-processing and use log-deviations from background frequencies instead.

In heatmaps such as the one in Fig. 1, each column corresponds to a topic. For the g-th topic, instead of plotting its raw log-deviations (encoded in $B_g\in \mathbb {R}^d$), which are harder to interpret, we exponentiated each word’s log-deviation to get the word’s multiplicative ratio from its background frequency (i.e., we compute $\exp (B_g)$); the color bar intensity values are precisely these multiplicative ratios of how often a word appears relative to the word’s background frequency.

To highlight features that distinguish topics from one another, we also sort rows in the heatmap by descending differences between the largest and smallest values in a row. Thus, features whose deviations vary greatly across topics tend to show up on the top. A technical detail is that we sorted with respect to the original features, rather than the one-hot encoded or binned features. Therefore, as an example, all bins under mean blood pressure stay together. For features associated with multiple rows in the heatmap, we computed the difference between the largest and smallest values for each row, and used the largest difference (across rows) for sorting.