Swin Transformer升級版來了！30億參數，刷榜多項視覺任務，微軟亞研原班人馬打造

編輯：杜偉、陳萍

微軟亞洲研究院升級了 Swin Transformer，新版本具有 30 億個參數，可以訓練解析度高達 1,536×1,536 的圖像，並在四個具有代表性的基準上刷新紀錄。

在不久之前公佈的 ICCV 2021 論文獎項中，來自微軟亞洲研究院的研究者憑藉論文《Swin Transformer: Hierarchical Vision Transformer using Shifted Windows》斬獲 ICCV 2021 馬爾獎（最佳論文）。這篇論文的作者主要包括中國科學技術大學的劉澤、西安交通大學的林宇桐、微軟的曹越和胡瀚等人。該研究提出了一種新的 vision Transformer，即 Swin Transformer，它可以作為計算機視覺的通用骨幹。

相比之前的 ViT 模型，Swin Transformer 做出了以下兩點改進：其一，引入 CNN 中常用的層次化構建方式構建分層 Transformer；其二，引入局部性（locality）思想，對無重合的窗口區域內進行自注意力計算。在 Swin Transformer 論文公開沒多久之後，微軟官方也在 GitHub 上開源了程式碼和預訓練模型，涵蓋圖像分類、目標檢測以及語義分割任務。

近日，該團隊又提出一種升級版 SwinTransformer V2。

論文地址：https://arxiv.org/pdf/2111.09883.pdf

通常來講，Transformer 適用於擴展視覺模型，但它還沒有像 NLP 語言模型那樣得到廣泛的探索，部分原因是因為在訓練和應用方面存在以下困難：

視覺模型經常面臨擴展不穩定問題；
許多下游視覺任務需要高解析度的圖像或窗口，目前尚不清楚如何有效地將在低解析度下預訓練的模型轉換為更高解析度的模型；
當圖像解析度較高時，GPU 記憶體消耗也是一個問題。

為了解決上述問題，該團隊將 SwinTransformer 作為基線提出了幾種改進技術，具體表現在：

提出後歸一化（post normalization）技術和縮放餘弦注意力（scaled cosine attention）方法，來提高大型視覺模型的穩定性；
提出 log-spaced 連續位置偏差技術，可有效地將在低解析度圖像和窗口中預訓練的模型遷移到更高解析度對應模型。

此外，該研究還介紹了關鍵實現細節，這些細節可顯著節省 GPU 記憶體消耗，使得常規 GPU 訓練大型視覺模型成為可能。使用這些技術和自監督預訓練，該團隊訓練了一個具有 30 億參數的 Swin Transformer 模型，並將其有效地遷移到高解析度圖像或窗口的各種視覺任務中，在各種基準上實現了 SOTA 性能。

通過擴展容量和解析度，Swin Transformer V2 在四個具有代表性的基準上刷新紀錄：在 ImageNet-V2 圖像分類任務上 top-1 準確率為 84.0%，COCO 目標檢測任務為 63.1 / 54.4 box / mask mAP，ADE20K 語義分割為 59.9 mIoU，Kinetics-400 視訊動作分類的 top-1 準確率為 86.8%。

部分刷榜截圖。圖源：https://paperswithcode.com/sota

Swin Transformer V2

研究者觀察到 Swin Transformer 在擴展模型容量和窗口解析度時存在以下兩個問題。

其一，擴展模型容量的不穩定問題。如下圖 2 所示，當我們將原始 Swin Transformer 模型從小到大擴展時，更深層的激活值急劇增加。具有最高和最低振幅的層之間的偏差達到了 10^4 的極值。

當我們進一步將其擴展到一個巨大的規模（6.58 億參數）時，Swin Transformer 無法完成訓練，如下圖 3 所示。

其二，跨窗口解析度遷移模型時性能下降。如下表 1 第一行所示，當我們通過雙三次插值方法，在更大的圖像解析度和窗口大小直接測試預訓練 ImageNet-1K 模型（解析度 256 × 256，窗口大小 8 × 8）的準確率時，發現準確率顯著下降。這可能值得去重新檢查原始 Swin Transformer 中的相對位置偏差方法。

擴展模型容量

在本節內容中，研究者介紹了上述兩個問題的解決方法，包括如下：

使用後歸一化和縮放餘弦注意力解決不穩定性問題；
使用 log-spaced 連續位置偏差方法解決跨窗口解析度遷移問題。

一方面，為了緩解擴展模型容量時的不穩定問題，研究者使用了後歸一化方法，如下圖 1 所示。在這種方法中，每個殘差塊的輸出在合併回主分支之前被歸一化，並且當層數越深時，主分支的振幅不會積聚。

又如上圖 2 所示，這種方法的激活幅度變得比原始預歸一化配置溫和得多。在研究者最大的模型訓練中，他們每 6 個 Transformer 塊都會在主分支上額外引入一個層歸一化單元，以進一步穩定訓練和振幅。

另一方面，在原始的自注意力計算中，像素對的相似項被計算為查詢向量和關鍵向量的點積。研究者發現將這種方法用於大型視覺模型時，特別是在 post-norm 配置中，一些塊和頭部學得的注意力圖經常由幾個像素對主導。為了緩解這個問題，他們提出了一種縮放餘弦注意力方法，它通過縮放餘弦函數計算像素對 i 和 j 的注意力對數：

擴展窗口解析度

在本節中，研究者介紹了一種 log-spaced 連續位置偏差方法，以使得相對位置偏差可以在窗口解析度之間平滑地遷移。連續位置偏差方法不是直接最佳化參數化偏差，而是在相對座標上引入一個小的元（meta）網路：

元網路

為任意相對座標生成偏差值，因此可以自然地遷移到具有任意變化窗口大小的微調任務。對於推理任務，每個相對位置的偏差值可以預先計算並儲存為模型參數，這樣在推理時與原始參數化偏差方法一樣方便。

當在變化很大的窗口大小之間遷移時，將有很大一部分相對座標範圍需要外推。為了緩解這個問題，研究者提出使用 log-spaced 座標替代原始 linear-spaced 座標：

節省 GPU 記憶體的實現

另一個問題在於當容量和解析度都很大時，常規實現的 GPU 記憶體消耗難以承受。為了解決記憶體問題，研究者採用以下幾種實現方法：

零冗餘最佳化器（Zero-Redundancy Optimizer, ZeRO）
激活檢查點（Activation check-pointing）
順序自注意力計算（Sequential self-attention computation）

通過這些實現，研究者成功地使用 Nvidia A100-40G GPU 訓練了一個 3B（30 億參數）模型，既可以用於輸入圖像解析度為 1,536×1,536 的 COCO 目標檢測，也可用於輸入解析度為 320 × 320 × 8 的 Kinetics-400 動作分類。

實驗結果

該團隊在 ImageNet-1K 圖像分類（V1 和 V2）、COCO 目標檢測和 ADE20K 語義分割進行了實驗。此外，對於 30 億參數模型實驗，該研究還報告了 Swin Transformer V2 在 Kinetics400 視訊動作識別上的準確率。

SwinV2-G 實驗設置：預訓練採用 192×192 圖像解析度，以節省訓練成本，實驗採用 2-step 預訓練方法：首先，在 ImageNet-22K-ext 資料集上使用自監督方法對模型進行 20epoch 的預訓練。其次，在 ImageNet-1K V1 和 V2 分類任務上，繼續將模型進行 30epoch 預訓練。

ImageNet-1K 圖像分類結果：表 2 將 SwinV2-G 模型與之前在 ImageNet-1K V1 和 V2 分類任務上的最大 / 最佳視覺模型進行了比較。SwinV2-G 是之前所有密集（dense）視覺模型中最大的。它在 ImageNet V2 基準測試中達到了 84.0% 的 top-1 準確率，比之前最好的 ViT-G (83.3%) 高 0.7%。但是，SwinV2-G 在 ImageNet-1K V1 上的準確率比 CoAtNet-7 略低（90.17% 比 90.88%）。

COCO 目標檢測結果：表 3 將 SwinV2-G 模型與之前在 COCO 目標檢測和實例分割任務上取得最佳性能模型進行了比較。SwinV2-G 在 COCO test-dev 上實現了 63.1/54.4 box/max AP，比 SoftTeacher(61.3/53.0) 提高了 + 1.8/1.4。這表明擴展視覺模型有利於目標檢測任務中的密集視覺識別任務。

ADE20K 語義分割結果：下表 4 將 SwinV2-G 模型與之前在 ADE20K 語義分割基準上的 SOTA 結果進行了比較。Swin-V2-G 在 ADE20K val 集上實現了 59.9 mIoU，比之前的 SOTA 結果（BEiT）58.4 高了 1.5。這表明擴展視覺模型有益於像素級視覺識別任務。在測試時使用更大的窗口大小還可以帶來 +0.2 的增益，這可能歸功於有效的 Log-spaced CPB 方法。

Kinetics-400 視訊動作分類結果：下表 5 將 SwinV2-G 模型與之前在 Kinetics-400 動作分類基準上的 SOTA 結果進行了比較。可以看到，Video-SwinV2-G 實現了 86.8% 的 top-1 準確率，比之前的 SOTA （TokenLearner）85.4% 高出 +1.4%。這表明擴展視覺模型也有益於視訊識別任務。在這種場景下，在測試時使用更大的窗口大小也可以帶來額外增益 ( +0.2% )，這也要歸功於有效的 Log-spaced CPB 方法。

post-norm 和縮放餘弦注意力的消融實驗：下表 6 展示了 post-norm 和縮放餘弦注意力方法應用於原始 Swin Transformer 方法的性能表現。可以看到，這兩種方法都提高了 Swin-Tiny、Swin-Small 和 Swin-Base size 的準確率，整體提升分別為 +0.2%、+0.4% 和 +0.5%，表明它們對更大的模型更有益。