For faster parallel token generation in diffusion LMs, a small learned policy (856K params) using confidence + hidden states can reliably decide which predictions to commit, improving throughput without sacrificing accuracy—and the policy transfers across different base models.
This paper studies how to decide which token predictions to commit to in parallel-decoding diffusion language models. The authors systematically analyze what makes a good unmasking policy—a learned model that decides when predictions are confident enough to finalize—by testing different supervision signals, inputs, and model sizes.