SlowGuess/NAACL_2022_OCR huggingface dataset:
Aakanksha Naik  $^{1*}$  Sravanthi Parasa $^{2}$  Sergey Feldman $^{3}$  Lucy Lu Wang $^{3}$  Tom Hope $^{3,4}$
$^{2}$ Swedish Medical Group  $^{3}$ Allen Institute for Artificial Intelligence
<sup>4</sup> Paul G. Allen School for Computer Science & Engineering, University of Washington, anaik@andrew.cmu.edu {tomh, lucyw, sergey}@allenai.org
We present BEEP (Biomedical Evidence-Enhanced Predictions), a novel approach for clinical outcome prediction that retrieves patient-specific medical literature and incorporates it into predictive models. Based on each individual patient's clinical notes, we train language models (LMs) to find relevant papers and fuse them with information from notes to predict outcomes such as in-hospital mortality. We develop methods to retrieve literature based on noisy, information-dense patient notes, and to augment existing outcome prediction models with retrieved papers in a manner that maximizes predictive accuracy. Our approach boosts predictive performance on three important clinical tasks in comparison to strong recent LM baselines, increasing F1 by up to 5 points and precision@Top-K by a large margin of over  $25\%$ .
Predicting the medical outcomes of hospitalized patients holds the promise of enhancing clinical decision making. With the advent of electronic health records (EHRs), more clinical data has become available to train AI models for outcome prediction (Rajkomar et al., 2018; Hashir and Sawhney, 2020). In particular, language models pretrained on biomedical and/or clinical text are demonstrating increasing proficiency when fine-tuned for the task of predicting outcomes such as in-hospital mortality or length of stay (van Aken et al., 2021).
In this work, we explore a novel approach for improving clinical outcome prediction by dynamically retrieving relevant medical literature for each patient, and incorporating this literature into language models (LMs) trained for outcome prediction from clinical notes. This is in contrast to existing outcome prediction work that uses
Admission Note
49-year-old male... refractory hypoglycemia... type 1 diabetic...
Literature Retrieval
patient desc. + in-hospital mortality Severe hypoglycemia... not associated with increased risk of mortality in adults with Type 1 diabetes...
BEEP system integrates notes + literature
More accurate prediction: Patient survives!
Figure 1: Overview of BEEP. We retrieve literature relevant to the patient description and an outcome of interest, in-hospital mortality in this example. We combine both sources of information to train a model to predict the outcome with better accuracy.
only clinical notes (Boag et al., 2018; Hashir and Sawhney, 2020). Recent LM-based approaches van Aken et al. (2021) have designed pretraining schemes over corpora of clinical notes and general biomedical literature. This is in contrast to our work, where we directly incorporate a literature retrieval mechanism into our outcome prediction model, by finding papers relevant to specific patient cases. Our approach, named BEEP (Biomedical Evidence-Enhanced Predictions), is broadly inspired by Evidence Based Medicine (EBM)—a leading paradigm in modern medical practice which calls for finding the "current best evidence" to support optimal clinical decisions for each individual patient (Sackett et al., 1996).
Our setting presents unique challenges. First, our approach requires retrieving literature based on noisy EHR notes containing multitudes of information (e.g., medical history, ongoing treatments), unlike orthogonal efforts on extracting and summarizing scholarly information related to well-formed questions (e.g., the efficacy of ACE inhibitors in adult patients with type-2 diabetes) (Wallace, 2019; Lehman et al., 2019; DeYoung et al., 2020, 2021).
*Work done during internship at AI2.
<sup>1</sup>Our code is available at https://github.com/allenai/BEEP.
Findings of the Association for Computational Linguistics: NAACL 2022, pages 438 - 453
In addition, as our end task is predicting patient outcomes, another challenge lies in aggregating the retrieved literature in a way that maximizes prediction accuracy. Toward these challenges, we make the following key contributions:
- Literature-Augmented Model. As illustrated in Figure 1, for each ICU patient and each target outcome to be predicted (e.g., mortality), our model retrieves papers from PubMed, encoded and fused together with the ICU admission note for making a final prediction. We present several architectures for retrieving papers and for aggregating and combining them with clinical notes. We make our code, cohort selection, paper identifiers and models publicly available.
- Adding Literature Boosts Results. For evaluation, we measure both overall performance and precision/recall@Top-K, to account for the real-world scenario where "alarms" are only raised for high-confidence predictions to avoid alarm fatigue (Sendelbach and Funk, 2013). BEEP provides substantial improvements over baselines, with strong gains in overall classification performance and precision@Top-K. For example, we improve F1 by up to 5 points and precision@Top-K by a large margin of over  $25\%$ .
- Exploring Patient-Specific Retrieval. We explore a range of sparse and dense retrieval approaches, including language models, for the complex and underexplored task of retrieving relevant literature based on a patient's noisy, information-dense clinical note. Our final retrieval module employs a retrieve-erank approach that effectively retrieves helpful literature, as shown in our analysis (section 5).
We hope our work opens new research directions for automatically scanning literature for patient-specific evidence, and combining it with EHR information to boost accuracy of medical predictive models. Finally, our work raises the more general prospect of building predictive models that can dynamically learn to retrieve literature for optimizing task accuracy, in medicine and other related areas.
Patient-Specific Literature Retrieval. Since 2014, the Text REtrieval Conference (TREC) has organized a series of challenges to advance research in this area. The TREC Clinical Decision Support (CDS) tracks focused on evaluating systems on the
task of retrieving biomedical articles relevant for answering generic clinical questions about patient medical records (e.g., identifying potential diagnoses, treatments, and tests) (Simpson et al., 2014; Roberts et al., 2015, 2016). TREC CDS 2014 and 2015 used short case reports as idealized representations of medical records due to the lack of available de-identified records. TREC 2016 shifted to using real-world medical records from the Medical Information Mart for Intensive Care (MIMIC) database (Johnson et al., 2016). In our work, our focus is on predicting clinical outcomes using ICU admission notes and patient-specific retrieved literature.
Ueda et al. (2021) use contextualized representations on more structured retrieval tasks not involving clinical notes (Voorhees et al., 2021), leaving open the question of how large pretrained language models (LMs) would fare on long, noisy EHR text. We explore this by experimenting with LMs for retrieval based on EHR text.
Clinical Outcome Prediction. The idea of using automated outcome prediction for assisting clinical triage, workflow optimization, and hospital resource management has received much interest recently, especially given the conditions of the COVID-19 pandemic (Li et al., 2020). Predictive models based on structured (e.g., lab results) and unstructured (e.g., nursing notes) information have been built for key clinical outcomes including mortality (Jain et al., 2019; Feng et al., 2020), length of hospital stay (van Aken et al., 2021), readmission (Jain et al., 2019), sepsis (Feng et al., 2020), prolonged mechanical ventilation (Huang et al., 2020), and diagnostic coding (Jain et al., 2019; van Aken et al., 2021). Increasingly, models have leveraged unstructured text from notes since they can contain key information for outcome prediction (Boag et al., 2018; Jin et al., 2018). Most recently, van Aken et al. (2021) attempted this using large pretrained LMs. Our work compares the performance of a broader range of state-of-the-art pretrained language models on outcome prediction tasks.
3 BEEP: Literature-Enhanced Clinical Predictive System
Task & Approach Overview. Our goal is to improve models for clinical outcome prediction
<sup>2</sup>Since 2017, the focus has switched to TREC-PM (precision medicine) tracks where articles are retrieved based on short structured queries with attributes such as patient condition and demographics, a less realistic scenario.
Figure 2: Complete system pipeline, unpacking the high-level overview seen in Figure 1. For a given patient ICU admission note, the literature retrieval module first retrieves relevant biomedical abstracts from a clinical outcome-specific index, then reranks a top-ranked subset of abstracts. The outcome prediction module aggregates information from these reranked abstracts and fuses it with the admission note to make the final prediction
from EHR notes by augmenting them with relevant biomedical literature. BEEP consists of two main stages: (i) literature retrieval, and (ii) outcome prediction. We also briefly experiment with a formulation that trains both jointly (details in section 4). Given a patient EHR note  $Q$  and a clinical outcome of interest  $y$ , the first stage is to identify a set of biomedical abstracts  $Docs(Q) = \{D_1, \dots, D_n\}$  from PubMed that may be helpful in assessing the likelihood of the patient having that outcome. The next stage is to augment the input to an EHR-based outcome prediction model with these retrieved abstracts  $(Q \cup Docs(Q))$  and predict the final outcome. Figure 1 provides a high-level illustration of BEEP, and Figure 2 unpacks it with more detail. Next, we describe our system's main components.
3.1 Literature Retrieval Module
Our literature retrieval module consists of three components: (i) an index of biomedical abstracts pertaining to the outcome of interest, (ii) a retriever that retrieves a ranked list of abstracts relevant to the patient note from the index, and (iii) a reranker that reranks retrieved abstracts using a stronger document similarity computation model. For the retriever, we experiment with both sparse and dense models. We follow the standard retrieve-erank approach, which has been shown to achieve good balance between efficiency and retrieval performance (Dang et al., 2013), and has recently also proved useful for large-scale biomedical literature search (Wang et al., 2021). In the retrieval step, we prioritize efficiency, using models that scale well to large document collections but are not as accurate, to return a set of top documents. In the reranker step, we prioritize retrieval performance by running a computationally expensive but more accurate model on the smaller set of retrieved documents.
3.1.1 Outcome-Specific Index Construction
Since we are interested in identifying information related to a specific outcome for a patient, we begin by constructing an index of all abstracts from PubMed relevant to that outcome to limit search scope. To gather all abstracts relevant to a clinical outcome, we first identify MeSH (Medical Subject Heading) terms associated with the outcome by performing MeSH linking on the outcome descriptions using scispaCy (Neumann et al., 2019). These associated MeSH terms are then used as queries to retrieve abstracts. For some MeSH terms that are too broad (e.g., "mortality"), we include additional qualifiers (e.g., "human") to make sure we do not gather articles that are not relevant to our overall patient cohort. Appendix A lists the final set of queries used for all clinical outcomes considered in this work. Abstracts retrieved via this process are used to construct the outcome-specific index.
3.1.2 Sparse Retrieval Model
The sparse retrieval model returns top-ranked abstracts based on cosine similarity between TF-IDF vectors of MeSH terms for the query (clinical note) and the documents (outcome-specific abstracts). MeSH terms from abstracts are extracted by running scispaCy MeSH linking over the abstract text. PubMed MeSH tagging is done only at the abstract level, and does not reflect actual term frequency in the text, requiring our extraction step. However, extracting MeSH terms from clinical notes requires a more elaborate pipeline, due to two major issues:
- Entity type and boundary issues: Off-the-shelf entity extractors like scispaCy and cTAKES (Savova et al., 2010) extract some entity types that are uninformative for relevant literature retrieval, e.g., hospital names, references
<sup>3</sup>https://pubmed.ncbi.nlm.nih.gov
4https://www.ncbi.nlm.nih.gov/books/NBK25499/
to family members, etc. They also have a tendency to ignore important qualifiers. For example, given a sentence containing the entity "right lower extremity pain", both extractors returned "extremity" and "pain" as separate entities.
- Negated entities: Clinical notes have a high density of negated entities (up to  $50\%$  of (Chapman et al., 2001)). These entities must be identified and discarded prior to literature retrieval to avoid retrieving articles about symptoms and conditions that are not exhibited by the patient.
To handle these issues, we train an entity extraction model that focuses on problems, tests, and treatments with empirically good coverage of important qualifiers (Uzuner et al., 2011). We then filter negated entities with negation detection (Harkema et al., 2009) and perform entity linking to MeSH terms. For more information and implementation details see Appendix B.
3.1.3 Dense Retrieval Model
We add a dense retrieval model to complement the sparse retriever, an approach that has shown promise in recent work (Gao et al., 2021). Our dense retrieval model maps clinical notes (queries) and biomedical abstracts (documents) to a shared dense low-dimensional embedding space. Computing similarity between these encoded vectors allows for softer matching beyond surface form. For dense retrieval, we use a BERT-based bi-encoder model. We use a bi-encoder to support scaling to large document collections, as opposed to cross-encoder models which are much slower (e.g., (Gu et al., 2021)). We use PubmedBERT (Gu et al., 2021) as the encoder and train our bi-encoder using the dataset from the TREC 2016 clinical decision support task (Roberts et al., 2016). For more details, see Appendix B. Our bi-encoder achieves mean precision@10 score of 45.67 on TREC 2016 data in 5-fold cross-validation, comparable to state-of-the-art results (Das et al., 2020).
3.1.4 Reranker Model
The reranker model takes a subset of top-ranked documents from both the sparse and dense retrieval models and rescores them. We use a BERT-based cross-encoder model for reranking, prioritizing ranking performance over efficiency on this smaller subset. Given a query clinical note  $Q$  and an abstract document  $D_{i}$ , we run a PubmedBERT-based encoder over the concatenation of both ([CLS]  $Q$  [SEP]  $D_{i}$  [SEP]) to compute an embedding
$E_{QD_i}$ . This embedding is run through a linear layer to produce a relevance score, trained using cross-entropy loss with respect to document relevance labels from the TREC 2016 dataset. Our cross-encoder achieves a mean precision@10 score of 48.33 on TREC 2016 in 5-fold cross-validation, which is also comparable to state-of-the-art performance on TREC CDS 2016 (Das et al., 2020).
From the top-ranked documents returned by the reranker, the top  $k$  are selected to be passed alongside the patient clinical note to the outcome prediction module, which we describe next.
3.2 Outcome Prediction Module
The goal of this module is to compute an aggregate representation from the set of top  $k$  abstracts relevant to the clinical note, and then predict the outcome of interest using this aggregate representation and the note representation.
3.2.1 Aggregation Strategies
Let  $Dcs(Q) = D_{1}, \ldots, D_{k}$  be the set of relevant abstracts retrieved for clinical note  $Q$  and BERT  $X$  be the encoder function that returns an embedding  $E_{X}$  given a document  $X$ . We experiment with four different strategies to compute an aggregate literature representation for  $Dcs(Q)$ , which we denote by  $LR(Q)$ .
Averaging. Averaging encoder representations:
$$
L R (Q) = \frac {1}{k} \sum_ {i = 1} ^ {k} \operatorname {B E R T} \left(D _ {i}\right) \tag {1}
$$
Weighted Averaging. Weighted average of encoder representations:
$$
L R (Q) = \frac {1}{\sum_ {i = 1} ^ {k} w _ {i}} \sum_ {i = 1} ^ {k} w _ {i} \cdot \operatorname {B E R T} \left(D _ {i}\right) \tag {2}
$$
where weights  $w_{i}$  are the relevance scores computed by the reranker. The final outcome is computed by concatenating note representation  $\mathrm{BERT}(Q)$  with  $LR(Q)$  and running this through a linear layer.
We also concatenate the note embedding with each abstract  $(E_{QD_i} = [\mathrm{BERT}(Q);\mathrm{BERT}(D_i)])$ , run outcome prediction and aggregate output probabilities as follows.
Soft Voting. Averaging per-class probabilities from  $k$  outcome prediction runs:
$$
p (y = c) = \frac {1}{k} \sum_ {i = 1} ^ {k} p (y = c \mid E _ {Q D _ {i}}) \tag {3}
$$
<sup>5</sup>We treat  $k$  as a hyperparameter, see appendix C.
Outcome	0	1	2	3
PMV	3,776	3,335	-	-
MOR	43,609	5,136	-	-
LOS	5,596	16,134	13,391	8,488
(a) Class distribution for all outcomes. For PMV, classes 0 and 1 refer to cases that don't/do require prolonged ventilation. For MOR, classes 0 and 1 refer to patients that don't/do die in admission. For LOS, classes 0-3 refer to stay lengths of  $< 3$  days, 3-7 days, 1-2 weeks, and  $>2$  weeks respectively.
Outcome	Train	Dev	Test	#Articles
PMV	5,691	712	708	81,311
MOR	33,997	4,918	9,830	90,125
LOS	30,421	4,391	8,797	93,594
(b) Training, development and test splits, and total number of PubMed articles in our outcome-specific index for each clinical outcome.
Table 1: Data statistics per outcome
Weighted Voting. Weighted average of per-class probabilities from  $k$  outcome predictions runs:
$$
p (y = c) = \frac {1}{\sum_ {i = 1} ^ {k} w _ {i}} \sum_ {i = 1} ^ {k} w _ {i} \cdot p (y = c | E _ {Q D _ {i}}) \tag {4}
$$
We test our system on the task of predicting clinical outcomes from patient admission notes. Predicting outcomes from admission notes can help with early identification of at-risk patients and assist hospitals in resource planning by indicating how long patients may require hospital/ICU beds, ventilators etc. (van Aken et al., 2021).
4.1 Clinical Outcomes
We evaluate our system on three clinical outcomes:
PMV: Prolonged mechanical ventilation prediction, identifying whether a patient will require ventilation for  $>7$  days (Huang et al., 2020).
- MOR: In-hospital mortality prediction, identifying whether a patient will survive their current admission (van Aken et al., 2021).
- LOS: Length of stay prediction is the task of identifying how long a patient will need to stay in the hospital. We follow van Aken et al. (2021) and group patients into four major categories based on clinician recommendations:  $< 3$  days, 3-7 days, 1-2 weeks, and  $>2$  weeks.
PMV and MOR are binary classification tasks, while LOS is a multi-class classification task. We
predict these outcomes from patient admission notes extracted from the MIMIC III v1.4 database (Johnson et al., 2016), which contains de-identified EHR data including clinical notes in English from the Intensive Care Unit (ICU) of the Beth Israel Deaconess Medical Center in Massachusetts between 2001 and 2012. Admission notes are constructed by filtering discharge summary documents from MIMIC to only retain the following sections typically known at admission: Chief complaint, (History of) Present illness, Medical history, Admission medications, Allergies, Physical exam, Family history and Social history. Notes that do not contain any of these sections are excluded. For PMV, we follow the cohort selection process from Huang et al. (2020), and include all patients who were above 18 years of age and were on mechanical ventilation for at least 2 days with more than 6 hours each day. Patients transferred from other hospitals, organ donors, and patients with neuromuscular disease, head and neck cancer, and extensive burns, which always lead to PMV and may act as confounds, were excluded. For MOR and LOS, we follow the same cohort selection process as van Aken et al. (2021), and include all patients except newborns and remove duplicate admissions. Following these cohort selection processes results in the data splits shown in Table 1b. Table 1b also shows the numbers of relevant PubMed articles for all three clinical outcomes.
4.2 Selecting the Encoder Language Model
Since the encoder used for outcome prediction needs to produce representations for both clinical notes and relevant abstracts, we choose language models that have been pretrained on both biomedical and clinical text. We evaluate the following models on outcome prediction (without literature augmentation) to choose a suitable encoder:
- ClinicalBERT (Alsentzer et al., 2019): ClinicalBERT further pretrains BioBERT (Lee et al., 2020), a biomedical language model, on EHR notes from MIMIC III. We evaluate both versions: one trained on discharge summary notes only, and one trained on both discharge summaries and nursing notes.
- CORe (van Aken et al., 2021): CORe further pretrains BioBERT with a next sentence prediction objective on sentences describing admissions and outcomes. CORe jointly trains on EHR notes and biomedical articles.
Model	PMV	MOR	LOS
AUROC	Micro F1	Macro F1	AUROC	Micro F1	Macro F1	AUROC	Micro F1	Macro F1
BLUEBERT	54.27	53.25	51.64	81.49	89.11	62.69	73.22	45.66	44.18
+Avg	57.21	54.66	52.32	83.90	90.52	61.62	71.66	45.22	40.66
+SVote	58.16	56.07	52.63	84.21	90.60	61.00	72.54	46.02	42.46
+WVote	57.71	57.91	56.67	84.00	90.45	61.02	71.49	44.82	39.55
+WAvg	57.59	55.65	52.21	84.26	90.44	60.49	72.58	45.90	42.39
UMLSBERT	56.44	56.07	54.97	83.34	87.93	66.93	72.19	43.12	42.20
+Avg	58.36	56.50	54.62	84.02	90.41	60.28	72.25	45.61	41.58
+SVote	55.92	54.66	50.94	83.30	84.82	67.23	72.14	45.55	42.12
+WVote	59.43	56.07	54.26	84.65	90.62	62.93	72.71	46.44	42.71
+WAvg	59.30	56.50	53.70	83.59	90.35	59.61	71.02	44.58	39.95
- BLUEBERT (Peng et al., 2019): BLUEBERT further pretrains BERT (Devlin et al., 2019) jointly on EHR notes and PubMed abstracts.
- UMLSBERT (Michalopoulos et al., 2021): UMLSBERT further pretrains ClinicalBERT on EHR notes from MIMIC, with tweaks to the architecture and pretraining objective to incorporate conceptual knowledge from the Unified Medical Language System (UMLS) Metathesaurus (Schuyler et al., 1993).
Note that in this experiment, we predict clinical outcomes from patient admission notes only, without incorporating literature. We also use weighted cross-entropy loss to manage class imbalance (see Appendix B). Table 5 in the Appendix shows the performance of the above language models on the validation sets for all clinical outcomes. We select the top-performing language models BLUEBERT and UMLSBERT for our remaining experiments.
4.3 Literature Augmentation Results
We provide two sets of results: for overall performance, and for high-confidence predictions.
Overall Performance. Table 2 shows the overall performance of our literature-augmented outcome prediction system on all three clinical outcomes. We test our system using both UMLSBERT and BLUEBERT as encoders, as well as all four literature aggregation strategies. We report three metrics for each setting: (i) area under the receiver operating characteristic (AUROC), (ii) micro-averaged F1 score, and (iii) macro-averaged F1 score. From Table 2, we observe that incorporating literature leads to performance improvements on two of three