[논문 리뷰] HiViT: Hierarchical Vision Transformer Meets Masked Image Modeling

Xiaosong Zhang, Yunjie Tian외 5명|arXiv (Cornell University)|2022. 05. 30.

Advanced Neural Network Applications인용 수 12

한 줄 요약

HiViT는 계층적인 비전 트랜스포머 아키텍처를 제안하며, 이는 이웃 단위 간의 국소적 연산(예: 시프트된 윈도우 자기주의)을 제거하면서도 다중 척도 특징 맵을 유지함으로써, 평범한 ViT처럼 토큰을 순차적으로 처리함으로써 효율적인 마스킹된 이미지 모델링(MIM)을 가능하게 한다. 이 설계는 MAE 사전학습 하에 ViT-B보다 +0.6% ImageNet-1K 상위-1 정확도를 달성하고, Swin-B보다 1.9배 빠른 학습 속도를 보인다.

ABSTRACT

Recently, masked image modeling (MIM) has offered a new methodology of self-supervised pre-training of vision transformers. A key idea of efficient implementation is to discard the masked image patches (or tokens) throughout the target network (encoder), which requires the encoder to be a plain vision transformer (e.g., ViT), albeit hierarchical vision transformers (e.g., Swin Transformer) have potentially better properties in formulating vision inputs. In this paper, we offer a new design of hierarchical vision transformers named HiViT (short for Hierarchical ViT) that enjoys both high efficiency and good performance in MIM. The key is to remove the unnecessary "local inter-unit operations", deriving structurally simple hierarchical vision transformers in which mask-units can be serialized like plain vision transformers. For this purpose, we start with Swin Transformer and (i) set the masking unit size to be the token size in the main stage of Swin Transformer, (ii) switch off inter-unit self-attentions before the main stage, and (iii) eliminate all operations after the main stage. Empirical studies demonstrate the advantageous performance of HiViT in terms of fully-supervised, self-supervised, and transfer learning. In particular, in running MAE on ImageNet-1K, HiViT-B reports a +0.6% accuracy gain over ViT-B and a 1.9$ imes$ speed-up over Swin-B, and the performance gain generalizes to downstream tasks of detection and segmentation. Code will be made publicly available.

연구 동기 및 목표

• 공간적으로 의존적인 국소적 연산으로 인해 계층적 비전 트랜스포머의 마스킹된 이미지 모델링(MIM)에서 효율성이 떨어지는 문제를 해결한다.
• MIM 효율성을 해치는 '국소적 단위 간 연산'을 제거함으로써 계층적 트랜스포머에서 토큰 순차화를 가능하게 한다.
• 계층적 특징 표현의 성능 이점을 유지하면서도 ViT와 유사한 계산 유연성을 확보한다.
• 자기지도 학습, 완전지도 학습, 그리고 후행 태스크 검출/분할 과제에서 최신 기술 성능을 입증한다.
• 평범한 ViT와 복잡한 계층적 모델(예: Swin 트랜스포머)의 간소하고 효율적이며 MIM 우호적인 대안을 제공한다.

제안 방법

• 계층적 비전 트랜스포머의 연산을 내부 단위, 전역 단위 간, 국소 단위 간 연산으로 분류한다.
• '국소 단위 간 연산'(예: 시프트된 윈도우 자기주의, 패치 병합)이 MIM에서 토큰 순차화를 방해하는 주요 장애물임을 규명한다.
• Swin 트랜스포머의 주 단계에서 모든 국소 단위 간 연산을 제거하고, 전역 자기주의와 계층적 특징 맵만 유지한다.
• 모델의 FLOPs와 능력을 유지하기 위해 초기 단계의 국소 연산을 등가의 내부 단위 MLP로 대체한다.
• 최종 단계를 주 단계와 병합하여 FLOPs를 유지하면서 아키텍처를 단순화한다.
• 마스킹된 단위를 버리고 오직 가시 단위만 처리함으로써 MIM 사전학습 중에 가시 토큰의 완전한 순차화를 가능하게 한다.

실험 결과

연구 질문

• RQ1 성능을 희생시키지 않고 계층적 비전 트랜스포머를 평범한 ViT만큼 MIM에서 효율적으로 만들 수 있는가?
• RQ2 계층적 트랜스포머의 '국소 단위 간 연산'이 인식 정확도에 실제로 기여하는가, 아니면 계층적 설계 자체가 핵심 요소인가?
• RQ3 Swin 트랜스포머에 국소 주의를 제거하는 최소한의 수정만으로도 ViT와 Swin 모두를 능가하는 MIM 효율성과 정확도를 달성할 수 있는가?
• RQ4 결과로 도출된 아키텍처인 HiViT는 자기지도 학습, 완전지도 학습, 후행 검출/분할 과제에서 강력한 성능을 유지하는가?
• RQ5 공간 인식 주의 메커니즘을 요구하지 않고도 MIM에서 계층적 구조를 효과적으로 활용할 수 있는가?

주요 결과

• HiViT-B는 ImageNet-1K에서 83.8%의 상위-1 정확도를 달성하여 ViT-B보다 +0.6% 향상되었고, Swin-B보다 +0.3% 향상되었다.
• MAE 사전학습 하에 HiViT-B는 Swin-B보다 1.9배 더 빠른 학습 속도를 보였으며, 정확도 면에서도 ViT-B를 능가했다.
• HiViT-B의 자기지도 사전학습은 ImageNet-1K에서 84.2%의 상위-1 정확도를 달성하여 지도 학습 기반 모델과 경쟁 가능했다.
• 향상된 특징 피라미드 활용으로 HiViT-B는 COCO 객체 검출에서 51.2% AP^box, 인스턴스 분할에서 44.2% AP^mask를 달성했다.
• ADE20K 세분화 분할 과제에서 HiViT-B는 51.2% mIoU를 기록하여 동일한 사전학습 프rotocol 하에서 ViT-B를 능가했다.
• 절단 실험 결과, 국소 단위 간 연산을 제거해도 성능에 미미한 영향을 미치지 않아 핵심 설계 선택의 타당성을 입증했다.