Proxy-based Zero-Shot Entity Linking by Effective Candidate Retrieval (2022)

Abstract: A recent advancement in the domain of biomedical Entity Linking is the development of powerful two-stage algorithms – an initial candidate retrieval stage that generates a shortlist of entities for each mention, followed by a candidate ranking stage. However, the effectiveness of both stages are inextricably dependent on computationally expensive components. Specifically, in candidate retrieval via dense representation retrieval it is important to have hard negative samples, which require repeated forward passes and nearest neighbour searches across the entire entity label set throughout training. In this work, we show that pairing a proxy-based metric learning loss with an adversarial regularizer provides an efficient alternative to hard negative sampling in the candidate retrieval stage. In particular, we show competitive performance on the recall@1 metric, thereby providing the option to leave out the expensive candidate ranking step. Finally, we demonstrate how the model can be used in a zero-shot setting to discover out of knowledge base biomedical entities.

## introduction the defining challenge in biomedical entity linking (el) is performing classification over a large number of entity labels with limited availability of labelled mention data, in a constantly evolving knowledge base. for instance, while the unified medical language system (umls) knowledge base @xcite ) contains millions of unique entity labels, the el training data in the biomedical domain as a whole is notoriously scarce, particularly when compared to the general domain -wikipedia, for instance, is powerful as both a knowledge base and a source of matching entities and mentions. furthermore, biomedical knowledge bases are evolving rapidly with new entities being added constantly. given this knowledge base evolution and scarcity of training data it is crucial that biomedical entity linking systems can scale efficiently to large entity sets, and can discover or discern entities outside of the knowledge base and training data. recent methods in the general entity linking domain @xcite address the data issue with zero-shot entity linking systems that use entity descriptions to form entity representations and generalise to entities without mentions. a particularly powerful architecture was initially proposed by @xcite and further improved by @xcite . it consists of a two-stage approach: 1) candidate retrieval in a dense space performed by a bi-encoder @xcite which independently embeds the entity mention and its description, and 2) candidate ranking performed by a cross-encoder which attends across both the mention and entity description @xcite . in this work we focus on the former, which is traditionally optimised with the cross-entropy (ce) loss and aims to maximise the similarity between the entity mention and its description relative to the similarities of incorrect mention-description pairs. in practice, the large number of knowledge base entities necessitates the use of negative sampling to avoid the computational burden of comparing each mention to all of the entity descriptions. however, if the sampled distribution of negatives is not reflective of the model distribution, the performance may be poor. recently, @xcite showed that using hard negatives -the highest scoring incorrect examples -results in bias reduction through better approximation of the model distribution. collecting hard negatives is computationally expensive, as it requires periodically performing inference and retrieving approximate nearest neighbours for each mention. at the ranking stage, negative sampling is not required, as the number of candidates usually does not exceed 64. however, the state-of-the-art crossencoder model used for ranking is very expensive to run, scaling quadratically with the input sequence length. this highlights the need for efficient and performant candidate retrieval models capable of disambiguating mentions without the need for the expensive ranking step. in this paper, we propose and evaluate a novel loss for the candidate retrieval model, which breaks the dependency between the positive and negative pairs. our contributions are: (1) a novel loss which significantly outperforms the benchmark cross-entropy loss on the candidate retrieval task when using random negatives, and performs competitively when using hard negatives. (2) we design and apply an adversarial regularization method, based on the fast gradient sign method @xcite , which is designed to simulate hard negative samples without expensively mining them. (3) we construct a biomedical dataset for out of knowledge base detection evaluation using the medmentions corpus and show that our model can robustly identify mentions that lack a corresponding entry in the knowledge base, while maintaining high performance on the retrieval task. our main testing ground is the biomedical entity linking dataset medmentions @xcite , which utilizes umls as its knowledge base. additionally, to confirm that our method works also in the general, non-biomedical domain, we evaluate it on the zero-shot entity linking (zeshel) dataset proposed in @xcite . we focus on the retrieval task with the recall@1 metric, because we are aiming to predict the entity directly without requiring the additional expensive ranking stage. our results show that both the proposed loss and regularization improve performance, achieving state-of-the-art results on recall@1 and competitive performance on recall@64 on both datasets. finally, we demonstrate that our model can robustly identify biomedical out of knowledge base entities, without requiring any changes to the training procedure. ## related work zero-shot entity linking there is a plethora of work on zero-shot entity linking methods leveraging the bi-encoder architecture @xcite for candidate retrieval. these include novel scoring functions between the input and the label @xcite @xcite , cross-domain pretraining methods @xcite , training and inference optimisation techniques @xcite and effective entity representation methods @xcite . our work instead focuses on optimising the candidate retriever's loss function. the impact of hard negatives on the entity linking model performance has also been investigated @xcite . notably, @xcite develop analytical tools to explain the role of hard negatives and evaluate their model on the zero-shot entity linking task. we draw on this work, but move away from the ce loss towards a novel contrastive proxy-based loss. finally, there is a body of work on zero-shot entity linking in the biomedical domain using clustering @xcite . our method does not consider the affinities between mentions directly and links them independently. therefore, we do not study entity discovery. an important aspect of biomedical entity linking systems is the detection of "unlinkable" mentions that lack a corresponding entry in the knowledge base -referred to as nil detection. methods for this task can be grouped into four main strategies @xcite : (1) label a mention as nil when the corresponding candidate retriever does not return any candidate entities @xcite , (2) assign the nil label to mentions whose corresponding top-ranked entity does not exceed some score threshold @xcite @xcite , (3) train a classifier that predicts whether the top-ranked entity for a given mention is correct @xcite , (4) explicitly introduce a nil class to the candidate ranking model @xcite . a downside of the final approach is that knowledge of the nil mention distribution is required at training time. in this work we tune a nil score threshold (2) on a validation set. detecting unlinkable mentions is particularly important in the biomedical domain, where the knowledge bases are rapidly evolving. proxy-based losses state-of-the-art entity linking models such as blink @xcite leverage metric learning loss during training to make mentions similar to its assigned entity representations. metric learning losses could be divided into two categories, pair-based and proxy-based losses @xcite . pair-based losses can leverage semantic relations between data points, here mentions. however, training them can be highly computationally expensive. on the other hand, proxy-based losses are significantly less computationally complex. this is done by establishing a proxy for each class and trying to increase the similarity between data points and its assigned proxies. therefore, avoiding comparing the mentions to each other in favour of comparing the mentions to their proxies. we draw heavily on proxy-based losses @xcite ) from metric learning by treating entity descriptions as the proxies. we establish a proxy for each entity, creating mention-proxy (i.e. entity) pairs, and optimise the model to embed the mention close to its assigned proxy. the loss proposed here is similar to the proxy-nca loss of @xcite . our modification is the use of the softplus function, similar to @xcite , to avoid a vanishing gradient for the true mention-proxy pair. adversarial regularization entity linking systems often rely on careful mining of hard negative examples to boost their performance @xcite at the expense of increased computational complexity. the model needs update hard negatives for each mention periodically. a potential alternative to hard negative mining is training on adversarial examples @xcite -synthetic data points designed to induce the model to making incorrect predictions, such that they are more challenging. adversarial training can be seen as data augmentation and can help reduce overfitting. @xcite introduced a simple method for generating adversarial examples, called fast gradient sign method (fgsm), which we build upon in this work. fgsm creates adversarial examples by applying small perturbations to the original inputs -often the word embeddings for nlp problems. fgsm has been used successfully as a regulariser in supervised and semi-supervised nlp tasks @xcite . here, we follow a similar approach and use fgsm to augment our training pairs with adversarial positive and negative examples. ## task formulation in the entity linking task we are provided with a list of documents @xmath0, where each document has a set of mentions m @xmath1, m 2 , . . . , m n d }. the task is to link each mention m i to an entity e i , where each entity belongs to the knowledge base (kb) e. in this work we focus specifically on the problem of biomedical zero-shot entity linking. the setup for the zero-shot task is the same as for entity linking introduced above, except that the set of entities present in the test set is not present in the training set, i.e. e train ∩ e @xmath2. we focus specifically on the candidate retrieval task, where the goal is given a mention m i , reduce the pool of potential candidate entities from a kb to a smaller subset. candidate retrieval is crucial for biomedical entity linking because of the large size of knowledge bases. in this work we use the bi-encoder architecture for candidate retrieval. finally, in addition to the in-kb entity linking task, where you only consider entities inside the kb, we also consider an out of kb scenario, where the task is to map mentions to the augmented set of labels e ∪ nil, with nil indicating the absence of a corresponding kb entity. in this section, we review the categorical ce loss, used by current state-of-the-art models, in the context of entity linking @xcite . we then compare it to our proposed proxy-based loss. finally, we describe and motivate our regularization approach. ## results we present the results for candidate retrieval and benchmark our models against suitable methods. we name our method proxy-based entity linking (pel-pb). we also report the results of a version of our model which uses the ce (pel-ce) loss on all experiments for comparison. ## discussion & future work we have proposed and evaluated a novel proxybased loss for biomedical candidate retrieval. additionally, we have adopted an adversarial regularization technique designed to simulate hard negatives, and shown that both our loss and regularization boost performance on the recall@1 metric. we have also constructed a biomedical dataset for nil detection and demonstrated that our candidate retrieval model can robustly identify biomedical nil entities, while maintaining high overall performance. these are important advances towards closing the gap between the two-stage approach that include an expensive cross-encoder and a candidate retriever-only setup. notably, our work highlights the importance of hard negative sampling when optimising the candidate generator with the ce loss. random negative sampling together with ce loss can result in the problem becoming trivial, for example the randomly sampled negative entity having a different type. however, accessing hard negative examples during model training can be challenging, particularly when the knowledge base is large and entity representations are frequently updated. considering this, we recommend to employ our proxy-based loss for the candidate retrieval task in three different scenarios: (1) training with random negatives, (2) optimising for recall@1, (3) detecting nil entities. moreover, we also recommend leveraging fgsm regularisation in any setup and both retrieval and ranking tasks. an interesting approach would be to attempt to approximate hard negatives without frequent updates of the entity representations. this could potentially be done by keeping the entity encoder frozen, or exploring alternative relatedness measures which does not require frequently running the model over the whole knowledge base. finally, there is a plethora of work on proxy-based @xcite and pair-based losses @xcite @xcite , usually discussed in the computer vision and metric learning literature. improving the candidate retrieval is a crucial step towards high-performing and efficient entity linking systems that can be easily applied in real-world settings.