论文笔记：Linear Attention Mechanism: An Efficient Attention for Semantic Segmentation
由于经典的 Dot product Attention 具有较大的空间和时间复杂度，虽然效果很好，但由于图像和视频输入的数据量经常比文字大很多，这种经典的 Attention 限制了在高分辨率图像中的应用。本文提出了一种线性的 Attention 模块，希望能够解决这样的问题同时在性能上不会有明显的损耗。
本文主要贡献点有：
1）提出了一种新的线性化 Attention 方式，能够将计算复杂度从 \(𝑂(𝑁^2)\) 降低到 \(𝑂(𝑁)\)
2）该 Attention 机制使得 Attention 与网络的结合更加通用和灵活；
3）在语义分割任务中，该 Attention 机制的引入在多种 baseline 上获得了性能提升。
1 METHODOLOGY
A. Definition of Dot-Product Attention
给定一个输入的特征向量 \(X=\left[\boldsymbol{x}_{1}, \cdots, \boldsymbol{x}_{N}\right] \in \mathbb{R}^{N \times D_{x}}\)，其中 \(N\) 代表特征长度，\(D_x\) 代表特征维度。那么 Dot-Product Attention 通过点乘变换矩阵 \(\boldsymbol{W}_{q} \in \mathbb{R}^{D_{x} \times D_{k}}\)、\(\boldsymbol{W}_{k} \in \mathbb{R}^{D_{x} \times D_{k}}\) 以及 \(\boldsymbol{W}_{v}= \mathbb{R}^{D_{x} \times D_{v}}\) 分别生成 Query Matrix，Key Matrix 和 Value Matrix ，其公式如下：
\(\begin{aligned}
&\boldsymbol{Q}=\boldsymbol{X} \boldsymbol{W}_{q} \in \mathbb{R}^{N \times D_{k}} \\
&\boldsymbol{K}=\boldsymbol{X} \boldsymbol{W}_{k} \in \mathbb{R}^{N \times D_{k}} \\
&\boldsymbol{V}=\boldsymbol{X} \boldsymbol{W}_{\boldsymbol{v}} \in \mathbb{R}^{\boldsymbol{N} \times \boldsymbol{D}_{v}}
\end{aligned}\tag{1}\)
其中 Q 和 K 的维度必须相同。
在此基础上，引入一个归一化函数 \(\rho\) 来度量 \(\boldsymbol{q}_{i}^{T} \in \mathbb{R}^{D_{k}}\) 和 \(\boldsymbol{k}_{j} \in \mathbb{R}^{D_{k}}\) 的相似度 \(\rho\left(\boldsymbol{q}_{i}^{T} \boldsymbol{k}_{j}\right) \in \mathbb{R}^{1}\)。而 Scaled Dot-Product Attention 就是相当于以 \(\rho\left(\boldsymbol{q}_{i}^{T} \boldsymbol{k}_{j}\right) \in \mathbb{R}^{1}\) 为权重对 \(v_j\) 做加权平均。Dot-Product Attention 对于整个矩阵的定义如下：
\(D(\boldsymbol{Q}, \boldsymbol{K}, \boldsymbol{V})=\rho\left(\boldsymbol{Q} \boldsymbol{K}^{T}\right) \boldsymbol{V}\tag{2}\)
其中最常用的一种归一化函数就是 softmax：
\(\rho\left(\boldsymbol{Q}^{T} \boldsymbol{K}\right)=\operatorname{softmax}_{\text {row }}\left(\boldsymbol{Q} \boldsymbol{K}^{T}\right)\tag{3}\)
对于 \(\boldsymbol{Q} \in \mathbb{R}^{N \times D_{k}}\) 和 \(\boldsymbol{K}^{T} \in \mathbb{R}^{D_{k} \times N}\) ，\(\boldsymbol{Q}\boldsymbol{K}^{T} \in \mathbb{R}^{N \times N}\) 因此 \(\rho\) 的时间和空间复杂度均为 \(𝑂(𝑁^2)\)。
B. Generalization of Dot-Product Attention Based on Kernel
基于 Kernal 对 Dot-Product Attention 进行一般化定义。对于 softmax 形式归一化函数来说，可以写成：
\(D(\boldsymbol{Q}, \boldsymbol{K}, \boldsymbol{V})_{i}=\frac{\sum_{j=1}^{N} e^{\boldsymbol{q}_{i}{ }^{T} \boldsymbol{k}_{j}} \boldsymbol{v}_{j}}{\sum_{j=1}^{N} e^{\boldsymbol{q}_{i}{ }^{T} \boldsymbol{k}_{j}}}\tag{4}\)
\(e^{\boldsymbol{q}_i^{\top}\boldsymbol{k}_j}\) 就是对于 \(v_j\) 做加权平均的权重， 该公式可以更进一步一般化定义为：
\(\begin{gathered}
D(\boldsymbol{Q}, \boldsymbol{K}, \boldsymbol{V})_{i}=\frac{\sum_{j=1}^{N} \operatorname{sim}\left(\boldsymbol{q}_{i}, \boldsymbol{k}_{j}\right) \boldsymbol{v}_{j}}{\sum_{j=1}^{N} \operatorname{sim}\left(\boldsymbol{q}_{i}, \boldsymbol{k}_{j}\right)} \\
\operatorname{sim}\left(\boldsymbol{q}_{i}, \boldsymbol{k}_{j}\right) \geq 0
\end{gathered}\tag{5}\)
显然当定义 \(\operatorname{sim}\left(\boldsymbol{q}_{i}, \boldsymbol{k}_{j}\right)=e^{\boldsymbol{q}_{i}{ }^{T} \boldsymbol{k}_{j}}\) 时，该式和标准 softmax 的 Attention 公式等价。这种形式的 Attention 在 CV 中一般也称作 Non-Local 网络。
显然如果直接去掉指数定义 \(\text{sim}(\boldsymbol{q}_i, \boldsymbol{k}_j) = \boldsymbol{q}_i^{\top}\boldsymbol{k}_j\) 不满足非负性质，为了完成这一点可以考虑加入核函数，更改定义为：
\(\text{sim}(\boldsymbol{q}_i, \boldsymbol{k}_j) = \phi(\boldsymbol{q}_i)^{\top} \varphi(\boldsymbol{k}_j)\tag{6}\)
那么 Attention 公式又可以重新定义为：
C. Linear Attention Mechanism
与之前很多线性化的方式不同，本文利用一阶泰勒展开来进行线性化：
\(e^{\boldsymbol{q}_{i}^{T} \boldsymbol{k}_{j}} \approx 1+\boldsymbol{q}_{i}{ }^{T} \boldsymbol{k}_{j}\tag{8}\)
为了保证 \(\boldsymbol{q}_{i}{ }^{T} \boldsymbol{k}_{j} \geq-1\) 我们将 \(\boldsymbol{q}_{i}\) 和 \(\boldsymbol{k}_{j}\) 进行 l2 norm：
\(\operatorname{sim}\left(\boldsymbol{q}_{i}, \boldsymbol{k}_{j}\right)=1+\left(\frac{\boldsymbol{q}_{i}}{\left\|\boldsymbol{q}_{i}\right\|_{2}}\right)^{T}\left(\frac{\boldsymbol{k}_{j}}{\left\|\boldsymbol{k}_{j}\right\|_{2}}\right) .\tag{9}\)
则 Attention 公式又可以重写为：
\(D(\boldsymbol{Q}, \boldsymbol{K}, \boldsymbol{V})_{i}=\frac{\sum_{j=1}^{N}\left(1+\left(\frac{\boldsymbol{q}_{i}}{\left\|\boldsymbol{q}_{i}\right\|_{2}}\right)^{T}\left(\frac{\boldsymbol{k}_{j}}{\left\|\boldsymbol{k}_{j}\right\|_{2}}\right)\right) \boldsymbol{v}_{j}}{\sum_{j=1}^{N}\left(1+\left(\frac{\boldsymbol{q}_{i}}{\left\|\boldsymbol{q}_{i}\right\|_{2}}\right)^{T}\left(\frac{\boldsymbol{k}_{j}}{\left\|\boldsymbol{k}_{j}\right\|_{2}}\right)\right)}=\frac{\sum_{j=1}^{N} \boldsymbol{v}_{j}+\left(\frac{\boldsymbol{q}_{i}}{\left\|\boldsymbol{q}_{i}\right\|_{2}}\right)^{T} \sum_{j=1}^{N}\left(\frac{\boldsymbol{k}_{j}}{\left\|\boldsymbol{k}_{j}\right\|_{2}}\right) \boldsymbol{v}_{j}^{T}}{N+\left(\frac{\boldsymbol{q}_{i}}{\left\|\boldsymbol{q}_{i}\right\|_{2}}\right)^{T} \sum_{j=1}^{N}\left(\frac{\boldsymbol{k}_{j}}{\left\|\boldsymbol{k}_{j}\right\|_{2}}\right)}\tag{10}\)
使用向量形式的写法：
其中 \(\sum_{j=1}^{N}\left(\frac{\boldsymbol{k}_{j}}{\left\|\boldsymbol{k}_{j}\right\|_{2}}\right) \boldsymbol{v}_{j}^{T}\) 和 \(\sum_{j=1}^{N}\left(\frac{\boldsymbol{k}_{j}}{\left\|\boldsymbol{k}_{j}\right\|_{2}}\right)\) 均可以预先计算，并在在每个 query 进行重用。
2 EXPERIMENTAL RESULTS
作者在常见的几类图像分割框架中引入本文提出的 Linear Attention，得到了相比基线均能够涨点的结果，说明了本文提出的结构具有普适的有效性。
论文 & 源码
论文
1636966378-Linear Attention Mechanism- An Efficient Attention for Semantic Segmentation
源码
https://github.com/lironui/Linear-Attention-Mechanism
https://github.com/lironui/MAResU-Net
https://github.com/lironui/MACU-Net
参考文献
[1] https://kexue.fm/archives/7546
