当前位置:首页 > 科技资讯 > 正文

策略蒸馏:强化学习与密集奖励的完美结合

最近,一向以发布博客著称而非论文的 Thinking Machines Lab(简称 TML)再次更新,发表了一篇题为《在策略蒸馏》的博客。

在策略蒸馏(on-policy distillation)是一种结合强化学习(RL)的纠错特性和 SFT 的奖励密度优势的训练方法。在将其应用于数学推理和内部聊天助手时,TML 发现策略蒸馏可以在极低成本下超越其他方法。

TML 的 CEO Mira Murati 表示,这种方法适用于小模型,使其具备强大的领域性能和持续学习能力。

策略蒸馏:强化学习与密集奖励的完美结合 策略蒸馏 强化学习 密集奖励 小型模型 第1张

值得注意的是,在这篇新博客中,TML 明确表示这项新成果受到了 Qwen 团队研究的启发,并且在实验过程中大量使用了 Qwen3 系列模型。实际上,在原英文博客中,「Qwen」这个关键词一共出现了 38 次之多!

策略蒸馏:强化学习与密集奖励的完美结合 策略蒸馏 强化学习 密集奖励 小型模型 第2张

作为一家明星创业公司,TML 的更新也吸引了广泛关注。有人总结其优势:

策略蒸馏:强化学习与密集奖励的完美结合 策略蒸馏 强化学习 密集奖励 小型模型 第3张

更是有网友盛赞,TML 才是真 Open AI。

策略蒸馏:强化学习与密集奖励的完美结合 策略蒸馏 强化学习 密集奖励 小型模型 第4张

策略蒸馏:强化学习与密集奖励的完美结合 策略蒸馏 强化学习 密集奖励 小型模型 第5张

博客地址:https://thinkingmachines.ai/blog/on-policy-distillation/

这篇博客的主要作者是 Thinking Machines Lab 研究者 Kevin Lu。他之前曾在 OpenAI 工作,领导了 4o-mini 的发布,并参与过 GPT-5 series、GPT-oss、o3 & o4-mini、4.1-nano & 4.1-mini、o1-mini、o3-mini 等模型的研发工作。

下面我们就来详细看看这篇博客的内容。

大型语言模型(LLM)能够在特定领域展现出专家级的水平。这是几种能力共同作用的结果,包括:对输入的感知、知识检索、规划选择和可靠执行。

要实现这一点,需要一系列的训练方法。我们可以将其大致分为三个阶段:

预训练(Pre-training):教授通用能力,例如语言使用、宽泛的推理和世界知识。

中训练(Mid-training):传授领域知识,例如代码、医疗数据库或公司内部文件。

后训练(Post-training):引导出目标行为,例如遵循指令、解决数学问题或聊天。

在特定专业领域,经过强化训练的小型模型的表现往往优于那些大型的通用模型。使用小型模型有很多好处:

出于隐私或安全考虑,它们可以进行本地部署。

它们可以更轻松地持续训练和更新。

它们还能节省推理成本。

想要利用这些优势,就需要为训练的后续阶段选择正确的方法。

后训练「学生」模型的方法可以分为两种:

在策略(On-policy)训练:从学生模型自身采样轨迹(rollouts),并为这些轨迹分配某种奖励。

离策略(Off-policy)训练:依赖于某个外部来源的目标输出,学生模型需要学习模仿这些输出。

例如,我们可能希望训练一个紧凑模型来解决如下的数学问题:

策略蒸馏:强化学习与密集奖励的完美结合 策略蒸馏 强化学习 密集奖励 小型模型 第6张

我们可以通过强化学习(RL)来进行在策略训练。具体做法是根据学生模型的每个轨迹是否解决了问题来为其评分。这个评分可以由人工完成,也可以由一个能可靠给出正确答案的「教师」模型来完成。

策略蒸馏:强化学习与密集奖励的完美结合 策略蒸馏 强化学习 密集奖励 小型模型 第7张

在策略训练的优势在于,学生通过在自己的样本上训练,能更直接地学会避免错误。

但 RL 有一个主要缺点:它提供的反馈非常稀疏(sparse feedback)。无论使用多少 token,它在每个训练回合(episode)中教授的比特数是固定的。

在我们上面的例子中,学生只知道「21」是错误答案,并更新模型以避免产生这个轨迹。但它并没有学到究竟错在哪里 —— 是搞错了运算顺序,还是算术本身出了错。这种反馈的稀疏性使得 RL 在许多应用中效率低下。

离策略训练通常通过监督微调(SFT)来完成,即在一组精心策划的、针对特定任务的有标注示例上进行训练。这些有标注示例的来源可以是一个在当前任务上表现出色的教师模型。

我们可以使用一种称为蒸馏(distillation)的机制:训练学生模型来匹配教师模型的输出分布。我们在教师的轨迹上进行训练,这些轨迹是生成的 token 的完整序列,包括中间的思考步骤。

在每一步,我们既可以使用教师完整的「下一个 token 分布」(常被称为 「logit distillation」),也可以只采样给定的序列。实践证明,采样序列提供了对教师分布的无偏估计,并能达到相同的目标。学生模型会根据自己生成该 token 的概率有多低,来相应地更新对序列中每个 token 的学习(在下例中用深色表示):

策略蒸馏:强化学习与密集奖励的完美结合 策略蒸馏 强化学习 密集奖励 小型模型 第8张

事实证明,蒸馏大型教师模型在训练小型模型方面非常有效,使其能够:

遵循指令

进行数学和科学推理

从医疗笔记中提取临床信息

以及参与多轮聊天对话

在策略蒸馏 —— 集两者之长