由於經典的 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