详细介绍Stand-Alone Self-Attention in CV中的公式

论文：Stand-Alone Self-Attention in Vision Models
因为我看这个文章的时候跳过了很多之前self-attention的知识，对公式的理解花了一段时间，所以这里主要写公式的理解，权当笔记。

前言：

卷积神经网络(CNN)通常学习小范围(kernel sizes)的局部特征，对于输入$x\in {\mathbb{R}}^{h\times w\times d_{in}}$，定义局部像素集 ${\mathcal{N}}_k$ ,为像素$x_{i,j}$​周围$k$区域的像素，大小为 $k\times k\times d_{in}$ ，如图1：

对于学习到的权重 $W\in {\mathbb{R}}^{k\times k\times d_{ou}\times k_{in}}$ ， 位置 $ij$ 的输出 $y_{ij}\in {\mathbb{R}}^{d_{out}}$ ,通过公式(1)计算所得, 其中 ${\mathcal{N}}_k(i,j)=\{a,b||a-i|\le k/2,|b-j\mid \le k/2\}$, CNN 使用权重共享， $W$ 用于所有像素位置 $ij$ 的输出，权重共享使得特征具有平移不变性以及降低卷积的参数量。目前有一些卷积的变种用 以提高预测的表现, 比如深度分离卷积

讲这些的原因主要是为了记住公式中各个符号的意义，从而更好地理解后面的自注意力机制公式计算方法。

Self-Attention

与传统的attention不同，self-attention应用于单个context而不是多个context间，能够直接建模context内长距离的交互信息，论文提出stand-alone self-attention layer用来替代卷积操作，并且构建full attention模型，这个attention layer主要是对之前的工作的一个简化

与卷积类似，对于像素 $x_{ij}\in {\mathbb{R}}^{d_{in}}$，首先会取$x_{ij}$的$k$范围内的局部区域像素 $ab\in {\mathcal{N}}_k(i,j)$ ($ab$是像素的位置，和$ij$一样) ，称为memory block。与之前的all-to-all attention不同，这个attention只在局部区域${\mathcal{N}}_k(i,j)$进行attention操作，全局attention只有在特征大小大幅减少后才能使用，不然会带来很大的计算开销

single-headed attention计算如公式2，输出像素$y_{ij}\in {\mathbb{R}}^{d_{out}}$，首先对输入向量进行三种变化得到3个值,：查询矩阵：$Q=W_{Q}x$，关键词矩阵：$K=W_{K}x$，值矩阵：$V=W_{V}x$
那么就有：查询像素queries $q_{ij}=W_Q{x}_{ij}$, 关键词像素keys $k_{ab}=W_K{x}_{ab}$ 以及值像素values $v_{ab}=W_V{x}_{ab}$ 都是像素 $ij$ 和其附近像素的线性变化， $softma{x}_{ab}$ 应用于所有 $q_{ij}^{\top }{k}_{ab}$ ，$W_Q,W_K,W_V\in {\mathbb{R}}^{d_{out}\times d_{in}}$ 为学习到的变化。与公式1的卷积类似, local self-attention通过结合混合权重( ${\mathrm{softmax}}_{ab}(\cdot )$ ) 与值向量进行输出，每个位置 $ij$ 都重复上述步骤。

Self-Attention公式分析

在这里解释一下公式：

1. 先解释一下$W_Q,W_K,W_V$的维数为什么是$d_{out}\times d_{in}$，以 $W_Q$ 为例，对于模型的输入$x$,它的通道数是$d_{in}$,维度是 $d_{in}\times w\times h\times batch$,我们对输入$x$进行降维处理，即用矩阵$W_Q$ 和 $x$(需要先对 $x$ 实施reshape变成二维矩阵)做矩阵乘法：
   $(d_{out}\times d_{in})\times (d_{in}\times (w\ast h\ast batch))$ ，
   然后再做reshape,从而得到输出的查询矩阵： $Q$ 的维度为 $d_{in}\times w\times h\times batch$。
   注：实际在代码上都是使用卷积层做降维处理，直接用$1\times 1$的卷积核做处理，就可以得到输出，输出的$d_{out}$就是卷积核的channel数，这样就避免了很多reshape操作。

2. 因为考虑的是一个像素点$ij$在周围一小块区域的被关注的程度，所以取了一个 ${\mathcal{N}}_k(i,j)=\{a,b||a-i|\le k/2,|b-j\mid \le k/2\}$
   ${\mathcal{N}}_k(i,j)$ 就是一个区域，和卷积核一样。
   比如现在令$k=3,i=2,j=2$,所以 ${\mathcal{N}}_k(i,j)=\{a,b||a-2|\le 1.5,\mid b-2\mid \le 1.5\}=\{(1,1)，(1,2)，(1,3)，(2,1)，(2,2)，(2,3)，(3,1)，(3,2)，(3,3)\}$

3. $q_{ij}^T{k}_{ab}$的矩阵乘法操作—— $q_{ij}$就是查询像素矩阵中位置为$ij$的像素向量，它由一个像素点$ij$在所有通道数上的那个值组成，所以对于每个$batch$，$q_{ij}$ 的维度就是$d_{out}\times 1\times 1$，同理 $k_{ab}$ 的维度就是$d_{out}\times k\times k$ 。
   那么$q_{ij}^T{k}_{ab}$就是(先reshape把尺寸长×宽变成一个维度)：
   $((1\ast 1)\times d_{out})\times (d_{out}\times (k\ast k))$,然后经过$softmax$处理，最终得到的就是$1\times (k\ast k)$。

4. 与值元素$v_{ab}$再做矩阵乘法的操作——其实上一步得到的就相当于一个卷积核，这一步就相当于一个卷积操作。只不过根据论文里的实现方式可以进行如下描述：因为$v_{ab}$的维度是$d_{out}\times k\times k$ (这里也是对于每个$batch$而言的)，先reshap成$(k\ast k)\times d_{out}$,然后与 $v_{ab}$ 做矩阵乘法，就得到了最终结果：$1\times d_{out}$。

5. 这里的每一个矩阵相乘都是在每个$batch$上操作的，实现的话其实不需要手动遍历$batch$，可以使用torch.bmm()函数，输入就是两个三维的矩阵，三个维度的意义分别为1是batch,2和3是矩阵的长和宽。函数功能顾名思义，在b(batch)上进行mm(matrix multiplication)操作。

相对位置编码

在实际中，使用multiple attention heads来学习输入的多个独立表达，将像素特征 $x_{ij}$ 分为 $N$ 组 $x_{ij}^n\in {\mathbb{R}}^{d_{in}/N}$, 每个head用不同的变化参数 $W_Q^n,W_K^n,W_V^n\in R^{d_{\text{out }}/N\times d_{\text{in }}/N}$ 进行single-headed attention计算, 最后将结果concatenate成最终的输出 $y_{ij}\in {\mathbb{R}}^{d_{out}}$ 。

公式2中没有使用位置信息，而之前的研究指出相对位置编码能为 self-attention 带来明显的精度提升。因此，使用二维相对位置编码(relative attention)，相对注意力首先定义$ij$到每个位置的相对距离 $ab\in N_k(i,j)$。 相对距离是跨维度分解的，因此每个元素 $ab\in N_k(i,j)$ 接收两个距离：行偏移 $a-i$ 和列偏移 $b-j$（见图 4）。 行和列偏移分别与嵌入 $r_{a-i}$ 和 $r_{b-j}$ 相关联，每个都具有 $\frac{1}{2}{d}_{out}$ 维，然后将行偏移 $a-i$ 和列偏移 $b-j$ 的编码进行concatenate成 $r_{a-i,b-j}$ ，最后将公式2变成公式3的spatial-relative attention。

相对位置编码公式分析

那么$r_{a-i,b-j}$是如何得到的呢，我们在此不用复杂的公式来定义它，根据前人的工作，我们知道这个相对位置编码是可以学习到的，所以我们仅仅使用一组随机数初始化就好了，然后让模型自己学习它。注意：这组随机数应该和Q,K,V的值没有任何关系，只和相对位置有关系。
这里涉及到代码了，简单介绍一下思路，首先要可以定义一个nn.Embedding()或者nn.Parameter()来为位置信息随机编码，原理可以看我的那个文章：nn.Embedding()。
我们这里就以nn.Embedding()为例，当令$k=3,i=2,j=2$,所以 ${\mathcal{N}}_k(i,j)=\{a,b||a-2|\le 1.5,\mid b-2\mid \le 1.5\}=\{(1,1)，(1,2)，(1,3)，(2,1)，(2,2)，(2,3)，(3,1)，(3,2)，(3,3)\}$ 。即每个kernel的大小为3×3，那么我们先实例化nn.Embedding()

1. 令embed = nn.Embedding(3， $\frac{1}{2}{d}_{out}$)，用来生成3个维度为 $\frac{1}{2}{d}_{out}$ 的嵌入向量以供编码。
2. 然后，根据论文里说的，行偏移和列偏移是分开计算并且编码然后再concatenate得到最终的编码信息的，所以我们现在先看横坐标，用 $(a,b)$ 的每个横坐标 $a$ 去减 $i(=2)$,得到的结果就是：
   [-1, -1, -1, 0, 0, 0, 1, 1, 1] ，因为在用nn.Embedding()编码时，这个位置信息实际上充当的索引值，然后再对每个元素加上$int(\frac{k}{2})$,这里就是加1，得到 [0, 0, 0, 1, 1, 1, 2, 2, 2] ，这样这些值都是大于等于0的整数了。
3. 接下来就可以刚刚实例化的 embed 来是生成行编码信息了。
   同理我们可以得到列编码信息。
   然后再进行concatenate,就可以生成最终的相对位置信息。
4. 注意把这个位置编码添加上模型的参数model.parameters()中，这样就可以训练了。

计算量分析

同时考虑query和key的内容间的相似性以及相对位置，相对于绝对位置编码，使用相对位置信息从某种角度来说变相地维持了类似卷积的平移不变性。另外, 与使用典型的卷积层相比，注意力层的计算成本随着注意力区域的增加增长很慢，只稍微 $d_{in}$ 和 $d_{\text{out }}$ 有关。例如，如果 $d_{in}=d_{\text{out }}$=128，k=3 的卷积层与 k=19 的注意力层具有差不多相同的计算成本。

这里再解释一下：
我们这里使用MACC作为评估计算成本的指标，MACC的意思是：乘法累加运算，也称为MADD，就是把一次乘法和一次加法合起来看成一次。
这里简单定性分析一下：

卷积层

1. 对于卷积层，生成的单通道 feature map 中每个像素位置$ij$都是经过了 $K\times K\times Channe{l}_{in}$ 次计算
2. 而一个feature map的尺寸为$H_{out}\times W_{out}$所以单通道 feature map 的计算就为：
   $K\times K\times Channe{l}_{in}\times H_{out}\times W_{out}$。
3. 生成的feature map 共有 $Channe{l}_{out}$个通道，那么这个卷积层的计算量就为：
   $K\times K\times Channe{l}_{in}\times Channe{l}_{out}\times H_{out}\times W_{out}$

注意力层

1. 首先是使用卷积层生成了三个矩阵，这个过程就是卷积层的过程，但是使用是$1\times 1$的卷积层，所以一共的计算量就是：
   $3\times 1\times 1\times Channe{l}_{in}\times Channe{l}_{out}\times H_{out}\times W_{out}$。
2. 然后要进行注意力的计算，我们知道每一个 $q_{ij}$ 都会生成一个 $y_{ij}$，而生成的过程是一个$q_{ij}$向量和两块大小为$K\times K$的区域矩阵(这个矩阵的每一个元素都是一个维度为$d_{out}的向量$)的每一个向量做计算，那么计算量为$2\times K\times K\times d_{out}$，因为输出的尺寸为$H_{out}\times W_{out}$,所以注意力的计算量共有：
   $2\times K\times K\times d_{out}\times H_{out}\times W_{out}$
3. 把两者相加，关于注意力层的计算量就为：
   $3\times 1\times 1\times Channe{l}_{in}\times Channe{l}_{out}\times H_{out}\times W_{out}+2\times K\times K\times channe{l}_{out}\times H_{out}\times W_{out}$
   注：$channe{l}_{out}=d_{out}$
4. 那么看看论文给的例子：
   k=3的卷积层计算量为：
   $3\times 3\times 128\times 128\times H_{out}\times W_{out}$
   k=19的注意力层计算量为：
   $3\times 128\times 128\times H_{out}\times W_{out}+2\times 19\times 19\times 128\times H_{out}\times W_{out}$
   两者做一下除法：
   $$\frac{3\times 3\times 128\times 128\times H_{out}\times W_{out}}{3\times 128\times 128\times H_{out}\times W_{out}+2\times 19\times 19\times 128\times H_{out}\times W_{out}}=\frac{3\times 3\times 128}{3\times 128\times +2\times 19\times 19}\approx 1.04$$
   果然差不多。。。。。

剩下的看论文应该啥问题了，就不写了，如有写错，再更正。

参考：
实战级Stand-Alone Self-Attention in CV | NeurIPS 2019