Pivot-SD：高效自蒸馏用于掩码扩散语言模型
来自 arXiv · 期号 2026-10-03
摘要
扩散语言模型（dLM）可以并行生成文本，但训练时很难判断哪个去噪步骤真正决定了最终答案。 Pivot-SD 提出只监督那些对输出影响最大的“关键承诺”（pivot）：用信息增益指标找出每一步中不确定性下降最多的 token，只对这些 token 做训练。成功的轨迹用交叉熵，失败的轨迹用定向不似然，其余 token 不动。 只用 200 道题、每道题 4 次 rollout，就在数学和代码 benchmark 上超过了全序列 SFT 和扩散 RL 基线。 论文附了开源代码。
深读
导读 扩散语言模型（dLM）能并行生成文本，但训练时面临一个关键问题：如何判断众多去噪步骤中，哪一步真正决定了最终答案。这篇论文提出 Pivot-SD，一种只监督那些对输出影响最大的“关键承诺”（pivot）的自蒸馏框架，用极少的训练数据就在数学和代码任务上超越了全序列监督微调（SFT）和扩散强化学习基线。对于关注高效训练方法和扩散模型后训练策略的研究者，这是一篇值得一读的简洁工作。 关键信息 - 扩散语言模型（dLM）在去噪过程中，少数几个“承诺”（commitments）会大幅降低剩余掩码位置的不确定性，从而决定最终回答的大部分内容。 - Pivot-SD 使用信息增益指标（衡量剩余掩码位置的不确定性下降程度）来选出这些关键承诺（pivot）。 - 对成功轨迹中的 pivot 使用交叉熵损失训练，对失败轨迹中的 pivot 使用定向不似然损失（targeted unlikelihood），其余 token 不参与训练。 - 仅用 200 道问题和每道题 4 次 rollout（共 800 次生成），Pivot-SD 就在数学和代码 benchmark 上超过了全序列 SFT 和预算匹配的扩散 RL 基线。 - 实验基于 LLaDA-8B-Instruct 模型。 逐节讲解 问题：扩散语言模型的后训练困境 掩码扩散语言模型（dLM）能并行生成文本，在复杂推理任务上是有潜力的自回归模型替代方案。但它在训练时面临一个独特的信用分配难题：去噪过程中，少数几个“承诺”（即对某些位置的 token 做出确定性的选择）会急剧降低剩余掩码位置的不确定性，从而塑造了整个回答。然而，现有的后训练方法（如对最终文本做 SFT，或对整个去噪步骤分配奖励）都没有利用这一信号来决定哪些 token 应该被训练，导致训练效率低下。 方法：Pivot-SD 框架 Pivot-SD 是一个高效的离线自蒸馏框架，核心思想是只监督那些高影响力的承诺（称为 pivot）。具体做法是：首先，用信息增益指标来衡量每一步去噪中，当前承诺使剩余掩码位置的不确定性下降了多少，据此选出 pivot。然后，对成功轨迹（最终答案正确）中的 pivot 使用交叉熵损失进行训练；对失败轨迹中的 pivot 使用定向不似然损失（即降低该 token 被生成的概率），而失败轨迹中其余非 pivot 的 token 则完全不动，不参与训练。 结果与结论 在 LLaDA-8B-Instruct 模型上的实验表明，Pivot-SD 仅使用 200 道训练问题和每道问题 4 次 rollout（共 800 次生成），就在数学和代码 benchmark 上超过了全序列 SFT 和预算匹配的扩散 RL 基线。论文附有开源代码。 值得注意 原文为论文摘要，信息有限。方法的具体实现细节（如信息增益的计算方式、pivot 的选择阈值、…