Post

[논문 리뷰] BEiT: BERT Pre-Training of Image Transformers

Posted
By Luna Lee
12 min read



Microsoft Research ICML 2022

Introduction

Computer Vision에서 Transformer가 뛰어난 성능을 달성했지만, 일반적으로 CNN에 비해 많은 학습 데이터를 필요로 한다.

BERT는 NLP 분야에서 큰 성공을 거두었다. 텍스트 내의 일부 토큰을 무작위로 마스크한 다음, 이렇게 손상된 텍스트의 Transformer 인코딩 결과를 기반으로 마스크된 토큰을 복구하는 MIM(masked language modeling) 방법을 사용했다. 본 논문에서는 BERT에서 영향을 받아 Vision Transformer를 사전 학습하기 위한 denoising auto-encoding idea를 제안한다.

BEIT_1.png 이미지 출처:https://wikidocs.net/115055

Language에는 단어라는 사전적으로 정의되는 명확한 token이 존재하는데 비해 이미지 데이터에는 patch에 대한 사전적인 정의가 존재하지 않으므로, BERT의 방식을 그대로 적용하는 것은 어렵다. 이에 대한 대안책은 마스크된 patch의 raw 픽셀을 예측하는 regression problem으로 문제를 수정하는 것이다. 하지만 이러한 픽셀 수준의 recovery task는 modeling capability를 short-range dependency와 high-frequency details에만 낭비하는 문제가 있다(row-frequency structure와 같은 부분이 일반적으로 더 중요함).

본 논문에서는 self-supervised vision representation model, BEiT(Bidirectional Encoder representation from Image Transformers)를 제안한다. BERT에서 영감을 받아 pre-training task, 즉 MIM(Masked Image Modeling)을 제안한다. Self-supervised learning을 수행한 다음 두 가지 down-stream task(classification, segmentation)에 대해 pretrain된 BEIT에 대해 fine-tuning했다고 한다.





Method

1. Image Representation

저자는 Image representation을 2가지 방식으로 표현했다. Pre-training에서 input으로 사용되는 Image patch, output representation인 Visual token이다.

Image Patch

BEIT_2.png

2D 이미지는 Transformer의 입력으로 사용할 수 있도록 patch로 분할된다(ViT와 동일한 방법을 사용했다). 수식적으로는 위의 그림과 같이 재구성된다고 볼 수 있다. 분할된 각각의 이미지 patch는 vector로 flatten되고, linear하게 project된다(마찬가지로 ViT와 동일).

논문에서는 224 × 224 이미지를 각각의 patch size(resolution, P)가 16 × 16인 14 × 14 patch grid로 분할했다고 한다.


Visual Token

BEIT_3.png

NLP와 동일하게 이미지를 raw 픽셀 대신 Image Tokenizer에서 생성한 discrete tokens으로 표현한다. Tokenizer는 DALL-E에서 사용한 dVAE(Discrete Variational Autoencoder)를 사용했다.

dVAE는 Tokenizer(Encoder)와 Decoder로 구성되어 있고, tokenizer $q_\phi(z|x)$는 이미지 픽셀 $x$를 visual codebook에 따라 discrete tokens $z$ 로 매핑하는 역할을 한다. Decoder $p_\psi(x|z)$는 visual token $z$를 이용하여 input 이미지 $x$를 재구성하는 역할을 한다.(DALL-E paper에서는 $p_\theta$로 decoder에 대한 notation을 조금 달리했다. 헷갈리지 않기..)

자세한 내용은 DALL-E 리뷰 포스트를 참고해보자.


각각의 이미지는 마찬가지로 14 × 14 grid의 visual token으로 tokenize되었다(앞의 이미지 token 개수와 동일하다). Codebook(vocabulary) size는 $|\mathcal{V}| = 8192$로 설정했다.




2. Backbone Network: Image Transformer

Backbone Network는 ViT와 동일하게 standard Transformer를 사용했다.

Transformer 입력으로 사용되는 sequence of image patches(위에서 언급했던,) ${{x^p_i}}^N_{i=1}$는 linearly project되고 patch embedding $Ex^p_i$가 된다$(E\inℝ^{(P^2C)\times D})$. 입력 시퀀스 앞에 special token [S]가 추가되며, 학습 가능한 1D position embeddings $E_{pos}\inℝ^{N \times D}$가 더해진다(ViT랑 동일하다).

  • Input Vector: $H_0 = [e_{[S]}, Ex^p_i, …, Ex^p_N] + E_{pos}$
  • Output Vector: $H_L = [h^L_{[S]}, h^L_1, …, h^L_N]$




3. Pre-Training BEIT: Masked Image Modeling

본 논문에서는 masked image modeling (MIM) task를 제안했다. 이미지 patch 중 일부를 랜덤하게 마스킹하고, 마스킹된 patch에 해당하는 위치의 visual token을 예측한다.

BEIT_4.png

전체적인 방법에 대한 그림은 위와 같다.

  1. 입력 이미지 $x$가 주어지면 N개의 이미지 patch $({x^p_i}^N_{i=1})$, N개의 visual token $({z_i}^N_{i=1})$이 생성된다.
  2. 이미지 patch의 약 40%를 랜덤하게 마스킹한다. (마스킹된 포지션: $\mathcal{M} \in {1, …, N}^{0.4N}$)
  3. 마스킹 된 patch를 learnable embedding $e_{[M]} \in ℝ^D$로 대체한다.
  4. 손상된 이미지 patch들$(x^{\mathcal{M}}={x^p_i: i \in \mathcal{M}}^N_{i=1} \bigcup \, {e_{[M]}: i \in \mathcal{M}}^N_{i=1})$ 은 L-layer Transformer(2에서 언급했던)에 입력으로 들어간다.
  5. 마지막 hidden vectors ${h^L_i}^N_{i=1}$은 입력 patch의 encoded representation으로 간주된다. 이 중 마스크된 각 위치 ${h^L_i: i \in \mathcal{M}}^N_{i=1}$ 에 대해 softmax classifier를 사용하여 각 위치에 대응하는 visual token $p_{MIM}(z’∣x^{\mathcal{M}}) = softmax_{z’}(W_ch^L_i + b)$을 예측한다.
    ($x^{\mathcal{M}}$는 손상된 이미지, $W_c \in ℝ^{∣\mathcal{V}\times D∣}$, $b_c \in ℝ^{∣\mathcal{V}∣}$ )

Pre-training objective는 손상된 이미지가 주어졌을 때 visual token $z_i$에 대한 maximize log-likelihood이다.

BEIT_5.png


마스크 위치를 패치에서 랜덤하게 뽑기보다는 blockwise masking 방법을 적용한다. 이미지 patch를 블록단위로 뭉친 다음 단계별로 블록에서 최소단위(여기서는 16)만큼은 마스킹이 되도록 한다(각 블록 내에서 랜덤). Algorithm 1과 같이 마스크 선택 과정을 반복하여 마스크 위치를 설정한다.

BEIT_6.png




4. From the Perspective of Variational Autoencoder

BEIT pre-training은 VAE 학습으로 볼 수 있다. $x$는 original image, $\tilde{x}$는 masked image, $z$는 visual token일 때, log-likelihood $p(x∣\tilde{x})$에 대한 ELBO식은 아래와 같다(= 손상된 이미지에서 원본 이미지 복구).

BEIT_7.png

학습은 2단계에 걸쳐 진행된다. 첫 번째 단계에서 dVAE를 통해 image tokenize를 학습하고(minimize reconstruction loss), 두 번째 단계에서 $q_\phi$ 와 $p_\psi$를 fix하고 prior $p_\theta$를 학습한다. 여기서 visual token에 대해 가장 가능성이 높은 token$(\hat{z_i} = \arg \max_z q_\phi(z∣x_i))$을 사용하여, $q_\phi(z∣x_i)$를 one-point distribution로 단순화할 수 있다. 그러면 위의 식은 아래와 같이 정리할 수 있다.

BEIT_8.png




5. Pre-Training Setup

  • ViT-Base 구조를 따름. 12-layer Transformer, 12 attention heads
  • 1.2M개의 이미지를 포함한 ImageNet-1K 데이터셋으로 pre-training
  • Augmentation: random resized cropping, horizontal flipping, color jittering
  • Transformer를 대규모 데이터셋에 대해 안정적으로 학습하기 위해서는 적절한 initialization이 중요하다고 함




6. Fine-Tuning BEIT on Downstream Vision Tasks

BEIT를 pre-train한 후 transformer에 task layer를 추가하고 down-stream task에 대해 fine-tuning을 진행했다.


  • Image classification

Task layer로 간단한 linear classifier를 추가했다. Average pooling → softmax classifier를 사용했다.


  • Semantic segmentation

Task layer로 SETR-PUP 구조를 사용헀다고 한다. 해당 논문(SETR)을 살펴보니 transformer output에 Decoder를 붙여주는 형태인데, 그 중에서도 PUP(Progressive UPsampling)구조는 4개의 Upsampling layer를 거쳐 점진적으로 upsampling하는 형태라고 한다.


  • Intermediate fine-tuining

Self-supervised pre-training이 끝난 뒤에, 조금 더 큰 데이터셋에 대해 추가로 학습을 진행한 다음 다시 downstream task에 대해 fine-tuning 할 수 있다. (이러한 중간 fine-tuning은 BERT의 일반적인 방법이라고 함)





Experiments

1. Image Classification

Image classification 성능 평가를 위해 1,000개의 클래스와 1.3M개의 이미지를 포함 ILSVRC-2012 ImageNet 데이터 세트를 사용했다.

BEIT_9.png

위의 표는 classification에 대한 top-1 accuracy를 비교한 표이다. 결과는 위와 같다.

추가적으로 224 × 224 이미지로 fine-tunig 한 후, 384 × 384 resolution 이미지에 대해 10 Epoch 더 fine-tuning을 진행했다. 이미 패치 크기는 resolution에 상관없이 동일하게 유지했으므로 resolution이 높아질수록 input sequence lengt가 길어진다. 위의 결과는 해상도가 높을수록 ImageNet에서 BEIT 결과가 1점 이상 향상됨을 보여주고 있다.

또한 모델 사이즈를 키워 실험을 진행했다. BEIT를 기본에서 대규모로 확장하면 ImageNet-22K를 사용한 supervised pre-training보다 더 큰 이점을 얻을 수 있었다.


BEIT_10.png

위의 수렴 곡선은 training-from-scratch 방법과 pre-training-then-fine-tuning 방법의 수렴 곡선을 비교한 그래프이다. BEIT를 fine-tuning하면 더 나은 성능을 얻을 수 있다고 한다.




2. Semantic Segmentation

Segmentation task를 위해 25K 이미지와 150개의 의미 카테고리를 갖춘 ADE20K benchmark에 대해 평가를 진행했다. 모든 semantic category에 대해 mean Intersection of Union (mIoU)를 측정했다.

BEIT_11.png




3. Ablation Studies

BEIT의 각 구성 요소의 기여도를 분석하기 위해 ablation study를 수행했다. 모델은 image classification (ImageNet), semantic segmentation (ADE20K)에 대해 평가된다. 평가 결과는 아래와 같다.

BEIT_12.png




4. Analysis of Self-Attention Map

저자는 manual annotation 없이도 self-attention mechanism이 객체를 잘 분리해낼 수 있음을 보여준다. 아래 그림은 MS COCO 데이터셋을 이용하여 다양한 reference points에 대한 self-attention map 결과를 보여주고 있다. visualization은 마지막 레이어의 query-key product를 통해 계산된 attention score에 의해 생성되었다.

BEIT_13.png



AI, Paper Review
Image Generation Transformer
This post is licensed under CC BY 4.0 by the author.