EANet解读
最近關于MLP的工作還是蠻多的,先是MLP-Mixer為視覺開了個新思路,接著EANet(即External Attention)和RepMLP橫空出世,乃至最新的《Do You Even Need Attention? A Stack of Feed-Forward Layers Does Surprisingly Well on ImageNet》質疑Transformer中attention的必要性。由于我個人對自注意力是比較關注的,因此EANet通過線性層和歸一化層即可替代自注意力機制還是很值得關注的。
論文作者最近發布了一篇解讀在知乎上,補充了多頭注意力的實現,感興趣的可以參考一下,鏈接給出。
簡介
自注意力其實早已在計算機視覺中有所應用,從較早的Non-Local到最近的Transformer,計算機視覺兜兜轉轉還是有回到全局感知的趨勢。相比于卷積這種局部感知的操作,自注意力計算每個位置與其他所有位置信息的加權求和來更新當前位置的特征,從而捕獲長程依賴(這在NLP中至關重要)獲取全局信息。但是自注意力的劣勢也是非常明顯的,一方面自注意力對每個位置都要計算與其他位置的關系,這是一種二次方的復雜度,是非常消耗資源的;另一方面,自注意力對待每個樣本同等處理,其實忽略了每個樣本之間的潛在關系。針對此,清華計圖團隊提出了一種新的external attention(外部注意力),僅僅通過兩個可學習的單元即可取代現有的注意力并將復雜度降到和像素數目線性相關,而且由于這個單元全數據集共享所以它能隱式考慮不同樣本之間的關系。實驗表明,由external attention構建的EANet在圖像分類、語義分割、圖像生成、點云分類和分割等任務上均達到甚至超越自注意力結構的表現,并且有著小得多的計算量和內存開銷。
-
論文標題
Beyond Self-attention: External Attention using Two Linear Layers for Visual Tasks
-
論文地址
https://arxiv.org/abs/2105.02358v1
-
論文源碼
https://github.com/MenghaoGuo/EANet
EA
論文中關于自注意力機制的發展就不多贅述了,這一節按照論文的思路來闡述一下External Attention。
Self-Attention
首先,不妨來回顧一下原始的自注意力機制,如下圖a所示,給定輸入特征圖F∈RN×dF \in \mathbb{R}^{N \times d}F∈RN×d,這里的NNN表示像素數目而ddd則表示特征的維度也就是通道數目。自注意力首先通過三個線性投影將輸入變換為query矩陣Q∈RN×d′Q \in \mathbb{R}^{N \times d^{\prime}}Q∈RN×d′、key矩陣K∈RN×d′K \in \mathbb{R}^{N \times d^{\prime}}K∈RN×d′和value矩陣V∈RN×dV \in \mathbb{R}^{N \times d}V∈RN×d,從而按照下面的式子計算自注意力的輸出,這兩個式子分別表示根據Query和Key計算相似性權重,然后將權重系數作用于Value得到加權求和的注意力結果。 至于這里的點積相似性只是相似性計算的一種方式而已。
A=(α)i,j=softmax?(QKT)Fout?=AV\begin{aligned} &A =(\alpha)_{i, j}=\operatorname{softmax}\left(Q K^{T}\right) \\ &F_{\text {out }} =A V \end{aligned} ?A=(α)i,j?=softmax(QKT)Fout??=AV?
在上面這個式子中,A∈RN×NA \in \mathbb{R}^{N \times N}A∈RN×N表示注意力矩陣并且αi,j\alpha_{i, j}αi,j?就是第iii個像素和第jjj個像素之間的逐對相似性。
既然QKV的計算都是線性變換,因此可以省略這個過程直接對輸入特征進行點積相似度的計算,從而得到注意力圖,因而簡化版的自注意力如下式,這對應上圖的b。然而,這種算法雖然簡化了,但是O(dN2)\mathcal{O}\left(d N^{2}\right)O(dN2)的計算復雜度大大限制了自注意力的使用,特別是用于圖像這種像素眾多的數據上,因此很多的工作都是patch之間計算自注意力而不是在pixel上計算。
A=softmax?(FFT)Fout?=AF\begin{aligned} &A =\operatorname{softmax}\left(F F^{T}\right) \\ &F_{\text {out }} =A F \end{aligned} ?A=softmax(FFT)Fout??=AF?
External Attention
論文作者通過可視化注意力圖,發現大部分像素其實只和少量的其他像素密切相關,因此構建一個N-to-N的注意力矩陣其實是非常冗余的。因此自然產生一個想法,能否通過少量需要的值來精煉原始的輸入特征呢?這就誕生了external attention模塊來替代原始的注意力,它在輸入特征和一個外部單元M∈RS×dM \in \mathbb{R}^{S \times d}M∈RS×d計算注意力。
A=(α)i,j=Norm?(FMT)Fout?=AM\begin{aligned} &A =(\alpha)_{i, j}=\operatorname{Norm}\left(F M^{T}\right) \\ &F_{\text {out }} =A M \end{aligned} ?A=(α)i,j?=Norm(FMT)Fout??=AM?
上面的式子表示EA的計算過程,這里的M是一個獨立于輸入的可學習參數,它相當于整個數據集的memory,因此它是所有樣本共享的,可以隱式考慮所有樣本之間的影響。A是通過先驗推斷出的注意力圖,它類似于自注意力那樣進行標準化,然后根據A和M對輸入特征更新。
實際上,為了增強EA的表達能力,與自注意力相似,論文采用了兩個記憶單元分別是MkM_kMk?和MvM_vMv?分別作為key和value,因此External Attention可以表示為下面修改后的式子。
A=Norm?(FMkT)Fout?=AMv\begin{aligned} &A =\operatorname{Norm}\left(F M_{k}^{T}\right) \\ &F_{\text {out }} =A M_{v} \end{aligned} ?A=Norm(FMkT?)Fout??=AMv??
從這個式子和上面的圖c不難發現,EA的計算過程其實就是線性運算和Norm層,源碼中線性層采用的是1x1卷積實現的,Norm層的實現下面會提到。整個EA的計算復雜度其實為O(dSN)\mathcal{O}(d S N)O(dSN),S和d都是超參數,這個復雜度因此和像素數目線性相關了,通過修改S的大小可以方便地控制計算的復雜度,事實上,實驗中,作者發現S設置為64效果就已經很好了。因此,EA通過線性層就替代了自注意力機制并且非常高效輕量,非常適用于大尺度的輸入特征圖。
Normalization
在自注意力中,Softmax常被用于注意力圖的標準化以實現∑jαi,j=1\sum_{j} \alpha_{i, j}=1∑j?αi,j?=1,然而注意力圖常常通過矩陣乘法計算得到,不同于余弦相似度,這種算法的相似度是尺度敏感的。因此作者對double-attention優化得到了下面這個分別行列進行標準化的方法。
(α~)i,j=FMkTαi,j=exp?(α~i,j)∑kexp?(α~k,j)αi,j=αi,j∑kαi,k\begin{aligned} (\tilde{\alpha})_{i, j} &=F M_{k}^{T} \\ \alpha_{i, j} &=\frac{\exp \left(\tilde{\alpha}_{i, j}\right)}{\sum_{k} \exp \left(\tilde{\alpha}_{k, j}\right)} \\ \alpha_{i, j} &=\frac{\alpha_{i, j}}{\sum_{k} \alpha_{i, k}} \end{aligned} (α~)i,j?αi,j?αi,j??=FMkT?=∑k?exp(α~k,j?)exp(α~i,j?)?=∑k?αi,k?αi,j???
下圖是EANet的結構示意圖,事實上就是在head之前添加了一個EA模塊,關于EA模塊的實現可以參考下面官方開源的Jittor的實現代碼。
class External_attention(Module):'''Arguments:c (int): The input and output channel number.'''def __init__(self, c):super(External_attention, self).__init__()self.conv1 = nn.Conv2d(c, c, 1)self.k = 64self.linear_0 = nn.Conv1d(c, self.k, 1, bias=False)self.linear_1 = nn.Conv1d(self.k, c, 1, bias=False)self.linear_1.weight = self.linear_0.weight.permute(1, 0, 2) self.conv2 = nn.Sequential(nn.Conv2d(c, c, 1, bias=False),nn.BatchNorm(c)) self.relu = nn.ReLU()def execute(self, x):idn = xx = self.conv1(x)b, c, h, w = x.size()n = h*wx = x.view(b, c, h*w) # b * c * n attn = self.linear_0(x) # b, k, nattn = nn.softmax(attn, dim=-1) # b, k, nattn = attn / (1e-9 + attn.sum(dim=1, keepdims=True)) # # b, k, nx = self.linear_1(attn) # b, c, nx = x.view(b, c, h, w)x = self.conv2(x)x = x + idnx = self.relu(x)return x實驗
為了證明方法的有效性,作者進行了大量任務的實驗,包括圖像分類、語義分割、圖像生成、點云分類和點云分割,我這里就只列一下圖像分類數據集ImageNet上的結果,在其他任務上EANet均能達到SOTA效果。
此外,關于自注意力及其變種也進行了計算消耗的比較,EA真的是非常輕量。
總結
論文提出了一種新的輕量級視覺注意力結構External Attention,克服了自注意力的超大計算量和缺乏樣本多樣性考慮的缺點,通過簡單的線性層即可實現自注意力,并且效果相當不錯,是很值得關注的一個方法。本文也只是我本人從自身出發對這篇文章進行的解讀,想要更詳細理解的強烈推薦閱讀原論文。最后,如果我的文章對你有所幫助,歡迎一鍵三連,你的支持是我不懈創作的動力。
總結
- 上一篇: SL解读
- 下一篇: RelationTrack解读