本文目錄

1 MAE
1.1 Self-supervised Learning
1.2 Masked AutoEncoder (MAE) 方法概述
1.3 MAE Encoder
1.4 MAE Decoder
1.5 自監(jiān)督學習目標函數(shù) Reconstruction Target
1.6 具體實現(xiàn)方法
1.7 ImageNet 實驗結(jié)果
1.8 masking ratio 對性能的影響
1.9 觀察到的一些實驗現(xiàn)象
1.10 訓練策略
1.11 結(jié)果對比
1.12 Partial Fine-tuning

Self-Supervised Learning超詳細解讀 (目錄)：

https://zhuanlan.zhihu.com/p/381354026

Self-Supervised Learning，又稱為自監(jiān)督學習，我們知道一般機器學習分為有監(jiān)督學習，無監(jiān)督學習和強化學習。而 Self-Supervised Learning 是無監(jiān)督學習里面的一種，主要是希望能夠?qū)W習到一種通用的特征表達用于下游任務 (Downstream Tasks)。其主要的方式就是通過自己監(jiān)督自己。作為代表作的 kaiming 的 MoCo 引發(fā)一波熱議， Yann Lecun也在 AAAI 上講 Self-Supervised Learning 是未來的大勢所趨。所以在這個系列中，我會系統(tǒng)地解讀 Self-Supervised Learning 的經(jīng)典工作。

1 MAE

論文名稱：Masked Autoencoders Are Scalable Vision Learners

論文地址：

https://arxiv.org/abs/2111.06377

計算機視覺領(lǐng)域的 GPT-3 要有了？Kaiming 在這篇論文中告訴我們：強大的 AI 技術(shù) idea 應該非常簡單，實現(xiàn)起來也很快捷。我們首先對 Self-supervised Learning 做個簡單的回顧：

1.1 Self-supervised Learning

在預訓練階段我們使用無標簽的數(shù)據(jù)集 (unlabeled data)，因為有標簽的數(shù)據(jù)集很貴，打標簽得要多少人工勞力去標注，那成本是相當高的，太貴。相反，無標簽的數(shù)據(jù)集網(wǎng)上隨便到處爬，它便宜。在訓練模型參數(shù)的時候，我們不追求把這個參數(shù)用帶標簽數(shù)據(jù)從初始化的一張白紙給一步訓練到位，原因就是數(shù)據(jù)集太貴。于是 Self-Supervised Learning 就想先把參數(shù)從一張白紙訓練到初步成型，再從初步成型訓練到完全成型。注意這是2個階段。這個訓練到初步成型的東西，我們把它叫做 Visual Representation。預訓練模型的時候，就是模型參數(shù)從一張白紙到初步成型的這個過程，還是用無標簽數(shù)據(jù)集。等我把模型參數(shù)訓練個八九不離十，這時候再根據(jù)你下游任務 (Downstream Tasks) 的不同去用帶標簽的數(shù)據(jù)集把參數(shù)訓練到完全成型，那這時用的數(shù)據(jù)集量就不用太多了，因為參數(shù)經(jīng)過了第1階段就已經(jīng)訓練得差不多了。

第一個階段不涉及任何下游任務，就是拿著一堆無標簽的數(shù)據(jù)去預訓練，沒有特定的任務，這個話用官方語言表達叫做：in a task-agnostic way。第二個階段涉及下游任務，就是拿著一堆帶標簽的數(shù)據(jù)去在下游任務上 Fine-tune，這個話用官方語言表達叫做：in a task-specific way。

以上這些話就是 Self-Supervised Learning 的核心思想，如下圖1所示，后面還會再次提到它。

Self-Supervised Learning 不僅是在NLP領(lǐng)域，在CV, 語音領(lǐng)域也有很多經(jīng)典的工作，如下圖2所示。它可以分成3類：Data Centric, Prediction (也叫 Generative) 和 Contrastive。

其中的主流就是基于 Generative 的方法和基于 Contrative 的方法。如下圖 3 所示這里簡單介紹下。基于 Generative 的方法主要關(guān)注的重建誤差，比如對于 NLP 任務而言，一個句子中間蓋住一個 token，讓模型去預測，令得到的預測結(jié)果與真實的 token 之間的誤差作為損失。基于 Contrastive 的方法不要求模型能夠重建原始輸入，而是希望模型能夠在特征空間上對不同的輸入進行分辨。

1.2 Masked AutoEncoders (MAE) 方法概述

這篇論文使用掩碼自編碼器 (masked autoencoders (MAE)) 進行自監(jiān)督學習。根據(jù)1.1節(jié)的介紹，它屬于 Generative (Predictive) pre-training 的類型。這種類型自監(jiān)督學習的另一個著名的例子就是 BERT。

對于 BERT 模型而言，一個 sentence 中間蓋住一些 tokens，讓模型去預測，令得到的預測結(jié)果與真實的 tokens 之間的誤差作為損失。它告訴了我們直接 reconstruct sentence 也可以做到很 work。

對于 MAE 模型而言，一個 image 中間蓋住一些 patches，讓模型去預測，令得到的預測結(jié)果與真實的 image patches 之間的誤差作為損失。它告訴了我們直接 reconstruct image 原圖也可以做到很 work。

上面的對比看上去很容易，但是問題在于：

為什么 BERT (2018) 提出這么久以后，直到 BEIT (2021.6) 和 MAE (2021.11) 之前，一直在 CV 領(lǐng)域都沒有一個很類似的 CV BERT 出現(xiàn)？

這里 Kaiming 提出了3條看法：

1. CV 和 NLP 主流架構(gòu)不同： 直到 ViT (2020.12) 出現(xiàn)之前，CV 的主流架構(gòu)一直是以卷積網(wǎng)絡為主，NLP 的主流架構(gòu)一直是以 Transformer 為主。卷積核作用在一個個的 grid 上面，直觀來講沒法產(chǎn)生像 Transformer 一樣的 token 的概念，也就是說如果我們只使用卷積網(wǎng)絡，那么 image token 概念的建立就不那么直觀。所以，像 Transformer 那樣在 token 的基礎(chǔ)上進行自監(jiān)督學習就不太適用，這是第一個難點。
2. 語言和圖片 (視頻) 的信息密度不同： 語言是人類造就的信號，它 highly semantic ， information-dense。而圖片 (視頻) 是自然產(chǎn)生的信號，它 heavy spatial redundancy。即擋住圖片的一部分 patches，可以很容易地通過看它周圍的 patches 而想象出它的樣子來。所以，語言和圖像，一個信息密度高，一個信息密度低，這是第二個難點。解決的辦法是什么呢？作者提出了一個簡單的策略：即擋住圖片的 patches 的比例高一些。比如之前你擋住一張圖片的 30% 的 patches，能夠輕松通過周圍的 patches 預測出來；那現(xiàn)在如果擋住圖片的 90% 的 patches，還能夠輕松通過周圍的 patches 預測出來嗎？
3. AutoEncoder 里面的 Decoder 部分 (就是將映射得到的中間特征重建為 input 的模塊) 在 CV 和 NLP 中充當?shù)慕巧煌?/strong> 在 CV 領(lǐng)域，Decoder 的作用是重建 image pixels，所以 Decoder 的輸出語義級別很低。在 NLP 領(lǐng)域，Decoder 的作用是重建 sentence words ，所以 Decoder 的輸出語義級別很豐富。

這3條看法非常準確地展示了：為什么在 CV 領(lǐng)域搞出一個像 BERT 一樣類似的 CV BERT 這件事比較難。

基于以上分析，作者提出了 MAE 方法，如下圖4所示就是 MAE 的架構(gòu)。MAE 的方法很簡單：Mask 掉輸入圖像的隨機的 patches 并重建它們。它基于兩個核心理念：研究人員開發(fā)了一個非對稱編碼器 - 解碼器架構(gòu)，其中一個編碼器只對可見的 patch 子集進行操作 (即沒有被 mask 掉的 token)，另一個簡單解碼器可以從潛在表征和被 masked 掉的 token 重建原始圖像。Decoder 的架構(gòu)可以是十分輕量化的模型，且具體的架構(gòu)對模型性能影響很大。研究人員進一步發(fā)現(xiàn)，Mask 掉大部分輸入圖像 (例如 75%) 會產(chǎn)生重要且有意義的自監(jiān)督任務。結(jié)合這兩種設計，我們就能高效地訓練大型模型：提升訓練速度至 3 倍或更多，并提高準確性。

MAE 方法嚴格來講屬于一種去噪自編碼器 (Denoising Auto-Encoders (DAE))，去噪自動編碼器是一類自動編碼器，它破壞輸入信號，并學會重構(gòu)原始的、未被破壞的信號。MAE 的 Encoder 和 Decoder 結(jié)構(gòu)不同，是非對稱式的。Encoder 將輸入編碼為 latent representation，而 Decoder 將從 latent representation 重建原始信號。

MAE 和 ViT 的做法一致，將圖像劃分成規(guī)則的，不重疊的 patches。然后按照均勻分布不重復地選擇一些 patches 并且 mask 掉剩余的 patches。作者采用的 mask ratio 足夠高，因此大大減小了 patches 的冗余信息，使得在這種情況下重建 images 不那么容易。

1.3 MAE Encoder：

MAE Encoder 采用 ViT 架構(gòu)，但只會作用于 unmasked images。和 ViT 思路一樣，MAE Encoder 會先通過 Linear Projection 編碼圖片，再加上位置編碼，隨后送入一堆連續(xù)的 Transformer Block 里面。但是編碼器只對整個圖片 patches 集合的一個小子集 (例如25%)進行操作，而刪除 masked patches。這里和 BERT 做法不一樣，BERT 使用對于 mask 掉的部分使用特殊字符，而 MAE 不使用掩碼標記。

1.4 MAE Decoder：

MAE Decoder 采用 Transformer 架構(gòu)，輸入整個圖片 patches 集合，不光是 unmasked tokens (圖4中藍色色塊)，還有被 mask 掉的部分 (圖4中灰色色塊)。每個 mask tokens 都是一個共享的、學習的向量，它指示了這里有一個待預測的 tokens。作者還將位置嵌入添加到這個完整 image patch 集合中的所有 tokens 中，位置編碼表示每個 patches 在圖像中的位置的信息。

MAE Decoder 僅用于預訓練期間執(zhí)行圖像重建任務。因為自監(jiān)督學習的特點就是只用最后預訓練好的 Encoder 完成分類任務。因此，可以靈活設計與編碼器設計無關(guān)的解碼器結(jié)構(gòu)。作者用比編碼器更窄更淺的很小的解碼器做實驗。在這種非對稱的設計下，tokens 就可以由輕量級解碼器處理，這大大縮短了預訓練的時間。

1.5 自監(jiān)督學習目標函數(shù) Reconstruction Target：

Decoder 的最后一層是一個 Linear Projection 層，其輸出的 channel 數(shù)等于圖像的像素 (pixel) 數(shù)。所以 Decoder 的輸出會進一步 reshape 成圖像的形狀。損失函數(shù)就是 MSE Loss，即直接讓 reconstructed image 和 input image 的距離越接近越好。

作者還嘗試了另外一種損失函數(shù)，就是先計算出每個 patch 的像素值的 mean 和 deviation，并使用它們?nèi)w一化這個 patch 的每個像素值。最后再使用歸一化的像素值進行 MSE Loss 計算。但是發(fā)現(xiàn)這樣做的效果比直接 MSE Loss 好。

1.6 具體實現(xiàn)方法：

MAE 的具體實現(xiàn)方法是：

1. 首先通過 Linear Projection 和位置編碼得到 image tokens。
2. 隨機 shuffle 這些 tokens，按照 masking ratio 扔掉最后的一部分。
3. 把 unmasked patches 輸出到 Encoder 中，得到這些 tokens 的表征。
4. 把 Encoder 的輸出，結(jié)合 masked tokens (可學習的向量)，執(zhí)行 unshuffle操作恢復順序，再一起輸入到 Decoder 中。
5. shuffle 和 unshuffle 操作的時間開銷可忽略不計。

1.7 ImageNet 實驗結(jié)果：

第1個實驗使用 ViT-Large (ViT-L/16) 作為 Encoder 的 Backbone，作者首先對比了3個結(jié)果，如下圖5所示：

• 76.5 是 ViT 論文 report 的 ViT-Large 使用 ImageNet 訓練的結(jié)果。
• 82.5 是作者自己實現(xiàn)的，ViT-Large 使用 ImageNet 訓練的結(jié)果，這里不得不佩服 Kaiming 強大的實驗能力啊，單訓 ViT 就能比原作者高出6個點來，作者認為要使用較大的 weight decay = 0.3。82.5的實驗設置如下圖6所示。
• 84.9 是作者自己實現(xiàn)的，ViT-Large 使用 MAE pretrain 之后，再在 ImageNet 上做完整 Fine-tuning 的結(jié)果，可以看到優(yōu)于有監(jiān)督的 82.5。

1.8 masking ratio 對性能的影響：

如下圖7所示代表了 masking ratio 對結(jié)果的影響。最優(yōu)的 masking ratio 竟然出乎意料地高，且當 masking ratio=75% 時，linear probing 和 fine-tuning 的性能依舊很高。這種行為與 BERT 相反，BERT 的典型mask ratio 只有15%。從這個結(jié)果可以得出結(jié)論：MAE 自監(jiān)督訓練好的模型可以推斷出缺失的補丁。作者假設這種類似推理的行為，與學習有用的表征有關(guān)。

圖7還顯示了linear probing 和 fine-tuning 的結(jié)果隨著的不同的變化趨勢。對于 fine-tuning，結(jié)果是40-80%的 masking ratio 都能 work well。

1.9 觀察到的一些實驗現(xiàn)象：

如下圖8所示，作者做了一些不同實驗設置下的對比實驗。

圖8 (a)： 足夠深的 Decoder 對于 linear probing 很重要。原因：pixel reconstruction task 和 recognition task 本質(zhì)不同，AutoEncoder 的末尾幾層對于重建 pixel 更重要。而且有趣的是，只有1層 Block 的 Decoder 在 fine-tuning 之后就能達到84.8%的 accuracy。默認配置是8層。

圖8 (b)： 更窄的 Decoder 也能夠在 fine-tuning 時獲得較好的性能。默認配置是 Embedding dimension=512。

計算量只有 ViT-L 的9%。

圖8 (c)： 研究有沒有必要在輸入 Encoder 的時候一并輸入 masked patches。結(jié)果發(fā)現(xiàn)，如果也一起給 Encoder 輸入了 mask tokens，效果其實是變差的：Linear Probing 性能降低了14%。作者認為，這可能是因為：在這種情況下，在預訓練和部署之間存在差距。 即：編碼器在預訓練的輸入中有很大一部分是 mask tokens，這在未損壞的圖像中是不存在的。這個 gap 可能會降低部署的準確性。

而且，通過跳過編碼器中的 mask token，大大減少了訓練計算量 (約3.3倍)。此外，GPU memory 大大減少，這可以使訓練更大的模型或加快 large-batch training。時間和 GPU memory 效率使我們更適合訓練大型模型。

圖8 (d)： 不同損失函數(shù)的影響。除了 MSE Loss 以外，作者還嘗試了另外一種損失函數(shù)，就是先計算出每個 patch 的像素值的 mean 和 deviation，并使用它們?nèi)w一化這個 patch 的每個像素值。最后再使用歸一化的像素值進行 MSE Loss 計算。但是發(fā)現(xiàn)這樣做的效果比直接 MSE Loss 好。

圖8 (e)： 不同數(shù)據(jù)增強策略的影響。MAE 只使用 cropping-only 就比較好，而進一步添加 color jittering 會影響性能。即使沒有使用數(shù)據(jù)增強 (only center-crop, no flipping) ，MAE 模型仍然能夠表現(xiàn)得體。作者也給出了原因，即：在 MAE 中，數(shù)據(jù)增強的角色其實是由 random masking 來扮演的。每個 iteration 的 mask 都不同，所以就相當于是產(chǎn)生了新的訓練樣本。

圖8 (9)，圖9： 不同 mask 采樣策略的影響。block-wise masking 在 mask 50%時的效果還行，但是當 mask ratio 達到75%時性能就大幅下降了。grid-wise sampling 使得任務變得更簡單，但相應的，圖像重建質(zhì)量也下降了。簡單的隨機抽樣最適合 MAE 模型。

1.10 訓練策略：

下圖10展示了不同 training epochs 的影響。在 training epochs=1600 時 MAE 還沒有達到飽和。

1.11 結(jié)果對比：

如下圖11所示為不同自監(jiān)督方法性能對比，對于 ViT-B 模型，所有的方法性能相似，但是對于更大的 ViT-L 模型，性能差距就拉開了，證明了 MAE 對于大模型的泛化性能。

使用 ViT-H 結(jié)構(gòu)的 Encoder，經(jīng)過 448 fine-tuning 之后可以得到 87.8% 的 Accuracy，而整個過程只使用了 ImageNet 數(shù)據(jù)集。

與 BEiT 相比，MAE 更快，訓練方式更簡單，MAE 重建 pixel，而 BEiT 重建 tokens。圖11中的模型是為了更高的精度預先訓練的 1600 epochs 的模型。即便如此，如果他們在同樣的硬件環(huán)境下訓練，我們總的訓練前時間比其他任何方法都短。例如，對于 ViT-L，MAE 的訓練時間是1600個 epochs，31小時。moco v3的訓練時間是300個 epochs，36小時，使用相同的128個 TPU-v3 核心。

1.12 Partial Fine-tuning

最后，作者提出了一種 Partial Fine-tuning 的新套路，它不同于之前人們常用的 Linear Probing (只訓練最后一層線性分類器的參數(shù)) 和 Fine-tuning (訓練所有層的參數(shù))。

Partial Fine-tuning 是指 只訓練最后模型的若干層的參數(shù)。如下圖12所示，值得注意的是，F(xiàn)ine-tuning 1個 Transformer Block 可以把 Accuracy 從73.0%提升至81.0%。此外，如果我們只Fine-tuning 1個 Transformer Block 的 MLP block，我們可以得到79.1% ，比 Linear Proing 要好得多。

總結(jié)

用 MAE 做 pre-training 只需 ImageNet-1k 就能達到 87.8% 的 Top-1 準確度，超過了所有在 ImageNet-21k pre-training 的 ViT 變體模型。而從方法上看，MAE 選擇直接重建原圖的元素，而且證明了其可行性，改變了人們的認知，又幾乎可以覆蓋 CV 里所有的識別類任務，看起來像是開啟了一個新的方向。直接重建原圖的元素是非常重要的，因為通過這個形式，作者就用最最直觀的方式完成了 MIM 任務，使得 MIM的潛力逐步被證實。從 MLM 到 MIM 的過渡已被證明，由此觀之比肩 GPT3 的 CV 預訓練大模型已不遠矣。

如果覺得有用，就請分享到朋友圈吧！