IBOT : IMAGE BERT PRE-TRAINING WITH ONLINETOKENIZER 리뷰

Zhou, Jinghao, Chen Wei, Huiyu Wang, Wei Shen, Cihang Xie, Alan Yuille, and Tao Kong. "ibot: Image bert pre-training with online tokenizer." arXiv preprint arXiv:2111.07832 (2021). 

https://arxiv.org/abs/2111.07832 

iBOT: Image BERT Pre-Training with Online Tokenizer

Abstract

이 연구에서는 visual tokenizer을 이용한 Masked Image Modeling에서의 이점과 그에대한 도전들을 포함한다. 우리는 online tokenizer을 이용해 masked prediction을 수행하는 자기지도 학습 형태의 프레임워크 iBOT을 제안한다. masked patch token에 대해 self-distillation을 적용하고 teacher network를 online tokenizer로 채택하였다, 또한 class token에 self-distillation하여 visual semantic을 획득한다. online tokenizer은 학습 가능하며 이는 선행학습 되지 않아도 된다.

 

Introduction

Masked Language Modeling은 자연어 모델을 사전 훈련시키기 위한 유명한 방법이다. 최근 이미지 처리를 위한 비지도 형태의 사전 훈련 모델들은 대부분 이미지의 internal-structure을 무시하고 global view만을 다루기 때문에 local token을 다루는 MLM과는 상반된다. 본 연구에서는 NLP에서와 같이 ViT의 훈련을 향상시키기 위한 Masked Image Modeling의 기본 요소를 탐구하는 것이 목적이다. MLM에서는 WordPiece와 같은 lingual tokenizer가 사용되었는데 이와 같이 MIM의 핵심은 Visual Tokenizer이다.

 

우리는 Knowledge Distillation 기반의 MIM 모델 iBOT을 제안한다. 이는 tokenizer로부터 knowledge distillation되어 학습되고 twin teacher를 통해 self-distillation 된다. target network로는 마스킹된 이미지가 입력되고 online network로는 원본 이미지가 입력된다. target network는 tokenizer output에 상응하는 masked patch token을 원상태로 복구하는 것을 목표로 한다. online tokenizer은 두가지의 주요 목표를 가진다. 1) 이미지의 유사성을 교차적으로 판단하여 클래스토큰에 주입함으로써 고품질의 시각적 정보를 주입한다. 2) tokenizer은 별도의 선행 훈련과정 없이 momentum update를 통해 최적화 된다. (lingual tokenizer같은 경우에는 단어를 token화 하기 위한 추가적인 훈련과정이 필요하다)

 

 

Preliminaries

1. Masked Image Modeling as Knowledge Distillation

MIM은 이미지를 \(x=\{{x_i}\}_{i=1}^N\) 인 이미지 토큰으로 나눈 후 N개 토큰에 대해 r 비율로 마스킹하여 샘플링 한다. 이후 마스킹된 토큰에 대해 원본 이미지를 복원 하도록 훈련된다. BEiT에서는 파라미터들이 categorical distribution으로 모델링되어 각 패치토큰에 해당하는 category를 예측하는 것인데 여기서 사용된 loss가 knowledge distillation과 비슷하다.

 

2. DINO에서 제안된 Self-distillation은 사전 정의된 posterior distribution으로부터 distillation되는 것이 아닌 모델의 과거 iteration에서 distillation받는 자기지도적인 방법이다. teacher, student network로 서로 다른 방법으로 증강된 이미지가 입력되어 클래스 토큰을 이용해 categorical distribution을 예측하고, cross-entropy를 최소화하도록 훈련된다. teacher, student는 동일한 backbone구조를 가지며 student network의 지수이동평균을 이용해 teacher network가 업데이트 된다.

 

IBOT

 

1. Framework

첫번째로 우리는 증강된 이미지 \(u, v\)에 대해 blockwise masking을 적용한 \(\hat{u}, \hat{v}\)을 획득한다. \(\hat{u}\)를 예로 들면, student는 masked view \(\hat{u}\)의 projection인 \(\hat{u}_s^{patch}=P_\theta^{patch}(\hat{u})\)를 출력하고, teacher은 non-masked view \(u\)의 projection인 \(u_t^{patch}=P_\theta\prime^{patch}(u)\) 를 출력한다.

 

iBOT의 MIM 목적함수는 (3)과 같다.

Fig3와 같이 teacher의 backbone인 \(f_t\)와 projection head \(h_t^{patch}\)은 마스킹 패치토큰에 대한 online token distribution을 생성하는 visual tokenizer에 해당한다. 따라서 tokenizer은 별도의 사전 학습이 필요 없이 현재 데이터셋에서의 지식을 증류 받는다.

 

online tokenizer가 의미있도록 보장하기 위해 자기주도 학습에서 사용된 것과 같이 cls 토큰에 대해 cross-view image를 이용한 self-distillation을 진행하여 시각 정보를 획득한다. 우리는 DINO에서와 달리 student에 \(u_s^{[CLS]}\) 대신 \(\hat{u}_s^{[CLS]}\)을 입력한다. 우리는 self-distillation 된 cls token으로부터 capability of semantics abstraction을 빌리기 위해 cls token, patch token에 공유된 projection head를 사용한다. 이는 각각 분리된 head를 사용한 경우보다 더 좋은 성능을 보인다고 한다.

 

정확한 의미를 내포하여 tokenize되는 단어와 달리 image token은 모호한 의미를 담는다. 때문에 우리는 one-hot token 대신 softmax를 이용한 토큰 분포를 사용하였다.

 

2. Implementation

- ViT, Swin Transformers를 사용하였고, 224크기의 이미지를 사용하여 pre-train, fine-tune 하였다. projection head는 3-layers MLP이며 DINO와 같이 l2-norm bottleneck을 사용하였고 8192-d를 출력한다.

 

이 이후에 매우 많은 downstream task results, ablation study가 있기 때문에 자세한 내용은 논문을 참고하길 바란다.