将线性注意力引入扩散语言模型,显著提升推理速度且不损失性能。
Retrofitting Linear Attention into Diffusion Language Models

- 用混合注意力机制,仅对历史块使用线性注意力,保持当前块精确计算。
- 在160亿参数模型上仅替换6层注意力,60小时完成改造,性能几乎不变。
- 实测推理吞吐量提升1.7倍,适合需要高并发的部署场景。
扩散语言模型(dLLMs)通过并行解码加速推理,但现有方法在每步去噪时仍需关注所有历史块,造成重复的前缀注意力开销。为此,本文提出块混合注意力:在活跃去噪块内保留精确softmax注意力,对历史块采用线性注意力。该方法可轻量级嵌入预训练模型——LLaDA-Hybrid仅替换20层中的6层,基于LoLCAT方案,耗时约60小时。性能保持良好:HumanEval达72.0%(原75.6%),MBPP+为63.0%(原57.7%),CMATH为86.7%(原88.3%)。使用Triton实现后,解码吞吐量最高提升1.7倍,支持更多并发请求,验证了预训练扩散模型高效线性化的可行性。代码已开源。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。