What Makes a Good Unmasking Policy? A Controlled Study of Learned Unmasking in Diffusion Language Models

Jinya Sakurai; Xun Xu; Jaehong Yoon

cs.AIOct 22, 2026

What Makes a Good Unmasking Policy? A Controlled Study of Learned Unmasking in Diffusion Language Models

Authors: Jinya Sakurai, Xun Xu, Jaehong Yoon

Organizations: NTU Singapore · Institute for Advanced Intelligence and Computing (IAIC), A*STAR · Centre for Frontier AI Research (CFAR), A*STAR

Abstract

Masked diffusion language models generate tokens in parallel, but their efficiency depends on which predictions are committed. Learned unmasking policies improve the speed--accuracy trade-off, yet differ in supervision, inputs, and capacity, obscuring which choices matter. We view supervised policies as distilling a decoding procedure through a teacher and examine independently what the teacher labels as committable, on which states the policy imitates it, what the policy observes, and how much the number of function evaluations (NFE) must decrease to offset policy overhead. We find that labeling predictions matching the final output is overly permissive at low NFE, and hidden states are weak alone but can complement confidence. About a 10% NFE reduction suffices to offset policy overhead for policies up to 7.5M parameters across tested settings. We therefore train a policy on trajectory-contraction targets, which accept only commitments staying on a reference trajectory, using confidence and hidden states. On LLaDA-8B-Instruct, the resulting 856K-parameter policy achieves the best accuracy--NFE frontier and highest throughput on GSM8K, is competitive at moderate targets on MATH-500 and HumanEval, and largely transfers to Dream-7B-Instruct.

Figures & tables

Method	#Params	Policy input	Commitment target	Training states
Learn2PD ( Bao et al., 2026 )	2.1K	confidence	terminal agreement	teacher-induced
TraceLock ( Sun et al., 2026 )	7.5M †	hidden states	terminal agreement	reference
LUP-conf ( Jazbec et al., 2026 )	300K	confidence	reinforcement learning	on-policy rollouts
LUP-hidden ( Jazbec et al., 2026 )	300M	hidden states	reinforcement learning	on-policy rollouts
Candidates studied	331K–856K	$c$ , $h$ , $[c,m]$ , $[c,h]^{\ast}$	terminal agreement, contraction ∗	reference, teacher-induced

Table 1: Design choices of existing learned unmasking policies and the candidates compared in our study. Existing methods differ simultaneously in their policy inputs, policy capacity, commitment targets, and training states. ∗ Candidates introduced in this work. † Estimated from the released configuration; not reported in the original paper.

Figure 1: Design choices in supervising unmasking policies. (a) Terminal agreement labels every prediction matching the final output, whereas trajectory contraction labels only positions in the furthest jump that stays on the reference trajectory; they disagree when a correct token is committed before a mispredicted one. (b) Following the teacher skips

s_{2}

s_{4}

, and

s_{5}

, which reference-trajectory supervision includes but teacher-induced supervision omits.

Figure 2: Accuracy–NFE frontiers for the

2\times 2

supervision designs. Faded markers denote dominated operating points.

Figure 3: Accuracy–NFE frontiers for different inputs.

		No policy	Policy overhead
Prompt	Cache	(ms/step)	331K	856K	7.5M	300M
128	None	16.9	+10%	+8%	+9%	+38%
	Prefix	16.0	+8%	+8%	+9%	+39%
512	None	28.0	+5%	+6%	+6%	+23%
	Prefix	16.5	+7%	+10%	+8%	+37%
1024	None	45.6	+3%	+3%	+4%	+15%

Table 2: Per-step latency on LLaDA-8B-Instruct without a policy, and the relative overhead of policies of different sizes. Absolute latencies are reported in Table 8 .

Figure 4: Accuracy versus NFE (top) and TPS (bottom) on LLaDA-8B-Instruct. Dotted lines mark the reference accuracy and the 98%, 95%, and 90% targets. Numerical results are provided in Tables 17 and 18 .

	Zero-shot				Four-shot
	98%		95%		98%		95%
Method	No	Pre.	No	Pre.	No	Pre.	No	Pre.
OnTrack (Ours)	331	249	381	423	121	225	149	241
Fast-dLLM	232	276	323	349	98	193	122	234
Learn2PD	242	253	282	299	92	183	110	247
EB-Sampler	221	234	313	329	120	171	120	231

Table 3: Highest measured TPS (tokens/s) on GSM8K reaching 98% and 95% of the uncached reference accuracy, without (No) and with (Pre.) prefix caching. Bold marks the best method in each column.

Figure 5: Accuracy versus NFE on Dream-7B-Instruct. Dotted lines mark the reference accuracy and 98%, 95%, and 90% targets. Faded markers denote dominated operating points. Numerical results are in Tables 21 and 22 .

Figure 6: Accuracy-NFE trade-off at

L=128

(top) and

L=512

(bottom), using a policy trained at

L=256

. Results on HumanEval and MBPP are provided in Section C.6 ; numerical results in Tables 23 and 25 .

Appendix figures & tables29 assets

Supplementary material from the paper’s appendix.

Appendix

Component	LLaDA	Dream
Policy width $d$	128	128
DiT depth	1	1
Attention heads	2	2
Confidence pathway + DiT	331,297	331,297
Hidden projection	524,416	458,880
Total policy parameters	855,713	790,177

Appendix

Table 4: Configuration of the unmasking policy for each backbone. The hidden projection maps the 4096-dimensional (LLaDA) or 3584-dimensional (Dream) hidden state to

d

dimensions.

Figure 7: Architecture of our unmasking policy. Following Jazbec et al. (2026) , confidence is projected into the policy representation space and processed by a DiT-style transformer. Our policy adds a projected backbone hidden representation to the confidence representation.

Section	Varied	Fixed
Section 3.3	label, training states	input $[c,h]$
Section 3.4	policy input	contraction, reference

Appendix

Table 5: Factors varied and fixed in the controlled comparisons. All other settings follow Section C.3 .

Training states	#States	Ratio
Reference	686,052	100%
Induced (contraction)	233,309	34.0%
Induced (terminal)	174,390	25.4%

Appendix

Table 6: Number of training states for each training-state configuration, over all 9,169 reference trajectories (74.8 states per trajectory on average).

Input	#Params
$c_{t,i}$	331,297
$h_{t,i}$	855,457
$[c_{t,i},\Delta_{t,i}]$	331,425
$[c_{t,i},h_{t,i}]$	855,713

Appendix

Table 7: Parameter counts of the policies in Section 3.4 . The DiT block accounts for 297,376 parameters and the timestep embedding for 33,024; the hidden-state projection adds 524,416.

		No policy	With policy (ms/step)
Prompt	Cache	(ms/step)	331K	856K	7.5M	300M
128	None	16.9	18.5	18.3	18.4	23.3
	Prefix	16.0	17.2	17.3	17.4	22.2
512	None	28.0	29.4	29.6	29.6	34.5
	Prefix	16.5	17.7	18.1	17.9	22.7
1024	None	45.6	47.1	47.1	47.3	52.3

Appendix

Table 8: Per-step latency (ms) on LLaDA-8B-Instruct with policies of different sizes (H200, generation length 256, block length 32, batch size 1). Each step is the backbone forward plus the policy forward and per-step bookkeeping; the 7.5M and 300M policies are randomly initialized.

Figure 8: Accuracy–NFE frontiers for the

2\times 2

supervision designs on MATH-500 and HumanEval. Faded markers denote dominated operating points.

Figure 9: Accuracy–NFE frontiers for different policy inputs on MATH-500 and HumanEval. Faded markers denote dominated operating points.

Dataset	Training	Evaluation
GSM8K	7,473	1,319
MATH	7,485	–
MATH-500	–	500
HumanEval	–	164
MBPP	–	257
Total	14,958	2,240

Appendix

Table 9: Training and evaluation dataset sizes.

Hyperparameter	Value
Optimizer	AdamW
Learning rate	$3\times 10^{-5}$
Adam $(\beta_{1},\beta_{2})$	$(0.9,0.99)$
Adam $\epsilon$	$10^{-8}$
Weight decay	$0.1$
LR scheduler	Cosine decay

Appendix

Table 10: Optimization hyperparameters used for policy training.

Setting	Value
Training data	GSM8K + MATH train
Reference decoder	Fast-dLLM
Reference threshold	$0.9$
Generation length	256
Block length	32
Training temperature	0

Appendix

Table 11: Training-trajectory and numerical configuration.

Method	Score / commitment rule	Param.
OnTrack (Ours)	learned policy probability $\sigma(\ell_{t,i})>\gamma$	$\gamma$
Learn2PD	learned filter probability $>\gamma$	$\gamma$
Fast-dLLM	top-1 token probability $>\gamma$	$\gamma$
EB-Sampler	entropy-bounded prefix of confidence-sorted positions	$\epsilon$

Appendix

Table 12: Commitment rules and swept parameters.

Figure 10: Accuracy versus throughput on GSM8K with zero-shot (top) and four-shot (bottom) prompting. Solid and dashed lines indicate decoding without and with prefix caching, respectively. Horizontal lines mark the uncached one-by-one reference accuracy and its 98%, 95%, and 90% retention thresholds.

Figure 11: Accuracy versus NFE on HumanEval and MBPP at generation lengths

L=128

(top) and

L=512

(bottom) for a policy trained at

L=256

GSM8K				MBPP
Method	Param.	Acc.	NFE	Method	Param.	Acc.	NFE
Teacher + terminal	0.99	79.2	256.0	Teacher + terminal	0.99	52.5	256.0
	0.98	78.8	100.1		0.98	52.5	187.7
	0.96	78.4	66.0		0.96	52.5	127.6
	0.94	77.8	55.9		0.94	52.5	100.5
	0.9	75.1	46.8		0.9	51.4	70.4

Appendix

Table 13: All operating points for the supervision study on GSM8K and MBPP ( Figure 2 ).

GSM8K				MBPP
Method	Param.	Acc.	NFE	Method	Param.	Acc.	NFE
Conf.	0.99	78.1	256.0	Conf.	0.99	53.3	256.0
	0.98	78.2	158.4		0.98	53.3	162.2
	0.96	78.1	72.5		0.96	52.5	65.8
	0.94	77.9	62.0		0.94	52.5	58.8
	0.9	77.0	52.9		0.9	51.4	51.9

Appendix

Table 14: All operating points for the policy-input study on GSM8K and MBPP ( Figure 3 ).

MATH-500				HumanEval
Method	Param.	Acc.	NFE	Method	Param.	Acc.	NFE
Teacher + terminal	0.99	38.0	256.0	Teacher + terminal	0.99	42.1	256.0
	0.98	38.2	137.6		0.98	42.1	246.0
	0.96	36.6	93.9		0.96	42.1	213.3
	0.94	37.0	81.3		0.94	41.5	175.4
	0.9	36.0	68.6		0.9	39.6	126.3

Appendix

Table 15: All operating points for the supervision study on MATH-500 and HumanEval ( Figure 8 ).

MATH-500				HumanEval
Method	Param.	Acc.	NFE	Method	Param.	Acc.	NFE
Confidence	0.99	37.6	256.0	Confidence	0.99	42.7	256.0
	0.98	37.8	170.0		0.98	42.7	182.9
	0.96	38.2	98.4		0.96	42.7	111.9
	0.94	37.6	86.5		0.94	40.9	100.1
	0.9	38.2	76.7		0.9	42.1	86.5

Appendix

Table 16: All operating points for the policy-input study on MATH-500 and HumanEval ( Figure 9 ).

GSM8K					MATH-500
Method	Param.	Acc.	NFE	TPS	Method	Param.	Acc.	NFE	TPS
OnTrack (Ours)	0.99	78.5	256.0	54	OnTrack (Ours)	0.99	37.2	256.0	53
	0.98	78.4	109.0	127		0.98	37.2	131.2	103
	0.96	78.8	70.6	196		0.96	37.6	88.7	153
	0.94	77.9	60.6	230		0.94	37.8	77.8	177
	0.9	77.8	51.7	274		0.9	38.0	68.0	200

Appendix

Table 17: All operating points on GSM8K and MATH-500 with LLaDA-8B-Instruct ( Figure 4 ).

HumanEval					MBPP
Method	Param.	Acc.	NFE	TPS	Method	Param.	Acc.	NFE	TPS
OnTrack (Ours)	0.99	43.3	256.0	39	OnTrack (Ours)	0.99	53.7	256.0	38
	0.98	43.3	247.1	40		0.98	53.7	196.9	48
	0.96	43.3	212.9	46		0.96	53.7	131.9	72
	0.94	43.9	176.5	57		0.94	53.7	103.3	93
	0.9	43.3	130.0	76		0.9	54.1	73.4	129

Appendix

Table 18: All operating points on HumanEval and MBPP with LLaDA-8B-Instruct ( Figure 4 ).

No cache					Prefix cache
Method	Param.	Acc.	NFE	TPS	Method	Param.	Acc.	NFE	TPS
OnTrack (Ours)	0.99	78.5	256.0	54	OnTrack (Ours)	0.99	77.3	256.0	58
	0.98	78.4	109.0	127		0.98	77.3	111.5	132
	0.96	78.8	70.6	196		0.96	78.2	72.0	203
	0.94	77.9	60.6	230		0.94	77.9	61.9	239
	0.9	77.8	51.7	274		0.9	77.9	52.7	249

Appendix

Table 19: Full GSM8K results with and without prefix caching under zero-shot prompting ( Table 3 ).

No cache					Prefix cache
Method	Param.	Acc.	NFE	TPS	Method	Param.	Acc.	NFE	TPS
OnTrack (Ours)	0.99	80.8	256.0	22	OnTrack (Ours)	0.99	79.3	256.0	50
	0.98	81.1	112.6	50		0.98	79.4	117.3	106
	0.96	80.9	72.6	78		0.96	78.8	76.1	154
	0.94	80.5	62.1	91		0.94	78.8	65.1	178
	0.9	79.2	52.8	107		0.9	77.2	55.1	203

Appendix

Table 20: Full GSM8K results with and without prefix caching under four-shot prompting ( Table 3 ).

GSM8K					MATH-500
Method	Param.	Acc.	NFE	TPS	Method	Param.	Acc.	NFE	TPS
OnTrack (Ours)	0.99	81.7	256.0	52	OnTrack (Ours)	0.99	46.0	256.0	54
	0.98	81.6	154.4	85		0.98	46.0	158.7	86
	0.96	81.5	98.4	131		0.96	46.6	103.5	132
	0.94	81.2	83.7	154		0.94	46.8	89.7	152
	0.9	80.4	70.6	184		0.9	46.0	77.0	175

Appendix

Table 21: All operating points on GSM8K and MATH-500 with Dream-7B-Instruct ( Figure 5 ).

HumanEval					MBPP
Method	Param.	Acc.	NFE	TPS	Method	Param.	Acc.	NFE	TPS
OnTrack (Ours)	0.99	58.5	256.0	27	OnTrack (Ours)	0.99	66.9	256.0	12
	0.98	58.5	140.5	49		0.98	66.9	65.7	44
	0.96	58.5	121.3	56		0.96	66.9	54.6	52
	0.94	58.5	104.6	65		0.94	66.9	47.3	61
	0.9	59.8	84.1	81		0.9	65.8	38.5	73

Appendix

Table 22: All operating points on HumanEval and MBPP with Dream-7B-Instruct ( Figure 5 ).

GSM8K					MATH-500
Method	Param.	Acc.	NFE	TPS	Method	Param.	Acc.	NFE	TPS
OnTrack (Ours)	0.99	72.5	128.0	64	OnTrack (Ours)	0.99	35.0	128.0	58
	0.98	72.3	67.3	124		0.98	35.0	86.2	92
	0.96	72.3	45.8	181		0.96	35.2	57.8	128
	0.94	71.9	39.7	210		0.94	33.0	50.2	153
	0.9	72.5	34.1	245		0.9	34.0	43.3	181

Appendix

Table 23: All operating points on GSM8K and MATH-500 at

L=128

( Figure 6 ).

HumanEval					MBPP
Method	Param.	Acc.	NFE	TPS	Method	Param.	Acc.	NFE	TPS
OnTrack (Ours)	0.99	43.9	128.0	53	OnTrack (Ours)	0.99	51.4	128.0	47
	0.98	43.9	126.3	53		0.98	51.4	110.3	58
	0.96	43.9	117.2	59		0.96	51.4	87.4	73
	0.94	43.9	101.7	67		0.94	51.4	71.6	86
	0.9	43.3	78.7	87		0.9	50.6	51.6	113

Appendix

Table 24: All operating points on HumanEval and MBPP at

L=128

( Figure 11 ).

GSM8K					MATH-500
Method	Param.	Acc.	NFE	TPS	Method	Param.	Acc.	NFE	TPS
OnTrack (Ours)	0.99	82.2	348.5	39	OnTrack (Ours)	0.99	39.0	445.9	40
	0.98	82.4	136.2	102		0.98	39.4	208.3	85
	0.96	82.9	86.6	160		0.96	38.4	134.9	130
	0.94	82.3	74.3	186		0.94	37.6	117.0	150
	0.9	82.3	63.1	217		0.9	39.8	101.1	174

Appendix

Table 25: All operating points on GSM8K and MATH-500 at

L=512

( Figure 6 ).

HumanEval					MBPP
Method	Param.	Acc.	NFE	TPS	Method	Param.	Acc.	NFE	TPS
OnTrack (Ours)	0.99	40.9	419.6	34	OnTrack (Ours)	0.99	50.6	222.3	38
	0.98	40.9	418.1	34		0.98	50.6	220.9	37
	0.96	40.9	352.6	40		0.96	50.6	157.0	53
	0.94	40.9	285.0	49		0.94	50.2	124.0	68
	0.9	40.2	198.9	71		0.9	50.6	85.7	100

Appendix

Table 26: All operating points on HumanEval and MBPP at

L=512

( Figure 11 ).

		NFE speedup (target: % of reference acc.)
Dataset	Baseline	90%	95%	98%
GSM8K	Fast-dLLM	1.20 $\times$ [1.11, 1.47]	1.25 $\times$ [1.11, 1.75]	1.56 $\times$ [1.08, 2.10]
	Learn2PD	1.10 $\times$ [1.08, 1.20]	1.40 $\times$ [1.01, 1.65]	1.45 $\times$ [1.00, 5.00]
	EB-Sampler	1.13 $\times$ [1.04, 1.49]	1.28 $\times$ [1.05, 1.94]	1.59 $\times$ [0.95, 2.52]
MATH-500	Fast-dLLM	1.25 $\times$ [1.05, 2.02]	1.50 $\times$ [1.00, 2.02]	n.r.
	Learn2PD	1.18 $\times$ [0.98, 1.80]	1.38 $\times$ [0.96, 1.84]	–

Appendix

Table 27: NFE speedup of OnTrack over each LLaDA-8B baseline, defined as the baseline’s NFE divided by OnTrack’s NFE to reach 90%, 95%, or 98% of the one-by-one highest-confidence reference accuracy (GSM8K 78.0, MATH-500 38.8, HumanEval 45.1, MBPP 52.9; the same target for all methods). NFE-to-target is the smallest measured mean NFE whose accuracy meets the target, on the common threshold grid without interpolation. Brackets give paired bootstrap 95% confidence intervals over evaluation problems (10,000 resamples); bold marks intervals that exclude 1. n.r.: OnTrack does not reach the target; –: neither method does.

Diffusion Language Model Decoding

Diffusion language model decoding generates text through iterative parallel refinement, balancing fewer denoising steps with accuracy, coherence, and control.

See all 249 papers

TIPS: A Tiny Inference-time Policy for Scaling in Diffusion Language Models
Oct 15, 2026Ya-Shi Zhang, Jian Tang, Pierre-André Noël +1Diffusion Language Model Decoding Inference-Time Adaptation of Language Models
ServiceNow AI Research · Mila - Quebec AI Institute · HEC Montréal
Mask-Aware Policy Gradients for Diffusion Language Models
Jul 16, 2026Haran Raajesh, Kulin Shah, Adam Klivans +1Masked Diffusion Models RL for Language Model Reasoning
Department of Computer Science The University of Texas at Austin
Enhanced Masked Diffusion Unlearning
Oct 18, 2026Tomer Garber, Tom TirerMachine Unlearning Diffusion Model Unlearning
Faculty of Engineering, Bar-Ilan University, Ramat Gan, Israel · Red Hat

Oct 15, 2026J/K move · Enter open · S save

TIPS: A Tiny Inference-time Policy for Scaling in Diffusion Language Models

Ya-Shi Zhang, Jian Tang, Pierre-André Noël, Oleksiy Ostapenko

ServiceNow AI Research · Mila - Quebec AI Institute · HEC Montréal

Diffusion Language Model Decoding Inference-Time Adaptation of Language Models Inference-Time Scaling Masked Diffusion Models

Open paper AskPDF arXiv

Masked diffusion language models expose an inference-time choice: which positions become visible after each denoising pass. Many existing decoding methods use hand-designed rules or require visibility to grow monotonically. We present a

\textbf{T}

iny

\textbf{I}

nference-time

\textbf{P}

olicy for

\textbf{S}

caling (

\textbf{TIPS}

), an approximately

600

K-parameter controller that learns both unmasking and remasking while keeping the underlying language model frozen. TIPS uses visibility-aware confidence features and a shallow transformer to select which tokens should be visible, while a fixed schedule controls the expected visible fraction. We train the controller through supervised warmup followed by reinforcement learning from terminal correctness rewards. On GSM8K and MATH-500 with LLaDA-8B-Instruct, the full-canvas and two-block TIPS variants achieve

58.46\%

and

57.97\%

micro-averaged accuracy at a budget of

32

decoder evaluations, exceeding the strongest baseline in our main comparison by

12.1

and

11.6

percentage points, respectively. Both variants are trained at a budget of

32

evaluations and remain competitive at higher budgets without retraining. These results demonstrate that learning when to reveal and reconsider tokens can improve the accuracy-compute trade-off without updating the decoder backbone.

Abstract

Diffusion Language Model Decoding

Explore similar work

TIPS: A Tiny Inference-time Policy for Scaling in Diffusion Language Models

Mask-Aware Policy Gradients for Diffusion Language Models

Enhanced Masked Diffusion Unlearning