MixMAE: 계층적 비전 트랜스포머의 효율적인 사전학습을 위한 혼합 및 마스킹 오토인코더

이 논문에서는 다양한 계층적 비전 트랜스포머(Visual Transformers)에 적용 가능한 간단하면서도 효율적인 사전 훈련 방법인 혼합 및 마스킹 오토인코더(MixMAE)를 제안한다. 기존의 계층적 비전 트랜스포머를 위한 마스킹 이미지 모델링(MIM) 방법은 입력 토큰 중 무작위 부분을 특수한 [MASK] 기호로 대체하고, 손상된 이미지에서 원본 이미지 토큰을 재구성하는 것을 목표로 한다. 그러나 우리는 [MASK] 기호를 사용할 경우, 특히 SimMIM에서와 같이 높은 마스킹 비율(예: 60%)을 적용할 경우 훈련 속도가 크게 저하되며, 사전 훈련과 미세 조정 사이의 일관성 문제를 야기함을 발견하였다. 반면, MAE는 인코더 단계에서 [MASK] 토큰을 전혀 도입하지 않지만, 계층적 비전 트랜스포머에는 적용할 수 없다. 이러한 문제를 해결하고 계층적 모델의 사전 훈련 속도를 가속화하기 위해, 한 이미지의 마스킹된 토큰을 다른 이미지의 가시적 토큰으로 대체하여 혼합된 이미지를 생성한다. 이후 이 혼합 입력에서 두 원본 이미지를 동시에 재구성하는 이중 재구성(dual reconstruction)을 수행함으로써 훈련 효율성을 크게 향상시킨다. MixMAE는 다양한 계층적 트랜스포머에 적용 가능하지만, 본 논문에서는 대규모 창 크기와 거대한 모델 규모(최대 6억 파라미터)를 갖춘 Swin 트랜스포머를 사용하여 성능을 확장하였다. 실증 결과에 따르면, MixMAE는 고품질의 시각적 표현을 효율적으로 학습할 수 있음을 보여주며, 특히 Swin-B/W14 아키텍처를 사용하여 ImageNet-1K에서 600 에포크 사전 훈련 시 85.1%의 Top-1 정확도를 달성하였다. 또한, 다른 6개의 데이터셋에 대한 전이 성능 평가 결과, 기존의 인기 있는 MIM 방법들에 비해 FLOPs 대비 성능의 균형이 더 우수함을 확인하였다. 코드는 https://github.com/Sense-X/MixMIM 에서 공개되어 있다.