Beyond Self-attention: External Attention using Two Linear Layers for Visual Tasks
Abstract
注意力機制，尤其是自注意力（self-attention），在視覺任務的深度特征表征中起着越來越重要的作用。自注意力通過在所有位置上使用成對的affinities計算特征的加權和來更新每個位置上的特征，以捕獲單個樣本中的long-range依賴關系。然而，自注意力具有二次復雜度，且忽略了不同樣本之間的潛在相關性。本文提出了一種新穎的外部注意力機制，即外部注意力機制（external attention），它基於兩個外部的、小的、可學習的共享memories，只需簡單地使用兩個級聯線性層和兩個歸一化層即可實現;使用起取代現有流行結構中的自注意力是很方便的。外部注意力具有線性復雜性，隱式考慮了所有數據樣本之間的相關性。我們進一步將multi-head機制納入到外部注意力中，以提供用於圖像分類的all-MLP體系結構，即外部注意力MLP ( external attention MLP，EAMLP)。大量關於圖像分類、目標檢測、語義分割、實例分割、圖像生成和點雲分析的實驗表明，我們的方法提供的效果等同或優於自注意力機制及其某些變量，而且計算和存儲成本更低。
1 INTRODUCTION
由於其捕獲長期依賴的能力，自注意力（self-attention）機制有助於提高各種自然語言處理[1]、[2]和計算機視覺[3]、[4]任務的性能。自注意力的工作原理是通過聚合單個樣本中所有其他位置的特征來細化每個位置的表征，這導致了一個樣本中關於位置數量的二次冪計算復雜度。因此，一些變體試圖以較低的計算成本來近似自注意力方法[5]，[6]，[7]，[8]。
此外，自注意力關注的是單個樣本內不同位置之間的親和度（affinities），而忽略了與其他樣本之間的潛在相關性。很容易看出，合並不同樣本之間的相關性有助於更好地表示特征。例如，在語義分割任務中，屬於同一類別但分布在不同樣本中的特征應該得到一致的處理，在圖像分類和其他各種視覺任務中也有類似的觀察結果。
本文提出了一種新的輕量級注意機制，我們稱之為外部注意力(external attention，見圖1c)。如圖1a)所示，計算自注意力需要首先通過計算self query vectors和self key vectors之間的親和度來計算一個注意力map，然后通過該注意力map對self value vectors進行加權來生成一個新的特征map。外部注意力的工作原理是不同的。我們首先通過計算self query vectors與一個external learnable key memory之間的親和度來計算注意力map，然后通過將該注意力map與另一個external learnable value memory相乘來得到一個精細的特征map。
在實踐中，這兩個memories是用線性層實現的，因此可以通過端到端反向傳播的方式進行優化。它們獨立於單個樣本，並在整個數據集中共享，起到了很強的正則化作用，提高了注意力機制的泛化能力。外部注意力的輕量級本質的關鍵在於，memories中的元素數量比輸入特征中的元素數量要少得多，這就產生了關於輸入元素數量的線性計算復雜度。外部memories的目的是學習整個數據集中最具區別性的特征，捕捉信息最豐富的部分，以及排除其他樣本的干擾信息。在稀疏編碼（sparse coding）[9]或字典學習（dictionary learning）[10]中也可以找到類似的思想。然而，與這些方法不同的是，我們既沒有嘗試重建輸入特征，也沒有對注意力map應用任何顯式稀疏正則化。
雖然所提出的外部注意力方法很簡單，但它對各種視覺任務都是有效的。由於它的簡單性，它可以很容易地合並到現有的流行的基於自注意力的架構中，如DANet [4]， SAGAN[11]和T2T-Transformer[12]。圖3展示了一個典型的架構，它用我們的外部注意力來代替自注意力來完成圖像語義分割任務。我們使用不同的輸入模式(圖像和點雲)，在諸如分類、目標檢測、語義分割、實例分割和生成等基本視覺任務上進行了廣泛的實驗。結果表明，我們的方法以更低的計算代價獲得了與原自注意力機制及其某些變體相當或更好的結果。
我們將multi-head機制引入到外部注意力中，以提高外部注意力的能力。受益於提出的multi-head外部注意力，我們設計了一個名為EAMLP的全新的all-MLP體系結構，它可以媲美CNNs和原始的Transformers來完成圖像分類任務。
本文的主要貢獻總結如下:
- 提出一個新的復雜度為O(n)的注意力機制 - 外部注意力；它可以取代現有架構中的自注意力。它可以挖掘整個數據集的潛在關系，具有很強的正則化作用，提高了注意力機制的泛化能力。
- multi-head外部注意力，這有利於我們構建一個all MLP體系結構;它在ImageNet-1K數據集上實現了79.4%的top-1准確度。
- 利用外部注意力對圖像分類，目標檢測，語義分割，實例分割，圖像生成，點雲分類和點雲分割進行了廣泛的實驗。在計算工作量必須保持在較低水平的情況下，它能獲得比原始的自注意力機制及其某些變體更好的結果
2 Related work
由於對注意力機制的全面回顧超出了本文的范圍，我們只討論視覺領域中與之最密切相關的文獻。
2.1 The attention mechanism in visual tasks
注意力機制可以看作是一種根據激活的重要性來重新分配資源的機制。它在人類視覺系統中起着重要的作用。在過去十幾年中，該領域得到了蓬勃的發展[3] [13] [14] [15] [16] [17] [18]。Hu等人提出了SENet[15]，表明注意力機制可以降低噪聲，提高分類性能。隨后，許多其他論文將其應用於視覺任務。Wang等人提出了用於視頻理解的non-local網絡[3],Hu等人[19]將注意力用於目標檢測,Fu等人提出用於語義分割的DANet[4],Zhang等人[11]演示在圖像生成中注意力機制的有效性,和Xie等人提出A-SCN[20]進行點雲處理。
2.2 Self-attention in visual tasks
自注意力是注意力的一種特殊情況，許多論文[3]，[4]，[11]，[17]，[21]都考慮了視覺的自注意力機制。自注意力的核心思想是計算特性之間的親和度，以捕獲長期依賴關系。但是，隨着特征map的增大，計算和內存開銷會呈冪次增長。為了減少計算和內存開銷，Huang等人提出了criss-cross注意力，它依次考慮行注意和列注意來捕獲全局上下文。Li等[6]采用期望最大化(EM)聚類優化自注意力。Yuan等人提出使用object-contextual向量來處理注意力;然而，它依賴於語義標簽。Geng等[8]研究表明，在語義分割和圖像生成中，矩陣分解是一種更好的全局上下文建模方法。其他研究[22]，[23]也探索了使用自注意力機制提取局部信息。
與通過計算self queries和self keys之間的親和度來獲得注意力map的自注意力不同，我們的外部注意力計算self queries和一個更小的可學習key memory之間的關系，后者捕獲數據集的全局上下文。外部注意力不依賴於語義信息，可以通過反向傳播算法進行端到端的優化，而不需要迭代算法。
2.3 Transformer in visual tasks
基於transformer的模型在自然語言處理[1]、[2]、[16]、[24]、[25]、[26]、[27]方面取得了巨大的成功。最近，他們也展示了其在視覺任務的巨大潛力。Carion等人[28]提出了一種端到端detection transformer，它以CNN特征作為輸入，用transformer生成bounding boxes。Dosovitskiy[18]提出了基於patch encoding和transformer的ViT，表明在有足夠的訓練數據的情況下，transformer比傳統的CNN具有更好的性能。Chen等人提出了基於transformer的iGPT圖像生成方法。
隨后，transformer方法已經成功應用於許多視覺任務，包括圖像分類[12]、[30]、[31]、[32]、目標檢測[33]、低級視覺(lower-level vision)[34]、語義分割[35]、跟蹤[36]、視頻實例分割[37]、圖像生成[38]、多模態學習(multimodal learning)[39]、目標再識別(object re-identification)[40]、圖像字幕[41]，點雲學習[42]，自監督學習[43]。讀者可以參考最近的調查[44]，[45]，以更全面地了解transformer方法在視覺任務中的使用。
3 METHODOLOGY
在這一節中，我們將從分析最初的自注意力（self-attention）機制開始。然后我們詳細介紹了我們定義的新方法:外部注意力(external attention)。只需使用兩個線性層和兩個歸一化層，就可以很容易地實現它，如算法1所示。
3.1 Self-Attention and External Attention
我們首先回顧一下自注意力機制(可見圖1a)。給定一個輸入特征map 
，其中N是元素數量（或圖像像素數量），d是特征維度數，自注意力線性投射輸入為一個query矩陣
、一個key矩陣
和一個value矩陣
[16]。自注意力公式為：
其中
是注意力矩陣，
是第i個和第j個元素的成對親和度值。
自注意力的一個常用簡化版變體（圖1b）直接從輸入特征
計算了一個注意力map，式子如下：
這里通過在特征空間中計算像素級相似度得到注意力map，輸出是輸入的精細化特征表征。
然而，即使經過簡化，O(dN2)的高計算復雜度也給自注意力的使用帶來了一個重大缺陷。輸入像素數的二次復雜度使得將自注意力直接應用於圖像是不可行的。因此，之前的工作[18]利用對patches而不是像素的自注意力來減少計算工作量。
自注意力可以看作是self values的線性組合來細化輸入特征。然而，在這個線性組合中，我們是否真的需要N × N的自注意矩陣和N個元素的self value矩陣，這一點並不明顯。此外，自注意力只考慮了數據樣本中元素之間的關系，而忽略了不同樣本中元素之間的潛在關系，潛在地限制了自注意力的能力和靈活性。
因此，我們提出了新的注意力模塊，命名為外部注意力（external attention），用於計算輸入像素和一個external memory單元
的注意力，公式為：
和自注意力不同，等式(5)中的
是第i個像素和M第j行的相似度，其中M是一個獨立於輸入的可學習參數，其作為整個訓練數據集的memory。A是從該已學好的dataset-level先驗知識推導得到的注意力map；以相同的方法歸一化為自注意力（可見3.2節）。最后,我們通過A的相似度更新M的輸入特性。
實際上，我們使用兩個不同的memory單元
和
作為key和value，去增強網絡的能力。因此外部注意力的計算變為：
外部注意力的計算復雜度是O(dSN)；d和S是超參數，該算法在像素數上是線性的。實際上，我們發現一個小的S，如64，就能在實驗中得到好效果。因此，外部注意力是比自注意力高效的，允許其直接應用到大尺度輸入中。我們還注意到外部注意力的計算負載粗略等價於1x1卷積。
3.2 Normalization
自注意力中使用Softmax去歸一化注意力map，因此
。然而，注意力map是通過矩陣乘法計算出來的。與余弦相似度不同，注意力map對輸入特征的scale非常敏感。為了避免這個問題，我們選擇[42]中提出的double-normalization，它分別歸一化列和行。這種double-normalization 公式為:
外部注意力的python版偽代碼為：
3.3 Multi-head external attention
在Transformer[16]中，自注意在不同的輸入channels上計算多次，這稱為multi-head注意力。multi-head注意力可以捕獲tokens之間的不同關系，提高了single head注意力的能力。我們對
multi-head外部注意力可以寫成:
其中hi是第i個head，H是heads的數量，Wo是一個用來保持輸入輸出維度一致的線性轉換矩陣。
和
是不同heads共享的memory單元
這種架構的靈活性允許我們在共享memory單元中平衡head H的數量和元素S的數量。例如，我們可以用H乘以k，同時用S除以k。
4 EXPERIMENTS
我們在圖像分類、目標檢測、語義分割、實例分割、圖像生成、點雲分類和點雲分割任務上進行了實驗，以評估我們提出的外部注意方法的有效性。所有實驗都是用Jittor[86]和/或Pytorch[87]深度學習框架上實現的。
4.1 Ablation study
為了在我們的完整模型中驗證所提出的模塊，我們在PASCAL VOC分割數據集上進行了實驗[88]。圖3描述了用於消融研究的架構，該架構以FCN[46]為特征backbone。batch size和總迭代次數分別設置為12和30000次。我們關注memory單元的數量、自注意力與外部注意力、backbone、歸一化方法和backbone的輸出步幅。如表1所示，在Pascal VOC數據集上，我們可以觀察到外部注意力比自注意力具有更好的准確性。選擇合適數量的memory單元對結果的質量很重要。歸一化方法可以對外部注意力產生巨大的積極影響，並對自注意力效果進行改善。
4.2 Visual analysis
使用外部注意力進行分割的注意力map(見圖3)和使用multi-head 外部注意力進行分類的注意力map(見章節4.3)分別如圖4和圖5所示。我們從一個層的memory單元Mk中隨機選擇一行
。然后通過計算對輸入特征的注意力來繪制注意力map。我們觀察到，在分割任務中，學習到的注意力map集中在有意義的物體或背景上，如圖4所示。圖5中最后兩行說明Mk的不同行注意的區域不同。multi-head外部注意力的每個head都能不同程度地激活興趣區域，如圖5所示，提高了外部注意力的表征能力。
4.3 Image classification
ImageNet-1K[89]是一種廣泛用於圖像分類的數據集。我們將T2T-ViT[12]中的Performer[90]和multi-head自注意力塊替換為外部注意力和multi-head外部注意力。為公平起見，其他超參數設置與T2T-ViT相同。表2的實驗結果表明，外部注意力的結果比Performer[90]更好，比multi-head注意力的結果低約2%。我們發現multi-head機制對自注意力和外部注意力都是必要的。我們還嘗試了MoCo V3[43]提出的策略，以BatchNorm(BN)[92]替代T2T-ViT backbone的MLP塊(非外部注意力塊)中的LayerNorm(LN)[91]。我們觀察到我們的EAMLP-7有1%的改善。然而，它在我們的大型模型EAMLP-14和EAMLP-19中產生了失敗的效果。
4.4 Object detection and instance segmentation
MS COCO數據集[93]是一個流行的目標檢測和實例分割基准。它包含超過200,000張圖片，其中包含來自80個類別的500,000多個帶注釋的對象實例。
MMDetection[47]是一個被廣泛使用的用於對象檢測和實例分割的工具包。我們使用基於RestNet-50 backbone的MMDetection對COCO數據集進行了目標檢測和實例分割實驗。我們只是在Resnet stage 4的末尾才增加了外部注意力。表3和表4的結果顯示，外部注意力使目標檢測和實例分割任務的准確率都提高了1%。
4.5 Semantic segmentation
在本次實驗中，我們采用圖3中的語義分割架構，將其稱為EANet，並將其應用於Pascal VOC[88]、ADE20K[94]和cityscapes[95]數據集。
Pascal VOC包含10,582張用於訓練的圖像，1,449張用於驗證的圖像，1,456張用於測試的圖像。在分割中，它有20個foreground對象類和一個的背景類。與所有比較方法一樣，我們使用輸出stride為8的dilated ResNet-101作為backbone;它是在ImageNet-1K上進行預訓練的。在訓練過程中采用了多元學習率策略（poly-learning rate policy）。初始學習率、batch size和輸入大小分別設置為0.009、16和513 × 513。我們首先在訓練集上訓練了45k次迭代，然后在trainval集上微調了15k次迭代。最后，我們在測試集上使用了multi-scale和flip測試。可視化結果如圖4所示，定量結果如表5所示:我們的方法可以達到與現有方法相當的性能。
ADE20K是一個更具挑戰性的數據集，它有150個類，以及20K、2K和3K張分別用於訓練、驗證和測試的圖像。我們采用dilated ResNet-101作為backbone，輸出stride為8。實驗配置與mmsegmentation[60]相同，訓練ADE20K進行160k次迭代。表6的結果顯示，我們的方法在ADE20K val集上優於其他方法。
Cityscapes在19個用於城市場景理解的語義類中包含了5000個高質量且具有像素級精細注釋的標簽。每個圖像都是1024×2048像素。分為2975張、500張和1525張，分別用於訓練、驗證和測試。(它還包含2萬張粗略標注的圖像，我們在實驗中沒有使用它們)。我們采用dilated ResNet-101作為所有方法的backbone，輸出stride為8。實驗配置與mmsegmentation相同，用80k次迭代訓練Cityscapes。表7的結果顯示，在cityscapes val集合上，我們的方法獲得了與最先進的方法(即DANet[4])相當的結果。
4.6 Image generation
自注意力是圖像生成中常用的方法，典型的方法是SAGAN[11]。我們將SAGAN生成器和鑒別器中的自注意力機制替換為我們的外部注意力方法，從而得到我們的EAGAN模型。所有實驗都基於流行的PyTorch-StudioGAN repo[96]。超參數使用SAGAN的默認配置。我們使用Frechet Inception Distance (FID)[97]和Inception Score (IS)[98]作為我們的評估指標。部分生成的圖像如圖6所示，定量結果在表8和表9中給出:外部注意力為SAGAN和其他的GANs提供了更好的結果。
4.7 Point cloud classification
ModelNet40[99]是一種流行的3D形狀分類基准，包含40個類別的12,311個CAD模型。它有9843個訓練樣本和2468個測試樣本。我們的EAT模型取代了PCT[42]中的所有自注意力模塊。我們在每個形狀上取樣1024個點，並使用如PCT[42]中的隨機平移、anisotropic縮放和dropout 增強輸入。表11表明，我們的方法優於其他所有方法，包括其他基於注意力的方法，如PCT，我們提出的方法為2D和3D視覺提供了一個出色的backbone。
4.8 Point cloud segmentation
我們在ShapeNet部分數據集上進行了點雲分割實驗[100]。訓練集中有14006個3D模型，評估集中有2874個。每個形狀被分割成parts，共16個對象類別和50個part標簽。我們遵循PCT[42]的實驗設置。EAT在這個數據集上取得了最好的結果，如表10所示。
4.9 Computational requirements
關於輸入大小的線性復雜度帶來了效率上的顯著優勢。我們在輸入大小為1 × 512 × 128 × 128的情況下，將外部注意力(EA)模塊與標准自注意力(SA)[16]及其幾個變體在參數數量和推理操作方面進行了比較，結果如表12所示。外部注意力只需要自注意力所需參數的一半，而且速度要快32倍。與最好的變體相比，外部注意力的速度仍然是前者的兩倍。
5 CONCLUSIONS
本文提出了一種新型的、輕量級的、有效的、適用於各種視覺任務的注意機制——外部注意力。外部注意力所采用的兩個外部memory單元可以看作是整個數據集的字典，能夠學習更多的輸入的代表性特征，同時降低計算代價。我們希望外部注意力能夠激發它在其他領域的實際應用和研究，比如自然語言處理。