Source-linked AI summary
Retaining by Doing: The Role of On-Policy Data in Mitigating Forgetting
Howard Chen, Noam Razin, Karthik Narasimhan, Danqi Chen
TL;DR
Post-training language models for new tasks can degrade existing capabilities, and the paper asks how SFT and RL differ in susceptibility to this forgetting. Through systematic experiments and simplified mixture-of-distributions analysis, it finds that RL forgets less while maintaining strong target performance, with on-policy data identified as the central factor and approximately on-policy data as an efficient mitigation.
Problem
Post-training language models for new tasks risks catastrophic forgetting, while comparative evidence on how SFT and RL differ in susceptibility remains limited.
Method
The paper systematically compares SFT and RL across tasks, model families, and scales, then uses mixture-mode analysis and ablations to examine why RL is more robust.
Results
RL consistently achieves strong target performance with substantially less forgetting than SFT across tasks, scales, and model families.
Takeaways & Limitations
On-policy data is the main supported explanation for RL’s robustness, and approximately on-policy data can mitigate forgetting more efficiently than fully on-policy data.
Takeaways & Limitations
The study does not establish the theoretical role of on-policy data and leaves behavior at larger model and dataset scales for future work.
Abstract
from arXiv · showhide
Adapting language models (LMs) to new tasks via post-training carries the risk of degrading existing capabilities -- a phenomenon classically known as catastrophic forgetting. In this paper, toward identifying guidelines for mitigating this phenomenon, we systematically compare the forgetting patterns of two widely adopted post-training methods: supervised fine-tuning (SFT) and reinforcement learning (RL). Our experiments reveal a consistent trend across LM families (Llama, Qwen) and tasks (instruction following, general knowledge, and arithmetic reasoning): RL leads to less forgetting than SFT while achieving comparable or higher target task performance. To investigate the cause for this difference, we consider a simplified setting in which the LM is modeled as a mixture of two distributions, one corresponding to prior knowledge and the other to the target task. We identify that the mode-seeking nature of RL, which stems from its use of on-policy data, enables keeping prior knowledge intact when learning the target task. We then verify this insight by demonstrating that the use on-policy data underlies the robustness of RL to forgetting in practical settings, as opposed to other algorithmic choices such as the KL regularization or advantage estimation. Lastly, as a practical implication, our results highlight the potential of mitigating forgetting using approximately on-policy data, which can be substantially more efficient to obtain than fully on-policy data.
1. Introduction
The paper compares SFT and RL for post-training, finding that RL achieves strong target performance with less forgetting across tasks, models, and scales. It attributes this robustness primarily to on-policy data and identifies approximately on-policy data as an efficient mitigation.
- Empirical comparison: RL consistently forgets less than SFT while achieving strong target-task performance across tasks, model families, and model scales.The study covers instruction following, general knowledge, and arithmetic reasoning with Qwen 2.5 and Llama 3 models up to 8B scale.
- Empirical comparison: SFT exhibits a tradeoff between target-task performance and forgetting, whereas RL improves target performance without noticeable non-target-task drops.
- Mechanistic intuition: In the mixture-of-modes model, forward-KL updates stretch the new mode and move probability mass from the old mode, leading to forgetting.The old mode represents prior knowledge, while the new mode represents the target task.
- Mechanistic intuition: Reverse-KL updates preserve the old mode’s shape while shifting the new mode to cover the target distribution.This mode-seeking behavior is associated with RL in the paper’s simplified analysis.
- Mechanism and practical implication: On-policy data underlies RL’s robustness to forgetting, rather than KL regularization or advantage estimation.
- Mechanism and practical implication: Approximately on-policy data generated at the start of each epoch can substantially reduce forgetting at lower computational cost.The paper contrasts this with generating data only from the initial policy, which is not enough for SFT.
2. Forgetting in LM Post-Training
The paper defines target-task gain and non-target-task drop to measure the tradeoff between adaptation and forgetting, then compares SFT variants with RL across tasks and models. SFT generally forgets more than RL: matching target gains can incur larger non-target drops, while RL improves target performance without noticeable drops.
- Experimental Setup: Figure 2 compares target-task gain and non-target-task drop for Self-SFT, SFT, and RL (GRPO).Gain is higher-is-better; drop is lower-is-better.
- Metrics: Target-task gain measures accuracy improvement, while non-target-task drop measures average accuracy decrease on other tasks.Post-training aims to maximize gain while minimizing drop.
- Experimental Setup: The study evaluates instruction following, general knowledge, arithmetic reasoning, MATH, and safety capabilities across Llama 3 and Qwen 2.5 instruct models.The compared models range from 1B to 8B parameters.
- Results: Self-SFT achieves similar target accuracy gains to RL but induces a significantly larger drop on non-target tasks.Self-SFT trains on responses generated by the initial model, retaining only correct responses.
- Results: SFT can achieve higher gain than RL on instruction following, but causes an even larger non-target drop than Self-SFT.Thus, higher target-task performance does not eliminate the forgetting tradeoff.
- Results: A high learning rate is typically required for SFT to reach high target performance, often causing severe forgetting; smaller rates reduce forgetting but fail to match performance even with more epochs.Figure 3 compares learning rates 1e−5 and 1e−4 across 2 and 10 epochs.
- Results: Both SFT variants show a consistent target-performance–forgetting tradeoff, whereas RL improves target performance without noticeable non-target drops.This pattern appears across datasets, model families, and model sizes.
3. Understanding Forgetting Dynamics Through the Lens of KL
SFT minimizes forward KL, whereas RL minimizes reverse KL, producing different forgetting dynamics that depend on whether the training policy is uni-modal or multi-modal. In the simplified mixture setting, forward KL forgets less with a uni-modal policy, while reverse KL is more robust with a bi-modal policy.
- 3.1. SFT and RL as KL Minimization: SFT corresponds to forward KL minimization, while RL corresponds to reverse KL minimization with respect to the optimal policy.Forward KL is mode-covering; reverse KL is mode-seeking.
- 3.2. Forward KL Forgets Less in a Uni-Modal Setting: The uni-modal analysis models the optimal policy as a mixture of old and new Gaussian modes, representing prior knowledge and the target task.The training policy is initialized to cover the old mode before optimizing toward the new mode.
- 3.2. Forward KL Forgets Less in a Uni-Modal Setting: In the uni-modal setting, reaching target task gain 0.9 yields non-target task drops of 0.64 for forward KL and 0.7 for reverse KL.The training policy is a univariate Gaussian initialized to cover the old mode, and overlap changes define gain and drop.
- 3.3. Reverse KL Forgets Less in a Multi-modal Setting: The contrasting results arise because a bi-modal training policy lets reverse KL match the new target mode without redistributing probability mass from prior-knowledge modes.The bi-modal policy contains separate trainable old and new Gaussian components with mixture weighting α.
- 3.3. Reverse KL Forgets Less in a Multi-modal Setting: KL regularization is not a major contributor to RL’s lower forgetting: except for Llama on IFEval, non-regularized GRPO has a similar gain-drop tradeoff.The comparison uses β = 0.05 with regularization and β = 0.0 without it.
- 3.3. Reverse KL Forgets Less in a Multi-modal Setting: In the bi-modal setting, achieving target task gain 0.9 with forward KL causes the overlap with the old mode to drop by 0.12.Reverse KL instead shifts the new component toward the target while largely preserving the old component.
4. Learning from On-Policy Data Mitigates Forgetting
The paper attributes RL’s lower forgetting primarily to on-policy data rather than KL regularization or advantage estimation. Approximately on-policy SFT data can substantially reduce forgetting while retaining comparable target-task performance.
- 4.1. On-Policy Data is the Primary Contributor for Mitigating Forgetting: On-policy data, rather than KL regularization or advantage estimation, is identified as the primary contributor to RL’s robustness against forgetting.Ablations compare these algorithmic distinctions and support on-policy data as the core factor.
- 4.1. On-Policy Data is the Primary Contributor for Mitigating Forgetting: KL-regularized and non-regularized GRPO achieve a similar target-task gain and non-target-task-drop tradeoff across considered models and datasets.These results suggest KL regularization does not explain RL’s lower forgetting.
- 4.1. On-Policy Data is the Primary Contributor for Mitigating Forgetting: REINFORCE lags behind GRPO in target-task accuracy but maintains a similar low level of forgetting.The comparison suggests advantage estimation mainly affects performance-gain magnitude, not forgetting mitigation.
- 4.2. Approximately On-Policy Data Suffices: Iterative-SFT reaches target accuracy higher than or comparable to SFT while exhibiting mild to no forgetting.It generates data at the start of each epoch, unlike Self-SFT, which uses data from only the initial policy.
- 4.2. Approximately On-Policy Data Suffices: SFT on data generated during an RL run also reduces forgetting, supporting approximately on-policy data as an efficient mitigation strategy.Self-SFT, which generates data only from the initial policy, suffers from severe forgetting.
5. Related Work
Related work has studied catastrophic forgetting in connectionist models and LM post-training, while recent studies similarly report that RL forgets less than SFT. This paper distinguishes its explanation by identifying on-policy data as the crucial factor.
- Catastrophic Forgetting: Catastrophic forgetting research began with connectionist models and later examined how forgetting patterns differ in language-model post-training.Earlier mitigation strategies focused on preventing drastic parameter changes.
- LM Post-Training: SFT uses ground-truth demonstrations, whereas RL generates model responses and learns from reward signals.The related-work discussion frames both as widely used methods for adding behaviors or improving target-task performance.
- LM Post-Training: Prior methods such as RAFT and STaR perform SFT in multiple rounds and can be viewed as approximately on-policy RL.The paper presents its contribution as a forgetting-centric comparison of SFT and RL.
- Concurrent Work: Concurrent studies also find less forgetting with RL, but this paper provides evidence against attributing robustness primarily to advantage estimation and instead emphasizes on-policy data.The cited concurrent explanations differ in whether advantage estimation or on-policy data is considered central.
6. Conclusion
Across tasks, model families, and scales, RL achieves strong target performance with substantially less forgetting than SFT. The paper attributes this robustness primarily to on-policy data and highlights approximately on-policy data as a practical direction.
- 6. Conclusion: Across tasks, scales, and model families, RL consistently achieves strong target performance with substantially less forgetting than SFT.The conclusion summarizes the paper’s systematic comparison of catastrophic forgetting in LM post-training.
- 6. Conclusion: RL’s robustness primarily stems from on-policy data rather than advantage estimation or KL regularization.This conclusion is supported by simplified analysis and practical experiments.
- 6. Conclusion: Approximately on-policy data, sampled asynchronously or at the start of each epoch, may efficiently mitigate forgetting.The paper presents this as a potential practical implication rather than a fully established theoretical result.
- 6. Conclusion: The study does not establish how forgetting patterns change at larger model and dataset scales or theoretically prove on-policy data’s role.Both questions remain directions for future research beyond the paper’s compute budget and current analysis.
Acknowledegment
The authors acknowledge comments, feedback, and financial support from academic, industry, and philanthropic organizations.
- Acknowledegment: The work received support from the National Science Foundation, Cisco Research, Schmidt Sciences, and the Zuckerman STEM Leadership Program.The acknowledgments also thank the Princeton Language Intelligence Group and Kevin Lu for feedback.
Impact Statement
The paper studies catastrophic forgetting during language-model post-training and examines how training-data choices may help preserve existing capabilities. Its findings are relevant to retaining safety behaviors and other desirable capabilities as models are fine-tuned and deployed as continually learning agents.
- The paper investigates catastrophic forgetting in language model post-training.
- On-policy methods are more robust to forgetting than supervised fine-tuning.
- Capability retention has implications for preserving safety behaviors and building more stable continually learning agents.
A.1. Simulation Details
The simulations compare forward-KL and reverse-KL updates in synthetic Gaussian-mixture settings. They show that reverse-KL updates can preserve the prior mode while learning the target, but forgetting increases when the target is far away.
- The synthetic task models forward KL as an SFT analogue and reverse KL as an RL analogue using a univariate Gaussian mixture.Updates use n = 1000 samples and are evaluated on a uniform density grid every 100 iterations.
- Forward KL with a high learning rate causes catastrophic forgetting, whereas a lower rate fails to learn the target mode.Reverse KL covers the target mode without sacrificing the prior mode.
- As the distance between qnew and pnew increases from 4.0 to 6.0, reverse-KL updates also begin to suffer forgetting.The simulation therefore does not show immunity when the target distribution is drastically far from the starting modes.
A.2. Connection Between Area Overlap and Total Variation Distance
The appendix connects area-overlap measures with total variation distance and documents the experimental implementation and evaluation setup. It also reports that approximate on-policy SFT nearly matches RL’s forgetting robustness, while KL divergence is not a fully monotonic forgetting indicator.
- A.2. Connection Between Area Overlap and Total Variation Distance: The non-target task drop equals the normalized increase in total variation distance from the scaled old component of the optimal policy.
- A.2. Connection Between Area Overlap and Total Variation Distance: The target task gain equals the normalized decrease in total variation distance from the scaled new component of the optimal policy.
- The experiments evaluate instruction following, general knowledge, and arithmetic reasoning across IFEval, MMLU, and Countdown, with additional non-target evaluations.
- SFT on RL-generated data trails full RL marginally in gains but has only slightly larger forgetting.The RL data is on-policy at each optimization step, while later SFT on that data is approximately rather than fully on-policy.
- A larger KL divergence from the initial policy does not necessarily imply greater forgetting.The relationship between KL divergence and forgetting remains not fully understood.
- Across AlpacaEval runs, SFT suffers much more degradation than RL after training on the three target tasks.Win rate is reported against GPT-4 for Llama-3.1-8B-Instruct and Qwen-2.5-7B-Instruct.