# Replaying pre-training data improves fine-tuning

> Replaying generic pre-training data during fine-tuning improves target-domain performance up to 2x data efficiency, especially when target data is scarce during pre-training.

- **Source:** [arXiv](https://arxiv.org/abs/2603.04964)
- **Published:** 2026-08-15
- **Permalink:** https://picx.dev/p/jb1hbm
- **Whiteboard:** https://picx.dev/p/jb1hbm/image

## Summary

# 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**:

| Domain | Optimal Data Efficiency Gain |
|--------|------------------------------|
| Flan | **1.87×** |
| FineMath | **1.49×** |
| StarCoder | **1.09×** |

**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:**

| Domain | Data Efficiency Gain |
|--------|---------------------|
| StarCoder | 9.92× |
| FineMath | 6.37× |
| Flan | 2.77× |

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$:

| Strategy | StarCoder | FineMath | Flan |
|----------|-----------|----------|------|
| Replay only (α=1) | 1.53× | 1.85× | 2.06× |
| Full space (optimal ρ, α) | 1.53× | 2.49× | 4.80× |

**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):

| Domain | Fraction of Optimal Gains Achieved by Replay Only |
|--------|---------------------------------------------------|
| FineMath | 67.4% |
| Flan | 46.0% |
| StarCoder | 100% |

> "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}$$

---

_Markdown view of https://picx.dev/p/jb1hbm, served by PicX — AI-generated visual whiteboard summaries of research papers._