Source-linked AI summary
Self-Alignment Pretraining for Biomedical Entity Representations
Fangyu Liu, Ehsan Shareghi, Zaiqiao Meng, Marco Basaldella, Nigel Collier
TL;DR
Biomedical entity linking needs representations that capture synonymy despite heterogeneous names and limitations of standard biomedical MLMs. SAPBERT self-aligns entity representations with UMLS-based metric learning, achieving state-of-the-art results across six benchmarks, including without task-specific supervision in scientific language.
Problem
Biomedical entity linking remains difficult because heterogeneous names and existing pretrained MLM representations do not reliably capture fine-grained synonymy.
Method
SAPBERT pretrains a BERT-based encoder on UMLS synonyms using hard-pair mining and a scalable metric-learning objective that clusters names of the same concept.
Results
SAPBERT achieves new state-of-the-art performance across all six medical entity-linking benchmarks and outperforms prior scientific-language state of the art without task-specific fine-tuning.
Takeaways & Limitations
SAPBERT provides a simple one-model approach for medical entity linking, requiring only nearest-neighbor search at test time and integrating into existing BERT-based systems.
Abstract
from arXiv · showhide
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.
1 Introduction
Biomedical entity representations underpin medical text-mining applications, but heterogeneous names and existing MLM representations make fine-grained concept separation difficult. SAPBERT addresses this by self-aligning synonymous entity names using UMLS-scale pretraining.
- Biomedical entity representations support literature search, clinical decision making, and relational knowledge discovery.
- Heterogeneous naming makes biomedical concept representation challenging, as Hydroxychloroquine appears as Oxichlorochine, HCQ, and Plaquenil.
- Medical entity linking maps entity mentions to unified concepts in a medical knowledge graph, with representation quality as its main bottleneck.
- Existing domain-specific MLMs such as PubMedBERT do not produce a well-separated biomedical entity representation space.
- SAPBERT pretrains on UMLS synonyms, clustering names of the same concept through hard-pair sampling and scalable metric learning.UMLS 2020AA contains 4M+ concepts and 10M+ synonyms from over 150 controlled vocabularies.
- SAPBERT reaches new state-of-the-art performance across six medical entity-linking benchmarks and improves accuracy over domain-specific pretrained models by up to 20%.It also outperforms previous scientific-language state of the art without task-specific fine-tuning.
2 Method: Self-Alignment Pretraining
SAPBERT represents biomedical names with a BERT encoder and aligns synonymous names using hard-pair mining and a metric-learning objective. The framework emphasizes informative positive and negative pairs so the representation space better reflects synonymy.
- Framework: The framework supports both UMLS pretraining and fine-tuning on task-specific datasets, starting from an existing BERT model.
- Formal Definition: SAPBERT uses a BERT model’s [CLS] output as the representation of each biomedical name.Cosine similarity between representations estimates resemblance, with high similarity for synonyms and low similarity otherwise.
- Online Hard Pairs Mining: Online hard-triplet mining selects difficult positive and negative pairs from each mini-batch for efficient training.A triplet contains an anchor, a same-label positive, and a different-label negative; disabling mining causes a drastic performance drop.
- Loss Function: The adapted Multi-Similarity loss learns from pairwise cosine similarities among BERT-produced name representations.The similarity matrix contains cosine similarities for all names in a mini-batch.
- Loss Function: The loss uses temperature scales, a similarity offset, and positive and negative sample indices for each anchor.
- Loss Function: The loss pushes negative pairs apart and pulls positive pairs together while re-weighting informative samples for stronger gradient signals.
3 Experiments and Discussions
Experiments evaluate SAPBERT across six MEL datasets and show consistent gains over BERT-based baselines and published SOTA systems. Ablations indicate that online hard-pair mining is important, while SAPBERT can also improve existing MEL systems.
- Experimental Setup: The UMLS pretraining data contains 9,712,959 deduplicated lowercase name–CUI entries, and training runs for one epoch with batch size 512.The reported setup uses UMLS 2020AA data, AdamW, mixed precision, and approximately 50k iterations.
- Evaluation Data and Protocol: SAPBERT is evaluated on six English MEL datasets using Acc@1 and Acc@5, with nearest-neighbour search at test time and no task fine-tuning unless stated.The datasets cover scientific and social-media language domains; reported results generally average five random-seed runs.
- Main Results: SAPBERT consistently improves seven BERT-based models across all datasets, with gains of up to 31.0% absolute Acc@1 in the social-media domain.PUBMEDBERT+SAPBERT performs best after SAPBERT pretraining among the compared BERT-based models.
- Main Results: Measured by Acc@1, SAPBERT reaches statistically significant new SOTA on 5 of 6 datasets and is comparable to SOTA on BC5CDR-c, 96.5 versus 96.6.On scientific datasets, unsupervised SAPBERT outperforms supervised SOTA; on social-media datasets, fine-tuning enables significant SOTA improvements.
- Ablation Analysis: Switching off online hard-pair mining causes a large Acc@1 drop and a smaller but significant Acc@5 drop.Easy, already-separated samples can dominate gradients and harm learning when hard-pair mining is removed.
- Integration with Existing Systems: Initializing BIOSYN with SAPBERT pretrained weights boosts performance across all evaluated datasets.This demonstrates that SAPBERT can be inserted into an existing BERT-based MEL system.
4 Conclusion
The paper concludes that SAPBERT self-aligns biomedical entity representations and consistently improves MEL performance. It achieves SOTA across six benchmarks, including without task-specific fine-tuning on multiple scientific-language datasets.
- Conclusion: SAPBERT is presented as a self-alignment pretraining scheme for learning biomedical entity representations.The approach is designed to improve representation quality for biomedical entity-level tasks.
- Conclusion: SAPBERT obtains new SOTA results on all six widely used MEL benchmarks and outperforms previous supervised SOTA on multiple scientific-language datasets without task-specific fine-tuning.The conclusion also identifies general-domain self-alignment, additional relations, and sentence-level representation learning as future directions.
A Evaluation Datasets Details
The evaluation datasets are divided into scientific-language and social-media-language categories according to their source text. Scientific datasets come from papers, whereas social-media datasets come from forums such as Reddit.
- Dataset Categories: The experiments divide datasets into scientific-language and social-media-language categories.The categories distinguish datasets by the source of their language data.
- Dataset Categories: Scientific-language datasets are extracted from scientific papers, while social-media-language datasets come from forums such as Reddit.com.The paper refers readers to Table 3 for an overview of key dataset statistics.
A.1 Scientific Language Datasets
The scientific-language evaluation includes NCBI, BC5CDR, and MedMentions, covering disease, chemical, and broad UMLS-linked entity mentions. MedMentions is especially large and challenging for conventional MEL systems.
- NCBI: NCBI contains 793 annotated PubMed abstracts and 6,881 mentions mapped to the MEDIC dictionary.The experiments denote this dataset as NCBI.
- BC5CDR: BC5CDR contains 1,500 PubMed articles with 4,409 chemical mentions, 5,818 disease mentions, and 3,116 chemical-disease interactions.Disease mentions use MEDIC, while chemical mentions use the CTD chemical dictionary.
- MedMentions: MedMentions contains over 4,000 abstracts and over 350,000 mentions linked to UMLS 2017AA.Its more than 3M concepts and massive mention count make many MEL systems infeasible, whereas SAPBERT can be applied with minimal effort.
A.2 Social-Media Language Datasets
The study includes two social-media MEL datasets: AskAPatient and COMETA, covering adverse drug reactions and health-related Reddit discussions.
- AskAPatient: AskAPatient contains 17,324 adverse drug reaction annotations mapped to 1,036 concepts from SNOMED-CT and AMT.The evaluation follows the original paper’s 10-fold protocol.
- Dataset scope: The social-media datasets represent distinct sources and evaluation settings within medical entity linking.AskAPatient uses blog posts, whereas COMETA uses Reddit discussions.
- COMETA: COMETA contains around 20k medical mentions from health-related Reddit discussions, mapped to SNOMED-CT.The study uses COMETA’s stratified (general) split and follows the original evaluation protocol.
B.1 The Choice of Base Models
The study evaluates BERT-based models and reports dataset statistics and baseline comparisons across scientific and social-media MEL datasets.
- Base-model setup: All BERT models in the study use the BERTBase architecture.The models and their specific versions are listed in Table 5.
- Dataset coverage: The evaluation covers six MEL datasets spanning scientific and social-media language domains.Table 3 reports concepts, surface forms or synonyms, and mentions or queries for these datasets.
- Baseline comparisons: The baseline table includes scientific and social-media MEL systems, with the final row collecting reported results from the best-performing models.Asterisks mark results produced using official released code; dashes indicate unreported results, and OOM denotes out-of-memory.
B.2 Comparing Loss Functions
The study compares multiple metric-learning losses for SAPBERT pretraining using COMETA’s zeroshot-general split and also documents related experimental configurations.
- Evaluation setting: Online mining is enabled for all models in the COMETA zeroshot-general comparison.The split is described as challenging and distinct from the stratified-general split used elsewhere.
- Loss comparison: Loss-function comparisons are reported on COMETA’s zeroshot-general split.Table 6 specifically compares losses used for SAPBERT pretraining.
- Compared objectives: The compared objectives include cosine, max-margin triplet, hinge-style, InfoNCE, Lifted-Structure, NCA, Multi-Similarity, and Circle losses.These objectives are situated among prior metric-learning and contrastive-learning approaches.
- Configuration details: The paper documents parameter counts for ADAPTER variants and full-model tuning, alongside hardware and hyper-parameter search details.The reported hyper-parameters balance memory limits and model performance but may be sub-optimal on individual datasets.
C.3 A High-Resolution Version of Fig. 1
The paper provides a higher-resolution version of the t-SNE visualization shown in Figure 1.
- Visualization: Figure 3 is a clearer, higher-resolution version of Figure 1’s t-SNE embedding visualization.The figure is presented to improve visualization clarity rather than introduce a new analysis.
- Hyper-parameters: The associated hyper-parameter table lists the search space used to obtain the reported results.Asterisks identify the hyper-parameters selected for reporting.