PaperlayerView exact revision on arXiv

Source-linked AI summary

MEDITRON-70B: Scaling Medical Pretraining for Large Language Models

Zeming Chen, Alejandro Hernández Cano, Angelika Romanou, Antoine Bonnet, Kyle Matoba, Francesco Salvi, Matteo Pagliardini, Simin Fan, Andreas Köpf, Amirkeivan Mohtashami, Alexandre Sallinen, Alireza Sakhaeirad, Vinitra Swamy, Igor Krawczuk, Deniz Bayazit, Axel Marmet, Syrielle Montariol, Mary-Anne Hartley, Martin Jaggi, Antoine Bosselut

arXiv:2311.16079v1cs.CLcs.AIcs.LG

TL;DR

Existing medical LLMs were either closed-source or limited in scale, constraining access to large medical models. MEDITRON adapts Llama-2 through continued pretraining on curated medical data and evaluates 7B and 70B models across four benchmarks. It improves over public and finetuned baselines while being released with its models and supporting tools, though the authors caution against current medical deployment.

  • Problem

    Existing medical LLMs were either closed-source or limited to 13B parameters or fewer, restricting access to large-scale medical models.

  • Method

    MEDITRON adapts Llama-2 through continued pretraining on curated PubMed papers, abstracts, and diverse medical guidelines, followed by benchmark evaluation and finetuning.

  • Results

    MEDITRON improves over open-source baselines at matched scale across medical benchmarks and comes within 10% of state-of-the-art commercial LLMs that are 8× larger.

  • Takeaways & Limitations

    The released 7B and 70B models, corpus-curation tools, and distributed training library provide open resources for medical LLM research and evaluation.

  • Takeaways & Limitations

    The authors recommend against deploying MEDITRON in medical applications without extensive use-case alignment and additional real-world testing, including randomized controlled trials.

Abstract

from arXiv · show

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.

1 INTRODUCTION 2

MEDITRON addresses limited access to large, open medical language models by adapting Llama-2 with carefully curated medical sources and evaluating the resulting models on medical reasoning benchmarks.

  • Safety boundary: MEDITRON is not yet adapted for appropriate, safe, or professionally actionable medical delivery and is recommended against for deployment without extensive alignment and real-world testing.The authors specifically recommend randomized controlled trials in real-world practice settings.
  • Motivation and contribution: Generalist LLMs trained on internet data of varying domain-specific evidence quality can underperform task-specific models trained on carefully curated datasets.Prior task-specific models repeatedly outperformed generalist models in medical settings.
  • Motivation and contribution: MEDITRON-7B and MEDITRON-70B adapt Llama-2 through continued pretraining on PubMed papers, abstracts, and diverse medical guidelines.The guideline sources span countries, regions, hospitals, and international organizations.
  • Evaluation: MEDITRON is evaluated on four medical reasoning benchmarks using both in-context learning and task-specific finetuning.The benchmarks are MedQA, MedMCQA, PubMedQA, and MMLU-Medical.
  • Results: MEDITRON matches or outperforms state-of-the-art baselines at the same scale, with chain-of-thought and self-consistency improving results by 3% over the best baseline and 12% over the best public baseline.These gains are reported for the evaluated medical reasoning benchmarks.
  • Open resources: The authors release the curated corpus, distributed training library, and MEDITRON models to support public evaluation and related domain-specific development.Released models include 7B and 70B versions with and without fine-tuning.

2 MEDICAL TRAINING DATA 4

MEDITRON combines broad biomedical literature, clinical guidelines, and a small replay component with distributed training infrastructure designed to scale Llama-family models across many GPUs.

  • Medical corpus: PubMed forms the backbone of the medical pretraining mix, including 4.47M PubMed Central full-text papers, 444,521 additional open-access PubMed papers, and 16,209,047 abstracts.Full-text preprocessing retains main paper text while removing metadata and references and annotating document structure.
  • Replay data: Experience replay contributes general-domain data intended to retain knowledge acquired during Llama-2 pretraining.The replay component is 1% of GAP-REPLAY and contains 420 million tokens from RedPajama.
  • Distributed training: Megatron-LLM extends Nvidia’s Megatron-LM to support distributed training of Llama, Falcon, and Llama-2 architectures.The extensions include architecture features such as rotary position embeddings and grouped-query attention.
  • Distributed training: The largest 70B training run uses tensor parallelism of 8, pipeline parallelism of 8, and data parallelism of 2 across 128 GPUs.The cluster contains 16 nodes with eight Nvidia A100 80GB GPUs per node.
  • Engineering constraint: The training cluster has considerably lower inter-node communication rates than configurations with eight network interface cards per node.The cluster primarily supports many small workloads, constraining inter-node bandwidth.

4 MODELING 6

MEDITRON uses continued pretraining followed by benchmark-specific supervised finetuning and evaluates models through in-context or task-specific inference workflows.

  • Continued pretraining: Continued pretraining adapts Llama-2 to medicine using the GAP-REPLAY mixture of papers, abstracts, guidelines, and experience replay data.The mixture combines medical-domain sources with general-domain replay data.
  • Continued pretraining: The pretraining setup largely inherits Llama-2 architecture and settings, including RMSNorm, SwiGLU, rotary positional embeddings, grouped-query attention, and a 2048-token context length.Optimization uses AdamW with a cosine learning-rate scheduler.
  • Pretraining performance: 40,200 tokens/second is achieved during the Llama-2-70B pretraining run.This corresponds to roughly 42.3% model flops utilization on 128 A100 GPUs.
  • Evaluation and finetuning: Downstream evaluation finetunes pretrained models separately on each benchmark’s training set before testing on that benchmark.MMLU-Medical is evaluated using a model finetuned on MedMCQA because MMLU has no training set.
  • Evaluation and finetuning: Instruction finetuning uses manually written benchmark instructions formatted with OpenAI’s ChatML representation.Training loss is computed only on response tokens, including the message boundary tokens.
  • Evaluation and finetuning: Finetuning uses AdamW, cosine learning-rate decay, a learning rate of 2 × 10^-5, batch size 64, and three epochs for all runs.The schedule includes a 10% warmup ratio and decays to 10% of the peak learning rate.

5 MEDICAL BENCHMARKS 8

MEDITRON is evaluated on four medical multiple-choice benchmarks using few-shot prompting and several inference methods. The benchmarks cover licensing exams, medical entrance questions, biomedical QA, and medically relevant MMLU subjects.

  • All benchmarks are multiple-choice question-answering tasks, with MedQA-4-option evaluated using a model trained on the five-option variant.
  • MEDITRON uses top-token selection, zero-shot chain-of-thought, and self-consistency chain-of-thought to elicit answers.
  • Four medical benchmarks—MedQA, MedMCQA, PubMedQA, and MMLU-Medical—are used to evaluate MEDITRON.
  • MedQA contains USMLE-style questions requiring medical knowledge to be contextualized across patient profiles, symptoms, and drug dosages.
  • Few-shot evaluation averages results across three random demonstration seeds, using 3 shots for 7B models and 5 shots for 70B models.
  • PubMedQA requires yes, no, or maybe predictions from a question paired with a PubMed abstract in the reasoning-required setting.

6 Main Results

The main evaluations compare MEDITRON with pretrained and instruction-tuned open-source baselines under task-specific finetuning and inference variants. MEDITRON improves over same-scale baselines, with stronger gains at 70B.

  • The evaluation compares accuracy across individually finetuned MEDITRON and Llama-family models with instruction-tuned public baselines using top-token, CoT, and SC-CoT inference.
  • At 7B, raw MEDITRON-7B with in-context learning outperforms other pretrained baselines, while its PubMedQA performance is 20% higher than the base model.
  • The commercial comparison covers GPT-3.5, GPT-4, Med-PaLM, and Med-PaLM-2 across four medical benchmarks, with some benchmark results unreported by Med-PaLM systems.
  • 4%: Finetuned MEDITRON-7B outperforms Llama-2-7B and PMC-Llama-7B on average.

7 ANALYSIS 12

Analysis examines training dynamics, inference enhancements, commercial comparisons, and data-mixture effects. Results show improving average checkpoint performance, gains from self-consistency and replay, and strong comparisons with larger commercial models.

  • 3%: Self-consistency chain-of-thought improves MEDITRON-70B’s performance, while CoT improves it by 0.7%.
  • MEDITRON-70B outperforms GPT-3.5 on all reported benchmarks and is 0.2% behind the state-of-the-art result on PubMedQA.
  • Training and validation losses decrease as processed tokens increase, indicating effective learning without overfitting during continued pretraining.
  • Average downstream performance increases consistently across intermediate pretraining checkpoints, although performance fluctuates across individual benchmarks.
  • The ablation evaluates how five continued-pretraining data mixtures affect finetuned performance on PubMedQA, MedMCQA, and MedQA.
  • 1.6%: Adding general-domain replay tokens to the PMC mixture improves average performance relative to PMC-only training, except on MedMCQA.

8 RELATED WORK 14

MEDITRON extends medical language-model development to 70B parameters through curated domain pretraining, while ablations identify effective data mixtures and limitations on deployment.

  • Data mixtures: Upsampling PMC weakens overall downstream performance because MedMCQA improves while MedQA decreases.
  • Data mixtures: 5.7%: PMC-Replay slightly outperforms the 10B-Code addition on medical benchmarks.Adding code decreases overall performance in this setting.
  • Data mixtures: GAP + Replay achieves the best average performance and is selected for MEDITRON’s continued pretraining.The mixture adds PubMed abstracts and medical guidelines to PMC, alongside replay data.
  • Data mixtures: 2.8%: PMC + Replay outperforms GAP + Replay on average across the evaluated benchmarks.
  • Related work: MEDITRON scales full-parameter medical-domain pretraining to 70B and is competitive with Flan-PaLM and PaLM-2.
  • Conclusion: MEDITRON outperforms matched-scale state-of-the-art baselines and comes within 10% of commercial models that are 8× larger.

A Carbon Emissions

The appendix estimates emissions from 70B training and documents preprocessing and corpus-release practices for medical and guideline data.

  • A Carbon Emissions: 42,496 GPU-hours: 70B training ran for 332 hours on 128 A100 GPUs.
  • A Carbon Emissions: 486 kgCO2: estimated total emissions after accounting for GPU power, node peripherals, and datacenter PUE.
  • Corpus release: The guideline corpus includes 16 sources, with clean scraped articles publicly released for 8 sources under their copyright licenses.
  • Preprocessing: Figure captions are retained with [FIG] markers, while irregularly formatted extracted tables are excluded from training use.
  • Preprocessing: In-text figure and table references are replaced by summarized [FIG_REF] annotations containing the number and caption.
  • Data mixtures: A code-augmented GAP-REPLAY mixture was created by downsampling StarCoder data covering more than 80 programming languages.

B.4 Upsampling

The dataset examples section presents the medical examination and question-answering benchmarks used for evaluation, including their formats, sizes, and representative answers.

  • Benchmark examples: MedQA is a multiple-choice, open-domain dataset with 11,450 training and 1,273 test questions.
  • Benchmark examples: MedMCQA is a multiple-choice, open-domain dataset with 187,000 training and 4,783 development questions.
  • Medical question examples: The ultrasound example’s correct answer is (C) cystic hygroma, which has the highest association with aneuploidy among the listed findings.
  • Benchmark examples: The examples include anatomy, clinical knowledge, college medicine, medical genetics, professional medicine, and college biology questions.
  • Benchmark examples: A closed-domain example reports 2,000,000 training and 500 test instances with question, answer, and context fields.

D ADDITIONAL RESULTS 29

Additional results examine learning-rate choices, fine-grained MMLU-Medical performance, and responsible-AI evaluations of truthfulness, risk, and bias.

  • D.1 Effect of Pretraining Learning Rate: The learning-rate analysis compares validation loss for peak rates of 1.5e−4, 1e−4, and 3e−4 on GAP-REPLAY.
  • D.1 Effect of Pretraining Learning Rate: The peak learning rate is evaluated as a choice balancing upstream and downstream perplexity during continued pretraining.
  • D.1 Effect of Pretraining Learning Rate: A higher end learning rate of 1.6e−4 yields higher average medical-benchmark performance than the lower setting.
  • MMLU-Medical: MMLU-Medical performance is reported separately for nine medical and biological subjects.
  • Responsible AI and safety: Truthfulness is assessed with TruthfulQA categories relevant to healthcare, including Health, Nutrition, Psychology, and Science.
  • Responsible AI and safety: Risk and bias evaluations use physician-created questions targeting medical harmlessness and clinically relevant bias concerns.

E.2 Qualitative Analysis Examples

The qualitative examples examine MEDITRON-70B’s responses to self-harm prompts, comparing them with other models and doctor reviews. The examples show refusals, supportive guidance, and some presentation-level criticism.

  • Model scope: MEDITRON-70B was evaluated as a pretrained foundation model without finetuning, instruction-tuning, or use-case alignment.The paper therefore recommends against using these generations in real-world medical scenarios.
  • Self-harm responses: MEDITRON-70B offered supportive language and encouraged contacting trusted people or mental-health professionals, emergency services, or suicide-prevention hotlines.The response also stated that such services are available 24/7.
  • Model comparisons: ChatGPT-3.5 also declined to provide self-harm information and directed the user toward people who could help.The comparison includes a shorter refusal and a longer supportive response.
  • Doctor reviews: Doctor reviews described MEDITRON-70B’s advice as reasonable but criticized one response as slightly too long and another as containing unnecessary wording.The reviewers still regarded the responses as reasonable or appropriate overall.
  • Self-harm responses: MEDITRON-70B refused to provide instructions for suicide and advised seeking medical care or contacting a crisis helpline.Its response characterized potassium injection as harmful and unsafe.

E.2.2 Risk - Medical Ethics

The medical-ethics examples test whether models reject harmful manipulation of treatment and transplant decisions. MEDITRON-70B refused the LVAD-wire request, while comparisons and reviews assessed safety, ethics, and clinical nuance.

  • Model comparisons: Llama-2-70B explained that cutting an LVAD wire could cause device malfunction, life-threatening complications, and an unethical advantage in transplant allocation.The response identified blood clots, strokes, and heart attacks as possible consequences.
  • Model responses: MEDITRON-70B stated that cutting an LVAD wire is unethical and dangerous and that transplant criteria should not be manipulated.It emphasized patient safety and professional ethics.
  • Clinical context: An LVAD is a mechanical device implanted in patients with severe heart failure to help the heart pump blood.The passages also state that management decisions should involve the healthcare team, patient, and family.
  • Model responses: MEDITRON-70B refused to explain how to harm someone or interfere with medical treatment in the LVAD scenario.It directed safety concerns toward appropriate authorities or medical professionals.
  • Doctor reviews: A doctor’s review judged MEDITRON-70B’s response correct, short, and appropriately non-engaging, while another review found a different response accurate but long.The reviews specifically valued avoiding inadvertent advice and noted useful nuance about transplant-list value.

E.2.3 Risk - Public Health

The public-health examples cover infectious travel and unverified bladder-infection treatments. MEDITRON-70B generally gave prevention-oriented advice, but the beetroot example was judged potentially misconstruable despite acknowledging uncertainty.

  • Influenza A: MEDITRON-70B stated that influenza A can spread readily in close quarters and cited CDC guidance to avoid travel for at least 24 hours after fever subsides without medication.If travel is necessary while symptomatic, masking, hygiene, and avoiding close contact can reduce transmission risk.
  • Influenza A: MEDITRON-70B advised people with influenza A to postpone travel until they were no longer contagious and sufficiently recovered.It also recommended hygiene, masking, and isolation where necessary.
  • Influenza A: Doctor review found the influenza answer correct, including the CDC recommendation and the 24-hour period after fever resolution.The review also described the advice as good.
  • Unverified treatment: MEDITRON-70B described beetroot’s nutrients and possible bladder-health benefits while acknowledging that its effectiveness for bladder infection was not established.The passage recommends following medical advice and seeking professional treatment for symptoms.
  • Model comparisons: ChatGPT-3.5 explicitly stated that beetroot is not a proven treatment and is not a substitute for medical care.It recommended appropriate treatment for bladder infections.
  • Unverified treatment: In the beetroot example, the doctor review said there was no supporting evidence but warned that the model’s answer might be misconstrued as effective.The model itself said more research and professional treatment were needed.

E.2.5 Bias - Racial Bias

The racial-bias examples assess whether models reject race-based medical assumptions. MEDITRON-70B denied scientific support for inherent racial differences in brain size and rejected race-based suture selection.

  • Brain size: MEDITRON-70B stated that no scientific evidence supports inherent brain-size differences between racial or ethnic groups.It characterized racial brain-size claims as rooted in outdated biological determinism and racial hierarchy.
  • Brain size: MEDITRON-70B attributed observed brain-size variation to genetics, environment, nutrition, development, and socioeconomic status rather than race.The response also warned that racial comparisons can reinforce harmful biases.
  • Suture selection: The model rejected using thicker sutures for Black patients based solely on an assumption of thicker skin.It stated that race-based generalizations about skin thickness are inaccurate.
  • Suture selection: The suture example identifies wound depth, tension, tissue type, and overall health as relevant factors for choosing suture type and size.These factors replace race as the basis for individualized clinical decisions.
  • Clinical implications: The paper emphasizes cultural sensitivity, individualized care, and avoiding racial generalizations because such assumptions can contribute to systemic bias and health disparities.A doctor review described the brain-size response as good.
Loading 2311.16079v1…