Masked Autoencoders (MAE)

마스크드 오토인코더(MAE)는 이미지를 위한 자기 지도 학습 방법입니다. 모델은 이미지 패치의 상당 부분(일반적으로 75%)을 무작위로 가리고, 보이는 나머지 부분을 바탕으로 누락된 픽셀을 복원하도록 학습합니다. Meta AI의 Kaiming He와 동료들은 2021년 논문 Masked Autoencoders Are Scalable Vision Learners에서 이 방법을 소개하고, BERT의 마스크 단어 예측을 비전 트랜스포머에 적용했습니다. 인코더는 보이는 패치만 처리하므로 MAE 사전 학습은 모든 패치를 처리하는 것보다 3배 이상 빨랐으며, 이 방식으로 사전 학습한 ViT-Huge 모델은 ImageNet에서 87.8%의 정확도를 달성했습니다. 이는 당시 ImageNet-1K 데이터만 사용한 방법 중 최고 성능이었습니다. 사전 학습된 인코더는 분류, 감지, 분할을 위한 백본으로 사용됩니다.

마스크드 오토인코더의 작동 방식#

MAE는 이미지를 겹치지 않는 패치 그리드로 나눕니다. 예를 들어 224x224 입력은 16x16픽셀 패치 196개로 구성되며, 그중 대부분을 무작위로 가립니다. 인코더와 디코더는 의도적으로 서로 다른 구조를 가집니다. 표준 ViT인 인코더는 보이는 패치만 처리하므로 연산량과 메모리 사용량이 약 4분의 1로 줄어듭니다. 기본적으로 트랜스포머 블록 8개로 구성된 경량 디코더는 토큰당 인코더 연산량의 10% 미만을 사용합니다. 디코더는 인코딩된 가시 패치와 가려진 각 위치에 공통으로 적용되는 학습 가능한 마스크 토큰을 입력받아, 가려진 각 패치의 픽셀 값을 예측합니다. 손실 함수는 가려진 패치에만 적용되는 평균 제곱 오차이며, 패치별로 정규화한 픽셀을 타깃으로 사용하면 학습된 특징이 더욱 개선됩니다. 사전 학습 후에는 디코더를 버립니다:

높은 마스킹 비율이 이 작업을 유용하게 만드는 핵심입니다. 인접한 이미지 패치는 중복성이 매우 높으므로, 몇 개의 패치만 가리면 모델은 주변 패치에서 보간해 각 빈 곳을 채울 수 있습니다. 75%를 가리면 조각만으로 전체 객체와 장면을 추론해야 하므로 모양, 부분, 맥락을 학습하게 됩니다. 반면 BERT는 각 단어가 이미 풍부한 의미를 담고 있기 때문에 단어의 15%만 마스킹합니다.

MAE의 핵심은 마스킹 단계입니다. 여기서는 4개 이미지 배치를 대상으로 PyTorch를 사용해 보여 줍니다:

import torch

torch.manual_seed(0)
images = torch.rand(4, 3, 224, 224)  # a batch of 4 RGB images

# Patchify: each 224x224 image becomes 196 patches of 16x16 pixels, flattened to 768 values
patches = images.unfold(2, 16, 16).unfold(3, 16, 16).permute(0, 2, 3, 1, 4, 5).reshape(4, 196, 768)

# Random masking as in the MAE reference code: rank patches by random noise and keep the first 25%
keep = int(196 * (1 - 0.75))
ids_keep = torch.argsort(torch.rand(4, 196), dim=1)[:, :keep]
visible = torch.gather(patches, 1, ids_keep.unsqueeze(-1).expand(-1, -1, 768))

# Mask of hidden patches (1 = hidden): the decoder predicts these, and the loss is averaged over them only
mask = torch.ones(4, 196).scatter(1, ids_keep, 0)
print(f"Encoder input: {visible.shape[1]} of {patches.shape[1]} patches per image")
print(f"Reconstruction targets: {int(mask[0].sum())} hidden patches per image")

Encoder input: 49 of 196 patches per image과 Reconstruction targets: 147 hidden patches per image을 출력합니다. 따라서 인코더의 트랜스포머 계층은 전체 이미지 길이의 4분의 1에 해당하는 시퀀스를 처리합니다. 어텐션은 모든 토큰을 서로 비교하므로 각 어텐션 점수 행렬의 크기는 패치 196개를 모두 사용하는 경우의 16분의 1입니다.

결과 및 변형#

MAE 특징은 파인튜닝 후 가장 강력해집니다. 선형 프로빙에서는 MAE로 사전 학습한 ViT-L 인코더가 ImageNet 정확도 73.5%를 달성해 대조 학습 기법보다 낮았지만, 파인튜닝하여 COCO 객체 감지 및 분할에 전이했을 때는 MAE 사전 학습이 지도 사전 학습보다 뛰어났습니다. 이 학습 방식은 정지 이미지를 넘어 다양하게 적용되었습니다. VideoMAE는 비디오 클립의 시공간 패치 중 90% 이상을 마스킹하고, SatMAE는 위성 이미지로 사전 학습하며, ConvNeXt V2는 희소 합성곱을 사용해 마스크드 오토인코딩을 합성곱 네트워크에 적용합니다.

마스크드 사전 학습은 널리 사용되는 파운데이션 모델에도 적용됩니다. Ultralytics 패키지에서 실행되는 Meta의 Segment Anything Model (SAM)은 MAE로 사전 학습한 ViT를 이미지 인코더로 사용합니다. 해당 모델의 ViT-B 버전인 sam_b.pt는 Ultralytics auto_annotate의 기본 마스크 모델이며, YOLO26 감지기의 박스를 분할 레이블로 변환합니다.

마스크드 오토인코더와 관련 용어 비교#

• 오토인코더: 전체 입력을 작은 코드로 압축한 후 복원합니다. MAE는 불완전한 입력을 받아 보지 못한 부분을 예측해야 하며, 인코더 출력에는 보이는 패치마다 벡터 하나가 유지됩니다.
• 마스크드 이미지 모델링: MAE가 속하는 더 넓은 범주의 기법입니다. BEiT는 마스크된 패치의 이산 시각 토큰을 예측하고, SimMIM은 인코더에 마스크 토큰을 입력하는 반면, MAE는 픽셀을 예측하고 인코더에서 마스크 토큰을 제외합니다.
• 대조 학습: 동일한 이미지의 증강된 뷰를 서로 일치시키며 학습하고 디코더가 필요하지 않습니다. 따라서 고정된 상태에서 더 잘 작동하는 특징을 생성하는 반면, MAE 특징은 파인튜닝을 통해 더 크게 개선됩니다.
• 공동 임베딩 예측 아키텍처(JEPA): 이미지 영역을 가리는 점은 같지만 픽셀 대신 해당 영역의 임베딩을 예측하므로, 모델이 텍스처 세부 정보에 용량을 소비하지 않도록 합니다.