SlowGuess/NAACL_2021_OCR huggingface dataset:
Yogarshi Vyas Miguel Ballesteros
{yogarshi, ballemig}@amazon.com
In entity linking, mentions of named entities in raw text are disambiguated against a knowledge base (KB). This work focuses on linking to unseen KBs that do not have training data and whose schema is unknown during training. Our approach relies on methods to flexibly convert entities with several attribute-value pairs from arbitrary KBs into flat strings, which we use in conjunction with state-of-the-art models for zero-shot linking. We further improve the generalization of our model using two regularization schemes based on shuffling of entity attributes and handling of unseen attributes. Experiments on English datasets where models are trained on the CoNLL dataset, and tested on the TAC-KBP 2010 dataset show that our models are  $12\%$  (absolute) more accurate than baseline models that simply flatten entities from the target KB. Unlike prior work, our approach also allows for seamlessly combining multiple training datasets. We test this ability by adding both a completely different dataset (Wikia), as well as increasing amount of training data from the TAC-KBP 2010 training set. Our models are more accurate across the board compared to baselines.
Entity linking consists of linking mentions of entities found in text against canonical entities found in a target knowledge base (KB). Early work in this area was motivated by the availability of large KBs with millions of entities (Bunescu and Paşa, 2006). Most subsequent work has followed this tradition of linking to a handful of large, publicly available KBs such as Wikipedia, DBPedia (Auer et al., 2007) or the KBs used in the now decade-old TAC-KBP challenges (McNamee and Dang, 2009; Ji et al., 2010). As a result, previous work always assumes complete knowledge of the schema of the target KB that entity linking models are trained for, i.e. how many and which attributes are used to
represent entities in the KB. This allows training supervised machine learning models that exploit the schema along with labeled data that link mentions to this a priori known KB. However, this strong assumption breaks down in scenarios which require linking to KBs that are not known at training time. For example, a company might want to automatically link mentions of its products to an internal KB of products that has a rich schema with several attributes such as product category, description, dimensions, etc. It is very unlikely that the company will have training data of this nature, i.e. mentions of products linked to its database.
Our focus is on linking entities to unseen KBs with arbitrary schemas. One solution is to annotate data that can be used to train specialized models for each target KB of interest, but this is not scalable. A more generic solution is to build entity linking models that work with arbitrary KBs. We follow this latter approach and build entity linking models that link to target KBs that have not been observed during training.<sup>1</sup> Our solution builds on recent models for zero-shot entity linking (Wu et al., 2020; Logeswaran et al., 2019). However, these models assume the same, simple KB schema during training and inference. We generalize these models to handle different KBs during training and inference, containing entities represented with an arbitrary set of attribute-value pairs.
This generalization relies on two key ideas. First, we convert KB entities into strings that are consumed by the models for zero-shot linking. Central to the string representation are special tokens called attribute separators, which represent frequently occurring attributes in the training KB(s), and carry over their knowledge to unseen KBs during inference (Section 4.1). Second, we generate more flexible string representations by shuffling entity attributes before converting them to strings,
<sup>1</sup> "Unseen KBs" refers to scenarios where we neither know the entities in the KB, nor its schema.
Proceedings of the 2021 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, pages 834-844 June 6-11, 2021. ©2021 Association for Computational Linguistics
Generic EL	Zero-shot EL
(Logeswaran et al., 2019)	Linking to any DB
(Sil et al., 2012)	This work
Test entities not seen during training	✓	✓	✓
Test KB schema unknown	✓
Out-of-domain test data	✓	✓
Unrestricted Candidate Set	✓	✓	✓
Table 1: This table compares the entity linking framework in the present work with those in previous work.
and by stochastically removing attribute separators to generalize to unseen attributes (Section 4.2).
Our primary experiments are cross-KB and focus on English datasets. We train models to link to one KB during training (viz. Wikidata), and evaluate them for their ability to link to an unseen KB (viz. the TAC-KBP Knowledge Base). These experiments reveal that our model with attribute-separators and the two generalization schemes are  $12 - 14\%$  more accurate than the baseline zero-shot models. Ablation studies reveal that all components individually contribute to this improvement, but combining all of them yields the most accurate models.
Unlike previous work, our models also allow seamless mixing of multiple training datasets which link to different KBs with different schemas. We investigate the impact of training on multiple datasets in two sets of experiments involving additional training data that links to (a) a third KB that is different from our original training and testing KBs, and (b) the same KB as the test data. These experiments reveal that our models perform favorably under all conditions compared to baselines.
Conventional entity linking models are trained and evaluated on the same KB, which is typically Wikipedia, or derived from Wikipedia (Bunescu and Paşa, 2006; Ling et al., 2015). This limited scope allows models to use other sources of information to improve linking, including alias tables, frequency statistics, and rich metadata.
Beyond Conventional Entity Linking There have been several attempts to go beyond such conventional settings, e.g. by linking to KBs from diverse domains such as the biomedical sciences (Zheng et al., 2014; D'Souza and Ng, 2015) and music (Oramas et al., 2016) or even being completely domain and language independent (Wang et al., 2015; Onoe and Durrett, 2020). Lin et al. (2017) discuss approaches to link entities to a KB
that simply contains a list of names without any other information. Sil et al. (2012) use database-agnostic features to link against arbitrary databases. However, their approach still requires training data from the target KB. In contrast, this work aims to train entity linking models that do not rely on training data from the target KB, and can be trained on arbitrary KBs, and applied to a different set of KBs. Pan et al. (2015) also do unsupervised entity linking by generating rich context representations for mentions using Abstract Meaning Representations (Banarescu et al., 2013), followed by unsupervised graph inference to compare contexts. They assume a rich target KB that can be converted to a connected graph. This works for Wikipedia and adjacent resources but not for arbitrary KBs. Logeswaran et al. (2019) introduce a novel zero-shot framework to "develop entity linking systems that can generalize to unseen specialized entities". Table 1 summarizes differences between our framework and those from prior work.
Contextualized Representations for Entity Linking Models in this work are based on BERT (Devlin et al., 2019). While many studies have tried to explain the effectiveness of BERT for NLP tasks (Rogers et al., 2020), the work by Tenney et al. (2019) is most relevant as they use probing tasks to show that BERT encodes knowledge of entities. This has also been shown empirically by many works that use BERT and other contextualized models for entity linking and disambiguation (Broscheit, 2019; Shahbazi et al., 2019; Yamada et al., 2020; Févry et al., 2020; Poerner et al., 2020).
3.1 Entity Linking Setup
Entity linking consists of disambiguating entity mentions  $\mathcal{M}$  from one or more documents to a target knowledge base,  $\kappa \mathcal{B}$ , containing unique entities. We assume that each entity  $e\in \kappa \mathcal{B}$  is represented using a set of attribute-value pairs
$\{(k_i,v_i)\}_{i = 1}^n$  . The attributes  $k_{i}$  collectively form the schema of  $\mathcal{KB}$  . The disambiguation of each  $m\in \mathcal{M}$  is aided by the context  $c$  in which  $m$  appears.
Models for entity linking typically consist of two stages that balance recall and precision.
1. Candidate generation: The objective of this stage is to select  $K$  candidate entities  $\mathcal{E} \subset \mathcal{KB}$  for each mention  $m \in \mathcal{M}$ , where  $K$  is a hyperparameter and  $K << |\mathcal{KB}|$ . Typically, models for candidate generation are less complex (and hence, less precise) than those used in the following (re-ranking) stage since they handle all entities in  $\mathcal{KB}$ . Instead, the goal of these models is to produce a small but high-recall candidate list  $\mathcal{E}$ . Ergo, the success of this stage is measured using a metric such as recall@ $K$  i.e. whether the candidate list contains the correct entity.
2. Candidate Reranking: This stage ranks the candidates in  $\mathcal{E}$  by how likely they are to be the correct entity. Unlike candidate generation, models for re-ranking are typically more complex and oriented towards generating a high-precision ranked list since the objective of this stage is to identify the most likely entity for each mention. This stage is evaluated using precision@1 (or accuracy) i.e. whether the highest ranked entity is the correct entity.
In traditional entity linking, the training mentions  $\mathcal{M}_{train}$  and test mentions  $\mathcal{M}_{test}$  both link to the same KB. Even in the zero-shot settings of Logeswaran et al. (2019), while the training and target domains and KBs are mutually exclusive, the schema of the KB is constant and known. On the contrary, our goal is to link test mentions  $\mathcal{M}_{test}$  to a knowledge base  $\kappa \mathcal{B}_{test}$  which is not known during training. The objective is to train models on mentions  $\mathcal{M}_{train}$  that link to  $\kappa \mathcal{B}_{train}$  and directly use these models to link  $\mathcal{M}_{test}$  to  $\kappa \mathcal{B}_{test}$ .
3.2 Zero-shot Entity Linking
The starting point (and baselines) for our work are the state-of-the-art models for zero-shot entity linking, which we briefly describe here (Wu et al., 2020; Logeswaran et al., 2019).<sup>2</sup>
Candidate Generation Our baseline candidate generation approach relies on similarities between mentions and candidates in a vector space to identify the candidates for each mention (Wu et al., 2020) using two BERT models. The first BERT model encodes a mention  $m$  along with its context  $c$  into a vector representation  $\mathbf{v}_m$ .  $\mathbf{v}_m$  is obtained from the pooled representation captured by the [CLS] token used in BERT models to indicate the start of a sequence. In this encoder, a binary (0/1) indicator vector is used to identify the mention span. The embeddings for this indicator vector (indicator embeddings) are added to the token embeddings of the mention as in Logeswaran et al. (2019). The second unmodified BERT model (i.e. not containing the indicator embeddings as in the mention encoder) independently encodes each  $e \in \mathcal{KB}$  into vectors. The candidates  $\mathcal{E}$  for a mention are the  $K$  entities whose representations are most similar to  $\mathbf{v}_m$ . Both BERT models are fine-tuned jointly using a cross-entropy loss to maximize the similarity between a mention and its corresponding correct entity, when compared to other random entities.
Candidate Re-ranking The candidate re-ranking approach uses a BERT-based cross-attention encoder to jointly encode a mention and its context along with each candidate from  $\mathcal{E}$  (Logeswaran et al., 2019). Specifically, the mention  $m$  is concatenated with its context on the left  $(c_l)$ , its context on the right  $(c_r)$ , and a single candidate entity  $e \in \mathcal{E}$ . An [SEP] token, which is used in BERT to separate inputs from different segments, is used here to separate the mention in context, from the candidate. This concatenated string is encoded using  $\mathrm{BERT}^3$  to obtain,  $\mathbf{h}_{m,e}$  a representation for this mention/candidate pair (from the [CLS] token). Given a candidate list  $\mathcal{E}$  of size  $K$  generated in the previous stage,  $K$  scores are generated for each mention, which are subsequently scored using a dot-product with a learned weight vector  $(\mathbf{w})$ . Thus,
$$
\begin{array}{l} \mathbf {h} _ {m, e} = \operatorname {B E R T} ([ \mathrm {C L S} ] c _ {l} m c _ {r} [ \mathrm {S E P} ] e [ \mathrm {S E P} ]), \\ s c o r e _ {m, e} = \mathbf {w} ^ {T} \mathbf {h} _ {m, e}. \\ \end{array}
$$
The candidate with the highest score is chosen as the correct entity, i.e.
$$
e^{*} = \operatorname *{arg  max}_{i = 1}^{K}score_{m,e_{i}}.
$$
3This BERT model also contains the indicator embeddings as described in candidate generation.
2We re-implemented these models and verified them by comparing results with those in the original papers.
4 Linking to Unseen Knowledge Bases
The models in Section 3 were designed to operate in settings where the entities in the target KB were only represented using a textual description. For example, the entity Douglas Adams would be represented in such a database using a description as follows: "Douglas Adams was an English author, screenwriter, essayist, humorist, satirist and dramatist. He was the author of The Hitchhiker's Guide to the Galaxy."
However, linking to unseen KBs requires handling entities with an arbitrary number and type of attributes. The same entity (Douglas Adams) can be represented in a different KB using attributes such as "name", "place of birth", etc. (top of Figure 1). This raises the question of whether such models, that harness the power of pre-trained language models, generalize to linking mentions to unseen KBs, including those without such textual descriptions. This section presents multiple ideas to this end.
4.1 Representing Arbitrary Entities using Attribute Separators
One way of using these models for linking against arbitrary KBs is by defining an attribute-to-text function  $f$ , that maps arbitrary entities with any set of attributes  $\{k_i, v_i\}_{i=1}^n$  to a string representation  $e$  that can be consumed by BERT, i.e.
$$
e = f (\{k _ {i}, v _ {i} \} _ {i = 1} ^ {n}).
$$
If all entities in the KB are represented using such string representations, then the models described in Section 3 can directly be used for arbitrary schemas. This leads to the question: how can we generate string representations for entities from arbitrary KBs such that they can be used for BERT-based models? Alternatively, what form can  $f$  take?
A simple answer to this question is concatenation of the values  $v_{i}$ , given by
$$
f (\{k _ {i}, v _ {i} \} _ {i = 1} ^ {n}) = v _ {1} v _ {2} \ldots v _ {n}.
$$
We can improve on this by adding some structure to this representation by teaching our model that the  $v_{i}$  belong to different segments. As in the baseline candidate re-ranking model, we do this by separating them with [SEP] tokens. We call this [SEP]-separation. This approach is also used by Logeswaran et al. (2019) and Mulang' et al. (2020)
Figure 1: Shown here are three ways of representing an entity with arbitrary attribute-values (Section 4.1). Concatenation simply concatenates all values, [SEP]-separation separates values using [SEP] tokens, and attribute separation introduces special tokens based on frequently occurring attributes (which in this example are "name" and "occupation").
to separate the entity attributes in their respective KBs.
$$
f (\{k _ {i}, v _ {i} \} _ {i = 1} ^ {n}) = [ \mathrm {S E P} ] v _ {1} [ \mathrm {S E P} ] v _ {2} \dots [ \mathrm {S E P} ] v _ {n}
$$
The above two definitions of  $f$  use the values  $v_{i}$ , but not the attributes  $k_{i}$ , which also contain meaningful information. For example, if an entity seen during inference has a capital attribute with the value "New Delhi", seeing the capital attribute allows us to infer that the target entity is likely to be a place, rather than a person, especially if we have seen the capital attribute during training. We capture this information using attribute separators, which are reserved tokens (in the vein of [SEP] tokens) corresponding to attributes. In this case,
$$
f (\{k _ {i}, v _ {i} \} _ {i = 1} ^ {n}) = [ K _ {1} ] v _ {1} [ K _ {2} ] v _ {2} \dots [ K _ {n} ] v _ {n}.
$$
These  $[K_i]$  tokens are not part of the default BERT vocabulary. Hence, we augment the default vocabulary with these new tokens and introduce them during training the entity linking model(s) based on the most frequent attribute values seen in the target KB of the training data, and randomly initialize their token embeddings. During inference, when faced with an unseen KB, we use attribute separators for only those attributes that have been observed during training, and use the [SEP] token for the remaining attributes.
Figure 1 illustrates the three instantiations of  $f$ . In all cases, attribute-value pairs are ordered in descending order of the frequency with which they appear in the training KB. Finally, since both
the candidate generation and candidate re-ranking models we build on use BERT, the techniques discussed here can be applied to both stages, but we only focus on re-ranking.
4.2 Regularization Schemes for Improving Generalization
Building models for entity linking against unseen KBs requires that such models do not overfit to the training data by memorizing characteristics of the training KB. This is done by using two regularization schemes that we apply on top of the candidate string generation techniques discussed in the previous section.
The first scheme, which we call attribute-OOV, prevents models from overtly relying on individual  $[K_i]$  tokens and generalize to attributes that are not seen during training. Analogous to how out-of-vocabulary tokens are commonly handled (Dyer et al., 2015, inter alia), every  $[K_i]$  token is stochastically replaced with the [SEP] token during training with probability  $p_{drop}$ . This encourages the model to encode semantics of the attributes in not only the  $[K_i]$  tokens, but also in the [SEP] token, which is used when unseen attributes are encountered during inference.
The second regularization scheme discourages the model from memorizing the order in which particular attributes occur. Under attribute-shuffle, every time an entity is encountered during training, its attribute/values are randomly shuffled before it is converted to a string representation using the techniques from Section 4.1.
Our held-out test bed is the TAC-KBP 2010 data (LDC2018T16) which consists of documents from English newswire, discussion forum and web data (Ji et al., 2010). The target KB ( $\mathcal{KB}_{test}$ ) is the TAC-KBP Reference KB and is built from English Wikipedia articles and their associated infoboxes (LDC2014T16). Our primary training and validation data is the CoNLL-YAGO dataset (Hoffart et al., 2011), which consists of documents from the CoNLL 2003 Named Entity Recognition task (Tjong Kim Sang and De Meulder, 2003) linked
Number of mentions	Size of target KB
CoNLL-YAGO (train)	18.5K	5.7M
CoNLL-YAGO (val.)	4.8K
Wikia (train)	49.3K	0.5M
Wikia (val.)	10.0K
TAC KBP 2010 (test)	1.7K	0.8M
Table 2: Number of mentions in our training, validation, and test sets, along with the number of entities in their respective KBs.
to multiple KBs. $^{6}$  To ensure that our training and target KBs are different, we use Wikidata as our training KB. $^{7}$  Specifically, we use the subset of entities from Wikidata with a Wikipedia page. We ignore all mentions without a corresponding entity in the KB, both during training and inference, leaving the task of handling such NIL entities to future work. Finally, we use the Wikia dataset (Logeswaran et al., 2019) for experiments that investigate the impact of multiple datasets (Section 5.5). $^{8}$  Table 2 describes the sizes of these various datasets along with the number of entities in their respective KBs.
While covering similar domains, Wikidata and the TAC-KBP Reference KB have different schemas. Wikidata is more structured and entities are associated with statements represented using attribute-value pairs, which are short snippets rather than full sentences. The TAC-KBP Reference KB contains both short snippets like these, along with the text of the Wikipedia article of the entity. The two KBs also differ in size, with Wikidata containing almost seven times the number of entities in TAC KBP.
Both during training and inference, we only retain the 100 most frequent attributes in the respective KBs. The attribute-separators (Section 4.1) are created corresponding to the 100 most frequent attributes in the training KB. Candidates and mentions (with context) are represented using strings of 128 sub-word tokens each, across all models.
4https://catalog.ldc.upenn.edu/ LDC2018T16 5https://catalog.ldc.upenn.edu/ LDC2014T16
$^{6}$ http://resources.mpi-inf.  
mpg.de/yago-naga/aida/download/  
aida-yago2-dataset.zip  
 $^{7}$ Retrieved from https://dumps.wikimedia.org/wikidatawiki/entities/in March, 2020.  
 $^{8}$ https://github.com/lajanugen/zeshel
5.2 Hyperparameters
All BERT models are uncased BERT-base models with 12 layers, 768 hidden units, and 12 heads with default parameters, and trained on English Wikipedia and the BookCorpus. The probability  $p_{drop}$  for attribute-OOV is set to 0.3. Both candidate generation and re-ranking models are trained using the BERT Adam optimizer (Kingma and Ba, 2015), with a linear warmup for  $10\%$  of the first epoch to a peak learning rate of  $2 \times 10^{-5}$  and a linear decay from there till the learning rate approaches zero. Candidate generation models are trained for 200 epochs with a batch size of 256. Re-ranking models are trained for 4 epochs with a batch size of 2, and operate on the top 32 candidates returned by the generation model. Hyperparameters are chosen such that models can be run on a single NVIDIA V100 Tensor Core GPU with 32 GB RAM, and are not extensively tuned. All models have the same number of parameters except the ones with attribute-separators which have 100 extra token embeddings (of size 768 each).
Candidate generation Since the focus of our experiments is on re-ranking, we use a fixed candidate generation model for all experiments that combines the architecture of Wu et al. (2020) (Section 3) with [SEP]-separation to generate candidate strings. This model also has no knowledge of the test KB and is trained only once on the CoNLL-Wikidata dataset. It achieves a recall@32 of 91.25 when evaluated on the unseen TAC-KBP 2010 data.
5.3 Research Questions
We evaluate the re-ranking model (Section 3) in several settings to answer the following questions:
1. Do the attribute-to-text functions (Section 4.1) generate useful string representations for arbitrary entities? Specifically, can these representations be used with the re-ranking model (Section 3) to link to the unseen  $\mathcal{KB}_{test}$ ?
2. Do all three key components—attribute-separators (Section 4.1), attribute-shuffling, and attribute-OOV (Section 4.2)—contribute equally to the final model?
3. Does training on more than one KB with different schemas help models in more accurately linking to  $\mathcal{KB}_{test}$ ?