[OUTTA Alpha팀 논문 리뷰] Part 11-5. MAE: Masked Autoencoders Are Scalable Vision Learners

논문 링크: https://arxiv.org/abs/2111.06377

Masked Autoencoders Are Scalable Vision Learners

깃허브 주소: mae/models_mae.py at main · facebookresearch/mae

mae/models_mae.py at main · facebookresearch/mae

---

Abstract

본 논문은 Masked Autoencoder (MAE)가 Computer Vision을 위한 확장 가능한 Self-supervised 학습 방법임을 보여준다.

핵심 아이디어:

• MAE는 단순한 구조를 가진다:
  입력 이미지의 일부 패치(75%)를 무작위로 마스킹하고, 누락된 픽셀을 복원한다.
• MAE는 두 가지 핵심 설계에 기반한다:
  1. 비대칭 인코더-디코더 구조 (asymmetric encoder-decoder)
     • Encoder는 마스킹되지 않은 일부 패치에만 작동한다.
     • Mask token은 사용하지 않음.
     • Lightweight Decoder는 latent representation과 mask tokens를 이용하여 전체 이미지를 복원한다.
  2. 높은 비율의 마스킹이 효과적임을 발견
     • 75% 이상 마스킹 → 더 어려운 학습 과제 → 더 좋은 표현 학습

성능 요약:

• 학습 속도는 3배 이상 향상, 정확도도 증가
• ViT-Huge 모델을 사용하면, ImageNet-1K 기준 87.8% 정확도 달성 (기존 self-supervised 방식 중 최고 성능)
• Transfer Learning에서도 supervised pre-training을 능가하는 성능을 보임

---

1. Introduction

배경

딥러닝은 하드웨어의 발전과 함께 지속적으로 더 큰 모델과 더 많은 데이터를 요구하게 되었음.

• 하지만 수억 개의 레이블 이미지 확보는 현실적으로 어렵다.

이 문제는 자연어 처리(NLP) 분야에서 Self-supervised learning으로 해결됨.

• 예: GPT (Autoregressive), BERT (Masked Autoencoder)
• 데이터를 일부 제거하고, 제거된 내용을 예측하는 방식
• 100B+ 파라미터 모델까지 일반화 가능

비전 분야에서 MAE가 어려운 이유

1. 과거 아키텍처의 차이점
   • 이전에는 CNN이 주류 → Mask token, positional embedding 통합이 어려움
   • Vision Transformer (ViT) 등장으로 이 문제는 해결됨
2. 언어 vs. 이미지의 정보 밀도 차이
   • 텍스트는 정보가 밀집되고 의미 중심 → 단어 몇 개만 예측해도 고차원 의미 학습
   • 이미지: 공간적 중복이 심함, 패치 몇 개로 대부분 예측 가능 → 학습이 쉬움
   • 해결책: 매우 높은 비율의 무작위 마스킹 (예: 75% 이상)
     → 더 어려운 과제 → 고차원 표현 학습 유도
3. Decoder의 역할 차이
   • NLP: decoder는 의미 있는 단어를 예측
   • Vision: decoder는 픽셀을 복원 → 낮은 수준의 의미
   • 따라서 Vision에서는 decoder 설계가 매우 중요

MAE 설계 및 특징

구조 개요

• Encoder: 마스킹되지 않은 패치만 처리 (mask token 없이)
• Decoder: latent vector + mask token으로 픽셀 복원
• 비대칭 구조로 계산량 감소 및 효율적인 학습 가능

→ Mask token을 decoder로만 보내면서 encoder 계산량을 75% 줄임

이점

• 학습 속도 3배 향상
• 메모리 소비 감소
• 대용량 모델로의 확장 가능 (e.g., ViT-Large, ViT-Huge)

성능 및 확장성

• ViT-Huge: ImageNet-1K 기준 87.8% 정확도 달성
• 다양한 다운스트림 태스크 (Object Detection, Segmentation 등)에서도 supervised pre-training보다 우수

→ NLP에서 BERT/GPT가 했던 역할을 Vision에서도 수행 가능함을 시사

MAE architecture

• MAE 구조를 요약한 그림
• 학습 시 75% 패치를 mask → decoder가 전체 이미지 복원
• 학습 후 decoder는 버리고, encoder만 사용

• ImageNet/COCO Validation 이미지 복원 결과
• 75~95% 마스킹에도 불구하고 의미 있는 복원 가능
• 일부 복원 결과는 실제와 다르지만 의미론적으로 타당함

---

2. Related Work

(1) Masked Language Modeling (MLM)과 Autoregressive Modeling

• BERT, GPT 등은 NLP에서 매우 성공적인 사전학습 방식으로 자리잡음.
• 입력 시퀀스 일부를 제거하고, 이를 예측하도록 학습.
• 확장성 뛰어남 → 사전학습된 표현이 다양한 다운스트림 태스크에 잘 일반화됨.

(2) Autoencoder 계열

• Autoencoder: 입력 → latent representation → 다시 원래 입력으로 복원
  • 예: PCA, k-means도 autoencoder의 일종
• Denoising Autoencoder (DAE): 입력을 의도적으로 손상시키고, 원래 신호를 복원하도록 학습
  • 예: 픽셀 마스킹, 컬러 채널 제거
• MAE도 DAE의 일종이지만, 전통적인 DAE와는 설계상 큰 차이가 있음

(3) Masked Image Encoding

• 이미지에 마스킹을 적용하여 표현을 학습
  • [59]: DAE에서 마스킹을 일종의 노이즈로 간주
  • Context Encoder: CNN 기반으로 큰 결손 영역을 복원
  • iGPT: 픽셀 시퀀스 기반 예측
  • ViT: patch 단위 마스킹 적용
  • BEiT: discrete token 예측 방식 채택

(4) Self-Supervised Learning

• Pretext task를 통해 학습 (예: Rotation 예측, Jigsaw Puzzle, etc.)
• 최근에는 Contrastive Learning이 큰 인기를 얻음
  • 대표적 연구: [SimCLR, MoCo, BYOL]
  • 데이터 증강에 크게 의존함
• Contrastive 방식과 달리, Autoencoding은 다른 방식의 표현 학습 경로를 제시

---

3. Approach

Masked Autoencoder (MAE) 개요

• 입력의 일부만 관찰한 상태에서 전체를 복원하는 간단한 autoencoder 방식
• 비대칭 구조 (Asymmetric Design):
  • Encoder: 마스킹되지 않은 패치만 처리 (mask token 없음)
  • Decoder: latent vector + mask token으로 전체 이미지 복원

3.1 Masking

• ViT [16]처럼 이미지를 patch 단위로 분할
• 무작위로 패치를 샘플링하여 나머지 패치를 마스킹
• Uniform sampling으로 중심 영역 마스킹 편향 방지
• 고마스킹 비율 (예: 75%) → 인접 패치로 예측이 불가능한 어려운 학습 과제 형성
  → 고차원적 표현 학습 유도

3.2 MAE Encoder

• ViT 구조를 기반으로 하되, visible patches만 입력
• 각 패치는 linear projection + positional embedding 후 transformer로 처리
• 전체 패치의 일부만 사용하므로, 연산량과 메모리 절약
  • 예: 전체 중 25%만 사용 → 대형 모델 학습 가능

3.3 MAE Decoder

• 입력:
  (i) 인코딩된 visible patches
  (ii) shared mask tokens (학습 가능한 벡터, 위치 정보는 positional embedding 추가)
• mask token은 예측 대상 위치를 알려주는 역할
• 경량화된 Transformer 구조
  • Encoder보다 얕고 좁은 구조 사용 (연산량 < 10%)
• Decoder는 pre-training에서만 사용, downstream task에는 사용하지 않음

3.4 Reconstruction Target

• 복원 대상: 마스킹된 패치의 픽셀 값
  • decoder의 출력 벡터는 각 patch의 픽셀 벡터
  • 마지막 layer는 linear projection (출력 채널 수 = patch pixel 수)
• Loss: MSE (mean squared error)
  • 마스킹된 패치에만 적용 (visible patch는 loss에 포함 X)
  • → BERT 방식과 유사함
  • 모든 픽셀에 loss를 주면 오히려 정확도 0.5% 감소
• Normalized pixel 값을 예측 대상으로 하는 변형도 실험
  • patch 내부 픽셀 평균과 표준편차로 정규화
  • → 표현 품질 향상

3.5 간단한 구현 방식

• MAE는 희소 연산(sparse ops) 없이도 효율적인 구현 가능

1. 모든 patch → token 생성 (linear + positional embedding)
2. 무작위로 토큰 리스트 섞고, 일부 제거 (마스킹 비율만큼)
3. Encoder는 이 작은 subset만 사용
4. Encoding 후 → mask token 붙이고, 원래 순서대로 unshuffle
5. 전체 토큰 리스트를 decoder에 입력

• shuffling / unshuffling 연산은 매우 빠르며, 추가 오버헤드는 거의 없음

Masking 비율에 따른 성능 변화

• Fine-tuning, Linear probing 모두에서 75% 마스킹 비율이 최적
  • 너무 낮거나 높은 비율은 성능 저하
  • 적절한 마스킹 비율 선택이 중요

---

4. ImageNet Experiments

Setup

• Pre-training: ImageNet-1K (IN1K) 학습 세트를 사용하여 self-supervised 학습
• 평가:
  1. End-to-end fine-tuning
  2. Linear probing
• 측정 지표: Top-1 validation accuracy (224×224 crop)

4.1 주요 특성 분석 (Ablation Study)

기준 모델: ViT-Large (ViT-L/16)

• ResNet-50보다 훨씬 큰 모델이며 오버피팅 경향 존재
• 비교 결과:

→ MAE 사전학습은 50 에폭만 fine-tuning 해도 성능이 뛰어남 (from scratch 학습은 200 에폭 필요)

Masking Ratio (마스킹 비율)

• 최적의 마스킹 비율은 75%로 매우 높음
• BERT의 마스킹 비율(15%)보다 훨씬 높고, 기존 Vision 방법(20~50%)보다도 큼
• linear probing 성능은 마스킹 비율에 민감하고 75%에서 최대 성능 도달
• fine-tuning은 다양한 비율(40~80%)에서 안정적

→ 높은 비율의 마스킹이 추론 기반의 표현 학습을 유도함

Decoder 구조 변화 실험

(a) Decoder 깊이

• Linear probing 성능은 깊이에 따라 최대 8% 향상
• Reconstruction task는 low-level 정보를 다루기 때문에, decoder가 깊을수록 효과적
• 하지만 fine-tuning에서는 영향이 적음

(b) Decoder 너비

• default: 512-d (ViT-L은 1024-d)
• narrow decoder도 충분히 성능 유지
• → decoder는 연산량 9%에 불과 (전체 compute 중)

Mask token을 Encoder에서 제거

(c) 실험 결과:

• mask token을 encoder에서 제거하면 성능이 크게 향상됨
• 실제로 추론 시에는 mask token이 없기 때문
• 학습 시간 2.8~4.1배 단축, 메모리 사용량 절감

Reconstruction Target 실험

(d) 다양한 예측 대상 비교:

→ 정규화된 픽셀이 가장 좋은 성능
→ BEiT 방식의 token 예측은 오히려 linear probing 성능 하락

Data Augmentation 실험

(e) 실험 결과:

• Color jitter는 오히려 성능 하락
• MAE는 강한 augmentation 없이도 효과적
• contrastive learning처럼 augmentation에 크게 의존하지 않음

Mask Sampling 전략

(f) 결과 비교:

• Random masking이 가장 좋음
• Block-wise는 재구성은 어렵지만 표현 학습 품질은 떨어짐
• Grid는 재구성은 쉽지만 representation 품질 낮음

Training Schedule

• 학습 에폭이 길어질수록 성능 향상
• 1600 epochs까지도 성능이 계속 증가 (saturation 없음)

4.2 기존 방법과의 성능 비교

ImageNet-1K fine-tuning 성능 비교:

→ MAE는 external data 없이 최고 성능 (ViT-H + 448 입력)

4.3 Partial Fine-tuning (일부 layer만 학습)

• MAE는 linear probing 성능은 낮지만, 일부 block만 fine-tuning 해도 성능 급상승
  • 예: 마지막 block 1개만 학습 → 73.5% → 81.0%
  • MLP sub-block만 학습해도 79.1%
• MoCo v3보다 부분 fine-tuning에서 일관되게 우수
• 결론: MAE는 비선형 표현 학습에 강함

즉,

• MAE는 간단한 구조, 고효율 연산, 우수한 표현력을 가짐
• contrastive 학습보다 학습 비용 낮고, 데이터 증강 의존도 적음
• 확장성 뛰어나며, transfer 학습에서도 강력한 성능

---

5. Transfer Learning Experiments

본 장에서는 ImageNet-1K에서 MAE로 pre-training한 모델을 다양한 다운스트림 태스크에 fine-tuning하여 transfer 성능을 평가하였다.

MAE Pre-training vs. Supervised Pre-training (ImageNet-1K, 224px)

 

• MAE가 같은 데이터만 사용해도 supervised 방법을 능가함
  • 특히 ViT-H/14에서는 MAE가 supervised (ImageNet-1K)의 모든 결과보다 높은 정확도를 보임
• JFT-300M처럼 수억 개 이미지로 사전학습한 supervised 모델과도 비슷하거나 뛰어난 성능
  • MAE는 외부 데이터 없이도 높은 성능 도달 (87.8%)
• 모델 크기가 커질수록 MAE의 강점이 더욱 뚜렷
  • 이는 MAE의 스케일링 성질을 강조하는 부분

5.1 COCO Object Detection & Segmentation

• 방법: Mask R-CNN을 이용해 COCO 데이터셋에서 end-to-end fine-tuning 수행
• Backbone: ViT-B / ViT-L, FPN에 맞게 조정
• 지표: APbox (Object Detection), APmask (Instance Segmentation)

 

→ MAE는 Supervised 및 MoCo v3보다 성능 우수
→ BEiT보다 간단하고 빠르면서 비슷하거나 더 좋은 성능

5.2 의미론적 분할 (ADE20K)

• 방법: UperNet 사용
• 성능 지표: mIoU (mean Intersection over Union)

→ MAE는 ViT-L에서 supervised보다 3.7% 더 높음

5.3 분류 (iNaturalist / Places)

→ iNat에서는 모든 모델 크기에서 기존 최고 기록을 초과
→ Places 데이터에서도 수십억 이미지로 pre-training한 기존 모델보다 우수

5.4 Pixel vs. Token 기반 복원 비교

→ 정규화된 픽셀과 token 방식은 성능 차이가 거의 없음
→ 복잡한 tokenization 과정 없이도 픽셀 기반 MAE로 충분히 고성능 달성

---

6. Discussion and Conclusion

• MAE는 간단하면서도 확장 가능한 self-supervised 방식
• ImageNet 및 다양한 다운스트림 태스크에서 supervised pre-training을 능가
• Contrastive 방식보다 연산량, 메모리 사용량, 구현 난이도 모두 낮음
• MAE는 정형화된 object 제거가 아닌 무작위 패치 제거라는 점에서 NLP와 차이 존재
• 하지만, 픽셀 복원만으로도 의미 있는 시맨틱 표현을 학습한다는 점에서 흥미로운 가능성 제시

Future Work

• 비전에서의 self-supervised learning도 NLP의 BERT/GPT처럼 스케일업을 통한 효과를 기대할 수 있음
• 실제 실험에서도 모델 크기가 커질수록 성능이 선형적으로 향상
• 추후 MAE 방식은 초대규모 모델 학습 및 전이학습의 기반이 될 수 있음

Boarder Impacts

• MAE는 학습 데이터 통계에 기반해 내용을 예측하므로, 데이터의 편향이 그대로 반영될 수 있음
• 가짜 또는 존재하지 않는 내용 생성 가능성 존재
• 추후 응용 시 적절한 통제와 책임 있는 사용 필요