Masked Autoencoders Are Scalable Vision Learners

--- title: 'Masked Autoencoders Are Scalable Vision Learners' tags: 'AI' --- # Masked Autoencoders Are Scalable Vision Learners(2021) [**原論文網址**](https://arxiv.org/abs/2111.06377) [**回到目錄**](https://jacksonchen890301.github.io/Jackson-Domain/) 撰寫日期: 2022/4/20 ## 概述 Masked Autoencoder(MAE)是以Vision Transformer為基礎所設計的Self-Supervised模型。由非對稱的encoder-decoder架構組成，encoder負責處理沒被刪去的patches之訊息，decoder則負責從latent space的特徵重建回原圖。另外，作者經由實驗發現，遮蓋75%的patches能讓MAE的學習效果最佳化，並且同時可以將計算速度加快約3倍(因為資料量被遮蓋的過程所減少了)。以此方法訓練的ViT-H模型，在不使用額外資料集的條件下，於ImageNet1K上取得了87.8%的最高準確率。 <center>下圖是模型的大致結構</center> <center><img src="https://i.imgur.com/mhqXDzM.png" width="70%" alt="img01"/></center> ## Intriduction Masked Autoencoding是透過遮蓋部分訊息，讓模型學習生成這些被刪除的資訊。此方法在NLP領域的模型上成效斐然，然而，在vision相關任務上卻沒有甚麼成果。主要原因可能是因為：(1)在CV上，CNN在近幾年有著主宰地位，但CNN顯然在結構上跟NLP的主流模型*Transformer*，有著非常大的差距。因為CNN沒有辦法套用Masked Token這種Indicator，以及Position Embedding這種封裝位置資訊的方法。不過，在ViT問世後，這些結構上的問題都可以被迎刃而解。(2)語言的資訊密度顯然比圖像高出許多，因此，預測一段語句中缺失的單詞是需要對語言有充足的了解。但是，若是要升成一個圖像中缺失的patch，只要一點點對於物件跟場景的理解，就能從鄰近的patch判斷出缺失的部份。所以在實作上，必須隨機遮蓋住大面積的圖像，才能讓這個生成任務足夠困難，讓模型學到足夠整體性的圖像資訊才能成功還原缺失的patch。 ## Approach MAE模型的任務是從片段的資訊，還原回原始的圖像。跟普通的Autoencoder一樣，encoder將輸入映射到潛在特徵空間(latent representation)，而decoder能夠將這些潛在特徵重建回原本的輸入。為了能夠讓模型只讀到部分的圖像資訊，MAE的設計使用了非對稱的結構。 * Masking：使用了uniform distribution來遮蓋大比例圖像面積，一方便避免被遮蓋區域不夠平均，另一方面則是讓這個圖像生成任務有足夠的困難性，不會簡單的被外推(extrapolation)或內插(interpolation)出來。 <center><img src="https://i.imgur.com/CGvSmva.png" width="70%" alt="img01"/></center> * MAE encoder：encoder輸入的patches會先經由linear projection再加上position embedding。之後，隨機遮蓋掉75%的patches，因此在encoder運算時只需要處理25%的剩餘資訊，可有效加速訓練時間。 * MAE decoder：decoder的輸入包含了被encode的visible patches還有mask tokens，並且都會加上position embedding。每個mask token都是共享權重，而且是在訓練過程中學習出來的vector，所以對decoder來說，讀到mask token就代表這個位置是被遮蓋的。decoder只會在self-supervised learning的訓練中使用到，而且經由實驗，decoder的設計對於後續encoder在downstream task的表現影響甚微，所以論文中是把decoder盡量設計的輕量化，以加速訓練流程。 * Reconstruction target：每個decoder輸出的token各代表一個patch，最後透過linear projection將token的channel數映射成跟原始圖像pixel values相等，就完成了圖像的還原。誤差的計算則是使用mean square error(MSE)，將原始的patch與重建出來的patch做比較。 ## ImageNet Experiments 使用ImageNet-1K(IN1K)在ViT Large模型上做self-supervised pre-training，並且以(1)end-to-end fine-tuning 或(2)linear probing去評估模型的性能。下圖為3種不同訓練方式的比較，可見MAE效果較佳。 <center><img src="https://i.imgur.com/IeMLe26.png" width="70%" alt="img01"/></center> 以下是對MAE進行的一系列實驗 * Masking ratio：一般來說，在NLP上masking的比例只有約15%，在linear probing實驗中，可以看到模型的性能隨著masking ratio逐漸上升，直到75%的ratio才停止。雖然在finetune上，ratio與性能的變化較不明顯，但都可以大致得出75%是一個相當合適的比例。而且，不管ratio如何變動，結果都比train from scratch(82.5% accuracy)來的好。 <center><img src="https://i.imgur.com/Kdju3w1.png" width="70%" alt="img01"/></center> * Decoder design：雖然在linear probing時，較深的decoder表現較佳(可以理解為，夠深的decoder可以把reconstruction任務與encoder找特徵的任務分開)，然而只要進行finetune，權重少的decoder訓練出來的模型也能有著很高的accuracy。所以，考慮到計算速度，使用比較輕量的decoder會是比較好的選擇。 <center><img src="https://i.imgur.com/PtCw0EM.png" width="90%" alt="img01"/></center> * Mask token：若將encoder的輸入以mask代替，而不是直接移除的話，會發現反而模型性能變差。可能是因為如果使用mask取代被移除的patches，會讓encoder看到一些根本不屬於圖像的資訊(而且是有75%的資訊都是被mask取代)，因此造成encoder無法歸納出圖像的特徵。所以綜合上述，把不使用mask不但可以增進模型的表現，還能夠同時減少資算量，可以算是一舉兩得。 <center><img src="https://i.imgur.com/t5ePlHK.png" width="50%" alt="img01"/></center> * Reconstruction target：如果分別對每個patch做normalization，會有比較好的訓練效果。估計是因為對每個patch做norm.可以增加對比度，讓模型更容易學習。另外，也嘗試了讓模型去重建出降維過後的圖像(基本上就是壓縮後的圖)，反而效果變差還增加了不必要的計算開銷(壓縮也需要計算量)。 <center><img src="https://i.imgur.com/ObOfuhQ.png" width="50%" alt="img01"/></center> * Data augmentation：發現pretain時若增加augmentation反而降低模型性能。與其相反的，contrasitive learning其實相當吃重這些augmentation手段(ex.SimCLR, BYOL)這邊論文歸納出的原因是，隨機masking本身就是一種augmentation手段，讓每次iteration的圖像都有所差距，而且MAE主要是透過這些隨機遮蓋去做學習，而不是透過那些augmentation的差異性。 <center><img src="https://i.imgur.com/wG6ddns.png" width="50%" alt="img01"/></center> * Mask sampling strategy：這邊使用了三種方式做masking，分別是隨機(random)、連續區塊(block)、以及網格(grid)。很直覺的，隨機的做法比較好，因為block的難度太高，像是下圖，中央區域被全部挖空，根本沒有足夠的資訊量模型腦補出中間的圖案。而grid則是太過簡單，沒有隨機性幫助模型學到更多不同的狀況。 <center><img src="https://i.imgur.com/acIQCtc.png" width="70%" alt="img01"/></center> <center><img src="https://i.imgur.com/nIRAUy9.png" width="50%" alt="img01"/></center> * Training schedule：發現pretrain的epoch越長，downstream task的表現越好，甚至到1600個epoch都還沒有飽和(saturate)的現象。 <center><img src="https://i.imgur.com/DoNGZ1m.png" width="80%" alt="img01"/></center> ## Comparisons with self-supervised methods 可以發現，相較於先前的其他ViT self-supervised learning方法，MAE不但表現的更好，而且隨著模型的規模越大，與其他方法的差距也跟著放大。因此，MAE會是一種能夠有效訓練高參數模型的做法，而且在執行效率以及方法設計上都比較簡單能夠實作。 <center><img src="https://i.imgur.com/tbGuW4e.png" width="80%" alt="img01"/></center> ## Comparisons with supervised pre-training 根據ViT原論文，使用IN1K訓練出來的ViT模型，在參數量過大時反而性能會下降，必須使用更大規模的資料集(JFT-300M)，才能訓練好高參數量的ViT-L。而本文使用的MAE方法，用較小規模的IN1K資料集，就能夠有效使高參數模型之性能能夠提升(下圖可見MAE訓練出的模型和JFT-300M這種超大資料集pretrain出來的模型，同樣能夠受益於參數量的提升)，且明顯比supervised的訓練方式來的好。 <center><img src="https://i.imgur.com/Qr3xo2Q.png" width="80%" alt="img01"/></center>