[控场AI]
· 5 分钟阅读· 2,874 字

递归稀疏注意力RecSA:百万级上下文外推,对数级成本

递归稀疏注意力RecSA:百万级上下文外推,对数级成本

RecSA用树结构压缩KV缓存,以对数级成本实现从32K到2M的上下文外推。

RecSA(递归稀疏注意力)针对百万级上下文下稠密注意力开销过大的问题,在KV缓存上构建一棵树,每个节点用key和均值池化value概括其子树,各层共享一个从稠密注意力蒸馏得到的pooler。每层最多选择k个候选,使每个查询的选择与读取成本均为O(log L)。其位置编码设计让每个条目的RoPE位置取其在同层候选中的索引而非原始token位置,从而把相对距离始终限制在训练分布内,实现从32K训练长度外推到2M。在RULER基准上,131K至2M各长度均优于基线,1M长度得分60.5对比最佳基线40.0;单张H200上524K长度预填充加速3.2倍、解码加速3.8倍。方法从Qwen3-4B初始化,仅用0.62B token在32K长度训练。

递归稀疏注意力RecSA:百万级上下文外推,对数级成本

当上下文长度扩展到百万token级别,稠密注意力在预填充和解码阶段的开销变得难以承受,而训练所需的长文档随着目标长度增加也愈发稀缺。现有稀疏注意力方法大多需要对所有前序块或token进行打分,导致每个查询的选择成本随上下文长度L线性增长;分层选择虽能将成本降至对数级,但此前可训练的分层方法仅在训练长度内报告了准确率。

核心思路:用树结构压缩KV缓存

据原论文描述,RecSA(Recursive Sparse Attention)在KV缓存上构建一棵树,每个节点用一个key和一个均值池化的value来概括其子树。这些key来自每层一个pooler,所有层级共享同一个pooler,并从稠密注意力中蒸馏得到。

这种设计的关键在于:树中每一层最多选择k个候选,因此每个查询的选择成本为O(log L),读取的条目数同样为O(log L)。未被选中候选的摘要与精确token共享同一个softmax,从而在保持稀疏性的同时不丢失全局信息。

RecSA架构示意图

KV缓存(KV Cache) 是Transformer解码阶段的核心优化:自回归生成时,已计算过的Key和Value向量会被缓存复用,避免每生成一个新token都重算全部历史。但缓存大小随上下文长度线性增长,百万token时单条序列的KV缓存可达数十GB,显存成为硬约束。稀疏注意力的思路正是只让每个查询读取缓存中的一小部分条目,从而同时降低显存占用与计算量。RecSA的树结构本质上是对这份缓存做多级摘要索引:底层是精确token,上层是子树均值摘要,查询先在上层粗筛、再逐层下钻,因此读取的条目数与选择成本都只随树高增长,即O(log L)。

均值池化(mean pooling) 指把一组向量逐维取平均得到一个概括向量。在RecSA中,每个树节点的value是其子树内所有token value的均值,key则由pooler生成。均值池化让上层节点以固定维度承载任意长度的子树信息,是树能压缩缓存、且层数随L对数增长的前提。需要注意的是,摘要向量会损失细粒度信息,因此RecSA让未选中候选的摘要与精确token共享同一个softmax,保证全局信息不被完全丢弃。

位置编码的巧思:让外推成为可能

RecSA的一个精妙设计在于RoPE位置的处理。每个条目的RoPE位置是其在所在层级候选中的索引,而非原始token位置。这样一来,无论上下文多长,所有相对距离都保持在训练时见过的范围内。

这解释了为何RecSA能够实现长度外推:模型在训练时只见过32K长度,但推理时可以处理远超此长度的输入,因为位置编码的相对距离始终处于训练分布内。这一思路与当前主流的长上下文外推方法形成对比,后者往往需要复杂的插值或缩放策略。

RoPE(旋转位置编码) 通过对Query和Key向量施加与位置成正比的旋转来注入相对位置信息,是目前主流大模型的位置编码方案。其关键性质是:注意力打分只依赖两个token的位置差(相对距离),而非绝对位置。训练时模型只见过有限范围内的相对距离,一旦推理长度远超训练长度,相对距离落入未见区间,注意力分布就会失真,这正是长度外推困难的根源。常见缓解手段包括位置插值(PI)、NTK-aware缩放、YaRN等,它们通过压缩或重映射位置索引把超长距离拉回训练分布,但往往需要额外调参或微调。RecSA走的是另一条路:不修改RoPE本身,而是改变被编码的位置含义——每个树条目记录的是它在同层候选中的序号,而非原始token位置,因此无论真实上下文多长,参与打分的相对距离天然被限制在每层候选数k的量级内。

实验结果:全面超越基线

RecSA从Qwen3-4B初始化,在原模型的32K训练长度上使用0.62B token进行训练。在RULER基准上,从131K到2M的每一个长度,RecSA都优于所有基线。

在1M长度下,RecSA得分为60.5,而最佳基线Dense + DCA仅为40.0。这一差距表明,递归稀疏注意力在超长上下文场景下具有显著优势。

更重要的是效率表现。借助优化的kernel,RecSA从64K到2M的端到端预填充和解码均快于稠密注意力。在单张H200 GPU上,524K长度时预填充加速3.2倍,解码加速3.8倍。这意味着百万级上下文不再是实验室里的奢侈品,而具备了实际部署的可行性。

RULER 是评估长上下文能力的合成基准,由多个子任务构成,包括单针/多针大海捞针(在长文本中定位特定信息)、变量追踪、聚合计数和长文问答等,并按目标长度(如4K、32K、128K、1M)分别报告得分。它之所以被广泛采用,是因为其长度可控、答案可自动判定,能系统检验模型在超出训练长度后的检索与理解衰减情况。需要留意的是,RULER整体仍以“定位与提取”类能力为主,对需要跨段落全局推理的任务覆盖有限,这也是文中指出RecSA在复杂整合任务上尚缺公开数据的原因。

Dense + DCA 是文中的对照基线之一:Dense指标准稠密注意力,DCA(Dual Chunk Attention)是一类通过分块并调整位置编码来实现训练长度外推的推理期方法,无需重新训练即可扩展上下文。RecSA在1M长度以60.5对40.0大幅领先该基线,说明可训练的稀疏选择在外推场景下比纯推理期的位置重映射更有效。

为什么这对AI应用意义重大

长上下文能力是当前大模型竞争的关键方向之一,但稠密注意力的二次复杂度始终是瓶颈。RecSA的价值在于,它用可训练的分层稀疏选择同时解决了两个问题:一是选择成本随长度对数增长,二是通过位置编码设计实现了训练长度之外的外推。

对于需要处理超长文档、代码仓库或大规模对话历史的AI应用而言,这种对数级成本的长上下文方案提供了新的工程路径。尤其是其从32K训练长度外推到2M的能力,大幅降低了对超长训练数据的依赖——而长文档的稀缺正是当前长上下文训练的痛点之一。

局限与待观察之处

原论文主要报告了RULER基准上的结果,这一基准侧重长上下文检索与理解能力。RecSA在需要全局推理或跨段落复杂整合的任务上表现如何,目前尚缺乏公开数据。此外,pooler的蒸馏质量对最终性能的影响,以及在不同基础模型上的泛化能力,也有待进一步验证。

从工程角度看,优化kernel的可用性是RecSA实用化的前提。论文提到“借助优化的kernel”才能实现加速,这意味着实际部署时对底层实现有较高要求。

总体而言,RecSA为百万级上下文提供了一条对数成本的技术路线,其在长度外推上的设计思路尤其值得关注。随着长上下文需求从 niche 走向主流,这类方法可能会成为下一代模型架构的重要组成部分。

#RecSA#递归稀疏注意力#长上下文#长度外推#稀疏注意力#KV缓存#RULER基准#百万token
分享:

相关推荐