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

最近,Transformer在各種視覺任務(wù)中表現(xiàn)出了良好的性能。為了降低全局自注意力引起的二次計(jì)算復(fù)雜度,各種方法都限制局部區(qū)域內(nèi)的注意力范圍,以提高全局自注意力的計(jì)算效率。因此,它們?cè)趩蝹€(gè)注意力層中的感受野不夠大,導(dǎo)致上下文建模不足。
為了解決這個(gè)問題,本文提出了一種Pale-Shaped self-Attention(PS-Attention),它在一個(gè)Pale-Shaped的區(qū)域內(nèi)進(jìn)行自注意力的計(jì)算。與全局自注意力相比,PS-Attention能顯著減少計(jì)算量和內(nèi)存開銷。同時(shí),利用已有的局部自注意力機(jī)制,可以在相似的計(jì)算復(fù)雜度下獲取更豐富的上下文信息。
在PS-Attention的基礎(chǔ)上,作者開發(fā)了一個(gè)具有層次結(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。對(duì)于下游的任務(wù),Pale Transformer Backbone在ADE20K語義分割和COCO目標(biāo)檢測(cè)和實(shí)例分割上比最近的最先進(jìn)的CSWin Transformer表現(xiàn)得更好。
1簡(jiǎn)介
受Transformer在自然語言處理(NLP)的廣泛任務(wù)上取得成功的啟發(fā),Vision Transformer(ViT)首次采用純Transformer架構(gòu)進(jìn)行圖像分類,這顯示了Transformer架構(gòu)在視覺任務(wù)方面的良好性能。

然而,全局自注意力的二次復(fù)雜度導(dǎo)致了昂貴的計(jì)算成本和內(nèi)存占用,特別是在高分辨率場(chǎng)景下,使其無法用于下游各種視覺任務(wù)。
一種典型的提高效率的方法是用局部自注意力取代全局自注意力。如何增強(qiáng)局部環(huán)境下的建模能力是一個(gè)關(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以實(shí)現(xiàn)Cross Window交換信息。Axial Self-Attention將局部注意力區(qū)域視為特征圖的單行或單列(圖1(c))。

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

在這項(xiàng)工作中提出了一種Pale-Shaped self-Attention(PS-Attention)來有效地捕獲更豐富的上下文依賴。具體來說:
首先將輸入特征圖在空間上分割成多個(gè)Pale-Shaped的區(qū)域。每個(gè)Pale-Shaped區(qū)域(縮寫為Pale)由特征圖中相同數(shù)量的交錯(cuò)行和列組成。相鄰行或列之間的間隔對(duì)于所有的Pale是相等的。例如,圖1(e)中的粉色陰影表示其中一種淡色。
然后,在每個(gè)Pale區(qū)域內(nèi)進(jìn)行Self-Attention計(jì)算。對(duì)于任何Token,它都可以直接與同一Pale中的其他Token交互,這使得本文的方法能夠在單個(gè)PS-Attention層中捕獲更豐富的上下文信息。
為了進(jìn)一步提高效率,作者開發(fā)了一個(gè)更高效的PS-Attention并行實(shí)現(xiàn)。得益于更大的感受野和更強(qiáng)的上下文建模能力,PS-Attention顯示出了圖1中所示的現(xiàn)有局部自注意力機(jī)制的優(yōu)越性。
本文基于提出的PS-Attention設(shè)計(jì)了一個(gè)具有層次結(jié)構(gòu)的通用Vision T燃燒former Backbone,命名為Pale Transformer。作者將本文的方法擴(kuò)展到一系列模型,包括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)檢測(cè))準(zhǔn)確率為47.4,mask mAP(實(shí)例分割)準(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工作主要集中在以下兩個(gè)方面:
增強(qiáng)Vision Transformer的局部性;
在性能和效率之間尋求更好的平衡。
2.1 局部增強(qiáng)ViTs
與CNN不同的是,原始Transformer沒有涉及到局部連接的感應(yīng)偏置,這可能會(huì)導(dǎo)致對(duì)線、邊緣、顏色結(jié)合點(diǎn)等局部結(jié)構(gòu)的提取不足。
許多工作致力于增強(qiáng)ViT的局部特征提取。最早的方法是用分層結(jié)構(gòu)取代單一尺度結(jié)構(gòu),以獲得多尺度特征。這樣的設(shè)計(jì)在以后的作品中也有很多。
另一種方法是把CNN和Transformer結(jié)合起來。Mobile-Former、Conformer和DSNet通過精心設(shè)計(jì)的雙分支結(jié)構(gòu)集成了CNN和Transformer功能。而Local ViT、CvT和Shuffle Transformer只在Transformer的某些組件中插入了幾個(gè)卷積。此外,一些作品通過融合不同尺度的多分支或配合局部自注意力以獲得更豐富的特征。
2.2 高效ViTs
目前主流的研究方向?yàn)樘岣咝蔞ision Transformer的主干有兩層:
通過剪枝策略減少冗余計(jì)算 設(shè)計(jì)更高效的自注意力機(jī)制
1、ViT的剪枝
對(duì)于剪枝,現(xiàn)有的方法可以分為三類:
Token剪枝:DVT提出了一種級(jí)聯(lián)Transformer結(jié)構(gòu),可以根據(jù)輸入圖像自適應(yīng)調(diào)整Token數(shù)量??紤]到Token含有不相關(guān)甚至混淆的信息可能不利于圖像分類,一些研究提出了通過可學(xué)習(xí)采樣和強(qiáng)化學(xué)習(xí)策略來定位有區(qū)別的區(qū)域,逐步減少信息Token的丟棄。然而,這種非結(jié)構(gòu)化的稀疏性導(dǎo)致了與密集預(yù)測(cè)任務(wù)的不兼容性。通過Token Pooling和Slow-Fast更新實(shí)現(xiàn)了一些結(jié)構(gòu)保持的Token選擇策略。
Channel剪枝:VTP為去除Reductant Channel提供了一個(gè)簡(jiǎn)單而有效的框架。
Attention Sharing:基于連續(xù)塊間注意力映射高度相關(guān)的觀測(cè)結(jié)果,提出了PSViT算法來重用相鄰層間的注意力計(jì)算過程。
2、高效Self-Attention機(jī)制
考慮到二次計(jì)算的復(fù)雜度是由Self-Attention引起的,許多方法致力于在避免性能衰減的同時(shí)提高二次計(jì)算的效率。
一種方法是減少key和value的序列長(zhǎng)度:
PVT提出了一種Spatial Reduction Attention,在計(jì)算注意力之前對(duì)key和value的尺度進(jìn)行下樣。
Deformable Attention使用一個(gè)線性層從全集合中選擇幾個(gè)關(guān)鍵字,可以看作是全局自注意力的稀疏版本。但是,過多的下采樣會(huì)導(dǎo)致信息混淆,Deformable Attention嚴(yán)重依賴于CNN學(xué)習(xí)的High-Level特征圖,可能不能直接用于原始輸入圖像。
另一種方法是將全局自注意力替換為局部自注意力,將每個(gè)自注意層的范圍限制在一個(gè)局部區(qū)域內(nèi)。如圖1(b)所示,首先將特征映射劃分為幾個(gè)不重疊的方形規(guī)則窗口(用不同的顏色表示),在每個(gè)窗口內(nèi)分別進(jìn)行自注意力計(jì)算。
設(shè)計(jì)局部自注意力機(jī)制的關(guān)鍵挑戰(zhàn)是如何彌合局部和全局感受野之間的差距?
一種典型的方式是在規(guī)則的方形窗口之間構(gòu)建連接。例如,在連續(xù)塊中交替使用常規(guī)窗口和另一種新設(shè)計(jì)的窗口分區(qū)方式(圖1(b)中Shift Window或Shuffle Window),并操作Messenger Tokens來跨窗口交換信息。
此外,Axial Attention通過對(duì)特征圖的每一行或每一列進(jìn)行自注意力計(jì)算,分別在水平和垂直方向上獲得較長(zhǎng)的距離依賴。
CSWin提出了一個(gè)包含多行多列的Cross-Shaped Self-Attention區(qū)域?,F(xiàn)有的這些局部注意力機(jī)制雖然在一定程度上提供了突破局部感受野的機(jī)會(huì),但它們的依賴關(guān)系不夠豐富,不足以在單個(gè)自注意力層中捕獲足夠的上下文信息,從而限制了整個(gè)網(wǎng)絡(luò)的建模能力。
與本文的工作最相關(guān)的是CSWin,CSWin開發(fā)了一個(gè)Cross-Shaped Self-Attention用于計(jì)算水平和垂直條紋中的自注意力,而本文提出的PS-Attention計(jì)算Pale形狀區(qū)域的自注意力。此外,本文的方法中每個(gè)Token的感受野比CSWin寬得多,這也賦予了本文的方法更強(qiáng)的上下文建模能力。
3本文方法
在本節(jié)中,首先介紹Pale-Shaped Self-Attention(PS-Attention)及其高效的并行實(shí)現(xiàn)。然后,給出了Pale Transformer block的組成。最后,將描述Pale Transformer Backbone的總體架構(gòu)和變體配置。
3.1 Pale-Shaped Attention
為了捕獲從短期到長(zhǎng)期的依賴關(guān)系,提出了Pale-Shaped Attention(PS-Attention),它在一個(gè)Pale-Shaped區(qū)域(簡(jiǎn)稱pale)中計(jì)算自注意力。如圖1(e)的粉色陰影所示,一個(gè)pale包含個(gè)交錯(cuò)的行和個(gè)交錯(cuò)的列,它覆蓋了包含個(gè)Token的區(qū)域。

定義為pale size。給定一個(gè)輸入特征圖,首先將其分割成多個(gè)相同大小的pale ,其中。Pale的個(gè)數(shù)等于,可以通過填充或插補(bǔ)操作來保證。對(duì)于所有pale,相鄰行或列之間的間隔是相同的。然后在每個(gè)pale中分別進(jìn)行自注意力計(jì)算。如圖1所示,PS-Attention的感受野比之前所有的局部自注意力機(jī)制都要廣泛和豐富得多,能夠?qū)崿F(xiàn)更強(qiáng)大的上下文建模能力。
3.2 高效的并行實(shí)現(xiàn)
為了進(jìn)一步提高效率,將上面提到的普通PS-Attention分解為行注意和列注意,它們分別在行Token組和列Token組內(nèi)執(zhí)行自注意力。

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

其中,包含個(gè)的交錯(cuò)行,包含個(gè)的交錯(cuò)列。
然后,分別在每個(gè)行Token組和列Token組中執(zhí)行自注意力。使用3個(gè)可分離的卷積層、、生成Query、Key和Value。

式中,、MSA表示多頭自注意力。
最后,行方向和列方向的注意力輸出沿著通道尺寸連接,得到最終的輸出:

與PS-Attention在整個(gè)系統(tǒng)內(nèi)的普通實(shí)現(xiàn)相比,這種并行機(jī)制具有較低的計(jì)算復(fù)雜度。此外,填充操作只需要確保能被整除,能被整除,而不是。因此,也有利于避免填充過多。
3.3 復(fù)雜度分析
給定尺寸為h×w×c的輸入特征,Pale-size為,標(biāo)準(zhǔn)全局自注意力的計(jì)算復(fù)雜度為:

然而,本文提出的PS-Attention并行實(shí)現(xiàn)下的計(jì)算復(fù)雜度為:

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

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

其中,LN(·)為層歸一化。CPE被實(shí)現(xiàn)為一個(gè)簡(jiǎn)單的深度卷積,它在以前的工作中廣泛使用,因?yàn)樗嫒萑我獯笮〉妮斎?。Eq.(7)中定義的PS-Attention模塊是通過將Eq.(1)對(duì)Eq.(3)依次執(zhí)行來構(gòu)造的。Eq.(8)中定義的MLP模塊由2個(gè)線性投影層組成,依次展開和收縮嵌入維數(shù)。
3.5 整體框架

如圖2(a)所示,Pale Transformer通過遵循CNN和Transformer的設(shè)計(jì),由四個(gè)層次級(jí)階段組成,用于捕獲多尺度特征。每個(gè)階段包含一個(gè)Patch合并層和多個(gè)Pale Transformer Block。patch merge層的目標(biāo)是對(duì)輸入特征進(jìn)行一定比例的空間下采樣,并將通道維數(shù)擴(kuò)展2倍以獲得更好的表示能力。
為了便于比較,使用重疊卷積來進(jìn)行Patch合并。其中,第一階段空間下采樣比為4,后三個(gè)階段空間下采樣比為2,分別采用7×7與stride=4的卷積和3×3與stride=2的卷積實(shí)現(xiàn)。Patch合并層的輸出被輸入到后續(xù)的Pale Transformer Block中,Token的數(shù)量保持不變。只需在最后一個(gè)塊的頂部應(yīng)用平均池化操作,以獲得最終分類頭的代表Token,該分類頭由單個(gè)線性投影層組成。
具體變體如下:


4實(shí)驗(yàn)
4.1 消融實(shí)驗(yàn)
1、Pale-size的影響

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

這里比較了3種PS-Attention的實(shí)現(xiàn)。普通的PS-Attention直接在整個(gè)Pale區(qū)域內(nèi)進(jìn)行自注意力計(jì)算,它可以近似為兩種更有效的實(shí)現(xiàn),順序和并行。
順序算法在連續(xù)塊中交替計(jì)算行方向和列方向上的自注意力,而并行算法在每個(gè)塊內(nèi)并行執(zhí)行行方向和列方向的注意力。如表8所示,并行PS-Attention在所有任務(wù)上都取得了最好的結(jié)果。作者將此歸因于在普通的PS-Attention中,對(duì)于非方形輸入大小的過多填充將導(dǎo)致輕微的性能下降。
為了直接比較PS-Attention與最相關(guān)的Axial-based Attention,分別用軸向自注意力和cross-shaped window self-attention替代了Pale-T的PS-Attention。如表8所示,PS-Attention明顯優(yōu)于這兩種機(jī)制。
3、不同的位置嵌入

位置編碼可以引入空間位置感知來實(shí)現(xiàn)自注意力的特征聚合,在Transformer中起著重要作用。本文比較了幾種常用的位置編碼方法,如無位置編碼(no pos.)、絕對(duì)位置編碼(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驗(yàn)證集上的性能。在相同的計(jì)算復(fù)雜度下,與先進(jìn)的CNN相比,Pale模型分別比RegNet模型好+3.4%、+2.6%和+2.0%。
與此同時(shí),Pale Transformer的性能優(yōu)于最先進(jìn)的基于Transformer的Backbone,在類似型號(hào)尺寸和FLOPs下的所有變種中,其性能比最相關(guān)的CSWin Transformer高0.7%。
請(qǐng)注意,LV-ViT和VOLO,使用額外的MixToken增強(qiáng)和token丟失進(jìn)行訓(xùn)練,似乎與本文方法相同。為了進(jìn)行比較,在Pale模型上使用了這兩種技巧,與lv-vi-s相比,Pale-T獲得了+0.9%的增益,且計(jì)算成本更低。Pale-S和Pale-B分別達(dá)到85.0%和85.8%,比VOLO高出0.8%和0.6%。
4.1 COCO目標(biāo)檢測(cè)

如表3所示,對(duì)于目標(biāo)檢測(cè)Pale-T、Pale-S和Pale-B的目標(biāo)檢測(cè)box mAP分別為47.4、48.4和49.2,超過了之前的最佳水平CSWin Transformer +0.7,+0.5,+0.6。
此外,Pale Transformer變體在實(shí)例分割上也有一致的改進(jìn),分別比之前的最佳Backbone提高了+0.5、+0.5、+0.3 mask mAP。
5參考
[1].Pale Transformer:A General Vision Transformer Backbone with Pale-Shaped Attention
輔助模塊加速收斂,精度大幅提升!移動(dòng)端實(shí)時(shí)的NanoDet-Plus來了!
SSD的torchvision版本實(shí)現(xiàn)詳解
? ??? ? ? ? ? ? ? ? ? ? ????????? ??一個(gè)用心的公眾號(hào)


