自監督學(xué)習新范式CAE：為什么 MIM 比對比學(xué)習更適合下游任務(wù)？

來(lái)自北京大學(xué)、香港大學(xué)和百度的研究者近日提出了一種名為CAE的新型 MIM 方法。

掩碼建模方法，在 NLP 領(lǐng)域 (例如 BERT) 得到了廣泛的應用。隨著(zhù) ViT 的提出和發(fā)展，人們也嘗試將掩碼圖像建模（MIM）應用到視覺(jué)領(lǐng)域并取得了一定進(jìn)展。在此之前，視覺(jué)自監督算法主要沿著(zhù)對比學(xué)習（contrastive learning）的思路去設計，而 MIM 無(wú)疑打開(kāi)了新的大門(mén)。

來(lái)自北京大學(xué)、香港大學(xué)和百度的研究者近日提出了一種名為CAE的新型 MIM 方法。該方法通過(guò)對 “表征學(xué)習” 和 “解決前置任務(wù)（pretext task）” 這兩個(gè)功能做完全分離，使得編碼器學(xué)習到更好的表征，從而在下游任務(wù)上實(shí)現了更好的泛化性能。

論文地址：https://arxiv.org/abs/2202.03026
該研究回答了如下幾個(gè)問(wèn)題：
1.MIM 方法中，網(wǎng)絡(luò )結構的哪個(gè)部分是學(xué)習表征的，哪個(gè)部分是解決前置任務(wù)的？2. 為什么之前典型的對比學(xué)習方法，在下游任務(wù) (例如檢測、分割) 上只能取得跟監督預訓練方法類(lèi)似的性能？3.MIM 方法為什么優(yōu)于目前的對比學(xué)習方法？
1. 背景
MIM 是一種自監督表征學(xué)習算法。它的主要思路是，對輸入圖像進(jìn)行分塊和隨機掩碼操作，然后對掩碼區域做一些預測。預測的目標可以是 Token ID (BEiT)，也可以是 RGB 的值 (MAE)。編碼器能夠通過(guò) MIM 學(xué)得一個(gè)好的表征，從而在下游任務(wù)上取得良好的泛化性能。
近期 MIM 有兩個(gè)代表性工作：BEiT 和 MAE。

• BEiT 使用一個(gè)編碼器做兩件事：(1) 學(xué)習一個(gè)好的圖像表征；(2) 解決前置任務(wù)：預測掩碼 patch 的 Token ID。編碼器的潛力并沒(méi)有完全被挖掘，只有部分被用來(lái)學(xué)習表征。
• MAE 使用了編碼器-****架構，編碼器負責對可見(jiàn) patch 進(jìn)行表征學(xué)習，****將可見(jiàn)和掩碼patch的表征（使用一個(gè)可學(xué)習的向量）作為輸入，預測掩碼 patch 的 RGB 值。但是，MAE 在****中也會(huì )對可見(jiàn) patch 的表征進(jìn)行更新，實(shí)際上****也負責了一部分學(xué)習表征的功能。

以上兩種方法，都沒(méi)有充分挖掘編碼器的潛力，限制了預訓練學(xué)習到的表征質(zhì)量。
2. Context Autoencoder (CAE)
CAE 設計的核心思想是對 “表征學(xué)習” 和 “解決前置任務(wù)” 這兩個(gè)功能做分離。研究者希望在預訓練時(shí)，編碼器只負責表征學(xué)習，****只負責解決前置任務(wù)，這樣可以盡可能大地挖掘編碼器的潛力。CAE 包括 4 個(gè)部分：(1) Encoder; (2) Latent contextual regressor; (3) Decoder; (4) Alignment模塊。

輸入圖像通過(guò)隨機掩碼被劃分成可見(jiàn) patch 和掩碼 patch 兩個(gè)部分。具體來(lái)說(shuō)：

• 編碼器（Encoder）是一個(gè) ViT 模型，負責學(xué)習可見(jiàn) patch 的表征。

• Latent contextual regressor 通過(guò)預測掩碼 patch 的表征。Latent contextual regressor 由一系列交叉注意力（cross-attention）模塊組成，query 是掩碼 patch 的表征，key 和 value 是全部 patch 的表征。在計算 query-key 相似度時(shí)，該方法會(huì )引入每個(gè) patch 對應的位置編碼。在這個(gè)階段，不斷更新、變得更加準確，而不會(huì )更新，對圖像特征的提取這個(gè)任務(wù)完全交給編碼器。

• ****（Decoder）只拿和對應的位置編碼作為輸入，通過(guò)預測掩碼 patch 的某些性質(zhì)，比如 Token ID，或者 RGB 的值。該研究的實(shí)驗與 BEiT 類(lèi)似，使用 DALL-E tokenizer 對輸入圖像 token 化，得到****的目標。

• 潛在表征對齊（Latent representation alignment）通過(guò)對添加約束，希望Latent contextual regressor 的輸出和編碼器的輸出在同一編碼空間中。該方法將圖像的掩碼 patch 也輸入到編碼器，獲得這部分的表征。將作為學(xué)習的目標。計算的過(guò)程不會(huì )計算梯度。

• 損失函數。損失函數由兩部分組成：(1) 對****預測的監督，使用交叉熵損失; (2) 對和的對齊的監督，使用 MSE損失。

3. 分析
3.1 CAE 關(guān)注每個(gè) patch 的表征
CAE 基于可見(jiàn) patch 的表征，從隨機采樣的掩碼 patch 中做一些預測，這要求 CAE 關(guān)注每個(gè) patch 的語(yǔ)義。這不同于典型的對比學(xué)習方法 (例如 MoCo v3, SimCLR)，不是只關(guān)注圖像的全局語(yǔ)義而忽略圖像的細節和非主體區域 (比如背景)。
3.2 Latent contextual regressor 的輸出和編碼器的輸出在同一編碼空間中
該研究對 Latent contextual regressor 的輸出做了約束，希望它能和編碼器的輸出盡可能接近。這樣，****會(huì )基于編碼器學(xué)到的編碼空間做預測，將對圖像的特征提取的重任完全交到了編碼器手上，驅使編碼器學(xué)習到好的表征。
為了驗證這一點(diǎn)，該研究用 RGB 值作為****目標 (考慮到 Token ID 難以可視化，這里使用 RGB)，訓練 CAE。在測試的時(shí)候，該研究將全部 patch 輸入到編碼器，然后跳過(guò) Latent contextual regressor，直接將編碼器的輸出送進(jìn)****，預測全部 patch 的 RGB 的值。下圖展示了預測結果，第一行是原圖，第二行是預測，研究者發(fā)現僅使用編碼器和****就可以將圖片重建出來(lái)，說(shuō)明編碼器的輸出和 Latent contextual regressor 的輸出屬于同一編碼空間。

如果訓練時(shí)不做對齊約束，那么就無(wú)法重建，如下圖所示，輸出都是亂碼，說(shuō)明編碼器輸出和 Latent contextual regressor 的輸出不在一個(gè)編碼空間中。這使得編碼器學(xué)到的表征質(zhì)量有所欠缺，在消融實(shí)驗部分也有驗證。

3.3 CAE 學(xué)到的表征可以區分不同類(lèi)別的對象/stuff
CAE 基于可見(jiàn) patch 的表征，在掩碼 patch 區域做預測，這要求 CAE 對可見(jiàn) patch 的內容有比較好的理解。舉例來(lái)說(shuō)，人們看到一只狗的頭部，可以預測出它的身體部分；看到一小片天空，也能預測出它的周?chē)蟾怕室彩且黄炜?。因此，研究者認為 CAE 學(xué)到的表征可以區分不同類(lèi)別的對象/stuff。為了驗證這一點(diǎn)，研究者從 ADE20K 數據集隨機采樣一些圖片輸入到編碼器。因為 ADE20K 提供了每個(gè)像素的類(lèi)別標簽 (150 類(lèi))，因此該研究可以使用 t-SNE 對編碼器輸出的表征進(jìn)行可視化。如下圖所示，每個(gè)顏色代表一個(gè)類(lèi)別，左圖是 CAE，右圖是隨機初始化的編碼器。研究者發(fā)現 CAE 可以有效區分不同類(lèi)別的對象/stuff (因為是在 ImageNet-1K 進(jìn)行預訓練，所以區分得不夠完美)，而隨機初始化的編碼器無(wú)法做到這一點(diǎn)。

3.4 典型的對比學(xué)習為什么在下游任務(wù)只能取得跟監督預訓練差不多的結果？
在對比學(xué)習中，隨機剪裁（random crop）是一個(gè)非常重要的數據增強策略。典型的 對比學(xué)習（比如 MoCo v3）希望最大化來(lái)自同一圖像的 2 個(gè)不同剪裁之間的全局語(yǔ)義相似度，而最小化來(lái)自不同圖像的剪裁之間的相似度。
這樣為什么能奏效呢？研究者首先分析了隨機剪裁的性質(zhì)。在 SimCLR 論文中提到，隨機剪裁是對比學(xué)習方法中非常重要的數據增強策略。在 ImageNet-1K 數據集中，圖像的主體對象大多處于圖像的中心區域，而對圖像進(jìn)行隨機剪裁，中心區域有很大的概率被囊括進(jìn)去，例如下圖展示的幾個(gè)例子，幾次剪裁基本都包括了圖像的主體對象。

對同一圖像的不同剪裁提取全局語(yǔ)義，實(shí)際上學(xué)到的是原始圖像中主體對象的特征，正因如此，同一圖像的不同剪裁之間才可能相似。在監督預訓練中，受到圖像分類(lèi)標簽的約束，網(wǎng)絡(luò )學(xué)習到的也是圖像主體區域的特征，這和對比學(xué)習學(xué)到的知識有很大的相似之處，因此在下游任務(wù)表現類(lèi)似。
3.5 MIM 和對比學(xué)習的區別
MIM 方法 (例如 CAE) 基于可見(jiàn) patch 的表征，對掩碼 patch 區域做預測。在做隨機掩碼時(shí)，圖像的每個(gè) patch (例如背景區域的對象/stuff) 都有可能被考慮到，而不僅僅是圖像的主體區域。為了做好掩碼 patch 的預測，CAE 會(huì )學(xué)好每個(gè) patch 的表征。
該研究對 CAE 以及 MoCo v3 的注意力圖做了可視化。如下圖所示，第一行是原圖，第二行是 MoCo v3，第三行是 CAE。紅色表示注意力值更高，藍色表示注意力值低。處于藍色邊界內部的區域，通過(guò)這樣的原則篩選：將注意力值從大到小排序后，保留累計和達到所有位置注意力值總和的 50% 的部分?？梢钥吹?，MoCo v3 的注意力圖主要在圖像的主體區域有高響應，而 CAE 能考慮到幾乎所有 patch。

4. 實(shí)驗
該研究使用 ViT-small 和 ViT-base 在 ImageNet-1K 上進(jìn)行實(shí)驗。輸入圖像的分辨率是 224 X 224，patch 大小是 16 X 16，一張圖會(huì )被劃分成 14 X 14 個(gè) patch。每次有 75 個(gè) patch 被隨機掩碼。
4.1 預訓練評估
自監督學(xué)習廣泛使用線(xiàn)性探測（linear probing）去評測預訓練表征的好壞：將編碼器的參數固定住，在之后加一個(gè)線(xiàn)性分類(lèi)器進(jìn)行圖像分類(lèi)。研究者認為線(xiàn)性探測不適合 MIM 方法，因為 MIM 方法通常會(huì )學(xué)到每個(gè) patch 的表征，不僅包含主體對象的信息，還學(xué)到了背景等若干知識，這是多而雜的，不適合直接進(jìn)行線(xiàn)性分類(lèi)。因此，研究者提出了一種新的測試指標：注意力探測（attentive probing）。該研究在固定參數的編碼器后加上一個(gè)簡(jiǎn)單的交叉注意力模塊（沒(méi)有 FFN）和一個(gè)線(xiàn)性分類(lèi)器，通過(guò)注意力機制動(dòng)態(tài)地選擇適合做圖像分類(lèi)的信息。
該研究對注意力探測階段使用的交叉注意力模塊做注意力圖可視化，發(fā)現可以關(guān)注到主體對象。

微調、線(xiàn)性探測、注意力探測的結果見(jiàn)下表。

研究者發(fā)現一些有趣的現象。(1) 對比學(xué)習方法 (MoCo v3, DINO) 的線(xiàn)性探測和注意力探測結果類(lèi)似。這說(shuō)明這類(lèi)方法在預訓練時(shí)已經(jīng)將注意力放到了圖像的主體對象上面，無(wú)需進(jìn)一步動(dòng)態(tài)篩選即可做好圖像分類(lèi)，這也與之前研究者對對比學(xué)習的分析一致。(2) MIM 方法 (例如 CAE) 的注意力探測相比線(xiàn)性探測有很大的提升。這說(shuō)明 MIM 方法學(xué)到了每個(gè) patch 的特征，而不僅僅是圖像主體對象的，因此需要做一些篩選才利于圖像分類(lèi)。
4.2 消融實(shí)驗
該研究對****和對齊模塊進(jìn)行消融實(shí)驗，見(jiàn)下表。單加一個(gè)****能改進(jìn)注意力探測的結果，但在下游任務(wù) (分割、檢測) 上的提升不明顯。使用對齊模塊之后能顯著(zhù)提升下游任務(wù)的性能，說(shuō)明約束編碼器的輸出和 Latent contextual regressor 的輸出在同一編碼空間非常重要，能提升編碼器學(xué)到的表征質(zhì)量。

4.3 語(yǔ)義分割
該研究在 ADE20K 上進(jìn)行語(yǔ)義分割的實(shí)驗。網(wǎng)絡(luò )使用 UperNet，迭代次數為 160K，輸入圖像分辨率為 512 X 512，使用單尺度測試。對比學(xué)習方法和監督預訓練方法（DeiT）的結果類(lèi)似，而 CAE 能取得明顯更好的結果。跟其他 MIM 方法相比，CAE 的結果也更好，說(shuō)明預訓練階段編碼器被充分利用，學(xué)到的表征更好。

4.4 目標檢測、實(shí)例分割該研究使用 Mask-RCNN 和 Cascade-RCNN 兩種網(wǎng)絡(luò )結構進(jìn)行目標檢測和實(shí)例分割的實(shí)驗。其中，使用多尺度訓練 12 epoch，測試階段僅使用單尺度測試。實(shí)驗結果和語(yǔ)義分割類(lèi)似：對比學(xué)習方法和監督預訓練方法的結果類(lèi)似且較差，CAE 的結果更好。

5 總結
該研究提出了 CAE，設計的核心有兩點(diǎn)：(1) 對 “表征學(xué)習” 和 “解決前置任務(wù)” 這兩個(gè)功能做完全分離; (2) 在可見(jiàn) patch 學(xué)習到的表征空間中對掩碼 patch 做預測。以上兩點(diǎn)都是為了驅使編碼器學(xué)習更好的表征，從而在下游任務(wù)取得良好的泛化能力。
此外，該研究對監督預訓練方法、對比學(xué)習和 MIM 方法進(jìn)行了分析，認為對比學(xué)習和監督預訓練主要關(guān)注圖像的主體區域 (例如 ImageNet-1K 標簽集中的對象)，而 MIM 會(huì )關(guān)注圖像的全部 patch，更有利于下游任務(wù)。