Beyond Self-attention: External Attention using Two Linear Layers for Visual Tasks
- 2021.5.12
- Under review
- https://arxiv.org/abs/2105.02358
Introduction
- self-attention通過計算所有位置上成對特征關系的加權和，來捕獲在單個圖片內long-range各部分的依賴關系，來更新每個位置的特征。但是它有一個很高的復雜度並且忽略了不同樣本之間的關系
- 本文提出了一種新的attention機制基於兩個額外的很小的共享memory，只需要用簡單的兩個級聯線性層和BN層就行 可以方便的代替attention，只有線性的復雜度並可以考慮樣本之間關系
- 對於之前的self-attention，一般的做法是通過計算query向量和key向量逐漸之間的關系來生成attention map,再加權乘在value上
- 本方法不同之處在於，通過計算query向量和一個可學習的key內存來得到attention map，生成一個細化的特征圖，然后通過該attention map乘上另一個memory所存儲的value，來生成特征圖
- 兩個memory通過線性層的方式實現，獨立於單個樣本，並且在整個數據集中共享參數，具有強大的正則作用並且提高了注意力的泛化機制。其之所以輕量化的核心在於memory中的元素數量比輸入特征中少很多，只具有線性復雜度。這個額外的memory設計來學習整個數據集上最有區分度的特征，捕獲最有信息的部分
Method
- 可以看到attention的結構對比圖如下
-
我們是在self-attention的基礎上來實現external-attention的，對於一個輸入特征圖\(F\in\mathbb{R}^{N \times d}\)，\(N\)是像素數量，\(d\)是特征維數，self-attention首先通過線性變化將其變成
-
query matrix \(Q \in \mathbb{R}^{N \times d^{'}}\)
-
key matrix \(K \in \mathbb{R}^{N \times d^{'}}\)
-
value matrix \(V\in\mathbb{R}^{N \times d}\)
-
self-attention可以描述成
\[\begin{array}{rl} A&= (\alpha)_{i,j}=softmax(QK^T) \\ F_{out}&= AV \end{array} \]
-
-
對於簡化版的self-attention，就直接利用特征圖來計算attention，可以描述成
\[\begin{array}{rl}A&= (\alpha)_{i,j}=softmax(FF^T) \\F_{out}&= AF\end{array} \] -
對attention進行可視化之后可以發現大多數像素點只和一部分像素點有關，沒有必要計算\(N \times N\)的attention map，所以我們提出了額外attention模塊，通過計算像素與額外的存儲單元\(M\in\mathbb{R}^{S \times d}\)得到attention，用和self-attention相似的方法進行歸一化，M是一個獨立於輸入可學習的參數，作為整個數據集的存儲單元，並且我們用不同的兩個M分別作為key和value，所以整個external attention可以描述成
\[\begin{array}{rl}A&= (\alpha)_{i,j}=Norm(FM_k^T) \\F_{out}&= AM_v\end{array} \] -
d和S都屬於超參，一個較小的S在實驗中表現就很不錯，所以external attention比self-attention更有效，應用在了多種任務上，論文中給出的偽代碼如下圖
- 計算量對比上，和之前的一些attention工作進行了對比，但是並沒有和ViT對比，這很奇怪
Conclusion
文章中做了多個實驗來證明其有效性，包括分類、分割、目標檢測、圖像生成，點雲等等，從結果來看確實是有效的，但是沒有任何的ablation study，和最新的transformer之間的對比還不夠，猜測最近幾天眾多與線性層相關的文章都會出現在今年的NeurIPS上，不知道專業的reviewer怎么看。我是感覺用一個線性層來存儲整個數據集中的所有信息來得到這個attention是不是有點太草率了，但它結果確實也nb，在可視化上來看也是work的，畢竟是大佬的組出來的文章，但也還是希望能夠看到和最近眾多視覺上transformer相關工作的結果進行更全面的對比。感謝該paper帶給我的啟發