PaperlayerView exact revision on arXiv

Source-linked AI summary

Efficient Continual Pre-training by Mitigating the Stability Gap

Yiduo Guo, Jie Fu, Huishuai Zhang, Dongyan Zhao, Yikang Shen

arXiv:2406.14833v2cs.CL

TL;DR

Continual pre-training is used to adapt LLMs to new domains, but their behavior during this process has received limited study and can include an early performance drop followed by recovery. The paper examines this stability gap and proposes three strategies to address it. Its Llama-3-8B model reports the strongest performance among similarly scaled open-source baselines and outperforms GPT-3.5.

  • Problem

    Few studies focus on LLM behavior during continual pre-training, where performance can initially drop before recovering.

    Source: §1 · pp. 1–2

  • Method

    The paper proposes training a suitable corpus subset for multiple epochs, selecting high-quality tokens, and matching the pre-training data mixture.

    Source: §1 · p. 2

  • Results

    The Llama-3-8B model achieves the strongest performance among open-source baselines of similar scale and outperforms GPT-3.5.

    Source: §6 · p. 11

  • Takeaways & Limitations

    The authors conclude that mitigating the stability gap can improve domain performance and reduce computational costs.

    Source: §6 · p. 11

    The released model is not considered ready for real-world medical use.

    Source: §6 · p. 11

Abstract

from arXiv · show

Continual pre-training has increasingly become the predominant approach for adapting Large Language Models (LLMs) to new domains. This process involves updating the pre-trained LLM with a corpus from a new domain, resulting in a shift in the training distribution. To study the behavior of LLMs during this shift, we measured the model's performance throughout the continual pre-training process. we observed a temporary performance drop at the beginning, followed by a recovery phase, a phenomenon known as the "stability gap," previously noted in vision models classifying new classes. To address this issue and enhance LLM performance within a fixed compute budget, we propose three effective strategies: (1) Continually pre-training the LLM on a subset with a proper size for multiple epochs, resulting in faster performance recovery than pre-training the LLM on a large corpus in a single epoch; (2) Pre-training the LLM only on high-quality sub-corpus, which rapidly boosts domain performance; and (3) Using a data mixture similar to the pre-training data to reduce distribution gap. We conduct various experiments on Llama-family models to validate the effectiveness of our strategies in both medical continual pre-training and instruction tuning. For example, our strategies improve the average medical task performance of the OpenLlama-3B model from 36.2% to 40.7% with only 40% of the original training budget and enhance the average general task performance without causing forgetting. Furthermore, we apply our strategies to the Llama-3-8B model. The resulting model, Llama-3-Physician, achieves the best medical performance among current open-source models, and performs comparably to or even better than GPT-4 on several medical benchmarks. We release our models at \url{https://huggingface.co/YiDuo1999/Llama-3-Physician-8B-Instruct}.

1 Introduction

Continual pre-training can produce an early performance drop followed by recovery, a stability gap that delays improvement. The paper proposes three strategies to mitigate it and improve performance under a fixed compute budget.

1 Introduction

  • The initial performance drop followed by recovery occurs across LLM scales and training corpora, despite improving medical-corpus perplexity.The paper identifies this pattern as the stability gap.Source: §1 · p. 2
  • The authors hypothesize that an initially insufficient stability gradient fails to maintain instruction-following ability against the plasticity gradient for learning domain knowledge.They relate recovery to a later rise in the stability gradient.Source: §1 · p. 2
  • Training a properly sized corpus subset for multiple epochs is proposed instead of training once on a large corpus.The authors associate the one-epoch approach with a prolonged large plasticity gradient.Source: §1 · p. 2
  • Selecting the highest-quality tokens is proposed to accelerate performance recovery and raise peak performance.Source: §1 · p. 2
  • Matching the pre-training data mixture is proposed to reduce distribution shift and mitigate forgetting of general instruction-following ability.Source: §1 · p. 2
  • The proposed strategies improve Llama-3-8B performance on diverse medical tasks, outperforming open-source baselines and achieving performance comparable to GPT-4.The authors also report that their strategies accelerate performance improvement and improve peak performance on OpenLlama-3B.Source: §1 · p. 2

2 Related work

The related work covers pre-training scaling and data reuse, methods for continual pre-training, and the stability gap previously observed in vision models. The paper focuses on that gap in LLM continual pre-training.

2 Related work

  • Scaling studies examine model and dataset size, while limited high-quality tokens have motivated multi-epoch data reuse despite reported degradation after four epochs.Source: §2 · p. 3
  • Continual pre-training methods include knowledge-protecting architecture changes, learning-rate schedules and replay, token selection, mixture-of-experts expansion, and distillation.The text notes that some architecture-based techniques consume substantial compute for LLMs.Source: §2 · p. 3
  • Prior vision work reports an initial loss and gradual recovery on previous tasks; this paper reports a similar stability gap in LLM domain performance and general ability.The cited vision studies are De Lange et al. and Caccia et al.Source: §2 · p. 3

3 Identifying the stability gap in continual pre-training

The experiments find that continual pre-training can initially reduce target-task performance before it recovers, even as medical perplexity improves. The authors interpret this stability gap through stability and plasticity gradients and examine task performance and layer updates.

3 Identifying the stability gap in continual pre-training

  • The experiments measure models at regular intervals and report an initial target-task performance drop followed by a rise.Source: §3 · p. 3

3.1 Investigating the behavior of LLMs during continual pre-training

  • The medical corpus uses 50B Refined-Web tokens selected for lowest perplexity under a KenLM model trained on a high-quality medical reference corpus.OpenLlama3B-v2 is the default model, with a 50B-token compute budget.Source: §3.1 · pp. 3–4
  • Medical task performance and Pythia common-sense performance both initially fall and then rise during continual pre-training.The medical-task pattern is also observed with TinyLlama, and the paper describes the pattern as general across LLM sizes.Source: §3.1 · p. 4, Figure 2 · p. 4
  • Medical Wikipedia perplexity steadily declines, which the authors interpret as medical domain knowledge continuing to improve throughout training.Source: §3.1 · p. 4

3.2 Stability Gap: A conceptual explanation for the initial performance drop and then following recovery.

  • Lange et al. [2022] describe the stability gap as an initial drop on previous tasks followed by recovery, associated with an imbalance between plasticity and stability gradients.The plasticity gradient learns the new task; the stability gradient maintains previous-task performance.Source: §3.2 · pp. 4–5
  • The authors hypothesize that a small initial stability gradient weakens instruction-following ability, while later recovery accompanies increased stability and learned domain knowledge.Their verification finds a similar V-shaped general-task curve and initially larger relative updates in bottom than top layers, at a rate above 1.35.Source: §3.2 · p. 5, Figure 3 · p. 5

4 Efficient continual pretraining strategies for mitigating the stability gap

The section proposes three strategies for mitigating the stability gap: train a corpus subset for multiple epochs, prioritize high-quality samples, and preserve a pre-training-like data mixture. The mixture strategies improve first-epoch performance, with dynamic sampling attaining a higher peak.

Strategy I

  • Training a randomly selected corpus subset across multiple epochs accelerates performance recovery, with peak performance reached at epoch four.This contrasts with training on the entire large corpus for one epoch.Source: §4 · pp. 5–6

Strategy II

  • Selecting the lowest-perplexity subset as the highest-quality data enables faster recovery and stronger medical-domain performance.Perplexity is calculated with the KenLM trained in Section 3.1; lower perplexity indicates greater closeness to the medical reference distribution.Source: §4 · p. 6

Strategy III

  • The mixture strategy follows the pre-training data rate, replacing CC and C4 tokens—82% of the initial five billion—with high-quality medical tokens sampled either once or anew each epoch.The two variants are rate-fixed-data-fixed and rate-fixed-data-dynamic.Source: §4 · p. 6
  • Both mixture-rate methods improve first-epoch medical performance; dynamic sampling reaches a higher peak, and the strategies also improve commonsense performance (Figure 4).The figure reports medical performance in panels (a) and (b), and commonsense performance in panel (c).Source: §4 · p. 6, Figure 4 · p. 6

5 Evaluation

The evaluation compares continual pre-training strategies, analyzes learning factors, and deploys them in Llama-3-8B pre-training and instruction tuning.

5.1 Comparison with other continual pre-training techniques

  • Our strategies raise the base model’s average medical task performance by 4.5% using 20 billion training tokens.Mixture-rate-matched replay is more effective than randomly sampled replay, while freezing the bottom five layers hinders medical learning.Source: Table 1 · p. 7, §5.1 · p. 7

5.2 Analysis

  • The optimal learning rate varies with model scale: TinyLlama-1.1B benefits from 3e-4, while OpenLlama-3B benefits from 3e-5.Rates that are too low impede domain learning, while rates that are too high can harm general instruction following.Source: §5.2 · p. 7, Figure 5 · p. 8
  • A 5-billion-token subset balances early recovery, peak performance, and compute efficiency; smaller subsets can overfit, while larger ones recover more slowly.The subset is selected from medical tokens ranked by KenLM perplexity.Source: §5.2 · p. 8, Figure 5 · p. 8

5.3 Deploying our strategies into the Llama-3 Model

  • Llama-3-Physician-8B beats similar-scale baselines across four medical benchmarks and outperforms many 70B models on average.Continual pre-training alone slightly lowers average medical performance, despite significantly reducing medical perplexity.Source: §5.3 · pp. 8–9, Table 2 · p. 9

5.4 Deploying our strategies into the instruction tuning process

  • With 25% instruction data, Llama-3-Physician-8B achieves the best diverse-task performance, beats open-source QA baselines, approaches GPT-4 on average, and exceeds GPT-4 on four other medical task types.The 25% mixture combines high-quality selected medical instructions with general instructions.Source: §5.4 · pp. 9–10, Table 3 · p. 10, Figure 6 · p. 10

6 Conclusion

The paper identifies the stability gap in continual pre-training and proposes strategies to improve domain performance while reducing computational costs. The authors also report strong Llama-3-8B results, while noting limits on model scale and real-world medical use.

6 Conclusion

  • The strategies address the stability gap and improve domain performance while reducing computational costs.The paper reports that Llama-3-8B achieves the strongest performance among open-source baselines of similar scale and outperforms GPT-3.5.Source: §6 · p. 11
  • The authors do not consider their model ready for real-world medical use and plan to verify their conclusions on larger LLMs.They also note that the pre-training data mixture is unavailable for most strong open-source LLMs, although their Llama-3-8B experiment improves performance in this scenario.Source: §6 · p. 11

A The Details of Pre-training

The pre-training setup specifies models, training data, and optimization choices, with training performed on 192 V100 GPUs using a 2048-token context.

A The Details of Pre-training

  • OpenLLaMa-3B and TinyLLaMa-1B were trained for one epoch on 50 billion medical tokens, while Pythia-410m used 100 billion RefinedWeb tokens from 2021–2022.The RefinedWeb subset was selected because the Pile contains only pre-2021 data and those tokens were considered unseen by Pythia-410m.Source: A The Details of Pre-training · p. 18
  • Training used a 2048-token context, 192 V100 GPUs, and batches of about 340 million tokens with four-step gradient accumulation.The task was next-token prediction.Source: A The Details of Pre-training · p. 18
  • The setup used AdamW with β1 = 0.9, β2 = 0.95, weight decay 0.01, and a learning rate of 3e-4.A cosine scheduler used a 0.1 warmup ratio, and training used bf16 precision.Source: A The Details of Pre-training · p. 18
  • The pre-training code was based on transformers and added FlashAttention-2 support for more efficient inference and long-context decoding.The passage attributes FlashAttention-2 to Dao, 2023.Source: A The Details of Pre-training · p. 18

B Task Information

The appendix describes the medical and commonsense task datasets used for evaluation, alongside dataset statistics and a zero-shot evaluation setup.

B Task Information

  • The appendix lists four medical tasks and ten commonsense tasks, reports dataset statistics in Table 5, and evaluates test sets zero-shot with lm-eval-harness.The passage identifying the number of medical tasks does not name them.Source: B Task Information · pp. 18–20, Table 5 · p. 19
  • Medical evaluation includes MedMCQA, a diverse medical entrance-exam question set, and MedQA, questions from professional US medical board exams.MedMCQA contains over 194,000 questions across 2,400 healthcare topics and 21 subjects; MedQA contains 12,723 questions sourced from 18 English medical textbooks.Source: B Task Information · p. 18
  • Other medical tasks cover selected MMLU medical reasoning subjects, PubMed abstract question answering, cancer-hallmark classification, drug-interaction relations, biomedical inference, and radiology impression generation.The MMLU subset includes clinical knowledge, college medicine, medical genetics, and professional medicine.Source: B Task Information · pp. 18–19
  • Commonsense evaluation includes science-exam questions, yes/no reading questions, causal-choice questions, grounded next-event prediction, and elementary science questions.These task descriptions correspond to ARC, BoolQ, COPA, HellaSWAG, and OpenBookQA.Source: B Task Information · p. 19
  • The remaining described commonsense datasets test physical reasoning, reading comprehension, science questions, and binary fill-in-the-blank commonsense reasoning.These are PIQA, RACE, SciQ, and WinoGrande; WinoGrande is a 44k-problem collection adjusted for scale and robustness against dataset-specific bias.Source: B Task Information · pp. 19–20

C The Perplexity and relative parameter update rate of the LLM using our strategies

Figure 7 reports that the strategies lower OpenLLaMa-3B's average medical perplexity over training and bring its bottom-to-top layer update ratio closer to 1.

C The Perplexity and relative parameter update rate of the LLM using our strategies

  • By the fourth epoch, the strategies' average medical perplexity is lower than that of the OpenLLaMa-3B baseline trained on 50 billion medical tokens.Figure 7(a) tracks average medical perplexity; the prose says it gradually decreases as the model acquires medical knowledge.Source: Figure 7 · p. 20, C The Perplexity and relative… · p. 20
  • The strategies bring the ratio of average relative updates in the bottom five versus top five layers closer to 1 than the 50-billion-token baseline.The paper interprets this as more balanced plasticity and stability gradients.Source: Figure 7 · p. 20, C The Perplexity and relative… · p. 20

D Deploying our strategies into the general continual pre-training setting

In Pythia-410m continual pre-training on RefinedWeb, multi-epoch subset training mitigated the stability gap, and strategies I and III achieved higher performance.

D Deploying our strategies into the general continual pre-training setting

E The Training Details of Deploying our Strategies into the Llama-3 Model

The Llama-3 training details specify the pre-training, task-specific fine-tuning, and instruction-tuning configurations.

E The Training Details of Deploying our Strategies into the Llama-3 Model

  • Pre-training predicts the next token with context size 8192, using 16 H100 80GB GPUs and an AdamW optimizer.The learning rate is 3e-5, with cosine scheduling, a 0.1 warmup ratio, bf16 precision, 12-step gradient accumulation, and batches of about 340 million tokens.Source: E The Training Details of Deploying… · p. 21
  • Task-specific fine-tuning uses AdamW with weight decay 0.01 and learning rate 3e-5 for 3 epochs.It uses cosine scheduling with a 10% warmup ratio and decays the final learning rate to 10% of peak.Source: E The Training Details of Deploying… · p. 21
  • Instruction tuning combines seven medical-task training sets and runs for 3 epochs with global batch size 96 and maximum sequence length 1024.For MedQA, PubMedQA, and MedMCQA, training samples containing test-sample options are removed to avoid potential contamination.Source: E The Training Details of Deploying… · p. 21

NeurIPS Paper Checklist

The checklist records the authors’ disclosures on reproducibility, experimental reporting, ethics, societal impact, and research assets, alongside standard guidance for these topics.

NeurIPS Paper Checklist

  • The authors report that experimental details are in the evaluation section and appendixes, and say they will release the models publicly.They state that fine-tuning data is open-source, pre-training data is too large to include, and code and relevant data will be released soon.Source: NeurIPS Paper Checklist · pp. 23–24
  • The authors state that training and test details are provided in the evaluation section and appendixes, and direct readers there for compute resources.The checklist marks both experimental-setting details and compute-resource information as Yes.Source: NeurIPS Paper Checklist · p. 25
  • The checklist response marks experimental statistical significance as No, with its justification left as [TODO].The checklist guidance asks for error bars, confidence intervals, or significance tests and explanations of their variability and calculation.Source: NeurIPS Paper Checklist · p. 25
  • The authors report compliance with the NeurIPS Code of Ethics and say societal impacts and safeguards are discussed after the conclusion.The checklist marks these ethics, impact, and safeguard questions Yes.Source: NeurIPS Paper Checklist · p. 26
  • The authors say they properly credit existing assets and document their new models in the main paper; human-subjects research is marked Not Applicable.The checklist also marks crowdsourcing and human-subjects research as not involved.Source: NeurIPS Paper Checklist · pp. 27–28
Loading 2406.14833v2…