機器之心專欄
作者：陳小康
來自北京大學、香港大學和百度的研究者近日提出了一種名為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 的結果也更好，說明預訓練階段編碼器被充分利用，學到的表征更好。
4.4 目标檢測、執行個體分割
該研究使用 Mask-RCNN 和 Cascade-RCNN 兩種網絡結構進行目标檢測和執行個體分割的實驗。其中，使用多尺度訓練 12 epoch，測試階段僅使用單尺度測試。實驗結果和語義分割類似：對比學習方法和監督預訓練方法的結果類似且較差，CAE 的結果更好。
5 總結
該研究提出了 CAE，設計的核心有兩點：(1) 對 “表征學習” 和 “解決前置任務” 這兩個功能做完全分離; (2) 在可見 patch 學習到的表征空間中對掩碼 patch 做預測。以上兩點都是為了驅使編碼器學習更好的表征，進而在下遊任務取得良好的泛化能力。
此外，該研究對監督預訓練方法、對比學習和 MIM 方法進行了分析，認為對比學習和監督預訓練主要關注圖像的主體區域 (例如 ImageNet-1K 标簽集中的對象)，而 MIM 會關注圖像的全部 patch，更有利于下遊任務。