用学习的调度策略优化文本去掩码生成，提升模型性能
Improving Discrete Diffusion Unmasking Policies Beyond Explicit Reference Policies
- 将去噪过程建模为带参考策略的强化学习问题，实现可优化的调度
- 在4个数据集上超越传统方法，苏格拉底任务提升11.2%
- 适合研究生成模型调度机制或改进语言建模的开发者
掩码扩散模型（MDMs）是近期出现的一种新型语言建模范式，通过逐步去噪掩码序列来生成句子。尽管支持任意顺序采样，但生成性能对下一步解码位置的选择极为敏感。以往工作依赖规则调度（如最大置信度、最大间隔），仅提供经验性改进。本文提出用学习的调度器替代这些启发式方法，将去噪建模为带有显式参考策略的KL正则化马尔可夫决策过程，在标准假设下保证策略优化与收敛。理论证明该框架生成样本更接近真实数据分布。实验表明，在4个基准测试中，所提方法持续优于最大置信度策略；例如在苏格拉底任务中，相比随机策略提升20.1%，相比最大置信度提升11.2%。代码已开源。
原文摘要 · Abstract (English)
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. Code is available at https://github.com/chunsanHong/UPO.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。