ViT — Vision Transformer : Convolution is dead, long live Transformers!

AppleHank
13 min readDec 18, 2020

--

在CV領域的Benchmark當中,大多仍是由CNN沿伸出的架構 — EfficientNet (2019)拿下了 State of the art (SOTA)的位置。然而一篇在今年十月提出的匿名論文「An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale」轟動了 CV 領域的ML社群,號稱依賴卷積的是沒有必要的,提出了不同於CNN的架構,也就是將NLP領域的 Transformer(2017)應用至 CV 界,並且在所有的Benchmark當中超越了EfficientNet以及基於ResNet來訓練的BiT。(備註:但在下筆當天(2020/12/11),EfficientNet 又超前了0.1%左右。並且揭發作者為Google)

Vision Tansformer(下稱ViT)的核心作法是將一張圖片,例如256(Width)x 256(Height) x 3(Channel)依照指定的 Patch Size 切成數個 Patch。例如標題的 Patch Size = 16,就會將圖片切為 n = 16(張),每張 Patch 變成16 x 16 x 3,再將這些 Patch 經過處理後直接丟入 Transformer。以 NLP 的概念來說,每一個Patch就是一個字,這16張圖片就是一個 sequence,會一起丟進Transformer。

一、Related Word

ViT 並不是第一個嘗試將 Transformer 融入 CV 領域的論文。在 Transformer 震撼 NLP 領域後,許多論文也嘗試將 Transformer 應用到 CV 領域,但這件事情並不簡單。

這個任務的阻礙是什麼呢?我們可以從Transformer based model中最知名的 BERT 的角度說起。在 BERT 中,我們可以 input 512個字,也就是512個的「embedding」。如果我們要將圖片放到BERT會怎麼做呢?最簡單的方式就是將每一個pixel當作一個embedding直接放進去,但我們是放不進去的。以MNIST的圖片來說,一張圖總共有28 x 28 = 784 個pixel,多出了兩百多個pixel;如果是ImageNet,則是256 x 256 = 65536 個pixel,這是不可能丟入的。如果我們將input維度擴大則會耗費龐大的訓練時間,試想input是512維的BERT的訓練時間都要那麼久了,如果是 128倍大的維度,論文可能要留給兒子發了(可能還會被reject),所以將Transformer融入CV的主要障礙在此,我們沒有辦法直接將原始圖片放入最原始的Transformer。

所以將Transformer融入CV領域的研究致力於降低運算成本,有兩種方法:

  1. 改變input:結合Convolution,將原始圖片轉為低維度的表達方式。但會丟失原始圖片上的部份資訊。例如FB在2020年發表的DERT,使用Convolution抽取特徵後將feature map做為input。
  2. 改變attention機制:修改Transformer的attention機制,降低運算成本。保留原始圖片的資訊,但特殊的attention機制無法在表現上或時間上勝過原始Transformer的global attention。例如Axial-DeepLab在2020年發表的 Axial Attention,將 global attention 改為注重row與column的local attention。

論文中將這方面的研究分為結合Convolution架構以及捨棄Convolution架構,但我認為分為以上兩種比較能理解研究方向。ViT採用改變了input,也就是將圖片切割為 Patch,並捨棄了Convolution。

二、Architecture

Architecture of ViT

整體架構如上,左側是整體架構,右邊只是其中的 Transformer Encoder 拉出來而已,接下來分為幾個步驟來看。

  1. Split image

假設左邊的原始圖片大小為48 x 48 x 3,Patch Size 就等於16,所以會切成9張 Patch,每一張 Patch 大小為16 x 16 x 3

目前圖片的狀況如上圖,為了方便理解所以先將Patch表示為矩陣。有9個16 x 16的矩陣,每個矩陣內的元素是一個 Pixel,每一個Pixel以(R,G,B)來表示,例如第一個元素代表R = 20, G = 1 , B = 197。第一張Patch稱為x¹ₚ。

2.Linear Projection

接著要將每個Patch轉變為D維的向量來放進 Transformer,因此將 x¹ₚ 的表達先轉換為向量,通過一個透過訓練得到的 Linear Projection,稱為EE是一個(P x P x C) x D的矩陣。因此[1張 x (16 x 16 x 3) x ((16 x 16 x 3) x D)],變為1 x D的向量,這個向量在論文中稱為 Patch Embedding

我們將9張 Patch(x¹ₚ, x², …, x)都通過E,得到9個 Patch Embedding。

3.Add position embedding

就像 Transformer 要加上 Position Embedding 給予位置資訊一樣,這邊也要將 Patch Embedding 加上一個透過訓練得到的1維 Position Embedding。圖中的紫框就是 Position Embedding,粉框則是 Patch Embedding。

特別注意,在最左邊有一個*的 Patch Embedding 是透過訓練得到的[CLS] Embedding,稱為 x꜀ₗₐₛₛ,代表的是原始圖片的整張資訊。就跟 Bert 的[CLS] token 一樣,當他和其他Embedding一起被丟進 encoder 後,他對應的 Output 就代表了整個句子的表示,並且可以使用這個Output來分析任務目標。

這個步驟結束後我們就得到了 Transformer 的 input:z₀

the input of transformer

4.Transformer

一個Transformer由多個 Multiheaded Self-Attention(MSA)與 MultiLayer Perceptron(MLP)區塊組成,一個區塊的運作如上圖,做法如下。將 z₀ 丟入Transformer,首先會先經過一層 Layer Normalization(LN),接著進入MSA,將這層的結果加上 input : z₀(Add),也就是使用 residual connection,MSA 的部分就完成了,得到z'ℓ。接下來再通過一層 LN 以及MLP,最後實施 residual connection,MLP 的部分就完成了,得到第ℓ層的輸出zℓ。論文中的公式如下右圖,最底部的公式中的z⁰ℓ代表 [CLS] embedding 經過Transformer後的輸出。

ℓ=1時,如左圖

5.Output

最後這一小部分在論文中似乎沒有講到,但我認為就跟大部分 Transformer based model一樣,將[CLS] token 丟入MLP,再接 softmax 產生出每個class的機率,去 predict 出結果。

三、Experiments

Google 提出了幾個不同大小,以及在不同資料集 pretrain 的 ViT 來實驗,如下圖。Table2 中最左邊的 Outs-JFT(ViT-H/14) 代表的是 Table1 中的 ViT-Huge/Patch Size = 14、在 JFT 資料集 pretrain。而旁邊的 ViT-L/16 的訓練集 I21K 代表 ImageNet-21K。JFT擁有18K個 class,303M 個高解析度的圖片;I21K 擁有 21K 個 class,14M 個圖片。

當 ViT 在較小的資料集上(I21K) pretrain 時,表現不比 EfficientNet 來的好。不過一旦在大資料集(JFT)訓練時,就可以超越現今的 SOTA ,實現取代 Convolution based model的野心(但在12/11查看 ImageNet的 SOTA 時,EfficientNet 硬擠出了高於 ViT 0.1%左右的 Accuracy)。看完了 Accuracy 後,來關心時間方面。

SOTA(2020/12/18),左:ImageNet ReaL 右:ImageNet

配上 EfficientNet 資料來看,ViT 的參數多出了 150M 左右,但運算時間卻快了將近五倍,歸功於 Self-Attention 的平行運算。

論文中還有一些解釋隨著 pretrain 資料量提升,不同Model之間的表現,有興趣的話可以再去看看。

四、Why Transformer works?

大家可能都會覺得Transformer可以拿來做 CV,但是,為什麼它能夠拿來做CV 呢?將圖片切割成Patch後,不會將特徵也切開,影響判斷嗎?

論文給出的解釋是,他們認為Self-attention讓 ViT 能夠整合不同 Patch 之間的資訊,為了分析這件事情,他們基於 Transformer 過程中產生的 Attention Weight 計算了每一個 Head 對於其他 Head 的「Attention Distance」,這個 Attention Distance 類似於CNN中的 receptive field size。receptive field size簡單來說就是每一張 feature map 在原始圖片上能夠被啟動的區域,並且給出了下圖。

Mean attention distance between 16 heads

如果將Attention Distance理解成receptive field size,可以看到有些Head在訓練一開始就能Attend到原始圖片中大多的區塊,而且隨著Layer越深,越多的Head 就能Attend到更多的區域,證明了Transformer所產生的資訊確實是有考慮到其他許多的Patch,是Global的。

看到這邊可能會想到,既然這樣,我把Patch size切得更小,照理來說Patch之間就能考慮到更細節的資訊吧?但需要注意的是,隨著Patch size越來越小,input的shape就會越來越大,就會花費更多的pretrain時間。在論文中提到的Patch size有32、16、14,其中ViT-H只有14的資料,而ViT-B與L有32與16兩種size。以有兩種size的ViT-B和ViT-L來看,Patch size = 16的表現都勝過Patch size = 32。所以也許Patch size越小的確是可以有更好的performance的。至於如果切割成pixel level的話結果會如何,我認為有可能會類似於將Convolution的kernel size設置為1 x 1而難以捕捉特徵,但畢竟Self-Attention的機制能夠讓每個Pixel相互Attend,所以可能也得要實驗後才能知道。

五、觀點分享

這篇Paper標榜的是完全取代Convolution,而Convolution的價值在於它能夠透過共享變數來大幅減少運算時間,同時擷取到圖片上的特徵,但我認為這篇Paper使用Self-Attention的方法也是類似的原理。使用Self-Attention共享參數以及透過切割Patch減少運算時間,同時也能夠擷取到不同Patch之間的資訊,老實說有點類似Convolution的概念,只是做法不同而已,而且在提出不久後,EfficientNet就再次得到了SOTA的位置。目前來看,我認為ViT能夠將Patch size設置更低,架構更大而再次贏過EfficientNet,同時有更短的訓練時間。

另外有人認為這篇論文帶來了一個證明。在理論上,我們能夠透過單純的MLP學習到所有的function。我們之所以產生了CNN、LSTM等等架構是因為我們需要「bias」MLP,這裡的bias指的是Inductive bias中的bias。當我們從一個現象中觀察到規律,將這個現象用觀察到的規律來表達的過程就稱為Inductive,例如在ML中可以理解成我們找到一組function來表達一個input和output的關係;但由於我們觀察到的現象有限,因此我們需要對input加上一些限制,這個限制就是bias。他認為Transformer在inductive bias上並沒有給出太多的限制,因此他可以fit許多的task。而CNN與LSTM給出限制是因為當時我們的資料量與運算能力還不足夠拿掉這些限制,但如今我們有能力慢慢捨去這些舊有的限制了,相關的想法我將影片附在下方的Reference。

六、結語

其實ViT的作法很類似於將原始圖片丟進Transformer,這種作法應該也有人想過,但如果將Patch切分為接近Pixel會提高運算成本,切成太大塊又會降低準確率。所以仔細想想這篇論文好像只厲害在他們有成本可以將Patch切的小塊又可以負擔運算成本而已XD。

不過論文中還有不少有趣的實驗,例如使用2D的Position Embedding的效果、運算時間上的分析、隨著層數加深帶來的影響…,有興趣的話可以去看看原文論文。而且這篇論文仍在修改當中,也許等到正式發表時又會修改許多內容。

Reference

--

--