在自回归（Autoregressive）模型逐字生成文本的范式之外，扩散语言模型（Diffusion Language Models, DLMs）正作为一股新兴力量崛起。它们通过类似图像生成的“去噪”过程，可以并行、无序地生成一个文本块，为复杂的推理任务带来了新的可能性。然而，这种潜力长期被一个核心瓶颈所限制：在处理长文本时，其注意力机制的计算成本和内存访问开销极其高昂。
ArXiv URL：arxiv.org/abs/2604.12…
一个直观的解决方案是借鉴标准大模型中的稀疏注意力（Sparse Attention）技术，即让每个查询（Query）只关注一小部分最相关的键值（Key-Value）。但这在块状扩散模型中却意外“失灵”了。原因在于一个被称为 “KV 膨胀”（KV Inflation） 的问题：在一个生成块中，尽管每个 Token（作为 Query）只稀疏地选择了少量上下文信息，但由于不同 Token 的选择各不相同，GPU 必须加载所有被选中信息的 并集。这个并集可能非常庞大，从而抵消了稀疏化带来的速度优势。
为了破解这一难题，来自加州大学伯克利分校、FuriosaAI 等机构的研究者们提出了 LoSA（Locality-aware Sparse Attention，局部性感知的稀疏注意力）。其核心洞察极为精妙：在扩散模型一轮轮的去噪迭代中，并非所有 Token 的状态都在剧烈变化。恰恰相反，大部分 Token 的内部表示趋于稳定，只有一小部分“活跃”Token 在进行实质性更新。
基于此，LoSA 提出了一种“动静分离”的策略：对那些“稳定”的 Token，直接复用上一步计算好的注意力结果；而只为少数“活跃”的 Token 运行稀疏注意力计算。这一改动极大地缩小了需要从内存中加载的 KV 数据并集，从而在保证（甚至提升）精度的同时，实现了显著的效率飞跃。实验显示，LoSA 在 RTX A6000 GPU 上带来了高达 4.14 倍的注意力计算加速，并在高稀疏度下将模型在长文本基准上的平均准确率提升了 9 个百分点。
扩散模型的“隐形”瓶颈：KV 膨胀问题
块状扩散语言模型（Block-wise DLMs）的生成方式很特别。它们将文本分成块（Block），块与块之间自回归地生成，但在每个块内部，所有 Token 是并行去噪、双向感知的。在每一轮去噪步骤中，块内的每个 Token 都需要与整个已经生成的上下文（即 Prefix KV Cache）进行注意力计算。当上下文变得很长时，反复加载这个庞大的 KV 缓存就成了性能的噩梦。
图 1：KV 膨胀问题示意图。在块状扩散中，每个查询（Query）只选择一小部分前缀 KV 位置（左侧），但实际计算成本取决于所有查询选择的 KV 位置的并集（右侧），这导致了 KV 访问的膨胀。LoSA 通过只对活跃 Token 计算稀疏注意力，显著减小了并集的大小和延迟。
稀疏注意力本是解决这个问题的良药，它允许每个查询只关注成千上万个上下文 Token 中的一小部分，比如 128 或 256 个。但在块状扩散模型中，问题出在“块”这个并行单元上。假设一个块有 16 个 Token，每个 Token 都独立地选出了自己最关心的 128 个上下文位置。由于大家的“兴趣点”不同，为了服务这 16 个 Token，系统最终需要加载的独立上下文位置的并集，可能远不止 128 个，最坏情况下甚至可能接近 。这个现象就是“KV 膨胀”，它使得内存访问的开销依然居高不下，稀疏化的优势大打折扣。
破局点：去噪过程中的“动静分离”
LoSA 的作者们发现，扩散模型独特的迭代去噪过程，提供了一个自回归模型所不具备的结构性机会：表征变化的局部性（Locality of Representation Changes）。
在一个去噪步骤到下一个步骤之间，模型可能只将块中的少数几个 [MASK] 标记转变为具体的词元。这些新生成的词元及其邻近位置的内部表征（hidden states）会发生剧烈变化，研究者将它们定义为 “活跃”Token（Active tokens）。而块中其他大部分 Token 的表征则几乎保持不变，被称为 “稳定”Token（Stable tokens）。
图 2：去噪步骤间表征变化的局部性可视化。左图和中图分别显示了 和 时刻的查询向量 L2 范数，右图则展示了两个时刻之间查询向量的均方误差（MSE）。可以看出，只有一小部分 Token 发生了大的变化，这启发了为稳定 Token 复用缓存前缀注意力的想法。
研究者通过计算连续两个去噪步骤之间，每个 Token 的查询向量（Query Vector）的均方误差（MSE）来量化这种变化。如下图所示的层级热力图清晰地揭示了这一现象：在每一层，将 Token 按 MSE 变化从大到小排序后，绝大部分变化都集中在顶部的少数 Token 上。
图 4：跨去噪步骤的局部性可视化热力图。图中每一行代表一个 Token，按 MSE 变化降序排列。可以看到，大部分变化量（图中显示了总变化量 50% 的阈值线）都由少数 Token 贡献，验证了“动静分离”的可行性。
这一发现是 LoSA 方法的基石。既然大部分 Token 的查询向量都没怎么变，那么它们对长文本上下文的注意力模式和最终的注意力输出，也理应是稳定的。因此，完全没有必要为它们一遍又一遍地重新计算。
LoSA 方法：为“稳定”Token 缓存结果，为“活跃”Token 精算注意力
基于“动静分离”的洞察，LoSA 的工作流程被设计得非常高效，其核心思想是区别对待、按需计算。
图 5：LoSA 工作流程概览。① 识别活跃 Token；② 仅为活跃 Token 运行稀疏选择器并取并集；③ 为活跃 Token 计算新的前缀注意力，为稳定 Token 复用缓存；④ 计算块内注意力；⑤ 合并结果并为下一步缓存。
整个流程可以分解为以下几个步骤：
-
局部性剪枝（Locality Pruning）：在每个去噪步骤，计算所有块内 Token 的查询向量与上一步的 MSE，将变化最大的 Top-K 个 Token（例如 K=5）识别为“活跃”集合 ，其余则为“稳定”集合。
-
为稳定 Token 复用结果：对于所有稳定 Token，LoSA 不执行任何新的前缀注意力计算。它直接从上一步缓存中取出已经计算好的前缀注意力输出 和对数归一化项 。这不仅极快，还保留了完整的注意力信息。
-
为活跃 Token 计算稀疏注意力：只对“活跃”集合 中的 Token 运行稀疏注意力算法（如 QUEST）。由于参与计算的 Token 数量从整个块的大小 锐减到 ，它们所选出的 KV 索引的并集 也随之大幅缩小。
-
合并与更新：加载缩减后的 KV 缓存数据，为活跃 Token 计算新的前缀注意力结果。然后，将所有 Token（包括复用结果的稳定 Token 和新计算的活跃 Token）的前缀注意力结果与块内（Suffix）注意力结果通过 online-softmax 方法合并，得到最终输出。最后，更新并缓存所有 Token 的前缀注意力统计数据，以备下一步使用。
值得一提的是，LoSA 在处理一个新块的第一次去噪迭代时，会执行一次完整的稠密前缀注意力。这确保了初始缓存的质量极高，后续的稳定 Token 可以从中受益，继承精确的全局上下文信息。正是这种“缓存高质量结果”与“只为必要部分计算”的结合，让 LoSA 在加速的同时，甚至比天真地对所有 Token 都做稀疏注意力的方案精度更高。
实验结果：更高精度、更快速度
LoSA 的有效性在多个模型和基准测试中得到了验证。
长文本处理能力（LongBench）
在 LongBench 基准上，LoSA 的优势尤为突出。与同样采用稀疏注意力的基线方法 QUEST 相比，LoSA 在所有预算配置下都取得了更高的准确率。特别是在非常激进的稀疏预算下（例如每个查询只关注 128 个位置），LoSA 的平均准确率比 QUEST 高出超过 10 个百分点。这证明了其为稳定 Token 保留完整注意力信息的核心机制，有效避免了稀疏化带来的信息损失。
推理速度
在性能方面，LoSA 展现了显著的加速效果。在 RTX A6000 GPU 上，针对 64K 上下文长度和 16 的块大小，LoSA 的前缀注意力计算速度是稠密注意力的 4.14 倍，也明显优于 QUEST 的 3.26 倍。
图 6：前缀注意力延迟分解比较。无论是在 64K 上下文（左）还是 32K 上下文（右）的设置下，LoSA 的延迟都最低，其速度优势主要源于通过局部性感知识别和复用，大幅减少了内存密集型的注意力操作。
这种加速趋势也成功延续到了更新的硬件上（如 RTX 5090），证明了该方法的普适性。速度提升的根本原因在于，注意力计算是内存带宽受限的操作，LoSA 通过大幅减少需要加载的 KV 数据量，直接命中了性能要害。
方法总结
LoSA 并非又一个通用的稀疏注意力算法，而是为块状扩散语言模型量身定做的优化方案。它敏锐地捕捉到了 DLM 迭代去噪过程中“表征变化局部性”这一独特现象，并以此为基础设计了“动静分离”的计算策略。
通过为绝大多数“稳定”的 Token 复用高质量的缓存结果，同时仅为少数“活跃”的 Token 精准计算稀疏注意力，LoSA 不仅成功破解了“KV 膨胀”这一棘手难题，更在大幅提升推理速度的同时，带来了令人意外的准确率增益。这项工作为提升扩散语言模型在长上下文场景下的实用性铺平了道路，也再次证明了深入理解特定模型架构的内在机理是通往极致优化的关键所在。
