[Paper Review] BioMedLM: A 2.7B Parameter Language Model Trained On Biomedical Text

Elliot Bolton, Abhinav Venigalla+9 more|arXiv (Cornell University)|Mar 27, 2024

Biomedical Text Mining and Ontologies31 citations

TL;DR

BioMedLM is a 2.7B-parameter GPT-style model trained solely on PubMed abstracts and articles, achieving competitive biomedical QA performance after fine-tuning while enabling on-device inference and open data provenance.

ABSTRACT

Models such as GPT-4 and Med-PaLM 2 have demonstrated impressive performance on a wide variety of biomedical NLP tasks. However, these models have hundreds of billions of parameters, are computationally expensive to run, require users to send their input data over the internet, and are trained on unknown data sources. Can smaller, more targeted models compete? To address this question, we build and release BioMedLM, a 2.7 billion parameter GPT-style autoregressive model trained exclusively on PubMed abstracts and full articles. When fine-tuned, BioMedLM can produce strong multiple-choice biomedical question-answering results competitive with much larger models, such as achieving a score of 57.3% on MedMCQA (dev) and 69.0% on the MMLU Medical Genetics exam. BioMedLM can also be fine-tuned to produce useful answers to patient questions on medical topics. This demonstrates that smaller models can potentially serve as transparent, privacy-preserving, economical and environmentally friendly foundations for particular NLP applications, such as in biomedicine. The model is available on the Hugging Face Hub: https://huggingface.co/stanford-crfm/BioMedLM.

Motivation & Objective

• Motivate the development of a domain-specific, smaller LLM to address privacy, cost, and transparency concerns of large models in biomedicine.

Proposed method

• Autoregressive decoder-only Transformer (GPT-2 style) with 2.7B parameters.
• Domain-specific Byte-Pair Encoding tokenizer trained on PubMed abstracts to improve tokenization of biomedical terms.
• Pre-training on PubMed abstracts and articles (34.6B tokens; 8.67 passes; ~300B tokens explored) using mixed-precision, bf16 for final training, and Decoupled AdamW optimizer.
• Fine-tuning for downstream biomedical QA tasks with architecture specialized for multiple-choice prompts (per-task prompt shaping and a final linear classifier over answer scores).
• Generation-style fine-tuning for consumer-health question answering (long-form responses) using web-derived QA pairs.

Figure 1: Train and Validation Loss after 100k Batches

Experimental results

Research questions

• RQ1 Can a compact, domain-specialized model (2.7B parameters) match or approach the performance of larger models on biomedical QA tasks?
• RQ2 Does training exclusively on PubMed data and using a biomedical tokenizer improve downstream task performance relative to general-domain baselines?
• RQ3 What are the trade-offs in privacy, cost, and accessibility when deploying a small, open biomedical LLM compared to closed, large models?

Key findings

• BioMedLM achieves competitive results on multiple biomedical QA benchmarks after fine-tuning, approaching or matching larger models in several tasks (e.g., MedMCQA 57.3%, MMLU Medical Genetics 69.0%).
• Domain-specific pre-training on PubMed with a specialized tokenizer yields noticeable gains over GPT-2/tokenizer baselines (e.g., MedQA improvement from 33.05 to 34.98 at 125M scale).
• Compared to GPT-Neo 2.7B trained on general English data, BioMedLM substantially outperforms on select QA tasks (e.g., 27 percentage point improvement on BioASQ).
• BioMedLM supports on-device inference and can be fine-tuned on modest hardware while providing transparency about training data and architecture.

Figure 2: Comparison of GPT-Neo 2.7B and BioMedLM on Select QA Tasks