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: The paper proposed to use autoregressive approach to solve entity-based problems. They proposed a uniform framework and showed that their model achieved the state of the art performance on 3 different types of tasks (~20 datasets). The GENRE model also significantly reduced the memory usage compared to previous models that stored a big memory table. It's also capable of linking novel entities at inference time. This paper is clearly written. The experiment results are convincing.