提出LoSA机制，让扩散语言模型更高效生成长文本。
LoSA: Locality Aware Sparse Attention for Block-Wise Diffusion Language Models
- 只对变化大的活跃词用稀疏注意力，稳定词复用之前缓存结果。
- 在高稀疏度下准确率提升9个百分点，注意力密度降低1.54倍。
- 适合需要高效长文本生成的研究者和工程落地场景。
块级扩散语言模型（Block-wise DLMs）可无序生成多个词元，是自回归解码的有力替代方案。然而，在长上下文场景中仍受制于内存密集型注意力计算。朴素稀疏注意力因存在KV膨胀问题而失效：不同查询选择不同前缀位置，导致访问的键值页并集过大。我们观察到，连续去噪步骤间仅一小部分活跃词状态显著变化，其余多数词元几乎保持不变。基于此，提出局部感知稀疏注意力（LoSA），对稳定词复用缓存的前缀注意力结果，仅对活跃词应用稀疏注意力。该方法大幅减少需加载的键值索引数量，显著提升速度与精度。在多个块级DLM模型与基准测试中，LoSA在保持接近稠密模型精度的同时大幅提升效率：在激进稀疏度下平均准确率提升高达9点，注意力密度降低1.54倍；在RTX A6000 GPU上实现最高4.14倍注意力加速，验证了方法的有效性。
原文摘要 · Abstract (English)
Block-wise diffusion language models (DLMs) generate multiple tokens in any order, offering a promising alternative to the autoregressive decoding pipeline. However, they still remain bottlenecked by memory-bound attention in long-context scenarios. Naive sparse attention fails on DLMs due to a KV Inflation problem, where different queries select different prefix positions, making the union of accessed KV pages large. To address this, we observe that between consecutive denoising steps, only a small fraction of active tokens exhibit significant hidden-state changes, while the majority of stable tokens remain nearly constant. Based on this insight, we propose LOSA (Locality-aware Sparse Attention), which reuses cached prefix-attention results for stable tokens and applies sparse attention only to active tokens. This substantially shrinks the number of KV indices that must be loaded, yielding both higher speedup and higher accuracy. Across multiple block-wise DLMs and benchmarks, LOSA preserves near-dense accuracy while significantly improving efficiency, achieving up to +9 points in average accuracy at aggressive sparsity levels while maintaining 1.54x lower attention density. It also achieves up to 4.14x attention speedup on RTX A6000 GPUs, demonstrating the effectiveness of the proposed method.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。