오늘은 IBOT이라는 논문에서 제안한 self-supervised learning model에 대해서 살펴보려고 한다.
https://arxiv.org/abs/2111.07832
iBOT: Image BERT Pre-Training with Online Tokenizer
The success of language Transformers is primarily attributed to the pretext task of masked language modeling (MLM), where texts are first tokenized into semantically meaningful pieces. In this work, we study masked image modeling (MIM) and indicate the adv
arxiv.org
Motivation
MAE와 DINO는 각각 다른 방향에서 representation learning을 잘 해내지만, 그만큼 분명한 한계도 가지고 있다.
먼저 MAE를 보면, 학습 시 reconstruction loss가 patch token에만 직접적으로 적용된다. 이 때문에 patch-level 표현은 잘 학습되지만, CLS token은 학습 목표에서 상대적으로 소외된다. 실제로 MAE의 CLS token은 classification과 같은 global task에서 성능이 기대보다 낮게 나오는 경우가 많다.
반대로 DINO는 self-distillation loss를 CLS token에만 직접 적용한다. 이 덕분에 global representation은 매우 잘 학습되지만, patch token들은 attention을 통해서만 간접적으로 학습되기 때문에 local / dense한 표현을 배우는 데에는 한계가 있다.
정리하면,
- MAE는 patch는 강하지만 global 표현이 약하고
- DINO는 global 표현은 강하지만 patch 표현이 약하다
이러한 상반된 한계를 보면 자연스럽게 다음과 같은 질문이 생긴다.
“그렇다면 CLS token과 patch token 둘 다에 직접 loss를 걸면 안 될까?”
iBOT은 바로 이 질문에서 출발한다.
CLS token에는 global semantic을, patch token에는 local semantic을 담당하도록 하고, 두 종류의 토큰 모두에 self-distillation 기반의 학습 목표를 명시적으로 부여함으로써, 전역적 표현과 국소적 표현을 동시에 잘 학습하는 representation을 만들고자 한다.
Method
1️⃣ View 생성
먼저 하나의 이미지 x로부터 두 개의 augmentation view u와 v를 생성한다.
그 다음, image masking 기법을 적용하여
각각의 view에 대해 masking된 버전인 u′와 v′를 생성한다.
즉, 총 네 개의 view가 만들어진다.
- u, v (masking 없음)
- u′, v′ (masking 적용)
2️⃣ Teacher–Student ViT 통과
u, v, u′, v′를 모두 ViT에 통과시켜 representation을 얻는다.
- u, v → teacher ViT
- u′, v′ → student ViT
각 view로부터
- CLS token
- patch tokens
의 representation이 각각 생성되며,
총 8개의 representation이 만들어진다.
- u_cls, u_patches
- v_cls, v_patches
- u′_cls, u′_patches
- v′_cls, v′_patches
3️⃣ Teacher 업데이트 방식
student 네트워크는 gradient를 통해 직접 학습되며,
teacher 네트워크는 student 네트워크의 EMA(Exponential Moving Average)로 업데이트된다.
이러한 teacher–student 비대칭 구조는
이전 self-supervised 학습 방법들에서 자주 사용되었으며,
representation collapse를 방지하는 데 중요한 역할을 한다.
중요한 점은,
loss는 같은 네트워크에서 나온 representation 사이에 걸리지 않는다는 것이다.
즉,
- student output ↔ student output
- teacher output ↔ teacher output
사이에는 loss를 걸지 않고,
항상 teacher에서 나온 representation을 기준으로 student를 학습시킨다.
4️⃣ Patch token에 대한 distillation loss
masked region에 대해,
u_patches와 u′_patches 사이에 distillation loss를 적용한다.
각 patch token은
C개의 class로 분류되는 확률 분포로 해석되며,
student가 teacher의 분포를 따라가도록 loss를 건다.
이를 통해 patch representation이
의미 있는 local semantic 정보를 학습하도록 유도한다.
같은 방식의 loss를
- v_patches와 v′_patches
사이에도 동일하게 적용한다.
5️⃣ CLS token에 대한 distillation loss
CLS token에 대해서도 distillation loss를 적용한다.
- u_cls와 v′_cls
- v_cls와 u′_cls
사이에서 loss를 걸어,
서로 다른 view에서 나온 CLS token이 같은 이미지를 같은 class로 분류하도록 학습한다.
여기서 중요한 점은,
- u_cls와 v_cls 사이에는 loss를 걸지 않는다.
이 경우 collapse가 발생할 수 있기 때문이다.
이러한 설계를 통해 CLS token은
이미지 전체를 대표하는 global representation을 학습하게 된다.
6️⃣ 전체 방법 정리
iBOT은 latent space에서
MAE의 masking 전략과 DINO의 self-distillation 방식을 결합한 방법이라고 볼 수 있다.
- patch token → masking + distillation → local 정보 학습
- CLS token → view-level distillation → global 정보 학습
이를 통해 iBOT은
CLS token과 patch token 모두가 유의미한 정보를 담도록 학습한다.
Experiments
두 loss는 서로 보완적으로 작용하여, CLS token이 중요한 downstream task와 patch token이 중요한 downstream task 모두에서 좋은 성능을 보인다.