PulseCol：周期性刷新的列稀疏注意力机制，用于加速扩散语言模型

# PulseCol：周期性刷新的列稀疏注意力加速扩散语言模型

来源：https://arxiv.org/html/2605.20813
Yanyi Lyu∗Letian Chen∗Futing Sun Miao Zhang†Weili Guan Liqiang Nie† 哈尔滨工业大学（深圳） \{lvyanyi, chenlltt, sunfuting\}@stu\.hit\.edu\.cn \{zhangmiao, guanweili\}@hit\.edu\.cnnieliqiang@gmail\.com *共同第一作者。†通讯作者。

###### 摘要

扩散大语言模型（dLLMs）的推理计算成本高昂，因为在去噪过程的每一步都必须重复执行完整的自注意力，而无法使用 KV 缓存。最近针对 dLLMs 的稀疏注意力方法通过块稀疏计算来缓解这一成本，但仅应用于后期迭代（此时模型性能对粗粒度稀疏近似的敏感度较低），因此在计算效率和加速方面带来的改进有限。这激发了一种更细粒度的稀疏化策略，该策略可以从更早的迭代开始应用，并利用可重用的稀疏模式，从而实现进一步的效率提升。在这项工作中，我们引入了 PulseCol，一种用于加速扩散语言模型的周期性刷新列稀疏注意力方法。PulseCol 用更细粒度的列稀疏结构取代了粗粒度的块级稀疏性，从而能够更精确地保留重要的注意力交互，同时暴露更大的稀疏性。基于这种列级公式，PulseCol 进一步在早期去噪步骤中识别稀疏模式，并在后续迭代中重用这些模式，仅在少量中间步骤刷新它们，以跟踪去噪过程中稀疏注意力模式的演变。实验表明，与先前针对 dLLMs 的稀疏注意力方法相比，PulseCol 实现了更高的稀疏性和更大的实际加速，同时保持了模型质量。借助针对列稀疏注意力优化的 GPU 内核，PulseCol 在各种上下文长度上实现了相对于 FlashAttention 最高 1.95× 的端到端加速。

参见标题图1：PulseCol 改善了 dLLM 推理的稀疏性-效率权衡。(a) 在 GSM8K 上，随着稀疏度增加，PulseCol 比 SparseD 保持了更高的准确率和 oracle 每 token 前 k 个注意力召回率，表明列级稀疏性能更好地保留重要的注意力交互。(b) 在使用 1024 步去噪的 LLaDA-1.5 模型上，PulseCol 实现了比基于 FlashAttention 的全注意力和 SparseD 更低的推理延迟，在 64K 上下文窗口下达到了 1.95× 的加速。

## 1 引言

扩散语言模型最近作为自回归语言建模的补充方向引起了越来越多的兴趣[1（https://arxiv.org/html/2605.20813#bib.bib23），14（https://arxiv.org/html/2605.20813#bib.bib24），31（https://arxiv.org/html/2605.20813#bib.bib21），16（https://arxiv.org/html/2605.20813#bib.bib25），21（https://arxiv.org/html/2605.20813#bib.bib26），19（https://arxiv.org/html/2605.20813#bib.bib15），9（https://arxiv.org/html/2605.20813#bib.bib27），18（https://arxiv.org/html/2605.20813#bib.bib2），28（https://arxiv.org/html/2605.20813#bib.bib1）]。通过使用双向上下文迭代地精炼序列，dLLMs 提供了一种超越严格从左到右解码的灵活生成范式[18（https://arxiv.org/html/2605.20813#bib.bib2），28（https://arxiv.org/html/2605.20813#bib.bib1）]。然而，这种灵活性伴随着巨大的推理成本，因为在每个去噪步骤都必须重新计算完整的注意力[17（https://arxiv.org/html/2605.20813#bib.bib18），25（https://arxiv.org/html/2605.20813#bib.bib20），15（https://arxiv.org/html/2605.20813#bib.bib33），24（https://arxiv.org/html/2605.20813#bib.bib34）]。这种注意力成本在重复的去噪步骤中累积，给高效的 dLLM 推理带来了实际挑战。

鉴于这种累积的注意力成本，利用注意力中的稀疏性是提高 dLLM 推理效率的一种自然方式。然而，为自回归语言模型开发的稀疏注意力方法不能直接转移到 dLLMs，因为它们的推理过程本质不同[6（https://arxiv.org/html/2605.20813#bib.bib28），29（https://arxiv.org/html/2605.20813#bib.bib30），27（https://arxiv.org/html/2605.20813#bib.bib10），30（https://arxiv.org/html/2605.20813#bib.bib32），13（https://arxiv.org/html/2605.20813#bib.bib9），24（https://arxiv.org/html/2605.20813#bib.bib34），26（https://arxiv.org/html/2605.20813#bib.bib35）]。自回归模型通过随时间扩展前缀并重用持久化的 KV 缓存来生成 token[12（https://arxiv.org/html/2605.20813#bib.bib31），27（https://arxiv.org/html/2605.20813#bib.bib10），30（https://arxiv.org/html/2605.20813#bib.bib32）]，而 dLLMs 则迭代地更新序列级的去噪状态[18（https://arxiv.org/html/2605.20813#bib.bib2），28（https://arxiv.org/html/2605.20813#bib.bib1），17（https://arxiv.org/html/2605.20813#bib.bib18），25（https://arxiv.org/html/2605.20813#bib.bib20）]。因此，dLLMs 中的稀疏注意力不能仅由静态的 token 位置或单步重要性估计来确定，因为有用的注意力连接可能会随着去噪的进行而改变[22（https://arxiv.org/html/2605.20813#bib.bib19），24（https://arxiv.org/html/2605.20813#bib.bib34），26（https://arxiv.org/html/2605.20813#bib.bib35）]。这些差异需要稀疏注意力设计能够明确表征注意力模式在去噪步骤间的演变方式。

最近关于 dLLM 推理的研究进一步观察到，相邻去噪步骤之间的注意力模式表现出显著的相似性[17（https://arxiv.org/html/2605.20813#bib.bib18），15（https://arxiv.org/html/2605.20813#bib.bib33），22（https://arxiv.org/html/2605.20813#bib.bib19），24（https://arxiv.org/html/2605.20813#bib.bib34），26（https://arxiv.org/html/2605.20813#bib.bib35）]。这一观察表明，不需要在每次迭代都从头重新计算稀疏注意力结构，这促使了一类通过在整个去噪过程中重用稀疏模式来加速 dLLM 推理的方法。尽管这些方法有效，但它们通常仅在推理的后期阶段应用稀疏化。其基本假设是，早期的去噪步骤对最终的生成质量影响很大，因此应保留完整的注意力[24（https://arxiv.org/html/2605.20813#bib.bib34）]。这种设计选择揭示了一个关键且未解决的问题。为什么在 dLLMs 中，早期阶段的注意力稀疏化如此脆弱？

在这项工作中，我们研究了早期去噪迭代中的注意力模式。我们观察到，推理早期阶段的注意力并非真正密集。相反，高得分的注意力始终集中在少数关键列上，揭示了显著的列向稀疏模式。这一发现为现有稀疏 dLLM 推理方法在早期步骤的脆弱性提供了一个可能的解释。它们的块级近似施加了粗暴的稀疏粒度，与注意力的固有列向结构不匹配，因此可能在早期去噪过程中丢弃重要的注意力连接。

受这一观察启发，我们提出了 PulseCol，一种用于高效 dLLM 推理的周期性刷新列稀疏注意力方法。PulseCol 不采用粗粒度的块级稀疏性来近似注意力[24（https://arxiv.org/html/2605.20813#bib.bib34），26（https://arxiv.org/html/2605.20813#bib.bib35）]，而是直接针对早期去噪步骤中显著的列向稀疏结构。通过用列稀疏性取代块稀疏性，PulseCol 更好地近似了底层的注意力模式，使得从更早的去噪阶段引入稀疏注意力变得更加安全。由于注意力模式在去噪过程中仍然会演变，PulseCol 在相邻步骤间重用列稀疏索引，同时定期刷新它们以适应这种演变。这种设计使得在早期阶段重用稀疏性成为可能，而无需在整个推理过程中假设固定的稀疏模式，从而扩大了 dLLM 推理的加速窗口。

实验表明，PulseCol 在生成质量和推理效率之间实现了良好的权衡。在相同的稀疏度预算下，PulseCol 在平均性能上优于先前的稀疏注意力基线。在 80% 稀疏度下，它仅略逊于全注意力，同时明显优于现有的稀疏方法。借助我们针对列稀疏注意力优化的 GPU 内核，PulseCol 进一步将其算法稀疏性转化为实际的推理加速，在不同上下文长度上实现了相对于 FlashAttention[8（https://arxiv.org/html/2605.20813#bib.bib12）] 最高 1.95× 的加速。

## 2 相关工作

### 2.1 扩散语言模型

扩散语言模型最近被探索为自回归生成的一种补充方法[18（https://arxiv.org/html/2605.20813#bib.bib2），32（https://arxiv.org/html/2605.20813#bib.bib3），3（https://arxiv.org/html/2605.20813#bib.bib4），28（https://arxiv.org/html/2605.20813#bib.bib1）]。最近的工作如 LLaDA[18（https://arxiv.org/html/2605.20813#bib.bib2）] 表明，这种范式可以在标准的预训练和指令微调下扩展到大型模型，同时支持对整个序列进行并行精炼，并具有竞争力的性能。这些优势使得 dLLMs 成为一个越来越有吸引力的语言生成框架。与此同时，它们的迭代去噪过程在多个步骤中反复调用双向全注意力，使得推理效率成为一个核心挑战[25（https://arxiv.org/html/2605.20813#bib.bib20），24（https://arxiv.org/html/2605.20813#bib.bib34）]。

现有加速 dLLM 推理的努力主要分为两类。第一类利用 KV 缓存和跨步骤重用。诸如 Fast-dLLM[25（https://arxiv.org/html/2605.20813#bib.bib20）]、FlashDLM[11（https://arxiv.org/html/2605.20813#bib.bib36）] 和 dLLM-Cache[15（https://arxiv.org/html/2605.20813#bib.bib33）] 等方法通过近似缓存、选择性更新或更有效地重用中间状态，减少了相邻去噪步骤间的冗余计算。第二类侧重于采样和解码策略。这些方法通过置信度或熵感知调度[25（https://arxiv.org/html/2605.20813#bib.bib20），2（https://arxiv.org/html/2605.20813#bib.bib37）]、流水线生成[23（https://arxiv.org/html/2605.20813#bib.bib38）] 或针对并行解码设计的蒸馏方法[5（https://arxiv.org/html/2605.20813#bib.bib39）] 来提高推理效率。虽然有效，但这些方法主要减少了有效迭代次数或跨步骤重用计算，而不是解决每个去噪步骤内注意力计算的成本问题。

### 2.2 用于 dLLMs 的稀疏注意力

稀疏注意力长期以来一直是降低注意力成本的常用方法，并已在自回归模型中得到了广泛研究。相比之下，其在 dLLMs 中的应用仍然鲜有探索[24（https://arxiv.org/html/2605.20813#bib.bib34），26（https://arxiv.org/html/2605.20813#bib.bib35）]。SparseD 是最早探索用于 dLLMs 的稀疏注意力的工作之一。其核心思想是在去噪步骤间重用稀疏注意力模式，同时在早期迭代中保留全注意力，并采用块级稀疏性来平衡稳定性和效率。除了稀疏注意力方法本身，最近对 dLLM 注意力动态的分析表明，dLLMs 中的注意力表现出独特的动态特性。特别是，[20（https://arxiv.org/html/2605.20813#bib.bib22）] 显示，dLLMs 中的 sink 位置在去噪过程中会发生位移，并且通常与语义或结构 token 相关，这与自回归模型中常见的更稳定的 sink 模式形成对比。总之，这些研究表明，dLLMs 中的注意力既表现出跨步骤的可重用结构，又在精炼过程中表现出动态变化，这使得稀疏化的时机、更新频率和粒度尤为重要。

我们的工作研究如何在迭代的 dLLM 推理过程中有效地稀疏化注意力。与先前在后期迭代中引入稀疏性并依赖块级约束的方法不同，我们利用更细粒度的稀疏模式，在去噪过程的早期阶段开掘可重用的注意力结构。

## 3 方法

### 3.1 预备知识和问题设定

#### 离散扩散语言模型。

dLLMs 将文本生成建模为离散 token 上的迭代去噪过程。给定提示 c = (c₁, …, cₘ) 和目标生成长度 N，响应被初始化为完全掩码状态，形成 x_T = [c₁, …, cₘ, [M], …, [M]]，其中 [M] 表示掩码 token。一个预训练的去噪模型然后通过 T 次反向去噪步骤，逐步将这个完全掩码的响应转换为完整的序列。在步骤 t 时，设 M_t = {i: x_t^i = [M]} 为掩码位置的集合。模型并行地为所有 i ∈ M_t 预测 token 分布 p_θ(x₀^i | x_t)，然后解码策略选择一个子集 U_t ⊆ M_t 进行解掩码，这通常根据生成计划和 token 置信度进行。例如，位置 i 的置信度可以衡量为 max_v p_θ(v | x_t)，因此高置信度的位置会更早被确定，而剩余位置保持掩码状态并在后续步骤中精炼。当所有响应位置都被解掩码时，过程终止。与一次生成一个 token 的自回归语言模型不同，dLLMs 可以在同一个去噪步骤内精炼多个位置。这种并行精炼使得 dLLMs 对于高效生成具有吸引力。然而，由于去噪网络必须在多个步骤中对不断演变的序列进行重复评估，总的推理成本仍然很大，尤其是在序列较长或去噪步骤数量较多时。

#### 用于 dLLM 推理的稀疏注意力。

在去噪步骤 t 时，考虑模型层 l 和头 h 的注意力计算。给定长度为 n = m + N 的隐藏状态，注意力 logits 计算为

Z_t^{l,h} = (Q_t^{l,h} (K_t^{l,h})^⊤) / √d_h, (1)

其中 Q_t^{l,h}, K_t^{l,h}, V_t^{l,h} ∈ ℝ^{n×d_h}。标准全注意力随后产生 A_t^{l,h} = softmax(Z_t^{l,h}) V_t^{l,h}。由于 dLLMs 在 T 个去噪步骤上反复执行全序列双向注意力，总的注意力成本缩放为 O(T L H n² d_h)，其中 L 是层数，H 是头数。这种重复的二次计算使得注意力成为 dLLM 推理中的主要瓶颈。

稀疏注意力通过限制 query-key 交互的集合来降低这一成本。对于每个步骤、层和头，设 S_t^{l,h} ∈ {0,1}^{n×n} 为稀疏注意力掩码，其中 S_t^{l,h}(i,j)=1 表示 query 位置 i 可以关注 key 位置 j。稀疏注意力输出计算为

Ã_t^{l,h} = softmax( Z_t^{l,h} + M(S_t^{l,h}) ) V_t^{l,h}, (2)