💬 NLP

LoSA: Locality Aware Sparse Attention for Block-Wise Diffusion Language Models

本文针对块状扩散语言模型在长上下文场景下因 KV 膨胀导致的注意力瓶颈，提出了利用稳定 Token 隐藏状态变化微小这一特性的 LOSA 方法，通过复用缓存结果并仅对活跃 Token 进行稀疏注意力计算，在显著降低注意力密度并提升推理速度的同时，实现了甚至优于密集注意力的准确率。

原作者： Haocheng Xi, Harman Singh, Yuezhou Hu, Coleman Hooper, Rishabh Tiwari, Aditya Tomar, Minjae Lee, Wonjun Kang, Michael Mahoney, Chenfeng Xu, Kurt Keutzer, Amir Gholami

发布于 2026-04-15

📖 1 分钟阅读☕ 轻松阅读

CC BY 4.0

原作者： Haocheng Xi, Harman Singh, Yuezhou Hu, Coleman Hooper, Rishabh Tiwari, Aditya Tomar, Minjae Lee, Wonjun Kang, Michael Mahoney, Chenfeng Xu, Kurt Keutzer, Amir Gholami

原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 ✨ 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。阅读完整免责声明

这篇论文介绍了一种名为 LoSA 的新方法，旨在让一种叫做“扩散语言模型”（Diffusion Language Models, DLMs）的 AI 写手在长文章生成时变得更快、更聪明。

为了让你轻松理解，我们可以把 AI 写文章的过程想象成一群画家在合作绘制一幅巨大的长卷画。

1. 背景：AI 是怎么画画的？

传统 AI（自回归模型）： 像是一个按部就班的工匠。他必须一笔一划地写，写完第一个字才能写第二个，写完第二个才能写第三个。虽然慢，但很稳。
扩散模型（DLMs）： 像是一群有创造力的画家。他们不一定要按顺序画，可以一次性把整块区域（比如一个段落）都画出来，甚至可以先画结尾再画开头。这让他们在处理复杂逻辑（比如写故事、做推理）时更有优势。

问题出在哪？
当这幅画变得非常长（比如几万字）时，画家们需要不断回头参考之前画好的部分（这叫"KV 缓存”）。

传统做法： 每次画新的一笔，画家都要把之前所有的画稿都翻一遍，看看哪部分有用。这就像在图书馆里，每写一个字，都要把整个书架的书都搬出来翻一遍，太慢了，而且内存（书架）不够用。

2. 之前的尝试：为什么“偷懒”行不通？

为了加速，研究人员想：“能不能只翻几页书，而不是翻整本书？”这就是稀疏注意力（Sparse Attention）。

** naive 的做法（Naïve Sparse Attention）：** 让每个画家只挑几页书看。
遇到的大坑（KV Inflation）： 想象一下，虽然每个画家只挑几页，但画家 A 挑了第 1 页，画家 B 挑了第 50 页，画家 C 挑了第 100 页。
- 为了照顾所有人，系统不得不把第 1、50、100 页全部搬出来。
- 结果就是：虽然每个人只看了几页，但系统搬运的总页数（并集）并没有减少多少，甚至可能因为大家挑的都不一样，搬运量反而更大了。这就叫 "KV 膨胀”问题。

3. LoSA 的灵感：观察画家的“变化”

论文作者发现了一个有趣的现象：在连续作画的过程中，大部分画家的状态其实没怎么变！

活跃画家（Active Tokens）： 只有少数几个画家正在画关键的新内容，他们的笔触、颜色变化很大。
静止画家（Stable Tokens）： 大多数画家只是在旁边看着，或者画一些背景，他们的状态几乎和上一秒一模一样。

LoSA 的核心思想：

“既然大部分画家的状态没变，我们为什么还要让他们每次都重新去翻书呢？直接复用他们上一轮的结果不就行了吗？”

4. LoSA 是怎么工作的？（三步走）

想象一下 LoSA 是一个聪明的项目经理：

识别谁在动（Locality Pruning）：
项目经理先扫一眼，发现只有 20% 的画家（活跃画家）在剧烈修改画作，剩下 80% 的画家（静止画家）几乎没动。
区别对待：
- 对静止画家： 直接说：“你们别动了，直接复用上一轮的结果！”（这就省去了去图书馆搬书的时间）。
- 对活跃画家： 只有这 20% 的画家需要去翻书。而且，因为人数少了，他们挑出来的书页重叠率更高，系统只需要搬运很少的书页就能满足所有人。
合并结果：
最后把“复用”的结果和“新计算”的结果拼在一起，完成这一轮作画。

5. 效果如何？

速度飞起： 因为不需要搬运那么多书（KV 缓存），在显卡（如 RTX A6000）上，速度提升了 4 倍多！
更聪明： 以前的“偷懒”方法（稀疏注意力）因为搬运了太多不相关的书，导致画家看漏了重要信息，画错了。而 LoSA 让静止画家保留了完整的“记忆”（全量信息），只有活跃画家才用“精简版”记忆。结果就是：既快，又准。
数据证明： 在长文本测试中，LoSA 比之前的方法准确率提高了 9 个百分点，同时内存占用减少了 1.5 倍。

总结

LoSA 就像是一个懂得“抓大放小”的超级管家。

它不再让 AI 对每一个字都重新做一遍繁琐的“查资料”工作，而是敏锐地察觉到：大部分时候，大家只是在“维持现状”，只有少数人在“搞创新”。

对维持现状的，直接给“旧答案”（复用缓存）。
对搞创新的，才给“新任务”（稀疏计算）。

这种方法巧妙地解决了“人多手杂导致效率低”的难题，让 AI 在处理超长文章时，既跑得快，又记得牢。

以下是关于论文 LoSA: Locality Aware Sparse Attention for Block-Wise Diffusion Language Models 的详细技术总结：

1. 研究背景与问题 (Problem)

背景：
块式扩散语言模型（Block-wise Diffusion Language Models, DLMs）通过迭代去噪生成多个 Token，支持块内任意顺序生成，是传统自回归（Autoregressive）解码的有力替代方案，尤其适用于复杂推理任务。

核心痛点：
尽管 DLMs 具有灵活性，但在长上下文场景下，其推理速度仍受限于内存绑定的注意力机制（Memory-bound Attention）。

KV 膨胀问题 (KV Inflation Problem)： 传统的稀疏注意力方法（如 QUEST）是为自回归模型设计的，即每个 Query 独立选择少量的 KV 位置。然而，在 DLM 的块式去噪过程中，同一个块内的不同 Query 往往会选择不同的前缀 KV 位置。
后果： 由于注意力计算是内存绑定的，实际加载的 KV 数据量取决于所有 Query 选中位置的并集（Union）。在块大小为 $B$ 的情况下，这个并集的大小可能膨胀至 $B \times k$ （ $k$ 为单 Query 预算），导致内存访问开销剧增，完全抵消了稀疏性带来的加速效果。

2. 核心观察 (Key Observation)

作者发现块式扩散解码具有自回归模型所不具备的独特结构：表示变化的局部性（Locality of Representation Changes）。

现象： 在连续的去噪步骤（ $t-1$ 到 $t$ ）之间，只有一小部分 Token（称为Active Tokens）的隐藏状态（Hidden States）发生了显著变化（通常是从 [MASK] 变为实际 Token 或其周围区域）。
稳定性： 绝大多数 Token（称为Stable Tokens）的隐藏状态变化极小，其 Query 向量几乎保持不变，因此它们的注意力模式和输出也基本保持恒定。

3. 方法论：LoSA (Methodology)

基于上述观察，作者提出了 LoSA (Locality-aware Sparse Attention)，旨在减少 KV Cache 的内存操作。

核心流程：

局部性剪枝 (Locality Pruning)：
- 计算每个 Token 在连续去噪步骤间的 Query 向量变化（使用均方误差 MSE 作为局部性分数）。
- 根据分数排序，仅保留变化最大的 Top- $k$ 个 Token 作为 Active Tokens 集合 $A$ ，其余视为 Stable Tokens。
缓存复用 (Cache Reuse)：
- 对于 Stable Tokens：直接复用上一轮去噪步骤中计算好的前缀注意力输出（Prefix Attention Outputs）和统计量（Log-normalizers），不再重新计算，也不加载对应的 KV 数据。
- 对于 Active Tokens：仅对这些 Token 应用稀疏注意力选择器（如 QUEST），计算它们需要加载的前缀 KV 索引。
缩小 KV 并集 (Reducing KV Union)：
- 由于只有 $|A|$ 个 Active Tokens 参与稀疏选择，实际加载的 KV 位置并集 $I = \bigcup_{i \in A} I_i$ 远小于传统方法中所有 $B$ 个 Token 的并集。
- 这直接降低了内存带宽需求，从而提升速度。
初始化策略：
- 在块生成的第一个去噪步骤，由于没有缓存，执行一次稠密注意力计算以初始化统计量，后续步骤开始应用 LoSA。

技术细节：

利用 Online Softmax 技术将前缀（Prefix）和当前块（Block）的注意力计算解耦并合并。
额外内存开销极小（仅缓存每个头的 $B \times (d+1)$ 个标量），远小于 KV Cache 本身。

4. 主要贡献 (Key Contributions)

发现 KV 膨胀问题： 首次明确指出在块式 DLM 中直接应用稀疏注意力会导致 KV 索引并集膨胀，从而失效。
提出表示变化局部性： 揭示了去噪步骤间 Token 状态变化的非均匀性，区分了 Active 和 Stable Tokens。
设计 LoSA 算法： 提出了一种结合局部性剪枝和缓存复用的稀疏注意力机制，显著减少了 KV 加载量。
全面的实验验证： 在多个模型（Trado-8B/4B, SDAR-8B）和基准测试（LongBench, 推理任务）上验证了有效性。

5. 实验结果 (Results)

实验在 NVIDIA RTX A6000 和 RTX 5090 GPU 上进行，对比基线包括 QUEST、SparseD 和 Sparse-dLLM。

精度提升 (Accuracy)：
- 在 LongBench 基准测试中，LoSA 在激进稀疏度（Per-Query Budget = 128）下，相比 QUEST 平均提升了 +10.43% 的准确率（Trado-8B 上）。
- 即使在低预算下，LoSA 也能保持接近稠密注意力（Dense）的精度，而 QUEST 精度下降明显。
- 在 SDAR-8B 和 Trado-4B 上均表现出一致的优势。
效率提升 (Efficiency)：
- KV 密度降低： LoSA 的 KV 缓存加载密度比 QUEST 平均低 1.54 倍。
- 加速比： 在 RTX A6000 上，LoSA 实现了 4.14 倍 的注意力加速（相比稠密注意力），而 QUEST 仅为 3.26 倍。
- 硬件扩展性： 在更新的 RTX 5090 上，LoSA 仍保持 3.67 倍 的加速，证明该方法受限于内存带宽而非计算，具有跨代硬件的通用性。
短上下文表现： 在短上下文（<1K tokens）任务中，由于 KV 膨胀问题不明显，LoSA 与 QUEST 表现相当，未出现性能下降。

6. 意义与结论 (Significance)

解决 DLM 落地瓶颈： LoSA 有效解决了块式扩散模型在长上下文推理中的内存瓶颈问题，使其在保持高生成质量的同时具备实用化的推理速度。
新范式： 该方法证明了利用去噪过程中的“状态稳定性”来优化注意力机制是可行的，为扩散语言模型的推理加速提供了新的思路（即从“静态稀疏”转向“动态局部性稀疏”）。
通用性： 该方法不依赖于特定的稀疏选择器（如 QUEST），可以与其他稀疏注意力算法正交结合，具有广泛的适用性。

总结： LoSA 通过识别并利用去噪过程中 Token 状态的局部稳定性，巧妙地规避了块式扩散模型中的 KV 膨胀问题，实现了精度与速度的双重优化，是扩散语言模型长上下文推理领域的重要进展。