SlowGuess/NAACL_2021_OCR huggingface dataset:
Rico Angell<sup>1</sup>, Nicholas Monath<sup>1</sup>, Sunil Mohan<sup>2</sup>, Nishant Yadav<sup>1</sup>, and Andrew McCallum<sup>1</sup>
<sup>1</sup> College of Information and Computer Sciences
{rangell, nmonath, nishantyadav, mccallum}@cs.umass.edu
2Chan Zuckerberg Initiative
smohan@chanzuckerberg.com
Due to large number of entities in biomedical knowledge bases, only a small fraction of entities have corresponding labelled training data. This necessitates entity linking models which are able to link mentions of unseen entities using learned representations of entities. Previous approaches link each mention independently, ignoring the relationships within and across documents between the entity mentions. These relations can be very useful for linking mentions in biomedical text where linking decisions are often difficult due mentions having a generic or a highly specialized form. In this paper, we introduce a model in which linking decisions can be made not merely by linking to a knowledge base entity but also by grouping multiple mentions together via clustering and jointly making linking predictions. In experiments we improve the state-of-the-art entity linking accuracy on two biomedical entity linking datasets including on the largest publicly available dataset.
Ambiguity is inherent in the way entities are mentioned in natural language text. Grounding such ambiguous mentions to their corresponding entities, the task of entity linking, is critical to many applications: automated knowledge base construction and completion (Riedel et al., 2013; Surdeanu et al., 2012), information retrieval (Meij et al., 2014), smart assistants (Balog and Kenter, 2019), question answering (Dhingra et al., 2020), text mining (Leaman and Lu, 2016; Murty et al., 2018).
Consider the excerpt of text from a biomedical research paper in Figure 1, the three highlighted mentions (expression, facial expressions, and facially expressive) all link to the same entity, namely C0517243 - Facial Expression in the leading biomedical KB, Unified Medical Language System (UMLS) (Bodenreider, 2004).
The mention expression is highly ambiguous and easily confused with the more prevalent en
Figure 1: Biomedical Entity Linking. All three highlighted mentions refer to the same entity. The mention expression is clearly related to the other two highlighted mentions which are much less ambiguous. If considered independently expression is more closely related to an incorrect entity.
tity, Gene expression. This linking decision may become easier with sufficient training examples (or sufficiently rich structured information in the knowledge-base). However, in biomedical (Mohan and Li, 2019) and other specialized domains (Logeswaran et al., 2019), it is often the case that the knowledge-base information is largely incomplete. Furthermore, the scarcity of training data leads to a setting in which most entities have not been observed at training.
State-of-the-art entity linking methods which are able to link entities unseen at training time make predictions for each mention independently (Logeswaran et al., 2019; Wu et al., 2019). In this way, the methods may have difficulty linking mentions which, as in the example above, have little lexical similarity with the entities in the knowledge-base, as well as mentions for which the context is highly ambiguous. These mentions cannot directly use information from one mention (or its linking decision) to inform the prediction of another mention. On the other hand, entity linking methods that do
Proceedings of the 2021 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, pages 2598-2608 June 6-11, 2021. ©2021 Association for Computational Linguistics
jointly consider entity linking decisions (Ganea and Hofmann, 2017; Le and Titov, 2018) are designed for cases in which all of the entities in the knowledge-base to have example mentions or metadata at training time (Logeswaran et al., 2019).
In this paper, we propose an entity linking model in which entity mentions are either (1) linked directly to an entity in the knowledge-base or (2) join a cluster of other mentions and link as a cluster to an entity in the knowledge-base. Some mentions may be difficult to link directly to their referent ground truth entity, but may have very clear coreference relationships to other mentions. So long as one mention among the group of mentions clustered together links to the correct entity the entire cluster can be correctly classified. This provides for a joint, tranductive-like, inference procedure for linking. We describe both the inference procedure as well as training objective for optimizing the model's inference procedure, based on recent work on supervised clustering (Yadav et al., 2019).
It is important to note that our approach does not aim to do joint coreference and linking, but rather makes joint linking predictions by clustering together mentions that are difficult to link directly to the knowledge-base. For instance, in Figure 1, the mention expression may be difficult to link to the ground truth Facial expression entity in the knowledge-base because the mention can refer to a large number of entities. However, the local syntactic and semantic information of the paragraph give strong signals that expression is coreferent with facial expression, which is easily linked to the correct entity.
We perform experiments on two biomedical entity linking datasets: MedMentions (Mohan and Li, 2019), the largest publicly available dataset as well as the benchmark BC5CDR (Li et al., 2016). We find that our approach improves over our strongest baseline by 2.3 points of accuracy on MedMentions and 0.8 points of accuracy on BC5CDR over the baseline method (Logeswaran et al., 2019). We further analyze the performance of our approach and observe that (1) our method better handles ambiguous mention surface forms (as in the example shown in Figure 1) and (2) our method can correctly link mentions even when the candidate generation step fails to provide the correct entity as a candidate.
Each document  $D \in \mathcal{D}$ , has a set of mentions  $\mathcal{M}^{(D)} = \{m_1^{(D)}, m_2^{(D)}, \ldots, m_N^{(D)}\}$ . We denote the set of all mentions across all documents as plainly  $\mathcal{M}$ . The task of entity linking is to classify each mention  $m_i$  as referent to a single entity  $e_i$  from a KB of entities. We use  $\mathcal{E}(m_i)$  to refer to the ground truth entity of mention  $m_i$  and  $\hat{e}_i$  to refer to the predicted entity.
Knowledge-bases. We assume that we are given a knowledge-base corresponding to a closed world of entities. These KBs are typically massive: English Wikipedia contains just over 6M entities<sup>1</sup> and the 2020 release of the UMLS contains 4.28M entities<sup>2</sup>. We describe in Sections 5.1 & 5.2 the details of the KBs used in each of the experiments.
Candidate Generation. Given the massive number of entities to which a mention may refer, previous work (Logeswaran et al., 2019, inter alia) uses a candidate generation step to reduce the restrict the number of entities considered for a given mention,  $m$ , to a candidate set  $\Gamma(m)$ . The recall of this step is critical to the overall performance of entity linking models.
In this section, we describe our clustering-based approach for jointly making entity linking predictions for a set of mentions. Our proposed inference method builds a graph where the nodes are the union of all of the mentions and entities and the edges have weights denoting the affinities between the endpoints. To make linking decisions, we cluster the nodes of the graph such that each cluster contains exactly one entity, following which each mention is assigned to the entity in its cluster.
3.1 Clustering-based Entity Linking
Let  $\varphi : \mathcal{M} \times \mathcal{M} \to \mathbb{R}$  and  $\psi : \mathcal{M} \times \mathcal{E} \to \mathbb{R}$  be parameterized functions which compute symmetric mention-mention and mention-entity affinities, respectively. The exact parameterizations of these functions are detailed in Section 3.2.
Define the undirected, weighted graph  $G = (V, E, w)$  where  $V = \mathcal{M} \cup \mathcal{E}$  and  $E = \mathcal{M} \times$
<sup>1</sup>number of content pages as of May 20, 2020, https://en.wikipedia.org/wiki/Special:Statistics
$^{2}$ https://www.nlm.nih.gov/research/uml/s/knowledge_sources/metathesaurus/release-notes.html
Graph w/ mention-mention and mention-entity affinities
Independently linking mentions to entities
Linking mention to entities with our proposed inference
Figure 2: Clustering-based Inference for Entity Linking. Mentions are shown in circles and entities in squares. Color families indicate ground-truth cluster assignments. The left figure shows the graph  $G$  that is the basis of the clustering task, the center figure show predictions under independent linking model, and the right figure shows our proposed inference linking mentions to entities by running our proposed constrained clustering inference procedure over  $G$  that assigns at most one entity per cluster.
$\mathcal{M} \cup \{(m, e) : e \in \Gamma(m)\}$ . The weight of each edge,  $w(v_i, v_j)$  for  $v_i, v_j \in V$ , is determined by  $\varphi$  or  $\psi$  depending on the vertices of the edge:  $w(m_i, m_j) = \varphi(m_i, m_j)$  and  $w(m_i, e_l) = \psi(m_i, e_l)$ . Linking decisions for each mention are determined by clustering the vertices of  $G$  under the constraint that every entity must appear in exactly one cluster.
Given the graph  $G$ , we start with every node in their own individual cluster. We define affinity between a pair of clusters as the strongest cross-cluster edge between nodes in the two clusters. Iteratively, we greedily merge clusters by choosing a pair of clusters with largest affinity between them under the constraint that we cannot merge two clusters which both contain an entity. When every cluster contains exactly one entity, this process can no longer merge any clusters, and thus terminates<sup>3</sup>. Each mention is linked to the entity present in its cluster at the end of inference. Algorithm 1 describes this process of constructing the graph and clustering nodes to make linking decisions more formally.
Figure 2 shows the proposed inference in action
on five entities and six mentions. Initially, every mention and entity start in a singleton cluster. In the first round, clusters  $\{m_1\}$  and  $\{m_2\}$  are merged, followed by merger of  $\{e_3\}$  and  $\{m_6\}$  in the second round, and so on. Note that in fifth round, clusters  $c_{1} = \{m_{4},e_{2}\}$  has higher affinity with  $c_{2} = \{m_{1},m_{2},m_{3},e_{1}\}$  than with  $c_{3} = \{m_{5}\}$ , yet  $c_{1}$  and  $c_{3}$  are merged instead of  $c_{1}$  and  $c_{2}$  due to the constraint that we cannot merge two clusters which both contain an entity. At the end, every mention is clustered together with exactly one entity, and there could be entities present as singleton clusters such as  $\{e_4\}$  and  $\{e_5\}$ . Note that  $m_{3}$  correctly links to its gold entity  $e_1$  as a result of being clustered with mentions  $m_{1},m_{2}$  even though it has higher affinity with entity  $e_3: w(m_3,e_3) > w(m_3,e_1)$ .
3.2 Affinity Models
We parameterize  $\psi (\cdot ,\cdot)$  and  $\phi (\cdot ,\cdot)$  using two separate deep transformer encoders (Vaswani et al., 2017) for our mention-mention affinity model and mention-entity affinity model — specifically we use the BERT architecture (Devlin et al., 2019) initialized using the weights from BioBERT (Lee et al., 2019).
This process is equivalent to single-linkage hierarchical agglomerative clustering with the constraint that two entities cannot be in the same cluster.
3.2.1 Mention-Mention Model
The mention-mention model is also a cross-encoder, taking as input a pair of mention in context and producing a single scalar affinity for every pair. The input tokens take the form:
$$
[ \mathrm {C L S} ] <   m _ {i} > [ \mathrm {S E P} ] <   m _ {j} > [ \mathrm {S E P} ]
$$
where  $<  m_{i} > \coloneqq c_{l}$  [START]  $m_{i}$  [END]  $c_{r}$
where  $m_{i}$  is the mention tokens and  $c_{l}$  and  $c_{r}$  are the left and right context of the mention in the text, respectively. The [START] and [END] tokens are special tokens fine-tuned to signify the start and end of the mention in context, respectively. We restrict the length of each input sequence to have a maximum of 256 tokens. A representations of each mention is computed using the average of the encoder's output representations corresponding to the mention's input tokens. The affinity for a mention pair is computed by concatenating their mention representations and passing it through a linear layer with a sigmoid activation. We make this affinity symmetric by averaging the two possible orderings of a pair of mentions in the cross-encoder input sequence.
3.2.2 Mention-Entity Model
The mention-entity affinity model is a cross-encoder model (Vig and Ramea, 2019; Wolf et al., 2019; Humeau et al., 2019, inter alia) and takes as input the concatenation of the mention in context with the entity description. The input tokens take the form:
$$
[ \mathrm {C L S} ] c _ {l} [ \mathrm {S T A R T} ] m [ \mathrm {E N D} ] c _ {r} [ \mathrm {S E P} ] e [ \mathrm {S E P} ]
$$
where the mention in context is the same as in the mention-mention model and  $e$  is the description of the entity. We restrict the length of this input sequence to 256 tokens. After passing the input sequence through BERT, we transform the output representation corresponding to the [CLS] token with a linear layer with one output unit. This value is finally passed through the sigmoid function to output affinity between the mention and the entity.
In this section, we explain the training procedure for the affinity models  $\varphi (\cdot ,\cdot)$  and  $\psi (\cdot ,\cdot)$  used by the clustering inference procedure. We train the mention-mention and mention-entity models independently in a way that allows the affinities to be comparable when performing inference.
We use triplet max-margin based training objectives to train both models. The most important aspect of our procedure is how we pick negatives during training. For the mention-entity model, we restrict our negatives to be from the candidate set. For the mention-mention model, we restrict our negatives to come from mentions within the same document. From these sets of possible negatives we choose the top- $k$  most offending ones according the instantaneous state of the model - i.e. the negatives with highest predicted affinities according to the model at that point during training. The following sections detail the training procedures for both models.
4.1 Mention-Mention Affinity Training
To train the mention-mention affinity model we use a variant of the maximum spanning tree (MST) supervised single linkage clustering algorithm presented in Yadav et al. (2019). Let  $\mathcal{M}_{e_l}^{(D)} = \{m\in$ $\mathcal{M}^{(D)}|\mathcal{E}(m) = e_l\}$  be the set of mentions referring to entity  $e_l$  in any one document and the set of ground truth clusters be represented by
$\mathcal{C}^* = \{\mathcal{M}_{e_l}^{(D)}|e_l\in \mathcal{E}\}$ . Let  $P$  be the set of positive training edges: the edges of the MST of the complete graph on the cluster  $C\in \mathcal{C}^*$ . Let  $N_{\varphi}(m_{*})$  be the  $k$ -nearest within document negatives to the anchor point  $m_{*}\in C$  according to the current state of the model during training. The objective of this training procedure is to minimize the following triplet max-margin loss with margin  $\mu$  for each cluster  $C\in \mathcal{C}^*$ :
$$
\mathcal {L} _ {\varphi} (\theta ; C) = \sum_ {m _ {*}, m _ {+} \in P} \sum_ {m _ {-} \in N _ {\varphi} (m _ {*})} \ell_ {\varphi , \mu} (m _ {*}, m _ {+}, m _ {-}),
$$
where  $\ell_{\varphi ,\mu}(a,p,n) = [\varphi (a,n) - \varphi (a,p) + \mu ]_{+}$
4.2 Mention-Entity Affinity Training
For the mention-entity model, we use a triplet max-margin based objective with margin  $\mu$  where the anchor is a mention  $m$  in the training set, the positive is the ground truth entity  $e_{+} = \mathcal{E}(m)$ , and the negatives are chosen from the candidate set  $\Gamma(m)$ . Denote the  $k$  most offending negatives according to the current state of the model during training as  $N_{\psi}(m) \subseteq \Gamma(m) \setminus \{\mathcal{E}(m)\}$ . Formally, the loss is
$$
\mathcal {L} _ {\psi} (\theta ; \mathcal {M}) = \sum_ {m, e _ {+}} \sum_ {e _ {-} \in N _ {\psi} (m)} \ell_ {\psi , \mu} (m, e _ {+}, e _ {-}),
$$
where  $\ell_{\psi ,\mu}(a,p,n) = [\psi (a,n) - \psi (a,p) + \mu ]_{+}$
We evaluate on biomedical entity linking using the MedMentions (Mohan and Li, 2019) and BC5CDR (Li et al., 2016) datasets. We compare to state-of-the-art methods. We then analyze the performance of our method in more detail and provide qualitative examples demonstrating our approaches' ability to use mention-mention relationships to improve candidate generation and linking.
5.1 MedMentions
MedMentions is a publicly available $^5$  dataset consisting of the titles and abstracts of 4,392 PubMed articles. The dataset is hand-labeled by annotators and contains labeled mention spans and entities linked to the 2017AA full version of UMLS. Following the suggestion of Mohan and Li (2019), we use the ST21PV subset, which restricts the entities linked in documents to a set of 21 entity types
MedMentions	BC5CDR
Train	Dev	Test	Train	Dev	Test
|M|	120K	40K	40K	18K	934	10K
|E(M)|	19K	9K	8K	2K	281	1K
% seen	100	57.7	57.5	100	80.1	64.8
Table 1: Linking Datasets. Statistics of each dataset, including the percent of ground truth entities seen during training (\% seen).
that were deemed most important for building scientific knowledge-bases. We refer the readers to Mohan and Li (2019) for a complete analysis of the dataset and provide a few important summary statistics here. The train/dev/test split partitions the PubMed articles into three non-overlapping groups. This means that some entities seen at training time will appear in dev/test and other entities will appear in dev/test but not at training time. In fact, a large number of entities that appear in dev/test time are unseen at training, about  $42\%$  of entities. See Table 1 for split details and statistics.
Previous work has evaluated on MedMentions using unfairly optimistic candidate generation settings such as using only 10 candidates including the ground truth (Zhu et al., 2019) or restricting candidates to entities appearing somewhere in the MedMentions corpus (Murty et al., 2018). We instead work in a much more general setting where all entities in UMLS are considered at candidate generation time and the generated candidates might not include the ground truth entity.
5.2 BC5CDR
BC5CDR (Li et al., 2016) is another entity linking benchmark in the biomedical domain. The dataset consists of 1,500 PubMed articles annotated with labeled disease and chemical entities. Unlike MedMentions, which contains 21 types of entities, this dataset contains just two types. These chemical and disease mentions are labeled with entities from  $\mathrm{MeSH}^6$ , a much smaller biomedical KB than UMLS. See Table 1 for split details and statistics.
5.3 Preprocessing
The MedMentions ST21PV corpus is processed as follows: (i) Abbreviations defined in the text of each paper are identified using AB3P (Sohn et al., 2008). Each definition and abbreviation instance is then replaced with the expanded form. (ii) The text of each paper in the corpus is tokenized and
4Define  $[x]_{+} = \max (x,0)$
<sup>5</sup>https://github.com/chanzuckerberg/
MedMentions
<sup>6</sup>https://www.nlm.nih.gov/mesh
Figure 3: Example predictions on Ambiguous Mentions. Here we show two example outputs for highly ambiguous mention surface forms (inhibition and activity). The independent model incorrectly makes predictions on these surface forms. The clustering-based model is able to have each ambiguous mention link to a less ambiguous mention in the same abstract and thereby make correct predictions.
split into sentences using CoreNLP (Manning et al., 2014). (iii) Overlapping mentions are resolved by preferring longer mentions that begin earlier in each sentence, and mentions are truncated at sentence boundaries. This results in 379 mentions to be dropped from the total of 203,282. (iv) Finally, the corpus is saved into the IOB2 tag format. The same preprocessing steps are used for BC5CDR, except overlapping mentions are not dropped.
5.4 Candidate Generation
For both datasets, we use a character  $n$ -gram TF-IDF model to produce candidates for all of the mentions in all splits. The candidate generator utilizes the 200k most frequent character  $n$ -grams,  $n \in \{2 \dots 5\}$  and the 200k most frequent words in the names in  $\mathcal{E}$  to produce sparse vectors for all of the mentions and entity descriptions (which in our case is the canonical name, the type, and a list of known aliases and synonyms). Table 5 provides candidate generation results for each dataset. The results report the average recall@  $K$  at different numbers of candidates  $(K)$ , i.e., whether or not the gold entity is top  $K$  candidates for a given mention.
5.5 Training and Inference Details
Our model contains 220M parameters, the majority of which are contained within the two separate BERT-based models. We optimize both the models with mini-batch stochastic gradient descent using the Adam optimizer (Kingma and Ba, 2014) with recommended learning rate of 5e-5 (Devlin et al.,
2019) with no warm-up steps. We accumulate gradients over all of the triples for a batch size of 16 within document clusters. We compute the top- $k$  most offending negatives on-the-fly for each batch by running the model in inference mode proceeding each training step. Training and inference are done on a single machine with 8 NVIDIA 1080 Ti GPUs. We train our model on MedMentions for two epochs and BC5CDR for four epochs. Training takes approximately three days for MedMentions and one day for BC5CDR. Clustering-based inference takes about three hours for MedMentions and one hour for BC5CDR. Code and data to reproduce experiments will be made available.
5.6 Results
We compare our clustering-based inference procedure, which we refer to our approach as CLUSTERING-BASED, to a state-of-the-art independent inference procedure, INDEPENDENT, which is the zero-shot architecture proposed by Logeswaran et al.. This same model is used as the mentionentity affinity model used in our approach. We also compare to an n-gram tfidf model (our candidate generation model), TAGGERONE (Leaman and Lu, 2016), BIOSYN (Sung et al., 2020), and SAPBERT (Liu et al., 2020) on both MedMentions and BC5CDR.
Table 2 shows performance of the baseline models, INDEPENDENT, and CLUSTERING-BASED inference procedure on MedMentions and BC5CDR. We report results using the gold mention segmen
MedMentions	BC5CDR
Overall Acc.	Acc. on Seen	Unseen	Overall Acc.	Acc. on Seen	Unseen
N-GRAM TF-IDF	50.9	50.9	51.0	86.9	89.2	74.6
BIOSYN (Sung et al., 2020)	72.5†	76.5†	58.7†	87.8†	89.0†	81.1†
SAPBERT (Liu et al., 2020)	69.8†	72.9†	58.9†	85.2†	85.8†	82.0†
INDEPENDENT (Logeswaran et al., 2019)	72.8	75.9	61.9	90.5	94.0	73.6
CLUSTERING-BASED (ours)	74.1	77.3	62.9	91.3	94.9	73.8
w/ Gold Types
N-GRAM TF-IDF	67.9	69.0	64.0	87.8	90.2	76.1
TAGGERONE (Leaman and Lu, 2016)	73.8	78.2	58.8	89.8	91.8	79.9
BIOSYN (Sung et al., 2020)	77.0†	80.7†	64.1†	87.9†	89.1†	81.3†
SAPBERT (Liu et al., 2020)	74.1†	77.0†	63.8†	86.0†	86.8†	82.0†
INDEPENDENT (Logeswaran et al., 2019)	76.8	79.2	68.4	90.6	94.1	73.6
CLUSTERING-BASED (ours)	79.1	81.5	70.5	91.4	94.9	74.0
Table 2: Entity Linking Results. We report linking accuracy on MedMentions and BC5CDR datasets with gold mentions spans and gold mention spans and types. We observe that CLUSTERING-BASED inference provides improved accuracy in each setting with additional improvements seen when gold entity types are provided. ( $\dagger$ Hits at one synonym — multiple entities could be predicted)