機器之心 - 自監督學習新範式CAE：為什麼 MIM 比對比學習更適合下游任務？－鑽石舞台

機器之心專欄

作者：陳小康

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

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

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

論文地址：https://arxiv.org/abs/2202.03026

該研究回答了如下幾個問題：

1.MIM 方法中，網絡結構的哪個部分是學習表徵的，哪個部分是解決前置任務的？

2. 為什麼之前典型的對比學習方法，在下游任務 (例如檢測、分割) 上只能取得跟監督預訓練方法類似的性能？

3.MIM 方法為什麼優於目前的對比學習方法？

1. 背景

MIM 是一種自監督表徵學習算法。它的主要思路是，對輸入圖像進行分塊和隨機掩碼操作，然後對掩碼區域做一些預測。預測的目標可以是 Token ID (BEiT)，也可以是 RGB 的值 (MAE)。編碼器能夠通過 MIM 學得一個好的表徵，從而在下游任務上取得良好的泛化性能。

近期 MIM 有兩個代表性工作：BEiT 和 MAE。

BEiT 使用一個編碼器做兩件事：(1) 學習一個好的圖像表徵；(2) 解決前置任務：預測掩碼 patch 的 Token ID。編碼器的潛力並沒有完全被挖掘，只有部分被用來學習表徵。

MAE 使用了編碼器-解碼器架構，編碼器負責對可見 patch 進行表徵學習，解碼器將可見和掩碼patch的表徵（使用一個可學習的向量）作為輸入，預測掩碼 patch 的 RGB 值。但是，MAE 在解碼器中也會對可見 patch 的表徵進行更新，實際上解碼器也負責了一部分學習表徵的功能。

以上兩種方法，都沒有充分挖掘編碼器的潛力，限制了預訓練學習到的表徵質量。

2. Context Autoencoder (CAE)

CAE 設計的核心思想是對「表徵學習」和「解決前置任務」這兩個功能做分離。研究者希望在預訓練時，編碼器只負責表徵學習，解碼器只負責解決前置任務，這樣可以儘可能大地挖掘編碼器的潛力。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 對應的位置編碼。在這個階段，不斷更新、變得更加準確，而不會更新，對圖像特徵的提取這個任務完全交給編碼器。

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

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

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

3. 分析

3.1 CAE 關注每個 patch 的表徵

CAE 基於可見 patch 的表徵，從隨機採樣的掩碼 patch 中做一些預測，這要求 CAE 關注每個 patch 的語義。這不同於典型的對比學習方法 (例如 MoCo v3, SimCLR)，不是只關注圖像的全局語義而忽略圖像的細節和非主體區域 (比如背景)。

3.2 Latent contextual regressor 的輸出和編碼器的輸出在同一編碼空間中

該研究對 Latent contextual regressor 的輸出做了約束，希望它能和編碼器的輸出儘可能接近。這樣，解碼器會基於編碼器學到的編碼空間做預測，將對圖像的特徵提取的重任完全交到了編碼器手上，驅使編碼器學習到好的表徵。

為了驗證這一點，該研究用 RGB 值作為解碼器目標 (考慮到 Token ID 難以可視化，這裡使用 RGB)，訓練 CAE。在測試的時候，該研究將全部 patch 輸入到編碼器，然後跳過Latent contextual regressor，直接將編碼器的輸出送進解碼器，預測全部 patch 的 RGB 的值。下圖展示了預測結果，第一行是原圖，第二行是預測，研究者發現僅使用編碼器和解碼器就可以將圖片重建出來，說明編碼器的輸出和Latent contextual regressor的輸出屬於同一編碼空間。

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

3.3 CAE 學到的表徵可以區分不同類別的對象/stuff

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

3.4 典型的對比學習為什麼在下游任務只能取得跟監督預訓練差不多的結果？

在對比學習中，隨機剪裁（random crop）是一個非常重要的數據增強策略。典型的對比學習（比如 MoCo v3）希望最大化來自同一圖像的 2 個不同剪裁之間的全局語義相似度，而最小化來自不同圖像的剪裁之間的相似度。

這樣為什麼能奏效呢？研究者首先分析了隨機剪裁的性質。在 SimCLR 論文中提到，隨機剪裁是對比學習方法中非常重要的數據增強策略。在 ImageNet-1K 數據集中，圖像的主體對象大多處於圖像的中心區域，而對圖像進行隨機剪裁，中心區域有很大的概率被囊括進去，例如下圖展示的幾個例子，幾次剪裁基本都包括了圖像的主體對象。

對同一圖像的不同剪裁提取全局語義，實際上學到的是原始圖像中主體對象的特徵，正因如此，同一圖像的不同剪裁之間才可能相似。在監督預訓練中，受到圖像分類標籤的約束，網絡學習到的也是圖像主體區域的特徵，這和對比學習學到的知識有很大的相似之處，因此在下游任務表現類似。

3.5 MIM 和對比學習的區別

MIM 方法 (例如 CAE) 基於可見 patch 的表徵，對掩碼 patch 區域做預測。在做隨機掩碼時，圖像的每個 patch (例如背景區域的對象/stuff) 都有可能被考慮到，而不僅僅是圖像的主體區域。為了做好掩碼 patch 的預測，CAE 會學好每個 patch 的表徵。

該研究對 CAE 以及 MoCo v3 的注意力圖做了可視化。如下圖所示，第一行是原圖，第二行是 MoCo v3，第三行是 CAE。紅色表示注意力值更高，藍色表示注意力值低。處於藍色邊界內部的區域，通過這樣的原則篩選：將注意力值從大到小排序後，保留累計和達到所有位置注意力值總和的 50% 的部分。可以看到，MoCo v3 的注意力圖主要在圖像的主體區域有高響應，而 CAE 能考慮到幾乎所有 patch。

4. 實驗

該研究使用 ViT-small 和 ViT-base 在 ImageNet-1K 上進行實驗。輸入圖像的分辨率是 224 X 224，patch 大小是 16 X 16，一張圖會被劃分成 14 X 14 個 patch。每次有 75 個 patch 被隨機掩碼。

4.1 預訓練評估

自監督學習廣泛使用線性探測（linear probing）去評測預訓練表徵的好壞：將編碼器的參數固定住，在之後加一個線性分類器進行圖像分類。研究者認為線性探測不適合 MIM 方法，因為 MIM 方法通常會學到每個 patch 的表徵，不僅包含主體對象的信息，還學到了背景等若干知識，這是多而雜的，不適合直接進行線性分類。因此，研究者提出了一種新的測試指標：注意力探測（attentive probing）。該研究在固定參數的編碼器後加上一個簡單的交叉注意力模塊（沒有 FFN）和一個線性分類器，通過注意力機制動態地選擇適合做圖像分類的信息。

該研究對注意力探測階段使用的交叉注意力模塊做注意力圖可視化，發現可以關注到主體對象。

微調、線性探測、注意力探測的結果見下表。

研究者發現一些有趣的現象。(1) 對比學習方法 (MoCo v3, DINO) 的線性探測和注意力探測結果類似。這說明這類方法在預訓練時已經將注意力放到了圖像的主體對象上面，無需進一步動態篩選即可做好圖像分類，這也與之前研究者對對比學習的分析一致。(2) MIM 方法 (例如 CAE) 的注意力探測相比線性探測有很大的提升。這說明 MIM 方法學到了每個 patch 的特徵，而不僅僅是圖像主體對象的，因此需要做一些篩選才利於圖像分類。

4.2 消融實驗

該研究對解碼器和對齊模塊進行消融實驗，見下表。單加一個解碼器能改進注意力探測的結果，但在下游任務 (分割、檢測)上的提升不明顯。使用對齊模塊之後能顯著提升下游任務的性能，說明約束編碼器的輸出和 Latent contextual regressor 的輸出在同一編碼空間非常重要，能提升編碼器學到的表徵質量。

4.3 語義分割

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