BEiT (2022) 논문 리뷰

BEiT 논문의 제목은 BEiT: BERT Pre-Training of Image Transformers다. (링크)

 

저자는 Hangbo Bao, Li Dong, Songhao Piao, Furu Wei다. 

 

Github: 링크

 

 

Abstract

본 논문에서는 self-supervised vision representation model인 BEiT (Bidirectional Encoder representation from Image Transformers 를 제시한다. 자연어 분야의 BERT를 따라서, 저자들은 masked image modeling task를 활용해서 vision Transformers를 pretrain한다. 구체적으로, 각각의 이미지는 pre-training에서 2가지 시점을 가진다. 하나는 Image patches 16 x 16 픽셀이고 다른 하나는 discrete tokens인 visual tokens다. 우선, 원본 이미지는 비주얼 토큰으로 토크나이즈한다. 그 다음 몇개의 이미지 패치를 마스킹한 다음 백본 트랜스포머로 보낸다. 사전학습 objectives는 corrupted image patches로 부터 원본 비주얼 토큰을 복원한다. BEiT를 사전학습한 다음, 저자들은 pretrained encoder에 task layers를 덧붙이고 전체 모델에 대해서 곧바로 다운스트림 태스크에 대한 파인 튜닝을 수행한다. Image classifiction과 semantic segmentation에서 경쟁력있는 성과를 달성함을 실험적으로 보인다. 

 

 

2. Methods

입력 이미지 $x$가 주어지면 BEiT는 이를 contextualized vector representation로 인코딩한다. Figure 1에서 볼 수 있듯이 BEiT는 masked image modeling (MIM)라는 self-supervised 학습의 측면에서 사전학습된다. MIM은 인코딩 벡터에 기반한 마스킹된 이미지 패치의 복원을 목적으로 한다. Image classifiction과 semantic segmentation과 같은 downstream tasks에 대ㅐㅎ서는 tak layers를 pretrained BEiT에 덧붙인 다음 각각의 특정한 데이터에 대해서 파라미터를 파인튜닝한다. 

 

 

 

2.1. Image Representations

 

2.1.1. Image Patch

 

2D 이미지는 patches의 시퀀스로 쪼개져서 표준 트랜스포머가 이미지를 직접 받아들일 수 있도록 한다. 수식으로 표현하면 이미지 $x \in \mathbb{R}^{H \times W \times C}$를 ($N = HW / P^2$) $x^p \in \mathbb{R}^{N \times (P^2 C)}$로의 reshape다. $C$는 채널의 수, $(H, W)$는 입력 이미지의 해상도, $(P, P)$는 개별 패치의 해상도다. 이미지 패치들 $\{ x_i^p \}_{i=1}^N$은 벡터로 flatten된 다음 linearly projection이 적용되여 BERT의 word embedding처럼 표현된다. 이미지 패치들은 raw pixels를 보존하고 BEiT에 입력 피쳐로 사용된다. 본 논문에서는 224 x 224 이미지를 14 x 14 그리드의 이미지 패치로 쪼개는데, 개별 패치의 크기는 16 x 16이 된다. 

 

 

2.1.2. Visual Token

 

자연어와 유사하게, 저자들은 raw pixels 대신 image tokenizer를 통해서 얻을 수 있는 dicrete tokens의 sequence로 이미지를 표현한다. 구체적으로 입력 이미지 $x \in \mathbb{R}^{H \times W \times C}$를 $z = [z_1, ..., z_N] \in \mathcal{V}^{h \times w}$로 토크나이징한다. 이때 $\mathcal{V} = \{ 1, ..., |\mathcal{V}| \}$는 discrete token indices를 보유하고 있는 vocabulary다. 

 

DALL-E (Zero-Shot Text-to-Image Generation) 논문에서 제시되었듯이, 저자들은 dVAE (discrete variational autoencoder)에 의해 학습된 이미지 토크나이저를 사용한다. Visual token learning 과정에는 2가지 모듈, tokenzier와 decoder가 존재한다. 토크나이저 $q_{\phi}(z|x)$는 이미지 픽셀 $x$를 discrete tokens $z$로 매핑하는데 visual codebook (i.e., vocabulary)에 따라서 수행한다. 디코더 $p_{\psi}(x|z)$는 비주얼 토큰 $z$에 기반해서 입력 이미지 $x$를 복원한다. Recontruction objective는 $\mathbb{E}_{z ~ q_{\phi}(z|x)} [ log p_{\psi}(x|z) ] $로 표기할 수 있다. Latent visual tokens는 이산형이기 때문에 모델 학습이 non-differentiable이다. Gumbel-softmax relaxation을 적용해서 모델의 파라미터를 학습한다. 추가적으로 dVAE의 학습 과정에서 uniform prior를 $q_{\phi}$에 적용한다. 이미지 토크나이저의 자세한 내용은 DALL-E를 참고하면 된다.

 

각각의 이미지를 14 x 14 크기의 이미지 패치들의 그리드로 토크나이징을 한다. 하나의 이미지에 대한 비주얼 토큰의 수와 이미지 패치의 수는 동일하다. Vocabulary 크기는 $|\mathcal{V}|$는 8192다. DALL-E에서 소개된 공개된 이미지 토크나이저를 그대로 사용하였으며, re-implementation은 Appendix C에 소개되어 있다. 

 

2.2. Backbone Network: Image Transformer

ViT를 따라서 표준적인 트랜스포머 구조를 백본 네트워크로 사용한다. 

 

입력은 이미지 패치들의 시퀀스인 $\{ x_i^p \}_{i=1}^N$다. 이는 linearly projected 되어서 patch embeddings $E_{x_i^p}$로 표기한다. 이때 $E \in \mathbb{R}^{N \times (P^2 C)}$의 차원을 가진다. 추가적으로 special token인 [S]를 입력 시퀀스의 맨앞에 덧붙인다. 또한 standard learnable 1D position embeddings  $E_{pos} \in \mathbb{R}^{N}$를 패치 임베딩에 더한다. Input vectors $H_0 = [ e_{[S]}, E_{x_i^p}, ..., E_{x_N^p} ] + E_{pos}$는 트랜스포머에 입력된다. Encoder는 L개의 레이어로 이루어지며, Transformer blocks $H^l$을 가지고 있다. $H^l = Transformer( H^{l-1}  )$이고 이때 $l = 1, ..., L$이다. 마지막 레이어의 output vectors $H_L = [ h_{[S]}^L, h_1^L, ..., h_N^L ] $는 이미지 패치의 encoded representations다. 이때 $h_i^L$는 $i$ 번째 이미지 패치의 벡터다. 

 

2.3. Pre-Training BEiT: Masked Image Modeling

Masked image modeling (MIM)에서는 이미지 패치의 임의의 퍼센트를 랜덤하게 마스킹한다. 그 다음 masked pathces에 해당하는 부분의 비주얼 토큰을 예측한다. 

 

자세한 과정은 다음과 같다. 우선 이미지 패치의 40% 정도를 랜덤하게 마스킹한다. Masked positions는 $\mathcal{M} \in \{ 1, ..., N \}^{0.4N}$로 표기한다. Masked patches를 learnable embedding $e_{[M]} \in \mathbb{R}^{D}$로 대체한다. Corrupted image patches $x^{\mathcal{M}} = \{ x_i^p : i \notin \mathcal{M} \}_{i=1}^N \cup \{ e_{[M]}: i \in \mathcal{M} \}$을 L개의 레이어로 구성된 트랜스포머에 주어진다. Final hidden vectors $\{ h_i^L \}_{i=1}^N$를 이미지 패치의 encoded representations로 간주한다. 각각의 마스킹된 포지션 $\{ h_i^L : i \in \mathcal{M} \}_{i=1}^N$에 대해서 softmax classifier를 활용해서 예측한다. 예측된 비주얼 토큰은 $p_{\text{MIM}}(z' | x^{\mathcal{M}})$ = $\text{softmax}_{z'} (W_c h_i^L + b_c)$다. 이때 $x^{\mathcal{M}}$은 corrupted image이며, 이때 $W_c \in \mathbb{R}^{ | \mathcal{V} | \times D }, b_c \in \mathbb{R}^{ | \mathcal{V} |  }$다. 아래 (1)의 Pre-training objective 는 주어진 corrupted image에 대해서 올바른 비주얼 토큰 $z_i$를 예측하는 log-likelihood의 최대화다.

 

max $\sum_{x \in D} \mathbb{E}_{\mathcal{M}} [ \sum_{i \in \mathcal{M}} log p_{\text{MIM}}(z' | x^{\mathcal{M}}) ]$ ... (1).  

 

이때 $D$는 training corpus고, $\mathcal{M}$은 randomly masked positions이며, $\mathcal{M}$에 해당하는 corrupted image는 $x^{\mathcal{M}}$다. 

 

 

 

2.4. From the Persepctive of Variational Autoencoder

BEiT의 pre-training은 variational autoencoder의 학습으로 볼 수도 있다. $x$를 원래 이미지, $\hat{x}$를 마스킹된 이미지, $z$는 비주얼 토큰이다. Log-likelihood $p(x | \hat{x})$의 evidence lower bound (ELBO)는 아래의 (2)처럼 표기가 가능하다. 이는 corrupted image에서 original image의 복원으로 볼 수 있다.

 

$\sum_{(x_i, \hat{x}_i) \in D} p(x_i | \hat{x}_i) \geq \sum_{(x_i, \hat{x}_i) \in D} \mathbb{E}_{z_i ~ q_{\phi}(z|x_i)}  [ log p_{\psi}(x_i | z_i) ] - D_{KL} [ q_{\phi}(z|x_i), p_{\theta}(z | \hat{x}_i)  ] $ ... (2)

 

 

이때 $q_{\phi}(z|x_i)$는 이미지 토크나이저로 비주얼 토큰을 생성한다. $p_{\psi}(x | z)$는 주어진 비주얼 토큰에 대해 원본 이미지를 디코드한다.  $p_{\theta}(z | \hat{x})$는 주어진 masked image에 기반해서 비주얼 토큰을 복원한다. 

 

Two-stage procedure로 모델을 학습한다.

 

첫 번째 단계에서는 uniform prior와 함께 $\mathbb{E}_{z_i ~ q_{\phi}(z|x_i)}  [ log p_{\psi}(x_i | z_i) ]$로 표기되는 recoonstruction loss를 최소화한다. 

 

두 번째 단계에서는 $q_{\phi}$와 $p_{\psi}$를 고정한 상태에서 prior인 $p_{\theta}$를 학습한다. 저자들은 $q_{\phi}(z | x_i)$를 one-point distribution로 단순화하는데 바로 most likely visual tokens를 반환하는 $\hat{z}_i = \text{argmax}_z q_{\phi}(z | x_i)$다. 

 

Equation (2)를 (3)처럼 다시 쓸 수 있다. 

 

 

 

2.5. Pre-Training Step

공정한 비교를 위해서 ViT-Base를 따라서 12-레이어, 768 hidden size, 12 attention heads로 설정한다. Intermediate size of feed-forward networks는 3072다. 16 x 16 크기의 이미지 패치를 적용했으며, DALL-E의 이미지 토크나이저를 그대로 사용하고, Vocab size는 8192다.

 

BEiT를 ImageNet-1K로 학습하는데 총 1.2 M개의 이미지를 가진다. Augmentation은 random resized cropping, horizontal flipping, color jittering을 사용했다. Self-supervsied learning이기에 labels을 사용하지 않으며, 224 x 224 해상도를 사용한다. 총 이미지 패치와 비주얼 토큰 모두 그리드의 크기는 14 x 14다. 랜덤하게 마스킹을 하는데 최대 총 75개까지만 한다. 이는 약 전체 패치의 40%에 해당한다. 

 

사전학습에서 총 500k steps (800 epochs)이며, 2k batch size를 사용한다. $\beta_1$ = 0.9, $\beta_2$ = 0.999을 가지는 Adam 옵티마이저를 사용하였으며 10 epochs를 warmup으로, learning rate는 1.5e-3이고 cosine learning rate decay, 그리고 weight decay는 0.0.5다. Stochastic depth를 0.1 rate로, 그리고 dropout은 사용하지 않는다. Nvidia의 Tesla V100 32GB GPU 16개를 사용해서 5일 동안 500k 스텝을 학습했다. 

 

안정적인 트랜스포머의 학습을 위해서 weights의 초기화는 [-0.02, 0.02] 구간에서 수행했다. $l$ 번째 레이어에서의 Self-attention과 feed-forward의 출력에 대해서 $\frac{1}{\sqrt{2l}}$로 rescaling을 수행한다. 

 

2.5. Fine-Tuning BEiT on Downstream Vision Tasks

Image Classification

Linear classifiier를 task layer로 활용하여 Softmax를 적용한다.

 

Semantic Segmentation

SETR-PUP 논문 (링크)에서 제시한 방법대로 task layer를 설정한다.

 

Intermediate fine-tuning

Self-supervised pre-training 이후 Intermediate-task transfer learning with pretrained language models: When and why does it work? 논문 (링크)에서 제시한 대로 intermediate fine-tuning을 학습한 다음, 다운스트림 태스크에 대해서 학습한다. 

 

 

3. Experiments

 

 

Table 1, 2, 3을 보면 BEiT가 ViT나 MoCo, DINO 같은 SOTA 모델들과 견줄만한 성능임을 알 수 있다.

 

 

 

Table 4에서는 Ablation 스터디의 결과인데 blockwise masking와 visual tokens를 빼면 성능이 감소함을 알 수 있다. 결국 두 가지 모두 필요하다. 

 

 

 

 

Figure 2에서는 Self-supervised가 사람이 매긴 라벨이 아니더라도 물체를 잘 탐지함을 알 수 있다.

 

 

 

Table 6에서는 B, L, H의 다양한 크기의 ViT와 BEiT의 성능 비교를 보여준다.