Abstract
BEIT는 Bidirectional Encoder representation form Image Transfomers의 약자로 self-supervised vision representation model이다. NLP에서의 BERT처럼 vision Transformers의 pretrain task로 masked image modeling을 제안한다. 각 image는 image patches (예를들어 16×16 pixels)와 visual tokens 두 가지로 분류된다. 우선, image를 visual token으로 "tokenize"한다. 이후 임의로 몇몇 image patch를 mask한 후 Transformer backbone에 넣어줍니다. Pre-training objective는 masking된 image patch를 원래 visual token으로 복구하는 것이다. BEIT를 pre-training한 후 encoder에 task layers를 추가하여 downstream task의 model parameter를 직접 finetune한다. 이미지 분류 및 의미론적 분할에 대한 실험 결과는 우리 model이 이전 pre-training 방법 대비 좋은 결과를 달성했음을 보여준다.
1 Introduction
Transformer는 computer vision 영역에서 많은 성공을 거뒀다. 그러나 vision Transformer는 convolutional nerural network보다 훨씬 많은 학습 데이터를 요구했다. 관련 문제를 해결하기위해 self-supervised pre-training이 대규모 데이터를 레버리지 하는 좋은 방안이 될 수 있다. (예를 들어 Contrastive learning, self-distillation)
NLP에서의 BERT에 영감을 받아 denoising auto-encoding 방식을 vision Transformer pretrain에 적용하였다.
해당 방식을 바로 적용하기에는 두 가지 문제가 존재한다.
- Pre-existing vocabulary가 존재하지 않는다.
- Mask 된 patch의 원본 pixel을 복구하기위한 예측은 회귀 문제이다. 그러나 이러한 pixel 수준 복구 작업은 short-range dependencies 및 high-frequency detail을 pre-train하는데 modeling 기능을 낭비하는 경향이 있다.
*Short-range dependencies: 인접한 pixel의 정보
*High-frequency detail: 이미지의 frequency는 intensity 값의 변화율이다. 따라서 high-frequency 이미지는 intensity 값이 한 pexel에서 다음 pixel로 빠르게 변경되는 이미지입니다.
해당 논문의 기여도:
- Self-supervised learning으로 vision Transformer를 pre-train하기 위한 Masked Image Modeling(MIM) task를 제안한다. 우리는 또한 Variational Autoencoder의 관점에서 이론적 설명을 제공한다.
- 우리는 BEIT를 pre-train하고 이미지 분류 및 의미론적 분할과 같은 다운스트림 작업에 대한 광범위한 fine-tuning 실험을 수행한다.
- 우리는 self-supervised BEIT의 self-attention 메커니즘이 사람의 주석을 사용하지 않고도 의미론적 영역과 객체 경계를 구별하는 방법을 학습한다는 것을 제시한다.
2 Methods
주어진 image 입력 $x$에 대해 BEIT는 contextualized vector representations로 변환한다. Figure 1에서 보이듯이 BEIT는 self-supervised learning 같은 MIM로 pretrained 한다. MIM은 encoding vectors에 기반하여 masked image patch를 복원한다. 이후 downstream task에 대해서는 layer를 추가하여 fine-tune한다.
2.1 Image Representations
Image는 image patch 와 visual tokens 두 가지 representation이 존재하며 각각 pre-training 동안 input과 output으로 적용한다.
2.1.1 Image Patch
2D image를 patch의 sequence로 변환하여 Transformer에 직접적으로 image data를 넣을 수 있게 변환한다.
$x \in \mathbb{R}^{H\times W\times C}$의 image를 $N=HW/P^2$로 reshape하여 $x^P \in \mathbb{R}^{N\times P^2 C}$ patch를 만든다. 이 때 $C$는 channels, $(H, W)$는 image resolution, $(P, P)$는 패치의 resolution이다. 이후 image patches $\left\{ x^p_i \right\}^N_{i=1}$는 flatten 한 후 BERT와 유사하게 linearly projected 된다.
해당 논문에서 $224 \times 224$ image를 $14 \times 14$ grid를 기준으로 image patch를 생성했다. (즉, 각각의 patch는 $16 \times 16$의 resolution을 가진다)
img_size=(224,224), patch_size=(16,16), in_chans=3, embed_dim=768
self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size)
x = self.proj(x) #(1,768,14,14)
x = x.flatten(2).transpose(1, 2) #(1,196,768)2.1.2 Visual Token
자연어와 같이 image를 discrete token의 sequence로 표현한다. Image $x \in \mathbb{R}^{H\times W\times C}$를 $z=[z_1,...,z_N] \in V^{h \times w}$로 tokenize 하며, 이 때 vocabulary $V=\left\{ 1,...,|V| \right\}$는 discrete token을 포함한다. Image tokenier 는 discrete variational autoencoder (dVAE)로 학습된다. Token 학습에는 tokenizer 와 decoder가 존재한다. Tokenizer $q_\phi(z|x)$는 image pixel $x$를 discrete token $z$로 vocabulary에 따라 매핑한다. Decoder $P_\psi(x|z)$ 는 image pixel $x$를 discrete token $z$로부터 복원해낸다. 이 때 복원에 관련된 objective는 $\mathbb{E}_{z \sim q_\phi(z|x)}[logp_\psi (x|z)]$ 이다. Vocabulary size 는 $|V|=8192$이다.
사용된 image tokenizer: https://github.com/openai/DALL-E
2.2 Backbone Network: Image Transformer
ViT에 따라 일반적인 Transformer를 backbone으로 사용했다. Transformer는 image patch의 input $x \in \mathbb{R}^{P^2 \times C}$ 을 받아 $Ex^p_i, E \in \mathbb {R}^{(P^2C)\times D}$와 같이 embedding한다. 또한 special token [S] 를 sequence의 맨 앞에 추가한 후 학습가능한 1D position embedding $E_{pos}\in \mathbb{R}^{N \times D}$를 더해준다. $$\texttt{Input vector:}\quad H_0=[e_{[S]}, Ex^p_i, ... Ex^p_N] + E_{pos}$$
$$\texttt{L layers of Transformer blocks:}\quad H^l = \texttt{Transformer(H^{l-1})}$$
$$\texttt{Output vector:}\quad H^L=[e^L_{[S]}, h^L_i, ... h^L_N]$$
이 때 $h^L_i$는 i번째 patch의 hidden vector다.
2.3 Pre-Training BEIT: Masked Image Modeling
논문에서 제안하는 Masked Image Modeling(MIM)은 임의로 image patch의 일부를 mask한 후 그에 해당하는 visual token을 예측하는 것이다. 약 40%의 image patch를 임의로 mask 하였으며 masked position은 $M\in \left\{ 1,...,N \right\}^{0.4N}$ 이다. Masked patch는 학습 가능하게 embedding 되어 $e_{[m]}\in \mathbb{R}^D$, $L$-layer Transformer에 $x^M=\left \{ x^p_i:i \notin M \right\}^N_{i=1} \cup \left \{ e_{[M]}:i \in M \right\}^N_{i=1}$ 가 입력된다. 마지막 hidden vector $\left \{ h_i^L\right \}^N_{i=1}$은 input patch의 encoded representation 으로 간주된다. 이 때 각각의 masked position $\left \{ h^L_i:i \in M \right\}^N_{i=1}$에서 softmax 분류기를 사용하여 오염된 image $x^M$의 visual token $p_{MIM}(z'|x^M)=\texttt{softmax}_{z'}(W_ch_i^L+b_c)$, $W_c \in \mathbb{R}^{|V|\times D}$와 $b_c \in \mathbb{R}^{|V|}$ 을 예측한다.
Pre-training objective: $\texttt{Max} \sum_{x\in D} {\mathbb{E}_M} [\sum_{i\in M} \texttt{log}p_{MIM}(z_i|x^M)]$
이 때, $D$는 training corpus, $M$은 임의의 masked position, $x^M$은 M으로 mask된 image이다.
이 때 단순 임의로 masking position을 선정하지 않고 blockwise로 masking했다.
한 번에 block 단위로 masking이 진행되었으며, block은 최소 16개의 patch로 구성되며 임의의 aspect ratio을 가지게 된다.
import random
import math
import numpy as np
class MaskingGenerator:
def __init__(
self, input_size, num_masking_patches, min_num_patches=4, max_num_patches=None,
min_aspect=0.3, max_aspect=None):
if not isinstance(input_size, tuple): # input이 tuple이 아닌경우 tuple로 변환
input_size = (input_size, ) * 2
self.height, self.width = input_size # (14, 14)
self.num_patches = self.height * self.width # 196
self.num_masking_patches = num_masking_patches # 118 (118/196 = 0.60) 살려놓을 pixel
self.min_num_patches = min_num_patches # 16
self.max_num_patches = num_masking_patches if max_num_patches is None else max_num_patches # 118
max_aspect = max_aspect or 1 / min_aspect # 3.33
self.log_aspect_ratio = (math.log(min_aspect), math.log(max_aspect)) # (-1.20, 1.20)
def __repr__(self):
repr_str = "Generator(%d, %d -> [%d ~ %d], max = %d, %.3f ~ %.3f)" % (
self.height, self.width, self.min_num_patches, self.max_num_patches,
self.num_masking_patches, self.log_aspect_ratio[0], self.log_aspect_ratio[1])
return repr_str
def get_shape(self):
return self.height, self.width
def _mask(self, mask, max_mask_patches):
delta = 0
for attempt in range(10):
target_area = random.uniform(self.min_num_patches, max_mask_patches) # ex) 20.16, 면적 즉 patch 갯수
aspect_ratio = math.exp(random.uniform(*self.log_aspect_ratio)) # ex) 2.67, 종횡비
h = int(round(math.sqrt(target_area * aspect_ratio))) # 7, mask 안 할 height
w = int(round(math.sqrt(target_area / aspect_ratio))) # 3, mask 안 할 width
if w < self.width and h < self.height: # mask할 w, h가 원본의 w, h를 넘으면 안됨
top = random.randint(0, self.height - h) # mask 시작 h 선정
left = random.randint(0, self.width - w) # mask 시작 w 선정
num_masked = mask[top: top + h, left: left + w].sum() # mask 안할 면적 중 mask 안 된 면적 확인
# Overlap
if 0 < h * w - num_masked <= max_mask_patches: # 면적단위로 한번에 mask 안할 면적이 max값보다 작은지 점검
for i in range(top, top + h):
for j in range(left, left + w):
if mask[i, j] == 0:
mask[i, j] = 1
delta += 1
if delta > 0:
break
return delta
def __call__(self):
mask = np.zeros(shape=self.get_shape(), dtype=np.int32) #input shape에 맞게 mask 행렬 생성
mask_count = 0
while mask_count < self.num_masking_patches: # 최대 mask 안할 횟수
max_mask_patches = self.num_masking_patches - mask_count
max_mask_patches = min(max_mask_patches, self.max_num_patches) # 1회 masking 안할 면적 단위 점검용
delta = self._mask(mask, max_mask_patches)
if delta == 0: # masking 실패
break
else:
mask_count += delta
# maintain a fix number {self.num_masking_patches}
if mask_count > self.num_masking_patches: # mask 안하는 횟수가 넘어간 경우
delta = mask_count - self.num_masking_patches
mask_x, mask_y = mask.nonzero() # mask 안된 위치 반환
to_vis = np.random.choice(mask_x.shape[0], delta, replace=False) # mask 안된 애들 중 넘치는 갯수 만큼을 선택
mask[mask_x[to_vis], mask_y[to_vis]] = 0 # mask 수행
elif mask_count < self.num_masking_patches: # mask가 부족한 경우
delta = self.num_masking_patches - mask_count
mask_x, mask_y = (mask == 0).nonzero()
to_mask = np.random.choice(mask_x.shape[0], delta, replace=False)
mask[mask_x[to_mask], mask_y[to_mask]] = 1
assert mask.sum() == self.num_masking_patches, f"mask: {mask}, mask count {mask.sum()}"
return mask
if __name__ == '__main__':
import pdb
generator = MaskingGenerator(input_size=14, num_masking_patches=118, min_num_patches=16,)
for i in range(10000000):
mask = generator()
if mask.sum() != 118:
pdb.set_trace()
print(mask)
print(mask.sum())Blockwise mask와 MIM은 pixel-level auto-encoding이 short-range dependecies와 high-frequency detail에 모델의 학습을 집중하는 것의 문제를 극복할 수 있게 되었다.
2.4 From the Perspective of Variational Autoencoder
BEIT pre-training 은 variational autoencoder training으로 볼 수 있다. $x$를 원본 image, $\tilde{x}$를 mask된 image, $z$를 visual tokens라고 하자. Log-likelihood $p(x|\tilde{x})$의 evidence lower bound (ELBO)를 고려할 때, 원본 이미지 복원은
여기서
- $q_{\phi}(z|x)$는 image tokenizer
- $p_{\psi}(x|z)$는 decoder
- $p_{\theta}(z|\tilde{x})$는 masked image를 visual token으로 복원하는 것이다. (MIM task)
*$D_{KL}$은 Kullback–Leibler divergence로 두 확률분포의 차이를 계산하는 데에 사용하는 함수로, 어떤 이상적인 분포에 대해, 그 분포를 근사하는 다른 분포를 사용해 샘플링을 한다면 발생할 수 있는 정보 엔트로피 차이를 계산한다.
위 식을 두 단계로 나누어 minimize the reconstruction loss 하고 $p_{\theta}$를 학습한다고 하면 위 수식은 다음과 같이 기술할 수 있다.
2.5 Pre-Training Setup
- 12-layer Transformer, 768 hidden size, 12 attention heads, 3072 feed-forward networks
- 16 $\times$ 16 input patch size
- 8192 vocabulary size of visual tokens
- Pretrain BEIT on ImageNet-1K
: This dataset spans 1000 object classes and contains 1,281,167 training images, 50,000 validation images and 100,000 test images. (예시: https://www.kaggle.com/c/imagenet-object-localization-challenge/overview/description) - Data augmentation (random resized cropping, horizontal flipping, color jittering)
from torchvision import datasets, transforms
from transforms import RandomResizedCropAndInterpolationWithTwoPic, _pil_interp
# oringinal beit data augmentation
self.common_transform = transforms.Compose([
transforms.ColorJitter(0.4, 0.4, 0.4),
transforms.RandomHorizontalFlip(p=0.5),
RandomResizedCropAndInterpolationWithTwoPic(
size=args.input_size, second_size=args.second_input_size, scale=(args.min_crop_scale, 1.0),
interpolation=args.train_interpolation, second_interpolation=args.second_interpolation,
),
])- 224 $\times$ 224 해상도, 14 $\times$ 14 patch들로 분할, 75개 patch를 mask (약 40%)
- 500k step (800 epochs) with 2k batch size
- Adam $\beta _1 = 0.9$, $\beta _2 = 0.999$
- lr = 1.5e-3 with warmup of 10 epoch, and cosine learning rate decay
- Weight decay 0.05
- Stochastic depth with 0.1 rate, disalbe dropout
- 500k training steps : 16ea Nvidia Tesla V100 32GB GPU 로 5일 소요
- Random initialization within small range $[-0.02, 0.02]$, 이후 $l$-th Transformer layer에서 $\frac{1}{\sqrt{2l}}$로 rescale
2.6 Fine-Tuning BEIT on Downstream Vision Tasks
BEIT를 pre-training한 이후 task layer를 Transformer 위에 올려서 downstream task에 대해 fine-tune했다.
Image classification
Linear classifier를 추가하여 분류했다. 구체적으로 average pooling을 통해 representation을 aggregate 후 softmax classifier에 넣었다. $\texttt{softmax}(\texttt{avg}(\left \{ h_i^L \right \} ^N_{i=1} W_c))$ 여기서 $h_i^L$ 는 $i$ -th image patch의 최종 encoding vector, $W_c \in \mathbb{R}^{D\times C}$ 는 parameter 행렬, $C$는 라벨의 수이다. Labelled data의 likelihood를 maximize하는 방향으로 parameter를 update했다.
Semantic segmentation
Semantic segmentation을 위해서 SETRPUP에서 사용된 task layer를 사용했다. 구체적으로 BEIT에 몇 개의 deconvolution layer를 segmentation 생성을 위한 decoder로 넣어주었다. 해당 모델 역시 Image classification처럼 end-to-end로 fine-tune 되었다.
Intermediate fine-tuning
Self-supervised pre-training 이후, ImageNet-1K로 intermediate 학습한 후 target downstream task로 finetuning 한다.
3 Experiments
Image classification과 semantic segmentation에 대해 실험했다.
3.1 Image Classificaton
ILSVRC-2012 ImageNet dataset with 1k classes and 1.3M image에 대해 DeiT의 hyperparameter를 그대로 따라서 fine-tuning을 진행하여 직접 비교를 수행했다. Fine-tuning 시의 epoch는 scratch로부터 학습할 때보다 적게 진행했다.
Table 1은 image classification의 top-1 accuracy를 보여준다. 세 가지 다른 방식으로 학습을 진행했으며 거의 동등한 수준에서 비교가 되지만 iGPT의 경우는 1.36B parameter를 가지고 있고 ViT-JFT300M의 경우 Google 자체 300M image로 pretrain되었다.
- Random initialization과 비교해 보았을 때 pre-train된 BEIT의 성능이 좋은 성능을 보여주었다.
- 과거 SOTA 모델들의 성능을 넘어선다.
Fine-tuning to $384 \times 384$ resolution
$224 \times 224$ 로 fine-tuning한 후에 추가적으로 $384 \times 384$에 대해 10 epochs 더 fine-tuning했다. DeiT의 higher-resolution setting을 따랐다.
참고해야할 점은 patch size는 resolution에 상관없이 동일하게 사용하였기 때문에 patch 수가 늘어났다. 즉, sequence length of Transformer가 길어졌다. (어떻게 처리했는가?)
Higher resolution에서 BEIT는 1+point 성능 향상이 있었으며 이 값은 supervised pre-training ViT384 (ImageNet-22K) 값을 상회하는 값이다.
Scaling up to larger size
BEIT의 규모를 키워서 학습을 진행해보았다. ViT384-L의 경우 ViT385보다 성능이 떨어지는 것이 확인되는데, 이는 vision Transformer의 data-hungry issue을 입증한다. 이 문제는 ImageNet-22K로 supervised pre-training 하여 1.2 point 상승하면서 어느 정도 해소가 되었다. 이에 비해 BEIT-L의 경우 2 point BEIT384의 경우 1.7 point 상승했다. 이를 통해 BEIT이 적은 데이터로 모델 규모를 키울 때 더 효과적이다.
Convergence curves
Figure 2에서 convergence curve를 통해 BEIT의 fine-tuning scratch부터 학습하는 DeiT보다 빠르게 수렴하고 금방 좋은 성능을 내는 것을 확인했다.
3.2 Semantic Segmentation
Semantic segmentation은 각 pixel의 class를 예측하는 것이다. 해당 논문에서는 ADE20K benchmark (25K images and 150 semantic categories)를 이용했다. 논문에서는 semantic categories에 대해 평균낸 metric of mean Intersection of Union (mIoU) 값을 비교했다.
해당 평가에서 hyperparameter는 SETR-PUP을 따랐으며 Adam optimizer로 lr은 1e-3 with layer-wise decay를 사용했다. fine-tuning은 160K step에서 진행했으며 batch size 는 16이다.
Table3에서 BEIT를 ImageNet데이터에 대해 supervised pre-training 한 것과 비교한다. 해당 결과에서 self-supervised가 0.3 point 좋은 결과를 보여준다. 추가적으로 intermediate fine-tuning을 수행하였다. 이는 ImageNet으로 fine-tuning을 한 후 ADE20K로 추가 fine-tune한 것이다. Intermediate finetuning은 좋은 결과를 보여주었다.
3.3 Abalation Studies
Abalation study를 image classification과 semantic segmentation에 대해 수행하였으며 전체 step의 37.5%인 300 epochs를 default pre-training step으로 잡았다.
- Blockwise masking은 양쪽 task 모두에서 효과적이였으며 특히 semantic segmentation에서 더 효과적
- Visual tokens 사용을 masked patch의 원본 pixel예측 task로 변경하여 pixel regression 문제로 변경하여 비교
: Pixel-level autoencoding을 outperform
: Table 1의 ViT (from scratch) 보다 성능이 낮은 것으로 볼 때 vision token이 BEIT의 핵심 - Blockwise masking과 visual token을 둘 다 사용하지 않을 것과 비교해 볼 때 blockwise masking은 pixel-level regression에서 더 효과적이다.
- 모든 visual token을 복원하는 것은 성능을 저하 (decoder를 태워서 복원한 값을 이용?)
- 다른 방식으로 학습
- 800 epochs 까지 길게 학습하면 더 성능이 향상
3.4 Analysis of Self-Attention Map
Pre-training 단계에서 특별한 annotation이 없더라도 BEIT는 object를 self-attention mechanism을 이용해 완벽히 분리할 수 있다.
Figure 2에서 확인할 수 있는 것은 image 내의 self-attention map이 서로 다른 reference point를 가리킨다. Visualizations 는 마지막 layer에서의 query-key prodect 값으로 연산된 attention score이다. 각각의 reference point에 해당하는 patch를 query로 하여 어떤 patch가 관여하는지 보여준다. 해당 figure를 통해 BEIT이 downstream task에 대해 왜 일반적으로 성능이 적은 데이터로도 좋았는지 보여준다.
4 Related Work
[skip]
5 Conclusion
Self-supervised pre-training를 vision Transformer에 도입.
Image classification, semantic segmentation과 같은 downstream task에서의 효과적인 fine-tuning 결과.
Multimodal pre-training을 일관된 방식으로 진행할 수 있다.
