Swin Transformer遇DCN,可變形注意力Transformer模型優(yōu)于多數ViT
Transformer 近來在各種視覺任務上表現(xiàn)出卓越的性能,感受野賦予 Transformer 比 CNN 更強的表征能力。然而,簡單地擴大感受野會引起一些問題。一方面,使用密集注意力(例如 ViT)會導致過多的內存和計算成本,并且特征可能會受到超出興趣區(qū)域的無關部分的影響;另一方面,PVT 或 Swin Transformer 中采用的稀疏注意力與數據無關,可能會限制對遠程(long range)關系建模的能力。
為了緩解這些問題,清華大學、AWS AI 和北京智源人工智能研究院的研究者提出了一種新型可變形自注意力模塊,其中以數據相關的方式選擇自注意力中鍵值對的位置。這種靈活的方案使自注意力模塊能夠專注于相關區(qū)域并捕獲更多信息特征。
在此基礎上,該研究提出了可變形注意力 Transformer(Deformable Attention Transformer,DAT),一種具有可變形注意力的通用主干網絡模型,適用于圖像分類和密集預測任務。該研究通過大量基準測試實驗證明了該模型的性能提升。
論文地址:https://arxiv.org/abs/2201.00520v1
可變形注意力 Transformer
現(xiàn)有的分層視覺 Transformer,特別是 PVT 和 Swin Transformer 試圖解決過多注意力的挑戰(zhàn)。前者的下采樣技術會導致嚴重的信息損失,而后者的 Swin 注意力導致感受野的增長要慢得多,這限制了對大型物體進行建模的潛力。因此,需要依賴于數據的稀疏注意力來靈活地對相關特征進行建模,從而導致首先在 DCN [9] 中提出可變形機制。
然而,在 Transformer 模型中實現(xiàn) DCN 是一個不簡單的問題。在 DCN 中,特征圖上的每個元素單獨學習其偏移量,其中 H ×W ×C 特征圖上的 3 × 3 可變形卷積具有 9 HWC 的空間復雜度。如果在注意力模塊中直接應用相同的機制,空間復雜度將急劇上升到 N_qN_kC,其中 N_q、N_k 是查詢和鍵的數量,通常與特征圖大小 HW 具有相同的比例,帶來近似于雙二次的復雜度。
盡管 Deformable DETR [54] 已經設法通過在每個尺度上設置較少數量的 N_k = 4 的鍵來減少這種開銷,并且可以很好地作為檢測頭,但由于不可接受的信息丟失(參見附錄中的詳細比較),在骨干網絡中關注如此少的鍵效果不佳。與此同時,[3,52] 中的觀察表明,不同的查詢在視覺注意力模型中具有相似的注意力圖。因此,該研究選擇了一個更簡單的解決方案,為每個查詢共享移位鍵和值,以實現(xiàn)有效的權衡。
模型架構
該研究在 Transformer(等式 (4))中的可變形注意力替換了 vanilla MHSA,并將其與 MLP(等式 (5))相結合,以構建一個可變形的視覺 transformer 塊。在網絡架構方面, DAT 與 [7, 26, 31, 36] 共享類似的金字塔結構,廣泛適用于需要多尺度特征圖的各種視覺任務。如下圖 3 所示,形狀為 H × W × 3 的輸入圖像首先被步長為 4 的 4 × 4 非重疊卷積嵌入,然后一個歸一化層獲得
補丁嵌入。
為了構建分層特征金字塔,主干包括 4 個階段,步幅逐漸增加。在兩個連續(xù)的階段之間,有一個步長為 2 的非重疊 2×2 卷積,對特征圖進行下采樣,將空間大小減半并將特征維度加倍。
在分類任務中,該研究首先對最后階段輸出的特征圖進行歸一化,然后采用具有池化特征的線性分類器來預測對數;在對象檢測、實例分割和語義分割任務中,DAT 在集成視覺模型中扮演主干的角色,以提取多尺度特征。該研究為每個階段的特征添加一個歸一化層,然后將它們輸入到以下模塊中,例如對象檢測中的 FPN [23] 或語義分割中的解碼器。
實驗
該研究在 3 個數據集上進行了實驗,以驗證提出的 DAT 的有效性。該研究展示了在 ImageNet-1K [10] 分類、COCO 目標檢測和 ADE20K 語義分割任務上的結果。此外,該研究提供了消融研究和可視化結果,以進一步展示該方法的有效性。
ImageNet-1K 分類
ImageNet-1K [10] 數據集有 128 萬張用于訓練的圖像和 5 萬張用于驗證的圖像。研究者在訓練分割上訓練 DAT 的三個變體,并報告驗證分割上的 Top-1 準確度,并與其他 Vision Transformer 模型進行比較。
該研究在下表 2 中給出了有 300 個訓練 epoch 的結果。與其他 SOTA 視覺 Transformer 模型相比, DAT 在具有相似計算復雜性的情況下在 Top-1 精度上實現(xiàn)了顯著提高。DAT 在所有三個尺度上都優(yōu)于 Swin Transformer [26]、PVT [36]、DPT [7] 和 DeiT [33]。沒有在 Transformer 塊 [13, 14, 35] 中插入卷積,或在補丁嵌入 [6, 11, 45] 中使用重疊卷積,DAT 比 Swin Transformer [26] 實現(xiàn)了 +0.7、+0.7 和 +0.5 的增益。在 384 × 384 分辨率下進行微調時,該模型繼續(xù)比 Swin Transformer 性能好 0.3。
COCO 目標檢測
COCO 目標檢測和實例分割數據集有 118K 的訓練圖像和 5K 的驗證圖像。該研究使用 DAT 作為 RetinaNet [24]、Mask R-CNN [17] 和 Cascade Mask R-CNN [2] 框架中的主干,以評估該方法的有效性。該研究在 ImageNet-1K 數據集上對該模型進行 300 個 epoch 的預訓練,并遵循 Swin Transformer [26] 中類似的訓練策略來公平地比較該方法。該研究在 1x 和 3x 訓練計劃中報告在 RetinaNet 模型上的 DAT。如下表 3 所示,在微型和小型模型中,DAT 的性能優(yōu)于 Swin Transformer 1.1 和 1.2 mAP。
當在兩階段檢測器(例如 Mask R-CNN、Cascade Mask R-CNN)中實現(xiàn)時,DAT 模型在不同尺寸的 Swin Transformer 模型上實現(xiàn)了一致的改進,如下表 4 所示。
下表 5 給出了在驗證集上各種方法的 mIoU 分數。
消融實驗
為了驗證 DAT 模型中關鍵組件設計的有效性, 該研究進行了消融實驗,報告了基于 DAT-T 的 ImageNet-1K 分類結果。對于幾何信息開發(fā),該研究首先評估了所提可變形偏移和可變形相對位置嵌入的有效性,如下表 6 所示。
對于不同階段的可變形注意力,該研究用不同階段的可變形注意力替換了 Swin Transfomer [26] 的移位窗口注意力。如下表 7 所示,僅替換最后階段的注意力提高了 0.1,替換最后兩個階段導致性能增益為 0.7(達到 82.0 的整體準確度)。然而,在早期階段用更多可變形注意力替換會略微降低準確性。
可視化
該研究在 DAT 中可視化學習變形位置的示例,以驗證該方法的有效性。如下圖 4 所示,采樣點描繪在對象檢測框和實例分割掩碼的頂部,從中可以看到這些點已轉移到目標對象。