Enhancing Entity Representations with Prompt Learning for Biomedical Entity Linking

1 Introduction

With the advancements in the healthcare domain, the volume of biomedical text, such as electronic health records (EHRs) or biomedical literature, is explosively growing daily, creating an urgent need to utilize the useful information contained in these records. Entity linking is a necessary step that builds on the output mentions of named entity recognition (NER), which refers to the process of automatically linking mentions in plain text, to a standardized list of entities in a knowledge

* Corresponding authors.

Mention with Context

base (KB). Although there are numerous entity linking approaches in the general domain, they cannot be directly transferred to biomedical domain due to the huge differences in language characteristics and knowledge bases. Thus, it is critical to have accurate biomedical entity linking techniques, which would better assist the understanding of biomedical text. Furthermore, biomedical entity linking can be useful in many downstream applications such as medical-related decision making, predictive modeling, medical information retrieval, information extraction, and question answering.

Generally, there are two main problems, namely variety and ambiguity, which challenge the accuracy of biomedical entity linking systems. The variety problem refers to a named entity may appear in different surface forms in a given text. The ambiguity problem refers to two named entities with the same surface form may have different semantic meanings. The latest research tends to address the biomedical entity linking task in a unified framework that considers only the biomedical entities themselves but without their contexts, ignoring the ambiguity issue[Liu et al., 2021; Sung et al., 2020]. However, the ambiguous instances account for a high proportion of the overall entities and mentions, and failure in disambiguation may lead to misinterpretation of the entire context. For example, there are 79, 609 synonyms in Unified Medical Language System (UMLS) share the same surface but have distinct Concept Unique IDs (CUIs), which relies on extra information to distinguish the correct entities. Moreover, the mentions can be ambiguous and it is difficult to associate a mention to the correct entity based on surface-level features alone, which requires a model to have a good semantic understanding of the mention and its context.

Consider the example in Figure 1, where nodules is a mention of the entity Thyroid Nodules, and others are the top candidates retrieved from UMLS by SAPBERT [Liu et al., 2021]. In this case, although the top candidates that ranked by SAPBERT are similar to the mention in surface, they are not the golden entities. However, if the linking model has the context of mentions and the type information of entities, the golden concept may have a greater chance of being recalled, which will make the entity disambiguation more precise.

Motivated by the above observations, in this paper, we propose a two-stage biomedical entity linking algorithm to enhance the entity representations based on prompt learning. To address the variety challenge and comprehensively consider the inference speed, we utilize a coarser-grained retrieval from a representation space defined by a bi-encoder that independently embeds the mention and entity's surface forms to retrieve candidates in the first stage. Afterwards, each candidate is re-ranked with a finer-grained encoder based on prompt-tuning that utilizes the mention context and entity type information, mainly aiming to solve ambiguity and improve linking accuracy. In detail, prompt-tuning is a new paradigm that bridges the gap of objective forms between pre-training and fine-tuning. And it can further inject and stimulate the knowledge in biomedical pre-trained language models (PLMs), thus boosts the model performance. Finally, we evaluate the performance of our model using the MedMentions [Mohan and Li, 2019] and the NCBI disease [Dogan et al., 2014] datasets, and against several state-of-the-art (SOTA) baselines. From the extensive experimental results, the proposed prompt-tuning strategy is demonstrated to be effective in recalling the correct entities that missed by the coarser-grained retrieval and enhancing the representations of the bio-entities. For example, in Figure 1, the entity Thyroid Nodules can be recalled by the proposed finer-grained encoder but missed by SAPBERT. The code is available at https://github.com/TiantianZhu110/BioPRO.

To sum up, the contributions of this paper are as follows: (1) We propose a two-stage algorithm to address both the variety and the ambiguity challenges of the entity linking task; (2) For the first time, we propose a finer-grained encoder based on prompt-tuning to re-rank the candidates output by the first stage; (3) Extensive experiments show that our model achieves promising performance improvements on two public datasets, which also demonstrates the effectiveness of prompt-tuning strategy in recalling golden entities that even failed in the first stage.

2 RelatedWork

2.1 Entity Linking

Entity linking focuses on mapping an input mention from biomedical text to its associated entity in a curated KB. In the biomedical domain, UMLS is often used as the KB for entities, as it is a rich conceptual ontology differentiated by CUIs and relations between them. There are many existing entity linking tools, such as TaggerOne [Leaman and Lu, 2016], QuickUMLS [Luca and Goharian, 2016], which rely on rulebased approaches. However, such approaches struggle when there is a large discrepancy between the morphologies of mentions and concepts. Recently, the field has shifted toward machine learning methods, which can be divided into classification [Niu et al., 2019] and learning to rank [Sung et al., 2020; Liu et al., 2021; Zhu et al., 2020] categories. Among them, classification approaches have the disadvantage that the output space must be the same as the number of concepts. However, ranking methods rank the similarities between input mention and candidate entities, which would reduce the output space. BIOSYN [Sung et al., 2020] encodes input and target terms by both TF-IDF and BioBERT [Lee et al., 2020] and uses synonym marginalization to maximize similarities between synonyms. SAPBERT [Liu et al., 2021] proposes self-alignment pretraining to fine-tune BERT on synonyms extracted from the UMLS and currently holds SOTA across all major English biomedical entity linking datasets. LATTE [Zhu et al., 2020] uses an attention-based mechanism to rank candidates for a given mention based on their semantic representations along with the latent-types.

In our framework, we propose a two-stage algorithm by first generating candidates and then ranking them, which takes the context of mentions as well as the type information of entities into consideration to obtain better representations.

2.2 Prompt Learning

Recently, PLMs provide a new trend to utilize massive unlabeled corpora and have been demonstrated their effectiveness on various natural language processing (NLP) tasks [Devlin et al., 2019; Raffel et al., 2020]. However, a fine-tuning process is still needed for applying the rich lexical, syntactic, and factual knowledge of the PLMs to different downstream NLP tasks. Specifically, pre-training is usually formalized as a cloze-style task, yet conventional fine-tuning often adds extra classifiers on the top of PLMs. And by analyzing the objective forms of the pre-training and the fine-tuning processes, we can find a significant gap between them, which would restrict taking full advantage of knowledge in PLMs.

Since the emergence of GPT-3 [Brown et al., 2020], prompt-tuning has attracted widespread attention from the NLP community. In contrast with fine-tuning, prompt-tuning is a new paradigm that leverages language prompts as contexts to adapt different tasks, which is also termed in-context learning and makes downstream tasks expressed as some objectives similar to pre-training objectives. Prompting means adding instructions and demos before input and output predictions to stimulate knowledge from PLMs. Recently, handcrafted prompts have achieved promising results in knowledge probing [Davison et al., 2019], natural language inference [Schick and Sch¨utze, 2021], and so on. Furthermore, to avoid handcrafting prompts for different tasks, a series of work focus on automatic prompt search [Jiang et al., 2020; Gao et al., 2021]. Moreover, unlike these discrete prompts, continuous prompts have also been used to steer PLMs [Li and Liang, 2021].

In this paper, we aim to steer PLMs with prompt-learning to capture the contextual information of both mentions and entities. We take biomedical entity linking, a crucial task in the biomedical domain, as the foothold to design the promptlearning strategy. In our work, we mainly emphasize that using prompt-learning to perform linking task can not only improve experimental performance, but also learn deeper biomedical entity representations.

3 Method

3.1 Problem Statement

Given a biomedical document of l tokens D = {w1,w2, . . . ,wl} with a list of n mentions M = {m1,m2, . . . ,mn}, and a KB of g entities K = {e1, e2, . . . , eg}, the entity linking system will output the target entity et in the KB for each mention m. Note that each mention has its own golden CUI and each CUI has at least one entity synonym in the KB. We assume that the type information of the entities is available, which is a common setting in biomedical KB such as UMLS.

3.2 Model Architecture

Following previous work [Wu et al., 2020], we adopt a twostage approach to perform biomedical entity linking and the overall architecture is illustrated in Figure 2. In the first stage, a coarser-grained retrieval model based on the bi-encoder is used to produce the mention and entity's vector representations and further generate the top candidates. The retrieved candidates along with contextual information are then passed to the finer-grained ranking model in the second stage for reranking. In the following sections, we describe in detail the coarser-grained retrieval model and the finer-grained ranking model used in this work.

3.3 Coarser-grained Retrieval Model

Considering the inference speed, we utilize a bi-encoder architecture similar to the work of [Wu et al., 2020] to independently embed the mention and entity's surface forms in the first stage. We use pre-trained SAPBERT to encode representations and share the same SAPBERT parameters for both mentions and entities. Specifically, the surface representation of the mention m is calculated as follows:

where rsm ∈ Rdh and dh denotes the hidden dimension of SAPBERT. [CLS] denotes the special token that BERT-style models use to compute a single vector representation of the input. The entity surface representation rs e is computed similarly. Then, the score of entity candidate ei is given by the scoring function:

Finer-grained Ranking Model

[Men_Context]. [Men] [Soft] and the[Ent_Type] [Ent] [Soft]have the same meaning, is it correct? [MASK].

where we use cosine similarity for f . Note that for each mention, we compute S(m, ei) for all entities ei ∈ K and return the top k nearest candidates. In addition, the entity representations can be cached at the inference time to reduce computational cost.

3.4 Finer-grained RankingModel

The input of the finer-grained ranking model is the mention with context and a set of top k candidate entities C = {c1, c2, . . . , ck} with type information, the model computes a relevance score for each entity in C. In this stage, we apply a prompt learning method based on PubMedBERT to perform the task, which transforms the original task into a masked language modeling problem. By incorporating the mention context and entity type information into the prompt template, the proposed finer-grained ranking model can reduce the semantic drift problem and achieve better performance.

Generally, a prompt consists of a template and a verbalizer. Among them, the template T(·) wraps the input sequence x into T(x) and there is at least one [MASK] in T(x) for the language model M to fill the label words. Moreover, for each label y ∈ Y, we define a corresponding label word set Vy, which is a subset of the vocabulary V of the M. And the verbalizer will map the label word in Vy to the label in Y. Thus, in prompt learning, the task can be transformed into a masked language modeling problem:

Templates

In this work, we construct two mention-entity prompts for the entity linking task. Specifically, one is a hard-encoding template and the other is a mixed template with both hardand soft-encoding.

In hard-encoding template, we fuse the labels into answers via question answering. Given the mention with context and the entity with semantic type, the template is constructed as: Thard =[Men Context]. [Men] and the [Ent Type] [Ent] have the same meaning, is it correct? [MASK].

Among them, [Men] and [Ent] denote the mention and entity names, respectively. [Men Context] denotes the context of the mention and here we take the shortest sentence containing the mention. [Ent Type] is the semantic type of the entity and we use "unknown type" for unknown entity types. The predicted words in the position [MASK] will be mapped to the final label. Here, we set "yes" as the positive label word and "no" as the negative label word. Thus, the probability of the label word w can fill the masked position, which is calculated as follows:

where vw is the embedding of the token w in the PLMM and h[MASK] is the hidden vector of [MASK].

In order to prevent the long-distance forgetting problem, we also add the trainable soft tokens to enhance the influence of the mention and the entity on the [MASK] position in the mixed template. The mixed template is constructed as follows: Tmixed =[Men Context]. [Men] [Soft] and the [Ent Type] [Ent] [Soft] have the same meaning, is it correct? [MASK]. Among them, [Soft] denotes the soft token, which is randomly initialized as a word in the PLM dictionary. And the soft token embeddings are free parameters which are initialized as a trainable matrix Pθ (parametrized by θ). Thus, the input embeddings of PLM can be calculated as follows:

where Softidx denotes soft token indices and Pθ ∈ R|Softidx|×d˜h. |Softidx| denotes the length of the soft tokens and d˜h denotes the hidden dimension of PLM.

Training

For training the ranking model, golden entities are selected as the positive samples, which may have multiple semantic types. We construct multiple positive samples by looping on different types of entities. Negative samples come from unmatched candidates generated by the SAPBERT model as well as negative instances randomly sampled from the full dataset. The loss function of the finer-grained ranking model is defined as follows:

where ω is the label weight to balance the positive and the negative instances, w+ denotes the label word "yes" and wdenotes "no".

Prompt for Candidates Retrieval

Since the vector representations produced by the coarsergrained model does not consider the extra information of mentions and entities, we use the trained finer-grained ranking model to generate deep representations. Furthermore, the candidates is retrieved by scoring both surface and deep representations in the first stage of our experiments. When generating deep representations, in order to maintain consistency between training and inference, we utilize the similar data format as prompt-tuning as the input. Specifically, when generating the entity deep representations, we replace the [Men Context] and [Men] with [Ent Type] and [Ent] in the mixed template, respectively. Then we extract the average representations of the corresponding entity tokens before the second soft token as the entity deep representations. Similarly, we can replace the [Ent Type] and [Ent] with "unknown type" and [Men] in the template and get the mention deep representations in the same way.

4 Experiments

4.1 Datasets

We used two datasets to evaluate our proposed model: MedMentions [Mohan and Li, 2019] and NCBI Disease [Dogan et al., 2014] datasets. MedMentions is the largest biomedical entity linking dataset consisting of 4, 392 abstracts from PubMed, with over 350, 000 mentions linked to the 2017AA full version of UMLS concepts. The size of the concept set in UMLS is 3, 415, 665, and the size of all surface forms / synonyms set is 14, 815, 318. And the concepts in UMLS belongs to 127 semantic types. NCBI Disease corpus contains 793 manually annotated PubMed abstracts and 6, 881 mentions with each CUI mapped to the July 6, 2012 version of MEDIC dictionary [Davis et al., 2012]. The size of the concept set in MEDIC dictionary is 11, 915, and the size of all surface forms / synonyms set is 71, 923. Descriptions of the datasets and their statistics are provided in Table 1.

4.2 EvaluationMetrics

Since picking the correct target entity among candidates is a ranking problem, we use the top k accuracy as an evaluation metric following the previous work [Sung et al., 2020; Liu et al., 2021]. For all datasets, we report Acc@1 and Acc@5 for evaluating performance of the proposed model against the baseline models, where Acc@k represents the accuracy when only k entities are retrieved.

4.3 Implementation Details

Table 2 lists the hyper-parameter search space for obtaining the set of used numbers. Note that all parameters are selected based on the best performance on the development set.

4.4 Baselines

For the evaluation of the proposed model, we use the following SOTA baseline methods for comparison.

• Sieve-based [D'Souza and Ng, 2015] is a multi-pass sieve approach to the under-studied task of normalizing disorder mentions in the biomedical domain.
• TaggerOne [Leaman and Lu, 2016] is a machine learning-based system that jointly performs disease NER and normalization by utilizing semi-Markov models.
• NormCo [Wright et al., 2019] is a deep coherence model which considers the semantics of an entity mention, as well as the topical coherence of the mentions within a single document.
• BNE [Phan et al., 2019] proposes a new framework for learning robust representations of biomedical terms by learning to encode names of the same concepts into similar representations, while preserving their conceptual and contextual meanings.
• BERTRANK [Ji et al., 2020] applies and evaluates the pre-trained BERT / BioBERT / ClinicalBERT models for candidate concept ranking by transforming the ranking task as a sentence-pair classification task, which is a point-wise learning to rank method.
• BIOSYN [Sung et al., 2020] utilizes the synonym marginalization technique and the iterative candidate retrieval for learning biomedical entity representations.
• SAPBERT [Liu et al., 2021] is a self-alignment pretraining schema for learning biomedical entity representations.

4.5 Experimental Results

Main Results

To evaluate our approach, we compared the proposed model with several SOTA models on the two datasets, and listed the performance of each model in Table 4. As shown in Table 4, our model outperformed all other models with a large margin, which improved 2.3 and 13.4 Acc@1 points over the baselines on NCBI Disease and MedMentions test sets, respectively. The improvement of Acc@1 is significantly greater than that of Acc@5, which demonstrates the effectiveness of prompt-tuning based strategy in improving the precision of the entity linking task. In addition, we list the performance comparison of the proposed model and its variants. As shown in Table 4, mixed template is better than hard template in prompt-tuning and the model that integrates the deep representations has the best performance. Note that [Mohan and Li, 2019] first reported that training TaggerOne on a subset of MedMentions requires > 900 GB of RAM due to the large search space. In comparison, our method can avoid the outof-memory problem on the full-set of MedMentions.

Resolution of the Ambiguity Problem

In order to study the resolution of different models on the ambiguity problem, we compared the top 1 candidates prediction's performance of our proposed model with SAPBERT on the MedMentions dataset. As mentioned before, the ambiguous instances account for a high proportion of the overall entities and mentions, increasing the challenge of the task. Table 3 lists the number of wrongly predicted top 1 candidates, the number of ambiguous examples in wrongly predicted top 1 candidates and the proportion of ambiguous instances in the overall top 1 errors. It is shown that our model can greatly reduce the proportion of ambiguous examples (↓ 26.6%) among the total misclassified examples, which demonstrates the effectiveness of the proposed prompt-tuning strategy in addressing the ambiguity challenge and understanding deep representations of mentions and entities.

Performance of the Coarser-grained Model

Although we measure the overall performance of our twostage entity linking algorithm in Table 4, the ranking results of the second step are based on the recall candidates output by the first step. Since the recall of the coarser-grained model will limit the performance of the finer-grained encoder, we further report the recall value of different candidate retrieval methods, namely SAPBERT and SAPBERT fusion prompt representations, in Table 6. As can be seen from Table 6, after incorporating the prompt representations in the first step, the recall rate has improved greatly. This may be due to the fact that the SAPBERT model tends to recall entities with similar surface forms, and misses instances of dissimilar forms but semantically matching. In comparison, prompt-tuning based representations preserve the contextual information well and are complementary to the SAPBERT representations, which can promote the recall improvement of the coarser-grained model.

Representations Analysis

From Table 6, we observe that the prompt-tuning based representations can improve the recall value of the coarser-grained model, and assume that the prompt representations may be complementary to the SAPBERT representations. To further study the differences between the prompt representations and the SAPBERT representations, we reported their Euclidean distance and cosine similarity to PubMedBERT's vector space representations in Table 7, respectively. Both SAP- BERT and our finer-grained encoder utilize PubMedBERT as a base model for further training, yet their tuning strategies differ. As shown in Table 7, our prompt-tuning based finergrained encoder is more similar to PubMedBERT in terms of vector space representation compared to fine-tuning based SAPBERT, which may better preserve the original biomedical knowledge of pre-trained language model and is suitable for modeling deeper contextual information.

Case Study

To illustrate the effectiveness of the prompt-tuning based finer-grained encoder, we listed the top 1 candidates retrieved by SAPBERT and our prompt encoder in Table 5, respectively. We can infer from Table 5 that SAPBERT prefers entities with similar surface form as the mentions, but may miss real matching entities. However, it clearly indicates that the proposed prompt encoder can make up for this disadvantage and make comprehensive judgments by leveraging the contextual information of entities and mentions, which will further improve the accuracy of the coarser-grained model.

5 Conclusion

In this paper, we focus on both the variety and the ambiguity challenges in biomedical entity linking. We propose a two-stage algorithm, including a coarser-grained retrieval and a finer-grained encoder based on prompt-tuning. Empirical results demonstrate that the prompt-tuning strategy of our model could considerably improve the recall of the coarsergrained model and meanwhile gain high performance.

Acknowledgments

This work is supported by the National Natural Science Foundation of China (Grant No. 62006061, 61872113, 62106115), Stable Support Program for Higher Education Institutions of Shenzhen (No.GXWD2020123015542700320200824155011001), Strategic Emerging Industry Development Special Funds of Shenzhen (No. XMHT20190108009), and Fundamental Research Fund of Shenzhen (No. JCYJ20190806112210067).