扩散语言模型的动态分块

# 动态分块扩散语言模型  
来源：https://arxiv.org/html/2605.15676  

Yichen Zhu CSE, HKUST yc\_zhu@zju\.edu\.cn  
& Xiaoming Shi†‡ Xiaohongshu Inc\. sxm728@hotmail\.com  
&Peng Zhao Alibaba group zhuyun\.zp@alibaba\-inc\.com  
Weiyu Chen CityUHK weiyu\.chen@cityu\.edu\.hk  
& Debing Zhang Xiaohongshu Inc\. dengyang@xiaohongshu\.com  
& James Kwok CSE, HKUST jamesk@cse\.ust\.hk  

###### 摘要  

块离散扩散语言模型通过固定大小的位置块自回归地分解序列，将块内并行去噪与块间条件依赖解耦。我们认为这种刚性划分浪费了序列中已有的结构：由位置而非内容定义的块会将语义相关的token拆开，并将无关的token组合在一起。我们提出了**动态分块扩散模型（DCDM）**，它用内容定义的语义块取代了位置块。其核心是**分块注意力**（Chunking Attention），一个可微分的路由层，将token分配到由可学习子空间参数化的 $K$ 个簇中，并通过扩散目标端到端地塑形。由此产生的簇分配诱导出一种块因果注意力掩码，在该掩码下，离散扩散去噪器在语义块上自回归地分解序列似然，严格推广了块离散扩散。在参数规模高达1.5B的下游基准测试中，DCDM一致优于非结构化和位置块扩散基线，其优势在不同规模下保持稳定，并在训练早期即可显现。

## 1 引言  

22footnotetext:†\dagger 项目负责人。  
33footnotetext:‡\ddagger 通讯作者：[email protected]  

扩散大语言模型（dLLMs）由于能够并行解码多个token，最近已成为一种有竞争力的文本生成范式。开源掩码扩散语言模型（MDLMs）[34](https://arxiv.org/html/2605.15676#bib.bib3)，如LLaDA[30](https://arxiv.org/html/2605.15676#bib.bib13)和Dream[44](https://arxiv.org/html/2605.15676#bib.bib14)，在相似规模下取得了与自回归模型相当的性能，而闭源模型如Gemini Diffusion[16](https://arxiv.org/html/2605.15676#bib.bib17)和Mercury[22](https://arxiv.org/html/2605.15676#bib.bib18)则展示了更高的生成吞吐量。这一进展背后的关键因素是**块扩散**[2](https://arxiv.org/html/2605.15676#bib.bib4)，它已成为可扩展扩散语言建模的主流设计。  

块扩散结合了自回归模型和扩散模型的优势。它在块上自回归地分解序列，保留了token组之间的因果条件依赖，同时在每个块内部双向并行地去噪。这种设计在自回归建模的质量与并行扩散采样的效率之间提供了一个实用的折中方案。然而，现有的块扩散语言模型（BDLMs）[2](https://arxiv.org/html/2605.15676#bib.bib4)通过固定的位置规则来定义块：一个序列被划分为等长的连续片段。这种选择施加了一个与序列内容无关的强结构先验。我们认为固定的位置块是语言建模中一种有局限的抽象。决定一个token的依赖关系往往不与局部的连续性对齐：一个实体可能支配远处的提及，一个数学推导可能依赖于前面的前提，一个代码token可能受到几行之外的作用域或语法的约束。因此，位置划分可能将应该联合去噪的token分开，同时将弱相关的相邻token置于同一个扩散过程中。在这种情况下，模型继承了块自回归分解，但应用粒度与序列的语义结构不匹配。  

为了解决这种不匹配，我们提出了**动态分块扩散模型（DCDM）**，它用学习的语义块取代了固定的位置块。DCDM不是根据位置对序列进行分割，而是根据模型内部产生的表示对token进行聚类。产生的块可以是不连续的、大小可变且依赖于序列的。它们起着与块扩散中的块相同的作用：块内的token被双向并行去噪，而块则通过块因果掩码以自回归方式排序。因此，DCDM保留了块扩散的计算结构，同时使并行去噪的单元具有内容自适应性。  

DCDM的核心组件是**分块注意力**（Chunking Attention），这是一个可微分的路由层，将token分配到 $K$ 个块中的一个。在高维语言模型表示中，直接的点-质心聚类是不稳定的，因为少数几个簇可能在早期占据主导地位，导致其余簇缺乏梯度信号。我们转而用可学习的低维子空间表示每个块，并根据子空间对齐度[32](https://arxiv.org/html/2605.15676#bib.bib7)来路由token。一条软注意力路径将分块几何结构直接置于扩散目标的梯度路径上，而由此产生的硬分配则定义了一个语义块因果注意力掩码。这种构造推广了位置块扩散，当学习的块恰好与固定的连续块重合时，即可恢复为位置块扩散。  

本文的贡献如下：  
- • 我们引入了分块注意力，一种基于子空间的可微分路由机制，为扩散语言建模诱导语义块因果掩码。  
- • 我们开发了DCDM，一种严格推广BDLM的扩散语言模型，并在扩散目标下端到端地训练分块机制和去噪器。  
- • 我们提供了大量实证证据，表明在通用推理、数学和代码生成等下游任务中，语义分块在0.5B和1.5B两种规模下均优于其位置对应物，验证了块扩散从内容自适应粒度中获得了显著收益。  

## 2 相关工作  

#### 扩散大语言模型。  
语言生成领域长期以来一直由自回归模型主导[33](https://arxiv.org/html/2605.15676#bib.bib24), [7](https://arxiv.org/html/2605.15676#bib.bib26), [1](https://arxiv.org/html/2605.15676#bib.bib25), [4](https://arxiv.org/html/2605.15676#bib.bib28), [39](https://arxiv.org/html/2605.15676#bib.bib27), [25](https://arxiv.org/html/2605.15676#bib.bib30), [12](https://arxiv.org/html/2605.15676#bib.bib29)]。虽然这些模型以高质量输出著称，但它们从根本上受到顺序的、逐token解码过程的限制[19](https://arxiv.org/html/2605.15676#bib.bib32), [38](https://arxiv.org/html/2605.15676#bib.bib33)]。为了缓解这些延迟瓶颈，dLLMs——一类专门为离散数据领域设计的扩散框架——已成为一种引人注目的替代方案。通过引入吸收状态（例如 `[MASK]`）来表示噪声，Austin等人[3](https://arxiv.org/html/2605.15676#bib.bib2)为掩码扩散建模奠定了基础。该框架随后被一系列近期工作[34](https://arxiv.org/html/2605.15676#bib.bib3), [37](https://arxiv.org/html/2605.15676#bib.bib12), [30](https://arxiv.org/html/2605.15676#bib.bib13), [44](https://arxiv.org/html/2605.15676#bib.bib14), [15](https://arxiv.org/html/2605.15676#bib.bib15)]扩展。值得注意的是，MDLM[34](https://arxiv.org/html/2605.15676#bib.bib3)是应用最广泛的模型之一，提供了简单高效的训练目标。LLaDA[30](https://arxiv.org/html/2605.15676#bib.bib13)系列将扩散语言模型扩展到超过80亿参数，展示了与同等规模自回归模型相当甚至更优的性能。  

#### 自回归-扩散混合语言模型。  
近期工作探索了将自回归模型的计算效率集成到扩散框架中，特别是在视频合成等复杂任务中。一种代表性方法BDLM[2](https://arxiv.org/html/2605.15676#bib.bib4)以自回归方式建模跨块的语义依赖，同时在每个块内独立执行去噪过程。Fast-dLLMs[42](https://arxiv.org/html/2605.15676#bib.bib16)采用诸如块级前缀缓存等技术，实现了远超过AR模型的生成效率，且不牺牲生成质量。另一项工作尝试在**推理时**放松固定位置块的刚性。AdaBlock-dLLM[26](https://arxiv.org/html/2605.15676#bib.bib22)在采样过程中使用局部去噪置信度信号自适应地调整块边界。  

## 3 预备知识  

### 3.1 掩码扩散模型  

掩码扩散语言模型（MDLMs）[34](https://arxiv.org/html/2605.15676#bib.bib3)是一类离散扩散模型，其中前向过程的吸收分布 $\pi$ 是特殊掩码token $m$ 上的点质量。令 $x \in V^L$ 表示从数据分布中抽取的长度为 $L$ 的干净序列，$z_t$ 表示其在时刻 $t \in [0,1]$ 的损坏潜变量。前向过程在连续时间 $t \in [0,1]$ 上操作，并以概率 $1 - \alpha_t$ 独立地将每个token替换为 $m$，其中 $\alpha_t$ 是一个预定义的噪声调度，从 $\alpha_0 = 1$（干净数据）严格递减到 $\alpha_1 = 0$（完全掩码）。逆过程由一个去噪器 $x_\theta(z_t, t)$ 参数化，该去噪器被训练以从掩码状态预测干净数据，其证据下界（ELBO）简化为每个样本的加权交叉熵损失：

$$L(x, \theta) = \mathbb{E}_{q(z_t \mid x)} \int_0^1 \frac{\alpha_t'}{1 - \alpha_t} \sum_{\ell: z_t^\ell = m} \log \langle x_{\theta, \ell}(z_t^{1:L}, t), x_\ell \rangle dt, \quad (1)$$

其中 $\alpha_t' = d\alpha_t / dt$，$x_{\theta, \ell}$ 是位置 $\ell$ 处的预测类别分布。

### 3.2 块扩散模型  

块扩散语言模型（BDLMs）[2](https://arxiv.org/html/2605.15676#bib.bib4)通过将长度为 $L$ 的token序列划分为 $K$ 个固定长度 $B$（$L = K \cdot B$）的连续块，结合了自回归和扩散建模，在每个块内执行离散扩散，同时保持跨块的自回归依赖。似然在块上自回归分解：

$$\log p_\theta(x) = \sum_{b=1}^K \log p_\theta(x_b \mid x_{<b}), \quad (2)$$

其中 $p_\theta(x_b \mid x_{<b})$ 由使用因果掩码的块内扩散过程建模。在训练过程中，BDLM 通过交替优化每个块的掩码扩散损失来训练去噪器，同时在推理时通过逐块采样生成序列：每个块 $x_b$ 使用去噪器 $p_\theta(x_b \mid x_{<b})$ 从完全掩码状态开始采样。

## 4 方法  

我们提出DCDM，它用一个可学习的、内容自适应的分块方案取代位置块。DCDM的核心是一种新颖的可微分路由机制，称为分块注意力，它根据token与可学习子空间的对齐度，将token灵活地组织成语义块。一旦分块确定，DCDM就根据这些语义块应用块因果去噪。  

### 4.1 通过子空间对齐进行分块  

令隐藏表示集合 $H = \{h_1, \dots, h_L\}$ 由语言模型的某一层产生。分块的目标是将每个token $\ell$ 分配到 $K$ 个语义块之一。直接学习 $K$ 个可学习质心 $\{c_k\}_{k=1}^K$ 并根据到质心的距离 $\|h_\ell - c_k\|$ 分配token在高维空间中是有问题的：少数簇倾向于捕捉大部分token，导致其余簇在训练早期饥饿并接收不到有意义的梯度。我们通过将每个块 $k$ 与一个低维子空间相关联来解决这个问题，该子空间由一个可学习投影矩阵 $W_k \in \mathbb{R}^{d \times r}$（$r \ll d$）参数化。  

**软路由：** token $\ell$ 到块 $k$ 的软对齐分数定义为：

$$s_{\ell k} = \|h_\ell^T W_k\|_2^2 \cdot \tau, \quad (3)$$

其中 $\tau$ 是一个温度超参数。这测量了 $h_\ell$ 在由 $W_k$ 张成的子空间上的投影范数：投影范数大的表示与子空间紧密对齐。然后，软分配由 $a_{\ell k} = \text{softmax}_k(s_{\ell k})$ 给出。  

**硬路由：** 为了因果掩码，我们需要每个token的硬块分配 $z_\ell \in \{1, \dots, K\}$，通过 $z_\ell = \arg\max_k s_{\ell k}$ 获得。为了保留可微性，我们使用直通估计器：前向传播使用 $z_\ell$ 进行硬分配，而梯度通过复制软分配 $a_{\ell k}$ 来反向传播。  

**子空间的可学习性：** 重要的是，$W_k$ 是通过扩散目标的梯度来训练的，因为对齐分数 $s_{\ell k}$ 的分支通过残差连接和后续层影响损失。剩余子空间随机初始化，并通过反向传播进行端到端优化。

### 4.2 语义块因果掩码  

给定硬块分配 $\{z_1, \dots, z_L\}$，我们通过定义以下因果掩码来施加块因果依赖：

$$\text{Mask}(\ell, j) = 
\begin{cases} 
1 & \text{如果 } z_\ell = z_j \text{ 且 } g_\ell < g_j, \\
1 & \text{如果 } z_\ell \neq z_j \text{ 且 } z_\ell < z_j, \\
0 & \text{否则},
\end{cases} \quad (4)$$

其中 $g_\ell$ 是token $\ell$ 在其块内的顺序索引。这强制执行块间自回归因果依赖和块内双向可见性。  

### 4.3 通过历史缓冲进行高效推理  

在推理期间并行解码所有块需要一次性获得完整分块方案，但在块被生成之前，真实的块分配是未知的。我们通过使用去噪器在每一步预测的token来近似最终块分配。具体来说，在扩散过程的第 $m$ 步，我们维护一个历史缓冲区 $\hat{x}^{(0)} = [MASK]^L$，并在每次去噪步骤后更新它：$\hat{x}^{(m)} \leftarrow \text{denoise}(\hat{x}^{(m-1)})$。块分配根据 $\hat{x}^{(m)}$ 的表示动态调整。  

### 4.4 分块注意力  

我们引入分块注意力，这是一个可微分的路由层，通过子空间对齐对token进行分块。核心组件是每个块 $k$ 的参数化投影矩阵 $W_k$。对于每个token，我们计算其与每个子空间的对齐分数，然后通过软注意力和硬路由机制获得硬块分配。此外，我们引入一个负载平衡偏差 $\mathbf{b} \in \mathbb{R}^K$ 以防止空块。该偏差在学习到的分配分数上添加一个控制循环修正：

$$\tilde{s}_{\ell k} = s_{\ell k} + b_k, \quad z_\ell = \arg\max_k \tilde{s}_{\ell k}, \quad \text{其中 } b_k \gets b_k - \eta_b \cdot (\text{count}_k - L/K). \quad (5)$$

$\eta_b > 0$ 是控制偏差对负载不平衡反应强度的步长。我们在整个训练过程中使用固定的 $\eta_b$。此更新降低了过载簇的偏差，提高了未充分利用簇的偏差。由于 $\mathbf{b}$ 仅进入离散的 $\arg\max$ 分支，没有梯度流经它：训练 $\{\mu_k\}$ 的软路径不受影响，而 $\mathbf{b}$ 纯粹作为叠加在结果硬分配之上的控制循环修正。

## 5 实验  

### 5.1 实验设置  

**数据集。** 所有扩散模型在 OpenWebText[14](https://arxiv.org/html/2605.15676#bib.bib34) 上以统一的训练协议进行预训练；完整训练细节见附录 D (https://arxiv.org/html/2605.15676#A4)。然后我们在标准基准测试套件上评估预训练的 LLM，该套件分为三类：通用推理与知识（ARC-C[9](https://arxiv.org/html/2605.15676#bib.bib35)、MMLU[18](https://arxiv.org/html/2605.15676#bib.bib36)、[17](https://arxiv.org/html/2605.15676#bib.bib37)、HellaSwag[45](https://arxiv.org/html/2605.15676#bib.bib38)、TruthfulQA[24](https://arxiv.org/html/2605.15676#bib.bib41)、WinoGrande[35](https://arxiv.org/html/2605.15676#bib.bib40)、PIQA[6](https://arxiv.org/html/2605.15676#bib.bib39)）、数学推理（MATH[23](https://arxiv.org/html/2605.15676#bib.bib42)、GSM8K[10](https://arxiv.org/html/2605.15676#bib.bib43)）和代码生成（HumanEval[5](https://arxiv.org/html/2605.15676#bib.bib44)）。密集扩散模型的训练动态总结在图 2 (https://arxiv.org/html/2605.15676#S5.F2) 中。附录 E.2 (https://arxiv.org/html/2605.15676#A5.SS2) 报告了在七个保留语料库上的补充零样本语言建模评估。

**基线。** 我们将 DCDM 与离散扩散语言建模的主流范式进行比较：(i) MDLM[34](https://arxiv.org/html/2605.15676#bib.bib3) 是一种掩码离散扩散语言模型，在没有块结构的情况下并行去噪 token。(ii) BDLM[2](https://arxiv.org/html/2605.15676#bib.bib4) 是一种块离散扩散模型，在固定大小的**位置**块上施加自回归分解，是离我们工作最近的先前研究。(iii) AdaBlock-dLLM[26](https://arxiv.org/html/2605.15676#bib.bib22) 是一种无训练推理时调度器，在解码过程中基于使用固定位置块训练的模型自适应调整块大小。我们在匹配的参数规模下评估所有基线，此外还包括一个稀疏的 DCDM (MoE) 变体，其活跃参数预算（0.4B / 1.2B）略低于密集模型，以评估 DCDM 的语义块结构是否与条件计算兼容。所有密集扩散模型在两种规模（0.5B 和 1.5B 参数）下进行评估，使用相同的训练数据和分词器，并在我们的统一评估协议下进行。

**指标。** 对于下游基准测试，我们遵循每个任务的标准指标[15](https://arxiv.org/html/2605.15676#bib.bib15)：多项选择任务（ARC-C、MMLU、HellaSwag、TruthfulQA、WinoGrande、PIQA）的准确率，MATH 和 GSM8K 的精确匹配准确率，以及 HumanEval 的 pass@1。每个基准测试所用的上下文示例数在表 1 (https://arxiv.org/html/2605.15676#S5.T1) 的括号内报告；没有括号的条目为零样本评估。

**表 1：** 在 0.5B 和 1.5B 规模下预训练 LLM 的下游基准测试结果。括号内的数字表示上下文示例的数量；没有括号的条目为零样本。† 标记了模型已在相应训练分割上进行微调的基准测试。对于 DCDM (MoE)，`Active Params` 行报告每个 token 的激活参数，括号内为总参数；所有其他模型在每个 token 处激活所有参数。在每个规模内，每行最佳结果以**粗体**显示，次佳结果以下划线显示。* 平均值计算了所有九个报告的基准测试。

### 5.2 主要结果  

**下游基准测试。** 表 1 (https://arxiv.org/html/2605.15676#S5.T1) 报告了在 0.5B 和 1.5B 两种规模下整个下游套件的零样本和少样本结果。三个模式在不同规模下保持一致。首先，DCDM 在每个基准测试上的表现都优于非结构化扩散基线 MDLM，除了 WinoGrande，在该测试中两者在两种规模下相差在 1 点以内（0.5B 下 DCDM 领先 0.84，1.5B 下落后 0.40）。最大的收益集中在...