北大博士生提出CAE，下游任務泛化能力優於何愷明MAE
楊淨 發自 凹非寺
量子位 | 公眾號 QbitAI
上回說道，何愷明時隔兩年發一作論文，提出了一種視覺自監督學習新正規化——
用掩蔽自編碼器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，以及對比學習、有監督預訓練方法的表徵結果更好。
在物體檢測、例項分割的結果也是如此。
本次研究由北京大學、香港大學、百度共同完成。
第一作者是在讀博士生陳小康，來自北京大學機器感知與智慧（教育部）重點實驗室。
通訊作者是百度計算機視覺首席架構師王井東，同時也是IEEE Fellow。
在加盟百度之前，曾在微軟亞研院視覺計算組擔任首席研究員。
感興趣的旁友，可戳下方連結進一步檢視論文~
論文連結：https://arxiv.org/abs/2202.03026
— 完 —
量子位 QbitAI · 頭條號簽約