SlowGuess/NAACL_2022_OCR huggingface dataset:
Dhruv Agarwal, Rico Angell, Nicholas Monath* Andrew McCallum
Manning College of Information and Computer Sciences
{dagarwal, rangell, nmonath, mccallum}@cs.umass.edu
Learning representations of entity mentions is a core component of modern entity linking systems for both candidate generation and making linking predictions. In this paper<sup>1</sup>, we present and empirically analyze a novel training approach for learning mention and entity representations that is based on building minimum spanning arborescences (i.e., directed spanning trees) over mentions and entities across documents to explicitly model mention coreference relationships. We demonstrate the efficacy of our approach by showing significant improvements in both candidate generation recall and linking accuracy on the Zero-Shot Entity Linking dataset and MedMentions, the largest publicly available biomedical dataset. In addition, we show that our improvements in candidate generation yield higher quality re-ranking models downstream, setting a new SOTA result in linking accuracy on MedMentions. Finally, we demonstrate that our improved mention representations are also effective for the discovery of new entities via cross-document coreference.
Natural language corpora, such as biomedical research papers (Leaman and Lu, 2016), news articles (Milne and Witten, 2008; Hoffart et al., 2011), and, more generally, web page text (Gabrilovich et al., 2013; Lazic et al., 2015a), often contain ambiguous mentions of entities. Resolving this ambiguity requires mentions to either be linked to a knowledge base (KB) of entities or discovered as a new KB concept if no suitable entry exists. Grounded entity mentions are beneficial for tasks such as question-answering (Das et al., 2019), semantic search (Leaman and Lu, 2016), recommendation ranking (Noia et al., 2016), and KB construction (Ling et al., 2015). The task is made particularly challenging in zero-shot settings, where not every
entity has labeled training data (Lin et al., 2017; Logeswaran et al., 2019). In such settings, a common approach is to make use of entity descriptions, types, and aliases to form entity representations, which can then be used for making predictions.
Learned vector representations of entity mentions are an integral part of modern linking systems (Gillick et al., 2019; Wu et al., 2020, inter alia). These representations are used for (a) retrieving a short-list of entity candidates for a mention to use with a re-ranker (Wu et al., 2020), (b) making linking predictions directly (Zhang et al., 2021; Liu et al., 2020; Sung et al., 2020), and (c) performing coreference by clustering mentions to form entities (Logan IV et al., 2020).
In this work, we present a new objective and training procedure for learning mention and entity representations that explicitly model mention coreference relationships. Our proposed method uses a supervised clustering training objective based on forming a directed minimum spanning tree, or arborescence, over mentions and entities. We hypothesize that such coreference links provide a useful inductive bias because the two tasks are inherently related (Angell et al., 2021; FitzGerald et al., 2021). We thoroughly analyze the performance of the proposed procedure in each of the aforementioned use cases on MedMentions (Mohan and Li, 2019) and ZeShEL (Logeswaran et al., 2019), two challenging datasets that require zero-shot generalization at inference.
Retrieving Candidates We illustrate that our approach yields mention and entity representations useful for candidate retrieval. We show improvements over baselines that use similarly parameterized models, achieving gains of at least 7.94 and 0.93 points in recall@64 over two standard dual-encoder training procedures on MedMentions and ZeShEL, respectively. We also consider the linking capacity of our learned embeddings without re-ranking and find that their performance (i.e re
*Now at Google.
1Code, datasets, and models are available at https://github.com/dhdhagar/arboEL.
Human Language Technologies, pages 4644 - 4658
call@1) indeed improves upon our baselines. Our best performing models show gains of 13.61 & 15.46 points in linking accuracy on MedMentions and 12.06 & 1.52 points on ZeShEL.
Linking Predictions We further consider the improvement in downstream training of full cross-attention re-ranker models using higher quality candidates generated by our approach. We show consistent gains in linking accuracy on MedMentions, setting a new state-of-the-art with a 1.63 point gain over the previous best model. We also note that our proposed approach shows mixed results on ZeShEL, with one variant outperforming all compared models by at least 1.19 points, while the other two underperform the baselines. We analyze this behavior in a later section and discuss the characteristics of the data distribution sufficient to make our approach effective.
Cross-Document Coreference Finally, we illustrate that the learned representations can be used to perform coreference of mentions across documents. This indicates that our approach could be used to discover entities in settings where there is limited or no existing knowledge base of entities.
2 Arborescence-based Training for Mention & Entity Representations
In this section, we describe our approach for constructing training objectives for dual-encoders that model mention coreference relationships.
Each document  $d$  of a corpus  $\mathcal{D}$  contains a set of entity mention spans  $\mathcal{M}^d = \{m_1^d, m_2^d, \ldots, m_N^d\}$ . All mentions in the corpus are given by  $\mathcal{M} = \bigcup_{d \in \mathcal{D}} \mathcal{M}^d$ . Following (Logeswaran et al., 2019; Angell et al., 2021), we assume that these mentions are pre-identified spans of text.
Entity Linking Formally, we define the task of entity linking as follows: given a knowledge base of entities  $\mathcal{E}$  and a set of mentions  $\mathcal{M}$ , predict an entity  $e_i^d \in \mathcal{E}$  for each mention  $m_i^d$ . We use  $e_i^{\star d}$  to refer to the ground truth entity label for  $m_i^d$ .
Zero-Shot Linking The zero-shot task refers to the setting where there are entities in the knowledge base that do not have any labeled mentions in the training data. Linking decisions must instead rely on the provided information for entities, such as descriptions, aliases, and/or entity types.
Coreference We also consider a setting in which the KB of entities is not known in advance and entities must be discovered. For this task, we map every entity mention  $m_i^d$  to a cluster and assign a coreference label  $c_i^d \in \mathcal{C}$  that is independent of the entity labels in the KB.
2.2 Coreference-based Similarity
In order to jointly train both the mention and entity encoders, we define a similarity measure and an analogous procedure for sampling positive training examples that intersperses the selection of coreferent mentions and gold entities based on a single-linkage structure formed by the representations generated by the model snapshot. We construct  $k$ -nearest neighbor graphs over coreferent mention and entity clusters, followed by the application of a pruning algorithm to generate arborescence (directed MST) structures rooted at entity nodes. The resultant edges after pruning the graphs represent the pairs of positive examples used for training.
Graph-based Dissimilarity Let  $G$  be a graph with nodes  $V = \mathcal{M} \cup \mathcal{E}$  and directed edges  $E \subset V \times V$ . Each edge  $(x, y)$  of the graph has an associated weight  $w_{x, y}$ . We define a dissimilarity function  $f$  between two nodes  $u, v \in V$  to be the weight of the minimax path between the nodes, i.e.
$$
f (u, v) = \left\{ \begin{array}{l l} \min  _ {p \in u \leftrightarrow v} \max  _ {(x, y) \in p} w _ {x, y}, & \text {i f c o n n e c t e d} (u, v) \\ \infty & \text {o t h e r w i s e} \end{array} \right. \tag {1}
$$
where  $\operatorname{connected}(u, v)$  is true if there exists a directed path from node  $u$  to  $v$  in  $G$ , and  $u \rightsquigarrow v$  is the set of all paths between  $u$  and  $v$ . In words, the dissimilarity between  $u$  and  $v$  is the minimum of the largest edge weights in all paths between the two nodes, and this is often referred to as the "bottleneck edge". This measure has the property of emitting low dissimilarities between nodes even when their direct edge weight  $w_{u,v}$  is high by connecting them through a chain of low-weight edges, providing an inductive bias well-suited for coreference, i.e. not all pairs of points in a cluster are nearby (Figure 1). This inductive bias is not achieved if we sum edge weights and simply find the minimum path.
Edge Weights With this definition of dissimilarity, we now define how edge weights are calculated. We use two models: a mention-pair affinity model,  $\phi :\mathcal{M}\times \mathcal{M}\to \mathbb{R}$ , and a mention-entity affinity model,  $\psi :\mathcal{E}\times \mathcal{M}\to \mathbb{R}$ . An edge between two
Figure 1: Arborescence-based Training Objective for Mention & Entity Representations. Shown above is an illustrative example of our proposed training objective for a dual-encoder  $(\mathsf{Enc}_{\mathsf{M}},\mathsf{Enc}_{\mathsf{E}})$  on real mentions and entities from the MedMentions dataset. Mentions are highlighted in context and entities from UMLS are represented using grey boxes with the name and unique identifier. First, each mention and entity is encoded into a dense vector representation ([????????]) using the respective transformer encoder. Mention-mention and mention-entity pairs are then selected using our arborescence-based procedure as described in §2.3. The embeddings of these pairs are encouraged to be pulled closer together if both endpoints are contained in the pruned arborescence structure (represented by the shaded regions), or encouraged to be pushed farther apart if the endpoints are sampled as hard negatives.
mentions  $m_{i}$  and  $m_j$  has weight:
$$
w _ {m _ {i}, m _ {j}} = - \phi (m _ {i}, m _ {j}), \tag {2}
$$
and the weight of the edge from entity  $e$  to  $m_{i}$  is:
$$
w _ {e, m _ {i}} = - \psi (e, m _ {i}) \tag {3}
$$
Each of  $\phi (\cdot ,\cdot)$  and  $\psi (\cdot ,\cdot)$  are independently parameterized by dual-encoder transformer models (Gillick et al., 2019; Humeau et al., 2019), one for mentions  $(\mathsf{Enc}_{\mathbf{M}})$ , and one for entities  $(\mathsf{Enc}_{\mathbf{E}})$ . The affinity models are simply the inner products of the associated encoded representations:
$$
\phi \left(m _ {i}, m _ {j}\right) = \operatorname {E n c} _ {\mathrm {M}} \left(m _ {i}\right) ^ {T} \operatorname {E n c} _ {\mathrm {M}} \left(m _ {j}\right) \tag {4}
$$
$$
\psi (e, m _ {i}) = \operatorname {E n c} _ {\mathrm {E}} (e) ^ {T} \operatorname {E n c} _ {\mathrm {M}} (m _ {i})
$$
For the mention encoder,  $\mathsf{Enc}_{\mathsf{M}}$  , the transformer input is the surrounding mention context with the mention span marked by special tokens [START] and [END]:
$$
[ \mathrm {C L S} ] c _ {\text {l e f t}} [ \text {S T A R T} ] m _ {i} [ \text {E N D} ] c _ {\text {r i g h t}} [ \text {S E P} ]
$$
where  $c_{\mathrm{left}}$  and  $c_{\mathrm{right}}$  are the left and right contexts of the mention  $m_i$  in the document. For the entity encoder,  $\mathsf{Enc}_{\mathsf{E}}$ , the transformer takes as input the title and description of the entity:
$$
[ \mathrm {C L S} ] e _ {\text {t i t l e}} [ \mathrm {T T I L E} ] e _ {\text {d e s c}} [ \mathrm {S E P} ]
$$
In this input,  $e_{\mathrm{desc}}$  is the token sequence corresponding to the description of the entity, which could include natural language text related to the entity, such as a "wiki" entry, a list of entity aliases, or other available features useful in forming an entity representation. We use a special token [TITLE] to separate the title text from the description.
2.3 Training Procedure
We now define our approach for training the affinity models,  $\phi (\cdot ,\cdot)$  and  $\psi (\cdot ,\cdot)$ , and their associated encoders,  $\mathsf{Enc}_{\mathsf{M}}$  and  $\mathsf{Enc}_{\mathsf{E}}$ . Our objective is to optimize the dissimilarity function  $f(\cdot ,\cdot)$  such that the procedure infers a set of clusters that each contain exactly one entity, and every mention is assigned to the cluster containing its ground truth entity. We optimize  $f(\cdot ,\cdot)$  using mini-batch gradient descent by sequentially building batches of mentions  $B\subset \mathcal{M}$  over the training data, where each  $m_{i}\in B$  has its gold entity defined by  $e_i^\star$ . We then build a graph  $G_{B}$  with nodes consisting of (a) each  $m_{i}\in B$ , (b) each mention coreferent to  $m_{i}\in B$ , and (c) the set of gold entities for each  $m_{i}\in B$ . For every  $m_{i}$ , we build a set of directed edges defined by
$$
E _ {m _ {i}} = \left\{\left(e _ {i} ^ {\star}, m _ {\ell}\right) \mid m _ {\ell} \in \mathcal {M} _ {e _ {i} ^ {\star}} \right\} \tag {5}
$$
$$
\cup \left\{\left(m _ {\ell}, m _ {p}\right) \mid m _ {\ell}, m _ {p} \in \mathcal {M} _ {e _ {i} ^ {\star}} \right\}
$$
where  $\mathcal{M}_{e_i^\star}$  is the set of coreferent mentions with  $e_i^\star$  as the ground-truth entity. The complete set of edges in graph  $G_B$  for a mini-batch  $B$  is then given by  $E(G_B) = \bigcup_{m_i\in B}E_{m_i}$ . Observe that the resultant edges ensure that each connected component contains exactly one entity (namely, the gold entity for the mentions in that connected component).
**Forming Clusters for Positive Sampling** The graph  $G_B$  is input to a constrained clustering procedure that partitions a graph  $G$  into disjoint clusters  $\mathcal{C} = \{C_1, \ldots, C_M\}$  such that each cluster contains at most one entity. There are three constraints that every  $C \in \mathcal{C}$  must satisfy:
(i)  $|C\cap \mathcal{E}|\leq 1$
(ii)  $\forall u,v\in C$  , connected  $(u,v)\Rightarrow f(u,v)\leq \lambda$
(iii)  $\forall u,v\in C$  , connected  $(u,v)\lor$  connected  $(v,u)$
where  $\lambda$  is a hyperparameter representing the dissimilarity threshold over which edges between nodes are dropped. We set  $\lambda = \infty$  during training. These constraints ensure that (i) there is at most one entity in each cluster, (ii) if  $u$  is reachable from  $v$  then every edge in the path from  $v$  to  $u$  has a weight  $\leq \lambda$ , and (iii) each node in the cluster has a path connecting itself with every other node in the cluster.
We solve this constrained clustering problem, i.e., partition graph  $G$ , using a process similar to Angell et al. (2021). Specifically, we first remove all edges in graph  $G$  with weight greater than threshold  $\lambda$ . We then evaluate each edge  $(u,v) \in E$  in descending order of dissimilarity and check if its presence violates any of the three constraints defined above, removing the edge from  $E$  if it does. If not, we evaluate whether there is an entity in the connected component of node  $u$ , i.e.  $|C_u \cap \mathcal{E}| = 1$ . If  $|C_u \cap \mathcal{E}| = 1$ , we temporarily drop edge  $(u,v)$  and check whether  $v$  can still be reached by an entity node. If reachable, we permanently drop  $(u,v)$ , maintaining the validity of constraint (i) as well as our minimax dissimilarity function  $f(\cdot, \cdot)$ . If an entity cannot reach  $v$ , we retain edge  $(u,v)$ , preserving the connectivity of the cluster, and iterate further. Our predicted clusters are the resultant connected components in the partitioned graph  $G$ .
Using this procedure on each  $E_{m_i}$  to generate a pruned set of edges  $E_{m_i}^{\star}$ , we construct a partitioned target graph  $G_B^\star = \{E_{m_i}^\star \mid m_i \in B\}$ , which is used to optimize the parametric encoder models. Note that each mention node in a target edge set  $E_{m_i}^{\star}$  has only one incoming edge originating from either
an entity or a mention, and the selection of  $E_{m_i}^{\star}$  was done in a way to minimize the dissimilarity function  $f(\cdot ,\cdot)$  between mentions and entities with coreferent labels on the subgraph of the mini-batch.
For every cluster with an entity node, the edge structure is a directed analogue of the minimum spanning tree, where there exists a directed path from the entity node to every other node in the cluster. This structure is often referred to as the minimum spanning arborescence, thus lending its name to our method, i.e. ARBORESCENCE-based linking.
Negative Sampling Akin to the graph embedding objectives used by Nickel and Kiela (2018) and others, we construct our objective by sampling hard negative edges. For each mention  $m_{i} \in B$ , the set of negative edges  $\mathbb{N}(m_i)$  is the  $k/2$  lowest-weight incoming edges from  $\mathcal{E} \setminus \{e_i^{\star}\}$  and the  $k/2$  lowest-weight incoming edges from  $\mathcal{M} \setminus \mathcal{M}_{e_i^{\star}}$ , where  $k$  is a tuned hyperparameter. In other words, we sample negative mention and entity edges that are most similar to the gold edge.
Loss Function We define  $\Gamma(m_i) = \{u \mid (u, m_i) \in E_{m_i}^*\} \cup \{u \mid (u, m_i) \in \mathbb{N}(m_i)\}$  to be the set of all neighbors with an outgoing edge to  $m_i$  in the training graph. Let  $\mathbb{I}_{u,m_i}$  be the indicator variable such that  $\mathbb{I}_{u,m_i} = 1$  if  $(u, m_i) \in E_{m_i}^*$  and  $\mathbb{I}_{u,m_i} = 0$  otherwise. Our loss function with respect to each mention  $m_i \in B$  is then defined as follows:
$$
\begin{array}{l} \mathcal {L} \left(m _ {i}\right) = \sum_ {u \in \Gamma \left(m _ {i}\right)} \left(\mathbb {I} _ {u, m _ {i}} \log \left(\sigma_ {u} \left(w _ {u, m _ {i}}\right)\right) \right. \tag {6} \\ \left. + \left(1 - \mathbb {I} _ {u, m _ {i}}\right) \log \left(1 - \sigma_ {u} \left(w _ {u, m _ {i}}\right)\right)\right), \\ \end{array}
$$
where  $\sigma (\cdot)$  is the softmax function over all edges in  $\Gamma (m_i)\times \{m_i\}$ . The loss for the entire batch  $B$  is the mean of losses over all mentions in  $B$ . Optimizing this loss function requires simultaneously increasing the likelihood of the positive edges and decreasing the likelihood of the negative edges. This objective and training routine are inspired by the supervised single-linkage clustering proposed by Yadav et al. (2019), but differs in the choice of loss function and selection of negative examples. We also experimented with the standard cross-entropy formulation, but found its performance subpar.
We are interested in investigating the following empirical research questions:
Training Method	Recall@	MedMentions	ZeShEL
1	2	4	8	16	32	64	1	2	4	8	16	32	64
IN-BATCH NEGATIVES	58.70	69.01	75.87	80.03	83.14	85.54	87.73	39.27	53.02	62.98	70.32	75.97	80.27	84.04
K-NN NEGATIVES	56.85	65.96	71.68	76.50	80.31	83.51	86.11	49.81	60.59	68.24	74.11	78.07	81.53	84.77
TF-IDF ‡	50.8	63.8	73.4	79.2	82.3	84.6	85.3	-	-	-	-	-	-	-
IN-BATCH NEGATIVES ‡‡	-	-	-	-	-	-	-	-	-	-	-	-	-	82.06
ARBORESCENCE†	72.31	80.88	86.09	89.86	92.36	94.31	95.67	50.31	61.04	68.34	74.26	78.40	82.02	85.11
1-NN ARBORESCENCE†	71.99	80.78	86.10	89.61	91.92	93.75	95.23	51.33	62.00	69.03	74.67	78.86	81.97	85.13
1-RAND ARBORESCENCE†	71.27	80.17	85.44	89.09	91.65	93.34	94.88	50.86	62.09	69.36	75.05	78.78	82.50	85.70
Table 1: Dual-Encoder Retriever Results: Recall@k (†Proposed methods; ‡Angell et al. (2021); †Wu et al. (2020))
MedMentions	ZeShEL
|M|	Train	120K	49K
Dev	40K	10K
Test	40K	10K
|\(\mathcal{E}\)|	Train	19K	26K
Dev	9K	7K
Test	8K	7K
|\(\mathcal{E}\setminus\mathcal{E}_{\text{Train}}\)|	Dev	4K	7K
Test	4K	7K
Table 2: Dataset Statistics.  $|\mathcal{M}|$  is the number of mentions.  $|\mathcal{E}|$  is the number of unique entities in the labeled partition (not the total KB size).  $|\mathcal{E} \setminus \mathcal{E}_{\text{Train}}|$  is the number of zero-shot entities. The total KB size of MedMentions and ZeShEL is 2.3M and 492K, respectively.
- Does our proposed approach improve the recall of candidate generators?
- Do improvements in candidate generation at training lead to improvements in downstream re-ranking models?
- Does our approach result in better learned mention embeddings that can be used for coreference or discovering entities when a KB does not exist?
Experiment Details Our experiments are run on top of BLINK (Wu et al., 2020), a PyTorch (Paszke et al., 2019) implementation of dual- and cross-encoder architectures for entity linking, with model fine-tuning performed over only BERT-base, since gains from pre-trained LM size are unrelated to our approach. More details are provided in Appendix §A.1.
We run experiments on two entity linking datasets that both require generalization to unseen entities at test time. Each document in the datasets contains a set of entity mention spans, which are pre-defined using common mention-detection heuristics. KB
entities are composed of two metadata attributes - an entity title and description, which are natural language sequences of text. ZeShEL, additionally, contains a fine-grained type specification, which is needed due to the diverse disjoint domains contained in the dataset. The statistics for both datasets are reported in Table 2.
MedMentions (Mohan and Li, 2019) is a collection of titles and abstractions of bio-medical research papers. The KB that is used for this dataset is the 2017AA full-version of UMLS. The validation and test sets contain both entities that are present in the training set as well as entities that are zero-shot (never seen at training time). We use the author-recommended ST21pv subset.
ZeShEL (Logeswaran et al., 2019) is a collection of crowd-sourced wikis, which are divided into train, validation, and test splits such that no Fandom topic overlaps across the sets. In this way, all entities that appear at validation and test time are not seen during training.
3.2 Dual-Encoder Retrieval
In order to robustly evaluate the benefit of modeling coreference relationships for learning representations, we construct three variants of our proposed dual-encoder training objective, which jointly train both the mention-mention similarity function  $\phi(\cdot, \cdot)$  and the mention-entity similarity function  $\psi(\cdot, \cdot)$ . We compare to baselines that only explicitly train  $\psi(\cdot, \cdot)$  and rely on the structure of  $\phi(\cdot, \cdot)$  sharing representations with  $\psi(\cdot, \cdot)$  to provide meaningful mention-mention similarities. Our proposed objectives differ in how the positive training pairs are constructed, thus providing a way to analyze the general idea of using coreference rather than any one specific target structure for training. Our baselines are identical to each other except in how negatives are sampled.
Re-ranker	Candidate Retriever	Accuracy	Oracle
Overall	Seen	Unseen	Self	Union
K-NN NEGATIVES	Dual (IN-BATCH NEGATIVES)	73.31	77.58	58.47	80.78	47.96
K-NN NEGATIVES	Dual (K-NN NEGATIVES)	70.76	77.05	48.85	79.90	21.12
MST & K-NN (Angell et al., 2021)	TF-IDF (Angell et al., 2021)	74.1	77.3	62.9	-	-
K-NN NEGATIVES	Dual (ARBORESCENCE)†	75.73	79.97	60.99	76.09	75.64
K-NN NEGATIVES	Dual (1-NN ARBORESCENCE)†	74.73	78.91	60.19	75.48	74.71
K-NN NEGATIVES	Dual (1-RAND ARBORESCENCE)†	74.89	79.39	59.22	75.75	74.95
Table 3: MedMentions: Cross-Encoder Linking Results: We report the re-ranker accuracy trained using the candidates generated by each retriever variant. (†Proposed methods)
Arborescence In the first training variant, for each mention query, we begin by constructing a fully-connected graph of the ground truth coreferent mention cluster along with the gold entity. We then apply the pruning procedure described in the previous section to compute an arborescence rooted at the entity node. From the resultant graph, each pair of a mention and its incoming-edge node (which can either be a coreferent mention or the gold entity) is treated as a positive example for training. Following previous work (Gillick et al., 2019), we use hard negative mining with  $k = 10$  negatives composed of equal number of mention and entities.