[논문 리뷰] FastViT: A Fast Hybrid Vision Transformer using Structural Reparameterization

Pavan Kumar Anasosalu Vasu, James Gabriel외 3명|arXiv (Cornell University)|2023. 03. 24.

Advanced Neural Network Applications인용 수 45

한 줄 요약

FastViT는 RepMixer를 도입하여 빠른 하이브리드 비전 트랜스포머를 재매개변수화 가능한 토큰 믹서로 구현하고, 여러 비전 작업을 지원하는 동시에 모바일 및 GPU 플랫폼 전반에서 지연-정확도 트레이드오프에서 우수한 성능을 달성합니다.

ABSTRACT

The recent amalgamation of transformer and convolutional designs has led to steady improvements in accuracy and efficiency of the models. In this work, we introduce FastViT, a hybrid vision transformer architecture that obtains the state-of-the-art latency-accuracy trade-off. To this end, we introduce a novel token mixing operator, RepMixer, a building block of FastViT, that uses structural reparameterization to lower the memory access cost by removing skip-connections in the network. We further apply train-time overparametrization and large kernel convolutions to boost accuracy and empirically show that these choices have minimal effect on latency. We show that - our model is 3.5x faster than CMT, a recent state-of-the-art hybrid transformer architecture, 4.9x faster than EfficientNet, and 1.9x faster than ConvNeXt on a mobile device for the same accuracy on the ImageNet dataset. At similar latency, our model obtains 4.2% better Top-1 accuracy on ImageNet than MobileOne. Our model consistently outperforms competing architectures across several tasks -- image classification, detection, segmentation and 3D mesh regression with significant improvement in latency on both a mobile device and a desktop GPU. Furthermore, our model is highly robust to out-of-distribution samples and corruptions, improving over competing robust models. Code and models are available at https://github.com/apple/ml-fastvit.

연구 동기 및 목표

• 모바일 기기와 데스크톱에서 정확도와 지연 사이의 균형을 맞춘 효율적인 비전 모델을 고안합니다.
• 합성곱과 트랜스포머의 강점을 활용하는 하이브리드 아키텍처를 개발합니다.
• 스킵 연결의 구조적 재매개변화를 통해 메모리 액세스 비용을 감소시킵니다.
• 학습 시 과매개화와 대형 커널을 사용하여 지연 페널티 없이 모델 용량을 확장합니다.
• 분류, 검출, 분할 및 3D 손 메쉬 추정과 같은 다양한 태스크에서 견고성과 일반화를 입증합니다.

제안 방법

• 스킵 연결을 제거하고 추론 시 재매개변수화 가능한 토큰 믹서인 RepMixer를 도입합니다.
• 밀집 k×k 합성곱을 분해된 깊이별과 포인트와이즈 합성곱으로 대체하고, 선형 학습 시 과매개화로 이를 보조합니다.
• 초기 자기-주 의(attention)를 대체하기 위해 FFN과 패치 임베딩 계층에 대형 커널 합성곱을 도입합니다.
• 깊이별 합성곱으로 생성된 조건부 위치 인코딩을 사용합니다.
• 스탬, 패치 임베딩, 그리고 프로젝션 계층에서 과매개화로 학습하여 용량을 증가시킵니다.

실험 결과

연구 질문

• RQ1 정밀 재매개변수화 가능한 토큰 믹서가 정확도를 손실하지 않으면서 메모리 액세스 비용과 지연을 줄일 수 있을까요?
• RQ2 선형 학습 시 과매개화가 분해된 합성곱 설계에서 정확도를 향상시킬까요?
• RQ3 초기 단계의 대형 커널 합성곱이 하이브리드 아키텍처에서 자기 주의(attention)와 비교해 지연 친화적 정확도 향상을 제공합니까?
• RQ4 실제 지연 제약 하에서 이미지 분류, 검출, 분할 및 3D 손 메쉬 추정에서 FastViT의 성능은 어떠합니까?
• RQ5 모델이 경쟁 아키텍처에 비해 오염 및 분포 밖 입력에 대해 견고합니까?

주요 결과

• FastViT는 모바일(iPhone 12 Pro) 및 데스크탑 GPU(RTX-2080Ti)에서 우수한 지연-정확도 트레이드오프를 달성하면서도 경쟁력 있는 정확도를 유지합니다.
• ImageNet-1k에서 Top-1 83.9%인 FastViT-S12는 모바일에서 유사 FLOPs 기준으로 EfficientNet-B5보다 4.9배, ConvNeXt-B보다 1.9배 더 빠르고 GPU에서 EfficientNetV2-S보다 1.6배 빠릅니다.
• FastViT-S12는 iPhone에서 MobileOne-S4보다 26.3% 빠르고 GPU에서 26.9% 빠르며 유사 정확도를 보이는 반면; FastViT-MA36은 매개변수 및 FLOPs를 감소시킨 상태에서 여러 최신 모델과 동등하거나 이를 능가합니다.
• RepMixer는 스킵 연결 제거로 메모리 액세스 비용을 감소시키고 입력 해상도가 클수록 특히 낮은 지연을 제공합니다(예: 384×384 및 1024×1024).
• 스탬, 패치 임베딩, 그리고 프로젝션 계층에서의 학습 시 과매개화는 정확도 향상을 가져오며(예: ImageNet에서 Top-1 최대 0.9% 증가), 학습 시간 부담은 작습니다.
• FFN 및 패치 임베딩 계층의 대형 커널 합성곱은 지연 영향이 미미한 범위에서 견고성 및 정확도 향상을 제공합니다.