小模型推理后训练:分阶段数据策略提升SFT与RL效果
速览
该研究针对小语言模型推理后训练,提出SFT与RL分阶段数据策略。SFT阶段聚焦未掌握技能,RL阶段巩固已有技能,并引入Bridge机制和Critique Fine-Tuning优化数据监督。实验表明该方法在多个基准上优于基线,证实了协调数据难度对提升模型推理能力的重要性。
AI 深度解读
Learning What to Learn: Stage-Specific Data Sets for SFT-then-RL in Small Language Model Reasoning
背景
在大型语言模型(LLM)向小型语言模型(Small Language Models, SLMs)延伸的过程中,推理能力(Reasoning)的提升已成为核心挑战。目前,对 SLMs 进行后训练(Post-training)以提升推理能力的标准范式通常是“监督微调(SFT)然后强化学习(RL)”(SFT-then-RL)。
然而,现有的研究工作往往忽视了一个关键问题:在每个训练阶段,模型究竟应该学习什么数据?
大多数现有方法倾向于在 SFT 和 RL 阶段使用相似或随机采样的数据分布,未能充分考虑到 SFT 和 RL 在技能习得过程中扮演的不同角色。这种数据策略与训练阶段目标的错位,可能导致训练效率低下或性能瓶颈。本文指出,数据策略应当与 SFT 和 RL 的独特作用相一致:SFT 更适合获取尚未掌握的新推理技能,而 RL 更适合巩固模型已经能够部分访问的技能。
核心内容
为了解决上述问题,作者提出了一种难度感知的 SFT-then-RL 框架,该框架将训练数据组织为特定阶段的集合(Stage-Specific Data Sets),并根据数据难度动态调整训练策略。
1. 核心原则:阶段特定的数据策略
- SFT 阶段(监督微调): 目标是获取尚未掌握的推理技能。因此,SFT 阶段应侧重于那些模型目前无法解决、但具有学习价值的“困难样本”。
- RL 阶段(强化学习): 目标是巩固模型已经具备一定基础的技能。因此,RL 阶段应侧重于那些模型能够部分解决、通过反馈可以进一步优化的样本,而非从头学习全新的高难度问题。
2. 关键技术机制
为了实现上述原则,作者设计了两个关键机制来处理不同阶段中的“困难样本”:
A. SFT 阶段:Bridge 机制
在 SFT 阶段,直接让 SLMs 学习由教师模型(Teacher Model)生成的原始推理轨迹(Reasoning Traces)往往效果不佳,因为 SLMs 的能力有限,难以直接模仿复杂的推理过程。
- 问题: 原始的教师推理轨迹对于 SLMs 来说可能过于复杂,导致监督信号难以被有效吸收。
- 解决方案: 引入 Bridge 机制。该机制将原始的、复杂的教师推理轨迹转化为更适合 SLMs 学习的、更易于理解的监督信号。这相当于在教师的高阶推理和 SLM 的低阶能力之间搭建了一座“桥梁”,使 SLMs 能够逐步习得这些新技能。
B. RL 阶段:Critique Fine-Tuning(批判性微调)
在 RL 阶段,模型可能会遇到一些始终无法解决的“硬样本”。如果这些样本在 RL 中持续失败(获得零奖励),传统方法可能会将其丢弃或重复训练,但效果有限。
- 问题: RL 中的零奖励失败样本缺乏明确的改进方向。
- 解决方案: 应用 Critique Fine-Tuning。对于在 RL 阶段仍未解决的硬样本,系统会将其转化为诊断性(Diagnostic)、修复性(Repair)以及新的推理轨迹监督信号。这些信号将被用于下一个 SFT 阶段的训练。
- 闭环流程: 这意味着 RL 阶段的失败经验不会丢失,而是被转化为具体的、可学习的 SFT 数据,反馈到后续的 SFT 循环中,形成“SFT -> RL -> 失败分析 -> 新 SFT 数据”的闭环优化。
3. 实验验证
作者在两个不同的 SLMs 上,针对五个推理基准测试(Reasoning Benchmarks)进行了实验。
- 对比基线: 代表性的 SFT、知识蒸馏(Distillation)和 RL 基线方法。
- 结果: 提出的方法在所有测试中均一致地优于现有基线。
关键要点
- 数据与阶段对齐: SFT 应专注于“学习新技能”(未掌握的技能),RL 应专注于“巩固已有技能”(部分掌握的技能)。
- Bridge 机制: 在 SFT 阶段,通过 Bridge 机制简化教师推理轨迹,使其更易于 SLMs 学习,解决“太难学不会”的问题。
- Critique Fine-Tuning: 在 RL 阶段,将失败的零奖励样本转化为诊断和修复信号,作为下一轮 SFT 的训练数据,解决“失败无反馈”的问题。
- 难度感知框架: 整个训练流程是动态的,根据数据难度和模型当前能力,自动分配数据到 SFT 或 RL 阶段,或进行跨阶段转化。
- 性能提升: 在多个 SLMs 和推理基准上,该方法显著优于传统的 SFT、蒸馏和 RL 方法。
意义与影响
这项研究对小型语言模型的推理能力训练具有重要的理论和实践意义:
- 重新定义后训练数据策略: 它挑战了“SFT 和 RL 使用相同数据分布”的传统假设,强调了数据难度与训练阶段匹配的重要性。这一原则可以推广到其他需要多阶段训练的场景。
- 提升 SLMs 的推理效率: 通过 Bridge 机制和 Critique Fine-Tuning,该方法有效地解决了 SLMs 在模仿复杂推理时的“认知负荷”问题,以及 RL 训练中“失败样本利用率低”的问题。
- 闭环优化范式: 提出的“SFT-then-RL”闭环框架,特别是将 RL 失败转化为 SFT 数据,为持续改进模型能力提供了一条可操作的路径。这对于资源受限的 SLMs 尤为重要,因为它们无法像大模型那样通过海量数据“暴力”学习。
- 推动 SLMs 在复杂任务中的应用: 随着 SLMs 在边缘设备、实时推理等场景中的应用增加,如何以较低成本提升其推理能力是关键。该方法提供了一种高效、低成本的后训练策略,有助于缩小 SLMs 与 LLMs 在推理能力上的差距。
总之,本文不仅提出了一种新的训练框架,更深刻地揭示了数据策略在 SLMs 推理能力发展中的核心作用,为后续研究提供了重要的方向指引。
