Autoregressive Entity Retrieval

1 INTRODUCTION . The ability to retrieve the correct entity from large Knowledge Bases ( KBs ) given a textual input is a fundamental building block for several applications ( Ferrucci , 2012 ; Slawski , 2015 ; Yang et al. , 2018a ) . Most commercial recommendation systems , for instance , include in their pipelines components to detect and disambiguate entity mentions in open text , in order to isolate relevant concepts from non-meaningful data ( Slawski , 2015 ; Yang et al. , 2018a ) . Another example are chat-bots and question answering systems , that are often equipped with retrieval components to surface specific KB entries ( e.g. , Wikipedia articles ) to find knowledge for sustaining a conversation or answering a question ( Ferrucci , 2012 ; Chen et al. , 2017 ; Lewis et al. , 2020b ; Roller et al. , 2020 ) . Although there has been extensive previous work on entity retrieval ( e.g . Hoffart et al. , 2011 ; Piccinno & Ferragina , 2014 ; Huang et al. , 2015 ; Le & Titov , 2018 ; Logeswaran et al. , 2019 ; Broscheit , 2019 ; Wu et al. , 2020 , to name just a few ) there is a common design choice to most current solutions : entities are associated with a unique atomic label and the retrieval problem can be interpreted as multi-class classification across these labels . The match between input and label is calculated through a bi-encoder ( Wu et al. , 2020 ; Karpukhin et al. , 2020 ) : a dot product between dense vector encodings of the input and the entity ’ s meta information ( such as title and description ) . ∗ Work done during internship with Facebook AI Research . Critically , this formulation enables sub-linear search using modern maximum-inner-product-search libraries ( Johnson et al. , 2019 ) and hence supports retrieving from large entity databases . Unfortunately , the classifier approach to entity retrieval also has several shortcomings . First , unless a costly cross-encoder is used for re-ranking ( Wu et al. , 2020 ) , the dot-product can miss fine-grained interactions between input and entity meta information ( Humeau et al. , 2020 ) . Second , storing dense vectors for the whole KB requires a large memory footprint , especially in real-world scenarios ( i.e. , ∼24GB to store 1024-dimensional vectors for all of the∼6M Wikipedia pages ) , and the size linearly grows with the addition of new entities . Third , computing an exact softmax over all entities is very expensive , hence current solutions need to subsample negative data ( Logeswaran et al. , 2019 ; Karpukhin et al. , 2020 ) at training time . Tuning an appropriately hard set of negative instances can be challenging and time-consuming . Finally , existing systems can suffer from a cold-start problem since they can not represent entities about which they have not yet gathered sufficient information , in the form , for instance , of a textual description or a set of relations with the existing entities . The treatment of entity identifiers as atomic labels in a classifier ignores the fact that we often have unambiguous , highly structured and compositional entity names . Wikipedia , for instance , associates unique titles to articles,1 that may be the name of the subject or a description of its topic , as well as potential distinctive information to disambiguate 2 ( see Figure 1 for some examples ) . These entity names often interact with mention contexts in a predictable and regular fashion . For example , often entity names are identical with the mention strings that refer to them ( e.g. , Fig . 1f ) . When this is not possible , they might be composed of tokens in the context ( e.g. , Fig . 1b ) , include a type specification that can inferred ( e.g. , Fig . 1a ) , be the translation of the string mention ( e.g. , Fig . 1c ) , require ‘ normalization ’ such as referring to the correct alias of a mention ( e.g. , Fig . 1d ) , or require factual knowledge that might be stored in the parameters of a model ( e.g. , Fig . 1e ) . These observations suggest that inputs could be translated into unique entity names , word by word , instead of being classified among a huge set of options . In this paper , we propose GENRE ( for Generative ENtity REtrieval ) , the first entity retriever that exploits a sequence-to-sequence architecture to generate entity names in an autoregressive fashion conditioned on the context . Concretely , GENRE uses a transformer-based architecture , pre-trained with a language modeling objective ( i.e. , we use BART weights from Lewis et al . ( 2020a ) ) and fine-tuned to generate entity names . This architecture has been shown to retain factual knowledge 1We use entity name to refer to the corresponding Wikipedia article title throughout the rest of the paper . 2often in the form of a description in parentheses after the name . Wikipedia naming conventions are de- scribed in https : //en.wikipedia.org/wiki/Wikipedia : Article_titles . to some extent ( Petroni et al. , 2019 ) and language translation skills ( Radford et al. , 2019 ) among other things , both desirable properties for an entity retriever . Naturally , the generated output might not always be a valid entity name . To solve this problem , GENRE employs a constrained decoding strategy that forces each generated name to be in a predefined candidate set . The autoregressive formulation allows us to directly capture the aforementioned relations between context and entity name , effectively cross encoding both . Also , the memory footprint required is orders of magnitude smaller than current systems , since the parameters of a sequence-to-sequence model scale linearly with the vocabulary size , not entity count . Moreover , the exact softmax can be computed efficiently for each output token ( i.e. , all non-gold tokens are considered negative ) , thereby eliminating the need for negative data downsampling . Finally , our model never accesses any explicit meta-information about the entity beyond their title , hence new entities can be added by simply appending their unambiguous name to the candidate set ( e.g. , Fig . 1b refers to an entity added after training ) . We empirically evaluate the performance of GENRE on more than 20 datasets , spanning three families of tasks : ( i ) entity disambiguation , using popular datasets and settings ( both in and out-of– domain ) ; ( ii ) end-to-end entity linking , with the GERBIL benchmarking tool ( Röder et al. , 2018 ) , by using a novel dynamically markup-constrained decoding strategy ; ( iii ) document retrieval , with the recently proposed KILT benchmark ( Petroni et al. , 2020b ) which spans 5 different sub-tasks . Our models achieve state-of-the-art or very competitive results on nearly all datasets , often with substantial improvement ( +13.7 precision points on KILT for retrieval on average ) . Further , we show that compared with recent models , GENRE requires substantially less memory ( ∼20 times smaller footprint on average ) . Finally , we demonstrate that our model can be applied in scenarios where the only entity information available is its name . We organize the paper as follows : in Section 2 we describe our problem formulation . Then , in Section 3 we present GENRE and eventually in Section 4 we extensively evaluate our method on the aforementioned settings . We will release code and pre-processed data to reproduce our experiments . 2 ENTITY RETRIEVAL . We assume to have a collection of entities E ( e.g. , Wikipedia articles ) where each entity is an entry in a Knowledge Base ( KB ) such as Wikipedia . We want to approach the following retrieval problem : given a textual input source x ( e.g. , question ) , a model has to return the most relevant entities from E with respect to x . We assume that each e ∈ E is uniquely assigned to a textual representation ( i.e. , its name ) : a sequence of tokens y ( e.g. , Wikipedia pages are identified by their titles ) . A particular instance of this problem is Entity Disambiguation ( ED ) ( see Figure 1 for an example ) where an input x is annotated with a mention and a system has to select either its corresponding entity from E , or to predict that there is no corresponding entry in the KB . Another instance is pagelevel Document Retrieval ( DR ) where the input x is intended as a query and E as a collection of documents identified by their unique titles ( e.g. , Wikipedia articles ) . 3 METHOD . We address the retrieval problem with an sequence-to-sequence model that generates textual entity identifiers ( i.e. , entity names ) . Concretely , GENRE ranks each e ∈ E by computing a score with an autoregressive formulation : score ( e|x ) = pθ ( y|x ) = ∏N i=1 pθ ( yi|y < i , x ) where y is the set of N tokens in the identifier of e , and θ the parameters of the model . We take advantage of fine-tuning the BART ( Lewis et al. , 2020a ) pre-trained language model . We train GENRE using a standard seq2seq objective , i.e. , maximizing the output sequence likelihood with teacher forcing ( Sutskever et al. , 2011 ; 2014 ) and regularized with dropout ( Srivastava et al. , 2014 ) and label smoothing ( Szegedy et al. , 2016 ) . Concretely , we use the objective that is typically used for neural machine translation ( NMT , Wu et al. , 2016 ) , that is maximizing log pθ ( y|x ) with respect to model ’ s parameters θ which , due to the factorized formulation , can be calculated exactly . We do not need negative sampling to approximate the loss normalizer . 3.1 INFERENCE WITH CONSTRAINED BEAM SEARCH . Naturally , at test time , we could compute a score for every element in E and then sort them . Unfortunately , this might be prohibitively expensive when E is very large ( e.g. , Wikipedia has ∼6M entities ) . Hence , we exploit Beam Search ( BS , Sutskever et al. , 2014 ) , an established approximate decoding strategies to efficiently navigate the search space . Instead of explicitly scoring all entities in E , we search for the top-k entities in E decoding from our model using BS with k beams . Note that using BS implies that the time cost of our retriever does not depend on the size of E , but only on the size of the beams and the average length of entity representations as we do autoregressive generation . The average length of entity representations is tractable ( e.g. , Wikipedia titles have 6 BPE tokens on average ) and we follow standard NMT settings where k is small ( e.g. , 10 ) . Since we want to output only entities from E we can not use traditional BS while decoding . Indeed , allowing to generate any token from the vocabulary at every decoding step might lead the model to generate output strings that are not valid identifiers . Hence , we resort to Constrained BS , forcing to only decode valid entity identifiers . BS only considers one step ahead during decoding so we can only constrain the generation of a single next token conditioned on the previous ones . Thus , we define our constrain in terms of a prefix tree T ( aka trie ) ( Cormen et al. , 2009 ) where nodes are annotated with tokens from the vocabulary . For each node t ∈ T , its children indicate all the allowed continuations from the prefix defined traversing the trie from the root to t. See Figure 9 in Appendix C for an exampled of a trie . When the number of allowed outputs is tractable ( e.g. , generating a Wikipedia title among ∼6M ) the trie is relatively small it can be precomputed and stored into memory ( e.g. , constraining on Wikipedia titles using the BART tokenizer produces a trie with ∼6M leaves , ∼17M internal nodes that occupied ∼600MB of disk space ) . We employed the constraints masking the log-probabilities of the invalid tokens and not their logits ( i.e. , we do not re-normalize the probability over the vocabulary ) .3

Summary: This paper proposes to tackle the entity linking task using a sequence-to-sequence neural model, trained by producing unique entity names, in autoregressive fashion. The paper makes a case that this approach can scale better with larger entity vocabularies than previous methods with dedicated entity representations both in terms of memory as well as computation costs. The model is studied under a number of tasks including entity disambiguation, entity linking and document retrieval for question answering.