Pivot-SD：面向掩码扩散语言模型的高效枢轴自蒸馏 Pivot-SD: Efficient Self-Distillation for Masked Diffusion Language Models
按信息增益挑出去噪中少数关键承诺，仅用200题离线自蒸馏即胜过全序列SFT与在线RL
前置知识
掩码扩散语言模型（Masked Diffusion LM, dLM）
以离散扩散方式生成文本的语言模型，代表如 LLaDA-8B-Instruct 与 Dream-7B。生成时从完全掩码的响应出发，每步模型对所有 $[\text{MASK}]$ 位置预测干净 token，采样器挑选部分位置写入，经 $T$ 步迭代得到完整回复。与自回归模型从左到右固定顺序生成不同，dLM 可以乱序填充位置，理论上支持并行解码，本文用它做复杂推理。
全文方法都建立在 dLM 特有的去噪轨迹之上：每个 token 在哪一步、哪个部分掩码状态中被提交。不理解 dLM 的生成机制就无法理解枢轴（pivot）这一核心概念。
置信度引导采样与去噪轨迹
实践中 dLM 常用 confidence-based sampler 解码：每步评估各掩码位置的预测置信度，优先填充最自信的位置。本文设置 block length 64、每步恰好提交 1 个 token，故 256 token 的回复对应 256 步去噪，形成轨迹 $\tau=\{(M_t, B_t, U_t)\}_{t=1}^{T}$，完整记录了每个 token 的提交顺序与当时的状态。
信息增益是对相邻两步的熵做差，而这两次前向传播是采样器本来就要做的，因此选枢轴不增加任何额外模型调用——这是方法近乎零开销的原因。
信用分配（credit assignment）与关键 token
后训练的核心问题是哪些 token 值得学：少数关键选择约束了后续大量内容，多数 token 由上下文基本确定。自回归侧已有工作用最终文本的统计量（如高熵少数 token、critical tokens）挑关键位置；但在 dLM 中，同一个 token 早提交能重塑大量未解码位置、晚提交几乎无影响，这种时序效应在最终文本中丢失。
Pivot-SD 的动机正是 dLM 的信用分配问题；消融显示选错监督位置（随机步）会让 MATH 从 37.47 跌到 29.93、低于未训练基线的 31.40。
Unlikelihood 训练
Welleck 等人（2019）提出的负向目标：对不希望出现的 token $y_p$，最小化 $\ell^-_p = -\log(1 - P_\theta(y_p \mid M_t))$，直接压低该 token 在该状态下的概率，而非只靠相对概率。相比对整条失败序列施加交叉熵惩罚，它可以精确作用于个别位置而不伤及邻近内容。
Pivot-SD 用 unlikelihood 只惩罚失败轨迹中被选中的枢轴 token，保留失败轨迹中正确的中间步骤；把它摊到整个状态（Pos+Neg All-token）会在四个基准上掉 2.5 到 5.1 分。
离线自蒸馏与拒绝采样（STaR 式）
由冻结基模型一次性采样轨迹，按结果（验证器成败）过滤或赋符号后做监督微调，训练期间不再生成新数据。与在线 RL 相比，无需训练中持续 rollout，不会出现零奖励导致的梯度停滞，超参数可离线调定，工程上简单稳定，代价是上限受基模型能力约束。
Pivot-SD 属于这一范式：800 条轨迹一次采完并缓存枢轴，训练仅 0.3 小时；而 5000 步 diffu-GRPO 需 22.9 小时，且附录 I 显示在线 RL 在 200 题预算下会完全停滞。
信息增益（平均熵下降）
对去噪步 $t$，比较提交前后剩余掩码位置的总熵 $H_{pre}(t)=\sum_{i\in \mathcal{A}_t\setminus\mathcal{U}_t} h(i\mid M_t)$ 与 $H_{post}(t)$，并除以剩余位置数归一化：$g(t)=\frac{H_{pre}(t)-H_{post}(t)}{|\mathcal{A}_t\setminus\mathcal{U}_t|}$，度量该步承诺平均使每个仍未解码位置的不确定性下降多少。
$g(t)$ 是选枢轴的唯一标准；除以剩余位置数是为了避免早期步因剩余位置多而天然占优。此前信息增益只被用于推理期选解码位置（Info-Gain Sampler、KLASS），本文首次用它挑训练目标。
研究动机
掩码扩散 LM 的后训练大多照搬自回归配方，产生三类具体问题。其一，标准掩码 SFT（式 3）对每个掩码位置等权交叉熵，琐碎的收尾填充与决定整条响应走向的关键承诺得到相同权重；轨迹级 RL（如 diffu-GRPO）则把结果奖励摊给整条轨迹。其二，失败轨迹中正确的中间推导与错误一起被惩罚，浪费了监督信号。其三，在线 RL 在小数据预算下信号稀疏且昂贵：附录 I 显示 wd1++ 在 200 题 AceCode 子集上整批样本单元测试全挂、组奖励与优势恒为零、正负两支梯度抵消，训练在 500 到 1000 步之间完全停滞；算力上 2000 题、5000 步的 diffu-GRPO 在双 L40 上需 22.9 小时。而 LLaDA-8B-Instruct 未训练时在 MATH500 上仅 31.40，改进空间与预算约束同时存在。
本文的目标是本文目标是给掩码扩散 LM 设计一种数据与算力都极小、且能稳定超过在线 RL 的后训练方法：只监督真正塑造了响应的少数关键去噪承诺（pivots），用任务验证器的成败决定每个更新的方向——成功轨迹的枢轴用交叉熵提升概率，失败轨迹的枢轴用 unlikelihood 压低概率，其余位置零损失。工程目标是把它做成纯离线自蒸馏：200 道题、每题 4 条轨迹（共 800 条）、双 L40 总共 2.8 小时，在 MATH500、GSM8K、HumanEval+、MBPP+ 四个基准上平均精度超过全序列 SFT 与预算匹配乃至 10 倍数据、5 倍步数的在线 RL，并迁移到第二个骨干 Dream-7B。
与已有工作不同的是，自回归模型的生成顺序由位置固定，只能事后从最终文本的统计量近似找关键 token；而 dLM 的置信度采样轨迹天然记录了哪个 token、在哪个位置、于哪一步被提交。本文的独特切入是利用这个 AR 模型不具备的信号：以每步承诺对剩余掩码位置造成的归一化熵下降（信息增益）为枢轴打分，因为同一个 token 在还剩很多掩码位置时提交与在最后一步提交的重要性截然不同，这一信息只存在于轨迹中。与把中间去噪步用于在线 RL 的 d1、d2、AGRPO、DiSPO 等不同，Pivot-SD 用该信号构造离线、局部化且带负目标（pivot 级 unlikelihood）的监督，既不需要过程奖励模型，也不需要步骤级人工标签。
核心方法
直觉是一条去噪轨迹的熵曲线只在少数几步陡降：模型写下某个 token 后，剩余所有掩码位置的不确定性骤减，说明这个承诺基本定下了响应的走向。技术路线分三步。第一步，从冻结基模型 $\theta_0$ 为每道训练题采 $G=4$ 条完整轨迹 $\tau=\{(M_t, B_t, U_t)\}_{t=1}^{T}$，$T=256$。第二步，对每步计算信息增益 $g(t)=\frac{H_{pre}(t)-H_{post}(t)}{|\mathcal{A}_t\setminus\mathcal{U}_t|}$，取 top-$K=10$ 步，每条轨迹得到 10 个枢轴三元组 $(t, p, y_p)$。第三步，用任务验证器（数学用精确答案匹配、代码用单元测试执行）给最终答案打标 $s(\tau)\in\{+1,-1\}$，把训练元组 $(M_t, y, p, y_p, s(\tau))$ 存入离线数据集 $D_{pivot}$。训练时重放枢轴被写下那一刻的掩码状态 $M_t$ 作为输入，只对枢轴 token 计算交叉熵或 unlikelihood 损失，其余位置不产生任何梯度。
核心创新有三层。第一，监督单位是已实现的去噪承诺 $(t,p,y_p)$ 而非文本推理步骤或最终序列位置：训练输入是枢轴被写下那一刻的部分掩码状态 $M_t$，让模型在完全相同的语境里重新预测该 token，这既不同于 SFT 的全序列监督，也不同于 RL 的轨迹级奖励。第二，定位信号是承诺时刻对剩余掩码位置的熵下降——GIFT 用的是最终序列位置的局部熵，Info-Gain Sampler 与 KLASS 只在推理期用信息增益选解码位置，都不在承诺时刻的状态上训练。第三，把轨迹成败转成局部化正负目标：成功枢轴做 CE $\ell^+_p=-\log P_\theta(y_p\mid M_t)$，失败枢轴做 pivot 级 unlikelihood $\ell^-_p=-\log(1-P_\theta(y_p\mid M_t))$，失败轨迹中正确的中间部分完全不碰。与已有方法的本质区别在于：SFT 等权一切 token、RL 奖励整条轨迹，Pivot-SD 只训练少数真正约束了后续内容的决定，并用离线缓存换取稳定性。
方法步骤详情
(1) 采样：每题用置信度采样器（块长 64、每步提交 1 个 token、温度 0.9）从冻结 $\theta_0$ 采 4 条 256 token 轨迹。(2) 构造候选承诺 $C(\tau)=\{(t,p,y_p):p\in U_t\}$，剔除块边界步与 EOS/轮末 token。(3) 逐步计算 $g(t)$：对 $\mathcal{A}_t\setminus\mathcal{U}_t$ 上每位置熵求和做差归一化，复用采样器已有前向，零额外开销。(4) 取 top-$K=10$ 步（K 取 5 至 20 稳定），把 $(M_t,y,p,y_p,s(\tau))$ 存入 $D_{pivot}$。(5) 训练：$\mathcal{L}_{pivot}=\mathbb{1}[s=+1]\ell^+_p+\lambda_{neg}\mathbb{1}[s=-1]\ell^-_p$，UL 概率剪裁防溢出；$\lambda_{neg}$ 按正负轨迹比设定：MATH 1、GSM8K 5、代码 0.02。(6) 从 $\theta_0$ 初始化，1000 步、batch size 8 微调。
技术新颖性
论文 Table 6 的信用分配版图清晰定位了新颖性。STaR 与拒绝采样只训成功轨迹的整条 rationale，失败样本被过滤；过程奖励模型（Lightman 2024、Math-Shepherd）需要人工或单独训练的过程标签；GIFT 用最终序列位置的局部熵做加权 SFT，丢失了提交时序；去噪过程奖励、ATPO 属于在线 RL 且没有 token 级负目标；AGRPO 在均匀或熵采样的每个去噪步上做轨迹级优势；DiSPO 有状态条件掩码填充但无缓存的已实现枢轴 UL 目标；Info-Gain Sampler 与 KLASS 把信息增益只用于推理期选位置，不参与训练；answer-aware 选择只盯末端答案块。Pivot-SD 是版图中唯一同时做到三件事的方法：只在选定的已实现承诺上训练、从失败轨迹构造局部化 pivot 级负目标、纯离线缓存——且全程不需要过程奖励模型或步骤级标签。
实验结果
主结果（Table 1，LLaDA-8B，256 token）Pivot-SD 四个基准全第一：MATH 37.47（超最强基线 SFT-GT 35.47 达 2.0 分）、GSM8K 79.51（超 5000 步 diffu-GRPO 77.63 达 1.9 分）、HumanEval+ 40.43（超 SFT-SD 35.78 达 4.7 分）、MBPP+ 47.12（超 SFT-GT 0.8 分），并胜过 2000 题、5000 步的扩展 RL。消融：CE/UL 摊到选定状态全部 token 掉 2.5 至 5.1 分；随机选步使 MATH 跌至 29.93、低于未训练 base 的 31.40；引入 UL 后随机步/随机 token 在 MATH 损失 7.5/2.6 分；K 增到 40 时 MATH 跌至 22.8。算力上 2.8 小时约 8 倍加速、FLOPs 省 1.65 倍、每轨迹仅监督 10 个 token（3.9%）。Dream-7B 迁移四项全升（MATH 42.20 对 37.97），OOD 微平均 48.61 最佳。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| 数学推理 MATH500 | 准确率（256 token，3 次重随机化均值） | 37.47 ±1.21 | 最强基线 SFT-GT 35.47 ±0.23；预算匹配 diffu-GRPO 33.47；扩展 RL（2000 题/5000 步）35.20 | +2.0 分 |
| 数学推理 GSM8K | 准确率（256 token） | 79.51 ±1.30 | 扩展 diffu-GRPO 77.63；预算匹配 wd1++ 76.93 | +1.9 分 |
| 代码生成 HumanEval+ | 准确率（256 token） | 40.43 ±2.75 | SFT-SD 35.78 ±1.54；预算匹配 wd1++ 33.54 | +4.7 分 |
| 代码生成 MBPP+ | 准确率（256 token） | 47.12 ±1.20 | SFT-GT 46.32 ±0.45 | +0.8 分（一个标准差内，计为平手） |
| 训练算力（2×NVIDIA L40 墙钟） | 总训练时长 | 2.8 小时（2.5 小时离线采样 + 0.3 小时训练） | 扩展 diffu-GRPO 22.9 小时；预算匹配 RL 3.6 至 5.4 小时 | 约 8 倍加速；总 FLOPs 省 1.65 倍（1875 对 3093 PFLOPs） |
局限与改进
作者承认的局限：若干超参数（$K$、$\lambda_{neg}$、学习率）按配置手工设定而非学习获得，跨域跨骨干需重新调整，例如代码域因分布移位必须把 $\lambda_{neg}$ 从 5 降到 0.02；信息增益公式只在单个去噪块内定义，长上下文推理需要跨块边界的步级信用分配扩展；实验只覆盖两个约 8B 的骨干，更大规模未验证。我的补充观察：其一，强依赖自动验证器（精确匹配与单元测试），开放任务要么无信号要么被奖励劫持；其二，离线自蒸馏上限受冻结基模型约束，按结果过滤无法引入基模型没有的能力，还可能放大其偏见与幻觉；其三，小基准统计功效不足——HumanEval+ 仅 164 题、一个问题值 0.61 分、标准差 2.75，MBPP+ 的 0.8 分优势在方差内；其四，HumanEval+ 在 512 token 时反而低于 base，长度迁移并不完全；其五，熵下降度量的可能是采样器置信度排序的副作用而非对结果的因果影响，$K$ 增大后选中无关步导致灾难性退化也提示这一准则并非稳健的重要性度量。
独立分析的弱点
第一，超参数脆弱：$K$ 固定为 10 时很好，但 $K=40/80$ 时 MATH 灾难性跌到 22.8/17.6；$\lambda_{neg}$ 跨域相差两个数量级（0.02 到 5），说明选多少、罚多重都对结果敏感。改进方向：按熵崩溃曲线自适应确定 $K$，或以轨迹正负比动态调 $\lambda_{neg}$。第二，无迭代自举：数据只来自冻结基模型的一次采样，没有 STaR 式的训后再采样再训练，能力上限受限。改进：用训后模型重采轨迹做第二轮自蒸馏。第三，依赖每步恰好提交一个 token 的采样器：在一步多 token 的块并行解码下，$g(t)$ 无法归属到单个 token。改进：按 token 级边际熵贡献分解。第四，验证器瓶颈：数学精确匹配与单元测试之外的任务（摘要、开放问答）没有可靠成败信号，pivot 方向无法确定。改进：接入 LLM 评审器并研究标签噪声下的鲁棒性。第五，熵下降不等价于因果重要性：随机步对照的高方差（HumanEval+ 标准差 5.51）与 $K$ 增大时的失效都说明该准则有盲区，可辅以反事实重放验证枢轴的真实影响。
未来方向
作者提出的方向：设计基于逐步熵崩溃的自适应规则，让枢轴数量与惩罚强度自动跟随任务复杂度，免去逐配置调参；把步条件信用分配扩展到跨块边界，以支持长上下文推理；在持续增大的 dLM 上检验 pivot 局部监督的可扩展性。基于本文成果可自然延伸的：其一，迭代式 Pivot-SD——用训后模型重采轨迹构建第二轮数据集，逼近 STaR 自举上限并缓解基模型能力封顶问题；其二，与在线 RL 组合，用缓存枢轴重放做暖启动，再放开有限探索，兼得离线稳定性与在线上限；其三，把局部化 unlikelihood 用于对齐——只惩罚诱发不良行为的枢轴 token 而保留其余正确内容，替代会误伤正常片段的全序列负训练；其四，与过程奖励模型结合，在无精确验证器的开放任务上用 PRM 分数替代 $\pm 1$ 结果标签；其五，理论层面刻画信息增益与结果因果影响的关系，解释 $K$ 增大时的失效机制并给出更稳健的选择准则。
复现评估
复现条件较好。论文承诺把生成的轨迹与评测脚本以 MIT 许可开源（项目页 sunwoohong.github.io/pivot-sd），训练与评测数据全部公开（GSM8K、MATH、HumanEval 为 MIT，MBPP 为 CC-BY-4.0），两个骨干 LLaDA-8B-Instruct 与 Dream-7B 均有公开权重。算力门槛低：全部实验在 2 张 NVIDIA L40 上完成，单次完整训练仅需 2.8 小时（离线采样 2.5 小时 + 训练 0.3 小时），方法本身实现简单——一次采样、复用采样器前向做信息增益打分（无额外调用）、标准 CE/UL 微调。附录 C 给出全部配置（块长 64、每步 1 token、训练采样温度 0.9、1000 步、batch size 8、$\lambda_{neg}$ 逐域设定规则），附录 D 规范了 3 次完全重随机化协议（重采 200 题、重采轨迹池、换训练种子）。主要成本在于严格对齐论文的对比实验需要复跑多个基线与随机种子，但核心方法一天内即可在单机复现，属于可复现性优秀的工作。
论文图表
一道 MATH 题目上去噪轨迹的预测熵热图（展示前 128 个 token）：x 轴为 token 位置，y 轴为自上而下推进的去噪步，颜色深浅表示该步各活跃掩码位置的预测熵。红色虚线与星标标出信息增益选中的关键去噪步及对应枢轴 token。可以清楚看到熵只在少数几步陡降、其余步几乎不变，而被选中的步恰好落在熵崖处。
这是枢轴这一核心概念存在的实证依据：不确定性在少数承诺处骤降的现象直接 motivating 了用信息增益挑训练目标的做法，也解释了为什么只监督约 4% 的 token 就足够。