Supervised topic models for clinical interpretability

Abundant count data—procedures, diagnoses, meds—are produced during clinical care. An important question is how such data can assist treatment decisions. Standard pipelines usually involve some dimensionality reduction—there are over 14,000 diagnostic ICD9-CM codes alone—followed by training on the task of interest. Topic models such as latent Dirichlet allocation (LDA) (Blei, 2012) are a popular tool for such dimensionality reduction (e.g. Paul and Dredze (2014) or Ghassemi et al. (2014)). However, especially given noise and irrelevant signal in the data, this two-stage procedure may not produce the best predictions; thus many efforts have tried to incorporate observed labels into the dimensionality reduction model. The most natural extension is supervised LDA (McAuliffe and Blei, 2007), though other attempts exist (Zhu et al., 2012, Lacoste-Julien et al., 2009).

Unfortunately, a recent survey by Halpern et al. (2012) finds that many of these approaches have little benefit, if any, over standard LDA. We take inspiration from recent work (Chen et al., 2015) to develop an optimization algorithm that prioritizes document-topic embedding functions useful for heldout data and allows a penalized balance of generative and discriminative terms, overcoming problems with traditional maximum likelihood point estimation or more Bayesian approximate posterior estimation. We extend this work with recognition network that allows us to scale to a data set of over 800,000 patient encounters via an approximation to the ideal but expensive embedding required at each document.

We consider models for collections of $D$ documents, each drawn from the same finite vocabulary of $V$ possible word types. Each document consists of a supervised binary label $y_{d}\in\{0,1\}$ (extensions to non-binary labels are straightforward) and $N_{d}$ observed word tokens $x_{d}=\{x_{dn}\}_{n=1}^{N_{d}}$, with each word token an indicator of a vocabulary type. We can compactly write $x_{d}$ as a sparse count histogram, where $x_{dv}$ indicates the count of how many words of type $v$ appear in document $d$.

To understand how different training objectives impact both predictive performance and interpretability of topic-word parameters, we consider a version of the toy bars dataset inspired by Griffiths and Steyvers (2004), but changed so the optimal $\phi$ parameters are distinct for unsupervised LDA and supervised LDA objectives. Our dataset has 144 vocabulary words visualized as pixels in a square grid in Fig. 1. To generate the observed words $x$, we use 6 true topics: 3 horizontal bars and 3 vertical bars. However, we generate label $y_{d}$ using an expanded set of 10 topics, where the extra topics are combinations of the 6 bars. Some combinations produce positive labels, but no single bar does. We train multiple initializations of each possible training objective and penalty weight setting, and show the best run of each method in Fig. 1. Our conclusions are listed below:

Standard training that instantiates $\pi$ can either ignore labels or overfit. Fig. 1’s first column shows two problematic behaviors with the optimization objective in Eq. (2). First, when $w_{x}=1$, the topic-word parameters are basically identical whether labels are ignored ($w_{y}=0$) or included ($w_{y}=1$). Second, when the observed data is weighted very low ($w_{x}=0.01$), we see severe overfitting, where the learned embeddings at training time are not reproducible at test time.

Ideal end-to-end training can be more predictive but has expensive runtime. In contrast to the problems with standard training, we see in the middle column of Fig. 1 that using the ideal test-time embedding function $f^{*}$ also during training can produce much lower error rates on heldout data. Varying the data weight $w_{x}$ interpolates between interpretable topic-word parameters $\phi$ and good predictions. One caveat to ideal embedding is its expensiveness: Completing 100 sweeps through this 1000 document toy dataset takes about 2.5 hours using our vectorized pure Python with autograd.

Approximate end-to-end training is much cheaper and often does as well. We see in the far right column of Fig. 1 that using our proposed approximate embedding $f^{\lambda}$ often yields similar predictive power and interpretable topic-word parameters when $w_{x}>0$. Furthermore, it is about 3.6X faster to train due to avoiding the expensive embedding iterations at every document.

We study a cohort of 875080 encounters from 49322 patients drawn from two large academic medical centers with at least one ICD9 diagnostic code for major depressive disorder (ICD9s 296.2x or 3x or 311, or ICD10 equivalent). Each included patient had an identified successful treatment: a prescription repeated at least 3 times in 6 months with no change.

We extracted all procedures, diagnoses, labs, and meds from the EHR (22,000 total codewords). For each encounter, we built $x_{d}$ by concatenating count histograms from the last three months and all prior history. To simplify, we reduced this to the 9621 codewords that occurred in at least 1000 distinct encounters. The prediction goal was to identify which of 16 common anti-depressants drugs would be successful for each patient. (Predicting all 25 primaries and 166 augments is future work).

Table 1 compares each method’s area-under-the-ROC-curve (AUC) with $K=50$ topics on a held-out set of 10% of patients. We see that our training algorithm using the ideal embedding $f^{*}$ improves its predictions over a baseline unsupervised LDA model as the weight $w_{x}$ is driven to zero. Our approximate embedding $f^{\lambda}$ is roughly 2-6X faster, allowing a full pass through all 800K encounters in about 8 hours, yet offers competitive performance on many drug tasks except for those like desipramine or imipramine for which less than 1% of encounters have a positive label. Unfortunately, our best sLDA model is inferior to simple bag-of-words features plus a logistic regression classifier (rightmost column “BoW”), which we guess is due to local optima. To remedy this, future work can explore improved data-driven initializations.