Recursive Sparse Attention：以对数成本实现百万词元长度外推

arXiv:2611.07232v1 公告类型：交叉发布 摘要：在约 100 万词元的上下文长度下，稠密注意力在预填充和解码阶段的成本高昂，而且训练通常依赖长文档，而随着目标长度增加，长文档也愈发稀缺。大多数稀疏注意力方法会对所有先前的块或词元进行评分，因此，在上下文长度为 $L$ 时，每个查询的选择成本为 $O(L)$；分层选择可将其降至 $O(\log L)$，但此前可训练的分层方法仅报告了训练长度范围内的准确率。我们提出 RecSA（Recursive Sparse Attention），它在 KV 缓存上构建一棵树，其中每个节点使用一个键和一个均值池化的值来概括其子树。这些键由每个模型层中的一个池化器生成，该池化器由所有树层级共享，并通过稠密注意力蒸馏得到。词元之上的每个层级最多选择 $k$ 个候选项，因此每个查询的选择成本为 $O(\log L)$，并读取 $O(\log L)$ 个条目；未选中候选项的摘要与精确的词元共同参与同一次 softmax 运算。每个条目的 RoPE 位置是其在所属层级候选项中的索引，因此，无论长度如何，所有相对距离都保持在训练期间见过的范围内。RecSA 以 Qwen3-4B 初始化，并在原模型的 32K 训练长度下使用 6.2 亿词元进行训练，在 RULER 上从 131K 到 2M 的所有长度下均优于全部基线；在 1M 长度下，其得分为 60.5，而最佳基线 Dense + DCA 的得分为 40.0。借助我们优化的计算内核，在 64K 至 2M 的长度范围内，RecSA 的端到端预填充和解码速度均快于稠密注意力；在单张 H200 GPU 上、长度为 524K 时，其预填充速度提升至 3.2 倍，解码速度提升至 3.8 倍。