What Makes a Good Unmasking Policy? A Controlled Study of Learned Unmasking for Masked Diffusion LMs

Forum topic · 小凯 · 2026-10-25

Summary

A 2026 arXiv paper (2610.31576) by Jinya Sakurai, Xun Xu, and Jaehong Yoon studies how to decide which token predictions to commit when generating text with masked diffusion language models. The authors show that labeling predictions as committable simply because they match the final output is overly permissive at low numbers of function evaluations (NFE). Hidden-state features are weak on their own but complement confidence-based signals. Their analysis finds that roughly a 10% reduction in NFE is enough to offset the inference overhead of policies up to 7.5M parameters. Training an 856K-parameter policy on trajectory-contraction targets using both confidence and hidden states yields the best accuracy-NFE frontier and highest throughput on GSM8K with LLaDA-8B-Instruct, and the policy largely transfers to Dream-7B-Instruct without retraining.

Paper Overview

Field: Machine Learning
Authors: Jinya Sakurai, Xun Xu, Jaehong Yoon
Published: 2026-10-22
arXiv: 2610.31576

Abstract

Masked diffusion language models generate tokens in parallel, but efficiency depends on which predictions are committed. The authors find that labeling predictions matching the final output is overly permissive at low NFE (number of function evaluations), and hidden states are weak alone but can complement confidence-based signals. About a 10% NFE reduction suffices to offset policy overhead for policies up to 7.5M parameters.

Method and Results

The team trains a learned unmasking policy on trajectory-contraction targets, using both token confidence and hidden states.
The resulting 856K-parameter policy achieves the best accuracy–NFE frontier and the highest throughput on GSM8K with LLaDA-8B-Instruct.
The policy largely transfers to Dream-7B-Instruct without task-specific retraining, demonstrating generalization across masked diffusion models.

Takeaway

Carefully designed lightweight unmasking policies can meaningfully improve the speed–accuracy trade-off of masked diffusion language models, with overhead costs that are easily repaid by modest NFE savings.

---

*Auto-collected on 2026-10-26*