真不是煉丹,務實敢爲的 MoCo v3

作者 | 黃掛  

編輯 | NewBeeNLP

這次分享的 paper 是 An Empirical Study of Training Self-Supervised Visual Transformers[1] ,前不久剛掛出來了的,可謂新鮮滾熱辣了。

這篇 paper 是 FAIR 何凱明的 MoCo-v3,行文還是很務實的,開篇第一句:“This paper does not describe a novel method”,明講沒有新方法,只是手把手教你如何訓練 visual transformer 的。所以這篇 paper 更像是一篇實驗報告。

Take Away

因爲我也經常看別人的論文閱讀筆記,其實不喜歡那種整篇順一遍幾乎沒有提煉的筆記。但其實會議通貨膨脹,很容易遇到水文,我想每篇閱讀筆記最前面都寫一些 take away,大家可以最方便快捷的瞭解到自己需不需要繼續瞭解這篇文章。

不過我總是建議大家多讀讀 paper,畢竟作者能把實驗思想寫成 paper 發表,再水也總是有可取之處的。下面是這篇文章的 take away

  1. 水文指數(滿分 5):⭐️(十分值得一讀

  2. ViT 這個結構最近很火,但自監督訓練會有不穩定的情況,可以通過凍結住第一層的 patch 層(隨機參數)來緩解。

  3. 如何 debug 自監督模型:kNN curve + 監控 gradient magnitude。

  4. 拋開文章提出的 idea 本身,我覺得作者在 paper 裏體現的思考過程是更可學習的,體現出頂級研究者的不同之處,也是我對這篇 paper 評價很高的原因。畢竟授人以魚不如授人以漁。大家可以往下看細品,具體下面 “如何思考” 的小節。

簡介

這篇文章就講一個事情:「如何訓練一個 visual transformer」

迴歸最基本的一些東西:learning rate、batch size、optimizer 這些。他們是從 self-supervised 的角度來研究 ViT 的訓練的,「發現關鍵因素是 instability」。也就是 ViT 的訓練可能會容易不穩定。

而且有趣的是,不穩定的 ViT 並不會崩掉,而是隻會在準確率上有一個微小的降低(比如 1-3% 個點),所以不好發現。CNN 訓練的時候,不太會有這種問題。他們爲了解決不穩定的問題,「將 ViT 的最開始的給圖片分塊的層 patch projction 固定住了」,並且實驗上看是確實會更好的。

這點我在想是不是 transformer 本身就是一個比較 fragile 的結構,我們最近在訓 ViT 的時候,也發現很容易炸。再往前追溯,在 19 年復現 Bert 的時候,也花了很多精力來調參,「可能 Transformer 是一個超參敏感的結構」

MoCo v3

作者把這個工作稱爲 MoCo v3 了,其實沒有新的東西,就是有個名字好聽點?可以方便稱呼這個工作。用的還是一樣的套路,一張圖片,做兩個 random data augmentation,然後過兩個 encoder,然後算一個 InfoNCE 的 loss,這部分有興趣的朋友可以找一下 SimCLR、MoCo v1 v2 等 paper 來看看。

InfoNCE 的公式:

MoCo v3 的 僞代碼:

MoCo v3 在 Resnet 的 setting 下,也還是比 MoCo v1 v2 好一些的。「主要原因是什麼呢?」 下圖是在 resnet setting 下,MoCo v3 和 v2 的對比

自監督訓練 ViT 的穩定性


這部分開始比較有意思了,開始講怎麼訓練一個穩定的 ViT。作者的說法是,ViT 的不穩定,不會體現 accuracy 上,實際上不穩定的 ViT 也能拿到一個看起來還湊合的準確率。這就比較難辦了。然而頂級研究者還是 nb,「他們用 kNN curve[2] 來監控訓練是否穩定」

Batch size & learning rate

Batch Size 和 Learning Rate 是一對好夥伴,要一起看。因爲一般 learning rate 都會根據 batch size 做一些 scale 的操作。這裏也不例外,用的是最經典的 setting:「lr x batch_size / 256」

在 base LR 一定的情況下,batch size 越大,反而越挫。4096 和 6144 會比 2048 這些要差。在 batch size 一定的情況下(4096),learning rate 太大不行,太小也不行。這個就是 overfit 和 underfit 的問題了。

lr 一定,效果最好的是 batch=2048,越大越挫

batch size 一定,Lr 也比較敏感,在 1e-4 的時候最好

optimizer

一般都是會選 AdamW 來做優化器,作者還比較了一下 LARS,在好好調了一下,能微弱漲 0.1 個點。「但是 LARS 對 LR 太敏感了,調一下很容易挫」,所以最後作者還是選了 AdamW。我們之前還試過 Adam 和 SGD,在 ViT 的 setting 下也是比 AdamW 挫一些。值得一提的是,Resnet 的 setting 下,SGD 一直是 yyds。

kNN Curve

這一部分作者沒有細講,但我挺感興趣的,翻了一下對應的論文。其實就是一個無需訓練的分類指標。

做法就是,每個 epoch 結束之後,用當前的參數對整個 test set 進行編碼得到 embedding,然後每條數據和所有數據都算 cosine 相似度,召回 top k,做一個 weighted 的投票來預測當前這條數據的類別。

「用 kNN classifier Curve 的好處就是,可以極其高效的看模型的效果」。可能會問,爲什麼不直接看 train 的指標呢?因爲 train 的指標容易過擬合。其實在有監督的 setting 下,是沒有這個問題的,因爲有監督的 setting 下,每個 epoch 結束可以直接看 val 集的 acc。但這裏是自監督的 setting,沒法直接看下游的 acc。

Tricks on Stability

OK,講完 ViT 訓練不穩定之後,一個自然的問題就是,怎麼讓他穩定啊。一般遇到這麼虛的問題,感覺都挺不知所措的。這裏開始展現作者的功力了,作者將訓練時每一個 iteration 的網絡的每一層的 gradient 都記錄了下來,如下圖。「發現在訓練的過程中,梯度會有一些尖刺。而且往往是第一層先出現,然後過幾十個 iter 之後,最後一層會出現。」

ViT 訓練時 第一層和最後一層的 gradient 可視化。發現第一層出現尖刺後,幾十個 iter 之後最後一層也會出現尖刺。

在這個觀察下,作者提出了一個假設:「訓練不穩定是發生在訓練時前面幾層淺層網絡上的,然後才慢慢衍生到後面層。並提出一個大單的做法:將第一層 patch projection 層固定住!」 看到這我還是挺喫驚的,畢竟在沒訓練的情況下,第一層 patch projection 是隨機初始化的,這是要把一個隨機初始化的參數凍結住,讓網絡學習一個隨機層映射出來的表示??好大膽的想法。結果證明,確實凍結住第一層會可以穩定訓練,如下圖:

凍結住隨機參數第一層,訓就十分穩定了

作者還嘗試了其他的自監督方法,SimCLR、BYOL,fix random patch 這個 trick 也是 work 的。還在 SwAV 上試了一下,SwAV 的好處是,他的不穩定是直接 nan loss(這個挺好的),然後這個 trick 依舊 work。總而言之,「fix random patch 對所有自監督框架 + ViT setting 都有效」

作者還試了可能能增加穩定性的方法,比如 BatchNorm、WeightNorm、gradient clip。前兩個 norm 沒有用,gradient clip 在很小的情況下有用,但很小的情況下就等價於 fixed 住了。

如何思考

fix random patch 這個 trick 看起來,確實不太直觀而且很大膽,畢竟是把一個隨機參數凍結住了... 你說你先隨便怎麼訓練一下,然後再凍結一個訓過的參數也好啊... 當然了,作者也說了,「fix random patch 不能解決訓練不穩定的問題,僅僅只是緩解。」 而 fix 第一層,也僅僅只是因爲他是唯一的非 transformer 的層。看來還是任重而道遠。

拋開 fix random patch 這個 idea 本身,讓我感悟比較深的是,這個 idea 作者是如何想到的,因爲我在細品這個 idea 之後,覺得自己決對想不出這樣的點子,就算想出來了也不會實施,因爲凍結一個隨機參數聽起來實在不靠譜。這裏就有認知差距了,我們捋一下作者是如何提出這個觀點的:

  1. 「如何 debug 模型」:因爲是自監督訓練 ViT,沒法直接看效果不好 debug。所以想到了用 kNN classifier Curve 的方法來看中間效果。發現訓練中間確實會有偶爾 acc 猛掉的情況。

  2. 「如何觀測」:那我們看看訓練的 gradient 怎麼樣,於是把訓練中網絡各個層的 gradient 可視化,發現確實會偶然炸毛一下。發現第一層炸毛之後幾十個 iter 後最後一層會炸毛。

  3. 「提出假設」:那麼自然的假設就是,第一層的參數可能是導致訓練不穩定的原因。

  4. 「設計實驗」:既然第一層是 “罪魁禍首”,那就試試把他凍結住會怎麼樣了。於是就有了這篇文章。

把整個推理過程捋順,你會發現,這個蹊蹺的 trick 其實是很 solid 的嘗試。可謂心思縝密,膽大心細,務實敢爲,不設限,令人佩服啊。讓我不禁感慨,深度學習真不是煉丹,是有邏輯的。

其他細節

還有一些實現細節,把一些有啓發的點記錄一下。

  1. 作者對 lr 和 weight decay 做了超參搜索,用 100 個 epoch 的結果,搜到最好的再訓更久。

  2. 用的 40 epoch warmup + cosine LR decay。

  3. cosine position embedding 比 learned position embedding 好 0.4 個點,比不加只有 1.6 個點

  4. ViT-B 的 setting,128TPU 跑 100 個 epoch 僅需 2.1h。有錢真好。他們的 Pytorch 多 GPU 實現,24 小時 + 128GPU 訓練 100 個 epoch。ViT-B/16 的設定。是一個驗證代碼訓練性能的 benchmark 了。而且這麼看,TPU 比 GPU 真的高效很多?

  5. 小模型 ViT-S 訓練更久效果會好,大模型 ViT-B 訓練更久效果提升不那麼顯著

  6. 模型越大效果提升不明顯,和 NLP 上不一樣。一個解決辦法當然是增加數據,但作者也覺得,有可能是因爲當前的 instance level 的 pretext task 太容易了。雖然已經論證了當前這種 Constractive 的任務效果要比 auto encoding 的好,但可能還有更好的 vision 任務呢。

寫在最後

這篇 paper 整體看下來,比較值得一提的點就是 fix random patch 了。第一次初略看完 paper 的感慨是,現在發 paper 真簡單啊,就這麼一個簡單的 trick 就能寫成 paper 了?第二次細看完 paper 的感慨是,這樣的 trick 也不是一般人敢想能做出來的,還是厚積薄發。而且看看人家實現 setting 裏動輒 256 個 TPU 的數量,還是很貴的。

不過不管怎麼說,這篇 paper 還是很值得一讀的。「務實敢爲,務實在推論過程嚴謹,敢爲在凍結住隨機參數這種騷操作都幹做。」

順着這個思路再延伸一下,其實看 paper 看到後面,一是從中汲取一些有啓發的 idea,第二也應該學習一下這些頂級研究者們是如何思考問題的。

拋開文章本身,「更有趣的是作者是如何嚴謹的一步步進行觀測做出假設設計實驗驗證觀點的」。授人以魚不如授人以漁,汲取文章觀點只是魚,學習別人如何思考問題是漁啊。

idea 總會被新的 idea 打敗,但找 idea 的方法總不過時。

一起交流

想和你一起學習進步!『**NewBeeNLP』**目前已經建立了多個不同方向交流羣(**機器學習 / 深度學習 / 自然語言處理 / 搜索推薦 / 圖網絡 / 面試交流 / **等),關注公衆號回覆『入羣』加入吧!

本文參考資料

[1]

An Empirical Study of Training Self-Supervised Visual Transformers: https://link.zhihu.com/?target=https%3A//arxiv.org/abs/2104.02057v1

[2]

kNN curve: https://arxiv.org/pdf/1805.01978.pdf

本文由 Readfog 進行 AMP 轉碼,版權歸原作者所有。
來源https://mp.weixin.qq.com/s/PdmiySjVSN7n0DLJBCNffg