cs.CLMar 5, 2026
SaveReplaying pre-training data improves fine-tuning
Organizations: Stanford University
Abstract
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 for fine-tuning and 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 and Basque question-answering accuracy by .
Figures & tables
Figure 1: Replaying the generic distribution can improve target performance. Standard fine-tuning trains on all target data (blue) after all generic data (purple). We find that replaying generic data during fine-tuning can surprisingly improve performance on the target domain, both for fine-tuning and mid-training (e.g. and for FineMath, respectively). We find that replay is most helpful when there is less target data present during pre-training.
Figure 2: Data scaling law for reference algorithm. We run a reference training strategy with different target data budgets. To estimate how effectively an algorithm is using the data, we invert the reference strategy’s scaling law to recover “effective data” for this loss and compare the data efficiency improvement between two strategies. All of our data efficiency estimates only need to interpolate this scaling law.
Figure 3: Controlled fine-tuning visualization. We systematically explore the benefit of replaying generic data while fine-tuning on the target data. On the right, we show standard fine-tuning for steps where fraction of the steps are on target data. On the left, we show fine-tuning with replay fraction (where we shorten pre-training to keep the total number of steps fixed). We use (independently tuned) cosine learning rate schedules for each stage, with an optimizer state reset between the stages to simulate standard practice for fine-tuning open-weight models.
Figure 5: Tuning learning rate cooldown. We tune how long we should cool down the learning rate for WSD. The above plot shows the optimal cooldown period is between 0.05 and 0.1; we use 0.1 for consistency across domains and being fair to changing data schedules.
Figure 6: Controlled mid-training visualization. We explore the space of data schedules when training on tokens where a fraction of the steps are on target data. A data schedule allocates an fraction to Stage 2 where Stage 2 has a replay fraction . Standard fine-tuning puts all target data at the end with no replay ( , ). We use a WSD learning rate schedule across both stages.
Figure 7: Full data schedule sweep. We sweep over data schedules, parameterized by their replay fraction and fraction of target data allocated to Stage 2. Standard fine-tuning (no replay and all target data in Stage 2, top right corner) achieves the worst loss for FineMath and Flan. This can be improved by adding replay data (right column). This can also be solved by adding some target data to Stage 1, in which case replay becomes less important.
Figure 8: Importance of replay fraction depends on rarity. When all the target data is seen during fine-tuning, tuning the replay fraction becomes critical to improve loss (blue line). When we change pre-training to see some of the target data, tuning the replay fraction is not important and can sometimes hurt loss (purple line).
Figure 9: Weblinx. We fine-tune Llama 3.1-8B Instruct on Weblinx demonstrations. Without any replay data, we get 32.86% accuracy following the original hyper-parameters. We find that mixing generic instruction following data (OpenHermes, UltraChat) improves accuracy by up to 4.5%. This is even better than replaying demonstrations from an alternative web agent task (Mind2Web).
Figure 10: Basque. We fine-tune Llama 3.1-8B on 200M Basque tokens from the Latxa training corpus and measure accuracy on Basque COPA. We find that replaying generic pre-training data from SlimPajama improves accuracy by up to 2%.
Appendix figures & tables15 assets
Supplementary material from the paper’s appendix.
Appendix
Figure 11: Train loss spike. We notice a large loss spike during the first few steps of training across our settings. This represents one barrier to training with a limited number of samples. Replay might help either because (1) there is less distribution shift between Stage 1 and Stage 2 or (2) there is more time to recover from the spike.
Figure 12: Linear regression with replay. We plot the loss of a linear regression model as a function of the number of fine-tuning points for different values of the replay fraction . We see that for , the model overfits to the noise in the fine-tuning data, resulting in higher error than random guessing. We plot the gray line to track the best possible error achievable by the model for a given by using any sample count up to . Replay significantly reduces the overfitting, resulting in better MSE than random guessing. We track the best possible error as a function of for the best as the red line. For a regime of some but not too many target points, replay improves over standard fine-tuning.
Figure 13: Linear regression with ridge regularization. We follow the same setup as in Figure 12 , but instead of tuning replay fraction, we tune the ridge regularization parameter . We see that the optimal ridge parameter is non-zero and that the best ridge parameter achieves much better loss than tuned replay count.
Figure 14: Model scaling. We take our standard 4B token training setup and scale the parameter count to 4x parameter count. For a given size, we sweep for the optimal replay fraction while reserving all the target data for Stage 2 ( ). We find that larger models still require replay data to get lower loss.
| Parameter | 150M | 300M | 600M | 1.9B |
|---|---|---|---|---|
| Context Length | 4096 | 4096 | 4096 | 4096 |
| Hidden Dimension | 512 | 768 | 1024 | 2048 |
| Intermediate Dimension | 1792 | 2688 | 3584 | 7168 |
| Attention Heads | 8 | 12 | 16 | 16 |
| KV Heads | 8 | 12 | 8 | 8 |
| Layers | 6 | 12 | 24 | 24 |
Table 1: Model architecture configurations for different model sizes. All models use the Llama architecture with a standardized context length of 4096 tokens. We default to the 150M model if not specified.
Figure 15: Mid-training tuning repetitions. We show that across our mid-training domains, we can tolerate maximum 32 repetitions of the original data before overfitting to the target data. Note that the x-axis is more comprehensive compared to 18 .
Figure 16: Tuning weight decay. We tune the weight decay for the fine-tuning baseline fixed to our above repetition count and learning rate cooldown. The effect is small and noisy, so we default to 0.1, at the upper range of weight decays used in the literature.
Figure 17: Full data schedule sweep for forgetting. We take the same data schedules in 7 and instead plot loss on the generic domain (C4) instead of loss on the target domain, quantifying the amount of forgetting. We find that replay and introducing target data early both significantly reduce forgetting.
Figure 18: Fine-tuning repetitions. We show that across our fine-tuning domains, we can tolerate 64 repetitions of the original data before overfitting to the target data.
Figure 19: Power law fit. We fit a power law to the loss of the reference algorithm (uniform mixing) as it receives more target data. We note that we use these laws in the interpolation regime since we train models with hypothetically more data than we actually train with.
Figure 20: Learning rate schedule. This figure shows the shape of a WSD learning rate schedule in contrast to a cosine learning rate schedule (both without warmup). Figure is taken from Schaipp et al. [2025] , Figure 2 left.
Figure 21: Loss curves. This figure shows the loss curves for a WSD learning rate schedule and a cosine learning rate schedule. Note that the loss of a WSD schedule initially makes slower progress than a cosine schedule and then makes up for it with much faster loss improvement at the end. Figure is taken from Schaipp et al. [2025] , Figure 1 left.
Figure 22: River valley landscape. This shows the intuitive picture of the river valley landscape. WSD makes progress along the river direction while making all the hill progress at the very end. Figure is taken from Wen et al. [2024] , Figure 2 left.
Figure 23: Weblinx weight decay ablation. We fine-tune Llama 3.1-8B Instruct on Weblinx demonstrations. We find that without replay, tuning weight decay gives little gains, improving by less than over the baseline.
Figure 24: Basque training with different token counts. We try Basque training with 40M and 200M tokens to confirm that the gain in accuracy is real.