滑动窗口注意力优于线性注意力
由于二次注意力的本质,大型语言模型 (LLM) 消耗大量内存和能量。每个新 token 的代价都比前一个更高,且所有键和值都必须无限期存储,这不可持续。 为修复二次缩放问题,已有多种替代方案,其中之一是将 LLM 改造为使用线性注意力。该思路因其以低成本实现先进性能的承诺而备受关注,但这一研究方向尚未与更简单的基线进行适当比较。 本文证明,带 sinks 的滑动窗口注意力(SWA) 在后训练线性注意力模型上表现相当或更优。我们在多个 LLM 和多种下游任务上观察到这一点。在长上下文推理任务 (Needle-in-a-Haystack 和 BABILong) 上,SWA 的性能大幅领先 (比线性注意力高 2 到 10 倍)。SWA 无需后训练,速度极快,内存占用低,因此是极其廉价且可靠的解决方案。 为降低推理内存成本,我们强烈建议改用 SWA,而非后训练线性模型。线性注意力模型虽然显示出一定潜力,但很可能需要从头训练或大量后训练才能匹敌 SWA。
论文精读
TL;DR 带 attention sinks 的滑动窗口注意力(SWA)无需后训练,在长上下文推理任务上性能达线性注意力后训练模型的 2-10 倍,且速度更快、显存更低,是更简单可靠的推理降本方案。
问题
问题背景
当前 LLM 推理的核心瓶颈是 二次复杂度注意力 带来的显存与能耗开销:每新增一个 token,键值缓存(KV cache)无限增长,长序列场景下不可持续。
现有方法局限
业界将 线性注意力(Linear Attention)作为 retrofit 方案改造已有 LLM,期望以低成本将复杂度降至线性。但这类后训练线性模型存在关键局限:
- 在长上下文推理任务(Needle-in-a-Haystack、BABILong)上性能大幅落后于简单的 滑动窗口注意力(Sliding Window Attention, SWA),差距可达 2–10 倍。
- 线性注意力需要大量后训练数据或从头训练才能勉强匹配 SWA,否则 MMLU 平均恢复率(Recovery)显著下降,且未经过与更朴素基线的公平对比。
为什么这个问题难/重要
长上下文推理要求模型在极长序列中精准检索与关联信息,线性注意力通过固定大小的状态压缩全局历史,容易丢失细粒度位置与内容信息;而 SWA 通过局部窗口 + attention sinks 保留关键全局锚点,在没有额外训练的情况下就能保持甚至超越线性注意力性能。业界高度关注推理成本与部署可行性,若盲目采用线性注意力替代方案而忽略简单基线,会浪费大量工程资源。
行业类比
类似实时视频流分析中,用固定滑动窗口缓存最近帧与关键帧,而非在线压缩全部历史特征——窗口策略在低延迟、低内存场景下往往更可靠。
核心洞察
- 滑动窗口注意力加注意力沉没是长期被忽视的强基线。该工作首次系统对比了后训练线性注意力模型与简单的 SWA+sinks 配置,发现后者无需任何后训练即可在通用知识、推理等任务上达到甚至超过线性注意力,在长上下文推理(Needle-in-a-Haystack、BABILong)上性能高出 2 至 10 倍,揭示此前线性注意力研究普遍缺乏与朴素掩码方案的公平对照,SWA 的零训练成本、高速和低内存特性使其成为推理部署的更务实选择。
- 线性注意力的“廉价适配”叙事存在验证缺口。已有后训练方法(SUPRA、Hedgehog、LoLCATs 等)消耗数十亿至数百亿 token 进行阶段式训练,才勉强恢复部分原始性能,而 SWA 仅通过改变注意力掩码和引入全局沉没 token 即可达到相似或更优效果,说明线性注意力若不从头训练或大规模后训练,难以在实际场景中匹敌 SWA,这为推理成本优化提供了更直接、低风险的工程路径。
方法
本研究不提出新架构,而是系统评估一个被长期忽视的简单基线:带 attention sinks 的滑动窗口注意力(Sliding Window Attention with sinks, SWA-sinks)。输入为预训练 LLM 的 Transformer 层,原始全注意力随序列长度呈二次复杂度。方法直接替换注意力机制,无需任何 post-training。
关键模块:
- 滑动窗口注意力:每个 token 仅关注前
w个 token(窗口大小,如 4096),将感受野限制在局部,使计算与缓存复杂度降至线性。 - Attention sinks:保留前
s个 token(通常 4 个)作为全局 sink,所有 token 始终关注这些 sink token,以稳定 softmax 注意力分布,避免信息丢失。
输出为轻量化的注意力层:KV 缓存只需保留窗口内 token 和 sink token,推理时内存占用大幅下降;长上下文任务中无需重新训练即可保持较强推理能力。实验在多个模型和任务(MMLU、Needle-in-a-Haystack、BABILong)上验证,SWA 在长上下文推理上比后训练线性注意力高 2-10 倍。
与同类方法差异:线性注意力通过改变注意力计算方式近似全注意力,需要大量 post-training 恢复性能;SWA 保持 softmax 注意力,仅限制上下文窗口,不引入近似误差,因此无需训练、即插即用,且长上下文表现更优。
实验
实验设计
- 基线对比:将 Sliding Window Attention (SWA) with attention sinks 与近期 post-trained Linear Attention 模型(如 SUPRA、Hedgehog、LoLCATs、Liger-GLA、MOHAWK)对比。
- 任务覆盖:通用知识与推理(MMLU)、长上下文推理(Single Needle-in-a-Haystack、BABILong)。
- 评估维度:下游任务性能、推理速度与内存。
关键发现
- 在长上下文推理上,SWA 性能大幅超过 linear attention,提升 2–10 倍,且无需任何 post-training。
- 在通用知识(MMLU)上,SWA 与大部分线性注意力模型持平或更好;部分线性模型如 LoLCATs 恢复率较高,但需两阶段 post-training,成本更高。
- SWA 推理速度极快、内存占用低,是更廉价可靠的方案。
与基线对比解读
- Linear attention 虽被宣传为低复杂度 SOTA,但多数 post-trained 变体在长上下文任务中明显退化,说明仅靠替换注意力机制并短期 post-train 难以保留原模型能力。
- SWA 利用窗口局部注意力 + sinks 保留显式记忆,避免全局 KV cache 增长,且不需要额外训练;工程上可立即替换,风险低。
- 论文建议若追求推理降本,应优先选择 SWA,而非盲目迁移到线性注意力;后者需要从头训练或更长 post-training 才可能匹敌。
行业影响
落地场景
SWA with sinks 可直接应用于需要处理长上下文的在线服务:客服对话系统(多轮会话)、法律/金融文档审阅、医疗记录摘要、代码库问答、内容平台的长视频/文章摘要 等。例如,电商平台对海量用户评价进行长文本情感分类,或在线教育平台对整章教材内容进行问答,都能在无需重新训练的情况下获得更低内存占用与更高长程推理准确性。
商业价值
采用 SWA 替代 post-trained linear attention 直接降低推理成本:
- 无需额外 post-training 数据与算力,缩短模型上线周期;
- 推理时 KV cache 内存从 O(n^2) 降至固定窗口 O(n),大幅减少长序列服务的显存开销;
- 在 Needle-in-a-Haystack 和 BABILong 等基准上性能是 linear attention 的 2-10 倍,避免信息丢失,提升用户信任与付费转化。 整体对降本(基础设施与训练开销)和体验提升(长上下文准确性)均有贡献。
与现有产品/工作流的接口
集成方式简单:
- 在现有 transformer 推理框架(如 vLLM、Hugging Face Transformers)中把 attention 实现替换为 sliding window attention with sinks(例如 Mistral、Gemma 等模型已内置窗口注意力);
- 无需修改模型权重,只需在推理配置中启用 sliding window 参数(window size、sink tokens),即可用于已部署模型;
- 可与量化、KV cache 压缩等推理优化技术叠加,进一步降低延迟。 对算法团队而言,可作为 A/B 测试的 baseline,快速验证长上下文任务的收益。
局限
- 本工作主要对比了**post-trained Linear Attention**模型(如 SUPRA、Hedgehog、LoLCATs 等),但未与从头训练的线性注意力架构(如 RetNet、Mamba 等)进行系统比较,因此结论是否适用于所有线性注意力变体仍待验证。此外,评估集中在长上下文推理(Needle-in-a-Haystack、BABILong),在一般知识和推理任务上 SWA 仅与线性注意力相当,并未展现压倒性优势,这可能限制了其普适性。
- SWA 依赖固定窗口,虽然使用 attention sinks 保留部分全局信息,但窗口之外的 token 交互被截断。对于需要全局依赖的任务(如长文档摘要、跨段落推理),SWA 可能无法与全局注意力或理论上可捕获全局的线性注意力相比。论文未评估此类任务,因此不能完全替代所有场景,尤其是在需要模型理解全文结构的应用(如代码库级问答)中可能失效。
- 线性注意力的一大优势是理论上的线性时间复杂度和恒定的 KV 缓存,在极长序列(如百万 token)下内存优势显著。而 SWA 虽然内存低,但仍是窗口内二次复杂度,当窗口大小或 batch size 较大时,其计算和内存可能并不优于线性注意力。论文未在超大窗口或超长序列下进行压力测试,因此 SWA 的可扩展性边界尚不清晰,无法判断在极端场景下是否仍具有优势。