Patch the Distribution Mismatch: RL Rewriting Agent for Stable Off-Policy SFT
This addresses a practical issue for LLM practitioners by improving adaptation to new data without losing prior capabilities, though it is an incremental improvement over existing data-rewriting approaches.
The paper tackles the problem of catastrophic forgetting in supervised fine-tuning (SFT) of large language models when downstream data has a distribution shift, by proposing an RL-based data-rewriting agent that optimizes for distributional alignment and diversity. The method reduces forgetting on non-downstream benchmarks by 12.34% on average while achieving comparable downstream gains to standard SFT.
Large language models (LLMs) have made rapid progress, yet adapting them to downstream scenarios still commonly relies on supervised fine-tuning (SFT). When downstream data exhibit a substantial distribution shift from the model's prior training distribution, SFT can induce catastrophic forgetting. To narrow this gap, data rewriting has been proposed as a data-centric approach that rewrites downstream training data prior to SFT. However, existing methods typically sample rewrites from a prompt-induced conditional distribution, so the resulting targets are not necessarily aligned with the model's natural QA-style generation distribution. Moreover, reliance on fixed templates can lead to diversity collapse. To address these issues, we cast data rewriting as a policy learning problem and learn a rewriting policy that better matches the backbone's QA-style generation distribution while preserving diversity. Since distributional alignment, diversity and task consistency are automatically evaluable but difficult to optimize end-to-end with differentiable objectives, we leverage reinforcement learning to optimize the rewrite distribution under reward feedback and propose an RL-based data-rewriting agent. The agent jointly optimizes QA-style distributional alignment and diversity under a hard task-consistency gate, thereby constructing a higher-quality rewritten dataset for downstream SFT. Extensive experiments show that our method achieves downstream gains comparable to standard SFT while reducing forgetting on non-downstream benchmarks by 12.34% on average. Our code is available at https://anonymous.4open.science/r/Patch-the-Prompt-Gap-4112 .