愷明大神在2021年提出掩碼自編碼器（MAE），一種簡單的自監(jiān)督學習方法，通過隨機掩蓋輸入圖像的部分區(qū)域并重建缺失的像素。主要研究動機如下：

• 數(shù)據(jù)需求與過擬合問題：深度學習模型，尤其是視覺領域的模型，需要大量數(shù)據(jù)來避免過擬合。自然語言處理（NLP）領域通過自監(jiān)督預訓練（如BERT）成功解決了這一問題，但在計算機視覺領域，類似的自監(jiān)督學習方法發(fā)展相對滯后。作者希望找到一種類似的有效方法來推動計算機視覺的發(fā)展。
• 信息密度差異：語言信號是高信息密度的，而圖像信號則具有較高的空間冗余。這導致直接將NLP領域的自監(jiān)督學習方法應用于視覺領域時效果不佳。
• 解碼器角色差異：在NLP中，解碼器預測的是具有豐富語義信息的詞匯，而在視覺領域，解碼器重建的是像素級別的圖像，語義層次較低。因此需要精心設計解碼器以提升模型性能。

下面來簡單看下這個算法：

01、方法介紹

Mask 策略

圖像分塊：首先，將輸入圖像劃分為規(guī)則的、不重疊的patches(ViT 中是 16x16 大小)。

隨機采樣：然后，使用均勻分布(uniform distribution)的采樣策略隨機采樣一部分塊，并mask剩余塊掩碼。被 mask 掉的 patches 占所有 patches 的大部分(實驗效果發(fā)現(xiàn)最好的比例是 75%)，它們不會輸入到 Encoder。

這種方式的優(yōu)點：服從均勻分布采樣，能夠避免潛在的中心歸納偏好，即避免 patch 的位置大多都分布在靠近圖像中心的區(qū)域；另外，采用高掩碼比例(mask 掉圖中大部分 patches)能夠防止模型輕易地根據(jù)鄰近的可見 patches 推斷出這些掩碼塊；最后，這種策略還造就了稀疏的編碼器輸入，因為 Encoder 只處理可見的 patches，計算量和內(nèi)存占用都減少，能夠以更低的代價訓練較大規(guī)模的 Encoder。

高掩碼比率：高掩碼比率（即移除的塊的比例）可以大幅減少冗余，從而創(chuàng)建一個挑戰(zhàn)性的任務，這個任務不能簡單地通過從可見的鄰近塊進行外推來解決。

作者通過實驗對比了不同mask 比例和不同 mask 策略進行訓練后模型的表現(xiàn)，無論是 fine-tune 還是 linear-probe 下，75% 左右的 mask 比例都表現(xiàn)最好，隨機(服從均勻分布)采樣策略下模型的表現(xiàn)最好。

Encoder

編碼器結構：MAE的編碼器基于Vision Transformer（ViT），但僅應用于可見的、未被掩碼的小塊。編碼器將小塊通過線性投影和位置嵌入進行嵌入，然后通過一系列Transformer塊進行處理。

計算效率：由于編碼器僅處理全部塊的一小部分（例如25%），這允許使用更少的計算和內(nèi)存資源來訓練非常大的編碼器。

劃分patch策略：先將圖像從 (B,C,H,W) reshape 成 (B,N,PxPxC)，其中 N 和 P 分別為 patch 數(shù)量 和 patch 大小()，也就是將3通道的圖像轉換成 N 個 維度大小為 PxPxC 的向量；然后，通過線性映射(linear projection，可以是全連接層)將其嵌入(embed)到指定的維度空間大小，記為'dim'(從PxPxC project到dim)，轉換成為 token(B,N,dim)；最后再加上位置嵌入(position embedding)，從而為各個 patch 添加位置信息。位置嵌入是所有圖像共享的、可學習的，shape 與 每張圖的 token 相對應，即：(N,dim)。

Decoder

解碼器輸入：MAE解碼器的輸入包括Encoder 編碼的 un-masked 的 tokens和掩碼標記（mask tokens）。掩碼標記masked token 并非由之前 mask 掉的 patch 經(jīng)過 embedding 轉換而來，而是可學習的、所有 masked patches 都共享的1個向量！

位置嵌入：為了給掩碼標記提供位置信息，對所有標記添加位置嵌入，每個 masked patch 對應1個，shape 是 (N',dim)，其中 N' 是 masked patch 的數(shù)量，從而也能區(qū)分各個 masked patch 所對應的 token。

解碼器結構：解碼器擁有一系列Transformer塊， 僅僅是在預訓練任務為了重建圖像而存在，下游任務形式多種多樣，因此實際應用時可能不再使用。所以，Decoder 的設計和 Encoder 是解耦的，Decoder 可以設計得簡單、輕量(比 Encoder 更窄、更淺。窄：對應通道數(shù)；淺：對應深度)，畢竟主要學習潛在特征表示的是 Encoder。

任務目標：重建像素值

像素重建：MAE的目標是通過預測每個掩碼塊（ masked patches）的像素值來重建輸入，解碼器的輸出是代表小塊像素值的向量，最后一個線性層的輸出通道數(shù)等于patch中的像素值數(shù)量。

均方誤差（MSE）：使用均方誤差作為重建和原始圖像之間的損失函數(shù)，僅對 mask 掉的部分計算 loss。這里作者也對全部patches計算loss進行了實驗，發(fā)現(xiàn)準確率會降低0.5%。

模型是如何去預測 masked patches 的像素值并計算 loss 的呢？具體來說，就是：

在 Decoder 解碼后的所有 tokens 中取出 masked tokens(在最開始 mask 掉 patches 的時候可以先記錄下這些 masked 部分的索引)，將這些 masked tokens 送入全連接層，將輸出通道映射到1個 patch 的像素數(shù)量(PxPxC)，也就是輸出的 shape 是：(B,N',PxPxC)，其中的每個值就代表預測的像素值。最后，以之前 mask 掉的 patches 的像素值作為 target，與預測結果計算 MSE loss。

另外，作者提到使用歸一化的像素值作為 target 效果更好，能夠提升學到的表征的質(zhì)量。這里的歸一化做法是：計算每個 patch 像素值的均值與標準差，然后用均值與標準差去歸一化對應的 patch 像素。

Pipeline

1. 將圖像劃分成 patches：(B,C,H,W)->(B,N,PxPxC)；
2. 對各個 patch 進行 embedding(實質(zhì)是通過全連接層)，生成 tokens，并加入位置信息(position embeddings)：(B,N,PxPxC)->(B,N,dim)；
3. 根據(jù)預設的掩碼比例(paper 中提倡的是 75%)，使用服從均勻分布的隨機采樣策略采樣一部分 tokens 送給 Encoder，另一部分“扔掉”(mask 掉)；
4. 將 Encoder 編碼后的 tokens 與 加入位置信息后的 masked tokens 按照原先在 patch 形態(tài)時對應的次序拼在一起，然后喂給 Decoder 玩(如果 Encoder 編碼后的 token 的維度與 Decoder 要求的輸入維度不一致，則需要先經(jīng)過 linear projection 將維度映射到符合 Decoder 的要求)；
5. Decoder 解碼后取出 masked tokens 對應的部分送入到全連接層，對 masked patches 的像素值進行預測，最后將預測結果與 masked patches 進行比較，計算 MSE loss

02、實驗結果

首先看重建效果，還是很驚艷的

Mask比例

mask 比例較高才能形成具有挑戰(zhàn)性的預訓練任務，模型才更有機會學到更好的潛在特征表示。由實驗結果可以看到，無論是在 fine-tune 還是 linear probe 中，mask 比例逐漸升高時，模型性能都會更好。但是，fine-tune 和 linear probe 的結果還是有所區(qū)別的：linear probe 幾乎是線性增漲的趨勢，而 fine-tune 則是 mask 比例在 30%~40% 之間激增，而后就傾向于飽和了。

這里的解釋：linear probe 之所以沒有那么快飽和，與其本身僅調(diào)整模型最后的幾層分類頭相關。因此，mask 比例越高，在預訓練時得到的 Encoder 就越強，但這部分在下游任務中不再被訓練，所以其性能就隨著 mask 比例的增加呈線性增漲的趨勢。相對地，fine-tune 還繼續(xù)訓練 Encoder 的參數(shù)去適配下游任務，因此在 mask 比例超過一定程度后，對于下游任務的性能提升就不那么明顯了。

Mask采樣策略

通過實驗比較，最終選擇了服從均勻分布的隨機采樣策略，下面是詳細的實驗結果?？梢杂^察出，block-wise 策略由于掩蓋掉的圖像塊區(qū)域太大了，因此在高于 50% 的 mask 比例下效果就不好。grid 策略在訓練時能夠對數(shù)據(jù)擬合得很好，但實際學到的特征表示泛化性其實是比較弱的。由此可以說明，代理任務設計得太困難(對應 block-wise)或太簡單(對應 grid)都不行，要適當(對應 random)才好。

Decoder 設計

Decoder 的深度和寬度對于 linear probe 有較為明顯的影響，但對于 fine-tune 的影響卻不那么突出。

因為當 Decoder 更深/寬時，它本身會擁有更強的重建能力，這樣就使得在預訓練時 Encoder 能更專注于提取抽象語義層級的特征。linear probe 是完全繼承預訓練的 Encoder，而 fine-tune 在下游任務中仍能夠繼續(xù)調(diào)整 Encoder 的參數(shù)適配下游任務，因此預訓練對其影響程度就沒那么大。究其本質(zhì)，其實是預訓練任務(圖像重建)與下游任務(圖像識別)之間存在著 gap！

Masked token

實驗結果顯示，如果編碼器使用掩碼標記，它的性能會更差。因為在下游任務中并不存在這些 masked tokens，上、下游任務之間存在 gap(參考BERT)。如果 Encoder 也對 masked tokens 進行編碼，會進一步將這種 gap 的影響“擴散”到下游任務。

重建目標比較

MAE 的重建目標是 masked patches 的像素值。同時，作者在 paper 中還提到，如果預測的是歸一化的像素值，那么效果會更好。另外，作者還和 BEiT 預測 token 的方式以及PCA 的方式(對 patch 空間實施 PCA 并預測最大的因子)進行了比較：

結果顯示，預測歸一化像素值的方式最強。

這里歸一化像素值的做法是分別針對每個 patch 使用它們獨立統(tǒng)計出來的均值與方差去歸一化的，這就會將各個 patch 歸一化到不同的表示空間，從而分成不同的“簇”，于是各個 patch 之間的差異性就更強，形成了高頻信息，相當于將各個 patch 構造成了邊緣與紋理，從整體圖像看來，對比度更高。從而使得模型更有針對性地學習各個 patch 的特征模式。同時，數(shù)值上由于做了歸一化，因此又不會使得模型在這方面有所偏倚。

至于 token 的方式是照搬 NLP 的玩法，是高度離散化和語義化的，一個字的差異也可能導致詞語之間的含義發(fā)生重大變化，本身就是高頻東西。

因此，究其本質(zhì)：高頻性質(zhì)的目標能夠“迫使”特征提取器(Encoder)編碼出來的各類特征更有差異性和區(qū)分性，它們之間的 margin 更大。

數(shù)據(jù)增強

不做隨機縮放(fixed size)和隨機縮放(rand size)的效果差不多，而采用色彩擾動(color jit)比簡單的 crop(vs fixed size）效果差 。這里應該是 MAE 對圖像進行 mask 的做法本身就已經(jīng)是一種數(shù)據(jù)增強手段，因此不需要過度的額外數(shù)據(jù)增強就能取得較好的效果。

局部微調(diào)

通過上面的實驗結果可以發(fā)現(xiàn)，linear probe與 fine-tune 之間總是存在著“不協(xié)同”的結果，比如前面說到的 Decoder 的深度和寬度對 linear probe 的影響挺大但對于 fine-tune 來說卻并不那么事關緊要。

作者“權衡”了 linear probe 和 fine-tune 之間的做法，設計出partial fine-tuning'：僅調(diào)整 Encoder 的最后幾層但 fix 住其它部分。如圖所示，調(diào)整 0 個 block 相當于是 linear probe，而調(diào)整所有24個 blocks 就是 fine-tuning。

可以看到，對于 MAE，僅調(diào)整1個 block 就可以將 acc 從73.5%(linear probe)漲到81%，并且對于 MOCO v3 也一樣可以漲點。

另外，MAE 在 partial fine-tuning 的方式下優(yōu)于 MOCO v3，這也說明 MAE 學到的特征非線性更強，于是當可以調(diào)整非線性頭部時效果就更好。

這些現(xiàn)象都表明：linear probe 并非是唯一的、正確地評估模型學到的表征質(zhì)量的方式。并且，作者后續(xù)還進行了 detection 與 segmentation 相關的實驗，從而在 linear probe 中學到的特征也并非是和遷移學習性能強相關的。

03、總結

第一篇把自編碼器在圖像上做出效果的文章。MAE為計算機視覺領域帶來了高效的自監(jiān)督學習方法，縮小了計算機視覺與自然語言處理在自監(jiān)督學習方面的差距。通過對圖像進行大量隨機掩碼，減少圖像空間冗余，使模型能夠學習到超越低級圖像統(tǒng)計的整體理解，學到更高級有效的特征。同時，MAE 的出現(xiàn)也讓研究者更加重視解碼器在視覺任務中的關鍵作用，推動了計算機視覺的發(fā)展，為后續(xù)的研究和應用（如在圖像分類、目標檢測等任務）提供了新的思路和方法。致敬大神。