北大博士生提出CAE，下游任務(wù)泛化能力優(yōu)于何愷明MAE

白交 2022-02-22 12:52:52 來源：量子位

楊凈 發(fā)自 凹非寺

量子位 | 公眾號 QbitAI

上回說道，何愷明時隔兩年發(fā)一作論文，提出了一種視覺自監(jiān)督學(xué)習(xí)新范式——

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

這一次，北大博士生提出一個新方法CAE，在其下游任務(wù)中展現(xiàn)的泛化能力超過了MAE。

來看看這是一項什么樣的研究？

這是一項什么研究？

自何愷明提出MAE以來，基于MIM，Masked Image Modeling，這一自監(jiān)督學(xué)習(xí)表征算法就越來越引發(fā)關(guān)注。

它的主要思路，就是對輸入圖像進行分塊和隨機掩碼操作，然后對掩碼區(qū)域做預(yù)測。

預(yù)測的目標(biāo)可以是Token ID（如微軟提出的BEiT），也可以是RGB值（如MAE）。

通過MIM這一方法，編碼器學(xué)習(xí)到好的表征，從而在下游任務(wù)中取得良好的泛化性能。

以往這一方法常見于NLP領(lǐng)域，但隨著ViT的提出和發(fā)展，這一方法在視覺領(lǐng)域也取得了一些進展。

團隊認為，近期兩個代表性工作，BEiT和MAE，沒有充分挖掘編碼器encoder的潛力，限制了預(yù)訓(xùn)練學(xué)習(xí)的表征質(zhì)量。

簡單來說，BEiT的編碼器只有一部分負責(zé)表征學(xué)習(xí)，還有一部分在做“前置/代理任務(wù)”（pretext task）。

到了MAE則是另一種情況，解碼器也做了一部分表征學(xué)習(xí)，可能讓編碼器學(xué)會“偷懶”。

基于這一背景，團隊提出了Context Autoencoder，簡稱CAE。核心設(shè)計思想是對“表征學(xué)習(xí)”和“前置/代理任務(wù)（pretext task）這兩個功能做分離。

在預(yù)訓(xùn)練時，編碼器只負責(zé)表征學(xué)習(xí)，解碼器只負責(zé)解決前置任務(wù)，兩者分工合作，將編碼器表征能力最大化。

CAE包含四個部分。

1、編碼器是一個ViT模型，負責(zé)學(xué)習(xí)圖像可見patch的表征，提取圖像特征Zv。

2、Latent contextual regressor（隱式上下文回歸器）則是在此基礎(chǔ)上預(yù)測出掩蔽patch的表征Zm。

3、解碼器以Zm和對應(yīng)的位置編碼作為輸入，通過Zm預(yù)測掩蔽patch的某些性質(zhì)，比如RGB值、Token ID。這過程中Zv不會更新，表征學(xué)習(xí)任務(wù)全交給編碼器。

4、Latent representation alignment對 Zm添加約束，希望 latent contextual regressor 的輸出和編碼器的輸出在同一空間。具體來說，圖像的掩蔽patch也將輸入到編碼器中（此過程不參與梯度反傳），獲得的這部分表征，作為Zm的學(xué)習(xí)目標(biāo)。

Alignment很重要，如果想對掩碼部分做好預(yù)測，要求latent contextual regressor 的輸出（也是解碼器的輸入）含有比較好的語義信息。通過對齊操作，可以鼓勵編碼器的操作也含有好的語義信息，提高編碼器的表征質(zhì)量。

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

如果訓(xùn)練時不做alignment約束，那么輸出的結(jié)果將是這樣…嗯，都是亂碼。

這種設(shè)計的編碼器學(xué)到的表征也相對更差，下游任務(wù)結(jié)果也會變差。

損失函數(shù)由兩部分組成，一個是對解碼器預(yù)測的監(jiān)督，使用的是cross-entropy loss；一個是對alignment的監(jiān)督，使用MSE損失函數(shù)。

除此之外，也進一步驗證了以CAE為代表的MIM方法，要比Moco v3、DINO為代表的對比學(xué)習(xí)方法更適合下游任務(wù)。

該論文從隨機裁剪操作的性質(zhì)分析，認為隨機裁剪有很大概率包含圖像的中心區(qū)域。

而ImageNet-1K這種數(shù)據(jù)集中，中心區(qū)域通常是1000類標(biāo)簽集中的物體（如下圖）。因此，對比學(xué)習(xí)方法主要提取圖像中主體物體的特征。

而MIM方法能學(xué)到每個patch的特征，包括圖像的背景區(qū)域，而不僅僅是圖像主體物體，這讓MIM學(xué)到的表征更適合下游檢測分割任務(wù)。

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

實驗結(jié)果

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

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

本文參照BEiT，使用DALL-E tokenizer對輸入圖像token化，得到預(yù)測目標(biāo)。

最終結(jié)果顯示，在語義分割任務(wù)中，跟其他MIM方法，比如MAE、BEiT，以及對比學(xué)習(xí)、有監(jiān)督預(yù)訓(xùn)練方法的表征結(jié)果更好。

在物體檢測、實例分割的結(jié)果也是如此。

百度CV大牛領(lǐng)銜

本次研究由北京大學(xué)、香港大學(xué)、百度共同完成。

第一作者是在讀博士生陳小康，來自北京大學(xué)機器感知與智能（教育部）重點實驗室。

通訊作者是百度計算機視覺首席架構(gòu)師王井東，同時也是IEEE Fellow。

在加盟百度之前，曾在微軟亞研院視覺計算組擔(dān)任首席研究員。

感興趣的旁友，可戳下方鏈接進一步查看論文~

論文鏈接：
https://arxiv.org/abs/2202.03026