PreDiff-LM: 预训练离散掩码扩散语言建模与混合注意力

# 预训练离散掩码扩散语言建模与混合注意力

语言模型工具用于编辑辅助。作者审阅并承担所有声明、分析和文本的责任。
来源：https://arxiv.org/html/2607.25157
姚正涛¹，李润豪¹,²²脚注标记:²陈旭鹏²，程嘉懿²，乐晨茜² Michael Yue³，王杰森¹，王思恒⁴，杨光⁵，徐浩岩¹，魏晨皓⁵，袁正清⁶，沈悠然⁴，叶艳芳⁶，董俊豪⁷
¹南加州大学²纽约大学³哥伦比亚大学
⁴加州大学伯克利分校⁵史蒂文斯理工学院
⁶圣母大学⁷南洋理工大学

###### 摘要

离散掩码扩散语言模型支持双向生成和填充，但适配预训练的自回归（AR）Transformer需要协调因果预训练与双向去噪。我们在注意力层面研究这个问题，而非声称AR权重重用本身是新的。
PreDiff-LM在观察到的提示内保留因果注意力，同时允许掩码目标内的完全双向注意力。在匹配的GPT-2 Medium、WikiText-103、90K步设置下，与使用相同AR初始化的均匀双向注意力相比，这种混合掩码将无条件困惑度从34.1降低到28.7，MAUVE从0.71提高到0.78。注意力适配还与DiffuGPT风格的目标适配结合，达到了26.9的困惑度。预训练初始化将困惑度降至50以下所需的步数从约350K减少到8K，尽管计算匹配的微调AR模型在同等规模下仍然更强（18.9对比28.7）。除了困惑度，PreDiff-LM在重复性、分布质量、四个零样本下游任务以及人工偏好方面优于先前的扩散基线。这些结果将混合注意力定位为适配预训练因果主干的补充机制，同时明确了与优化AR模型之间仍存在的质量和推理效率差距。

## 引言

自回归（AR）语言模型通过逐个预测标记生成高质量文本（Radford等人，2019 (https://arxiv.org/html/2607.25157#bib.bib17)；Brown等人，2020 (https://arxiv.org/html/2607.25157#bib.bib2)）。其因果分解方式对从左到右的生成有效，但本身不支持任意顺序的细化或填充。离散扩散语言模型则是对整个序列进行破坏和去噪（Austin等人，2021 (https://arxiv.org/html/2607.25157#bib.bib1)；Sahoo等人，2024 (https://arxiv.org/html/2607.25157#bib.bib18)；Lou等人，2023 (https://arxiv.org/html/2607.25157#bib.bib12)）。这赋予了它们双向生成能力，但许多早期系统是从头开始训练的，并且性能仍然明显落后于同等规模的AR模型。

最近的研究表明，这种差距并非仅仅是扩散目标本身的特性。DiffuGPT和DiffuLLaMA通过持续扩散训练来适配预训练的AR模型（Gong等人，2024 (https://arxiv.org/html/2607.25157#bib.bib7)）；Dream从一个AR预训练模型初始化离散扩散LM，并引入了上下文自适应噪声重调度（Ye等人，2025 (https://arxiv.org/html/2607.25157#bib.bib23)）。这些结果使AR初始化成为一种既定策略。我们研究的未解决问题更为狭窄：*当因果权重被重用于掩码去噪时，信息流应如何变化？*

难点在于注意力模式的不匹配。AR权重是在严格的因果注意力下学习的，而掩码目标标记受益于左右两侧的上下文。在所有地方用双向注意力掩码替换因果掩码会改变每一层看到的表示，包括观察到的条件提示。PreDiff-LM改用不对称的混合掩码：提示标记保留其预训练的因果计算，而目标标记则双向关注提示和彼此。通过一个具有相同AR初始化的受控实验，与均匀双向注意力相比，分离出了5.4个点的困惑度改进。

我们的贡献是：

1. 我们将注意力适配形式化为AR到扩散迁移中的一个独立组件，并引入了一种混合因果-双向掩码，该掩码在保留提示侧因果结构的同时，实现了目标侧的去噪。
2. 我们将这种机制与Dream风格的均匀注意力控制进行分离，并展示了注意力适配可以与DiffuGPT风格的目标适配结合（结合后困惑度为26.9）。
3. 我们添加了计算匹配的微调AR控制、对重复敏感的度量、四个下游任务以及对偏好评估的三名人工审核。这些检查支持非AR的改进，同时不声称与匹配的AR模型性能相当。
4. 我们区分了强大的训练效率结果和推理延迟：在某种情况下，优化的AR解码速度更快，而扩散对于少步并行细化、填充和约束生成仍然有用。

参见图注图1：问题、机制和受控证据。*左：*预训练初始化在8K步内达到PPL低于50，而MDLM约需350K步，并提高了生成文本的质量。*中：*混合掩码在迭代去噪过程中保留因果提示计算，同时开放双向目标注意力。*右：*PreDiff-LM在评估的扩散基线中拥有最低的PPL，尽管外部结果在模型大小和训练预算上有所不同。匹配的微调AR控制在18.9 PPL时仍然更强，而匹配的AR差距在Large规模下缩小。
## 相关研究

#### 从头训练的掩码扩散。

D3PM引入了离散转移核（Austin等人，2021 (https://arxiv.org/html/2607.25157#bib.bib1)），而SEDD和MDLM改进了评分熵和掩码扩散目标（Lou等人，2023 (https://arxiv.org/html/2607.25157#bib.bib12)；Sahoo等人，2024 (https://arxiv.org/html/2607.25157#bib.bib18)）。ARDM和MAC将排序或自回归结构与扩散风格生成相结合（Hoogeboom等人，2021 (https://arxiv.org/html/2607.25157#bib.bib9)；Shi等人，2024 (https://arxiv.org/html/2607.25157#bib.bib20)）。LLaDA证明，从头训练的掩码扩散LM可以扩展到数十亿参数并执行下游语言任务（Nie等人，2025 (https://arxiv.org/html/2607.25157#bib.bib14)）。其基于置信度的重新掩码与我们推断规则密切相关；因此我们将置信度感知去掩码视为一个支持组件而非主要新颖点。

#### 适配预训练AR模型。

Diffusion-LM和CDCD研究嵌入空间中的连续扩散（Li等人，2022 (https://arxiv.org/html/2607.25157#bib.bib11)；Dieleman等人，2022 (https://arxiv.org/html/2607.25157#bib.bib5)）。DiffuGPT和DiffuLLaMA连接了AR和扩散目标，并持续适配预训练模型，包括逐步改变注意力结构（Gong等人，2024 (https://arxiv.org/html/2607.25157#bib.bib7)）。Dream在大型离散扩散LM中重用AR预训练权重，并将其与上下文自适应标记级噪声重调度相结合（Ye等人，2025 (https://arxiv.org/html/2607.25157#bib.bib23)）。这些工作在我们之前就进行了AR权重重用。我们的贡献是一个注意力层面的受控研究：我们保留条件提示的因果计算，仅对去噪目标开放双向注意力，并在其他条件匹配的情况下测试这种变化与均匀双向注意力的对比。目标层面和注意力层面的观点是互补而非竞争的解释；它们的组合表现优于单独任何一种。

#### 连续状态与路径统一生成。

离散随机定位（DSL）将预训练的掩码扩散检查点微调为连续状态、时间不变的去噪器，可以支持掩码细化、混合连续-离散采样以及使用单个模型进行随机顺序AR生成（Wu等人，2026 (https://arxiv.org/html/2607.25157#bib.bib22)）。DSL改变了破坏表示并统一了采样路径，而我们研究的是当起始检查点是因果AR Transformer时注意力应如何变化；因此这两个方向是互补的。

#### 生成质量与解码。

BERT使用双向掩码预测，但不是一个迭代生成模型（Devlin等人，2019 (https://arxiv.org/html/2607.25157#bib.bib4)）。最近的工作记录了扩散解码中的重复和长窗口失效模式，并提出了卷积解码和拒绝性微调（Seo等人，2025 (https://arxiv.org/html/2607.25157#bib.bib19)）。因此，除了平均困惑度之外，我们还评估MAUVE、Distinct-n、重复率、Self-BLEU、基于离散度的困惑度诊断、下游任务和盲法偏好。

## 方法

#### 设计原理。

每个组件都解决AR到扩散迁移中的不同失效模式。混合注意力保留了提示表示预训练时所依据的信息流模式，同时使掩码目标暴露于双向上下文。掩码率嵌入将一个单一共享网络的去噪噪声水平标识出来。自条件作用将模型之前的软预测反馈到下一轮，以减少同时掩码的标记之间的不一致性。置信度感知去掩码在解码预算小时限制错误传播。第一个机制是我们的核心贡献；其余组件构成了用于评估它的训练和推理方案。

### 3.1 背景：离散掩码扩散

令 \( \mathbf{x} = (x_1, \ldots, x_L) \in \mathcal{V}^L \) 为来自词汇表 \( \mathcal{V} \) 的长度为 \( L \) 的标记序列。离散掩码扩散定义了一个正向过程，该过程以速率 \( t \in [0,1] \) 随机用 [MASK] 标记替换标记：

\[
q(\mathbf{x}_t|\mathbf{x}_0,t) = \prod_{i=1}^{L}\bigl[(1-t)\cdot\delta_{x_i^t=x_i^0} + t\cdot\delta_{x_i^t=\texttt{[MASK]}}\bigr]
\]
(1)

其中 \( t=0 \) 是干净数据，\( t=1 \) 是完全掩码。逆向过程学习从损坏序列预测原始标记：
\[
p_{\theta}(\mathbf{x}_0|\mathbf{x}_t,t) = \prod_{i:x_i^t=\texttt{[MASK]}} p_{\theta}(x_i^0|\mathbf{x}_t,t)
\]
。训练最小化掩码位置上的期望交叉熵：

\[
\mathcal{L}_{\text{diff}} = \mathbb{E}_{t\sim p(t),\, \mathbf{x}_t\sim q(\cdot|\mathbf{x}_0,t)}\left[-\frac{1}{|\mathcal{M}_t|}\sum_{i\in\mathcal{M}_t}\log p_{\theta}(x_i^0|\mathbf{x}_t,t)\right]
\]
(2)

其中 \( \mathcal{M}_t = \{i:x_i^t=\texttt{[MASK]}\} \) 是掩码位置的集合，\( t \) 从噪声调度 \( p(t) \) 中抽取。

### 3.2 混合因果-双向注意力

核心适配问题是不对称的。提示标记是观察到的条件变量，所以保留它们的因果计算就保留了AR权重学习时所处的内部表示机制。目标标记是潜在的重建变量，所以每个目标标记受益于两侧的证据以及模型对其他目标的当前信念。对两个区域进行相同处理要么会移除有用的目标上下文（完全因果注意力），要么会扰乱预训练的提示计算（均匀双向注意力）。

我们通过一个混合的4D注意力掩码 \( \mathbf{M} \in \{0,1\}^{(L_p+L_x)\times(L_p+L_x)} \) 来编码这种不对称性。它将输入划分为长度为 \( L_p \) 的提示前缀 \( \mathbf{p} \) 和长度为 \( L_x \) 的目标区域 \( \mathbf{x}_t \)：

\[
M_{ij} = \begin{cases}
\mathbf{1}[j\le i] & i,j\in\text{prompt}\\
1 & i\in\text{target},\;j\in\text{prompt}\\
1 & i,j\in\text{target}\\
0 & i\in\text{prompt},\;j\in\text{target}
\end{cases}
\]
(3)

第一种情况保留了提示侧的因果模式。第二种和第三种情况允许每个目标标记使用完整的提示和当前的目标状态。第四种情况防止来自损坏目标的信息改变提示表示。这种设计并非关于因果提示通用性的上界声明：它是针对重用因果权重的迁移假设。第4.3节 (https://arxiv.org/html/2607.25157#S4.SS3) 在匹配的初始化和训练条件下，测试了它与均匀双向注意力的对比。这种掩码是否有助于完全从头训练的模型仍是开放问题。

### 3.3 自条件化

掩码扩散后验是在掩码位置上分解的。在高掩码率下，这种边际预测可能产生相互不一致的标记选择，因为共同掩码的标记彼此观察不到对方的干净值。自条件化为缺失的联合上下文提供了一个软近似。在训练期间，以概率 \( \rho=0.5 \)，我们首先计算初步预测 \( \hat{\mathbf{x}}_0 = f_{\theta}(\mathbf{x}_t,t) \)，然后在此基础上进行细化的传递：

\[
p_{\theta}(\mathbf{x}_0|\mathbf{x}_t,t,\hat{\mathbf{x}}_0) = f_{\theta}\!\left(\mathbf{x}_t + \text{proj}(\hat{\mathbf{x}}_0),\;t\right)
\]
(4)

其中 \( \text{proj}:\mathbb{R}^{|\mathcal{V}|}\to\mathbb{R}^d \) 将之前的softmax分布映射到一个嵌入残差。这使每个位置都能访问模型对其邻居的当前信念，而无需提交硬标记。当许多标记被掩码时，效果最大：移除自条件化将PPL@50%从69.0提高到214.0，无条件困惑度提高32.1个点（表7 (https://arxiv.org/html/2607.25157#S4.T7)）。

### 3.4 掩码率时间嵌入和噪声调度

掩码率 \( t \) 通过正弦编码（Vaswani等人，2017 (https://arxiv.org/html/2607.25157#bib.bib21)）后跟一个MLP注入：\( \mathbf{e}_t = \text{MLP}(\text{SinEmbed}(t)) \in \mathbb{R}^d \)，并添加到所有标记位置。我们从余弦调度的 \( t \) 中采样，该调度将训练集中到信息量较大的中等掩码率：

\[
t = 1 - \cos\!\left(\tfrac{\pi}{2}u\right), \quad u \sim \text{Uniform}(0,1)
\]
(5)

### 3.5 置信度感知去掩码 (CAU)

在推理时，我们在 \( T \) 步内迭代地按 *容易优先* 策略去掩码标记，该策略与LLaDA中的基于置信度的重新掩码密切相关（Nie等人，2025 (https://arxiv.org/html/2607.25157#bib.bib14)）。我们将其作为一个解码组件包含在内，并且不声称通用的置信度排序原则是新颖的。在步骤 \( s \)，我们计算预测 \( \hat{\mathbf{x}}_0 = f_{\theta}(\mathbf{x}_{t_s}, t_s) \)，按置信度 \( c_i = \max_v p_{\theta}(x_i=v|\mathbf{x}_{t_s}, t_s) \) 对掩码位置进行排序，然后去掩码前 \( k_s \) 个位置，其中：

\[
k_s = \left\lceil |\mathcal{M}_{t_s}| \cdot \right.
\]