微软新研究：无需后训练，只加4个Token，滑动窗口注意力反超线性注意力

随着大语言模型（LLM）上下文长度不断增加，推理阶段的内存占用和能源消耗正逐渐成为实际部署中的关键瓶颈。

标准全注意力机制会为每个新 token 保留此前所有 token 的 Key 和 Value，KV Cache 随上下文持续增长，解码成本也随之攀升。

尽管线性注意力方法试图将时间和内存复杂度从二次降为线性，用固定大小的递归状态替代不断扩张的 KV Cache，但这类方法通常需要额外的后训练，且必须学会哪些信息该保留、哪些可以遗忘，因而仍存在表达能力下降、训练成本高以及软硬件生态不成熟等问题。

在这项工作中，微软团队提出，相比于后训练的线性注意力模型，“滑动窗口注意力（Sliding Window Attention，SWA）长期被忽视，其性能实际上可以达到甚至超过前者，这一结论在多个 LLM、多种下游任务上均得到了验证。

论文链接：https://arxiv.org/pdf/2608.28444

值得一提的是，与以往的滑动窗口注意力方法相比，该方法只需在窗口之外额外保留 4 个 token。

研究方法

滑动窗口注意力的基本思路是限制每个 token 的注意力范围：模型不再访问整个历史序列，而只关注前面固定窗口内的 token。随着上下文不断增长，模型只需维护窗口内的 Key 和 Value，KV Cache 在达到窗口大小后便保持稳定，推理内存占用不再随上下文长度增加而上升。

然而，单纯的滑动窗口存在一个问题：当最早的 token 被窗口移出后，模型可能出现灾难性失效。研究团队表示，许多此前工作与线性注意力对比时，使用的是不带 sinks 的普通滑动窗口注意力，这导致对滑动窗口注意力性能的估计并不准确。

他们的做法是，在窗口之外额外保留序列最开始的 4 个 token。这些固定保留的 token 被称为“注意力汇聚”（attention sinks）。由此，滑动窗口注意力的注意力范围由两部分组成：当前 token 前方最近的若干 token，以及序列开头固定保留的 4 个 token。

Attention sinks 的作用并非简单地保存最有用的语义内容，而是为注意力分布提供稳定的“汇聚位置”。在标准 Transformer 中，模型往往会将部分注意力集中到序列开头的 token 上。即使这些 token 本身并不携带重要信息，它们也可能承担着稳定注意力分布的功能。保留最前面的几个 token，可以避免这些位置被完全移除，使模型在有限窗口下仍能维持接近原始模型的计算行为。

他们重点比较了预训练模型直接使用滑动窗口注意力，与纯线性注意力、线性注意力加滑动窗口注意力等后训练线性化方法间的差异。为保持公平，比较对象涵盖多种线性化模型、原始预训练模型，以及原始模型直接接入滑动窗口注意力后的结果，部分层保留全注意力的混合模型则不在比较范围内。

图｜不同类型的注意力掩码：全注意力（FA）、LoLCATs/Liger-GLA（线性注意力与滑动窗口注意力的混合）、以及带有 sinks 的滑动窗口注意力。

实验结果

研究团队将现有线性方法（纯线性方法或带有滑动窗口注意力的线性方法）与基础预训练模型及其对应的滑动窗口注意力变体进行了对比。

出于简单与公平性考虑，他们没有与包含部分全注意力层的混合模型进行对比。对比测试涉及通用知识和推理、长上下文推理，以及速度和内存三个方面。

在通用知识和推理实验中，评测指标包括 MMLU、ARC-C、ARC-E、HellaSwag、PIQA 和 Winogrande。

结果显示，SWA(64, 4) 在 11 组主要比较中有 9 组取得了最佳平均下游成绩。其中，SWA(64, 4) 将 MMLU 性能恢复到原始模型的 93.2%，QRWKV6 为 92.4%；如果看 6 个知识与推理基准的平均恢复率，SWA(64, 4) 为 99.0%，QRWKV6 为 99.1%。此外，在 MMLU 上，SWA(64, 4) 几乎始终是表现最好的非教师方案。

两种方法的关键区别在于训练成本。滑动窗口注意力无需任何后训练 token，而 LoLCATs 需要使用约 4000 万个 token 进行微调，才能达到 83.2% 的 MMLU 得分和 97.5% 的平均恢复率。就“以较低训练成本换取较高性能”而言，滑动窗口注意力在这组实验中更具优势。

图｜蒸馏的线性化模型、教师模型以及带有滑动窗口注意力（包含 4 个 Sink、滑动窗口大小为 64 或 128）的教师模型的基准测试得分（%）。

长上下文实验以 Llama 3.1 8B 为基线模型，比较对象包括滑动窗口注意力、LoLCATs 和 Liger-GLA。研究团队首先在单大海捞针（S-NIAH）任务中，于 0.5K、1K、2K 和 4K 上下文长度下测试模型从长文本中找回关键信息的能力，窗口大小覆盖 128、256 和 512。

结果显示，在所有窗口大小和上下文长度组合中，滑动窗口注意力的得分均不低于另外两种方法。当上下文长度达到 4K 时，滑动窗口注意力在三个 S-NIAH 子任务上的准确率处于12.6%-23.0%区间，而 LoLCATs 的最高准确率只有 5.8%，Liger-GLA 最高仅为 0.8%。

图｜不同上下文长度与窗口大小下的 S-NIAH 测试准确率。

他们还在 BABILong 基准上比较了滑动窗口注意力和 LoLCATs。该实验选取 QA1 至 QA5 五个任务的平均结果，同样以 Llama 3.1 8B 为基线模型，窗口大小为 256。

结果显示，在较短上下文下，LoLCATs 略占优势；然而，随着上下文长度增大，滑动窗口注意力的得分更高，上下文长度为 2K 时达到 19%，4K 时达到 15%，而 LoLCATs 分别对应为 10% 和 3%。

以全注意力结果为参照，在 4K 上下文下，滑动窗口注意力恢复了约 25% 的基线性能，LoLCATs 仅恢复约 5%。可见，在需要从不断增加的历史上下文中检索信息时，带 sinks 的滑动窗口注意力表现得更加稳定。

图｜在不同上下文长度下的 BABILong 基准测试分数。

研究团队在一个 4 层 Transformer 上开展了速度和内存实验。其中，全注意力和滑动窗口注意力使用 FlashAttention 作为后端，线性注意力使用 ThunderKittens，LoLCATs 使用融合后的滑动窗口注意力与线性注意力 Kernel。

结果显示，当上下文长度超过 1K 后，全注意力的解码速度开始下降，而其他方法的速度相对平稳。滑动窗口注意力是速度最快的方法，且窗口大小为 64 的配置快于窗口大小为 512 的配置。

在内存方面，全注意力的开销会随上下文长度线性增长，滑动窗口注意力的内存开销先增加，在达到窗口大小后便保持稳定，窗口大小为 64 的滑动窗口注意力内存成本最低，然后是线性注意力、LoLCATs，最后是窗口大小为 512 的滑动窗口注意力。

因此，在窗口大小小于 512 时，滑动窗口注意力能同时具备更高的解码速度，以及与线性注意力相近或更低的内存成本。

图｜不同注意力类型在不同上下文长度（128 至 256K）下的速度与内存开销。

下一步？

当然，这项研究仍存在一些局限。

例如，他们的工作聚焦于无需训练的滑动窗口注意力，但后训练仍有可能进一步提升滑动窗口注意力的性能。未来的工作可以比较滑动窗口注意力和线性方法在不同后训练 token 数量下的表现差异。

此外，当前实验尚未覆盖仅让部分 token 或部分层使用全注意力的混合模型，全注意力与滑动窗口注意力、线性注意力之间应采用何种比例，仍有待研究。他们表示，后续工作可以扩展到更大规模的模型、更复杂的 Agentic 任务，以及多模态大语言模型（MLLM）和视频生成模型。值得一提的是，针对视频扩散模型，已有强有力的证据表明免训练滑动窗口注意力是一个强大的基线方法。

上述结果表明，降低 LLM 推理内存未必需要先进行大规模模型改造或后训练。对于已完成预训练的 Transformer，带 sinks 的滑动窗口注意力提供了一条实现简单、速度较快、内存稳定的路径。这一方法无需任何后训练即可同时实现更快的解码速度和更低的内存开销。

在线性注意力仍需解决长期信息保留和模型表达能力问题的当下，滑动窗口注意力或许是许多实际部署场景中值得优先验证的方案。

更多技术细节，详见原论文。

作者：学术君

如需转载或投稿，请直接在本文章评论区内留言