通过重掩码修复扩散语言模型中的不一致生成问题
Remask, Don't Replace: Token-to-Mask Refinement in Diffusion Language Models

- 用当前词概率识别低置信度位置,重掩码后在后续步骤重构
- 重掩码比直接替换更有效,提升多任务性能
- 无需训练,可跨模型通用,适合追求生成质量的场景
扩散语言模型(dLLMs)通过迭代去噪生成文本,每步填充多个掩码位置。由于同一步骤内各位置预测互不依赖新填值,易产生相互矛盾;一旦保留,这些矛盾会成为后续预测的上下文。我们提出无需训练的推理阶段修正方法Token-to-Mask(T2M),利用模型对当前词的概率识别低置信度位置,重新掩码并延迟重构。在配备修正机制的dLLMs上,单一T2M配置无需微调即可跨任务、跨模型通用,显著优于各模型原生修正机制。控制实验表明,修正方法可分解为检测器与修正动作:固定检测器下,重掩码优于替换;以当前词概率为检测器、重掩码为动作组合表现最佳。相比直接编辑,T2M更高效地将额外推理计算转化为性能提升,在多数任务中仍保持序列步骤优势,优于自回归逐词解码。
原文摘要 · Abstract (English)
Diffusion language models (dLLMs) generate text through iterative denoising, filling multiple masked positions at each step. Positions filled in the same step are predicted without conditioning on one another's newly filled values and can therefore be mutually inconsistent; once retained, these inconsistencies become context for later predictions. We introduce \emph{Token-to-Mask} (T2M), a training-free inference-time correction method that identifies low-confidence positions using the model's probability of the current token, remasks them, and reconstructs them in later denoising steps. On dLLMs equipped with correction mechanisms, a single T2M configuration transfers across tasks and models without retuning and broadly improves task metrics over each model's native correction mechanism. In controlled experiments, we decompose correction methods into a detector that identifies suspicious tokens and an action that determines how to revise them. Holding the detector fixed, remasking yields higher task metrics than replacement; across the tested detector--action combinations, current-token-probability detection paired with remasking performs best. Compared with direct editing, T2M converts additional inference compute into performance gains more effectively and, on most tasks, retains a sequential-step advantage over autoregressive token-by-token decoding.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。