模態(tài)編碼器 | 愷明大神的MAE--掩碼自編碼器
愷明大神在2021年提出掩碼自編碼器（MAE），一種簡單的自監(jiān)督學(xué)習(xí)方法，通過隨機(jī)掩蓋輸入圖像的部分區(qū)域并重建缺失的像素。主要研究動機(jī)如下：
- 數(shù)據(jù)需求與過擬合問題：深度學(xué)習(xí)模型，尤其是視覺領(lǐng)域的模型，需要大量數(shù)據(jù)來避免過擬合。自然語言處理（NLP）領(lǐng)域通過自監(jiān)督預(yù)訓(xùn)練（如BERT）成功解決了這一問題，但在計算機(jī)視覺領(lǐng)域，類似的自監(jiān)督學(xué)習(xí)方法發(fā)展相對滯后。作者希望找到一種類似的有效方法來推動計算機(jī)視覺的發(fā)展。
- 信息密度差異：語言信號是高信息密度的，而圖像信號則具有較高的空間冗余。這導(dǎo)致直接將NLP領(lǐng)域的自監(jiān)督學(xué)習(xí)方法應(yīng)用于視覺領(lǐng)域時效果不佳。
- 解碼器角色差異：在NLP中，解碼器預(yù)測的是具有豐富語義信息的詞匯，而在視覺領(lǐng)域，解碼器重建的是像素級別的圖像，語義層次較低。因此需要精心設(shè)計解碼器以提升模型性能。
下面來簡單看下這個算法：
01、方法介紹
Mask 策略
圖像分塊：首先，將輸入圖像劃分為規(guī)則的、不重疊的patches(ViT 中是 16x16 大小)。
隨機(jī)采樣：然后，使用均勻分布(uniform distribution)的采樣策略隨機(jī)采樣一部分塊，并mask剩余塊掩碼。被 mask 掉的 patches 占所有 patches 的大部分(實驗效果發(fā)現(xiàn)最好的比例是 75%)，它們不會輸入到 Encoder。
這種方式的優(yōu)點：服從均勻分布采樣，能夠避免潛在的中心歸納偏好，即避免 patch 的位置大多都分布在靠近圖像中心的區(qū)域；另外，采用高掩碼比例(mask 掉圖中大部分 patches)能夠防止模型輕易地根據(jù)鄰近的可見 patches 推斷出這些掩碼塊；最后，這種策略還造就了稀疏的編碼器輸入，因為 Encoder 只處理可見的 patches，計算量和內(nèi)存占用都減少，能夠以更低的代價訓(xùn)練較大規(guī)模的 Encoder。
高掩碼比率：高掩碼比率（即移除的塊的比例）可以大幅減少冗余，從而創(chuàng)建一個挑戰(zhàn)性的任務(wù)，這個任務(wù)不能簡單地通過從可見的鄰近塊進(jìn)行外推來解決。
作者通過實驗對比了不同mask 比例和不同 mask 策略進(jìn)行訓(xùn)練后模型的表現(xiàn)，無論是 fine-tune 還是 linear-probe 下，75% 左右的 mask 比例都表現(xiàn)最好，隨機(jī)(服從均勻分布)采樣策略下模型的表現(xiàn)最好。
Encoder
編碼器結(jié)構(gòu)：MAE的編碼器基于Vision Transformer（ViT），但僅應(yīng)用于可見的、未被掩碼的小塊。編碼器將小塊通過線性投影和位置嵌入進(jìn)行嵌入，然后通過一系列Transformer塊進(jìn)行處理。
計算效率：由于編碼器僅處理全部塊的一小部分（例如25%），這允許使用更少的計算和內(nèi)存資源來訓(xùn)練非常大的編碼器。
劃分patch策略：先將圖像從 (B,C,H,W) reshape 成 (B,N,PxPxC)，其中 N 和 P 分別為 patch 數(shù)量 和 patch 大小(
)，也就是將3通道的圖像轉(zhuǎn)換成 N 個 維度大小為 PxPxC 的向量；然后，通過線性映射(linear projection，可以是全連接層)將其嵌入(embed)到指定的維度空間大小，記為'dim'(從PxPxC project到dim)，轉(zhuǎn)換成為 token(B,N,dim)；最后再加上位置嵌入(position embedding)，從而為各個 patch 添加位置信息。位置嵌入是所有圖像共享的、可學(xué)習(xí)的，shape 與 每張圖的 token 相對應(yīng)，即：(N,dim)。
Decoder
解碼器輸入：MAE解碼器的輸入包括Encoder 編碼的 un-masked 的 tokens和掩碼標(biāo)記（mask tokens）。掩碼標(biāo)記masked token 并非由之前 mask 掉的 patch 經(jīng)過 embedding 轉(zhuǎn)換而來，而是可學(xué)習(xí)的、所有 masked patches 都共享的1個向量！
位置嵌入：為了給掩碼標(biāo)記提供位置信息，對所有標(biāo)記添加位置嵌入，每個 masked patch 對應(yīng)1個，shape 是 (N',dim)，其中 N' 是 masked patch 的數(shù)量，從而也能區(qū)分各個 masked patch 所對應(yīng)的 token。
解碼器結(jié)構(gòu)：解碼器擁有一系列Transformer塊， 僅僅是在預(yù)訓(xùn)練任務(wù)為了重建圖像而存在，下游任務(wù)形式多種多樣，因此實際應(yīng)用時可能不再使用。所以，Decoder 的設(shè)計和 Encoder 是解耦的，Decoder 可以設(shè)計得簡單、輕量(比 Encoder 更窄、更淺。窄：對應(yīng)通道數(shù)；淺：對應(yīng)深度)，畢竟主要學(xué)習(xí)潛在特征表示的是 Encoder。
任務(wù)目標(biāo)：重建像素值
像素重建：MAE的目標(biāo)是通過預(yù)測每個掩碼塊（ masked patches）的像素值來重建輸入，解碼器的輸出是代表小塊像素值的向量，最后一個線性層的輸出通道數(shù)等于patch中的像素值數(shù)量。
均方誤差（MSE）：使用均方誤差作為重建和原始圖像之間的損失函數(shù)，僅對 mask 掉的部分計算 loss。這里作者也對全部patches計算loss進(jìn)行了實驗，發(fā)現(xiàn)準(zhǔn)確率會降低0.5%。
模型是如何去預(yù)測 masked patches 的像素值并計算 loss 的呢？具體來說，就是：
在 Decoder 解碼后的所有 tokens 中取出 masked tokens(在最開始 mask 掉 patches 的時候可以先記錄下這些 masked 部分的索引)，將這些 masked tokens 送入全連接層，將輸出通道映射到1個 patch 的像素數(shù)量(PxPxC)，也就是輸出的 shape 是：(B,N',PxPxC)，其中的每個值就代表預(yù)測的像素值。最后，以之前 mask 掉的 patches 的像素值作為 target，與預(yù)測結(jié)果計算 MSE loss。
另外，作者提到使用歸一化的像素值作為 target 效果更好，能夠提升學(xué)到的表征的質(zhì)量。這里的歸一化做法是：計算每個 patch 像素值的均值與標(biāo)準(zhǔn)差，然后用均值與標(biāo)準(zhǔn)差去歸一化對應(yīng)的 patch 像素。
Pipeline
- 將圖像劃分成 patches：(B,C,H,W)->(B,N,PxPxC)；
- 對各個 patch 進(jìn)行 embedding(實質(zhì)是通過全連接層)，生成 tokens，并加入位置信息(position embeddings)：(B,N,PxPxC)->(B,N,dim)；
- 根據(jù)預(yù)設(shè)的掩碼比例(paper 中提倡的是 75%)，使用服從均勻分布的隨機(jī)采樣策略采樣一部分 tokens 送給 Encoder，另一部分“扔掉”(mask 掉)；
- 將 Encoder 編碼后的 tokens 與 加入位置信息后的 masked tokens 按照原先在 patch 形態(tài)時對應(yīng)的次序拼在一起，然后喂給 Decoder 玩(如果 Encoder 編碼后的 token 的維度與 Decoder 要求的輸入維度不一致，則需要先經(jīng)過 linear projection 將維度映射到符合 Decoder 的要求)；
- Decoder 解碼后取出 masked tokens 對應(yīng)的部分送入到全連接層，對 masked patches 的像素值進(jìn)行預(yù)測，最后將預(yù)測結(jié)果與 masked patches 進(jìn)行比較，計算 MSE loss
02、實驗結(jié)果
首先看重建效果，還是很驚艷的
Mask比例
mask 比例較高才能形成具有挑戰(zhàn)性的預(yù)訓(xùn)練任務(wù)，模型才更有機(jī)會學(xué)到更好的潛在特征表示。由實驗結(jié)果可以看到，無論是在 fine-tune 還是 linear probe 中，mask 比例逐漸升高時，模型性能都會更好。但是，fine-tune 和 linear probe 的結(jié)果還是有所區(qū)別的：linear probe 幾乎是線性增漲的趨勢，而 fine-tune 則是 mask 比例在 30%~40% 之間激增，而后就傾向于飽和了。
這里的解釋：linear probe 之所以沒有那么快飽和，與其本身僅調(diào)整模型最后的幾層分類頭相關(guān)。因此，mask 比例越高，在預(yù)訓(xùn)練時得到的 Encoder 就越強(qiáng)，但這部分在下游任務(wù)中不再被訓(xùn)練，所以其性能就隨著 mask 比例的增加呈線性增漲的趨勢。相對地，fine-tune 還繼續(xù)訓(xùn)練 Encoder 的參數(shù)去適配下游任務(wù)，因此在 mask 比例超過一定程度后，對于下游任務(wù)的性能提升就不那么明顯了。
Mask采樣策略
通過實驗比較，最終選擇了服從均勻分布的隨機(jī)采樣策略，下面是詳細(xì)的實驗結(jié)果。可以觀察出，block-wise 策略由于掩蓋掉的圖像塊區(qū)域太大了，因此在高于 50% 的 mask 比例下效果就不好。grid 策略在訓(xùn)練時能夠?qū)?shù)據(jù)擬合得很好，但實際學(xué)到的特征表示泛化性其實是比較弱的。由此可以說明，代理任務(wù)設(shè)計得太困難(對應(yīng) block-wise)或太簡單(對應(yīng) grid)都不行，要適當(dāng)(對應(yīng) random)才好。
Decoder 設(shè)計
Decoder 的深度和寬度對于 linear probe 有較為明顯的影響，但對于 fine-tune 的影響卻不那么突出。
因為當(dāng) Decoder 更深/寬時，它本身會擁有更強(qiáng)的重建能力，這樣就使得在預(yù)訓(xùn)練時 Encoder 能更專注于提取抽象語義層級的特征。linear probe 是完全繼承預(yù)訓(xùn)練的 Encoder，而 fine-tune 在下游任務(wù)中仍能夠繼續(xù)調(diào)整 Encoder 的參數(shù)適配下游任務(wù)，因此預(yù)訓(xùn)練對其影響程度就沒那么大。究其本質(zhì)，其實是預(yù)訓(xùn)練任務(wù)(圖像重建)與下游任務(wù)(圖像識別)之間存在著 gap！
Masked token
實驗結(jié)果顯示，如果編碼器使用掩碼標(biāo)記，它的性能會更差。因為在下游任務(wù)中并不存在這些 masked tokens，上、下游任務(wù)之間存在 gap(參考BERT)。如果 Encoder 也對 masked tokens 進(jìn)行編碼，會進(jìn)一步將這種 gap 的影響“擴(kuò)散”到下游任務(wù)。
重建目標(biāo)比較
MAE 的重建目標(biāo)是 masked patches 的像素值。同時，作者在 paper 中還提到，如果預(yù)測的是歸一化的像素值，那么效果會更好。另外，作者還和 BEiT 預(yù)測 token 的方式以及PCA 的方式(對 patch 空間實施 PCA 并預(yù)測最大的因子)進(jìn)行了比較：
結(jié)果顯示，預(yù)測歸一化像素值的方式最強(qiáng)。
這里歸一化像素值的做法是分別針對每個 patch 使用它們獨(dú)立統(tǒng)計出來的均值與方差去歸一化的，這就會將各個 patch 歸一化到不同的表示空間，從而分成不同的“簇”，于是各個 patch 之間的差異性就更強(qiáng)，形成了高頻信息，相當(dāng)于將各個 patch 構(gòu)造成了邊緣與紋理，從整體圖像看來，對比度更高。從而使得模型更有針對性地學(xué)習(xí)各個 patch 的特征模式。同時，數(shù)值上由于做了歸一化，因此又不會使得模型在這方面有所偏倚。
至于 token 的方式是照搬 NLP 的玩法，是高度離散化和語義化的，一個字的差異也可能導(dǎo)致詞語之間的含義發(fā)生重大變化，本身就是高頻東西。
因此，究其本質(zhì)：高頻性質(zhì)的目標(biāo)能夠“迫使”特征提取器(Encoder)編碼出來的各類特征更有差異性和區(qū)分性，它們之間的 margin 更大。
數(shù)據(jù)增強(qiáng)
不做隨機(jī)縮放(fixed size)和隨機(jī)縮放(rand size)的效果差不多，而采用色彩擾動(color jit)比簡單的 crop(vs fixed size）效果差 。這里應(yīng)該是 MAE 對圖像進(jìn)行 mask 的做法本身就已經(jīng)是一種數(shù)據(jù)增強(qiáng)手段，因此不需要過度的額外數(shù)據(jù)增強(qiáng)就能取得較好的效果。
局部微調(diào)
通過上面的實驗結(jié)果可以發(fā)現(xiàn)，linear probe與 fine-tune 之間總是存在著“不協(xié)同”的結(jié)果，比如前面說到的 Decoder 的深度和寬度對 linear probe 的影響挺大但對于 fine-tune 來說卻并不那么事關(guān)緊要。
作者“權(quán)衡”了 linear probe 和 fine-tune 之間的做法，設(shè)計出partial fine-tuning'：僅調(diào)整 Encoder 的最后幾層但 fix 住其它部分。如圖所示，調(diào)整 0 個 block 相當(dāng)于是 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 學(xué)到的特征非線性更強(qiáng)，于是當(dāng)可以調(diào)整非線性頭部時效果就更好。
這些現(xiàn)象都表明：linear probe 并非是唯一的、正確地評估模型學(xué)到的表征質(zhì)量的方式。并且，作者后續(xù)還進(jìn)行了 detection 與 segmentation 相關(guān)的實驗，從而在 linear probe 中學(xué)到的特征也并非是和遷移學(xué)習(xí)性能強(qiáng)相關(guān)的。
03、總結(jié)
第一篇把自編碼器在圖像上做出效果的文章。MAE為計算機(jī)視覺領(lǐng)域帶來了高效的自監(jiān)督學(xué)習(xí)方法，縮小了計算機(jī)視覺與自然語言處理在自監(jiān)督學(xué)習(xí)方面的差距。通過對圖像進(jìn)行大量隨機(jī)掩碼，減少圖像空間冗余，使模型能夠?qū)W習(xí)到超越低級圖像統(tǒng)計的整體理解，學(xué)到更高級有效的特征。同時，MAE 的出現(xiàn)也讓研究者更加重視解碼器在視覺任務(wù)中的關(guān)鍵作用，推動了計算機(jī)視覺的發(fā)展，為后續(xù)的研究和應(yīng)用（如在圖像分類、目標(biāo)檢測等任務(wù)）提供了新的思路和方法。致敬大神。
