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× for fine-tuning and 2.06× 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%.
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. 1.87× and 2.06× 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 T 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 4
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 T 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 ( α=1 , ρ=0 ). 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 n for different values of the replay fraction ρ . We see that for n<d , 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 n by using any sample count up to n . Replay significantly reduces the overfitting, resulting in better MSE than random guessing. We track the best possible error as a function of n 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 ( α=1.0 ). 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
Appendix
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 16× 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 2% 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.
Real-world model deployments demand strong performance on narrow domains where data is often scarce. Typically, practitioners finetune models to specialize them, but this risks overfitting to the domain and forgetting general knowledge. We study a simple strategy, specialized pretraining (SPT), where a small domain dataset, typically reserved for finetuning, is repeated starting from pretraining as a fraction of the total tokens. Across three specialized domains (ChemPile, MusicPile, and ProofPile), SPT improves domain performance and preserves general capabilities after finetuning compared to standard pretraining. In our experiments, SPT reduces the pretraining tokens needed to reach a given domain performance by up to 1.75x. These gains grow when the target domain is underrepresented in the pretraining corpus: on domains far from web text, a 1B SPT model outperforms a 3B standard pretrained model. Beyond these empirical gains, we derive overfitting scaling laws to guide practitioners in selecting the optimal domain-data repetition for a given pretraining compute budget. Our observations reveal the finetuner's fallacy: while finetuning may appear to be the cheapest path to domain adaptation, introducing specialized domain data during pretraining stretches its utility. SPT yields better specialized domain performance (via reduced overfitting across repeated exposures) and better general domain performance (via reduced forgetting during finetuning), ultimately achieving stronger results with fewer parameters and less total compute when amortized over inference. To get the most out of domain data, incorporate it as early in training as possible.
Christina Baek, Ricardo Pio Monti, David Schwab +31
Continued pretraining enables language models to adapt to new domains and knowledge, but often at the cost of forgetting previously acquired capabilities. Replay can mitigate this trade-off, but fixed replay mixtures allocate training independently of the model's actual retention needs. We introduce Replay on Demand (RoD), which instead derives the replay allocation from the model's learning dynamics. RoD jointly prioritizes adaptation samples by their remaining learning potential and replay samples by their observed forgetting. Their competition for a shared training budget yields an online curriculum that determines what to train on at each step. Across models, scales, and adaptation domains, RoD reaches or improves upon the adaptation-forgetting frontier of tuned fixed-replay baselines and model merging without prescribing a replay allocation in advance. Replay concentrates on sources that are more vulnerable to forgetting and dynamically increases and redistributes as forgetting emerges during training. Together, our results show that replay can be allocated online from the model's evolving state, targeting what is needed, when it is needed.
Lukas Thede, Shengzhuang Chen, Stefan Winzeck +3
University of Tübingen, Tübingen AI Center · Helmholtz Munich · Munich Center for Machine Learning (MCML) +3
How can we train models whose post-trained capabilities survive subsequent fine-tuning? Rather than focusing on downstream interventions to mitigate forgetting of upstream capabilities, we study how upstream training choices - that is, the manner in which a capability is acquired - shape how robustly that capability is retained. We investigate this question in a controlled three-stage language-model pipeline: pretraining, post-training to acquire a target capability, and downstream fine-tuning on a new objective. Across 135M and 1B models, two post-training domains, and two downstream fine-tuning tasks, we find that immediate post-training performance does not reliably predict retention after subsequent fine-tuning: training recipes that look equivalent immediately after post-training can retain the target capability very differently after subsequent fine-tuning. In particular, early exposure - mixing post-training data into pretraining - consistently improves the frontier between retained upstream performance and downstream performance. In compute-matched experiments, where the target data must be allocated between pretraining and post-training, we find that the optimum lies at neither extreme. Together with our other empirical and theoretical findings, this supports the view that post-training drives immediate specialization while early exposure improves robustness to later forgetting. Replay and dropout, typically used to mitigate forgetting as it occurs during fine-tuning, provide complementary gains to early exposure when applied during post-training. Our findings suggest that robustness to subsequent fine-tuning should be treated as a first-class objective of upstream training, addressed preventatively through choices like early exposure rather than reactively during fine-tuning itself.
Lawrence Feng, Gaurav R. Ghosal, Jacob Mitchell Springer +2