[논문 리뷰]Masked Auto-Encoder

terms

• patch / token: (이미지) 파편
• 잠재 변수(Latent representations): 특정 데이터를 가장 잘 설명할 수 있는 값
• Auxiliary dummy token: ViT의 Class token 역할

요약

MAE is Scalable Self-supervised learners.
MAE의 Concept는 패치를 랜덤으로 Masking하고, Masking된 Missing Pixels를 재 구성 하는 것이다.
MAE의 Design은 다음 두 가지를 따르며, 무거운 모델에 대한 efficiently and effectively하게 학습이 가능하다.

1. 비대칭 Encoder-Decoder 구조이다.
   Encoder는 without Mask patches에 대해서 동작하며, Decoder는 latent representations와 mask patches를 통해 원본 이미지를 재구성한다.
2. 75%와 같은 매우 큰 영역에 대해 마스킹을 처리하게 될 시 의미있는 supervisory task를 반환할 것으로 예상한다.

해당 Design은 3배 이상의 속도와 함께 정확도 향상을 도출한다.

초록

NLP 모델의 Solutions: 데이터의 일부를 제거하고, 제거된 영역을 예측

• GPT
  • Auto Regressive
• BERT
  • Masked Auto-Encoding

대략 1000억개의 Parameter를 소비

NLP vs Vision

1. 아키텍쳐가 다르다.
   • Vision은 CNN이 압도적으로 널리 활용
   • Masked patch이나 Positional embedding 같은 indicators를 결합하는 데에 장애
     • ViT에서 다룸으로써 어느정도 해결
2. 정보 밀집도가 다르다
   • 언어: 인간이 만들어낸 일종의 신호로, 매우 큰 정보를 포함
     • 문장에서 Masking 후 예측 할 때 Sophisticated languages understanding을 유도
   • 이미지: 공간적 중복성(정보)을 보유한 자연 신호
     • 패치가 없어져도 주변 패치의 parts, objects, scenes를 high-level로 이해하여 회복 가능
     • 차이를 극복하기 위해 이미지의 매우 많은 영역을 랜덤으로 Masking
3. Decoder
   • 언어의 Decoder는 찾아 낼 Missing Words의 Semantic이 풍부
   • BERT의 Decoder는 MLP로 보잘 것 없다!
   • 이미지 Decoder는 Pixel을 형성하므로 Recognition Task에 비해 낮은 Semantic을 보유
   • 이미지의 Decoder은 학습된 잠재 변수의 Semantic level을 결정

MAE

Random patch에 대해 마스킹 후 재구성

• Encoder와 Decoder가 비대칭 구조
• Win-Win Scenario
  1. 인코더는 데이터의 25%만을 이용해서 최적화
  2. 전반적으로 3배 이상의 시간과 메모리를 절약
• ViT-Huge: 87.8%(ImageNet 1K)

관련 연구

Bert & GPT: Masked language model

Auto Encoder: Representations를 학습하는 가장 classic한 방법

• Encoder: Input을 latent representations으로 맵핑
• Decoder: latent representations로 부터 Input을 재구성
• Ex) PCA, K-means

Denoising Auto Encoder(DAE)

• Auto encoder의 일종
• Input signal을 오염시킨 후, original signal을 재구성하는 학습
• MAE에 비해 무지한 방법

Stacked DAE

• DAE에 대한 일종의 후속작으로 Masking을 통해 Noise를 추가
• Context Encoder: 넓은 missing regions를 CNN을 통해 Inpainting
• Motivations: Transformers, NLP
  • iGPT: pixel의 시퀀스로 동작하며, unknown pixels를 예측
  • ViT: self-supervised learning을 위한 masked patch 예측
  • BEiT: 개별 patch를 예측 (Recent research)

Self-supervised learning

• Pre-training을 위해 다른 pre-text tasks에 초점
• 최근, 두 개 이상의 view로 이미지 유사도와 비유사도(혹은 유사도만)를 모델링하는 Contrastive learning이 유행

  즉, 입력 샘플 간의 비교를 통해 학습을 하는 것

  • Feature representations의 우수한 일반화
  • 새로운 class에 대해 대응이 용이

제안

MAE Architecture

• 비대칭 구조의 Auto encoder
  • Encoder: Visible patch에 대한 Signal을 Latent representations로 맵핑
  • Decoder: Encoder에서 추출한 Latent representations로 부터의 전체 Signal과 Masked patch를 재구성하는 가벼운 Decoder
  • Figure 1 참조
• Masking
  • 중복되지 않은 패치로 분리
  • 샘플을 분리하여 masking 처리
  • Sampling 방법: 위치 변경 없이 특정 확률에 대해 랜덤으로 Patch 샘플들을 선택하여 중앙 영역의 Biases를 억제
  • Masking에 높은 비중을 주어 대부분의 공간 정보를 제거
• MAE Encoder
  • Visible patches에 대해 ViT의 Encoder를 적용
  • 기존 Transformer와 달리 25% 가량의 적은 Subset에 대해서만 연산하며, Masked patches는 연산하지 않고 제거
  • 연산량 및 메모리 효율을 대폭 향상
  • Patch의 Representations 생산만을 목적
• MAE Decoder
  • 입력 값: Fullset(Visible patches의 Encoding 정보와 Mask token)
  • 각 Masked patch는 스스로를 학습하는 벡터
  • 공간 정보를 식별하기 위해 Fullset patch 별 Positional embeddings를 부여
  • Patch의 재구성만을 목적으로 하며, Encoder와 독립적이므로 유연하게 디자인 가능
  • 본 연구에서는 Encoder에 비해 10% 미만의 연산량을 확보한 매우 가볍고 작은 Decoder로 설계
• Target 재구성
  • Decoder의 최종 레이어가 Output channels와 Pixel 수가 동일한 Linear projection으로 구성
  • Reshape를 통한 이미지 재구성
  • 본 연구에서 Loss 함수로 MSE를 사용하며, Loss는 BERT와 유사하게 Masked patches에 대해서만 적용
  • Representation의 퀄리티를 높이기 위해 평균과 표준편차를 고려한 Normalized pixel로 이미지를 재구성

Simple implementations

1. 매 Input에 대해 patch를 생성(Linear projection with positional embedding)
2. Shuffle 후 Masking ratio에 따른 뒷 부분 제거
3. 인코딩 후 리스트에 Masked token 추가
4. Full list에 대해 Unshuffle(Inverting) 후 Patch 정렬
5. Decoder에서 해당 Full list를 적용

※ Shuffling과 Unshuffling의 overhead는 적다.

실험

ImageNet Expreiments

Baseline: ViT-Large

• ViT 기본 구조
  • Multi-head self-attention bloack(Layer Norm) + MLP block(Layer Norm)
  • Class token 보유
• 구조 변형
  • Encoder와 Decoder의 크기가 다르므로 Encoding 후 Linear projection layer를 도입
  • Encoder와 Decoder 각 Input에서 Sine-Cosine Positional embeddings 적용
• 학습 방법
  • Pre-training: Encoder와 Decoder를 학습
  • (end-to-end) Fine-tuning: Encoder(trainable) + Classifier
  • Linear probing: Encoder(frozen) + Classifier
    • Linear layer만 학습
    • Mixup, Cutmix, Drop path, Color jittering와 같은 Regularization이 큰 방해 요소로 취급
    • Weight decay를 0으로 초기화
    • Pre-trained features를 정규화하는 것이 유용하여 Affine 변환을 하지 않은 Batch Norm 레이어를 Linear classifier 앞 단에 적용
  • Scratch
    • ViT-L/H의 학습이 불안정한 점을 완화
    • Original < Implemented
• 특징 추출
  • Encoder의 output
  • Auxiliary dummy token 없이도 원활한 작동을 지원

학습 방법 별 Settings
	Pre-training	Fine-tuning	Linear probing	Scratch
optimizer	AdamW	AdamW	LARS	AdamW
optimizer momentum	β1, β2=0.9, 0.95	β1, β2=0.9, 0.999		β1, β2=0.9, 0.95
base learning rate	1.5e-4	1e-3	0.1	1e-4
weight decay	0.05	0.05	0	0.3
batch size	4096	1024	16384	4096
layer-wise lr decay		0.75
learning rate schedule	cosine decay	cosine decay	cosine decay	cosine decay
warmup epochs	40	5	10	20
training epochs		100 (B), 50 (L/H)	90	300 (B), 200 (L/H)
augmentation	RandomResizedCrop	RandAug (9, 0.5)	RandomResizedCrop	RandAug (9, 0.5)
label smoothing		0.1		0.1
mixup		0.8		0.8
cutmix		1.0		1.0
drop path		0.1 (B/L) 0.2 (H)		0.1 (B), 0.2 (L/H)
exp. moving average (EMA)				0.9999

Properties

Masking ratio 별 validation 정확도 비교: Fine-tuning model에 대한 결과(상)와 Linear probing에 대한 결과(하) Reconstructions 결과

• Masking ratio 별 비교
  • Fine-tuned와 Linear probing 모두 75% 가량의 Mask ratio가 이상적
  • BERT는 15%의 ratio가 이상적이며 그 외의 관련 연구에서도 20%~50%에 비해 월등히 높은 비율
  • Fine-tuned model에서는 ratio와의 연관성이 낮으며 scratch에 대한 결과(82.5%)보다 우수한 성능을 도출
  • Linear probed model에서는 ratio에 비례하는 추세

Decoder depth	Decoder width	Mask token

Reconstruction target	Data augmentation	Mask sampling

Masking 방법	Epochs에 따른 Accuracy

• Decoder design
  • Decoder의 Depth와 Width에 대해 성능 차이
  • Mask tokens 유무에 따른 속도 및 성능 차이 (FLOPS vs FLOPs)
  • Reconstruction 방법에 따른 성능 차이
  • Data augmentation에 따른 성능 차이
  • Masking 방법에 따른 성능 차이

MAE vs Self-supervised pretrained models	MAE vs Supervised pretrained models	Partial fine-tuning

Object detection and segmentation	Semantic segmentation

Object detection and segmentation	Semantic segmentation

결론

• 도메인 별 Self-supervised learning
  • NLP: 기하급수적으로 확장되는 models에 대해 유용
  • CV: Self-supervised learning이 성장하더라도 해당 task에서는 기존 CNN이나 ViT가 적용
• ImageNet과 Transfer learning을 통해 Auto-Encoder가 Scalable 모델에서 유용
• Object를 Masking하는 대신 Patch를 제거하였음에도 불구하고, Semantics와 같은 visual concepts를 학습하여 Patch 재구성
  • 가설: Hidden representations