LLM 优化技术(3)——sliding_window_attention

哦豁灬

于 2025-03-31 09:01:29 发布

阅读量2.1k

点赞数 39

分类专栏： LLM 大模型深度学习文章标签： LLM 大模型 transformers self-attention

本文链接：https://blog.csdn.net/qq_38342510/article/details/146714992

版权

深度学习同时被 3 个专栏收录

59 篇文章

订阅专栏

LLM

36 篇文章

订阅专栏

大模型

19 篇文章

订阅专栏

一种（理论有损）提升注意力计算效率的方法：SWA（sliding window attention）。Qwen系列和Mistral就使用了SWA。

对于原始的 causal attention，其注意力矩阵是一个下三角矩阵，这样每个 token 都能看到自己和在自己前面的所有的 token。这样随着输入长度 $s$ 增大，这个下三角矩阵中 1 的元素数量以 $s_2$ 的速度增长，计算量以序列长度平方的速度增长，缓存和长度成线性关系增长。

1 SWA

CNN 中的感受野：

在这里插入图片描述

如上图，假设模型有 3 层，每层卷积核大小为 $[3, 3]$ （实际上CNN里卷积操作就是一个 sliding window）。

对于 layer3，每一个像素能看到 layer2 中的一个 $[3, 3]$ 的区域，layer2 中其他较远的像素就看不到了。再往前推，layer2 里的每个像素也可以看到 layer1 中的一个 $[3, 3]$ 区域，那么 layer2 中的 $[3, 3]$ 区域就可以看到 layer1 中一个 $[5, 5]$ 的区域，相当于 layer3 中一个像素可以间接看到 layer1 中一个 $[5, 5]$ 的输入。以此类推，如果再增加一层 layer4，那么 layer4 中一个像素就能获取 layer1 一个 $[7, 7]$ 区域的信息。

每层只能多看周围一格的信息，但是只要层数够多，理论上靠近输出端的层想看多远就能看多远。

一般认为模型低层部分提取比较基础的特征，高层会提取高级的语义特征。在 CNN 里，前几层提取的可能更多是关于简单的边界、颜色、形状等基础特征，后面的层则提取较复杂的语义特征，比如在分类任务中会是和分类类别相关的花纹、物体大小、风格等特征。

如果我们把模型设计成最后一层的一个像素刚好要到第一层才能接收到全局信息（在其它层都只能看到局部），那对于图像边缘的语义特征识别能力可能会受到一些限制。假设我们做猫和狗的图像分类任务，如果这个时候决定性的特征出现在图像最边缘几个像素里，那这种情况下的错误率会比特征出现在图像中间时要高。

对于语言模型，一般情况下，越远距离的信息，对当前位置的重要性越低，因此只要我们的窗口大小不要太过极限小，问题应该都还不大。

1.1 SWA 的具体做法

在这里插入图片描述

左边是正常的 causal attention，每个位置能看到自己和自己前面的所有位置，attention mask 是个下三角矩阵；
中间则是 SWA 的 attention mask，这里的窗口大小为 3。包括自己在内，每个位置只能往前看 3 个输入。

同 CNN 的感受野一样，随着层数的堆叠，模型理论上能处理的最远距离也逐层线性递增。LLM 里递增的方向是单向的，只能往前。

Mistral 7B 使用了 4096 的窗口大小，模型层数为 32，则最终输出的感受野为 $4096 \times 32 = 128K$ 。

Attention 的计算量与 sequence length 的二次方成正比，而使用 SWA 的理论计算量与 sequence length 的一次方成正比。而缓存和上下文长度 sequence length 成线性关系。

SWA 在上下文长度在 4k 以下时，和普通 causal attention 一样；当上下文长度超过4k时，则相对节省资源，长度越大，节省的比例越高。实际使用中，Mistral 通过把 SWA 实现在 FlashAttention 和 xFormers 中，对于 16k 的上下文长度，获得了 2 倍的速度提升。