MoCo的基本原理,包括其歷史來龍去脈在前文中[1,2,3]中已經介紹的比較充足了,本文就不再進行贅述。本文主要介紹下MoCo v3 [4]中的一些新發現。MoCo v3中并沒有對模型或者MoCo機制進行改動,而是探索基于Transformer的ViT(Visual Transformer)模型[5,6]在MoCo機制下的表現以及一些訓練經驗。作者發現ViT在采用MoCo機制的訓練過程中,很容易出現不穩定的情況,并且這個不穩定的現象受到了學習率,batch size和優化器的影響。如Fig 1.所示,在batch size大于4096的時候已經出現了明顯的劇烈抖動,如Table 1.所示,我們發現在bs=2048時候取得了最好的測試性能,繼續增大batch size反而有很大的負面影響,這個結論和MoCo v1里面『batch size越大,對比學習效果越好』相悖,如Fig 2.所示。這里面的大幅度訓練抖動肯定是導致這個結論相悖的罪魁禍首。這個抖動并不容易發現,因為在bs=4096時候,模型訓練最終能收斂到和bs=1024,2048相同的水平,但是泛化效果確實會存在差別。
Fig 2. 在MoCo v1中,隨著batch size的增大,對比學習的結果也逐漸變好。不僅僅是batch size,學習率也會導致ViT訓練的不穩定,如Fig 3.所示,我們發現在較大的學習率下訓練曲線存在明顯的抖動,而最終的訓練收斂位置卻差別不大。在測試結果上看,則會受到很大的影響。如果將優化器從AdamW更換到LAMB優化器,那么結果也是類似的,如Fig 4.所示,只是可以采用更大的學習率進行訓練了。
Fig 3. ViT以MoCo機制訓練,在不同學習率下的訓練曲線和對應測試結果。
Fig 4. ViT以MoCo機制訓練,在采用LAMB優化器的情況下,不同學習率下的訓練曲線和對應測試結果。這種出現訓練時的劇烈抖動,很可能是梯度劇變導致的,因此作者對ViT的第一層和最后一層的梯度的無窮范數進行了統計。注意到無窮范數相當于求所有梯度值中絕對值的最大值,也即是如(1)所示。結果如圖Fig 5.所示,我們發現的確會存在有梯度的驟變,而且總是第一層先發生,然后經過約數十個step之后傳遞給了最后一層。因此,導致訓練曲線劇烈抖動的原因可能是ViT的Transformer的第一層梯度不穩定導致。
Fig 5. ViT在訓練過程中第一層和最后一層的梯度無窮范數。
考慮到在ViT中的第一層是將patch映射到visual token,也就是一層FC全連接層,如圖Fig6.所示。作者在MoCo v3里面的做法也很直接,直接將ViT的第一層,也即是從Patch到Visual Token的線性映射層隨機初始化后固定住,不參與訓練。
Fig 6. 在ViT中通過FC層將圖片patch線性映射到了visual token,從而輸入到Transformer。
這個做法挺奇怪的,但是實驗結果表明在固定住了線性映射層之后,的確ViT的訓練穩定多了,如Fig 7.所示,訓練曲線的確不再出現詭異的劇烈抖動,最主要的是其測試結果也能隨著學習率的提高而增大了,并且同比learned path proj.的情況還更高。
Fig 7. 在固定住了patch映射到visual token的線性映射層之后,訓練曲線不再出現明顯的劇烈抖動。
這種現象還是蠻奇怪的,也就是說即便不訓練這個patch projection layer,模型的性能也不會打折,而且還會更加穩定。作者給出的解釋就是目前這個映射是完備的(complete),甚至是過完備(over-complete)的,以的patch,768的visual token為例子,那么這個映射矩陣就是
的。也就是說對于所有可能的patch來說,可能在隨機的M \mathbf{M}M?????中就已經有著近似的唯一輸出對應,即便這個映射可能不保留太多的視覺語義信息,但是也保留了原始的視覺信息,不至于損失原始信息。但是正如作者最后所說的,這個『trick』只是緩解了問題,但是并沒有解決問題,顯然這個問題出現在了優化階段,而固定FC層減少了解空間提高了其穩定性。在更大的學習率下,還是會受到相同的不穩定現象,對該現象的研究值得繼續深究。
筆者在大規模的對比學習訓練過程中也遇到過類似的訓練曲線抖動(雖然沒有那么劇烈),但是筆者發現可能是溫度系數的劇烈變化導致的,我們以后再繼續討論下溫度系數的影響。
Reference
[1]. https://fesian.blog.csdn.net/article/details/119515146
[2]. https://fesian.blog.csdn.net/article/details/120039316
[3]. https://fesian.blog.csdn.net/article/details/120364242
[4]. Chen, Xinlei, Saining Xie, and Kaiming He. “An empirical study of training self-supervised vision transformers.” arXiv preprint arXiv:2104.02057 (2021).
[5]. https://blog.csdn.net/LoseInVain/article/details/116031656
[6]. Dosovitskiy, Alexey, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani et al. “An image is worth 16x16 words: Transformers for image recognition at scale.” arXiv preprint arXiv:2010.11929 (2020).