機器之心 - 只需2040張圖片，訓練視覺Transformer：南大吳建鑫團隊提出IDMM－鑽石舞台

機器之心報道

機器之心編輯部

預訓練模型是否也能只需要很少的樣本就能完成任務？

Transformer 架構的強大性能不僅在 NLP 領域成為了主流，也在代替卷積神經網絡 CNN，成為視覺識別的一個方向（ViT）。一些 transformer 模型已經取得了有競爭力的結果，但因為缺乏典型的卷積歸納偏差使得它們比普通 CNN 需要更多的訓練數據。

在南京大學吳建鑫團隊近日提交的一篇論文中，研究者研究了如何使用有限數據訓練 ViT，由於可以捕獲特徵對齊和實例相似性，新方法展現了優勢。

在各種 ViT 主幹下的 7 個小型數據集上從頭開始訓練之後，該研究獲得了最先進的結果。研究者還討論了小數據集的遷移能力，發現從小數據集學習的表示甚至可以改善大規模 ImageNet 的訓練。

論文鏈接：https://arxiv.org/abs/2201.10728

Transformer 近來已廣泛用於視覺識別，替代了卷積神經網絡（CNN）。視覺 Transformer（ViT）是一種直接繼承自自然語言處理的架構，但適用於以原始圖像 patch 作為輸入的圖像分類。ViT 及其變體獲得了可與 CNN 媲美的結果，但卻需要更多的訓練數據。

例如，在 ImageNet（128 萬張圖像）上訓練時，ViT 的性能比具有相似容量的 ResNet 差。一個可能的原因可能是 ViT 缺乏 CNN 架構中固有的某些理想屬性，這使得 CNN 非常適合解決視覺任務，例如局部性、平移不變性和層次結構。因此，ViT 通常需要比 CNN 更大量的數據進行訓練。

為了緩解這個問題，很多工作都嘗試將卷積引入 ViT。這些架構具有兩種範式的優點，注意力層對遠程依賴進行建模，而卷積則強調圖像的局部屬性。實驗結果表明，這些在 ImageNet 上訓練的 ViT 在該數據集上優於類似大小的 ResNet。

然而，ImageNet 仍是一個大規模數據集，當在小數據集（例如 2040 張圖像）上訓練時，這些網絡的行為仍不清楚。該研究從數據、計算和靈活性的角度進行了分析，證明不能總是依賴如此大規模的數據集。

該論文探究了如何使用有限的數據從頭開始訓練 ViT。

該研究首先執行自監督預訓練，然後對同一目標數據集進行監督微調，與（Cao et al.,2021）等人的方法類似。該研究重點關注自監督的預訓練階段，方法基於參數實例判別（parametric instance discrimination）。

從理論的角度分析，參數實例判別不僅可以捕獲 positive pair 之間的特徵對齊，還可以找出實例之間的潛在相似性，這要歸功於最終可學習的全連接層 W。實驗結果進一步驗證了研究者的分析，該研究的方法比非參數方法實現了更好的性能。

眾所周知，在大規模數據集上，高維全連接層的實例判別會受到大量 GPU 計算、內存過載和收斂速度慢的影響。由於該研究專注於小型數據集，因此不需要針對大型數據集的複雜策略。相反，該研究對小數據設置採用小分辨率、多裁剪和 CutMix 的方法，並且研究者還從理論和實驗的角度對其進行了分析。

該研究將這種方法命名為帶有 Multi-crop 和 CutMix 的實例判別（Instance Discrimination with Multi-crop and CutMix，IDMM）。實驗結果表明，在 7 個小型數據集上從頭開始訓練多種 ViT 主幹網絡，實現了 SOTA 結果。例如，該研究在 flowers 數據集（含 2040 張圖像）上從頭開始訓練模型，結果達到 96.7% 的準確率，這表明使用小數據集訓練 ViT 是完全可以的。

此外，該研究首先分析了小數據集的遷移能力，並發現：即使在小型數據集上進行預訓練，ViT 也具有良好的遷移能力，甚至可以促進對大規模數據集（例如 ImageNet）的訓練。（Liu et al.,2021）也研究了使用小型數據集訓練 ViT，但他們專注於微調階段，而南大的這項研究專注於預訓練階段，並且南大的方法取得了更好的結果，其中在 flowers 數據集上獲得的最佳準確率為 56.3%。

總的來說，該研究的主要貢獻包括：

提出了用於自監督 ViT 訓練的 IDMM，即使在 7 個小型數據集上對各種 ViT 主幹網絡從頭開始訓練也能獲得 SOTA 結果；

對於在處理小數據時為什麼應該更注重參數實例判別，該研究給出了損失角度的理論分析。此外，展示了 CutMix 等策略應如何從梯度的角度緩解不頻繁更新的問題；

該研究通過實驗表明，投影 MLP 頭對於非參數方法（例如，SimCLR）是必不可少的，但對於參數實例判別而言卻不是，這要歸功於實例判別中最終可學習的全連接層；

分析了小數據集的遷移能力，發現即使在小數據集上進行預訓練，ViT 也具有良好的遷移能力。

具體方法如上圖所示，輸入圖像 x_i (i = 1, · · · , N) 被傳至神經網絡 f(·) 中，得到輸出表徵 z_i = f(x_i) ∈ R^d ，其中 N 表示實例總數。然後使用全連接層 W 進行分類，類數等於用於參數實例判別的訓練圖像總數。隨後將 w_j ∈ R^d 表示第 j 類的權重，W = [w_1| . . . |w_N ] ∈ R^(d×N) 包含所有 n 個類的權重。

因此有

，其中第 j 類的輸出是

。最後，將 O^(i) 被傳到 softmax 層以獲得有效的概率分布 P^(i)。

損失函數為：

其中上標 i 代表對實例求和，下標 c 代表對類求和。

實驗結果

研究者在實驗中使用了 7 個小型數據集，具體如下表 1 所示。

為什麼要從頭開始訓練

為什麼要直接在目標數據集上從頭開始訓練呢？研究者從以下三個方面進行了解答

首先是數據。當前的 ViT 模型往往是在大規模數據集（如 ImageNet 或更大）上預訓練，然後在各種下游任務中微調。典型卷積歸納偏差的缺失也使得這些模型比常見 CNN 更需要數據。因此，探究是否可以針對某個可用圖像有限的任務從頭開始訓練 ViT 模型非常重要；

其次是計算。大規模數據集、大量 epoch 和複雜的骨幹網絡，這些組合在一起意味着 ViT 訓練的計算成本非常高。這種現象導致 ViT 模型成為少數機構研究人員才能使用的「特權」；

最後是靈活性。下游微調範式之前的預訓練有時會變得非常麻煩。例如，我們可能需要為同一個任務訓練 10 個不同的模型，並將它們部署到不同的硬件平台，但在大規模數據集上預訓練 10 個模型是不切實際的。

從頭開始訓練的結果

研究者提供了從頭開始訓練的結果。

在下表 2 和下圖 4 中，研究者首先將 IDMM 與用於 CNN 和 ViT 的流行 SSL 方法進行了比較。公平起見，所有方法都預訓練了 800 個 epoch，然後微調了 200 個 epoch。可以看出，即使從頭開始訓練，SSL 預訓練也很有用，並且所有 SSL 方法的性能都比隨機初始化好。

表 2：不同預訓練方法之間的比較。

圖 4：在花數據集上與不同 SSL 方法進行比較。所有模型都經過了同樣 epoch 數的預訓練和微調。

接着，研究者微調模型更長的 epoch，以獲得更好的結果。具體來講，他們使用 IDMM 初始化權重，先在 224x224 分辨率下微調了 800 個 epoch，然後在 448x448 分辨率下微調了 100 個 epoch。如下表 3 所示，當在這 7 個數據集上從頭訓練所有這些 ViT 模型時，IDMM 實現了 SOTA 結果。

小型數據集上的遷移能力

在下表 5 中，研究者評估了在不同數據集上預訓練模型的遷移準確率。可以看到，即使在小型數據集上進行預訓練，ViT 也具有良好的遷移能力。與 SimCLR 和 SupCon 相比，IDMM 在所有這些數據集上的遷移準確率也更高。即使預訓練數據集和目標數據集不在同一個域中，研究者也可以獲得非常好的結果。