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 θ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 Mt 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 ‡ 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 ‡ 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
MATH
GSM8K
AceCode
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)
Table 6: Credit-assignment landscape. Among the methods listed, Pivot-SD is the only one that trains on selected realized commitments (t,p,yp) alone and builds a localized negative target from failed trajectories.
Method
MATH
GSM8K
HumanEval+
MBPP+
SFT-GT
35.47 ± 0.23
76.94 ± 1.05
33.94 ± 2.32
46.32 ± 0.45
SFT-SD
31.73 ± 0.58
73.90 ± 0.89
35.78 ± 1.54
45.92 ± 1.46
diffu-GRPO
33.47 ± 1.03
76.04 ± 0.80
32.52 ± 1.27
43.39 ± 0.27
wd1++
33.27 ± 0.50
76.93 ± 0.59
33.54 ± 2.80
43.12 ± 0.27
Pivot-SD (ours)
37.47 ± 1.21
79.51 ± 1.30
40.43 ± 2.75
47.12 ± 1.20
Appendix
Table 7: Main comparison with per-run standard deviations. Mean ± standard deviation over three fully re-randomized runs, 256-token evaluation. Rows correspond to the ‡ -marked entries of Table 1 . RL baselines are the budget-matched variants.
Method
MATH
GSM8K
HumanEval+
MBPP+
Pos+Neg All-token
35.00 ± 1.25
74.44 ± 0.70
36.80 ± 0.35
44.54 ± 2.37
Random-Step Pivot
29.93 ± 3.13
76.51 ± 0.74
39.00 ± 5.51
44.80 ± 0.85
Random-Token Pivot
34.87 ± 1.94
76.47 ± 1.68
34.78 ± 1.07
45.24 ± 0.70
Entropy pivot
32.27 ± 1.10
74.75 ± 0.82
35.57 ± 1.76
44.80 ± 1.10
Pivot-SD (ours)
37.47 ± 1.21
79.51 ± 1.30
40.43 ± 2.75
47.12 ± 1.20
Appendix
Table 8: Ablation with per-run standard deviations. Mean ± 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× )
2,819
Training fwd+bwd (PFLOPs)
197
274
274
Total Compute (PFLOPs)
1,875
3,093 ( ∼1.65× )
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 × L40)
2.8 h
5.4 h
3.6 h
Appendix
Table 11: Compute accounting at L=512 and N=8×109 . Pivot-SD uses about 1.65× 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
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.
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.
Diffusion large language models (dLLMs) are a promising alternative to autoregressive generation. However, reasoning-oriented post-training for dLLMs remains challenging. Supervised fine-tuning (SFT) for dLLMs requires dense but often off-policy masked states, while reinforcement learning (RL) relies on sparse rewards or value modeling. This paper proposes trace-based on-policy distillation (TOPD), a teacher-supervised framework that transfers reasoning ability to a target dLLM without reward estimation. The key idea is to supervise a dLLM on its own denoising trajectory, focusing on the trace-aligned token decisions that form the final response. Specifically, TOPD samples on-policy diffusion trajectories from the target dLLM, obtains teacher token distributions from a teacher model on the corresponding partially denoised states, and updates the target dLLM with a token-level Reverse Kullback-Leibler (Reverse-KL) objective. This design preserves dense teacher supervision while aligning training with the model's own denoising states. On mathematical reasoning benchmarks, TOPD enables SDAR-4B-Chat to match the MATH500 accuracy of its RL-trained counterpart TraDo-4B-Instruct, with gains of +5.7 under static evaluation and +4.5 under dynamic evaluation. Compared with the RL-trained counterpart, TOPD achieves this with 4× fewer rollout rounds, corresponding to an estimated 96.0× to-accuracy model-compute speedup.
Haolin Ren, Ziyang Huang, Chenhao Yuan +2
Institute of Automation, Chinese Academy of Sciences, Beijing, China · University of Chinese Academy of Sciences, Beijing, China
Masked diffusion language models (dLMs) generate text through bidirectional denoising, yet this capability remains locked for infilling prompts. This limitation is an artifact of the current supervised finetuning (SFT) convention of applying response-only masking. To unlock this capability, we extend full-sequence masking during SFT, where both prompts and responses are masked jointly. Once unlocked, the model infills masked portions of a prompt template conditioned on few-shot examples. We show that such model-infilled prompts match or surpass manually designed templates, transfer effectively across models, and are complementary to existing prompt optimization methods. Our results suggest that training practices, not architectural limitations, are the primary bottleneck preventing masked diffusion language models from infilling effective prompts