全面超越Swin Transformer ！Facebook用ResNet思想升級多尺度視覺Transformer

論文鏈接：https://arxiv.org/abs/2112.01526

在本文中研究了多尺度Vision Transformer(MViT)作為圖像和視頻分類以及目標(biāo)檢測的統(tǒng)一架構(gòu)。作者提出了一個(gè)MViT的改進(jìn)版本，它結(jié)合了分解的相對位置嵌入和池化殘差連接。

作者以五種大小實(shí)例化了這種架構(gòu)，并在ImageNet分類，COCO檢測和動(dòng)力學(xué)視頻識別方面進(jìn)行評估，它們優(yōu)于之前的工作。作者還進(jìn)一步比較了MViT的pooling attention和window attention，前者在準(zhǔn)確性/計(jì)算方面優(yōu)于后者。

MViT在3個(gè)領(lǐng)域具有最先進(jìn)的性能:ImageNet分類準(zhǔn)確率為88.8%，COCO目標(biāo)檢測準(zhǔn)確率為56.1AP%，Kinetics-400視頻分類準(zhǔn)確率為86.1%。

1 簡介

為不同的視覺識別任務(wù)設(shè)計(jì)架構(gòu)一直以來都很困難，而采用最廣泛的架構(gòu)是那些結(jié)合了簡單和高效的架構(gòu)，例如VGGNet和ResNet。最近，Vision Transformers(ViT)已經(jīng)展現(xiàn)出了有前途的性能，并可以與卷積神經(jīng)網(wǎng)絡(luò)競爭，最近也有很多研究提出了很多的改進(jìn)工作，將它們應(yīng)用到不同的視覺任務(wù)。

雖然ViT在圖像分類中很受歡迎，但其用于高分辨率目標(biāo)檢測和時(shí)空視頻理解任務(wù)仍然具有挑戰(zhàn)性。視覺信號的密度對計(jì)算和內(nèi)存需求提出了嚴(yán)峻的挑戰(zhàn)，因?yàn)樵诨赩ision Transformer的模型的Self-Attention Block中，這些信號的復(fù)雜性呈二次型。

采用了兩種不同的策略來解決這個(gè)問題:

1. Window Attention：在 Window 中進(jìn)行局部注意力計(jì)算用于目標(biāo)檢測;
2. Pooling Attention：在 Self-Attention 之前將局部注意聚合在一起。

而Pooling Attention為多尺度ViT帶來了很多的啟發(fā)，可以以一種簡單的方式擴(kuò)展ViT的架構(gòu):它不是在整個(gè)網(wǎng)絡(luò)中具有固定的分辨率，而是具有從高分辨率到低分辨率的多個(gè)階段的特性層次結(jié)構(gòu)。MViT是為視頻任務(wù)設(shè)計(jì)的，它具有最先進(jìn)的性能。

在本文中，作者做了兩個(gè)簡單的改進(jìn)以進(jìn)一步提高其性能，并研究了MViT作為一個(gè)單一的模型用于跨越3個(gè)任務(wù)的視覺識別：圖像分類、目標(biāo)檢測和視頻分類，以了解它是否可以作為空間和時(shí)空識別任務(wù)的一般視覺Backbone(見圖1)。

本文改進(jìn)MViT體系結(jié)構(gòu)主要包含以下內(nèi)容:

1. 創(chuàng)建了強(qiáng)大的Baseline，沿著2個(gè)軸提高pooling attention:使用分解的位置距離將位置信息注入Transformer塊;通過池化殘差連接來補(bǔ)償池化步長在注意力計(jì)算中的影響。上述簡單而有效的改進(jìn)帶來了更好的結(jié)果；
2. 將改進(jìn)的MViT結(jié)構(gòu)應(yīng)用到一個(gè)帶有特征金字塔網(wǎng)絡(luò)(FPN)的Mask R-CNN，并將其應(yīng)用于目標(biāo)檢測和實(shí)例分割；

作者研究MViT是否可以通過pooling attention來處理高分辨率的視覺輸入，以克服計(jì)算和內(nèi)存成本。

實(shí)驗(yàn)表明，pooling attention比 local window attention(如Swin)更有效。

作者進(jìn)一步開發(fā)了一個(gè)簡單而有效的Hybrid window attention方案，可以補(bǔ)充pooling attention以獲得更好的準(zhǔn)確性和計(jì)算折衷。

3. 在5個(gè)尺寸的增加復(fù)雜性(寬度，深度，分辨率)上實(shí)例化了架構(gòu)，并報(bào)告了一個(gè)大型多尺度ViT的實(shí)踐訓(xùn)練方案。該MViT變體以最小的改進(jìn)，使其可以直接應(yīng)用于圖像分類、目標(biāo)檢測和視頻分類。

實(shí)驗(yàn)表明，Improved MViT在ImageNet-21K上進(jìn)行預(yù)處理后，準(zhǔn)確率達(dá)到了88.8%(不進(jìn)行預(yù)處理的準(zhǔn)確率為86.3%)；僅使用Cascade Mask R-CNN在COCO目標(biāo)檢測AP可以達(dá)到56.1%。對于視頻分類任務(wù)，MViT達(dá)到了前所未有的86.1%的準(zhǔn)確率；對于Kinetics-400和Kinetics-600分別達(dá)到了86.1%和87.9%，并在Kinetics-700和Something-Something-v2上也分別達(dá)到了79.4%和73.7%的精度。

2 回顧一下多尺度ViT

MViT的關(guān)鍵思想是為low-level和high-level可視化建模構(gòu)建不同的階段，而不是ViT中的單尺度塊。如圖2所示MViT從輸入到輸出的各個(gè)階段緩慢地?cái)U(kuò)展通道D，同時(shí)降低分辨率L(即序列長度)。

為了在Transformer Block內(nèi)進(jìn)行降采樣，MViT引入了Pooling Attention。

具體來說, 對于一個(gè)輸入序列 , 對它應(yīng)用線性投影, 然后是Pool運(yùn)算(P)，分別用于query、key和value張量:

其中, 的長度 可以通過 來降低, 和 的長度可以通過 和 來降 低。

隨后, pooled self-attention可以表達(dá)為:

計(jì)算長度靈活的輸出序列 。注意, key和value的下采樣因子 和 可能與 應(yīng)用于query序列 的下采樣因子不同。

Pooling attention可以通過Pooling query Q來降低MViT不同階段之間的分辨率，并通過Pooling key K和value V來顯著降低計(jì)算和內(nèi)存復(fù)雜度。

Pooling Attention的Pytorch實(shí)現(xiàn)如下：

#?pool通常是MaxPool3d或AvgPool3d??
def?attention_pool(tensor,?pool,?thw_shape,?has_cls_embed=True,?norm=None):??
????if?pool?is?None:??
????????return?tensor,?thw_shape??
????tensor_dim?=?tensor.ndim??
????if?tensor_dim?==?4:??
????????pass??
????elif?tensor_dim?==?3:??
????????tensor?=?tensor.unsqueeze(1)??
????else:??
????????raise?NotImplementedError(f"Unsupported?input?dimension?{tensor.shape}")??

????if?has_cls_embed:??
????????cls_tok,?tensor?=?tensor[:,?:,?:1,?:],?tensor[:,?:,?1:,?:]??
??????
????B,?N,?L,?C?=?tensor.shape??
????T,?H,?W?=?thw_shape??
????tensor?=?(tensor.reshape(B?*?N,?T,?H,?W,?C).permute(0,?4,?1,?2,?3).contiguous())??
????#?執(zhí)行pooling操作??
????tensor?=?pool(tensor)??
??????
????thw_shape?=?[tensor.shape[2],?tensor.shape[3],?tensor.shape[4]]??
????L_pooled?=?tensor.shape[2]?*?tensor.shape[3]?*?tensor.shape[4]??
????tensor?=?tensor.reshape(B,?N,?C,?L_pooled).transpose(2,?3)??
????if?has_cls_embed:??
????????tensor?=?torch.cat((cls_tok,?tensor),?dim=2)??
????if?norm?is?not?None:??
????????tensor?=?norm(tensor)??
????#?Assert?tensor_dim?in?[3,?4]??
????if?tensor_dim?==?4:??
????????pass??
????else:??#??tensor_dim?==?3:??
????????tensor?=?tensor.squeeze(1)??
????return?tensor,?thw_shape??

3 Improved MViT

3.1 改進(jìn)Pooling Attention

1、分解相對位置嵌入

雖然MViT在建模Token之間的交互方面已經(jīng)顯示出了潛力，但它們關(guān)注的是內(nèi)容，而不是結(jié)構(gòu)。時(shí)空結(jié)構(gòu)建模僅依靠絕對位置嵌入來提供位置信息。這忽略了視覺的平移不變性的基本原則。也就是說，即使相對位置不變，MViT建模兩個(gè)patch之間交互的方式也會(huì)隨著它們在圖像中的絕對位置而改變。為了解決這個(gè)問題，作者將相對位置嵌入(相對位置嵌入只依賴于Token之間的相對位置距離)納入到Pooled Self-Attention 計(jì)算中。

這里將 2 個(gè)輸入元素 和 之間的相對位置編碼為位置嵌入 , 其中 和 表示元素 和 的空間(或時(shí)空)位置, 然后將兩兩編碼嵌入到Self-Attention模塊中:

然而, 在 中possible embeddings的 數(shù)量規(guī)模, 計(jì)算起來比較復(fù) 雜。為了降低復(fù)雜度, 作者將元素 和 之間的距離計(jì)算沿時(shí)空軸分解為:

其中 為沿高度、寬度和時(shí)間軸的位置嵌入， 和 分別表示 Token 的垂直、水平和時(shí)間位置。注意, 這里的 是可選的, 僅在視頻情況下需要 支持時(shí)間維度。相比之下, 分解嵌入將學(xué)習(xí)嵌入的數(shù)量減少到 , 這對早期的高分辨率特征圖有很大的影響。

2、池化殘差連接

pooled attention對于減少注意力塊中的計(jì)算復(fù)雜度和內(nèi)存需求是非常有效的。MViT在K和V張量上的步長比Q張量的步長大，而Q張量的步長只有在輸出序列的分辨率跨階段變化時(shí)才下采樣。這促使將殘差池化連接添加到Q（pooled后）張量增加信息流動(dòng)，促進(jìn)MViT中pooled attention Block的訓(xùn)練和收斂。

在注意力塊內(nèi)部引入一個(gè)新的池化殘差連接，如圖4所示。具體地說，將pooled query張量添加到輸出序列z中，因此將式(2)重新表述為:

注意，輸出序列Z與pooled query張量的長度相同。

消融實(shí)驗(yàn)表明，對于池化殘差連接，query的pool運(yùn)算符和殘差路徑都是必需的。由于上式中添加池化的query序列成本較低，因此在key和value pooling中仍然具有大跨步的低復(fù)雜度注意力計(jì)算。

3.2 MViT用于目標(biāo)檢測

1、改進(jìn)版MViT融合FPN

MViT的層次結(jié)構(gòu)分4個(gè)階段生成多尺度特征圖可以很自然地集成到特征金字塔網(wǎng)絡(luò)中(FPN)為目標(biāo)檢測任務(wù)，如圖5所示。在FPN中，帶有橫向連接的自頂向下金字塔為MViT在所有尺度上構(gòu)建了語義的特征映射。通過使用FPN與MViT Backbone將其應(yīng)用于不同的檢測架構(gòu)(例如Mask R-CNN)。

2、Hybrid window attention

Transformers中的self-attention具有與token數(shù)量的二次復(fù)雜度。這個(gè)問題對于目標(biāo)檢測來說更加嚴(yán)重，因?yàn)樗ǔＰ枰叻直媛实妮斎牒吞卣鲌D。在本文中研究了兩種顯著降低這種計(jì)算和內(nèi)存復(fù)雜度的方法:

1. 引入Pooling Attention
2. 引入Window Attention

Pooling Attention和Window Attention都通過減少計(jì)算Self-Attention時(shí)query、key和value的大小來控制Self-Attention的復(fù)雜性。

但它們的本質(zhì)是不同的:Pooling Attention通過局部聚合的向下采樣匯集注意力池化特征，但保持了全局的Self-Attention計(jì)算，而Window Attention雖然保持了分辨率，但通過將輸入劃分為不重疊的window，然后只在每個(gè)window中計(jì)算局部的Self-Attention。

這兩種方法的內(nèi)在差異促使研究它們是否能夠在目標(biāo)檢測任務(wù)中結(jié)合。

默認(rèn)Window Attention只在Window內(nèi)執(zhí)行局部Self-Attention，因此缺乏跨window的連接。與Swin使用移動(dòng)window來緩解這個(gè)問題不同，作者提出了一個(gè)簡單的Hybrid window attention(Hwin)來增加跨window的連接。

Hwin在一個(gè)window內(nèi)計(jì)算所有的局部注意力，除了最后3個(gè)階段的最后塊，這些階段都饋入FPN。通過這種方式，輸入特征映射到FPN包含全局信息。

消融實(shí)驗(yàn)顯示，這個(gè)簡單的Hwin在圖像分類和目標(biāo)檢測任務(wù)上一貫優(yōu)于Swin。進(jìn)一步，將證明合并pooling attention和Hwin在目標(biāo)檢測方面實(shí)現(xiàn)了最好的性能。

3.3 MViT用于視頻識別

1、從預(yù)訓(xùn)練的MViT初始化

與基于圖像的MViT相比，基于視頻的MViT只有3個(gè)不同之處:

1. patchification stem中的投影層需要將輸入的數(shù)據(jù)投影到時(shí)空立方體中，而不是二維的patch;
2. 現(xiàn)在，pool運(yùn)算符對時(shí)空特征圖進(jìn)行池化;
3. 相對位置嵌入?yún)⒖紩r(shí)空位置。

由于1和2中的投影層和池化操作符默認(rèn)由卷積層實(shí)例化，對CNN的空洞率初始化為[8,25]。具體地說，作者用來自預(yù)訓(xùn)練模型中的2D conv層的權(quán)重初始化中心幀的conv kernel，并將其他權(quán)重初始化為0。

對于3利用在Eq.4中分解的相對位置嵌入，并簡單地從預(yù)訓(xùn)練的權(quán)重初始化空間嵌入和時(shí)間嵌入為0。

4.4 MViT架構(gòu)變體

本文構(gòu)建了幾個(gè)具有不同數(shù)量參數(shù)和FLOPs的MViT變體，如表1所示，以便與其他vision transformer進(jìn)行公平的比較。具體來說，通過改變基本通道尺寸、每個(gè)階段的塊數(shù)以及塊中的head數(shù)，為MViT設(shè)計(jì)了5種變體(Tiny、Small、Base、Large和Huge)。

遵循MViT中的pooling attention設(shè)計(jì)，本文在所有pooling attention塊中默認(rèn)采用Key和Value pooling，pooling attention blocks步長在第一階段設(shè)置為4，并在各個(gè)階段自適應(yīng)地衰減stride。

4 實(shí)驗(yàn)

4.1 消融實(shí)驗(yàn)

1、不同的注意力機(jī)制

通過上表關(guān)于ImageNet分類和COCO目標(biāo)檢測的實(shí)驗(yàn)可以看出，本文提出的HSwin注意力與Pooling可以得到最好性能的結(jié)果。

2、位置嵌入

上表比較了不同位置嵌入?？梢杂^察到:

• 比較(2)與(1)，絕對位置僅比無位置略有提高。這是因?yàn)閜ool操作符已經(jīng)建模了位置信息。
• (2)比較(3,4)和(1,2)，相對位置通過引入轉(zhuǎn)移不變性先驗(yàn)來Pooling Attention，從而帶來收益。
• 最后，在COCO上分解的相對位置嵌入序列比joint相對位置快3.9倍。

3、池化殘差連接

上表研究了池化殘差連接的重要性。可以看到：

• (2)簡單地添加殘差路徑可以改善這兩種情況的結(jié)果IN-1K(+0.3%)和COCO(+0.8 APbox)的成本可以忽略不計(jì)。
• (3)使用池化殘差連接，并在所有其他層中添加Q pooled(stride=1)，性能得到顯著的提升，特別是在COCO(+1.4 APbox)上。這表明在MViT中，Q pooled塊和殘差路徑都是必要的。
• (4)僅僅添加(沒有殘差)更多的Q pooled層stride=1沒有幫助。

4、運(yùn)行時(shí)間比較

5、FPN消融

如表所示，F(xiàn)PN顯著提高了兩種Backbone的性能，而MViT-S始終優(yōu)于ViT-B。MViT-S的FPN增益(+2.9)比ViT-B(+1.5 APbox)大得多，這表明了分層多尺度設(shè)計(jì)對于密集目標(biāo)檢測任務(wù)的有效性。

4.2 ImageNet-1K

4.3 COCO目標(biāo)檢測

4.4 視頻識別

參考

[1].Improved Multiscale Vision Transformers for Classification and Detection

如果覺得有用，就請分享到朋友圈吧！