本文目錄

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

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

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

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

1 MAE

論文名稱(chēng)：Masked Autoencoders Are Scalable Vision Learners

論文地址：

https://arxiv.org/abs/2111.06377

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

1.1 Self-supervised Learning

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

第一個(gè)階段不涉及任何下游任務(wù)，就是拿著一堆無(wú)標(biāo)簽的數(shù)據(jù)去預(yù)訓(xùn)練，沒(méi)有特定的任務(wù)，這個(gè)話(huà)用官方語(yǔ)言表達(dá)叫做：in a task-agnostic way。第二個(gè)階段涉及下游任務(wù)，就是拿著一堆帶標(biāo)簽的數(shù)據(jù)去在下游任務(wù)上 Fine-tune，這個(gè)話(huà)用官方語(yǔ)言表達(dá)叫做：in a task-specific way。

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

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

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

1.2 Masked AutoEncoders (MAE) 方法概述

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

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

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

上面的對(duì)比看上去很容易，但是問(wèn)題在于：

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

這里 Kaiming 提出了3條看法：

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

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

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

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

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

1.3 MAE Encoder：

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

1.4 MAE Decoder：

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

MAE Decoder 僅用于預(yù)訓(xùn)練期間執(zhí)行圖像重建任務(wù)。因?yàn)樽员O(jiān)督學(xué)習(xí)的特點(diǎn)就是只用最后預(yù)訓(xùn)練好的 Encoder 完成分類(lèi)任務(wù)。因此，可以靈活設(shè)計(jì)與編碼器設(shè)計(jì)無(wú)關(guān)的解碼器結(jié)構(gòu)。作者用比編碼器更窄更淺的很小的解碼器做實(shí)驗(yàn)。在這種非對(duì)稱(chēng)的設(shè)計(jì)下，tokens 就可以由輕量級(jí)解碼器處理，這大大縮短了預(yù)訓(xùn)練的時(shí)間。

1.5 自監(jiān)督學(xué)習(xí)目標(biāo)函數(shù) Reconstruction Target：

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

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

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

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

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

1.7 ImageNet 實(shí)驗(yàn)結(jié)果：

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

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

1.8 masking ratio 對(duì)性能的影響：

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

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

1.9 觀(guān)察到的一些實(shí)驗(yàn)現(xiàn)象：

如下圖8所示，作者做了一些不同實(shí)驗(yàn)設(shè)置下的對(duì)比實(shí)驗(yàn)。

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

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

計(jì)算量只有 ViT-L 的9%。

圖8 (c)： 研究有沒(méi)有必要在輸入 Encoder 的時(shí)候一并輸入 masked patches。結(jié)果發(fā)現(xiàn)，如果也一起給 Encoder 輸入了 mask tokens，效果其實(shí)是變差的：Linear Probing 性能降低了14%。作者認(rèn)為，這可能是因?yàn)椋?strong style="color: black;">在這種情況下，在預(yù)訓(xùn)練和部署之間存在差距。 即：編碼器在預(yù)訓(xùn)練的輸入中有很大一部分是 mask tokens，這在未損壞的圖像中是不存在的。這個(gè) gap 可能會(huì)降低部署的準(zhǔn)確性。

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

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

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

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

1.10 訓(xùn)練策略：

下圖10展示了不同 training epochs 的影響。在 training epochs=1600 時(shí) MAE 還沒(méi)有達(dá)到飽和。

1.11 結(jié)果對(duì)比：

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

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

與 BEiT 相比，MAE 更快，訓(xùn)練方式更簡(jiǎn)單，MAE 重建 pixel，而 BEiT 重建 tokens。圖11中的模型是為了更高的精度預(yù)先訓(xùn)練的 1600 epochs 的模型。即便如此，如果他們?cè)谕瑯拥挠布h(huán)境下訓(xùn)練，我們總的訓(xùn)練前時(shí)間比其他任何方法都短。例如，對(duì)于 ViT-L，MAE 的訓(xùn)練時(shí)間是1600個(gè) epochs，31小時(shí)。moco v3的訓(xùn)練時(shí)間是300個(gè) epochs，36小時(shí)，使用相同的128個(gè) TPU-v3 核心。

1.12 Partial Fine-tuning

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

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

總結(jié)

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

如果覺(jué)得有用，就請(qǐng)分享到朋友圈吧！