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
With 200 questions, Pivot-SD trains 10 pivots per trajectory and lifts LLaDA-8B to 37.47 MATH and 40.43 HumanEval+, above full-sequence SFT and 10×-data diffusion RL.
A masked diffusion language model (dLM) does not write left to right. The response starts fully masked, and the sampler fills whichever position the model is confident about next. Writing a concrete token into a position is a commitment.
A few commitments cut the entropy of the positions that are still masked, and much of the rest of the answer follows from them. Figure 1 is an entropy heatmap of the base model on a MATH problem. Most denoising steps barely move the remaining uncertainty. The cliffs sit on a handful of steps. The tokens written on those steps are the pivots.
Post-training rarely uses that cliff to choose tokens. Supervised fine-tuning borrowed from autoregressive models weights the finished text uniformly, so a token that locks the answer early counts the same as a comma filled in at the end. On a failed trajectory, correct intermediate algebra is penalized along with the actual mistake. Online RL spreads the reward across a whole denoising step and does not isolate the commitment that reshaped the distribution over the masks that were left. Autoregressive order is fixed by position, so statistics on the finished text still see that order. A dLM chooses its order while sampling, and the finished text does not record how many masks remained when a token was written.
Pivot-SD is offline self-distillation. Trajectories are sampled once from a frozen base model. Training only replays the selected commitments.
Each question gets four trajectories. The sampler writes one token per step. Information gain for a step is the drop in average predictive entropy over positions that stay masked, divided by the number of those positions. Without that division, early steps look stronger simply because more masks remain. Steps that change the active mask set are dropped, so the before and after entropies cover the same positions. The score only selects steps. It does not weight the loss. Both entropy reads come from forwards the sampler already runs, so selection adds no model calls. Each trajectory keeps the 10 highest-gain steps. K was set to 10 before the experiments. The appendix reports little change for K from 5 to 20.
A 256-token response therefore supervises 10 tokens, under 4% of the budget. The training input is the partially masked state from just before that pivot was written. The model has to predict the same token in that state.
A verifier sets the sign. Math uses exact match. Code uses unit tests. Pivots from correct trajectories get cross-entropy, which raises the token probability. Pivots from incorrect trajectories get unlikelihood, -log(1-p), which pushes that probability down. No other position in the failed trajectory enters the loss. Failed samples often still contain correct algebra, and a full-sequence penalty would push that algebra down too. Probabilities are clipped away from 1. The dataset is frozen after sampling, so hyperparameters can be tuned offline, and training avoids the all-zero reward batches that online RL hit at this budget.
The main backbone is LLaDA-8B-Instruct. GSM8K, MATH, and AceCode each supply 200 training questions, four trajectories each, 800 per dataset, and 1,000 optimizer steps. Evaluation is GSM8K, MATH500, HumanEval+, and MBPP+. Methods on the controlled budget are fully re-randomized three times at 256 tokens. SFT-GT fits ground-truth solutions with random masking. SFT-SD runs ordinary masked SFT on the correct rollouts from the same pool. Budget-matched diffu-GRPO and wd1++ use the same 200 questions and 1,000 steps. The extended runs use 2,000 questions and 5,000 steps.
| Method | MATH | GSM8K | HumanEval+ | MBPP+ |
| Base | 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, budget-matched | 33.47 | 76.04 | 32.52 | 43.39 |
| diffu-GRPO, 2000 q / 5000 steps | 35.20 | 77.63 | 34.15 | 42.59 |
| wd1++, budget-matched | 33.27 | 76.93 | 33.54 | 43.12 |
| wd1++, 2000 q / 5000 steps | 35.40 | 76.65 | 34.15 | 44.97 |
| Pivot-SD | 37.47 | 79.51 | 40.43 | 47.12 |
Against the strongest baseline on each benchmark, MATH rises from 35.47 (SFT-GT) to 37.47, a gain of 2.0. GSM8K rises from 77.63 (5,000-step diffu-GRPO) to 79.51, a gain of 1.9. HumanEval+ rises from 35.78 (SFT-SD) to 40.43, a gain of 4.7. MBPP+ rises from 46.32 (SFT-GT) to 47.12, a gain of 0.8 that the paper places inside one standard deviation. SFT-SD saw the same rollouts, so the gap is pivot selection plus the negative update.
At 512 tokens, Pivot-SD still leads MATH (42.6), GSM8K (84.38), and MBPP+ (48.15). HumanEval+ falls to 41.46, behind the base model and budget-matched diffu-GRPO, both at 42.07.
With both signs of update, supervising only the pivot beats supervising every masked token in the selected state on all four benchmarks, by 2.5 to 5.1 points. With positive updates alone, the two scopes have no stable winner. The same positive-and-negative loss on a random step drops MATH to 29.93, below the base score of 31.40 and 7.5 behind Pivot-SD. Replacing the information-gain token with a random token inside an information-gain state costs 2.6 MATH points. In the positive-only setting, local-entropy selection and information-gain selection stay within 2.3 points on every benchmark.
On two GPUs, Pivot-SD spends 2.5 hours sampling and 0.3 hours training, 2.8 hours total. Budget-matched diffu-GRPO takes 5.4 hours and wd1++ takes 3.6. The 5,000-step runs take 22.9 and 17.3 hours. Average accuracy is more than 3 points above those 5,000-step RL runs, at about one eighth of the wall-clock of 5,000-step diffu-GRPO. In FLOPs, Pivot-SD uses about 40% less compute than budget-matched RL. From 50 to 2,000 training questions it stays ahead on HumanEval+, and at 2,000 questions it leads all four benchmarks.
The same recipe on Dream-v0-Instruct-7B:
| Method | MATH | GSM8K | HumanEval+ | MBPP+ |
| Dream base | 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 |
Pivot-SD is above the Dream base and both SFT runs on all four. SFT-GT drops GSM8K from 79.55 to 61.49.
Out of domain, the micro-average is taken over source-target pairs, and the three training columns do not cover the same benchmarks. Pivot-SD scores 48.61, ahead of SFT-GT at 47.83. It ranks first among GSM8K-trained models (39.82), second among AceCode-trained models (55.42 versus 55.49 for budget-matched diffu-GRPO), and third among MATH-trained models (52.87 versus 53.57 for SFT-GT and 53.02 for extended diffu-GRPO).
On math and code, where a verifier exists, post-training a masked dLM can skip online rollouts. Two hundred questions, four trajectories, ten tokens each: of the 2.8 hours, 2.5 are the one-time sample and 0.3 are the update.
The architecture stays the same. The gradient moves onto a single commitment, and a failed sample is penalized only there. The requirements are specific. The sampler must keep the step, the position, and the token, and the outcome must be automatically checkable. Open-ended generation has no sign to attach. Process reward models need a solution sliced into steps. Pivot-SD does not. It localizes supervision with the model's own entropy drop.
The paper lists three open points. A few hyperparameters are set by hand per configuration, and a new domain or backbone is expected to need retuning. Credit assignment stays inside one denoising block, so long contexts need the signal to cross block boundaries. The experiments cover two masked diffusion backbones at similar scale, LLaDA-8B and Dream-7B.
The tables already narrow a four-benchmark sweep. The MBPP+ gain of 0.8 sits inside one standard deviation. At 512 tokens HumanEval+ falls below the base score of 42.07, so the main claim is tied to decoding at the training length. On Dream, SFT-GT itself cuts GSM8K from 79.55 to 61.49, so part of the margin over full-sequence SFT is an unstable baseline. Against the Dream base, MBPP+ moves only from 60.58 to 61.64.
Information gain is never checked against human process labels, so there is no evidence that a pivot is a derivation step a person would mark. The main table does not include GIFT, which reweights tokens by entropy on the finished sequence. Entropy-based selection appears only in the positive-only ablation. Trajectories are drawn once from the frozen model, and the paper never resamples to see whether the trained model still commits in the same places. Every task has a hard verifier. Where correctness is not automatic, what unlikelihood would suppress is unanswered.