Replaying Pre-training Data Improves Fine-tuning

Authors: Suhas Kotha & Percy Liang (Stanford University)

---

Summary (Overview)

• Key finding: Replaying generic pre-training data during fine-tuning improves performance on the target domain, contrary to the standard intuition that replay is only useful for preventing catastrophic forgetting.
• Quantified gains: In controlled experiments with 150M parameter models (4M target tokens, 4B total tokens), generic replay increases target data efficiency by up to 1.87× for fine-tuning and 2.06× for mid-training.
• Interaction with pre-training: Replay is most beneficial when the target data is scarce during pre-training; when target data is already present in Stage 1 (α < 1.0), the benefit of replay diminishes significantly.
• Practical validation at scale: Fine-tuning 8B parameter Llama 3 models with replay improves web agent navigation success by 4.5% and Basque question-answering accuracy by 2%.
• Practical recommendation: For settings where only fine-tuning can be modified, replaying an approximation of the pre-training distribution can improve target performance, especially for low-resource domains.

---

Introduction and Theoretical Foundation

Background and Motivation

The standard paradigm for building domain-specific language models (e.g., math, code, instruction following) involves:

1. Pre-training on a vast amount of generic web text (e.g., C4, Common Crawl)
2. Fine-tuning on a relatively limited amount of target-domain data

The conventional assumption is that mixing generic data during fine-tuning serves only to prevent catastrophic forgetting of the generic domain [French, 1999; Rolnick et al., 2019]. The authors challenge this assumption by asking a fundamental question:

Can different data schedules—specifically, replaying generic data during fine-tuning—actually improve performance on the target domain itself?

Theoretical Context

The work sits at the intersection of several research areas:

• Mid-training: Recent practice of annealing the learning rate while training on high-quality data [Grattafiori et al., 2024; Li et al., 2025; OLMo et al., 2025]
• Data mixture optimization: Online algorithms that greedily upweight beneficial data components [Chen et al., 2023; Xie et al., 2023]—which the authors note are myopic and miss that the most relevant data should appear at the end of training
• Continual learning: Traditionally focused on reducing catastrophic forgetting [Kirkpatrick et al., 2017] rather than improving target task performance
• Robust fine-tuning: Techniques like weight averaging [Ilharco et al., 2023; Wortsman et al., 2022]—replay is qualitatively similar but operates in data distribution space rather than parameter space

Key Insight

The authors formalize a two-stage data schedule framework where:

• Stage 1: Pre-training (generic data, possibly with some target data)
• Stage 2: Fine-tuning/mid-training (target data, potentially mixed with replayed generic data)

The central surprising finding is that replay improves in-distribution target performance, departing from the standard continual learning intuition that replay only helps out-of-distribution generalization.

---

Methodology

2.1 Controlled Pre-training Setup

Data Pools:

• Generic domain: C4 (web text)
• Target domains (4M tokens each):
  • FineMath (math)—moderate overlap with C4
  • StarCoder (coding)—furthest from C4 (C4 is filtered for code)
  • Flan (instruction following)—closest to C4 (most natural language)

Training Configuration:

• 150M parameter Llama-style models
• Total training budget: 4B tokens (compute-matched comparisons)
• Target data constraint: 4M tokens
• Optimizer: AdamW

2.2 Evaluation via Data Efficiency

The authors evaluate using target validation loss (which scales smoothly and correlates with downstream performance) and define data efficiency as follows:

Given a training strategy $S$ accepting $D$ target tokens and producing loss $\mathcal{L}(S(D))$, the effective target data is estimated by inverting a reference strategy's scaling law:

$$\hat{\mathcal{L}}_{\text{ref}}^{-1}(\mathcal{L}(S(D)))$$

The relative data efficiency of strategy $S_2$ over $S_1$ is:

$$\frac{\hat{\mathcal{L}}_{\text{ref}}^{-1}(\mathcal{L}(S_2(D)))}{\hat{\mathcal{L}}_{\text{ref}}^{-1}(\mathcal{L}(S_1(D)))}$$

A data efficiency improvement of k× means "S₁ would require k times more target data to match the loss of S₂ at D tokens."

2.3 Experimental Designs

Fine-tuning experiments (Section 3):

• Two stages with separate cosine learning rate schedules
• Optimizer state reset between stages (simulating standard practice for fine-tuning open-weight models)
• Replay fraction $\rho$: fraction of Stage 2 steps on generic data
• Total steps conserved by shortening Stage 1

Mid-training experiments (Section 4):

• Single Warmup-Stable-Decay (WSD) learning rate schedule [Hu et al., 2024] across both stages
• No optimizer state reset
• Two degrees of freedom:
  • Replay fraction $\rho$: generic data fraction in Stage 2
  • Target Stage 2 allocation $\alpha$: fraction of total target data allocated to Stage 2

---

Empirical Validation / Results

3.1 Fine-tuning: Replay Improves Target Loss

Baseline: Standard fine-tuning (γ fraction of steps on target data, cosine LR, optimizer reset).

Key Result: For each target domain, a non-zero replay fraction minimizes loss:

Observation: Code (which C4 explicitly filters out) tolerates less replay than higher-overlap domains (math, instruction following). The loss improvement is relatively constant across domains despite different optimal replay fractions.

4.1 Mid-training Baseline

WSD Learning Rate Schedule: The authors find that annealing the learning rate over only 10% of training (cooldown period) dramatically improves data efficiency:

WSD with 10% cooldown increases data efficiency by 28.47× relative to annealing over all of training for FineMath.

Mid-training baseline gains over standard fine-tuning:

The improvement is attributed to joint training without optimizer state reset and learning rate rewarmup.

4.3 Two-Stage Data Schedule Search

Sweeping over replay fraction $\rho$ and target Stage 2 allocation $\alpha$:

Key finding: Introducing target data in Stage 1 offers additional gains for FineMath and Flan, but not for StarCoder.

4.4 Interaction Between Replay and Pre-training

Critical discovery: Replay matters most when target data is least present during pre-training:

• α = 1.0 (all target data in Stage 2): Replay is critical for improving loss
• α = 0.25 (75% of target data in Stage 1): Replay is no longer helpful and can sometimes hurt

This suggests replay serves as a substitute for having target data during pre-training.

5.1 Web Agents (Weblinx)

Fine-tuning Llama 3.1 8B Instruct on web navigation demonstrations:

• Baseline (no replay): 32.86% accuracy
• With replay (OpenHermes or UltraChat): Up to +4.5% accuracy
• Replay outperformed even replaying demonstrations from an alternative web agent task (Mind2Web)

5.2 Basque Language

Continual pre-training of Llama 3.1 8B on 200M Basque tokens (0.035% of Common Crawl):

• With replay (SlimPajama as proxy for pre-training data): Up to +2% accuracy on Basque COPA
• A large range of replay fractions offer benefit, making tuning easy

---

Theoretical and Practical Implications

When is Replay Necessary?

The authors quantify how much of the optimal data schedule's gains can be achieved by only modifying Stage 2 (replay):

"It is encouraging that we can get away with not using data early since it is prohibitive or impossible to change pre-training for many applications."

Hypotheses for Why Standard Fine-tuning Underperforms

1. Training instability: A few steps of fine-tuning cause optimization instability, which replay partially mitigates (Appendix B.1)
2. Statistical overfitting: A toy model shows failure arises from a small number of noisy data points, leveraging classical double-descent intuition (Appendix B.2)

Practical Recommendations

1. For fine-tuning practitioners: Replay an approximation of the pre-training distribution when the target domain is scarce in pre-training
2. For model developers: Release the model and optimizer state before cooldown (pre-annealed checkpoint), as this is more useful for downstream fine-tuning applications
3. Compute considerations: Using replay fraction $\rho$ requires $\frac{1}{1-\rho}$ times more training steps, which is generally permissible for fine-tuning since it is rarely compute-constrained

Connection to Existing Literature

• Contradicts necessity of pre-training: Prior work [Allen-Zhu and Li, 2024; Jiang et al., 2024b] argued target data must be seen during pre-training. This work shows that for some tasks, data may not need to appear in pre-training if optimal adaptation procedures are followed.
• Limitation of online mixture optimization: Greedy online algorithms miss that the most relevant data should appear at the end of training.

---

Conclusion

Main Takeaways

1. Replay improves target performance—not just prevents forgetting. This challenges a core assumption in fine-tuning practice.
2. The benefit is largest when target data is scarce in pre-training, offering a clear heuristic for when to apply replay.
3. Mid-training with WSD and no optimizer reset dramatically outperforms standard fine-tuning (up to 9.92× data efficiency), suggesting practitioners should fine-tune from pre-annealed checkpoints.
4. Results scale to 8B models with meaningful downstream gains (web agents: +4.5%, Basque QA: +2%).

Limitations

• Assumes only two data distributions, whereas real pre-training involves much higher diversity
• Simple two-stage schedules preclude studying continuous annealing or sample-level orderings
• Validation loss may not perfectly correlate with all downstream metrics
• Replay increases compute, which may be limiting outside standard fine-tuning

Future Directions

• Understanding the optimization and statistical mechanisms underlying replay's benefit
• Extending to more complex data schedules and multi-task settings
• Developing better metrics for quantifying distribution similarity between target and pre-training data
• Exploring whether the necessity of pre-training can be further relaxed with optimal adaptation procedures

---

Key Equations Summary

Data efficiency metric:

$$\frac{\hat{\mathcal{L}}_{\text{ref}}^{-1}(\mathcal{L}(S_2(D)))}{\hat{\mathcal{L}}_{\text{ref}}^{-1}(\mathcal{L}(S_1(D)))}$$

Compute cost of replay:

$$\text{Steps with replay } \rho = \frac{1}{1-\rho} \times \text{ steps without replay}$$