Study isolates what improves learned unmasking in diffusion language models
Original title: What Makes a Good Unmasking Policy? A Controlled Study of Learned Unmasking in Diffusion Language Models
An 856K-parameter policy on LLaDA-8B-Instruct leads the accuracy–NFE frontier on GSM8K and largely transfers to Dream-7B-Instruct.
The paper studies supervised unmasking policies for masked diffusion language models as distillation of a decoding procedure. It varies teacher labels, imitation states, observed inputs, and how much the number of function evaluations must fall to offset policy cost. Labeling any prediction that matches the final output is overly permissive at low NFE, while hidden states help mainly alongside confidence. A trajectory-contraction policy using both signals is competitive on MATH-500 and HumanEval at moderate targets.
- About a 10% NFE cut offsets overhead for policies up to 7.5M parameters.
- The selected policy has 856K parameters.
- It uses confidence and hidden states with trajectory-contraction targets.
- Tested models are LLaDA-8B-Instruct and Dream-7B-Instruct.
- AI models
- LLaDA-8B-InstructDream-7B-Instruct
From the original article
1 min read at huggingface.co
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.
By Jinya Sakurai, Xun Xu, Jaehong Yoon · huggingface.co
A short excerpt; the article and its images belong to the publisher. Summary, key facts and assessment above are AI-written from the article and can be wrong.