模糊窗口注意力

# 模糊窗口注意力

来源：https://arxiv.org/html/2606.09862

SSMState\-Space Model  
LALinear Attention  
BLABlurry Window Attention  
GLAGated Linear Attention  
GSAGated Slot Attention  
ABCAttention with Bounded\-memory Control  
FLAFlash Linear Attention  
GDNGated DeltaNet  
SWASliding Window Attention  
MQARMulti\-Query Associate Recall  
AI人工智能  
LMLanguage Model  
RNNRecurrent Neural Network  

（1华为，瑞士苏黎世 2华为先进计算与存储实验室，中国深圳）

###### 摘要

Transformer语言模型中的Softmax注意力操作具有序列长度的二次复杂度，且状态大小以KV缓存形式不断增长，这在长上下文场景中成为瓶颈。为克服这一局限，人们引入了具有线性复杂度和有限状态大小的替代架构，如状态空间模型（https://arxiv.org/html/2606.09862#id1.1.id1）、线性注意力（https://arxiv.org/html/2606.09862#id2.2.id2）（LA（https://arxiv.org/html/2606.09862#id2.2.id2））和带有限记忆控制的注意力（https://arxiv.org/html/2606.09862#id6.6.id6）（ABC（https://arxiv.org/html/2606.09862#id6.6.id6））。尽管线性模型在语言困惑度上与Transformer相当，但在需要检索或回忆特定信息的任务中仍落后。本文提出模糊窗口注意力（https://arxiv.org/html/2606.09862#id3.3.id3）（BLA（https://arxiv.org/html/2606.09862#id3.3.id3）），一种受SSM（https://arxiv.org/html/2606.09862#id1.1.id1）启发的新型ABC（https://arxiv.org/html/2606.09862#id6.6.id6）方法。BLA（https://arxiv.org/html/2606.09862#id3.3.id3）存储一个频率窗口，通过使用狄利克雷核进行插值来重建模糊的KV历史。根据狄利克雷核的分辨率，BLA（https://arxiv.org/html/2606.09862#id3.3.id3）可理解为滑动窗口注意力（https://arxiv.org/html/2606.09862#id9.9.id9）（SWA（https://arxiv.org/html/2606.09862#id9.9.id9））的推广，或门控槽注意力（https://arxiv.org/html/2606.09862#id5.5.id5）（GSA（https://arxiv.org/html/2606.09862#id5.5.id5））的特例，其中衰减因子通过狄利克雷核实现。我们详细描述了BLA（https://arxiv.org/html/2606.09862#id3.3.id3）的理论和高效实现。在多查询关联回忆（https://arxiv.org/html/2606.09862#id10.10.id10）（MQAR（https://arxiv.org/html/2606.09862#id10.10.id10））合成任务上，我们展示了BLA（https://arxiv.org/html/2606.09862#id3.3.id3）的状态效率比SWA（https://arxiv.org/html/2606.09862#id9.9.id9）高8×，且与流行的线性注意力模型竞争力相当；在RegBench合成任务中，在我们测试的线性模型中，只有BLA（https://arxiv.org/html/2606.09862#id3.3.id3）和SWA（https://arxiv.org/html/2606.09862#id9.9.id9）的性能随状态大小增长而提升。

## 1 引言

Transformer架构[1 (https://arxiv.org/html/2606.09862#bib.bib1)]及其注意力机制是大型语言模型（https://arxiv.org/html/2606.09862#id12.12.id12）的主要支柱之一。注意力的优势在于其能够沿序列长度并行化，并在token之间实现全连接路径，使远距离时间点之间能够直接交互。然而，这种交互的计算成本随序列长度呈二次增长，当上下文长度超过模型维度时成为主要瓶颈，这在智能体AI或长思维链等场景中很常见。此外，Transformer在推理过程中需要不断增长的KV缓存，每个新token需要相应增加计算量。

滑动窗口注意力（https://arxiv.org/html/2606.09862#id9.9.id9）（SWA（https://arxiv.org/html/2606.09862#id9.9.id9））通过将KV历史截断为有限时间窗口来克服二次复杂度。虽然堆叠层数原则上可以将感受野扩展到窗口大小之外，但这种效果并非可加性[2 (https://arxiv.org/html/2606.09862#bib.bib2)]，仍需要全注意力层来维持长距离交互。为缓解注意力的二次复杂度瓶颈同时允许长距离性能，人们设计了具有线性序列复杂度的替代架构。最突出的替代架构包括状态空间模型（https://arxiv.org/html/2606.09862#id1.1.id1）（SSM（https://arxiv.org/html/2606.09862#id1.1.id1））[3 (https://arxiv.org/html/2606.09862#bib.bib3),4 (https://arxiv.org/html/2606.09862#bib.bib4),5 (https://arxiv.org/html/2606.09862#bib.bib5)]、线性注意力（https://arxiv.org/html/2606.09862#id2.2.id2）（LA（https://arxiv.org/html/2606.09862#id2.2.id2））[6 (https://arxiv.org/html/2606.09862#bib.bib6),7 (https://arxiv.org/html/2606.09862#bib.bib7),8 (https://arxiv.org/html/2606.09862#bib.bib8)]和带有限记忆控制的注意力（https://arxiv.org/html/2606.09862#id6.6.id6）（ABC（https://arxiv.org/html/2606.09862#id6.6.id6））[9 (https://arxiv.org/html/2606.09862#bib.bib9),10 (https://arxiv.org/html/2606.09862#bib.bib10)]。与Transformer一样，这些神经网络可沿序列并行化，但与Transformer不同，它们的线性序列混合操作使用有限状态而非不断增长的KV缓存。这使得线性LM（https://arxiv.org/html/2606.09862#id12.12.id12）在长上下文场景中相比Transformer具有计算优势。然而，近期研究指出，线性LM（https://arxiv.org/html/2606.09862#id12.12.id12）在需要长距离信息回忆的特定任务中不如注意力变体[11 (https://arxiv.org/html/2606.09862#bib.bib11),12 (https://arxiv.org/html/2606.09862#bib.bib12)]，这使纯线性LM（https://arxiv.org/html/2606.09862#id12.12.id12）在文本处理中的长期可行性受到质疑。

本文提出模糊窗口注意力（https://arxiv.org/html/2606.09862#id3.3.id3）（BLA（https://arxiv.org/html/2606.09862#id3.3.id3）），一种新型线性注意力架构，旨在结合SWA（https://arxiv.org/html/2606.09862#id9.9.id9）的精确检索与SSM（https://arxiv.org/html/2606.09862#id1.1.id1）和LA（https://arxiv.org/html/2606.09862#id2.2.id2）模型的长距离依赖能力。线性注意力的状态以外积格式存储键值关联，而BLA（https://arxiv.org/html/2606.09862#id3.3.id3）维护分离的键状态和值状态，这使得BLA（https://arxiv.org/html/2606.09862#id3.3.id3）更类似于ABC（https://arxiv.org/html/2606.09862#id6.6.id6）和SWA（https://arxiv.org/html/2606.09862#id9.9.id9）。但与ABC（https://arxiv.org/html/2606.09862#id6.6.id6）方法不同，BLA（https://arxiv.org/html/2606.09862#id3.3.id3）的写入机制可视为SWA（https://arxiv.org/html/2606.09862#id9.9.id9）的推广。这是通过在一组有限的傅里叶模式上独立乘积累加输入键和值来实现的，类似于S4D[13 (https://arxiv.org/html/2606.09862#bib.bib13)]等SSM（https://arxiv.org/html/2606.09862#id1.1.id1）。这种状态空间表示允许在时域上使用狄利克雷核进行有损插值，直至一个周期。然后使用当前查询对插值后的键和值计算softmax注意力。接下来，我们首先介绍BLA（https://arxiv.org/html/2606.09862#id3.3.id3）的理论，然后评估其在回忆密集型合成任务上的性能。我们展示了BLA（https://arxiv.org/html/2606.09862#id3.3.id3）在MQAR（https://arxiv.org/html/2606.09862#id10.10.id10）任务上的状态效率比SWA（https://arxiv.org/html/2606.09862#id9.9.id9）高8×，且接近流行的线性模型。此外，与门控线性注意力（https://arxiv.org/html/2606.09862#id4.4.id4）（GLA（https://arxiv.org/html/2606.09862#id4.4.id4））和门控DeltaNet（https://arxiv.org/html/2606.09862#id8.8.id8）（GDN（https://arxiv.org/html/2606.09862#id8.8.id8））不同，BLA（https://arxiv.org/html/2606.09862#id3.3.id3）在RegBench任务上达到了与全注意力相当的性能，并且在状态大小较小时优于SWA（https://arxiv.org/html/2606.09862#id9.9.id9）。

## 2 背景

我们首先简要回顾原始因果Softmax注意力[1 (https://arxiv.org/html/2606.09862#bib.bib1)]及其线性变体的操作，为简单起见考虑单头单样本的情况。

### 2.1 Softmax注意力

给定一个长度为 \(L\) 的 \(d\) 维向量序列 \(\mathbf{X} \in \mathbb{R}^{L \times D}\)，Softmax注意力使用投影矩阵 \(\mathbf{W}_q, \mathbf{W}_k, \mathbf{W}_v \in \mathbb{R}^{D \times D}\) 将输入投影为查询、键和值序列：\(\mathbf{Q} = \bm{W}_q \mathbf{X}\), \(\mathbf{K} = \bm{W}_k \mathbf{X}\), \(\mathbf{V} = \bm{W}_v \mathbf{X} \in \mathbb{R}^{L \times D}\)。输出由以下公式给出：

\[
\mathbf{O} = \mathrm{Softmax}\left( \frac{\mathbf{Q} \mathbf{K}^\top}{\sqrt{D}} + \mathbf{M} \right) \mathbf{V} \quad \in \mathbb{R}^{L \times D},
\tag{1}
\]

其中softmax按行应用。\(\mathbf{M} \in \{-\infty, 0\}^{L \times L}\) 是因果掩码，防止查询 \(\bm{q}_t\) 查询未来的键向量 \(\bm{k}_{t' > t}\)。softmax项是一个 \(L \times L\) 矩阵，称为注意力掩码，它导致了标准注意力在序列长度上的 \(O(L^2 D)\) 二次复杂度。在滑动窗口注意力（https://arxiv.org/html/2606.09862#id9.9.id9）（SWA（https://arxiv.org/html/2606.09862#id9.9.id9））中，窗口大小为 \(w\)，查询 \(\bm{q}_t\) 仅关注滑动窗口内的键 \(\bm{k}_{t'}\)，其中 \(t' \in [t - w, t]\)，这将复杂度降至 \(O(L w D)\)，但代价是丢失了向量间的长距离交互。

### 2.2 带有限记忆控制的注意力

带有限记忆控制的注意力（https://arxiv.org/html/2606.09862#id6.6.id6）（ABC（https://arxiv.org/html/2606.09862#id6.6.id6））[9 (https://arxiv.org/html/2606.09862#bib.bib9)]引入了累积的softmax写入门控 \(\bm{\phi}_t\)，允许将多个token存储到固定大小的记忆槽中：

\[
\widetilde{\mathbf{K}}_t = \widetilde{\mathbf{K}}_{t-1} + \bm{\phi}_t \otimes \mathbf{k}_t, \quad \widetilde{\mathbf{V}}_t = \widetilde{\mathbf{V}}_{t-1} + \bm{\phi}_t \otimes \mathbf{v}_t.
\tag{2}
\]

\(\bm{\phi}_t\) 通过对token特征进行归一化指数得到，提供了一种数据依赖的、类似FIFO的记忆更新方式，同时保留了对槽的softmax注意力。该公式可以表示为两遍线性注意力，使得能够对较小的循环状态进行硬件高效的分块训练。

门控槽注意力（https://arxiv.org/html/2606.09862#id5.5.id5）（GSA（https://arxiv.org/html/2606.09862#id5.5.id5））[10 (https://arxiv.org/html/2606.09862#bib.bib10)]在ABC机制基础上，为每个记忆槽添加了一个数据依赖的门控标量 \(\alpha_i \in [0,1]\)。每一步，键和值槽通过门控循环进行更新：

\[
\widetilde{\mathbf{K}}_t = \operatorname{Diag}(\bm{\alpha}_t) \widetilde{\mathbf{K}}_{t-1} + (1 - \bm{\alpha}_t) \otimes \mathbf{k}_t
\tag{3}
\]

（\(\widetilde{\mathbf{V}}_t\) 类似），这使得模型能够遗忘过时信息并引入近因偏差，解决了ABC无法丢弃旧token以及对早期token有偏差的问题。该更新可以写成两遍门控线性注意力，从而能够使用与线性注意力相同的硬件高效分块训练，同时提供紧凑的循环状态和改进的推理效率。

### 2.3 状态空间模型

状态空间模型（https://arxiv.org/html/2606.09862#id1.1.id1）（SSM（https://arxiv.org/html/2606.09862#id1.1.id1））文献可追溯到勒让德记忆单元[14 (https://arxiv.org/html/2606.09862#bib.bib14)]和Hippo理论[3 (https://arxiv.org/html/2606.09862#bib.bib3)]。SSM（https://arxiv.org/html/2606.09862#id1.1.id1）最初解决的问题可概括为：给定一个输入的1维连续信号 \(x(t)\) 和一个有限的 \(N\) 维存储空间，如何保留关于信号的最多信息？SSM（https://arxiv.org/html/2606.09862#id1.1.id1）理论表明，给定关于信号的某种期望度量，我们可以将其投影到一个基上，维护一组坐标，从而可以近似重构信号。忽略离散化步骤，离散SSM（https://arxiv.org/html/2606.09862#id1.1.id1）的方程如下：

\[
\begin{split}
\bm{h}(t+1) &= A \bm{h}(t) + B x(t), \\
y(t) &= C \bm{h}(t) + D x(t).
\end{split}
\tag{4}
\]

这里 \(\bm{h}(t)\) 是信号 \(x(t)\) 的 \(N\) 维状态空间表示。矩阵 \(A, B, C, D\) 是SSM（https://arxiv.org/html/2606.09862#id1.1.id1）的参数。从方程中可以看出，状态更新是线性循环，只要 \(A\) 矩阵可对角化，就可以沿序列长度高效并行化。早期SSM（https://arxiv.org/html/2606.09862#id1.1.id1）将输入信号投影到勒让德多项式或截断傅里叶模式上[14 (https://arxiv.org/html/2606.09862#bib.bib14),3 (https://arxiv.org/html/2606.09862#bib.bib3)]，这对应使用特定的参数矩阵。早期SSM（https://arxiv.org/html/2606.09862#id1.1.id1）和线性RNN（如LRU或S4-FouT）曾使用具有复特征值的 \(A\) 矩阵[15 (https://arxiv.org/html/2606.09862#bib.bib15),3 (https://arxiv.org/html/2606.09862#bib.bib3)]，因为任何实矩阵 \(A\) 在 \(\mathbb{C}\) 中几乎必然可对角化，这有助于更好的表达性。这一趋势后来转变为使用从数据中学习的对角实值 \(A\) 矩阵，并在之前添加短卷积以改善回忆[16 (https://arxiv.org/html/2606.09862#bib.bib16),13 (https://arxiv.org/html/2606.09862#bib.bib13),5 (https://arxiv.org/html/2606.09862#bib.bib5),17 (https://arxiv.org/html/2606.09862#bib.bib17)]。有趣的是，Mamba 3又回到了使用复特征值[18 (https://arxiv.org/html/2606.09862#bib.bib18)]。

## 3 理论

现在描述我们提出的模糊窗口注意力（https://arxiv.org/html/2606.09862#id3.3.id3）（BLA（https://arxiv.org/html/2606.09862#id3.3.id3））框架的理论。首先，以类似于传统SSM（https://arxiv.org/html/2606.09862#id1.1.id1）的方式介绍，突出BLA（https://arxiv.org/html/2606.09862#id3.3.id3）与SSM（https://arxiv.org/html/2606.09862#id1.1.id1）文献的相似性。然后，展示一种更高效的实现方式，该方式利用softmax注意力的置换不变性，不需要卷积，更类似于ABC（https://arxiv.org/html/2606.09862#id6.6.id6）。最后，展示如何在BLA（https://arxiv.org/html/2606.09862#id3.3.id3）中实现类似于GSA（https://arxiv.org/html/2606.09862#id5.5.id5）的状态衰减，使其看起来像SWA（https://arxiv.org/html/2606.09862#id9.9.id9）的更一般版本。

请参考图注
图1：模糊窗口注意力机制概览。
左图：BLA（https://arxiv.org/html/2606.09862#id3.3.id3）的状态是键和值对一组 \(M\) 个傅里叶模式的余弦和正弦分量进行卷积的结果，参数