Token Sparse Attention: Efficient Long-Context Inference with Interleaved Token Selection

以下是论文《Token Sparse Attention: Efficient Long-Context Inference with Interleaved Token Selection》的详细技术总结。

1. 问题陈述

大型语言模型（LLM）在长上下文推理中面临根本性瓶颈，这是由于预填充（prefill）阶段注意力机制的**二次方复杂度（$O(L^2)$）**所致。尽管 FlashAttention 等硬件优化减少了内存 I/O 开销，但它们并未改变算法复杂度。

现有的加速方法存在两个主要局限性：

• 结构化稀疏性：Minference 或 FlexPrefill 等方法应用块级稀疏性。如果无关 token 与重要 token 被分组在同一块中，这些无关 token 往往会被保留，导致稀疏性次优。
• Token 驱逐：FastKV 或 GemFilter 等方法在早期层永久丢弃 token。这假设 token 重要性是静态的，忽略了 token 相关性在不同层和注意力头之间的动态性质。一旦 token 被驱逐，便无法重新考虑，从而导致精度下降。

2. 方法论：Token Sparse Attention

作者提出了Token Sparse Attention，这是一种轻量级、动态且可逆的 token 级稀疏化机制。与基于驱逐的方法不同，它不永久移除 token，而是压缩计算空间，随后恢复序列。

核心架构

该方法在每个注意力头中分两个阶段运行：

1. 压缩（QKV 选择）：
   • 对于每个注意力头 $h$，根据重要性分数选择一组 token 索引 $S_h$。
   • 收集查询（$Q$）、键（$K$）和值（$V$）张量，形成压缩张量 $\hat{Q}, \hat{K}, \hat{V}$，其序列长度缩减为 $L' \ll L$。
   • 关键特性：选择是按头独立执行的，允许不同的头关注不同的 token 子集，从而解决头间的异质性。
   • 压缩后的张量保持密集且连续，使得可以使用 FlashAttention 等优化内核。
2. 解压缩（输出恢复）：
   • 注意力输出 $\hat{O}$（形状为 $L' \times d$）利用特定于头的索引，散播回原始序列长度 $L \times d$ 的零初始化张量中。
   • 关键特性：未选择的位置被置零，但来自前一层的残差连接保留了被跳过 token 的信息。这使得模型能够在后续层中“重访”并重新评估 token 的重要性。

动态 Token 覆盖

为了确定保留哪些 token 以及保留多少，论文引入了动态 Token 覆盖（Dynamic Token Coverage）：

• 评分：使用最近的查询对所有键计算轻量级代理注意力图。Token 重要性分数通过对注意力权重垂直求和得出。
• 预算分配：该方法不是保留前 $k$ 个 token，而是按升序对 token 进行排序，并修剪“尾部”（最不重要的 token），直到被修剪 token 的累积质量超过覆盖阈值 $\tau$。这将修剪视为针对注意力噪声的结构化正则化。
• 按头选择：在确定全局预算后，每个头根据其特定分数独立选择其前 $k$ 个 token。

稀疏层选择

作者观察到，将稀疏性应用于所有层会降低性能。他们引入了层间表示漂移（Inter-Layer Representation Drift）（$R_\ell$）来衡量 token 表示在层与层之间的稳定性。

• 低漂移（表示稳定）的层被选中进行稀疏化。
• 高漂移（表示发生显著变化）的层保持密集，以保留模型行为。
• 该选择是每模型的一次性预处理步骤。

3. 主要贡献

1. 可逆的 Token 级稀疏化：一种新颖的“压缩与解压缩”设计，在不永久驱逐 token 的情况下降低计算成本，允许跨层动态重新评估 token 重要性。
2. 头间异质性：该方法支持每个注意力头的独立 token 选择，捕捉了不同头优先处理不同语义特征的事实。
3. 无缝兼容性：该方法完全兼容现有的密集（FlashAttention）和稀疏（Minference, FlexPrefill）内核，作为一种互补的加速器而非替代品。
4. 动态预算：一种数据驱动的方法，根据 token 重要性分布和层间稳定性确定稀疏度级别，而非固定比例。

4. 实验结果

实验在 LLaMA-3.1-8B-Instruct 和 Mistral-Nemo-12B-Instruct 上进行，使用了 RULER、InfiniteBench 和 Needle-in-a-Haystack 等基准测试。

• 精度 - 延迟权衡：
  • Token Sparse Attention 始终改善速度 - 精度权衡。
  • 在 128K 上下文下，与 FlashAttention 相比，它实现了高达 $\times 3.23$ 的注意力加速，且精度下降小于 1%。
  • 当与 FlexPrefill 结合使用时，它将加速比从 $\times 2.44$ 提升至 $\times 2.76$，同时保持相同的平均精度（87.27%）。
• 与驱逐方法的比较：
  • 显著优于 token 驱逐方法（PyramidInfer, FastKV, GemFilter）。例如，在 128K 下，Token Sparse Attention 实现了 86.84% 的精度，而 PyramidInfer 在相似加速比下仅为 78.49%。
• 可扩展性：
  • 加速增益随上下文长度增加而增加（例如，4K 时增益最小，128K 时显著），因为注意力稀疏性随序列长度增加而增加。
• 开销：
  • 评分和索引管理的开销极小，即使在高度稀疏的情况下，也仅占总注意力延迟的 <11%。

5. 意义

本文通过将不可逆的 token 驱逐转向动态、交错的稀疏化，解决了长上下文 LLM 推理中的一个关键差距。

• 理论洞察：它证明了 token 重要性在层和头之间高度动态，需要灵活的选择机制而非静态修剪。
• 实际影响：它提供了一种“即插即用”的优化，可以叠加在现有的最先进注意力内核（FlashAttention、Minference 等）之上，在无需重新训练模型或牺牲显著精度的情况下，实现长上下文场景中的巨大加速。
• 未来方向：它为可扩展推理确立了新的设计点，表明与永久 token 移除相比，可逆压缩是处理 $O(L^2)$ 瓶颈的更优策略。