SlowGuess/ACL_2020_OCR huggingface dataset:
Zhen Wang $^{1}$ , Jennifer Lee $^{2,3}$ , Simon Lin $^{4}$ , Huan Sun $^{1}$
<sup>1</sup>The Ohio State University
$^{2}$ Department of Family Medicine, The Ohio State University Wexner Medical Center
$^{3}$ Department of Physician Informatics, Nationwide Children's Hospital
Abigail Wexner Research Institute at Nationwide Children's Hospital
{wang.9215, sun.397}@osu.edu
{Jennifer.Lee2, Simon.Lin}@nationwidechildrens.org
Nowadays, the interpretability of machine learning models is becoming increasingly important, especially in the medical domain. Aiming to shed some light on how to rationalize medical relation prediction, we present a new interpretable framework inspired by existing theories on how human memory works, e.g., theories of recall and recognition. Given the corpus-level statistics, i.e., a global cooccurrence graph of a clinical text corpus, to predict the relations between two entities, we first recall rich contexts associated with the target entities, and then recognize relational interactions between these contexts to form model rationales, which will contribute to the final prediction. We conduct experiments on a real-world public clinical dataset and show that our framework can not only achieve competitive predictive performance against a comprehensive list of neural baseline models, but also present rationales to justify its prediction. We further collaborate with medical experts deeply to verify the usefulness of our model rationales for clinical decision making<sup>1</sup>.
Predicting relations between entities from a text corpus is a crucial task in order to extract structured knowledge, which can empower a broad range of downstream tasks, e.g., question answering (Xu et al., 2016), dialogue systems (Lowe et al., 2015), reasoning (Das et al., 2017), etc. There has been a large amount of existing work focusing on predicting relations based on raw texts (e.g., sentences, paragraphs) mentioning two entities (Hendrickx et al., 2010; Zeng et al., 2014; Zhou et al., 2016; Mintz et al., 2009; Riedel et al., 2010; Lin et al., 2016; Verga et al., 2018; Yao et al., 2019).
Figure 1: Our intuition for how to rationalize relation prediction based on the corpus-level statistics. To infer the relation between the target entities (red nodes), we recall (blue dashed line) their associated entities (blue nodes) and infer their relational interactions (red dashed line), which will serve as assumptions or model rationales to support the target relation prediction.
In this paper, we study a relatively new setting in which we predict relations between entities based on the global co-occurrence statistics aggregated from a text corpus, and focus on medical relations and clinical texts in Electronic Medical Records (EMRs). The corpus-level statistics present a holistic graph view of all entities in the corpus, which will greatly facilitate the relation inference, and can better preserve patient privacy than raw or even de-identified textual content and are becoming a popular substitute for the latter in the research community for studying EMR data (Finlayson et al., 2014; Wang et al., 2019).
To predict relations between entities based on a global co-occurrence graph, intuitively, one can first optimize the graph embedding or global word embedding (Pennington et al., 2014; Perozzi et al., 2014; Tang et al., 2015), and then develop a relation classifier (Nickel et al., 2011; Socher et al., 2013; Yang et al., 2015; Wang et al., 2018) based on the embedding vectors of the two entities. However, such kind of neural frameworks often lack the desired interpretability, which is especially important for the medical domain. In general, despite
1Our code and datasets are available at: https://github.com/zhenwang9102/X-MedRELA
8078
Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics, pages 8078-8092 July 5 - 10, 2020. ©2020 Association for Computational Linguistics
their superior predictive performance in many NLP tasks, the opaque decision-making process of neural models has concerned their adoption in high stakes domains like medicine, finance, and judiciary (Rudin, 2019; Murdoch et al., 2019). Building models that provide reasonable explanations and have increased transparency can remarkably enhance user trust (Ribeiro et al., 2016; Miller, 2019). In this paper, we aim to develop such a model for our medical relation prediction task.
To start with, we draw inspiration from the existing theories on cognitive processes about how human memory works, e.g., two types of memory retrieval (recall and recognition) (Gillund and Shiffrin, 1984). Basically, in the recall process, humans tend to retrieve contextual associations from long-term memory. For example, given the word "Paris", one may think of "Eiffel Tower" or "France", which are strongly associated with "Paris" (Nobel and Shiffrin, 2001; Kahana et al., 2008; Budiu, 2014). Besides, there is a strong correlation between the association strength and the co-occurrence graph (Spence and Owens, 1990; Lundberg and Lee, 2017). In the recognition process, humans typically recognize if they have seen a certain piece of information before. Figure 1 shows an example in the context of relation prediction. Assume a model is to predict whether Aspirin may treat Headache or not (That "Aspirin may treat Headache" is a known fact, and we choose this relation triple for illustration purposes). It is desirable if the model could perform the aforementioned two types of memory processes and produce rationales to base its prediction upon: (1) Recall. What entities are associated with Aspirin? What entities are associated with Headache? (2) Recognition. Do those associated entities hold certain relations, which can be leveraged as clues to predict the target relation? For instance, a model could first retrieve a relevant entity Pain Relief for the tail entity Headache as they co-occur frequently, and then recognize there is a chance that Aspirin can lead to Pain Relief (i.e., formulate model rationales or assumptions), based on which it could finally make a correct prediction (Aspirin, may treat, Headache).
Now we formalize such intuition to rationalize the relation prediction task. Our framework consists of three stages, global association recall (CogStage-1), assumption formation and representation (CogStage-2), and prediction decision making (CogStage-3), shown in Figure 2. CogStage-1
Figure 2: A high-level illustration of our framework.
models the process of recalling diverse contextual entities associated with the target head and tail entities respectively, CogStage-2 models the process of recognizing possible interactions between those recalled entities, which serve as model rationales (or, assumptions²) and are represented as semantic vectors, and finally CogStage-3 aggregates all assumptions to infer the target relation. We jointly optimize all three stages using a training set of relation triples as well as the co-occurrence graph. Model rationales can be captured through this process without any gold rationales available as direct supervision. Overall, our framework rationalizes its relation prediction and is interpretable to users³ by providing justifications for (i) why a particular prediction is made, (ii) how the assumptions of the prediction are developed, and (iii) how the particular assumptions are relied on.
On a real-life clinical text corpus, we compare our framework with various competitive methods to evaluate the predictive performance and interpretability. We show that our method obtains very competitive performance compared with a comprehensive list of various neural baseline models. Moreover, we follow recent work (Singh et al., 2019; Jin et al., 2020) to quantitatively evaluate model interpretability and demonstrate that rationales produced by our framework can greatly help earn expert trust. To summarize, we study the important problem of rationalizing medical relation prediction based on corpus-level statistics and propose a new framework inspired by cognitive theories, which outperforms competitive baselines in terms of both interpretability and predictive performance.
Different from existing work using raw texts for relation extraction, we assume a global co-occurrence graph (i.e., corpus-level statistics) is given, which was pre-constructed based on a text corpus  $\mathcal{D}$ , and denote it as an undirected graph  $\mathcal{G} = (\mathcal{V},\mathcal{E})$ , where
2We use the two terms interchangeably in this paper.
3Following Murdoch et al. (2019), desired interpretability is supposed to provide insights to particular audiences, which in our case are medical experts.
8079
Figure 3: Framework Overview.
each vertex  $v \in \mathcal{V}$  represents an entity extracted from the corpus and each edge  $e \in \mathcal{E}$  is associated with the global co-occurrence count for the connected nodes. Counts reflect how frequent two entities appear in the same context (e.g., co-occur in the same sentence, document, or a certain time frame). In this paper, we focus on clinical co-occurrence graph in which vertices are medical terms extracted from clinical notes. Nevertheless, as we will see later, our framework is very general and can be applied to other relations with corpus-level statistics.
Our motivation for working under this setting lies in three folds: (1) Such graph data is stripped of raw textual contexts and thus, has a better preserving of patient privacy (Wang et al., 2019), which makes itself easier to be constructed and shared under the HIPPA protected environments (Act, 1996) for medical institutes (Finlayson et al., 2014); (2) Compared with open-domain relation extraction, entities holding a medical relation oftentimes do not co-occur in a local context (e.g., a sentence or paragraph). For instance, we observe that in a widely used clinical co-occurrence graph (Finlayson et al., 2014), which is also employed for our experiments later, of all entity pairs holding the treatment relation according to UMLS (Unified Medical Language System), only about  $11.4\%$  have a co-occurrence link (i.e., co-occur in clinical notes within a time frame like 1 day or 7 days); (3) As suggested by cognitive theories (Spence and Owens, 1990), lexical co-occurrence is significantly correlated with association strength in the recall memory process, which further inspires us to utilize such statistics to find associations and form model rationales for relation prediction.
Finally, our relation prediction task is formulated as: Given the global statistics  $\mathcal{G}$  and an entity pair, we predict whether they hold a relation  $r$  (e.g., MAY_TREAT), and moreover provide a set of model rationales  $\mathcal{T}$  composed of relation triples for the
prediction. For the example in Figure 1, we aim to build a model that will not only accurately predict the MAY_TREAT relation, but also provide meaningful rationales on how the prediction is made, which are crucial for gaining trust from clinicians.
Following a high-level framework illustration in Figure 2, we show a more detailed overview in Figure 3 and introduce each component as follows.
3.1 CogStage-1: Global Association Recall
Existing cognitive theories (Kahana et al., 2008) suggest that recall is an essential function of human memory to retrieve associations for later decision making. On the other hand, the association has been shown to significantly correlate with the lexical co-occurrence from the text corpus (Spence and Owens, 1990; Lund and Burgess, 1996). Inspired by such theories and correlation, we explicitly build up our model based on recalled associations stemming from corpus-level statistics and provide global highly-associated contexts as the source of interpretations.
Given an entity, we build an estimation module to globally infer associations based on the corpus-level statistics. Our module leverages distributional learning to fully explore the graph structure. One can also directly utilize the raw neighborhoods in the co-occurrence graph, but due to the noise introduced in the preprocessing of building the graph, it is a less optimal choice in real practice.
Specifically, for a selected node/entity  $e_i \in \mathcal{E}$ , our global association recall module estimates a conditional probability  $p(e_j | e_i)$ , representing how likely the entity  $e_j \in \mathcal{E}$  is associated with  $e_i^4$ . We formally define such conditional probability as:
$$
p \left(e _ {j} \mid e _ {i}\right) = \frac {\exp \left(\boldsymbol {v} _ {e _ {j}} ^ {\prime T} \cdot \boldsymbol {v} _ {e _ {i}}\right)}{\sum_ {k = 1} ^ {| \mathcal {V} |} \exp \left(\boldsymbol {v} _ {e _ {k}} ^ {\prime T} \cdot \boldsymbol {v} _ {e _ {i}}\right)} \tag {1}
$$
4We assume all existing entities can be possible associations for the given entity.
where  $\pmb{v}_{e_i} \in \mathbb{R}^d$  is the embedding vector of node  $e_i$  and  $\pmb{v}_{e_j}' \in \mathbb{R}^d$  is the context embedding for  $e_j$ .
There are many ways to approximate  $p(e_j|e_i)$  from the global statistics, e.g., using global log-bilinear regression (Pennington et al., 2014). To estimate such probabilities and update entity embeddings efficiently, we optimize the conditional distribution  $p(e_j|e_i)$  to be close to the empirical distribution  $\hat{p}(e_j|e_i)$  defined as:
$$
\hat {p} \left(e _ {j} \mid e _ {i}\right) = \frac {p _ {i j}}{\sum_ {(i , k) \in \mathcal {E}} p _ {i k}} \tag {2}
$$
where  $\mathcal{E}$  is the set of edges in the co-occurrence graph and  $p_{ij}$  is the PPMI value calculated by the co-occurrence counts between node  $e_i$  and  $e_j$ . We adopt the cross entropy loss for the optimization:
$$
\mathcal {L} _ {n} = - \sum_ {\left(e _ {i}, e _ {j}\right) \in \mathcal {V}} \hat {p} \left(e _ {j} \mid e _ {i}\right) \log \left(p \left(e _ {j} \mid e _ {i}\right)\right) \tag {3}
$$
This association recall module will be jointly trained with other objective functions to be introduced in the following sections. After that, given an entity  $e_i$ , we can select the top- $N_c$  entities from  $p(\cdot | e_i)$  as  $e_i$ 's associative entities for subsequent assumption formation.
3.2 CogStage-2: Assumption Formation and Representation
As shown in Figure 3, with the associative entities from CogStage-1, we are ready to formulate and represent assumptions. In this paper, we define model assumptions as relational interactions between associations, that is, as shown in Figure 1, the model may identify (Caffeine, MAY_TREAT, Migraine) as an assumption, which could help predict Aspirin may treat Headache (Caffeine and Migraine are associations for Aspirin and Headache respectively). Such relational rationales are more concrete and much easier for humans to understand than the widely-adopted explanation strategy (Yang et al., 2016; Mullenbach et al., 2018; Vashishth et al., 2019) in NLP that is based on pure attention weights on local contexts.
One straightway way to obtain such rationales is to query existing medical knowledge bases (KBs), e.g., (Caffeine, MAY_TREAT, Migraine) may exist in SNOMED CT and can serve as a model rationale. We refer to rationales acquired in this way as the Closed-World Assumption (CWA) (Reiter, 1981) setting since only KB-stored facts are considered and trusted in a closed world. In contrast
to the CWA rationales, considering the sparsity and incompleteness issues of KBs that are even more severe in the medical domain, we also propose the Open-World Assumptions (OWA) (Ceylan et al., 2016) setting to discover richer rationales by estimating all potential relations between associative entities based on a seed set of relation triples (which can be regarded as prior knowledge).
In general, the CWA rationales are relatively more accurate as each fact triple has been verified by the KB, but would have a low coverage of other possibly relevant rationales for the target prediction. On the other hand, the OWA rationales are more comprehensive but could be noisy and less accurate, due to the probabilistic estimation procedure and the limited amount of prior knowledge. However, as we will see, by aggregating all OWA rationales into the whole framework with an attention-based mechanism, we can select high-quality and most relevant rationales for prediction. For the rest of the paper, by default we adopt the OWA setting in our framework and describe its details as follows.
Specifically, given a pair of head and tail entities,  $e_h, e_t \in \mathcal{V}$ , let us denote their association sets as  $\mathcal{A}(e_h) = \{a_h^i\}_{i=1}^{N_h}$  and  $\mathcal{A}(e_t) = \{a_t^j\}_{j=1}^{N_t}$ , where  $N_h, N_t$  are the number of associative entities  $a_h, a_t$  to use. Each entity has been assigned an embedding vector by the previous association recall module. We first measure the probability of relations holding for the pair. Given  $a_h^i \in \mathcal{A}(e_h), a_t^j \in \mathcal{A}(e_t)$  and a relation  $r_k \in \mathcal{R}$ , we define a scoring function as Bordes et al. (2013) to estimate triple quality:
$$
s _ {k} ^ {i j} = f \left(a _ {h} ^ {i}, r _ {k}, a _ {t} ^ {j}\right) = - \left| \left| \boldsymbol {v} _ {a _ {h} ^ {i}} + \boldsymbol {\xi} _ {k} - \boldsymbol {v} _ {a _ {t} ^ {j}} \right| \right| _ {1} \tag {4}
$$
where  $\pmb{v}_{a_h^i}$  and  $\pmb{v}_{a_t^j}$  are embedding vectors, relations are parameterized by a relation matrix  $R\in$ $\mathbb{R}^{N_r\times d}$  and  $\xi_{k}$  is its  $k$  -th row vector. Such a scoring function encourages larger value for correct triples. Additionally, in order to filter unreliable estimations, we define an NA relation to represent other trivial relations or no relation as the score,  $s_{\mathrm{NA}}^{ij} = f(a_h^i,\mathrm{NA},a_t^j)$  , which can be seen as a dynamic threshold to produce reasonable rationales.
Now we formulate OWA rationales by calculating the conditional probability of a relation given a pair of associations as follows (we save the superscript  $ij$  for space):
$$
p \left(r _ {k} \mid a _ {h} ^ {i}, a _ {t} ^ {j}\right) = \left\{ \begin{array}{l l} \frac {\exp \left(s _ {k}\right)}{\sum_ {s _ {k} \geq s _ {\mathrm {N A}}} \exp \left(s _ {k}\right)}, & s _ {k} > s _ {\mathrm {N A}} \\ 0, & s _ {k} \leq s _ {\mathrm {N A}} \end{array} \right. \tag {5}
$$
5https://www.snomed.org/
For each association pair,  $(a_h^i,a_t^j)$ , we only form an assumption with a relation  $r_k^*$  if  $r_k^*$  is top ranked according to  $p(r_{k}|a_{h}^{i},a_{t}^{j})$ .
To represent assumptions, we integrate all relation information per pair into a single vector representation. Concretely, we calculate the assumption representation by treating  $p(r_k | a_h^i, a_t^j)$  as weights for all relations as follows:
$$
\mathsf {a} _ {i j} = \rho (a _ {h} ^ {i}, a _ {t} ^ {j}; \mathcal {R}) = \sum_ {k ^ {\prime} = 1} ^ {N _ {r}} p (r _ {k ^ {\prime}} | a _ {h} ^ {i}, a _ {t} ^ {j}) \cdot \pmb {\xi} _ {k ^ {\prime}} \quad (6)
$$
Finally, we combine the entity vectors as well as the relation vector to get the final representation of assumptions for association pair  $(a_h^i,a_t^j)$ , where  $c_{i}\in \mathcal{A}(e_{h})$  and  $c_{j}\in \mathcal{A}(e_{t})$ :
$$
\mathrm {e} _ {i j} = \tanh  \left(\left[ \boldsymbol {v} _ {a _ {h} ^ {i}}; \boldsymbol {v} _ {a _ {t} ^ {j}}; \boldsymbol {a} _ {i j} \right] \boldsymbol {W} _ {p} + \boldsymbol {b} _ {p}\right) \tag {7}
$$
where  $[\cdot ;\cdot ]$  represents vector concatenation,  $W_{p}\in$ $\mathbb{R}^{3d\times d_p},b_p\in \mathbb{R}^{d_p}$  are the weight matrix and bias in a fully-connected network.
3.3 CogStage-3: Prediction Decision Making
Analogical to human thinking, our decision making module aggregates all assumption representations and measures their accountability for the final prediction. It learns a distribution over all assumptions and we select the ones with highest probabilities as model rationales. More specifically, we define a scoring function  $g(\mathbf{e}_{ij})$  to estimate the accountability based on the assumption representation  $\mathbf{e}_{ij}$  and normalize  $g(\mathbf{e}_{ij})$  as:
$$
g \left(\mathbf {e} _ {i j}\right) = \boldsymbol {v} ^ {T} \cdot \tanh  \left(\boldsymbol {W} _ {a} \mathbf {e} _ {i j} + \boldsymbol {b} _ {a}\right) \tag {8}
$$
$$
p _ {i j} = \frac {\exp (g (\mathbf {e} _ {i j}))}{\sum_ {m = 1} ^ {N _ {h}} \sum_ {n = 1} ^ {N _ {t}} \exp (g (\mathbf {e} _ {m n}))} \tag {9}
$$
where  $W_{a},b_{a}$  are the weight matrix and bias for the scoring function. Then we get the weighted rationale representation as:
$$
\mathbf {r} = \psi \left(e _ {h}, e _ {t}\right) = \sum_ {i = 1} ^ {N _ {h}} \sum_ {j = 1} ^ {N _ {t}} p _ {i j} \mathbf {e} _ {i j} \tag {10}
$$
With the representation of weighted assumption information for the target pair  $(e_h, e_t)$ , we calculate the binary prediction probability for relation  $r$  as:
$$
p (r \mid e _ {h}, e _ {t}) = \sigma \left(\boldsymbol {W} _ {r} r + \boldsymbol {b} _ {r}\right) \tag {11}
$$
where  $\sigma (x) = 1 / (1 + \exp (-x))$  and  $W_{r},b_{r}$  are model parameters.
Rationalizing relation prediction. After fully training the entire model, to recover the most contributing assumptions for predicting the relation between the given target entities  $(e_h, e_t)$ , we compute the importance scores for all assumptions and select those most important ones as model rationales. In particular, we multiply  $p_{ij}$  (the weight for association pair  $(a_h^i, a_t^j)$  in Eqn. 9) with  $p(r_k | a_h^i, a_t^j)$  (the probability of a relation given the pair  $(a_h^i, a_t^j)$  in Eqn. 5) to score the triple  $(a_h^i, r_k, a_t^j)$ . We rank all such triples for  $a_h^i \in \mathcal{A}(e_h), a_t^j \in \mathcal{A}(e_t), r_k \in \mathcal{R}$  and select the top- $K$  triples as model rationales for the final relation prediction.
We now describe how we train our model efficiently for multiple modules. For relational learning to estimate the conditional probability  $p(r_k | a_h^i, a_t^j)$ , we utilize training data as the seed set of triples for all relations as correct triples denoted as  $(h, r, t) \in \mathcal{P}$ . The scoring function in Eqn. 4 is expected to score higher for correct triples than the corrupted ones in which we denote  $\mathcal{N}(?, r, t)$  ( $\mathcal{N}(t, r, ?)$ ) as the set of corrupted triples by replacing the head (tail) entity randomly. Instead of using margin-based loss function, we adopt a more efficient training strategy from (Kadlec et al., 2017; Toutanova and Chen, 2015) with a negative log likelihood loss function as:
$$
\begin{array}{l} \mathcal {L} _ {r} = - \sum_ {(h, r, t) \in \mathcal {P}} \log p (h | t, r) \tag {12} \\ - \sum_ {(h, r, t) \in \mathcal {P}} \log p (t | h, r) \\ \end{array}
$$
where the conditional probability  $p(h|t,r)$  is defined as follows ( $p(t|h,r)$  is defined similarly):
$$
p (h \mid t, r) = \frac {\exp (f (h , r , t))}{\sum_ {h ^ {\prime} \in \mathcal {N} (\ref {e q : 1})} \exp (f (h ^ {\prime} , r , t))} \tag {13}
$$
For our binary relation prediction task, we define a binary cross entropy loss function with Eqn. 11 as follows:
$$
\begin{array}{l} \mathcal {L} _ {p} = - \sum_ {i = 1} ^ {M} \left(y _ {i} \cdot \log \left(p \left(r \mid e _ {h} ^ {i}, e _ {t} ^ {i}\right)\right) \right. \tag {14} \\ + \left(1 - y _ {i}\right) \cdot \log \left(1 - p \left(r \left| e _ {h} ^ {i}, e _ {t} ^ {i}\right)\right)\right) \\ \end{array}
$$
where  $M$  is the number of samples,  $y_{i}$  is the label showing whether  $e_h, e_t$  holds a certain relation.
The above three loss functions, i.e.,  $\mathcal{L}_n$  for global association recall,  $\mathcal{L}_r$  for relational learning and  $\mathcal{L}_p$  for relation prediction, are all jointly optimized. All three of them share the entity embeddings and  $\mathcal{L}_p$  will reuse the relation matrix from  $\mathcal{L}_r$  to conduct the rationale generation.
<sup>6</sup>We remove the target relation to predict if it exists in the assumption set.
In this section, we first introduce our experimental setup, e.g., the corpus-level co-occurrence statistics and datasets used for our experiments, and then compare our model with a list of comprehensive competitive baselines in terms of predictive performance. Moreover, we conduct expert evaluations as well as case studies to demonstrate the usefulness of our model rationales.