LOSA：基于局部性的稀疏注意力用于块状扩散语言模型
计算与语言
2026-04-15 v1 机器学习
摘要
块状扩散语言模型（DLMs）以任意顺序生成多个标记，为自回归解码管道提供了有前景的替代方案。然而，他们仍在长情境场景中受限于内存受限注意力。由于KV膨胀问题，朴素的稀疏注意力在DLMs上无法正常工作——不同查询选择不同的前缀位置，导致访问KV页面的并集变大。为解决此问题，我们观察到在连续去噪步骤之间，只有少数活跃标记显示出显著的隐藏状态变化，而大多数稳定标记几乎保持不变。基于这一洞察，我们提出了LOSA（局部感知稀疏注意力），它为稳定标记复用已缓存的前缀注意力结果，只对活跃标记应用稀疏注意力。这显著缩小了必须加载的KV索引数量，实现更高的加速速度和更高准确率。在多种块状DLMs和基准测试中，LOSA在保持接近稠密准确率的同时显著提高效率，在激进稀疏水平下平均准确率提升最高达9分，同时保持1.54倍更低的注意力密度。它还在RTX A6000 GPU上实现了最高4.14倍的注意力加速，展示了所提出方法的有效性。
引用
@article{arxiv.2604.12056,
title = {LoSA: Locality Aware Sparse Attention for Block-Wise Diffusion Language Models},
author = {Haocheng Xi and Harman Singh and Yuezhou Hu and Coleman Hooper and Rishabh Tiwari and Aditya Tomar and Minjae Lee and Wonjun Kang and Michael Mahoney and Chenfeng Xu and Kurt Keutzer and Amir Gholami},
journal= {arXiv preprint arXiv:2604.12056},
year = {2026}
}
备注
16 pages, 11 figures, 6 tables