Papers
arxiv:2610.31576

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

Published on Oct 22
· Submitted by
Jaehong Yoon
on Oct 23
Authors:
,
,

Abstract

Optimized unmasking policies for masked diffusion language models improve efficiency by using trajectory-contraction targets with confidence and hidden states, achieving better speed-accuracy trade-offs with minimal overhead.

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.

Community

Paper submitter

.

Before the speedup matters, I want the cheap baseline in the table: unmask whichever token the denoiser is most confident about, zero extra parameters, no training run. If a learned policy doesn't clearly beat that, the gain is a rounding error with a GPU bill attached. The other thing that nags me is coupling — the policy is fit to one denoiser's confidence landscape, so swap the checkpoint or fine-tune it and the ordering it learned is some other model's opinion. A scheduler you have to retrain per checkpoint isn't a scheduler. Maybe it clears both bars; I just can't tell from the headline.

Sign up or log in to comment

Models citing this paper 1

Datasets citing this paper 0

No dataset linking this paper

Cite arxiv.org/abs/2610.31576 in a dataset README.md to link it from this page.

Spaces citing this paper 0

No Space linking this paper

Cite arxiv.org/abs/2610.31576 in a Space README.md to link it from this page.

Collections including this paper 0

No Collection including this paper

Add this paper to a collection to link it from this page.