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

cs.LGOct 2, 2026

Pivot-SD: Efficient Self-Distillation for Masked Diffusion Language Models

Authors: Seo Hyun Kim, Sunwoo Hong, Younwoo Choi, Chen-Hao Chao, Se-Young Yun, Rahul G. Krishnan

Organizations: KAIST AI · University of Toronto & Vector Institute

Abstract

Masked diffusion language models (dLMs) offer a promising parallel alternative to autoregressive models for complex reasoning. However, they face a distinct credit-assignment challenge, since a few commitments during denoising sharply reduce the uncertainty over the remaining masked positions and shape much of the response. Most post-training recipes for dLMs do not use this signal to decide which tokens to train on: they typically train on the final text or assign rewards to whole denoising steps, rather than selecting the individual commitments that shape the response. We introduce Pivot-SD, an efficient offline self-distillation framework that supervises only these high-impact commitments (pivots). Pivot-SD selects pivots using an information-gain metric measuring uncertainty reduction over the remaining masked positions. Pivots from successful trajectories are trained with cross-entropy, and pivots from failed trajectories with targeted unlikelihood, leaving the rest of the failed trajectory untouched. Using only 200 questions and four rollouts each, Pivot-SD improves LLaDA-8B-Instruct over full-sequence SFT and budget-matched diffusion RL baselines across math and code benchmarks.

Figures & tables

Figure 1: Predictive Entropy and Pivot Selection. An entropy heatmap of a masked denoising trajectory for a MATH problem (showing the first 128 tokens). The x-axis represents the token position, and the y-axis denotes the denoising step progressing downward. The color scale indicates the predictive entropy of the active masked positions. Red dashed lines and stars denote the critical denoising steps and the specific pivot tokens selected by our Information Gain metric, respectively.

Figure 2: Illustrative Overview of Pivot-SD. (a) In a trajectory sampled from

\theta_{0}

, steps whose commitment most reduces the average entropy of the remaining masked positions are selected as pivots (yellow). (b) The final answer determines the sign: for each pivot, its masked state

M_{t}

is replayed and only the pivot token is trained, with cross-entropy if the answer is correct and unlikelihood if it is wrong. All other positions receive no loss.

Method	MATH		GSM8K		HumanEval+		MBPP+
	256	512	256	512	256	512	256	512
Base	31.40	38.8	75.13	83.40	32.32	42.07	43.12	44.44
SFT-GT ‡	35.47	38.6	76.94	83.17	33.94	36.59	46.32	42.86
SFT-SD ‡	31.73	38.6	73.90	81.96	35.78	40.85	45.92	44.97
diffu-GRPO (budget-matched) ‡	33.47	38.2	76.04	82.11	32.52	42.07	43.39	45.24
diffu-GRPO	35.20	40.6	77.63	83.62	34.15	40.85	42.59	43.65

Table 1: Accuracy of LLaDA-8B-Instruct on math and code benchmarks at generation lengths of 256 and 512. Budget-matched RL variants use the same 200 questions and 1,000 optimizer steps as Pivot-SD. Methods marked with

{\ddagger}

are averaged over three fully re-randomized runs at 256 tokens. Standard deviations are in Appendix D .

Method	State	Target	Pos.	Neg.	MATH	GSM8K	HumanEval+	MBPP+
Positive All-token	IG	All	CE	-	31.60	74.37	36.59	39.68
Positive Pivot	IG	Pivot	CE	-	32.20	72.48	35.37	43.65
Entropy pivot ‡	Entropy	Pivot	CE	-	32.27	74.75	35.57	44.80
Pos+Neg All-token ‡	IG	All	CE	State-UL	35.00	74.44	36.80	44.54
Random-Step Pivot ‡	Random	Pivot	CE	Pivot-UL	29.93	76.51	39.00	44.80
Random-Token Pivot ‡	IG	Random	CE	Pivot-UL	34.87	76.47	34.78	45.24

Table 2: Ablation study of pivot selection and credit assignment. “IG” denotes denoising states selected by our information-gain criterion. “All” applies the loss to all supervised masked tokens in the selected state, while “Pivot” applies the loss only to the selected pivot token. CE denotes cross-entropy and UL denotes unlikelihood. Methods marked with

{\ddagger}

are averaged over three fully re-randomized runs. Standard deviations are in Appendix D .

Method	Steps	Gen. (h)	Train (h)	Total (h)
diffu-GRPO (matched)	1000	0	5.4	5.4
diffu-GRPO	5000	0	22.9	22.9
wd1++ (matched)	1000	0	3.6	3.6
wd1++	5000	0	17.3	17.3
Pivot-SD (ours)	1000	2.5	0.3	2.8

Table 3: Compute accounting. Wall-clock hours on two GPUs. Pivot-SD generates its trajectories once before training. The RL baselines generate rollouts during training. RL times are averaged across training domains.

Figure 3: Accuracy vs. compute cost. Average accuracy versus total wall-clock hours on two GPUs. Pivot-SD is more than 3 points above the 5,000-step RL runs while using about one-eighth of the training time of 5,000-step diffu-GRPO.

Method	MATH	GSM	HE	MBPP
Base (Dream-Instruct-7B)	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 (ours)	42.20	81.88	58.54	61.64

Table 4: Dream generalization results. We keep the same baseline categories as in the LLaDA experiments and report the same evaluation metrics.

Method	Training domain			Overall
Method	MATH	GSM8K	AceCode	Overall
Base	50.19	35.61	53.27	45.49
SFT-GT	53.57	39.46	51.80	47.83
SFT-SD	51.13	35.44	51.21	45.27
diffu-GRPO (matched)	50.89	36.45	55.49	46.62
diffu-GRPO	53.02	37.99	54.66	47.79

Table 5: Average out-of-domain (OOD) performance. Each column reports the mean accuracy over the benchmarks not used for training: MATH- and GSM8K-trained models are evaluated on the other three benchmarks, AceCode-trained models on MATH and GSM8K. Column composition therefore differs, and “Overall” is the micro-average over all source-target pairs rather than the mean of the three columns.

Appendix figures & tables8 assets

Supplementary material from the paper’s appendix.

Appendix

Method	Local unit	Localization signal	Training regime	Localized Failure Target
STaR / rejection SD ( Zelikman et al., 2022 )	Full successful rationale	Final correctness	Offline SFT	No (failed traces are filtered or retried)
Process supervision (PRMs) ( Lightman et al., 2024 ; Wang et al., 2024 )	Textual reasoning step	Human or separately constructed process label	PRM / SFT / RL	Possible, but requires process-label construction
GIFT ( Xu et al., 2025 )	Final-response token position	Local predictive entropy	Offline weighted SFT	No
Denoising Process Reward / ATPO ( Xie et al., 2026 ; Chen et al., 2025 )	Denoising interval or step	Outcome-progress reward or trajectory uncertainty	Online RL	No explicit token-level UL target
AGRPO ( Zhan, 2025 )	Sampled unmasking step	Trajectory-level advantage; steps sampled uniformly or by entropy	Online RL	No (trajectory-level advantage on every sampled step)
DiSPO ( Oba et al., 2026 )	State-conditioned mask filling / newly filled tokens	Branched terminal rewards	Policy optimization	No cached realized-pivot UL target

Appendix

Table 6: Credit-assignment landscape. Among the methods listed, Pivot-SD is the only one that trains on selected realized commitments

(t,p,y_{p})

alone and builds a localized negative target from failed trajectories.

Method	MATH	GSM8K	HumanEval+	MBPP+
SFT-GT	35.47 $\pm$ 0.23	76.94 $\pm$ 1.05	33.94 $\pm$ 2.32	46.32 $\pm$ 0.45
SFT-SD	31.73 $\pm$ 0.58	73.90 $\pm$ 0.89	35.78 $\pm$ 1.54	45.92 $\pm$ 1.46
diffu-GRPO	33.47 $\pm$ 1.03	76.04 $\pm$ 0.80	32.52 $\pm$ 1.27	43.39 $\pm$ 0.27
wd1++	33.27 $\pm$ 0.50	76.93 $\pm$ 0.59	33.54 $\pm$ 2.80	43.12 $\pm$ 0.27
Pivot-SD (ours)	37.47 $\pm$ 1.21	79.51 $\pm$ 1.30	40.43 $\pm$ 2.75	47.12 $\pm$ 1.20

Appendix

Table 7: Main comparison with per-run standard deviations. Mean

\pm

standard deviation over three fully re-randomized runs, 256-token evaluation. Rows correspond to the

{\ddagger}

-marked entries of Table 1 . RL baselines are the budget-matched variants.

Method	MATH	GSM8K	HumanEval+	MBPP+
Pos+Neg All-token	35.00 $\pm$ 1.25	74.44 $\pm$ 0.70	36.80 $\pm$ 0.35	44.54 $\pm$ 2.37
Random-Step Pivot	29.93 $\pm$ 3.13	76.51 $\pm$ 0.74	39.00 $\pm$ 5.51	44.80 $\pm$ 0.85
Random-Token Pivot	34.87 $\pm$ 1.94	76.47 $\pm$ 1.68	34.78 $\pm$ 1.07	45.24 $\pm$ 0.70
Entropy pivot	32.27 $\pm$ 1.10	74.75 $\pm$ 0.82	35.57 $\pm$ 1.76	44.80 $\pm$ 1.10
Pivot-SD (ours)	37.47 $\pm$ 1.21	79.51 $\pm$ 1.30	40.43 $\pm$ 2.75	47.12 $\pm$ 1.20

Appendix

Table 8: Ablation with per-run standard deviations. Mean

\pm

standard deviation over three fully re-randomized runs, 256-token evaluation.

	MATH				GSM8K				HumanEval+				MBPP+
Method	50	200	400	2000	50	200	400	2000	50	200	400	2000	50	200	400	2000
SFT-GT	36.0	35.6	34.2	35.0	76.7	75.8	77.9	75.1	32.9	32.9	32.9	33.5	46.6	46.6	45.8	46.8
SFT-SD	31.6	31.4	33.0	31.6	72.4	74.9	72.7	73.2	37.8	35.4	37.2	37.8	45.5	45.8	44.7	45.5
diffu-GRPO	34.0	34.6	34.6	35.2	76.4	75.7	76.4	77.6	31.7	31.1	32.9	34.2	42.3	43.7	43.9	42.6
wd1++	32.4	33.8	39.2	35.4	75.9	77.1	79.0	76.7	32.3	32.9	31.1	34.2	44.7	43.1	45.8	45.0
Pivot-SD (ours)	36.2	38.6	36.0	37.2	75.0	79.5	80.1	78.9	42.1	40.9	40.2	38.4	46.0	46.0	46.6	47.1

Appendix

Table 9: Accuracy across training budgets (256-token evaluation). Column headers give the number of training questions

q

. Step budgets are fixed within each group (

q=50

200

: 1,000 steps;

q=400

: 2,000;

q=2000

: 5,000), matching the extended RL baselines. All entries are single runs. Multi-seed means for

q=200

are in Table 7 .

	MATH			GSM8K			HumanEval+			MBPP+
$K$	Random Step	Random Token	Ours	Random Step	Random Token	Ours	Random Step	Random Token	Ours	Random Step	Random Token	Ours
5	22.8	34.2	37.8	74.3	75.1	79.6	39.0	34.8	39.0	42.6	44.2	45.0
10	26.6	33.2	38.6	76.9	75.7	79.5	35.4	37.8	40.9	46.0	45.8	46.0
20	17.6	36.4	38.2	73.8	76.0	76.8	39.6	37.2	41.5	41.5	45.5	45.5
40	22.0	34.6	22.8	73.0	77.0	79.5	39.0	43.3	39.6	41.0	44.4	39.7
80	20.4	33.0	17.6	69.8	77.6	76.7	37.2	42.7	30.5	34.7	41.8	37.6

Appendix

Table 10: Effect of the pivot budget

K

(

q=200

, 256-token evaluation).

K

is the number of pivots supervised per trajectory. Random-Step and Random-Token are the two random-selection controls of Table 2 . We use

K=10

K=40

and

80

show what happens when selection is no longer selective. Entries are single runs on a shared trajectory pool. The multi-seed mean for

K=10

is in Table 7 .

Metric	Pivot-SD (ours)	diffu-GRPO	wd1++
Generation (PFLOPs)	1,678 (offline)	2,819 ( $1.68\times$ )	2,819
Training fwd+bwd (PFLOPs)	197	274	274
Total Compute (PFLOPs)	1,875	3,093 ( $\sim$ $1.65\times$ )	3,093
Model Calls (NFE)	204,800 (offline)	358,080 (online)	358,080 (online)
Supervised Tokens / Trajectory	10 (3.9%)	256 (100%)	256 (100%)
Wall-Clock Time (2 $\times$ L40)	2.8 h	5.4 h	3.6 h

Appendix

Table 11: Compute accounting at

L=512

and

N=8\times 10^{9}

. Pivot-SD uses about

1.65\times

fewer total FLOPs than the budget-matched RL runs. Its generation happens once before training, while the RL runs generate rollouts throughout training.

Method	MATH Trained				GSM8K Trained				AceCode Trained
Method	GSM8K	HumanEval+	MBPP+	Avg.	MATH	HumanEval+	MBPP+	Avg.	MATH	GSM8K	Avg.
SFT-GT	76.35	37.80	46.56	53.57	35.40	37.20	45.77	39.46	30.20	73.39	51.80
SFT-SD	75.74	32.93	44.71	51.13	30.60	30.49	45.24	35.44	29.40	73.01	51.21
diffu-GRPO (budget-matched)	75.21	33.54	43.92	50.89	33.80	31.10	44.44	36.45	34.40	76.57	55.49
diffu-GRPO	76.35	37.20	45.50	53.02	35.80	32.93	45.24	37.99	33.80	75.51	54.66
wd1++ (budget-matched)	76.42	32.32	43.39	50.71	35.00	36.59	42.59	38.06	32.60	74.91	53.76

Appendix

Table 12: Out-of-Domain (OOD) generalization results. Models are trained on one domain and evaluated on others. “budget-matched” denotes models trained with the 1,000-step budget configuration matching Pivot-SD.

	Math		Code (256)		Code (512)
Variant	MATH	GSM	MBPP+	HE+	MBPP+	HE+
w/o bootstrap
Step 500	32.60	74.91	43.12	32.93	43.92	40.24
Step 1000	-	-	43.12	32.93	-	-
w/ bootstrap
Step 500	33.20	74.53	44.71	31.10	46.56	39.63

Appendix

Table 13: Effect of the bootstrap reward on wd1++ trained on 200 AceCode prompts. GSM and HE+ denote GSM8K and HumanEval+. Without bootstrap, training stalls between steps 500 and 1000 because the zero-advantage gradients cancel. With the bootstrap reward, training resumes on code and MBPP+ gains 3.97 points by step 1000, while MATH does not improve. The wd1++ numbers in the main tables (Tables 1 and 7 ) come from the standard runs without bootstrap shaping. The

47.09

here is the single-run peak after bootstrap shaping and is outside the multi-seed protocol.

Diffusion Language Model Distillation

Diffusion language model distillation transfers knowledge or decoding behavior to make diffusion-based text generation faster and more capable.

See all 34 papers

Revise, Don't Freeze: Sampler-Matched Training for Self-Correcting Masked Diffusion Language Models
May 31, 2026Longxuan Yu, Shaorong Zhang, Yu Fu +3Masked Diffusion Models Diffusion Model Fine-Tuning
University of California, Riverside · Microsoft
Trace-Based On-Policy Distillation for Masked Diffusion Language Models
Jul 18, 2026Haolin Ren, Ziyang Huang, Chenhao Yuan +2Masked Diffusion Models Diffusion Model Distillation
Institute of Automation, Chinese Academy of Sciences, Beijing, China · University of Chinese Academy of Sciences, Beijing, China
Unlocking Prompt Infilling Capability for Diffusion Language Models
Apr 4, 2026Yoshinari Fujinuma, Keisuke SakaguchiDiffusion Language Models Masked Diffusion Models
Patronus AI · Tohoku University, RIKEN

May 31, 2026J/K move · Enter open · S save

Revise, Don't Freeze: Sampler-Matched Training for Self-Correcting Masked Diffusion Language Models

Longxuan Yu, Shaorong Zhang, Yu Fu, Hui Liu+2

University of California, Riverside · Microsoft

Masked Diffusion Models Diffusion Model Fine-Tuning Diffusion Language Model Decoding Diffusion Language Models Diffusion Language Model Training

Open paper AskPDF arXiv

Masked diffusion language models (MDLMs) re-predict every position at each denoising step, but standard samplers commit tokens once revealed, leaving this revision capability unused. Existing approaches either add heuristic or learned mechanisms to revise committed tokens, or remask them back to [MASK] before re-predicting; a principled sampler that directly revises visible tokens without auxiliary modules remains underexplored. We introduce D3IM, a parameter-free sampler derived as a corrector-style reverse update that permits direct visible-to-visible revision without additional modules or auxiliary passes. D3IM also reveals a model-side obstacle we term preservation bias: the model tends to reproduce its own wrong committed tokens rather than correct them. We address this with SCOPE (Self-Conditioned On Prediction Errors), a lightweight post-training procedure that simulates D3IM's sampling process. On LLaDA-8B at 64 denoising steps, SCOPE+D3IM improves over the original LLaDA-8B with standard unmasking by +13.0 on GSM8K (68.3%), +4.8 on MATH-500 (23.6%), +15.3 on HumanEval (29.3%), and +10.4 on MBPP (30.8%), with gains that increase as more denoising steps are used on math and HumanEval.

Abstract

Diffusion Language Model Distillation

Explore similar work

Revise, Don't Freeze: Sampler-Matched Training for Self-Correcting Masked Diffusion Language Models

Trace-Based On-Policy Distillation for Masked Diffusion Language Models

Unlocking Prompt Infilling Capability for Diffusion Language Models