[컴퓨터 비전] iBOT - ViT에서 [MASK] 사용하는 방법

"iBOT: Image BERT Pre-Training with Online Tokenizer"
ICLR 2022, 27 Jan 2022

 

1. Vision Transformer도 BERT처럼 Mask로 학습할 수 있을까?

NLP에서는 BERT와 같은 Transformer 모델이 MLM (Masked Language Modeling)를 통해 엄청난 성과를 냈습니다. 즉 모델에게 가려진 부분(mask)이 무엇인지 맞추는 방식으로 학습하여 성능을 높혔습니다.

 

ViT (Vision Transformer)가 이미지 분야에서 떠오르면서 BERT와 비슷하게 Mask를 이용한 학습 방법에 아이디어를 차용했습니다. 

“이미지의 일부 조각을 가리고 맞추는 식으로 학습하면 더 잘하지 않을까?”

이런 시도가 MIM (Masked Image Modeling) 입니다. ViT는 원래 이미지를 패치로 나눠서 처리하기 때문에, 패치를 가리는 것은 어렵지 않습니다. 문제는 가려진 부분을 어떻게 맞힐 건란 겁니다.

 

이에 MIM에서 중요한 게 Visual Tokenizer 입니다. Visual Tokenizer는 각 패치를 의미 있는 벡터(토큰)으로 변환해주는 역할 입니다. 가려진 부분을 맞출 때 해당 벡터를 출력하는 방식으로 맞히면 되는 겁니다.

 

이전 BEiT 모델에서는 사전학습된 토크나이저를 활용해서 MIM 학습을 진행했습니다. 하지만 사전학습된 토크나이저는 낮은 수준의 특징만 학습했고, 다른 도메인에 적용하기가 힘들었습니다.  

 

이에 iBOT (image BERT pre-training with Online Tokenizer)은 이름에서 보듯이 Online Tokenizer 를 사용했습니다. BEiT과 달리 Tokenizer를 사전학습하지 않고, 모델과 함께 학습합니다. 이에 데이터에 맞게 높은 수준의 이미지 특징을 학습할 수 있었습니다. 

 

iBOT의 핵심 출발점은 여기서부터 시작됩니다.

 

이미지의 일부를 가리고, 이를 맞추면서 학습하자. 그리고 그 의미는 직접 스스로 Tokenizer를 학습해서 데이터에서 이미지 의미를 배우자

 

2. 사전 지식: BERT의 MLM

iBOT을 이해하기 위해선 우선 BERT의 MLM에 대해서 알아보도록 한다. 

 

BERT의 핵심 학습 방식은 **Masked Language Modeling (MLM)**이라는 기법입니다. 예를 들어, 이런 문장이 있다고 해봅시다:

"나는 [MASK]를 먹는 걸 좋아해."

 

BERT는 이 문장에서 [MASK] 자리에 어떤 단어가 들어갈지를 맞히는 방식으로 학습합니다. 정답이 ‘바나나’일 수도 있고, ‘축구’일 수도 있습니다. 하지만 우리는 이 두가지 중에서 바나나가 들어갈 것 알수 있습니다. 이는 우리가 이걸 맞히기 위해선 앞뒤 문맥 전체를 이해해야 하고 있다는 점입니다. 

 

이와 비슷하게 AI 모델도 [MASK]에 있는 내용을 맞추어서 앞뒤 문맥을 파악하게 만들어 자연스럽게 "언어의 의미와 구조"를 배우게 하는 것입니다. 

 

이미지도 마찬가지 입니다. 가려진 이미지 패치가 뭔지 알려면 전체 이미지가 뭔지 알아야 되고, 각 구성들의 고차원 의미를 알아야 됩니다. iBOT에서도 이런 Mask를 통한 방식으로 학습을 진행합니다.

 

 

3. iBOT의 구조  

 

 

3-1. Self-Distillation

iBOT은 위 구조처럼 DINO 처럼 teacher 모델과 student 모델로 구성되어 있습니다.

학습 방법도 DINO 처럼 student 모델이 teacher 모델이 되도록 학습 시킨 후, student의 모델의 weight를 모먼트로 teacher 모델을 업데이트하는 방식을 사용합니다. 

 

이를 Momentum Encoder라고 합니다. 자세한 내용은 DINO에 설명하고 있으니 아래 링크를 참고하면 좋을 것 같습니다.

 

DINO : https://whale-yoon.tistory.com/8 

 

3-2. DINO와 차이점: Mask된 토큰 입력

DINO와 다른 점은 "Mask된 토큰을 사용한다는 점"입니다.

 

구성 요소	DINO	iBOT
Teacher 입력	전체 패치	전체 패치
Student 입력	전체 패치	일부 패치가 [MASK] 처리된 이미지

 

이 때 Teacher 모델에는 원본을 넣고, Student 모델에만 Masked 된 토큰을 넣어줍니다.

그리고 Student가 일부 패치를 [MASK]로 처리한 상태로 입력 받고, Teacher가 본 "원래 패치"와 똑같은 출력을 만들도록 학습합니다.

 

이 구조 덕분에 Teacher는 자연스럽게 Online Tokenizer 역할을 하게 됩니다. 미리 학습된 Tokenizer가 없어도 되고, 데이터에서 의미 있는 토큰 표현을 표현할 수 있게 만들어 줍니다.

 

3-3. 두가지 Loss : L[MIM] + L[CLS]

iBOT에서는 두가지 Loss를 사용합니다.

 

1. L[MIM] : 패치 복원 학습

 

L[MIM]은 Student의 Masked 패치 출력을 Teacher의 패치 출력과 비슷하도록 하는 학습합니다. 이 Loss를 통해서 BERT처럼 Mask로 이미지의 의미를 학습할 수 있게 해주는 Loss로 Cross entrophy을 사용해서 학습합니다.

 

수식은 아래와 같습니다. 

 

 

2. L[CLS] : 이미지 전체 의미 학습

 

iBOT도 DINO와 마찬가지로 Multi-Crop 전략을 이용해서 학습을 진행합니다. Teacher와 Student에 각각 다른 두 이미지 뷰를 넣고, Class 토큰([CLS])가 같은 결과가 나오도록 하여 이미지의 시맨틱 정보를 학습하도록 합니다. 이 Loss도 Cross Entropy를 이용해서 학습을 진행합니다. 

 

수식은 아래와 같습니다.

 

3-4. Projection head 공유 

그 외의 iBOT에서 사용한 구조는 패치 토큰을 생성하는 projection head와 [CLS] 토큰을 생성하는 projection head의 파라미터를 공유한다는 점입니다.

 

[CLS] 토큰으로 학습한 추상화 능력을 MIM에서 잘 활용되기 위해 이런 구조를 택했다고 합니다. 이런 구조는 별도의 head를 사용하는 것보다 실험적으로 좋은 성능을 냈다고 합니다.

 

4. 실험 결과

논문에서는 frozen된 feature에 대한 k-NN 분류, 리니어 프로빙과 Downstream Task로 Object Detection, Segmentation등의 성능 실험을 보여줍니다. 당연하겠지만 모든 실험에서 iBOT는 좋은 성능을 보여줍니다.

 

정량적으로 얼마나 좋아졌는 지 보기 위해서 ImageNet-1K를 통한 분류 실험 결과만 보겠습니다.

 

위 표는 실험 결과 입니다. 이 실험은 대규모 이미지 분류 벤치마크인 ImageNet-1K를 사용했습니다. 실험으로 출력한 feature가 얼마나 잘 의미가 담겼는지 보기 위해서 k-NN, 리니어 프로빙, Fine Tuning을 진행했습니다.

 

참고로 리니어 프로빙은 출력된 feature에서 뒤 단에 선형 분류기만 학습하여 분류가 잘되는 지 보는 평가 방법입니다.

 

ViT-B/16 기준 DINO와 성능을 보면 모두 iBOT이 성능이 좋다는 것을 알 수 있습니다.

방법	모델	Accuracy
k-NN	DINO	76.1%
	iBOT	77.1%
Linear Probing	DINO	78.2%
	iBOT	79.5%
Fine-tuning	DINO	83.6%
	iBOT	84.0%

 

이 결과 말고도 논문에서는 iBOT이 더 많은 파라미터를 가진 큰 모델에서 DINO 보다 좋은 결과를 냈다고도 합니다. 

 

5. iBOT의 패턴 구조

5-1. 높은 수준의 의미 학습 

iBOT은 사전 학습된 tokenizer를 사용한 BEiT과 달리 높은 수준의 의미를 학습한다고 합니다. 이를 위해서 비슷한 특징을 보이는 패치를 모아서 보았습니다. 

 

아래의 iBOT의 비슷한 특징의 패치들을 보면, 높은 수준의 의미 (자동차, 개 귀)와 낮은 수준의 의미 (줄무늬) 가 잘 학습된 걸 볼 수 있습니다.

 

iBOT의 비슷한 특징의 패치

 

하지만 아래처럼 BEiT의 결과를 보면 높은 수준의 의미가 없고, 단순한 패턴들만 학습된 것을 확인할 수 있습니다.

BEiT의 비슷한 패치 결과

 

5-2. 의미가 풍부한 패치 토큰

iBOT이 더 풍부한 의미를 학습한 토큰임을 보여주기 위해서 일부 패치만 사용해 리니어 프로빙을 수행했습니다.

더 자세히 말하면 self-attention score가 높은 패치 k를 뽑아서 패치 토큰의 평균을 내고, 이를 리니어 프로빙을 수행했습니다. 

 

DINO와 비교해서 결과를 보면 iBOT이 DINO보다 성능이 좋았다. 특히나 패치의 갯수가 적을수록 iBOT의 성능이 DINO보다 월등히 좋은 모습을 보여줍니다. 

 

 

5-3. 지역 세부 정보 손실에 견고함

iBOT은 MASK된 토큰을 맞추는 방법으로 학습했기 때문에 이미지의 일부 정보가 손실되어도 이미지 특징을 담을 수 있습니다.

 

이를 정량적으로 보여주기 위해서 배경 변화, 가림, out-of-distribution에 의한 선형 분류 결과를 보았습니다. 

DINO와 비교하면 일부만 제외하고 대부분에서 성능이 좋아지는 것을 볼 수 있습니다.

 

 

 

6. 정리

iBOT은 BERT의 Mask를 활용한 자가 학습을 ViT에서 적용할 수 있는 방법을 알려주었습니다. 이를 위해서 BEiT처럼 사전 학습된 Tokenizer를 사용하지 않고, 학습 데이터로 직접 학습하는 방법 (Online-Tokenizer)에 대해서 설명했습니다.

 

ViT의 Mask 활용 방법

• Self-Distillation + Masked Image Modeling 결합
• 두가지 Loss 사용 : L[MLM], L[CLS]
• CLS, 패치 Projection Head의 파라미터 공유

iBOT 학습시 나타나는 특징

• 데이터에서 직접 학습해 고수준의 의미 학습
• 패치 토큰에 더 많은 의미를 학습
• 일부 정보가 손실되어도 의미를 잘 보존

“이미지 일부만 보고 전체를 이해하는 능력, iBOT은 진짜 ‘이미지 BERT’가 되고 있어요.”