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

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

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

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

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

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

輸入圖像通過隨機掩碼被劃分成可見 patch 和掩碼 patch 兩個部分。具體來說：

• 編碼器（Encoder）是一個 ViT 模型，負責學習可見 patch 的表征。

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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