Retaining by Doing: The Role of On-Policy Data in Mitigating Forgetting

AI summary

Overview

Research area: Language model post-training, specifically catastrophic forgetting in supervised fine-tuning (SFT) versus reinforcement learning (RL), and the role of on-policy data.

Technical level: Intermediate. The paper combines large-scale empirical comparisons with a synthetic mixture-of-Gaussians analysis, but it is written to be readable with a basic background in language model training.

Scope (one sentence): A systematic comparison of how much SFT and RL each degrade a model's existing capabilities, an analysis attributing RL's relative robustness to its use of on-policy data, and a demonstration that approximately on-policy data can recover the same benefit more cheaply.

What This Paper Is About

Post-training a language model on a new task often damages capabilities it already had, a problem known as catastrophic forgetting. The authors ask a simple but underexplored question: do supervised fine-tuning and reinforcement learning forget differently, and if so, why? They find that RL consistently forgets far less than SFT, and then trace that advantage to the fact that RL trains on data generated by the model itself rather than data supplied from outside.

Key Contributions

1. A systematic empirical comparison of forgetting between SFT and RL across tasks (instruction following, general knowledge, arithmetic reasoning), model families (Llama, Qwen), and model sizes up to 8B, establishing that RL forgets significantly less than SFT while matching or exceeding it on the target task.

2. Identification of on-policy data as the core factor behind RL's robustness to forgetting. The authors rule out alternative explanations such as KL regularization and advantage estimation, and provide intuition for why RL's mode-seeking updates can counterintuitively preserve prior knowledge.

3. A practical demonstration that approximately on-policy data suffices. Data regenerated at the start of each epoch, rather than at every training step, can still substantially mitigate forgetting at lower computational cost.

4. A mechanistic account via a simplified setting in which the language model is modeled as a mixture of distributions, one for prior knowledge and one for the target task, explaining when mode-seeking behavior helps and when it does not.

Main Findings

• RL forgets less than SFT across the board. Across the Llama and Qwen families, model sizes from 1B to 8B, and the three target tasks, SFT variants showed a consistent tradeoff between target performance and forgetting, while RL improved target performance "without noticeable drops on non-target tasks."

• Self-SFT achieves similar target gain to RL but a much larger drop. Self-SFT uses data generated by the initial model and produces a similar accuracy gain to RL on the target task while inducing a significantly larger drop on non-target tasks.

• Standard SFT can gain more on instruction following but forget more. SFT using data from Llama-3.3-70B-Instruct can achieve a higher gain than RL on instruction following, but induces an even larger drop than Self-SFT.

• SFT's learning-rate tradeoff. A high learning rate (default 1e-4) is typically needed to reach high target performance for SFT, often at the cost of severe forgetting; a smaller learning rate (1e-5) reduces forgetting but fails to reach the same target performance even with more epochs (2 versus 10 were tested).

• In a uni-modal simulation, forward KL (SFT) forgets less. With a uni-modal training policy modeled as a Gaussian, reaching a target task gain of 0.9 caused a non-target drop of 0.64 under forward KL versus 0.7 under reverse KL.

• In a multi-modal simulation, reverse KL (RL) forgets less. When the training policy is modeled as a bi-modal mixture, reaching a target task gain of 0.9 with forward KL caused the area overlap with the old mode to drop by 0.12, whereas reverse KL shifted the new mode toward the target while largely keeping the old mode intact.

• The mode-seeking behavior comes from on-policy data. The authors trace reverse KL minimization, and therefore mode-seeking behavior, to the use of data sampled from the current policy.

• KL regularization is not the cause. GRPO with KL regularization (beta = 0.05) and without it (beta = 0.0) achieved a similar gain-drop tradeoff across the considered models and datasets, except for Llama models trained on IFEval.

• The advantage estimator is not the cause. REINFORCE, which does not use an advantage estimator, maintained a similar low level of forgetting to GRPO while lagging in target task accuracy. On Llama-3.1-8B-Instruct, for example, REINFORCE's drop was 7.7 on IFEval, -0.1 on MMLU, and -0.8 on Countdown, versus GRPO's 3.4, -0.2, and -0.5; SFT's drops on the same model were 27.8, 38.5, and 36.4.

• Approximately on-policy data works. Iterative-SFT, which trains on data generated at the start of each epoch, reached target accuracy higher than or comparable to SFT while exhibiting mild to no forgetting, in experiments with Qwen 2.5 1.5B Instruct and Qwen 2.5 7B Instruct on IFEval and MMLU. Applying SFT to data generated during an RL run also reduced forgetting.

• RL is not immune to forgetting. As the distance between the training policy's new mode and the target mode grows (distances of 4.0, 5.0, and 6.0 were simulated), RL begins to suffer from forgetting as well.

Methodology in Plain English

The authors start by defining two metrics. Gain is the increase in accuracy on the target task after training. Drop is the average decrease in accuracy on a collection of non-target tasks, including MATH and two safety datasets (WildJailbreak and WildGuardTest). Good post-training means high gain and low drop.

They take four instruct models — Llama-3.2-1B-Instruct, Llama-3.1-8B-Instruct, Qwen-2.5-1.5B-Instruct, and Qwen-2.5-7B-Instruct — and train each on one of three target tasks: IFEval for instruction following, MMLU for general knowledge, and Countdown for arithmetic reasoning. After each run they evaluate on all the other tasks. They compare three methods: SFT using responses from Llama-3.3-70B-Instruct, Self-SFT using responses from the initial model itself (filtered to keep only correct answers), and RL using GRPO with a reward of 1 for correct generations and 0 for incorrect ones. All models are trained for two epochs.

To understand why the difference exists, the authors build a stripped-down mathematical model. Instead of a full language model over sequences, they use mixtures of one-dimensional Gaussians: one "old" mode standing in for prior knowledge and one "new" mode standing in for the target task. Training means moving the training distribution toward the target, and forgetting is measured as the loss of overlap between the training distribution and the old mode. They run this simulation in two versions — a training policy with a single Gaussian, and one with two Gaussians — and compare forward KL minimization (the SFT analogue) against reverse KL minimization (the RL analogue) using 1000 samples per gradient update over up to 1000 iterations.

Finally, they isolate which algorithmic ingredient of GRPO actually matters. They test GRPO without KL regularization, compare GRPO against REINFORCE (which has no advantage estimator), and replace fully on-policy data with progressively staler data: Self-SFT (data from the initial policy only), Iterative-SFT (data regenerated at each epoch), and SFT applied to data produced during an RL run.

Why This Matters

Impact on research. The paper reframes a common intuition. Forward KL is mode-covering and reverse KL is mode-seeking, so the conventional expectation is that SFT should preserve old capabilities better than RL. The authors show this holds only for a uni-modal starting policy, and reverses once the policy has multiple modes — which is arguably the realistic case for language models. It also directly challenges concurrent work (Lai et al., 2025) that attributed RL's robustness to an advantage estimator, and reports that Shenfeld et al.'s (2025) hypothesis linking KL divergence from the initial policy to forgetting does not always hold in their setting. The work gives a forgetting-centric lens on the well-documented behavioral differences between SFT and RL.

Real-world applications:

• Preserving safety behavior during fine-tuning. The paper notes that safety capabilities are often eroded through fine-tuning and uses WildJailbreak and WildGuardTest to measure this. On-policy or approximately on-policy training could help practitioners keep safety behavior intact while adapting a model to a new task.

• Efficient post-training pipelines. Because regenerating data once per epoch (Iterative-SFT) largely recovers the forgetting benefit, teams can reduce the compute overhead of fully on-policy RL while still avoiding severe capability loss.

• Continually learning agents. The authors argue that as the community moves toward agents that continually learn from experience, deciding which data to consume becomes a stability question, and that on-policy data better preserves existing capabilities than off-policy data such as internet content or experience from other agents.

• Test-time training. The paper states its insights lay groundwork for understanding forgetting in the emerging paradigm of test-time training.

Industry relevance. Post-training is standard practice for deploying useful models, and unintended capability regressions — sometimes called the "alignment tax" — are a direct operational cost. A guideline that says "train on data your own model produced, refreshed periodically" is cheap to implement relative to full RL infrastructure, and it applies directly to instruction tuning and domain adaptation workflows.

Future Directions

• Scaling studies. The authors state that investigating how forgetting patterns vary as model and dataset sizes are scaled further, beyond their compute budget, remains a valuable direction.

• A theoretical account. The paper provides intuition from a mixture-of-Gaussians simulation and empirical evidence identifying on-policy data as a main cause, but states that additional research is needed to theoretically establish the role of on-policy data in mitigating forgetting.

• Forgetting under distribution shift. The simulation shows that RL begins to forget when the target task is drastically far from the starting modes, raising the open question of how to handle far-away targets.

• Data selection for continual learning. With agents that learn from experience, deciding which data to consume is consequential to stability, and the paper's results suggest on-policy sources are safer than off-policy ones — but the precise policy for mixing on-policy and off-policy data is not settled here.

Target Audience

Researchers and engineers working on language model post-training, including those who design SFT and RL pipelines and those concerned with capability regression after fine-tuning. It is also relevant to safety researchers who track the erosion of safety behaviors during adaptation, and to anyone building systems that continually update a model from its own experience. Readers need only a working familiarity with supervised fine-tuning, policy gradient methods, and KL divergence directions; the synthetic Gaussian analysis is self-contained and explained step by step.

Authors’ abstract

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.