PaperlayerView exact revision on arXiv

Source-linked AI summary

Simple and Scalable Strategies to Continually Pre-train Large Language Models

Adam Ibrahim, Benjamin Thérien, Kshitij Gupta, Mats L. Richter, Quentin Anthony, Timothée Lesort, Eugene Belilovsky, Irina Rish

arXiv:2403.08763v4cs.LGcs.AIcs.CL

TL;DR

To avoid costly retraining when new data arrives, the paper studies continual pre-training with learning-rate re-warming, re-decaying, and replay. This combination matches average retraining performance across tested shifts and model sizes while using less compute.

  • Problem

    Retraining LLMs on old and newly collected data for every update is costly, motivating the question of whether continual strategies can match union-data training.

    Source: §1 · p. 2

  • Method

    The paper combines learning-rate re-warming and re-decaying with replay of previous data during continual pre-training.

    Source: §1 · p. 3

  • Results

    Continual pre-training achieved average performance on par with retraining from scratch on all data across two shifts and two model sizes, using significantly less compute.

    Source: §1 · p. 3, §9 · p. 24

  • Takeaways & Limitations

    The reported strategies provide a simple, scalable way to update LLMs while matching the retraining baseline in the studied settings.

    Source: §1 · p. 3, §9 · p. 24

    The study covers only two model sizes, primarily two sequential tasks, and infinite schedules only at 405M scale without distribution shift.

    Source: §8 · p. 24

Abstract

from arXiv · show

Large language models (LLMs) are routinely pre-trained on billions of tokens, only to start the process over again once new data becomes available. A much more efficient solution is to continually pre-train these models, saving significant compute compared to re-training. However, the distribution shift induced by new data typically results in degraded performance on previous data or poor adaptation to the new data. In this work, we show that a simple and scalable combination of learning rate (LR) re-warming, LR re-decaying, and replay of previous data is sufficient to match the performance of fully re-training from scratch on all available data, as measured by the final loss and the average score on several language model (LM) evaluation benchmarks. Specifically, we show this for a weak but realistic distribution shift between two commonly used LLM pre-training datasets (English$\rightarrow$English) and a stronger distribution shift (English$\rightarrow$German) at the $405$M parameter model scale with large dataset sizes (hundreds of billions of tokens). Selecting the weak but realistic shift for larger-scale experiments, we also find that our continual learning strategies match the re-training baseline for a 10B parameter LLM. Our results demonstrate that LLMs can be successfully updated via simple and scalable continual learning strategies, matching the re-training baseline using only a fraction of the compute. Finally, inspired by previous work, we propose alternatives to the cosine learning rate schedule that help circumvent forgetting induced by LR re-warming and that are not bound to a fixed token budget.

1 Introduction

Because retraining on each update is costly and naive continuation can adapt poorly or forget previous data, the paper tests scalable continual-pretraining strategies against union-data retraining. Across two model sizes and distribution shifts, learning-rate re-warming, re-decaying, and replay match average retraining performance with less compute.

1 Introduction

  • Naive continuation can under-adapt to new data or cause catastrophic forgetting, while retraining on each update is expensive.The study asks whether continual strategies can match the baseline trained on the union of existing and new data.Source: §1 · p. 2
  • For cosine-pretrained decoder-only transformers, re-warming and re-decaying the learning rate is necessary for adaptation during continual pre-training.Source: §1 · p. 2
  • Across two model sizes and distribution shifts, LR re-warming, re-decaying, and compute-equivalent replay match union-data retraining’s average performance with significantly less compute.Figure 1 illustrates similar average performance at substantially lower compute for its reported settings.Source: Figure 1 · p. 3, §1 · p. 3
  • With an appropriate replay amount, forgetting can be significantly mitigated even when updating decoder-only transformers on hundreds of billions of new tokens.Source: §1 · p. 3
  • The paper proposes infinite learning-rate schedules as a way to smooth transitions across datasets and address optimization difficulties associated with re-warming.Source: §1 · p. 3

2 Main Findings and Takeaways and Examples our Method’s Practicality

The paper reports that continual pre-training can match retraining across weak and strong distribution shifts and two model sizes under its experimental setting. Its practical guidance covers learning-rate choices, replay fractions, and alternative schedules.

2 Main Findings and Takeaways and Examples our Method’s Practicality

  • With sequential pre-training phases of at least 100B tokens, appropriate LR re-warming, re-decaying, and replay match retraining across weak and strong shifts and two model sizes.Source: §2 · pp. 3–4

Learning rate schedule:

  • For cosine-decayed models, re-warming and re-decaying improves adaptation; lowering the schedule maximum can reduce forgetting, while raising it can improve adaptation.Source: Learning rate schedule · p. 4
  • Infinite LR schedules transition to a high constant rate across tasks, avoiding between-task re-warming and fixed token budgets.A final exponential decay can be used to train to convergence at any point.Source: Learning rate schedule · p. 4
  • Even small replay amounts can mitigate forgetting; experimenting with fractions early can help find a balance between adapting and retaining prior performance.Source: Learning rate schedule · p. 4
  • Glorioso et al. (2024) and DeepSeek-AI et al. (2024) applied combinations of these techniques at scale, reporting improved performance without catastrophic forgetting and gains in code generation while retaining most natural-language abilities.Their examples use 60% replay on 50B tokens and 30% replay on 6T tokens, respectively.Source: Learning rate schedule · p. 4

3 Related Work

Related work frames continual learning as balancing adaptation to novel data with retention of prior knowledge. The paper situates its general-purpose, large-scale setting alongside work on model scale, domain adaptation, task-specific continual pre-training, and learning-rate schedules.

3.1 Continual learning

  • Continual learning balances adapting to novel data with retaining prior knowledge, and may reduce large-scale retraining costs; this work combines learning-rate scheduling with replay for LLM pretraining.Source: §3.1 · p. 5

3.2 Pre-training, Model Scale, and Continual Learning

  • Prior studies report that pretraining and larger model scale can reduce forgetting, while transfer to new tasks improves with parameter count.Source: §3.2 · p. 5

3.3 Domain Adaptive Continual Pre-training (DACPT)

  • Unlike domain-specific DACPT, this work uses general-purpose datasets; LLaMA Pro instead grows the model and adjusts only new weights.Source: §3.3 · pp. 5–6

3.4 Continual Learning for LMs Applied to Specific Domains

  • Most related task- or domain-specific studies use smaller models and datasets; Gogoulou et al. (2023) approach the data scale, while Yang et al. (2024) use a similar setup without controlled systematic evaluation.Source: §3.4 · p. 6

3.5 Learning Rate Schedules

  • Cosine schedules are widely adopted for LLM pretraining, while inverse-square-root and infinite schedules offer flexibility in training duration.The proposed infinite schedules for LLMs are inspired by vision-transformer work on indefinite training.Source: §3.5 · p. 6

4 Background & Methodology

This section introduces the learning-rate schedule background and defines how replay comparisons keep training compute equivalent. It describes the warmup and cosine-annealing phases, then explains replay’s token-budget adjustment.

4.1 Linear Warmup and Cosine Decay Schedule

  • The cosine recipe warms the learning rate linearly, then anneals it; the illustrated 10% warmup is illustrative, since common warmups span 0.1%–0.5% of training steps.Source: §4.1 · pp. 6–7, Figure 2 · p. 7

4.2 Compute-equivalent Replay

  • Compute-equivalent replay replaces some new-dataset tokens with replayed prior data while keeping total training tokens and compute fixed; the example uses 5% replay.The example allocates 95B unique tokens to D1 and 5B replayed tokens from D0, for 200B total tokens.Source: §4.2 · p. 7
  • Replay samples follow their original pretraining order; preliminary tests found no noticeable reshuffling differences, while sample-selection methods remain future work.Source: §4.2 · p. 7

5 Experimental Setup

The experiments compare continual pre-training strategies with baselines in a controlled setting, without seeking state-of-the-art performance. They cover multiple datasets, sequential-training settings, model sizes, and English and German evaluations.

5 Experimental Setup

  • The experiments aim for a controlled comparison with baselines, not state-of-the-art performance or comparisons to models outside the study's scope.Source: §5 · p. 8

5.1 Datasets

  • SlimPajama, German Common Crawl, and Pile provide the training and validation data; SlimPajama is subsampled to 300B tokens, with domain sizes and batch-sampling percentages reported in Table 1.Source: §5.1 · p. 8, Table 1 · p. 8

5.2 Continual Learning Settings

  • The two-dataset settings assume sequential pre-training phases with at least 100B tokens per dataset and reset optimizer states; they compare Pile-to-SlimPajama English data with Pile-to-German Common Crawl.The English transition is described as a weak but realistic shift, while the language change to German constitutes a stronger shift.Source: §5.2 · pp. 8–9
  • Additional settings test three sequential 100B-token SlimPajama splits and domain-ordered SlimPajama data; the latter is described as unusually difficult compared with mixing domains.Source: §5.2 · p. 9

5.3 Training Setup

  • The study trains decoder-only transformers at 405M and 9.6B parameters using GPT-NeoX, AdamW, and a shared Pile-trained tokenizer.Source: §5.3 · pp. 9–10

5.4 German and English LM Evaluation Benchmark

  • English evaluations cover commonsense reasoning, world knowledge, reading comprehension, math, and MMLU.Source: §5.4 · p. 10
  • German evaluations include MMLU-DE, commonsense reasoning benchmarks, and TriviaQA-DE; the German benchmarks translate English counterparts using GPT 3.5 API.Source: §5.4 · p. 10

6 Results

The results evaluate continual pre-training with LR re-warming, LR re-decaying, and replay across data shifts and at 10B scale, reporting final loss and benchmark performance comparable to joint training or union-data retraining. They also examine replay’s effect on forgetting and adaptation cost.

6.1.2 The effect of re-warming, re-decaying, and varying ηmax for Weak and Strong Distribution Shifts

  • Shorter warmups initially accelerate adaptation and forgetting, but differences diminish by 50B tokens; re-warming and re-decaying improve adaptation, with higher ηmax increasing both adaptation and forgetting.The stronger shift also exacerbates forgetting and adaptation; warmup duration ultimately has little impact on either outcome.Source: §6.1.1 · p. 12, Figure 4 · p. 12, §6.1.2 · p. 13

6.2 The Effect of Replay

  • 1% replay significantly reduces forgetting on Pile across both shifts with little downstream-performance impact, while 50% replay worsens D1 adaptation but attains or surpasses joint-training average validation performance.The 50%-replay models saw 150B fewer SlimPajama tokens and 100B fewer German tokens than their respective joint-training baselines.Source: §6.2 · pp. 14–15, Figure 5 · p. 14, Table 2 · p. 15

6.3.1 Final Performance Evaluated by Loss

  • Continual models with re-warming, re-decaying, and replay achieve nearly identical or identical average final validation loss to joint training on weak and strong shifts.For the stronger shift, 25% replay reduces Pile forgetting by 1.23 in final loss; replay minimally affects downstream performance compared with no replay.Source: §6.3.1 · pp. 15–17, Table 3 · p. 16

6.3.2 Final Performance Evaluated by Zero-shot and Few-shot Results on Popular LM Benchmarks

  • On English evaluations, the continual method improves over joint training, while German-trained continual models outperform German-only and joint baselines on German HellaSwag.English evaluations show a 1.19-point range across English models, which the authors suspect may reflect noise; qualitative German generations also improve over the Pile baseline despite poor benchmark performance at 405M scale.Source: §6.3.2 · p. 17

6.4.1 Final Performance Evaluated by Loss

  • 0.02 was the average final validation-loss difference from joint training for 5%-replay continual models at both 10B and 405M parameters.At 10B and 405M, 5% replay reduced Pile forgetting by 0.19 and 0.21, respectively; no-replay models forgot 0.23 and 0.27 nats, respectively.Source: §6.4.1 · p. 19, Table 4 · p. 19

6.4.2 Final Performance Evaluated by Zero-shot and Few-shot Results on Popular LM Benchmarks

  • At 405M and 10B, continual pre-training with LR re-warming, LR re-decaying, and replay exceeds single-dataset baselines and matches expensive union-data retraining on average.The reported comparison includes average final validation loss and evaluation accuracy.Source: Table 5 · p. 19, §6.4.2 · p. 20
  • With 5% replay, 405M models surpass union-trained models in average accuracy, while 10B models approach them; task-level winners alternate between replay and union models.At 10B, OpenBookQA favors replay and MMLU favors union training; at 405M, task-level differences are not notable. The authors suspect the 10B MMLU gap reflects limited training data.Source: §6.4.2 · p. 20
  • Across the reported benchmark results, 10B models outperform 405M models, models trained on more tokens perform better, and English-only 405M models outperform German-trained counterparts.HellaSwag uses normalized accuracy, NaturalQuestions and TriviaQA use exact match, and other tasks use unnormalized accuracy.Source: Table 5 · p. 19, §6.4.2 · p. 20
  • Glorioso et al. (2024) later applied these techniques without MMLU degradation and reported improved MMLU during continual pre-training.Source: §6.4.2 · p. 20

7 Understanding and Circumventing the Pathologies of Re-warming

Re-warming can increase loss even without a distribution shift, motivating alternatives to cosine schedules that avoid re-warming while supporting continued pre-training. In the reported single-dataset and same-distribution experiments, infinite learning rate schedules perform similarly to cosine decay and show negligible forgetting across dataset boundaries.

7.1 Re-warming on the Same Data

  • Re-warming increased Pile validation loss even without distribution shift: by 0.1 at ηmax = 3 · 10−4 and by 0.2 at ηmax = 6 · 10−4.The models did not recover quickly from the performance hit, motivating alternatives to schedules that require re-warming.Source: §7.1 · p. 20, Figure 8 · p. 21

7.2 Infinite Learning Rate Schedules

  • Infinite schedules keep the learning rate constant across future tasks after an initial warm-up and cooldown, avoiding re-warming and allowing annealing when pre-training is ready to end.Their four phases are warm-up, cooldown to a constant rate, constant-rate training, and final annealing; the constant-phase checkpoint is used to continue on a new dataset.Source: §7.2 · p. 21

7.3 Comparing Cosine Decay to Variants of our Infinite Schedules

  • All three 405M models reached similar final validation loss on 300B SlimPajama tokens, showing infinite schedules can match cosine decay in single-dataset pre-training.A pre-annealing checkpoint can be resumed, and annealing can begin at any time during the constant phase.Source: §7.3 · p. 22, Figure 9 · p. 23

7.4 Infinite Learning Rate Schedules: Scaling to Infinite Future Updates

  • Across three IID 100B-token SlimPajama splits, infinite schedules showed negligible forgetting at dataset boundaries and immediately recovered from loss increases after optimizer-state resets.The experiments cover same-distribution updates; the authors leave distribution shifts and stability over long constant-rate phases for future study.Source: §7.4 · pp. 22–23
  • The experiments leave open whether infinite schedules work under distribution shifts and remain stable over long constant-rate phases.The authors identify both questions as future work.Source: §7.4 · pp. 23–24

8 Limitations

The study is limited in model scale, task sequence length, and experimental coverage; some dataset and schedule settings also leave uncertainty.

8 Limitations

  • Experiments cover only 405M- and 10B-parameter models, not another order of magnitude such as 100B.The authors cite computational limitations.Source: §8 · p. 24
  • The study covers transitions between at most two subsequent tasks, rather than longer task sequences.Source: §8 · p. 24
  • The German validation data was not deduplicated against training data, so some contamination may remain.The authors used distinct shards but acknowledge possible overlap.Source: §8 · p. 24
  • Experiments were not run over multiple seeds, leaving a stochastic element in some results.The authors expect relatively stable results given large batches and consistent data order.Source: §8 · p. 24
  • Infinite schedules were tested only at 405M scale without distribution shift; their behavior over more tokens or shifted datasets remains unclear.The authors also flag stability at larger scale, particularly during constant learning-rate training without µP.Source: §8 · p. 24

9 Conclusion

The conclusion reports that continual pre-training with learning-rate re-warming, re-decaying, and replay matches retraining performance across the studied shifts and scales. The accompanying sections discuss practical and energy implications, future directions, and extended evaluation materials.

9 Conclusion

  • Continual pre-training with learning-rate re-warming, re-decaying, and replay matches retraining performance across two distribution shifts and 405M- and 10B-parameter models.The conclusion also identifies re-warming and re-decaying as important for adaptation and replay as a way to mitigate forgetting.Source: §9 · p. 24
  • The authors describe continual pre-training as an efficient alternative for updating decoder-only transformer LLMs on new data, including higher-quality datasets.They suggest the strategies may also apply to pre-training curricula and identify infinite schedules, model growth, tokenizer adaptation, and multimodal settings for future study.Source: §9 · pp. 24–25

Broader Impact Statement

  • Continual pre-training is presented as a way to reduce compute and energy associated with updating foundation models, although the experiments themselves are environmentally costly.Source: Broader Impact Statement · p. 25

A Extended results

  • The extended results section introduces domain-incremental results and German-language models, and reports aggregated evaluation and final-loss tables for all models.Source: A Extended results · p. 36

A.1 Domain Incremental continual pre-training

This appendix studies sequential domain-by-domain continual pre-training, using replay and learning-rate schedule changes across successive data domains. The authors describe this setting as difficult and of limited practical use compared with consuming data as a mixture.

A.1 Domain Incremental continual pre-training

  • The domain-incremental setup processes distinct datasets sequentially, with Pile as D0 and SlimPajama as D1, to model data arriving from different domains over time.The authors characterize distribution shifts at domain transitions as a particular challenge.Source: §A.1 · p. 36
  • For multiple tasks, discrete reservoir sampling updates the replay buffer at each dataset transition, retaining data from prior datasets according to the replay ratio.The buffer is updated when training moves to a new dataset.Source: §A.1 · p. 36
  • Figure 11’s 300B-token SlimPajama experiment processes domains from largest to smallest and uses re-warming, re-decaying, and replay; total data processed matches the 300B-token schedule baseline, but order and schedule differ.Source: Figure 11 · p. 36, §A.1 · p. 37
  • The authors hypothesize that repeated learning-rate warmups and strong distribution shifts make it difficult to achieve strong performance across domains, especially when adapting to new data.They suggest less frequent warmup and decay or infinite schedules may help, but say further research is needed.Source: §A.1 · p. 37
  • The authors consider sequential domain ingestion of limited practical use compared with consuming the data as a mixture.Source: §A.1 · p. 37

A.2 Model Merging v.s. Continual Pre-training

Preliminary TIES experiments compare model merging with continual pre-training. The authors report that continual pre-training generally performs better at the same compute, while merging continual-pretraining checkpoints does not improve performance.

A.2 Model Merging v.s. Continual Pre-training

  • Continual pre-training generally outperforms merging separately trained models at the same compute, while merging continual-pretraining checkpoints does not increase performance.The comparison is described as preliminary; the English evaluation table reports normalized accuracy for HellaSwag and exact match for NaturalQuestions and TriviaQA, with unnormalized accuracy for other tasks.Source: §A.2 · p. 37, Table 6 · p. 37
  • The authors caution that the merging experiments are not comprehensive and that other techniques may perform better.Source: §A.2 · p. 37

A.3 Replay for Different Dataset Sizes

The LR re-warming, re-decaying, and replay strategy remains effective across continued-training dataset sizes of 100B, 200B, and 300B tokens. More training tokens consistently reduce final loss.

A.3 Replay for Different Dataset Sizes

  • The combination of LR re-warming, LR re-decaying, and replay remains effective across dataset sizes of 100B, 200B, and 300B tokens.The schedules are fitted to each altered training length.Source: §A.3 · p. 37
  • Final SlimPajama loss consistently decreases as the continued-training token count increases.Figure 12 compares models with and without 5% replay at each dataset size, starting from a 300B-token Pile checkpoint.Source: §A.3 · p. 37, Figure 12 · p. 38

A.4 The effect of resetting optimizer states at dataset transitions

In the tested Pile-to-SlimPajama transition, retaining optimizer states does not change the final loss relative to discarding them. The starting moment estimates quickly decay during optimization.

A.4 The effect of resetting optimizer states at dataset transitions

  • The authors drop optimizer states between dataset transitions to reflect open-source LLM checkpoints, which rarely include them.The experiment tests whether retaining them can resolve instabilities during continual pre-training.Source: §A.4 · p. 38
  • After 40B tokens, models that kept or dropped optimizer states reached similar loss values, with differences attributed to randomness.The comparison used 405M-parameter models starting from a 300B-token Pile checkpoint and trained on 40B SlimPajama tokens.Source: §A.4 · p. 38
  • Keeping optimizer states did not affect continual pre-training relative to discarding them in this setting.The authors report that after 1,000 optimization steps, pre-training moment estimates contribute less than 0.0044% of current moment estimates.Source: §A.4 · p. 38

A.5 Qualitative evaluation of German models

The German qualitative evaluation finds an observable generative-quality effect from German Common Crawl training, though its small, low-quality sample set limits robust model-to-model conclusions. Aggregated evaluations report differing patterns across English and German benchmarks, and the appendix summarizes final losses and hyperparameters.

A.5 Qualitative evaluation of German models

  • German Common Crawl models generally produced grammatically correct sentences, but outputs were repetitive and the small, low-quality sample set precluded robust comparisons among them.The Pile baseline more often had grammatical errors and failed to respect prompt context; the authors conclude German Common Crawl training had an observable effect on German generative quality.Source: §A.5 · p. 39

A.6 Aggregated LM Evaluation Results

  • English-only 405M models outperformed German-trained counterparts on English tasks, while English models matched German counterparts on average across German tasks.For German benchmarks, performance was near chance on MMLU and ARC-C; German models did better on HellaSwag and English models on Truthful QA.Source: Table 11 · p. 42, Table 10 · p. 42

A.7 Aggregated average final accuracy

  • The 50%-replay model achieved nearly identical loss values while using 150B fewer SlimPajama tokens.The table reports that more replay improves the adaptation-forgetting trade-off in average loss.Source: Table 12 · p. 43

B Model hyperparameters

  • Tables 13 and 14 present the 405M and 10B model hyperparameters and the default learning-rate schedule hyperparameters, respectively.Source: Table 13 · p. 44, Table 14 · p. 45
Loading 2403.08763v4…