论文筆記: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}
其中 QK 的維度必須相同。

在此基礎上,引入一個歸一化函數 \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 公式又可以重新定義為:

D(\boldsymbol{Q}, \boldsymbol{K}, \boldsymbol{V})_{i}=\frac{\sum_{j=1}^{N} \phi\left(\boldsymbol{q}_{i}\right)^{T} \varphi\left(\boldsymbol{k}_{j}\right) \boldsymbol{v}_{j}}{\sum_{j=1}^{N} \phi\left(\boldsymbol{q}_{i}\right)^{T} \varphi\left(\boldsymbol{k}_{j}\right)}=\frac{\phi\left(\boldsymbol{q}_{i}\right)^{T} \sum_{j=1}^{N} \varphi\left(\boldsymbol{k}_{j}\right) \boldsymbol{v}_{j}^{T}}{\phi\left(\boldsymbol{q}_{i}\right)^{T} \sum_{j=1}^{N} \varphi\left(\boldsymbol{k}_{j}\right)}\tag{7}

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}
使用向量形式的寫法:

D(\boldsymbol{Q}, \boldsymbol{K}, \boldsymbol{V})=\frac{\sum_{j} \boldsymbol{V}_{i, j}+\left(\frac{\boldsymbol{Q}}{\|\boldsymbol{Q}\|_{2}}\right)\left(\left(\frac{\boldsymbol{K}}{\|\boldsymbol{K}\|_{2}}\right)^{T} \boldsymbol{V}\right)}{N+\left(\frac{\boldsymbol{Q}}{\|\boldsymbol{Q}\|_{2}}\right) \sum_{j}\left(\frac{\boldsymbol{K}}{\|\boldsymbol{K}\|_{2}}\right)_{i, j}^{T}}\tag{11}

其中 \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

Add a Comment

發佈留言必須填寫的電子郵件地址不會公開。 必填欄位標示為 *