Pivot-SD: Efficient Self-Distillation for Masked Diffusion Language Models
Seo Hyun Kim, Sunwoo Hong, Younwoo Choi, Chen-Hao Chao, Se-Young Yun, Rahul G. Krishnan
EMNLP 2026 Main (Oral)
cs.LG, cs.CL
2026-10-03
Pivot-SD按信息增益只监督每条轨迹十个关键承诺，做对的用交叉熵、做错的只做unlikelihood，二百题就让LLaDA-8B在数学和代码上超过全序列SFT与十倍数据的扩散RL。
Masked diffusion language model（dLM，遮罩扩散语言模型）不从左往右写。回答从全 [MASK] 开始，模型按置信度决定先填哪一格。某个位置被写成具体 token，就是一次 commitment（承诺）。
少数承诺会把剩余遮罩位置的熵打下去，后面的内容大体被这几步定住。图 1 里，基座解 MATH 题的熵热力图多数步几乎不动，断崖只出现在少数几步，这些步上的 token 就是 pivot。
后训练很少用这个断崖来选 token。借自回归的 SFT 在最终文本上均匀加权，先写下的关键 token 和最后补的标点一样重。失败轨迹里，做对的中间式子会和错误一起被罚。在线 RL 把奖励摊到整个去噪步，不单独处理改写了后文分布的那一次承诺。自回归顺序由位置固定，在最终文本上统计不会丢掉顺序。dLM 的填写顺序是采样当场决定的，最终文本不记录写下时还剩多少遮罩。
Pivot-SD 做离线自蒸馏。轨迹从冻结基座采一次，之后只回放被选中的承诺。
每道题采 4 条。采样器每步写一个 token。一步的信息增益，是它对剩余遮罩位置的平均预测熵的下降，再除以这些位置的个数。不归一化时，早步遮罩更多，熵差天然偏大。遮罩集合发生变化的步被丢掉，避免前后比的不是同一组位置。分数只决定留哪些步，不作为 loss 权重。这两次熵来自采样器本就要做的前向，挑选不再额外调用模型。每条轨迹留信息增益最高的 10 步。K 在实验前定成 10，附录里 5 到 20 差别不大。
256 token 的回答因此只监督 10 个 token，不到响应预算的 4%。训练输入是该 pivot 写下去之前的部分遮罩状态，模型要在那个状态下重预测同一个 token。
对错来自验证器。数学精确匹配，代码过单元测试。做对的轨迹，pivot 用交叉熵，提高该 token 概率。做错的轨迹，只在 pivot 上做 unlikelihood，损失是 -log(1-p)，把该 token 概率压低，其余位置不加损失。失败样本里常有写对的代数，整段回推会把它们一起压掉。概率接近 1 时会裁剪。数据采完即固定，超参可以离线调，也避开了小预算在线 RL 里整批奖励为零的情况。
主骨干是 LLaDA-8B-Instruct。GSM8K、MATH、AceCode 各 200 题，每题 4 条轨迹，各 800 条，优化 1000 步。评测是 GSM8K、MATH500、HumanEval+、MBPP+。受控预算的方法在 256 token 上重随机三次取平均。SFT-GT 在真值解答上做随机遮罩。SFT-SD 只用同一批做对的 rollout，做普通 masked SFT。diffu-GRPO 和 wd1++ 的对齐版同样是 200 题、1000 步，加长版是 2000 题、5000 步。
| 方法 | MATH | GSM8K | HumanEval+ | MBPP+ |
| 基座 | 31.40 | 75.13 | 32.32 | 43.12 |
| SFT-GT | 35.47 | 76.94 | 33.94 | 46.32 |
| SFT-SD | 31.73 | 73.90 | 35.78 | 45.92 |
| diffu-GRPO 预算对齐 | 33.47 | 76.04 | 32.52 | 43.39 |
| diffu-GRPO 2000题/5000步 | 35.20 | 77.63 | 34.15 | 42.59 |
| wd1++ 预算对齐 | 33.27 | 76.93 | 33.54 | 43.12 |
| wd1++ 2000题/5000步 | 35.40 | 76.65 | 34.15 | 44.97 |
| Pivot-SD | 37.47 | 79.51 | 40.43 | 47.12 |
相对各榜最强对照，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 的 46.32 只高 0.8，论文写明这个差距在一个标准差以内。SFT-SD 看到的是同一批 rollout，空档来自挑选和负向监督。
512 token 上，MATH 42.6、GSM8K 84.38、MBPP+ 48.15 仍然最高。HumanEval+ 为 41.46，低于基座和预算对齐 diffu-GRPO，两边都是 42.07。
正负损失都用时，只监督 pivot 比监督该状态里全部遮罩 token，四榜高 2.5 到 5.1。只保留正损失时，两种范围没有稳定胜负。同一套正负损失放到随机步上，MATH 跌到 29.93，低于基座 31.40，也比 Pivot-SD 低 7.5。信息增益选中的状态里若换成随机 token，MATH 低 2.6。只做正样本时，局部熵挑选和信息增益挑选的差距不超过 2.3。
两张 GPU 上，Pivot-SD 采样 2.5 小时、训练 0.3 小时，合计 2.8 小时。预算对齐的 diffu-GRPO 要 5.4 小时，wd1++ 要 3.6 小时。5000 步分别要 22.9 小时和 17.3 小时。平均精度比这些 5000 步 RL 高出 3 分以上，用时大约是 5000 步 diffu-GRPO 的八分之一。按 FLOPs 计，比预算对齐的 RL 少大约 40%。题量从 50 增到 2000 时，HumanEval+ 上一直最好，2000 题时四榜都最好。
Dream-v0-Instruct-7B 用同一套配方。
| 方法 | MATH | GSM8K | HumanEval+ | MBPP+ |
| Dream 基座 | 37.97 | 79.55 | 54.27 | 60.58 |
| SFT-GT | 35.40 | 61.49 | 54.27 | 59.26 |
| SFT-SD | 40.80 | 70.74 | 48.78 | 61.38 |
| Pivot-SD | 42.20 | 81.88 | 58.54 | 61.64 |
四项都高于基座和两条 SFT。SFT-GT 把 GSM8K 从 79.55 打到 61.49。
域外按源-目标对做微平均，三列的榜单组成并不相同。Pivot-SD 的总平均是 48.61，次高 SFT-GT 是 47.83。GSM8K 上训练的模型里它排第一（39.82）。AceCode 上训练的里排第二（55.42，预算对齐 diffu-GRPO 为 55.49）。MATH 上训练的里排第三（52.87，SFT-GT 为 53.57，加长 diffu-GRPO 为 53.02）。
数学和代码有自动验证器，masked dLM 的后训练可以先不做在线采样。200 题、每题 4 条、每条 10 个 token，2.8 小时里大头是那一次 2.5 小时的采样，更新只占 0.3 小时。
模型结构没有换，梯度改落在单次填写上，失败样本也只罚这一次。条件很具体：采样过程要留下步、位置和 token，结局要能自动判对错。开放生成没有验证器，正负号没有来处。过程奖励模型得先把推导切成步，这里不需要那层标注，定位用的是模型自己的熵降。
论文自己列了三条。超参按配置手调，换领域、换骨干还要调。信用分配停在单个去噪块内，长上下文得把信号传到块外。实验只有 LLaDA-8B 和 Dream-7B，规模接近。
表里的数字把四榜全胜收窄了。MBPP+ 的 0.8 分在一个标准差内。512 token 的 HumanEval+ 低于基座的 42.07，主结论绑在与训练同长的解码上。Dream 上 SFT-GT 把 GSM8K 从 79.55 打到 61.49，相对全序列 SFT 的优势里有一部分来自对照自己掉点。对 Dream 基座，MBPP+ 只从 60.58 到 61.64。
信息增益没有和人工过程标注对过，pivot 是不是人会标成关键推导，没有证据。主表没有纳入 GIFT。GIFT 在最终序列上按熵给 token 加权。熵挑选只出现在纯正样本的消融里。轨迹只从冻结模型采一次，训完之后新轨迹里的承诺会不会挪地方，没有再测。任务全都有硬验证器。没有自动对错时，unlikelihood 会压低什么，这篇没有回答。