Empirical Evaluation of Pretraining Strategies for Supervised Entity Linking

1 . Introduction . Traditionally , entity linking approaches have relied on knowledge bases , complicated modelling and task-specific hand-engineered features to achieve high performance . More recently , Broscheit 2019 , Ling et al . 2020 and Wu et al . 2019 show that using large-scale pretrained language models like BERT [ Devlin et al. , 2018 ] , pretraining on Wikipedia entity links , and fine-tuning on a specific entity linking corpus leads to state-of-the-art performance without relying on such features . However , Ling et al . 2020 focused mainly on constructing general-purpose entity representations , Wu et al . 2019 on building strong zero-shot entity linking systems , and Broscheit 2019 on end-to-end linking , so that the limits of pretraining for entity disambiguation have not been fully explored . In this paper we present a thorough study of pretraining strategies for supervised entity linking . We establish new upper bounds for performance on the widely studied CoNLL and TAC-KBP 2010 entity linking tasks . We also show that our pretraining approach yields a very competitive entity linking system without any further domain specific tuning . We present a detailed analysis of significant design choices including the choice of negative candidates used during training , and the document context encoded for each mention . We find that the optimal choice of negative candidates is dependent on whether or not the final linking system has access to an alias table . For a system that will use an alias table at inference , it is helpful to pretrain the models with lexically similar candidates . However , when no alias table is used for the downstream task , ensuring candidates are random improves the model ability to distinguish the right entity among all possible entities . ∗ . Denotes equal contribution . † . Work conducted during Google AI Residency . In our studies we found two surprising results : ( 1 ) it is possible to achieve optimal entity linking results with a four layer transformer , which is one third of the size of BERT-base , and ( 2 ) given the abundance of supervision from Wikipedia links , we did not get any gains in performance from training with an auxiliary language modeling loss . However , we did find that the input perturbations introduced by [ Devlin et al. , 2018 ] themselves increase the quality of our pretraining approach and we present an analysis of how these perturbations add robustness to the model . Finally , to demonstrate the generality of our model , we present results on the end-toend entity linking task in which both mention location and identity are predicted . On this task , our model outperforms all but the tailored methods introduced by Kolitsas et al . 2018 , Broscheit 2019 . We argue that our model is more practical , and easier to integrate , than heavily engineered existing approaches , and we believe that downstream tasks such as information extraction and question answering can benefit from this robust standalone end-to-end entity linker ( e.g . [ Févry et al. , 2020 ] ) . 2 . Related Work . Early entity linking systems [ Bunescu and Pasca , 2006 , Mihalcea and Csomai , 2007 ] focused on matching the context of the mention with that of the entity page . In addition to context features , systems have relied on P ( e|m ) , the prior probability that mention m refers to entity e , computed from Wikipedia mention counts . The set of entities in a document should be globally coherent , and several approaches have introduced sophisticated global disambiguation methods [ Globerson et al. , 2016 , Cheng and Roth , 2013 , Sil and Yates , 2013 ] that consider all mentions in a document to make predictions . In contrast , we do not model document-level disambiguation explicitly . However , our long context windows contain several mentions which should allow such disambiguation . Other approaches have also sought to incorporate types and knowledge base information in their modelling , such as Radhakrishnan et al . 2018 , Raiman and Raiman 2018 . We pretrain our model by learning distributed representation of entities directly from Wikipedia text , similarly to Yamada et al . 2016 , 2017 , Gillick et al . 2019 , Ling et al . 2020 . Unlike Wu et al . 2019 , Gupta et al . 2017 , our embeddings are learned directly rather than generated through entity descriptions . In contrast with Yamada et al . 2016 , 2017 , we do not use additional features , such as prior probabilities or string match features . Our method is therefore most similar to Ling et al . 2020 and Broscheit 2019 who also use a transformer . We differ from Ling et al . 2020 by simultaneously considering all mentions and entities in a context and , in contrast to both , only use a four layer , randomly-initialized transformer instead of twelve layers initialized from large scale language modelling pretraining . In addition , we experiment with end-to-end entity linking [ Sil and Yates , 2013 , Luo et al. , 2015 , Kolitsas et al. , 2018 ] , where instead of predicting the entity for gold spans , a system must both predict the span and its label . A closely related task is multilingual entity linking . Approaches have used multilingual embeddings to link text in several languages [ Sil et al. , 2018 , Tsai and Roth , 2016 ] . Zero-shot entity linking [ Logeswaran et al. , 2019 , Wu et al. , 2019 ] is another relevant task . In that setting , entities predicted at test time are not seen in training . Instead , the system relies on the entity name and description . 3 . Model . 3.1 Task Definition . Let E = { e1 . . . eN } be a predefined set of entities , and let V = { [ mask ] , w1 . . . wM } be a vocabulary of words . A context x = [ x0 . . . xt ] is a sequence of words xi ∈ V. A span s = ( sstart , send ) , is a tuple with 0 ≤ sstart , send < t which defines a contiguous sequence of tokens in a given context . A mention label l = ( sk , ek ) consists of a span si and an entity label , ei ∈ E ∪ ∅ . We use l to denote a set of such mention labels . The NULL-symbol ∅ indicates a span that is labeled as a mention , but without an entity linking label . Our training data , D = { ( x0 , l0 ) . . . ( xN , lN ) } , is a corpus of contexts , each paired with a set of mention labels , one for each mention in the context . Given an input context xi , our goal is to predict the set of entity mentions li . In Entity Disambiguation , we are given the set of spans , and predict the entity linked by each span . In End-to-End Entity Linking , we must predict both the set of mention spans , and their linked entities . 3.2 Contextual Language Representation . Our model is built using the now-standard Transformer-based architecture [ Vaswani et al. , 2017 ] . The model computes a matrix representation Ĥ ∈ Rt×d of a text sequence through successive application of a Transformer block to the output of the previous layer : Hi = TransformerBlock ( Hi−1 ) = MLP ( MultiHeadAttention ( Hi−1 , Hi−1 , Hi−1 ) ) H0 is a sequence of context-independent token embeddings and Ĥ = Hn , where n is the number of Transformer layers . 3.3 Entity Disambiguation . Each entity e ∈ E is mapped directly onto a dedicated vector in Rd via a |E|×d dimensional embedding matrix . In our experiments , we have a distinct embedding for every concept that has an English Wikipedia page , resulting in approximately 5.7m entity embeddings . In order to perform entity linking for a particular span with word-piece token indices ( i , j ) , we ( following Lee et al . 2016 ) first obtain a representation of the span by concatenating the representation at the span start and end , and pass this through a multi-layer perceptron which projects the span representation into the same space as the entity embeddings . ŝsi = MLP ( [ Hn , sstart , Hn , send ] ) ( 1 ) Our model scores each span-entity pair by taking the dot-product between the projected span representation and the embedding of ec . Thus , the conditional probability that the span si refers to entity ec is defined as : P ( ec|si ) = exp ( ŝsi · ec ) ∑ c′∈E exp ( ŝsi · ec′ ) ( 2 ) In practice , this is expensive to compute for large |E| . Therefore , for every l we select a set Cl of k candidates , which contains the entity labels for all l ∈ l as well as a set of negative candidates . We do not have an entity linking loss on mentions that do not have a label . Therefore , our per-example entity linking loss is : llinking ( l ) = ∑ li∈l exp ( ŝsi · ei ) 1ei ! =∅∑ c∈Cl exp ( ŝsi · ec ) ( 3 ) We will discuss further how we select Cl in Section 4.1 . 3.4 Mention Detection . For many entity linking tasks , the target spans are provided . In order to be able to do endto-end entity linking , we additionally train our model to predict mentions , independently of entity linking . One way to do this would be to score every possible span-entity pair , and either use a score threshold to filter spans where no entity link achieves a sufficiently high score , or to additionally score a special NULL-link embedding . However , enumerating all spans for the long contexts we use in our model would be prohibitively expensive . We take the approach of encoding mentions as a BIO sequence , and train an MLP on the context representation to predict this sequence with a standard cross-entropy loss . Our final loss sums the mention detection loss and the linking loss . 4 . Experimental Setup . 4.1 Wikipedia Pretraining . We build a training corpus of contexts paired with entity mention labels from the 2019-04- 14 dump of English Wikipedia . We first divide each article into chunks of 1000 unicode characters , resulting in a corpus of over 17.5 million contexts with over 17 million entity mentions covering over 5.7 million entities . These are processed with the BERT tokenizer , limited to 256 word-piece tokens . In addition to the Wikipedia links , we annotate each sentence with unlinked mention spans using a state-of-the-art named entity recognizer . These are used as additional signal for our mention detection component . Entity Candidates Selection Training the model with a full softmax over all 5.7 million entities for every mention is computationally expensive . A common solution is to use a noise contrastive loss [ Gutmann and Hyvärinen , 2012 , Mnih and Kavukcuoglu , 2013 ] and sample candidates according to their relative frequency , as in Ling et al . [ 2020 ] . In this work , we experiment with other approaches to candidate generation that might provide better negatives in training . In addition to negatives selected uniformly at random from the entire entity vocabulary , we define two types of hard negatives : 1 . Page candidates , which is the set of all entities linked to in the article from which the given context was taken . This is meant to capture semantically related concepts . 2 . Phrase table candidates , the set of lexically related entities for each mention candidate , obtained from the Phrase Table provided by SLING [ Ringgaard et al. , 2017 ] . Throughout the paper , we will use |Cl| = 768 . In our base setup , we use up to 256 page candidates , and 384 phrase table candidates , equally divided between each mention in the example . Any remaining room in the set of 768 is filled with random candidates sampled uniformly from the entity vocabulary ( meaning a minimum of 128 random candidates per example ) . We will study the impact of different candidate selection methods in Section 6.2 . In addition to those candidates , for every example in a batch , we use the candidates of other examples as additional negatives . Input Noising We also add noise to the input data . We apply the same noise function as is used in Devlin et al . 2018 : 15 % of the tokens are chosen to be modified . 80 % of those tokens are changed to the [ mask ] token , 10 % are changed to a random token and 10 % are left unmodified . Pretraining hyperparameters We use Adam [ Kingma and Ba , 2014 ] with a learning rate of 1e-4 to optimize our model . We use a linear warmup schedule for the first 10 % of training , decay the learning rate afterwards and use gradient clipping with a norm of 1.0 . We train from scratch for up to a million steps and use a large batch size of 8192 for pretraining . We follow BERT [ Devlin et al. , 2018 ] base for many of our model parameters , though we do not use large-scale language-modeling pretraining and only use four layers , as we did not find more to layers to further improve performance . We use the same wordpiece vocabulary as the lowercase version of BERT . We use entity embeddings of size 256 unless mentioned otherwise . We weight both the entity disambiguation loss and the mention detection loss to 1 . We use a context window of 256 tokens .

Summary: This paper investigates the use of a simple architecture for entity disambiguation: encode the mention and its context with BERT, use an MLP over the mention's fenceposts to compute an embedding, then compare that embedding with embeddings of entity candidates and take the one with the highest dot product. Notably, it uses a transformer pre-trained on Wikipedia to do entity resolution, but does *not* use the BERT model or its pre-trained parameters directly.  The paper deals with several design decisions along the way: how to pre-train this model on Wikipedia, how to generate candidates at train and test time, whether or not to mask the input as in BERT, and other hyperparameters. Results show state-of-the-art performance on CoNLL (with a good candidate set) and TAC-KBP, as well as good performance on end-to-end entity linking (detecting and linking mentions).