Generative Biomedical Entity Linking via Knowledge Base-Guided Pre-training and Synonyms-Aware Fine-tuning
Abstract
Entities lie in the heart of biomedical natural language understanding, and the biomedical entity linking (EL) task remains challenging due to the ﬁne-grained and diversiform concept names.
Generative methods achieve remarkable performances in general domain
EL with less memory usage while requiring expensive pre-training.
Previous biomedical
EL methods leverage synonyms from knowledge bases (KB) which is not trivial to inject into a generative method.
In this work, we use a generative approach to model biomedical EL and propose to inject synonyms knowledge in it.
We propose KB-guided pre-training by constructing synthetic samples with synonyms and deﬁnitions from KB and require the model to recover concept names.
We also propose synonyms-aware ﬁne-tuning to select concept names for training, and propose decoder prompt and multi-synonyms constrained preﬁx tree for inference.
Our method achieves state-of-the-art results on several biomedical EL tasks without candidate selection which displays the effectiveness of proposed pre-training and ﬁne-tuning strategies.
The source code is available at https:// github.com/
Introduction
Biomedical entity linking (EL) refers to mapping a biomedical mention (i e., entity) in free texts to its concept in a biomedical knowledge base (KB)
(e.g., UMLS (Bodenreider, 2004)).
This is one of the most concerned tasks of research in medical natural language processing which is highly related to high-throughput phenotyping (Yu et al., 2015), relation extraction (Li et al., 2016), and automatic diagnosis (Yuan and Yu, 2021).
Recent methods in biomedical EL mainly used neural networks to encode mentions and each concept name into the same dense space, then linked mentions to corresponding concepts depending on embedding similarities (Sung et al., 2020; Lai et al., 2021; Bhowmik et al., 2021; Ujiie et al., 2021; Agarwal et al., 2021).
Synonyms knowledge has been injected into these similarity-based methods by contrastive learning (Liu et al., 2021a; Yuan et al., 2022).
For example in UMLS, concept C0085435 has synonyms: Reiter syndrome , Reactive arthritis and ReA which help models to learn different names of a concept entity.
However, similarity-based methods requires large memory footprints to store representation for each concept which are hard to deploy.
In general domain, GENRE (Cao et al., 2021a) viewed EL as a seq2seq task which inputs mentions with contexts and outputs concept names tokenby-token.
Mentions, contexts and concepts can be mingled due to the power of transformers, and the model does not need to store representations for each concept during inference.
GENRE pretrained on Wikipedia EL datasets to boost performances.
However, directly implementing GENRE on biomedical EL cannot harvest satisfying results.
The gap occurs in two aspects: (1) There are no such large-scale human-labeled biomedical EL datasets for pre-training.
(2) Biomedical concepts may have multiple synonyms. we ﬁnd selecting synonyms for training is sensitive to the results, We ﬁnd the results are sensitive to the synonyms selection for training, and simply using a 1-tomapping between names and concepts as Cao et al.
(2021a,b) may hurt performances.
To address the above issues, we propose KBguided pre-training and synonyms-aware ﬁnetuning to improve generative EL.
For pre-training , we construct pre-training samples using synonyms and deﬁnitions collected from KBs and sentence templates. KB-guided pre
-training has the same format as seq2seq EL, which ﬁlls the gap of loss of pre-training corpus.
Compared to the method introduced in Cao et al.
(2021a), ours performs better in biomedical EL with fewer resources.
For ﬁnetuning , we propose decoder prompts to highlight mentions.
We ﬁnd the model tends to generate textually similar names to mentions.
Hence textual similar criterion is proposed for selecting concept names during ﬁne-tuning.
During inference, we propose a multi-synonyms constrained preﬁx tree, which results in signiﬁcantly improved performance.
The overview of our approach is illustrated in Figure 1.
We conduct experiments on various biomedical EL datasets and achieve SOTA results on
COMETA, BC5CDR, and AskAPatient (AAP) even without candidate selection.
Extensive studies show the effectiveness of our proposed pre-training and ﬁne-tuning schemes.
Approach
Deﬁne a set of concepts E as target concepts (i.e. concepts from target KBs).
For each concept e ∈ E , we have a set of synonyms names f ( e ) =
{ s ie | i ∈ { 1 , . . . , n e }} .
All the synonyms forms a name set S =
(cid:83) e ∈E { s ie | i ∈ { 1 , . .
n e }} .
Names-toconcept mappings can be deﬁned by: σ ( s ie ) = e where σ = f − 1 .
For mention m with left and right contexts c l and c r which gold label is e m ∈ E , we need to ﬁnd the target concept ˆ e m ∈ E . m , c and s comprise a sequence of tokens.
2. Seq2seq EL
Our model applies an encoder-decoder transformer architecture following Cao et al.
(2021a).
The encoder input is: [BOS] c l [ST] m [ET] c r [EOS] , where [ST] and [ET] are the special tokens marking m 1 .
For the decoder side, unlike GENRE decoding target names directly, we use simple preﬁx prompts P m = <
m is > to strengthen the interaction between mentions and make the decoder side output resemble a natural language sentence: [BOS] m is s, where s is a target name belong to label concept e .
The training objective of
Seq2Seq EL is to maximize the likelihood:
[Equation]
where N s is the number of tokens of s and y i indicates the i th token.
The inference of Seq2seq EL applies beam search (Sutskever et al., 2014) with targets constrained to the name set S by a preﬁx tree (constructed by name set S ).
Unlike mGENRE Cao et al. (2021b) using provided candidates to decrease the size of the preﬁx tree, we use the whole name set S instead.
2. KB-Guided Pre-training
As the training data of EL is tiny compared to the vast number of concepts in KB, thus it makes
EL for some mentions zero-shot problems.
We want to leverage synonyms knowledge from KB to enhance the model’s performance.
Injecting synonyms knowledge to the encoder-only models can be done by contrastive learning (Liu et al., 2021a;
Yuan et al., 2022).
However, such a paradigm cannot directly apply to encoder-decoder architecture as entities are not represented by dense embeddings.
To mitigate this problem, we construct a pre-training task that shares a similar form as Section 2.1.
We manually deﬁne a set of clause templates to splice with synonyms and deﬁnitions in KB to form input synthetic language discourses.
Concretely, we select two synonyms s ae and s be and deﬁnition c e of a concept e ∈ E .
Then we randomly pick a template to concatenate them to form the encoder input, here we give two examples: [BOS][ST] s ae [ET] is deﬁned as c e [EOS] [BOS] c e describes [ST] s ae [ET][EOS]
For the decoder: [BOS] s ae is s be .
c e is the simulated context and s be is for model to predict.
If deﬁnitions are absent in KB, we will use other synonyms to construct c e .
All templates we used can be found in Appendix E.2.
2. Synonyms-Aware Fine-tuning
We propose and validate by experiments in Section 4 that seq2seq EL is profoundly inﬂuenced by the textual similarity between mentions and concept names.
It tends to generate textually similar names.
We select the target name by calculating the character 3-gram TF-IDF similarity (Neumann et al., 2019) between mention m and all synonyms { s ie } and choosing the most similar one as
[Equation]
By the textual similarity criterion, we manually reduce the difﬁculty of ﬁne-tuning.
We do not use this criterion for pre-training since we want it to learn various synonyms to improve generalization.
[Table]
Different from GENRE using only one canonical name for each concept for preﬁx tree, we use multisynonym names (i.e. S ) to construct preﬁx tree.
During inference, we apply preﬁx tree constrained beam search to decode the name ˆ s m , and map to the concept ˆ e m = σ (ˆ s m ) via N-to-1 names-to-concept mapping σ .
Experiments 3. Datasets and KBs Pre-training
We use a subset of UMLS st21pv (Mohan and Li, 2019) as E for pre-training.
It contains 2.37M concepts, where 160K concepts contain deﬁnitions and 1.11M concepts have multiple synonyms.
While pre-training, we iterate concepts and synonyms to construct inputs and outputs
. Fine-tuning
We evaluate our model on BC5CDR (Li et al., 2016), NCBI (Do˘gan et al., 2014),
COMETA (Basaldella et al., 2020) and AAP (Limsopatham and Collier, 2016).
These benchmarks focus on different entity types, including disease, chemicals, and colloquial terms.
For fair comparison, we follow Varma et al.
(2021) in BC5CDR,
Lai et al.
(2021) in NCBI and COMETA, and Limsopatham and Collier (2016) in AAP to construct dataset splits and target KB concepts E .
Name set
S is constructed by synonyms from UMLS and [Table] original KB which is detailed in Appendix E.1.
Datasets summaries are shown in Appendix A.
We use recall @1 / @5 as metrics for performance illustration, and recall @5 results are listed in Appendix C.
3. Implementation Details
We use BART-large (Lewis et al., 2020) as the model backbone.
We pre-train and ﬁne-tune our model using teacher forcing (Williams and Zipser, 1989) and label smoothing (Szegedy et al., 2016) which are standard in seq2seq training.
The hyperparameters can be found in Appendix E.3.
3. Main Results
Table 1 and 2 compare the recall @1 with stateof-the-art(SOTA) methods.
Our method with KBguided pre-training exceeds the previous SOTA on
BC5CDR, COMETA and AAP, which is up to 1. on BC5CDR, 1.3 on COMETA and 0.3 on AAP.
Our method also shows superiority over previous SOTA on BC5CDR and COMETA without pretraining.
Besides, our method shows competitive results on NCBI compared to SOTA, and we further analyse the results of NCBI through case studies in Appendix D.
Discussion
Does pre-training help?
On all datasets, KBguided pre-training improves ﬁne-tuning consis- [Table] tently, which is 0.7 on BC5CDR, 0.3 on NCBI, 0.7 on COMETA, and 0.5 on AAP.
To better understand KB-guided pre-training, we conduct ablation studies.
We compare different pre-trained models without ﬁne-tuning in Table 3.
BART fails to link mentions due to the mismatch of the pretraining task.
GENRE has been pre-trained on the large-scale BLINK dataset (Wu et al., 2020), and it obtains a decent ability to disambiguate biomedical mentions.
Our pre-trained model shows improvement on BC5CDR and COMETA compared to GENRE with fewer pre-training resources (6 GPU days vs. 32 GPU days).
We then conduct synonyms-aware ﬁne-tuning on different pretrained models in Table 3.
Our pre-trained model outperforms BART (+0.8 on BC5CDR, +0.5 on
COMETA) and GENRE (+0.4 on BC5CDR, +0. on COMETA) which proves the effectiveness of pre-training.
Selection of Names
We ablate the selection of target names s for ﬁne-tuning: (a) proposed TFIDF similarity; (b) the shortest name in a concept; (c) randomly sampled name in a concept.
We also compare how to construct S for inference: (i) using all synonyms from UMLS and target KB; (ii) using the shortest name for a concept; (iii) using randomly sampled name for a concept.
We note (i) establish an N-to-1 mapping from synonym name to concept, while (ii) and (iii) establish a 1-tomapping.
We conduct experiments on available combinations.
From Table 4, we conclude that (1) N-to-1 mapping performs better than 1-to-1 mappings during inference, which means synonyms can boost performances.
Using one synonym like GENRE degrades performances.
(2) Textual similarity criterion performs better than shortest or sampled names when training.
We also check the accuracies of different TF-IDF [Table] similarity sample groups on COMETA trained with using all synonyms as S and TF-IDF for selecting target names s .
From Figure 2, we ﬁnd the distribution of TF-IDF similarity between mentions and selected names is polarized, and the accuracy increases along with textual similarities which prove textually similar targets are easy to generate.
This phenomenon validates the advantage of selecting textually similar names for ﬁne-tuning.
Decoder Prompting
One difference between
GENRE and ours is using prompt tokens on the decoder side.
Prompting has shown improvement on various NLP tasks (Liu et al., 2021b).
Here, prompt tokens serve as informative hints by providing additional decoder attention queries and making the outputs resemble language models’ pre-training tasks.
We test dropping the prompt tokens, and Tableshows degraded performances (-0.3 on BC5CDR and -0.5 on COMETA).
The results illustrate the improvement brought by decoder prompting.
Sub-population Analysis
We list several subpopulations of the BC5CDR benchmark to illustrate the model’s performance on different ﬁnegrained categories of mentions.
The details of subpopulations are shown in Appendix C. Our model’s performance on different subpopulations of BC5CDR is shown in Table 9. Through the results, we have several ﬁndings: Compared with Varma et al.
(2021), our method shows superiority over most of the sub-populations of BC5CDR.
Our method without pre-training outperforms the dataaugmented version of Varma et al.
Our method surpasses Varma et al.
(2021) by the largest margin on Unseen Concepts .
One possible explanation is that our generative method learns the linkage between mentions and contextual information better, thus gaining superior zero-shot performance.
[Table]
[Figure]
Our KB-guided pre-training gains improvement universally on most subsets.
This reﬂects the effectiveness of the pre-training task.
Popular concepts and single-word mentions are more easily resolved compared to unseen mentions or concepts and multi-word mentions, respectively.
Unsurprisingly, mentions with less training resources and longer textual forms are more difﬁcult.
The lengths of mentions may challenge the generative model.
Related Work
Biomedical EL is an important task in biomedical NLP.
Classiﬁcation-based methods used a softmax layer for classiﬁcation (Limsopatham and Collier, 2016;
Miftahutdinov and Tutubalina, 2019) which consider concepts as category factors and lost information of concept names.
Recent methods (Sung et al., 2020; Liu et al., 2021a; Lai et al., 2021;
Bhowmik et al., 2021; Ujiie et al., 2021; Agarwal et al., 2021; Yuan et al., 2022) encoded mentions and names into a common space and disambiguated mentions by nearest neighbors.
Angell et al.
(2021) and Varma et al.
(2021) adopted a retrieve-andrerank framework to boost performances.
Varma et al.
(2021) emphasized the lack of training samples in EL and augmented data using Wikipedia and PubMed, while our pre-training corpus constructed by KB and templates can serve as good supplementary training data.
Generative EL
Cao et al.
(2021a) proposed to view EL as a seq2seq problem that got rid of hard negative sampling during training and required less memory at inference.
The shortage is it demanded vast training sources (11 GB training data and GPU for 24 hours) to achieve competitive performance.
Cao et al.
(2021b) explored dealing synonyms in multilingual generative EL by adding language identiﬁers that cannot be directly implemented in biomedical EL.
Conclusion
To the best of our knowledge, our work is the ﬁrst to explore generative EL in the biomedical domain.
We inject synonyms and deﬁnition knowledge into the generative language model by KGguided pre-training.
We emphasize the synonym selection issue and propose synonyms-aware ﬁnetuning by considering the textual similarity.
Decoding prompts are also introduced to improve the model’s performance.
Our model sets new stateof-the-art on different biomedical EL benchmarks.
GENRE shows that well-selected candidate sets can improve seq2seq EL, and we believe this will further boost our performances. Acknowledgements
We thank Shengxuan Luo and Chuanqi Tan for fruitful discussions and advises, and Xixi Mo for drawing the overview ﬁgure.
We appreciate the anonymous reviewers for their helpful comments and suggestions.
This work was supported by the National Natural Science Foundation of China (Grant No. 12171270), the Natural Science Foundation of Beijing Municipality (Grant No. Z190024), and the International Digital Economy Academy.
A Dataset Summary and Statistics
We pre-process the datasets by the following procedures: (1) the abbreviations in the texts are expanded using AB3P (Sohn et al., 2008); (2) the texts are lower-cased, and the beginning and ending of a mention are marked by two special tokens [ST] and [ET] ; (3) the overlapping mentions and mentions absent from the target KB are discarded.
BC5CDR (Li et al., 2016) is a benchmark for biomedical entity recognition and disambiguation.
The dataset annotates 1500 PubMed article abstracts with 4409 chemicals, 5818 diseases entities, and 3116 chemical-disease interactions.
All the annotated entities are mapped to MeSH ontology, which is a smaller medical vocabulary that comprises a subset of UMLS (Bodenreider, 2004).
In this work, in consideration of fairness, we follow two most recent works (Angell et al., 2021; Varma et al., 2021) that use MeSH contained in UMLS 2017 AA release to construct the target knowledge base.
NCBI (Do˘gan et al., 2014) contains a corpus of 793 PubMed abstracts.
It consists of 6892 annotated disease mentions of 790 unique disease concepts.
All the mentions are labeled with concepts in MELIC ontology (Davis et al., 2012).
MELIC is a medical dictionary that merges the diseases concepts, synonyms, and deﬁnitions in MeSH and
OMIM and is composed of 9700 unique diseases.
In our work, we used the processed data and the target ontology provided by BioSyn (Sung et al., 2020) and ResCNN (Lai et al., 2021).
We followed their works to construct our training, developing, and testing data.
(Basaldella et al., 2020) consists of 20k English biomedical entity mentions from publicly available and anonymous health discussions on Reddit.
All the mentions are expert-annotated with concepts from SNOMEL CT.
We use the “stratiﬁed (general)” split and follow the evaluation protocol of SapBert (Liu et al., 2021a) and
ResCNN (Lai et al., 2021).
AskAPatient
(Limsopatham and Collier, 2016) is a dataset containing 8,662 phrases of social media language.
Each phrase can be mapped to one of the 1,036 medical concepts from SNOMEL-CT and AMT (the Australian Medicines Terminology).
The samples in AskAPatient do not include contextual information and mentions can only be disambiguated by phrases per se.
We follow the experimental settings from works of Sung et al.
(2020) and Limsopatham and Collier (2016)
and apply the 10-fold evaluation protocol.
Statistics of above-mentioned datasets are listed in Table 6. [Table] [Table]
B License and Availability of Resources
BC5CDR, NCBI, COMETA, and AskAPatients are all publicly available datasets on the Internet.
Their target KBs MeSH, MELIC, and SNOMEL
CT are covered by UMLS Metathesaurus License.
One can require such a license by signing up for a UMLS terminology services account to access
KBs mentioned above.
C Additional Experiment Results C. Recall@5 Results
We show the Recall@5 result of our method with and without pre-training in Table 7.
C. Sub-population Analysis
Following Varma et al.
(2021), we split test samples into different sub-populations.
The details of different sub-population categories we use is shown in Table 8.
D Case Study
We provide case studies on the NCBI-disease benchmark to give an insight and justiﬁcation of the performance of our method.
For the case of mention colorectal adenomas which is annotated with
D018256 adenomatous polyp , the mention exists for 6 times in the test set (account for 0.63 score of Recall@1).
Our model fails to correctly disambiguate all such mentions while linking the mention to other concepts D000236 colorectal adenomas or D003123 hereditary nonpolyposis colorectal cancer .
In the training set, the mention colorectal adenomas exists for two times and is annotated with D003132 and D000236 respectively.
Thus, through this case, we can see (1) our model learns the information contained in the training set; (2) such inconsistent test samples are hard to disambiguate correctly.
E Implementation Details
In this section, we provide more details of our experiments.
E. Knowledge Base Pre-processing
Given different knowledge bases, we pre-process their content by the following procedure:
For each concept, we include all its synonyms from the original target KB.
We also expand the synonym set for each concept using UMLS.
We use the 2017 AA Active Release of UMLS.
For synonyms in the expanded name set S , we lowercase them and remove the symbols (e.g., dash line - or comma ,).
There may exist a name as a synonym for multiple concepts.
We de-duplicate these overlapped synonyms by removing the synonym from the concept with more other synonyms to avoid the unbalanced number of synonyms in each concept.
It is worth noticing that we do not de-duplicate the synonyms in the target KB of NCBI in consideration of comparison fairness.
In the previous works, a mention link to multiple concepts, and correct disambiguation is claimed if the target concept is hit by one of the predicted concepts in NCBI.
E. Pre-training Clause Templates
We list pre-training clause templates we used in
Table 10.
For those concepts containing onlysynonyms, s ae and s be are the two synonyms respectively and c e is the same as s be .
For those concepts containing only 1 sole synonym, s ae , s be and c e are the same. [Table] [Table]
[Table]
same concept as
[Table] s ae [Table] has synonyms such
[Table]
E. Experiment Parameters
Our model contains 406M parameters with 12layer transformer encoders and 12-layer transformer decoders.
We list the hyper-parameters of our model for KB-guided pre-training and synonyms-aware ﬁne-tuning on different benchmarks in Table 11.
For pre-training, we heuristically select our parameters.
For ﬁne-tuning, we tune training steps among { 20000 , 30000 , 40000 } on development set.
For BC5CDR and COMETA, we use learning rate as 1 e − 5 and warmup steps as 500.
For AskAPatient and NCBI, we search learning rate among { 5 e − 6 , 8 e − 7 , 3 e − 7 } , and do not use warmup.
We only evaluate our model at the end of training.
For each benchmark, we run three times to calculate means and standard deviations.
E. Computational Resource
For our KB-guided pre-training, we implement our model on 6 A100 GPU with 40 GB memory with the help of DeepSpeed ZeRO 2 (Rajbhandari et al., 2020) and train for 1 day.
For ﬁne-tuning on different benchmarks, we implement our model on
A100 GPU.
[Table]