[논문 리뷰] Uniform Masking: Enabling MAE Pre-training for Pyramid-based Vision Transformers with Locality

Xiang Li, Wenhai Wang외 2명|arXiv (Cornell University)|2022. 05. 20.

Advanced Neural Network Applications인용 수 36

한 줄 요약

Uniform Masking을 통해 MAE-style 사전 학습을 지역성 기반 피라미드 ViT와 통합하여 효율적 사전 학습을 가능하게 하고 다양한 작업에서 강한 파인튜닝 성능을 보존합니다.

ABSTRACT

Masked AutoEncoder (MAE) has recently led the trends of visual self-supervision area by an elegant asymmetric encoder-decoder design, which significantly optimizes both the pre-training efficiency and fine-tuning accuracy. Notably, the success of the asymmetric structure relies on the "global" property of Vanilla Vision Transformer (ViT), whose self-attention mechanism reasons over arbitrary subset of discrete image patches. However, it is still unclear how the advanced Pyramid-based ViTs (e.g., PVT, Swin) can be adopted in MAE pre-training as they commonly introduce operators within "local" windows, making it difficult to handle the random sequence of partial vision tokens. In this paper, we propose Uniform Masking (UM), successfully enabling MAE pre-training for Pyramid-based ViTs with locality (termed "UM-MAE" for short). Specifically, UM includes a Uniform Sampling (US) that strictly samples $1$ random patch from each $2 imes 2$ grid, and a Secondary Masking (SM) which randomly masks a portion of (usually $25\%$) the already sampled regions as learnable tokens. US preserves equivalent elements across multiple non-overlapped local windows, resulting in the smooth support for popular Pyramid-based ViTs; whilst SM is designed for better transferable visual representations since US reduces the difficulty of pixel recovery pre-task that hinders the semantic learning. We demonstrate that UM-MAE significantly improves the pre-training efficiency (e.g., it speeds up and reduces the GPU memory by $\sim 2 imes$) of Pyramid-based ViTs, but maintains the competitive fine-tuning performance across downstream tasks. For example using HTC++ detector, the pre-trained Swin-Large backbone self-supervised under UM-MAE only in ImageNet-1K can even outperform the one supervised in ImageNet-22K. The codes are available at https://github.com/implus/UM-MAE.

연구 동기 및 목표

• 로컬 윈도우를 사용하는 피라미드 기반 ViT에 대해 MAE 스타일의 자체지도 사전 학습을 고무하고 가능하게 한다.
• 효율성을 유지하면서 로컬 윈도우 전반에 걸쳐 균일한 입력 구조를 유지하도록 Uniform Masking을 설계한다.
• UM-MAE가 다운스트림 성능을 보존하거나 향상시키면서 사전 학습 시간과 GPU 메모리를 감소시킴을 입증한다.
• ImageNet-1K 분류, ADE20K 분할, COCO 객체 탐지와 같은 다운스트림 작업에서 UM-MAE가 기존 MIM 방법과 어떻게 비교되는지 조사한다.

제안 방법

• Uniform Sampling (US)는 2x2 격자마다 임의의 패치를 하나 선택하여 25% 가시 패치 세트를 만든다.
• Secondary Masking (SM)은 이미 샘플링된 영역의 일부(예: 25%)를 무작위로 마스킹하고 학습 가능한 마스크 토큰을 사용한다.
• 균일 샘플링된 패치를 피라미드형 ViT 인코더에 입력되는 촘촘한 2D 입력으로 재구성한다.
• 디코더는 MAE의 경량 ViT로 남아 남겨진 패치의 평균 제곱 오차를 통해 원래 이미지 픽셀을 재구성한다.
• 인코더 입력은 토큰의 25%로 축소되며, 피라미드 백본의 해상도를 회복하기 위해 픽셀 셔플이 사용된다.
• 학습은 IN1K, ADE20K, COCO 전반에 걸쳐 UM-MAE를 SimMIM 및 MAE 기준선과 비교하고, 가끔 중간 파인튜닝이 논의된다.

실험 결과

연구 질문

• RQ1 로컬 윈도우를 가진 피라미드 기반 ViT에 MAE 스타일의 사전 학습을 과도한 계산 부담 없이 효과적으로 적용할 수 있는가?
• RQ2 피라미드 기반 아키텍처에 대해 어떤 샘플링 및 마스킹 전략이 전달 가능한 표현을 가장 잘 보존하거나 향상시키는가?
• RQ3 기존 MIM 방법과 비교했을 때 사전 학습 효율성과 다운스트림 작업 정확도 측면에서 UM-MAE의 성능은 어떠한가?
• RQ4 밀집 예측 작업에서 UM-MAE의 전이 이득에 중간 파인튜닝이 영향을 미치는가?

주요 결과

• UM-MAE는 피라미드 기반 ViT의 사전 학습을 크게 가속화(≈2×)하고 GPU 메모리도 감소시킨다(≥2×) SimMIM과 비교해서.
• 25%의 Secondary Masking 비율을 가진 Uniform Sampling은 강력한 트레이드오프를 제공하며 다운스트림 작업에서 MAE 기준선과 동일하거나 이를 상회한다.
• Swin-T의 경우, UM-MAE는 IN1K 82.04 Top-1, ADE20K 45.96 mIoU, COCO 47.7 AP를 설정 간에 달성하며 SimMIM에 비해 메모리/시간 개선이 있다.
• 대형 모델(Swin-L)에서는 IN1K에서 사전 학습된 UM-MAE가 더 적은 사전 학습 에폭으로 감독학습 IN22K 기준선을 능가할 수 있다.
• MIM 하에서 피라미드 기반 ViT의 우수한 다운스트림 성능을 위해 IN1K에서의 중간 파인튜닝이 중요하며, 종종 직접 파인튜닝보다 이득을 준다.
• UM-MAE는 강력한 MIM 기준선에 비해 프리-트레이닝 자원을 줄이면서도 경쟁력 있거나 향상된 다운스트림 성능을 유지한다.