AB
AiBoss
Wiki

什么是滑动窗口注意力(Sliding Window Attention)?

滑动窗口注意力(Sliding Window Attention,SWA)是一种稀疏注意力模式,让每个词元只关注固定窗口内的邻近词元,从而把自注意力的计算复杂度从二次降为线性。两篇 arXiv 论文分别讨论了把全注意力模型适配到 SWA 的方法,以及 SWA 与线性注意力在长上下文任务上的对比。

滑动窗口注意力(Sliding Window Attention,SWA)是 Transformer 类模型中一种稀疏注意力(sparse attention)模式:每个词元(token)在计算注意力时,只关注它前面固定长度窗口内的邻近词元,而不是像全注意力(Full Attention,FA)那样关注此前所有词元。它要解决的问题是自注意力的二次复杂度——随着上下文变长,全注意力的计算量与键值缓存(KV cache)占用会迅速膨胀,而 SWA 把每个词元的注意力范围限制在常数大小的窗口内,使长上下文推理的代价随序列长度近似线性增长。

需要说明的是,SWA 本身是一种结构上的取舍,而不是无代价的等价替换。arXiv 论文 SWAA 把 SWA 称为「最简单的稀疏注意力模式」,同时指出它存在「灾难性的长上下文性能崩塌」。另一篇 arXiv 论文则把 SWA 与注意力汇(attention sinks)的组合,与改造为线性注意力(Linear Attention)的路线做了直接对比。以下内容分别标注了各自主张的来源。

为什么重要

在 SWA 这类稀疏模式被系统讨论之前,长上下文推理主要依赖全注意力。全注意力的特点是:每生成一个新词元,都要与之前所有词元计算注意力,因此第 n 个词元的代价高于第 n−1 个;同时,所有历史词元的键(key)和值(value)必须一直保存在显存中。arXiv 论文《Sliding-window beats linear attention》把这种性质描述为「不可持续」——随着上下文增长,内存与能耗持续上升。

针对这一问题,该论文归纳出两条主要技术路线。第一条是压缩键值缓存,例如淘汰(evicting)或量化(quantizing)键和值。第二条是把已有模型改造(retrofit)为使用线性注意力,用固定大小的状态替代键值缓存。该论文指出,改造路线因为「以低成本达到先进性能」的承诺而受到很多关注,但它并没有被与「最简单的键值缓存淘汰形式」——即带注意力汇的滑动窗口注意力——做过恰当的比较。

另一篇论文 SWAA 则从另一个角度说明 SWA 的重要性:它把 SWA 视为降低长上下文推理成本的最简方案,但强调直接把用全注意力预训练好的模型套用 SWA 会失败,原因有两个。其一是训练与推理不匹配(training inference mismatch):模型在预训练阶段见到的注意力范围与推理时不同。其二是结构上的固有限制:如果所有模块、所有时刻都只用 SWA,模型就无法访问远距离信息。这两个因素共同导致了长上下文性能崩塌。

工作机制

SWA 的核心做法可以拆成几个层面来理解。

  • 窗口化注意力:每个词元只与窗口内的邻近词元计算注意力,窗口大小固定,因此单个词元的注意力计算量不随序列长度增长。这是 SWA 获得线性复杂度的直接来源。
  • 注意力汇:SWA 与注意力汇配合使用时,会保留一部分「汇」词元(sink tokens),使模型始终能访问到这些位置的信息。arXiv 论文《Sliding-window beats linear attention》讨论的正是「带汇的 SWA」这一形式。
  • 全注意力解码:SWAA 论文提出的策略之一,是在解码阶段保留全注意力,以缓解纯 SWA 无法访问远距离信息的结构缺陷。
  • FA 与 SWA 层交错:SWAA 的另一项策略是把全注意力层与滑动窗口注意力层交错排布,选择性地允许模型访问远距离信息,而不是在所有层、所有时刻都使用 SWA。
  • 轻量微调:SWAA 还使用轻量级微调来缓解训练与推理不匹配的问题,其目标是在不做昂贵预训练的前提下完成适配。

SWAA 论文强调,上述策略单独使用时都不足以解决问题,只有特定的协同组合才能有效恢复长上下文性能。该论文把 SWAA 描述为一个即插即用(plug and play)的方法集合,用于把全注意力模型适配到 SWA,而无需重新进行昂贵的预训练。论文还报告了性能与效率的权衡分析,称其识别出适用于不同场景的最优 SWAA 配置,在可接受的质量保持下为长上下文推理带来 30% 到 100% 的加速。

另一篇论文的机制主张则更直接:带汇的 SWA 不需要额外训练,速度极快,内存需求很小。该论文称,在训练预算有限的情况下,切换到 SWA 比改造为线性注意力更能有效降低推理内存成本。

典型例子

两篇论文给出了若干具体的评测与对比场景。

  • 长上下文推理任务:arXiv 论文《Sliding-window beats linear attention》在「大海捞针」(Needle-in-a-Haystack)与 BABILong 这两项长上下文推理任务上做了比较,称 SWA 的表现大幅高于线性注意力,达到后者的 2 到 10 倍。该论文还表示,在多个模型与多种下游任务上,带汇的 SWA 表现与大多数改造后的线性注意力模型相当或更好,在长上下文与生成类任务上优势最大。
  • 与线性注意力的路线对比:同一篇论文指出,线性注意力模型虽然展现出潜力,但要获得其架构优势、接近 SWA 的水平,需要从头训练或进行大规模改造。
  • SWAA 的适配实验:SWAA 论文的实验显示,孤立的适配策略效果不足,而特定的协同组合能够恢复长上下文性能;论文据此给出了不同场景下的配置建议,并报告了 30% 至 100% 的长上下文推理加速。论文称其代码、数据与模型权重均已公开。

需要注意的是,上述数字与结论均来自对应论文自身的实验设置,属于论文作者的主张,并非独立复现的结果,也不代表所有模型与任务上的普遍表现。

边界与常见误解

SWA 容易被误解为「全注意力的等价替代」,但两篇论文都指出了它的代价与适用边界。

  • 不是免费的降本手段:SWAA 论文明确指出,把用全注意力预训练的模型直接套用 SWA,会导致灾难性的长上下文性能崩塌,原因包括训练与推理不匹配,以及所有模块、所有时刻都使用 SWA 时无法访问远距离信息。因此直接替换并不可行,需要配合全注意力解码、层交错、保留汇词元或轻量微调等手段。
  • 单独策略往往不够:SWAA 论文称,孤立的策略不足以解决问题,只有特定的协同组合才有效。这意味着 SWA 的收益依赖于具体的配置选择,而不同配置的计算开销也不同。
  • 与线性注意力并非同一路线:线性注意力用固定大小的状态替代键值缓存,而 SWA 属于键值缓存淘汰的一种形式。arXiv 论文《Sliding-window beats linear attention》主张在训练预算有限时 SWA 更划算,但这一比较是在该论文选定的模型与任务上做出的,不能推广为线性注意力在所有场景下都不如 SWA。
  • 窗口大小与汇的取舍:窗口越小,单步计算与内存开销越低,但可访问的历史信息也越少;注意力汇保留了多少信息、保留哪些位置,同样会影响效果。这些都属于需要按场景调优的超参数,论文并未给出普适取值。
  • 结论的来源限定:关于「SWA 表现优于线性注意力」「30% 到 100% 加速」等说法,均应理解为对应论文在其实验条件下的报告结果,而非行业共识。

参考资料