BEIT: BERT Pre-Training of Image Transformers

Oct 29, 2021

BEIT: BERT Pre-Training of Image Transformers

Contents

Table of Content 1. Introduction 2. Method 2.1. Image Representations 2.2. Backbone Network: Image Transformer 2.3. Pre-Training BEIT: Masked Image Modeling 2.4. From the Perspective of VAE 2.5. Pre-Training Setup 2.6. Fine-Tuning BEIT on Downstream Vision Tasks 3. Experiments 3.1. Image Classification 3.2. Semantic Segmentation 3.3. Ablation Studies 3.4. Analysis of Self-Attention Map 4. Conclusion

Table of Content

Table of Content 1. Introduction 2. Method 2.1. Image Representations 2.1.1. Image Patch 2.1.2. Visual Token 2.2. Backbone Network: Image Transformer 2.3. Pre-Training BEIT: Masked Image Modeling 2.4. From the Perspective of VAE 2.5. Pre-Training Setup 2.6. Fine-Tuning BEIT on Downstream Vision Tasks 2.6.1. Image Classification 2.6.2. Semantic Segmentation 2.6.3. Intermediate Fine-tuning 3. Experiments 3.1. Image Classification 3.1.1. Fine-Tuning to 384 x 384 Resolution 3.1.2. Scaling Up to Large Size 3.1.3. Convergence Curves 3.2. Semantic Segmentation 3.3. Ablation Studies 3.4. Analysis of Self-Attention Map 4. Conclusion

1. Introduction

NLP에서 Transformer가 한 차례 센세이션을 불러일으킨 후, 다음 주자는 BERT였다는 것을 모두 기억하실 겁니다!

현재 CV에서도 ViT의 등장으로 같은 역사가 반복되고 있는데요, BERT, 혹은 Masked Language Modeling(MLM)을 접목시킨 모델이 등장하는 것은 시간문제일 뿐, 필연적인 일이었습니다. 그리고 그 필연이 불과 두 달 전인 6월에 일어났습니다!

Microsoft Research 에서 발표한 BEIT는 Bidirectional Encoder representation from Image Transformers의 약어로, 제안하는 Masked Image Modeling(MIM) 기법으로 Pre-Training한 모델입니다. OpenAI 에서 발표한 DALL-E 논문과 겹치는 부분이 제법 있으니 함께 보시면 도움이 많이 될 겁니다!

2. Method

Figure 1: Overview of BEIT pre-training. Before pre-training, we learn an “image tokenizer” via autoencoding-style reconstruction, where an image is tokenized into discrete visual tokens according to the learned vocabulary. During pre-training, each image has two views, i.e., image patches, and visual tokens. We randomly mask some proportion of image patches (gray patches in the figure) and replace them with a special mask embedding [M]. Then the patches are fed to a backbone vision Transformer. The pre-training task aims at predicting the visual tokens of the original image based on the encoding vectors of the corrupted image

BEIT는 입력 이미지 x를 Contextualized Vector Representation 으로 인코딩하는 역할을 합니다. Encoder은 MIM을 통해 Masking된 이미지를 원본으로 복원하는 방법을 Self-Supervised Learning하게 됩니다. 그렇게 학습된 Encoder의 뒷단에 Task Layer를 추가함으로써 Downstream Task에서의 활용이 가능해집니다. 아래에서 더 자세히 살펴보겠습니다.

2.1. Image Representations

저자들은 이미지를 Image Patch 와 Visual Token 두 가지 개념으로 처리하였습니다.

2.1.1. Image Patch

<AN IMAGE IS WORTH 16X16 WORDS: TRANSFORMERS FOR IMAGE RECOGNITION AT SCALE>

ViT에서 사용되는 개념과 동일하게, 차원의 2D 이미지를 N개의 (P x P) 패치로 Split하여 차원으로 Reshape합니다. 후에 차원의 Linear Transform을 거쳐 D차원의 Representation Vector로 변환합니다. 이렇게 변환된 Vector를 Image Patch라 칭하며, Image Patch는 Standard Transformer의 입력으로 바로 사용할 수 있습니다.

2.1.2. Visual Token

BERT의 철학을 유지하며 Masked Image Modeling을 진행하기 위함인지, 저자들은 Encoder가 Masking된 부분을 추론할 때 Raw Pixel 대신 Discrete한 Token으로 추론하게끔 설계하였습니다. Raw Pixel을 직접적으로 생성하게 하는 것은 모델의 시야를 편협하게 만들며 자주 등장하는 패턴에 집중하게 만듦으로 바람직하지 않습니다. 이 과정은 Tokenizer와 Decoder로 이루어진 dVAE(discrete Variational AutoEncoder)를 활용함으로써 해결하였습니다. dVAE는 DALL-E에서 사용된 것과 동일하므로 아래 인용으로 설명을 대체합니다. 저자들은 OpenAI 측에서 공개한 dVAE를 튜닝 없이 그대로 사용했습니다.

<Zero-Shot Text-to-Image Generation> 논문 정리글 인용

먼저 dVAE(discrete Variational AutoEncoder)를 훈련합니다. dVAE는 256 x 256 이미지를 32 x 32의 Image Tokens로 변환하는 역할을 합니다. 각 Token들은 8192개의 Codeword가 되어 Visual Codebook에 업데이트됩니다. 이를 이 방법을 통해 이미지의 큰 손실 없이 Transformer 학습에 필요한 Context를 192배 줄일 수 있습니다. 필연적으로 이미지는 손상되나 그 정도가 심하지 않다는 것을 아래 그림에서 확인할 수 있습니다.

Comparison of original images (top) and reconstructions from the discrete VAE (bottom). The encoder downsamples the spatial resolution by a factor of 8. While details (e.g., the texture of the cat’s fur, the writing on the storefront, and the thin lines in the illustration) are sometimes lost or distorted, the main features of the image are still typically recognizable. We use a large vocabulary size of 8192 to mitigate the loss of information.

단, dVAE가 Codebook에서 Image Tokens을 Discrete하게 처리한다는 특성을 가지고 있기 때문에, 일반적인 방법으로는 Gradient가 흐를 수 없어 Back Propagation이 불가능합니다. 이에 저자들은 Gumbel Softmax를 적용하였습니다. Gumbel Softmax는 Gumbel-max Trick과 Softmax의 결합으로 tau값에 따라 One-hot과 Uniform을 오가는 값으로 Reparameterization할 수 있게 하는 함수입니다. tau가 0에 가까울수록 One-hot에 가까운 값이 나오고 tau가 커질수록 Uniform에 가까워집니다.

2.2. Backbone Network: Image Transformer

저자들은 ViT와 직접적인 비교를 위해 2017년에 발표된 Standard Transformer를 Backbone Network로 채택했습니다. 앞서 언급한 테크닉들을 통해 자연어의 Embedding과 완전히 동일한 형태로 Image를 변환할 수 있기 때문에 가능한 일입니다.

Input Sequence 앞에는 [S] 토큰을 추가하여 문장(?)의 시작을 알리고, 1D Position Embedding을 Patch Embedding에 더해주었습니다.

2.3. Pre-Training BEIT: Masked Image Modeling

저자들은 전체 Sequence에서 약 40%의 비율로 Image Patch를 Masking을 진행했습니다. Masking은 Masking할 Image Patch를 학습되는 Embedding인 [M] 으로 대체함으로써 이루어집니다.

조금 특이한 것은 Random Index에 대해 Masking을 하는 대신에 Blockwise Masking을 진행했다는 점입니다. Blockwise Masking의 알고리즘은 아래와 같습니다.

2D Image Patches에서 한 Patch를 선택한 후, 해당 Patch로부터 Random Aspect Ratio를 결정하여 Masking을 하는 방식입니다. 자연어에서도 동일한 이름 또는 N-gram Masking 이라 칭해지며 그 효과가 입증되었다고 합니다.

2.4. From the Perspective of VAE

BEIT의 Pre-Training 과정은 어쩌면 Variable Autoencoder의 학습 과정으로 볼 수도 있다고 저자들은 말합니다. 원본 이미지를 , 마스킹 된 이미지를 , Visual Token을 라 하면 아래 수식처럼 마스킹 된 이미지로부터 원본 이미지를 복원해내는 ELBO(Evidence Lower BOund)로 나타낼 수 있다는 겁니다.

😅

필자는 Autoencoder에 대한 이해가 부족해 자세한 설명은 생략하겠습니다. 자세한 내용이 궁금하신 분은 논문을 직접 참고하시길 바랍니다!

2.5. Pre-Training Setup

BEIT 는 ViT 와의 공정한 비교를 위해 ViT-Base 와 완전히 동일한 세팅으로 설계하였습니다.

12 Layer, 768 Hidden Size, 12 Attention Heads 3072 Feed-forward Hidden, 16 x 16 Patch Size

Tokenizer는 DALL-E 의 것을 빌려왔기 때문에, 그 셋팅과 동일한 8192의 Vocab Size를 적용하였습니다.

Pre-Training은 ImageNet-1K의 1.2M개의 Image에 대해 진행되었습니다. Augmentation은 Random Resize Cropping, Horizontal Flipping, Color Jittering을 적용하였고 입력 크기는 224 x 224 를 사용하였습니다. 총 14 x 14 개의 Patch로 분할되었으며 그중 75개를 Masking하여 약 40%의 Masking 비율을 적용하였습니다.

그 외 Hyperparameters Adam Optimizer() Learning Rate: 1.5e-3 with warmup 10 epochs apply cosine learning rate decay Weight Decay: 0.05 Stochastic Depth with 0.1 rate Disable Dropout Training 500K Step(800 epochs) with 2K Batch Size 5 Days using 16 Nvidia Tesla V100 32GB GPUs

2.6. Fine-Tuning BEIT on Downstream Vision Tasks

Pre-Trained BEIT 는 BERT처럼 Task Layer를 모델에 붙여 Fine-Tuning을 합니다. 저자들은 예시로써 Image Classification과 Semantic Segmentation Task에 대한 Fine-Tuning을 소개합니다.

2.6.1. Image Classification

단 하나의 Linear Layer를 추가함으로써 Fine-Tuning이 가능합니다. Representations를 Aggregate하는 데에는 Average Pooling을 사용했습니다.

2.6.2. Semantic Segmentation

SETR-PUP와 동일한 형태의 Task Layer를 추가합니다. Pre-Trained BEIT를 Backbone Encoder로 사용하고 그 뒤에 몇 개의 Deconvolution Layer를 추가하여 Segmentation을 생성하는 Decoder를 완성합니다.

2.6.3. Intermediate Fine-tuning

Self-supervised Pre-training을 마친 후에, 저자들은 Intermediate Fine-tuning을 진행하였습니다. 이름 그대로 중간 Fine-tuning인데, Self-supervised Learning 후, Downstream Task에 적용하기 전, 그 중간에 수가 많은 데이터셋(본 논문에선 ImageNet-1K)에 대해 사전 Fine-tuning을 진행합니다. Fine-tuning 과정은 BERT가 NLP에서 학습하는 방법과 완전히 동일합니다.

3. Experiments

3.1. Image Classification

ImageNet 과 CIFAR-100 에 대해 성능을 측정하였습니다. 공정한 비교를 위해 DeiT 의 하이퍼파라미터를 적극적으로 따랐으며 Fine-Tuning 모델은 Scratch 모델보다 더 적은 Epoch을 훈련하였습니다.

Table 1: Top-1 accuracy of image classification on CIFAR-100 and ImageNet-1K. The models are at resolution 224 × 224, except ViT384 and ViT384-JFT300M use 384 × 384. The results, unless otherwise indicated, are all obtained by base-size models. *: result is taken from (Chen et al., 2021). †: iGPT-1.36B contains 1.36 billion parameters, while others are base-size models. ‡: ViT-JFT300M is pretrained with the “masked patch prediction” task on Google’s in-house 300M images, while others use ImageNet.

주목할 만한 점은 CIFAR-100 데이터셋에 대해 Scratch ViT 가 48.5% Accuracy 밖에 달성하지 못한 것에 비해, BEIT 는 Self-Supervised Training 만으로 90.1%를 달성했다는 점입니다. 이는 데이터셋의 Annotation에 Dependency가 있는 ViT에 비해 BEIT가 효과적일 수 있음을 보여줍니다. 게다가 ImageNet에서도 더 높은 성능을 보이고 있으므로 Resource가 충분한 경우에도 BEIT가 더 효과적임을 보여줍니다.

이전의 Transformer 계통의 Self-Supervised 모델인 DINO나 MoCo v3 모두 ImageNet Fine-Tuning된 BEIT가 능가하는 모습을 확인할 수 있습니다.

3.1.1. Fine-Tuning to `384 x 384` Resolution

저자들은 224 x 224 에 대해 Fine-Tuning을 마친 후에, 추가적으로 384 x 384 입력에 대해 10 Epoch 학습을 했습니다. 모든 세팅은 DeiT 의 High-Resolution 세팅을 적극적으로 따랐으며, Epoch 수는 오히려 감소시켰습니다. Patch Size는 동일하게 유지하기 때문에 결론적으로 Input Sequence의 길이가 길어지게 되었습니다. 아래는 결과입니다.

Table 2: Top-1 accuracy on ImageNet-1K. We evaluate base- (“-B”) and large-size (“-L”) models at resolutions 224 × 224 and 384 × 384. †: iGPT-1.36B contains 1.36 billion parameters, while others are base-size models. ‡: ViT384-B-JFT300M is pretrained with the “masked patch prediction” task on Google’s in-house 300M images, while others use ImageNet.

눈여겨볼 점은 ImageNet-22K 를 학습한 모델보다 ImageNet-1K 로 학습한 의 성능이 더 뛰어나다는 점입니다!

3.1.2. Scaling Up to Large Size

Scratch부터 학습한 경우에 모델이 모델보다 성능이 좋은 것은 ViT의 Data-Hungry Issue를 여실히 보여줍니다. BEIT의 경우 모델의 사이즈가 커질 수록 성능이 좋아지는 것을 확인할 수 있습니다. Supervised에서 ViT는 모델 확장을 통해 1.2% 상승을 보였지만, Self-Supervised에서 BEIT-B를 BEIT-L로 확장했을 때 2.0% 이상 상승하는 것을 확인할 수 있는데, 이것은 Extremely Bigger Model을 만드는 데에 BEIT가 적합한 모델이 될 수 있음을 시사합니다.

3.1.3. Convergence Curves

Figure 2: Convergence curves of training DeiT from scratch and fine-tuning BEIT on ImageNet-1K.

위 그래프에서 확인할 수 있듯이, BEIT는 성능이 좋을 뿐만 아니라 DeiT에 비해 수렴되는 속도도 더 빠릅니다. 게다가 Pre-Training BEIT 를 Fine-Tuning 하는 경우라면 심지어 단 몇 번의 Epoch 만으로도 수렴이 가능합니다.

3.2. Semantic Segmentation

Semantic Segmentation에 대해서는 ADE20K 데이터셋에 대해 성능을 측정하였습니다. ADE20K 데이터셋은 25K개의 이미지와 150개의 클래스를 가지고 있습니다. 저자들은 모든 카테고리의 IoU(Intersection of Union) 값을 평균 낸 mIoU(mean IoU)를 Metric으로 사용했습니다. 결과는 아래와 같습니다.

Table 3: Results of semantic segmentation on ADE20K. We directly use the same task layer of SETR-PUP (Zheng et al., 2020) for a fair comparison. Only single-scale inference is used.

표에서 확인할 수 있듯이, 제안하는 BEIT는 ImageNet에 Supervised Pre-Training된 모델보다 좋은 성능을 보입니다. 나아가, Intermediate Fine-tuning은 ADE20K와 관련 없는 ImageNet 데이터셋에 대해 진행되었음에도 2% 가량의 성능 향상을 보여주어 유효한 방법론임을 입증하였습니다.

3.3. Ablation Studies

Table 4: Ablation studies for BEIT pre-training on image classification and semantic segmentation.

저자들은 모델의 다양한 세팅을 변경해가며 비교 실험을 진행했습니다. Blockwise Masking은 두 가지 Task에서 모두 이점이 있음을 확인할 수 있었는데, 특히 Sementic Segmentation에서 좋은 성능 향상을 확인할 수 있었습니다.

Visual Token 대신 Raw Pixel을 Regression하게 학습을 진행하면 성능이 하락하는 것을 확인할 수 있었습니다. 이 결과를 통해 Visual Token이 BEIT의 중요한 Key 중 하나임을 알 수 있습니다. Blockwise Masking도 함께 제거했을 때 더 큰 성능 하락을 보이는 것을 보면 Blockwise Masking이 모델의 시야가 편협해지는 걸 방지하는 데에 도움이 되지 않았을까 저자들은 말합니다.

😅

Recover 100% visual tokens는 무슨 실험인지 언급이 되어있지를 않습니다...(?)

3.4. Analysis of Self-Attention Map

Figure 3: Self-attention map for different reference points. The self-attention mechanism in BEIT is able to separate objects, although self-supervised pre-training does not use manual annotations.

저자들은 BEIT의 Self-Attention Mechanism이 객체를 분리해낼 수 있다고 말합니다. Pre-Training 중에 한 번도 노출되지 않은 MS COCO Dataset 에 대해 Self-Attention Map을 각기 다른 Reference Point에 대해 시각화하였고, 그 결과는 위와 같습니다.

Reference Point에 해당하는 Patch를 Query로 사용하여 어떤 Patch가 그것에 Attend되는지를 확인하였습니다(?). BEIT 는 어떤 Supervision도 없이 스스로 Self Attention Heads를 통해 객체를 구분해내고 있었으며 이는 BEIT 가 Downstream Task에서 잘 동작할 수 있는 이유가 될 수 있을 것입니다.

4. Conclusion

이로써 논문은 마무리가 되었습니다. 말미에서 저자들이 제안하는 Future Works가 흥미로운데요, 아래와 같은 두 가지 Future Works를 제시하고 있습니다.

우리는 BEIT 를 어마어마하게 Scale Up 할 것이다!

우리는 Multi-Modal Pre-Training을 할 것이다. Text와 Image를 혼합해서!

NLP에서 BERT 이후가 GPT 시리즈의 등장이었던 만큼 1번을 통해 CV계의 GPT를 만들어 낼 수 있으리라 기대가 되며, 2번은 개인적으로 관심을 가지고 있는 본질적인 Embedding에 대해 한 발짝 다가갈 것 같아 흥미가 돋네요. 향후 연구가 기대가 됩니다!