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

[2111.06377] Masked Autoencoders Are Scalable Vision Learners

Masked Autoencoders Are Scalable Vision Learners

 

1. Intro

배경 : 모델이 점점 커지면서 → 100만 장 라벨 데이터로는 부족 → 수억 장 필요 → 대부분 접근 불가

• NLP ⇒ self-supervised pretraining으로 해결
  • 라벨 없는 데이터에서 일반화 가능한 표현을 학습하는 것!
  • GPT의 autoregressive
  • BERT의 masked autoencoding
• Vision ⇒ NLP에서보다 autoencoding 적용이 어려움

왜?

1. 구조 차이

• Vision에서는 CNN이 주류
  • CNN은 regular grid 기반이라 mask token, position embedding을 넣기 어려움 → ViT로 이 문제 해결

1. 정보 밀도 차이

• NLP : 단어 하나가 많은 의미 포함 → 일부 단어만 mask해도 어려운 문제
• Vision : 공간 중복 높음 → 인접 패치로 쉽게 복원 가능
  • 따라서 높은 mask 비율 필요 → 복원이 어려워지고, 의미 있는 feature 학습 가능

1. Decoder 역할 차이

• NLP : Decoder 예측 대상 = 단어
  • high-level semantic = 단어 자체가 문장 의미에 직결되니까 복잡한 구조 필요 X
  • 단순 MLP decoder로도 표현력 +
• Vision: Decoder 예측 대상 = pixel
  • low-level semantic
  • decoder 설계가 학습 표현의 질에 영향
  • 단순 MLP로는 부족, decoder 구조 최적화 필요

제안 : visual representation learning을 위한 MAE

• Asymmetric encoder-decoder
  • 인코더 : 마스킹되지 않은 패치(25%)만 처리 → 연산 감소
  • 디코더: latent + mask token을 입력받아 픽셀 복원
  • 마스킹 토큰 처리를 decoder 쪽에 몰아줌 → encoder 계산량 크게 감소
• High masking ratio (75%) → redundancy 줄이고, global context 학습 유도
• 결과: 학습 속도 3배+, 메모리 절감, 대규모 모델로 쉽게 확장 가능

 

2. Related Work

Masked language modeling

• GPT의 autoregressive
• BERT의 masked autoencoding
  • scaling +
  • downstream task에 일반화 + (pretraining → fine-tuning 구조가 다양한 과제에 잘 적용됨)
• ⇒ NLP의 사전학습 +

Autoencoding

• 표현 학습
• 입력 → 인코더 → 잠재공간(latent representation) → 디코더 → 원본 복원
• DAE
  • 훼손하고 → 깨끗한 신호 복원
  • MAE도 DAE의 일종이지만, 여러 면에서 다름

Masked Image Encoding

• iGPT [6]: 픽셀 시퀀스로 예측
• ViT [16]: 마스크된 패치 예측
• BEiT [2]: discrete token 예측

Self-supervised learning

• pretext task를 중심으로 발전
• 최근에는 contrastive learning ↑

 

3. Approach

 

MAE

• 부분적인 정보로부터 원래 신호 복원 : autoencoding
• encoder : obcerved signal → latent 표현
• decoder : 복원
• asymmetric이라는 점! → 그래서 encoder는 mask 없는 부분만 다루게 됨

 

 

Masking

• ViT처럼 이미지를 패치로 나눔
• 그 중 75%는 마스킹
  • 모델이 단순한 저수준 픽셀 연속성만 보고 복원하는 것을 방지 → 더 의미 있는 표현 학습 유도 가능
  • encoder의 계산량 감소
  • uniform distribution, random sampling : 이미지 중앙에 더 마스킹 되는 현상 (potential center bias) 방지

MAE encoder

• 마스킹 안된 것 처리
• ViT encoder 처럼 패치를 linear projection해서 임베딩하고, 위치 정보를 더한 후 (⇒ token으로 변환) transformer 블록에 입력
  • token list random shuffle
  • 마스킹 비율(예: 75%)만큼 리스트 뒤쪽 토큰 제거
  • → 남은 25%만 인코더 입력

MAE decoder

• encoded visible patches + mask tokens 입력
  • 제거했던 위치에 mask token 넣어 다시 전체 토큰 개수(N개)로 복원
  • 셔플했던 순서를 Unshuffle해서 다시 원래 패치 순서로 정렬해서 입력
• 모든 token에 위치 임베딩
• 사전 학습 때만 사용 (reconstruction) : 인코더 + 디코더
• Recognition 목적의 이미지 표현을 생성시에는 : 인코더만
• ⇒ 디코더는 인코더와 independent, 작고 연산량 적은 디코더 사용 가능! : 사전 학습 시간 크게 단축

Reconstruction target

• 마스킹된 각 패치의 픽셀값을 예측
• 디코더 출력 (embedding) → Linear layer → 픽셀 값 벡터 (패치 내 픽셀 수)
• 디코더 output → img로 복원
• pixel 공간에서 loss 계산 : 마스킹된 패치에 대해서만

Simple implementation

• 구현상 shuffle/unshuffle만 사용
• 희소 행렬 연산(sparse operation) 없이 순서만 바꿔 처리하므로 구현 단순

 

4. ImageNet Experiments

• Self-supervised pre-training
  • ImageNet-1K 데이터 라벨 없이 MAE로 사전 학습
  • 여기서는 encoder가 좋은 representation을 배우게 하는 단계
  • 분류 같은 downstream task는 풀지 않음
• Supervised evaluation
  • 사전학습된 모델의 성능을 downstream task(여기서는 ImageNet 분류)에서 평가
  1. End-to-end fine-tuning → pre-trained encoder + classifier 전부 학습
  2. Linear probing → pre-trained encoder freeze, 마지막 classifier만 학습
• Baseline: ViT-Large
   

  학습 방식	Epoch 수	Accuracy	특징
  ViT-L scratch (original)	200	76.5%	기본 세팅, 성능 낮음
  ViT-L scratch (strong reg.)	200	82.5%	정규화로 성능 향상
  MAE pretrain → fine-tune	50	84.9%	더 짧은 학습, 더 높은 성능

  • 모두 ViT-Large 모델을 ImageNet-1K에 대해 학습한 결과
  1. Vision Transformer 논문([16])에서 사용한 원래 구현 그대로, 완전히 무작위 초기화로 시작해서 supervised learning만 진행
     • 성능이 낮은 이유: ViT-L은 파라미터 수가 많아서, 데이터 양이 충분해도 처음부터 학습하면 최적화가 어렵고 overfitting이 심해짐
  2. 여전히 무작위 초기화로 시작하지만, ViT에서 사용한 레시피 대신 이 논문에서 새로 조정한 학습 세팅을 사용
     • 같은 scratch 학습이지만, 원래 구현보다 성능이 크게 향상(76.5 → 82.5%)
  3. 무작위 초기화 대신, MAE(self-supervised)로 사전학습된 encoder 가중치를 불러와서 supervised fine-tuning 진행
     • 사전학습 덕분에 모델이 이미지 구조와 패턴을 이미 잘 이해하고 있어서, supervised 학습 시 더 빠르고 효율적으로 최적화 가능
• scratch ↔ fine-tuning : 사전 학습이 얼마나 도움이 되는지 확인
  • scratch → 사전 학습 없이 처음부터 supervised로 학습
• fine-tuning ↔ linear probing : 사전 학습 representation이 얼마나 좋은지

 

4.1. Main Properties

 

Masking ratio

[Fig 5] y축 : ImageNet-1K validation accuracy (%)

• lin & fin 모두 75%가 + (BERT는 15%)
• 전체 맥락(gestalt)을 이해해야만 복원 가능 → 더 의미 있는 표현 학습
• lin → 마스킹 비율이 높아질수록 성능이 증가하다가 75%에서 최고점 (acc gap이 54.6% → 73.5%, 약 20% 향상)
• fin → 마스킹 비율에 덜 민감 (40~80%에서 모두 좋은 성능), 모두 scratch (82.5%)보다 높음

 

Decoder design

 

 

• depth(Transformer block 개수)
  • lin → decoder depth 중요
    • autoencoder의 마지막 layer : reconstruction에 특화 vs recognition(분류)
    • 디코더 깊으면, reconstruction 처리를 디코더에서 해결, 인코더 표현은 더 추상적이고 recognition 친화적으로 유지 가능
    • 성능 최대 8% 향상
  • fin → 영향이 작음
    • 인코더 마지막 layer를 recognition task에 맞게 재학습 가능
  • 1-block 디코더도 fin에서는 높은 성능(84.8%)
    • mask token → visible token 정보 전달만 하면 되기 때문(최소 요건)
• width(채널 수)
  • 기본값: 512
  • fin & lin 둘 다 +
  • 더 좁은 디코더도 fin에서는 +

 

Mask token

 

 

• 인코더에서 mask token 사용하면
  • lin에서는 14% -
    • 사전 학습과 실제 적용 사이 gap 때문
    • 실제 사용(Downstream task)에서는 mask token이 없음
  • 결국 인코더에서 real 패치만 보도록 → 성능 +
• 속도↑
  • self-attention 연산량이 대폭 감소
  • FLOPs 3.3배 절감 → 실시간 속도 2.8배 향상
  • masking 비율 높으면 속도 차이가 4배 이상 나기도 함
• 메모리 절감
  • token 개수 감소 → 메모리 사용량 감소 → 더 큰 모델이나 배치 학습 가능

 

Reconstruction target

 

 

1. 기본적으로는 per-patch norm 없이 각 패치의 픽셀 값 복원
2. Per-patch norm
   • 각 패치 내부 픽셀을 평균 0, 표준편차 1로 정규화한 뒤 복원
   • 패치 내 contrast 증가 → 정확도 향상
   ⇒ “밝기 정보”라는 쉽게 복원 가능한 저주파 신호를 제거하고, 복원 난이도가 높은 디테일 신호를 강조해서 representation을 더 잘 학습하도록!

참고..

3. PCA 수행 후 복원

• 고차원 데이터를 저차원으로 변환하여 데이터의 주요 패턴을 찾고 시각화하는 데 사용되는 통계 기법
• 각 패치를 PCA로 변환하고, 상위 96개 주성분만 복원 목표로 사용.
• PCA는 저주파(large-scale structure) 성분이 먼저, 고주파(디테일)는 뒤에
  • 상위 성분은 데이터의 큰 분산(저주파/큰 구조/평균 밝기 변화)
  • 하위 성분은 경계·질감 같은 고주파 디테일
• 상위 주성분만 쓰면 → 디테일 손실 → MAE 성능 저하
• 따라서 MAE는 고주파 정보가 중요!

4. Token 예측 (BEiT 방식)

• 픽셀 대신 token 복원

1. DALL·E 사전 학습된 dVAE로 이미지를 토큰화
2. MAE 디코더가 마스크된 위치의 토큰 ID를 예측 (CE loss)

• 결과:
  • fin → 0.4% 향상 (token ↔ unnormalized pixels)
    • 비정규화 픽셀은 저주파(밝기, 평균 색) 성분이 그대로 있어서, 토크나이저로 변환하면 그 영향을 줄이고 semantic 단서를 강화 → 미세하게 성능 향상
  • 성능 차이 없음 (token ↔ normalized pixels)
  • lin에서는 오히려 성능 하락
• 추가 단점:
  • dVAE 자체 사전 학습 필요(추가 데이터 필요: 2.5억 이미지)
  • dVAE 인코더가 크고(FLOPs의 40%) 연산 부담 큼
  • 구현 복잡성 증가

 

Data augmentation

 

 

• Cropping-only (고정 or 랜덤 크기) : 성능+
• Cropping + Color jittering : 성능 -
• No augmentation (센터 crop만, flip 없음) : 꽤 좋은 성능
  • contrastive learning에서는 절대 불가능
    • BYOL, SimCLR 같은 contrastive learning → augmentation 필수
      • Crop only로 바꾸면 → 정확도 13% ~ 28% 감소
      • trivial solution(특징이 너무 쉬워서 의미 없는 학습)로 빠짐
      • contrastive는 같은 이미지의 다른 뷰끼리 embedding을 가깝게 만들고, 다른 이미지끼리는 멀게 만드는 게 목표.
      • Augmentation 없으면 "다른 뷰"가 사실상 같은 이미지라 학습이 무너짐.

• MAE가 Augmentation에 덜 의존하는 이유
  • 무작위 masking 자체가 augmentation 역할
  • 각 iteration마다 mask 패턴이 다르기 때문에, 같은 이미지여도 매번 다른 입력이 되는 것

 

Mask sampling strategy

 

 

1. Block-wise masking
   1. 75%로 가면 성능 하락
   2. 큰 덩어리로 가려서 복원 난이도가 급상승, 학습 loss↑, blurry해짐
2. Grid-wise masking
   1. 복원 난이도 낮음, loss↓, 재구성 결과도 sharp해서 나쁘지 않지만~?
   2. but representation 품질 낮음(즉, downstream task 성능↓)

⇒ random masking이 가장 효율적이다!

 

Training schedule

 

 

• 기본 setting : 800 epoch 사전 학습
• 학습 길이가 길어질수록 lin 성능 계속 증가
  • 1600 epoch에서도 성능 포화(saturation) x
• MoCo v3(Contrastive learning)는 약 300 epoch에서 이미 성능 포화
• 이유:
  • MAE encoder는 한 epoch에 전체 패치의 25%만 봄
  • contrastive는 200% 이상 봄 (2 crop, multi-crop)
  • 즉, MAE는 동일 epoch 수라도 실제 본 데이터 양이 훨씬 적기 때문에 더 긴 학습 필요

 

4.2. Comparisons with Previous Results

 

 

Comparisons with self-supervised methods

• ViT-B
  • 대부분의 self-supervised 방법들이 비슷한 fine-tuning 성능을 보임.
• ViT-L
  • 방법 간 성능 차이가 커짐 → 큰 모델일수록 overfitting될 위험 높음, 좋은 pre-training 방법을 쓰면 모델의 큰 용량이 진짜 일반화 성능으로 이어짐
• MAE 특징
  • scale up well
  • ViT-H(224×224) fine-tuning → 86.9%
  • ViT-H(448×448) fine-tuning → 87.8%
    • IN1K 데이터만 사용
  • 기존 SOTA(512, 87.1%)보다 0.7% 향상
  • • SOTA는 복잡한 고급 네트워크 사용, MAE는 vanilla ViT로 달성
• BEiT와 비교
  • 정확도는 MAE가 더 높음.
  • BEiT는 token prediction인데, ViT-B로 픽셀 재구성을 시도하면 성능이 1.8% 떨어짐
    • 추가적인 tokenizer(dVAE) 사전학습 필요함
  • MAE는 애초에 픽셀 기반 재구성이라 성능 저하 없음
  • 속도
    • MAE가 BEiT보다 epoch당 3.5배 빠름
    • 이유: Table 1c에서 설명한 것처럼 encoder에서 mask token을 제거해 연산량과 메모리 사용량을 크게 줄였기 때문.

Comparisons with supervised pre-training

• Supervised pre-training → 라벨 필요
• MAE → 라벨 없이 self-supervised pre-training

1. ViT-L degrades when trained in IN1K
   1. 데이터 양(1.28M 이미지)이 ViT-Large의 큰 용량을 일반화할 만큼 충분하지 않음 → overfitting 발생
2. ViT-L supervised 학습을 자체적으로 더 잘 구현했으나 성능이 일정 수준 이상 올라가지 않고, plateau

⇒ ViT-L, ViT-H 같이 파라미터 수가 많으면 supervised는 overfitting 경향이 있고 성능이 정체됨

1. 반면 MAE는(self-supervised pre-training → fine-tuning)
   1. 일반화 성능이 supervised보다 더 높음
   2. 모델이 커질수록 성능이 계속 올라감
2. JFT-300M
   1. ViT 원 논문에서 대규모 데이터셋 JFT-300M(3억 장)으로 supervised pre-training 후 fine-tuning하면 모델 크기가 클수록 성능이 더 크게 향상됨
   2. MAE도 데이터 크기를 늘리지 않고 IN1K만으로 비슷한 scaling-up 효과를 냄

 

정리!

Vision 분야는 지금까지 대부분 감독 학습 기반이었고, 자기지도 학습은 제한적으로만 쓰였음. 그러나 NLP처럼 간단한 자기지도 학습 → 대규모 모델 확장 → 지속적인 성능 향상의 길로 갈 가능성이 있음.