model.eval()

딥러닝 모델들은 학습과 추론에서 다르게 작동하는 Layer를 가지고 있는 경우가 있습니다. 대표적으로 BatchNorm, Dropout가 있으며 이런 Layer 들이 추론 과정에서 학습시와는 다르게 작동하게 해주는 것이 .eval() 입니다

#train
...
model.train() #train mode
...

#eval 
model.eval() #-> eval mode
...
with torch.no_grad():
    ...

torch.no_grad()

위 코드에서 모델 추론 부분에서는 model.eval() 말고도 torch.no_grad()도 함께 사용했습니다. torch.no_grad()로 감싸진(with) 부분에서는 gradient 계산을 하지 않아 메모리/속도가 개선됩니다. 공식 문서에서는 loss.backward()를 호출하지 않는다면 사용할 것을 권장하고 있습니다.

requires_grad=False

보통 아래와 같이 많이 사용합니다.

for p in model.parametes():
    p.requires_grad=False

requires_grad=False 를 적용하면 모델의 특정 부분에 그라디언트 계산을 멈출 수 있습니다. torch.no_grad() 와 가장 큰 차이는 그라디언트를 저장은 한다는 것입니다. 따라서 모델의 특정 부분은 freeze 하고 나머지는 학습시키는 등의 전략을 사용할 때 사용합니다.

torch.no_grad() VS requires_grad=False 의 설명이 약간 부실한것 같은데 아래 링크로 가시면 더 자세한 설명이 있습니다.

LINK

처음에는 wheel 문제로 생각하고 wheel을 재설치해봤지만 여전히 똑같은 문제가 발생했습니다. 그래서 다시 살펴보니 gcc를 인식못하고 있는 것으로 보였습니다.

파란 박스 부분을 보시면 gcc를 실행시키지 못하고 있습니다. 아래 코드로 gcc를 설치하고 다시 pycocotools를 설치하니까 문제가 해결되었습니다!

apt update apt install gcc pip install pycocotools

Abstract/Intro

본 논문은 BERT, BEIT, DINO 논문들에서 나오는 내용을 모두 알고 읽으시는게 이해가 쉬울 것 같습니다. 위 논문들의 링크입니다. BERT | BEIT | DINO

BERT에서 사용되었던 MLM을 최근 Computer Vision에 적용하기 위한 연구들이 있다. DALL-E tokenizer를 사용해서 Image를 Token으로 만들었던 BEIT가 대표적이다. 본 논문은 BEiT의 Tokenizer를 사전에 학습해서 사용하는 Multi-stage training 방식에서 벗어나 Tokenizer를 target model과 함께 훈련시키는 online tokenizer 방식을 제시한다. Online Tokenizer를 학습하기 위해서 DINO에서 사용했던 Self-Distillation 방식을 채택했다.

Masked Language Model(이하 MLM)은 BERT에서 사용했던 사전학습 방식으로, 입력값의 일부에 마스킹을 적용한뒤 이를 최종 출력에서 예측한다. Computer Vision 분야에서도 ViT의 등장과 함께 MLM을 이미지 인식모델 사전학습에 사용하려고 했지만 MLM에서 필수인 Tokenizer의 부재로 사용에 어려움이 있었다. 최근 논문인 BEiT에서는 DALL-E pre-trained Tokenzier로 이미지를 토큰화해서 이를 해결했다.

BEiT의 학습 순서는 다음과 같다.

1. 이미지를 사전에 학습된 DALL-E Tokenizer를 이용해서 토큰화
2. 이미지 패치의 일부에 마스킹을 적용
3. 모델이 마스킹된 패치의 token값을 예측 하도록 학습(MIM Head에서 예측)
4. Downstream task로 fine-tune

BEiT는 결국 DALL-E Tokenizer의 결과를 학습하는 것이기 때문에 DALL-E의 지식을 target 모델에 Distillation한다고 봐도 될것 같다.

IBOT에서 지적한 부분은 1번의 사전 학습된 토크나이저를 사용한 것이다. 토크나이저를 학습하기 위한 데이터가 추가로 필요하고, end to end 방식이 아닌 토크나이저 학습 - target 모델 학습 이라는 mulit-stage 방식으로 학습을 해야한다는 것이다.

본 논문에서 Tokenizer는 이미지의 semantic 정보를 잘 추출할 수 있어야하고, target 모델의 목표또한 이미지의 semantic 정보 추출에 있기 때문에 이 둘을 동시에 학습할 수 있다고 가정하고 IBOT 모델을 실험했다.

BEiT의 한계점이었던 사전에 학습된 Tokenizer를 사용했다는 부분을 DINO의 Self-Distillatio으로 해결했다. BEiT + DINO = IBOT

Preliminaries

Masked Image Modeling as Knowledge Distillation

BEiT에서는 DALL-E의 tokenizer를 사용해서 MIM을 구현했습니다. BEiT는 아래의 수식을 최소화합니다.

$$ \begin{aligned} -\sum_{i=1}^{N} m_{i} \cdot P_{\phi}(x_{i})^{T} \log P_{\theta}(\hat{x_{i}}) \end{aligned} $$

$P(\cdot)$은 사전에 정의된 K(논문에선 8192) 차원의 확률 분포를 나타내며 $m_{i}$는 마스킹 여부, $x_{i}, \hat{x_{i}}$는 각각 원본 이미지, 마스킹된 이미지를 나타냅니다. 위 식을 살펴보면 사전 확률분포(DALL-E tokenizer) $P_{\phi}(x_{i})$ 와 모델의 출력인 $P_{\theta}(\hat{x_{i}})$가 서로 비슷해 지도록 합니다. DALL-E tokenizer는 사전 학습되어 있으며 모델의 학습 동안 파라미터가 바뀌지 않기 때문에 BEiT는 DALL-E tokenizer의 지식을 모델에 증류하는 것으로 볼 수 있습니다.

Self-Distillation

DINO 논문에서 제시된 자가 증류는 사전에 학습된 모델 $P_{\phi}$가 아닌 모델의 전 iteration인 $P_{{\theta}^{'}}$ 의 지식을 증류한다. Training Set 에서 랜덤하게 추출된 이미지 $x$를 서로 다른 Augmentation을 적용해서 $u, v$를 생성하고 각각 teacher, student 모델의 입력값으로 사용한다. 아래의 식으로 지식 증류가 일어난다.

$$ \begin{aligned} L_{[CLS]} = -P_{{\theta}^{'}}^{[CLS]}(v)^{T} \log P_{\theta}^{[CLS]}(u) \end{aligned} $$

[CLS]는 CLS토큰을 의미한다. Teacher, student 모델은 동일한 backbone을 사용한다. Student 모델은 위 식을 통해서 업데이트 되고 teacher 모델은 Student 모델의 파라미터를 이용해서 Exponential Moving Average로 업데이트 된다.

DINO 논문에서 Teacher에는 원본 이미지를 Global crop한 입력값을, student에는 local crop한 입력값을 사용한다.

Student는 이미지의 일부분만 보기 때문에 이를 마스킹된 이미지라고 해석할 수 있다고 생각한다. Student는 이미지의 일부만 보고 이미지 전체를 본 Teacher와 비슷한 CLS를 출력할 수 있게 학습되는 것이다.

Methodology

본 논문에서 제시하는 IBOT은 DALL-E Tokenzier와 같이 사전에 학습된 tokenizer를 사용하는 것이 아닌 self-distillation을 적용한 online tokenzier를 Target 모델과 함께 학습하는 MIM 모델을 제시합니다.

위 그림에서 볼 수 있듯이 DINO와 마찬가지로 같은 원본 이미지를 2개의 서로다른 Augmentation을 적용합니다($u,v$). $u,v$에 서로 다른 blockwise masking(BEiT 방법)을 적용한 $\hat{u}, \hat{v}$를 student의 입력값으로, $u,v$를 teacher의 입력값으로 사용합니다.

Student와 Teacher는 각각 head를 통해서 CLS토큰과 Patch토큰을 예측하게 됩니다. Student의 경우 마스킹을 적용 되었기 때문에 마스킹된 부분의 Patch토큰을 예측하게 되고 Teacher는 마스킹 없이 예측합니다. BEiT에서 DALL-E Tokenizer가 했던 역할을 그대로 Teacher가 수행합니다. Teacher와 Student의 Patch 예측값을 비교해서 첫번째 로스인 $L_{MIM}$이 나옵니다. 이는 아래와 같습니다.

$$ \begin{aligned} L_{MIM} = -\frac{1}{2}\sum_{i=1}^{N} m_{i} \cdot [ P_{{\theta}^{'}}^{patch}(u_{i}) \log P_{\theta}^{patch}(\hat{u}{i}) + P{{\theta}^{'}}^{patch}(v_{i}) \log P_{\theta}^{patch}(\hat{v}_{i})] \end{aligned} $$

IBOT에서는 tokenizer를 더 semantically-meaningful(본문 인용)하게 만들기 위해서 DINO에서처럼 CLS 토큰을 이용한 자가 증류를 합니다. 이를 수식으로 나타내면 아래와 같습니다.

$$ \begin{aligned} L_{CLS} = -\frac{1}{2} \sum_{i=1}^{N} P_{{\theta}^{'}}^{[CLS]}(u_{i}) \log P_{\theta}^{[CLS]}(\hat{v}{i}) + P{{\theta}^{'}}^{[CLS]}(v_{i}) \log P_{\theta}^{[CLS]}(\hat{u}_{i}) \end{aligned} $$

위 식은 2개의 패치이미지 쌍을 사용했다는 점을 제외하면 DINO와 완전히 동일합니다.

CLS 토큰과 Patch 토큰 예측에는 같은 head를 사용하며, BEiT에서 토큰 임베딩 차원으로 사용한 8192를 그대로 사용했습니다.

DINO에서는 6만 정도의 큰 차원으로 CLS토큰을 임베딩 시켰습니다.

또한 BEiT에서는 One-hot 인코딩된 토큰을 사용했지만 본 논문에서는 softmax값을 그대로 사용했습니다. 이는 이미지 패치가 문장의 단어와는 다르게 매우 애매한 semantic meaning을 가지고 있기 때문이라고 설명합니다.

DINO에서도 temparature가 적용된 Softmax값을 사용했습니다.

추가적으로 실험에서 사용한 setting은 DINO의 그것과 거의 유사합니다. ViT, SwinTransformer를 기본 모델로 사용했으며, CLS와 패치 projection head는 3개의 MLP를 사용했으며, 모델이 자명한 답을 내지 않게 하기 위해서 L2 정규화를 MLP블락의 bottleneck에 적용했습니다.

정리

1. DINO와 마찬가지로 2개의 서로다른 Augmentation이 적용된 이미지를 준비한다. -> $u,v$
2. 각 이미지에 blockwise 마스킹을 적용한다. -> $\hat{u}, \hat{v}$
3. Teacher에는 $u,v$를 Student에는 $\hat{u}, \hat{v}$를 입력값으로 사용한다.
4. Student가 예측한 마스킹된 패치의 값과 Teacher가 예측한 패치의 값을 크로스 엔트로피를 적용해서 Teacher -> Student로 지식을증류한다.
5. 서로 다른 이미지의 CLS 토큰값을 이용해서 Semantic meaning을 증류한다. $u \leftrightarrow \hat{v} , v \leftrightarrow \hat{u}$
6. CLS토큰과 패치 예측은 같은 head를 사용한다.

   Experiment/Result

Result

평가방식은 DINO에서 사용했던 방식인 KNN, Linear Probing, Fine-tuning 을 사용했습니다.

결과를 보면 알 수 있듯이 대부분의 경우 이전 SOTA였던 DINO, EsViT에 비해서 높은 성능을 내고 있습니다. 특히 ImageNet 1k 훈련을 시켰을 때 kNN성능이 뛰어난것을 확인할 수 있습니다.

추가적으로 Semi-supervised learning, Unsupervised learning, Object detection, Semantic segmentation 등에서도 기존의 자가지도 학습방식 모델들보다 높은 성능을 보이고 있습니다.

Pattern layout

본 논문에서는 자가 증류를 통해서 MIM이 무엇을 학습하는지 밝히기 위해서 ImageNet 1k validation set에서 가장 높은 신뢰도(??)를 가진 패치들을 시각화 했습니다. 위 이미지에서 확인 할 수 있듯이 자동차의 헤드라이트, 동물의 귀, 사람의 머리처럼 특정 군집으로 묶을 수 있었습니다. 논문에서는 MIM을 통해서 high-level semantic meaning과 low-level detail을 모두 포착했다고 설명했습니다.

MIM & Recognition

본 논문에서는 MIM이 이미지 인식을 어떻게 향상시키는지 밝히고자 2가지 실험을 했습니다. 첫번째는 CLS 토큰을 이용한 linear classfication이고, 두번째는 k 개의 높은 self-attention을 가지는 패치들을 선택해서 평균내서 linear classification을 하는 것입니다.

첫번째 실험에서는 DINO와 IBOT의 성능이 77.9% v.s. 77.0%로 DINO가 소폭 앞섰습니다. 하지만 두번째 실험 즉, 패치 representation을 직접적으로 사용했을 경우에는 IBOT이 더 높은 성능이 나왔습니다. 아래 그림 참고

논문에서는 IBOT이 MIM을 통해서 각 패치들이 더 풍부한 Semantic information을 가지고 있기 때문에 가능하다고 설명했습니다.

Conclusion

본 논문에서 제시하는 IBOT은 이미지에서 BERT 방식의 학습을 하기위해서 필요했던 Image Tokenzier를 사전에 학습하지 않고 자가 증류를 이용해서 online으로 학습하는 방법을 제시했습니다. BEiT에서 부족하다고 평가받았던 pre-trained tokenizer 사용을 DINO의 자가 증류를 통해서 해결했다는 기여점이 있는 것 같습니다.

저는 NLP에 대한 지식이 거의 없기 때문에 본 논문에서 나오는 NLP에서 자주 사용하는 단어의 뜻을 잘 못 이해하고 있을 수 있습니다. 따끔한 지적해주시면 감사하겠습니다.

Abstract/Intro

본 논문은 NLP분야에서 가장 큰 영향력을 미쳤다고 평가받는 BERT의 훈련방식인 Masked language modeling을 이미지 모델에 적용했습니다. 기존의 이미지 모델들이 BERT의 훈련방식을 채택할 수 없었던 이유는 크게 두가지라고 밝힙니다.

1.NLP와 다르게 이미지는 사전에 정의된 사전이 존재하지 않아서 token화 시키기 어렵다 2.토큰을 쓰지 않고 픽셀 수준의 예측을 할 순 있지만 그렇게 되면 모델이 사전학습에서 short-range dependencies와 high-frequencey details만 학습할 수 있다.

본 논문에서 밝히진 않았지만 추가적으로 NLP와 Vision에서 주로 사용하는 기본 모델이 Transformer VS CNN 으로 달랐기 때문에 BERT와 같은 방식으로 학습하기 어려웠을 것으로 생각하고 있습니다.

본 논문에서는 OpenAI에서 발표한 논문에서 공개된 사전학습된 DALL-E tokenizer를 이용해서 이미지 패치들을 tokenizer 해서 위 두 문제들을 해결하고 BERT방식으로 모델을 훈련했습니다. 결과는 본 논문 이전까지 나왔던 대부분의 지도학습 및 비지도학습 기반 ViT 관련 논문에 비해서 성능이 좋았습니다.

BERT에서 영감을 받았다는 사실을 논문에서도 여러번 언급하고 이와 같은 학습 방식도 BERT의 Masked Language Modeling(MLM)의 이름을 모방하여 Masked Image Modeling(MIM)이라고 지었습니다.

위 설명에서도 알 수 있듯이 Tokenizer는 DALL-E에서 공개한것을 사용했고, 훈련방식은 BERT와 완전히 동일합니다. 전에 없던 새로운 방법론을 제시한건 아니지만 Masked Image를 사용해서 ViT를 사전훈련했을 때 높은 성능이 나온다는 것을 확인할 수 있다는 기여점이 있는 것 같습니다.

Method

Backbone/Image Patch

BEIT는 이미지를 ViT의 논문에서 제시했던 16x16패치로 이미지를 일정하게 나누는 방식을 사용했습니다. 개별 패치들은 flatten과 linear projection 과정을 거치고 Position embedding이 더해져서 최종 입력값이 됩니다. Backbone network도 ViT논문에서 제시했던 ViT모델을 변경없이 그대로 사용했습니다.

Visual Token

본 논문에서는 이미지 모델을 훈련하는데 BERT의 학습 방식인 Masked modeling을 적용했습니다. BERT는 문장의 단어들을 token화 하고 이중 일부를 masking해서 지웠습니다. Masking되지 않은 토큰과 masking된 토큰을 모델의 입력값으로 사용하고 masking된 토큰을 예측하면서 문장을 전체적으로 이해하는 학습방식 입니다.

BEIT에서 첫번째로 지적한 문제점이 여기서 발생합니다. NLP에서는 사전에 정의된 tokenizer를 이용해서 단어들을 토큰으로 만들 수 있지만 이미지의 경우 사전에 정의된 tokenizer가 없었기 때문에 이미지를 일정하게 토큰으로 만들 수 없었는데 OpenAI에서 발표한 DALL-E 논문에서 발표한 이미지 tokenizer를 사용하면서 이 문제를 해결했습니다.

위 그림에서 빨간 박스안에 있는 부분이 DALL-E 에서 공개한 Tokenizer를 사용한 부분입니다. DALL-E의 Tokenizer를 이용해서 이미지를 discrete한 visual token으로 만들 수 있습니다.

DALL-E의 Tokenizer는 discrete Variational AutoEncoder 방식으로 학습되었는데 자세한 내용은 Paper에서 확인하실 수 있습니다.

Pre-Training

DALL-E Tokenizer를 이용해서 이미지를 discrete token들로 변환하고 image patch들 중에서 40%에 masking을 적용합니다. Masking된 패치들은 학습 가능한 임베딩으로 교체한뒤 masking되지 않은 패치와 함께 ViT의 입력값으로 사용됩니다.

ViT모델의 최종 히든 벡터들은 입력값의 encoded representation으로 취급하며, softmax classifier(MIM Head)를 통과해서 최종출력을 하게 됩니다.

masking되었던 패치들의 최종출력이 DALL-E를 통해서 미리 생성되었던 discrete token을 예측하도록 모델을 훈련합니다(log-likelihood사용).

위 내용만 놓고 보면 자료형태가 이미지인것을 제외하면 BERT와 거의 동일하다.

주목할만한 점은 masking할 패치들을 임의 선택한것이 아닌 blockwise하게 했다. 논문에서 밝힌 방법은 랜덤하게 block의 사이즈를 정하고 block의 위치를 정하는 과정을 masked patch가 전체의 40%가 될때까지 반복했다.

위 과정에서 이미지의 라벨을 전혀 사용하지 않았기 때문에 사전학습 과정은 온전히 자가지도학습 방법론이라고 주장한다.

Fine-Tuing

사전학습이 끝난 BEIT는 downstream task로 fine-tuning되어야 하는데 그중에서 Image classification Global Average Pooling을 이용해서 최종 출력층의 벡터를 모으고 linear classifier를 통과시키게 됩니다.

Experiment

결과를 보면 똑같이 자가지도학습 방식을 채택했던 DINO에 비해서 더 좋은 결과가 나온것을 확인할 수 있습니다. 또한 JFT300M이라는 구글 내부의 거대한 데이터셋으로 지도학습시킨 모델보다도 ImageNet1k 만 사용한 BEIT가 높은 성능을 내주었습니다.

하지만 저는 BEIT도 사전 학습된 DALL-E를 사용했기 때문에 실제로는 ImageNet 1k 보다 많은 데이터를 사용했다고 보고 있습니다.

이외의 결과는 다른 많은 논문들과 마찬가지로 대부분의 downstream task에서 높은 성능이 나왔다고 주장합니다.

한가지 주목할만한 점은 BEIT도 DINO모델과 마찬가지로 지도 학습을 하지 않았음에도 attention map을 시각화하면 이미지에 존재하는 오브젝트가 활성화됩니다. ViT모델을 자가지도학습 방식으로 학습시키면 나타나는 하나의 특징인 것 같습니다.

결론/생각

본 논문은 DALL-E에서 공개한 이미지 tokenizer를 이용해서 BERT방식의 discrete token을 예측하는 masked image modeling을 제시했습니다.

사실 tokenizer는 DALL-E에서 공개했고, Backbone은 ViT를 그대로 사용했고, 학습방식은 BERT의 그것을 그대로 사용했습니다. 새로운 방법론을 제시했다고 보기는 어려울 것 같지만 ViT를 masking patch prediction 방식으로 학습하면 어떤 결과가 나오는지 보여준 논문이라고 생각합니다.

추가적으로 이후에 나온 Masked Autoencoders are scalable vision learner는 BEIT에서 좋지 않다고 주장한 pixel 수준의 재구성을 통해서 BEIT보다 높은 성능을 내주었습니다. 이 논문도 읽어보면 많은 도움이 될것 같습니다.

Abstract

본 논문은 두가지 간단한 디자인을 통해서 ImageNet-1K 데이터만을 이용한 학습에서 높은 정확도(87.8%)를 달성했습니다. 첫번째는 비대칭 Encdoer-Decoder구조이며 두번째는 높은 비율(75%)로 input 이미지를 마스킹하는 것입니다. 비대칭적 구조로 Decoder가 Encoder에 비해서 상대적으로 작은 scale이며, input의 대부분을 마스킹 했기 때문에 적은 메모리 사용량으로 학습을 가속할 수 있으며, scale을 변경하기 쉬워서 높은 capacity의 모델을 학습할 수 있습니다. 자세한 내용은 방법론에서 다루겠습니다.

Intro

NLP에서 가장 성공적인 모델을 꼽으라면 단연 GPT와 BERT입니다. 두 모델은 다른 구조를 가지고 있지만 본 논문에서는 두 모델이 데이터의 일부를 지우고 지워진 부분을 예측하기 위해서 학습하는 같은 개념을 가지고 있다고 하며 이를 masked auto encoding이라고 합니다. 또한 이렇게 데이터의 일부분을 지우고 예측하는 것을 Denoising Autoencoder와 유사하다고 합니다.

하지만 BERT의 학습 방법을 Computer Vision(이하 CV)에 적용했을 때 그 성능이 주목할 만큼 좋지는 않았습니다. 여기서 논문이 해결하고자하는 질문이 나옵니다.

What makes masked autoencoding different between vision and language?

본 논문에서는 총 3가지의 이유로 차이가 생긴다고 했습니다.

1. NLP와 CV에서 주로 사용하는 모델의 구조(Attention VS CNN)가 근본적으로 다르다.
2. Text가 Image에 비해서 정보의 밀도가 높다.
3. Autoencoder의 디코더가 Text와 Image에서 수행하는 역할이 다르다.

1번의 경우 ViT 모델의 등장 이후로 CV의 많은 분야에서 Transformer 구조를 채택했기 때문에 NLP와 CV간의 모델 구조의 차이점이 많이 사라졌습니다. 하지만 2번과 3번은 Text와 Image의 데이터 특성이 다르기 때문에 발생합니다. Text의 경우 단어 하나하나가 모두 Semantic 정보를 가지고 있지만 Image의 경우 pixel하나가 가지고 있는 semantic 정보는 많지 않습니다.

예를 들어서 Text에서 단어 하나를 지우고 이를 예측하려면 문장 전체의 맥락과 문장을 이루고 있는 단어들의 의미를 잘 이해해야 하지만 Image에서 pixel하나를 지우고 예측한다고 하면 단순히 근처 픽셀들의 평균값을 넣어도 훌륭한 예측이 됩니다. -> 이미지는 공간적 중복성이 높은 정보를 가지고 있다.

본 논문에서는 위 문제점들을 해결하기 위해서 2가지 방법을 제시합니다.

1. 높은 비율(75% 이상)으로 mask를 적용해서 패치들간의 공간적 중복성을 제거한다.
2. 비대칭적 구조를 가진 encoder-decoder 구조를 통해서 계산 효율성을 높여서 scale을 크게 키운다.

Masked AutoEncoder(이하 MAE)는 입력 이미지를 높은 비율로 mask 하기 때문에 실제 계산에서는 이미지의 일부(25%)만 사용해서 비교적 큰 모델을 특징 추출기로 사용해도 메모리 사용량이 높지 않다. 또한 인코더에 비해서 작은 디코더를 사용해서 계산량을 줄였다.

결과적으로 MAE는 ViT-Large/Huge 모델을 ImageNet-1K 데이터만 사용해서 87.8% 라는 높은 성능을 보여주였고 다른 task로의 전이학습에서도 SOTA에 준하는 성능을 보여주어서 일반화도 잘되는 좋은 이미지 분류 모델임을 증명했습니다. 본 논문의 큰 기여점으로 자가 지도학습(ImageNet 1K 데이터사용)을 통해서 지도학습보다 높은 성능을 내주었으며, 모델의 scale을 키울수록 높은 성능이 나오는 것을 확인해서 NLP에서의 자기 지도학습과 결이 같은 모습을 보여주었다.

Methodology

MAE 는 기존의 autoencoder와 마찬가지로 인코더가 입력 이미지를 latent representation으로 만들고 디코더가 이를 이용해서 입력 이미지를 생성합니다. 하지만 비대칭적 구조를 채택해서 인코더가 입력 이미지의 일부분(mask된 부분을 제외한 이미지)만 latent representation으로 만들고 인코더에 비해서 경량화된 디코더를 사용해서 입력 이미지를 재구성합니다.

Maksing

많은 Vision Transformer 모델이 그러하듯 겹치지 않게 격자로 이미지를 나누고 Linear projection하는 Patch Embedding으로 입력 이미지를 패치로 만들어줍니다. 이후 패치의 집합에서 부분집합을 샘플링 해주고 선택되지 않은 패치들은 mask해줍니다. 샘플링 방식은 균등 분포를 따르는데 논문에서는 이를 random sampling(이하 랜덤 샘플링) 이라고 합니다. 높은 비율로 masking하는 랜덤 샘플링은 중복성(redundancy)를 제거해서 인접한 패치를 보고 문제를 쉽게 해결할 수 없게 만듭니다. 균등 분포 이외의 분포를 사용하게 되면 중앙 부분만 높게 샘플링 되는 등의 bias가 생길 수 있기 때문에 균등 분포를 사용했다고 밝혔습니다.

MAE encoder

ViT를 인코더로 사용했으며 위에서 설명했듯이 입력 이미지는 Patch Embedding을 거쳐서 패치로 변환되며 이 패치들중 일부만 선택되어 인코더의 입력값으로 사용됩니다. Mask된 패치들은 일련의 과정에서 사용되지 않기 때문에 제거해서 메모리 효율성을 높였습니다.

MAE decoder

디코더는 인코더와 마찬가지로 Transformer block으로 이루어져 있지만 인코더 보다는 가벼운 구조를 채택합니다. 위 그림과 같이 디코더는 인코딩된 패치들과 mask된 패치 모두를 입력값으로 받게 됩니다. 이 때 mask된 패치에도 positional embedding을 더해주게 됩니다. 이렇게 하지 않으면 mask된 패치가 이미지에서 어느 위치에 있는지 나타내는 정보가 없기 때문에 디코더가 작동하지 않게 됩니다.

디코더는 pre-training 단계에서 입력 이미지를 latent representation으로 부터 복원하기 위해서 사용되며 실제로 이미지 분류에는 인코더의 latent representation만 사용됩니다. 따라서 디코더는 인코더의 구조와는 별개로 사용이 가능하다는 것을 의미합니다. 본 논문에서는 굉장히 가볍고 얕은 구조를 사용했습니다.

인코더의 경우 입력 이미지의 25%만 연산하기 때문에 계산량이 적지만 디코더의 경우 모든 패치에 대해서 연산을 하기 때문에 인코더와 같은 구조(깊이와 넓이)를 채택하게 되면 계산량이 한번에 너무 커지게 됩니다.

Reconstruction target

MAE는 mask된 패치의 pixel값을 예측합니다. 따라서 디코더의 출력값은 입력 이미지의 크기와 형태로 reshape 됩니다. 또한 BERT와 마찬가지로 mask되었던 패치들의 예측에 대해서만 MSE Loss 를 적용합니다.

저자들은 원본 이미지를 그대로 복원하는 것 보다는 mean, std를 이용해서 정규화된 입력값을 복원하는 것이 성능 향상에 기여한다고 밝혔습니다.

Experiments

MAE의 기본 구조자체가 워낙 간단해서 방법론에 대한 설명은 길지 않았습니다. 본 논문에서도 1페이지 정도로 모든 설명을 마쳤습니다. 본 논문에서 제시하는 방법론은 CV에서 기존의 분류 모델을 학습시키는 그것보다는 NLP와 비슷하기 때문에 다양한 실험을 통해서 제시하는 방법론을 구체화 했습니다.

실험에 사용한 모델은 ViT-Large 16이며 비교에 사용한 방법론은 ViT논문에서 제시했던 scratch 방식이 아닌 DeiT입니다. 제가 중요하다고 생각하는 실험의 결과와 주관적 해석은 아래와 같습니다.

여기부터는 부족한 저의 의견이 많이 있습니다. 적당히 걸러서 읽어주시길 바랍니다.

Masking ratio

본 논문에서는 Masking 비율을 찾기위해서 1090%까지 비율을 바꾸면서 실험을 했습니다. 위 그래프를 보면 4080% 구간에서 높은 성능이 나오는 것을 확인할 수 있으며, 저자들은 높은 비율의 masking이 spatial redundancy를 줄여서 더 높은 일반화 성능을 얻을 수 있다고 기술했습니다.

하지만 실제 10% mask의 결과도 DeiT에서 제시하는 82.5% 보다 높은것을 확인할 수 있습니다. Linear probing 방식에서는 mask 비율 10~75% 구간에서 성능이 선형적으로 상승하지만 fine-tuning에서는 이미 10%에서도 충분히 높은 수치입니다. 따라서 높은 masking 비율보다 더 크리티컬하게 작용하는 무언가가 있다는 사실을 내포하고 있을 수 있다고 생각합니다.

인코더에 비해서 굉장히 작은 디코더를 사용했기 때문에 mask 비율이 높을 때 좋은 성능이 나온것은 아닌가 생각합니다. 인코더와 비슷한 크기의 디코더를 사용하고 mask비율을 조정하면서 실험을 해보면 굉장히 큰 컴퓨팅 파워가 필요하겠지만 어떤 결과가 나올지 궁금합니다.

Decoder design

디코더의 디자인에 따라서 Fine-tuning과 Linear probing의 성능이 차이를 보입니다. 위 표에서 알 수 있듯이 디코더의 깊이가 깊어질수록 linear probing의 성능은 올라가지만 Fine-tuning의 성능은 큰 차이를 보이지 않습니다. 극단적으로 디코더의 깊이가 1일때와 8일때의 fine-tuning 성능도 차이를 보이지 않습니다.

본 논문에서는 이렇게 차이가 발생하는 이유를 인코더의 마지막 부분과 디코더가 recognition보다 reconstruction에 집중하기 때문이라고 해석했습니다. 따라서 디코더가 깊을수록 latent representation이 더 추상적인 수준(논문에서는 abstract level이라고 표현)으로 올라가기 때문에 linear probing에서 높은 성능이 나오지만, fine-tuning시에는 인코더의 마지막 부분이 recognition task에 더 집중하도록 조정되기 때문에 성능에 큰 차이가 없다고 합니다.

MAE가 만들어내는 latent representation을 이용해서 contrasive learning이나 DINO에서 적용했던 distillation을 이용한 학습을 적용해 보면 어떤 결과가 나올지 궁금해지기도 합니다.

깊이가 큰 차이를 만들어내지 않듯이 dimenstion 또한 128과 1024가 큰 차이를 만들지 않기 때문에(오히려 감소) 512라는 상대적으로 작은 차원을 디코더에 적용한것을 알 수 있습니다.

기타

MAE는 위에서 설명했다 싶이 mask토큰을 제거하고 입력값으로 mask되지 않은 토큰만을 이용합니다. mask 토큰을 입력값으로 넣으면 오히려 성능을 떨어지고(linear probing -14%) 계산량만 늘어났다고 합니다.

패치별로 정규화된 값을 예측하게 하면 성능이 더 올라간다고 합니다. 이러한 방법이 지역적인 contrast를 올린다고 해석했습니다. 또한 PCA를 패치별로 적용해보았는데 이 경우에는 성능이 떨어진다고 합니다.

BEiT에서 적용되었던 token을 예측하는 방식을 적용했다고 합니다. MAE는 pixel을 예측하는 방식인데 DALLE pre-trained dVAE를 토크나이저로 이용해서 token 예측방식을 적용하니 fine-tunig에서는 성능이 아주 약간 올랐지만 linear probing에서는 성능이 떨어졌다고 합니다. 더 직관적인 pixel 방식이 났다고 주장합니다.

MAE는 데이터 증강기법을 적용하지 않는게 더 높은 성능이 나온다고 합니다. 기존 contrastive learning등의 연구에서는 다양한 증강기법을 적용했는데 비해서 특이한 결과인것 같습니다. 높은 비율로 mask를 적용한다는 것 자체가 강력한 데이터 증강기법이 사용된것이며 랜덤 샘플링은 이터레이션마다 다른 패치를 mask하기 때문에 추가적인 증강기법이 유의미하지 않다고 합니다.

결론/생각

NLP 에서 주로 사용되고 성공적인 결과를 가져왔던 Masked token prediction 방식을 이미지에 적용하기 위해서 Autoencoder 방식을 채택했다. 결과로 비교적 간단한 구조에 적은 자원을 사용하고도 기존의 지도학습 방식보다 높은 성능을 얻을 수 있었다. 본 논문에서 가장 중요하다고 생각하는 점은 높은 mask 비율을 적용해서 spatial redundancy를 줄여서 모델이 각 patch 들의 semantic 정보를 학습하게 한 것이라고 생각합니다. FAIR에서 나왔던 DINO논문에서도 small crop과 large crop을 이용해서 이미지의 일부만을 보게하는 학습방식을 사용했었는데 본 논문에서 적용했던 방식과 비교하면 MAE의 방식이 더 직관적이어서 이해하기 쉬웠던 것 같습니다. 더 자세한 내용은 코드리뷰에서 다뤄보겠습니다.

year : 2021 Code : https://github.com/microsoft/Swin-Transformer

Method

Overall Architecture

기존의 ViT와 동일하게 patch embedding으로 input을 만들어줍니다. 다른점이 있다면 패치 사이즈가 4x4라는 것입니다. ViT에서는 16x16의 패치가 디폴트값이었습니다. 패치사이즈가 작아지면 전체 패치의 갯수가 제곱으로 늘어나고 패치의 임베딩 차원이 동일하면 연산량도 그에 상응하게 늘어납니다. 따라서 4x4 사이즈는 상당히 작은 크기임을 알 수 있습니다.

아래는 과거에 만들어 두었던 발표자료에서 가져온 패치 임베딩(패치사이즈 16) 예시입니다. 만약 4x4 패치가 된다면 16x16에 비해서 패치수가 16배 많으므로 리니어 임베딩을 지나면 64 x C(임의의 차원)이 됩니다.

Swin transformer는 여러개의 Stage로 구성되어 있으며 각 Stage에서는 패치의 갯수가 변하지 않습니다. 따라서 Stage 1에서는 전체 토큰의 갯수 $\frac{H}{4} \times \frac{W}{4}$가 유지됩니다. 논문의 제목처럼 계층적 표현을 위해서 Stage가 진행됨에 따라서 패치는 합쳐지고 이 과정에서 패치의 갯수는 줄어들게 됩니다.

Stage 2에서는 Stage 1에 있던 패치를 합치는 Patch merging layer를 통과합니다. Patch mergin layer는 인전한 2x2패치를 concat하고 linear layer를 통해서 채널을 절반으로 줄입니다. 이를 통해서 패치의 갯수를 1/4로 줄일 수 있습니다.

간단하게 그림을 표현해보았는데 실제 연산에서는 아래 그림과 같이 2차원 벡터의 연산이 아닌 1차원으로 flatten된 벡터로 연산됩니다.

위와 같은 연산은 Stage를 지날때마다 동일하게 적용되어서 Stage 4까지 이어집니다. Stage 1에서는 4x4크기의 패치를 이용해서 Patch Embedding을 했기때문에 패치 갯수가 $\frac{H}{4} \times \frac{W}{4}$ 이지만, 위와 같이 4개의 패치를 하나로 묶었기 때문에 패치 갯수는 $\frac{H}{8} \times \frac{W}{8}$ 로 1/4이 됩니다. 이를 통해서 논문에서 제시하는 아래와 같은 계층적 구조가 완성됩니다.

위 그림과 같이 총 4개의 Stage 를 지나면서 한번 지날때마다 패치의 크기가 4배씩 늘어나고 패치의 갯수는 1/4로 줄어들게 됩니다. 본 논문에서는 Swin Transformer Block을 제시하였는데 기존의 ViT에서 사용했던 MSA(Multi Head Self Attention)을 대체하는 W-MSA(Window based MSA), SW-MSA(Shifted Window MSA)로 이루어져있습니다. 이외의 구조는 기존 ViT와 동일하게 Pre-Norm구조에 2개의 Linear layer와 GELU를 사용하는 MLP block이 Attention 연산 다음에 적용되고 각 연산마다 residual connection이 존재합니다. 

Shifted Window based Self-Attention

본 논문의 가장 큰 기여점이라고 생각하는 Shifted Window 입니다. Swin transformer는 4x4라는 아주 작은 Patch를 사용합니다. ViT의 기본 모델이 사용한 16x16에 비교하면 $(\frac{16}{4} \times \frac{16}{4})^2$배 많은 연산이 필요합니다. 이를 해결하기 위해서 겹치지 않는(non overlapping) MxM(default 7)개의 패치를 하나의 Window로 설정하고 Window안에서만 Attention 연산을 수행합니다.

여담이지만 처음에 7이라는 숫자를 보고 왜 굳이 저 숫자로 설정했는지 의문이었습니다. 아마 pre-train과정에서 가장 많이 사용하는 이미지 해상도인 224(32x7 -> (2^5)x7)가 7의 배수이기 때문으로 생각합니다. Stage를 지나면서 패치의 갯수는 (8x7) x (8x7) -> (4x7) x (4x8) -> (2x7) x (2x7) -> (1x7) x (1x7) 순으로 줄어드는데 계속해서 7의 배수이기 때문에 Window size를 7x7로 설정하면 계속해서 나누어 떨어집니다.

위 그림에서 왼쪽 예시가 4x4 Window 입니다.

기존의 Vanilla Transformer와 ViT가 채택한 Global Attention은 패치수에 Quadratic하게 연산량이 증가하지만 Window based Attention의 연산량은 패치의 수에 Linear하게 증가합니다. 이를 확인하기 위해서 임의로 Attention 연산을 3단계로 나누어서 행렬곱의 형태를 살펴보면 아래와 같습니다. 편의를 위해서 scale factor는 제외했으며, $N_{patch}$는 $\frac{H}{patch_size} \times \frac{W}{patch_size}$입니다.

-Input -> Q,K,V : $(N_{patch} \times C) * (C \times C) \rightarrow (N_{patch} \times C)$ -$Q * K^{T} * V$ : $(N_{patch} \times C) * (C \times N_{patch}) * (N_{patch} \times C) \rightarrow (N_{patch} \times C)$ -Attention -> Output : $(N_{patch} \times C) * (C \times C) \rightarrow (N_{patch} \times C)$

위 연산들의 연산량은 결국 행렬곱의 연산량을 구하면 알 수 있습니다. 간단한 예시로 Naive하게 행렬곱을 했다고 가정하면, (MxN)*(NxP)행렬곱의 연산량은 MNP 입니다. 이를 사용해서 위 연산들의 연산량을 구하면 아래와 같습니다.

-Input -> Q,K,V : $N_{patch}C^2$ -$Q * K^{T} * V$ : $2N_{patch}^{2}C$ -Attention -> Output : $N_{patch}C^2$

따라서 총 연산량은 $4N_{patch}C^2 +2 N_{patch}^{2}C$로 패치갯수의 제곱해서 연산량이 증가합니다.

하지만 본 논문에서는 Window 내에서만 Attention연산을 수행하기 때문에 총 연산량은 $4N_{patch}C^2 +2M^{2}hwC$로 $hw$에 대해서는 선형적으로 연상량이 증가합니다. 이를 통해서 작은 패치를 사용함에도 불구하고 적은 연산량을 유지할 수 있습니다.

Window based Attention은 분명 연산량을 선형적으로 줄임으로서 작은패치를 사용함에 있어서 부담을 덜 수 있지만 연산이 줄었다는건 그만큼 잃는 것도 있습니다. Global Attention과 같은 경우에 모든 패치들이 상호작용을 하지만 Window based Attention은 같은 Window 안에서만 상호작용이 이루어지기 때문에 멀리 떨어진 위치의 패치들과 상호작용이 불가능합니다.

Cyclic Shift

본 논문에서는 Window간의 연결을 추가하면서 동시에 적은 연산량을 유지하기 위해서 Shifted window partiotioning을 도입합니다. 위 그림에서 오른쪽이 Shifted 된 window의 예시입니다. $[\frac{M}{2}, \frac{M}{2}]$만큼 대각선으로 Window가 Shift된것을 볼 수 있습니다. 이런식으로 Window를 옮기게 되면 기존의 Window에서 연결성이 없었던 패치들간의 연결성이 생기지만 가장자리에는 원래 Window 크기보다 작은 Window들이 여러개 생기게 됩니다.

만약 패딩을 통해서 작은 window들을 크기를 원래 window 사이즈와 맞춰주면 전체 Window의 갯수가 늘어나기 때문에 계산 효율성이 떨어지게 됩니다.

따라서 위 그림과같이 짜투리(?) Window들을 모아서 기존과 동일한 크기의 Window를 만들고 서로 다른 짜투리 Window에서 온 패치들끼리는 Attention 연산이 수행되지 않도록 mask를 적용합니다. 동일한 말이지만 Window가 shift되었다기 보다는 패치들이 옮겨졌다는 것이 더 적절할것 같습니다. Cyclic shift에 대해서는 이후 코드 리뷰에서 더 자세히 다루겠습니다.

Relative position bias

본 논문에서는 Attention 연산에 아래와 같이 relative position bias 를 추가했습니다.

$$ Attention(Q,K,V) = \text{SoftMax}(QK^T / \sqrt{d} +B)V $$

M개의 패치를 가지는 window 내에서 상대적인 패치들간의 거리는 [-M+1, M-1]의 범위에 있습니다.

1번째 패치를 기준으로 M번째 패치의 거리는 M-1 이고 M번째 패치를 기준으로 1번째 패치는 -M+1입니다. 절대 거리는 같지만 방향이 다르므로 상대 거리를 사용합니다.

$B$는 $\hat{B} \in R^{(2M-1)\times(2M-1)}$에서 선택해서 사용합니다. 저자들은 이와 같은 상대적 위치를 사용하는 것이 성능향상에 큰 도움이 되었다고 밝혔습니다.