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

白交 2022-02-22 12:52:52 來(lái)源：量子位

楊凈 發(fā)自 凹非寺

量子位 | 公眾號(hào) QbitAI

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

用掩蔽自編碼器MAE，為視覺(jué)大模型開(kāi)路。

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

來(lái)看看這是一項(xiàng)什么樣的研究？

這是一項(xiàng)什么研究？

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

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

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

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

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

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

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

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

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

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

CAE包含四個(gè)部分。

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

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

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

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

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

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

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

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

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

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

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

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

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

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

實(shí)驗(yàn)結(jié)果

研究團(tuán)隊(duì)使用ViT-small和ViT-base在 ImageNet-1K 上進(jìn)行實(shí)驗(yàn)，輸入圖像的分辨率224*224，每張圖被分成14*14的patch，每個(gè)patch的大小為16*16。

每次將有75個(gè)patch被隨機(jī)掩碼，其余patch則為可見(jiàn)的。

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

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

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

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

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

第一作者是在讀博士生陳小康，來(lái)自北京大學(xué)機(jī)器感知與智能（教育部）重點(diǎn)實(shí)驗(yàn)室。

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

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

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

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