NTU Trains 856K Unmasking Policy, Tops GSM8K on LLaDA-8B

A new controlled study from Nanyang Technological University Singapore by Jinya Sakurai, Xun Xu, and Jaehong Yoon lands on a very small network — 856K parameters — as the learned unmasking policy that wins the speed-accuracy frontier on a masked diffusion LLM.

Masked diffusion language models predict tokens in parallel, but efficiency depends on which predictions each step commits. The paper separates the design choices in prior learned policies — the teacher's labels, the inputs the policy sees, the states it imitates, and the NFE budget it has to recoup. Two findings land flat. "Labeling predictions matching the final output is overly permissive at low NFE, and hidden states are weak alone but can complement confidence," the authors report. From there they adopt "trajectory-contraction targets, which accept only commitments staying on a reference trajectory," and feed the policy both confidence and hidden states.

On LLaDA-8B-Instruct, the 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." Across the tested settings, policies up to 7.5M parameters need only about a 10% NFE reduction to offset their overhead. The abstract publishes no per-benchmark numbers. The work arrives alongside HKUST's trace of 'Wait' tokens in reasoning models back to data, part of a run of controlled-study papers we tracked this week.