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

Abstract

在 CV 領域 transformer 表現有限,目前 attention 常常是和卷積神經網路一起用,或是用來把一些卷積層換成 self-attention,但整體架構不變。這篇論文想展現一個純 Transformer 可以直接在影像分類上表現很好。如果用大量資料作預訓練,再遷移到中小型的資料集,可以和 SOTA 的 CNN 表現得一樣好,還需要較少的訓練資源作訓練。

Introduction

self-attention-based 架構,特別是 Transformer,已經是 NLP 的重要選擇。主流的作法是在大型文字資料集上作訓練,再針對小型任務資料集作 fine-tune。由於 Transformer 的計算效率高,還有可擴展性,可以 train 一些很大的 model,隨著 model 和資料集增大,目前還沒看出飽和的現象。

然而在 CV,CNN 還是主流,一些工作嘗試用 self-attention 結合 CNN-like 的架構,比如把 feature map 當 transformer 的輸入,因為原始 pixel 太多,或甚至把卷積層全換成 self-attention,雖然後者理論上效率很高(原論文中有另外 cite 兩篇作法),但因為他們做法特殊,在現代硬體上很難加速,所以無法很有效地擴展。在 large-scale 的影像識別上, ResNet-like 的架構還是 SOTA。

該實驗直接把一個標準的 Transformer 作用於圖片上,只作最少的修改。把影像分成多個 patch,並把它們變成一系列的 linear embedding,當作 NLP 中的 tokens(words) 來處理。

當在中型大小的資料集(e.g. ImageNet)上訓練,如果沒有 strong regularization,ViT 會略輸同等大小的 ResNets

這篇論文在更大的資料集(14M-300M 的影像)上訓練,就打敗了 inductive bias。在大量資料上作預訓練就很讚。

Related Work

大型的 Transformer-based 模型常常是先在大資料集上預訓練然後根據任務 fine-tune,比如 BERT 和 GPT。

要把 self-attention 用在 CV 上,最簡單的做法就是把每個 Pixel 當一個元素,但 self-attention 是平方複雜度,在現實的圖片很難應用。一個應用 Transformer 的做法是只把 self-attention 用在 local neighborhood,另外一個是用 Sparse Transformer,還有一堆特殊的方法,雖然表現不錯,但要用硬體加速起來不容易。

另一個有關的模型是 iGPT,在 reduce image resolution 和 color space 後把 transformer 應用在 image pixels 上。它用非監督式訓練後,再 fine-tune 或做 linear probing(只更新最後的 linear layer) 分類任務,表現很好。

已經有類似的工作了,抽取 patches of size 2 * 2,最後再接 full self-attention,基本上和 ViT 非常像,這篇論文進一步證明了作大規模的預訓練可以讓 Transformer 和 SOTA 的 CNN 相比,而且 ViT 因為 patch 比較大,可以處理 medium-resolution 的圖片。這問題是可預期的,因為 Transformer 缺少了一些 inductive biases。

  • inductive biases
    • 一些假設
    • 比如 CNN 常有四個假設
      • locality
      • translation invariance with pooling layers
        • 平移不變性
      • translation equivariance
        • f(g(x)) = g(f(x))
        • 卷積和平移的先後順序沒差

Method

模型盡可能類似原始 Transformer,這樣可以把一些 NLP 上成功的 Transformer 架構拿來用,還可以用一些很有效率的 implementation

embedding 維度是 768 = 16 * 16 * 3 position embedding 的做法是 standard learnable 1D positional embeddings,就是 BERT 的做法,簡單來說就是生出一張可以訓練的表,(序列長度, embedding size),作者也有嘗試其他方法,但發現成效差不多,比如 2D positional embedding,概念就是從生出(序列長度, embedding size)變成生出 2 個(sqrt(序列長度), embedding size)。

[class] 的概念是 NLP 出來的,ResNet-like 的架構常見的做法也有通過 globally average-pooling (GAP)來生出向量,再接上分類器做預測。實驗發現直接在 transformer 的輸出做 GAP 和 [class] 都可以達到不錯的效果。

Conclusion

拿標準的 Transformer 來作 Image recognition,和以往用 self-attention 在 CV 的方法不一樣,除了一開始的 initial patch extraction,沒有引入其他影像特有的 inductive biases。直接把圖片當成是一系列的 patch,然後直接用 Transformer encoder 當一般 NLP 任務處理。在很多影像分類訓練集上表現得更好還在 pre-train 上相對便宜。

還有一些值得挑戰的地方,比如把 ViT 應用在其他 CV 任務,比如 detection 和 segmentation。另一個挑戰是探索自監督預訓練的方法。這篇論文其實有實驗自監督,表現 OK,但和監督式還是有很大的落差。擴大 ViT 可能有更好的結果。