变分掩码扩散模型：优化并发标记预测的依赖关系处理

变分掩码扩散(VMD)模型通过引入潜在变量，成功解决了标准掩码扩散在并发标记预测中无法有效捕捉依赖关系的问题。在合成数据、数独谜题和文本数据上的实验都证明了VMD的有效性，特别是在标记间依赖关系重要的场景中。

1. 研究背景与问题定义

基于扩散的大型语言模型(DLLMs)作为自回归模型(ARMs)的重要扩展，正在成为生成式AI领域的重要创新方向。与传统ARMs按预定顺序逐一生成标记的方式不同，DLLMs提供了并发标记生成、更高输出多样性、增强全局一致性以及更好的生成文本可控性等优势。近期的突破性模型如LLaDA、Mercury和Gemini Diffusion都凸显了DLLMs的潜力。

然而，当前的掩码扩散模型(MDM)存在一个关键限制：无法有效捕捉并发预测的标记之间的依赖关系。这导致在标记间依赖性较强的推理任务中性能下降。例如，在预测"A poker hand that consists of two English words is: a a"的后续两个词时，合适的预测应为"high card"、"two pair"、"full house"或"straight flush"。这些词对之间存在强依赖关系，但MDM在并发预测时会独立采样，无法考虑这种依赖性，从而可能产生不合理的组合。

2. 变分掩码扩散(VMD)模型

为解决上述问题，研究者提出了变分掩码扩散(Variational Masked Diffusion, VMD)框架，通过引入潜在变量来建模并发预测期间的联合标记分布。VMD的核心思想是：通过潜在变量模型捕捉任意联合分布，而不仅仅是可因式分解的分布。

图片

2.1 基本变分公式

VMD的基本公式为：

pθ(x0i|xt) = ∫pθ(x0i|xt,z)p(z)dz

其中z是全局潜在变量，不依赖于标记位置i。这使得模型能够在标记之间建立联合分布。条件于潜在变量z，标记可以独立采样，但边缘化潜在变量后，能够从正确的联合分布中获得样本。

训练目标函数(NELBO)为：

-log pθ(x0) ≤ ∫(0→1) (1/t)