SlowGuess/NAACL_2022_OCR huggingface dataset:
Tom Ayoola*
Joseph Fisher*
{tayoola, fshjos, apierleo}@amazon.com
Andrea Pierleoni
Recent work in entity disambiguation (ED) has typically neglected structured knowledge base (KB) facts, and instead relied on a limited subset of KB information, such as entity descriptions or types. This limits the range of contexts in which entities can be disambiguated. To allow the use of all KB facts, as well as descriptions and types, we introduce an ED model which links entities by reasoning over a symbolic knowledge base in a fully differentiable fashion. Our model surpasses state-of-the-art baselines on six well-established ED datasets by 1.3 F1 on average. By allowing access to all KB information, our model is less reliant on popularity-based entity priors, and improves performance on the challenging ShadowLink dataset (which emphasises infrequent and ambiguous entities) by 12.7 F1.
Entity disambiguation (ED) is the task of linking mentions of entities in text documents to their corresponding entities in a knowledge base (KB). Recent ED models typically use a small subset of KB information (such as entity types or descriptions) to perform linking. These models have strong performance on standard ED datasets, which consist mostly of entities that appear frequently in the training data.
However, ED performance deteriorates for less common entities, to the extent that many recent models are outperformed by outdated feature engineering-based ED systems on datasets that focus on challenging or rare entities (Provatorova et al., 2021). This suggests models over-rely on prior probabilities, which are either implicitly learned or provided as features, rather than make effective use of the mention context. One reason for this is that the subset of KB information used by the models is not enough to discriminate between
similar entities in all contexts, meaning the model has to fall back on predicting the most popular entity. Another explanation for the performance drop is that less common entities are prone to missing or inconsistent KB information (e.g. they may not have a description), which is problematic for models which rely on a single source of information. To illustrate, we find that  $21\%$  of the  $25\%$  least popular<sup>1</sup> entities in Wikidata have neither an English description nor any entity type<sup>2</sup>, leaving no mechanism for models which rely on these two sources of information alone to disambiguate them (other than their label).<sup>3</sup> Over half of these entities have at least one KB fact (e.g. [Cafe Gratitude], [headquarters location], [San Francisco]); so by including KB facts the percentage of the least popular entities with no information aside from a label drops from  $21\%$  to  $8\%$ .
Figure 1: Example of a sentence where fine-grained KB information is required for entity disambiguation.
In light of this, we introduce an ED model which has access to entity types and descriptions, and all KB facts. By using a larger variety of information, our model is more robust to missing KB information, and is able to disambiguate entities in a broader range of contexts without relying on entity priors. Figure 1 shows an example sentence where there is insufficient information in the entity
<sup>1</sup>We use the number of KB facts where the entity is the subject entity as a proxy for popularity, and only consider entities with an English Wikipedia page.
See for example Q5017238.
${}^{3}$  Conversely,  ${100}\%$  of the  ${25}\%$  most popular entities in Wikidata have either a description or type.
*Tom and Joseph contributed equally to this work.
Human Language Technologies, pages 2899 - 2912
descriptions and types to disambiguate the mention, Clinton. Fine-grained KB information, such as facts about the birthplace or education of candidate entities, is required.
To incorporate KB facts, our model begins by re-ranking candidate entities using descriptions (Wu et al., 2019) and predicted entity types (Raiman and Raiman, 2018). We then predict, using the document context, the relations which exist between every pair of mentions in the document. For example, given the sentence in Figure 1, the model may predict that the [place of birth] relation exists between the mention Clinton and the mention Hope, Arkansas. For this, we introduce a novel "coarse-to-fine" document-level relation extraction (RE) module, which increases accuracy and reduces inference time relative to the standard RE approach. Given the relation predictions, we query the KB (Wikidata in our case) for facts which exist between any of the candidate entities for the mention Clinton and for the mention Hope, Arkansas. In this case we would find the Wikidata fact [Bill Clinton], [place of birth], [Hope], and would correspondingly boost the scores of both the [Bill Clinton] and [Hope] entities. We implement this mechanism with the KB stored in a one-hot encoded sparse tensor, which makes the architecture end-to-end differentiable.
Our model surpasses state-of-the-art (SOTA) baselines on well-established ED datasets by 1.3 F1 on average, and significantly improves performance on the challenging ShadowLink dataset by 12.7 F1. In addition, the model predictions are interpretable, in that the facts used by the model to make predictions are accessible.
1. We empirically show that using KB facts for ED increases performance above SOTA methods, which generally rely on a single source of KB information.
2. We introduce a scalable method of incorporating symbolic information into a neural network ED model. To our knowledge, this is the first time an end-to-end differentiable symbolic KB has been used for ED.
3. We introduce a novel document-level relation extraction (RE) architecture which uses
coarse-to-fine predictions to obtain competitive accuracy with high efficiency.
Recent work on ED has primarily focused on feature-based approaches, whereby a neural network is optimised so that the representation of the correct KB entity is most similar to the mention representation, and each mention is resolved independently. The way in which the KB entities are represented varies between work. Initial work (Ganea and Hofmann, 2017) learned entity embeddings directly from training examples, which performed well for entities seen during training, but could not resolve unseen entities. More recent work improved performance on common datasets by enabling linking to entities unseen during training by using a subset of KB information to represent entities, such as entity descriptions (Logeswaran et al., 2019; Wu et al., 2020) or entity types (Raiman and Raiman, 2018; Onoe and Durrett, 2020).
2.1 ED with KB context
Mulang' et al. (2020) and Cetoli et al. (2019) incorporate KB facts into ED models by lexicalising KB facts and appending them to the context sentence, then using a cross-encoder model to predict whether the facts are consistent with the sentence. Our model differs from this approach as we resolve entities in the document collectively rather than independently; enabling pairwise dependencies between entity predictions to be captured. Another potential limitation of the cross-encoder method is the high computational cost of encoding the long sequence length of every fact appended to the document context. By accessing KB facts from sparse tensors, we are able to avoid this bottleneck and scale to a larger volume of facts (Cohen et al., 2020).
2.2 ED with knowledge graph embeddings
Graph neural networks (GNN) have been used to represent KB facts to inform ED predictions (Sevgili et al., 2019; Ma et al., 2021). These approaches can potentially access the information in all KB facts, but are reliant on the quality of the graph embeddings, which may struggle to represent many basic semantics (Jain et al., 2021) particularly for unpopular entities (Mohamed et al., 2020).
2.3 Global ED
There has been a series of papers which aim to optimise the global coherence of entity choices
4We use square brackets to denote relations and entities in the KB, and italics to represent mentions in the input text.
Figure 2: Our model architecture shown for a document with two mentions, England and 1966 World Cup. The model disambiguates all entity mentions in a single pass; making use of the KB facts connecting the candidates of each mention.
across the document (Hoffart et al., 2011; Cheng and Roth, 2013; Moro et al., 2014; Pershina et al., 2015). Our model differs from previous approaches in that the model predicts the relations which exist between mentions based on the document text and weights the coherence scores by these predictions, rather than considering coherence independently of document context. We also limit the model to pairwise coherence between mentions as opposed to global coherence for computational efficiency.
2.4 ED with multiple modules
The most similar work to ours is Orr et al. (2021), which achieves strong results on tail entities by introducing an ED model which uses entity embeddings, relation embeddings, type embeddings, and a KB module to link entities. A key difference to our model is the way in which KB facts are used for disambiguation. In their work, KB facts are encoded independently of the document context in which the candidate entities co-occur, whereas our model is able to leverage the relevant KB facts for the document context.
Given a document  $X$  with mentions,  $M = \{m_1, m_2, \ldots, m_{|M|}\}$ , a KB with a set of facts  $G = \{(s, r, o) \subset E \times R \times E\}$  which express relations  $r \in R$  between subject  $s \in E$  and object entities  $o \in E$ , and a description  $d_k$  for each KB entity  $e_k$ ,
the goal of ED is to assign each mention  $m \in M$  the correct corresponding KB entity  $e \in E$ .
3.2 Overview
Figure 2 shows a high-level overview of our model. We use a transformer model to encode all mentions in the document in a single-pass. We use these mention embeddings both to generate initial candidate entity scores for each mention using the entity types and descriptions of KB entities and to predict relations between every pair of mentions in the document. We retrieve KB facts for every pair of mentions in the document, for each combination of candidate entities. We weight the retrieved KB facts by multiplying the initial candidate entity score for the subject entity, the predicted score for the relation, and the initial candidate entity score for the object entity. Then we generate KB scores by summing the weighted facts for each candidate entity. The final score used for ranking entities is a weighted sum of the initial score and KB score.
3.3 Mention representation
We encode the tokens in the document  $X$  using a transformer-based model, giving contextual token embeddings  $H = \{\mathbf{h_1},\mathbf{h_2},\dots,\mathbf{h_N}\}$ . We obtain mention embeddings  $\mathbf{m_i}$  for each mention  $m_{i}$  by average pooling the contextualised token embeddings of the mention from the final transformer layer. This allows all mentions  $M$  in the document
<sup>5</sup>We use bold letters for vectors throughout our paper.
$X$  to be encoded in a single forward pass.
3.4 Initial entity score  $\psi_{\mathbf{a}}$
Initially, we score candidate entities using entity typing and description scores. We combine the two with a learned weighted sum  $\psi_{a}$ :
$$
\psi_ {a} \left(c _ {i k}\right) = w _ {1} \psi_ {t} \left(c _ {i k}\right) + w _ {2} \psi_ {d} \left(c _ {i k}\right) \tag {1}
$$
where  $c_{ik}$  is the mention-entity pair  $(m_i, e_k)$ ,  $\psi_t$  is a scoring function based on candidate entity types, and  $\psi_d$  is a scoring function based on candidate entity descriptions.
3.4.1 Entity typing score  $\psi_{\mathbf{t}}$
We construct a fixed set of types  $T = \{(r,o) \subset R \times E\}$  by taking relation-object pairs (r, o) from the KB  $G$ ; for example (instance of, song). We predict an independent unnormalised score for each type  $t \in T$  for every mention in the document by applying a linear layer  $FF_{1}$  to the mention embedding  $\mathbf{m_i}$ . To compute entity scores  $\psi_t$  using the predicted types, we calculate the dot product between the predicted types and the candidate entity's types binary vector  $\mathbf{t_k}$ . Additionally, we add a  $P(e|m)$  (PEM score) which expresses the probability of an entity given the mention text only, and is obtained from hyperlink count statistics as in previous work (Raiman and Raiman, 2018):
$$
\psi_ {t} \left(c _ {i k}\right) = \left(F F _ {1} \left(\mathbf {m} _ {\mathbf {i}}\right) \cdot \mathbf {t} _ {\mathbf {k}}\right) + P \left(e _ {k} \mid m _ {i}\right) \tag {2}
$$
3.4.2 Entity description score  $\psi_{\mathrm{d}}$
We use a bi-encoder architecture similar to (Wu et al., 2019) but altered to encode all mentions in a sequence in a single forward pass, as opposed to requiring one forward pass per mention. We represent KB entities as:
[CLS] label [SEP] description [SEP]
where "label" and "description" are the tokens of the entity label and entity description in the KB. We refer to this as  $d_{k}$ . To compute entity scores  $\psi_{d}$  using entity descriptions, we use a separate transformer model  $TR_{1}$  to encode  $d_{k}$ , taking the final layer embedding for the [CLS], and calculate the dot product between this embedding and the contextual mention embedding  $\mathbf{m}_{\mathrm{i}}$  projected by linear layer  $FF_{2}$ :
$$
\psi_ {d} \left(c _ {i k}\right) = F F _ {2} \left(\mathbf {m} _ {\mathbf {i}}\right) \cdot T R _ {1} \left(d _ {k}\right) \tag {3}
$$
3.5 Relation extraction
Our relation extraction layer outputs a relation score vector  $\hat{\mathbf{r}}_{\mathbf{ij}} \in \mathbb{R}^{|R|}$  for each pair of mentions  $m_i$  and  $m_j$  in the document, where  $R$  is the subset of relations chosen from the KB. To calculate  $\hat{\mathbf{r}}_{\mathbf{ij}}$  we begin by passing  $\mathbf{m_i}$  and  $\mathbf{m_j}$  through a bilinear layer  $B$  with output dimension 1, to predict the likelihood  $\hat{r}_{ij}^{coarse}$  that a relation exists between mentions  $m_i$  and  $m_j$ .
$$
\hat {r} _ {i j} ^ {\text {c o a r s e}} = \sigma \left(B \left(\mathbf {m} _ {\mathbf {i}}, \mathbf {m} _ {\mathbf {j}}\right)\right) \tag {4}
$$
Note that  $\hat{r}_{ij}^{\text{coarse}}$  is a scalar, denoting the likelihood that any relation exists between mention  $m_{i}$  and  $m_{j}$ .
Figure 3: The model component for document-level relation extraction.  $\hat{r}_{ij}^{coarse}$  denotes the predicted probability that any relation exists between mentions  $i$  and  $j$ .  $R$  denotes the number of relations we include in the model - set to 2 in the Figure for illustration purposes only.
We then take the top-k mention pairs with the highest values of  $\hat{r}_{ij}^{\text{coarse}}$  (in similar style to Lee et al. (2018) who introduce a coarse-to-fine approach for coreference resolution), illustrated with  $K = 2$  in Figure 3. These are the pairs of mentions which the model predicts have the highest likelihood of having a relation connecting them. For the surviving mention pairs, we pass each of the two mention embeddings individually through a linear layer,  $FF_{3}$ , to reduce their dimension by a factor of two. This ensures that when we concatenate the two representations back together we get a
<sup>6</sup>We use 1 to indicate the presence of an entity type and 0 the absence of an entity type for our binary vector. We also follow this convention for the KB facts binary vector.
representation of the mention pair  $\mathbf{m}_{\mathrm{ij}}^{*}$  of the same dimension as the contextual token embeddings  $H$
$$
\mathbf {m} _ {\mathbf {i j}} ^ {*} = c o n c a t (F F _ {3} (\mathbf {m _ {i}}), F F _ {3} (\mathbf {m _ {j}})) \quad (5)
$$
We then pass the resulting embedding  $\mathbf{m}_{\mathbf{ij}}^{*}$  through a series of transformer layers  $TR_{2}$ , where they can attend to the contextual embeddings of the original input tokens,  $H = \{\mathbf{h}_1,\mathbf{h}_2,\dots,\mathbf{h}_{\mathbf{N}}\}$ . The mention-pair embeddings from the final transformer layer are passed through a linear layer  $FF_{4}$  with output dimension  $|R|$  to give the score that each relation exists between this mention-pair,  $\hat{\mathbf{r}}_{\mathrm{ij}}^{\mathrm{fine}}$ .
$$
\hat {\mathbf {r}} _ {\mathbf {i j}} ^ {\text {f i n e}} = F F _ {4} \left(T R _ {2} \left(\mathbf {m} _ {\mathbf {i j}} ^ {*}, H\right)\right) \tag {6}
$$
Finally, to get  $\hat{\mathbf{r}}_{\mathrm{ij}}$  we multiply the coarse layer score  $\hat{r}_{ij}^{\text{coarse}}$  with the fine layer score  $\hat{\mathbf{r}}_{\mathrm{ij}}^{\text{fine}}$ , ensuring that gradients are propagated through the coarse layer during training, despite only the top-k mention pairs being passed to the fine layer.
$$
\hat {\mathbf {r}} _ {\mathbf {i j}} = \hat {r} _ {i j} ^ {\text {c o a r s e}} * \hat {\mathbf {r}} _ {\mathbf {i j}} ^ {\text {f i n e}} \tag {7}
$$
For all mention pairs outside the top-k pairs, we set  $\hat{\mathbf{r}}_{\mathrm{ij}}$  to a vector of 0s.
The relation extraction layer is trained end-to-end using the signal from the entity disambiguation loss only, and is not pretrained with any task-specific relation extraction data. To validate the effectiveness of the architecture, we include results with the RE module trained in isolation on the DO-CRED RE dataset in Appendix D.
3.6 KB score  $\psi_{\mathrm{b}}$
We retrieve KB facts for every mention-entity pair in the document and represent it as a 5-dimensional tensor  $\mathbf{r}$ , where  $\mathbf{r}_{\mathbf{ij},\mathbf{kn}}$  is a binary vector indicating the relations that exist in the KB between the two entities ( $e_k$  and  $e_n$ ) for mention-entity pair  $c_{ik}$  and  $c_{jn}$ . We weight KB facts  $\mathbf{r}$  based on initial entity scores  $\psi_a$  and relation predictions  $\hat{\mathbf{r}}$ , according their relevance to the document. To compute the KB score  $\psi_b$  for a mention-entity pair, we sum KB facts where the entity (from the mention-entity pair) is the subject entity to give score  $\psi_s$  and sum the KB facts where the entity is the object entity to give score  $\psi_o$ :
$^{7}$ Facts are efficiently retrieved by indexing into a sparse tensor.
The dimensions of tensor  $\mathbf{r}$  are: [nmentions (M), n_mentions (M), nCandidates, nCandidates, nRelations (R)]
$$
\psi_ {s} \left(c _ {i k}\right) = \dot {\psi} _ {a} \left(c _ {i k}\right) \sum_ {j = 1} ^ {j \leq | M |} \sum_ {n = 1} ^ {n \leq | E |} \left(\hat {\mathbf {r}} _ {\mathbf {i j}} \cdot \mathbf {r} _ {\mathbf {i j}, \mathbf {k n}}\right) \dot {\psi} _ {a} \left(c _ {j n}\right) \tag {8}
$$
$$
\psi_ {o} \left(c _ {i k}\right) = \dot {\psi} _ {a} \left(c _ {i k}\right) \sum_ {j = 1} ^ {j \leq | M |} \sum_ {n = 1} ^ {n \leq | E |} \left(\hat {\mathbf {r}} _ {\mathbf {j i}} \cdot \mathbf {r} _ {\mathbf {j i}, \mathbf {n k}}\right) \dot {\psi} _ {a} \left(c _ {j n}\right) \tag {9}
$$
where  $\dot{\psi}_a$  is the initial entity scoring function  $\psi_{a}$  followed by the softmax function applied over the candidate entities for the given mention. We then combine the two scores with a weighted sum giving  $\psi_{b}$ :
$$
\psi_ {b} \left(c _ {i k}\right) = w _ {3} \psi_ {s} \left(c _ {i k}\right) + w _ {4} \psi_ {o} \left(c _ {i k}\right) \tag {10}
$$
Note that for computational efficiency, this scoring mechanism considers the coherence of entity predictions between pairs of mentions only, in contrast to methods which consider global coherence Hoffart et al. (2011).
3.7 Optimisation and inference
To obtain final entity scores  $\psi_{f}$ , we add the KB scores  $\psi_{b}$  to the initial entity scores  $\psi_{a}$ .
$$
\psi_ {f} \left(c _ {i k}\right) = \psi_ {a} \left(c _ {i k}\right) + \psi_ {b} \left(c _ {i k}\right) \tag {11}
$$
We train our model on entity linked documents using cross-entropy loss. Our model is fully differentiable end-to-end, with the training signal propagating through all modules, including the relation extraction module. During ED inference, we take the candidate entity with the highest final entity score for each mention.
4.1 Standard ED
We evaluate our model on the following well-established standard ED datasets: AIDA-CoNLL (Hoffart et al., 2011), MSNBC (Cucerzan, 2007), AQUAINT (Milne and Witten, 2008), ACE2004 (Ratinov et al., 2011), CWEB (Gabrilovich et al., 2013) and WIKI (Guo and Barbosa, 2018). We train our model on Wikipedia hyperlinks and report InKB micro-F1 (which only considers entities with a non-NIL entity label). To ensure fair comparisons to baselines, we use the same method to generate candidates as previous work (Cao et al., 2021; Le and Titov, 2018). Concretely, we use the top-30 entities based on entity priors (PEM) obtained by mixing hyperlink count statistics from Wikipedia hyperlinks, a large Web corpus, and YAGO.
4.2 Long-tail and ambiguous ED
We use the ShadowLink ED dataset (Provatorova et al., 2021) to evaluate our model on long-tail and ambiguous examples. The dataset consists of 3 subsets. SHADOW where the correct entity is overshadowed by a more popular entity; TOP where the correct entity is the most popular entity; and TAIL where the correct entity is a long-tail entity. All examples in SHADOW and TOP are ambiguous, whereas TAIL has some unambiguous examples, as it is a representative sample of long-tail entities. The original dataset consists of short text snippets from Web pages, which often only include one or two mentions of entities. This limits the ability of our model to use its document-level RE module, and reason over the relationships between entities. We therefore also evaluate on the full-text version of the SHADOW and TOP subsets, referred to as SHADOW-DOC and TOP-DOC in the results tables. The dataset consists of 1 annotated entity per document, so we run spaCy ("en_core_web_lg" model) (Honnibal and Montani, 2017) to identify additional mentions to allow our model and baselines to utilise other mentions to disambiguate the annotated entity mention.
4.3 Model details
We use Wikidata (July 2021) as our KB, restricted to entities with a corresponding English Wikipedia page. This results in 6.2M entities. We use this data to generate lookups for entity types, entity descriptions, and KB facts. We select a fixed set of 1400 relation-object pairs, based on usefulness for disambiguation, to use as our entity types (Appendix A). For the KB facts, we represent the top 128 relations as separate classes and collapse the remaining relations into a single class we refer to as OTHER. Additionally, we add a special relation which exists between every entity and itself. We refer to this relation as the SAME AS relation, and the idea behind this is to enable the model to implicitly learn coreference resolution.
4.4 Training details
We use Wikipedia hyperlinks (July 2021) with additional weak labels as our training dataset, which
consists of approximately 100M labelled mentions. We limit candidate generation to top-30 entities based on entity priors obtained from Wikipedia hyperlink statistics. $^{12}$  Our model operates at the document-level and is trained using multiple mentions simultaneously. We initialise the mention embedding Transformer model weights from the RoBERTa (Liu et al., 2019) model and train our model for 1M steps with a batch size of 64 and a maximum sequence length of 512 tokens. This requires approximately 4 days when using 8 V100 GPUs. For additional details, see Appendix B.
5.1 Standard ED