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

In plain words

Large language models know a lot in general, but the strongest ones are closed, and the open ones that are small enough to run locally are weak on medicine. This paper takes an open general model (Llama 2) and keeps training it on a large pile of medical text: clinical practice guidelines, scientific papers and abstracts about medicine. A small slice of ordinary general text is mixed in so the model does not forget how to write normally.

After that continued training, the team trains the model separately for each exam-style test (medical licensing questions, Indian entrance-exam questions, a yes/no research question set, and a medical subset of a general exam). At answer time they test a few ways of asking: picking the single most likely answer, having the model reason step by step, and sampling several reasoning paths and taking the majority answer. The largest model with the most involved answer method comes out ahead of the open baselines on average.

The team also released code for building the medical corpus and the weights. The paper is explicit that the models are research artifacts and should not be used in medical settings without testing.

TL;DR

• Method: continued pretraining of Llama-2 at 7B and 70B on a medical mixture (GAP-Replay), then one supervised finetune per benchmark §4.1 §4.2.
• Data: 46K clinical guidelines, PubMed/PMC full texts and abstracts, and 420M replay tokens from RedPajama (1% of the mix) §2 tab. 1.
• Headline (author-reported): MEDITRON-70B with self-consistency CoT averages 72.0 versus 69.2 for finetuned Llama-2-70B tab. 5.
• Engineering: Megatron-LLM, 3D parallelism on 128 A100 80GB GPUs §3.
• Worth reading if you do domain continued pretraining of a Llama-2-class model and want the data pipeline, the trial-run ablations and the knobs that the code and paper set differently.

Problem

Closed medical LLMs were the strongest, while open medical models were limited to about 13B parameters or less. The paper's aim is an open suite at 7B and 70B that is adapted to medicine §1.

Method

The pipeline has four stages: curate a medical corpus §2, pretrain Llama-2 on it with a distributed trainer §3, finetune per benchmark and choose an inference mode §4 §4.3, then evaluate §5 §6.

1. Corpus: GAP-Replay (48.1B tokens including validation)

The mixture has four sources §2:

• Clinical Guidelines (G): 46,469 guideline articles from 16 sources §2.1 tab. 9. Text is cleaned of URLs, references, figures and table delimiters, headers are normalised, titles are used for deduplication, and short or non-English articles are dropped §2.1.
• Paper Abstracts (A): 16.2M abstracts from PubMed and PubMed Central, split into 15.7M train and 487K validation §2.2 tab. 1.
• Medical Papers (P): full texts of about 5M open-access PubMed/PMC papers, with metadata, references, tables and figures removed. Inline citations, figure and table references, and formulas are replaced by special tokens that carry a short summary §2.2 app. B.2.
• Replay (R): 420M tokens randomly sampled from RedPajama, roughly 1% of the mixture, to reduce forgetting of general-domain knowledge §2.3 tab. 1.

Train tokens by source: guidelines 107M, abstracts 5.48B, papers 40.7B, replay 420M; total 46.7B train plus 1.4B validation tab. 1. The final continued-pretraining mixture is GAP + Replay, 46.8B tokens in the trial-run table tab. 7.

Guideline articles from only 8 of the 16 sources are released (35,733 articles). Scrapers and preprocessing code are released for all 16 §2.1.

2. Engineering: Megatron-LLM

The team extended Nvidia's Megatron-LM to support Llama, Llama 2 and Falcon. They added rotary embeddings, grouped-query attention, FlashAttention and FlashAttention-2, and untied word-embedding and output weights for Llama §3. The trainer is a separate repository (see Paper vs code).

Training uses tensor parallelism of 8 (one node), and for the 70B run pipeline parallelism of 8 and data parallelism of 2 on 128 GPUs §3. The 70B run reached 40,200 tokens/s, about 42.3% of peak bf16 flops of the 128 GPUs §4.1. The 70B run took 332 hours, or 42,496 GPU-hours app. A.

3. Continued pretraining

• Start from Llama-2, keep its architecture (RMSNorm, SwiGLU, rotary embeddings, grouped-query attention), and keep the Llama tokenizer, extended with the special tokens §4.1 app. B.2.
• Context length 2048 for 7B and 4096 for 70B §4.1.
• Optimizer AdamW, cosine schedule, warmup 2000 steps, decay to 10% of the peak LR §4.1.
• Peak LR is 1.5e-4 for 70B and 3e-4 for 7B and 13B §4.1.
• The 70B run uses global batch 512 and micro-batch 2 §3.
• Training lasted about 23,000 iterations, which the intermediate evaluation table ties to 48B tokens tab. 6.

4. Supervised finetuning, one model per benchmark

• Each benchmark is finetuned on its own training set, then tested on its test set. MMLU-Medical has no training set, so the MedMCQA-finetuned model is used there §4.2 §5.
• Instruction text per benchmark is in tab. 2.
• Prompts use OpenAI's ChatML layout, with the system instruction, question and answer each wrapped in special tokens. Loss is computed only on the response tokens, including the closing token §4.2.
• Padding is inserted at the end of each document and masked from the loss §4.2.
• Hyperparameters: AdamW with beta1 0.9, beta2 0.95, eps 1e-5; cosine schedule with a 10% warmup ratio, decaying to 10% of the peak LR; LR 2e-5; weight decay 0.1; batch 64; 3 epochs §4.2.

# Continued pretraining on the GAP-Replay mixture  [§4.1]
# Input:  llama2_weights; token batches drawn from the GAP-Replay mixture
# Output: meditron_base, the pretrained checkpoint

def continued_pretraining(llama2_weights, token_batches):
    # Start from the Llama 2 weights and keep training with next-token loss  [§4.1]
    model = llama2_weights
    for batch in token_batches:              # (B, T), T = 4096 for 70B  [§4.1]
        loss = next_token_loss(model, batch)
        # AdamW (beta1 0.9, beta2 0.95, eps 1e-5), cosine LR, clip at 1.0  [§4.1]
        update(model, loss)
    return model

# Per-benchmark supervised finetuning  [§4.2]
# Input:  meditron_base; train_set of (instruction, question, answer)
# Output: a model finetuned for one benchmark

def finetune(meditron_base, train_set):
    # Build a ChatML prompt for each example  [§4.2]
    for example in train_set:
        text = chatml(system=example.instruction,
                      user=example.question,
                      assistant=example.answer)
        # The loss covers only the answer tokens  [§4.2]
        loss_mask = mask_non_response_tokens(text)
    # 3 epochs, batch 64, LR 2e-5, 10% warmup, cosine to 10% of peak  [§4.2]
    return train(meditron_base, train_set, loss_mask)

5. Inference modes

• Top-token: for single-label questions, take the vocabulary token with the highest log-probability and compare it with the expected answer. Pretrained models get few-shot demonstrations in the prompt; finetuned models get the question directly §4.3.
• Chain-of-thought (CoT): add the zero-shot cue "Let's think step-by-step" to the question, after Kojima et al. §4.3.
• Self-consistency CoT: decode at temperature 0.8, draw 5 generations, extract each answer, and take the majority vote §4.3.

# Self-consistency chain-of-thought  [§4.3]
# Input:  model, question with options
# Output: one option letter

def sc_cot_answer(model, question):
    # Append the zero-shot step-by-step cue  [§4.3]
    prompt = question + " Let's think step-by-step."
    answers = []
    for branch in range(5):                  # 5 generations  [§4.3]
        # Sample at temperature 0.8  [§4.3]
        text = sample(model, prompt, temperature=0.8)
        answers.append(extract_option(text))
    # Majority vote over the extracted options  [§4.3]
    return most_common(answers)

Hyperparameters

Results

Author-reported. Setup: five result columns over four benchmarks (MedQA, MedQA-4-option, MedMCQA, PubMedQA, MMLU-Medical). MedQA-4-option is the 4-option form of MedQA. Dataset sizes and test splits are in tab. 3 and §5. Models are compared to open-source baselines unless noted.

Few-shot, no finetuning tab. 4:

Finetuned per benchmark, with three inference modes tab. 5:

Reported baselines: Mistral-7B-Instruct, Zephyr-7B-β, Clinical-Camel-70B and Med42-70B are instruction-tuned models used without further finetuning tab. 5. Commercial comparisons (GPT-3.5, GPT-4, Med-PaLM, Med-PaLM-2) are taken from their papers, not rerun fig. 3.

Headline statements in the text: MEDITRON-7B gains 4% over the Llama-2-7B and PMC-Llama-7B fine-tuned baselines; MEDITRON-70B gains 1.8% over fine-tuned Llama-2-70B in top-token mode; the MedQA scores of all three MEDITRON-70B modes exceed the 60.0 human passing score cited in the paper §6.2 tab. 5.

Training curve and checkpoints: validation loss falls without overfitting fig. 4. The few-shot average rises across checkpoints, from 60.8 at Llama-2 to 63.3 at 23,000 iterations, though individual benchmarks fluctuate tab. 6.

Paper vs code

Code: epfLLM/meditron at a7c7cda3014e537f0df2ec58f836fbe920d6283b, Apache-2.0. The paper links this repository as a MEDITRON release alongside the weights §1. The distributed trainer is the separate epfLLM/Megatron-LLM repository, which is a submodule and was not present in the clone, so its code was not checked.

Discrepancies:

• Pretraining defaults are not the paper's 70B run. pretrain/finetune.sh defaults to 7B, pipeline parallelism 1, global batch 12, micro-batch 1, 5 warmup iterations, minimum LR 1e-6 and --train_iters 102 [code: pretrain/finetune.sh#L5-L9] finetune.sh [code: pretrain/finetune.sh#L173]. The paper's 70B run uses PP 8, global batch 512, 2000 warmup steps and about 23,000 iterations §3 §4.1 tab. 6.

• 70B finetuning LR. The paper sets 2e-5 for finetuning §4.2. In finetuning/finetune_sft.sh, the instruction branch sets 2e-5 [code: finetuning/finetune_sft.sh#L94-L96], but the Llama-2 branch then overrides it to 1.5e-4 peak and 1.5e-5 minimum for any model above 13B, including the instruction case finetune_sft.sh.

• Self-consistency branches. The paper uses 5 generations §4.3. The code's --sc_branch default is 10 inference.py.

• Sampling penalties. The paper does not state the penalties used for benchmark decoding. The vLLM call sets presence penalty 0.0, frequency penalty 1.0, top-k off and top-p 1.0 inference.py.

• Majority vote. Counter.most_common picks the first-seen answer on a tie. Answers outside the allowed labels are dropped from the denominator evaluate.py.

• Upsampling filters. The paper's upsampling table gives articles older than 10 years, bottom-quartile citations and journals not reviewed by UpToDate a score of 0 tab. 10. The code excludes these outright [code: gap-replay/pubmed/upsample.py#L132-L134] [code: gap-replay/pubmed/upsample.py#L147-L150] [code: gap-replay/pubmed/upsample.py#L158-L161]. This affects only the PMC Upsampled arm of the trial runs, not the final GAP mixture §7.2.

• Dropout is not stated in the paper but the code sets it to 0.0 [code: pretrain/finetune.sh#L159] finetune.sh.

Consistent with the code: AdamW betas and eps §4.1 [code: pretrain/finetune.sh#L162]; 4096 context for Llama-2 70B §4.1 [code: pretrain/finetune.sh#L136]; batch 64 and 3 epochs in SFT §4.2 [code: finetuning/sft.py#L33] sft.py; SC-CoT temperature 0.8 §4.3 inference.py.

Pitfalls

• Three different end-LR statements. Section 4.1 decays the LR to 10% of the peak. Appendix D.1 runs the end-LR study on the 7B model, and says the 2-epoch total-iteration setting chosen from it is used for MEDITRON pretraining. That setting ends at 1.6e-4 from a 3e-4 peak, about 53% of the peak, not 10% app. D.1 tab. 11. The paper states no 70B end LR. The code's default minimum LR is 1e-6 finetune.sh. Confirm the schedule before reproducing the 70B run.
• Guideline counts disagree. Section 2.1 gives 46,469 articles, the Table 9 total is 46,649, and Table 1 gives about 43K (41K train plus 2,284 validation) §2.1 tab. 9 tab. 1.
• Warmup. The paper uses 2000 warmup steps §4.1; the code default is 5 finetune.sh.
• Reversed claim in the data ablation. Section 7.2 says PMC + Replay beats GAP + Replay by 2.8%, but Table 8 shows GAP + Replay ahead at 58.9 against 56.2 on average §7.2 tab. 8. The table is consistent with the choice of GAP.
• Ablations are 7B only. Data-mixture choices come from 7B continued pretraining with SFT on three tasks (PubMedQA, MedMCQA, MedQA), not from the 70B model §7.2 tab. 8.
• Evaluation splits. MedMCQA's test labels are not public, so results are on its validation split §5. MedQA is trained on the 5-option version and evaluated on both 5- and 4-option forms §5. MMLU-Medical is out of distribution: it uses the MedMCQA-finetuned model §6.2.
• Few-shot setup. 7B uses 3-shot because the context window is limited to 2048 tokens; 70B uses 5-shot tab. 4.
• Corpus not fully released. Only 8 of 16 guideline sources are released. PubMed texts come from S2ORC and PMC, and replay data from RedPajama §2.1 §2.2 §2.3. Reproducing the mixture needs the released code and those upstream sources.
• Data cutoff. Papers and abstracts end in August 2023 §2.2.
• Preprocessing matters. Citations, figure and table references, and formulas are rewritten into special tokens with truncated titles or captions, and bibliographies and author blocks are dropped app. B.2. Models trained on raw PubMed text will not match this.
• Chain-of-thought data. The MedQA CoT training set is a separate set with human explanations. It was deduplicated against the test set with 8-gram collisions at an 80% threshold app. B.5.
• Cost. 332 hours on 128 A100s, 42,496 GPU-hours, and an estimated 486 kgCO2 for the 70B run app. A. The 40,200 tokens/s throughput and 42.3% utilisation are the paper's own numbers §4.1.
• Safety. The paper's Safety Advisory, placed after the abstract, says the models are not ready for clinical use and advises against deployment in medical applications without use-case alignment and real-world testing. A separate responsible-AI appendix evaluates truthfulness, risk and bias app. E.

Lineage

• Builds on: Llama 2 (arXiv:2307.09288, no card), with Megatron-LM for distributed training.
• Compared against: Llama-2 base models, PMC-Llama-7B, Mistral-7B-Instruct, Zephyr-7B-β, Clinical-Camel-70B, Med42-70B, and commercial models (GPT-3.5, GPT-4, Med-PaLM, Med-PaLM 2) §6.2 fig. 3.