全新Backbone | Pale Transformer完美超越Swin Transformer

最近，Transformer在各種視覺任務(wù)中表現(xiàn)出了良好的性能。為了降低全局自注意力引起的二次計算復(fù)雜度，各種方法都限制局部區(qū)域內(nèi)的注意力范圍，以提高全局自注意力的計算效率。因此，它們在單個注意力層中的感受野不夠大，導(dǎo)致上下文建模不足。

為了解決這個問題，本文提出了一種Pale-Shaped self-Attention(PS-Attention)，它在一個Pale-Shaped的區(qū)域內(nèi)進(jìn)行自注意力的計算。與全局自注意力相比，PS-Attention能顯著減少計算量和內(nèi)存開銷。同時，利用已有的局部自注意力機制，可以在相似的計算復(fù)雜度下獲取更豐富的上下文信息。

在PS-Attention的基礎(chǔ)上，作者開發(fā)了一個具有層次結(jié)構(gòu)的通用Vision Transformer Backbone，名為Pale Transformer，在224×224 ImageNet-1K分類中，模型尺寸分別為22M、48M和85M, Top-1準(zhǔn)確率達(dá)到83.4%、84.3%和84.9%，優(yōu)于之前的Vision Transformer Backbone。對于下游的任務(wù)，Pale Transformer Backbone在ADE20K語義分割和COCO目標(biāo)檢測和實例分割上比最近的最先進(jìn)的CSWin Transformer表現(xiàn)得更好。

1簡介

受Transformer在自然語言處理(NLP)的廣泛任務(wù)上取得成功的啟發(fā)，Vision Transformer(ViT)首次采用純Transformer架構(gòu)進(jìn)行圖像分類，這顯示了Transformer架構(gòu)在視覺任務(wù)方面的良好性能。

然而，全局自注意力的二次復(fù)雜度導(dǎo)致了昂貴的計算成本和內(nèi)存占用，特別是在高分辨率場景下，使其無法用于下游各種視覺任務(wù)。

一種典型的提高效率的方法是用局部自注意力取代全局自注意力。如何增強局部環(huán)境下的建模能力是一個關(guān)鍵而又具有挑戰(zhàn)性的問題。

例如，Swin Transformer和Shuffle Transformer分別提出了Shift Window和Shuffle Window(圖1(b))，并交替使用兩種不同的Window分區(qū)(即Regular Window和Proposed Window)在連續(xù)塊中構(gòu)建Cross Window連接。

MSG Transformer操縱Messenger Token以實現(xiàn)Cross Window交換信息。Axial Self-Attention將局部注意力區(qū)域視為特征圖的單行或單列(圖1(c))。

CSWin提出了Cross-Shaped Window Attention(圖1(d))，可以看作是Axial Self-Attention的多行、多列展開。雖然這些方法的性能很好，甚至優(yōu)于CNN的同類方法，但每個自注意力層的依賴關(guān)系不夠豐富，不足以捕捉足夠的上下文信息。

在這項工作中提出了一種Pale-Shaped self-Attention(PS-Attention)來有效地捕獲更豐富的上下文依賴。具體來說：

• 首先將輸入特征圖在空間上分割成多個Pale-Shaped的區(qū)域。每個Pale-Shaped區(qū)域(縮寫為Pale)由特征圖中相同數(shù)量的交錯行和列組成。相鄰行或列之間的間隔對于所有的Pale是相等的。例如，圖1(e)中的粉色陰影表示其中一種淡色。

• 然后，在每個Pale區(qū)域內(nèi)進(jìn)行Self-Attention計算。對于任何Token，它都可以直接與同一Pale中的其他Token交互，這使得本文的方法能夠在單個PS-Attention層中捕獲更豐富的上下文信息。

為了進(jìn)一步提高效率，作者開發(fā)了一個更高效的PS-Attention并行實現(xiàn)。得益于更大的感受野和更強的上下文建模能力，PS-Attention顯示出了圖1中所示的現(xiàn)有局部自注意力機制的優(yōu)越性。

本文基于提出的PS-Attention設(shè)計了一個具有層次結(jié)構(gòu)的通用Vision T燃燒former Backbone，命名為Pale Transformer。作者將本文的方法擴展到一系列模型，包括Pale-T(22M)、Pale-S(48M)和Pale-B(85M)，達(dá)到了比以前方法更好的性能。

Pale-T在ImageNet-1k上的Top-1分類準(zhǔn)確率為83.4%，在ADE20K上的單尺度mIoU(語義分割)準(zhǔn)確率為50.4%，在COCO上的box mAP(目標(biāo)檢測)準(zhǔn)確率為47.4，mask mAP(實例分割)準(zhǔn)確率為42.7，分別比目前最先進(jìn)的Backbone算法提高了+0.7%、+1.1%、+0.7和+0.5。

此外，最大的版本Pale-B在ImageNet-1K上的Top-1準(zhǔn)確率為84.9%，在ADE20K上的單尺度mIoU準(zhǔn)確率為52.2%，在COCO上的box mAP為49.3，mask mAP為44.2。

2相關(guān)工作

ViT將輸入圖像作為一系列的Patches，為多年來由CNN主導(dǎo)的視覺任務(wù)開辟了一條新的道路，并顯示出了良好的性能。為了更好地適應(yīng)視覺任務(wù)，之前的Vision Transformer Backbone工作主要集中在以下兩個方面:

1. 增強Vision Transformer的局部性；

2. 在性能和效率之間尋求更好的平衡。

2.1 局部增強ViTs

與CNN不同的是，原始Transformer沒有涉及到局部連接的感應(yīng)偏置，這可能會導(dǎo)致對線、邊緣、顏色結(jié)合點等局部結(jié)構(gòu)的提取不足。

許多工作致力于增強ViT的局部特征提取。最早的方法是用分層結(jié)構(gòu)取代單一尺度結(jié)構(gòu)，以獲得多尺度特征。這樣的設(shè)計在以后的作品中也有很多。

另一種方法是把CNN和Transformer結(jié)合起來。Mobile-Former、Conformer和DSNet通過精心設(shè)計的雙分支結(jié)構(gòu)集成了CNN和Transformer功能。而Local ViT、CvT和Shuffle Transformer只在Transformer的某些組件中插入了幾個卷積。此外，一些作品通過融合不同尺度的多分支或配合局部自注意力以獲得更豐富的特征。

2.2 高效ViTs

目前主流的研究方向為提高效率Vision Transformer的主干有兩層:

• 通過剪枝策略減少冗余計算
• 設(shè)計更高效的自注意力機制

1、ViT的剪枝

對于剪枝，現(xiàn)有的方法可以分為三類:

1. Token剪枝：DVT提出了一種級聯(lián)Transformer結(jié)構(gòu)，可以根據(jù)輸入圖像自適應(yīng)調(diào)整Token數(shù)量。考慮到Token含有不相關(guān)甚至混淆的信息可能不利于圖像分類，一些研究提出了通過可學(xué)習(xí)采樣和強化學(xué)習(xí)策略來定位有區(qū)別的區(qū)域，逐步減少信息Token的丟棄。然而，這種非結(jié)構(gòu)化的稀疏性導(dǎo)致了與密集預(yù)測任務(wù)的不兼容性。通過Token Pooling和Slow-Fast更新實現(xiàn)了一些結(jié)構(gòu)保持的Token選擇策略。

2. Channel剪枝：VTP為去除Reductant Channel提供了一個簡單而有效的框架。

3. Attention Sharing：基于連續(xù)塊間注意力映射高度相關(guān)的觀測結(jié)果，提出了PSViT算法來重用相鄰層間的注意力計算過程。

2、高效Self-Attention機制

考慮到二次計算的復(fù)雜度是由Self-Attention引起的，許多方法致力于在避免性能衰減的同時提高二次計算的效率。

一種方法是減少key和value的序列長度：

• PVT提出了一種Spatial Reduction Attention，在計算注意力之前對key和value的尺度進(jìn)行下樣。

• Deformable Attention使用一個線性層從全集合中選擇幾個關(guān)鍵字，可以看作是全局自注意力的稀疏版本。但是，過多的下采樣會導(dǎo)致信息混淆，Deformable Attention嚴(yán)重依賴于CNN學(xué)習(xí)的High-Level特征圖，可能不能直接用于原始輸入圖像。

另一種方法是將全局自注意力替換為局部自注意力，將每個自注意層的范圍限制在一個局部區(qū)域內(nèi)。如圖1(b)所示，首先將特征映射劃分為幾個不重疊的方形規(guī)則窗口(用不同的顏色表示)，在每個窗口內(nèi)分別進(jìn)行自注意力計算。

設(shè)計局部自注意力機制的關(guān)鍵挑戰(zhàn)是如何彌合局部和全局感受野之間的差距？

一種典型的方式是在規(guī)則的方形窗口之間構(gòu)建連接。例如，在連續(xù)塊中交替使用常規(guī)窗口和另一種新設(shè)計的窗口分區(qū)方式(圖1(b)中Shift Window或Shuffle Window)，并操作Messenger Tokens來跨窗口交換信息。

此外，Axial Attention通過對特征圖的每一行或每一列進(jìn)行自注意力計算，分別在水平和垂直方向上獲得較長的距離依賴。

CSWin提出了一個包含多行多列的Cross-Shaped Self-Attention區(qū)域?，F(xiàn)有的這些局部注意力機制雖然在一定程度上提供了突破局部感受野的機會，但它們的依賴關(guān)系不夠豐富，不足以在單個自注意力層中捕獲足夠的上下文信息，從而限制了整個網(wǎng)絡(luò)的建模能力。

與本文的工作最相關(guān)的是CSWin，CSWin開發(fā)了一個Cross-Shaped Self-Attention用于計算水平和垂直條紋中的自注意力，而本文提出的PS-Attention計算Pale形狀區(qū)域的自注意力。此外，本文的方法中每個Token的感受野比CSWin寬得多，這也賦予了本文的方法更強的上下文建模能力。

3本文方法

在本節(jié)中，首先介紹Pale-Shaped Self-Attention(PS-Attention)及其高效的并行實現(xiàn)。然后，給出了Pale Transformer block的組成。最后，將描述Pale Transformer Backbone的總體架構(gòu)和變體配置。

3.1 Pale-Shaped Attention

為了捕獲從短期到長期的依賴關(guān)系，提出了Pale-Shaped Attention(PS-Attention)，它在一個Pale-Shaped區(qū)域(簡稱pale)中計算自注意力。如圖1(e)的粉色陰影所示，一個pale包含個交錯的行和個交錯的列，它覆蓋了包含個Token的區(qū)域。

定義為pale size。給定一個輸入特征圖，首先將其分割成多個相同大小的pale ，其中。Pale的個數(shù)等于，可以通過填充或插補操作來保證。對于所有pale，相鄰行或列之間的間隔是相同的。然后在每個pale中分別進(jìn)行自注意力計算。如圖1所示，PS-Attention的感受野比之前所有的局部自注意力機制都要廣泛和豐富得多，能夠?qū)崿F(xiàn)更強大的上下文建模能力。

3.2 高效的并行實現(xiàn)

為了進(jìn)一步提高效率，將上面提到的普通PS-Attention分解為行注意和列注意，它們分別在行Token組和列Token組內(nèi)執(zhí)行自注意力。

具體而言，如圖2(c)，首先將輸入特征劃分為兩個獨立的部分和，然后將其分為多個組，以便分別按行和列進(jìn)行注意力計算。

其中，包含個的交錯行，包含個的交錯列。

然后，分別在每個行Token組和列Token組中執(zhí)行自注意力。使用3個可分離的卷積層、、生成Query、Key和Value。

式中，、MSA表示多頭自注意力。

最后，行方向和列方向的注意力輸出沿著通道尺寸連接，得到最終的輸出:

與PS-Attention在整個系統(tǒng)內(nèi)的普通實現(xiàn)相比，這種并行機制具有較低的計算復(fù)雜度。此外，填充操作只需要確保能被整除，能被整除，而不是。因此，也有利于避免填充過多。

3.3 復(fù)雜度分析

給定尺寸為h×w×c的輸入特征，Pale-size為，標(biāo)準(zhǔn)全局自注意力的計算復(fù)雜度為：

然而，本文提出的PS-Attention并行實現(xiàn)下的計算復(fù)雜度為：

由于始終保持不變，因此與全局算法相比，可以明顯減輕計算量和內(nèi)存負(fù)擔(dān)。

3.4 Pale Transformer Block

如圖2(b)所示，Pale Transformer Block由3個順序部分組成，用于動態(tài)生成位置嵌入的條件位置編碼(CPE)，用于捕獲上下文信息的PS-Attention模塊，以及用于特征投影的MLP模塊。第l塊的傳播可以表示為:

其中，LN(·)為層歸一化。CPE被實現(xiàn)為一個簡單的深度卷積，它在以前的工作中廣泛使用，因為它兼容任意大小的輸入。Eq.(7)中定義的PS-Attention模塊是通過將Eq.(1)對Eq.(3)依次執(zhí)行來構(gòu)造的。Eq.(8)中定義的MLP模塊由2個線性投影層組成，依次展開和收縮嵌入維數(shù)。

3.5 整體框架

如圖2(a)所示，Pale Transformer通過遵循CNN和Transformer的設(shè)計，由四個層次級階段組成，用于捕獲多尺度特征。每個階段包含一個Patch合并層和多個Pale Transformer Block。patch merge層的目標(biāo)是對輸入特征進(jìn)行一定比例的空間下采樣，并將通道維數(shù)擴展2倍以獲得更好的表示能力。

為了便于比較，使用重疊卷積來進(jìn)行Patch合并。其中，第一階段空間下采樣比為4，后三個階段空間下采樣比為2，分別采用7×7與stride=4的卷積和3×3與stride=2的卷積實現(xiàn)。Patch合并層的輸出被輸入到后續(xù)的Pale Transformer Block中，Token的數(shù)量保持不變。只需在最后一個塊的頂部應(yīng)用平均池化操作，以獲得最終分類頭的代表Token，該分類頭由單個線性投影層組成。

具體變體如下：

4實驗

4.1 消融實驗

1、Pale-size的影響

4個階段的Pale-size控制了上下文信息豐富度與計算成本之間的權(quán)衡。如表7所示，增加Pale-size(從1到7)可以持續(xù)提高所有任務(wù)的性能，而進(jìn)一步增加到9不會帶來明顯的、持續(xù)的改進(jìn)，但會帶來更多的FLOPs。因此，默認(rèn)情況下，對所有任務(wù)使用。

2、不同注意力對比

這里比較了3種PS-Attention的實現(xiàn)。普通的PS-Attention直接在整個Pale區(qū)域內(nèi)進(jìn)行自注意力計算，它可以近似為兩種更有效的實現(xiàn)，順序和并行。

順序算法在連續(xù)塊中交替計算行方向和列方向上的自注意力，而并行算法在每個塊內(nèi)并行執(zhí)行行方向和列方向的注意力。如表8所示，并行PS-Attention在所有任務(wù)上都取得了最好的結(jié)果。作者將此歸因于在普通的PS-Attention中，對于非方形輸入大小的過多填充將導(dǎo)致輕微的性能下降。

為了直接比較PS-Attention與最相關(guān)的Axial-based Attention，分別用軸向自注意力和cross-shaped window self-attention替代了Pale-T的PS-Attention。如表8所示，PS-Attention明顯優(yōu)于這兩種機制。

3、不同的位置嵌入

位置編碼可以引入空間位置感知來實現(xiàn)自注意力的特征聚合，在Transformer中起著重要作用。本文比較了幾種常用的位置編碼方法，如無位置編碼(no pos.)、絕對位置編碼(APE)和條件位置編碼(CPE)。如表9所示，CPE的性能最好。不使用任何位置編碼將導(dǎo)致嚴(yán)重的性能下降，這證明了位置編碼在Vision Transformer模型中的有效性。

4.2 ImageNet-1K分類

表2比較了Pale Transformer與最先進(jìn)的CNN和Vision Transformer Backbone在ImageNet-1K驗證集上的性能。在相同的計算復(fù)雜度下，與先進(jìn)的CNN相比，Pale模型分別比RegNet模型好+3.4%、+2.6%和+2.0%。

與此同時，Pale Transformer的性能優(yōu)于最先進(jìn)的基于Transformer的Backbone，在類似型號尺寸和FLOPs下的所有變種中，其性能比最相關(guān)的CSWin Transformer高0.7%。

請注意，LV-ViT和VOLO，使用額外的MixToken增強和token丟失進(jìn)行訓(xùn)練，似乎與本文方法相同。為了進(jìn)行比較，在Pale模型上使用了這兩種技巧，與lv-vi-s相比，Pale-T獲得了+0.9%的增益，且計算成本更低。Pale-S和Pale-B分別達(dá)到85.0%和85.8%，比VOLO高出0.8%和0.6%。

4.1 COCO目標(biāo)檢測

如表3所示，對于目標(biāo)檢測Pale-T、Pale-S和Pale-B的目標(biāo)檢測box mAP分別為47.4、48.4和49.2，超過了之前的最佳水平CSWin Transformer +0.7，+0.5，+0.6。

此外，Pale Transformer變體在實例分割上也有一致的改進(jìn)，分別比之前的最佳Backbone提高了+0.5、+0.5、+0.3 mask mAP。

5參考

[1].Pale Transformer：A General Vision Transformer Backbone with Pale-Shaped Attention

6推薦閱讀

激活函數(shù) | Squareplus性能比肩Softplus激活函數(shù)速度快6倍（附Pytorch實現(xiàn)）

YOLO-Z | 記錄修改YOLOv5以適應(yīng)小目標(biāo)檢測的實驗過程

LVT | ViT輕量化的曙光，完美超越MobileNet和ResNet系列

掃描下方二維碼關(guān)注【集智書童】公眾號，獲取更多實踐項目源碼和論文解讀，非常期待你我的相遇，讓我們以夢為馬，砥礪前行！