|論文閱讀|Swin Transformer V2: Scaling Up Capacity and Resolution

Elena
8 min readMay 20, 2024

--

上一篇聊到了Swin-Transformer v1可以達到甚至超越CNN-based的視覺模型,本篇要來聊聊Swin Transformer v2的改進。

儘管大型語言模型已在多個 NLP 任務中展現卓越的效能,但同等規模的視覺模型卻相對發展緩慢。 為了成功訓練大型視覺通用模型,需解決幾個問題如下所列:

  • Unstable Training 訓練不穩定:當模型容量增大時,層間的激活值差異變大,導致訓練過程不穩定。
  • Resolution Gaps 解析度差異:在預訓練和fine-tune階段之間解析度的不一致,使得模型難以從低解析度的預訓練有效過渡到高解析度的下游任務。
  • GPU 記憶體消耗的問題: 在模型擴展到更大的容量和解析度時面臨的額外技術挑戰
  • Hunger on labelled data 數據依賴:模型對大量標籤數據的依賴性高,這在實際應用中可能受到限制。

因此,Swin Transformer v2 為了解決上述問題,提出以下貢獻:

  1. Residual Post-Normalization 殘差後歸一化:這種新的歸一化配置將層歸一化(Layer Normalization, LN)從每個殘差單元的開始移動到末端,減少了不同層間激活值的累積和波動。
  2. Cosine Attention 余弦注意力機制:引入余弦相似度來計算注意力權重,取代了傳統的點積注意力機制。這種方法對於模型輸入的振幅不敏感,有助於減少極端注意力值的發生。
  3. Log-Spaced Continuous Position Bias 對數空間連續位置偏差:透過對數轉換的位置編碼來進行連續的位置偏差計算,使模型能夠更好地處理在預訓練和微調階段間解析度的變化。
  4. 自監督預訓練方法(SimMIM):採用自監督學習策略減少對大量有標籤數據的依賴,通過模擬遮蔽圖像建模(SimMIM)來預訓練模型,降低對標籤的需求。

模型架構

以下架構圖比較了v1及v2的差別:

Swin Transformer v2架構圖
  1. 殘差後歸一化(Res-Post-Norm)
  • 在 V1 中,Layer Norm 是位於每個殘差單元的開始。
  • 在 V2 中,將 Layer Norm 移至每個殘差單元的末端。這種配置改變有助於緩解在模型深層中激活值累積增大的問題,從而改善訓練的穩定性。

2. 余弦注意力機制(Cosine Attention)

  • V1 中採用的是標準的點積自注意力(dot product attention)。
  • V2 中引入了一種基於余弦相似度的注意力機制,這樣做可以使注意力分數的計算不再依賴於輸入塊的振幅,減少極端值的發生,從而有助於穩定訓練過程。

3. 對數間隔的連續相對位置偏差(Log-CPB)

  • V1 使用的是參數化的位置偏差。
  • V2 引入了一種新的連續位置偏差方法,它基於對數空間的座標輸入,這使得模型在面對不同窗口大小的輸入時能更靈活地調整,特別是從預訓練到微調階段的過渡。

以下將針對以上的貢獻再做詳細介紹:

Scaling Up Model Capacity

Swin Transformer V2 透過提高模型容量來增強其處理複雜視覺任務的能力。當模型的容量增大時,層間的激活值差異會變大,導致訓練不穩定。為解決這個問題,提出了以下方法:

Post normalization

引入了殘差後歸一化技術,將層歸一化放在每個殘差塊的輸出端。

Scaled cosine attention

在原始的自注意力計算中,是用傳統的點積注意力來計算。這可能會遇到一個問題,即學習到的注意力圖中某些區塊和頭部可能會由少數幾個像素對所主導。為了解決這一問題,研究者提出了 Scaled cosine attention 方法。這種方法通過計算像素對 𝑖和 𝑗之間的餘弦相似度來確定注意力分數,並進行尺度調整。具體的數學表達式如下:

  • 其中,cos(qi,kj) 表示查詢(query)向量 qi 和鍵(key)向量 kj 之間的餘弦相似度。
  • 𝐵𝑖𝑗是介於像素 ij 之間的相對位置偏差,是一個可學習的參數。

尺度化餘弦注意力的特點在於餘弦函數是自然歸一化的,即輸出範圍是固定的(從 -1 到 1)。這可避免訓練過程中由於極端的注意力值而導致的不穩定性問題。

Scaling Up Window Resolution

Swin Transformer V2 引入了一種新的方法來處理在不同視窗解析度之間轉移模型時的位置偏差問題,稱為對數空間連續位置偏差(Log-Spaced Continuous Position Bias)。這種方法的主要優勢是能夠平滑地適應不同的視窗大小,特別是在從預訓練到微調階段過渡時。

Continuous relative position bias

與傳統的參數化位置偏差不同,連續位置偏差使用一個小型的meta network (G),通常是一個包含 ReLU 激活函數的兩層 MLP(多層感知器),來生成相對坐標的偏差值,如下:

這允許模型針對任意的視窗大小生成位置偏差,從而更好地適應變化的輸入尺寸。

Log-spaced coordinates

在處理大幅度變化的視窗大小時,尤其是需要外推(extrapolate)坐標範圍時,直接使用線性空間(linear-spaced)坐標可能導致效率低下和精度問題。為了解決這一問題,研究者們提出使用對數空間(log-spaced)坐標:

Δx , Δy和Δ𝑥^, Δ𝑦^ 分別代表線性空間坐標和對數空間坐標。

這樣的坐標轉換能夠大幅減少在從低解析度到高解析度轉換時所需的外推比率,從而提高位置偏差的準確性和模型的適應性。

Self-Supervised Pre-training

由於大型模型通常對資料有很高的需求,傳統上需要依賴大規模的標記資料集來進行有效訓練,如 JFT-3B。然而,這不僅成本高昂,也限制了模型應用的靈活性。為了解決這一問題,Swin Transformer V2 採用了一種稱為 SimMIM 的自監督學習策略。

Implementation to Save GPU Memory

當訓練大規模和高解析度的模型時,GPU 記憶體消耗往往成為一個巨大的挑戰。本篇論文實施了幾種優化措施來降低記憶體需求,同時保持訓練速度的效率。以下是這些措施的詳細介紹:

Zero-Redundancy Optimizer (ZeRO)

ZeRO 是一種數據並行的優化技術,透過將模型參數和對應的優化狀態分割並分佈到多個 GPU 上,顯著減少了單個 GPU 的記憶體消耗。

Activation check-pointing

這是另一種用於降低 GPU 記憶體消耗的技術,特別是在處理高解析度影像和大視窗時。通過臨時存儲前向傳播中的中間激活值,然後在需要時重新計算它們,而不是一直保持它們在記憶體中。這種方法可以顯著減少記憶體的使用,但可能使訓練速度慢下來約 30%。

Sequential self-attention computation

當在非常大的解析度上訓練大型模型時,自注意力模塊可能會成為瓶頸。為此,研究團隊採用了Sequential self-attention computation,這與傳統的批量計算方法不同。這種優化主要應用於模型的前兩階段,對整體訓練速度的影響較小。

實驗結果

ImageNet-1K 分類結果比較

Swin Transformer V2在這些任務上相比其他大型視覺模型(如 ViT-G, V-MoE 等)表現出更高的精確度。特別是 SwinV2-G,在 ImageNet-1K V2 上達到了非常高的 top-1 正確率84.00%。

COCO 物件偵測和實例分割結果比較

此表格展示了在 COCO 資料集上進行物件偵測和實例分割任務的結果。SwinV2-G 在這些測試中顯示了優越的表現,特別是在更大的視窗尺寸和輸入解析度下,例如 1536x1536。

ADE20K 語義分割結果比較

SwinV2-G 特別在使用更大視窗尺寸的設置下,顯示了卓越的 mIoU成績,這再次證明了其在大規模視覺模型中處理大視窗和高解析度任務的能力。

結論

Swin Transformer V2 將模型擴展到 30 億參數,並使其能夠處理高達1536x1536 解析度的影像。包括使用 res-post-norm 和 scaled cosine attention來使模型更容易在容量上進行擴展,以及使用對數空間連續位置偏差方法,這讓模型能夠更有效地在不同窗口解析度間轉換。Swin Transformer V2 不僅在影像分類任務上表現出色,還在更複雜的視覺任務(如物體檢測、實例分割和語義分割)中達到了當前最佳表現。

--

--

Elena

AI Engineer | Data | GIS | 金魚腦女子,希望可以透過文字記錄以喚起被遺忘的記憶。