2040張圖片訓練出的ViT,準確率96.7%,連遷移性能都令人驚訝
本文經AI新媒體量子位(公眾號ID:QbitAI)授權轉載,轉載請聯系出處。
ViT在計算機視覺領域取得了巨大的成功,甚至大有取代CNN之勢。
但是相比CNN,訓練ViT需要更多的數據,通常要在大型數據集JFT-300M或至少在ImageNet上進行預訓練,很少有人研究少量數據訓練ViT。
最近,南京大學吳建鑫團隊提出了一種新方法,只需2040張圖片即可訓練ViT。
他們在2040張花(flowers)的圖像上從頭開始訓練,達到了96.7%的準確率,表明用小數據訓練ViT也是可行的。
另外在ViT主干下的 7 個小型數據集上從頭開始訓練時,也獲得了SOTA的結果。
而且更重要的是,他們證明了,即使在小型數據集上進行預訓練,ViT也具有良好的遷移能力,甚至可以促進對大規(guī)模數據集的訓練。
論文內容
在這篇論文中,作者提出了用于自我監(jiān)督 ViT訓練的IDMM(Instance Discrimination with Multi-crop and CutMix)。
我們先來看一下ViT圖像分類網絡的基本架構:
將圖像樣本x?(i = 1, 2, …, N; N為圖片數量)送入ViT中,得到一組輸出表征z?。w?為第j個分類的權重。
然后,使用全連接層W進行分類,當類的數量等于訓練圖像的總數N時,即參數化實例判別。
第j類的輸出為:
我們把O送入Softmax層,就得到一個概率分布P???。對于實例判別,損失函數為:
對于深度聚類,其損失函數為:
可以看出,只要適當設置權重(讓w? = ~w? ),就可以讓實例判別等價于深度聚類。
從下圖中可以看出,與其他方法相比,實例判別可以學習到更多的分布式表征,并能更好地捕捉到類內的相似性。
作者之所以選擇參數化的實例判別,還有一個重要的原因:簡單性和穩(wěn)定性。
不穩(wěn)定性是影響自監(jiān)督ViT訓練的一個主要問題。實例判別(交叉熵)的形式更穩(wěn)定,更容易優(yōu)化。
接下來開始梯度分析,損失函數對權重求導:
其中δ是指示函數,當k=i時值為1,否則為0。
需要注意的是,對于實例判別,類的數量N通常很大,而且存在對實例樣本訪問極稀少的問題。
對于稀少的實例k≠i,可以預計P????≈0,因此?L/?w?≈0,這意味著w?的更新頻率極低。
在小數據集問題上,作者使用CutMix和標簽平滑,來緩解此問題。
CutMix:
標簽平滑:
最后梯度變?yōu)椋?/p>
這樣通過直接修改單次標簽,來更頻繁地更新權重矩陣,也是ViT監(jiān)督訓練中常用的方法。
總之,作者使用了以下策略來加強小數據集上的實例判別。
- ?小分辨率:預訓練中的小分辨率對小數據集很有用。
- 多次裁剪:實例判別概括了對比損失,保證了在使用多種實例時獲取特征的對齊和統(tǒng)一性。
- CutMix和標簽平滑:有助于緩解使用實例判別時的過擬合和不經常訪問的問題。
至于為什么需要直接在目標數據集上從頭開始訓練,作者給出了3點原因:
1、數據
目前的ViT模型通常在一個大規(guī)模的數據集上進行預訓練,然后在各種下游任務中進行微調。由于缺乏典型的卷積歸納偏向,這些模型比普通的CNN更耗費數據。
因此從頭開始訓練ViT,能夠用圖像總量有限的任務是至關重要的。
2、算力
大規(guī)模的數據集、大量的耗時和復雜的骨干網絡的,讓ViT訓練的算力成本非常昂貴。這種現象使ViT成為少數機構研究人員的特權。
3、靈活性
預訓練后再進行下游微調的模式有時會很麻煩。
例如,我們可能需要為同一任務訓練10個不同的模型,并將它們部署在不同的硬件平臺上,但在一個大規(guī)模的數據集上預訓練10個模型是不現實的。
在上圖中,很明顯與從頭開始訓練相比,ImageNet預訓練的模型需要更多的參數和計算成本。
在小數據集上進行預訓練時的遷移能力。每個單元格和列中精度最高的元素分別用下劃線和粗體表示
最后,在下表中,作者評估了在不同數據集上預訓練模型的遷移精度。
對角線上的單元(灰色)是在同一數據集上進行預訓練和微調。對角線外的單元格評估了這些小數據集的遷移性能。
從這張表中,我們可以看到以下幾點:
- 即使在小數據集上進行預訓練,ViT也有良好的遷移能力。
- 與SimCLR和SupCon相比,該方法在所有這些數據集上也有更高的遷移精度。
- 即使預訓練的數據集和目標數據集不在同一領域,也能獲得令人驚訝的好結果。例如,在Indoor67上預訓練的模型在轉移到Aircraft上時獲得了最高的準確性。
作者簡介
本文第一作者是南京大學在讀博士曹云浩,通訊作者是南京大學人工智能學院吳建鑫教授。
吳建鑫本科和碩士畢業(yè)于南京大學計算機專業(yè),博士畢業(yè)于佐治亞理工。2013年,他加入南京大學科學與技術系,任教授、博士生導師,曾擔任ICCV 2015領域主席、CVPR 2017領域主席,現為Pattern Recognition期刊編委。