Masked Autoencoders Are Scalable Vision Learners

Language model에서 Transformer 이후 BERT가 등장한 이후 competer vision 분야에서도 Vision Transformer의 BERT급을 만들기 위한 논문이 이 Masked Autoencoders(MAE) 입니다.

Introduction

본 논문은 이 문제점에서 시작되었습니다.

" What makes masked autoencoding different between vision and language? "

NLP 에서 BERT 의 성공 이후에 나오는 Autoencoder 의 성능이 Computer Vision 에서 나오지 않자, vision 과 language 영역의 차이가 어디서 나오는것인지 궁금증을 가지고 3가지의 관점에서 서술하였습니다.

1. Architecture
   비전 분야에서 CNN 은 주된 아키텍쳐였었습니다. 이러한 CNN 은 구조는 혼자 단독적인 모듈로써 존재할 뿐, network 의 indicator 같은 역할(positional embeddings, mask token 등)을 수행하는 모듈이 존재하지 않았습니다. 이러한 아키텍처 간극은 Vision transformer 의 등장으로 극복할 수 있게 되었습니다.

2. Information density
   NLP 의 주 데이터인 언어는 인류가 만들어낸, 많은 정보가 내포되어있는 signal 입니다. 따라서 AI 에서 문장의 일부를 Masking 하고 이를 예측하는 방식으로 언어를 학습할 수 있었습니다. (이것이 바로 BERT 죠)
   하지만 비전 분야에서 메인 데이터인 이미지의 경우, 공간적 중복이 많은 natural signal 에 해당됩니다. 이 말은 어떠한 픽셀은 인접한 픽셀과 공간적으로 유사할 확률이 높다는 뜻입니다. 따라서 이를 고려해 논문의 저자들은 높은 비율의 random masking을 적용합니다. (BERT 의 masking 비율은 15%, MAE 의 masking 비율은 75%)

3. Autoencoder
   autoencoder의 decoder 는 latent representation 으로부터 input 과 동일한 값을 재건(똑같이 만들어내는)하는데, 여기서 NLP 와 비전에서 다른 양상을 보입니다.
   NLP 에서는 풍부한 의미 정보를 갖는 단어를 예측하기 때문에, decoder 로 단순 MLP로도 사용할 수 있지만, 의미 정보가 부족한 pixel 을 예측하는 비전에서 decoder 는 학습된 latent representation 의 의미 수준을 결정하는데 중요한 역할을 한다고 이야기합니다.

이 세가지 분석을 바탕으로 저자들은 이미지의 일부를 random masking 하고, 이를 다시 pixel level 로 복원하는 MAE 를 제안합니다.

Related Work

MAE 와 관련된 연구들을 정리해보았습니다.

Masked language modeling

NLP 에서 우선적으로 좋은 성능을 낸 pre-trained representation method 인 BERT , GPT 가 Masked language modeling 입니다.

Autoencoding

Autoencoding은 앞서 말한 learning representation 방식 중 가장 클래식한 방식으로, encoder 와 decoder로 이루어져 있는 구조입니다.

Encoder 는 input을 latent representation으로 mapping 하고, decoder는 latent representation 된 것을 다시 input 으로 reconstruct 하도록 진행됩니다.

즉, 인코더로 들어오는 정보를 압축하고, 디코더로 이들을 복원하는 구조로 학습이 진행된다고 보시면 됩니다.

Denoising autoencoder(DAE) 는 autoencoder의 일종으로, input에 noise 를 넣은 후 다시 uncorrupted input signal 로 복원하는 모델입니다. 저자들은 MAE 를 DAE 의 classical 한 방식이라고 말하고 있답니다.

Self-supervised learning

Self-supervised learning은 대표적인 비지도 학습으로, label 없이 학습하는 방법입니다. 비전 분야에서는 contrastive learning 과 masked image modeling 이라는 2가지 방식으로 발전하고 있습니다.

contrastive learning은 저차원 공간에서 이미지를 인코딩 하는 방식으로, 유사한 이미지들끼리는 가깝게, 다른 이미지들끼리는 멀리 떨어지는 과정으로 학습합니다.
즉, 유사한 양수 쌍 혹은 샘플들의 표현 사이의 거리를 최소화하고, 다른 음수쌍 혹은 샘플들의 표현 사이의 거리를 최대화하여 특징을 추출 학습 시키는 방법입니다.

이 방식은 label 없이 학습하기 위해 동일한 이미지에서 나온 것인지를 맞출 수 있도록 모델의 pre-trained objective를 설정하여 진행합니다.

MAE 는 masked-image modeling의 방식으로, 이미지의 일부 영역에 masking을 적용하고, 다시 원래의 이미지로 복원하는 pre-training objective 로 학습을 진행합니다.

Approach

Masked autoencoder 를 구현한 자세한 방식들이 설명되어있습니다.

figure 1 을 확인해보면, MAE의 특징들을 확인할 수 있습니다. MAE 는 encoder 에 대해 lightweight decoder 를 사용하는 asymmetric design 을 가지고 있고, encoder에서 masked patch 를 처리하지 않기 때문에 training time 과 memory 측면에서 좋으며 model scalability 가 좋은 특징을 지니고 있습니다.

Masking

MAE 는 Vision transformer 를 기반으로 제작되었으므로 이미지로부터 조각조각 나눠진 patch 를 얻습니다. 이후 uniform distribution 을 따르는 random sampling(단순 무작위 추출 방법) 을 통해서 patch(패치) 에 masking을 진행합니다.

masking을 어느정도 하는 것이 아닌 무려 75%나 masking 을 진행하는데요, 이렇게 높은 비율로 진행하면 갖게 되는 이점을 이야기하고 있습니다.

1. 공간적으로 중복이 많은 인접한 patch 정보를 제거함으로써, pre-training objective 의 난이도를 높혀 이미지를 잘 학습하게 만듭니다. (많이 제거하지 않으면, 너무 쉬워서 학습이 잘 안되기 때문 !)

2. Uniform distribution을 따르는 masking 방식으로 center bias 문제를 방지할 수 있습니다. (contrastive learning에서 view를 만드는데 중요한 역할을 하는 crop augmentation의 경우 이미지의 중심에 편향되는 center bias 문제가 있는데, 이를 고려하여 해결)

MAE Encoder

MAE의 encoder는 VIT 와 동일하지만 작은 과정만 작동하는데요, recognition을 위한 image representation을 추출합니다.

그리고 visible patch만 학습하기 때문에 pre-training cost를 줄이고 memory 또한 줄일 수 있게 됩니다. masked patches 를 처리하는 부분은 가벼워진 decoder 에서 처리한다고합니다.

MAE Decoder

MAE의 decoder는 pre-training 과정에서만 image 재구성(reconstruction) 을 수행하기 때문에, encoder design과 무관한 decoder design이 가능해집니다.

decoder 과정에서 positional embedding 을 추가하여 masked token 에 지역 정보(local information)를 넣어줍니다.

또한, encoder에서 나온 encoded visible token과 mask token을 같이 학습하는데, lightweight decoder를 사용함으로써 pre-training cost를 줄입니다.

MAE decoder의 forward 과정을 차례대로 확인해보면,

1) Embedded tokens

Encoded visible token의 차원과 decoder의 차원을 맞춰주기 위한 linear projection 을 진행합니다.

2) Append mask token + unshuffle

Unshuffle : masking 구현을 위해 적용했던 shuffle을 되돌리는 과정으로, 이를 위해 encoding에서 적용했던 masking 과정의 index 정보가 필요하게 됩니다. 따라서

3) Add positional embedding

과정을 진행하게 되는데요, mask token은 embedding이 적용되지 않았기 때문에, 추가적인 positional embedding을 진행하여 지역 정보를 넣어줌으로써 maked token 에도 information을 넣어줍니다.

4) Apply Transformer blocks

를 진행하고

5) Predictor projection

Loss 계산을 위해 target image와 shape을 맞춰주기 위한 linear projection 을 진행합니다.

6) Remove class token

7) Reconstruction target

MAE의 reconstruction target은 masked patch 안의 pixel이며, MSE loss를 사용합니다. 그

loss 계산 전에 target image에 patch 단위의 normalization을 적용했을 때 더 좋은 representation quality를 보였다고 말하는데, 이에 대해 별 다른 이야기는 없지만, transformer block의 layer normalization이 patch 단위로 적용되기 때문인 것으로 생각됩니다.

Simple Implementation

MAE는 앞서 언급했던 것처럼, encoder에서 masked patch를 학습하지 않고, lightweight decoder를 사용하기 때문에 효율적으로 학습이 가능한데요.

이뿐 아니라 Encoding의 masking 과정은 단순 random shuffle 이후 visible token만 encoder로 넣어주며, Decoder의 경우에도 mask token을 concat한 다음 random shuffle의 역 과정 (=unshuffle) 을 거치면 됩니다.

즉, masked image modeling을 위한 masking, shuffle의 과정에 sparse operation을 사용하지 않기 때문에 빠르다고 이야기합니다.

ImageNet Experiments

실험 결과들을 차례대로 확인해보겠습니다.
저자들은 ViT-L를 baseline으로 설정하였으며, original ViT 대비 strong regularization 을 적용하였습니다. 여기서의 MAE는 fine-tuning 성능을 기록하였는데, 200 epoch를 학습한 scratch ViT에 비해 MAE는 단 50epoch만으로 2.4%p 높은 성능을 달성했습니다. 다음으로는 다양한 ablation study를 바탕으로 MAE의 특성들을 밝히고 있습니다.

Main properties

Masking ratio

Masking ratio가 상대적으로 높은 수치인 75%에서 linear probing, fine-tuning 모두 좋은 성능을 보였습니다.

이는 BERT의 15%, iGPT의 20%, ViT의 50% 보다도 높은 수치입니다.

Linear probing은 75%까지 성능이 상승하다가 떨어지지만, fine-tuning은 40%부터 80%까지 괜찮은 성능을 내고 있습니다.

Decoder design

decoder의 depth가 일정 수준 이상이여야 linear probing 성능이 나오는 것을 확인할 수 있습니다. 이는 pre-training 단계에서 reconstruction 방식을 학습하기에 decoder의 depth가 충분하지 못하기 때문에, recognition을 담당하는 encoder의 일부 layer까지 영향을 줄 수 있기 때문입니다. 즉, encoder가 온전히 recognition에 집중하지 못해서 backbone을 freeze하는 linear probing의 성능이 낮아지게 된다고 이야기합니다.

Mask token

encoder는 mask token을 생략하였다 하는데, 이는 성능 면에서도 좋은 결과를 보여주고 있으며, 특히 linear probing 성능의 차이가 많이 나는 것을 확인할 수 있습니다. Mask token은 실제 이미지가 아니기 때문에 encoder의 학습을 방해하는 결과를 가져오고, 이것이 linear probing 성능 차이로 드러난다고 저자들은 이야기합니다. 당연히 FLOPs 측면에서도 효율적임을 알 수 있습니다.

Reconstruction target

Pixel with per-patch normalization 성능이 가장 좋은 것을 확인할 수 있습니다. 하지만 회색 음영을 보시면 ablation study에서 patch norm을 적용하지 않은 것을 default로 설정했습니다.

Pixel with per-patch normalization은 말 그대로, target 이미지를 patch 단위로 쪼갠 후 patch 단위로 normalization이 적용된 pixel을 예측하는 것입니다.

저자들은 per-patch normalization을 통해 local contrast를 향상시킬 수 있다고 이야기합니다.

또 다른 실험으로 PCA를 활용했습니다. 224 이미지에 16 patch로 쪼갤 경우, 1개의 patch는 196 size를 갖게 됩니다. 여기에 PCA를 적용해서 96 차원으로 embedding 시킨 결과를 target으로 학습하는 방식을 적용했습니다. 생각보다 성능 하락이 크지 않았다는 점은 의외인 것 같습니다.

저자들은 앞선 2개의 실험 결과를 바탕으로, MAE는 high-frequency가 필요하다고 이야기합니다.

그 다음으로, BEiT의 DALL-E tokenizer 실험을 합니다. 이때의 decoder는 cross-entropy loss를 사용하여 token indices를 예측하는 task로 학습했는데, patch-norm을 하기 전보다 fine-tuning 성능이 좋은 것을 볼 수 있습니다. 이에 대해 저자들은 그럼에도 patch-norm을 적용하면 MAE의 fine-tuning 성능이 개선되고, MAE는 tokenizer가 추가로 필요하지 않다는 점에서 BEiT보다 더 좋다고 주장합니다.

Data augmentation

None은 center-crop만 적용했다는 것을 의미하고, 나머지는 horizontalflip이 추가되어 있습니다. 여기서 중요한 부분은 center-crop만으로도 성능 하락이 크지 않았다는 점인데요, 이는 매 iteration 마다 적용되는 random masking 의 효과로, augmentation 없이도 충분히 training regularization이 가능하다는 것을 보여줍니다.

특히, augmentation에 크게 의존하는 contrastive learning과 비교하여 장점 이라고 할 수 있습니다.

Mask sampling stategy

BEiT에서 적용했던 block-wise masking 과 grid masking 을 대조군으로 실험하였습니다.

다른 masking 대비 random masking의 성능이 좋은 것을 알 수 있고, 저자는 block masking을 적용하였을 때 training loss가 불안정했고 reconstruction 결과가 blurry했다고 합니다.

Training schedule

Ablation study에서는 800 epoch를 baseline으로 실험하였습니다.

Results

정리하면, MAE는 알고리즘이 간단하고, 학습 속도가 빠르며, 모델 확장성이 좋은, vision 도메인에서 masking을 활용한 self-supervised method 입니다.

저자들이 말하길 이미지는 단어의 시각적 유사체일뿐, 의미론적인 분해가 불가능한, 기록된 빛에 불과하다고 합니다. 그렇기에 MAE는 높은 비율로 의미를 형성하지 않을 random patch를 제거했고, 의미를 형성하는 단위가 아닌 pixel level에서의 reconstruct를 하며 학습하도록 설계되었습니다. 이러한 학습 방식에도 불구하고 MAE는 복잡하고 전체적인 reconstruct를 추론해내는 능력을 가졌고, visual concept도 잘 학습하는 것으로 보입니다. 이는 앞으로 연구될 가치가 있는 MAE의 hidden representation 특성이므로 future work로 제안하고 있습니다.

Masked Autoencoders Are Scalable Vision Learners - 논문 보기