Pivot-SD: 面向掩码扩散语言模型的高效自蒸馏方法

摘要

掩码扩散语言模型（dLMs）为复杂推理提供了可并行的自回归替代方案，但存在信用分配难题：去噪过程中的少数几步承诺会大幅削减其余掩码位置的不确定性，并主导整个回答。现有后训练方案多在最终文本上训练，或把奖励归到整个去噪步骤，而不会挑选具体承诺。Pivot-SD 是一种高效离线自蒸馏框架，用信息增益指标衡量对剩余掩码位置的不确定性削减，只监督这些高影响的“枢轴”（pivots）：成功轨迹的枢轴用交叉熵训练，失败轨迹的枢轴用针对性的非似然训练，失败轨迹其余部分保持不动。仅用 200 个问题、每个 4 次采样，Pivot-SD 在数学与代码基准上均优于全序列 SFT 和预算相当的扩散 RL 基线，提升了 LLaDA-8B-Instruct。