Source-linked AI summary
Autoregressive Entity Retrieval
Nicola De Cao, Gautier Izacard, Sebastian Riedel, Fabio Petroni
TL;DR
Entity retrieval systems commonly classify among atomic labels, which can miss fine-grained context–entity interactions, require large dense indexes, and use negative sampling. GENRE instead autoregressively generates unique entity names with constrained decoding, achieving strong results across entity tasks while substantially reducing memory use. The paper also shows that new entities can be added by specifying their names.
Problem
Existing classifier-based entity retrieval can miss fine-grained interactions, require large memory for dense entity representations, and need subsampled negative data.
Method
GENRE uses a sequence-to-sequence model to generate unique entity names autoregressively from context, with constrained decoding restricting outputs to valid candidates.
Results
GENRE achieves state-of-the-art or very competitive results on more than 20 datasets spanning entity disambiguation, entity linking, and document retrieval.
Takeaways & Limitations
GENRE reduces the entity-index memory footprint and permits adding new entities by appending their unambiguous names to the candidate set.
Takeaways & Limitations
The formulation assumes each entity is uniquely assigned a textual representation, such as a name or title.
Abstract
from arXiv · showhide
Entities are at the center of how we represent and aggregate knowledge. For instance, Encyclopedias such as Wikipedia are structured by entities (e.g., one per Wikipedia article). The ability to retrieve such entities given a query is fundamental for knowledge-intensive tasks such as entity linking and open-domain question answering. Current approaches can be understood as classifiers among atomic labels, one for each entity. Their weight vectors are dense entity representations produced by encoding entity meta information such as their descriptions. This approach has several shortcomings: (i) context and entity affinity is mainly captured through a vector dot product, potentially missing fine-grained interactions; (ii) a large memory footprint is needed to store dense representations when considering large entity sets; (iii) an appropriately hard set of negative data has to be subsampled at training time. In this work, we propose GENRE, the first system that retrieves entities by generating their unique names, left to right, token-by-token in an autoregressive fashion. This mitigates the aforementioned technical issues since: (i) the autoregressive formulation directly captures relations between context and entity name, effectively cross encoding both; (ii) the memory footprint is greatly reduced because the parameters of our encoder-decoder architecture scale with vocabulary size, not entity count; (iii) the softmax loss is computed without subsampling negative data. We experiment with more than 20 datasets on entity disambiguation, end-to-end entity linking and document retrieval tasks, achieving new state-of-the-art or very competitive results while using a tiny fraction of the memory footprint of competing systems. Finally, we demonstrate that new entities can be added by simply specifying their names. Code and pre-trained models at https://github.com/facebookresearch/GENRE.
1 INTRODUCTION
GENRE reframes entity retrieval as autoregressive generation of structured entity names rather than classification over atomic labels. This design addresses interaction modeling, memory, and negative-sampling limitations, and performs strongly across broad evaluations.
- Entity retrieval supports recommendation, chatbots, question answering, and other applications that need relevant Knowledge Base entries.
- Classifier-based retrieval uses a bi-encoder dot product between input and dense entity representations, enabling sub-linear search but limiting fine-grained interaction modeling.
- Structured entity names encode compositional relations with contexts, motivating word-by-word generation instead of classification among a huge option set.
- GENRE uses a transformer sequence-to-sequence model to generate entity names autoregressively from context, with constrained decoding enforcing valid candidates.
- Across more than 20 datasets, GENRE achieves state-of-the-art or competitive results, including +13.7 precision points on average for KILT retrieval and substantially lower memory use.The reported average memory footprint is approximately 20 times smaller than recent models.
2 ENTITY RETRIEVAL
The paper formulates entity retrieval as selecting relevant entities from a Knowledge Base for a textual input. Each entity is represented by a unique textual name, supporting disambiguation and document retrieval.
- The task retrieves the most relevant entities from a collection E given textual input x.
- Each entity is assigned a unique textual representation, such as a Wikipedia article title, modeled as a token sequence.
- Entity disambiguation selects the corresponding Knowledge Base entity for a mention or predicts that no corresponding entry exists.
- Page-level document retrieval treats the input as a query and entities as documents identified by unique titles.
3 METHOD
GENRE scores candidate entities by autoregressively generating their names, trains with exact sequence likelihood, and uses constrained beam search to decode only valid identifiers. Dynamic constraints extend the framework to end-to-end entity linking.
- GENRE ranks each entity by the autoregressive probability of its name conditioned on the input and previously generated tokens.
- Training maximizes exact sequence likelihood with teacher forcing, dropout, and label smoothing, avoiding negative sampling because the factorized objective is computed exactly.
- INFERENCE WITH CONSTRAINED BEAM SEARCH: Beam search efficiently finds top-k entities without explicitly scoring every candidate in large collections.
- INFERENCE WITH CONSTRAINED BEAM SEARCH: Constrained beam search uses a prefix tree whose children specify valid next-token continuations, preventing invalid entity identifiers.
- End-to-end entity linking generates the source with marked mention spans and entity identifiers, using dynamic constraints because annotated outputs are exponentially large.
4 EXPERIMENTS
GENRE is evaluated across more than 20 datasets covering entity disambiguation, entity linking, and page-level document retrieval. It achieves competitive or leading results while substantially reducing retrieval memory, with additional evidence for cold-start and structured-name behavior.
- Experimental settings: GENRE is evaluated on more than 20 datasets across entity disambiguation, end-to-end entity linking, and page-level document retrieval.The experiments include KILT document-retrieval tasks and English-language settings.
- Overall results: GENRE achieves very competitive results across all three settings and is the best-performing system on average.The paper reports state-of-the-art or highly competitive performance across nearly all datasets.
- Entity linking: +13 F1 points on Derczynski and +4.7 on KORE50 are reported for out-of-domain entity linking, although OKE15 and OKE16 perform poorly.The weaker datasets contain coreference annotations that GENRE was not specifically trained to handle.
- Document retrieval: +13.7 R-precision points over the best baseline is achieved on average for page-level document retrieval, with GENRE best across five KILT tasks except Natural Questions.The comparison includes DPR, RAG, BLINK+flair, and other public-leaderboard systems.
- Memory footprint: 14 times less memory than BLINK and 34 times less than DPR is required by GENRE for the entity index.GENRE stores a prefix tree of entity names, using millions rather than billions of model/index parameters.
5 RELATED WORKS
Related work applies sequence-to-sequence modeling and constrained generation to structured NLP outputs, while autoregressive document reranking requires a limited candidate list for computational feasibility.
- Sequence-to-sequence formulations have been explored for semantic parsing, semantic role labeling, discourse representation parsing, and structured language generation.
- Constrained generation methods use copying or structural constraints to keep outputs lexically close to inputs or structurally well formed.
- The cited related-work passages include a note that pre-training and fine-tuning use 2019 data dumps.
- Autoregressive document reranking predicts relevance labels but requires a limited candidate-document list, unlike entity retrieval over the full entity space.
6 CONCLUSIONS
GENRE retrieves entities by autoregressively generating entity names, exploiting their compositional structure and predictable context interactions. It achieves strong performance across entity retrieval tasks with reduced memory and supports adding new entities by appending their names.
- GENRE generates entity names autoregressively, using their compositional structure and predictable interaction with context.
- The approach cross-encodes mention context and entity candidates, computes an exact softmax without negative subsampling, and reduces memory usage by a factor of twenty on average.
- GENRE achieves state-of-the-art performance across entity disambiguation, end-to-end entity linking, and page-level document retrieval datasets.
- New entities can be considered by appending their unambiguous names to the candidate set.
A EXPERIMENTAL DETAILS
GENRE is trained as a sequence-to-sequence model with Adam, linear learning-rate scheduling, and categorical cross-entropy with label smoothing.
- GENRE uses Adam with a learning rate of 3 · 10^-5, followed by linear warm-up for 500 steps and linear decay.
- Training optimizes sequence-to-sequence categorical cross-entropy loss with 0.1 label smoothing.
A.1 NAMED ENTITY DISAMBIGUATION
The experiments define entity disambiguation, entity linking, and document retrieval settings, with task-specific training data, candidate handling, and constrained beam-search inference.
- NAMED ENTITY DISAMBIGUATION: Entity disambiguation assigns each mention either a knowledge-base entity or NIL from a restricted candidate set.
- NAMED ENTITY DISAMBIGUATION: Constrained beam search is used at inference, with task-specific beam counts, decoding limits, context truncation or chunking, and length-normalized log-probabilities.
- NAMED ENTITY DISAMBIGUATION: Entity linking returns mention-entity tuples and uses Wikipedia pre-training enriched with string-matching co-reference examples before AIDA fine-tuning.
- NAMED ENTITY DISAMBIGUATION: Document retrieval ranks Wikipedia pages by relevance to a query and trains GENRE jointly on all KILT data.
- NAMED ENTITY DISAMBIGUATION: The Wikipedia training corpus is based on the 2019/08/01 dump pre-processed for KILT.
B.1 NAMED ENTITY DISAMBIGUATION
This section evaluates GENRE on WikilinksNED unseen mentions and reports additional AIDA results using established entity-disambiguation metrics.
- The section combines benchmark evaluation on WikilinksNED with literature comparisons on AIDA.
- Table 6 evaluates GENRE on WikilinksNED Unseen-Mentions data using accuracy, or precision at 1.The evaluation splits test examples by seen versus unseen entities and by three mention–gold-entity matching types.
- Table 7 reports additional AIDA results using Micro InKB F1 on the test sets.
B.2 DOCUMENT RETRIEVAL
GENRE’s retrieval ablation tests whether meaningful entity names matter and examines performance variation with title length and entity connectivity. Numerical identifiers substantially weaken performance, while accuracy is lower for long or rarely linked titles.
- Numerical entity IDs significantly lower GENRE performance compared with meaningful entity names in the retrieval ablation.The experiment tests whether entity names provide meaningful, compositional information beyond memorization.
- Unconstrained generation allows GENRE to produce entity names that are not present in the knowledge base.Constrained decoding is the corresponding alternative in the ablation description.
- 78.6% average accuracy is higher for titles shorter than 10 BPE tokens and lower for titles with at least 10 tokens.Most Wikipedia titles contain fewer than 15 BPE tokens, and the distribution mode is 5.
- Pages with no incoming Wikipedia links achieve accuracy 20% below the global average of 78.6%.Incoming-link count is used as a proxy for how often a page or entity has been observed.
C EXAMPLES
The examples show GENRE generating ranked entity titles for disambiguation and retrieval, markup strings for end-to-end linking, and constrained title sequences through a prefix tree.
- Named entity disambiguation: For named entity disambiguation, GENRE marks a mention with [START ENT] and [END ENT] and generates ranked entity candidates with log-likelihoods.The examples include predictions for Metropolis (comics), Bronzino, and Tool (band).
- Named entity disambiguation: In open-domain question answering, a painter query receives Bronzino as a ranked prediction, while a band statement receives Tool (band).
- Document retrieval: For document retrieval, GENRE takes a query as input and outputs a ranked list of Wikipedia article titles with log-likelihoods.
- End-to-end entity linking: For end-to-end entity linking, GENRE outputs a markup string whose links are Wikipedia titles and whose spans encode mention start, length, and title.Spans use the format ⟨s_i, l_i, t_i⟩.
- Constrained decoding: The prefix tree constrains generation to allowed entity identifiers, sharing prefixes and terminating valid sequences with EOS.The illustrated identifiers are English language, English literature, and France.