[论文笔记] PAI-Megatron 源码解读之Mistral的滑动窗口sliding window

Mistral滑动窗口

        滑动窗口注意力（Sliding Window Attention，SWA）利用了变压器（Transformer）堆叠层的特性，以关注超出窗口大小W的信息。位于层k的位置i的隐藏状态hi，会关注到前一层中从i-W到i之间的所有隐藏状态。递归地，能够访问到输入层上距离高达W×k个令牌的位置，如图1所示。在最后一层，使用窗口大小为W=4096时，我们理论上的注意力范围大约为131K个令牌。实际上，对于一个长度为16K的序列和W=4096的情况，对FlashAttention [11] 和 xFormers [18] 所做的修改带来了与普通注意力基线相比2倍的速度提升。

        滚动缓冲缓存。固定的注意力范围意味着我们可以使用滚动来限制我们的缓存大小缓冲区缓存。缓存的大小固定为W，并且存储时间步骤i的键和值在缓存的位置i mod W。因此，当位置i大于W时，过去的值在缓存中被覆盖，并且缓存的大小停止增加。我们提供了一个例子。在图2中，W = 3。在序列长度为32k令牌的情况下，这减少了缓存内存的使用，在不影响模型质量的