💬 NLP

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

本文提出了一种基于强化学习的重写智能体，通过将数据重写建模为策略学习问题，在确保任务一致性的同时优化重写数据与模型自然生成分布的对齐及多样性，从而在提升下游任务性能的同时显著缓解监督微调带来的灾难性遗忘。

原作者： Jiacheng Wang, Ping Jian, Zhen Yang, Zirong Chen, Keren Liao, Zhongbin Guo

发布于 2026-02-13

📖 1 分钟阅读☕ 轻松阅读

原作者： Jiacheng Wang, Ping Jian, Zhen Yang, Zirong Chen, Keren Liao, Zhongbin Guo

原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 ✨ 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。阅读完整免责声明

这篇论文讲述了一个关于如何让大语言模型（LLM）在“学习新技能”时，不会“忘记老本行”的故事。我们可以把它想象成一位全能学霸（大模型）去进修一门非常专业的技能（比如数学竞赛），但结果往往是他数学变强了，却把原本擅长的语文、历史等通用知识给忘了。

这篇论文提出了一种聪明的“补丁”方法，叫作RL 重写代理（RL Rewriting Agent）。

下面我用几个生活化的比喻来解释它的核心思想：

1. 问题的根源：水土不服的“教材”

想象一下，这位学霸（大模型）平时读的是**“自然对话风格”**的书籍（比如聊天、写故事），说话很自然、流畅。

现在，他要去学习**“数学解题”。但是，他拿到的“专家教材”（训练数据）写得非常生硬，充满了特定的格式和模板，就像是用一种“外语”**写的。

传统做法（SFT）： 老师直接让他背这些生硬的教材。结果他虽然学会了做题，但因为教材风格和他平时的说话方式（自然分布）太不一样，他为了适应这种生硬风格，把自己原本自然的说话方式（通用能力）给改坏了。这就叫**“灾难性遗忘”**。
现有的改进方法（数据重写）： 以前的方法试图把教材“翻译”一下，但翻译得还是很生硬，或者只是机械地套用几个固定的模板（比如“第一步、第二步、第三步”），导致教材虽然能看懂，但依然不够自然，而且千篇一律，缺乏多样性。

2. 核心创新：请了一位“金牌编辑”

这篇论文没有直接让模型去背生硬的教材，而是先请了一位**“金牌编辑”（这就是他们的RL 重写代理**）。

这位编辑的任务是：把那些生硬的专家教材，改写成学霸平时最喜欢读的那种“自然风格”，同时保证解题逻辑完全正确。

这位编辑是怎么工作的？（三个关键原则）

任务一致性（必须做对）：
- 比喻： 编辑改写的题目，答案必须是对的，解题思路必须通顺。如果改错了，直接扔掉，绝不使用。
- 技术点： 这是一个“硬门槛”，只有做对的改写才会被保留。
分布对齐（要像“自己人”）：
- 比喻： 编辑改写后的文字，必须读起来像学霸自己平时说话的样子（QA 风格），而不是像机器人生成的模板。如果读起来太生硬，编辑就会继续修改，直到它变得自然流畅。
- 技术点： 这叫做“分布对齐”，让训练数据更符合模型原本的自然生成习惯。
多样性（拒绝千篇一律）：
- 比喻： 如果所有题目都被改写成同一种格式（比如都变成“首先...其次...最后..."），学霸就会学傻了。编辑被要求：同一个题目，要能写出好几种不同的、自然的解法。
- 技术点： 防止“模式坍塌”，让数据丰富多彩。

3. 训练过程：如何培养这位“编辑”？

这位编辑（重写代理）不是靠死记硬背规则来工作的，而是通过**“强化学习”（RL）**来学习的。

奖励机制：
- 如果改写后的题目做对了 + 读起来很自然 + 和之前的改写不一样，编辑就得到高分奖励。
- 如果做错了，或者读起来很怪，或者太重复，编辑就没有奖励甚至被惩罚。
结果： 经过训练，这位编辑学会了如何把任何生硬的专家教材，都“润色”成学霸最舒服、最容易吸收的自然风格。

4. 最终效果：双赢

最后，学霸用这位编辑润色过的新教材进行学习：

数学成绩（下游任务）： 因为教材质量高、逻辑对，数学成绩提升很大，和直接背生硬教材的效果差不多。
通用能力（防止遗忘）： 因为教材读起来很自然，符合他原本的说话习惯，所以他没有为了学数学而丢掉原本的语言能力。

总结数据：
实验表明，这种方法在提升数学能力的同时，平均减少了 12.34% 的通用能力遗忘。也就是说，它既让学霸学会了新技能，又保护了他的老本行。

一句话总结

这篇论文就像是为大模型请了一位**“超级翻译官”，把那些生硬、难懂的专业教材，翻译成模型最擅长、最自然**的语言风格，让模型在学习新技能时，既能学得快，又不会“忘本”。

这是一篇关于大语言模型（LLM）在下游任务微调中解决分布不匹配（Distribution Mismatch）导致灾难性遗忘（Catastrophic Forgetting）问题的论文。论文提出了一种基于强化学习（RL）的数据重写代理（Rewriting Agent），旨在生成更符合模型自然生成分布的高质量监督数据。

以下是该论文的详细技术总结：

1. 研究背景与问题 (Problem)

核心痛点：将预训练好的指令微调模型（Instruction-tuned LLMs）适配到特定下游场景时，通常需要进行监督微调（SFT）。然而，当下游数据分布与模型先前的训练分布存在显著差异时，SFT 会导致灾难性遗忘，即模型在通用能力上的表现大幅下降。
现有方法的局限性：
- 分布不匹配：现有的数据重写方法（如 SDFT, Mind the Gap）通常使用固定的提示模板（Prompt Templates）引导模型重写数据。这导致重写后的数据服从于“提示诱导的条件分布” $\pi_0(\cdot | x, y^*, x_{prompt})$ ，而非模型自然的“问答式生成分布” $\pi_0(\cdot | x)$ 。这种不匹配使得重写数据无法完全消除分布差距。
- 多样性崩溃：固定模板容易导致输出格式单一、缺乏多样性（Mode Collapse），限制了模型学习多种有效解法的能力。
- 优化困难：分布对齐、多样性和任务一致性难以通过可微分的目标函数端到端优化。

2. 方法论 (Methodology)

作者提出了一种基于强化学习的数据重写代理（RL-based Data-Rewriting Agent），将数据重写建模为策略学习问题。

2.1 整体框架

两阶段流程：
1. 阶段 I（RL 训练）：训练一个轻量级的重写策略 $R_\phi$ ，使其生成的重写数据 $\tilde{y}$ 既符合任务要求，又贴近基座模型的自然 QA 分布，并保持多样性。
2. 阶段 II（下游 SFT）：利用训练好的代理生成重写数据集 $D_R$ ，并在该数据集上对目标模型进行标准的监督微调。

2.2 参数高效策略与优化

LoRA 微调：在冻结的指令微调基座模型 $\pi_0$ 上，使用 LoRA（Low-Rank Adaptation）适配器参数化重写策略 $R_\phi$ 。这被视为对基座模型的一个轻量级“补丁（Patch）”，用于局部修正生成行为，避免策略漂移。
GRPO 优化：采用 Group Relative Policy Optimization (GRPO) 进行策略优化。该方法无需额外的价值网络（Value Network），通过组内相对优势估计来更新策略。

2.3 统一目标函数 (Unified Objective)

为了同时优化任务一致性、分布对齐和多样性，设计了基于奖励信号的统一目标：
$r(x, y^*, \tilde{y}) = r_{task} + r_{task} \cdot (\lambda_{dist} r_{dist} + \lambda_{div} r_{div})$

任务一致性奖励 ( $r_{task}$ )：硬门控（Hard Gate）。只有当重写结果通过任务验证（答案正确且推理逻辑有效）时，才给予非零奖励。验证采用“粗到细”策略：先规则检查答案，再调用 LLM-as-a-Judge 检查推理。
分布对齐奖励 ( $r_{dist}$ )：鼓励生成符合基座模型自然 QA 分布的数据。计算重写结果 $\tilde{y}$ 在条件 $\pi_0(\cdot | x)$ 下的对数似然（NLL），并通过组内归一化转换为奖励。这确保重写数据更易于被模型在标准 SFT 下拟合。
多样性奖励 ( $r_{div}$ )：基于语义嵌入（Embedding）计算可行重写样本组内的语义距离。鼓励每个样本为集合贡献独特的语义内容，防止模板化导致的同质化。

2.4 数据集构建 (Generate-Verify-Fallback)

在构建最终训练集时，采用 生成 - 验证 - 回退 机制：

如果重写结果 $\tilde{y}$ 通过任务一致性检查，则使用 $\tilde{y}$ 作为监督目标。
如果检查失败，则回退（Fallback）使用原始的专家演示 $y^*$ ，确保不丢失监督信号。

3. 主要贡献 (Key Contributions)

分布不匹配的洞察：指出了现有数据重写流程中“提示诱导分布”与“自然 QA 分布”的不匹配问题，并提出了 RL 框架来解决这一根本矛盾。
统一优化目标：提出了一个联合优化任务一致性、分布对齐和多样性的统一目标，通过硬门控和奖励塑形，构建了高质量的重写数据集。
显著的性能提升：在多个模型规模（1B, 3B, 7B）和多个基准测试上验证了方法的有效性。

4. 实验结果 (Results)

实验在数学推理（下游任务）和通用领域（MMLU, MMLU-Pro 等，用于评估遗忘）上进行。

下游性能：该方法在数学推理任务上的提升幅度与标准 SFT 相当（例如在 Mistral-7B 上提升了 55.23%，与 SFT 的 54.19% 持平）。
遗忘缓解：显著减少了灾难性遗忘。在通用基准测试上的性能下降幅度比标准 SFT 小得多。
- 关键数据：在 Mistral-7B 上，标准 SFT 导致通用能力下降 19.17%，而该方法仅下降 7.35%。
- 平均遗忘减少：在多个基准上，平均减少了 12.34% 的遗忘。
训练稳定性：实验显示，使用该方法生成的数据进行 SFT 时，训练损失（Training Loss）起始更低且收敛更平稳，表明数据更符合模型的自然分布，优化更稳定。
消融实验：
- 移除分布对齐奖励 ( $r_{dist}$ ) 导致通用能力保留大幅下降，证明其关键性。
- 移除多样性奖励 ( $r_{div}$ ) 导致数学性能下降，证明多样性对任务学习同样重要。
- 硬门控（Hard Gate）比软塑形（Soft Shaping）效果更好，证明了过滤无效样本的重要性。

5. 意义与结论 (Significance)

数据为中心的解决方案：不同于传统的通过约束更新（如 DFT、PPO 约束）来防止遗忘，本文提出从数据源头入手，通过主动重写监督目标来缩小分布差距。
平衡增益与保留：该方法成功打破了“下游性能提升”与“通用能力保留”之间的权衡（Trade-off），实现了在大幅提升特定任务能力的同时，最大程度地保留模型的通用能力。
通用性：虽然实验主要集中在数学推理领域，但其提出的“分布对齐 + 多样性 + 任务一致性”的 RL 重写框架具有推广到其他下游领域（如代码、对话）的潜力。

总结：这篇论文通过引入强化学习代理来优化数据重写过程，有效地解决了 SFT 中的分布偏移问题，提供了一种稳定、高效且能显著减轻灾难性遗忘的下游模型适配方案。