【原】ECCV2022兼具線性復(fù)雜度和全局感受野的優(yōu)點(diǎn)！谷歌&UT-Austin提出MaxViT，在多個(gè)任務(wù)上實(shí)現(xiàn)SOTA性能！

我愛計(jì)算機(jī)視覺 2022-08-18 發(fā)布于江蘇

展開全文

詳細(xì)信息如下：

論文地址：https:///abs/2204.01697
代碼地址：未開源

摘要

Transformer最近在計(jì)算機(jī)視覺界引起了極大的關(guān)注。然而，自注意力機(jī)制在圖像大小方面缺乏可擴(kuò)展性，限制了其在最先進(jìn)的視覺主干中的廣泛采用。在本文中，作者提出了一種高效且可擴(kuò)展的注意力模型，稱之為多軸注意力，它包括兩個(gè)方面：塊局部（blocked local）注意力和擴(kuò)展全局（dilated global）注意力。這些設(shè)計(jì)選擇允許在任意輸入分辨率上進(jìn)行全局-局部空間交互，僅具有線性復(fù)雜度。

作者還提出了一個(gè)新的架構(gòu)元素，有效地將本文提出的注意力模型與卷積相結(jié)合，并通過簡單地在多個(gè)階段重復(fù)基本構(gòu)建塊提出了一個(gè)簡單的分層視覺主干，稱為MaxViT。值得注意的是，MaxViT能夠在整個(gè)網(wǎng)絡(luò)中“看到”全局，甚至在早期的高分辨率階段。作者證明了本文的模型在廣泛的視覺任務(wù)中的有效性。

在圖像分類方面，MaxViT在各種設(shè)置下都達(dá)到了最先進(jìn)的性能：在沒有額外數(shù)據(jù)的情況下，MaxViT達(dá)到了86.5%的ImageNet-1K top-1精度；通過ImageNet-21K預(yù)訓(xùn)練，本文的模型達(dá)到了88.7%的top-1精度。對于下游任務(wù)，MaxViT作為主干在目標(biāo)檢測和視覺美感評估方面提供了良好的性能。作者還表明，本文提出的模型在ImageNet上表現(xiàn)出強(qiáng)大的生成建模能力，證明了MaxViT塊作為通用視覺模塊的優(yōu)越潛力。

Motivation

自AlexNet以來，卷積神經(jīng)網(wǎng)絡(luò)（ConvNet）一直是計(jì)算機(jī)視覺的主導(dǎo)架構(gòu)設(shè)計(jì)選擇。ConvNet通過深入、拓寬、添加密集連接、高效可分離卷積、空洞卷積、使用編碼器-解碼器框架甚至引入現(xiàn)代微設(shè)計(jì)組件，繼續(xù)在眾多視覺問題上表現(xiàn)出色。

同時(shí)，受自然語言處理中Transformers等自注意力模型演變的啟發(fā)，許多研究人員已經(jīng)開始將注意力機(jī)制引入視覺。視覺Transformer（ViT）可能是第一個(gè)完全基于Transformer的視覺架構(gòu)，其中圖像patch僅被視為文字序列，Transformer編碼器應(yīng)用于這些視覺token。當(dāng)對大規(guī)模數(shù)據(jù)集進(jìn)行預(yù)訓(xùn)練時(shí)，ViT可以在圖像識別方面取得令人信服的結(jié)果。

然而，已經(jīng)觀察到，如果沒有廣泛的預(yù)訓(xùn)練，ViT在圖像識別方面表現(xiàn)不佳。這是因?yàn)門ransformer的模型容量很強(qiáng)，充滿了較少的歸納偏置，這會導(dǎo)致過配合。為了適當(dāng)?shù)卣{(diào)整模型容量并提高其可擴(kuò)展性，許多后續(xù)工作研究了為視覺任務(wù)（如局部注意）定制的稀疏Transformer模型。這些方法通常重新引入層次結(jié)構(gòu)以補(bǔ)償非局部性的損失。

Swin Transformer就是這樣一種成功的嘗試，通過在移位的非重疊窗口上應(yīng)用自注意力來修改Transformer。這種方法在ImageNet基準(zhǔn)上優(yōu)于純視覺Transformer和ConvNet。盡管與ViT中使用的完全注意力相比，基于窗口的注意力具有更大的靈活性和通用性，但由于非局部性的損失，已觀察到基于窗口的注意力的模型容量有限，因此在更大的數(shù)據(jù)區(qū)域（如ImageNet-21K和JFT）上的擴(kuò)展不利。

然而，由于注意力算子需要二次復(fù)雜度，因此在層次網(wǎng)絡(luò)的早期或高分辨率階段通過完全注意力獲取全局交互需要大量計(jì)算。如何有效地結(jié)合全局和局部交互，在計(jì)算預(yù)算下平衡模型容量和可推廣性仍然是一個(gè)挑戰(zhàn)。

在本文中，作者提出了一種新型的Transformer模塊，稱為多軸自注意力（Max SA），它可以作為基本的架構(gòu)組件，在單個(gè)塊中執(zhí)行局部和全局空間交互。

與完全自注意力相比，Max SA具有更大的靈活性和效率，即自然適應(yīng)線性復(fù)雜度的不同輸入長度；與窗口/局部注意力不同，Max SA通過提出一個(gè)全局接受域，允許更強(qiáng)的模型容量。此外，MaxSA僅具有線性復(fù)雜度，即使在早期的高分辨率階段，也可以在網(wǎng)絡(luò)的任何層中用作一般的獨(dú)立注意力模塊。

為了證明其有效性和通用性，作者進(jìn)一步設(shè)計(jì)了一種簡單但有效的視覺主干，稱為多軸視覺Transformer（MaxViT），通過分層堆疊由Max-SA和卷積組成的重復(fù)塊。雖然作者提出的模型屬于混合視覺Transformer的范疇，但MaxViT與以前的方法不同之處在于，我們力求簡單，通過設(shè)計(jì)一個(gè)基本塊來統(tǒng)一卷積、局部和全局注意力，然后簡單地重復(fù)它。

本文的實(shí)驗(yàn)表明，對于廣泛的視覺任務(wù)，MaxViT在所有數(shù)據(jù)模式下顯著提高了最先進(jìn)的（SOTA）性能。具體而言，如上圖所示，MaxViT在精度與浮點(diǎn)運(yùn)算以及精度與參數(shù)曲線方面均優(yōu)于所有最近基于Transformer的模型。

本文的貢獻(xiàn)是：

提出一種通用的Transformer架構(gòu)：MaxViT，可以在網(wǎng)絡(luò)的每個(gè)階段捕獲局部和全局空間交互。
提出一種新穎的獨(dú)立多軸注意力模塊，由塊局部注意力和擴(kuò)展全局注意力組成，以線性復(fù)雜度享受全局感知。
大量實(shí)驗(yàn)表明，MaxViT在各種數(shù)據(jù)模式下實(shí)現(xiàn)了SOTA結(jié)果，用于廣泛的任務(wù)，包括圖像分類、目標(biāo)檢測、圖像美學(xué)評估和圖像生成。

方法

本文引入了一種新型的注意力模塊，稱為多軸自注意力（MaxSA），通過將完全密集的注意力機(jī)制分解為兩種稀疏形式（窗口注意力和網(wǎng)格注意力）將普通注意力的二次復(fù)雜度降低到線性，而不失任何非局部性。

由于Max SA的靈活性和可擴(kuò)展性，可以通過在分層架構(gòu)中簡單地將Max SA的替代層和MBConv堆疊，來構(gòu)建一個(gè)視覺主干，稱之為MaxViT，如上圖所示。MaxViT在整個(gè)網(wǎng)絡(luò)中，從淺到深，從全局和局部接受域中受益，在模型容量和泛化能力方面表現(xiàn)出卓越的性能。

3.1 Attention

自注意力允許整個(gè)空間位置的空間混合，同時(shí)也受益于基于歸一化成對相似性的內(nèi)容相關(guān)權(quán)重。標(biāo)準(zhǔn)自注意力是位置不可感知的，即非平移等變性，這是一種重要的歸納偏見。相對自注意力被提出，通過在注意力權(quán)重中引入相對學(xué)習(xí)偏置來改善普通注意力，已證明在許多視覺任務(wù)中，相對學(xué)習(xí)偏置始終優(yōu)于原始注意力。在這項(xiàng)工作中，作者主要采用預(yù)歸一化相對自注意力作為MaxViT中的關(guān)鍵算子。

3.2 Multi-axis Attention

與局部卷積相比，全局交互是自注意力的關(guān)鍵優(yōu)勢之一。然而，直接沿整個(gè)空間應(yīng)用注意力在計(jì)算上是不可行的，因?yàn)樽⒁饬λ阕有枰螐?fù)雜度。為了解決這個(gè)問題，作者提出了一種多軸方法，通過簡單地分解空間軸，將全尺寸注意力分解為兩種稀疏形式——局部和全局。

設(shè)為輸入特征圖。作者沒有將注意力放在鋪平的空間維度HW上，而是將特征分為形狀張量，表示劃分為大小為P×P的非重疊窗口。將自注意力應(yīng)用于局部空間維度，即P×P，相當(dāng)于在一個(gè)小窗口內(nèi)進(jìn)行注意力。作者將使用這種塊注意力（block attention）來進(jìn)行局部交互。

盡管繞過了繁重的完全自注意力計(jì)算，但在大規(guī)模數(shù)據(jù)集上觀察到局部注意力模型不足。受塊注意力的啟發(fā)，作者提出了一種簡單但有效的方法來獲得稀疏的全局注意力，作者稱之為網(wǎng)格注意力（grid attention）。

作者使用固定的G×G均勻網(wǎng)格將張量網(wǎng)格化為形狀，而不是使用固定的窗口大小對特征圖進(jìn)行分區(qū)，從而使窗口具有自適應(yīng)大小。在分解的網(wǎng)格軸上使用自注意力，即G×G，對應(yīng)于token的擴(kuò)展、全局空間混合。通過使用相同的固定窗口和網(wǎng)格大小，可以充分平衡局部和全局運(yùn)算之間的計(jì)算，兩者在空間大小或序列長度方面僅具有線性復(fù)雜度。

注意，本文提出的Max SA模塊可以用相同數(shù)量的參數(shù)和FLOPs來替代Swin注意力模塊。然而，它具有全局交互能力，不需要mask、填充或循環(huán)移位，使其更易于實(shí)現(xiàn)，優(yōu)于移位窗口方案。

作者連續(xù)交替的采用了這兩種注意力機(jī)制，以獲取局部和全局建模。作者還在多軸注意力之前采用了MBConv，因?yàn)橹暗墓ぷ鞅砻鬟@個(gè)模塊增加模型的泛化性能。其中，MBConv中的深度卷積也可以看做是位置編碼，MaxViT的偽代碼如算法1所示。

本文提出的多軸注意力可以一起使用，也可以單獨(dú)使用，用于不同的目的——塊注意力用于局部交互，網(wǎng)格注意力用于全局混合。這些元素可以很容易地插入到許多視覺體系結(jié)構(gòu)中，尤其是在高分辨率任務(wù)中。

3.3 Architecture Variants

作者設(shè)計(jì)了一系列極其簡單的架構(gòu)變體，以探索本文提出的MaxViT塊的有效性。作者使用了類似于常見ConvNet實(shí)踐的層次主干，其中輸入首先在stem階段（S0）使用Conv3x3層進(jìn)行降采樣。網(wǎng)絡(luò)主體包含四個(gè)階段（S1-S4），每個(gè)階段的分辨率為前一階段的一半，通道數(shù)增加了一倍（隱藏維度）。

在本文的網(wǎng)絡(luò)中，作者在整個(gè)主干中使用相同的MaxViT塊。在每個(gè)階段的第一個(gè)MBConv塊的深度Conv3x3層中應(yīng)用下采樣。默認(rèn)情況下，反向瓶頸和SE模塊的膨脹率和收縮率分別為4和0.25。作者將所有注意力塊的注意力頭大小設(shè)置為32。通過增加每個(gè)階段B的塊數(shù)和通道維度C來擴(kuò)大模型。下表總結(jié)了MaxViT變體的架構(gòu)配置。

實(shí)驗(yàn)

作者在上表中展示了ImageNet-1K分類的性能比較。在基本224×224設(shè)置下，MaxViT在整個(gè)FLOPs中大大優(yōu)于最近的強(qiáng)混合模型CoAtNet，此外在性能上也有明顯的優(yōu)勢。

上表顯示了在ImageNet-21K上預(yù)訓(xùn)練的模型的結(jié)果。值得注意的是，MaxViT-B模型實(shí)現(xiàn)了88.38%的精度，僅使用43%的參數(shù)計(jì)數(shù)和38%的浮點(diǎn)運(yùn)算，就比之前的最佳模型CoAtNet-4高出0.28%。

上圖全面顯示了模型大小比較MaxViT量表明顯優(yōu)于先前具有類似復(fù)雜性的基于注意力的模型。此外，MaxViT XL模型實(shí)現(xiàn)了新的SOTA性能，在分辨率512×512下進(jìn)行微調(diào)時(shí)，精度達(dá)到88.70%。

本文的模型也可擴(kuò)展到大規(guī)模訓(xùn)練數(shù)據(jù)–MaxViT XL在4.75億個(gè)參數(shù)下實(shí)現(xiàn)了89.53%的高精度，在可比模型大小下優(yōu)于以前的模型。

上表展示了在目標(biāo)檢測和分割任務(wù)中使用的MaxViT主干模型在精度和效率方面大大優(yōu)于所有其他主干，包括各種模型大小的Swin、ConvNeXt和UViT。

為了將本文的模型與現(xiàn)有方法進(jìn)行評估和比較，作者在上表中總結(jié)了結(jié)果。對于相似的輸入分辨率，提出的MaxViT-T模型優(yōu)于現(xiàn)有的圖像美學(xué)評估方法。

在生成任務(wù)上，本文的MaxViT實(shí)現(xiàn)了更好的FID，并且參數(shù)數(shù)量明顯減少。這些結(jié)果證明了MaxViT塊對生成任務(wù)的有效性。

本文的主要貢獻(xiàn)之一是網(wǎng)格注意力模塊，它允許在線性時(shí)間內(nèi)進(jìn)行稀疏的全局交互，使模型能夠捕獲所有階段的全局信息。作者進(jìn)行了兩次消融以了解其收益：1）在每個(gè)階段完全消除了全局注意力；2）將網(wǎng)格注意力替換為塊注意力，以保持相同的參數(shù)計(jì)數(shù)和FLOPs。如上表所示，與僅使用局部注意力或卷積相比，在早期階段啟用全局注意力可以進(jìn)一步提高性能。

作者還通過刪除每個(gè)階段中的所有MBConv來消除MaxViT中MBConv層的使用。如上表所示，在MaxViT中使用MBConv層顯著提高了性能。

本文提出了三個(gè)不同的模塊來構(gòu)建MaxViT塊——MBConv、塊和網(wǎng)格注意力——它捕捉從局部到全局的空間交互。為了研究最有效的組合方式，作者使用所有6種排列評估了MaxViT-T模型，結(jié)果如上表所示。

作者還探究了順序和并行堆疊模塊的影響，從上表中可以看出順序方法在參數(shù)和計(jì)算量較少的情況下顯著優(yōu)于并行方法。原因可能是并行設(shè)計(jì)學(xué)習(xí)互補(bǔ)線索，它們之間的交互更少，而本文的順序堆棧能夠?qū)W習(xí)局部和全局層之間更強(qiáng)大的融合。

上圖展示了本文方法和Swin在不同模型大小下的實(shí)驗(yàn)結(jié)果。

總結(jié)

雖然這幾年的最新研究表明，ConvNet和視覺Transformer可以在圖像識別方面實(shí)現(xiàn)類似的性能，但本文的工作提出了一種統(tǒng)一的設(shè)計(jì)，它利用了兩個(gè)模型的優(yōu)點(diǎn)（高效卷積和稀疏注意力）并證明了一個(gè)建立在上面的模型，即MaxViT，可以在各種視覺任務(wù)上實(shí)現(xiàn)最先進(jìn)的性能，更重要的是，可以很好地?cái)U(kuò)展到大規(guī)模數(shù)據(jù)。

盡管作者在視覺任務(wù)的背景下提出了本文的模型，但所提出的多軸方法可以很容易地?cái)U(kuò)展到語言建模，以在線性時(shí)間內(nèi)捕獲局部和全局依存關(guān)系。此外，還可以在MLP-Mixer或gMLP模型上采用所提出的多軸方法，以提高可擴(kuò)展性和效率。