Patch the Distribution Mismatch: RL Rewriting Agent for Stable Off-Policy SFT

TL;DR

This paper introduces an RL-based data rewriting agent that improves the alignment and diversity of training data for fine-tuning large language models, reducing catastrophic forgetting and enhancing downstream performance.

Contribution

It proposes a novel reinforcement learning approach to learn a data rewriting policy that aligns with the model's QA distribution and maintains diversity, addressing limitations of existing methods.

Findings

Achieves comparable downstream performance to standard SFT.

Reduces forgetting on non-downstream benchmarks by 12.34%.

Effectively balances distributional alignment and diversity during data rewriting.

Abstract

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 .