Improving Discrete Diffusion Unmasking Policies Beyond Explicit Reference Policies
This work addresses the sensitivity of unmasking order in diffusion models for language generation, offering a learned approach that outperforms heuristic methods, though it is incremental in nature.
The paper tackled the problem of improving masked diffusion models for language modeling by replacing rule-based unmasking schedules with a learned scheduler, resulting in a 20.1% gain over random and 11.2% gain over max-confidence on the SUDOKU benchmark.
Masked diffusion models (MDMs) have recently emerged as a novel framework for language modeling. MDMs generate sentences by iteratively denoising masked sequences, filling in [MASK] tokens step by step. Although MDMs support any-order sampling, performance is highly sensitive to the choice of which position to unmask next. Prior work typically relies on rule-based schedules (e.g., max-confidence, max-margin), which provide ad hoc improvements. In contrast, we replace these heuristics with a learned scheduler. Specifically, we cast denoising as a KL-regularized Markov decision process (MDP) with an explicit reference policy and optimize a regularized objective that admits policy improvement and convergence guarantees under standard assumptions. We prove that the optimized policy under this framework generates samples that more closely match the data distribution than heuristic schedules. Empirically, across four benchmarks, our learned policy consistently outperforms max-confidence: for example, on SUDOKU, where unmasking order is critical, it yields a 20.1% gain over random and a 11.2% gain over max-confidence.