Review: An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale

Guan
工人智慧
Published in
13 min readAug 29, 2022

--

Introduction

本篇是 Vision Transformer 的濫觴,因著 Transformer 在 NLP 中的卓越表現,本文嘗試直接將 Transformer 應用於圖片;在此之前 Convolution 幾乎主宰電腦視覺領域,部分試圖整合 convolution 和 attention 機制的架構,部分試圖完全僅使用 attention,但因特化的 attention 運算的硬體支援不足外(註),效果皆也遠不及傳統大型 Resnet。而本文提供一種使用大型資料集搭配 Transformer 架構,取代 CNN 架構的可能性。

Method

Model Overview

要讀懂如何將圖片輸入 Transformer 的關鍵是理解每一個部分的 tensor shape,以下我們搭配 lucidrains/vit-pytorch 將每個步驟的 tensor shape 列出,並於此基礎上分析其物理意義。

Input and Output of Attention Layers

首先,假定輸入圖片為 (224, 224, 3) ,我們將圖片打散為 N 張 patches,設 N = 14 * 14,所以我們會得到 196 張 patches (Fig.1 範例為 9 張),每張維度 (1, 16, 16, 3);即 x1_p, …, x196_p,P =224 / 14 =16。

故設 N=196, P = 16 得出 (1) 中各項的維度
x1_p, …, x196_p : (196, 768)
E : (768, 1024)
x_class : (1, 1024)
E_pos : (197, 1024)

得出 z_0 之維度 (197, 1024)。

而 (2) (3) 之中的 Multi-head Self Attention MSA 和 Multilayer perceptron MLP 之輸入輸出維度相同,故 (4) 取 197 序列的第一個z0_L 長度為 1024,故得到 y 之長度 1024。

Inductive Bias

中文翻譯作 “歸納偏置”,白話文來說,為在做出決策前,由於演算法 缺乏特定資訊做出決策,必須要作出假設,而每種演算法所傾向的假設不同,譬如 Convolution 的特性,導致 CNN 有著 局部性 Locality 和 平移不變性 Translation Equivariance 這兩種假設傾向,而剛好這種假設非常符合圖片資料,故原文中有這樣的描述:

We note that Vision Transformer has much less image-specific inductive bias than CNNs. … In ViT, only MLP layers are local and translationally equivariant, while the self-attention layers are global.

描述了 ViT 先天在處理圖片資料時,缺乏正確的假設傾向導致效能的下降。

Hybrid Architecture

作者另外提出了一種混合架構,不單單是把圖片打散為 patch 後輸入,而是使用 CNN 提取出的 feature maps;但其後的實驗能看到雖然在相對小模型時有優勢,但可以透過加大 ViT 模型得到相似與 Hybrid 精度。

Fine-tune And Higher Resolution

在 Fine-tune 時,當以下兩點改變時,要改變原先訓練的模型架構:

  1. Number of classes changed
    當類別數量改變時,將 (3) 中,第 L 層原先訓練的 MLP 移除,並接上另一 fully connected layer,並以 0 初始化,維度為 (1024, new_num_classes) 即可。
  2. Resolution of image changed
    本文使用 (224, 224, 3) 圖片作為範例,當使用在圖片資料集 (384, 384, 3) 之上時,我們將固定 P = 16,因為 Transformer 可以處理變化的 N,但 E_pos 仍會受到不同 N 的影響,故這裡我們選擇對 E_pos 做 “2D interpolation” (我選擇理解為 Bilinear interpolation)。

故當圖片為 (384, 384, 3)
N = 24, P = 16 得到

x1_p, …, x196_p : (576, 768)
E : (768, 1024)
x_class : (1, 1024)
E_pos : (577, 1024)

於是我們將原先的 E_pos (197, 1024) 內插為 (577, 1024) 即可。

Experiments

此處使用了三種大型資料集作為 pre-train dataset,包括 ImageNet 1K , ImageNet 21K 和 JFT,再使用如 ImageNet, ImageNet ReaL, CIFAR 等較小的資料進行 Fine-tune 實驗。

Comparisons of models pre-train on JFT and ImageNet21k

上表可以看出 ViT-H 在 accuracy 和 training efficiency 都遠超現有模型,尤其是在與也使用了 JFT 的 Noisy Student 做比較時,其可說是 CNN 模型在分類任務上的標竿。

另外也提供了 VTAB benchmark ,VTAB (Visual Task Adaptation Benchmark) 是用於評估模型泛化在不同任務的能力,或我認為是能夠量化 model capacity;VTAB 提供一份訓練集使模型輸出 embedding representations,在使用訓練好的模型在 19 個不同的下游任務進行評估。

VTAB results

Pre-training Data Requirements

本節談論了 ViT 的弱點,也就是在沒有 Image Inductive Bias 的狀況下,需要多少資料集才能學習到與 CNN 相似的能力;眾所周知 CNN 透過增加資料集的大小對於表現的提升是有限的,所以首先才有 Residuals Block 打造更深但有效的模型,才有能夠有效率的使用 JFT 的 Noisy Student 出現。

而 ViT 在大資料集有更好的效率,本文設計兩種實驗討論預訓練資料集大小對於 ViT 表現的影響。

  1. 分別使用 ImageNet, ImageNet21k, JFT 作為預訓練資料集,並使用 weight decay, dropout 和 label smoothing:

可以看到隨著資料集的增大,ViT 的整體表性更好,但另外也有幾點觀察:

a. Large patch size 通常表現更差,可以看到 ViT-B/32 和 ViT-L/32 的表現都遜於它們 16 的版本。

b. 越小的 ViT 受到資料集增加的效益遞減非常大,雖然符合直覺,因較小的模型 model capacity 自然受限,但圖中從 ImageNet-21k 到 JFT-300M ,在資料集擴大 1000 多倍的狀況下,ViT-B 的表現幾乎沒有進步,甚至 ViT-B/32 有退化的跡象,有可能是受 regularization 影響。

c. CNN 也如同 ViT 一樣,在大資料集的狀況下,小型 CNN 一樣有退化的狀況。

2. 不使用任何 regularization ,將 JFT 分為四個大小 10M, 30M, 100M 及完整的 300M:

觀察與上圖相似,文中寫道

This result reinforces the intuition that the convolutional inductive bias is useful for smaller dataasets, but for larger ones, learning the relevant patterns directly from data is sufficient, even beneficial.

Scaling Study

本實驗皆使用 JFT 作為預訓練資料,在不同的資料集上 Fine-tune 固定 epoch (7 or 14),討論在受限的運算資源下,模型訓練的效率。

左邊是除 ImageNet 外五種資料集的平均準確度,右邊是 ImageNet,更詳細的數據請見 Table 6。本次實驗有兩個觀察:

  1. ViT 在效率上全面勝過 ResNet BiT,幾乎是 2 – 4 倍有效率
  2. Hybrids 在小資料時可能勝過 ViT,但大資料則否,符合前述實驗及直覺(CNN 在大資料集之限制)

Inspecting Vision Transformer

Center
我們前面以 (224, 224, 3) 為例,E_pos 是 (197, 1024) ,除去x_class每一個 patch 都有一個長度 1024 的 E_n_pos ,所以每個 patch 之間的 positional embedding 可以求出 cosine similarity。

再來看中間圖片,先看第四行第四列的 patch ,也就是 (4, 4),在上圖(4, 4)的小圖中,可以看到中間是最黃的,也就是最相關的,因為它自身與它自身的 positional embedding 當然是一模一樣的。

可以看出每個 patch 的 positional embedding 大致上會與周邊的,以及有垂直水平關係的 patch 的 positional embedding 有極大相關。

Left
左圖的意義是將前面所提的 (1) 式之 E 視覺化,可以看到這個 projection filter 產生了各種不同模式提取 patch 資訊。

Right
右圖所謂的 “attention distance” 很抽象,文中也沒有寫出計算式,但可以從文字中理解它的意義,首先

This “attention distance” is analogous to receptive field size in CNNs.

所以越大的 attention distance 意味著該 head 的 receptive field 越大,更能夠眾覽全圖;再來

This highly localized attention is less pronounced in hybrid models that apply a ResNet before the Transformer, suggesting that it may serve a similar function as early convolutional layers in CNNs.

當我們觀察 hybrid 模型的時候,由於使用了 Resnet,它提供了 CNN 的 inductive bias,所以更少的 head 去負責小的 recpetive field。如下圖右,幾乎沒有 head 是小於 40 的 mean attention distance 。

最後猜測一下 attention distance 是如何被計算的 :

Attention distance was computed for 128 example images by averaging the distance between the query pixel and all other pixels, weighted by the attention weight.

(Deprecated)依照這段文字,attention distance 與 MSA 中的 query 和權重有關,符合描述的應是 QW^QKW^K 相乘後得到一 (seq_len, seq_len) 的向量,這個向量的意義可理解為序列元素之間的相關性權重,原文中稱其為 attention weights,然後對其取平均,這是第一個可能。

(Deprecated)往下計算,得到 Attention(Q,K,V)即某個 head 的 self attention 矩陣 ( seq_len, num_feature) 計算出 (seq_len, seq_len) 的 cosine similarity 矩陣,取平均,這是第二種可能。

後來找到了關於 attention distance 的實作 (credit to sayakpaul) 如下

Attention Distance
Code to compute Attention Distance

得到 Attention_i(Q,K,V) (num_patches, num_patches) 之後,乘以 distance_matrix (num_patches, num_patches),自 0 到同樣行列的 num_patches 逐漸加大,譬如(0, 16, 32, 48…, 213.4666),在對其平均得到一 scalar 則為該 head 之 attention distance。

直覺來看,越遙遠的地方有越大的 feature value (of Attention_i(Q,K,V)) 就會被認為是有更大 receptive field 的 head。

Conclusion

本篇為 Vision Transformer 的濫觴,不引入任何 CNN 之 Image Inductive Bias ,且以序列方式理解圖片,也能夠有很好的效果。

然而,本文使用了未開源 JFT-300M,認為 ViT 必須要在大型的資料集上訓練才能夠有與 CNN 匹敵的表現,但之後我們知道了,這只是訓練參數問題, DeiT 和 Swin Transformer 都能夠僅使用 ImageNet 1K 訓練就有極佳表現;接下來我們會就這兩篇續作探討。

--

--