何愷明團隊新作解決訓練不穩(wěn)定性!自監(jiān)督學習+Transformer=MoCoV3
點擊上方“AI算法與圖像處理”,選擇加"星標"或“置頂”
重磅干貨,第一時間送達
導讀
本文是FAIR的愷明團隊針對自監(jiān)督學習+Transformer的一篇實證研究。針對Transformer在自監(jiān)督學習框架中存在的訓練不穩(wěn)定問題,提出了一種簡單而有效的技巧:Random Patch Projection
paper: https://arxiv.org/abs/2104.02057
本文是FAIR的愷明團隊針對自監(jiān)督學習+Transformer的一篇實證研究。針對Transformer在自監(jiān)督學習框架中存在的訓練不穩(wěn)定問題,提出了一種簡單而有效的技巧:Random Patch Projection,它不僅適用于MoCoV3框架,同樣適用于其他自監(jiān)督學習框架(比如SimCLR、BYOL);與此同時,從不同角度的對比分析得出:Transformer中的位置信息尚未得到充分探索,即Transformer仍有繼續(xù)改善的空間。
Abstract
本文并沒有提出一種新的方法,相反,鑒于最近計算機視覺的進展,我們研究了一個簡單、漸進、但必須知道的基線:用于視覺Transformer的自監(jiān)督學習。盡管標準卷積網(wǎng)絡的訓練方法已經(jīng)非常成熟且魯棒,然而ViT的訓練方案仍有待于構建,特別是自監(jiān)督場景下的訓練極具挑戰(zhàn)。
在這里,我們從基礎出發(fā),對訓練自監(jiān)督ViT的幾種基本組件的影響進行了分析調(diào)研。我們發(fā)現(xiàn):不穩(wěn)定性是影響精確下降的最主要問題,它會被表面上好的結果覆蓋(容易陷入局部最優(yōu))。我們通過實驗發(fā)現(xiàn):這些結果確實存在部分失?。划斢柧氉兊梅€(wěn)定時,這些結果可以進一步提升?;贛oCoV3以及其他自監(jiān)督框架,我們從不同角度對ViT進行了測試分析;我們對觀察到的積極面、挑戰(zhàn)性以及開放問題進行了討論,期望該工作可以為未來的研究提供有用的數(shù)據(jù)支撐和經(jīng)驗參考。
Introduction
本文主要聚焦于:采用視覺領域的自監(jiān)督框架進行Transformer的訓練。CNN的訓練方法已被進行充分的研究與論證,而ViT模型是新的,其訓練方法尚未完整構建。本文從基礎出發(fā),研究了影響深度網(wǎng)絡訓練的幾個基本模塊:batch size,learning rate以及optimizer。我們發(fā)現(xiàn):在不同場景下,不穩(wěn)定性均是影響自監(jiān)督ViT訓練的主要問題。
有意思的是,我們發(fā)現(xiàn):不穩(wěn)定的ViT訓練可能不會導致災難性結果(比如發(fā)散);相反,它可以導致精度的輕度退化(約下降1-3%)。除非有一個更穩(wěn)定的作為對比,否則這個程度的退化可能難以被注意到。據(jù)我們所知,該現(xiàn)象在卷積網(wǎng)絡訓練中鮮少發(fā)生,我們認為該問題及其隱含的退化值得注意。
為了證明不穩(wěn)定性的可能危害,我們研究了一個可以在實踐中提高穩(wěn)定性的簡單技巧。基于梯度變換的經(jīng)驗觀察,我們固化ViT中的塊投影層,即采用固定隨機塊投影。我們發(fā)現(xiàn)該trick可以緩解多種場景下的不穩(wěn)定問題并提升模型精度。
基于對比學習框架,自監(jiān)督Transformer可以取得非常好的結果。不同于ImageNet監(jiān)督的ViT(模型變大時精度反而變差),更大的自監(jiān)ViT可以取得更高的精度。比如ViT-Large的自監(jiān)督預訓練可以取得超越監(jiān)督預訓練版本的性能。
此外,本文所提自監(jiān)督ViT模塊可以取得與大的卷積網(wǎng)絡相當?shù)男阅?。一方面驗證了ViT的潛力;另一方面意味著自監(jiān)督ViT仍有進一步提升的空間。因為我們發(fā)現(xiàn):移除ViT中的position embedding僅僅造成了輕微的性能下降,這意味著:自監(jiān)督ViT無需位置信息即可學習很強的特征表達,同時也也暗示位置信息并未得到充分探索。
MoCoV3
我們先來看一下本文所提出的MoCoV3,它是對MoCo V1/2的一種改進,尋求在簡單性、精度以及可縮放等方面提供更好的均衡。MoCoV3的偽代碼實現(xiàn)如下:

類似MoCo、SimCLR,我們采用隨機數(shù)據(jù)增強從每個圖像中裁剪兩個圖像塊,并經(jīng)由兩個編碼器編碼為,我們采用InfoNCE損失函數(shù):
其中,表示q同源圖像的輸出,即正樣本;表示異源圖像輸出,即負樣本。
延續(xù)SimCLR處理方式,MoCoV3采用同一批次自然共存的密鑰,移除了Memory Queue(當batch足夠大時,其收益遞減)?;谶@種簡化,上面的對比損失可以簡化為上述偽代碼中ctr實現(xiàn)。我們采用了對稱損失:。
我們的編碼器由骨干網(wǎng)絡(如ResNet、ViT)、投影頭以及額外的預測頭構成;而編碼則由骨干網(wǎng)絡、投影頭構成,沒有預測頭。通過的滑動平均更新。
作為參考,我們以ResNet50作為骨干網(wǎng)絡,其在ImageNet上的性能如下,由于額外的預測頭與大的batch,MoCoV3具有更加性能。
Stability of Self-Supervised ViT Training
原則上來講,我們可以直接在對比自監(jiān)督學習框架中采用ViT骨干替換ResNet骨干網(wǎng)絡。但實際上,主要挑戰(zhàn)在于:訓練不穩(wěn)定。
我們發(fā)現(xiàn):不穩(wěn)定問題不僅僅由精度反映。實際上,即使存在不穩(wěn)定問題,它也可以“表面上很好”的進行訓練并提供一個不錯的結果。為揭示這種不穩(wěn)定性,我們在訓練過程中對kNN曲線進行了監(jiān)控,研究了它如何影響不穩(wěn)定性并提出了一種簡單的trick進行穩(wěn)定訓練,進而提升不同場景下的精度。
Empirical Observations on Basic Factors

按照上述方式,我們設計了不同計算量的ViT模型并采用大batch(它有助于提升自監(jiān)督學習方法的性能)進行訓練。在下面的分析中,我們采用ViT-B/16作為基準。

BatchSize。上圖給出了不同batch時的訓練曲線??梢钥吹剑?/section> 1k和2k的batch具有比較平滑的曲線,精度分別為71.5%和72.6%; 4k的曲線就開始變得不再穩(wěn)定,其最終精度為72.2%,要低于2k的72.6%。后面會提到:這里的性能下降是受不穩(wěn)定訓練導致。 6k的曲線的不穩(wěn)定性進一步加劇,最終仍有一個不錯的結果69.7%。我們猜測:訓練被部分重啟,跳出了當前局部最優(yōu)并重新尋找新的優(yōu)化軌跡。因此,訓練無法收斂,最終的精度會依賴于重啟的局部性能。 此外,我們還發(fā)現(xiàn)這種不穩(wěn)定不會導致一個顯著差的結果。實驗發(fā)現(xiàn):相同配置下,結果的差異較小,范圍內(nèi)。這就使得不穩(wěn)定導致了性能退化難以被注意到。

LearningRate。事實上,學習率通常會隨batch提升而縮放。在本文實驗中,我們采用了線性縮放規(guī)則:。上圖給出了學習率的影響曲線。當lr比較小時,訓練更為穩(wěn)定,但會欠擬合。比如0.5e-4訓練的模型精度要比1.0e-4訓練的模型精度低1.8%。當lr比較大時,訓練不再穩(wěn)定。如1.5e-4訓練曲線具有更多的下降波谷且精度更低。

Optimizer。我們默認采用AdamW作為訓練優(yōu)化器;另一方面,LARS是自監(jiān)督方法常用的一種優(yōu)化器。我們研究了LMAB優(yōu)化器,它是LARS的AdamW版,結果見上圖。給定合理的學習率,LMAB可以獲得比AdamW稍高的精度(72.5%)。但是當學習率比最優(yōu)學習率大時,模型精度會迅速下降。有意思的是:此時的訓練曲線仍然很平滑,但會緩慢下降。我們發(fā)現(xiàn):當學習率合適時,LAMB可以取得與AdamW相當?shù)木?。但是學習率的敏感性使其無法在進行學習率搜索的前提下適配不同架構設計。因此,我們?nèi)赃x擇AdamW作為默認優(yōu)化器。
A Trick for Improving Stability
前面的所有實驗均表明:不穩(wěn)定是主要問題。接下來,我們將提出一種簡單的trick提升不同場景下的穩(wěn)定性。

在訓練過程中,我們注意到梯度突變(見上圖突變波峰)會導致訓練曲線的“下沉”(見上圖)。通過比較所有層的梯度,我們發(fā)現(xiàn):梯度突變發(fā)生在(patch projection)第一層先發(fā)生,然后延遲一定迭代后最后一層再發(fā)生?;谠摪l(fā)現(xiàn),我們猜測:不穩(wěn)定性發(fā)生在淺層。受此啟發(fā),我們在訓練過程中對塊投影進行凍結。換句話說:我們采用固定的Random Patch Projection層進行塊嵌入,而非通過學習方式。

上圖對比了可學習與隨機塊投影的MoCoV3結果對比??梢钥吹剑?strong style="font-weight: bold;color: black;">隨機塊投影可以穩(wěn)定訓練,訓練曲線更為平滑、精度更高(精度提升約1.7%),進一步說明了訓練不穩(wěn)定是影響精度的主要問題。

我們發(fā)現(xiàn):除了MoCo外,其他相關方法(比如SimCLR、BYOL)同樣存在不穩(wěn)定問題。隨機塊投影同樣可以改善SimCLR與BYOL的性能(分別提升0.8%和1.3%),見上圖對比。不穩(wěn)定對于SwAV會導致?lián)p失發(fā)散(NaN),本文所提隨機塊投影可以穩(wěn)定訓練SwAV并將其精度由65.8%提升到66.4%。總而言之,本文所提t(yī)rick對于所有自監(jiān)督方案均有效。
Discussion:一個有意思的發(fā)現(xiàn):塊投影層的訓練并非必要。對于標準ViT塊尺寸,其塊投影矩陣時過完備的,此時隨機投影足以保持原始塊的信息。
我們注意到:凍結第一層并不會改變架構,但它會縮小解決方案空間。這意味著:根本問題與優(yōu)化相關。該trick可以緩解優(yōu)化問題,但不能解決它;當lr過于大時模型仍存在不穩(wěn)定問題。第一層不太可能時不穩(wěn)定的主要原因,相反,該問題與所有層相關。但第一層只是更易于分開處理,因為它是骨干網(wǎng)絡僅有的非Transformer層。
Implementation Details
在這里,我們ViT+MoCoV3的實現(xiàn)細節(jié)進行更詳細的描述。
Optimizer:默認選擇AdamW,batch為4096,通過100epoch訓練搜索lr和wd,然后選擇最優(yōu)訓練更長。學習率采用40epoch進行warmup,它同樣有助于緩解不穩(wěn)定性;warmup之后,學習率按照cosine衰減; MLP Head:投影頭是一個3層MLP;預測頭是一個2層MLP。MLP的隱含層維度均為4096,輸出層維度均為256。在MoCoV3中,參考SimCLR,MLP的所有Linear均后接BN。 Loss:參考BYOL,損失函數(shù)進行了縮放。盡管該方法可以通過調(diào)節(jié)lr和wd合并,但可以使得其對的敏感性降低,默認。 ViT Architecture:參考ViT一文,輸入塊為或者,經(jīng)過投影后它將輸出一個長度為196/256的序列,Position Embedding與該序列相加;所得序列與可學習類token拼接并經(jīng)由后續(xù)Transformer模塊編碼;最終所得類token視作輸出并送入MLP頭。 Linear probing。延續(xù)常規(guī)方案,我們采用線性方式評估特征表達質(zhì)量。完成自監(jiān)督與訓練后,移除MLP頭并采用監(jiān)督方式訓練一個線性分類器,此時訓練90epoch,且僅僅采用RandomResize、Flipping進行數(shù)據(jù)增廣。
Experiments Results

上表給出了不同模型的計算量、訓練時長等信息(上述結果為谷歌云平臺實驗結果)。ViT-B訓練100epoch花費2.1小時;ViT-H訓練100epoch花費9.8小時(512個TPU)。如果采用GPU的話,ViT-B需要24小時(128GPU)。TPU的擴展性要比GPU更優(yōu)哇。
Self-supervised Learning Framework

上表給出了四種不同自監(jiān)督學習框架的實驗對比??梢钥吹剑?1) 相比其他自監(jiān)督方法,在相同骨干網(wǎng)絡ViT-S與ViT-B時,MoCoV3具有更加的精度。(2) MoCoV3與SimCLR在ViT-B方面的性能要優(yōu)于ResNet50。
Ablations of ViT+MoCoV3
接下來,我們將從不同角度對本文所提框架進行消融實驗分析。

Position Embedding(PE)。上表比較了的不同PE的性能對比,可以看到:(1) 可學習方式表現(xiàn)比較好,并并不如sin-cos方式;(2) 移除掉PE后,模型仍有一個不錯的結果74.9%,也就是說PE僅僅貢獻了1.6%。該實驗揭示了當前模型的強處與局限性:一方面,模型可以僅僅通過塊集合即可學習很強的表達能力,類似于bag-of-word模型;另一方面,模型可能并未充分利用位置信息。

Class Token(CLS)。上表給出了CLS的影響性對比,可以看到:(1) 移除CLS,保留LN,此時性能比較差,僅有69.7%;(2) 移除LN和CLS,結果幾乎不變76.6%。這意味著:CLS對于該系統(tǒng)并非關鍵因子;同時也意味著:規(guī)范化層的選擇影響較大。

BatchNorm in MLP。上表比較了BN存在魚頭的影響性??梢钥吹剑篵atch=2048時,移除BN導致了2.1%的性能下降。這意味著:BN并非對比學習的必要因子,但合理的使用BN可以提升精度。

Prediction Head。上表對比了預測頭有無的性能對比。預測頭并非MoCo的必選項,但預測頭會帶來額外的精度提升(1%);而在BYOL與SimSiam中卻是必選項。

Momentum Encoder。上表對比了Momentum稀疏的影響對比。Momentum編碼可以帶來2.2%的性能提升。

Training length。上表對比了不同訓練時長的性能對比。輕量版ViT-S可以從更長周期的訓練中受益更多,比如精度提升0.9%;重量型ViT-B則受益很少。
Comparison with Prior Art

上表給出了MoCoV3框架下不同ViT模型的對比??梢钥吹剑?1) 相比iGPUT,無需額外數(shù)據(jù)預訓練,基于MoCo的ViT取得了更高的精度、更小的模型。(2) 隨著模型變大,所提方案的精度會逐漸提升;而在監(jiān)督學習方式中,基于ImageNet-1k/2k預訓練的ViT-L的精度要低于ViT-B。事實上,本文自監(jiān)督預訓練的ViT-L精度(77.6%)要比監(jiān)督方式(76.53%)的更高,這意味著:自監(jiān)督學習作為一種廣義表達學習工具不易于過擬合。

上圖給出了所提方案與ResNet系列+其他自監(jiān)督學習的性能對比??梢钥吹剑?/p>
在小模型方面,本文基線ViT MoCo(即ViT, MoCoV3)的性能要比SimCLRv2+ResNet更佳;在大模型方面則與之相當。 SimCLRV2+SK-ResNet的組合具有更高的性能; BYOL+Wider-ResNet具有更高的性能,配合R200-2x時可以得到更優(yōu)異的結果; 正如前面提到的規(guī)范化技術的影響,將LN替換為BN后模型的性能提升了1%; 將ViT默認塊尺寸替換為后,模型的性能進一步提升2-3%。MoCoV3+ViT-BN-L/取得了81.0%的top1精度,作為對比,SimCLRV2+SK-ResNet152-3x的最佳精度為79.8%,BYOL-ResNet200-2x的最佳精度為79.6%。
Transfer Leanring

最后,我們再看一下所提方案在下游任務遷移學習方面的性能。結果見上表,可以看到:
當模型大小從ViT-B提升到ViT-L,所提方案具有更加的遷移學習精度;當提升到ViT-H時則出現(xiàn)了過擬合問題。作為對比,ImageNet監(jiān)督ViT在ViT-L時就出現(xiàn)了過擬合。 相比ImageNet監(jiān)督方案,本文所提自監(jiān)督ViT取得了更佳的結果。 在這些小數(shù)據(jù)上,采用大的ViT模型從頭開始訓練時過擬合問題非常嚴重。這意味著:如果數(shù)據(jù)量不足會導致難以訓練ViT學習好的特征表達;而自監(jiān)督預訓練可以彌補這種差距,極大的避免小數(shù)據(jù)集上的過擬合問題。
個人微信(如果沒有備注不拉群!) 請注明:地區(qū)+學校/企業(yè)+研究方向+昵稱
下載1:何愷明頂會分享
在「AI算法與圖像處理」公眾號后臺回復:何愷明,即可下載。總共有6份PDF,涉及 ResNet、Mask RCNN等經(jīng)典工作的總結分析
下載2:終身受益的編程指南:Google編程風格指南
在「AI算法與圖像處理」公眾號后臺回復:c++,即可下載。歷經(jīng)十年考驗,最權威的編程規(guī)范!
下載3 CVPR2021 在「AI算法與圖像處理」公眾號后臺回復:CVPR,即可下載1467篇CVPR 2020論文 和 CVPR 2021 最新論文
點亮
,告訴大家你也在看

