论文

$x$-预测流:掩码扩散语言模型的高效连续解码

$x$-Prediction Flow: Efficient Continuous Decoding for Masked Diffusion Language Models

模型推理解码与生成控制

摘要

屏蔽扩散语言模型(MDLM)通过迭代地揭开词元来生成文本,但它们的标准解码器将每个步骤简化为二元操作:一个位置要么致力于单个词元,要么完全屏蔽,丢弃丰富的预测信息而不是继续下去,并强制过早的、不可撤销的承诺,从而导致在有限的解码预算下性能不佳。在本文中,我们将掩模预测重新解释为干净状态预测($x$-预测),并表明它可以用于在输入嵌入空间中诱导连续流。基于这个观点,我们提出了一个 MDLM 的连续解码框架,其中词元可以在每个扩散步骤中累积部分进度并保持可修改。为了匹配语言中不同位置的不均匀上下文约束,我们用基于置信度的异步更新替换图像扩散中的全局同步计划,其中扩散进度是按标记累积的。此外,我们引入了一个轻量级策略网络,并将其训练制定为强化学习问题。应用于预训练的 LLaDA,我们的解码器使用不到 15% 的扩散步骤保留了 83--97% 的全预算精度,在匹配预算下大大优于离散掩模预测解码。