跳到正文
北京时间
原文
Thinking Machines Lab:官方博客· Thinking Machines Lab·· 2025-10-27精选AI 评分66

Thinking Machines Lab 提出 On-Policy Distillation:以 RL 约 1/10 的成本达到相近性能

On-Policy Distillation

AI 导读

Thinking Machines Lab 发布博客介绍 on-policy distillation 方法,从学生模型采样轨迹并由教师模型对每个 token 按 reverse KL 打分,结合了 on-policy 训练的相关性与蒸馏的密集奖励信号。

推荐理由

原文给出从 SFT、RL 到 on-policy 蒸馏的对比实验与成本数据,读者可以据此评估后训练方法的算力取舍。

正文 · AI 翻译

LLM 能够在特定领域达到专家级表现,这是多种能力叠加的结果:对输入的感知、知识检索、方案选择和可靠执行。这需要一整套训练方法,我们可以将其大致分为三个阶段:

  • 预训练教会模型通用能力,如语言使用、广泛推理和世界知识。
  • 中期训练传授领域知识,例如代码、医学数据库或公司内部文档。
  • 后训练引导出特定行为,例如指令遵循、数学问题推理或对话。

经过更强训练的较小模型,在其受训的专业领域中往往优于更大的通用模型。使用较小模型有许多好处:出于隐私或安全考虑,它们可以本地部署;可以更容易地持续训练和更新;还能节省推理成本。要利用这些优势,就需要为后训练阶段选择正确的方法。

对“学生”模型进行后训练的方法可分为两类:

  • 同策略训练从学生模型自身采样轨迹,并为其分配某种奖励。
  • 异策略训练依赖来自某些外部来源的目标输出,让学生模型学习模仿。

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

我们可以通过强化学习进行同策略训练,对学生的每次轨迹是否解出题目进行评分。评分可以由人类完成,也可以由一个能可靠得出正确答案的“教师”模型完成。

同策略训练的优势在于,通过在自身采样上训练,学生能以更直接的方式学会避免错误。但强化学习有一个主要缺点:它提供的反馈非常稀疏,无论使用多少 token,每个训练回合只传授固定数量的比特。在上面的例子中,学生得知“21”是错误答案,并更新以避免产生它尝试过的轨迹。但它并不知道错误具体出在哪里,是运算顺序错了,还是算术本身出了错。这种反馈的稀疏性使强化学习在许多应用中效率低下。

异策略训练通常通过监督微调(SFT)进行:在一组精心挑选的任务特定标注样本上训练。这些标注样本的来源可以是一个已被证明在该任务上表现良好的教师模型。

我们可以使用一种称为蒸馏的机制:训练学生模型去匹配教师模型的输出分布。我们在教师的轨迹上训练:即包括中间思考步骤在内的完整生成 token 序列。我们可以使用教师在每一步的完整下一 token 分布(通常称为“logit 蒸馏”),也可以只对给定序列进行采样。在实践中,采样序列能对教师的分布给出无偏估计,并达到相同的目标。学生模型会针对序列中的每个 token 进行更新,更新幅度与该 token 自身生成它的可能性成比例,在下面的示例中用更深的颜色表示:

从大型模型教师进行蒸馏已被证明能有效训练小模型遵循指令,Alpaca: A Strong, Replicable Instruction-Following Model(Taori 等,2021)进行数学和科学推理,OpenThoughts: Data Recipes for Reasoning Models(Guha 等,2025)从医疗记录中提取临床信息Distilling Large Language Models for Efficient Clinical Information Extraction(Vedula 等,2025),以及参与多轮聊天对话。Enhancing Chat Language Models by Scaling High-quality Instructional Conversations(Ding 等,2023)用于这些及其他应用的蒸馏数据集通常会被开源并发布。

离策略训练的缺点在于,学生是在教师经常出现的上下文中学习,而不是学生自己经常会遇到的上下文。这可能导致误差累积:如果学生犯了教师从不犯的早期错误,它就会发现自己越来越偏离训练中观察到的状态。当我们关心学生在长序列上的表现时,这个问题尤为突出。为避免这种偏离,学生必须学会从自己的错误中恢复。

在离策略蒸馏中观察到的另一个问题是,学生可能学会模仿教师的风格和自信,但不一定学会其事实准确性。The False Promise of Imitating Proprietary LLMs(Gudibande 等,2023)

如果你正在学习下棋,在策略 RL 类似于在没有指导的情况下对弈。赢或输的反馈直接与你的下法相关,但每局只能收到一次,而且不会告诉你哪些走法对结果贡献最大。离策略蒸馏类似于观看大师下棋——你观察到极强的棋步,但它们出现在新手很少会遇到的棋盘状态中。

我们希望将 RL 的在策略相关性与蒸馏的密集奖励信号结合起来。对于学习下棋,这将是一位老师,对你自己的每一步棋从“失误”到“精彩”进行评分。对于 LLM 后训练,这就是在策略蒸馏。

来自chess.com的截图。每一步棋都由分析引擎进行颜色分级,将棋步标记为失误(红色)、错误(橙色)、不精确(黄色)或精彩(蓝色)。

在策略蒸馏——两全其美

在策略蒸馏的核心思想是从学生模型采样轨迹,并使用高性能教师对每条轨迹的每个 token进行评分。回到我们上面的数学示例,在策略蒸馏会对解决方案的每一步打分,惩罚导致学生得出错误答案的错误,同时强化正确执行的步骤。

在这篇文章中,我们探索了将同策略蒸馏应用于诸如训练数学推理模型以及训练结合领域知识与指令遵循的助手模型等任务。我们在已具备预训练和中期训练能力基础的模型上应用同策略蒸馏。我们发现,这是一种廉价而强大的后训练方法,兼具同策略训练与稠密奖励信号的优势。

方法采样奖励信号
监督微调离策略稠密
强化学习同策略稀疏
同策略蒸馏同策略稠密

我们的同策略蒸馏工作受到 DAGGER 的启发,A Reduction of Imitation Learning and Structured Prediction to No-Regret Online Learning(Ross 等,2010)这是一种迭代式 SFT 算法,包含教师对学生所访问状态的评估。它也与 过程奖励建模Let’s Verify Step by Step(Lightman 等,2023)相似,后者是一种对学生在思维链中每一步进行评分的 RL 方法。我们扩展了先前同策略蒸馏的工作,包括 Agarwal 等,On-Policy Distillation of Language Models: Learning from Self-Generated Mistakes(Agarwal 等,2023)Gu 等,MiniLLM: Knowledge Distillation of Large Language Models(Gu 等,2023)以及 Qwen3 团队Qwen3 Technical Report(Qwen 团队,2025)。使用 Tinker 训练 API,我们复现了 Qwen3 的结果,即以远低于 RL 的成本,通过同策略蒸馏在推理基准上达到同等性能。

实现

你可以在 Tinker cookbook 中跟随实现的每一步。

损失函数:反向 KL

同策略蒸馏可以使用多种损失函数来为学生轨迹评分。关于各种损失函数选择的分析,请参见 Agarwal 等。为简单起见,我们选择逐 token 反向 KL——即在相同先验轨迹条件下,学生(πθπθ​)与教师(πteacherπteacher​)分布之间每个 token 的散度:

KL(πθ∣∣πteacher)=Ex∼πθ[log⁡πθ(xt+1∣x1..t)−log⁡πteacher(xt+1∣x1..t)]KL(πθ​∣∣πteacher​)=Ex∼πθ​​[logπθ​(xt+1​∣x1..t​)−logπteacher​(xt+1​∣x1..t​)]

我们的奖励函数最小化反向 KL,这促使学生在它所处的每个状态下都去逼近教师的行为。当学生的行为与教师完全一致时,反向 KL 为零。为简单起见,我们使用零折扣因子:在任意给定时间步,尽管从数学上更正确,但我们在实践中并未发现折扣因子 > 0 能提升性能,因此为简单起见选择零。学生只优化紧接着的下一个 token,不考虑未来的 token。

反向 KL 与 RL 具有天然的协同性,因为 RL 通常优化的正是由奖励模型诱导的一种序列级反向 KL。然而,与实践中的大多数奖励模型不同,反向 KL 是“不可被钻空子”的,也就是说,从教师模型的视角来看,低 KL 总是对应着期望行为的高概率。反向 KL 的另外两个有用特性是:它具有“模式寻求”关于模式寻求行为的更多讨论,参见 Eric Jang 的文章。——它学习一种特定的行为(教师的行为),而不是把分布分散到若干次优选项上——并且它能减少暴露偏差。Scheduled Sampling for Sequence Prediction with Recurrent Neural Networks(Bengio 等,2015)。更多讨论参见 Gu 等。

这种方法能显著节省计算量。由于它不需要等待一次 rollout 完成采样来计算奖励,我们可以在训练中使用更短或部分的 rollout。查询教师的对数概率也只需要从较大模型进行一次前向传播,而轨迹则由更小、更廉价的学生模型生成。

我们也不需要单独的奖励模型或标注模型。将基于蒸馏的逐 token 奖励与序列级环境奖励相结合可能具有优势;这是一个值得未来研究的有趣方向。

示意图

下面我们看到一个由教师评分的错误学生轨迹的真实示例。该示例来自 SimpleBench,它依赖于模型做出一个关键观察,即问题的前提很重要:正确答案是“B. 0”,因为冰块在煎锅中会融化。学生模型 Qwen3-4B-Instruct-2507 错误地将其视为纯数学问题,没有考虑物理背景。

由教师模型评分的示例轨迹。深红色的 token 对应更高的反向 KL。

颜色越深表示该 token 从教师模型 Qwen3-235B-A22B-Instruct-2507 处受到的惩罚越高,而该教师模型能正确解决此问题。我们看到,它会惩罚那些开启引导学生误入歧途的短语的 token,直观上对应于引导推理的重要“分叉 token”。Beyond the 80/20 Rule: High-Entropy Minority Tokens Drive Effective Reinforcement Learning for LLM Reasoning(Wang 等,2025) 最终答案虽然是错误的,却并未受到惩罚——在给定整个前序序列的条件下,它是完全可以预测的。

伪代码

我们在 Tinker 中的 RL 脚本之上实现在策略蒸馏,该脚本已经实现了采样、奖励计算和策略梯度式训练。我们的实现实际上可以在使用 KL 正则化的 RL 实现之上做一行改动:我们只需替换正则化模型。

  1. 初始化教师客户端。 Tinker API 可以轻松地为不同模型创建不同的客户端,而无需担心模型引擎的利用率。我们使用采样客户端,因为我们不需要通过教师模型传播 logprobs。
  2. 采样轨迹。我们完全按照在 RL 中的方式从学生模型采样 rollout。在采样过程中,RL 已经计算了学生模型的 logprobs log⁡πθ(x)logπθ​(x),用作 重要性采样损失的一部分。
  3. 计算奖励。我们用 compute_logprobs 在采样得到的轨迹上查询教师客户端,它会返回教师模型在学生模型采样的 token xx 上的 logprobs log⁡πteacher(x)logπteacher​(x)。在本文中,我们的所有实验都不考虑 logit(top-k)蒸馏,而该方法本可用于进一步提升计算效率。随后我们用它来计算反向 KL。
  4. 用 RL 训练。我们将每个 token 的优势设为负的反向 KL,并调用 RL 重要性采样损失函数对学生模型执行训练更新。
# Initialize teacher client (main):
teacher_client = service_client.create_sampling_client(
    base_model=teacher_config.base_model,
    model_path=teacher_config.load_checkpoint_path,
)

# Sample trajectories (main):
trajectories = do_group_rollout(student_client, env_group_builder)
sampled_logprobs = trajectories.loss_fn_inputs["logprobs"]

# Compute reward (compute_teacher_reverse_kl):
teacher_logprobs = teacher_client.compute_logprobs(trajectories)
reverse_kl = sampled_logprobs - teacher_logprobs
trajectories["advantages"] = -reverse_kl

# Train with RL (train_step):
training_client.forward_backward(trajectories, loss_fn="importance_sampling")

在下面的实验中,我们通常对已经在特定领域知识上经过中期训练的模型应用同策略蒸馏。这种训练提高了学生模型生成落在教师分布内的 token 的概率,尽管这通常远不足以复现教师模型的性能。使用前向 KL 的 SFT 会为新 token 增加支撑集。反向 KL 方法随后可以在初始化支撑集内进行模式搜索。正如我们将在个性化示例中看到的,由于学生模型缺乏任何相关领域知识,生成相关 token 的概率往往从零开始。

我们将同策略蒸馏用于后训练,并将其与这一专家模型训练最后关键阶段的其他方法进行比较。

面向推理的蒸馏

我们使用蒸馏来训练 Qwen3-8B-Base 模型的数学推理能力,并以 Qwen3-32B 作为教师模型。教师模型(Qwen3-32B)和学生模型(Qwen3-8B-Base)目前都是 Tinker 上的受支持模型,因此你可以用 Tinker cookbook 复现我们的实验。更新(2026 年 6 月):此处使用的教师模型(Qwen3-32B)和学生模型(Qwen3-8B-Base)此后已从 Tinker 模型阵容中退役。Tinker cookbook 中的蒸馏配方已更新为使用 Qwen3.5-9B-Base 作为学生模型、Qwen3.5-9B 作为教师模型,因此这些实验仍可复现。

离策略蒸馏

如上所述,我们所有的实验都从离策略蒸馏形式的中期训练开始——在教师模型生成的示例数据集上进行监督微调。用于数学推理的数据集是 OpenThoughts-3,它是由 QwQ-32B(一个与 Qwen3-32B 类似的推理模型)生成的推理提示与回复的集合。

在 40 万条提示上对学生的模型(Qwen3-8B-Base)进行全参数微调,在数学题基准 AIME’24 上取得 60% 的分数。我们也可以用 LoRA 训练,LoRA: Low-Rank Adaptation of Large Language Models (Hu et al, 2021),不过在训练大规模数据集时,它落后于全参数微调。在所有情况下,我们都看到性能呈对数线性增长——最初的性能提升成本较低,但后期的提升成本高昂。

在离策略蒸馏(SFT)过程中 AIME’24 分数的变化。在最初的 5 万至 10 万条提示之后,性能遵循可预测的对数线性缩放曲线。正如 LoRA Without Regret 中所预测的,在使用大批量进行大规模 SFT 时,我们观察到 LoRA 的性能更差。

我们可以把在 40 万条提示上微调后的模型视为一个检查点,然后再尝试各种后训练方法来提升其性能。我们可以比较将 AIME’24 基准上的分数从 60% 提升到 70% 所需的投入。

默认方法是在更多提示上进行微调,延续离策略蒸馏的过程。外推对数线性趋势,我们估计该模型在约 200 万条提示时将在 AIME’24 上达到 70%。这一外推要求缩放定律在不出现停滞的情况下成立,而这并非易事。不过,也有大规模离策略蒸馏将 8B 模型性能提升至 70% 以上的例子,例如 OpenThoughts-3 和 DeepSeek-R1-0528-Qwen3-8B。DeepSeek-R1-0528-Qwen3-8B 在该基准上达到 86%,但未说明训练提示的数量。较旧的模型(Qwen2.5-7B、Qwen2.5-14B)在使用 DeepSeek-R1 的 80 万条蒸馏提示训练后,分别达到 55.5% 和 69.7% 的性能。我们可以将这一外推作为离策略蒸馏性价比的乐观估计。

强化学习

Qwen3 技术报告在类似的 SFT 初始化基础上,使用 17,920 GPU 小时的强化学习在该基准上达到 67.6% 的性能。很难将其直接与蒸馏的成本进行比较,但基于对 SFT 训练栈的一些合理假设,这与在 200 万条离策略蒸馏提示上训练的成本相近。

方法AIME’24GPQA-DiamondGPU 小时
离策略蒸馏55.0%55.6%未报告
+ 强化学习67.6%61.3%17,920
+ 同策略蒸馏74.4%63.3%1,800
来自 Qwen3 Technical Report,表 21。

Qwen 团队还报告称,以强化学习十分之一的成本通过同策略蒸馏在 AIME’24 上达到了 74.4 的更高分数,这启发了我们的工作。我们在下面的基础设置中尝试复现它。

同策略蒸馏

作为离策略蒸馏或强化学习的替代方案,我们按上述方式运行同策略蒸馏。我们实际上使用 Qwen3-8B 作为教师模型,因为它的表现略好。不过,为了比较计算量,我们仍然可以测量 32B 模型的 FLOPs。从 40 万条 SFT 检查点出发,同策略蒸馏在约 150 步内使 AIME’24 达到 70%。150 步约对应 7.7 万条提示;我们对每条提示使用 4 个样本进行训练。

AIME’24 在 on-policy 蒸馏过程中的表现。我们以训练 FLOPs 来衡量额外计算量(见下文)。On-policy 蒸馏在计算效率上显著优于 SFT,尤其是对 LoRA 模型而言。在 rank = 32 时,LoRA 在 SFT 后比全量微调落后 13%,但在 on-policy 蒸馏后仅落后 6%。

跨方法比较计算成本并非易事,因为训练、采样与 log-prob 计算成本之间的比例会因实现方式不同而差异显著。下面我们以 FLOPs 计算成本,这会不利于那些可在 GPU 上有效并行化的方法。特别是,它会高估计算 log-probs 的实际成本。

方法AIME’24教师 FLOPs学生 FLOPsCE 与 SFT-2M 对比
初始化:SFT-400K60%8.5 × 10203.8 × 1020–
SFT-2M(外推)~70%(外推)3.4 × 10211.5 × 10211×
强化学习68%--≈1×
On-policy 蒸馏70%8.4 × 10198.2 × 10199-30×

我们发现,当 SFT 数据集已给定(如我们使用 OpenThoughts-3 的示例),或在多次训练运行中摊销时,基线成本可降低 9 倍。CE = (学生 + 教师) / (学生) 在这种情况下,对于 off-policy 训练我们不计教师 FLOPs 的成本,但对于 on-policy 训练则计入,因为我们必须运行教师模型来计算学生轨迹的 log-probs。由于该计算可在 GPU 上低成本并行化,GPU 小时数上的成本降低接近 18 倍。

然而,我们常常希望在新任务上训练一个小模型,而该任务没有可用的 off-policy 蒸馏数据集。如果我们在 off-policy 蒸馏中计入教师模型的全部成本——即包括从教师模型采样的额外成本——总成本降低约为 30 倍。CE = (学生 + 教师) / (学生 + 教师)

用于个性化的蒸馏

除了训练小模型在常见任务上达到高性能外,蒸馏的另一个用例是个性化。示例包括在对话中遵循特定语气和输出格式,或工具使用和成本预算等能力。我们常常希望将这种行为与新的领域知识结合起来训练。

同时训练两者通常很困难,而轻量级微调往往不足以实现这一目标,Unfamiliar Finetuning Examples Control How Language Models Hallucinate(Kang 等,2024) 因此需要更大规模的 midtrain。在新知识之上学习后训练行为需要复杂的后训练技术栈,通常包括专有数据和奖励模型。虽然这种方法对前沿实验室来说可行,但其他从业者可能难以复制或成本过高。

在本节中,我们展示 on-policy 蒸馏可有效用于后训练专门行为。这种方法也适用于持续学习或“测试时训练”:在模型部署时更新模型而不使基础性能回退。我们使用一个在我们的内部公司文档上 midtrain 的模型作为示例应用。

训练内部助手

自定义模型的一个常见目标是充当助手:在某个领域拥有专家知识,此外还具备可靠的助手式行为。我们可能需要分别对两者进行训练,尤其是当专业领域无法仅从预训练数据中学习,或学习它会干扰行为时。

我们的示例是一个公司内部助手,对此我们有两个期望:

  1. 模型对领域(公司文档)知识渊博。预训练模型从未见过公司的任何内部文档,因此无论模型规模多大,都只能靠猜测。我们将使用内部知识回忆评估(“internal QA”)来衡量这一点。
  2. 模型展现出强大的后训练行为,即指令遵循。我们将使用常用的 IF-eval 来衡量这一点。Instruction-Following Evaluation for Large Language Models(Zhou 等,2023)

在新知识上训练会降低已学习的行为

我们将从 Qwen3-8B 开始,而不是基础模型。Qwen3-8B 已针对助手的有用技能进行了后训练,例如指令遵循和基于 RL 的推理。先前的研究表明,这种强化学习只训练原始模型的小型子网络,Reinforcement Learning Finetunes Small Subnetworks in Large Language Models(Mukherjee 等,2025)因此当网络在大量数据上进一步训练时,这种能力可能很脆弱。我们研究这种情况发生的程度,以及如何恢复期望的行为。

为了减少这种灾难性遗忘,中期训练中的一种常见做法是混入来自原始模型预训练分布的“背景数据”。Midtraining Bridges Pretraining and Posttraining Distributions(Liu 等,2025)在这种情况下,我们无法访问 Qwen3 的预训练分布。因此,我们考虑一个更强也更昂贵的基线:我们取 Tulu3Tulu 3: Pushing Frontiers in Open Language Model Post-Training(Ivison 等,2024)的提示——一个广泛的聊天和指令遵循数据集——并用 Qwen3-8B 重新采样它们,以充当聊天背景数据。

这种由 Qwen3-8B 采样的“同策略”背景数据充当正向 KL 正则化器,在整个中期训练过程中强化模型的原始行为。我们发现,在整个中期训练过程中,从 Qwen3-8B 采样比从 Qwen3-32B 采样更能保持聊天能力,这凸显了数据源的敏感性;Chen 等人也发现了类似的同策略 SFT 结果。Retaining by Doing: The Role of On-Policy Data in Mitigating Forgetting(Chen 等,2025)我们假设这种方法甚至可能比访问原始预训练数据分布更有效,代价是必须采样大规模数据集。

然后,我们在内部文档和聊天数据的不同混合比例上微调 Qwen3-8B。增加文档数据的比例会直接提升模型的知识。然而,尽管混入至少 30% 的聊天数据有助于保留大部分指令遵循能力,但没有任何权重配置能维持 IF-eval 上的原始性能。即使 SFT 数据集包含 100% 聊天数据也是如此。我们在关于持续学习的讨论中进一步探讨这一点。

扫描内部文档与中期训练期间背景聊天数据的比例。尽管混入少量聊天数据有助于防止灾难性退化,但没有任何权重能维持原始的 IF-eval 性能。

对于任何给定的数据混合,我们观察到 IF-eval 性能在微调期间下降。这削弱了我们利用更长训练来进一步专精模型的能力。从趋势上看,我们可能期望一个在某些数据集上训练的过参数化模型只在该数据集的上下文中更新其行为,而不影响其在其他上下文中的表现。然而,在实践中我们并未观察到这一点,因为即使在 QA 上下文中,在原始文档数据上训练也会导致性能回退。

在所有数据混合下,IF-eval 在中期训练期间都会下降。当我们使用线性学习率(如图所示)时,性能下降最终趋于平缓,并随着学习率衰减而缓慢开始恢复。然而,性能从未完全恢复。

另一种常用方法是使用 LoRA 来约束参数更新,从而降低灾难性遗忘的可能性。然而,这种方法仍不足以保持 IF-eval,而且 LoRA 学到的更少。LoRA Learns Less and Forgets Less (Biderman et al, 2024)

当用于在已后训练的 Qwen3-8B 之上进行个性化中期训练时,LoRA 学到的(知识)更少,并且仍然会遗忘其原有的后训练行为。

同策略蒸馏可恢复后训练行为

接下来,我们寻求在内部文档上微调后恢复指令遵循行为。该行为最初是通过 RL 训练的,成本高昂,而且正如我们所见,十分脆弱。取而代之,我们以早期版本的模型 Qwen3-8B 作为教师,在 Tulu3 提示上运行同策略蒸馏。请注意,这一训练阶段与内部文档数据无关,仅用于恢复指令遵循能力。

使用早期版本的模型作为教师来“重新唤起”微调期间丢失的能力,使同策略蒸馏在持续学习方面非常有前景。我们可以在新数据上微调的阶段与恢复行为的蒸馏阶段之间交替,使我们的模型能够随时间学习并保持知识更新。这种阶段交替方法此前已由 Cobbe 等人探索过。Phasic Policy Gradient (Cobbe et al, 2020)

在按 70-30 比例混合内部文档数据和聊天数据进行微调后,同策略蒸馏在 IF-eval 上恢复了几乎全部性能,且没有丢失任何知识;我们还观察到聊天能力与模型在内部 QA 评估中的“知识”表现之间存在一些正向迁移。

模型内部 QA 评估(知识)IF-eval(聊天)
Qwen3-8B18%85%
+ 中期训练(100%)43%45%
+ 中期训练(70%)36%79%
+ 中期训练(70%)+ 蒸馏41%83%
中期训练后特定领域(内部 QA 评估)和聊天(IF-eval)的性能。尽管中期训练遗忘了 Qwen3-8B 的后训练行为,但通过同策略蒸馏可以低成本地恢复这些行为,同时还能获得通过中期训练学到的额外知识。

本质上,我们将语言模型本身视为奖励模型,高概率行为即获得奖励。Direct Preference Optimization: Your Language Model is Secretly a Reward Model(Rafailov 等,2023)这与逆强化学习有关:高概率行为对应于假设的底层偏好模型中的优势奖励。Algorithms for Inverse Reinforcement Learning(Ng 和 Russell,2000)任何经过指令微调的开放权重模型都可以在这个意义上用作奖励模型;我们只需要访问 compute_logprobs 函数。

蒸馏作为整合行为与知识的工具,也已在混合推理模型(Qwen3)和专家蒸馏中得到探索。DeepSeek-V3.2-Exp: Boosting Long-Context Efficiency with DeepSeek Sparse Attention(DeepSeek-AI Team,2025)正如我们和 Chen 等人的结果所表明的,on-policy 学习可以成为增强类似基于蒸馏的“模型合并”设置的关键工具。

讨论

稠密监督大幅提升计算效率

强化学习和 on-policy 蒸馏都通过反向 KL 学习,剪除基础策略中存在的动作空间。区别在于奖励的密度。在 LoRA Without Regret 中,我们提出了信息论视角:强化学习每个 episode 只教授 O(1)O(1) 比特。相比之下,蒸馏每个 episode 教授 O(N)O(N) 比特,其中 NN 是 token 数量。我们能否量化通过更稠密奖励获得的训练效率提升?

我们进行了一项实验,对两者进行直接比较:

  1. 从 Qwen3-8B-Base 开始(无额外 SFT)。
  2. 在 DeepMath 上运行 RL,匹配我们在 LoRA Without Regret 中的流程。我们使用 LoRA rank 为 128。所得模型作为蒸馏的教师。
  3. 将 RL 训练的模型(2)on-policy 蒸馏回基础模型(1)。
从相同的初始化开始,on-policy 蒸馏可以在大约少 7-10 倍的梯度步数内学到 RL 训练的策略,这对应于 50-100 倍的计算效率。

我们看到,在匹配模型架构(LoRA rank 128)的情况下,蒸馏达到教师性能水平的速度比 RL 快约 7-10 倍。反向 KL 降至接近零,AIME 分数在不到 10 个梯度步内恢复,而 RL 需要 70 步才能达到该水平。

累计来看,所需计算的减少约为 50-100 倍:

  • 虽然 RL 需要在大约评估上下文长度下训练(以便策略学习上下文限制且不产生格式惩罚),但蒸馏在较短的上下文长度下也能合理学习,因为在已完成采样的轨迹和继续采样的轨迹之间,奖励没有急剧的截断。
  • 当 SFT 初始化较强时,即:教师策略处于学生策略的支持范围内。当这一点不成立时,如“Distillation for reasoning”中那样,我们需要显著更大的批量大小。on-policy 蒸馏在批量大小小得多的情况下也能有效工作,因为它每个 episode 提供显著更多的比特,从而减少梯度噪声。

尽管通常难以用过程监督训练强化学习模型,但这些结果表明,作为一个广泛的方向,过程监督和稠密奖励有潜力将学习效率提升一个数量级。这与 Lightman 等人早期在 RL 研究中的结果一致。

蒸馏可以有效地复用训练数据以提高数据效率

对于从业者来说,收集大型训练提示数据集可能既困难又耗时。因此,我们希望在训练中能够多次复用提示。使用 RL 时,对同一提示训练多个 epoch 往往会导致对最终答案的简单记忆,尤其是在大模型上。“Reinforcement Learning for Reasoning in Large Language Models with One Training Example”(Wang 等,2025)展示了一个正面结果,然而,仅在某些设置下。相比之下,on-policy 蒸馏通过最小化反向 KL 来学习近似教师的完整分布,而不是记忆单一答案。这使我们能够从同一提示中训练许多样本。

我们重复上述在数学上训练 Qwen3-8B-Base 的实验,但现在仅使用数据集中随机选择的一个提示。提示:“计算极限:lim⁡x→∞x(x+13−x−13)limx→∞​x​(3x+1​−3x−1​)”

我们在这个提示上连续训练 20 步,每步使用 256 个 rollout 的批量,总共 5120 个评分序列。我们以顺序方式对同一提示进行多步训练,这通常会导致过拟合。尽管这自然在计算效率上较低,但我们仅在一个提示上训练,却大致达到了教师模型的性能。

在这个例子中,在单个训练样本上进行多轮训练,就足以蒸馏出教师的 AIME'24 表现。我们的默认配置(也用于个性化实验)运行 on-policy 蒸馏,每批 64 个提示,每个提示 4 个样本。所示的所有方法均以每批 256 个样本进行训练。请注意,右图显示的是训练 KL,因此总共 1 个提示优于每批 1 个提示是自然的。

RL 在语义策略空间中搜索

我们已经看到,on-policy 蒸馏可以用少得多的训练步数复现 RL 所提供的学习效果。对这一结果的一种解释是,与预训练不同,RL 不会在梯度步本身上花费大量计算。我们应当把 RL 的大部分计算视为花在搜索上——展开策略并分配信用——而不是花在更新上。摘自The Bitter Lesson(Rich Sutton):“突破性进展最终会通过一种相反的思路到来,即通过搜索和学习来扩展计算”

通过随机梯度下降进行的预训练是在探索高维参数空间。预训练需要海量信息,且很难蒸馏,部分原因是参数空间在某种程度上对每个网络都是独特的。”The Lottery Ticket Hypothesis: Finding Sparse, Trainable Neural Networks”(Frankle and Carbin, 2018) 预训练所需的梯度步在计算上极其昂贵且耗时。

相比之下,我们应当把 RL 视为在语义策略空间中探索。请注意,对策略的探索与对结果的探索有微妙的不同;RL 要求基础模型一开始就有非零的成功率,因此它已经“找到了结果”,但策略可以在 RL 过程中被改进,使成功结果更可能出现。 每一步,RL 都会尝试对它过去找到的某个策略做小幅修改。它不是在参数空间中探索,而是靠运气“偶然”发现新策略——它从自己已有的权重集合中随机采样。

一旦找到好的策略,蒸馏就成为学习它的捷径:on-policy 蒸馏不需要对 RL 课程期间的中间策略建模,而只需对最终学到的策略建模。如果我们只关心最终策略(在生产用例中很常见),就无需花费计算去建模所有中间策略。

考虑一个类比:在科学研究中,我们花费大量时间和资源寻找答案、探索新想法。一旦发现某个结果,用自然语言表达出来教给别人就简单得多。我们可以将其与打运动等直觉性身体技能相对比。这些技能更难教给别人,因为知识以一种只有我们自己才能轻易理解的天生语言(例如肌肉记忆)存在。运动只能通过反复练习来学会。

On-policy 学习作为持续学习的工具

在关于个性化蒸馏的章节中,我们探讨了 on-policy 蒸馏将专门训练的行为重新引入模型的能力。这可以推广到更广泛的持续学习任务,这些任务需要在不降低先前能力的情况下获取新知识。

先前的工作发现,on-policy 学习(RL)比 off-policy 学习遗忘更少。“RL's Razor: Why Online Reinforcement Learning Forgets Less”(Shenfeld 等,2025)然而,RL 只能塑造行为——它无法很好地教授新知识,因此不足以支撑持续学习。

在上一节中,我们看到 SFT(包括 off-policy 蒸馏)在搭建持续学习方面会失败,因为它会损害行为。我们对此进行了更细致的考察,并用一个直接例子加以说明。与上文类似,我们构建了一个数据集:取 Tulu3 的提示,从 Qwen3-32B 在 temperature = 1.0 处采样,不做进一步修改。因此该数据集相对于 Qwen3-32B 的 KL 为零。“真正 on-policy”的 KL=0 数据的重要性也在我们之前的文章 Defeating Nondeterminism in LLM Inference 中探讨过。

当我们在这个由模型自身样本构成的数据集上运行 SFT 时会发生什么?我们看到,任何大于零的实用学习率都会导致指令遵循评估上的性能下降!

在 Qwen3-32B 自身样本之上运行 SFT 会降低性能。我们使用与个性化部分相同的学习率,该学习率是为实际性能考量而调优的。线性学习率可以防止前向 KL / IF-eval 无限退化,但在学习率衰减到零之前无法恢复性能。

对此的一种可能解释是:虽然 KL 散度在期望上为 0,但在实践中每个有限批次都会呈现略有不同的分布。在这些有限批次上训练会产生非零的梯度更新,进而使更新后模型的策略偏离其原始状态的策略。这一过程随时间推移将“在自身样本上训练”变成了 off-policy 训练,从而导致与 off-policy 训练相同的误差累积,以及在长序列上的发散。

On-policy 蒸馏始终保持 on-policy,且由于教师保持固定,学生会收敛到教师的理想行为,而不会像 SFT 那样在自蒸馏设置中出现性能退化。这使得 on-policy 蒸馏成为持续学习的一个非常有前景的工具。

结论

我们探索了 on-policy 蒸馏在诸如训练小型模型进行数学推理或构建持续学习的助手等应用中的使用。我们将 on-policy 蒸馏与另外两种后训练方法进行了比较:off-policy 蒸馏和 on-policy RL。我们发现 on-policy 蒸馏结合了两者的优点:既有 on-policy 训练的可靠性能,又有稠密奖励信号的成本效率。

后训练是达到前沿模型能力的关键部分。通过利用来自学生的 on-policy 采样以及来自教师的稠密监督,on-policy 蒸馏方案能以远低于前沿高算力 RL 运行的成本达到这些能力。

我们的实现可以在 Tinker cookbook 中找到。我们的工作探索了 on-policy 蒸馏简单直接的实例化方式,以清晰地展示其优势。我们希望继续研究蒸馏的新应用、改进教师监督的新方法,以及提升数据效率和持续学习的方式。

在 Thinking Machines,我们的使命是让人们能够使用将前沿性能与适应性和个性化相结合的人工智能模型。同策略蒸馏是实现这一目标的有力工具。

引用

请将此工作引用为:

Lu, Kevin and Thinking Machines Lab, "On-Policy Distillation",
Thinking Machines Lab: Connectionism, Oct 2025.

或使用 BibTeX 引用:

@article{lu2025onpolicydistillation,
  author = {Kevin Lu and Thinking Machines Lab},
  title = {On-Policy Distillation},
  journal = {Thinking Machines Lab: Connectionism},
  year = {2025},
  note = {https://thinkingmachines.ai/blog/on-policy-distillation},
  doi = {10.64434/tml.20251026},
}

来源:Thinking Machines Lab:官方博客 · thinkingmachines.ai