![[2510.05725] Improving Discrete Diffusion Unmasking Policies Beyond Explicit Reference Policies](https://arxiv.org/static/browse/0.3.4/images/arxiv-logo-fb.png){kind=logo}
[2510.05725] Improving Discrete Diffusion Unmasking Policies Beyond Explicit Reference Policies
Summary
This article presents a novel approach to improving masked diffusion models (MDMs) for language modeling by introducing a learned scheduler that outperforms traditional heuristic methods.
Why It Matters
The study addresses the limitations of existing unmasking strategies in MDMs, which are crucial for generating coherent language. By optimizing the unmasking process through a learned policy, the research contributes to advancements in natural language processing, enhancing model performance and applicability.
Key Takeaways
- Introduces a learned scheduler for masked diffusion models (MDMs).
- Demonstrates significant performance improvements over heuristic methods.
- Proves that optimized policies yield samples closer to the data distribution.
- Empirical results show a 20.1% gain in specific benchmarks.
- Provides code access for further experimentation and validation.
Computer Science > Machine Learning arXiv:2510.05725 (cs) [Submitted on 7 Oct 2025 (v1), last revised 26 Feb 2026 (this version, v2)] Title:Improving Discrete Diffusion Unmasking Policies Beyond Explicit Reference Policies Authors:Chunsan Hong, Seonho An, Min-Soo Kim, Jong Chul Ye View a PDF of the paper titled Improving Discrete Diffusion Unmasking Policies Beyond Explicit Reference Policies, by Chunsan Hong and 3 other authors View PDF HTML (experimental) Abstract: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...
