Efficient One-Pass End-to-End Entity Linking for Questions

Introduction

Entity linking (EL), the task of identifying entities and mapping them to the correct entries in a database, is crucial for analyzing factoid questions and for building robust question answering (QA) systems. For instance, the question “when did shaq come to the nba?” can be answered by examining Shaquille O’Neal’s Wikipedia article Min et al. (2019), or its properties in a knowledge graph Yih et al. (2015); Yu et al. (2017). However, real-world user questions are invariably noisy and ill-formed, lacking cues provided by casing and punctuation, which prove challenging to current end-to-end entity linking systems Yang and Chang (2015); Sorokin and Gurevych (2018). While recent pre-trained models have proven highly effective for entity linking Logeswaran et al. (2019); Wu et al. (2020), they are only designed for entity disambiguation and require mention boundaries to be given in the input. Additionally, such systems have only been evaluated on long, well-formed documents like news articles Ji et al. (2010), but not on short, noisy text. Also, most prior works have focused mainly on improving model prediction accuracy, largely overlooking efficiency.

In this work, we propose ELQ, a fast and accurate entity linking system that specifically targets questions. Following the Wikification setup Ratinov et al. (2011), ELQ aims to identify the mention boundaries of entities in a given question and their corresponding Wikipedia entity. We employ a biencoder based on BERT Devlin et al. (2019) as shown in Figure 1. The entity encoder computes entity embeddings for all entities in Wikipedia, using their short descriptions. Then, the question encoder derives token-level embeddings for the input question. We detect mention boundaries using these embeddings, and disambiguate each entity mention based on an inner product between the mention embeddings (averaged embedding over mention tokens) and the entity embeddings. Our model extends the work of Wu et al. (2020) but with one major difference: our system does not require pre-specified mention boundaries in the input, and is able to jointly perform mention detection and entity disambiguation in just one pass of BERT. Thus, at inference time, we are able to identify multiple entities in the input question efficiently.

We extend entity disambiguation annotations from Sorokin and Gurevych (2018) to create an end-to-end question entity linking benchmark. Evaluated on this benchmark, we are able to outperform previous methods in both accuracy and run-time. ELQ has much faster end-to-end inference time than any other neural baseline (by $2\times$), while being more accurate than all previous models we evaluate against, suggesting that it is practically useful for downstream QA systems. We verify the applicability of ELQ to practical QA models in a proof-of-concept experiment, by augmenting GraphRetriever (Min et al., 2019) to use our model, improving its downstream QA performance on three open-domain QA datasets (by up to $6\%$).

Related Work

Much prior work on entity linking has focused on long, grammatically coherent documents that contain many entities. This setting does not accurately reflect the difficulties of entity linking on questions. While there has been some previous work on entity linking for questions Sorokin and Gurevych (2018); Blanco et al. (2015); Chen et al. (2018); Tan et al. (2017), such works (mostly from the pre-BERT era) utilize complex models with many interworking modules. For example, Sorokin and Gurevych (2018) proposes a variable-context granularity (VCG) model to address the noise and lack of context in questions, which incorporates signals from various levels of granularity by using character-level, token-level, and knowledge-base-level modules. They also rely on external systems as a part of the modeling pipeline.

In this work, we take a much simpler approach that uses a biencoder. Biencoder models have been used in a wide range of tasks (Seo et al., 2019; Karpukhin et al., 2020; Wu et al., 2020). They enable fast inference time through maximum inner product search. Moreover, as we find, biencoders can be decomposed into reusable question and entity encoders, and we can greatly expedite training by training one component independently of the other.

Problem Definition & ELQ Model

We propose an end-to-end entity linking system that performs both mention detection and entity disambiguation on questions in one pass of BERT.

Given an input question $q=q_{1}\mathinner{\cdotp\mkern-2.0mu\cdotp\mkern-2.0mu\cdotp}q_{n}$ of length $n$, we first obtain question token representations based on BERT (Devlin et al., 2019):

where each $\mathbf{q_{i}}$ is a $h$-dimensional vector. We then obtain entity representations $\mathbf{x}_{e}$ for every $e_{i}\in\mathcal{E}$.

where ${\mathtt{[CLS]}}$ indicates that we select the representation of the $\mathtt{[CLS]}$ token. We consider candidate mentions as all spans $[i,j]$ ($i$-th to $j$-th tokens of $q$) in the text up to length $L$.

To compute the likelihood score of a candidate span $[i,j]$ being an entity mention, we first obtain scores for each token being the start or the end of a mention:

Entity Disambiguation

We obtain a mention representation for each mention candidate $[i,j]$ by averaging $\mathbf{q_{i}}\mathinner{\cdotp\mkern-2.0mu\cdotp\mkern-2.0mu\cdotp}\mathbf{q_{j}}$, and compute a similarity score $s$ between the mention candidate and an entity candidate $e\in\mathcal{E}$:

We then compute a likelihood distribution over all entities, conditioned on the mention $[i,j]$:

Training

We jointly train the mention detection and entity disambiguation components by optimizing the sum of their losses. We use a binary cross-entropy loss across all mention candidates:

whereby $y_{[i,j]}=1$ if $[i,j]$ is a gold mention span, and otherwise. $N$ is the total number of candidates we consider.If $n\geq L$, $N=L(L+1)/2+(n-L)L$. Otherwise, $N=n(n+1)/2$.

The entity disambiguation loss is given by

where $e_{g}$ is the gold entity corresponding to mention $[i,j]$.

To expedite training, we use a simple transfer learning technique: we take the entity encoder trained on Wikipedia by Wu et al. (2020) and freeze its weights, training only the question encoder on QA data. In addition, we mine hard negatives. As entity encodings are fixed, a fast search of hard negatives in real time is possible.

Inference

Figure 1 shows our inference process. Given an input question $q$, we use our mention detection model to obtain our mention set $\mathcal{M}=\{[i,j]:1\leq i\leq j\leq\min(i+L-1,n),p([i,j])>\gamma\}$, where $\gamma$ is our threshold (a hyperparameter). We then compute $p(e,[i,j])=p(e|[i,j])p([i,j])$ for each mention $[i,j]\in\mathcal{M}$, and threshold according to $\gamma$. In contrast to a two-stage pipeline which first extracts mentions, then disambiguates entities Févry et al. (2020), a joint approach grants us the flexibility to consider multiple possible candidate mentions for entity linking. This can be crucial in questions as it can be difficult to extract mentions from short, noisy text in a single step.

More implementation details can be found in Appendix D.

Experiments

Table 1 reports the statistics of the resulting datasets, WebQSP${}_{\text{EL}}$ and GraphQ${}_{\text{EL}}$. Following Sorokin and Gurevych (2018), we use WebQSP${}_{\text{EL}}$ for training and GraphQ${}_{\text{EL}}$ for zero-shot evaluation.

Using the rule defined by Carmel et al. (2014), a prediction is correct only if the groundtruth entity is identified and the predicted mention boundaries overlap with the groundtruth boundaries. (This is sometimes known as “weak matching”.) Specifically, let ${\mathcal{T}}$ be a set of gold entity-mention tuples and $\widehat{\mathcal{T}}$ be a set of predicted entity-mention tuples, we define precision ($p$), recall ($r$) and F1-score ($F_{1}$) as follows:

Baselines

We use the following baselines: (1) TAGME Ferragina and Scaiella (2012), a lightweight, on-the-fly entity linking system that is popular for many downstream QA tasks, being much faster than most neural models (Joshi et al., 2017; Sun et al., 2018; Min et al., 2019), (2) VCG (Sorokin and Gurevych, 2018), the current state-of-the-art entity linking system on WebQSP, and (3) biencoder from BLINK (Wu et al., 2020). As BLINK requires pre-specified mention boundaries as input, we train a separate, BERT-based span extraction model on WebQSP in order to predict mention boundaries (details in Appendix B).

2 Results

Table 2 show our main results. We find that BERT-based biencoder models far outperform the state-of-the-art (VCG) on both datasets, in performance and in runtime. Moreover, ELQ outperforms all other models trained in a comparable setting, and is much more efficient than every other neural baseline (VCG and BLINK). ELQ is also up to $2.3\times$ better than TAGME — in the case of WebQSP${}_{\text{EL}}$.

ELQ outperforms BLINK, suggesting that it is possible to train representations from a single model to resolve both entity references as well as mention boundaries of all entities in text, without restricting the model to focusing on a single marked entity as in BLINK.

Runtime

We record the inference speed on CPUs in number of questions processed per second for all models (Table 2). For BLINK, we report the combined speed of our span extraction model and the BLINK entity linker, in order to compare the end-to-end speeds. ELQ, which performs both detection and disambiguation in one pass of BERT, is approximately $2\times$ faster than BLINK, which performs multiple passes, while also outperforming BLINK in F1 score. Moreover, against TAGME Ferragina and Scaiella (2012), ELQ is only $1.5\times$ slower on WebQSP${}_{\text{EL}}$ and $2.0\times$ slower on GraphQ${}_{\text{EL}}$, despite TAGME being a completely non-neural model (with much lower accuracy).

QA Experiments

To demonstrate the impact of improved entity linking on the end QA accuracy, we experiment with the task of textual open-domain question answering, using GraphRetriever (GRetriever) (Min et al., 2019). GRetriever uses entity linking to construct a graph of passages in the retrieval step and deploys a reader model to answer the question. The original model uses TAGME for entity linking; we replace TAGME with ELQ and keep the other components the same, in order to isolate the impact of entity linking.Min et al. (2019) used two reader models, ParReader++ and GraphReader; for simplicity, we only use ParReader++ As an additional baseline, we also add the result of TF-IDF, implemented by Chen et al. (2017), a widely used retrieval system.

Results are shown in Table 3. Following literature in open-domain QA, we evaluate our approach on three datasets, WebQuestions (Berant et al., 2013), Natural Questions (Kwiatkowski et al., 2019) and TriviaQA (Joshi et al., 2017). In particular, WebQuestions (WQ) and Natural Questions (NQ) consist of short, noisy questions from Web queries, in line with the motivation of our work. We observe that simply replacing TAGME with ELQ significantly improves performance, including 5.9% and 3.9% absolute improvements on WQ and NQ, respectively. While ELQ trained on Wikipedia achieves good results overall, further fine-tuning on WebQSP${}_{\text{EL}}$ gives extra gains on WQ. This indicates that, if entity linking annotations in the same domain are available, using them to finetune ELQ can bring further gains.

Analysis

We set up experiments to disentangle the capability of ELQ’s entity linker and mention detector. First, to test just the mention detector (MD only), we measure just the mention boundary overlap between predicted and groundtruth mentions, ignoring the entity label. Next, to test just the entity linker (EL only), we give the entity linking component gold mention boundaries, and compute the resulting F1 score. We do this for both ELQ and BLINK. For comparability, we use the version of ELQ trained on Wikipedia. Results are reported in Table 4. Surprisingly, we find that both components of ELQ outperform BLINK, suggesting that the two tasks might mutually benefit from being trained jointly.

Runtime

To confirm that our biencoder’s main bottleneck is the BERT forward pass — and thus, investing in decreasing the number of BERT forward passes is valuable — we separately time each component of ELQ during inference. We run examples from WebQSP${}_{\text{EL}}$ test set one at a time through ELQ, on 1 CPU, and average runtimes across all examples. Indeed, we find that the BERT forward pass to be the slowest component of the model, taking $0.683$s, over $6\times$ slower than the next slowest component of the model, inner-product search (taking $0.107$s). Everything else takes a combined total of $5.08\times 10^{-3}$s.

Qualitative

We manually examine all our model’s errors on the WebQSP${}_{\text{EL}}$ and GraphQ${}_{\text{EL}}$ dev sets. We identify four broad error categories: (1) technically correct — where our model was technically correct but limitations in evaluation falsely penalized our model (i.e., we found a more or less precise version of the same entity), (2) not enough entities — where the model did not fully identify all entities in the question, (3) wrong entities — where our model linked to the wrong entity, (4) insufficient context — where the model made reasonable mistakes due to the lack of context (that even reasonable humans would make). Error type breakdowns can be found in Table 5.

Conclusion

We proposed an end-to-end model for entity linking on questions that jointly performs mention detection and disambiguation with one pass through BERT. We showed that it is highly efficient, and that it outperforms previous state-of-the-art models on two benchmarks. Furthermore, when applied to a QA model, ELQ improves that model’s end QA accuracy. Despite being originally designed with questions in mind, we believe ELQ could also generalize to longer, well-formed documents.

References

Appendix A Annotation Statistics

We show annotators a question and a gold entity in the question, and instruct them to annotate (i.e., put brackets around) the appropriate mention span.

For quality control, we created a shared annotation set by pulling a subset of examples from each dataset, and having all four annotators label that set. Inter-annotator agreement statistics on our shared set are shown in Table 6. Note that only 2 out of 41 shared examples did not have unanimous mention boundary agreement. The two conflicting examples, with the respective annotations, are shown below:

Appendix B Span-Extraction Model for Mention Boundary Detection

The BLINK Entity Linker requires mention boundaries to be marked in the input. In order to evaluate against BLINK for end-to-end Entity Linking, we train a span-extraction model to first obtain candidate mention boundaries, and subsequently use BLINK on these candidate mentions. Our span extraction model first represents every token $q_{i}$ in question $q=q_{1},\mathinner{\cdotp\mkern-2.0mu\cdotp\mkern-2.0mu\cdotp},q_{n}$ of length $n$ using a dense representation using BERTbase:

The model then computes a start span probability $p_{s}(q_{i}|q)$ and an end span probability $p_{e}(q_{i}|q)$ for every token $q_{i}$ using learnable vectors $\mathbf{w_{s}}$ and $\mathbf{w_{t}}$ respectively:

The model is trained to maximize the likelihood of $p_{s}(q_{s}|q)\times p_{e}(q_{e}|q)$ for each correct mention $[q_{s},q_{e}]$ in the training set of WebQSP, and similarly during inference, outputs the top-K scoring spans. These spans are used as mention boundary candidates, to evaluate BLINK in an end-to-end setting.

Appendix C Analysis

Table 7 presents the contributions of each component of our training scheme to our final result. We record model performance on WebQSP${}_{\text{EL}}$ (valid) after 20 (of 100) epochs of training on Wikipedia, having seen just 20% of the data. We note that both our transfer learning technique and our adversarial hard-negatives training expedites convergence.

Qualitative Error Analysis

We record specific examples for each of our four error categories (technically correct, not enough entities, wrong entities, and insufficient context), detailed in Section 6. Note there were 61 total mistakes for WebQSP${}_{\text{EL}}$ dev and 158 total mistakes for GraphQ${}_{\text{EL}}$ dev. Specific examples can be found in Table 9.

Appendix D Implementation Details and Hyperparameters

Following Wu et al. (2020), we use BERT${}_{\text{Large}}$ ($\sim$340M parameters) for the question and entity encoder. The span extraction model detailed in Appendix B, used for our BLINK baselines, is a BERT${}_{\text{Base}}$ model ($\sim$110M parameters).

We lowercase all inputs to ELQ, during both training and inference time, to make it case-insensitive. For both training and inference, the mention scorer considers all spans up to length $L=10$.

During training, we use FAISS Johnson et al. (2019) for fast inner product search when mining hard negatives. We do this in real time, inside the training loop. As we do not update the entity encoder, we were able to train the model with a single FAISS index, greatly increasing the training speed. For further speedup, we use a hierarchical index (IndexHNSWFlat), with $efConstruction=200$ and $efSearch=256$.

For $\mathcal{L}_{\text{ED}}$ at each iteration, computing $s(e,[i,j])$ for every $e\in\mathcal{E}$ is intractable. We thus approximate $\mathcal{L}_{\text{ED}}$ by replacing $\mathcal{E}$ with $\mathcal{E}^{\prime}$ for the softmax, where $\mathcal{E}^{\prime}$ is a set of hard negative entities, specifically, negative entities that have the highest 10 similarity scores with the mention representation.

For our WebQSP${}_{\text{EL}}$-trained model, we train for up to 100 epochs on WebQSP${}_{\text{EL}}$ data, using batch size 128 and context window size of 20 tokens. For our Wikipedia-trained model, we split the data evenly into 100 chunks and train on each (thus, making one pass through Wikipedia overall). For Wikipedia, we use batch size 32 and a context window size of 128 tokens. For Wikipedia + WebQSP${}_{\text{EL}}$ model, we take our Wikipedia-trained model and further fine-tune it on WebQSP${}_{\text{EL}}$ for up to 100 epochs (using the WebQSP${}_{\text{EL}}$ training settings). For all three training settings, we use the AdamW optimizer with learning rate 1e-5, coupled with a linear schedule with $10\%$ warmup. We clip gradients to max norm $1.0$.

Inference

During inference, we consider all mention candidates $[i,j]$ with mention score $\log p([i,j])\geq\gamma$. If no mention candidate has mention score $\geq\gamma$, we simply take the top-$50$-scoring mentions. $\gamma$ is a threshold we tune on each dataset’s dev data.

The linker then retrieves the $10$ closest entity candidates per mention boundary. We use the same hierarchical FAISS index as during training to expedite retrieval. Since the search is approximate, we expect some performance degradation. However, in practice, we found minimal performance degradation for significant speedup. On WebQSP${}_{\text{EL}}$ dev set, F1 score decreased from $92.5\to 91.9$, but run-time decreased from $127.0s\to 24.3s$ (for the entire dataset, with batch size 64).

As computing the softmax over all entities for $\log p(e|[i,j])$ is intractable, we simply softmax over our 10 retrieved candidates. At the end, we threshold the final joint score $\log p([i,j])+\log p(e|[i,j])$ based on $\gamma$.

We use manual tuning and binary search to find the best-performing hyper-parameters for the threshold $\gamma$. We optimize for F1-score on the development sets of WebQSP${}_{\text{EL}}$ and GraphQ${}_{\text{EL}}$. Best settings are reported in Table 8.

As mention overlaps are not allowed in the questions data, we have an additional global step of removing overlapping mention boundaries — in the case of multiple entities, we greedily choose the highest-scoring entity each time, and remove all entities which overlap with it.

Appendix E Infrastructure Details

We ran all training distributed across 8 NVIDIA TESLA V100 GPUs, each with 32 GB of memory. For 80-CPU inference, we run on 2 chips of Intel(R) Xeon(R) CPU E5-2698 v4 @ 2.20GHz with 20 cores (40 threads) each. For 1-CPU inference (reported in Table 2), we run only on a single core.