Abstract
본 논문은 Masked Autoencoders (MAE)가 컴퓨터 비전 분야에서 확장 가능한 self-supervised 학습 모델임을 보여준다. 제안하는 MAE 방식은 간단하다: 입력 이미지의 일부 패치를 무작위로 masking하고, 누락된 픽셀을 reconstruction한다.
이 방식은 두 가지 핵심 설계에 기반한다. 첫째, 인코더는 마스킹되지 않은 패치들만 처리하고, 디코더는 잠재 표현과 마스크 토큰을 이용해 원래 이미지를 복원하는 비대칭 인코더-디코더 구조를 설계한다.
둘째, 입력 이미지의 높은 비율(예: 75%)을 마스킹하는 것이 의미 있는 자기지도 학습 과제를 만들어냄을 발견했다.
이 두 설계를 결합함으로써 모델 학습을 더욱 빠르고 효율적으로 수행할 수 있으며, 학습 속도는 최대 3배 이상 빨라지고, 정확도도 향상된다. 이러한 확장 가능한 접근법을 통해 높은 용량의 모델을 우수한 일반화 성능으로 학습할 수 있으며, 예를 들어 기본 ViT-Huge 모델은 ImageNet-1K만을 사용한 방법 중 최고 정확도를 달성한다. 또한 다양한 다운스트림 작업에서의 전이 학습 성능도 지도학습 방식보다 우수하며, 확장성 있는 특성을 보인다.
Introduction
딥러닝은 최근 몇 년 동안 능력과 규모가 지속적으로 증가하는 다양한 아키텍처의 폭발적인 발전을 이루어왔다. 하드웨어 성능의 비약적 발전에 힘입어, 오늘날의 모델들은 백만 장 규모의 이미지 데이터셋에서는 쉽게 과적합해버리고, 점점 더 많은 수억 장의 라벨링된 이미지가 필요해지고 있다.
이러한 데이터에 대한 높은 요구는 NLP 분야에서 self-supervised pretraining을 통해 해결되어 왔다. GPT의 자기회귀 언어 모델링, BERT의 마스킹 기반 자동인코딩 같은 접근법은 개념적으로 매우 간단하다: 입력의 일부를 제거하고, 제거된 내용을 예측하도록 학습하는 방식이다. 이러한 방식은 현재 수백억 개의 파라미터를 가진 대규모 NLP 모델을 학습할 수 있게 만들고 있다.
마스킹 기반 자동인코더(masked autoencoder)라는 아이디어는 보다 일반적인 잡음제거 자동인코더(denoising autoencoder)의 한 형태로, 컴퓨터 비전에서도 자연스럽게 적용할 수 있다. 실제로 비전 분야에서는 BERT 이전부터 관련된 연구들이 진행된 바 있다. 하지만 BERT의 성공 이후 관심이 커졌음에도, 비전 분야의 자동인코딩 연구는 NLP에 비해 더디게 발전해왔다. 그러므로 우리는 이렇게 질문한다: 비전과 언어에서 masked autoencoding은 무엇이 다른가?
우리는 이 질문을 다음 세 가지 관점에서 살펴본다.
① 최근까지 비전과 NLP는 아키텍처가 달랐다. 비전에서는 지난 10년간 CNN이 절대적 우위를 점해왔다. 그러나 CNN은 규칙적인 그리드 위에서 연산하며, 마스크 토큰이나 위치 임베딩처럼 토큰 기반의 신호를 자연스럽게 통합하기가 어렵다. 이러한 아키텍처 격차는 Vision Transformer(ViT)의 등장으로 해소되었고, 더 이상 장애물이 아니다.
② 언어와 비전은 정보 밀도가 다르다. 언어는 사람이 생성한 매우 고도의 의미 정보를 담고 있어, 문장에서 몇 단어만 가려도 모델은 복잡한 언어 이해 능력을 학습하게 된다. 반면 이미지는 공간적 중복이 매우 크기 때문에, 한 패치가 사라져도 단순히 주변 패치만으로 충분히 복원되는 경우가 많다. 이를 극복하기 위해, 우리는 비전에서는 무작위 패치를 매우 높은 비율로(예: 75%) 마스킹하는 단순한 전략이 효과적임을 보였다. 이 전략은 중복을 크게 줄이고, 단순한 저수준 통계가 아닌 전체적 이해가 필요한 어려운 자기지도 과제를 만들어준다.
③ 자동인코더의 decoder는 언어와 비전에서 근본적으로 다른 역할을 한다. 비전에서는 디코더가 픽셀을 다시 복원하므로, 그 출력은 고수준 의미보다는 저수준 시각 정보에 더 가깝다. 반면 언어에서는 디코더가 의미 있는 단어를 직접 예측한다. BERT에서는 디코더가 단순한 MLP로도 충분하지만, 비전에서는 디코더 설계가 잠재 표현의 의미 수준을 결정하는 데 중요함을 발견했다.
이러한 분석을 바탕으로, 우리는 단순하면서도 효과적이며 확장 가능한masked autoencoder(MAE) 비전 표현학습 방법을 제안한다. MAE는 입력 이미지에서 무작위 패치를 제거하고, 제거된 패치를 픽셀 단위로 복원한다. MAE는 비대칭적 인코더–디코더 구조를 가진다. 인코더는 마스킹되지 않은 패치만 처리하고, 디코더는 잠재 표현과 마스크 토큰을 이용해 전체 이미지를 복원한다. 이러한 구조는 디코더로 계산 부담을 이동시켜 전체 연산량을 크게 줄인다.
이 설계 덕분에, 매우 높은 마스킹 비율(예: 75%)을 사용하면서도 높은 정확도를 유지할 수 있으며, 인코더는 전체 패치의 작은 부분(25% 정도)만 처리하면 된다. 이는 사전학습 시간을 3배 이상 줄이고, 메모리 사용량도 줄여 더 큰 모델로의 확장을 가능하게 한다.
MAE는 일반화 성능이 뛰어난 고용량 모델을 학습한다. MAE 사전학습을 통해 ViT-Large, ViT-Huge 같은 대규모 모델도 ImageNet-1K에서 더 나은 성능을 발휘한다. 예를 들어, 기본 ViT-Huge 모델은 ImageNet-1K만 사용한 방법 중 최고 정확도를 달성한다.
또한 객체 탐지, 인스턴스 분할, 의미 분할 등 다양한 다운스트림 작업에서도 지도학습 대비 더 나은 전이 성능을 보이며, 모델 규모가 커질수록 성능 향상이 더욱 뚜렷하게 나타나는 확장성을 확인했다. 이러한 경향은 NLP에서 자기지도 사전학습이 성공적으로 확장되어온 흐름과 매우 유사하다.
Related Work
masked language modeling과 autoregressive 기반 방식들(BERT, GPT 등)은 NLP 분야에서 매우 성공적인 사전학습 기법이다. 이러한 방법들은 입력 시퀀스의 일부를 가리고, 가려진 내용을 예측하도록 모델을 학습시킨다. 이 방식들은 매우 뛰어난 확장성을 보이는 것으로 알려져 있으며, 다양한 다운스트림 작업에서 우수한 일반화 성능을 보인다는 강력한 증거들이 존재한다.
autoencoding은 표현 학습에서 오래된 고전적 방법이다. 자동인코더는 입력을 잠재 표현으로 변환하는 인코더와, 이를 다시 원래 입력으로 복원하는 디코더로 구성된다. 예를 들어 PCA와 k-means조차도 자동인코더의 한 형태로 볼 수 있다. Denoising Autoencoder(DAE)는 입력 신호에 잡음을 넣고 원래 신호를 복원하도록 학습하는 자동인코더의 한 종류이다. 픽셀 마스킹, 색 채널 제거와 같은 다양한 종류의 손상을 적용한 방법들은 모두 DAE의 확장으로 볼 수 있다. 우리의 MAE는 이러한 denoising autoencoding의 한 형태이지만, 기존 DAE와는 여러 면에서 다르다.
Masked image encoding은 마스킹된 이미지로부터 표현을 학습하는 방식이다. DAE에서 마스킹을 일종의 잡음으로 활용하는 초기 연구가 제안됐다. Context Encoder는 CNN을 활용하여 큰 결손 영역을 인페인팅하는 방법을 제시했다. NLP의 성공에서 영감을 받아, 이후 연구들은 Transformer 기반 접근을 사용한다.
컴퓨터 비전 분야의 자기지도 학습은 오랫동안 다양한 프리텍스트 과제를 중심으로 연구되어 왔다. 최근에는 특히 contrastive learning이 큰 인기를 얻었다. 대조 학습은 두 개 이상의 시각적 뷰 간의 유사성 또는 비유사성을 모델링하는 방식이며, 경우에 따라서는 유사성만 모델링하기도 한다. 이러한 대조학습 계열 방법들은 강력한 데이터 증강에 크게 의존한다. 반면 자동인코딩 방식은 근본적으로 다른 방향을 취하며, 이후 논문 전반에서 보여주듯 전혀 다른 학습 특성을 나타낸다.
Approach
우리의 MAE(masked autoencoder)는 일부분만 관찰된 입력으로부터 원래 신호를 복원하는 단순한 자동인코딩 방식이다. 모든 자동인코더와 마찬가지로, 인코더는 관찰된 신호를 잠재 표현으로 변환하고, 디코더는 잠재 표현으로부터 원래 신호를 복원한다.
하지만 기존 자동인코더와 달리, 우리는 비대칭 구조를 채택한다. 즉, 인코더는 관찰 가능한 패치만 처리하며 마스크 토큰 없이 동작하고, 디코더는 잠재 표현과 마스크 토큰을 함께 입력받아 전체 신호를 복원한다.
ⓐ Masking
ViT처럼 이미지를 겹치지 않는 규칙적 패치로 나눈 후, 일부 패치는 선택하고 나머지는 마스킹한다. 샘플링 방식은 단순하다. 균일 분포로 무작위 패치를 선택하고 중복 선택은 없다. 이를 random sampling이라 한다.
높은 마스킹 비율(예: 75%)을 적용하면 이미지 내 중복 정보가 크게 줄어들고, 단순히 주변 패치를 기반으로 예측하기 어려운 도전적인 자기지도 학습 과제가 된다. 또한 균일 분포 샘플링은 이미지 중심에 마스크가 몰리는 중심 편향을 방지한다. 마지막으로, 입력이 크게 희소해지므로 효율적인 인코더 설계가 가능해진다.
ⓑ MAE encoder
인코더는 ViT와 동일한 구조를 따르지만, 마스크되지 않은 보이는 패치들만 처리한다. 패치는 선형 projection + positional embedding을 통해 토큰으로 변환된다. 이 토큰들은 Transformer block들을 통과해 처리된다. 그러나 인코더는 전체 패치 중 작은 비율(약 25%)만 처리한다. 마스크된 패치는 제거되며 마스크 토큰은 사용하지 않는다.
덕분에 매우 큰 인코더도 적은 연산량과 메모리로 학습할 수 있다. 전체 패치는 대신 경량 디코더가 처리한다.
ⓒ MAE decoder
디코더의 입력은 다음 두 종류의 토큰을 모두 포함한다: 인코더가 처리한 보이는 패치, 마스크 토큰
각 마스크 토큰은 학습되는 고정 벡터이며, 여기에 마스킹된 패치가 있음을 나타낸다. 모든 토큰에는 위치 정보를 주기 위해 positional embedding을 더한다. 디코더는 또 다른 Transformer block들의 집합으로 구성된다.
디코더는 pretraining 단계에서만 사용된다. 다운스트림에서는 인코더만 사용한다. 따라서 디코더는 인코더와 독립적으로 매우 작고 얕게 설계할 수 있다. 실제로 기본 디코더는 인코더 대비 토큰당 연산량이 10% 미만이다. 이 비대칭 구조 덕분에 전체 토큰을 디코더에서만 처리하여 사전학습 시간을 크게 단축한다.
ⓓ Reconstruction target
MAE는 각 마스킹된 패치의 픽셀 값을 예측하여 입력을 복원한다. 디코더의 출력 벡터는 하나의 패치를 구성하는 픽셀 값들의 집합이다. 마지막 선형 레이어의 출력 차원 수는 패치에 포함된 총 픽셀 수와 동일하다. 최종적으로 패치들을 재배치하여 복원된 이미지를 만든다. 손실 함수는 재구성된 이미지와 원본 이미지 간의 MSE(평균제곱오차)이며,
BERT와 같이 마스크된 패치에 대해서만 손실을 계산한다.
또 하나의 변형으로, 우리는 정규화된 픽셀 값을 복원 타겟으로 사용하는 방법도 연구했다. 패치 내부의 픽셀 평균과 표준편차를 구해 픽셀 값을 정규화하는 방식이다. 이 정규화된 픽셀을 목표로 하면 표현 품질이 더 좋아지는 것을 확인했다.
ⓔ Simple implementation
MAE는 복잡한 희소 연산 없이도 효율적으로 구현할 수 있다.
- 입력 이미지의 모든 패치를 토큰으로 변환한다.
- 토큰 리스트를 무작위로 섞은 뒤, 마스킹 비율만큼 뒷부분의 토큰을 제거한다.
- 인코더 출력에 마스크 토큰을 붙이고, 섞기 전 순서로 다시 되돌린다.
- 이 전체 토큰 리스트를 디코더에 입력한다.
이렇게 shuffle와 unshuffle 정도의 단순 연산만 필요하며, 오버헤드는 거의 없다.
ImageNet Experiments
우리는 ImageNet-1K 데이터셋에 대해 self-supervised pre-training을 수행하였다. 이후 end-to-end 파인튜닝과 linear probing을 통해 표현 품질을 평가한다. 모든 실험은 224×224 단일 crop의 top-1 validation accuracy로 보고한다.
ⓐ Baseline: ViT-Large
소거 실험을 위해 ViT-Large(ViT-L/16)을 사용한다. ViT-L은 매우 큰 모델이며, 쉽게 과적합되는 경향이 있다. 아래는 ViT-L을 처음부터 학습한 경우 vs MAE 사전학습 후 파인튜닝한 경우의 비교이다.
ViT-L을 ImageNet에서 처음부터 지도학습하는 것은 쉽지 않으며, 강한 정규화 기법이 필요하다. 그럼에도 불구하고 MAE 사전학습은 큰 향상을 보인다. 여기서 파인튜닝은 50 epoch만 수행되며(처음부터 학습은 200 epoch), 이는 성능의 상당 부분이 사전학습에 의존함을 의미한다.
ⓑ Main Properties
-Masking ratio
최적의 비율은 매우 높다(75%). 이는 일반적으로 15%만 마스킹하는 BERT와 대조적이다. 기존 비전 관련 마스킹 연구(20~50%)보다도 훨씬 높다. 모델은 누락된 패치를 다양한 방식으로, 하지만 말이 되는 형태로 복원한다. 이는 단순히 주변 텍스처를 확장하는 것이 아니라 사물과 장면의 전체적 구조를 이해하고 있음을 보여준다.
linear probing과 finetuning은 서로 다른 경향을 보인다. linear probing은 마스킹 비율 증가하면서 정확도가 sweet point 전까지 꾸준히 상승한다.
fine-tuning은 비율에 둔감한 편이며 40~80% 범위가 고르게 잘 작동한다.
-Decoder design
우리의 MAE 디코더는 유연하게 설계할 수 있다.
디코더의 깊이(Transformer 블록 수)를 변화시킨다. 충분히 깊은 디코더는 linear probing에서 중요하다. 이는 픽셀 복원 작업과 인식 작업 사이의 간극으로 설명할 수 있다. 자동인코더의 마지막 몇 개 층은 복원에 더욱 특화되지만, 인식에는 덜 관련된다. 적절히 깊은 디코더는 이러한 복원 특화를 담당하게 하여 잠재 표현을 보다 추상적인 수준에 유지하도록 한다. 이러한 설계는 linear probing에서 최대 8%의 향상을 가져올 수 있다. 그러나 finetuning이 사용될 경우, 인코더의 마지막 층들은 인식 작업에 맞도록 조정될 수 있다. 디코더의 깊이는 finetuning 성능 향상에는 덜 영향을 미친다.
흥미롭게도, 단일 블록 디코더를 가진 MAE도 finetuning에서 강력한 성능을 보일 수 있다. 단일 Transformer 블록은 보이는 토큰에서 마스크된 토큰으로 정보가 전달되기 위한 최소한의 요구사항임을 주목하라. 이렇게 작은 디코더는 학습 속도를 더욱 높일 수 있다.
디코더의 폭(channel 수)을 연구한다. 기본적으로 우리는 512-d를 사용하며, fine-tuning과 linear probing 모두에서 잘 작동한다. 더 좁은 디코더도 fine-tuning에서는 잘 작동한다.
전체적으로 우리의 기본 MAE 디코더는 가볍다. 그것은 8개의 블록과 512-d 폭을 가지며, ViT-L(24 블록, 1024-d)과 비교했을 때 토큰당 FLOPs는 9%에 불과하다. 따라서 디코더가 모든 토큰을 처리하더라도 전체 연산량에서 차지하는 비중은 여전히 작다.
-Mask token
우리의 MAE에서 중요한 설계 중 하나는 인코더에서 마스크 토큰을 생략하고, 그것을 가벼운 디코더에만 적용하는 것이다.
만약 인코더가 마스크 토큰을 사용하면 성능이 떨어진다. linear probing에서 정확도가 14% 감소한다. 이 경우, 사전학습과 실제 배포 간에 간극이 존재한다. 이 인코더는 사전학습 중에는 입력에 많은 마스크 토큰이 포함되지만, 실제 깨끗한 이미지에서는 그런 토큰이 존재하지 않는다. 이 간극은 배포 시 정확도를 떨어뜨릴 수 있다. 인코더에서 마스크 토큰을 제거함으로써 인코더는 항상 실제 패치만 보게 되며, 이는 정확도를 향상시킨다.
또한 인코더에서 마스크 토큰을 생략함으로써 학습 연산량을 크게 줄일 수 있다. 우리는 전체 학습 FLOPs를 3.3배 줄인다. 이는 실제 실행 시간에서 2.8배 속도 향상으로 이어진다. 이렇게 작은 디코더(1 블록), 큰 인코더(ViT-H) 또는 둘 모두를 사용하는 경우 실행 속도 향상은 더 크다.
-Reconstruction target
우리는 다양한 복원 목표를 비교한다. 지금까지의 결과는 정규화되지 않은 픽셀을 기반으로 한다. 정규화된 픽셀을 사용하면 정확도가 향상된다. 이러한 패치별 정규화는 국소적으로 대비를 높인다. 또 다른 변형으로, 우리는 패치 공간에서 PCA를 수행하고 가장 큰 PCA 계수들을 목표로 사용한다. 이렇게 하면 정확도가 떨어진다. 두 실험은 고주파 성분이 우리의 방법에서 유용함을 시사한다.
우리는 MAE의 변형으로 토큰을 예측하는 방식도 비교한다. 이는 BEiT 에서 사용되는 목표이다. 구체적으로, 우리는 DALLE 사전학습된 dVAE를 토크나이저로 사용한다. 이 토크나이제이션은 정규화되지 않은 픽셀에 비해 finetuning 정확도를 0.4% 향상시키지만, 정규화된 픽셀과 비교했을 때는 이점이 없다. 또한 linear probing 정확도를 감소시킨다.
우리의 픽셀 기반 MAE는 토크나이제이션보다 훨씬 단순하다. dVAE 토크나이저는 추가 사전학습 단계가 필요하며 2억 5천만 개의 이미지라는 추가 데이터에 의존한다. dVAE 인코더는 큰 CNN이며 비 trivial한 오버헤드를 추가한다. 픽셀을 사용하면 이러한 문제들을 피할 수 있다.
-Data augmentation
MAE는 크롭만 사용하는 증강에서도 잘 작동한다. color jittering을 추가하면 결과가 저하되므로 우리는 사용하지 않는다. 놀랍게도 MAE는 증강이 전혀 없을 때도 상당히 잘 작동한다. 이는 데이터 증강에 크게 의존하는 contrastive learning과 극명하게 다르다. MAE에서는 증강의 역할이 크게 무작위 마스킹에 의해 대체된다. 마스크는 매 반복마다 달라지므로 자연스러운 변화를 제공한다.
-Mask sampling strategy
Block-wise 마스킹은 큰 블록들을 제거한다. Block-wise 마스킹을 사용한 MAE는 50% 마스킹에서는 괜찮지만 75%에서는 성능이 저하된다. 작업이 너무 어려워지며 복원도 더 흐릿해진다. Grid-wise 샘플링은 네 개 중 한 개의 패치만 유지한다. 이는 더 쉬운 작업으로, 더 선명한 복원을 만들어내지만 표현 품질은 떨어진다. 단순한 random sampling이 가장 잘 작동한다. 이는 높은 마스킹 비율을 허용하며 속도와 정확도 모두에서 이점을 가진다.
-Training schedule
Linear probing은 1600 에포크에서도 포화되지 않는다. MAE 인코더는 한 에포크당 전체 패치의 25%만 보지만, 대조학습 방법들은 multi-crop 증강으로 한 에포크당 200% 이상을 본다.
Discussion and Conclusion
딥러닝의 핵심은 잘 확장되는 간단한 알고리즘들이다. NLP에서는 간단한 자기지도 학습 방법들이 기하급수적으로 확장되는 모델들의 이점을 가능하게 한다. 컴퓨터 비전에서는 자기지도 학습의 발전에도 불구하고 실용적인 사전학습 패러다임은 주로 감독 학습에 기반한다. 본 연구에서 우리는 ImageNet과 전이 학습에서 NLP의 기법과 유사한 간단한 자기지도 방식인 autoencoder가 확장 가능한 이점을 제공함을 관찰한다. 비전 분야의 자기지도 학습은 이제 NLP와 유사한 궤적을 시작할지도 모른다.
한편, 우리는 이미지와 언어는 본질적으로 다른 종류의 신호이며 이 차이를 신중하게 다뤄야 한다는 점도 주목한다. 이미지는 단지 기록된 빛이며, 단어에 해당하는 시각적 유사물로 의미론적 분해가 존재하지 않는다. 객체를 제거하려는 대신, 우리는 의미론적 세그먼트를 형성하지 않을 가능성이 높은 임의의 패치들을 제거한다. 마찬가지로, 우리의 MAE는 의미론적 엔티티가 아닌 픽셀을 재구성한다. 그럼에도 불구하고 우리는 MAE가 복잡하고 전체적인 재구성을 유추함을 관찰하며, 이는 수많은 시각 개념(시각적 의미론)을 학습했음을 시사한다. 우리는 이러한 행동이 MAE 내부의 풍부한 잠재 표현을 통해 발생한다고 가정한다. 이러한 관점이 향후 연구를 고무하길 바란다.
