跳到正文
Thinking Machines Lab:官方博客· Thinking Machines Lab·· 2025-10-27精选AI 评分68

Thinking Machines Lab 提出 On-Policy Distillation 训练方法

On-Policy Distillation

AI 导读

Thinking Machines Lab 发布技术博客,提出 on-policy distillation:从学生模型采样轨迹,由教师模型对每个 token 计算反向 KL 作为密集奖励。

推荐理由

原文给出反向 KL 损失、成本对比和可复现实验,读者可以据此评估用稀疏 RL 换成密集蒸馏监督的落地空间。

正文 · AI 翻译

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

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

训练更充分的小模型,往往在其专精领域胜过更大的通用模型。用小模型有很多好处:出于隐私或安全考虑可以本地部署,能持续训练、更容易更新,还能省下推理成本。要享受这些好处,就得为后训练阶段选对方法。

对“学生”模型做后训练,方法可分为两类:

  • 同策略训练从学生模型自身采样轨迹,再给这些轨迹打分作为奖励。
  • 异策略训练依赖外部来源给出的目标输出,让学生模型学着模仿。

举个例子,我们可能想训练一个小模型来解这样的数学题:

我们可以用强化学习做同策略训练,对学生的每条轨迹按是否解出题目来打分。打分可以由人来完成,也可以交给一个能稳定得出正确答案的“教师”模型。

同策略训练的优势在于,学生用自己产生的样本训练,能更直接地学会避开错误。但强化学习有个大缺点:反馈非常稀疏,每个训练回合只传递固定数量的比特信息,与用了多少 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)

如果你在学下棋,在策略强化学习就像在没有教练指导的情况下对弈。赢棋或输棋的反馈直接与你的走法挂钩,但每局只能收到一次,而且不会告诉你哪些走法对结果影响最大。离策略蒸馏则像是观看大师下棋——你看到的是极强的棋招,但它们所处的棋盘局面是新手很少会遇到的。

我们希望把强化学习的在策略相关性与蒸馏的稠密奖励信号结合起来。就学习下棋而言,这相当于有一位老师给你自己的每一步棋打分,从“严重失误”到“妙手”。对于大语言模型的后训练来说,这就是在策略蒸馏。

截图来自 chess.com。每一步棋都由分析引擎按颜色分级,将走法标记为严重失误(红色)、失误(橙色)、不精确(黄色)或妙手(蓝色)。

在策略蒸馏——两全其美

在策略蒸馏的核心思路是:从学生模型采样轨迹,再用一个高性能教师模型为每条轨迹的每个 token 打分。回到上面数学的例子,在策略蒸馏会给解题的每一步打分,惩罚那些导致学生模型得出错误答案的错误,同时强化那些执行正确的步骤。

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

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

我们的同策略蒸馏工作受到 DAGGERA 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 通常优化的正是由奖励模型诱导出的一种序列级反向 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 一起发给教师客户端,它返回教师模型在这些轨迹上的 logprobs log⁡πteacher(x)logπteacher​(x),对应的是学生模型采样出的 token xx。本文所有实验都不考虑 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),不过在数据量很大时,LoRA 的效果不如全参数微调。无论哪种方式,性能都呈对数线性增长——前期涨分很便宜,越往后越贵。

离策略蒸馏(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 在在线策略蒸馏过程中的表现。我们以训练 FLOPs 来衡量额外算力(见下文)。在线策略蒸馏的算力效率显著高于 SFT,对 LoRA 模型尤其如此。在 rank = 32 时,LoRA 在 SFT 后比全量微调落后 13%,但在在线策略蒸馏后只落后 6%。

跨方法比较算力成本并不简单,因为训练、采样与对数概率计算三者的成本比例会随实现方式大幅变化。下面我们以 FLOPs 计算成本,这种口径对能在 GPU 上高效并行的方法不利。具体来说,它会高估计算对数概率的实际成本。

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

我们发现,在 SFT 数据集已给定(如我们使用 OpenThoughts-3 的例子)或在多次训练运行间摊销的情况下,基线成本可降低 9 倍。CE =(学生 + 教师)/(学生) 这种情况下,离策略训练不计教师 FLOPs 成本,但在线策略训练要计,因为必须运行教师模型来计算学生轨迹的对数概率。由于这部分计算可以在 GPU 上低成本并行,按 GPU 小时计的成本降幅接近 18 倍。

不过,我们常常需要针对新任务训练一个小模型,而该任务没有现成的离策略蒸馏数据集。如果把离策略蒸馏中教师模型的全部成本都算进来——即包括从教师模型采样的额外成本——总成本降幅约为 30 倍。CE =(学生 + 教师)/(学生 + 教师)

面向个性化的蒸馏

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

同时训练两者通常很困难,轻量微调往往不足以实现这一目标,Unfamiliar Finetuning Examples Control How Language Models Hallucinate(Kang 等,2024)因此需要更大规模的中期训练。在新知识之上学习后训练行为,需要一套复杂的后训练流程,通常包含专有数据和奖励模型。前沿实验室固然能用这套方法,但其他从业者要复现可能很困难,或成本高得难以承受。

本节我们展示在线策略蒸馏可以有效地用于后训练专门行为。这一方法同样适用于持续学习或“测试时训练”:在模型部署后更新模型,同时不让基础性能回退。我们用一个在公司内部文档上做过中期训练的模型作为示例应用。

训练内部助手

定制模型的一个常见目标是充当助手:既掌握某个领域的专家知识,又具备可靠的助手式行为。我们可能需要分别训练这两者,尤其是当专业知识无法仅从预训练数据中学到,或学习它会干扰行为时。

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

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

在新知识上训练会损害已学到的行为

我们从 Qwen3-8B 开始,而不是基座模型。Qwen3-8B 已经过后训练,具备助手所需的有用技能,例如指令遵循和基于强化学习的推理。此前的研究表明,这类强化学习只训练了原模型中的小子网络,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 的表现会在微调过程中下降。这让我们无法通过延长训练来进一步让模型专精化。从趋势上看,我们或许会预期:一个过参数化的模型在某个数据集上训练后,只会更新其在该数据集语境下的行为,而不影响它在其他语境下的表现。但实际并非如此:用原始文档数据训练,即使在问答语境下性能也会回退。

在所有数据配比下,IF-eval 都会在中期训练期间下降。使用线性学习率时(如图所示),下降最终会趋于平缓,并随着学习率衰减而缓慢开始回升。但性能始终无法完全恢复。

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

在已完成后训练的 Qwen3-8B 之上做个性化中期训练时,LoRA 学到的(知识)更少,却依然会遗忘原有的后训练行为。

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

接下来,我们希望在用内部文档微调之后恢复指令遵循行为。这一行为原本是通过强化学习训练出来的,成本高昂,而且如前所述十分脆弱。于是我们改用同策略蒸馏,以早期版本的模型 Qwen3-8B 作为教师,在 Tulu3 提示词上训练。注意这一训练阶段与内部文档数据无关,唯一目的是恢复指令遵循能力。

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

在用内部文档数据与对话数据按 70 比 30 配比微调之后,同策略蒸馏几乎完全恢复了 IF-eval 的性能,且没有丢失任何知识;我们还观察到,对话能力与模型在内部问答评测中的“知识”表现之间存在一定的正向迁移。

模型内部问答评测(知识)IF-eval(对话)
Qwen3-8B18%85%
+ 中期训练(100%)43%45%
+ 中期训练(70%)36%79%
+ 中期训练(70%)+ 蒸馏41%83%
中期训练后模型在领域特定任务(内部问答评测)与对话任务(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 团队,2025)正如我们和 Chen 等人的结果所示,在策略学习可以成为增强此类基于蒸馏的“模型合并”方案的关键工具。

讨论

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

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

我们做了一项实验,直接对比两者:

  1. 从 Qwen3-8B-Base 开始(不做额外 SFT)。
  2. 在 DeepMath 上运行 RL,流程与 LoRA Without Regret 一致。我们使用 LoRA rank 128。得到的模型作为蒸馏的教师。
  3. 从 RL 训练出的模型(2)在策略蒸馏回基础模型(1)。
从相同初始化出发,在策略蒸馏学习 RL 训练出的策略所需的梯度步数约为原来的 1/7 到 1/10,对应 50-100 倍的计算效率。

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

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

  • RL 需要在接近评估上下文长度下训练(以便策略学会上下文限制、避免格式惩罚),而蒸馏在较短上下文长度下也能正常学习,因为已结束采样的轨迹与仍在继续的轨迹之间不存在奖励的急剧截断。
  • 当 SFT 初始化较强时,即教师策略处于学生策略的支撑集内。若这一条件不成立,就像“Distillation for reasoning”中那样,我们就需要显著更大的批大小。在线策略蒸馏在批大小小得多的情况下也能有效工作,因为每条轨迹能提供多得多的信息量,从而降低梯度噪声。

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

蒸馏能有效复用训练数据,提升数据效率

对从业者来说,收集大规模训练提示词数据集既困难又耗时。因此,我们希望能在训练中多次复用提示词。用强化学习时,对同一提示词训练多个轮次往往会导致模型简单记住最终答案,大模型尤其如此。“Reinforcement Learning for Reasoning in Large Language Models with One Training Example”(Wang 等,2025)在某些设置下给出了正面结果。相比之下,在线策略蒸馏通过最小化反向 KL 散度来学习逼近教师的完整分布,而不是记住单一答案。这让我们能对同一提示词训练许多样本。

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

我们在这个提示词上连续训练 20 步,每步一批 256 条 rollout,总共 5120 条带评分的序列。我们对同一提示词按顺序训练多步,这通常会导致过拟合。虽然这样在计算效率上自然较低,但我们仅用单个提示词训练,性能却大致追平了教师模型。

在这个例子中,在单个训练样本上做多轮训练,就足以蒸馏出教师的 AIME'24 表现。我们的默认配置(个性化实验中也使用)采用同策略蒸馏,每批 64 条提示词、每条提示词 4 个样本。图中所有方法均按每批 256 个样本训练。注意右图展示的是训练 KL,因此总共 1 条提示词优于每批 1 条提示词是自然的。

强化学习在语义策略空间中搜索

我们已经看到,同策略蒸馏能用少得多的训练步数复现强化学习带来的学习效果。对这一结果的一种解读是:与预训练不同,强化学习并不会把大量算力花在梯度步本身。我们应当把强化学习理解为把大部分算力花在搜索上——展开策略并分配信用——而不是花在更新上。摘自《苦涩的教训》(Rich Sutton):“突破性进展最终来自一种相反的方法,即通过搜索和学习来扩展计算。”

通过随机梯度下降进行的预训练是在探索高维参数空间。预训练需要海量信息,而且很难蒸馏,部分原因是参数空间对每个网络而言都多少有些独特。《彩票假说:寻找稀疏、可训练的神经网络》(Frankle 和 Carbin,2018) 预训练所需的梯度步在计算上极其昂贵,也非常耗时。

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

一旦找到好的策略,蒸馏就成了学习它的捷径:同策略蒸馏不需要建模强化学习课程中的中间策略,只需要建模最终学到的策略。如果我们只关心最终策略(生产用例中很常见),就不必花费算力去建模所有中间策略。

打个比方:在科学研究中,我们花费大量时间和资源寻找答案、探索新想法。一旦发现某个结果,用自然语言把它表达出来教给别人就简单得多。我们可以把它与打球这类直觉性的身体技能作对比。这类技能要教给别人难得多,因为知识以只可自身体会的内在语言(比如肌肉记忆)存在。运动只能靠反复练习学会。

把同策略学习用作持续学习的工具

在面向个性化的蒸馏一节中,我们探讨了同策略蒸馏把专门训练出的行为重新引入模型的能力。这可以推广到更广泛的持续学习任务,这些任务要求在不降低既有能力的前提下获取新知识。

此前的研究发现,同策略学习(RL)比离策略学习遗忘得更少。“RL的剃刀:为什么在线强化学习遗忘更少”(Shenfeld 等,2025)然而,RL 只能塑造行为——它无法很好地传授新知识,因此不足以支撑持续学习。

上一节我们看到,SFT(包括离策略蒸馏)在搭建持续学习基础时会失败,因为它会让行为退化。我们对此做了更细致的研究,并用一个直接例子加以说明。与上文类似,我们用 Tulu3 的提示词从 Qwen3-32B 在 temperature = 1.0 下采样构建数据集,不做任何进一步修改。因此该数据集相对 Qwen3-32B 的 KL 为零。“真正同策略”的 KL=0 数据的重要性,我们在此前的文章《击败 LLM 推理中的不确定性》中也探讨过。

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

在 Qwen3-32B 自身样本上跑 SFT 会导致性能退化。我们使用与个性化部分相同的学习率,该学习率是出于实际性能考虑调出来的。线性学习率可以防止前向 KL / IF-eval 无限恶化,但在学习率衰减到零之前无法恢复性能。

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

同策略蒸馏始终保持同策略,而且由于教师模型固定不变,学生模型会收敛到教师模型的理想行为,在自蒸馏场景下不会像 SFT 那样出现性能回退。这让同策略蒸馏成为持续学习的一个非常有前景的工具。

结论

我们探索了同策略蒸馏在训练小型数学推理模型或持续学习助手等场景中的应用。我们将同策略蒸馏与另外两种后训练方法做了比较:离策略蒸馏和同策略 RL。我们发现同策略蒸馏兼具两者之长:既有同策略训练的可靠性能,又有稠密奖励信号的成本效率。

后训练是达到前沿模型能力的关键一环。借助从学生模型进行同策略采样、并由教师模型提供稠密监督,同策略蒸馏方案能以远低于前沿高算力 RL 训练的成本达到这些能力。

我们的实现见 Tinker cookbook。我们的工作探索了同策略蒸馏简单直接的实现方式,以清晰展示其优势。我们希望继续研究蒸馏的新应用、改进教师监督的新方法,以及提升数据效率和持续学习的方式。

在 Thinking Machines,我们的使命是让人们用上兼具前沿性能、适应性与个性化的 AI 模型。同策略蒸馏正是实现这一目标的有力工具。

引用

引用本文时请使用:

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