Source-linked AI summary
BioMegatron: Larger Biomedical Domain Language Model
Hoo-Chang Shin, Yang Zhang, Evelina Bakhturina, Raul Puri, Mostofa Patwary, Mohammad Shoeybi, Raghav Mani
TL;DR
Biomedical language-model studies had not deeply established how vocabulary, labeling, model size, pre-training corpus, and domain transfer affect different applications. The paper evaluates these factors with BioMegatron across biomedical benchmarks and finds that targeted domain and application alignment matters more than model size alone, with gains across token classification, sequence classification, and question answering.
Problem
Prior work provided limited study of factors affecting domain-language applications, while biomedical datasets are often small and label-biased.
Method
The paper evaluates subword vocabulary, labeling method, model size, pre-training corpus, and domain transfer using BioMegatron and biomedical benchmarks.
Results
Targeted domain- and application-specific models perform best; model size is secondary to vocabulary for token classification, and larger size does not necessarily improve cross-domain benchmarks.
Takeaways & Limitations
There is no single model that performs all biomedical applications as well as a targeted model, although larger size may further help domain- and application-specific models.
Abstract
from arXiv · showhide
There has been an influx of biomedical domain-specific language models, showing language models pre-trained on biomedical text perform better on biomedical domain benchmarks than those trained on general domain text corpora such as Wikipedia and Books. Yet, most works do not study the factors affecting each domain language application deeply. Additionally, the study of model size on domain-specific models has been mostly missing. We empirically study and evaluate several factors that can affect performance on domain language applications, such as the sub-word vocabulary set, model size, pre-training corpus, and domain transfer. We show consistent improvements on benchmarks with our larger BioMegatron model trained on a larger domain corpus, contributing to our understanding of domain language model applications. We demonstrate noticeable improvements over the previous state-of-the-art (SOTA) on standard biomedical NLP benchmarks of named entity recognition, relation extraction, and question answering. Model checkpoints and code are available at [https://ngc.nvidia.com] and [https://github.com/NVIDIA/NeMo].
1 Introduction
Biomedical language-model research expanded rapidly after BERT-based models demonstrated strong domain performance. This paper studies vocabulary, labeling, model size, and domain transfer in greater detail across several biomedical tasks.
- 1 Introduction: BERT-based biomedical models such as BioBERT and SciBERT motivated many subsequent domain-specific language-model studies.Follow-up work added clinical text and evaluated models on biomedical and clinical benchmarks.
- 1 Introduction: Gu et al. studied pre-training corpus domain, masking method, and adversarial training across token, sequence, and regression datasets.
- 1 Introduction: This paper examines subword vocabulary, labeling method, model size, and domain transfer, reporting gains across token classification, sequence classification, and question answering.
2 Related Works
Prior biomedical language models adapted BERT to biomedical and clinical corpora, generally improving benchmark performance. Most studies used BERTBase and provided limited evidence about scaling beyond standard model sizes.
- 2 Related Works: BioBERT continued pre-training BERT on PubMed abstracts, while SciBERT trained BERT from scratch with a scientific-text vocabulary.
- 2 Related Works: Both BioBERT and SciBERT improved over previous non-BERT biomedical benchmark results, including named entity recognition and relation extraction.BioBERT also reported question-answering results.
- 2 Related Works: Follow-up models extended biomedical pre-training with clinical text, while PubMedBERT added full-text PubMed data and a PubMed-derived vocabulary.
- 2 Related Works: Most prior studies used BERTBase with 110M parameters, and BERTLarge produced mixed results across relation extraction and named entity recognition.These biomedical tasks showed less scaling benefit than standard NLP benchmarks such as GLUE or SQuAD.
3 Language Model Pre-training
BioMegatron applies scalable Megatron-LM architecture to biomedical pre-training using large PubMed and PMC corpora. The experiments vary model size, initialization, casing, and domain vocabulary.
- 3 Language Model Pre-training: BioMegatron uses Megatron-LM’s architecture, including layer-normalization and residual-connection ordering that supports scaling BERT-style models beyond 336M parameters.
- 3 Language Model Pre-training: Pre-training uses 4.5 billion PubMed abstract words and 1.6 billion words from the PMC full-text corpus.
- 3 Language Model Pre-training: The study trains BioMegatron models with 345M, 800M, and 1.2B parameters.
- 3 Language Model Pre-training: The 345M experiments compare BERT cased and uncased vocabularies, initialization strategies, and PubMed-derived vocabularies with 30K or 50K subword units.
- 3 Language Model Pre-training: Larger models use fewer pre-training variations: 800M models train from scratch on PubMed, while the 1.2B model starts from a general-domain checkpoint.
4 Downstream Benchmark Tasks
The paper evaluates BioMegatron on widely used biomedical benchmarks for named entity recognition, relation extraction, and question answering. These tasks cover entity tagging, sentence-level relation classification, and factoid answering.
- 4 Downstream Benchmark Tasks: The benchmark suite includes NER datasets BC5CDR and NCBI-disease, which annotate biomedical disease and chemical entities with IOB tagging.
- 4 Downstream Benchmark Tasks: ChemProt evaluates chemical-protein relation extraction by classifying PubMed sentences into five relation categories.
- 4 Downstream Benchmark Tasks: BioASQ-7b evaluates biomedical factoid question answering using context, question, and answer triplets.Metrics include strict accuracy, lenient accuracy, and mean reciprocal rank.
5 Results and Discussion
Results show that subword vocabulary strongly affects biomedical NER and RE, while larger models generally help QA but can plateau or overfit on small datasets. Domain transfer requires sufficient PubMed pre-training, after which general-domain models can approach domain-only performance.
- Named Entity Recognition: Subword vocabulary is the most critical factor for NER, with single-term entity representations generally more helpful than splitting entities into subtokens.PubMedBERT has a low entity break-out rate despite a smaller vocabulary, which may help its NER performance.
- Named Entity Recognition: Different NER labeling schemes can change F1-score by approximately 2%, and subtoken labeling is generally used except for NCBI-disease.Whole-entity labeling performs better for NCBI-disease across models.
- Named Entity Recognition and Relation Extraction: Larger models produce higher precision at lower recall for both NER and RE, while additional hyperparameter tuning may improve F1-scores with uncertain generalization.
- Question Answering: Biomedical vocabularies perform much worse when models are first fine-tuned on SQuAD, possibly because those vocabularies are less relevant to the initial task.The QA setup fine-tunes on SQuAD before BioASQ and uses repeated question-answer pairs in BioASQ.
- Question Answering: Larger models tend to improve QA, but the effect levels off after 345m parameters and is more evident when fine-tuning directly on BioASQ.Table 5 reports a consistent improvement pattern with model size, while other models and domain-vocabulary variants do not achieve comparable BioASQ results.
- Domain Transfer and Generalization: A general-domain BioMegatron-1.2b model needs sufficient PubMed pre-training to match a model trained only on domain text, with vocabulary remaining especially important for NER.On BioASQ, larger general-domain models do not perform better, possibly indicating overfitting on the small training set.
- Domain Transfer and Generalization: A large biomedical language model trained on a large corpus can outperform smaller general-domain models on SQuAD, even when those models were pre-trained on biomedical text.Biomedical datasets are generally smaller and more label-biased than their general-domain counterparts, including substantially smaller QA data.
6 Conclusion
The study identifies factors affecting domain language model performance and finds that domain- and application-targeted models perform best. Model size is secondary to vocabulary set for token classification, and larger models do not necessarily improve cross-domain benchmark performance.
- The study reviews and tests several factors that can affect domain language model performance.
- Targeting a language model to a specific domain and application produces the best performance.
- Vocabulary set is more important than model size for token classification tasks.
- Larger models do not necessarily perform better on cross-domain benchmark tasks.
- There may be no single model that performs well enough across all domains and applications compared with targeted models.