只需2040張圖片，訓(xùn)練視覺(jué)Transformer：南大吳建鑫團(tuán)隊(duì)提出IDMM

來(lái)源：機(jī)器之心
本文共3000字，建議閱讀5分鐘
本文研究了如何使用有限數(shù)據(jù)訓(xùn)練 ViT，由于可以捕獲特征對(duì)齊和實(shí)例相似性，新方法展現(xiàn)了優(yōu)勢(shì)。

預(yù)訓(xùn)練模型是否也能只需要很少的樣本就能完成任務(wù)？

Transformer 架構(gòu)的強(qiáng)大性能不僅在 NLP 領(lǐng)域成為了主流，也在代替卷積神經(jīng)網(wǎng)絡(luò) CNN，成為視覺(jué)識(shí)別的一個(gè)方向（ViT）。一些 transformer 模型已經(jīng)取得了有競(jìng)爭(zhēng)力的結(jié)果，但因?yàn)槿狈Φ湫偷木矸e歸納偏差使得它們比普通 CNN 需要更多的訓(xùn)練數(shù)據(jù)。

在南京大學(xué)吳建鑫團(tuán)隊(duì)近日提交的一篇論文中，研究者研究了如何使用有限數(shù)據(jù)訓(xùn)練 ViT，由于可以捕獲特征對(duì)齊和實(shí)例相似性，新方法展現(xiàn)了優(yōu)勢(shì)。

在各種 ViT 主干下的 7 個(gè)小型數(shù)據(jù)集上從頭開(kāi)始訓(xùn)練之后，該研究獲得了最先進(jìn)的結(jié)果。研究者還討論了小數(shù)據(jù)集的遷移能力，發(fā)現(xiàn)從小數(shù)據(jù)集學(xué)習(xí)的表示甚至可以改善大規(guī)模 ImageNet 的訓(xùn)練。

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

Transformer 近來(lái)已廣泛用于視覺(jué)識(shí)別，替代了卷積神經(jīng)網(wǎng)絡(luò)（CNN）。視覺(jué) Transformer（ViT）是一種直接繼承自自然語(yǔ)言處理的架構(gòu)，但適用于以原始圖像 patch 作為輸入的圖像分類。ViT 及其變體獲得了可與 CNN 媲美的結(jié)果，但卻需要更多的訓(xùn)練數(shù)據(jù)。

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

為了緩解這個(gè)問(wèn)題，很多工作都嘗試將卷積引入 ViT。這些架構(gòu)具有兩種范式的優(yōu)點(diǎn)，注意力層對(duì)遠(yuǎn)程依賴進(jìn)行建模，而卷積則強(qiáng)調(diào)圖像的局部屬性。實(shí)驗(yàn)結(jié)果表明，這些在 ImageNet 上訓(xùn)練的 ViT 在該數(shù)據(jù)集上優(yōu)于類似大小的 ResNet。

然而，ImageNet 仍是一個(gè)大規(guī)模數(shù)據(jù)集，當(dāng)在小數(shù)據(jù)集（例如 2040 張圖像）上訓(xùn)練時(shí)，這些網(wǎng)絡(luò)的行為仍不清楚。該研究從數(shù)據(jù)、計(jì)算和靈活性的角度進(jìn)行了分析，證明不能總是依賴如此大規(guī)模的數(shù)據(jù)集。

該論文探究了如何使用有限的數(shù)據(jù)從頭開(kāi)始訓(xùn)練 ViT。

該研究首先執(zhí)行自監(jiān)督預(yù)訓(xùn)練，然后對(duì)同一目標(biāo)數(shù)據(jù)集進(jìn)行監(jiān)督微調(diào)，與（Cao et al.,2021）等人的方法類似。該研究重點(diǎn)關(guān)注自監(jiān)督的預(yù)訓(xùn)練階段，方法基于參數(shù)實(shí)例判別（parametric instance discrimination）。

從理論的角度分析，參數(shù)實(shí)例判別不僅可以捕獲 positive pair 之間的特征對(duì)齊，還可以找出實(shí)例之間的潛在相似性，這要?dú)w功于最終可學(xué)習(xí)的全連接層 W。實(shí)驗(yàn)結(jié)果進(jìn)一步驗(yàn)證了研究者的分析，該研究的方法比非參數(shù)方法實(shí)現(xiàn)了更好的性能。

眾所周知，在大規(guī)模數(shù)據(jù)集上，高維全連接層的實(shí)例判別會(huì)受到大量 GPU 計(jì)算、內(nèi)存過(guò)載和收斂速度慢的影響。由于該研究專注于小型數(shù)據(jù)集，因此不需要針對(duì)大型數(shù)據(jù)集的復(fù)雜策略。相反，該研究對(duì)小數(shù)據(jù)設(shè)置采用小分辨率、多裁剪和 CutMix 的方法，并且研究者還從理論和實(shí)驗(yàn)的角度對(duì)其進(jìn)行了分析。

該研究將這種方法命名為帶有 Multi-crop 和 CutMix 的實(shí)例判別（Instance Discrimination with Multi-crop and CutMix，IDMM）。實(shí)驗(yàn)結(jié)果表明，在 7 個(gè)小型數(shù)據(jù)集上從頭開(kāi)始訓(xùn)練多種 ViT 主干網(wǎng)絡(luò)，實(shí)現(xiàn)了 SOTA 結(jié)果。例如，該研究在 flowers 數(shù)據(jù)集（含 2040 張圖像）上從頭開(kāi)始訓(xùn)練模型，結(jié)果達(dá)到 96.7% 的準(zhǔn)確率，這表明使用小數(shù)據(jù)集訓(xùn)練 ViT 是完全可以的。

此外，該研究首先分析了小數(shù)據(jù)集的遷移能力，并發(fā)現(xiàn)：即使在小型數(shù)據(jù)集上進(jìn)行預(yù)訓(xùn)練，ViT 也具有良好的遷移能力，甚至可以促進(jìn)對(duì)大規(guī)模數(shù)據(jù)集（例如 ImageNet）的訓(xùn)練。（Liu et al.,2021）也研究了使用小型數(shù)據(jù)集訓(xùn)練 ViT，但他們專注于微調(diào)階段，而南大的這項(xiàng)研究專注于預(yù)訓(xùn)練階段，并且南大的方法取得了更好的結(jié)果，其中在 flowers 數(shù)據(jù)集上獲得的最佳準(zhǔn)確率為 56.3%。

總的來(lái)說(shuō)，該研究的主要貢獻(xiàn)包括：

• 提出了用于自監(jiān)督 ViT 訓(xùn)練的 IDMM，即使在 7 個(gè)小型數(shù)據(jù)集上對(duì)各種 ViT 主干網(wǎng)絡(luò)從頭開(kāi)始訓(xùn)練也能獲得 SOTA 結(jié)果；
• 對(duì)于在處理小數(shù)據(jù)時(shí)為什么應(yīng)該更注重參數(shù)實(shí)例判別，該研究給出了損失角度的理論分析。此外，展示了 CutMix 等策略應(yīng)如何從梯度的角度緩解不頻繁更新的問(wèn)題；
• 該研究通過(guò)實(shí)驗(yàn)表明，投影 MLP 頭對(duì)于非參數(shù)方法（例如，SimCLR）是必不可少的，但對(duì)于參數(shù)實(shí)例判別而言卻不是，這要?dú)w功于實(shí)例判別中最終可學(xué)習(xí)的全連接層；
• 分析了小數(shù)據(jù)集的遷移能力，發(fā)現(xiàn)即使在小數(shù)據(jù)集上進(jìn)行預(yù)訓(xùn)練，ViT 也具有良好的遷移能力。

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

因此有，其中第 j 類的輸出是。最后，將 O^(i) 被傳到 softmax 層以獲得有效的概率分布 P^(i)。

損失函數(shù)為：

其中上標(biāo) i 代表對(duì)實(shí)例求和，下標(biāo) c 代表對(duì)類求和。

實(shí)驗(yàn)結(jié)果

研究者在實(shí)驗(yàn)中使用了 7 個(gè)小型數(shù)據(jù)集，具體如下表 1 所示。

為什么要從頭開(kāi)始訓(xùn)練

為什么要直接在目標(biāo)數(shù)據(jù)集上從頭開(kāi)始訓(xùn)練呢？研究者從以下三個(gè)方面進(jìn)行了解答

• 首先是數(shù)據(jù)。當(dāng)前的 ViT 模型往往是在大規(guī)模數(shù)據(jù)集（如 ImageNet 或更大）上預(yù)訓(xùn)練，然后在各種下游任務(wù)中微調(diào)。典型卷積歸納偏差的缺失也使得這些模型比常見(jiàn) CNN 更需要數(shù)據(jù)。因此，探究是否可以針對(duì)某個(gè)可用圖像有限的任務(wù)從頭開(kāi)始訓(xùn)練 ViT 模型非常重要；
• 其次是計(jì)算。大規(guī)模數(shù)據(jù)集、大量 epoch 和復(fù)雜的骨干網(wǎng)絡(luò)，這些組合在一起意味著 ViT 訓(xùn)練的計(jì)算成本非常高。這種現(xiàn)象導(dǎo)致 ViT 模型成為少數(shù)機(jī)構(gòu)研究人員才能使用的「特權(quán)」；
• 最后是靈活性。下游微調(diào)范式之前的預(yù)訓(xùn)練有時(shí)會(huì)變得非常麻煩。例如，我們可能需要為同一個(gè)任務(wù)訓(xùn)練 10 個(gè)不同的模型，并將它們部署到不同的硬件平臺(tái)，但在大規(guī)模數(shù)據(jù)集上預(yù)訓(xùn)練 10 個(gè)模型是不切實(shí)際的。

從頭開(kāi)始訓(xùn)練的結(jié)果

研究者提供了從頭開(kāi)始訓(xùn)練的結(jié)果。

在下表 2 和下圖 4 中，研究者首先將 IDMM 與用于 CNN 和 ViT 的流行 SSL 方法進(jìn)行了比較。公平起見(jiàn)，所有方法都預(yù)訓(xùn)練了 800 個(gè) epoch，然后微調(diào)了 200 個(gè) epoch?？梢钥闯觯词箯念^開(kāi)始訓(xùn)練，SSL 預(yù)訓(xùn)練也很有用，并且所有 SSL 方法的性能都比隨機(jī)初始化好。

表 2：不同預(yù)訓(xùn)練方法之間的比較。

圖 4：在花數(shù)據(jù)集上與不同 SSL 方法進(jìn)行比較。所有模型都經(jīng)過(guò)了同樣 epoch 數(shù)的預(yù)訓(xùn)練和微調(diào)。

接著，研究者微調(diào)模型更長(zhǎng)的 epoch，以獲得更好的結(jié)果。具體來(lái)講，他們使用 IDMM 初始化權(quán)重，先在 224x224 分辨率下微調(diào)了 800 個(gè) epoch，然后在 448x448 分辨率下微調(diào)了 100 個(gè) epoch。如下表 3 所示，當(dāng)在這 7 個(gè)數(shù)據(jù)集上從頭訓(xùn)練所有這些 ViT 模型時(shí)，IDMM 實(shí)現(xiàn)了 SOTA 結(jié)果。

小型數(shù)據(jù)集上的遷移能力

在下表 5 中，研究者評(píng)估了在不同數(shù)據(jù)集上預(yù)訓(xùn)練模型的遷移準(zhǔn)確率。可以看到，即使在小型數(shù)據(jù)集上進(jìn)行預(yù)訓(xùn)練，ViT 也具有良好的遷移能力。與 SimCLR 和 SupCon 相比，IDMM 在所有這些數(shù)據(jù)集上的遷移準(zhǔn)確率也更高。即使預(yù)訓(xùn)練數(shù)據(jù)集和目標(biāo)數(shù)據(jù)集不在同一個(gè)域中，研究者也可以獲得非常好的結(jié)果。

在下表 6 中，研究者將 IDMM 與各種 SSL 方法以及不同主干下的監(jiān)督基線方法進(jìn)行了比較。結(jié)果顯示，IDMM 比這些競(jìng)品方法有很大的優(yōu)勢(shì)，在 SIN-10k 上學(xué)到的表征可以在遷移到其他數(shù)據(jù)集時(shí)作為一個(gè)很好的初始化。

此外，研究者還探究了當(dāng)在 ImageNet 上訓(xùn)練時(shí)，是否可以從 10,000 張圖像的預(yù)訓(xùn)練中受益。從下表 7 可以看到，使用從 10,000 張圖像中學(xué)到的表征作為初始化，可以大大加快訓(xùn)練過(guò)程，最終在 ImageNet 上實(shí)現(xiàn)了更高的準(zhǔn)確率（提升約 1%）。

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

最后，研究者對(duì)不同的組件進(jìn)行了消融實(shí)驗(yàn)。所有實(shí)驗(yàn)使用 PyTorch 完成，并且在 ImageNet 實(shí)驗(yàn)中使用了 Titan Xp GPU，在小型數(shù)據(jù)集實(shí)驗(yàn)中使用了 Tesla K80.

如下表 8 所示，研究者提出的所有策略都很有用，策略的結(jié)合使用更是實(shí)現(xiàn)了 SOTA 結(jié)果。

在下表 9 中，所有方法在 SIN-10k 上預(yù)訓(xùn)練了 800 個(gè) epoch，然后在遷移到目標(biāo)數(shù)據(jù)集時(shí)微調(diào)了 200 個(gè) epoch。?

該工作中，研究者對(duì)于 IDMM 的局限性進(jìn)行了探討，在像 DeiT 這樣的架構(gòu)上從頭開(kāi)始訓(xùn)練這些小型數(shù)據(jù)集時(shí)，模型仍有改進(jìn)的空間。還有哪些屬性對(duì)于小型數(shù)據(jù)集預(yù)訓(xùn)練起到關(guān)鍵作用？這還有待未來(lái)的進(jìn)一步研究。

編輯：王菁