最近關(guān)于MLP的工作還是蠻多的，先是MLP-Mixer為視覺開了個新思路，接著EANet（即External Attention）和RepMLP橫空出世，乃至最新的《Do You Even Need Attention? A Stack of Feed-Forward Layers Does Surprisingly Well on ImageNet》質(zhì)疑Transformer中attention的必要性。由于我個人對自注意力是比較關(guān)注的，因此EANet通過線性層和歸一化層即可替代自注意力機制還是很值得關(guān)注的。

論文作者最近發(fā)布了一篇解讀在知乎上，補充了多頭注意力的實現(xiàn)，感興趣的可以參考一下，鏈接給出。

簡介

自注意力其實早已在計算機視覺中有所應(yīng)用，從較早的Non-Local到最近的Transformer，計算機視覺兜兜轉(zhuǎn)轉(zhuǎn)還是有回到全局感知的趨勢。相比于卷積這種局部感知的操作，自注意力計算每個位置與其他所有位置信息的加權(quán)求和來更新當(dāng)前位置的特征，從而捕獲長程依賴（這在NLP中至關(guān)重要）獲取全局信息。但是自注意力的劣勢也是非常明顯的，一方面自注意力對每個位置都要計算與其他位置的關(guān)系，這是一種二次方的復(fù)雜度，是非常消耗資源的；另一方面，自注意力對待每個樣本同等處理，其實忽略了每個樣本之間的潛在關(guān)系。針對此，清華計圖團隊提出了一種新的external attention（外部注意力），僅僅通過兩個可學(xué)習(xí)的單元即可取代現(xiàn)有的注意力并將復(fù)雜度降到和像素數(shù)目線性相關(guān)，而且由于這個單元全數(shù)據(jù)集共享所以它能隱式考慮不同樣本之間的關(guān)系。實驗表明，由external attention構(gòu)建的EANet在圖像分類、語義分割、圖像生成、點云分類和分割等任務(wù)上均達(dá)到甚至超越自注意力結(jié)構(gòu)的表現(xiàn)，并且有著小得多的計算量和內(nèi)存開銷。

• 論文標(biāo)題

  Beyond Self-attention: External Attention using Two Linear Layers for Visual Tasks

• 論文地址

  https://arxiv.org/abs/2105.02358v1

• 論文源碼

  https://github.com/MenghaoGuo/EANet

EA

論文中關(guān)于自注意力機制的發(fā)展就不多贅述了，這一節(jié)按照論文的思路來闡述一下External Attention。

Self-Attention

首先，不妨來回顧一下原始的自注意力機制，如下圖a所示，給定輸入特征圖$F\in {\mathbb{R}}^{N\times d}$，這里的$N$表示像素數(shù)目而$d$則表示特征的維度也就是通道數(shù)目。自注意力首先通過三個線性投影將輸入變換為query矩陣$Q\in {\mathbb{R}}^{N\times d^{\prime }}$、key矩陣$K\in {\mathbb{R}}^{N\times d^{\prime }}$和value矩陣$V\in {\mathbb{R}}^{N\times d}$，從而按照下面的式子計算自注意力的輸出，這兩個式子分別表示根據(jù)Query和Key計算相似性權(quán)重，然后將權(quán)重系數(shù)作用于Value得到加權(quán)求和的注意力結(jié)果。 至于這里的點積相似性只是相似性計算的一種方式而已。

$$\begin{array}{rl}& A=(\alpha {)}_{i,j}=\mathrm{softmax}\left(Q{K}^T\right)\\ & F_{\text{out?}}=AV\end{array}$$

在上面這個式子中，$A\in {\mathbb{R}}^{N\times N}$表示注意力矩陣并且${\alpha }_{i,j}$就是第$i$個像素和第$j$個像素之間的逐對相似性。

既然QKV的計算都是線性變換，因此可以省略這個過程直接對輸入特征進(jìn)行點積相似度的計算，從而得到注意力圖，因而簡化版的自注意力如下式，這對應(yīng)上圖的b。然而，這種算法雖然簡化了，但是$\mathcal{O}\left(d{N}^2\right)$的計算復(fù)雜度大大限制了自注意力的使用，特別是用于圖像這種像素眾多的數(shù)據(jù)上，因此很多的工作都是patch之間計算自注意力而不是在pixel上計算。

$$\begin{array}{rl}& A=\mathrm{softmax}\left(F{F}^T\right)\\ & F_{\text{out?}}=AF\end{array}$$

External Attention

論文作者通過可視化注意力圖，發(fā)現(xiàn)大部分像素其實只和少量的其他像素密切相關(guān)，因此構(gòu)建一個N-to-N的注意力矩陣其實是非常冗余的。因此自然產(chǎn)生一個想法，能否通過少量需要的值來精煉原始的輸入特征呢？這就誕生了external attention模塊來替代原始的注意力，它在輸入特征和一個外部單元$M\in {\mathbb{R}}^{S\times d}$計算注意力。

$$\begin{array}{rl}& A=(\alpha {)}_{i,j}=\mathrm{Norm}\left(F{M}^T\right)\\ & F_{\text{out?}}=AM\end{array}$$

上面的式子表示EA的計算過程，這里的M是一個獨立于輸入的可學(xué)習(xí)參數(shù)，它相當(dāng)于整個數(shù)據(jù)集的memory，因此它是所有樣本共享的，可以隱式考慮所有樣本之間的影響。A是通過先驗推斷出的注意力圖，它類似于自注意力那樣進(jìn)行標(biāo)準(zhǔn)化，然后根據(jù)A和M對輸入特征更新。

實際上，為了增強EA的表達(dá)能力，與自注意力相似，論文采用了兩個記憶單元分別是$M_k$和$M_v$分別作為key和value，因此External Attention可以表示為下面修改后的式子。

$$\begin{array}{rl}& A=\mathrm{Norm}\left(F{M}_k^T\right)\\ & F_{\text{out?}}=A{M}_v\end{array}$$

從這個式子和上面的圖c不難發(fā)現(xiàn)，EA的計算過程其實就是線性運算和Norm層，源碼中線性層采用的是1x1卷積實現(xiàn)的，Norm層的實現(xiàn)下面會提到。整個EA的計算復(fù)雜度其實為$\mathcal{O}(dSN)$，S和d都是超參數(shù)，這個復(fù)雜度因此和像素數(shù)目線性相關(guān)了，通過修改S的大小可以方便地控制計算的復(fù)雜度，事實上，實驗中，作者發(fā)現(xiàn)S設(shè)置為64效果就已經(jīng)很好了。因此，EA通過線性層就替代了自注意力機制并且非常高效輕量，非常適用于大尺度的輸入特征圖。

Normalization

在自注意力中，Softmax常被用于注意力圖的標(biāo)準(zhǔn)化以實現(xiàn)${\sum }_j{\alpha }_{i,j}=1$，然而注意力圖常常通過矩陣乘法計算得到，不同于余弦相似度，這種算法的相似度是尺度敏感的。因此作者對double-attention優(yōu)化得到了下面這個分別行列進(jìn)行標(biāo)準(zhǔn)化的方法。

$$\begin{array}{rl}(\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{array}$$

下圖是EANet的結(jié)構(gòu)示意圖，事實上就是在head之前添加了一個EA模塊，關(guān)于EA模塊的實現(xiàn)可以參考下面官方開源的Jittor的實現(xiàn)代碼。

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

實驗

為了證明方法的有效性，作者進(jìn)行了大量任務(wù)的實驗，包括圖像分類、語義分割、圖像生成、點云分類和點云分割，我這里就只列一下圖像分類數(shù)據(jù)集ImageNet上的結(jié)果，在其他任務(wù)上EANet均能達(dá)到SOTA效果。

此外，關(guān)于自注意力及其變種也進(jìn)行了計算消耗的比較，EA真的是非常輕量。

總結(jié)

論文提出了一種新的輕量級視覺注意力結(jié)構(gòu)External Attention，克服了自注意力的超大計算量和缺乏樣本多樣性考慮的缺點，通過簡單的線性層即可實現(xiàn)自注意力，并且效果相當(dāng)不錯，是很值得關(guān)注的一個方法。本文也只是我本人從自身出發(fā)對這篇文章進(jìn)行的解讀，想要更詳細(xì)理解的強烈推薦閱讀原論文。最后，如果我的文章對你有所幫助，歡迎一鍵三連，你的支持是我不懈創(chuàng)作的動力。

總結(jié)

以上是生活随笔為你收集整理的EANet解读的全部內(nèi)容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網(wǎng)站內(nèi)容還不錯，歡迎將生活随笔推薦給好友。