(论文）Local Attention

局部注意力方法 （Local attention，LA）

摘要：

主要研究了一些局部注意力的工作。

对于NLP，LA能够更好的外推；对于CV，能够更加高效。

作用不同，主要也是因为两者的关注目标，实现方法不同。(没有讨论多头，因为那是对最后一个维度，通道维度的操作，现在关注对注意力矩阵的操作)。

---

1. Full Self-Attention (SA)

对于输入 \(X \in \mathbb{R}^{n\times d}\)，首先经过线性变化提取 \(K,Q,V\)：

\[K = X W_k \in \mathbb{R}^{n\times d_k}, \\ Q = X W_q \in \mathbb{R}^{n\times d_k}, \\ V = X W_v \in \mathbb{R}^{n\times d_v}, \]

其中，\(n\) 表示 embedding 个数。全注意力为：

\[score = softmax(\frac{QK^{\top}}{\sqrt{d}}), \\ Attention(Q,K,V) = score * V. \]

此时，每次 query points 都可以访问到全部输入。

2. Mask Self-Attention

2.1 Causal attention

上述的注意力也可以被看作是全的双向注意力，因为每一个输入仅能访问前面的信息，也能访问后面的信息。

然而，在语言任务中（next token predictions），文本生成更需要符合因果关系：即后面生成的token，仅能看到前面的tokens，不应该看到后面未生成的tokens的信息。因此，需要对注意力矩阵添加一个 mask 矩阵 \(\mathcal{M}\)：

\[score = softmax(\frac{QK^{\top} + \mathcal{M}}{\sqrt{d}}), \]

这里 \(\mathcal{M}\) 是一个对角线及以下为0，以上为极大负数的矩阵。经过 \(softmax\) 之后，每个tokens 后面的注意力分数为0。\(\mathcal{M}\) 如下图的例子所示所示：

\[\left[\begin{array}{c} 0 & -inf & -inf \\ 0 & 0 & -inf \\ 0 & 0 & 0 \\\end{array}\right] \]

2.2 AliBi [1]

把 \(\mathcal{M}\) 替换成带状：

\[\left[\begin{array}{c} 0 & 0 & 0 & -inf & -inf & -inf\\ 0 & 0 & 0 & 0 & -inf & -inf \\ 0 & 0 & 0 & 0 & 0 & -inf \\ -inf & 0 & 0 & 0 & 0 & 0 \\ -inf & -inf & 0 & 0 & 0 & 0 \\ -inf & -inf & -inf & 0 & 0 & 0 \\\end{array}\right] \]

用于提升模型的外推能力。

3. Local attention

以上两类工作尽管也是利用了局部的信息（也可以看作一种local的注意力），但这不可以直接迁移到 image 中。在 CV 中，local注意力发展是为了提升效率。

3.1 Sparse global attention

对于所有的 query，随机采样得到 K 和 V的一个子集，在子集上计算注意力。

3.2 Window attention

以 swin transformer 为代表，将自注意力限制到一个局部区域；通过窗口移动来实现窗口外的交互。（实际上是特征偏移实现的。）

3.3 Local attention

将注意力完全限制到一个局部区域。邻居间的交互通过间接方式实现，因为是不重叠的划分。

但是会出现两个问题：

1. 需要使用 Im2Col，计算慢。
2. 构建新的cuda支持。

== 特别的，CV中local attention是在 local 中，每次只关注local中的一个tokens，所以需要将qv切片成 \(k^2\) 个（如下图），才能照顾到local中的所有点。这一步骤需要 Im2Col，会很慢。

3.3.1 Slide-transformer

这篇文章使用深度水平的卷积替代了Im2Col，还加入了一个可学习核增加表示力，无需手写cuda。

3.3.2 Neighborhood attention transformer

手写cuda，实现 query 仅仅与local范围内 K 、V 交互。

---

参考文献

1. Ofir Press, Noah A. Smith, Mike Lewis. Train short, test long: Attention with linear biases enables input length extrapolation. ICLR, 2022.