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

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

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

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

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

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

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

• 編碼器（Encoder）是一個(gè) ViT 模型，負(fù)責(zé)學(xué)習(xí)可見 patch 的表征。

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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