超越PVT、Swin，南大開(kāi)源高效Transformer：ResT?

paper: https://arxiv.org/abs/2105.13677

code: https://github.com/wofmanaf/ResT

本文是南京大學(xué)提出的一種高效Transformer架構(gòu):ResT，它采用了類似ResNet的設(shè)計(jì)思想：stem提取底層特征信息、stages捕獲多尺度特征信息。與此同時(shí)，為解決MSA存在的計(jì)算量與內(nèi)存占用問(wèn)題，提出了EMSA模塊進(jìn)一步降低計(jì)算量與內(nèi)存消耗。所提ResT在圖像分類、目標(biāo)檢測(cè)以及實(shí)例分割等任務(wù)均取得了顯著的性能提升，比如在ImageNet數(shù)據(jù)上，在同等計(jì)算量前提下，所提方法取得了優(yōu)于PVT、Swin的優(yōu)異性能，實(shí)乃一種強(qiáng)力骨干網(wǎng)絡(luò)。

Abstract

本文提出一種高效多尺度Vision Transformer：ResT，它可作為圖像中識(shí)別的通用骨干架構(gòu)。不同于現(xiàn)有采用固定分辨率+標(biāo)準(zhǔn)Transformer模塊的Transformer模型，它有這樣幾個(gè)優(yōu)勢(shì)：

• (1) 提出了一種內(nèi)容高效的多頭自注意力模塊，它采用簡(jiǎn)單的深度卷積進(jìn)行內(nèi)存壓縮，并跨注意力頭維度進(jìn)行投影交互，同時(shí)保持多頭的靈活性；
• (2) 將位置編碼構(gòu)建為空域注意力，它可以更靈活的處理任意分辨率輸入，且無(wú)需插值或者微調(diào)；
• (3) 并未在每個(gè)階段的開(kāi)始部分進(jìn)行序列化，我們把塊嵌入設(shè)計(jì)成重疊卷積堆疊方式。

我們?cè)趫D像分類與下游任務(wù)上對(duì)所提ResT進(jìn)行了性能驗(yàn)證，實(shí)驗(yàn)結(jié)果表明：所提ResT大幅優(yōu)于現(xiàn)有骨干架構(gòu)，比如，ResNet18(69.7%)、PVT-Tiny(75.1%)相似大小的模型下，所提方法取得了79.5%的top1精度，這說(shuō)明它是一種強(qiáng)有力的骨干網(wǎng)絡(luò)。

Method

上圖給出了ResT的架構(gòu)示意圖，它具有與ResNet相似的結(jié)構(gòu)方案，比如采用stem模塊提取底層特征，后接四個(gè)stage捕獲多尺度特征。每個(gè)stage包含三個(gè)成分：一個(gè)塊嵌入模塊，一個(gè)位置編碼模塊以及L個(gè)高效Transformer模塊。具體來(lái)說(shuō)，在每個(gè)stge開(kāi)始前，塊嵌入模塊用于降低輸入的分辨率并擴(kuò)展通道維度；位置編碼用于約束位置信息提升塊嵌入的特征提取能力；然后將所得送入到后續(xù)高效Transformer模塊中。

Rethinking of Transformer Block

標(biāo)準(zhǔn)的Transformer模塊由MSA與FFN以及殘差鏈接構(gòu)成，在MSA與FFN之前還采用LN。對(duì)于輸入token，每個(gè)Transformer模塊的輸出表示如下：

其中，MSA的單頭SA與FFN的定義分別如下：

MSA與FFN的計(jì)算復(fù)雜度分別為, 。

Efficient Transformer Block

如前所述，MSA有兩個(gè)缺點(diǎn)：(1) 計(jì)算量隨平方增長(zhǎng)，這會(huì)導(dǎo)致較大的訓(xùn)練與推理負(fù)載；(2) MSA的每個(gè)頭僅負(fù)責(zé)輸入的部分子集，這會(huì)影響模型的性能，尤其當(dāng)通道維度非常小時(shí)。

為解決上述問(wèn)題，我們提出了上圖所示的高效多頭自注意力模塊?？梢钥吹剑?/p>

• 類似MSA，EMSA首先采用投影集合得到Q；
• 為壓縮內(nèi)存，2D輸入將被reshap為3D形式，然后送入深度卷積以因子降低空域維度;
• 將上述所得特征reshape為2D形式并送入后兩個(gè)投影集合得到K與V；
• 然后采用下面公式計(jì)算注意力，注：Conv為卷積，用于對(duì)不同頭進(jìn)行信息交互。為補(bǔ)償Conv導(dǎo)致的多樣性素食，我們?cè)赟oftmax之后添加了IN。

• 最后，每個(gè)頭的輸出進(jìn)行拼接并線性 投影構(gòu)成最終的輸出。

EMSA的計(jì)算復(fù)雜度為，具有比MSA更低的計(jì)算量。此時(shí)，高效Transformer模塊定義如下：

Patch Embedding

標(biāo)準(zhǔn)的Transformer采用一序列token作為輸入，以ViT為例，3D圖像需要拆分為的塊，這些塊再平展為2D形式并映射為隱嵌入。然而，這種直接的tokenization難以捕獲底層特征信息(比如邊緣、角點(diǎn))。此外，ViT中的tokens長(zhǎng)度是固定的，這使其難以進(jìn)行下游任務(wù)(比如目標(biāo)檢測(cè)、實(shí)例分割)適配。

為解決上述問(wèn)題，我們構(gòu)建了一種高效多尺度骨干ResT用于稠密預(yù)測(cè)。正如前面所提到的，每個(gè)階段的高效Transformer模塊在同尺度同分辨率上跨通道、空域維度進(jìn)行處理。因此，塊嵌入模塊同樣需要漸進(jìn)的擴(kuò)展通道維度，同時(shí)降低空域分辨率。

類似于ResNet，我們采用stem模塊以倍率4收縮寬高維度。為高效捕獲底層特征信息，我們引入了一種簡(jiǎn)單而有效的方式：堆疊三個(gè)卷積，stride分別為212，前兩個(gè)后接BatchNorm與ReLU。在234階段，采用塊嵌入模塊下采樣空間分辨并提升通道維度，這與stride=2的卷積作用類似。

Position Encoding

位置編碼對(duì)于序列順序的探索非常關(guān)鍵，ViT一文將可學(xué)習(xí)參數(shù)加到輸入tokens中編碼位置信息。假設(shè)為輸入，表示位置參數(shù)，那么編碼后輸入表示如下：

然而，此時(shí)要求位置長(zhǎng)度與輸入tokens長(zhǎng)度相同，這無(wú)疑會(huì)限制了其應(yīng)用。

為解決上述問(wèn)題，我們需要設(shè)計(jì)一種新的變長(zhǎng)位置編碼，我們將上式修改為如下：

其中表示組線性操作，組數(shù)為c。

除了上述形式外，我們還可以采用更靈活的注意力機(jī)制得到像素級(jí)權(quán)值。因此，我們提出了一種簡(jiǎn)單且高效的像素注意力(Pixel-wise Attention,PA)模塊進(jìn)行位置編碼。具體來(lái)說(shuō)，PA采用采用深度卷積計(jì)算像素權(quán)值，然后采用sigmoid激活，那么帶PA的位置編碼可以描述如下：

由于每個(gè)stage的輸入token通過(guò)卷積得到，我們可以將位置編碼嵌入到塊嵌入模塊中，整體結(jié)果見(jiàn)上圖。注：這里的PA可以采用任意空域注意力替換，這使得ResT中的PE極為靈活。

Linear Head

分類頭采用全局均值池化+線性分類器的方式，ResT的架構(gòu)配置信息見(jiàn)下表。

Experiments

接下來(lái)，我們?cè)诔Ｓ没鶞?zhǔn)任務(wù)上進(jìn)行所提方案驗(yàn)證，包含ImageNet數(shù)據(jù)上的圖像分類、COCO數(shù)據(jù)上的目標(biāo)檢測(cè)與實(shí)例分割等。

Image Classification

上表給出了圖像分類任務(wù)上不同方案的性能對(duì)比，從中可以看到：

• 在小模型方面，ResT-small憑借相似的復(fù)雜度以79.6%精度大幅超過(guò)PVT-T的75.1%；
• 在中等模型方面，ResT-base憑借相似復(fù)雜度以81.6%超過(guò)Swin-T的81.3%；
• 在大模型方面，ResT-Large憑借相似復(fù)雜度以83.6%精度超過(guò)Swin-S的83.3%；
• 相比ConvNet，如RegNet，所提ResT憑借相似復(fù)雜度取得了更佳的性能；
• 總而 言之，在不同復(fù)雜度模型方面，ResT均顯著優(yōu)于現(xiàn)有模型。

Object Detection and Instance Segmentation

上表給出了RetinaNet架構(gòu)下的不同骨干模型在目標(biāo)檢測(cè)上的性能對(duì)比，可以看到：

• 在小模型方面，相比PVT-T，ResT-Small取得了2.8的指標(biāo)提升；
• 在大模型方面，相比PVT-S，ResT-Base取得了0.8的指標(biāo)提升。

上表給出了實(shí)例分割任務(wù)上的性能對(duì)比，可以看到：

• 在小模型方面，相比PVT-T，ResT-Small取得了1.8boxAP指標(biāo)提升，1.0MaskAP指標(biāo)提升；
• 在大模型方面，相比PVT-S，ResT-Base分別取得了2.1與1.9的指標(biāo)提升。

Ablation Study

接下來(lái)，我們對(duì)所提ResT進(jìn)行消融實(shí)驗(yàn)分析，主要從stem、EMSA、PE三個(gè)角度進(jìn)行對(duì)比分析。

從下圖的Table5可以看到：ResT中的stem比PVT、ResNet中的Stem更加高效，分別取得了0.92%、0.64%的性能提升。

從上圖的Table6可以看到：

• 當(dāng)移除卷積操作與IN后，模型性能下降1.16%，這說(shuō)明長(zhǎng)序列與靈活性的組合對(duì)于注意力非常重要；
• 當(dāng)移除IN后，模型同樣出現(xiàn)了大幅性能下降，我們將其歸因于不同頭之間的多樣性遭到了破壞。

上表對(duì)比了不同PE的性能對(duì)比，從中可以看到：

• 當(dāng)移除PA編碼后，模型性能從72.88%下降到71.54%，這說(shuō)明位置編碼對(duì)于ResT非常重要；
• LE與GL具有相似性能，而PA以0.84%精度優(yōu)于GL，這說(shuō)明：空域注意力可用于進(jìn)行位置編碼建模。

如果覺(jué)得有用，就請(qǐng)分享到朋友圈吧！