Self-Alignment Pretraining for Biomedical Entity Representations (2021)

Abstract: Despite the widespread success of self-supervised learning via masked language models (MLM), accurately capturing fine-grained semantic relationships in the biomedical domain remains a challenge. This is of paramount importance for entity-level tasks such as entity linking where the ability to model entity relations (especially synonymy) is pivotal. To address this challenge, we propose SapBERT, a pretraining scheme that self-aligns the representation space of biomedical entities. We design a scalable metric learning framework that can leverage UMLS, a massive collection of biomedical ontologies with 4M+ concepts. In contrast with previous pipeline-based hybrid systems, SapBERT offers an elegant one-model-for-all solution to the problem of medical entity linking (MEL), achieving a new state-of-the-art (SOTA) on six MEL benchmarking datasets. In the scientific domain, we achieve SOTA even without task-specific supervision. With substantial improvement over various domain-specific pretrained MLMs such as BioBERT, SciBERTand and PubMedBERT, our pretraining scheme proves to be both effective and robust.

## introduction biomedical entity foot_0 representation is the foundation for a plethora of text mining systems in the medical domain, facilitating applications such as literature search @xcite , clinical decision making @xcite and relational knowledge discovery (e.g. chemical-disease, drug-drug and protein-protein relations, @xcite . the heterogeneous naming of biomedical concepts poses a major challenge to representation learning. for instance, the medication hydroxychloroquine is often referred to as oxichlorochine (alternative name), hcq (in social media) and plaquenil (brand name). mel addresses this problem by framing it as a task of mapping entity mentions to unified concepts in a medical knowledge graph. 3 the main bottleneck of mel is the quality of the entity representations @xcite . prior works in this domain have adopted very sophisticated text pre-processing heuristics (d'souza and ng, 2015; @xcite @xcite which can hardly cover all the variations of biomedical names. in parallel, self-supervised learning has shown tremendous success in nlp via leveraging the masked language modelling (mlm) objective to learn semantics from distributional representations @xcite . domain-specific pretraining on biomedical corpora (e.g. biobert, @xcite ) have made much progress in biomedical text mining tasks. nonetheless, representing medical entities with the existing sota pretrained mlms (e.g. pubmedbert, @xcite ) as suggested in fig. 1 (left) does not lead to a well-separated representation space. to address the aforementioned issue, we propose to pretrain a transformer-based language model on the biomedical knowledge graph of umls @xcite , the largest interlingua of biomedical ontologies. umls contains a comprehensive collection of biomedical synonyms in various forms (umls 2020aa has 4m+ concepts and 10m+ synonyms which stem from over 150 controlled vocabularies including mesh, snomed ct, rxnorm, gene ontology and omim). 4 we design a selfalignment objective that clusters synonyms of the same concept. to cope with the immense size of umls, we sample hard training pairs from the knowledge base and use a scalable metric learning loss. we name our model as self-aligning pretrained bert (sapbert). being both simple and powerful, sapbert obtains new sota performances across all six mel benchmark datasets. in contrast with the current systems which adopt complex pipelines and hybrid components @xcite @xcite , sapbert applies a much simpler training procedure without requiring any pre-or post-processing steps. at test time, a simple nearest neighbour's search is sufficient for making a prediction. when compared with other domain-specific pretrained language models (e.g. biobert and scibert), sapbert also brings substantial improvement by up to 20% on accuracy across all tasks. the effectiveness of the pretraining in sap-bert is especially highlighted in the scientific language domain where sapbert outperforms previous sota even without fine-tuning on any mel datasets. we also provide insights on pretraining's impact across domains and explore pretraining with fewer model parameters by using a recently introduced adapter module in our training scheme. ## method: self-alignment pretraining we design a metric learning framework that learns to self-align synonymous biomedical entities. the framework can be used as both pretraining on umls, and fine-tuning on task-specific datasets. we use an existing bert model as our starting point. in the following, we introduce the key components of our framework. formal definition. let (x, y) ∈ x × y denote a tuple of a name and its categorical label. for the self-alignment pretraining step, x × y is the set of all (name, cui foot_3 ) pairs in umls, e.g. (remdesivir, c4726677); while for the finetuning step, it is formed as an entity mention and its corresponding mapping from the ontology, e.g. (scratchy throat, 102618009). given any pair of tuples (x i , y i ), (x j , y j ) ∈ x × y, the goal of the self-alignment is to learn a function f (•; θ) : @xmath0. then, the similarity f (x i ), f (x j ) (in this work we use cosine similarity) can be used to estimate the resemblance of x i and x j (i.e., high if x i , x j are synonyms and low otherwise). we model f by a bert model with its output [cls] token regarded as the representation of the input. 6 during the learning, a sampling procedure selects the informative pairs of training samples and uses them in the pairwise metric learning loss function (introduced shortly). online hard pairs mining. we use an online hard triplet mining condition to find the most informative training examples (i.e. hard positive/negative pairs) within a mini-batch for efficient training, fig. 2 . for biomedical entities, this step can be particularly useful as most examples can be easily classified while a small set of very hard ones cause the most challenge to representation learning. 7 we start from constructing all possible triplets for all names within the mini-batch where each triplet is in the form of (x a , x p , x n ). here x a is called anchor, an arbitrary name in the minibatch; x p a positive match of x a (i.e. y @xmath1.e. y @xmath2. among the constructed triplets, we select out all triplets that violate the following condition: ## conclusion we present sapbert, a self-alignment pretraining scheme for learning biomedical entity representations. we highlight the consistent performance boost achieved by sapbert, obtaining new sota in all six widely used mel benchmarking datasets. strikingly, without any fine-tuning on task-specific labelled data, sapbert already outperforms the previous supervised sota (sophisticated hybrid entity linking systems) on multiple datasets in the scientific language domain. our work opens new avenues to explore for general domain self-alignment (e.g. by leveraging knowledge graphs such as dbpedia). we plan to incorporate other types of relations (i.e., hypernymy and hyponymy) and extend our model to sentence-level representation learning. in particular, our ongoing work using a combination of sapbert and adapter is a promising direction for tackling sentence-level tasks. acknowledge funding from health data research uk as part of the national text analytics project. florian schroff, dmitry kalenichenko, and james philbin. 2015. facenet: a unified embedding for face recognition and clustering. in proceedings of the ieee conference on computer vision and pattern recognition, pages 815-823. elliot schumacher, andriy mulyar, and mark dredze. 2020. clinical concept linking with contextualized neural representations. in proceedings of the 58th annual meeting of the association for computational linguistics, pages 8585-8592. hoo-chang shin, yang zhang, evelina bakhturina, raul puri, mostofa patwary, mohammad shoeybi, and raghav mani. 2020. biomegatron: larger biomedical domain language model. in proceedings of the 2020 conference on empirical methods in natural language processing (emnlp), pages 4700-4706, online. association for computational linguistics. yifan sun, changmao cheng, yuhan zhang, chi zhang, liang zheng, zhongdao wang, and yichen wei. 2020. circle loss: a unified perspective of pair similarity optimization. in proceedings of the ieee/cvf conference on computer vision and pattern recognition, pages 6398-6407. mujeen sung, hwisang jeon, jinhyuk lee, and jaewoo kang. 2020. biomedical entity representations with synonym marginalization. in proceedings of the 58th annual meeting of the association for computational linguistics (acl), pages 3641-3650, online. association for computational linguistics. elena tutubalina, artur kadurin, and zulfat miftahutdinov. 2020. fair evaluation in concept normalization: a large-scale comparative analysis for bertbased models. in proceedings of the 28th international conference on computational linguistics (coling). dongfang xu, zeyu zhang, and steven bethard. 2020. a generate-and-rank framework with semantic type regularization for biomedical concept normalization. in proceedings of the 58th annual meeting of the association for computational linguistics, pages 8452-8464.