2040張圖片訓(xùn)練出的ViT，準(zhǔn)確率96.7%，連遷移性能都令人驚訝 | 南京大學(xué)

蕭簫 2022-02-02 13:26:48 來源：量子位

曉查 發(fā)自 凹非寺
量子位 | 公眾號 QbitAI

ViT在計算機(jī)視覺領(lǐng)域取得了巨大的成功，甚至大有取代CNN之勢。

但是相比CNN，訓(xùn)練ViT需要更多的數(shù)據(jù)，通常要在大型數(shù)據(jù)集JFT-300M或至少在ImageNet上進(jìn)行預(yù)訓(xùn)練，很少有人研究少量數(shù)據(jù)訓(xùn)練ViT。

最近，南京大學(xué)吳建鑫團(tuán)隊提出了一種新方法，只需2040張圖片即可訓(xùn)練ViT。

他們在2040張花（flowers）的圖像上從頭開始訓(xùn)練，達(dá)到了96.7%的準(zhǔn)確率，表明用小數(shù)據(jù)訓(xùn)練ViT也是可行的。

另外在ViT主干下的 7 個小型數(shù)據(jù)集上從頭開始訓(xùn)練時，也獲得了SOTA的結(jié)果。

而且更重要的是，他們證明了，即使在小型數(shù)據(jù)集上進(jìn)行預(yù)訓(xùn)練，ViT也具有良好的遷移能力，甚至可以促進(jìn)對大規(guī)模數(shù)據(jù)集的訓(xùn)練。

論文內(nèi)容

在這篇論文中，作者提出了用于自我監(jiān)督 ViT訓(xùn)練的IDMM（Instance?Discrimination with?Multi-crop and CutMix）。

我們先來看一下ViT圖像分類網(wǎng)絡(luò)的基本架構(gòu)：

將圖像樣本x?（i = 1, 2, …, N; N為圖片數(shù)量）送入ViT中，得到一組輸出表征z?。w?為第j個分類的權(quán)重。

然后，使用全連接層W進(jìn)行分類，當(dāng)類的數(shù)量等于訓(xùn)練圖像的總數(shù)N時，即參數(shù)化實例判別。

第j類的輸出為：

我們把O送入Softmax層，就得到一個概率分布P???。對于實例判別，損失函數(shù)為：

對于深度聚類，其損失函數(shù)為：

可以看出，只要適當(dāng)設(shè)置權(quán)重（讓w? = ~w? ），就可以讓實例判別等價于深度聚類。

從下圖中可以看出，與其他方法相比，實例判別可以學(xué)習(xí)到更多的分布式表征，并能更好地捕捉到類內(nèi)的相似性。

作者之所以選擇參數(shù)化的實例判別，還有一個重要的原因：簡單性和穩(wěn)定性。

不穩(wěn)定性是影響自監(jiān)督ViT訓(xùn)練的一個主要問題。實例判別（交叉熵）的形式更穩(wěn)定，更容易優(yōu)化。

接下來開始梯度分析，損失函數(shù)對權(quán)重求導(dǎo)：

其中δ是指示函數(shù)，當(dāng)k=i時值為1，否則為0。

需要注意的是，對于實例判別，類的數(shù)量N通常很大，而且存在對實例樣本訪問極稀少的問題。

對于稀少的實例k≠i，可以預(yù)計P????≈0，因此?L/?w?≈0，這意味著w?的更新頻率極低。

在小數(shù)據(jù)集問題上，作者使用CutMix和標(biāo)簽平滑，來緩解此問題。

CutMix：

標(biāo)簽平滑：

最后梯度變?yōu)椋?/p>

這樣通過直接修改單次標(biāo)簽，來更頻繁地更新權(quán)重矩陣，也是ViT監(jiān)督訓(xùn)練中常用的方法。

總之，作者使用了以下策略來加強(qiáng)小數(shù)據(jù)集上的實例判別。

1. 小分辨率：預(yù)訓(xùn)練中的小分辨率對小數(shù)據(jù)集很有用。
2. 多次裁剪：實例判別概括了對比損失，保證了在使用多種實例時獲取特征的對齊和統(tǒng)一性。
3. CutMix和標(biāo)簽平滑：有助于緩解使用實例判別時的過擬合和不經(jīng)常訪問的問題。

至于為什么需要直接在目標(biāo)數(shù)據(jù)集上從頭開始訓(xùn)練，作者給出了3點原因：

1、數(shù)據(jù)

目前的ViT模型通常在一個大規(guī)模的數(shù)據(jù)集上進(jìn)行預(yù)訓(xùn)練，然后在各種下游任務(wù)中進(jìn)行微調(diào)。由于缺乏典型的卷積歸納偏向，這些模型比普通的CNN更耗費數(shù)據(jù)。

因此從頭開始訓(xùn)練ViT，能夠用圖像總量有限的任務(wù)是至關(guān)重要的。

2、算力

大規(guī)模的數(shù)據(jù)集、大量的耗時和復(fù)雜的骨干網(wǎng)絡(luò)的，讓ViT訓(xùn)練的算力成本非常昂貴。這種現(xiàn)象使ViT成為少數(shù)機(jī)構(gòu)研究人員的特權(quán)。

3、靈活性

預(yù)訓(xùn)練后再進(jìn)行下游微調(diào)的模式有時會很麻煩。

例如，我們可能需要為同一任務(wù)訓(xùn)練10個不同的模型，并將它們部署在不同的硬件平臺上，但在一個大規(guī)模的數(shù)據(jù)集上預(yù)訓(xùn)練10個模型是不現(xiàn)實的。

在上圖中，很明顯與從頭開始訓(xùn)練相比，ImageNet預(yù)訓(xùn)練的模型需要更多的參數(shù)和計算成本。

在小數(shù)據(jù)集上進(jìn)行預(yù)訓(xùn)練時的遷移能力。每個單元格和列中精度最高的元素分別用下劃線和粗體表示

最后，在下表中，作者評估了在不同數(shù)據(jù)集上預(yù)訓(xùn)練模型的遷移精度。

對角線上的單元（灰色）是在同一數(shù)據(jù)集上進(jìn)行預(yù)訓(xùn)練和微調(diào)。對角線外的單元格評估了這些小數(shù)據(jù)集的遷移性能。

從這張表中，我們可以看到以下幾點：

1. 即使在小數(shù)據(jù)集上進(jìn)行預(yù)訓(xùn)練，ViT也有良好的遷移能力。
2. 與SimCLR和SupCon相比，該方法在所有這些數(shù)據(jù)集上也有更高的遷移精度。
3. 即使預(yù)訓(xùn)練的數(shù)據(jù)集和目標(biāo)數(shù)據(jù)集不在同一領(lǐng)域，也能獲得令人驚訝的好結(jié)果。例如，在Indoor67上預(yù)訓(xùn)練的模型在轉(zhuǎn)移到Aircraft上時獲得了最高的準(zhǔn)確性。

作者簡介

本文第一作者是南京大學(xué)在讀博士曹云浩，通訊作者是南京大學(xué)人工智能學(xué)院吳建鑫教授。

吳建鑫本科和碩士畢業(yè)于南京大學(xué)計算機(jī)專業(yè)，博士畢業(yè)于佐治亞理工。2013年，他加入南京大學(xué)科學(xué)與技術(shù)系，任教授、博士生導(dǎo)師，曾擔(dān)任ICCV 2015領(lǐng)域主席、CVPR 2017領(lǐng)域主席，現(xiàn)為Pattern Recognition期刊編委。

參考鏈接：
[1]https://arxiv.org/abs/2201.10728
[2]https://cs.nju.edu.cn/wujx/index.htm