Replaying pre-training data improves fine-tuning

TL;DR

In a controlled 150M-param / 4B-token / 4M-target-token pre-train→fine-tune setting (C4 as generic, FineMath / StarCoder / Flan as targets), mixing generic data back into the fine-tuning stage — “distributional replay” — improves target-domain validation loss, even though it moves the fine-tuning distribution further from the target. The improvement holds when the optimization process is unified into a single WSD schedule (i.e. mid-training rather than discrete fine-tuning), and replay matters most precisely when target data is scarce in pre-training. At 8B Llama 3 scale the same recipe lifts web-agent navigation accuracy and Basque QA accuracy, giving the paper its practical recommendation: when only fine-tuning is modifiable, replay generic data — especially for rare target domains.

Key claims

Replaying generic data during fine-tuning (Stage 2) improves target-domain validation loss in the controlled 150M setting; the loss-vs-replay-fraction curve has an interior minimum at a non-zero replay fraction for every target domain tested [§3.2, Fig. 4].
Replay’s benefit increases as the overlap between generic and target distributions shrinks: StarCoder (C4 explicitly filters code) tolerates less replay than FineMath, which tolerates less than Flan [§3.2, Fig. 4].
A unified WSD schedule with no optimizer-state reset between pre-training and fine-tuning (“mid-training baseline”) beats the standard separate-cosine fine-tuning baseline on target data efficiency; the gain comes from not re-warming the LR and not resetting AdamW moments [§4.1, Fig. 5].
A short cooldown (10% of total training) annealing the LR sharply at the end gives the largest target-loss drop and is the optimal cooldown fraction tested for all three domains [§4.1, Appendix I].
The benefit of replay also persists in the WSD/mid-training setting; pure replay alone (no target data in Stage 1) gives substantial data-efficiency gains over the mid-training baseline for all three target domains [§4.3, Fig. 7].
Replay is most beneficial precisely when less target data is present in Stage 1; with 75% of target data moved to Stage 1, replay no longer helps [§4.4, Fig. 8].
For applications where pre-training cannot be modified, just adding replay to Stage 2 captures the majority of the gain for FineMath and Flan, and captures all of the optimal-data-schedule gain for StarCoder [§7, §3.2 vs §4.3].
At 8B (Llama 3.1 Instruct), replaying instruction-following data (OpenHermes / UltraChat) during fine-tuning on web-agent trajectories improves WebLinx accuracy under the original Weblinx hyperparameters [§5.1, Fig. 9].
At 8B (Llama 3.1 base), replaying SlimPajama as a Llama-pre-training-distribution proxy during continual pre-training on 200M Basque tokens improves Basque-COPA accuracy across a broad band of replay fractions [§5.2, Fig. 10].
The authors flag a recommendation to model-release teams: ship the pre-annealed (pre-cooldown) checkpoint and its optimizer state alongside the final checkpoint, since downstream mid-training prefers to initialize there rather than at the fully-annealed model [§4.1].
Replay differs from continual-learning’s classical use: continual-learning literature replays old data to prevent catastrophic forgetting of previous tasks; this paper finds replay improves the new (target) task as well, departing from that standard intuition [§3.2, §6 “Continual learning”].
Two hypotheses are offered for why fine-tuning underperforms replay: (i) a brief training instability in the first few fine-tuning steps that replay mitigates, and (ii) a statistical small-sample overfitting barrier analogous to double-descent that even a perfect optimizer would hit [§7 “Hypotheses for inefficiency of fine-tuning”, Appendix B.1, Appendix B.2].

Method

The controlled setup is a 150M-param Llama-style language model trained for a total compute budget of 4B tokens, with a fixed data constraint of 4M tokens drawn from a “target” pool (FineMath, StarCoder, or Flan) and the rest drawn from a “generic” pool (C4). Two-stage data schedules are parameterized by (a) the replay fraction — what fraction of Stage 2 training steps is on generic rather than target data — and (b) the target Stage 2 allocation — what fraction of total target tokens is held back for Stage 2 vs. mixed into Stage 1. Standard fine-tuning corresponds to replay_fraction = 0, Stage2_allocation = 1, with a separate cosine LR schedule and optimizer-state reset; the mid-training variant unifies pre-training and Stage 2 under a single WSD LR schedule (linear warmup → constant stable phase → short linear decay over 10% of training) with no optimizer reset.

Performance is measured by held-out validation loss on the target distribution. To compare strategies independent of the reference’s data efficiency, the paper fits a reference scaling law L_ref(T) = (T/c)^{-α} and reports a relative data efficiency of a strategy A at T target tokens as T_eff(A; T) / T where T_eff is the data the reference would need to match A’s loss; a 2.0× data efficiency means the reference would need 2× more target data to match A. Search over (replay_fraction, Stage2_allocation) is the main experimental sweep (Figure 7), with the StarCoder result that replay-only attains the full optimum being the extremum.

At 8B scale the recipe is run in two settings: (1) Llama-3.1-8B-Instruct fine-tuned for the WebLinx web-agent benchmark, replaying OpenHermes / UltraChat instruction data as a proxy for the prior post-training distribution; and (2) Llama-3.1-8B continual-pre-trained on Latxa Basque (200M tokens), replaying SlimPajama as a Common-Crawl-Llama-pre-training proxy. Validation in both cases is downstream task accuracy (Weblinx offline scoring and Basque-COPA).

Results

150M / 4B tokens / FineMath, StarCoder, Flan: a non-zero replay fraction strictly improves target loss over replay=0 for every domain, with the curve shape varying by domain [§3.2, Fig. 4].
Mid-training baseline vs. standard fine-tuning: unifying the LR schedule and skipping optimizer-state reset gives a measurable data-efficiency win for all three target domains [§4.1, Fig. 5].
Replay-only mid-training: at the mid-training baseline (no target in Stage 1), adding replay alone outperforms the no-replay mid-training baseline by margins reported per domain [§4.3, Fig. 7 right column].
Full data-schedule search (replay + target reallocation): a sweep over (replay_fraction, Stage2_allocation) improves over the mid-training baseline for all three domains; for StarCoder the optimum sits at Stage2_allocation = 1 (i.e. no target in Stage 1 needed) [§4.3, Fig. 7].
Replay × target-in-pretraining interaction: with Stage2_allocation = 1 (all target in Stage 2), replay is critical; with Stage2_allocation = 0.25 (75% of target moved to Stage 1), replay no longer helps [§4.4, Fig. 8].
Llama 3.1 8B Instruct / WebLinx: replaying instruction-following data during agent fine-tuning gives a consistent accuracy gain at the offline scoring procedure of Weblinx [§5.1, Fig. 9].
Llama 3.1 8B base / Latxa Basque 200M tokens: replaying SlimPajama gives a wide band of replay fractions that improve Basque-COPA accuracy, making the ratio easy to tune in practice [§5.2, Fig. 10].

Why it’s interesting

The headline reframes the boundary between pre-training and fine-tuning as a single optimizer trajectory over a time-varying data mixture, and argues the right primitives are not “where does fine-tuning start” but “what’s the replay fraction in Stage 2” and “what fraction of target data sits before the cooldown.” Two consequences are load-bearing for Luma’s training stack: (i) the explicit recommendation to release pre-annealed checkpoints + optimizer state so downstream mid-training can resume into the cooldown rather than re-warm, and (ii) the finding that replay’s benefit grows with the rarity of the target domain in pre-training — which is exactly the regime Luma’s domain-specific fine-tunes sit in.

This connects directly to several papers already on the wiki. How to Set the Learning Rate for Large-Scale Pre-training? also uses the WSD stable phase as the unit of analysis and argues Fitting scaling-laws on small WSD-trained proxies beat µTransfer; both papers are treating the WSD schedule as the production primitive and building recipes on top of its stable-phase behavior. Decouple Searching from Training: Scaling Data Mixing via Model Merging for Large Language Model Pre-training (DeMix) and Nemotron-CLIMB: CLustering-based Iterative Data Mixture Bootstrapping for Language Model Pre-training (CLIMB) optimize pre-training data mixtures with predictor-loops; this paper attacks the same problem from the opposite end — the data mixture at the end of training — and arrives at the surprising result that fine-tuning-time mixing can recover most of the optimal-schedule gain without touching pre-training, which is exactly the regime DeMix and CLIMB struggle to serve. Shaping capabilities with token-level data filtering uses negative curation at pre-training time to remove a capability; this paper’s distributional replay is the positive mid-training analog. Worth pairing with the synthetic-data lineage (see synthetic-training-data) since “what to mix in” is one knob and “synthetically generated vs replayed real” is an orthogonal one — neither paper isolates that interaction, and Luma should.