paper: https://arxiv.org/abs/2104.02057

本文是FAIR的愷明團(tuán)隊(duì)針對(duì)自監(jiān)督學(xué)習(xí)+Transformer的一篇實(shí)證研究。針對(duì)Transformer在自監(jiān)督學(xué)習(xí)框架中存在的訓(xùn)練不穩(wěn)定問(wèn)題，提出了一種簡(jiǎn)單而有效的技巧：Random Patch Projection，它不僅適用于MoCoV3框架，同樣適用于其他自監(jiān)督學(xué)習(xí)框架(比如SimCLR、BYOL)；與此同時(shí)，從不同角度的對(duì)比分析得出：Transformer中的位置信息尚未得到充分探索，即Transformer仍有繼續(xù)改善的空間。

Abstract

本文并沒(méi)有提出一種新的方法，相反，鑒于最近計(jì)算機(jī)視覺(jué)的進(jìn)展，我們研究了一個(gè)簡(jiǎn)單、漸進(jìn)、但必須知道的基線：用于視覺(jué)Transformer的自監(jiān)督學(xué)習(xí)。盡管標(biāo)準(zhǔn)卷積網(wǎng)絡(luò)的訓(xùn)練方法已經(jīng)非常成熟且魯棒，然而ViT的訓(xùn)練方案仍有待于構(gòu)建，特別是自監(jiān)督場(chǎng)景下的訓(xùn)練極具挑戰(zhàn)。

在這里，我們從基礎(chǔ)出發(fā)，對(duì)訓(xùn)練自監(jiān)督ViT的幾種基本組件的影響進(jìn)行了分析調(diào)研。我們發(fā)現(xiàn)：不穩(wěn)定性是影響精確下降的最主要問(wèn)題，它會(huì)被表面上好的結(jié)果覆蓋(容易陷入局部最優(yōu))。我們通過(guò)實(shí)驗(yàn)發(fā)現(xiàn)：這些結(jié)果確實(shí)存在部分失?。划?dāng)訓(xùn)練變得穩(wěn)定時(shí)，這些結(jié)果可以進(jìn)一步提升?；贛oCoV3以及其他自監(jiān)督框架，我們從不同角度對(duì)ViT進(jìn)行了測(cè)試分析；我們對(duì)觀察到的積極面、挑戰(zhàn)性以及開放問(wèn)題進(jìn)行了討論，期望該工作可以為未來(lái)的研究提供有用的數(shù)據(jù)支撐和經(jīng)驗(yàn)參考。

Introduction

本文主要聚焦于：采用視覺(jué)領(lǐng)域的自監(jiān)督框架進(jìn)行Transformer的訓(xùn)練。CNN的訓(xùn)練方法已被進(jìn)行充分的研究與論證，而ViT模型是新的，其訓(xùn)練方法尚未完整構(gòu)建。本文從基礎(chǔ)出發(fā)，研究了影響深度網(wǎng)絡(luò)訓(xùn)練的幾個(gè)基本模塊：batch size，learning rate以及optimizer。我們發(fā)現(xiàn)：在不同場(chǎng)景下，不穩(wěn)定性均是影響自監(jiān)督ViT訓(xùn)練的主要問(wèn)題。

有意思的是，我們發(fā)現(xiàn)：不穩(wěn)定的ViT訓(xùn)練可能不會(huì)導(dǎo)致災(zāi)難性結(jié)果(比如發(fā)散)；相反，它可以導(dǎo)致精度的輕度退化(約下降1-3%)。除非有一個(gè)更穩(wěn)定的作為對(duì)比，否則這個(gè)程度的退化可能難以被注意到。據(jù)我們所知，該現(xiàn)象在卷積網(wǎng)絡(luò)訓(xùn)練中鮮少發(fā)生，我們認(rèn)為該問(wèn)題及其隱含的退化值得注意。

為了證明不穩(wěn)定性的可能危害，我們研究了一個(gè)可以在實(shí)踐中提高穩(wěn)定性的簡(jiǎn)單技巧。基于梯度變換的經(jīng)驗(yàn)觀察，我們固化ViT中的塊投影層，即采用固定隨機(jī)塊投影。我們發(fā)現(xiàn)該trick可以緩解多種場(chǎng)景下的不穩(wěn)定問(wèn)題并提升模型精度。

基于對(duì)比學(xué)習(xí)框架，自監(jiān)督Transformer可以取得非常好的結(jié)果。不同于ImageNet監(jiān)督的ViT(模型變大時(shí)精度反而變差)，更大的自監(jiān)ViT可以取得更高的精度。比如ViT-Large的自監(jiān)督預(yù)訓(xùn)練可以取得超越監(jiān)督預(yù)訓(xùn)練版本的性能。

此外，本文所提自監(jiān)督ViT模塊可以取得與大的卷積網(wǎng)絡(luò)相當(dāng)?shù)男阅?。一方面?yàn)證了ViT的潛力；另一方面意味著自監(jiān)督ViT仍有進(jìn)一步提升的空間。因?yàn)槲覀儼l(fā)現(xiàn)：移除ViT中的position embedding僅僅造成了輕微的性能下降，這意味著：自監(jiān)督ViT無(wú)需位置信息即可學(xué)習(xí)很強(qiáng)的特征表達(dá)，同時(shí)也也暗示位置信息并未得到充分探索。

MoCoV3

我們先來(lái)看一下本文所提出的MoCoV3，它是對(duì)MoCo V1/2的一種改進(jìn)，尋求在簡(jiǎn)單性、精度以及可縮放等方面提供更好的均衡。MoCoV3的偽代碼實(shí)現(xiàn)如下：

類似MoCo、SimCLR，我們采用隨機(jī)數(shù)據(jù)增強(qiáng)從每個(gè)圖像中裁剪兩個(gè)圖像塊，并經(jīng)由兩個(gè)編碼器編碼為，我們采用InfoNCE損失函數(shù)：

其中，表示q同源圖像的輸出，即正樣本；表示異源圖像輸出，即負(fù)樣本。

延續(xù)SimCLR處理方式，MoCoV3采用同一批次自然共存的密鑰，移除了Memory Queue(當(dāng)batch足夠大時(shí)，其收益遞減)?；谶@種簡(jiǎn)化，上面的對(duì)比損失可以簡(jiǎn)化為上述偽代碼中ctr實(shí)現(xiàn)。我們采用了對(duì)稱損失：。

我們的編碼器由骨干網(wǎng)絡(luò)(如ResNet、ViT)、投影頭以及額外的預(yù)測(cè)頭構(gòu)成；而編碼則由骨干網(wǎng)絡(luò)、投影頭構(gòu)成，沒(méi)有預(yù)測(cè)頭。通過(guò)的滑動(dòng)平均更新。

作為參考，我們以ResNet50作為骨干網(wǎng)絡(luò)，其在ImageNet上的性能如下，由于額外的預(yù)測(cè)頭與大的batch，MoCoV3具有更加性能。

Stability of Self-Supervised ViT Training

原則上來(lái)講，我們可以直接在對(duì)比自監(jiān)督學(xué)習(xí)框架中采用ViT骨干替換ResNet骨干網(wǎng)絡(luò)。但實(shí)際上，主要挑戰(zhàn)在于：訓(xùn)練不穩(wěn)定。

我們發(fā)現(xiàn)：不穩(wěn)定問(wèn)題不僅僅由精度反映。實(shí)際上，即使存在不穩(wěn)定問(wèn)題，它也可以“表面上很好”的進(jìn)行訓(xùn)練并提供一個(gè)不錯(cuò)的結(jié)果。為揭示這種不穩(wěn)定性，我們?cè)谟?xùn)練過(guò)程中對(duì)kNN曲線進(jìn)行了監(jiān)控，研究了它如何影響不穩(wěn)定性并提出了一種簡(jiǎn)單的trick進(jìn)行穩(wěn)定訓(xùn)練，進(jìn)而提升不同場(chǎng)景下的精度。

Empirical Observations on Basic Factors

按照上述方式，我們?cè)O(shè)計(jì)了不同計(jì)算量的ViT模型并采用大batch(它有助于提升自監(jiān)督學(xué)習(xí)方法的性能)進(jìn)行訓(xùn)練。在下面的分析中，我們采用ViT-B/16作為基準(zhǔn)。

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

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

• Optimizer。我們默認(rèn)采用AdamW作為訓(xùn)練優(yōu)化器；另一方面，LARS是自監(jiān)督方法常用的一種優(yōu)化器。我們研究了LMAB優(yōu)化器，它是LARS的AdamW版，結(jié)果見上圖。給定合理的學(xué)習(xí)率，LMAB可以獲得比AdamW稍高的精度(72.5%)。但是當(dāng)學(xué)習(xí)率比最優(yōu)學(xué)習(xí)率大時(shí)，模型精度會(huì)迅速下降。有意思的是：此時(shí)的訓(xùn)練曲線仍然很平滑，但會(huì)緩慢下降。我們發(fā)現(xiàn)：當(dāng)學(xué)習(xí)率合適時(shí)，LAMB可以取得與AdamW相當(dāng)?shù)木?。但是學(xué)習(xí)率的敏感性使其無(wú)法在進(jìn)行學(xué)習(xí)率搜索的前提下適配不同架構(gòu)設(shè)計(jì)。因此，我們?nèi)赃x擇AdamW作為默認(rèn)優(yōu)化器。

A Trick for Improving Stability

前面的所有實(shí)驗(yàn)均表明：不穩(wěn)定是主要問(wèn)題。接下來(lái)，我們將提出一種簡(jiǎn)單的trick提升不同場(chǎng)景下的穩(wěn)定性。

在訓(xùn)練過(guò)程中，我們注意到梯度突變(見上圖突變波峰)會(huì)導(dǎo)致訓(xùn)練曲線的“下沉”(見上圖)。通過(guò)比較所有層的梯度，我們發(fā)現(xiàn)：梯度突變發(fā)生在(patch projection)第一層先發(fā)生，然后延遲一定迭代后最后一層再發(fā)生?；谠摪l(fā)現(xiàn)，我們猜測(cè)：不穩(wěn)定性發(fā)生在淺層。受此啟發(fā)，我們?cè)谟?xùn)練過(guò)程中對(duì)塊投影進(jìn)行凍結(jié)。換句話說(shuō)：我們采用固定的Random Patch Projection層進(jìn)行塊嵌入，而非通過(guò)學(xué)習(xí)方式。

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

我們發(fā)現(xiàn)：除了MoCo外，其他相關(guān)方法(比如SimCLR、BYOL)同樣存在不穩(wěn)定問(wèn)題。隨機(jī)塊投影同樣可以改善SimCLR與BYOL的性能(分別提升0.8%和1.3%)，見上圖對(duì)比。不穩(wěn)定對(duì)于SwAV會(huì)導(dǎo)致?lián)p失發(fā)散(NaN)，本文所提隨機(jī)塊投影可以穩(wěn)定訓(xùn)練SwAV并將其精度由65.8%提升到66.4%?？偠灾?，本文所提t(yī)rick對(duì)于所有自監(jiān)督方案均有效。

Discussion：一個(gè)有意思的發(fā)現(xiàn)：塊投影層的訓(xùn)練并非必要。對(duì)于標(biāo)準(zhǔn)ViT塊尺寸，其塊投影矩陣時(shí)過(guò)完備的，此時(shí)隨機(jī)投影足以保持原始?jí)K的信息。

我們注意到：凍結(jié)第一層并不會(huì)改變架構(gòu)，但它會(huì)縮小解決方案空間。這意味著：根本問(wèn)題與優(yōu)化相關(guān)。該trick可以緩解優(yōu)化問(wèn)題，但不能解決它；當(dāng)lr過(guò)于大時(shí)模型仍存在不穩(wěn)定問(wèn)題。第一層不太可能時(shí)不穩(wěn)定的主要原因，相反，該問(wèn)題與所有層相關(guān)。但第一層只是更易于分開處理，因?yàn)樗枪歉删W(wǎng)絡(luò)僅有的非Transformer層。

Implementation Details

在這里，我們ViT+MoCoV3的實(shí)現(xiàn)細(xì)節(jié)進(jìn)行更詳細(xì)的描述。

• Optimizer：默認(rèn)選擇AdamW，batch為4096，通過(guò)100epoch訓(xùn)練搜索lr和wd，然后選擇最優(yōu)訓(xùn)練更長(zhǎng)。學(xué)習(xí)率采用40epoch進(jìn)行warmup，它同樣有助于緩解不穩(wěn)定性；warmup之后，學(xué)習(xí)率按照cosine衰減；
• MLP Head：投影頭是一個(gè)3層MLP；預(yù)測(cè)頭是一個(gè)2層MLP。MLP的隱含層維度均為4096，輸出層維度均為256。在MoCoV3中，參考SimCLR，MLP的所有Linear均后接BN。
• Loss：參考BYOL，損失函數(shù)進(jìn)行了縮放。盡管該方法可以通過(guò)調(diào)節(jié)lr和wd合并，但可以使得其對(duì)的敏感性降低，默認(rèn)。
• ViT Architecture：參考ViT一文，輸入塊為或者，經(jīng)過(guò)投影后它將輸出一個(gè)長(zhǎng)度為196/256的序列，Position Embedding與該序列相加；所得序列與可學(xué)習(xí)類token拼接并經(jīng)由后續(xù)Transformer模塊編碼；最終所得類token視作輸出并送入MLP頭。
• Linear probing。延續(xù)常規(guī)方案，我們采用線性方式評(píng)估特征表達(dá)質(zhì)量。完成自監(jiān)督與訓(xùn)練后，移除MLP頭并采用監(jiān)督方式訓(xùn)練一個(gè)線性分類器，此時(shí)訓(xùn)練90epoch，且僅僅采用RandomResize、Flipping進(jìn)行數(shù)據(jù)增廣。

Experiments Results

上表給出了不同模型的計(jì)算量、訓(xùn)練時(shí)長(zhǎng)等信息（上述結(jié)果為谷歌云平臺(tái)實(shí)驗(yàn)結(jié)果）。ViT-B訓(xùn)練100epoch花費(fèi)2.1小時(shí)；ViT-H訓(xùn)練100epoch花費(fèi)9.8小時(shí)(512個(gè)TPU)。如果采用GPU的話，ViT-B需要24小時(shí)(128GPU)。TPU的擴(kuò)展性要比GPU更優(yōu)哇。

Self-supervised Learning Framework

上表給出了四種不同自監(jiān)督學(xué)習(xí)框架的實(shí)驗(yàn)對(duì)比。可以看到：(1) 相比其他自監(jiān)督方法，在相同骨干網(wǎng)絡(luò)ViT-S與ViT-B時(shí)，MoCoV3具有更加的精度。(2) MoCoV3與SimCLR在ViT-B方面的性能要優(yōu)于ResNet50。

Ablations of ViT+MoCoV3

接下來(lái)，我們將從不同角度對(duì)本文所提框架進(jìn)行消融實(shí)驗(yàn)分析。

• Position Embedding(PE)。上表比較了的不同PE的性能對(duì)比，可以看到：(1) 可學(xué)習(xí)方式表現(xiàn)比較好，并并不如sin-cos方式；(2) 移除掉PE后，模型仍有一個(gè)不錯(cuò)的結(jié)果74.9%，也就是說(shuō)PE僅僅貢獻(xiàn)了1.6%。該實(shí)驗(yàn)揭示了當(dāng)前模型的強(qiáng)處與局限性：一方面，模型可以僅僅通過(guò)塊集合即可學(xué)習(xí)很強(qiáng)的表達(dá)能力，類似于bag-of-word模型；另一方面，模型可能并未充分利用位置信息。

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

• BatchNorm in MLP。上表比較了BN存在魚頭的影響性?？梢钥吹剑篵atch=2048時(shí)，移除BN導(dǎo)致了2.1%的性能下降。這意味著：BN并非對(duì)比學(xué)習(xí)的必要因子，但合理的使用BN可以提升精度。

• Prediction Head。上表對(duì)比了預(yù)測(cè)頭有無(wú)的性能對(duì)比。預(yù)測(cè)頭并非MoCo的必選項(xiàng)，但預(yù)測(cè)頭會(huì)帶來(lái)額外的精度提升(1%)；而在BYOL與SimSiam中卻是必選項(xiàng)。

• Momentum Encoder。上表對(duì)比了Momentum稀疏的影響對(duì)比。Momentum編碼可以帶來(lái)2.2%的性能提升。

• Training length。上表對(duì)比了不同訓(xùn)練時(shí)長(zhǎng)的性能對(duì)比。輕量版ViT-S可以從更長(zhǎng)周期的訓(xùn)練中受益更多，比如精度提升0.9%；重量型ViT-B則受益很少。

Comparison with Prior Art

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

上圖給出了所提方案與ResNet系列+其他自監(jiān)督學(xué)習(xí)的性能對(duì)比?？梢钥吹剑?/p>

• 在小模型方面，本文基線ViT MoCo（即ViT, MoCoV3）的性能要比SimCLRv2+ResNet更佳；在大模型方面則與之相當(dāng)。
• SimCLRV2+SK-ResNet的組合具有更高的性能；
• BYOL+Wider-ResNet具有更高的性能，配合R200-2x時(shí)可以得到更優(yōu)異的結(jié)果；
• 正如前面提到的規(guī)范化技術(shù)的影響，將LN替換為BN后模型的性能提升了1%；
• 將ViT默認(rèn)塊尺寸替換為后，模型的性能進(jìn)一步提升2-3%。MoCoV3+ViT-BN-L/取得了81.0%的top1精度，作為對(duì)比，SimCLRV2+SK-ResNet152-3x的最佳精度為79.8%，BYOL-ResNet200-2x的最佳精度為79.6%。

Transfer Leanring

最后，我們?cè)倏匆幌滤岱桨冈谙掠稳蝿?wù)遷移學(xué)習(xí)方面的性能。結(jié)果見上表，可以看到：

• 當(dāng)模型大小從ViT-B提升到ViT-L，所提方案具有更加的遷移學(xué)習(xí)精度；當(dāng)提升到ViT-H時(shí)則出現(xiàn)了過(guò)擬合問(wèn)題。作為對(duì)比，ImageNet監(jiān)督ViT在ViT-L時(shí)就出現(xiàn)了過(guò)擬合。
• 相比ImageNet監(jiān)督方案，本文所提自監(jiān)督ViT取得了更佳的結(jié)果。
• 在這些小數(shù)據(jù)上，采用大的ViT模型從頭開始訓(xùn)練時(shí)過(guò)擬合問(wèn)題非常嚴(yán)重。這意味著：如果數(shù)據(jù)量不足會(huì)導(dǎo)致難以訓(xùn)練ViT學(xué)習(xí)好的特征表達(dá)；而自監(jiān)督預(yù)訓(xùn)練可以彌補(bǔ)這種差距，極大的避免小數(shù)據(jù)集上的過(guò)擬合問(wèn)題。