ICLR 2022Citations: 3,000+

BEiT: BERT Pre-Training of Image Transformers

BEiT: 이미지 트랜스포머의 BERT 사전학습

Hangbo Bao, Li Dong, Songhao Piao, Furu Wei (2021)

BERT의 마스크 토큰 예측 방식을 비전 트랜스포머(ViT)에 적용하여, 이미지 패치를 마스킹하고 시각적 토큰을 예측하는 사전학습으로 강력한 시각 표현을 학습한 논문.

배경

BERT가 NLP에서 마스크 언어 모델링으로 획기적 성공을 거두었지만, 비전에서의 마스크 예측 사전학습은 효과적이지 않았다. 이미지 패치는 텍스트 토큰과 달리 연속적이어서 예측 타겟을 정의하기 어려웠다. ViT의 등장으로 이미지를 패치 시퀀스로 처리할 수 있게 되었지만, 효과적인 자기지도 사전학습 방법은 아직 확립되지 않았다.

핵심 아이디어

BEiT(Bidirectional Encoder representation from Image Transformers)는 두 단계로 구성된다. 먼저 dVAE(discrete VAE) 토크나이저를 학습하여 이미지 패치를 이산적 시각 토큰으로 변환한다. 이후 ViT를 사전학습할 때, 입력 이미지의 약 40% 패치를 마스킹하고, 마스킹된 위치의 시각 토큰을 예측하도록 학습한다. 핵심은 원시 픽셀이 아닌 이산 시각 토큰을 예측함으로써, 모델이 저수준 세부사항보다 의미적으로 풍부한 표현을 학습하게 한다는 것이다. 블록 단위 마스킹(blockwise masking)을 사용하여 연속적인 패치 영역을 가리는 것도 효과적이다.

방법론

DALL-E의 dVAE를 이미지 토크나이저로 사용하여 14x14 격자의 시각 토큰을 생성한다. ViT-B/16 인코더에서 약 40%의 패치를 블록 단위로 마스킹하고, 마스킹된 위치에 [MASK] 토큰을 삽입한다. 소프트맥스 분류기로 해당 위치의 시각 토큰(8192 클래스)을 예측한다. ImageNet-1K 300 에폭 사전학습 후 미세조정했다.

주요 결과

ImageNet 미세조정에서 ViT-B 기준 83.2% top-1 정확도를 달성하여, DeiT(81.8%)와 DINO(82.8%)를 앞섰다. ADE20K 시맨틱 세그멘테이션에서도 사전학습의 효과가 두드러졌으며, 특히 저데이터 레짐에서 지도학습 사전학습 대비 큰 우위를 보였다.

임팩트

BERT 스타일의 마스크 예측 사전학습이 비전에서도 효과적임을 최초로 대규모로 입증하여, MAE, data2vec 등 후속 마스크 이미지 모델링 연구의 직접적 선구자가 되었다. 이산 시각 토큰을 예측 타겟으로 사용하는 아이디어는 BEiT v2, PeCo 등으로 발전했다.

관련 Foundation 논문

관련 논문