本文中，來(lái)自清華大學(xué)、AWS AI 和北京智源人工智能研究院的研究者提出了一種新型可變形自注意力模塊，其中以數(shù)據(jù)相關(guān)的方式選擇自注意力中鍵值對(duì)的位置，使得自注意力模塊能夠?qū)Ｗ⒂谙嚓P(guān)區(qū)域，并捕獲更多信息特征。

Transformer 近來(lái)在各種視覺(jué)任務(wù)上表現(xiàn)出卓越的性能，感受野賦予 Transformer 比 CNN 更強(qiáng)的表征能力。然而，簡(jiǎn)單地?cái)U(kuò)大感受野會(huì)引起一些問(wèn)題。一方面，使用密集注意力（例如 ViT）會(huì)導(dǎo)致過(guò)多的內(nèi)存和計(jì)算成本，并且特征可能會(huì)受到超出興趣區(qū)域的無(wú)關(guān)部分的影響；另一方面，PVT 或 Swin Transformer 中采用的稀疏注意力與數(shù)據(jù)無(wú)關(guān)，可能會(huì)限制對(duì)遠(yuǎn)程（long range）關(guān)系建模的能力。

為了緩解這些問(wèn)題，清華大學(xué)、AWS AI 和北京智源人工智能研究院的研究者提出了一種新型可變形自注意力模塊，其中以數(shù)據(jù)相關(guān)的方式選擇自注意力中鍵值對(duì)的位置。這種靈活的方案使自注意力模塊能夠?qū)Ｗ⒂谙嚓P(guān)區(qū)域并捕獲更多信息特征。

在此基礎(chǔ)上，該研究提出了可變形注意力 Transformer（Deformable Attention Transformer，DAT），一種具有可變形注意力的通用主干網(wǎng)絡(luò)模型，適用于圖像分類和密集預(yù)測(cè)任務(wù)。該研究通過(guò)大量基準(zhǔn)測(cè)試實(shí)驗(yàn)證明了該模型的性能提升。

論文地址：https://arxiv.org/abs/2201.00520v1

可變形注意力 Transformer

現(xiàn)有的分層視覺(jué) Transformer，特別是 PVT 和 Swin Transformer 試圖解決過(guò)多注意力的挑戰(zhàn)。前者的下采樣技術(shù)會(huì)導(dǎo)致嚴(yán)重的信息損失，而后者的 Swin 注意力導(dǎo)致感受野的增長(zhǎng)要慢得多，這限制了對(duì)大型物體進(jìn)行建模的潛力。因此，需要依賴于數(shù)據(jù)的稀疏注意力來(lái)靈活地對(duì)相關(guān)特征進(jìn)行建模，從而導(dǎo)致首先在 DCN [9] 中提出可變形機(jī)制。

然而，在 Transformer 模型中實(shí)現(xiàn) DCN 是一個(gè)不簡(jiǎn)單的問(wèn)題。在 DCN 中，特征圖上的每個(gè)元素單獨(dú)學(xué)習(xí)其偏移量，其中 H ×W ×C 特征圖上的 3 × 3 可變形卷積具有 9 HWC 的空間復(fù)雜度。如果在注意力模塊中直接應(yīng)用相同的機(jī)制，空間復(fù)雜度將急劇上升到 N_qN_kC，其中 N_q、N_k 是查詢和鍵的數(shù)量，通常與特征圖大小 HW 具有相同的比例，帶來(lái)近似于雙二次的復(fù)雜度。

盡管 Deformable DETR [54] 已經(jīng)設(shè)法通過(guò)在每個(gè)尺度上設(shè)置較少數(shù)量的 N_k = 4 的鍵來(lái)減少這種開(kāi)銷，并且可以很好地作為檢測(cè)頭，但由于不可接受的信息丟失（參見(jiàn)附錄中的詳細(xì)比較），在骨干網(wǎng)絡(luò)中關(guān)注如此少的鍵效果不佳。與此同時(shí)，[3,52] 中的觀察表明，不同的查詢?cè)谝曈X(jué)注意力模型中具有相似的注意力圖。因此，該研究選擇了一個(gè)更簡(jiǎn)單的解決方案，為每個(gè)查詢共享移位鍵和值，以實(shí)現(xiàn)有效的權(quán)衡。

模型架構(gòu)

該研究在 Transformer（等式 (4)）中的可變形注意力替換了 vanilla MHSA，并將其與 MLP（等式 (5)）相結(jié)合，以構(gòu)建一個(gè)可變形的視覺(jué) transformer 塊。在網(wǎng)絡(luò)架構(gòu)方面， DAT 與 [7, 26, 31, 36] 共享類似的金字塔結(jié)構(gòu)，廣泛適用于需要多尺度特征圖的各種視覺(jué)任務(wù)。如下圖 3 所示，形狀為 H × W × 3 的輸入圖像首先被步長(zhǎng)為 4 的 4 × 4 非重疊卷積嵌入，然后一個(gè)歸一化層獲得補(bǔ)丁嵌入。

為了構(gòu)建分層特征金字塔，主干包括 4 個(gè)階段，步幅逐漸增加。在兩個(gè)連續(xù)的階段之間，有一個(gè)步長(zhǎng)為 2 的非重疊 2×2 卷積，對(duì)特征圖進(jìn)行下采樣，將空間大小減半并將特征維度加倍。

在分類任務(wù)中，該研究首先對(duì)最后階段輸出的特征圖進(jìn)行歸一化，然后采用具有池化特征的線性分類器來(lái)預(yù)測(cè)對(duì)數(shù)；在對(duì)象檢測(cè)、實(shí)例分割和語(yǔ)義分割任務(wù)中，DAT 在集成視覺(jué)模型中扮演主干的角色，以提取多尺度特征。該研究為每個(gè)階段的特征添加一個(gè)歸一化層，然后將它們輸入到以下模塊中，例如對(duì)象檢測(cè)中的 FPN [23] 或語(yǔ)義分割中的****。

實(shí)驗(yàn)

該研究在 3 個(gè)數(shù)據(jù)集上進(jìn)行了實(shí)驗(yàn)，以驗(yàn)證提出的 DAT 的有效性。該研究展示了在 ImageNet-1K [10] 分類、COCO 目標(biāo)檢測(cè)和 ADE20K 語(yǔ)義分割任務(wù)上的結(jié)果。此外，該研究提供了消融研究和可視化結(jié)果，以進(jìn)一步展示該方法的有效性。

ImageNet-1K 分類

ImageNet-1K [10] 數(shù)據(jù)集有 128 萬(wàn)張用于訓(xùn)練的圖像和 5 萬(wàn)張用于驗(yàn)證的圖像。研究者在訓(xùn)練分割上訓(xùn)練 DAT 的三個(gè)變體，并報(bào)告驗(yàn)證分割上的 Top-1 準(zhǔn)確度，并與其他 Vision Transformer 模型進(jìn)行比較。

該研究在下表 2 中給出了有 300 個(gè)訓(xùn)練 epoch 的結(jié)果。與其他 SOTA 視覺(jué) Transformer 模型相比， DAT 在具有相似計(jì)算復(fù)雜性的情況下在 Top-1 精度上實(shí)現(xiàn)了顯著提高。DAT 在所有三個(gè)尺度上都優(yōu)于 Swin Transformer [26]、PVT [36]、DPT [7] 和 DeiT [33]。沒(méi)有在 Transformer 塊 [13, 14, 35] 中插入卷積，或在補(bǔ)丁嵌入 [6, 11, 45] 中使用重疊卷積，DAT 比 Swin Transformer [26] 實(shí)現(xiàn)了 +0.7、+0.7 和 +0.5 的增益。在 384 × 384 分辨率下進(jìn)行微調(diào)時(shí)，該模型繼續(xù)比 Swin Transformer 性能好 0.3。

COCO 目標(biāo)檢測(cè)

COCO 目標(biāo)檢測(cè)和實(shí)例分割數(shù)據(jù)集有 118K 的訓(xùn)練圖像和 5K 的驗(yàn)證圖像。該研究使用 DAT 作為 RetinaNet [24]、Mask R-CNN [17] 和 Cascade Mask R-CNN [2] 框架中的主干，以評(píng)估該方法的有效性。該研究在 ImageNet-1K 數(shù)據(jù)集上對(duì)該模型進(jìn)行 300 個(gè) epoch 的預(yù)訓(xùn)練，并遵循 Swin Transformer [26] 中類似的訓(xùn)練策略來(lái)公平地比較該方法。該研究在 1x 和 3x 訓(xùn)練計(jì)劃中報(bào)告在 RetinaNet 模型上的 DAT。如下表 3 所示，在微型和小型模型中，DAT 的性能優(yōu)于 Swin Transformer 1.1 和 1.2 mAP。

當(dāng)在兩階段檢測(cè)器（例如 Mask R-CNN、Cascade Mask R-CNN）中實(shí)現(xiàn)時(shí)，DAT 模型在不同尺寸的 Swin Transformer 模型上實(shí)現(xiàn)了一致的改進(jìn)，如下表 4 所示。

下表 5 給出了在驗(yàn)證集上各種方法的 mIoU 分?jǐn)?shù)。

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

為了驗(yàn)證 DAT 模型中關(guān)鍵組件設(shè)計(jì)的有效性， 該研究進(jìn)行了消融實(shí)驗(yàn)，報(bào)告了基于 DAT-T 的 ImageNet-1K 分類結(jié)果。對(duì)于幾何信息開(kāi)發(fā)，該研究首先評(píng)估了所提可變形偏移和可變形相對(duì)位置嵌入的有效性，如下表 6 所示。

對(duì)于不同階段的可變形注意力，該研究用不同階段的可變形注意力替換了 Swin Transfomer [26] 的移位窗口注意力。如下表 7 所示，僅替換最后階段的注意力提高了 0.1，替換最后兩個(gè)階段導(dǎo)致性能增益為 0.7（達(dá)到 82.0 的整體準(zhǔn)確度）。然而，在早期階段用更多可變形注意力替換會(huì)略微降低準(zhǔn)確性。

可視化

該研究在 DAT 中可視化學(xué)習(xí)變形位置的示例，以驗(yàn)證該方法的有效性。如下圖 4 所示，采樣點(diǎn)描繪在對(duì)象檢測(cè)框和實(shí)例分割掩碼的頂部，從中可以看到這些點(diǎn)已轉(zhuǎn)移到目標(biāo)對(duì)象。