SWA:Sliding Window Attention 技术介绍
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 不只是一个单独的注意力技巧,它还会影响长上下文扩展策略。
和其他技术的关系
总结
SWA 是一种面向长上下文效率的注意力机制。它通过限制多数层的注意力窗口,把完整注意力的高成本转化为局部窗口内的可控成本。
在 Mellum 2 中,SWA 的设计不是简单地“少看上下文”,而是“多数层看局部,少数层看全局”。这种 3:1 模式让模型在 IDE 场景中既能处理较长输入,又能保持可部署的推理速度。
