Predicting Patient Readmission Risk from Medical Text via Knowledge Graph Enhanced Multiview Graph Convolution

For each document (e.g., clinical note), we construct a weighted and undirected four-view graph $\mathcal{G}=(\mathcal{N},\mathcal{E})$ with an associated adjacency matrix $\mathbf{A}$, where $\mathcal{N}$ and $\mathcal{E}$ refer to the vertex set and edge set respectively. We also denote the representation of vertices by $\mathbf{X}$. Instead of using unique words in the document as vertices, we first conduct entity linking over the text and link the entity mentions to UMLS¹¹1We use ScispaCy (Neumann et al., 2019) as the entity linker in this work.. Consequently, we consider two types of vertices in the document-level graph $\mathcal{G}$, i.e., the unique words $\mathcal{N}_{w}$ and the linked UMLS entities $\mathcal{N}_{e}$. The vertex set $\mathcal{N}$ is thus formed as the union of $\mathcal{N}_{w}$ and $\mathcal{N}_{e}$: $\mathcal{N}=\mathcal{N}_{w}\cup\mathcal{N}_{e}$. Four views are then designed to exploit intra-document, intra-UMLS, and document-UMLS interactions that will be combined to form the adjacency matrix as follows.

In this work, we incorporate a two-layer graph convolutional network (GCN) (Kipf and Welling, 2016) to encode the graph representation of clinical notes, as depicted in Figure 1. We include an attention layer after GCN, which serves as a decoder to decode the document-level representation $\mathbf{D}_{G}$ from node embeddings. The encoding process can be described as:

where $\mathbf{\hat{A}}=\mathbf{A}+\mathbf{I}$, and $\mathbf{I}$ is the identity matrix of $\mathbf{A}$. $\mathbf{\hat{D}}$ is the diagonal degree matrix of $\mathbf{\hat{A}}$, and $\mathbf{W}^{(l)}$ is the weight matrix for the $l$-th layer where $l=0,1,2$ in this work.

We incorporate a graph summation module (Li et al., 2015; Zhang et al., 2020b) to decode the document-level representation $\mathbf{D}_{G}$ from the constructed graph, by assigning different attention weights to the nodes. The decoding process can be described as:

where $\mathbf{X}^{(2)}$ is the output of the GCN encoder and $f_{1}$, $f_{2}$ are two feed-forward networks with sigmoid and leakyrelu activation, respectively. The $f_{1}$ network acts as a soft attention mechanism that indicates the relative importance of nodes, while $f_{2}$ serves as feature transformation. The operator $\odot$ denotes element-wise multiplication. Then the document-level representation $\mathbf{D}_{G}$ is summarized as the addition of the mean and maximum values of the attentive node embeddings.

We also use a two-layer bidirectional LSTM to directly encode the document and decode the document-level representation $\mathbf{D}_{T}$ with a linear decoder, where linear transformation and max-pooling are applied. Then the two document-level representations, i.e., $\mathbf{D}_{G}$ and $\mathbf{D}_{T}$, are concatenated and fed into a MLP classifier. The model is optimized with cross-entropy loss.

The experiment is conducted based on the MIMIC-III Critical Care (Medical Information Mart for Intensive Care III) Database, which is a large, freely-available database composed of de-identified EHR data (Johnson et al., 2016). For a fair comparison, we use the same data split with the baseline (Zhang et al., 2020a), where the discharge summaries are extracted from EHRs and the generated $48,393$ documents are split into training ($80\%$), validation ($10\%$), and testing ($10\%$).

We use three metrics for evaluation, i.e., the area under the receiver operating characteristics curve (AUROC), the area under the precision recall curve (AUPRC), and the recall at precision of $80\%$ (RP80). AUROC and AUPRC are widely used for evaluating patient outcome prediction tasks, including readmission prediction (Zhang et al., 2020a; Lu et al., 2019; Lin et al., 2019). RP80 is a clinically-relevant metric that helps minimize the risk of alarm fatigue, as introduced in ClinicalBERT (Huang et al., 2019), where we fix the precision at $80\%$ and calculate the recall rate.

The following baselines are used for comparison.

• •

  BioBERT. BioBERT is a domain-specific BERT variant pre-trained on large biomedical corpora, e.g., PubMed abstracts and PMC full-text articles (Lee et al., 2020). In the experiment, we use the latest version, i.e., BioBERT v1.1, with a classification head as the baseline. The last $512$ tokens of each note are used as input to the model.

• •

  ClinicalBERT. ClinicalBERT is initialized from BioBERT v1.0 and pre-trained on MIMIC notes (Alsentzer et al., 2019). Note that there is another ClinicalBERT (Huang et al., 2019) model which presents a similar idea.

• •

  CC-LSTM. Zhang et al. propose CC-LSTM that encodes UMLS knowledge into text representations and report state-of-the-art performance on readmission prediction on the MIMIC-III dataset (Zhang et al., 2020a). For a fair comparison, we use the same pre-trained word embeddings, i.e., BioWordVec (Zhang et al., 2019), in our model.

• •

  MedText-x. Specifically, we replace the Bi-LSTM encoder with ClinicalBERT and BioBERT to demonstrate the effectiveness of the proposed graph-based knowledge injection strategy. The last two baselines are denoted by MedText-ClinicalBERT and MedText-BioBERT, respectively.

The experimental results are presented in Table 1. Generally, the proposed method, i.e., MedText, compares favorably with all the other baselines and outperforms the state-of-the-art method. Besides, directly applying pre-trained language models, such as BioBERT and ClinicalBERT, to readmission prediction does not work well. It is most likely due to the long and noisy nature of clinical notes, and only the last $512$ tokens are taken as input in the experiment. However, by combining with MedText, the performance gets improved greatly, indicating the effectiveness of the proposed graph-based knowledge injection method.

Additionally, Lin et al. propose a readmission prediction model that takes numerical features, e.g., chart events, of patients as input, and claim a state-of-the-art AUROC of $0.791$ with AUPRC of $0.513$ on the same dataset (Lin et al., 2019). This is essentially not comparable as they are using numerical features instead of text, but it highlights the value of clinical notes in EHRs.

We present the ablation study in Table 2. As shown in the table, removal of the four views will cause the performance to drop greatly, indicating the effectiveness and necessity of the four views. It is also worth noticing that the model still performs on par with CC-LSTM if the Bi-LSTM module is removed, i.e., w/o $\mathbf{D}_{T}$, and it would be more efficient in training. We also show the AUROC score with different masking threshold in Figure 2, where AUROC reaches the peak when $\gamma=0.5$. To further assess the performance of the model in terms of precision and recall, we show the P-R curve in Figure 3.

Entity linking plays an important role in this method as it is the first step of graph construction and all four views either directly or indirectly depend on the linked entities. Since a relatively high linking precision can be achieved by setting appropriate parameters of the ScispaCy linker, we mainly focus on the missed entities in the text. After manually examining a subset of notes, we roughly estimate that $15\%$ to $25\%$ of entities are not recognized or linked, which may have negatively influenced the prediction model. Some example snippets of clinical notes include:

“this is a 69 year old man with a history of end stage cardiomyopathy ( nyha class 4 ) and severe chf with an ef of 15 ( ef of 20 on milrinone drip ) as well as severe mr p/w sob , doe , pnd , weight gain of 6lbs in a week , likely due to chf exacerbation . ”

“he has a history of v-tach which responded to amiodarone . patient also has icd in place . respiratory : sob and increased o2 requirement were likely secondary to chf exacerbation and resultant pulmonary edema”

“you were admitted for increasing shortness of breath and oxygen requirements on increasing doses of lasix”

The texts in bold refer to unrecognized entity mentions. Essentially they should be linked to UMLS entities C4086268 (Exacerbation), C0034063 (pulmonary edema) and C0013404 (Dyspnea), respectively. These uncovered entities might indicate the severity of patients’ conditions and thus are critical for predicting the readmission risk.