[Paper Review] MEDITRON-70B: Scaling Medical Pretraining for Large Language Models

Zeming Chen, A. Cano+18 more|arXiv (Cornell University)|Nov 27, 2023

Artificial Intelligence in Healthcare and EducationMedicine117 citations

TL;DR

MEDITRON-7B and MEDITRON-70B are open-source medical LLMs built by continued pretraining on curated medical data, achieving state-of-the-art results among open models and competitive performance versus some closed models on multiple medical benchmarks.

ABSTRACT

Large language models (LLMs) can potentially democratize access to medical knowledge. While many efforts have been made to harness and improve LLMs' medical knowledge and reasoning capacities, the resulting models are either closed-source (e.g., PaLM, GPT-4) or limited in scale (<= 13B parameters), which restricts their abilities. In this work, we improve access to large-scale medical LLMs by releasing MEDITRON: a suite of open-source LLMs with 7B and 70B parameters adapted to the medical domain. MEDITRON builds on Llama-2 (through our adaptation of Nvidia's Megatron-LM distributed trainer), and extends pretraining on a comprehensively curated medical corpus, including selected PubMed articles, abstracts, and internationally-recognized medical guidelines. Evaluations using four major medical benchmarks show significant performance gains over several state-of-the-art baselines before and after task-specific finetuning. Overall, MEDITRON achieves a 6% absolute performance gain over the best public baseline in its parameter class and 3% over the strongest baseline we finetuned from Llama-2. Compared to closed-source LLMs, MEDITRON-70B outperforms GPT-3.5 and Med-PaLM and is within 5% of GPT-4 and 10% of Med-PaLM-2. We release our code for curating the medical pretraining corpus and the MEDITRON model weights to drive open-source development of more capable medical LLMs.

Motivation & Objective

• Demonstrate how domain-specific continued pretraining improves medical reasoning in large LLMs.
• Release open-source medical pretraining data, training code, and model weights to enable reproducible development.
• Evaluate MEDITRON on key medical benchmarks with in-context learning and finetuning.
• Show how larger models (70B) compare to baselines and closed LLMs in medical tasks.

Proposed method

• Continued pretraining of Llama-2-based models on GAP-REPLAY: a medical data mix of PubMed/PubMed Central papers, abstracts, clinical guidelines, and an experience replay corpus.
• Engineering extensions to Megatron-LM to support Llama/Llama-2/Falcon architectures and efficient training (DP/PP/TP, FlashAttention).
• Hyperparameters: AdamW with cosine schedule; 2048/4096 context; 32k BPE vocab; domain-specific token handling.
• Supervised finetuning on benchmark training sets (PubMedQA, MedMCQA, MedQA) with task-specific instructions.
• Inference strategies including Top Token Selection, Chain-of-Thought prompting (CoT), and Self-Consistency CoT (SC-CoT).
• Evaluation on four medical benchmarks (MedQA, MedMCQA, PubMedQA, MMLU-Medical) under both in-context learning and finetuning regimes.

Experimental results

Research questions

• RQ1 How does continued pretraining on a curated medical corpus impact performance of large-scale LLMs in medical reasoning tasks?
• RQ2 What are the gains of scaling from 7B to 70B parameters in medical domains when using in-context learning versus finetuning?
• RQ3 Do advanced inference strategies (CoT, SC-CoT) provide substantial gains for medical reasoning with MEDITRON models?
• RQ4 How do open-source medical LLMs compare to closed models (GPT-3.5, GPT-4, Med-PaLM) on standard medical benchmarks?
• RQ5 What is the impact of domain-specific data (clinical guidelines, PubMed abstracts/papers) versus general data on medical pretraining efficacy?

Key findings

• MEDITRON-7B with in-context learning outperforms several open baselines; continued medical pretraining provides additional gains over base Llama-2-7B.
• MEDITRON-70B significantly improves over its base Llama-2-70B across benchmarks, including PubMedQA and MedMCQA.
• Self-Consistency CoT with MEDITRON-70B yields the highest average accuracy among evaluated configurations (72.0% Avg in Table 5).
• At 70B, MEDITRON-70B approaches performance of some commercial LLMs and outperforms GPT-3.5 and Med-PaLM on several tasks, while remaining below GPT-4/Med-PaLM-2 in some metrics.
• On MedQA-like tasks, MEDITRON achieves strong performance, with PubMedQA in-context results close to finetuned baselines (e.g., 79.8% in-context vs 80.0% finetuned).
• The release includes the curated medical pretraining corpus, training library, and model weights to support open development.