S p EL : Structured Prediction for Entity Linking (2023)

Abstract: Entity linking is a prominent thread of research focused on structured data creation by linking spans of text to an ontology or knowledge source. We revisit the use of structured prediction for entity linking which classifies each individual input token as an entity, and aggregates the token predictions. Our system, called SpEL (Structured prediction for Entity Linking) is a state-of-the-art entity linking system that uses some new ideas to apply structured prediction to the task of entity linking including: two refined fine-tuning steps; a context sensitive prediction aggregation strategy; reduction of the size of the model’s output vocabulary, and; we address a common problem in entity-linking systems where there is a training vs. inference tokenization mismatch. Our experiments show that we can outperform the state-of-the-art on the commonly used AIDA benchmark dataset for entity linking to Wikipedia. Our method is also very compute efficient in terms of number of parameters and speed of inference.

## introduction knowledge bases, such as wikipedia and yago @xcite , are valuable resources that facilitate structured information extraction from textual data. entity linking @xcite involves identifying text spans (mentions) and disambiguating the concept or knowledge base entry to which the mention is linked. entity linking can be viewed as three interlinked tasks @xcite @xcite : gr ace kelly by m ika reached ... w1 w2 w3 w4 w5 classification head of fixed candidate set size ... single_(music) united_kingdom fc_mika mika_solo mika_(singer) uk_singles_chart malika_ayane grace_kelly_(song) grace_kelly_(actress) o 0.0 0.0 1.0 0.01 0.83 0.16 0.01 0.80 0.19 0.01 0.85 0.14 1.0 0.0 0.0 0.0 0.0 0.01 0.52 0.99 0.48 pred: grace_kelly_(song) pred: mika_(singer) pred: o pred: o grace_kelly_ (actress) is not in the candidate set for this phrase, so will be ignored! prediction based on the highest average probability ... ... in this example, we demonstrate top 3 most probable entities for each tokenized subword. in the first phrase, the most likely entity identified for grace kelly is incorrect. we use a candidate set for this mention to filter out irrelevant entity links with a high probability. in the second phrase, mika is correctly narrowed down to the top 3 potential related entities. however, subword predictions do not consistently agree on the most probable prediction. in such cases, we calculate the average predicted probability for each entity across subwords and select the entity with the highest average probability. an end-to-end entity linking system does all three tasks and links text spans to concepts. the system can either have independently modelled components @xcite or jointly modelled components @xcite . similar to almost all nlp tasks, recent en-tity linking models use pre-trained representation learning methods that are based on transformers @xcite . these methods commonly utilize bidirectional language models like bert @xcite or auto-regressive causal language models such as gpt @xcite or bart @xcite , which are then finetuned on specific entity linking training datasets. in a number of such techniques, entity linking is framed as another well-studied problem, such as sequence-to-sequence translation @xcite or question answering @xcite . entity linking can be viewed as sequence tagging using structured prediction foot_0 , aiming to assign one of the finitely many classes to every input utterance. this approach involves using a pre-trained model to encode each input subword token into a multilayer context-aware dense vector representation. a classifier head calibrates each token representation to predict the entity for each subword token. due to the large number of possible entities and issues with consistency of entity prediction across multiple subword tokens, structured prediction for entity linking (surprisingly) has not been studied in-depth. our contributions in this paper are as follows: we introduce recent entity linking methods in section 2; explain our approach to structured prediction for entity linking and our context sensitive prediction aggregation strategy in section 3 with comprehensive experiments in section 4. ## related work entity linking can be framed as another wellstudied task and the best solution for that task is applied. autoregressive encoder-decoder sequenceto-sequence translation is one such approach. de cao et al. (2021b) consider the input text as the source for translation and the text is annotated with wikimedia markup containing the mention spans and the entity for each mention. instead of mapping the entity identifiers into a single id (this is the default in many techniques including this work because the approach can be easily ported to ontologies other than wikipedia), their model generates the entity label in a token-by-token basis (it generates the wikipedia url one subword at a time). the generation process follows a constrained decoding schema that prevents the model from producing invalid entity urls while limiting the generated output to the entities in a predefined candidate set (discussed in section 3.1). de cao et al. (2021a) use a bert-style bidirectional model fine-tuned to identify potential spans (mention detection) by learning spans using a begin probability and an end probability for each subword in the input text. for each potential span, they use a generative lstm-based @xcite language model to generate the entity identifiers (token-by-token), and limit the generation process to pre-defined candidate sets. mrini et al. (2022) frame entity linking as a sequence-to-sequence translation task using bart @xcite . they duplicate the bart decoder three times to fine-tune the model in a multitask setting. the two additional decoder modules are trained using auxiliary objectives of mention detection and re-ranking. while this training method increases the model size during training, they mitigate increased model size and speed at inference time by excluding the auxiliary decoder modules and employing sampling and re-ranking techniques on the generated target sequences. @xcite use question answering as a way to frame the entity linking task. they suggest a two-step entity linking model in which they use a fine-tuned transformer-based blink @xcite model to find all the potential entity records that might exist in the text and then utilize a finetuned question answering electra @xcite model to identify the matching occurrences of the potential entities discovered in the first step. this approach obtains high accuracy; however, it is very resource intensive and inference is slow. structured prediction (subword token multi-label classification) is the other well-studied problem. @xcite proposes a very simple entity linking model which places a classification head on top of a bert language model and directly classifies each subword representation using a softmax over all the entities known to the model. our work also uses structured prediction because it is one of the most lightweight techniques in terms of compute cost and inference speed. we extend the structured prediction model in this work by (1) utilizing a context sensitive prediction aggregation strategy to form meaningful span annotations (section 3.2), (2) addressing a training/inference tokenization mismatch issue (section 3.3), (3) examining the role of different types of candidate sets (section 3.1) in curating the predicted results, and (4) optimizing the implementation of the structured prediction model. we obtain more than 6.1% points micro-f1 improvement as well as more than 10 times reduction in required disk space, and close to 4 times reduction in the required gpu memory (in base case) compared to @xcite . our approach is also faster at inference time than previous transformer-based methods. a number of recent techniques focus on enhancing the entity linking knowledge in bert (or one of its variations) and utilize one or more such knowledge-enhanced models to perform the task of entity linking. peters et al. (2019) inject wikipedia and wordnet information into the last few layers of bert, poerner et al. (2020) inject pre-trained wikipedia2vec (yamada et al., 2020a) entity embeddings into the input layer of the language model while freezing the rest of its parameters, and martins et al. (2019) leverage a stack-lstm @xcite named entity recognition model to enhance entity linking performance using multitask learning to improve entity linking. these approaches are wikipedia-centric and while we also experiment on wikipedia exclusively, our approach can be used on any entity-linking dataset that has entities from other ontologies such as the medmentions dataset (mohan and li, 2019) which links to concepts in umls ontology @xcite . @xcite jointly model mention detection and mention disambiguation using an lstm-based architecture while reusing the candidate sets created by @xcite as a replacement for the candidate generation step, and kannan @xcite follow a similar framework while modeling each of mention detection and mention disambiguation using separate bert models. @xcite compute entity embeddings (instead of using pre-trained ones) using the average of the subword embeddings of the candidates and compare them to the average of the subword embeddings for the potential span (training a siamese network; @xcite . @xcite investigate pre-training strategies specifically tailored for transformer models to perform entity linking, diverging from the conventional use of pre-trained bert models. @xcite propose a modular configuration that composes mention detection, candidate generation, and mention disambiguation in a pipeline approach, incorporating the most promising components from prior research. in section 4.4, we conduct experiments to compare spel to the methods discussed in this section. ## entity linking formally, entity linking receives a passage (p) containing words {w 1 , ..., w n } and produces a list containing ℓ span annotations. each span annotation is expected to be a triplet of the form (span start, span end, entity identifier). the span start and span end values are expected to be character positions on the raw input p, and the entity identifier values are selected from a predefined vocabulary of entities (e.g. there would be approximately 6 million entities to choose from when entity linking to wikipedia). the massive entity vocabulary size increases the model's hardware requirements and in some cases renders the task intractable. ## conclusion we introduced several improvements to a structured prediction approach for entity linking. our experimental results on the aida dataset show that our proposed improvements to the structured prediction model for entity linking can achieve state-ofthe-art results using a commonly used evaluation toolkit providing head to head numbers for competing methods on the same dataset. we show that our approach has the best f1-score on this task compared to the state-of-the-art on this dataset. we are more compute efficient with many fewer parameters and our approach is also much faster at inference time, providing faster throughput, compared to previous methods. in future work, we plan to extend our work to other entity linking datasets in other domains such as biomedical research. we also plan to research multilingual applications of structured prediction for entity linking including benefits to projection of entity linking concepts from one language to another and using multilingual representation learning for our base model.