文 | happy
源 | 極市平臺

論文鏈接:
https://arxiv.org/abs/2104.02057

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

Abstract

本文并沒有提出一種新的方法，相反，鑒于最近計(jì)算機(jī)視覺的進(jìn)展，我們研究了一個簡單、漸進(jìn)、但必須知道的基線：用于視覺Transformer的自監(jiān)督學(xué)習(xí)。盡管標(biāo)準(zhǔn)卷積網(wǎng)絡(luò)的訓(xùn)練方法已經(jīng)非常成熟且魯棒，然而ViT的訓(xùn)練方案仍有待于構(gòu)建，特別是自監(jiān)督場景下的訓(xùn)練極具挑戰(zhàn)。在這里，我們從基礎(chǔ)出發(fā)，對訓(xùn)練自監(jiān)督ViT的幾種基本組件的影響進(jìn)行了分析調(diào)研。我們發(fā)現(xiàn)：不穩(wěn)定性是影響精確下降的最主要問題，它會被表面上好的結(jié)果覆蓋(容易陷入局部最優(yōu))。我們通過實(shí)驗(yàn)發(fā)現(xiàn)：這些結(jié)果確實(shí)存在部分失敗；當(dāng)訓(xùn)練變得穩(wěn)定時，這些結(jié)果可以進(jìn)一步提升。基于MoCoV3以及其他自監(jiān)督框架，我們從不同角度對ViT進(jìn)行了測試分析；我們對觀察到的積極面、挑戰(zhàn)性以及開放問題進(jìn)行了討論，期望該工作可以為未來的研究提供有用的數(shù)據(jù)支撐和經(jīng)驗(yàn)參考。

Introduction

本文主要聚焦于：采用視覺領(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)練的幾個基本模塊：batch size，learning rate以及optimizer。我們發(fā)現(xiàn)：在不同場景下，不穩(wěn)定性均是影響自監(jiān)督ViT訓(xùn)練的主要問題。有意思的是，我們發(fā)現(xiàn)：不穩(wěn)定的ViT訓(xùn)練可能不會導(dǎo)致災(zāi)難性結(jié)果(比如發(fā)散)；相反，它可以導(dǎo)致精度的輕度退化(約下降1-3%)。除非有一個更穩(wěn)定的作為對比，否則這個程度的退化可能難以被注意到。據(jù)我們所知，該現(xiàn)象在卷積網(wǎng)絡(luò)訓(xùn)練中鮮少發(fā)生，我們認(rèn)為該問題及其隱含的退化值得注意。為了證明不穩(wěn)定性的可能危害，我們研究了一個可以在實(shí)踐中提高穩(wěn)定性的簡單技巧。基于梯度變換的經(jīng)驗(yàn)觀察，我們固化ViT中的塊投影層，即采用固定隨機(jī)塊投影。我們發(fā)現(xiàn)該trick可以緩解多種場景下的不穩(wěn)定問題并提升模型精度。基于對比學(xué)習(xí)框架，自監(jiān)督Transformer可以取得非常好的結(jié)果。不同于ImageNet監(jiān)督的ViT(模型變大時精度反而變差)，更大的自監(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無需位置信息即可學(xué)習(xí)很強(qiáng)的特征表達(dá)，同時也也暗示位置信息并未得到充分探索。

MoCoV3

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

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

其中，表示q同源圖像的輸出，即正樣本；表示異源圖像輸出，即負(fù)樣本。延續(xù)SimCLR處理方式，MoCoV3采用同一批次自然共存的密鑰，移除了Memory Queue(當(dāng)batch足夠大時，其收益遞減)。基于這種簡化，上面的對比損失可以簡化為上述偽代碼中ctr實(shí)現(xiàn)。我們采用了對稱損失：。我們的編碼器由骨干網(wǎng)絡(luò)(如ResNet、ViT)、投影頭以及額外的預(yù)測頭構(gòu)成；而編碼則由骨干網(wǎng)絡(luò)、投影頭構(gòu)成，沒有預(yù)測頭。通過的滑動平均更新。作為參考，我們以ResNet50作為骨干網(wǎng)絡(luò)，其在ImageNet上的性能如下，由于額外的預(yù)測頭與大的batch，MoCoV3具有更加性能。

Stability of Self-Supervised ViT Training

原則上來講，我們可以直接在對比自監(jiān)督學(xué)習(xí)框架中采用ViT骨干替換ResNet骨干網(wǎng)絡(luò)。但實(shí)際上，主要挑戰(zhàn)在于：訓(xùn)練不穩(wěn)定。我們發(fā)現(xiàn)：不穩(wěn)定問題不僅僅由精度反映。實(shí)際上，即使存在不穩(wěn)定問題，它也可以“表面上很好”的進(jìn)行訓(xùn)練并提供一個不錯的結(jié)果。為揭示這種不穩(wěn)定性，我們在訓(xùn)練過程中對kNN曲線進(jìn)行了監(jiān)控，研究了它如何影響不穩(wěn)定性并提出了一種簡單的trick進(jìn)行穩(wěn)定訓(xùn)練，進(jìn)而提升不同場景下的精度。

Empirical Observations on Basic Factors

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

• BatchSize。上圖給出了不同batch時的訓(xùn)練曲線。可以看到：

• • 1k和2k的batch具有比較平滑的曲線，精度分別為71.5%和72.6%；

  • 4k的曲線就開始變得不再穩(wěn)定，其最終精度為72.2%，要低于2k的72.6%。后面會提到：這里的性能下降是受不穩(wěn)定訓(xùn)練導(dǎo)致。

  • 6k的曲線的不穩(wěn)定性進(jìn)一步加劇，最終仍有一個不錯的結(jié)果69.7%。我們猜測：訓(xùn)練被部分重啟，跳出了當(dāng)前局部最優(yōu)并重新尋找新的優(yōu)化軌跡。因此，訓(xùn)練無法收斂，最終的精度會依賴于重啟的局部性能。

  • 此外，我們還發(fā)現(xiàn)這種不穩(wěn)定不會導(dǎo)致一個顯著差的結(jié)果。實(shí)驗(yàn)發(fā)現(xiàn)：相同配置下，結(jié)果的差異較小，范圍內(nèi)。這就使得不穩(wěn)定導(dǎo)致了性能退化難以被注意到。

• LearningRate。事實(shí)上，學(xué)習(xí)率通常會隨batch提升而縮放。在本文實(shí)驗(yàn)中，我們采用了線性縮放規(guī)則：。上圖給出了學(xué)習(xí)率的影響曲線。當(dāng)lr比較小時，訓(xùn)練更為穩(wěn)定，但會欠擬合。比如0.5e-4訓(xùn)練的模型精度要比1.0e-4訓(xùn)練的模型精度低1.8%。當(dāng)lr比較大時，訓(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í)率大時，模型精度會迅速下降。有意思的是：此時的訓(xùn)練曲線仍然很平滑，但會緩慢下降。我們發(fā)現(xiàn)：當(dāng)學(xué)習(xí)率合適時，LAMB可以取得與AdamW相當(dāng)?shù)木取５菍W(xué)習(xí)率的敏感性使其無法在進(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)定是主要問題。接下來，我們將提出一種簡單的trick提升不同場景下的穩(wěn)定性。

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

上圖對比了可學(xué)習(xí)與隨機(jī)塊投影的MoCoV3結(jié)果對比。可以看到：隨機(jī)塊投影可以穩(wěn)定訓(xùn)練，訓(xùn)練曲線更為平滑、精度更高(精度提升約1.7%)，進(jìn)一步說明了訓(xùn)練不穩(wěn)定是影響精度的主要問題。

我們發(fā)現(xiàn)：除了MoCo外，其他相關(guān)方法(比如SimCLR、BYOL)同樣存在不穩(wěn)定問題。隨機(jī)塊投影同樣可以改善SimCLR與BYOL的性能(分別提升0.8%和1.3%)，見上圖對比。不穩(wěn)定對于SwAV會導(dǎo)致?lián)p失發(fā)散(NaN)，本文所提隨機(jī)塊投影可以穩(wěn)定訓(xùn)練SwAV并將其精度由65.8%提升到66.4%。總而言之，本文所提t(yī)rick對于所有自監(jiān)督方案均有效。Discussion：一個有意思的發(fā)現(xiàn)：塊投影層的訓(xùn)練并非必要。對于標(biāo)準(zhǔn)ViT塊尺寸，其塊投影矩陣時過完備的，此時隨機(jī)投影足以保持原始塊的信息。我們注意到：凍結(jié)第一層并不會改變架構(gòu)，但它會縮小解決方案空間。這意味著：根本問題與優(yōu)化相關(guān)。該trick可以緩解優(yōu)化問題，但不能解決它；當(dāng)lr過于大時模型仍存在不穩(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，通過100epoch訓(xùn)練搜索lr和wd，然后選擇最優(yōu)訓(xùn)練更長。學(xué)習(xí)率采用40epoch進(jìn)行warmup，它同樣有助于緩解不穩(wěn)定性；warmup之后，學(xué)習(xí)率按照cosine衰減；

• MLP Head：投影頭是一個3層MLP；預(yù)測頭是一個2層MLP。MLP的隱含層維度均為4096，輸出層維度均為256。在MoCoV3中，參考SimCLR，MLP的所有Linear均后接BN。

• Loss：參考BYOL，損失函數(shù)進(jìn)行了縮放。盡管該方法可以通過調(diào)節(jié)lr和wd合并，但可以使得其對的敏感性降低，默認(rèn)。

• ViT Architecture：參考ViT一文，輸入塊為或者，經(jīng)過投影后它將輸出一個長度為196/256的序列，Position Embedding與該序列相加；所得序列與可學(xué)習(xí)類token拼接并經(jīng)由后續(xù)Transformer模塊編碼；最終所得類token視作輸出并送入MLP頭。

• Linear probing。延續(xù)常規(guī)方案，我們采用線性方式評估特征表達(dá)質(zhì)量。完成自監(jiān)督與訓(xùn)練后，移除MLP頭并采用監(jiān)督方式訓(xùn)練一個線性分類器，此時訓(xùn)練90epoch，且僅僅采用RandomResize、Flipping進(jìn)行數(shù)據(jù)增廣。

Experiments Results

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

Self-supervised Learning Framework

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

Ablations of ViT+MoCoV3

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

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

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

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

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

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

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

Comparison with Prior Art

上表給出了MoCoV3框架下不同ViT模型的對比。可以看到：(1) 相比iGPUT，無需額外數(shù)據(jù)預(yù)訓(xùn)練，基于MoCo的ViT取得了更高的精度、更小的模型。(2) 隨著模型變大，所提方案的精度會逐漸提升；而在監(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í)工具不易于過擬合。

上圖給出了所提方案與ResNet系列+其他自監(jiān)督學(xué)習(xí)的性能對比。可以看到：

• 在小模型方面，本文基線ViT MoCo（即ViT, MoCoV3）的性能要比SimCLRv2+ResNet更佳；在大模型方面則與之相當(dāng)。

• SimCLRV2+SK-ResNet的組合具有更高的性能；

• BYOL+Wider-ResNet具有更高的性能，配合R200-2x時可以得到更優(yōu)異的結(jié)果；

• 正如前面提到的規(guī)范化技術(shù)的影響，將LN替換為BN后模型的性能提升了1%；

• 將ViT默認(rèn)塊尺寸替換為后，模型的性能進(jìn)一步提升2-3%。MoCoV3+ViT-BN-L/取得了81.0%的top1精度，作為對比，SimCLRV2+SK-ResNet152-3x的最佳精度為79.8%，BYOL-ResNet200-2x的最佳精度為79.6%。

Transfer Leanring

最后，我們再看一下所提方案在下游任務(wù)遷移學(xué)習(xí)方面的性能。結(jié)果見上表，可以看到：

• 當(dāng)模型大小從ViT-B提升到ViT-L，所提方案具有更加的遷移學(xué)習(xí)精度；當(dāng)提升到ViT-H時則出現(xiàn)了過擬合問題。作為對比，ImageNet監(jiān)督ViT在ViT-L時就出現(xiàn)了過擬合。

• 相比ImageNet監(jiān)督方案，本文所提自監(jiān)督ViT取得了更佳的結(jié)果。

• 在這些小數(shù)據(jù)上，采用大的ViT模型從頭開始訓(xùn)練時過擬合問題非常嚴(yán)重。這意味著：如果數(shù)據(jù)量不足會導(dǎo)致難以訓(xùn)練ViT學(xué)習(xí)好的特征表達(dá)；而自監(jiān)督預(yù)訓(xùn)練可以彌補(bǔ)這種差距，極大的避免小數(shù)據(jù)集上的過擬合問題。

后臺回復(fù)關(guān)鍵詞【入群】

加入賣萌屋NLP/IR/Rec與求職討論群

后臺回復(fù)關(guān)鍵詞【頂會】

獲取ACL、CIKM等各大頂會論文集！

總結(jié)

以上是生活随笔為你收集整理的MoCoV3：何恺明团队新作！解决Transformer自监督训练不稳定问题！的全部內(nèi)容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網(wǎng)站內(nèi)容還不錯，歡迎將生活随笔推薦給好友。