SWA:Sliding Window Attention 技术介绍

1 minute read

Published:

SWA,全称是 Sliding Window Attention,中文可以叫“滑动窗口注意力”。它的核心思想很直接:不是让每个 token 都关注完整上下文,而是让大部分 token 只关注附近的一段窗口,从而减少长上下文推理中的计算量和显存压力。

在 Mellum 2 Technical Report 中,SWA 是 Mellum 2 控制推理延迟的重要设计。Mellum 2 使用 3:1 的 SWA 模式:每 4 层中有 3 层使用窗口大小为 1,024 tokens 的滑动窗口注意力,剩下 1 层保留完整注意力。这个设计让模型既能降低大部分层的计算成本,又能通过全局注意力层保留跨长距离的信息流动。

为什么需要 SWA

标准 self-attention 的计算复杂度和上下文长度成平方关系。假设上下文长度是 n,完整注意力大致需要处理:

n x n

当上下文从 8K 扩到 32K、64K、128K 时,完整注意力的成本会快速增长。这对代码模型尤其明显,因为 IDE 场景经常需要塞入当前文件、相关文件、错误信息、工具输出和用户对话。

SWA 的做法是把注意力限制在局部窗口内。例如窗口大小是 1,024,那么每个 token 主要关注它附近的 1,024 个 token,而不是完整上下文。

完整注意力:每个 token 看全部历史
SWA:每个 token 只看局部窗口

这样计算规模就从近似 n x n,变成近似:

n x window_size

window_size 固定时,长上下文的成本增长会温和很多。

SWA 如何工作

可以把 SWA 想象成一个沿着文本滑动的窗口。模型生成或处理当前位置时,只允许它在局部范围内做注意力计算。

例如窗口大小为 4,当前位置是第 10 个 token,那么它可能只能关注:

token 6, token 7, token 8, token 9, token 10

真实大模型中的窗口通常更大。Mellum 2 使用的是 1,024 tokens 窗口。

窗口注意力适合捕捉局部依赖,例如:

  • 一段函数内部的变量关系。
  • 当前代码块附近的类型和注释。
  • 当前段落附近的语义衔接。
  • 短到中等距离的上下文线索。

但它不擅长单独处理非常远距离的依赖,因此通常会和全局注意力层结合。

Mellum 2 的 3:1 SWA 模式

Mellum 2 并没有让所有层都使用 SWA,而是使用 3:1 模式:

Layer 1: SWA
Layer 2: SWA
Layer 3: SWA
Layer 4: Full Attention
Layer 5: SWA
Layer 6: SWA
Layer 7: SWA
Layer 8: Full Attention
...

报告中的最终配置是:

配置项数值
SWA 层比例每 4 层中 3 层
全局注意力层比例每 4 层中 1 层
滑动窗口大小1,024 tokens
原生上下文8,192 tokens
长上下文扩展后131,072 tokens

这种设计的好处是:大部分层用 SWA 降低成本,少数全局注意力层负责跨窗口信息传递。模型既不会完全失去远距离感知能力,也不会承担所有层完整注意力的高成本。

为什么窗口是 1,024

窗口大小是 SWA 的关键超参数。

窗口太小,模型可能看不到足够上下文;窗口太大,节省的计算又不明显。Mellum 2 报告中提到,窗口大小 1,024 在质量评测上优于 512,因此最终选择 1,024。

这个选择符合代码场景的直觉。一个局部代码片段常常需要几百到上千个 token 的上下文,512 可能不够覆盖完整函数、类定义或周边注释,而 1,024 能提供更稳的局部信息。

SWA 的优点

1. 降低长上下文成本

SWA 避免每一层都对完整上下文做注意力计算。上下文越长,收益越明显。

2. 降低推理延迟

Mellum 2 报告中把 SWA 明确作为延迟优化。对于 IDE 里的代码补全和对话,用户对延迟很敏感,SWA 能帮助模型在更长输入下维持可用速度。

3. 保留局部建模能力

很多语言和代码依赖都是局部的。SWA 强化了模型在局部窗口内的注意力计算,适合代码编辑、补全和短距离推理。

4. 可以和全局层配合

SWA 不必单独使用。Mellum 2 的做法是混合 SWA 层和 full attention 层,让全局层承担远距离信息聚合。

SWA 的局限

SWA 的局限也来自它的优势:它限制了注意力范围。

如果所有层都只看局部窗口,模型可能难以处理:

  • 跨文件依赖。
  • 文档开头定义、结尾使用的长距离关系。
  • 长对话中早期约束对后续回答的影响。
  • repository 级别的代码理解。

因此,Mellum 2 没有采用“全层 SWA”,而是保留了周期性的全局注意力层。

和 YaRN 的关系

Mellum 2 把上下文从 8K 扩展到 128K 时,使用了 YaRN。一个关键细节是:报告中只对全局注意力层应用 YaRN 频率重映射,而没有改变 SWA 层的 RoPE 参数。

原因很直观:SWA 层一直只处理固定大小的局部窗口,不需要真正外推到 128K 的全局距离;需要适应超长位置的是 full attention 层。

这说明 SWA 不只是一个单独的注意力技巧,它还会影响长上下文扩展策略。

和其他技术的关系

  • GQA:减少 KV heads,降低 KV Cache 成本。
  • YaRN:扩展 RoPE 模型的上下文长度。
  • MTP:用于辅助训练和推测解码 draft。

总结

SWA 是一种面向长上下文效率的注意力机制。它通过限制多数层的注意力窗口,把完整注意力的高成本转化为局部窗口内的可控成本。

在 Mellum 2 中,SWA 的设计不是简单地“少看上下文”,而是“多数层看局部,少数层看全局”。这种 3:1 模式让模型在 IDE 场景中既能处理较长输入,又能保持可部署的推理速度。