Token Sparse Attention：通过交错Token选择实现高效长上下文推理 Token Sparse Attention: Efficient Long-Context Inference with Interleaved Token Selection
通过动态token级稀疏化，在128K上下文下实现3.23倍注意力加速，精度损失小于1%
前置知识
注意力机制 (Attention Mechanism)
Transformer架构的核心组件，通过计算查询(Query)与键(Key)之间的相似度来加权值(Value)。标准注意力的计算复杂度为O(L²d)，其中L是序列长度，d是维度。在长上下文场景下，这种二次复杂度成为主要瓶颈，导致计算时间和内存消耗随序列长度平方增长。
理解注意力机制的二次复杂度问题是本文动机的核心，Token Sparse Attention正是为了解决这一瓶颈而提出的。
稀疏注意力 (Sparse Attention)
通过跳过注意力图中低重要性区域的计算来加速注意力的方法。现有方法包括块稀疏（如FlexPrefill）和结构化模式（如Minference的Vertical-Slash模式）。这些方法通过预定义的稀疏模式或动态块选择来减少计算量，但存在粒度较粗的问题。
Token Sparse Attention是在稀疏注意力基础上提出的token级稀疏化方法，理解现有稀疏注意力的局限性有助于把握本文的创新点。
Token驱逐 (Token Eviction)
在模型的特定层选择重要token子集并永久丢弃其余token的方法。代表工作包括FastKV、GemFilter和PyramidInfer。这些方法在早期层估计token重要性后，在后续层中只保留选中的token，被丢弃的token无法被重新考虑。
Token驱逐是本文主要对比的基线方法，理解其"永久丢弃"的特性是理解Token Sparse Attention"可逆设计"优势的关键。
KV缓存 (KV Cache)
在自回归生成中，为了避免重复计算，将之前token的键(Key)和值(Value)缓存起来的技术。在长上下文场景下，KV缓存会占用大量内存，成为解码阶段的主要瓶颈。
虽然本文主要关注预填充阶段的加速，但KV缓存是长上下文推理的另一重要瓶颈，理解这一点有助于全面评估方法的适用场景。
FlashAttention
一种硬件感知的精确注意力实现，通过优化内存访问模式（如分块计算、减少HBM访问）来显著减少注意力计算的内存I/O开销。FlashAttention不改变注意力的数学形式，而是通过更好的工程实现来加速计算。
Token Sparse Attention是在FlashAttention基础上构建的，理解FlashAttention的工作原理有助于理解本文方法如何与现有高效注意力实现兼容。
研究动机
在长上下文推理场景中，标准注意力机制的O(L²d)计算复杂度成为主要瓶颈。当上下文长度L超过100K时，即使使用FlashAttention这样的硬件优化实现，注意力计算仍然占据推理总时间的绝大部分。现有加速方法主要分为两类：结构化稀疏注意力和token驱逐。结构化稀疏方法（如Minference、FlexPrefill）通过预定义的块级稀疏模式跳过部分计算，但由于稀疏粒度较粗（块级），当重要token与不重要token被分在同一个块时，仍然会保留不相关的计算。Token驱逐方法（如FastKV、GemFilter、PyramidInfer）在早期层选择重要token并永久丢弃其余token，但这种方法存在两个关键问题：一是token重要性在不同层之间会发生显著变化，早期层的选择决策可能在后续层变得不适用；二是所有注意力头被迫共享相同的token集合，忽略了不同头关注不同语义特征的异质性。
本文的目标是本文的目标是设计一种动态的、可逆的token级稀疏化机制，在保持模型精度的同时显著加速长上下文推理。具体而言，方法需要满足三个要求：一是支持token级粒度的选择，比块级稀疏更精细；二是允许不同层和不同头独立选择token子集，捕捉注意力模式的异质性；三是保持可逆性，被跳过的token在后续层可以被重新考虑。量化目标是在128K上下文下实现2-3倍的注意力加速，同时将精度损失控制在1%以内。
与已有工作不同的是，本文的独特切入点在于发现了token重要性的两个关键动态特性。通过分析LLaMA-3.1-8B-Instruct模型，作者发现：第一，token重要性在层间存在显著变化，相邻层的重要token重叠度仅为中等水平，随着层距离增加重叠度快速下降；第二，同一层内不同注意力头关注的token集合存在明显差异，不同头优先关注不同的token子集。这两个发现表明，token选择应该在层和头两个维度上保持灵活性，而不是像现有方法那样做出永久性、统一的选择。这一观察直接启发了"压缩然后解压缩"的设计理念。
核心方法
Token Sparse Attention的核心思想可以类比为"信息筛选器"：在每一层的每个注意力头中，动态选择最重要的token进行注意力计算，然后将计算结果"解压缩"回原始序列长度，确保被跳过的token在后续层仍有机会被重新考虑。这种设计类似于一个可逆的过滤过程：先压缩（选择重要token），再解压缩（恢复序列维度），使得token信息能够在层间流动。技术路线包括三个关键组件：Token Sparse Attention主体（压缩-解压缩机制）、Dynamic Token Coverage（动态token选择策略）和Sparse Layer Selection（基于表示漂移的层选择策略）。整个方法完全兼容FlashAttention等高效注意力实现，可以无缝集成到现有稀疏注意力框架中。
Token Sparse Attention最本质的创新在于"交错token选择"设计，即通过压缩-解压缩机制实现token级稀疏化的同时保持可逆性。与现有token驱逐方法的关键区别在于：驱逐方法在早期层做出永久性丢弃决策，被丢弃的token在后续层无法被重新考虑；而Token Sparse Attention在每个注意力操作后都将输出解压缩回原始序列维度，使得每个头在每一层都可以从完整的token序列中重新选择自己的重要token子集。这种设计直接解决了前文发现的两个动态性问题：层间动态性（不同层关注不同token）和头间动态性（不同头关注不同token）。更重要的是，由于压缩后的Q、K、V张量在内存中保持连续，可以直接使用FlashAttention等高效实现，无需修改底层kernel。
方法步骤详情
Token Sparse Attention的执行流程分为两个主要阶段和一个预处理步骤。第一阶段是压缩：对于每个注意力头h，根据Dynamic Token Coverage策略选择的token索引集合S_h，从原始Q、K、V张量中提取对应的行，形成压缩后的张量\hat{Q}、\hat{K}、\hat{V}，序列长度从L压缩到L'（L' << L）。由于选择是按头独立进行的，不同头可以选择不同的token子集。第二阶段是注意力计算：在压缩后的张量上执行标准注意力操作（使用FlashAttention），得到压缩维度的输出\hat{O}。第三阶段是解压缩：将\hat{O}通过scatter操作放回零初始化的L×d张量中，未被选中的位置保持零值，然后与残差连接相加。Dynamic Token Coverage策略包括：使用最近查询计算轻量级注意力分数、聚合头间分数得到层级重要性分布、按升序排序token并找到累积质量超过阈值 au的最小不重要token集合、按头选择top-k重要token。Sparse Layer Selection基于层间表示漂移指标R_l，选择漂移值低于阈值\delta的层应用稀疏化。
技术新颖性
Token Sparse Attention在技术新颖性上体现在三个方面。首先是压缩-解压缩的可逆设计：与驱逐方法的"不可逆丢弃"不同，本文通过在每个注意力操作后恢复序列维度，实现了token信息的持续流动。这种设计虽然增加了scatter操作的开销，但保留了后续层重新选择token的可能性。其次是头独立的token选择：现有方法通常在所有头间共享统一的token集合，而本文允许每个头独立选择，直接捕捉注意力模式的头间异质性。第三是与现有稀疏注意力方法的互补性：由于Token Sparse Attention在QKV级别操作，而块稀疏方法在注意力图级别操作，两者可以无缝组合，实现"异构粒度"的稀疏化策略。这种互补性在实验中得到了验证：与FlexPrefill组合后，在相同精度下实现了更高的加速比。
实验结果
实验结果表明Token Sparse Attention在多个维度上都表现出色。在RULER基准测试上，当与FlexPrefill组合应用于LLaMA-3.1-8B-Instruct时，平均精度保持在87.27%，与原始FlexPrefill完全一致，但128K注意力加速比从2.44倍提升到2.76倍。在更激进的稀疏设置下（τ=0.01），方法仍能保持86.84%的平均精度，实现1.51倍加速。与token驱逐方法对比时，在相似加速比下（约1.5倍），Token Sparse Attention达到86.84%的平均精度，而PyramidInfer为78.49%，FastKV为85.64%，GemFilter为85.12%。特别值得注意的是，在128K长度下，本文方法精度为73.25%，而PyramidInfer仅为69.35%，差距达到3.9个百分点。Dynamic Token Coverage相比固定稀疏策略也表现出明显优势：在相似稀疏度下，动态策略保持87.02%精度，而固定策略仅为86.91%；在更激进的稀疏下，动态策略86.84%对比固定策略85.43%。延迟分析显示，在128K上下文下，Token Sparse Attention引入的额外开销仅占总注意力延迟的不到11%，在A100 GPU上实现24.35秒的预填充延迟，相比FlashAttention的31.04秒节省了21.6%的时间。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| RULER基准测试（LLaMA-3.1-8B-Instruct） | 平均精度(%) | 87.02 | 87.01 (FlashAttention) | 精度基本不变，加速1.36倍 |
| RULER基准测试（128K长度） | 精度(%) | 73.68 | 74.15 (FlashAttention) | 精度损失0.47%，加速1.36倍 |
| 与FlexPrefill组合（LLaMA-3.1-8B-Instruct） | 平均精度(%) | 87.27 | 87.27 (FlexPrefill) | 精度不变，加速从2.44倍提升到2.76倍 |
| 与token驱逐方法对比（128K） | 精度(%) | 73.25 | 69.35 (PyramidInfer) | 精度提升3.9% |
| InfiniteBench基准测试 | 平均精度(%) | 50.88 | 50.86 (FlashAttention) | 精度基本不变 |
| 预填充延迟（A100, 128K） | 延迟(秒) | 24.35 | 31.04 (FlashAttention) | 延迟减少21.6% |
局限与改进
作者在论文中承认了几个局限性。首先，在短上下文场景下，由于注意力计算不占主导地位，Token Sparse Attention的加速效果有限。其次，过于激进的token稀疏化可能增加排除语义重要token的风险，导致精度下降。此外，本文主要关注预填充阶段的加速，解码阶段仍使用标准密集注意力。从独立观察来看，方法的额外开销虽然不大（<11%），但在短序列上可能变得相对显著。Dynamic Token Coverage策略依赖于最近查询的注意力分数，这可能不适用于所有任务类型。层选择阈值δ=0.5是通过实验确定的经验值，缺乏理论指导。最后，方法在Mistral-Nemo-12B-Instruct上的表现不如LLaMA-3.1-8B-Instruct稳定，表明对不同架构的泛化性可能有限。
独立分析的弱点
从独立分析来看，Token Sparse Attention存在几个值得关注的弱点。首先是token评分机制的局限性：当前方法使用最近查询的注意力分数作为重要性代理，但在某些任务中（如需要全局理解的任务），最近查询可能无法准确反映token的整体重要性。改进方向可以探索多尺度评分策略，结合局部和全局查询信息。其次是层选择策略的刚性：基于表示漂移的层选择是静态的，对所有输入使用相同的层子集。可以考虑输入自适应的层选择，根据具体输入动态决定哪些层应用稀疏化。第三是头独立选择的计算开销：虽然实现了头间异质性，但每个头独立执行top-k选择会增加索引操作的开销。可以探索分组选择策略，在相近的头间共享token集合。最后是稀疏度预算的分配：当前方法使用固定的token覆盖率阈值τ，可以考虑基于层深度或输入复杂度的自适应预算分配。
未来方向
论文作者提出了几个有前景的未来方向。首先是将Token Sparse Attention扩展到解码阶段，当前方法仅针对预填充加速，解码阶段的KV缓存压缩仍是开放问题。其次是探索多模态场景的应用，如视觉-语言模型中的token选择。基于本文成果还可以延伸几个方向：一是将token级稀疏化与其他优化技术（如量化、知识蒸馏）结合，探索多维度的效率优化；二是研究token重要性动态变化的根本原因，为设计更鲁棒的选择策略提供理论指导；三是探索Token Sparse Attention在不同任务类型（如检索、推理、生成）中的表现差异，开发任务自适应的稀疏化策略。此外，将方法扩展到分布式推理场景，研究跨设备的token选择协调也是有意义的方向。
复现评估
从复现角度来看，Token Sparse Attention具有较好的可复现性。作者提供了开源代码（https://github.com/dongwonjo/Token-Sparse-Attention），实现了基于Triton的自定义kernel。实验设置清晰：使用公开可用的LLaMA-3.1-8B-Instruct和Mistral-Nemo-12B-Instruct模型，在RULER、InfiniteBench、LongBench等标准基准上进行评估。算力需求方面，所有实验在单块NVIDIA A100 80GB GPU上完成，对于大多数研究团队来说是可接受的。复现难度中等：核心算法逻辑清晰，但需要理解Triton编程和FlashAttention的集成细节。token覆盖率参数τ（LLaMA用0.005，Mistral用0.008）和层选择阈值δ=0.5需要根据具体模型调整。总体而言，具备中等GPU资源和注意力机制背景的研究者应该能够复现本文结果。
论文图表
(a)展示重要token在层间的重叠度，显示token重要性随层距离增加而快速变化。(b)展示同一层内不同注意力头关注的token排名差异，说明token重要性是头特定的。
这张图提供了Token Sparse Attention设计动机的关键证据，解释了为什么需要交错token选择。