Source-linked AI summary
Scalable Zero-shot Entity Linking with Dense Entity Retrieval
Ledell Wu, Fabio Petroni, Martin Josifoski, Sebastian Riedel, Luke Zettlemoyer
TL;DR
Entity linking must handle millions of candidates, especially when zero-shot systems can use only short entity descriptions. The paper introduces a two-stage BERT system with dense bi-encoder retrieval and cross-encoder re-ranking, achieving state-of-the-art results while evaluating its accuracy-speed trade-off. It also shows fast large-scale retrieval and transfers cross-encoder gains to the bi-encoder through distillation.
Problem
Entity linking must efficiently rank millions of possible entities, including in zero-shot settings where short textual descriptions provide the available entity information.
Method
A bi-encoder independently embeds mention context and entity descriptions for dense retrieval, after which a cross-encoder concatenates and re-ranks the retrieved candidates.
Results
The approach achieves state-of-the-art performance across zero-shot and established evaluations, including nearly 6 absolute points on Wikia, more than 7 on WikilinksNED Unseen-Mentions, and over 30% relative error reduction on TACKBP-2010.
Takeaways & Limitations
Bi-encoder linking scales to 5.9 million candidates in 2 milliseconds, while knowledge distillation transfers much of the cross-encoder’s accuracy gain to the faster bi-encoder.
Takeaways & Limitations
The evaluation assumes in-KB linking with valid gold entities and leaves out-of-KB or nil prediction to future work.
Abstract
from arXiv · showhide
This paper introduces a conceptually simple, scalable, and highly effective BERT-based entity linking model, along with an extensive evaluation of its accuracy-speed trade-off. We present a two-stage zero-shot linking algorithm, where each entity is defined only by a short textual description. The first stage does retrieval in a dense space defined by a bi-encoder that independently embeds the mention context and the entity descriptions. Each candidate is then re-ranked with a cross-encoder, that concatenates the mention and entity text. Experiments demonstrate that this approach is state of the art on recent zero-shot benchmarks (6 point absolute gains) and also on more established non-zero-shot evaluations (e.g. TACKBP-2010), despite its relative simplicity (e.g. no explicit entity embeddings or manually engineered mention tables). We also show that bi-encoder linking is very fast with nearest neighbour search (e.g. linking with 5.9 million candidates in 2 milliseconds), and that much of the accuracy gain from the more expensive cross-encoder can be transferred to the bi-encoder via knowledge distillation. Our code and models are available at https://github.com/facebookresearch/BLINK.
1 Introduction
The paper addresses scalable zero-shot entity linking, where millions of entities must be considered using only short textual descriptions. It proposes a two-stage BERT approach and evaluates its accuracy-speed trade-off.
- Motivation: Millions of possible entities make scale a central challenge for entity linking, motivating efficient candidate filtering and ranking.Existing methods use external cues such as mention tables, link popularity, and entity categories.
- Method: The model first retrieves candidates with a bi-encoder that independently embeds mention context and entity descriptions in a shared dense space.A cross-encoder then concatenates mention and entity text to re-rank the retrieved candidates.
- Results: Over 30% relative error reduction on TACKBP-2010 establishes a new state-of-the-art result for the two-stage approach.The result is achieved using provided text descriptions despite prior methods using additional cues.
- Results: The method improves zero-shot benchmarks by nearly 6 absolute points on Wikia and more than 7 absolute points on WikilinksNED Unseen-Mentions.These gains are reported alongside the TACKBP-2010 result.
- Efficiency: The system evaluates accuracy-speed trade-offs at Wikipedia scale and supports approximate nearest-neighbor linking over 5.9 million candidates in 2 milliseconds.Knowledge distillation transfers much of the cross-encoder’s accuracy gain to the faster bi-encoder.
2 Related Work
Related work frames entity linking as candidate generation followed by ranking, with prior systems using lexical, frequency-based, neural, and external-knowledge signals. This paper combines dense retrieval with pre-trained encoders for scalable zero-shot linking.
- Prior entity linking: Entity linking systems typically separate candidate generation from candidate ranking, using frequency information, alias tables, TF-IDF, and neural models.Prior ranking models represent context words, spans, and entities with neural networks.
- Zero-shot linking: Recent zero-shot work used cross-encoders for entity ranking but relied on traditional information-retrieval methods for candidate generation.That work also did not evaluate large-scale benchmarks such as TACKBP.
- Scope: The authors leave exploration of end-to-end linking to future work, although they state that their techniques should apply to that setting.This scope boundary distinguishes the presented evaluation from end-to-end linking.
- Dense retrieval: Dense embeddings had been shown useful for candidate generation, but earlier work omitted pre-training and used external category labels tied to Wikipedia entities.The paper positions its approach as generalizing these dense-retrieval and zero-shot lines of work.
- Pre-trained architectures: The paper adapts bi-, poly-, and cross-encoder ideas from pre-trained transformer research on dialogue sentence selection to entity linking.It additionally demonstrates that a bi-encoder can serve as a strong retrieval model.
3 Definition and Task Formulation
The paper formulates entity linking as predicting mention-entity pairs from a document and a knowledge base, and separately defines zero-shot evaluation with disjoint training and test entities.
- Entity Linking: Given a document and its mentions, an entity-linking model outputs mention-entity pairs whose entities belong to a knowledge base such as Wikipedia.The formulation assumes entity titles and descriptions are available and evaluates mentions with valid in-KB gold entities.
- Entity Linking: The task assumes each mention has a valid gold entity in the knowledge base, making the setting in-KB evaluation.Out-of-KB or nil prediction is left to future work.
- Zero-shot Entity Linking: Zero-shot entity linking separates the training and test knowledge bases so that E_train ∩ E_test = ∅ and test entities are unseen during training.Documents, mentions, and entity dictionaries are also separated between training and test.
4 Methodology
The method combines independently encoded mention and entity representations for scalable retrieval with cross-encoder re-ranking for more detailed candidate scoring. It also uses batched negatives, hard negatives, cached representations, and knowledge distillation to balance training efficiency, inference speed, and accuracy.
- Bi-encoder: The bi-encoder independently encodes mention context and entity text into vectors, enabling candidate retrieval through dense representations.Mention context includes the surrounding text and marked mention, while entity representations use the title and description.
- Bi-encoder: Candidate entities are scored by the dot product of the mention and entity vectors, with representations reduced from transformer outputs.The reduction function uses the final-layer [CLS] representation.
- Bi-encoder: Training maximizes the correct entity's score against randomly sampled in-batch entities and additional hard negatives from the top 10 predicted entities.The hard negatives are added to the random in-batch negatives.
- Inference: At inference, cached entity representations reduce linking to maximum-dot-product search, with exact and approximate nearest-neighbor methods supporting large-scale retrieval.The paper evaluates efficiency and accuracy trade-offs using FAISS.
- Cross-encoder: The cross-encoder concatenates mention and entity inputs in one transformer, enabling deep cross-attention before a linear layer computes candidate scores.Because of its larger memory and compute footprint, it re-ranks at most 100 bi-encoder candidates rather than performing retrieval.
- Knowledge Distillation: Knowledge distillation trains a bi-encoder using cross-encoder logits softened by temperature alongside the ground-truth student loss.The combined objective mixes distillation and student losses through a coefficient α.
5 Experiments
Experiments evaluate the approach on zero-shot entity linking, TACKBP-2010, and WikilinksNED Unseen-Mentions, covering retrieval, ranking, scalability, and knowledge distillation. The model performs strongly across datasets while supporting full-Wikipedia retrieval and speed–accuracy trade-offs.
- Datasets: The evaluation covers three challenging datasets: Zero-shot EL, TACKBP-2010, and WikilinksNED Unseen-Mentions.Zero-shot EL tests entirely unseen domains and entities, while the other datasets evaluate established and unseen-mention settings.
- Experimental setup: The bi-encoder retrieves candidates in a dense space, and the cross-encoder re-ranks the retrieved candidates for end-to-end linking.The experiments use top-64 candidates for the Zero-shot EL ranker and top-100 candidates for TACKBP-2010 training examples.
- Zero-shot entity linking: The two-stage model achieves much better end-to-end accuracy, largely because of improved retrieval.On the Zero-shot EL dataset, the cross-encoder is evaluated over candidates generated by the bi-encoder.
- TACKBP-2010: The approach remains effective with the 5.9M-article full-Wikipedia knowledge base and without TACKBP-2010 fine-tuning.Both full-Wikipedia and no-fine-tuning variants are included as ablations.
- Inference efficiency: 9.2 ms is the average batch retrieval time for top-100 candidates across 5.9M entities using exact FAISS search.The entity embeddings require about 2.8 hours to compute on 8 Nvidia Volta V100 GPUs, but this preprocessing can be parallelized.
- Inference efficiency: 2.6 ms and 1.4 ms average query times are achieved with HNSW17 and HNSW28, with less than 1.2% and 2.1% accuracy drops, respectively.Approximate search reduces query time relative to exact retrieval while preserving most accuracy and recall.
- Knowledge distillation: Knowledge distillation from the cross-encoder improves the bi-encoder over a bi-encoder trained without distillation.This provides another point in the accuracy–speed trade-off between the two architectures.
6 Qualitative Analysis
The qualitative examples show that cross-encoders often exploit contextual cues more effectively than bi-encoders, though misleading context can also cause cross-encoder errors.
- Context use: Cross-encoders use contextual words such as “Juventus” to disambiguate mentions that bi-encoders link incorrectly.The bi-encoder selects the Brazilian football player for “Ronaldo,” while the cross-encoder identifies the contextually appropriate entity.
- Context use: Cross-encoders can distinguish the intended sense of an entity from contextual descriptions that bi-encoders fail to interpret.One example concerns recognizing that a sentence describes art rather than fiction.
- Error patterns: Cross-encoders are often better at using context, but misleading contextual cues can still produce incorrect links.The examples include a cross-encoder error linking “Ancient Greek” to “Ancient Greek philosophy” because “philosophers” appeared in context.
7 Conclusion
The paper concludes that a simple two-stage BERT approach provides accurate and scalable entity linking without task-specific heuristics or external entity knowledge. It also identifies efficiency and knowledge distillation as routes to practical deployment, while outlining several extensions.
- Conclusion: The two-stage BERT model achieves state-of-the-art results on zero-shot WikilinksNED Unseen-Mentions and the established TACKBP-2010 benchmark.The conclusion attributes these results to a model that does not use task-specific heuristics or external entity knowledge.
- Conclusion: The evaluation shows that efficient linking is possible with a modest loss of accuracy despite the trade-off between accuracy and speed.The paper frames this trade-off as inherent to large pre-trained models.
- Conclusion: Knowledge distillation can further improve bi-encoder performance.This transfers benefits from the more accurate cross-encoder to the retrieval-oriented model.
- Future work: Future work includes enriching entity representations with type and graph information and modeling coherence across mentions.The proposed extensions also cover other languages, other domains, and joint mention detection and entity linking.
A Training details and hyper-parameters Optimization
The appendix records the computational infrastructure, hyperparameter search bounds, model parameter references, and validation criterion used during experimentation.
- Infrastructure: The experiments use 8 Nvidia Volta V100 GPUs for model training.This specifies the training hardware reported by the authors.
- Hyperparameter optimization: Hyperparameter bounds include context/mention token limits of [32, 64, 128], 128 candidate tokens, α = 0.5, and T in [2, 5].The settings apply to bi-encoder token limits and knowledge-distillation experiments, with grid search over 24 trials.
- Model configuration: The appendix refers to Table 11 for the number of parameters in the models.The supplied appendix passage provides the reference but not the parameter values themselves.
- Hyperparameter optimization: Validation-set accuracy is used to select hyperparameters in all experiments.This is the stated model-selection criterion.
- Appendix tables: Table 10 lists the bounds of the models’ hyperparameters, while Table 11 lists their parameter counts.The table captions identify the contents of each appendix table.
A.1 Training on Wikipedia data
Training begins with Wikipedia data and then fine-tunes models on specific datasets, using separate bi-encoder and cross-encoder configurations.
- Training procedure: Wikipedia data is used to pretrain the models before dataset-specific fine-tuning.This procedure is applied in experiments on TACKBP-2010 and WikilinksNED Unseen-Mentions.
- Wikipedia data: The May 2019 English Wikipedia dump contains 5.9M entities and supplies hyperlink examples for training.Anchor text serves as the mention, and the bi-encoder training set contains 9M examples with a 10K-example holdout set.
- Cross-encoder training: The cross-encoder is trained from the bi-encoder’s top 100 retrieved Wikipedia results.The cross-encoder training data is further down-sampled to 1M examples.
- Model configurations: The large bi-encoder uses a learning rate of 1e−5, batch size 128, and four training epochs.Its reported average runtime is 17.5 hours per epoch with a maximum of 32 context tokens.
- Model configurations: The large cross-encoder uses a learning rate of 2e−5, batch size 1, and one training epoch.Its reported average runtime is 37.2 hours per epoch with a maximum of 32 context tokens.
A.2 Zero-shot Entity Linking Dataset
The zero-shot entity linking dataset contains 49K training, 10K validation, and 10K test examples. The reported configurations compare bi-encoder and cross-encoder models, with cross-encoders requiring substantially longer training per epoch.
- Dataset: 49K, 10K, and 10K examples comprise the train, validation, and test sets, respectively.The dataset is available at the cited project URL.
- Training configurations: Bi-encoder base training used a 2e−5 learning rate, batch size 128, 128 maximum context tokens, and 5 epochs.Average runtime was 28.2 minutes per epoch.
- Training configurations: Bi-encoder large training used a 1e−5 learning rate, batch size 128, 128 maximum context tokens, and 5 epochs.Average runtime was 38.2 minutes per epoch.
- Training configurations: Cross-encoder base training used batch size 1 and required 2.6 hours per epoch, compared with 28.2 minutes for the bi-encoder base.Both configurations used 128 maximum context tokens, while the cross-encoder used a 1e−5 learning rate and 2 epochs.
- Training configurations: Cross-encoder large training required 8.5 hours per epoch, compared with 38.2 minutes for the bi-encoder large.The cross-encoder used batch size 1, 128 maximum context tokens, a 1e−5 learning rate, and 2 epochs.
A.3 TACKBP-2010 Dataset
The TACKBP-2010 data include 1,074 training and 1,020 test examples, with 10-fold cross-validation on the training set. The reported configurations include bi-encoder, knowledge-distilled bi-encoder, and cross-encoder variants with distinct training costs.
- Dataset: 1,074 and 1,020 annotated examples comprise the TACKBP-2010 train and test sets, respectively.The training set is evaluated using 10-fold cross-validation.
- Training configurations: The bi-encoder large model uses a 2e−6 learning rate, batch size 128, 32 maximum context tokens, and 10 epochs at 9.0 minutes per epoch.
- Training configurations: Knowledge-distilled bi-encoder training uses a 2e−5 learning rate, batch size 128, 32 maximum context tokens, T = 2, α = 0.5, and 10 epochs.Average runtime is 11.2 minutes per epoch.
- Training configurations: The cross-encoder large model uses a 1e−5 learning rate, batch size 1, 128 maximum context tokens, and 10 epochs at 20.4 minutes per epoch.
- Dataset: The train, validation, and test sets contain 2.2M, 10K, and 10K examples, respectively, while fine-tuning uses a 100K-example subset.The authors report that using more than 100K examples does not help.
- Training configurations: The cross-encoder large model for the larger dataset uses batch size 5 and requires 4.2 hours per epoch.It uses a 2e−6 learning rate, 128 maximum context tokens, and 1 epoch.