提出稀疏注意力机制,让扩散语言模型推理更快且不损失质量。
SparseD: Sparse Attention for Diffusion Language Models
- 针对不同注意力头设计专用稀疏模式,一次预计算重复使用。
- 早期去噪步骤用全注意力保质量,后期转为稀疏加速。
- 在64k上下文下速度提升1.5倍,生成效果无损,适合长文本场景。
尽管扩散语言模型(DLMs)为自回归模型提供了有前景的替代方案,但现有开源DLMs存在高推理延迟问题,主要源于注意力机制随上下文长度呈二次复杂度。为降低复杂度,自然策略是限制注意力至稀疏模式,仅保留最相关连接。这类方法在自回归模型中已有成熟应用,但DLMs表现出独特稀疏行为:(1) 各注意力头的模式不同,(2) 每个头的模式在不同去噪步骤间高度一致,(3) 早期去噪步骤对生成至关重要。这使得为自回归模型设计的稀疏注意力方法难以直接应用于DLMs,因无法捕捉头特异性结构,且可能在早期步骤中损害生成质量。为此,我们提出SparseD,一种专用于DLMs的新型稀疏注意力方法。基于上述观察,SparseD仅需一次预计算各头的稀疏模式并跨所有步骤复用,避免每步重算;同时在早期去噪阶段采用全注意力以保障质量,后期切换至稀疏注意力以实现加速。实验表明,SparseD实现无损加速,在64k上下文长度、1,024个去噪步骤下,相比FlashAttention最高提速1.50倍。
原文摘要 · Abstract (English)
While diffusion language models (DLMs) offer a promising alternative to autoregressive models (ARs), existing open-source DLMs suffer from high inference latency. This bottleneck is mainly due to the attention's quadratic complexity with respect to context length in computing all query-key pairs. Intuitively, to reduce this complexity, a natural strategy is to restrict attention to sparse patterns that retain only the most relevant connections. Such approaches are well-established in ARs, where attention follows fixed and clearly defined sparse patterns. However, in DLMs, we observe distinct sparsity behaviors: (1) attention patterns vary across heads, (2) attention patterns in each head remain highly similar across denoising steps, and (3) early denoising steps are critical for generation. These findings render sparse attention methods designed for ARs largely incompatible with DLMs, as they fail to capture head-specific structures and risk degrading generation when applied in early denoising steps. To address these challenges, we propose SparseD, a novel sparse attention method for DLMs. Leveraging the observations, SparseD only requires pre-computing head-specific sparse patterns one time, and reuses them across all steps. This prevents recomputing sparse patterns at each denoising step. Meanwhile, SparseD uses full attention in the early steps, then switches to sparse attention later to maintain generation quality. Together, these establish SparseD as a practical and efficient solution for deploying DLMs in long-context applications. Experimental results demonstrate that SparseD achieves lossless acceleration, delivering up to $1.50\times$ speedup over FlashAttention at a 64k context length with 1,024 denoising steps.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。