掩码扩散解码作为$x$-预测流

# 掩码扩散解码作为 x 预测流
来源：https://arxiv.org/html/2606.29066
Weitian Wang1,2, Lianlei Shan3, Shubham Rai1, Cecilia De La Parra1, Akash Kumar2 1德国博世有限公司，2德国波鸿鲁尔大学，3中国科学院大学，中国

###### 摘要

掩码扩散语言模型（MDLM）通过逐步去掩码 token 来生成文本，但其标准解码器将每一步简化为一个二元动作：一个位置要么被承诺为一个单一 token，要么保持完全掩码状态，中间没有任何部分信念的表示。这种“全有或全无”的机制丢弃了丰富的预测信息，并迫使做出过早且不可撤销的承诺，导致在有限解码预算下性能不佳。在本文中，我们将掩码预测重新解释为干净状态预测（x 预测），并展示它可用于在输入嵌入空间中诱导一个连续流。基于这一观点，我们为 MDLM 提出一个连续解码框架，其中 token 可以在每个扩散步骤中累积部分进展，并且保持可修正性。为了匹配语言中不同位置间不均匀的上下文约束，我们用基于置信度的异步更新取代了图像扩散中的全局同步调度，其中扩散进展按 token 逐步累积。此外，我们引入一个轻量级策略网络，并将其训练公式化为一个强化学习问题。应用于预训练的 LLaDA，我们的连续解码器在 HumanEval 数据集上仅用 25% 的解码预算就达到了其 97% 的性能。

## 1 引言

参见图注
图 1：通过将掩码预测重新解释为嵌入空间中的干净状态预测，并定义一个速度场，将当前状态从掩码嵌入（噪声状态）移向预测的干净状态，我们的 x 预测流解码使*所有* token 在嵌入空间中*连续地*演化。

基于自回归（AR）范式构建的大型语言模型[1 (https://arxiv.org/html/2606.29066#bib.bib13);7 (https://arxiv.org/html/2606.29066#bib.bib14)]推动了自然语言处理的最新进展。然而，其严格的顺序分解在推理时施加了一个基本约束：token 必须一个接一个地生成，仅条件于先前生成的前缀。这种顺序依赖导致高解码延迟。这在现代推理模型中更加严重，其中思维链轨迹通常跨越数千个 token[9 (https://arxiv.org/html/2606.29066#bib.bib15)]。

扩散语言模型（DLM）最近作为一种有前景的替代方案出现[16 (https://arxiv.org/html/2606.29066#bib.bib1);3 (https://arxiv.org/html/2606.29066#bib.bib5)]。DLM 不是从左到右承诺 token，而是将生成建模为一个迭代去噪过程：一个完全损坏的响应在少量步骤中并行优化，每个位置都可以双向关注序列的其余部分。这种并行性使得生成更快，支持可控和非因果的生成模式，并提供了一个内置机制，用于根据后面的上下文优化早期的预测。

一个自然的设计是遵循图像扩散[11 (https://arxiv.org/html/2606.29066#bib.bib10);15 (https://arxiv.org/html/2606.29066#bib.bib9);14 (https://arxiv.org/html/2606.29066#bib.bib2)]，直接在连续嵌入空间中去噪。但这个方案不能直接迁移到语言。语言 token 是离散的且高度依赖上下文，一个 token 的含义由其邻居的身份决定，而不是由其在任何连续坐标中的位置决定。因此，向嵌入中注入随机噪声会产生一个扰动，其幅度在语义上与句法或词汇变化没有有意义的对应关系。因此，主流的 DLM 家族[17 (https://arxiv.org/html/2606.29066#bib.bib8);16 (https://arxiv.org/html/2606.29066#bib.bib1);3 (https://arxiv.org/html/2606.29066#bib.bib5)]采用了*掩码*扩散，基于 BERT[5 (https://arxiv.org/html/2606.29066#bib.bib4)]推广的掩码预测的长期成功，即一部分 token 被特殊的 [M] 符号损坏，并要求模型从周围上下文中恢复它们。这种形式化使得 MDLM 能够扩展到数十亿参数，并在标准基准上达到 AR 基线的水平，同时在受益于双向结构的任务（如代码生成和逆向推理）上也显示出优势[16 (https://arxiv.org/html/2606.29066#bib.bib1)]。

尽管具有这种可扩展性，MDLM 的标准解码过程对其计算预算的利用效率低下。在每一步，模型在每个掩码位置输出一个词汇表上的分类分布，但采样器将这个分布简化为一个二元动作：要么该位置被解除掩码为一个承诺的 token，要么保持为 [M] 并在下一步从头开始重新预测。从任何单个位置的角度来看，每步状态因此是“全有或全无”的，步骤之间没有部分信念的表示。这导致两个后果。首先，完整输出分布中包含的丰富预测信息，包括候选者及其相对似然度，一旦一个位置被保留就被丢弃了。其次，一旦一个 token 被解除掩码，承诺就是最终的，即使后来对相邻位置的更新本应支持不同的选择。模型被迫在过早的确定性和完全不确定性之间做出选择，这限制了固定解码步数的有效利用。这一局限性促使设计一种解码方案，其中对每个 token 的信念可以跨步骤*连续地*演化，使得置信度随着上下文逐步积累，并且已承诺的预测在生成收敛之前保持可修正性。

在这项工作中，我们通过将掩码预测重新解释为干净状态预测来实现这种连续解码方案。我们将模型在掩码位置的输出视为嵌入空间中干净状态的估计，并用它来定义一个速度场，将当前状态从掩码嵌入移向预测的干净状态。我们将由此产生的连续动力学称为 x 预测流。与从随机高斯噪声开始的图像扩散不同，该流从确定的掩码嵌入（MDLM 明确训练过的状态）初始化。基于这个流，我们进一步将调度适应于语言生成的不对称、依赖上下文的性质，其中一些 token 必须提前承诺以告知其他 token。我们的主要贡献是：

- •**一种用于预训练 MDLM 的连续解码范式。** 我们展示掩码预测解码可以重新表述为输入嵌入空间中的 x 预测流，其中每个轨迹从掩码嵌入开始，并迭代地向预测的干净状态移动。这将二元解除掩码转化为连续、可修正的更新，并且仅需几百步对齐训练即可在现成的 MDLM 上运行。
- •**一种异步 token 级扩散调度。** 我们用基于置信度的异步更新方案取代图像扩散中的全局同步调度，其中每个 token 承载自己的解码进度。这使得高置信度的位置可以更快地进展，从而为其他 token 提供更清晰的上下文。
- •**一种学习到的 token 级步长策略。** 我们将每个 token 的更新参数化为其剩余扩散距离的对数尺度分数，条件于 token 的置信度统计和当前解码进度，并使用 GRPO[18 (https://arxiv.org/html/2606.29066#bib.bib11)] 结合任务级奖励和完成正则化项来训练此策略。

## 2 掩码扩散语言模型

本节介绍掩码扩散语言模型（MDLM）——一类非自回归文本生成模型，从部分掩码序列中恢复干净 token——并建立本文后续使用的符号。我们描述 MDLM 的一般形式，其中 LLaDA[16 (https://arxiv.org/html/2606.29066#bib.bib1)] 是规模上的一个代表性实例。

令 \(\mathcal{V}\) 为一个离散词汇表，并增加一个特殊的掩码 token [M]。一个长度为 \(N\) 的语言序列为 \(\mathbf{x}_0 = (x_0^1, \ldots, x_0^N) \in \mathcal{V}^N\)。每个 token \(v \in \mathcal{V}\) 通过一个嵌入矩阵 \(\mathbf{W}_e \in \mathbb{R}^{|\mathcal{V}| \times E}\) 与一个学习到的嵌入相关联，我们将掩码嵌入记为 \(\mathbf{m} \triangleq \mathbf{W}_e[[\texttt{M}]] \in \mathbb{R}^E\)。我们假设存在一个损坏机制，给定掩码比例 \(\sigma \in [0,1]\)，生成一个部分掩码序列 \(\tilde{\mathbf{x}}\)，其中位置的一个子集 \(\mathcal{M} \subseteq \{1,\ldots,N\}\)（预期大小为 \(\sigma N\)）被替换为 [M]，而剩余位置 \(\mathcal{U} = \{1,\ldots,N\} \setminus \mathcal{M}\) 保留其原始 token。

#### 训练目标

MDLM 被训练来在给定周围上下文的情况下恢复掩码位置的干净 token。具体地，模型定义每个位置 \(i\) 的条件分布 \(p_\theta(x_0^i \mid \tilde{\mathbf{x}})\)，并通过最小化仅限于掩码位置的交叉熵损失来优化：

\[
\mathcal{L}_{\text{MDLM}}(\theta) = -\,\mathbb{E}_{\sigma,\, \mathbf{x}_0,\, \tilde{\mathbf{x}}}\left[\frac{1}{\sigma}\sum_{i=1}^N \mathbf{1}\!\left[i \in \mathcal{M}\right] \log p_\theta\!\left(x_0^i \mid \tilde{\mathbf{x}}\right)\right],
\]
其中 \(\sigma \sim \mathcal{U}[0,1]\)，且 \(\frac{1}{\sigma}\) 权重补偿掩码 token 的预期比例。该目标已被证明是数据分布负对数似然的一个变分上界[17 (https://arxiv.org/html/2606.29066#bib.bib8);16 (https://arxiv.org/html/2606.29066#bib.bib1)]，提供了一个基于似然的原理性训练准则。

#### 掩码预测器

条件分布 \(p_\theta(\cdot \mid \tilde{\mathbf{x}})\) 由一个*掩码预测器* \(f_\theta\) 参数化，通常实现为双向 Transformer[19 (https://arxiv.org/html/2606.29066#bib.bib17)]，以便每个掩码位置可以关注完整的周围上下文。给定 \(\tilde{\mathbf{x}}\)，预测器输出每个位置的 logits：

\[
\mathbf{z}_{\text{pred}}^i = f_\theta(\tilde{\mathbf{x}})^i \in \mathbb{R}^{|\mathcal{V}|},\qquad p_\theta(\cdot \mid \tilde{\mathbf{x}})^i = \operatorname{softmax}(\mathbf{z}_{\text{pred}}^i),
\]
并且位置 \(i\) 的预测 token 为 \(x_{\text{pred}}^i = \arg\max_v \, p_\theta(v \mid \tilde{\mathbf{x}})^i\)。关键的是，所有掩码位置在单次前向传递中并行预测，这与自回归语言模型的从左到右顺序分解形成鲜明对比。

#### 标准离散采样

在推理时，模型的输入是干净提示 \(\mathbf{p}\) 和选定长度的完全掩码响应的拼接，解码在固定步数内进行，直到每个响应位置都被填充。在每一步，\(f_\theta\) 基于 \(\mathbf{p}\) 和部分解码的响应同时预测所有当前掩码的响应 token；这些预测的一个子集随后被承诺（解除掩码），而其余的在下一步被重新掩码。提示 \(\mathbf{p}\) 永远不会被掩码。现有的 MDLM 通常每步承诺固定比例的 token，并依赖诸如低置信度重新掩码（保留最自信的预测，重新掩码最不自信的预测[16 (https://arxiv.org/html/2606.29066#bib.bib1)]）等启发式方法来决定哪些 token 被承诺。然而，这种离散的同步更新忽略了 token 嵌入空间的连续结构，并迫使做出后期步骤无法修正的过早硬决策——这促使我们在下一节中引入连续解码框架。

## 3 MDLM 的连续状态扩散

本文的目标是将连续状态扩散引入 MDLM。与标准 MDLM 解码中的离散更新（每个掩码位置每步要么被承诺为一个单一 token，要么保持完全掩码）不同，连续状态解码器可以跨步骤为每个 token 携带一个软的中间估计，并逐渐细化它。保留这个中间状态而不是将其坍缩为硬决策，可以更有效地利用每个扩散步骤，并在有限的扩散预算下产生更高质量的生成。为此，我们提出一个基于 token 嵌入空间中 x 预测流的连续解码框架，该流锚定于掩码嵌入而非高斯噪声，这使得预训练的 MDLM 能够仅通过轻量级的对齐训练就作为连续状态扩散模型运行。

### 3.1 将连续状态扩散迁移到语言的挑战

在本节中，我们首先确定将连续状态扩散从图像迁移到语言的两个基本挑战，这些挑战将在后续小节中由我们的方法解决。连续状态扩散在图像领域取得了显著成功[11 (https://arxiv.org/html/2606.29066#bib.bib10);15 (https://arxiv.org/html/2606.29066#bib.bib9);14 (https://arxiv.org/html/2606.29066#bib.bib2)]，其中通过一系列小的细化步骤逐步去噪一个损坏的连续状态来恢复干净状态。然而，天真地将这个方案移植到语言会遇到两个基本障碍，任何文本的连续解码器都必须面对。

#### 高斯噪声不是语言的良定义状态

在连续状态图像扩散中，损坏过程总是涉及向干净图像添加某种高斯噪声。向干净图像添加高斯噪声会产生另一个图像，虽然模糊但仍然是像素空间的一个有效元素，并且模型在这种状态的连续体上进行训练。语言不允许类似噪声的概念。Token 是离散的，其含义由词汇身份决定，而不是由任何连续坐标中的位置决定，因此 token 嵌入的高斯扰动不对应于任何词汇项，且扰动的幅度没有语义可解释的对应物。因此，一个随机加噪的嵌入不是预训练 MDLM 曾经条件过的状态，将其用作连续扩散轨迹的起点会使模型从一开始就在分布外运行。

#### 语言生成高度依赖上下文

连续状态图像扩散采用同步去噪调度，因为像素虽然局部相关，但可以大致并行细化：消除一个像素歧义所需的大部分信息与其邻居对称共享，并随着噪声去除而逐渐揭示。语言表现出更强且更不对称的上下文依赖形式。一个 token 的身份