Masked Autoencoders: 借鏡BERT與ViT的Self-Supervised Learners

Jia-Yau Shiau
AI Blog TW
Published in
22 min readNov 21, 2021

長期以來, deep learning在 NLP與 CV領域建模 (modeling) 擁有各自強大的慣例。在架構部分是比較顯而易見的, NLP領域使用 attention為基礎的Transformer,而 CV領域使用有更強的 inductive biases的 CNN。其實在此之外,還有一個隱性的差異在於預訓練 (pre-training) 模型的方法

在 NLP領域, self-supervised pre-training是行之有年的方法,如採用 MLM (Masked-Language Modeling) 的 BERT。然而, CV領域即使這兩三年出現許多 self-supervised方法,如 SimCLRMoCoBYOL等等,最常被採用的還是 supervised pre-training或是 joint training。

MAE (Masked Autoencoders)是 FAIR在 2021年提出,借鏡 BERT的 MLM概念、結合切割影像為 patch的 ViT,精煉出不需要 contractive架構的 self-supervised pre-training方法。論文中也討論了 image與 language的資料差異,為未來 CV領域 self-supervised pre-training點出可行方向。

Cover made with Canva (圖片來源)

文章難度:★★★★☆
閱讀建議: 這篇文章主要介紹 Masked Autoencoders與 self-supervised learning在 CV與 NLP領域的一些異同。文章會先簡單介紹一些地基知識,包含: (1) Attention與 Transformer的概念; (2) Masked-Language Modeling的訓練方法; (3) ViT (Vision Transformer) 的概念等等。才會開始討論 Masked Autoencoders的方法與實驗數據。當然也建議可以搭配一些 BERT與 self-supervised learning的文章相輔學習。
推薦背景知識:self-supervised learning, pre-training, CNN, RNN, attention, Transformer, ViT.

Attention & Transformer

深度學習在人工智慧的發展中,許多不同領域的應用都有自己主流的一套方法或是模型。在自然語言處理 (Natural Language Processing, NLP) 中,近幾年主流的模型是 2017年 Google團隊提出 multi-head self-attention為基礎的 Transformer [1]。

Attention

Attention的大方向很簡單,可以理解成利用網路預測的權重來組和多個 feature

以 text sequence為例,就是對這句話找出每個字的權重,然後組合起來。

Attention的基本概念 (資料來源)

Transformer

為了避免扯太遠,實際上 multi-head self-attention的細節先放著不管 (或是參考文末 [補充1] )。總之, attention可以說是透過計算權重,找出關注位置的方法,而 Transformer則是基於 attention的精神,在 encoder-decoder的架構下建構了一套幾乎完全都是以 attemtion組成的網路

Transformer架構圖 (資料來源)

Positional encoding

雖然這邊只是快速地帶過 Transformer,但要特別補充一下 positional encoding的觀念。

因為 Attention其實沒有辦法 modeling前後順序

更精確地,以語言文字為例-每個字的絕對與相對位置。因此, Transformer相關方法都會用一些方法將 positional資訊放進 latent features中。比如說常見的將文字位置的索引 (index) 放到弦函數中,組合成合理且連續的位置編碼資訊。

利用 sin/cos function的不可學習positional encoding。 (資料來源)

Masked-Language Modeling

以文字處理為例,Transformer可以理解成不斷地綜觀所有文字的網路,然後找出值得關注的部分來做出最佳的預測。但這樣看起來好像還是需要具體任務,比如說中翻英,才能夠訓練 Transformer。但在使用標註資料訓練外,

在 NLP的發展中,能夠 self-supervised pre-training (SSL),或者更具體地說,如何運用未標註或其他種類的標註資料,一直都是很關鍵的事情。

因此 Transformer也有許多不需要標註的 pre-training方法。其中最知名的就是 BERT的 Masked-Language Modeling (MLM) [2]。MLM的做法相當簡單,但是效果卻是出奇地好。假設今天的輸入句子為 ”the movie is very boring”,MLM的訓練會自動將某些詞遮蓋掉,比如說 boring被 masked。然後就將這樣一個帶有遮蔽的文字丟進 Transformer,訓練目標就是可以還原出原本的句子。

Masked-Language Modeling的做法。(資料來源)

這個遮蔽其實就是以某組數值取代 (以下以 [MASK]代稱),在 BERT論文中使用了 15%的 mask ratio。因為在訓練實際的下游任務中,並不會出現 [MASK]這樣的詞。為了讓訓練時不要 overfit [MASK],在 BERT訓練時會有一定機率把這些數值替換為其他詞,而不是替換為 [MASK]。

Masked Autoencoders (MAE)

以 NLP的成功案例來看, masked-language modeling這樣的 masking method其實也是有機會套用到 vision task上,但是一直以來直接套用效果都不甚理想。這其實也不是很奇怪的事情,畢竟影像與文字本身的特徵是不一樣的

Contractive SSL methods

以往在 CV上的 SSL pre-training大多與 NLP上有明顯的差距,如同前面所提及,在 NLP的 Transformer上,僅僅是簡單的 masked-language modeling pre-training就可以達到很傑出的效果。但以往 CV上則是加上一些 tricks的 contractive方法 [3, 4, 5, 6]。

很簡單的contrastive learning概念,來自SimCLR (資料來源)

Masking Method for Vision

這或許很大一部分可以回歸到影像與文字之間的差異。影像基本上是來自感光元件直接對環境採樣、記錄下來的資料,每個 pixel本身其實是不帶任何 semantic資訊的。文字本身的每個詞,比如說 movie或 happy,就已經帶有 semantic資訊。

如果直接將 masking的目標放成 pixel,以重建為目標來訓練,網路可能仍然無法理解 sematic資訊。再加上其實圖片上很多相近的 pixel數值其實都可以使用內插或外插逼近,那這樣 masking就更沒意義了。

因此,如果要將 masking method放到 CV任務上,重點就會是:

  1. 如何讓這個重建任務難到無法以簡單插值獲得。
  2. 如何在 self-supervised抽取特徵時得到 semantic特徵,而非一定程度 redundant的 pixel-wise資訊。

Concept of MAE

Masked Autoencoders (MAE) [7]是在 2021年11月初次放到 arXiv上的方法,是由 FAIR的 Kaiming He主導的一個研究。 Kaiming He在近年 CV領域的 self-supervised learning (SSL) 研究可以說是不遺餘力,很明顯的這位優秀研究員現在的主要精神都放在 SSL上。

在MAE的最終提出的 masking作法卻是非常簡單:

Masking a high proportion of the input image, e.g., 75%, yields a nontrivial and meaningful self-supervisory task.

換句話說,整體訓練任務還是那個熟悉的 autoencoder,只是輸入會搭配大面積的 mask,然後要求網路重建出那些被遮蔽的資訊。

左到右依序為輸入、網路重建結果、真實圖片。(資料來源)

但這邊有一個小細節是,因為MAE的論文預設使用基於 Transformer的 ViT [8],因此在 masking時也是一塊一塊的遮蔽,而非 pixel-wise。

論文雖然沒有特別在這邊討論,但回想 DropBlock [9] 當年的論點,八成這樣區塊性的 masking更等價於 BERT中遮蔽 semantic的文字。不過論文也有提到,這個如果更進一步地拓展 block大小,反而會得到較差的結果。所以論文建議的方法就是以 ViT的一個 patch為基本單位。

Vision Transformer (ViT)

2020年 Google團隊提出的 Vision Transformer (ViT) [8] 可以說是第一個有效真正 end-to-end將 Transformer搬到 vision任務,並且完全不使用任何 CNN的方法。

ViT的實際作法很單純,把整張影像分成多個 2D patch,每個 patch的解析度為 P×P,攤平可以形成一個 sequence,丟給 linear projection layer得到 embedding,就可以直接丟給 Transformer了。具體來說,假設輸入影像的 shape為 (H, W, C),那麼攤平後就是 (HW/P², CP²) 的序列資料了。

Vision Transformer的架構圖。(資料來源)

這也是為什麼 ViT的標題是用 “An Image is Worth 16x16 Words”。

其實這 ViT有很多有趣的細節,而後來騰空殺出的 Swin Transformer [10]更是精采之作,不過就不再這邊提太多,以後有機會再多做討論。

Pipeline of MAE

因此 MAE可以說是一定程度站在 ViT的基礎上,把 masked-language modeling的簡單思路帶到了 CV領域。基本上的流程如下:

  1. 對於輸入image,直接用 ViT的切 patch方法切成多個 patches。
  2. 使用 random uniform將很高比例 (如75%) 的 patches mask掉。
  3. 把沒masked的 patches丟到 ViT based的 encoder抽 features。
  4. 把之前被masked掉的部分,以 masked token的形式補到 features裡。
  5. 丟給 ViT based的 decoder還原原圖,直接在 pixel-level 計算 MSE loss。
Masked Autoencoders的訓練流程。(資料來源)

整體上 MAE的精神與流程很簡單,需要稍微補充的大概就是 positional embedding與所謂的 masked token是什麼。不過這邊其實與 Transformer與 BERT中的設計如出一轍。

Positional embedding & Masked token

positional embedding的部分就使用 1D embedding的方法,直接加在 encoder與 decoder的輸入。更精確地說, encoder的輸入是 patch通過 linear layer轉換的 embedding,再加上這個 patch在 sequence 對應的 1D位置的 positional embedding。而 decoder也是做一樣的操作。

至於 masked token其實就是一個共用的、學來的向量,指出這個 feature原本應該是來自 masked的部分,然後再加上 positional embedding將原本位置的資訊考慮進來。具體地說, masked token就是一組可訓練的參數,而在輸入 decoder前的每個應該來自 masked部分的 feature,都直接帶入這組參數。

Masked Autoencoders的實驗數據

MAE論文的實驗大致分成 fine-tune與 linear probing兩種形式來探討:

  • fine-tune: backbone (encoder) 參數來自 pre-training,加上 predict head後一起訓練。
  • linear probing: backbone (encoder)參數來自 pre-training,在 encoder後接簡單的 linear layer作為 predict head,直接預測目標任務。訓練下游任務時凍結 backbone的參數,換句話說,只訓練 linear layer。

通常 linear probing可以想像成更接近學到的 representation品質。但其實也有其不合理之處,畢竟誰訓練網路最後不 fine-tune的。因此論文也有少量地討論到所謂的 partial fine-tune,即 backbone最後幾個 layer是不凍結的。

MAE中需要大的mask ratio

如同論文一直強調的, mask ratio要夠大才有效果。這樣的情況自然在 linear probing的 schema下更為明顯。對比 fine-tuning大約40% -75%效果都是相近的,linear probing則是在70–75%明顯較好。

MAE相當依賴高 masking ratio。(資料來源)

Comparison with SOTA

在 self-supervised方面,論文中將 MAE與三個 vision領域的方法做了比,基本上都是 Vision Transformer方面,沒有使用 Transformer的都不在上列 (如 SimCLR、 BYOL等等)。

MAE與其他SOTA方法的比較數據。(資料來源)

不過參考 MoCo v3 [11] 的數據,將上述三者視為目前的 SOTA是公正的。

MoCo v3中與其他 self-supervised方法的比較數據。(資料來源)

論文也有拿 MAE與 supervised pre-training做比較,結果顯而易見地,在使用同樣的資料下, MAE是比一般 supervised pre-training來的有優勢。但使用大量的資料 (JET300M) supervised pre-training還是遠勝使用小量資料 (1N1K) 的 MAE。

MAE與 supervised pre-training的比較數據。(資料來源)

另一方面,前述中有提到 linear probing可能不是最好地評估 self-supervised pre-training結果的 scenario。論文也有提供 partial fine-tuning的數據,下圖越靠左情境越接近 linear probing,越靠右則是接近 full fine-tuning。

MAE與 MoCo v3以 partial fine-tuning評估的數據。(資料來源)

Transfer Learning

論文最後當然在 transfer learning結果上秀了一番肌肉,再 object detection與 sematic segmentation都展示了最好的水平。

Transfer to object detection. (資料來源)
Transfer to semantic segmentation. (資料來源)

如同一兩年前的文章中所預測,未來幾年的研究會持續地鑽研「使用無標註資料」這類的方法。在CV領域的 self-supervised learning方面,前些日子都是依靠 contractive learning為主,近期其他考慮 masking方法的也都較為繁雜。而 Masked Autoencoders就是 Kaiming He在這個領域沉澱並鑽研多時帶來的乾淨的方法,如果 BERT的 pre-training一樣單純。

然而,這類方法也有一些小問題,就是對於 ViT的依賴。畢竟,在主流的平台運行的 vision model還是以 convolution為主,這不僅僅是慣例,也是因為 CNN發展多時,在硬體平台都已有許多加速的方法,這部分是 ViT在應用上很大的考驗。當然這些方法也有可見的改進之處,如直接套用 Swin Transformer取代 inductive bias較弱的 ViT。相信過不了多久,就可以看到相關的論文了。

雖然說把 MAE的想法放到 CNN上也不見得沒有效果,不過論文也沒有提供相關的數據,就要留給其他研究員或開發者自行嘗試了。

好了~這篇文章就先到這邊。老話一句,Deep Learning領域每年都會有大量高質量的論文產出,說真的要跟緊不是一件容易的事。所以我的觀點可能也會存在瑕疵,若有發現什麼錯誤或值得討論的地方,歡迎回覆文章或來信一起討論 :)

Supplementary

[補充1] Attention

Transformer的論文 [1] 中系統性地解釋 attention (論文中的 scaled dot-product attention)。概念上非常接近 hash table,使用三組 input (Q for query、 K for key、 V for value),分別通過不同的參數乘法後,使用 query尋找對應的 key,再利用 key取 value。

scaled dot-product attention (圖片來源)

直覺看起來, Q、 K、 V三者的輸入應該會來自不同地方。不過所謂的 self-attention就是指 Q、 K、 V三者的輸入都是一樣的,只是通過不同的轉換。而 multi-head指的則是 Q、 K、 V會許多 attention head,再組合再一起。

[補充2] 更多MAE實驗數據

Asymmetric Encoder-Decoder Design

MAE比較特別的設計之一是不對稱 (asymmetric) 的 encoder-decoder設計,因為輸入只有沒有被 masked的 patches,輸出則是所有 patches。乍看之下這樣的設計是運算量的考量,而非一定要這樣才有效果。不過論文卻指出了這是讓 MAE效果極好的關鍵之一。

論文認為是這些被 masked的資訊,在真實 inference的時候不會發生,因此造成 degraded的表現。雖然這點很難被證明,不過其實對比 BERT中也有類似的處理(前述:在訓練實際的下游任務中,並不會出現 [MASK]這樣的詞)

Linear Probing需要較深的 decoder

雖然在實際的 inference時通常是不需要 MAE的 decoder,但在論文中花了一些篇幅討論了 decoder的架構。結論大致是 decoder的深度會大幅影響是否抽取到 semantic的 representation,因為就 linear probing的結果來看,過淺的 decoder效果是明顯差的。

Supervised Target直接對Pixel就足夠好

MAE使用的 (self-) supervised target是與熟悉的 autoencoder可以說是沒什麼差別,但可以針對每個 patch額外計算 mean與 variance來做 normalization得到額外的提升。然而,卻是不需要使用更 semantic的 loss,如 DALL-E dVAE的 image token。

Data Augmentation非主導地位

與 SimCLR、 MoCo、 BYOL這些基於 contractive的方法不同, MAE並不需要使用 data augmentation生成對應的 sample來得到足夠好的 self-supervised agreement。因此,即使不使用 data augmentation也是可以訓練的,而且效果還不差。

Reference

[1] Attention Is All You Need [NIPS 2017]
[2] BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding [NAACL 2019]
[3] A Simple Framework for Contrastive Learning of Visual Representations [ICML 2020]
[4] Momentum Contrast for Unsupervised Visual Representation Learning [CVPR 2020]
[5] Bootstrap your own latent: A new approach to self-supervised Learning [NeurIPS 2020]
[6] Exploring Simple Siamese Representation Learning [arXiv 2020]
[7] Masked Autoencoders Are Scalable Vision Learners [arXiv 2021]
[8] An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale [NeurIPS 2021]
[9] DropBlock: A regularization method for convolutional networks [NIPS 2018]
[10] Swin Transformer: Hierarchical Vision Transformer using Shifted Windows [ICCV 2021]
[11] An Empirical Study of Training Self-Supervised Vision Transformers [ICCV 2021]

--

--

AI Blog TW
AI Blog TW

Published in AI Blog TW

AI Blog Taiwan專注於分享、探討資料科學、機器學習與人工智慧。核心理念是盡量結合當代學術知識與業界經驗,以避免從單一論文或實作視角產生的盲點;不造神,也不做內容農場。另外,鑒於繁體中文相關文章相對稀缺,內容使用繁體中文為主要語言進行撰寫。

Jia-Yau Shiau
Jia-Yau Shiau

Written by Jia-Yau Shiau

Broadening Impact of Data Science @ Taiwan

No responses yet