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

TL;DR

Catastrophic forgetting during language-model post-training is addressed by comparing supervised fine-tuning (SFT) and reinforcement learning (RL) across tasks and model families. RL consistently achieves strong target-task gains with substantially less forgetting, with the authors arguing that on-policy data drives this robustness by enabling a mode-seeking update that preserves prior knowledge. A mixture-of-distributions analysis shows forward KL can forget less in uni-modal settings, but reverse KL forgets less in multi-modal scenarios, aligning with practical LM behavior. The authors further demonstrate that approximately on-policy data—collected at the start of epochs or via RL traces—can dramatically reduce forgetting, offering a practical, data-efficient guideline for LM post-training.

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.

Figures (9)

• Figure 1: Illustration of the forgetting dynamics for the forward KL objective, corresponding to SFT, and the reverse KL objective, corresponding to RL. Left: we model LM post-training as a mixture of two modes. The “old" mode represents prior knowledge and the “new" mode represents a target task. Initially, the old mode of the training policy $\pi_\theta$ roughly matches the old mode of the optimal policy $\pi^*$, but its additional “new" mode does not match the new target mode. The goal is for the training policy to match the optimal policy. Top right: minimizing forward KL first stretches the new mode of $\pi_\theta$ and then moves probability mass from the old mode to cover the target, leading to forgetting. Bottom right: in contrast, minimizing reverse KL maintains the shape of the old mode and covers the target distribution by shifting the new mode of $\pi_\theta$.
• Figure 2: SFT forgets more than RL across tasks and models. We compare the Gain (solid bar) / Drop (shaded bar) across models and datasets for: (1) Self-SFT, which uses data generated from the initial policy; (2) SFT, which uses data generated by Llama-3.3-70B-Instruct; and (3) RL (GRPO). Gain (higher better) represents the accuracy increase on the target task, while drop (lower better) represents the average accuracy decrease on non-target tasks.
• Figure 3: SFT exhibits a tradeoff between target task performance and forgetting. Comparison of Self-SFT runs with learning rates: $1e{-5}$ (small) and $1e{-4}$ (default), and training epochs ($2, 10$).
• Figure 4: Forward KL (SFT) with uni-modal training policy forgets less than reverse KL (RL). Learning and forgetting dynamics of forward KL (left) and reverse KL (right). Dashed lines represent the modes of the optimal policy: $p_\text{old}$ (left) and $p_\text{new}$ (right). For forward KL, the data is sampled from the target mode $p_\text{new}$; the curve goes from red to yellow as training progresses. For reverse KL, the data is sampled from $\pi_\theta$; the curve goes from blue to purple. Forgetting corresponds to the decrease of overlap on the left mode and learning a new target task corresponds to the increase in overlap on the right mode.
• Figure 5: Reverse KL (RL) with multi-modal training policy forgets less than forward KL (SFT). Learning and forgetting patterns of forward KL with different high ($0.15$) and low ($0.01$) learning rates (left and middle) and reverse KL (right). Dashed lines represent the modes of the optimal policy: $p_\text{old}$ (left) and $p_\text{new}$ (right). For forward KL, the data is sampled from the target mode $p_\text{new}$; the curve goes from red to yellow as training progresses. For reverse KL, the data is sampled from $\pi_\theta$; the curve goes from blue to purple. Forgetting corresponds to the decrease of overlap on the left mode and learning a new task is the increase in overlap on the right mode.