将线性注意力改造融入扩散语言模型
Retrofitting Linear Attention into Diffusion Language Models
浏览论文内容
中文总结 AI 辅助
本文提出分块混合注意力，将其融入预训练dLLM LLaDA~2.1得到LLaDA-Hybrid，在保持基准性能的同时提升解码吞吐量与并发请求支持能力。
中文摘要 AI 辅助
扩散语言模型（dLLMs）是自回归模型的有前景替代方案，可通过并行解码加速推理。近期的dLLMs通常采用分块半自回归解码，即自回归生成分块，同时在每个活跃分块内并行去噪标记。然而，即便使用KV缓存，每次去噪步骤仍会关注所有先前分块，反复产生前缀注意力成本。受此瓶颈启发，本文探究能否通过对先前分块应用线性化注意力进一步加速dLLM推理。本文提出分块混合注意力，在活跃去噪分块内保留精确的softmax注意力，同时对先前分块应用线性注意力。研究表明，该混合注意力可通过极少的后训练改造融入预训练dLLM：LLaDA-Hybrid替换了16B开源dLLM LLaDA~2.1的20个注意力层中的6个，大致遵循LoLCAT（Zhang等人，2024）的方法。转换仅需约60小时，同时保持基准性能：在HumanEval上为72.0%对比75.6%，在MBPP+上为63.0%对比57.7%，在CMATH上为86.7%对比88.3%。借助Triton实现，LLaDA-Hybrid的解码吞吐量最高提升1.7倍，且在内存耗尽前支持更多并发请求，表明预训练dLLMs可被高效线性化以实现更快推理。本文代码可访问：this https URL。
英文摘要
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.
发表机构
- Apple（苹果公司）
- Google DeepMind（谷歌DeepMind）
- Harvard University（哈佛大学）
- MIT（麻省理工学院）
机构由 AI 辅助整理，请以论文原文为准。
