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

1. Problem Statement

The paper addresses the issue of catastrophic forgetting that occurs when Large Language Models (LLMs) are adapted to downstream tasks via Supervised Fine-Tuning (SFT).

• The Core Conflict: When downstream data exhibits a significant distribution shift from the model's pre-training distribution, standard SFT acts as an off-policy learning signal. The expert demonstrations ($y^*$) often have low likelihood under the model's current policy ($\pi_\theta$), leading to high-variance gradient updates, training instability, and the erosion of general capabilities.
• Limitations of Existing Solutions:
  • Standard SFT: Directly fine-tunes on expert data, causing severe forgetting.
  • Existing Data Rewriting (e.g., SDFT, Mind the Gap): These methods attempt to rewrite downstream data to be more "in-distribution" before SFT. However, they typically rely on fixed prompt templates to induce a conditional distribution $\pi_0(\cdot | x, y^*, x_{prompt})$.
  • The Mismatch: The target distribution for SFT is naturally $\pi_0(\cdot | x)$ (QA-style completion). Rewrites generated via fixed prompts often fail to align with this natural distribution, introducing templated phrasing and causing diversity collapse (mode collapse), thereby only partially mitigating the distribution gap.

2. Methodology

The authors propose an RL-based Data Rewriting Agent ($R_\phi$) that learns to rewrite supervision targets to better match the backbone model's natural generation distribution while preserving task correctness and diversity.

A. Framework Overview

The process consists of two stages:

1. Stage I: RL Training of the Rewriting Agent.
   • Architecture: A lightweight LoRA adapter is trained on top of a frozen instruction-tuned base model ($\pi_0$). This acts as a "patch" to correct local deviations without overfitting or drifting the entire model.
   • Optimization: Uses Group Relative Policy Optimization (GRPO), a PPO-style algorithm that eliminates the need for a separate value network by estimating baselines from groups of sampled outputs.
2. Stage II: Dataset Construction & Downstream SFT.
   • The trained agent generates rewritten data.
   • A Generate-Verify-Fallback pipeline is used: If a rewrite passes a task-consistency check, it is used; otherwise, the original expert demonstration is retained to ensure no supervision is lost.
   • Standard SFT is performed on this constructed dataset.

B. Unified Objective (Reward Function)

The agent is optimized to maximize a unified reward signal that balances three criteria:

1. Task Consistency (Hard Gate): A binary reward ($r_{task} \in \{0, 1\}$) acting as a hard constraint. It verifies both the final answer correctness (rule-based) and reasoning validity (LLM-as-a-judge). Rewards for alignment and diversity are only computed if $r_{task}=1$.
2. Distributional Alignment ($r_{dist}$): Encourages rewrites to be high-likelihood under the QA-style distribution $\pi_0(\cdot | x)$ (conditioning only on the input $x$, not the prompt or reference solution).
   • Calculated via length-normalized Negative Log-Likelihood (NLL) of the rewrite under the frozen $\pi_0$.
   • Normalized via group-wise whitening to stabilize training.
3. Diversity ($r_{div}$): Prevents mode collapse by encouraging semantic diversity among feasible rewrites for the same input.
   • Computed as the marginal contribution of a rewrite to the pairwise semantic distance (using embeddings) of the set of valid rewrites.

3. Key Contributions

• Distribution Mismatch Insight: The authors identify that existing data-rewriting methods fail because they optimize for a prompt-induced conditional distribution rather than the model's natural QA-style generation distribution.
• RL-Based Rewriting Agent: They propose a novel framework casting data rewriting as a policy learning problem, using GRPO to optimize for alignment, diversity, and consistency simultaneously.
• Unified Objective with Hard Gating: A reward design that enforces task consistency as a hard constraint while using soft rewards for alignment and diversity, preventing the agent from learning to "cheat" by generating invalid but high-reward outputs.
• Parameter Efficiency: The use of LoRA on a frozen backbone ensures the rewriting policy is a targeted "patch" that avoids catastrophic drift of the base model's general capabilities.

4. Experimental Results

Experiments were conducted on Llama-3.2 (1B, 3B) and Mistral-7B backbones, using mathematical reasoning as the downstream task and general-domain benchmarks (MMLU, etc.) to measure forgetting.

• Performance vs. Forgetting Trade-off:
  • Standard SFT: Achieved the highest math gains but suffered the most severe forgetting (e.g., ~17-19% drop in general benchmarks).
  • Baseline Rewriting (SDFT, Mind the Gap): Reduced forgetting slightly but yielded lower math gains, indicating they failed to fully optimize the distribution gap.
  • Proposed Method (Rewriting-Agent): Achieved comparable math gains to Standard SFT (e.g., +55.23% vs +54.19% on Mistral-7B) while significantly reducing forgetting (reducing the general benchmark drop from ~19% to ~7.35%).
• Training Stability: Downstream SFT training loss curves showed that models trained on the agent's rewritten data started with lower loss and converged to a more stable plateau compared to baselines, confirming better distributional alignment.
• Ablation Studies:
  • Removing the distributional alignment reward caused the largest drop in generalization, proving its critical role in preventing forgetting.
  • Removing the diversity reward hurt in-domain math performance, confirming the need to prevent mode collapse.
  • Hard gating was superior to soft shaping, as applying auxiliary rewards to invalid rewrites introduced noise.

5. Significance

This work provides a robust solution to the stability-plasticity dilemma in LLM adaptation. By shifting the intervention point from the optimization algorithm (e.g., PPO constraints) to the data source (rewriting supervision targets), the method proactively narrows the distribution gap.

• Practical Impact: It allows practitioners to fine-tune models on specialized data (like math or code) without sacrificing their general conversational or reasoning abilities, a critical requirement for deploying LLMs in real-world scenarios.
• Theoretical Contribution: It demonstrates that optimizing the distribution of supervision data via RL is a more effective strategy for mitigating catastrophic forgetting than simply constraining the fine-tuning updates.

In summary, the paper introduces a data-centric, RL-driven approach that "patches" the distribution mismatch between expert demonstrations and the model's natural generation style, resulting in stable fine-tuning with minimal catastrophic forgetting.