T rain X – Named Entity Linking with Active Sampling and Bi-Encoders (2020)

Abstract: We demonstrate TrainX, a system for Named Entity Linking for medical experts. It combines state-of-the-art entity recognition and linking architectures, such as Flair and fine-tuned Bi-Encoders based on BERT, with an easy-to-use interface for healthcare professionals. We support medical experts in annotating training data by using active sampling strategies to forward informative samples to the annotator. We demonstrate that our model is capable of linking against large knowledge bases, such as UMLS (3.6 million entities), and supporting zero-shot cases, where the linker has never seen the entity before. Those zero-shot capabilities help to mitigate the problem of rare and expensive training data that is a common issue in the medical domain.

## introduction named entity linking is a well-studied task for decades @xcite . it includes recognizing and disambiguating mentions of entities in text against a catalogue or a knowledge base. however, training data is often missing and requires additional expensive labeling, especially in domains like medicine, where the availability of domain experts for rare diseases is limited. moreover, novel as well as uncommon entities such as rare diseases might not have been part of the training data; in that case, the linker must solve a zero-shot scenario by disambiguating a mention never seen before. existing easyto-use annotation interfaces like prodigy 1 either fail to support entity linking annotations or have limited support for an end-user to find the correct entity in a large knowledge base. further, they do not support the annotator by actively sampling relevant documents to save annotation time. active-learning-enabled annotation tools, like inception @xcite , overcome this problem, but they are optimized for annotating multiple layers of linguistic features, which makes their user interfaces very complex and crowded. using such tools leads to additional training costs for medical professionals. contribution we present trainx, a system that consists of state-of-the-art entity recognition and linking architectures combined with an easy-to-use interface for healthcare professionals. our system supports medical experts in annotating data and training models for medical named entity linking based on umls @xcite with more than 3.6 million entities. by using active sampling, we minimize labeling efforts. trainx uses transfer learning by leveraging bi-encoders @xcite @xcite for disambiguation and a knn-index to retrieve candidate entities within milliseconds. we mitigate issues caused by sparse training data by using zero-shot optimized techniques that can generalize beyond the labels seen in training. to our knowledge, this is the first named entity linking approach that combines an easy-to-use frontend with the transfer learning capabilities of recent bert models. the system is licensed under apache 2.0 and is available on github 2 . ## demonstrating medical named entity linking in this section, we demonstrate the usage of trainx for a medical entity linking scenario. although we show a domain-specific task, trainx is not limited to the medical domain but can be easily adapted to any entity linking use case with a knowledge base that contains names and short descriptions for every entity. figure 1 shows the usage of trainx in an example scenario where existing gold annotations are available. first, the annotator begins a new session or resumes an existing one (1). in the case of a new session, she uploads a new dataset (2). after the upload, she obtains sampled documents from the dataset, in order to annotate them (3). the samples view (4) allows her to add new user annotations (green) or to view/edit gold annotations provided in the dataset (yellow). by clicking on an annotation, she can examine details of the linked entity and make a correction if needed by using the annotation helper (5); the modified samples are henceforth user annotations and marked in green. a full-text search on the umls supports the annotator to interactively explore the knowledge base in order to speed up the annotation/examination/correction workflow. by clicking on the checkmark (6), she can mark the entire sample as correct, or she can use the arrows above to request as many further samples as she likes. when one round of annotation is finished, she uploads the annotated samples (7) and starts the training phase (8), while the system will apply the model to the newly adjusted data. she can query the training status at any time (9). when training is finished, she retrieves the newly processed samples (10) and is returned to the samples view, where the predictions of the newly trained model are shown as pred annotations in blue (11). now, she can further correct and/or add annotations and iterate the process. a video of this demonstration is available under https://youtu.be/xat94uneeq4 . ## named entity linking with active sampling and bi-encoders a system component overview is shown in figure 2 . prior to training, the system needs to be initialized with the knowledge base (umls in our case) and an optional set of pre-training documents (1). after the initialization, the user can upload her documents (2) and annotate (3) them using the support of the annotation helper (4). the user is supported by an active sampling of further samples to annotate and correct. next, the updated annotations and the current model are sent to training component where the named entity recognizer and linker are now (re-)trained on the supplied data (5). after the training succeeded, the newly trained model is used to recognize and link mentions in the uploaded documents to provide feedback to the user (6). named entity recognition and linking with bi-encoders the recognition step is the first step of an entity-linking pipeline. a high recall is crucial because the linker will not be able to disambiguate mentions that have not been found by the recognizer in the first place. we chose the flair-framework @xcite because it has proven to achieve state-of-the-art results @xcite . we implemented a bi-encoder based on the work by @xcite and @xcite . the bi-encoder uses fine-tuned bert models to project mentions and entities in a common dense vector space to allow retrieval based on vector similarity. the projection is enforced using a cross-entropy loss function on both encoders' output that rewards a high similarity between the mention and the matching entity representation. in contrast to other entity linking architectures such as the entity linker from the "scispacy" framework @xcite , the bi-encoder approach is fully focused on solving zero-shot problems @xcite . also, it allows us to apply confidence measurements needed for our active sampling mechanisms. we chose not to implement an additional cross-encoder as proposed by @xcite , @xcite and @xcite due to computational efficiency @xcite . the bi-encoder consists a mention encoder y @xmath0. t 1 and t 2 are bert models, and m and e are sequences of wordpiece tokens that encode mention and entity, respectively. the pooling method pool() aggregates the resulting tensor into single vector representations y m and y e by using the vector of the [cls] token as a representation of the whole token sequence. training the bi-encoder our system works with three types of labels: gold for labels from (optional) pre-training data, user labels that have been created or updated by the user, and pred labels that were predicted from the entity linker. we train on gold and user labels. for the mention encoder t 1 , the mention and its context are encoded by using two special tokens to mark the beginning and end of a mention i.e. ## evaluation and discussion evaluation on medmentions we chose the publicly available "medmentions-st21pv" dataset (mohan and li, 2019) containing 4,392 annotated abstracts of pubmed articles with 203,282 annotations of 25,419 unique concepts. we use the pre-defined train/test/dev split. the test set contains annotations for 3,590 concepts that have not been in the training set which allows us to evaluate zero-shot capabilities. as named entity recognizer, we selected the flair framework @xcite and trained it on bioes tagging until the early-stop mechanism of flair stopped the training process. the scores of this model resulted in a precision of 69.2, a recall of 69.0 and an f1 of 69.1. in umls @xmath1.9 63.2 63.2 50.3 50.0 21.7 44.4 43.4 26.5 25.3 r@20 73.7 86.2 86.1 76.3 76.0 50.7 73.6 71.7 55.6 53.1 end-to-end (ner + nel) precision -47.4 47.4 ---32.1 31.1 --recall -45.7 45.7 ---30.9 30.0 --f1 -46.5 46.5 ---31.5 30.6 -- the lower half provides an end-to-end evaluation of the whole pipeline.