Global Entity Disambiguation with BERT (2022)

Abstract: We propose a global entity disambiguation (ED) model based on BERT. To capture global contextual information for ED, our model treats not only words but also entities as input tokens, and solves the task by sequentially resolving mentions to their referent entities and using resolved entities as inputs at each step. We train the model using a large entity-annotated corpus obtained from Wikipedia. We achieve new state-of-the-art results on five standard ED datasets: AIDA-CoNLL, MSNBC, AQUAINT, ACE2004, and WNED-WIKI. The source code and model checkpoint are available at https://github.com/studio-ousia/luke .

## introduction entity disambiguation (ed) refers to the task of assigning mentions in a document to corresponding entities in a knowledge base (kb). this task is challenging because of the ambiguity between mentions (e.g., "world cup") and the entities they refer to (e.g., fifa world cup or rugby world cup). ed models typically rely on local contextual information based on words that co-occur with the mention and global contextual information based on the entity-based coherence of the disambiguation decisions. a key to improve the performance of ed is to effectively combine both local and global contextual information @xcite . in this study, we propose a global ed model based on bert @xcite . our model treats words and entities in the document as input tokens, and is trained by predicting randomly masked entities in a large entity-annotated corpus obtained from wikipedia. this training enables the model to learn how to disambiguate masked entities based on words and non-masked entities. at the inference time, our model disambiguates * work done at riken. ## related work transformer-based ed. several recent studies have proposed ed models based on transformer @xcite trained with a large entity-annotated corpus obtained from wikipedia @xcite @xcite @xcite . @xcite trained an ed model based on bert by classifying each word in the document to the corresponding entity. similarly, févry et al. @xcite addressed ed using bert by classifying mention spans to the corresponding entities. ling a in c word d 4 in a the c word d 5 the a world c word d 6 world a cup c word d 7 cup a [sep] c word d 8 [sep] blionel_messi c entity (e 1 + e 2 ) / 2 lionel messi b[mask] c entity (e 6 + e 7 ) / 2 [mask] words entities input position emb. token type emb. token emb. + + + + + + + + + + + + + + + + + + + + + + transformer figure 2: the input representation of our model with the text "messi played in the world cup" with mentions "messi" and "world cup". the entity corresponding to the mention "world cup" is replaced by the [mask] token. et al. (2020) trained bert by predicting entities using the document-level representation. cao et al. (2021) addressed ed by training bart @xcite to generate referent entity titles of target mentions in an autoregressive manner. @xcite formulated ed as a text extraction problem; they fed the document and candidate entity titles to bart and longformer @xcite and disambiguated a mention in the document by extracting the referent entity title of the mention. however, unlike our model, these models addressed the task based only on local contextual information. treating entities as inputs of transformer. recent studies @xcite @xcite have proposed transformerbased models that treat entities as input tokens to enrich their expressiveness using additional information contained in the entity embeddings. however, these models were designed to solve general nlp tasks and not tested on ed. we treat entities as input tokens to capture the global context that is shown to be highly effective for ed. ## model given a document with n mentions, each of which has k entity candidates, our model solves ed by selecting a correct referent entity from the entity candidates for each mention. ## experiments our experimental setup follows @xcite . in particular, we test the proposed ed models using six standard datasets: aida-conll (conll) @xcite , msnbc, aquaint, ace2004, wned-cweb (cweb), and wned-wiki (wiki) @xcite . we consider only the mentions that refer to valid entities in wikipedia. for all datasets, we use the kb+yago entity candidates and their associated p(e|m) @xcite , and use the top 30 candidates based on p(e|m). for the conll dataset, we also test the performance using pprforned entity candidates @xcite . we report the in-kb accuracy for the conll dataset and the micro f1 score (averaged per mention) for the other datasets. further details of the datasets are provided in appendix c. furthermore, we optionally fine-tune the model by maximizing the log likelihood of the ed predictions (ŷ ed ) using the training set of the conll dataset with the kb+yago candidates. we mask 90% of the mentions and fix the entity token embeddings (b and the model is trained for two epochs using adamw. additional details are provided in appendix b. ## conclusion and future work we propose a new global ed model based on bert. our extensive experiments on a wide range of ed datasets demonstrate its effectiveness. one limitation of our model is that, similar to existing ed models, our model cannot handle entities that are not included in the vocabulary. in our future work, we will investigate the method to compute the embeddings of such entities using a post-hoc training with an extended vocabulary @xcite .