Retrofitting Linear Attention into Diffusion Language Models

📄 arXiv: 2608.06628v1 📥 PDF

作者: Jinha Kim, Younghun Roh, Jaeyeon Kim

分类: cs.LG

发布日期: 2026-08-06

🔗 代码/项目: GITHUB

💡 一句话要点

提出块混合注意力以加速扩散语言模型推理

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)

关键词: 扩散语言模型 块混合注意力 推理加速 自然语言处理 模型优化

📋 核心要点

  1. 现有的扩散语言模型在推理过程中仍需处理前缀注意力成本，导致效率低下。
  2. 本文提出块混合注意力，结合活跃块的softmax注意力与先前块的线性注意力，以加速推理过程。
  3. LLaDA-Hybrid在解码吞吐量上提升了1.7倍，并在多个基准测试中保持了高性能，展示了其有效性。

📝 摘要（中文）

扩散语言模型（dLLMs）通过并行解码为自回归模型提供了有前景的替代方案。现有的dLLMs通常采用块式半自回归解码，尽管使用了KV缓存，但每个去噪步骤仍需关注所有先前块，导致前缀注意力成本的重复。为了解决这一瓶颈，本文提出了块混合注意力，保留活跃去噪块内的精确softmax注意力，同时对先前块应用线性注意力。实验表明，LLaDA-Hybrid在保留基准性能的同时，解码吞吐量提高了1.7倍，支持更多并发请求，展示了预训练dLLMs的高效线性化潜力。

🔬 方法详解

问题定义：现有的扩散语言模型在推理时，尽管采用KV缓存，仍需在每个去噪步骤中关注所有先前块，导致前缀注意力的重复计算，影响推理效率。

核心思路：本文提出的块混合注意力方法，通过在活跃去噪块中保留精确的softmax注意力，同时对先前块应用线性注意力，从而减少计算复杂度，提高推理速度。

技术框架：整体架构包括预训练的dLLM模型，采用块混合注意力替换部分注意力层。具体而言，LLaDA-Hybrid替换了LLaDA~2.1中的6个注意力层，保持了模型的整体性能。

关键创新：最重要的创新在于块混合注意力的引入，它有效结合了传统的softmax注意力与线性注意力，显著降低了计算成本，并保持了模型的性能。

关键设计：在参数设置上，LLaDA-Hybrid的转换过程仅需约60小时，使用Triton实现以提高解码吞吐量，确保在内存使用上支持更多的并发请求。实验结果显示，模型在HumanEval、MBPP+和CMATH等基准测试中表现优异。

🖼️ 关键图片

fig_0
fig_1

📊 实验亮点

LLaDA-Hybrid在解码吞吐量上实现了高达1.7倍的提升，同时在HumanEval、MBPP+和CMATH基准测试中，分别保持了72.0%对75.6%、63.0%对57.7%和86.7%对88.3%的性能，展示了其有效性和优越性。

🎯 应用场景

该研究的潜在应用领域包括自然语言处理、对话系统和代码生成等。通过提高扩散语言模型的推理效率，能够在实时应用中实现更快的响应时间，提升用户体验。此外，该方法的高效性也为大规模模型的部署提供了可能，具有重要的实际价值和未来影响。

📄 摘要（原文）

Diffusion language models (dLLMs) offer a promising alternative to autoregressive models by accelerating inference through parallel decoding. Recent dLLMs commonly use blockwise semi-autoregressive decoding, generating blocks autoregressively while denoising tokens within each active block in parallel. However, despite KV caching, each denoising step still attends to all previous blocks, repeatedly incurring prefix-attention cost. Motivated by this bottleneck, we ask whether dLLM inference can be further accelerated by linearizing attention over previous blocks. We introduce block-hybrid attention, which retains exact softmax attention within the active denoising block while applying linear attention over previous blocks. We show that this hybrid attention can be retrofitted into a pretrained dLLM with minimal post-training: LLaDA-Hybrid replaces 6 of the 20 attention layers in LLaDA~2.1, a 16B open-source dLLM, largely following LoLCAT (Zhang et al, 2024). The conversion takes only approximately 60 hours while preserving benchmark performance: 72.0% vs. 75.6% on HumanEval, 63.0% vs. 57.7% on MBPP+, and 86.7% vs. 88.3% on CMATH. With a Triton implementation, LLaDA-Hybrid achieves up to $1.7\times$ higher decoding throughput and supports more concurrent requests before exhausting memory, showing that pretrained dLLMs can be efficiently linearized for faster inference. Our code is available at: https://github.com/Diuven/LLaDA-Hybrid.