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

What Makes a Good Unmasking Policy? A Controlled Study of Learned Unmasking in Diffusion Language Models. By Jinya Sakurai and two others, from Nanyang Technological University, Singapore, the Institute for Advanced Intelligence and Computing at A-STAR, and one other institution. Published on the twenty-second of October. No code repository is linked to it.

Masked diffusion language models generate tokens in parallel, and how efficient they are depends on which predictions get committed at each step. Learned unmasking policies improve the trade-off between speed and accuracy. But the existing ones differ in supervision, in inputs, and in capacity all at once, which hides which of those choices actually matter. This study pulls them apart.

The authors treat a supervised policy as distilling a decoding procedure through a teacher. Then they vary four things independently: what the teacher labels as committable, on which states the policy imitates it, what the policy gets to observe, and how far the number of function evaluations, or NFE, has to fall to pay for the policy's own overhead.

Three findings come out. Labeling a prediction as committable simply because it matches the final output is overly permissive when NFE is low. Hidden states are weak as the only input, but can complement confidence. And a reduction of about ten percent in NFE is enough to offset the overhead, for policies up to seven point five million parameters.

Putting those together, the authors train a policy on trajectory-contraction targets, which accept only commitments that stay on a reference trajectory, using both confidence and hidden states. On LLaDA-8B-Instruct, this policy, with eight hundred and fifty-six thousand parameters, achieves the best frontier of accuracy against NFE and the highest throughput on GSM8K.

The wins are not uniform. On MATH-500 and HumanEval the policy is competitive at moderate targets rather than best. The overhead threshold is stated across the tested settings, not in general. And the transfer to Dream-7B-Instruct is described as largely holding, not complete.

What the controlled study buys is attribution. Instead of one more policy with a bundle of changes, it says which label to use, which inputs to give, and how small the speed-up can be before a learned policy stops paying for itself.