PaperlayerView exact revision on arXiv

Source-linked AI summary

Replaying pre-training data improves fine-tuning

Suhas Kotha, Percy Liang

arXiv:2603.04964v1cs.CLcs.LG

TL;DR

Standard training pre-trains on generic data before fine-tuning on limited target data, but the effect of replaying generic data during adaptation is unclear. This paper evaluates replay across controlled schedules and larger-scale applications, finding improved target-data efficiency and downstream performance, especially when target data is scarce during pre-training.

  • Problem

    Current practice typically pre-trains on generic web text and then fine-tunes on limited target data, leaving open whether alternative data schedules improve target-domain performance.

  • Method

    The paper evaluates generic-data replay in controlled fine-tuning and mid-training schedules, then tests the recommendation by fine-tuning an 8B-parameter Llama 3 model on downstream tasks.

  • Results

    Generic replay increases target-data efficiency by up to 1.87× for fine-tuning and 2.06× for mid-training, while improving web-navigation success by 4.5% and Basque question-answering accuracy by 2%.

  • Takeaways & Limitations

    Replay is especially helpful when less target data appears during pre-training, and it can improve target performance even when only fine-tuning is changed.

  • Takeaways & Limitations

    The controlled study simplifies pre-training to two distributions and uses validation loss, which might not perfectly correlate with downstream metrics; replay also requires increased compute.

Abstract

from arXiv · show

To obtain a language model for a target domain (e.g. math), the current paradigm is to pre-train on a vast amount of generic web text and then fine-tune on the relatively limited amount of target data. Typically, generic data is only mixed in during fine-tuning to prevent catastrophic forgetting of the generic domain. We surprisingly find that replaying the generic data during fine-tuning can actually improve performance on the (less related) target task. Concretely, in a controlled pre-training environment with 4M target tokens, 4B total tokens, and 150M parameter models, generic replay increases target data efficiency by up to $1.87\times$ for fine-tuning and $2.06\times$ for mid-training. We further analyze data schedules that introduce target data during pre-training and find that replay helps more when there is less target data present in pre-training. We demonstrate the success of replay in practice for fine-tuning 8B parameter models, improving agentic web navigation success by $4.5\%$ and Basque question-answering accuracy by $2\%$.

1 Introduction

The paper asks whether changing the generic-to-target data schedule can improve target-domain performance. Generic replay improves target data efficiency in controlled fine-tuning and mid-training, especially when little target data appeared during pre-training, and transfers to 8B-model applications.

  • The paper tests whether alternative data schedules can improve performance when generic web text precedes limited target-domain fine-tuning.
  • 1.87× target data-efficiency improvement is obtained by replaying generic data during controlled fine-tuning.
  • 2.06× target data-efficiency improvement is obtained from replay in the mid-training setting.
  • Replay helps more when less target data is present during the first training stage.
  • 4.5% higher web-navigation success and 2% higher Basque question-answering accuracy demonstrate replay benefits for fine-tuning 8B-parameter models.

2 Controlled pre-training setup

The controlled setup models generic pre-training followed by target-domain adaptation with compute-matched schedules. It evaluates strategies using target validation loss and a scaling-law-derived measure of effective target data.

  • Data and training: The experiments use C4 as generic data and FineMath, StarCoder, and Flan as target domains with varying overlap.
  • Data and training: The controlled environment uses 150M-parameter Llama-style models, 4M target tokens, and 4B total training tokens.
  • Evaluation: Validation loss on held-out target data is the primary performance measure because it scales smoothly at the studied model size.
  • Evaluation: Data efficiency estimates invert a reference strategy’s target-data scaling law to calculate effective data for a given loss.
  • Data and training: The controlled comparisons use compute-matched schedules with independently configured training stages and optimizer-state resets.

3 Modifying fine-tuning

The fine-tuning intervention mixes generic data into the target-training stage while preserving total training steps. Across target domains, a nonzero replay fraction improves target loss and data efficiency, with domain-dependent tolerance.

  • Modifying fine-tuning: The study compares standard two-stage fine-tuning with schedules that mix generic data into the final target-training stage.
  • Replay results: 1.87×, 1.49×, and 1.09× data efficiency are achieved for Flan, FineMath, and StarCoder, respectively, at domain-specific nonzero replay fractions.
  • Replay results: Code tolerates less replay than math and instruction-following data because C4 explicitly filters out code.
  • Interpretation: Replay improves the new target training task, contrasting with its usual role in preventing catastrophic forgetting of old tasks.

4 Modifying mid-training and pre-training

The paper expands mid-training beyond replay-only fine-tuning by varying when target data appears and how much generic data is replayed. Replay improves data efficiency, especially when target data is scarce during pre-training, while earlier target-data exposure can provide further gains.

  • Data schedule space: A unified WSD schedule supports two-stage experiments that vary Stage 2 replay and the allocation of target data between Stages 1 and 2.The schedule uses replay fraction ρ and Stage 2 allocation α as its two degrees of freedom.
  • Mid-training baseline: The mid-training baseline improves data efficiency over standard fine-tuning by 9.92× for Starcoder, 6.37× for FineMath, and 2.77× for Flan.The authors attribute this likely advantage to joint training without resetting optimizer state or rewarming the learning rate.
  • Searching over two-stage schedules: Replay alone improves data efficiency over the mid-training baseline by 1.53× for StarCoder, 1.85× for FineMath, and 2.06× for Flan.These results come from sweeping replay fraction while keeping all target data in Stage 2.
  • Searching over two-stage schedules: Allowing target data in Stage 1 raises data-efficiency improvements to 1.53× for StarCoder, 2.49× for FineMath, and 4.80× for Flan over the same baseline.The full sweep varies both replay fraction ρ and Stage 2 target-data allocation α.
  • Interaction between replay and pre-training: Replay matters most when target data is less present during pre-training, becoming critical when no target data appears in Stage 1 and less helpful when 75% is retained there.The pattern holds across StarCoder, FineMath, and Flan, and replay can become less important when α < 1.0.

5 Recommendations for post-training practice

The paper recommends replaying generic data during fine-tuning when pre-training cannot be changed, especially for scarce target domains. Tests with 8B models find gains in web navigation and Basque question answering.

  • 5 Recommendations for post-training practice: Replay can improve target performance when changing pre-training is impractical and the target distribution is scarce during pre-training.The recommendation applies to settings that can modify fine-tuning but not pre-training.
  • 5 Recommendations for post-training practice: A replay fraction requires proportionally more fine-tuning steps, but this is generally permissible because fine-tuning is rarely compute-constrained.Using replay fraction ρ requires 1/(1−ρ) times as many training steps.
  • 5.1 Web Agents: Replay generic instruction-following data improves web-agent accuracy by up to 4.5%.The comparison uses Llama 3.1 8B Instruct on Weblinx demonstrations and an offline scoring procedure.
  • 5.2 Basque: Basque replay uses SlimPajama as a proxy for unreleased Llama pre-training data, while web-agent replay uses OpenHermes or UltraChat instruction data.These replay distributions approximate the data used in the models’ previous training stages.
  • 5.2 Basque: Basque is especially low-resource, comprising only 0.035% of Common Crawl.The evaluation uses a professional Basque translation of COPA.
  • 5.2 Basque: Replay of generic pre-training-like data improves Basque accuracy by up to 2%.The experiment fine-tunes Llama 3.1 8B on 200M Basque tokens and evaluates Basque COPA accuracy.

6 Related work

Related work frames replay as connected to data-mixture optimization, continual learning, robust fine-tuning, and curriculum learning. The paper distinguishes its focus on how pre-training exposure affects the best fine-tuning recipe and target performance.

  • Optimizing data mixtures: Data-mixture optimization methods can be myopic because they greedily upweight promising data early, missing that relevant data may belong at the end.This motivates considering training order rather than only online mixture weights.
  • Continual learning: Unlike traditional continual-learning work focused on reducing catastrophic forgetting, this paper studies improving target-task performance.The related literature also includes methods for teaching models new facts.
  • Necessity of pretraining: The paper argues that some tasks may not need to appear during pre-training if adaptation follows optimal training procedures.This challenges the assumption that target skills must always be incorporated during pre-training.
  • Robust fine-tuning: Replay is described as analogous to weight averaging, but averaging occurs in data-distribution space rather than parameter space.The paper emphasizes that its analysis connects optimal fine-tuning to the model’s prior target-task exposure.
  • Robust fine-tuning: The paper’s distinction from robust fine-tuning is its explicit characterization of the interaction between pre-training and fine-tuning.It also notes that prior work emphasizing out-of-distribution performance under-focused on in-distribution improvement.
  • Curriculum learning: Curriculum learning studies sequences of training distributions, with theory suggesting tractable intermediate tasks can accelerate convergence.Recent work designs curricula using reference models or structure in the data distribution.

7 Discussion

The discussion asks whether pre-training must change to exploit task-relevant data and identifies both practical gains and methodological limits. Stage-2 replay alone captures all optimal-schedule gains for StarCoder but only part for FineMath and Flan.

  • Do we need to change pre-training?: Stage-2-only changes capture 67.4% of FineMath gains and 46.0% of Flan gains from the optimal data schedule.For StarCoder, the optimal schedule requires only adding replay data to Stage 2.
  • Hypotheses for inefficiency of fine-tuning: Standard fine-tuning may underperform replay because of early training instability and overfitting to small samples.The paper presents these as hypotheses, including a toy model involving noisy data points.
  • Limitations: The controlled setting simplifies pre-training to two distributions, whereas practical pre-training is a more diverse multi-task problem.The simple schedules also exclude continuous annealing, sample-level orderings, and advanced fine-tuning methods.
  • Limitations: Replay increases compute requirements and validation loss may not perfectly correlate with downstream metrics.These constraints limit how directly the controlled results transfer to practice.

8 Impact statement

The paper aims to improve data efficiency for low-resource domains. It acknowledges that replay may increase training compute and that other harms are broadly shared with language-model research.

  • Impact: The work targets improved data efficiency for low-resource domains that receive relatively less attention.The impact statement frames this as the intended benefit of the research.
  • Impact: Replay may increase the compute used in language-model training.The authors identify this as a potential harm and note that most other harms are common to language-model research.

A Data schedule equivalences

The paper represents two-stage data schedules using seven variables but shows that only two are freely chosen; the remaining quantities follow algebraically. It also examines fine-tuning instability, including an initial loss spike that replay may mitigate through reduced distribution shift or additional recovery steps.

  • A Data schedule equivalences: Seven schedule variables reduce to two degrees of freedom because total steps and target-step fraction are fixed, while the remaining quantities are derived.The paper sets T = 1 and treats the repetition count as fixed for this section.
  • A Data schedule equivalences: Replay fraction ρ determines the Stage 2 target weight as w2 = 1 −ρ.
  • A Data schedule equivalences: Target allocation α determines Stage 2 target steps, total Stage 2 duration δ, and the corresponding Stage 1 target weight w1.The derived weights satisfy w1(1 −δ) + w2δ = γ.
  • B.1 Instability of fine-tuning: Fine-tuning begins with a large loss spike, especially at higher learning rates, but later reaches lower loss despite the spike.
  • B.1 Instability of fine-tuning: Replay may reduce the spike by lowering Stage 1–Stage 2 distribution shift or providing more Stage 2 steps for recovery.The authors present these as hypotheses and say further experimentation is needed to determine when the spike matters.

B.2 Overfitting to target data

A linear-regression toy model frames fine-tuning as overfitting under limited target samples. Replay reduces this overfitting, while ridge regularization performs even better in the toy setting; larger models nevertheless still require replay for lower loss.

  • B.2 Overfitting to target data: The toy model tests whether unregularized fine-tuning overfits when the target sample count is small.It uses a 400-dimensional distribution with pre-training and fine-tuning parameters related by a Gaussian perturbation.
  • B.2 Overfitting to target data: Replay mixes pre-training data into fine-tuning and significantly reduces overfitting in the linear-regression experiment.
  • B.2 Overfitting to target data: The best non-zero ridge parameter achieves much lower loss than replay-count tuning in the toy model.This motivates seeking an appropriate regularization mechanism for language-model training.
  • B.2 Overfitting to target data: Larger models still require replay data to obtain lower loss in the tested joint-training setting.
  • B.2 Overfitting to target data: The authors suggest that optimal data schedules for large models may be tunable on smaller models, although the project’s experimental coverage was more comprehensive for mid-training.

D General training settings

The experiments use 150M-parameter Llama-style models trained on generic C4 data and target domains spanning math, coding, and instruction following. The study tunes repetition, learning-rate cooldown, weight decay, and data schedules, while measuring target efficiency and generic-domain forgetting.

  • D General training settings: The main controlled setup trains 150M-parameter Llama-style models for 4B tokens with context length 4096, Adam, and weight decay 0.1.
  • D General training settings: Generic data comes from C4, while target domains are FineMath, StarCoder, and Flan instruction following.
  • D General training settings: The study uses a 150M default model and checks whether results hold at larger scales, while varying schedule fractions and target-data allocations.The schedule uses replay fractions and target allocations selected on powers-of-two-like grids.
  • D General training settings: Mid-training tolerates up to 32 target-data repetitions before overfitting across domains.
  • D General training settings: A cooldown period is critical for WSD, while weight decay has a minimal and noisy effect; the experiments therefore use weight decay 0.1.
  • D General training settings: Replay and early target-data introduction both significantly mitigate forgetting on the generic domain.
  • D General training settings: Fine-tuning tolerates up to 64 target-data repetitions before overfitting across domains.
  • D General training settings: Data efficiency is measured by comparing each strategy’s loss with a reference power law and converting it into equivalent target-data points.

H.2 Training runs

The training runs compare cosine and WSD schedules, target-data placement, and replay in controlled and downstream settings. WSD accelerates loss improvement during decay, while replay produces gains for Weblinx and Basque evaluation, including accuracy improvements that persist across token counts.

  • H.2 Training runs: The main runs use cosine learning rates, 4B total tokens, and target-token counts from 4M to 64M, with up to 32 target-data repetitions.
  • H.2 Training runs: WSD has warmup, stable, and decay phases, with the decay period being critical to final loss.
  • H.2 Training runs: During WSD decay, loss improves much faster, unlike cosine schedules whose improvement slows near the end.
  • H.2 Training runs: The schedule places target data at the end of WSD training because faster late-stage loss reduction makes high-quality late data more consequential.With cosine scheduling, placing target data at the end instead of uniformly can hurt.
  • H.2 Training runs: Without replay, Weblinx weight-decay tuning improves performance by less than 2% over baseline.
  • H.2 Training runs: Basque gains remain real across evaluations using 40M and 200M tokens.
  • H.2 Training runs: Replay’s loss improvement decreases with more total tokens, but accuracy gains persist and may become more important for evaluation metrics.
  • H.2 Training runs: Repeating target data too many times eventually increases loss rather than producing the diminishing improvements predicted by prior scaling laws.
Loading 2603.04964v1…