卷!超越 Swin、ConvNeXt |Facebook提出Neighborhood Attention Transformer

本文提出了Neighborhood Attention Transformer(NAT),NAT是一種集高效、準(zhǔn)確和可擴展的分層Transformer,Neighborhood Attention是一種簡單而靈活的
Self Attention機制,它將每個query的感受野擴展到其最近的鄰近像素,并隨著感受野的增大而接近Self-Attention。在FLOPs和內(nèi)存使用方面,與相同的感受野大小帶有Shifted Window Attention的Swin-Transformer相同,但是NAT受到的約束更少。
此外,NA還包含了局部歸納偏差,這消除了的額外操作,如像素移動。
NAT-Tiny在ImageNet上達到83.2%的Top-1精度,只有4.3 GFLOPs和28M參數(shù),在MS-COCO上的mAP為51.4%,ADE20k上的mIoU為48.4%。
1事出緣由
自2020年以來,學(xué)術(shù)界已經(jīng)提出很多關(guān)于Vision Transformer的架構(gòu)和方法,也有很多關(guān)于高效Transformer的工作被提出。但是,不管是標(biāo)準(zhǔn)的ViT還是高效的ViT,雖然其中的Self Attention相對于嵌入維度具有線性復(fù)雜度(不包括線性投影),但相對于Token數(shù)量卻具有二次復(fù)雜度。而Token的數(shù)量通常又與圖像分辨率呈線性相關(guān)。因此,更高的圖像分辨率將會帶來復(fù)雜度和內(nèi)存的二次增長。
對于ViT的應(yīng)用來說,這是一個很嚴峻的問題,因為過高的復(fù)雜度和計算量會影響模型應(yīng)用于下游的視覺任務(wù),如目標(biāo)檢測和語義分割,因為這些任務(wù)中圖像分辨率通常比分類大得多。
另一個問題是,卷積的性能大多受益于歸納偏差,如局部性、平移等方差和二維鄰域結(jié)構(gòu),而 dot-product Self Attention是一個全局的操作。雖然Vision Transformer中的MLP層具有具有局部性和平移不變性,但其余的歸納偏差必須通過大量數(shù)據(jù)或煉丹技術(shù)和增強來彌補。
因此,有研究便采用Local attention modules來減輕這個問題。Swin-Transformer是第一個基于Local attention的分層Vision Transformer。Swin-Transformer的層次結(jié)構(gòu)設(shè)計和Shifted-Window Self Attention使其可以比較容易地應(yīng)用到下游任務(wù),同時也通過注入的額外偏差提高了性能。HaloNet探索了另一種Local Attention Block,并發(fā)現(xiàn)將Local Attention Block和卷積結(jié)合起來可以獲得最佳性能,因為該組合操作在內(nèi)存使用和平移不變性之間進行了最佳權(quán)衡。
基于以上描述,作者在本文中提出Neighborhood Attention(NA),并在此基礎(chǔ)上構(gòu)建Neighborhood Attention Transformer(NAT),在跨視覺任務(wù)上實現(xiàn)了具有競爭性的結(jié)果。
NA是dot-product Self Attention的一種局部化,將每個query token的感受野限制在key-value pair中對應(yīng)token周圍的一個固定大小的鄰域內(nèi)。較小的感受野帶來更多的局部信息,而較大的感受野帶來更多的全局信息。這種設(shè)計使感受野能夠很好的控制平移不變性和平移等方差的平衡。
NA是受到卷積的局部性的啟發(fā),以及它們?nèi)绾萎a(chǎn)生更多有利于視覺任務(wù)的局部性歸納偏差。它不同于將Self Attention應(yīng)用于Local Windows(Swin),可以將其視為與Content-Dependant Kernel的卷積。
2主要貢獻
Neighborhood Attention(NA):一種簡單、靈活的視覺注意力機制,它將每一個Token的感受野定位到Token的鄰域。并將該模塊的復(fù)雜性和內(nèi)存使用量與Self Attention、Window Self Attention和卷積進行了比較。構(gòu)建了
Neighborhood Attention Transformer(NAT),一種由Neighborhood Attention組成的高效、準(zhǔn)確、可擴展的新型分層Transformer。每一層之后都進行下采樣操作,將空間大小減少一半。類似的設(shè)計可以在許多最近的基于注意力的模型中看到,如Swin和Focal Transformer。與那些模型不同,NAT利用小內(nèi)核重疊卷積來嵌入和下采樣,而不是非重疊卷積。與Swin等現(xiàn)有技術(shù)相比,NAT還引入了一組更有效的體系結(jié)構(gòu)配置。展示了NAT在圖像分類和下游視覺任務(wù)(包括目標(biāo)檢測和語義分割)中的有效性。作者觀察到NAT的性能不僅優(yōu)于Swin Transformer,還優(yōu)于ConvNeXt。NAT-Tiny模型在ImageNet上僅用4.3 GFLOPs和28M參數(shù)就能達到83.2%的top-1精度,在MS-COCO上達到51.4%的Box mAP,在ADE20k上達到48.4%的多尺度mIoU,為這種簡單、小規(guī)模的
Transformer模型創(chuàng)造了水平。

3本文方法
與Self Attention相比,Neighborhood Attention不僅減少了計算成本,而且引入了類似于卷積的局部歸納偏差。該操作適用于鄰域大小L,當(dāng)每個像素最小時,它只關(guān)注自身周圍的1個像素鄰域(創(chuàng)建一個3×3方形窗口)。
作者還表明,當(dāng)鄰域大小達到最大值(即輸入的大小)時,Neighborhood Attention等于Self Attention。因此,如果鄰域大小超過或匹配feature map的大小,在相同的輸入下,Neighborhood Attention和Self Attention的輸出是相等的。
此外,NAT利用了一個多級分層設(shè)計,類似于Swin-Transformer,這意味著特征映射在級別之間被向下采樣,而不是一次性全部采樣。然而,與Swin-Transformer不同的是,NAT使用重疊卷積來向下采樣特征映射,而不是不重疊的映射。這略微增加了計算量和參數(shù),作者也通過提出計算成本較低的配置來彌補這一點。
3.1 Neighborhood Attention
Neighborhood Attention主要是受卷積如何引入鄰域偏差和局部性的啟發(fā)。Neighborhood Attention的目的是允許特征圖中的每個像素只關(guān)注其相鄰的像素。由于鄰域依賴于size,因此Neighborhood Attention機制也是如此。
用表示處的一個像素的鄰域,它是最接近的像素的一個固定長度的索引集。對于大小為。因此,在單個像素上的Neighborhood Attention可以定義如下:

其中,其中Q、K、V是X的線性投影,為相對位置偏差,根據(jù)相對位置將其添加到每個注意力權(quán)重中。這個操作可以進一步擴展到所有像素,從而形成一種局部注意力的形式。
然而,如果函數(shù)將每個像素映射到所有像素,這將等同于Self Attention(帶有額外的位置偏差)。這是因為當(dāng)鄰域邊界超過輸入大小時,將包括所有可能的像素。結(jié)果,和,通過去除偏置項,推導(dǎo)出Self Attention的表達形式:

Neighborhood Attention在計算上是很廉價的。它的復(fù)雜性相對于分辨率是線性的,不像Self Attention是二次的。此外,其復(fù)雜度也與鄰域大小呈線性關(guān)系。函數(shù)將一個像素映射到一組相鄰的像素,可以很容易地通過Raster-Scan Sliding Window操作產(chǎn)生,原理類似于卷積。

每個像素被簡單地映射到一組相鄰的像素和自身。圖2展示了該操作的示例。對特征圖中的每個像素重復(fù)此操作。對于不能居中的角像素,擴展鄰域以保持感受野的大小。這是一個關(guān)鍵的設(shè)計選擇,它允許NA隨著鄰域大小向特征圖分辨率增長而一般化到Self Attention。

擴展鄰域是通過簡單持續(xù)選擇L2中最接近原始鄰域的像素來實現(xiàn)的。例如,對于L=3,每個query將以圍繞它的9個key-value像素結(jié)束(query位于中心的3×3網(wǎng)格)。
對于角像素,鄰域是另一個3×3網(wǎng)格,但query沒有定位在中心。圖6展示了這個想法的說明。

通過上表可以看出,Neighborhood Attention的復(fù)雜度和內(nèi)存消耗和Swin相同。
3.2 Neighborhood Attention Transformer
NAT通過使用2個連續(xù)的3×3卷積(stride=2)來嵌入輸入,進而產(chǎn)生輸入空間大小1/4的輸入。這類似于使用patch和4×4 patch的嵌入層,但它使用的是重疊卷積而不是非重疊卷積。
另一方面,使用重疊卷積會增加成本,而2次卷積會引入更多的參數(shù)。然而,作者通過重新配置模型來處理這個問題,這將產(chǎn)生更好的權(quán)衡。
NAT由4個level組成,每個level后面都鏈接一個下采樣器(最后一個除外)。下采樣器將空間大小減少了為原來的一半,而通道數(shù)量增加了一倍。這里下采樣使用的是3×3卷積(stride=2)。由于tokenizer的降采樣倍數(shù)為4倍,因此模型生成了大小為H/4×W/4、H/8×W/8、H/16×W/16和H/32×W/32的特征圖.使得NAT可以更容易地將預(yù)訓(xùn)練過的模型遷移到下游任務(wù)中。

此外,在訓(xùn)練較大的模型時,使用LayerScale來提高穩(wěn)定性。圖4展示了整體網(wǎng)絡(luò)架構(gòu)的說明。在下表中總結(jié)了不同的NAT變體及其關(guān)鍵差異。

4實驗結(jié)果
4.1 分類實驗

4.2 目標(biāo)檢測

4.3 語義分割

4.4 可視化分析

5參考文獻
[1].Neighborhood Attention Transformer
