[Paper Review] Masked Autoencoders Are Scalable Vision Learners

Masked Autoencoders Are Scalable Vision Learners

He, Kaiming, et al. "Masked autoencoders are scalable vision learners." Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. 2022.

 

Abstract

해당 논문에서 제시한 Masked Autoencoder(이하 MAE)는 Autoencoder 방식과 Self-supervised Learning을 결합한 모델로, asymmetric encoder-decoder 모델 구조를 가진다. 기존 자연어처리 분야에서 많이 사용되던 Masked language modeling을 이미지 분야에 적용시킨 것으로, 입력 이미지 patch에 Mask token을 사용하여 masking을 한 후 이를 복원시키는 작업을 수행한다.

기본적으로 input signal을 encoder에 통과시켜 latent representation을 얻고 이를 다시 decoder를 통과시켜 original signal을 복원시키는 Autoencoder의 콘셉트를 따라가지만, mask 되지 않은 patch만을 사용하는 encoder 구조와 비교적 가벼운 decoder 구조를 사용한 asymmetric architecture design을 제안하였다는 것이 차이점이다.

해당 논문은 Facebook AI Research의 Kaiming He가 1 저자로 참여한 논문으로, CVPR 2022에 게재되었으며 리뷰일 기준 약 1,020회의 인용 수를 자랑한다.

 

Self-supervised Learning

Supervised Learning 방식의 대표적인 문제점은 'Lack of labeled data'이다. Supervised model을 학습시키기 위해서는 데이터와 함께 정답 label이 필요한데, 이를 구축하는 데에는 분야에 따라 많은 시간과 노력, 그리고 비용이 투자되어야 한다.

이러한 문제에서 벗어날 수 있는 방법으로 제시된 것이 Self-supervised Learning이다. 이름에서 알 수 있듯이, 모델은 미리 정의된 Pretext task에 따라 스스로 label을 만들고 학습을 수행한다.

하지만 이러한 Self-supervised 방식은 NLP와 달리 Vision 분야에서는 많이 활약하지 못하였는데, 그 이유를 저자들은 다음과 같이 설명하며, 각각에 대한 해결책을 제시하며 논문을 시작한다.

 

1. Language-Vision Architecture Gap

input의 순서 및 위치가 크게 작용했던 NLP 분야와 달리, input을 grid로 나누어 locally 하게 처리하는 CNN 계열이 지배적이었던 Vision field에서는 근본적인 아키텍처 구조 차이가 존재했다. 이러한 아키텍처의 차이로 mask token이나 positional embedding을 활용하기가 어려웠던 것이다.

하지만 이는 Vision Transformer(ViT)의 등장으로 해결되었다.

2. Information density 차이

NLP 분야에서 다루는 Language는 semantic / information density가 높은 human-generated signal인 반면, Vision 분야에서 다루는 Image는 heavy spatial redundancy를 가지는 natural signal이라는 차이가 존재한다.

저자들은 이러한 redundancy를 줄이기 위해 대부분의 patch를 masking 하는 High-proportional masking strategy를 활용하여 low-level image static 이상의 전반적인 understanding을 요구하는 Self-supervised task로 문제를 재정의하였다.

3. Decoder에서의 차이

Rich semantic information을 가지는 'Missing words'가 decoder를 거쳐 reconstruct 되는 NLP 분야의 task와 달리, Vision에서는 low semantic information을 가지는 'Pixels'가 reconstruct 된다. Decoder는 encoder로 얻어진 latent representation으로부터 semantic level을 결정하는 중요한 역할을 하기 때문에 low semantic information을 reconstruct 하는 Vision 분야에서는 decoder의 design이 중요한 역할을 한다.

따라서 저자들은 decoder를 encoder와 독립적으로 작동하도록 만들어 유연한 디자인이 가능하도록 구성하였다.

 

Masked Autoencoders (MAE)

MAE는 input의 partial observation(masked input)으로부터 original signal을 복원하는 모델이다. 기존 Autoencoder와 동일하게 encoder부와 decoder부로 이루어져 있으며, input을 latent representation으로 encoding 후 이를 다시 decoding 하는 프로세스를 거친다. 다만 MAE는 기존 Autoencoder와는 달리 다음과 같은 asymmetric 한 아키텍처 구조를 가진다.

MAE Architecture Overview

1) MAE encoder

ViT와 유사하게, MAE encoder는 input의 patch들을 linear projectiopn 시킨 후 positional embedding을 더해준 뒤 Transformer block을 거쳐 encoding을 수행한다. 이때 ViT와 다른 점은, MAE encoder의 경우 mask 되지 않은 visible 한 patch만을 input으로 활용한다는 점이다. 논문에서는 75%의 patch들을 random 하게 masking 하여 25%만을 사용하는데, 이때 uniform distribution을 따라 샘플링함으로써 mask token의 편향을 방지하였다. 이러한 masking을 통해, MAE는 signal의 spatial redundancy를 줄이는 동시에 적은 메모리와 컴퓨팅 파워로 연산을 수행할 수 있다.

이렇게 인코딩 된 patch들은 다시 mask token과 합쳐져 full set을 이루고 이는 이후 설명할 lightweight decoder에 의해 처리된다.

2) MAE decoder

decoder의 input은 encoder를 거친 encoded visible patch와 mask token이 합쳐진 full set이다. 이때 mask token에도 positional embedding을 더해주어 location 정보를 가지도록 한다.

MAE decoder는 image reconstruction task를 수행하기 위한 pre-training 과정에만 사용되고 recognition을 위한 image representation 과정에서는 사용되지 않는다. 따라서 decoder 아키텍처는 encoder와 독립적으로 디자인될 수 있다는 장점을 가진다.

논문에서는 token 당 연산량이 encoder 대비 10% 미만의 narrow 하고 shallow 한 lightweight decoder를 사용하였는데, 이로써 pre-training time을 크게 줄일 수 있었다고 한다.

MAE decoder의 마지막 layer는 patch의 픽셀 값 수와 동일한 output channel을 가지는 linear projection으로, 픽셀 값의 벡터를 출력하며 이는 reshape 되어 reconstructed image를 생성한다. 이렇게 출력된 reconstructed patch는 원본 이미지로부터의 Mean Squared Error(MSE) loss를 계산하게 되는데, 이때의 loss는 BERT와 유사하게, masked patch에 대해서만 계산된다.

 

Experimental Results

ImageNet-1K를 training set으로 하여 self-supervised pre-training을 수행하였다. Backbone으로는 ViT-Large(ViT-L/16)을 사용하였는데, MAE로 fine-tune 하여 학습을 수행한 결과 다음과 같이 Top-1 accuracy의 상승이 있었음을 확인할 수 있다.

 

Ablation Study

해당 논문에서는 Vision task의 특성을 고려하여 BERT(15%) 보다 월등히 높은 Masking ratio인 75%를 default로 성능 실험을 진행하였다. 그 결과 fine-tuning과 linear probing 모두에서 공통적으로 75% Masking ratio로 실험 시 성능이 가장 좋음을 확인할 수 있었다.

이 외에도 Decoder depth, Decoder width를 다르게 하며 decoder의 크기에 따른 성능 변화를 실험하였고 encoder input에 대한 mask token 포함 여부를 가르게 하여 실험을 진행하였으며, Reconstruction target, Data augmentation의 유무 및 종류 그리고 Masking strategy를 변경하여 실험을 진행하였다.

 

결론적으로, decoder의 크기는 depth 8, dimension 512일 때가 가장 성능이 우수하였으며, Random sampling strategy를 적용하여 visible 한 patch만을 encoding 하였을 때가 가장 성능이 좋았다.

 

또한 epoch 수는 클수록 좋았으며 1,600 epoch가 진행될 때 까지도 noticeable 한 improvement를 보였다.

 

기존 SOTA methods와 비교한 결과, IN1K 데이터만으로 pre-train 하였음에도 다른 방법론을 능가하는 성능을 보여주었다. 또한 두 번째로 성능이 좋았던 BEiT과 비교하여도 MAE가 훨씬 간결한 구조와 빠른 연산속도를 보여주었다.

 

Supervised pre-training과 비교하였을 때에도 MAE pre-training 방식의 일반화 성능이 높았으며 Transformer block을 partially fine-tuning 한 경우 기존에 인기 있던 protocol인 linear probing을 압도하는 성능을 보였다.

 

Conclusion

해당 논문은 NLP 분야의 SSL 방식이 Vision 분야에도 충분히 적용될 수 있다는 확장 가능성을 보여주었다. 하지만 동시에 저자들은 image와 language는 본질적으로 다른 signal 특성을 가지므로 다소 주의하여 다룰 필요가 있음을 강조하기도 하였다.