DPRM：一种用于扩散语言模型的Doob h变换诱导的Token排序模块
机器学习
2026-04-28 v1 人工智能
摘要
扩散语言模型无固定的从左到右顺序生成，这使得token排序成为核心算法选择：在每一步中，应该揭示、保留、修订或验证哪些token？现有系统主要使用随机遮蔽或基于置信度的排序。随机遮蔽导致训练-测试不匹配，而仅依赖置信度的规则虽高效，但可能盲目并抑制有用的探索。我们引入DPRM（Doob h变换过程奖励模型），作为扩散语言模型的插件token排序模块。DPRM保持主机架构、去噪目标和监督不变，仅改变排序策略。它从基于置信度的渐进式排序开始，逐渐转向Doob h变换过程奖励引导的排序。我们刻画了精确的DPRM策略作为奖励倾斜的Gibbs揭示法，证明了阶段性Soft-BoN近似的O(1/N)收敛，并展示了在经验-伯恩斯坦速率下，线上桶式控制器跟踪精确DPRM得分。在可行的优化假设下，DPRM还在随机和基于置信度的排序方面具有样本复杂度优势。DPRM在预训练、后训练、测试时扩展以及单细胞遮蔽扩散中均优于基于置信度的基线，在更难的推理子集上表现尤为突出。在蛋白质、分子生成和DNA设计中，效果更为多目标：排序感知的变体在所选结构或片段受限指标上显著改善，但并不在所有质量指标上都统一优于主机基线。这些结果表明，token排序是扩散语言模型中的一个基本控制轴，DPRM作为改进它的通用模块。代码可在https://github.com/DakeBU/DPRM-DLLM访问。
引用
@article{arxiv.2604.24357,
title = {DPRM: A Plug-in Doob h transform-induced Token-Ordering Module for Diffusion Language Models},
author = {Dake Bu and Wei Huang and Andi Han and Hau-San Wong and Qingfu Zhang and Taiji Suzuki and Atsushi Nitanda},
journal= {arXiv preprint arXiv:2604.24357},
year = {2026}
}