总成本降低30倍，告别昂贵后训练，在线策略蒸馏敲开大模型后训练的未来
一家明星公司 Thinking Machines，发了一篇论文，提出一种LLM后训练方法——在线策略蒸馏，让小模型更懂专业领域。
Thinking Machines成立于2024年，由前OpenAI首席技术官Mira Murati联合创立。专注于开发新一代人工智能基础设施和模型，目标是构建安全、可靠且可扩展的 AI 系统。团队还包括多位来自 OpenAI、Google 和 Meta 的资深工程师与研究人员。
在线策略蒸馏用十分之一的成本，达到了比强化学习更高的模型性能。
大型语言模型展现出的专家级能力，并非单一技术的产物，而是输入感知、知识检索、计划选择和可靠执行等多种能力的复杂叠加。
要铸就这些能力，一套完整的训练方法必不可少。这个过程通常被划分为三个宏观阶段。
预训练（pre-training）是第一步，它通过海量的无监督文本，教授模型语言使用、广泛推理和世界知识等通用能力。
中期训练（mid-training）是第二步，它向模型灌输特定领域的知识，比如代码、医疗数据库或一家公司的内部文档。
后训练（post-training）是最后一步，也是至关重要的一步，它负责引导模型产生目标行为，例如遵循指令、进行数学推理或展开对话。
Thinking Machines Lab的Kevin Lu及其合作者发表的这篇研究，就将焦点精准地对准了后训练阶段。
他们提出了一种名为在线策略蒸馏（On-policy distillation）的新方法，巧妙地将在线策略训练的相关性与蒸馏的密集奖励信号结合起来。
这种方法在保持甚至超越现有模型性能的同时，戏剧性地降低了训练成本，为大型语言模型的高效后训练开辟了一条全新的道路。
LLM后训练面临一个根本性的两难选择
大型语言模型训练方法的演进，是一部从简单到复杂的历史。
早期的语言模型几乎完全依赖预训练。但随着模型规模的爆炸式增长和应用需求的日益多样化，人们很快发现，仅靠预训练已无法满足特定任务的高性能要求。中期训练和后训练阶段应运而生。
后训练的方法，主要可以归为两大流派：在线策略训练（on-policy training）和离策略训练（off-policy training）。
在线策略训练，顾名思义，是从学生模型自身生成的轨迹（或称经验）中采样，并为其分配奖励。学生在自己走过的路上学习。
离策略训练，则依赖于一个外部来源提供的目标输出。学生模型需要学习的，是模仿这些现成的、高质量的输出。
强化学习是典型的在线策略训练方法。它通过奖励信号，像驯兽师一样引导模型学习正确的行为。但强化学习有一个致命的弱点：反馈极其稀疏。
无论模型生成了多少个token，在一个训练周期内，它能学到的信息量是固定的。
以解答一道数学题为例。
如果学生模型最终给出了错误答案，强化学习的奖励机制只能告诉它你错了。这个反馈是序列级别的，它无法指出错误究竟发生在哪一个具体的推理步骤。学生模型就像一个考砸了却拿不到批改试卷的孩子，只知结果，不知所以然。
离策略训练则试图解决这个问题。它通常采用监督微调（Supervised Fine-Tuning, SFT）的形式，在一个经过精心策划和标记的任务示例集上进行训练。
这些高质量的标记示例，往往来自于一个在任务上表现出色的教师模型。蒸馏（distillation）就是一种常用的离策略训练机制，它训练学生模型去匹配教师模型的输出分布，学习教师的思考过程。
传统蒸馏方法提供了密集的奖励信号，看似解决了强化学习的痛点。但它引入了一个新的、同样棘手的问题：学生模型学习的上下文，是教师模型经常遇到的上下文，而不是它自己会遇到的。
这会导致一种被称为复合错误（compounding error）的现象。
想象一下，学生模型在解题的早期步骤犯了一个小错误，一个教师模型绝不会犯的错误。这个错误会立刻将它带入一个未知领域，一个在教师提供的训练数据中从未出现过的状态。由于缺乏应对这种陌生局面的经验，学生模型很可能会一步错、步步错，最终离正确答案越来越远。
这就是后训练面临的两难困境：在线的强化学习反馈稀疏，学习效率低下；离策略的蒸馏存在分布不匹配问题，可能导致复合错误。人们似乎必须在学得慢和学不准之间做出抉择。
在线策略蒸馏找到了最优解
在线策略蒸馏的核心思想非常直观：从学生模型自己采样的轨迹中学习，同时使用一个高性能的教师模型，对这条轨迹中的每一个词元进行评分。
它像一位经验丰富的导师，看着学生亲手完成作业，并在每一个步骤旁边给出即时反馈。
这种方法完美结合了在线策略训练的相关性（学生在自己遇到的状态下学习）和蒸馏的密集奖励信号（每个步骤都有反馈），实现了鱼与熊掌兼得的效果。
回到数学解题的例子。在线策略蒸馏会对学生解答过程的每一步进行评分。它不仅会惩罚那些最终导致错误答案的步骤，还会强化那些被正确执行的步骤。
这提供了一个比传统强化学习密集得多的反馈信号，同时又完美规避了离策略学习的分布不匹配问题。
为了评估学生模型的轨迹，在线策略蒸馏可以采用多种损失函数。研究人员为了简化问题，选择了一种高效且强大的度量方式：每词元反向KL散度（reverse KL）。
它衡量的是，在给定相同历史轨迹的条件下，学生模型（π_θ）的预测分布与教师模型（π_teacher）的预测分布之间的差异。
奖励函数的目标就是最小化这个反向KL散度。这会推动学生模型，在它自己所处的每一个状态下，都尽可能地模仿教师模型的行为。
当学生模型的行为与教师模型完全一致时，反向KL散度为零。为了简化计算，研究者使用了零折扣因子，这意味着在任何给定的时间步，学生模型只关心如何优化紧接着的下一个词元，而不去考虑更遥远的未来。
反向KL散度与强化学习有着天然的亲和力，因为强化学习本质上也是在优化某种由奖励模型定义的序列级反向KL。
但与大多数奖励模型不同，反向KL是不可破解的。一个低的KL值总是对应着从教师模型角度看的高概率期望行为，不存在被学生模型利用以获得高奖励却产生无用输出的漏洞。
反向KL还有另外两个非常实用的特性。
它是模式寻求的。这意味着它会引导学生学习一种特定的、高质量的行为（教师模型的行为），而不是将概率分散在几个次优的选项上，让学习目标更加明确。
它还能有效减少暴露偏差（exposure bias）。暴露偏差指的是模型在训练时看到的是真实数据，但在推理时看到的是自己生成的数据，这种差异导致的性能下降。在线策略蒸馏通过让学生在自己生成的数据上学习，天然地缓解了这个问题。
实现这套方法的过程也相对简单，可以基于现有的强化学习脚本进行构建。Thinking Machines Lab利用其Tinker训练API，主要通过以下几个步骤实现：
初始化教师客户端。Tinker API使得为不同模型创建客户端变得轻而易举，无需担心底层模型引擎的资源利用率。这里使用采样客户端，因为只需要教师模型进行前向传播计算概率，而不需要反向传播梯度。
采样轨迹。像在标准强化学习中一样，从学生模型采样轨迹。在采样过程中，学生的对数概率已经被计算出来，这是后续重要性采样损失计算的一部分。
计算奖励。使用
compute_logprobs函数查询教师客户端，获取教师模型对学生采样出的那些词元的对数概率。然后，利用这个值计算出每一步的反向KL散度。使用强化学习进行训练。将每词元的优势函数（advantage）设置为负的反向KL散度，然后调用强化学习的重要性采样损失函数，对学生模型执行一次训练更新。
这个流程带来了显著的计算效率提升。
由于计算奖励的过程非常快，不必等待整个长轨迹采样完成，可以使用较短或部分的轨迹就开始训练，大大加快了迭代速度。
查询教师模型的对数概率，也仅仅需要一次从较大模型的前向传播，而生成轨迹的任务则由更小、更便宜的学生模型来完成。
更重要的是，这种方法不需要一个独立的奖励模型或标记模型。教师模型本身就扮演了奖励模型的角色，高明且廉价。未来一个有趣的研究方向，可能是将这种基于蒸馏的每词元奖励与传统的序列级环境奖励结合起来。
惊人的效率提升是其最直接的优势
理论上的优雅必须通过实践来检验。研究者使用在线策略蒸馏进行后训练，并将其与传统方法进行了直接比较，以此来展示其作为训练专家模型最后关键阶段的威力。
实验场景设定在数学推理能力的训练上。
研究者选择Qwen3-8B-Base模型作为学生，Qwen3-32B作为教师，目标是提升学生模型的数学推理能力。这两个模型都是Tinker平台目前支持的模型，使得实验可以被轻松复现。
所有实验都从一个共同的起点开始：一种离策略蒸馏形式的中期训练。
具体来说，就是在一个由教师模型生成的数据集上，对学生模型进行监督微调。这个数据集名为OpenThoughts-3，它是由一个类似Qwen3-32B的推理模型QwQ-32B生成的推理提示和响应集合。
在40万个提示上进行完全微调后，学生模型Qwen3-8B-Base在AIME'24（一个数学问题基准测试）上达到了60%的分数。研究者也尝试了LoRA（低秩适应）进行训练，但在这种高容量数据集上，LoRA的效果不如完全微调。
所有情况下的性能增长都呈现出一种对数线性关系——初期的性能提升成本低廉，但越往后，每提升一个百分点所需的成本就越高。
现在，这个经过40万提示微调、AIME'24分数达到60%的模型，成为了比较不同后训练方法效果的检查点。目标是：将AIME'24分数从60%提高到70%，看看哪种方法最省力。
默认的方法是继续进行监督微调，也就是延续离策略蒸馏的过程。根据之前观察到的对数线性趋势进行外推，研究者估计，模型大约需要在200万个提示上进行训练，才能在AIME'24上达到70%的分数。
这个外推的前提是缩放定律（scaling law）持续有效而不会停滞，这本身就是一个不小的挑战。尽管如此，确实存在通过大规模离策略蒸馏将8B模型性能提升到70%以上的先例。
Qwen3在一个类似的SFT初始化的基础上，通过强化学习，花费了17920个GPU小时，在AIME'24基准测试上达到了67.6%的性能。
直接比较这个成本与蒸馏的成本很困难，但根据对SFT训练堆栈的一些合理假设，这大致相当于训练200万个离策略蒸馏提示的成本。
真正引人注目的是，Qwen团队报告称，他们使用在线策略蒸馏，仅用了强化学习十分之一的成本（1800 GPU小时），就在AIME'24上达到了更高的74.4%的分数。这正是启发这项研究工作的关键发现。
研究者在自己的基础设置中尝试复现这一惊人的结果。
从60%分数的SFT检查点开始，在线策略蒸馏仅用了大约150个训练步骤，就将AIME'24的分数提升到了70%。
比较不同方法的计算成本是一件复杂的事情，因为训练、采样和对数概率计算的成本比例因具体实现而异。
为了得到一个更公平的比较，研究者以浮点运算次数（FLOPs）来计算成本。这种度量方式会惩罚那些可以在GPU上高效并行化的操作，因此它实际上高估了计算对数概率的真实成本。
结果是惊人的。
当SFT数据集是现成的（比如OpenThoughts-3），或者其生成成本可以在多次训练中分摊时，在线策略蒸馏的基线成本降低了9倍。
在这种情况下，不计入离策略训练的教师FLOPs成本，但必须计入在线策略的，因为教师模型需要实时运行以计算学生轨迹的对数概率。
考虑到这种计算可以在GPU上廉价地并行化，换算成GPU小时的成本降低幅度更接近18倍。
在更现实的场景中，常希望在没有现成离策略蒸馏数据集的新任务上训练模型。如果将离策略蒸馏中生成数据的教师模型成本也完全计算在内，那么总成本的降低幅度达到了约30倍。
密集监督带来的效率提升是显著的。强化学习和在线策略蒸馏都通过反向KL进行学习，本质都是在修剪基础策略中存在的动作空间。它们的根本区别在于奖励的密度
一篇名为LoRA Without Regret的研究提出了一个信息论的视角：强化学习每个周期只教授O(1)比特的信息。相比之下，蒸馏每个周期能教授O(N)比特的信息，其中N是序列中的词元数量。
为了量化这种密集奖励带来的训练效率提升，研究者进行了一项直接对比实验：
从Qwen3-8B-Base模型开始，不进行任何额外的SFT。
在DeepMath数据集上运行强化学习，得到的模型作为蒸馏的教师。
使用在线策略蒸馏，将从强化学习训练出的模型的知识，蒸馏回基础模型。
结果显示，在线策略蒸馏达到教师性能水平的速度，比强化学习快了约7-10倍。反向KL散度迅速降至接近零，AIME分数在不到10个梯度步内就恢复了，而强化学习则需要70步才能达到同样的水平。
累计来看，所需的总计算量减少了约50-100倍。
这其中的原因有两点。强化学习通常需要在接近评估上下文长度的序列上进行训练，以便策略能够学习到上下文的限制。
而蒸馏在较短的上下文长度下也能合理学习，因为奖励是密集的，不会因为轨迹是否完成而出现急剧的悬崖。
当SFT初始化较强时，在线策略蒸馏可以有效地使用更小的批量大小，因为它每个周期提供的信息比特更多，从而降低了梯度噪声。
这些结果有力地表明，过程监督和密集奖励有潜力将学习效率提高一个数量级。
它为持续学习和个性化打开了新大门
除了在通用任务上训练高性能小模型，蒸馏的另一个重要用例是个性化。
例子包括让模型在对话中遵循特定的语调和输出格式，或者实现工具使用和成本预算等特定功能。通常希望将这种定制化的行为与新的领域知识一起训练到模型中。
同时训练知识和行为通常非常困难。
轻量级的微调往往不足以实现目标，需要进行更大规模的中期训练。而在新知识的基础上学习后训练行为，则需要一套复杂的后训练堆栈，通常依赖于专有数据和奖励模型。这种方法对于顶尖实验室来说是可行的，但对于大多数从业者而言，成本过高且难以复制。
在线策略蒸馏在这里展现了其作为一种高效后训练工具的价值。
这种方法同样适用于持续学习或测试时训练：在模型部署后，不断用新数据更新它，同时不损害其基础性能。
研究者用一个在内部公司文档上进行中期训练的模型作为例子。
定制模型的一个常见目标是让它扮演一个助手角色：既在某个领域拥有专业知识，又具备可靠的助手行为。这两个目标可能需要分开训练，特别是当专业领域知识无法仅从预训练数据中学到，或者学习新知识的过程会干扰原有行为时。
实验中的内部公司助手有两个期望：
模型了解公司文档。预训练模型没有见过这些内部文档，因此只能猜测。研究者使用内部知识问答评估（内部QA）来衡量这一点。
模型表现出强大的后训练行为，即指令跟随能力。研究者使用常用的IF-eval基准来衡量这一点。
这次，研究者从Qwen3-8B开始，这是一个已经经过指令跟随和推理强化学习后训练的模型。
先前的研究表明，这种强化学习可能只训练了原始模型的一小部分子网络，因此当网络在大量新数据上进一步训练时，这些习得的能力可能非常脆弱。
研究者首先探究了灾难性遗忘发生的程度。
为了减少灾难性遗忘，中期训练的一个常见做法是混合来自原始模型预训练分布的背景数据。
但研究者无法访问Qwen3的预训练分布。因此，他们采用了一个更强但更昂贵的基线：使用Tulu3提示（一个广泛的聊天和指令跟随数据集），并用Qwen3-8B重新采样，作为聊天背景数据。
这种由Qwen3-8B采样的在线策略背景数据，充当了一个前向KL正则化器，在整个中期训练过程中强化模型的原始行为。
然后，研究者在内部文档和聊天数据的不同混合比例上微调Qwen3-8B。
结果显示，增加文档数据的比例直接提高了模型的知识水平。但同时，虽然混合至少30%的聊天数据有助于保留大部分指令跟随能力，但没有任何一个混合权重能够保持IF-eval上的原始性能不下降。
对于任何给定的混合比例，研究者都观察到IF-eval性能在微调期间持续下降。这严重损害了使用更长时间的训练来进一步专业化模型的能力。
接下来，研究者试图在内部文档微调之后，恢复模型的指令跟随行为。
这种行为最初是用昂贵且脆弱的强化学习训练出来的。这一次，他们选择在Tulu3提示上运行在线策略蒸馏，并巧妙地使用模型的早期版本——即原始的Qwen3-8B——作为教师。
这个训练阶段与内部文档数据完全无关，其唯一目的就是恢复被遗忘的指令跟随能力。
使用模型的一个早期版本作为教师，来重新激活在微调过程中丢失的能力，这使得在线策略蒸馏在持续学习方面非常有前景。
可以设想一种训练模式：在新数据上进行微调阶段，然后进入恢复行为的蒸馏阶段，两者交替进行。这使得模型能够随着时间的推移不断学习新知识，同时保持核心能力不退化。
结果非常理想。
在一个70%内部文档和30%聊天数据的混合微调之后，在线策略蒸馏几乎完全恢复了模型在IF-eval上的性能，并且没有丢失任何新学到的知识。研究者甚至观察到聊天能力和模型在内部QA评估上的知识性能之间存在一些正向迁移。
本质上，在线策略蒸馏将语言模型本身视为一个奖励模型，高概率的行为会得到奖励。
这与逆强化学习（Inverse Reinforcement Learning）有异曲同工之妙。
任何经过指令调优的开源模型，只要能访问其计算对数概率的函数，就可以在这种意义上被用作奖励模型。
蒸馏作为整合行为和知识的工具，其潜力远不止于此。在线学习可以是增强类似模型合并设置的关键工具。
更有趣的是，在线策略蒸馏还能极大地提高数据效率。
对于许多从业者来说，收集大规模的训练提示数据集既困难又耗时。因此，在训练中多次重复使用提示就显得尤为重要。
使用强化学习时，在同一个提示上训练多个周期，通常会导致模型简单地记住最终答案，而不是学会推理过程，特别是对于大模型。
相比之下，在线策略蒸馏通过最小化反向KL散度，学习的是近似教师的完整分布，而不仅仅是记忆单个答案。这允许在同一个提示上训练多个样本而不导致过拟合。
研究者重复了在数学上训练Qwen3-8B-Base的实验，但这次只使用数据集中随机选择的一个提示。
他们在这个唯一的提示上连续训练了20步，总共生成并评分了5120个序列。尽管这种方式计算效率较低，但结果令人惊讶：仅靠这一个训练提示，学生模型就成功匹配了教师模型的AIME'24性能。
在线策略蒸馏的成功，让人们重新思考强化学习的本质。
或许我们应该将强化学习看作是在一个语义策略空间中进行探索。
它的大部分计算并非花费在梯度更新上，而是花费在搜索上——推出不同的策略并为其分配信用。
它不是在参数空间中艰难跋涉，而是在已有的能力集合中进行随机组合，偶然发现新的有效策略。
一旦找到了好的策略，蒸馏就成了学习它的捷径。
在线策略蒸馏不需要建模强化学习课程中所有中间的、不成熟的策略，它直奔终点，只建模那个最终学到的、最优的策略。
这就像科学研究。我们花费大量时间和资源去探索、试错，最终找到答案。一旦结果被发现，用自然语言把它教给别人，就变得相对简单。
与之相对的是直觉性的身体技能，比如运动。这些技能更难教授，因为知识存在于一种内在的语言——肌肉记忆中。运动只能通过反复练习来学习。强化学习更像后者，而在线策略蒸馏则像前者中的教学环节。
在线策略蒸馏始终保持在线策略，并且由于教师是固定的，学生会稳定地收敛于教师的理想行为，而不会像自蒸馏的SFT那样性能退化。
这使其成为持续学习的一个极具前景的强大工具。
这项研究探索了在线策略蒸馏在数学推理和持续学习等应用中的巨大潜力。
通过利用来自学生的在线策略采样和来自教师的密集监督，在线策略蒸馏配方以一小部分成本，达到了过去需要高昂计算成本的强化学习才能达到的能力。
后训练是通往最前沿模型能力的关键一环，而在线策略蒸馏，无疑是这条路上最高效、最经济的路径之一。
参考资料：
https://thinkingmachines.ai/blog/on-policy-distillation/
https://arxiv.org/abs/2303.08774
