近日,Thinking Machines Lab(简称TML)——一个更倾向于通过博客而非学术论文分享进展的研究机构——发布了最新更新,内容围绕一篇题为《在策略蒸馏》的博客文章展开。
在策略蒸馏(on-policy distillation)是一种创新的训练方法,它巧妙地将强化学习(RL)的纠错相关性与监督微调(SFT)的奖励密度相结合。当应用于数学推理和内部聊天助手时,TML发现该方法能以极低成本超越其他训练策略。
TML的首席执行官Mira Murati指出,这种技术尤其适用于小型模型,能赋予它们强大的领域性能与持续学习能力。
值得注意的是,在这篇新博客中,TML明确承认其成果受到了Qwen团队研究的启发,并在实验过程中大量使用了Qwen3系列模型。实际上,在英文原博客中,“Qwen”一词出现了38次,甚至超过了小米17系列发布会中雷军提及“苹果”的37次。
作为一家备受瞩目的创业公司,TML的更新引发了广泛关注。有观察者总结了其优势:
更有网友高度赞扬,称TML是真正的Open AI。
博客地址:https://thinkingmachines.ai/blog/on-policy-distillation/
这篇博客的主要作者是Thinking Machines Lab的研究员Kevin Lu。他此前曾在OpenAI工作,领导了4o-mini的发布,并参与了GPT-5系列、GPT-oss、o3 & o4-mini、4.1-nano & 4.1-mini、o1-mini、o3-mini等模型的研发。
下面我们将深入解析这篇博客的核心内容。
大型语言模型(LLM)能够在特定领域展现出专家级水平,这得益于多种能力的协同:输入感知、知识检索、规划选择以及可靠执行。
实现这一目标需要一系列训练方法,大致可分为三个阶段:
预训练(Pre-training):教授通用能力,如语言运用、基础推理和世界知识。
中训练(Mid-training):传授领域知识,如代码、医疗数据库或公司内部文档。
后训练(Post-training):引导目标行为,如遵循指令、解决数学问题或进行对话。
在专业领域,经过强化训练的小型模型往往能超越大型通用模型。使用小型模型有多重好处:
出于隐私或安全考虑,它们可部署于本地环境。
它们能更轻松地持续训练和更新。
它们还能显著降低推理成本。
要充分利用这些优势,关键在于为后训练阶段选择正确的方法。
后训练“学生”模型的方法可分为两类:
在策略(On-policy)训练:从学生模型自身采样轨迹,并为这些轨迹分配奖励。
离策略(Off-policy)训练:依赖外部来源的目标输出,学生模型需模仿这些输出。
例如,我们可能希望训练一个紧凑模型来解决如下数学问题:
我们可以通过强化学习(RL)进行在策略训练,即根据学生轨迹是否解决问题来评分。评分可由人工完成,也可由能可靠给出答案的“教师”模型执行。
在策略训练的优势在于,学生通过自身样本训练,能更直接地学会避免错误。但RL有一个主要缺点:反馈稀疏。无论使用多少token,每个训练回合教授的比特数是固定的。在上例中,学生仅知“21”是错误答案,并更新模型以避免该轨迹,但未学到具体错误原因——是运算顺序混淆还是算术错误。这种稀疏性使得RL在许多应用中效率低下。
离策略训练通常通过监督微调(SFT)实现,即在精心策划的任务标注示例上训练。这些标注可来自在当前任务上表现优异的教师模型。我们可以使用蒸馏机制:训练学生模型以匹配教师模型的输出分布。我们在教师的完整轨迹上训练,包括中间思考步骤。在每一步,我们既可使用教师完整的“下一个token分布”(称为“logit蒸馏”),也可采样给定序列。实践表明,采样序列能提供对教师分布的无偏估计,并达到相同目标。学生模型会根据自身生成该token的概率高低来更新学习(如下例深色部分所示):
事实证明,蒸馏大型教师模型在训练小型模型方面非常有效,使其能够:
遵循指令
进行数学和科学推理
从医疗笔记中提取临床信息
以及参与多轮聊天对话
用于这些应用的蒸馏数据集通常是开源和公开发布的。离策略训练的缺点是,学生在教师常遇的上下文中学习,而非在学生自身未来常遇的上下文中。这可能导致复合错误:如果学生早期犯下教师从未犯过的错误,它会发现自己越来越偏离训练中观察到的状态。当关注学生在长序列上的表现时,此问题尤为突出。为避免偏离,学生必须学会从自身错误中恢复。离策略蒸馏的另一问题是,学生可能学会模仿教师的风格和自信,而非事实准确性。打个比方:在策略RL如无教练指导自学国际象棋,输赢反馈与自身下法直接相关,但每局仅一次反馈,且不指明关键棋步;离策略蒸馏则如观看特级大师对弈——观察到高超棋步,但这些棋步处于新手少遇的棋局状态。我们期望结合RL的在策略相关性与蒸馏的密集奖励信号。对于国际象棋学习,这好比有老师为每一步棋评分;对于LLM后训练,这便是在策略蒸馏。
在策略蒸馏的核心思想是:从学生模型中采样轨迹,并使用高性能教师模型为每个轨迹的每一个token评分。回到上述数学例子,在策略蒸馏会给解题每一步打分,惩罚导致错误答案的错误步骤,同时强化正确步骤。
本文探讨了在策略蒸馏在以下任务的应用:1. 训练模型进行数学推理;2. 训练兼具领域知识和指令遵循能力的助手模型。我们在已具备预训练和中训练基础能力的模型上应用在策略蒸馏。我们发现,这是一种廉价而强大的后训练方法,成功结合了在策略训练的优势和密集奖励信号。
我们的在策略蒸馏工作借鉴了DAGGER(Ross et al, 2010),这是一种迭代式SFT算法,包含教师对学生访问状态的评估。它也类似于过程奖励建模(Lightman et al, 2023),这是一种RL方法,会对学生思维链的每一步评分。我们扩展了Agarwal et al.(2023)和Qwen3团队(2025)之前的在策略蒸馏工作。使用Tinker训练API,我们复现了Qwen3的成果,即通过在策略蒸馏在推理基准上实现同等性能,而成本仅为RL的一小部分。
你可在此Tinker cookbook中跟随实现每一步:https://github.com/thinking-machines-lab/tinker-cookbook/tree/main/tinker_cookbook/recipes/distillation
损失函数:反向KL
在策略蒸馏可使用多种损失函数为学生轨迹评分。为简化,我们选择逐token的反向KL——即在给定相同先前轨迹下,学生(π_θ)和教师(π_teacher)在每个token上分布之间的散度:
我们的奖励函数最小化反向KL,促使学生在每种状态下近似教师行为。当学生行为与教师完全一致时,反向KL为零。为简化,我们使用折扣因子为零:在任何给定时间步,学生仅优化下一个token,不考虑未来token。反向KL与RL有天然协同作用,RL通常优化由奖励模型引导的序列级反向KL。然而,与大多数奖励模型不同,反向KL是“不可破解的”,因为从教师视角看,低KL总对应高概率期望行为。反向KL的另一有用特性是“寻找众数”——它学习特定行为(教师行为),而非将分布分散于多个次优选项。该方法可节省大量计算资源。因无需等待轨迹完成采样来计算奖励,我们可以使用更短或部分轨迹训练。查询教师的对数概率也仅需大型模型一次前向传播,而轨迹由更小、更廉价的学生模型生成。我们也不需要单独奖励或标注模型。将基于蒸馏的逐token奖励与序列级环境奖励结合可能有益;这是未来有趣的研究领域。
图解
下面是一个真实例子,展示错误的学生轨迹由教师模型评分。此例来自SimpleBench,要求模型做出关键观察:问题前提很重要。正确答案是“B. 0”,因为冰块在煎锅中会融化。而学生模型(Qwen3-4B-Instruct-2507)错误地将其视为纯数学问题,忽略物理背景。
颜色越深,代表该token受教师模型(Qwen3-235B-A22B-Instruct-2507)惩罚越高(教师模型正确解决此问题)。我们看到,它惩罚了引导学生误入歧途的短语起始token,直观对应引导推理的重要“分叉token”。最终答案(虽错误)未受惩罚——因为在给定前面所有序列下,此答案完全可预测。
伪代码
我们在Tinker的RL脚本上实现在策略蒸馏,该脚本已实现采样、奖励计算和策略梯度式训练。地址:https://thinkingmachines.ai/blog/on-policy-distillation/(https://github.com/thinking-machines-lab/tinker-cookbook/blob/main/tinker_cookbook/rl/train.py)
初始化教师客户端:Tinker API可轻松为不同模型创建客户端。我们使用采样客户端,因无需通过教师模型传播对数概率。
采样轨迹:我们像在RL中一样从学生模型采样轨迹。采样期间,RL已计算学生对数概率log π_θ(x),用作重要性采样损失部分。
计算奖励:我们用compute_logprobs函数在采样轨迹上查询教师客户端,返回教师在学生采样token x上的对数概率log π_teacher(x)。然后我们用此计算反向KL。
使用RL进行训练:我们将逐token优势设置为负反向KL,并调用RL的重要性采样损失函数对学生模型执行训练更新。
伪代码如下:
在以下实验中,我们通常将在策略蒸馏应用于已过特定领域知识中训练的模型。这种训练提高了学生生成教师分布范围内token的概率,尽管通常远不足以复刻教师性能。通常,正如个性化示例所示,生成相关token的概率开始时为零,因学生缺乏相关领域知识。我们将使用在策略蒸馏进行后训练,并将其与训练专家模型的其他最后关键阶段方法比较。
我们使用蒸馏训练Qwen3-8B-Base模型的数学推理能力,并使用Qwen3-32B作为教师模型。教师和学生均为Tinker上支持的模型,因此你可使用Tinker cookbook复现实验。
离策略蒸馏
如前所述,我们所有实验以离策略蒸馏(即在教师生成的示例数据集上进行监督微调)形式作为中训练起点。用于数学推理的数据集是OpenThoughts-3,这是一个由QwQ-32B(类似Qwen3-32B的推理模型)生成的推理提示和响应集合。在40万个提示上对学生模型(Qwen3-8B-Base)进行全参数微调,在AIME"24(数学问题基准测试)上获得60%分数。我们也可使用LoRA训练,但在高容量数据集上训练时,它落后于全参数微调。在所有情况下,我们看到性能呈对数线性增长——初始性能提升廉价,但后期提升成本高昂。
我们可将在此40万提示上微调的模型视为检查点,然后尝试各种后训练方法,将其在AIME’24基准分数从60%提高到70%。默认方法是在更多提示上微调,即继续离策略蒸馏过程。根据对数线性趋势推断,我们估计模型在大约200万个提示上能达到70%分数。此推断需scaling law持续有效而不停滞,这并不简单。
强化学习
Qwen3技术报告称,在类似SFT初始化基础上,通过17,920个GPU小时的RL,在基准测试上达到67.6%性能。这很难与蒸馏成本直接比较,但基于对SFT训练堆栈的合理假设,这与在200万个离策略蒸馏提示上训练成本相似。
Qwen团队还报告称,使用在策略蒸馏,能以RL成本的十分之一,在AIME’24上达到74.4%的更高分数。这也启发了我们的工作。
在策略蒸馏
作为替代方案,我们运行了在策略蒸馏。从40万SFT检查点开始,在策略蒸馏在大约150个步骤内就达到AIME’24 70%的成绩。
跨方法比较计算成本并非易事。下面,我们用FLOPs(浮点运算次数)计算成本。
我们发现,当SFT数据集是现成的(如OpenThoughts-3示例)或在多次训练中被摊销时,基线成本降低9倍。然而,我们常希望在没有现成离策略蒸馏数据集的新任务上训练小模型。如果我们将教师模型的全部成本(包括从教师模型采样的额外成本)计算在离策略蒸馏中,那么总成本可降低约30倍。
除了将小型模型训练到在通用任务上表现出色之外,蒸馏的另一用例是个性化。例子包括在对话中遵循特定语气和输出格式,或像工具使用和成本预算这样的能力。我们经常希望在传授新领域知识的同时训练这种行为。同时训练这两者通常很困难,轻量级微调(如LoRA)往往不足以实现此目标,因此需要更大规模的中训练。在掌握新知识的基础上学习后训练行为,需要一个复杂的后训练堆栈,通常由专有数据和奖励模型组成。虽然前沿实验室可以做到,但其他从业者要复刻可能很困难或成本高昂。本节展示了在策略蒸馏可有效用于后训练专业化行为。此方法也适用于持续学习或“测试时训练”:即在模型部署后更新它们,而不会导致基础性能下降。我们使用在公司内部文档上进行中训练的模型作为应用示例。
训练一个内部助手
定制模型的常见目标是充当助手:在某个领域拥有专家知识,并具有可靠的助手式行为。我们可能需要对这两者进行单独训练,尤其是当专业领域知识无法仅从预训练数据中学到,或学习它会干扰行为时。我们的例子是一个公司内部助手,我们有两个期望:1. 模型对该领域(公司文档)知识渊博。预训练模型未见过任何公司内部文档,因此无论模型规模多大,都只能猜测。我们将使用内部知识召回评估(“内部QA”)衡量此点。2. 模型表现出强大的后训练行为,即遵循指令。我们将使用常用的IF-eval衡量此点。
训练新知识会降低已学到的行为
我们将从Qwen3-8B(而非基础模型)开始。Qwen3-8B已通过RL进行后训练,掌握了作为助手有用的技能,如指令遵循和推理。先前研究表明,这种强化学习只训练了原始模型的一小部分子网络,因此当网络在大量数据上进一步训练时,可能会变得很脆弱。我们研究了这种情况发生的程度,以及如何恢复所需行为。为减少灾难性遗忘,中训练中常见做法是混入来自模型原始预训练分布的“背景数据”。在此情况下,我们无法访问Qwen3的预训练分布。因此,我们考虑一个更强、成本更高的基线:我们获取Tulu3提示(广泛的聊天和指令遵循数据集),并使用Qwen3-8B重新采样它们,以充当聊天背景数据。然后,我们在内部文档和聊天数据的不同混合比例上微调Qwen3-8B。提高文档数据比例直接提升模型知识水平。然而,尽管混入至少30%聊天数据有助于保留大部分指令遵循能力,但没有任何一种权重配比能维持在IF-eval上的原始性能。
对于任何给定混合比例,我们都观察到IF-eval性能在微调过程中下降。
另一种常用方法是使用LoRA来约束参数更新,从而减少灾难性遗忘可能性。然而,此方法仍不足以保留IF-eval,而且LoRA学到的知识也更少。
在策略蒸馏恢复后训练行为
接下来,我们试图在对内部文档进行微调后,恢复指令遵循行为。此行为最初用RL训练,成本高昂,而且正如我们所见,很脆弱。取而代之的是,我们在Tulu3提示上,使用模型的早期版本Qwen3-8B作为教师,来进行在策略蒸馏。请注意,此训练阶段与内部文档数据无关,其唯一目的是恢复指令遵循能力。使用模型的早期版本作为教师来“重新唤起”在微调过程中丢失的能力,这使得在策略蒸馏在持续学习方面非常有前景。我们可以交替进行“在新数据上微调”和“蒸馏以恢复行为”这两个阶段,使我们的模型能够随着时间的推移学习并保持知识的最新状态。在70-30混合的内部文档数据和聊天数据上微调后,在策略蒸馏几乎完全恢复了在IF-eval上的性能,且没有损失任何知识;我们还观察到聊天能力和模型在内部QA评估中的“知识”性能之间存在一些正向迁移。
从本质上讲,我们将语言模型本身视为一个奖励模型,高概率的行为会受到奖励。这与逆向RL有关:高概率的行为对应于假定的潜在偏好模型中的有利奖励。任何经过指令调优的开源权重模型都可在此意义上用作奖励模型;我们只需要能访问compute_logprobs函数。
密集监督可极大提高计算效率
强化学习和在策略蒸馏都通过反向KL进行学习,修剪基础策略中存在的动作空间。区别在于奖励的密度。在LoRA Without Regret中,我们提出了信息论观点,即强化学习每个回合只教授O(1)比特。相比之下,蒸馏每个回合教授O(N)比特,其中N是token数量。通过更密集的奖励,我们到底能获得多少训练效率的提升?我们做了一个实验来直接比较两者:从Qwen3-8B-Base(没有额外的SFT)开始。在DeepMath上运行RL。我们使用128的LoRA rank。生成的模型是蒸馏的教师。从RL训练的模型(2)在策略蒸馏回基础模型(1)。
从相同初始化开始,在策略蒸馏学习RL训练的策略所需的梯度步数大约少7-10倍,这对应于50-100倍的计算效率提升。我们看到,蒸馏达到教师性能水平的速度比RL快了大约7-10倍。反向KL下降到接近零,AIME分数在10个梯度步内就得以恢复,而RL则需要70步才能达到该水平。
蒸馏可以有效地重用训练数据以提高数据效率
对于从业者来说,收集大量的训练提示可能既困难又耗时。因此,我们希望能够在训练中多次重用提示。使用RL时,在同一个提示上训练多个轮次常常导致对最终答案的简单记忆。相比之下,在策略蒸馏学习的是近似教师的完整分布,而不是记忆单个答案。这使我们能够用同一个提示训练多个样本。我们重复了上述在数学上训练Qwen3-8B-Base的实验,但这次只使用数据集中随机选择的一个提示。我们在这个提示上连续训练了20步。尽管我们只在一个提示上训练,但我们确实达到了与教师模型相当的性能。
RL在语义策略空间中搜索
我们已经看到,在策略蒸馏可以用少得多的训练步骤来复刻RL提供的学习效果。一种解释是,与预训练不同,RL并未在梯度步骤本身上花费大量计算。我们应该认为RL将其大部分计算用在了搜索上——即推出一个策略并分配功劳——而不是进行更新。预训练通过随机梯度下降探索高维参数空间。预训练需要海量信息,并且非常难以蒸馏。相比之下,我们应该认为RL是在探索语义策略的空间。RL不是在参数空间中探索,而是靠运气“偶然”发现新策略——它只是从它已有的权重集合中随机抽样。一旦找到了一个好的策略,蒸馏就成了学习它的捷径:在策略蒸馏不需要对RL课程中的中间策略进行建模,而只需要学习最终的策略。打个比方:在科学研究中,我们花费大量时间和资源来寻找答案和探索新思想。一旦发现一个结果,用自然语言将其教给别人就简单得多。相比之下,像运动这样的直觉性身体技能,就很难教给别人,因为这些知识存在于一种天生的语言中(例如,肌肉记忆),只有我们自己才能轻易理解。运动只能通过反复练习来学习。
作为持续学习工具的在策略学习
在关于个性化的部分,我们探讨了在策略蒸馏将专业训练行为重新引入模型的能力。这可以推广到更广泛的持续学习任务,这些任务要求在不降低先前能力的情况下获取新知识。先前的工作发现,在策略学习(RL)比较少地遗忘。然而,RL只能塑造行为——它不能很好地教授新知识,因此不足以用于持续学习。在上一节中,我们看到SFT(包括离策略蒸馏)在支持持续学习方面是失败的,因为它会降低行为。我们更深入地研究了这个问题。当我们在模型自己的样本数据集上运行SFT时会发生什么?我们看到,任何大于零的实用学习率都会导致指令遵循评估的性能下降!
一个可能的解释是,虽然KL散度在期望上为0,但每个有限的批次在实践中都会表现出略微不同的分布。在这些有限批次上训练会导致非零的梯度更新,这会使更新后的模型策略偏离其原始状态。在策略蒸馏始终保持在在策略状态,并且由于教师保持不变,学生会收敛于教师的期望行为,而不会像SFT那样在自蒸馏设置中出现性能衰退。这使得在策略蒸馏成为一种非常有前景的持续学习工具。
我们探讨了在策略蒸馏在训练小型模型进行数学推理或持续学习助手等方面的应用。我们将在策略蒸馏与其他两种后训练方法进行了比较:离策略蒸馏和在策略RL。我们发现,在策略蒸馏结合了两者的优点:在策略训练的可靠性能以及密集奖励信号带来的成本效益。后训练是达到前沿模型能力的关键部分。通过利用来自学生的在策略采样和来自教师的密集监督,在策略蒸馏方案能够以前沿高计算量RL运行成本的一小部分,达到这些能力。我们的实现可以在Tinker cookbook中找到。我们的工作探索了在策略蒸馏的简单直接的实例化,以清晰地展示其优势。我们希望继续研究蒸馏的新应用、改进教师监督的新方法,以及提高数据效率和持续学习的方法。在Thinking Machines,我们的使命是为人们提供兼具前沿性能、适应性和个性化的AI模型。在策略蒸馏是实现这一目标的有力工具。
参考链接
https://x.com/thinkymachines/status/1982856272023302322
https://x.com/miramurati/status/1982856564970254772
本文由主机测评网于2026-01-14发表在主机测评网_免费VPS_免费云服务器_免费独立服务器,如有疑问,请联系我们。
本文链接:https://www.vpshk.cn/20260117683.html