[논문 리뷰] MixMAE: Mixed and Masked Autoencoder for Efficient Pretraining of Hierarchical Vision Transformers

Jihao Liu, Xin Huang외 3명|arXiv (Cornell University)|2022. 05. 26.

Advanced Neural Network Applications인용 수 7

한 줄 요약

MixMAE는 [MASK] 토큰을 사용하지 않고 두 이미지의 토큰을 혼합하여 마스킹된 패치를 대체함으로써 계층적 비전 트랜스포머의 효율적인 사전학습을 위한 새로운 마스킹된 자동에코딩 프레임워크를 제안한다. 이는 이중 재구성 기능을 가능하게 하며, FLOPs/성능 트레이드오프 측면에서 기존의 MIM 방법보다 뛰어나다. ImageNet-1K에서 85.1%의 top-1 정확도를 달성했으며, 사전학습에 단지 600 에포크만 소요된다.

ABSTRACT

In this paper, we propose Mixed and Masked AutoEncoder (MixMAE), a simple but efficient pretraining method that is applicable to various hierarchical Vision Transformers. Existing masked image modeling (MIM) methods for hierarchical Vision Transformers replace a random subset of input tokens with a special [MASK] symbol and aim at reconstructing original image tokens from the corrupted image. However, we find that using the [MASK] symbol greatly slows down the training and causes pretraining-finetuning inconsistency, due to the large masking ratio (e.g., 60% in SimMIM). On the other hand, MAE does not introduce [MASK] tokens at its encoder at all but is not applicable for hierarchical Vision Transformers. To solve the issue and accelerate the pretraining of hierarchical models, we replace the masked tokens of one image with visible tokens of another image, i.e., creating a mixed image. We then conduct dual reconstruction to reconstruct the two original images from the mixed input, which significantly improves efficiency. While MixMAE can be applied to various hierarchical Transformers, this paper explores using Swin Transformer with a large window size and scales up to huge model size (to reach 600M parameters). Empirical results demonstrate that MixMAE can learn high-quality visual representations efficiently. Notably, MixMAE with Swin-B/W14 achieves 85.1% top-1 accuracy on ImageNet-1K by pretraining for 600 epochs. Besides, its transfer performances on the other 6 datasets show that MixMAE has better FLOPs / performance tradeoff than previous popular MIM methods. Code is available at https://github.com/Sense-X/MixMIM.

연구 동기 및 목표

• 계층적 비전 트랜스포머를 위한 마스킹된 이미지 모델링(MIM)에서 사전학습의 비효율성과 사전학습-微조정 불일치 문제를 해결하기 위해.
• 에코더에서 [MASK] 토큰의 사용을 제거함으로써 학습 속도 저하와 미세조정 시 분포 불일치 문제를 해결하기 위해.
• [Swin Transformer]와 같은 계층적 ViT와 호환되며, 높은 표현 품질을 유지하면서도 학습 효율성을 향상시키는 방법을 개발하기 위해.
• 기존의 MIM 방법들인 SimMIM과 MAE와 비교해 더 나은 FLOPs/성능 트레이드오프를 달성하기 위해.

제안 방법

• MixMAE는 한 이미지의 마스킹된 패치를 다른 이미지의 가시 패치로 대체하여 혼합 입력을 생성함으로써 에코더에서 [MASK] 토큰의 사용을 피한다.
• 계층적 비전 트랜스포머 에코더는 혼합 입력을 처리하고 두 원본 이미지 모두에 대한 은닉 표현을 생성한다.
• 디코딩 이전에 은닉 표현을 다시 분리하고 [MASK] 토큰으로 채워 이중 재구성 기능을 가능하게 한다.
• 경량 디코더는 혼합 표현에서 두 원본 이미지를 재구성하며, 특징 품질 향상을 위해 이중 재구성 손실을 사용한다.
• 다양한 마스킹 비율을 지원하기 위해 여러 이미지를 혼합할 수 있으며, 75% 마스킹 비율(4개의 이미지 혼합)에서 최적의 성능을 달성한다.
• 디코더에서는 마스킹된 자기주의 주의를 사용하고 에코더에서는 [MASK] 토큰을 사용하지 않아 사전학습-微조정 일치성을 보장한다.

Figure 1 : Overview of MixMAE. For pretraining, two images are mixed with a random mixing mask to create a mixed image. MixMAE takes the mixed image as input and reconstructs the two original images. Right before decoding, the token embeddings are unmixed and filled with mask tokens for dual reconst

실험 결과

연구 질문

• RQ1 계층적 비전 트랜스포머를 위한 마스킹된 자동에코딩 접근법을 설계할 수 있는가? 이는 에코더에서 [MASK] 토큰의 사용을 피함으로써 학습 효율성과 일관성을 향상시킬 수 있는가?
• RQ2 다른 이미지의 패치로 마스킹된 패치를 대체하는 것이 [MASK] 토큰이나 학습 가능한 벡터를 사용하는 것보다 표현 학습에 어떤 영향을 미치는가?
• RQ3 높은 성능를 달성하면서도 최소한의 계산 비용으로 계층적 ViT에서 최적의 마스킹 비율과 혼합 전략은 무엇인가?
• RQ4 혼합 입력에서의 이중 재구성은 단일 이미지 재구성보다 더 나은 다운스트림 전이 성능을 제공하는가?
• RQ5 기존의 MIM 방법들보다 적은 사전학습 에포크 수로 MixMAE는 강력한 성능을 달성할 수 있는가?

주요 결과

• MixMAE는 단지 600 에포크의 사전학습만으로 ImageNet-1K에서 85.1%의 top-1 정확도를 달성하여 높은 효율성과 뛰어난 성능을 입증한다.
• 900 에포크의 사전학습을 거친 MixMAE는 ADE20K에서 51.0 mIoU를 기록하여 조밀한 예측 작업으로의 일반화 능력이 뛰어나다는 것을 보여준다.
• 4개의 이미지를 혼합하여 75% 마스킹 비율을 달성한 결과, 동일한 비율에서 [MASK] 토큰을 사용한 경우보다 성능이 뛰어나 혼합 전략의 효과를 확인한다.
• 이중 재구성은 특히 ADE20K에서 성능 향상을 크게 이끌어내며, 단일 재구성 대비 mIoU를 2.6 포인트 향상시킨다.
• 이미지 분류, 객체 검출, 세그멘테이션을 포함한 일곱 개인 다운스트림 벤치마크에서 MixMAE는 SimMIM과 MAE보다 FLOPs/성능 트레이드오프 측면에서 뛰어난 성능을 보인다.
• 에코더에서 [MASK] 토큰을 제거함으로써 사전학습 효율성이 뛰어나지며, 계산 낭비를 줄이고 사전학습-微조정 불일치 문제를 방지한다.

Figure 2 : Tradeoffs of FLOPs vs. (left) top-1 accuracy on ImageNet-1K, (middle) AP box on COCO, (right) and mIoU on ADE20K. All results are from various self-supervised pretraining methods followed by supervised finetuning. All entries on COCO [ 31 ] use Mask RCNN [ 22 ] framework. All entries on A