技术摘要：统一具有多种生成顺序的掩码扩散模型及超越

问题陈述

掩码扩散模型（MDMs）已成为自回归模型（ARMs）在语言生成领域的一个有前景的替代方案。然而，MDMs 的生成质量 critically 取决于生成顺序（即 token 被取消掩码的序列）。现有方法面临显著局限：

1. 固定顺序： 经典 MDMs 通常使用均匀随机掩码调度，导致对文本而言次优的随机生成顺序。
2. 硬编码调度器： 如 Block Diffusion (BD3LM) 等方法通过手工设计的调度器强加特定结构（例如分块从左到右），限制了灵活性。
3. 两阶段优化： 学习取消掩码策略的方法（例如通过后训练采样器）会产生额外的计算成本，并且由于扩散骨干网络与排序策略的解耦优化，往往收敛至次优解。
4. 缺乏统一性： 目前没有一个单一的理论框架能够表达多样化的范式（ARMs、标准 MDMs、分块扩散和词汇依赖型调度器），同时允许联合学习生成顺序和扩散骨干网络。

方法论

本文提出了一种统一框架和特定的模型架构来解决这些问题。

1. 顺序表达型掩码扩散模型 (OeMDM)

作者引入了OeMDM，这是一个广义框架，将噪声调度器视为一个可学习的、位置相关的组件，而非固定函数。

• 广义调度器： OeMDM 定义了一个“自由形式调度器”类 $\alpha_F$，它可以依赖于输入上下文（例如完整序列 $x$ 或当前状态 $z_t$）。这使得前向破坏过程能够对不同位置应用不同的噪声率，从而建立明确的生成优先级。
• 广义 NELBO： 作者为 OeMDM 推导了广义负证据下界（NELBO）。该目标分解为两项：
  • 重构损失 ($L_{main}$)： 标准的 token 预测损失，由“速度”$A$（取消掩码的速率）加权。
  • 速度失配损失 ($L_{velocity}$)： 量化真实前向后验速度 ($A$) 与参数化反向过程速度 ($\hat{A}$) 之间差距的项。
• 统一性： 该框架从理论上证明，标准 MDMs、ARMs、分块扩散 (BD3LM) 和词汇依赖型调度器 (GenMD4) 均为 OeMDM 的特例，通过特定选择调度器函数即可实现。

2. 可学习顺序掩码扩散模型 (LoMDM)

基于 OeMDM，作者提出了LoMDM，它通过单一目标从头开始联合学习生成顺序和扩散骨干网络。

• 上下文感知调度： 与之前学习受限调度器或后训练采样器的方法不同，LoMDM 参数化了一个基于完整序列上下文的、位置相关的调度器。
• 架构：
  • 扩散骨干网络 $\theta$ 充当特征提取器。
  • 两个轻量级网络，$\phi$（前向调度器）和 $\psi$（反向调度器），附加在骨干网络上。它们根据上下文预测每个 token 的生成优先级（速度）。
  • 调度器使用归一化的 sigmoid 激活函数，以确保速度保持在有效范围内，并正则化相对生成优先级。
• 训练算法：
  • 训练目标是最小化单一 NELBO。
  • 为了处理调度器的梯度（涉及从前向过程采样），作者采用了REINFORCE Leave-One-Out (RLOO) 估计器。与朴素 REINFORCE 相比，这降低了梯度估计的方差，实现了骨干网络与调度器的稳定联合训练。
  • 模型使用“停止梯度”机制，在计算调度器损失时分离骨干网络特征，在调度器更新步骤中将骨干网络视为固定的特征提取器。

主要贡献

1. 统一理论框架 (OeMDM)： 本文提供了一个广义的 NELBO 公式，将各种生成顺序（随机、从左到右、分块和词汇依赖型）统一在单一的掩码扩散框架下。它证明了 ARMs 可以被视为具有特定调度器的 OeMDM 的极限情况。
2. 顺序与骨干网络的联合学习 (LoMDM)： 作者提出了第一种从头开始使用单一目标联合学习生成顺序和扩散模型的方法，避免了以往工作中两阶段优化的陷阱。
3. 上下文依赖的生成： LoMDM 学习到一个能适应句子特定上下文的调度器，使模型能够动态决定“接下来生成什么”和“在哪里生成”，而不是遵循僵化或随机的顺序。
4. 高效训练机制： RLOO 估计器的集成使得调度器参数能够与扩散骨干网络一起稳定训练，克服了学习离散采样策略通常伴随的高方差问题。

实验结果

作者在标准语言建模基准（LM1B 和 OpenWebText）上评估了 LoMDM，对比了各种离散扩散基线（包括 MDLM、BD3LM 和 GenMD4）以及自回归 Transformer。

• 困惑度 (PPL)： LoMDM 在离散扩散模型中实现了最先进的测试困惑度。在 LM1B 和 OpenWebText 上，它显著优于 MDLM（例如 PPL 降低超过 3 点）。它也超越了 BD3LM（即使使用较小的块大小）和 GenMD4。
• 零样本泛化： LoMDM 在 7 个多样化数据集（包括 PTB、WikiText 和 Lambada）上表现出卓越的零样本泛化能力，在所有 7 个数据集上均优于 MDLM，甚至在其中的 4 个数据集上击败了自回归 Transformer。
• 生成质量： 就生成困惑度（由预训练的 GPT-2 Large 评估）而言，LoMDM 显著优于 MDLM 基线，表明学习到的速度 $\hat{A}$ 提高了实际文本生成质量，而不仅仅是似然下界。
• 训练效率： LoMDM 被证明是一个“快速且高效的学习者”。它仅需约 0.18M 步（大约看到 18% 的 token）即可达到 1M 步 MDLM 训练运行的性能，将墙钟训练时间减少了约 3 倍。
• 生成轨迹分析： 对生成过程的分析显示，LoMDM 采用了一种“由粗到细”的策略：它在解码过程的早期优先生成结构词（冠词、介词），而在后期填充语义内容（名词、动词）。这与倾向于早期生成语义内容的 MDLM 形成对比，也不同于 ARMs 严格的从左到右偏差。

意义与主张

本文声称，LoMDM 在离散扩散语言建模方面代表了重大进步，具体体现在：

• 解决顺序瓶颈： 它通过将顺序变为可学习的、上下文感知的组件，而非固定的超参数或事后添加项，解决了 MDM 质量对生成顺序的关键依赖问题。
• 理论统一： 它提供了一个单一的理论视角（OeMDM）来理解和比较不同的生成范式（自回归、分块扩散、标准扩散），阐明了它们之间的关系和局限性。
• 实际效率： 通过联合优化顺序和模型，LoMDM 实现了比需要单独训练阶段或固定调度器的方法更高的样本质量和更快的收敛速度。

作者总结道，他们的方法使扩散模型能够以比随机或僵化的从左到右顺序更有效的上下文依赖顺序生成文本，特别是在无条件生成任务中。他们指出，虽然当前结果集中在无条件生成上，但学习到的调度器的适应性表明其在条件任务中具有潜力，尽管这仍是未来研究的领域。