(논문리뷰) Masked Autoencoders Are Scalable Vision Learners[CVPR2022]

https://arxiv.org/abs/2111.06377

Masked Autoencoders Are Scalable Vision Learners

 

ConvNeXt v2 논문을 읽기까지 긴 여정을 시작하려고 한다. (마지막 CV 팔로업 ViT... 갈 길 구만리,ㅜㅜ)

 

일단 MAE는 ResNet의 창시자인 Kaiming He] 가 1저자인 논문으로 language에서 GPT, BERT가 가져온 혁신을 image에서도 가져올 수 있다는 가능성을 제기한 논문이라고 볼 수 있을 것 같다. 

 

이 논문에서 알아야할 핵심 내용은 아래와 같다. 

 

• - maksed token을 제외하고 입력받는 encoder와 latent representation과 maksed token을 입력받아 복원하는 경량화된 decoder로 구성된 모델을 제안하였다.
  •  " 비대칭 EncDec 구조 "
• 입력 이미지의 75%를 즉, 높은 비율을 making하고 복원하는 것이 의미있는 결과를 가져왔다. 
  • 모델의 학습 시간과 메모리 사용량을 크게 줄일 수 있어, MAE를 대형 모델로 쉽게 확장가능
• downstream task의 transfer performance와 high-capacity model에서 좋은 성능을 가져왔다. 
  • ImageNet-1K 데이터만으로 ViT Huge 모델을 이용해서 SOTA달성. 
  • NLP에서 Self Supervised Learning에서의 혁신을 보여줬던 것처럼 Vision에서도 가능성을 제기.
    • Supervised learning은 라벨과 관련이 있는 내용만을 학습하는데 Self Supervised Learning은 데이터의 본질을 학습할 수 있게된다.

 

이제 introduction부터 살펴보자. 

 

Introduction

BERT가 masking된 token을 복원하는 과정에서 language의 본질적인 의미를 학습할 수 있게 되었다. 하지만 이미지의 경우는 이러한 것이 어려웠는데, 이를 본 논문에서 3가지 이유로 language와 Vision의 차이점을 설명한다. 

 

1) ViT가 나오기 이전까지 이미지는 convolution기반의 연산이 주를 이뤘다. 이것은 masked token이나 positional embedding을 통합하기에는 어려운 연산방법이였다. 

 

2) Language랑 Vision은 정보의 밀도가 다르다. 

Language는 사람이 만들었기에 정보가 밀집되어있고 의미가 함축되어있는 것이다. 그래서 각 토큰마다의 하지만 이미지는 픽셀로 나타내고 그 픽셀마다의 의미정보가 크다기 보단 인접한 픽셀과의 관계를 통해 의미를 전달하기 때문에 parts, objects, scenes에 대한 이해가 조금이라도 있으면 인접한 픽셀을 통해 바로 복원을 할 수 있게 된다. BERT에서 사용했던 masking 방법론을 적용하기 위해 그래서 높은 비율의 random한 patch를 masking하는 전략을 선택했다. 

 

3)Vision에서는 decoder를 통해 pixel단위를 복원한다. 즉, 낮은 semantic 정보를 가지고 있다. 하지만 Language의 decoder에서는 missing word를 복원하고 그 자체로 높은 semantic 정보를 갖고 있다. BERT에서는 decoder가 단지 몇개의 MLP로 연결되어있어 그 의미가 사소하지만 Vision의 경우 Decoder단에서 학습된 latent representation을 통해 semantic level로 결정해야하므로 중요하다. 

 

결국 MAE 모델은 encoder에서는 가려지지 않은 패치만 처리하고, decoder에서는 가려진 패치를 포함해 입력 이미지를 복원하는 역할을 하는 것이다. 

 

Approach

먼저 이미지를 일정한 크기의 패치로 나눈 후, 무작위로 일부 패치(약 25%)를 선택해 나머지를 마스킹한다. 

이렇게 높은 비율로 패치를 제거하면 중복을 줄이고, 주변 패치로 추측하기 어려운 과제를 만든다. (마치 BERT가 TOKEN 15%를 없앤 것과 같은 느낌.)

 

MAE의 encoder는 ViT 기반으로 동작하지만, 마스킹되지 않은 패치만 처리하여 메모리와 연산 자원을 줄일 수 있었다. 

decoder는 masked patch와 latent representation을 모두를 처리해 이미지를 복원한다.

 

MAE는 복원 과정에서 마스킹된 패치의 픽셀 값(정규화된 픽셀 값)을 예측하는데, 이 과정에서 평균 제곱 오차(MSE) 손실 함수를 사용한다. 

 

Masking

 

위 figure 5와 figure 6을 관찰해보면 random하게 75%를 masking한 것이 finetuning과 linear probing기법 모두에서 성능이 좋은 것을 관찰했다. 

 

* linear probing: encoder 단을 사전에 학습하고 head에 MLP를 붙여서 이 부분만 학습하는 방법. 

 

위 표는 다양한 ablation study의 결과를 나타낸 것으로 하나씩 살펴보면, 

(a)와 (b)에서는 decoder가 lightweight해도 비교적 높은 성능을 보인다는 것을 보여주고 있다. (한토큰이 encoder를 거치는 것보다 decoder를 거치는 것이 10%의 burden도 안됨)

 

(c)에서는 encoder단에서 maked token을 making해서 함께 넣는 것보다 아예 없애는게 FLOPs 관점에서도 성능 관점에서도 더 좋은 것을 확인했다. 

 

(d)를 보면 dVAE token case는 DeiT의 경우 달리모델에서 toknizer를 활용해서 token을 예측하는 방법과 같은 것이다. 

즉, token을 patch별로 PCA를 진행해서 target으로 쓸 때와 toknizing을 통해 token을 target으로 썼을 때보다 원본 pixel을 정규화했을 때가 가장 좋은 성능을 보여주었다. 

MAE pre-training과 supervised pre-training을 비교해보았을 때, IN1K( ImageNet-1K(224size) )로 ViT를 사전학습했을 때는 정확도가 포화상태에 빠진 것을 확인했지만 MAE로 사전학습을 진행한 경우는 parameter가 증가할수록 좋은 성능을 나타내는 것을 확인했다.  

 

MOCO v3는 transformer기반의 contrastive learning을 한 모델인데 x축은 finetuning을 한 block개수로 0개면 linear probing이고 나머지는 점진적으로 block을 늘려간 것이다. 위 결과를 보면 linear probing만을 적용했을 떄는 MAE가 성능이 낮고 block을 하나만 finetuning해줘도 바로 성능이 우세한 것을 확인할 수 있다. 

 

이를 통해 linear하게 representation을 표현하는 것은 불충분하다.

 

하나의 nonlinear단만 추가해주더라도 바로 MAE의 성능이 보장되므로 MAE의 표현이 선형적으로 구분되기 어려운 대신, 비선형적인 특징에서 강력함을 보여줌. 

 

마지막 target을 pixel로 두는 것과 token으로 두는 것의 차이를 보면 미세하다고 생각할 수 있지만 token을 출력하기 위해 dalle를 학습하고 toknier를 사용하는 것도 연산과 메모리가 필요하기 때문에 결과적으로 pixel을 target으로 두는 것이 더 낫다.