完整自注意力会让每个 Token 查看前面的全部 Token,序列增长时,注意力计算量按长度的平方增长;逐词生成还要保留越来越多的 Key 和 Value。滑动窗口注意力给每层划定范围,一个位置只直接查看最近的 W 个位置。
我查了 Longformer 和 Mistral 7B 的原始论文。两者都使用局部窗口降低长序列成本,设计目的却有差别,Longformer 面向长文档编码,把局部窗口与少量任务指定的全局注意力结合;Mistral 7B 面向自回归生成,用窗口限制每层读取的历史,并配套固定大小的滚动 KV 缓存。
一层看近处,多层向前传
局部窗口不等于窗口外的信息立刻消失,前一层已经把邻近 Token 的信息混进当前表示,下一层可以再向前移动一个窗口。Mistral 论文给出的上限是,经过 k 层以后,信息最多传播 k 乘 W 个位置。
Mistral 7B 有 32 层,窗口大小为 4096,按这个传播关系,最后一层的理论范围接近 131K。这个数字描述信息可以沿层传到多远,不代表末尾 Token 能像完整注意力那样直接比较前面 131K 个位置。经过多次中转后,细节会怎样保留,仍由训练和任务决定。
计算和显存收益更容易量。论文在 16K 序列上修改 FlashAttention 与 xFormers 的实现,滑动窗口相对完整注意力快约两倍。滚动缓存只保留 W 个位置,第 i 个位置写入缓存的 i 对 W 取模位置,旧内容随后被覆盖;输入达到 32K 时,论文报告缓存显存减少八倍,模型质量没有因这项缓存实现改变。
很长的提示词还可以分块预填充。Mistral 以窗口长度作为块大小,每个新块关注当前块和缓存中仍在窗口内的内容。这样峰值内存更容易控制,块外信息只能通过已有隐藏表示继续传递。
论文表里的 Mistral 7B 配置还写着 8192 的上下文长度。4096 是每层窗口,8192 是该模型发布时采用的上下文配置,约 131K 则来自层数乘窗口的理论传播范围。三个数字描述不同东西。把理论传播范围直接写成可用上下文,会把模型配置与信息传递上限混在一起。
全局位置要留给真正需要的内容
Longformer 提供了另一种处理。普通 Token 使用局部窗口,问题中的关键位置或分类标记可以获得全局注意力,直接查看整篇文档。其注意力计算随序列长度近似线性增长,但哪些位置应当设为全局,需要按任务决定。全局位置放得太多,计算又会涨回去。
评估滑动窗口模型时,我会把关键信息放在窗口内、刚过窗口和远离窗口的多处位置,分别测试检索与跨段推理。还要记录缓存大小、预填充时间和逐词速度。窗口让长文本跑得动,远处信息能否准确到达答案,则是另一项必须实测的能力。