北大博士生提出CAE,下游任務泛化能力優於何愷明MAE

對白的演算法屋發表於2022-02-23

大家好,我是對白。

何愷明時隔兩年發一作論文,提出了一種視覺自監督學習新正規化——

用掩蔽自編碼器MAE,為視覺大模型開路。

這一次,北大博士生提出一個新方法CAE,在其下游任務中展現的泛化能力超過了MAE。

圖片

來看看這是一項什麼樣的研究?

這是一項什麼研究?

自何愷明提出MAE以來,基於MIM,Masked Image Modeling,這一自監督學習表徵演算法就越來越引發關注。

它的主要思路,就是對輸入影像進行分塊和隨機掩碼操作,然後對掩碼區域做預測。

預測的目標可以是Token ID(如微軟提出的BEiT),也可以是RGB值(如MAE)。

通過MIM這一方法,編碼器學習到好的表徵,從而在下游任務中取得良好的泛化效能。

以往這一方法常見於NLP領域,但隨著ViT的提出和發展,這一方法在視覺領域也取得了一些進展。

團隊認為,近期兩個代表性工作,BEiT和MAE,沒有充分挖掘編碼器encoder的潛力,限制了預訓練學習的表徵質量。

簡單來說,BEiT的編碼器只有一部分負責表徵學習,還有一部分在做“前置/代理任務”(pretext task)。

到了MAE則是另一種情況,解碼器也做了一部分表徵學習,可能讓編碼器學會“偷懶”。

基於這一背景,團隊提出了Context Autoencoder,簡稱CAE。核心設計思想是對“表徵學習”和“前置/代理任務(pretext task)這兩個功能做分離。

在預訓練時,編碼器只負責表徵學習,解碼器只負責解決前置任務,兩者分工合作,將編碼器表徵能力最大化。

CAE包含四個部分。

圖片

1、編碼器是一個ViT模型,負責學習影像可見patch的表徵,提取影像特徵Zv。

2、Latent contextual regressor(隱式上下文迴歸器)則是在此基礎上預測出掩蔽patch的表徵Zm。

3、解碼器以Zm和對應的位置編碼作為輸入,通過Zm預測掩蔽patch的某些性質,比如RGB值、Token ID。這過程中Zv不會更新,表徵學習任務全交給編碼器。

4、Latent representation alignment對 Zm新增約束,希望 latent contextual regressor 的輸出和編碼器的輸出在同一空間。具體來說,影像的掩蔽patch也將輸入到編碼器中(此過程不參與梯度反傳),獲得的這部分表徵,作為Zm的學習目標。

Alignment很重要,如果想對掩碼部分做好預測,要求latent contextual regressor 的輸出(也是解碼器的輸入)含有比較好的語義資訊。通過對齊操作,可以鼓勵編碼器的操作也含有好的語義資訊,提高編碼器的表徵質量。

論文對alignment做了視覺化:將全部patch輸入編碼器,然後將表徵直接輸入到解碼器中,進行RGB的重建。CAE可以將原圖重建出來 (第一行是原圖,第二行是重建結果),說明編碼器的輸出和latent contextual regressor 的輸出處於同一編碼空間。

圖片

如果訓練時不做alignment約束,那麼輸出的結果將是這樣…嗯,都是亂碼。

圖片

這種設計的編碼器學到的表徵也相對更差,下游任務結果也會變差。

損失函式由兩部分組成,一個是對解碼器預測的監督,使用的是cross-entropy loss;一個是對alignment的監督,使用MSE損失函式。

圖片

除此之外,也進一步驗證了以CAE為代表的MIM方法,要比Moco v3、DINO為代表的對比學習方法更適合下游任務。

該論文從隨機裁剪操作的性質分析,認為隨機裁剪有很大概率包含影像的中心區域。

ImageNet-1K這種資料集中,中心區域通常是1000類標籤集中的物體(如下圖)。因此,對比學習方法主要提取影像中主體物體的特徵。

圖片

MIM方法能學到每個patch的特徵,包括影像的背景區域,而不僅僅是影像主體物體,這讓MIM學到的表徵更適合下游檢測分割任務。

論文對CAE和MoCo v3的注意力圖做了視覺化。紅色表示注意力值更高,藍色表示注意力值更低。第一行是原圖,第二行是 MoCo v3,第三行是 CAE。可以看到,MoCo v3 的注意力圖主要在影像的主體區域有高響應,而 CAE 能考慮到幾乎所有patch。

圖片

實驗結果

研究團隊使用ViT-small和ViT-base在 ImageNet-1K 上進行實驗,輸入影像的解析度224*224,每張圖被分成14*14的patch,每個patch的大小為16*16。

每次將有75個patch被隨機掩碼,其餘patch則為可見的。

本文參照BEiT,使用DALL-E tokenizer對輸入影像token化,得到預測目標。

最終結果顯示,在語義分割任務中,跟其他MIM方法,比如MAE、BEiT,以及對比學習、有監督預訓練方法的表徵結果更好。

圖片

在物體檢測、例項分割的結果也是如此。

圖片

圖片

 

最後歡迎大家關注我的微信公眾號:對白的演算法屋duibainotes),跟蹤NLP、推薦系統和對比學習等機器學習領域前沿,日常還會分享我的創業心得和人生感悟。想進一步交流的同學也可以通過公眾號加我的微信,和我一同探討技術問題,謝謝!

相關文章