Representation-based Masked Diffusion Model

# 基于表征的掩码扩散模型
来源：https://arxiv.org/html/2609.12382

**作者**  
Yangrong Hu（单位：数据科学与人工智能系；邮箱：[yangrong\.hu@connect\.polyu\.hk](mailto:[email protected])）  
Ding Huang（单位：应用数学系，香港理工大学，中国香港特别行政区；邮箱：[ding\.huang@connect\.polyu\.hk](mailto:[email protected])）  
Xueyu Zhou（单位：数据科学与人工智能系；邮箱：[xueyu\.zhou@connect\.polyu\.hk](mailto:[email protected])）  
Jian Huang（单位：数据科学与人工智能系、应用数学系，香港理工大学，中国香港特别行政区；邮箱：[j\.huang@polyu\.edu\.hk](mailto:[email protected])）

###### 摘要
掩码扩散模型（MDMs）已成为语言建模中一种引人注目的范式，能够实现高效的并行文本生成。然而，现有的并行采样方法通常独立地更新多个掩码标记，并忽略了掩码标记之间复杂的相互依赖关系。这种独立更新机制缺乏全局协调，可能导致输出不连贯。为解决这一局限性，我们提出了**基于表征的掩码扩散模型（RMDM）**，该框架利用文本表征来显式编码全局语义，并帮助更精确地并行更新标记。具体而言，我们首先使用预训练的编码器将文本编码到连续的语义空间，并学习一个可逆变换，将表征分布规范化为高斯先验，以便在生成过程中进行高效采样。基于此潜在语义表征，我们训练一个掩码扩散模型来学习条件文本分布，其中该表征作为全局语义引导，协调并行标记更新并忠实地近似目标分布。实证结果表明，RMDM显著提高了生成质量，尤其是在激进的少步采样模式下。

††脚注：通讯作者：Jian Huang。  
‡‡脚注：现地址：字节跳动 Seed，北京，中国。邮箱：huanhgding@bytedance\.com\.

**第一阶段：表征映射**  
**第二阶段：生成建模**

编码器 \(g_{\psi}(\bm{x}_{0})\)  → \(\bm{z}_{1}\)  
MeanFlow \(u_{\phi}\) → \(\bm{z}_{0} \sim \mathcal{N}(\bm{0},\bm{I})\)  
掩码 \(\bm{x}_{t}\)  
RMDM \(p_{\theta}(\bm{x}_{0}|\bm{x}_{t},\bm{z}_{0})\) → \(\hat{\bm{x}}_{0}\)

**图 1：** 第一阶段将从预训练编码器 \(g_{\psi}(\bm{x}_{0})\) 得到的潜在空间 \(\bm{z}_{0}\) 与通过 MeanFlow \(u_{\phi}\) 引入的高斯先验对齐。第二阶段训练 RMDM \(p_{\theta}\) 以在潜在表征 \(\bm{z}_{0}\) 的引导下，从受损数据 \(\bm{x}_{t}\) 重构原始数据 \(\bm{x}_{0}\)。

## 1 引言

大型语言模型（LLMs）在建模文本数据方面取得了巨大成功（Brown 等人，2020 (https://arxiv.org/html/2609.12382#bib.bib1)；Radford 等人，2019 (https://arxiv.org/html/2609.12382#bib.bib35)；Achiam 等人，2023 (https://arxiv.org/html/2609.12382#bib.bib24)；Liu 等人，2024 (https://arxiv.org/html/2609.12382#bib.bib25)），主要通过自回归（AR）范式。虽然 AR 模型擅长通过顺序生成标记来捕获复杂依赖关系，但其推理速度本质上受限于长度为 \(L\) 的序列所需的 \(O(L)\) 顺序复杂度。这种延迟瓶颈激发了人们对非自回归或并行生成框架的极大兴趣（Gu 等人，2018 (https://arxiv.org/html/2609.12382#bib.bib45)）。

掩码扩散模型（MDMs）（Austin 等人，2021 (https://arxiv.org/html/2609.12382#bib.bib5)；Sahoo 等人，2024 (https://arxiv.org/html/2609.12382#bib.bib13)；Lou 等人，2024 (https://arxiv.org/html/2609.12382#bib.bib40)；Ou 等人，2025 (https://arxiv.org/html/2609.12382#bib.bib39)）最近作为一种引人注目的替代方案出现，它通过迭代去噪实现并行标记生成。与连续扩散模型相比，MDMs 使用一种*离散扩散*过程，其中前向过程通过将干净序列的子集替换为特殊的 \([MASK]\) 标记（而非高斯噪声）来逐渐添加噪声，并生成一个部分观测的上下文。然后训练一个去噪器来反转此损坏过程，从上下文中预测原始标记。在反向过程中，MDMs 利用去噪器通过反复优化一个完全由掩码标记组成的句子来生成整个序列。

在实际部署中，要求 MDMs 在单次迭代中高效地更新多个掩码标记。大多数并行方法（Nie 等人，2025 (https://arxiv.org/html/2609.12382#bib.bib11)；Zhu 等人，2025 (https://arxiv.org/html/2609.12382#bib.bib10)；Ye 等人，2025 (https://arxiv.org/html/2609.12382#bib.bib12)；Chang 等人，2022 (https://arxiv.org/html/2609.12382#bib.bib49)；Wu 等人，2025 (https://arxiv.org/html/2609.12382#bib.bib17)）在给定当前观测上下文的情况下，独立地对每个掩码位置进行采样，同时忽略了掩码标记之间的依赖性。这种独立采样将掩码位置上的联合分布近似为标记级边际分布的乘积。然而，这种分解近似通常与自然语言不一致，因为在自然语言中，掩码标记由于语法一致性、长程约束、多词表达和实体一致性等原因，常常表现出强烈的残余依赖性。因此，激进的并行化可能会产生在每个位置局部看似合理，但在真实联合分布下全局不一致的序列。

我们在第 3.1 节 (https://arxiv.org/html/2609.12382#S3.SS1) 形式化并量化了这种**条件依赖差距**，并在不牺牲并行性和效率的前提下，围绕缓解该差距构建了我们的方法。

##### 我们的方法。为了缓解这种不一致性，我们引入了一个全局潜在表征 \(\bm{z}\)，它编码了整个序列的语义信息。该表征捕获了共享的全局因素，如主题和意图，这些因素在很大程度上解释了掩码标记之间的残余依赖。因此，以 \(\bm{z}\) 为条件减少了残余的跨标记依赖，并使并行预测更接近真实的联合条件分布。我们在第 3.1 节 (https://arxiv.org/html/2609.12382#S3.SS1) 将这种直觉形式化为*条件依赖差距*的减少。

为了使潜在变量 \(\bm{z}\) 具有可操作性，我们使用预训练的编码器，如 BERT 模型（Devlin 等人，2019 (https://arxiv.org/html/2609.12382#bib.bib36)；Warner 等人，2024 (https://arxiv.org/html/2609.12382#bib.bib51)）和现代大规模嵌入模型（Zhang 等人，2025 (https://arxiv.org/html/2609.12382#bib.bib57)），来提取输入文本的上下文表征。该表征编码了底层序列的全局属性，并捕获了难以从部分揭示的上下文中推断出的高层语义因素。通过提供这种缺失的全局信息，该表征成为 \(\bm{z}\) 的自然候选者，并有助于在并行更新期间减少条件依赖差距。

一个关键的复杂性在于，这些表征在训练时可以从真实文本中轻易获得，但在推理时却是未观测的。直接以这种嵌入为条件会导致训练-测试不匹配，除非在生成过程中可以从一个定义良好的先验中采样 \(\bm{z}\)。为解决训练-推理不匹配问题，我们采用 MeanFlow（Geng 等人，2025 (https://arxiv.org/html/2609.12382#bib.bib3)）将预训练序列表征的经验分布映射到一个高斯先验。这种传输使我们能够在生成时直接采样一个全局潜在变量 \(\bm{z} \sim \mathcal{N}(0,I)\)，它作为 RMDM 的全局条件，以确保连贯的并行更新。

尽管变分推断（Kingma 和 Welling，2014 (https://arxiv.org/html/2609.12382#bib.bib38)；Xie 等人，2025 (https://arxiv.org/html/2609.12382#bib.bib6)）是处理高斯潜在变量的经典方法，但其重构和 KL 目标在高容量语言模型中常常冲突，导致优化不稳定或后验坍缩。诸如 CCDD（Zhou 等人，2025 (https://arxiv.org/html/2609.12382#bib.bib7)）等方法协同演化连续和离散扩散，而 RMDM 则采样一个固定的潜在变量并在整个离散去噪过程中重复使用它。预训练编码器和 MeanFlow 网络用于构建训练潜在变量，但在生成过程中不被调用。

实验结果表明，对于固定的采样步数，RMDM 在样本质量上始终优于标准 MDMs，并且在我们的 A800 评估中，在匹配的生成质量下，测量到的吞吐量提高了约 \(3.6\times\)。我们在附录 A (https://arxiv.org/html/2609.12382#A1) 中更详细地讨论了相关的离散扩散语言模型、连续扩散语言模型和潜在增强掩码扩散方法。

## 2 预备知识

### 2.1 连续扩散与 MeanFlow

连续扩散模型（Ho 等人，2020 (https://arxiv.org/html/2609.12382#bib.bib4)；Song 等人，2021 (https://arxiv.org/html/2609.12382#bib.bib2)）建立在这样一个原则之上：前向过程在连续时间 \(t \in [0,1]\) 上逐渐将数据 \(\bm{z}_{0} \in \mathbb{R}^{d}\) 损坏为高斯噪声 \(\bm{z}_{1} \sim \mathcal{N}(\bm{0},\bm{I})\)。生成过程则通过学习反转这种损坏来定义——本质上是“去噪”潜在变量——通过遵循一个引导噪声回到数据流形的速度场。

流匹配（Lipman 等人，2023 (https://arxiv.org/html/2609.12382#bib.bib8)）通过训练模型 \(\bm{v}_{\theta}\) 来匹配条件速度场，从而简化了这一过程：
\[
\mathcal{L}_{\text{CFM}}(\theta)=\mathbb{E}_{t,\bm{z}_{0},\bm{z}_{1}}\left[\|\bm{v}_{\theta}(\bm{z}_{t},t)-(\bm{z}_{1}-\bm{z}_{0})\|^{2}\right],
\]
其中 \(\bm{z}_{t}=(1-t)\bm{z}_{0}+t\bm{z}_{1}\) 表示数据和噪声之间的线性插值。

虽然标准的流匹配需要多步数值 ODE 积分，但 MeanFlow（Geng 等人，2025 (https://arxiv.org/html/2609.12382#bib.bib3)）通过建模平均速度 \(u(\bm{z}_{t},r,t)\triangleq\frac{1}{t-r}\int_{r}^{t}\bm{v}(\bm{z}_{\tau},\tau)d\tau\)，实现了高效的单步生成。模型 \(u_{\theta}\) 通过 MeanFlow 恒等式进行训练：
\[
\mathcal{L}_{\text{MF}}(\theta)=\mathbb{E}\left\|u_{\theta}(\bm{z}_{t},r,t)-\text{sg}\left[\bm{v}_{t}-(t-r)\frac{d}{dt}u_{\theta}(\bm{z}_{t},r,t)\right]\right\|^{2}
\]
其中 \(\text{sg}\) 是停止梯度算子。一旦优化完成，\(\bm{z}_{0}\) 可以在一步内从噪声 \(\bm{z}_{1}\) 恢复：\(\bm{z}_{0}=\bm{z}_{1}-u_{\theta}(\bm{z}_{1},0,1)\)。

### 2.2 掩码扩散模型

虽然连续扩散适合于连续分布，但语言建模需要一种处理离散分类变量的合理方法。掩码扩散模型（MDMs）（Austin 等人，2021 (https://arxiv.org/html/2609.12382#bib.bib5)；Lou 等人，2024 (https://arxiv.org/html/2609.12382#bib.bib40)）将扩散框架扩展到离散状态空间，定义了一个损坏过程，该过程逐渐将标记 \(\bm{x}_{0}\) 替换为吸收性的 \([MASK]\)，而完全由 \([MASK]\) 组成的序列表示为状态 \(\bm{m}\)。与连续情况类似，前向过程 \(q(\bm{x}_{t}|\bm{x}_{0})=\text{Cat}(\bm{x}_{t};\alpha_{t}\bm{x}_{0}+(1-\alpha_{t})\bm{m})\) 在干净数据和掩码先验之间插值，其中 \(\alpha_{t}\in[0,1]\) 是递减的噪声调度。

生成模型 \(p_{\theta}(\bm{x}_{s}|\bm{x}_{t})\) 通过近似封闭形式的后验来学习反转这个掩码过程：
\[
q(\bm{x}_{s}|\bm{x}_{t},\bm{x}_{0})=
\begin{cases}
\delta_{\bm{x}_{t}} & \text{if } \bm{x}_{t}\neq\bm{m},\\
\text{Cat}\left(\bm{x}_{s};\frac{\alpha_{s}-\alpha_{t}}{1-\alpha_{t}}\bm{x}_{0}+\frac{1-\alpha_{s}}{1-\alpha_{t}}\bm{m}\right) & \text{if } \bm{x}_{t}=\bm{m}
\end{cases}
\]
通过训练去噪模型 \(p_{\theta}(\bm{x}_{0}|\bm{x}_{t})\)，MDMs 可以通过一个源自数据对数似然 \(\log p_{\theta}(\bm{x}_{0})\) 的目标进行优化：
\[
\mathcal{L}(\theta)=\int_{0}^{1}\frac{1}{t}\mathbb{E}_{q(\bm{x}_{t}|\bm{x}_{0})}\left[\sum_{i\in\mathcal{I}_{t}}-\log p_{\theta}(\bm{x}_{0}^{i}|\bm{x}_{t})\right]dt\quad(1)
\]
其中 \(\mathcal{I}_{t}=\{i:\bm{x}_{t}^{i}=\texttt{[MASK]}\}\)。该损失使模型能够预测每个掩码位置上原始标记的分类分布，为上述连续去噪流提供了离散类比。

## 3 方法论

### 3.1 条件依赖差距

在本节中，我们记 \(\bm{x}_{M}\) 为掩码数据，\(\bm{x}_{U}\) 为在扩散过程中时间 \(t\) 的未掩码数据。掩码扩散模型（MDMs）的一个基本挑战源于其训练目标与效率所需的并行采样策略之间的差异。虽然标准目标方程（1）(https://arxiv.org/html/2609.12382#S2.E1) 优化的是单个标记 \(p_{\theta}(x_{i}|\bm{x}_{U})\) 的独立重构，但快速生成通常需要同时采样所有掩码标记 \(\bm{x}_{M}\)。现有的并行采样器隐式地通过因子化近似联合后验：
\[
p_{\theta}(\bm{x}_{M}|\bm{x}_{U})\triangleq\prod_{i\in M}p_{\theta}(x_{i}|\bm{x}_{U})\approx\prod_{i\in M}p_{\text{data}}(x_{i}|\bm{x}_{U})\neq p_{\text{data}}(\bm{x}_{M}|\bm{x}_{U}),
\]
其中 \(M\) 和 \(U\) 分别表示掩码和未掩码索引的集合。这种因子化假设掩码标记在给定上下文 \(\bm{x}_{U}\) 的条件下是独立的。然而，对于结构化数据，此假设常常被违反，导致**条件依赖差距**。我们使用条件总相关性来量化这个差距：
\[
\mathcal{T}(\bm{x}_{M}|\bm{x}_{U})\triangleq D_{\mathrm{KL}}\!\left(p_{\text{data}}(\bm{x}_{M}|\bm{x}_{U})\,\Big\|\,\prod_{i\in M}p_{\text{data}}(x_{i}|\bm{x}_{U})\right).
\]

**图 2：因果图。** 全局潜在变量 \(\bm{z}\) 作为掩码标记的共同原因。以 \(\bm{z}\) 为条件消除了 \(\bm{x}_{M}^{1}\) 和 \(\bm{x}_{M}^{2}\) 之间的相关性。