SHViT: Single-Head Vision Transformer with Memory Efficient Macro Design

Recently, efficient Vision Transformers have shown great performance with low latency on resource-constrained devices. Conventionally, they use $4\times 4$ patch embeddings and a 4-stage structure at the macro level, while utilizing sophisticated attention with multi-head configuration at the micro level. This paper aims to address computational redundancy at all design levels in a memory-efficient manner. We discover that using larger-stride patchify stem not only reduces memory access costs but also achieves competitive performance by leveraging token representations with reduced spatial redundancy from the early stages. Furthermore, our preliminary analyses suggest that attention layers in the early stages can be substituted with convolutions, and several attention heads in the latter stages are computationally redundant. To handle this, we introduce a single-head attention module that inherently prevents head redundancy and simultaneously boosts accuracy by parallelly combining global and local information. Building upon our solutions, we introduce SHViT, a Single-Head Vision Transformer that obtains the state-of-the-art speed-accuracy tradeoff. For example, on ImageNet-1k, our SHViT-S4 is $3.3\times, 8.l\times$, and $2.4\times faster$ than MobileViTv2 $\times l.0$ on GPU, CPU, and iPhone12 mobile device, respectively, while being 1.3% more accurate. For object detection and instance segmentation on MS COCO using Mask-RCNN head, our model achieves performance comparable to FastViT-SA12 while exhibiting $3.8\times and2.0\times lower$ backbone latency on GPU and mobile device, respectively.

한 줄 요약

SHViT는 메모리 효율적인 설계를 통해 ImageNet-1k에서 1.3% 더 정확하고 3.3× 빠른 추론 속도를 달성한 Single-Head Vision Transformer이다.

핵심 기여도

16×16 패치 임베딩과 3단계 구조를 도입하여 초기 단계의 공간적 중복을 줄이고 메모리 접근 비용을 16배 감소.
Single-Head Self-Attention(SHSA) 모듈로 헤드 중복을 제거하고, 전역/로컬 정보를 병렬 처리하여 정확도를 0.95% 이하로 저하시키지 않으면서 연산량을 줄임.
SHViT-S4는 GPU에서 MobileViTv2 ×1.0 대비 3.3×, CPU에서 8.1×, iPhone12에서 2.4× 빠른 추론 속도를 보임.
MS COCO에서 Mask-RCNN 헤드를 사용할 때 FastViT-SA12 대비 GPU와 모바일 기기에서 각각 3.8×, 2.0× 더 빠른 백본 지연 시간을 기록.

핵심 아이디어

기존의 효율적인 Vision Transformer(ViT)는 4×4 패치 임베딩과 4단계 구조를 사용하면서, 초기 단계에서 공간적 토큰 수가 많아 메모리 접근 비용이 높고, MHSA의 다중 헤드 구조가 연산적 중복을 유발한다는 문제를 제기한다. SHViT는 이에 대응하기 위해 16×16 패치 임베딩을 도입해 초기 단계에서 토큰 수를 3136개에서 196개로 줄이고, 공간적 중복을 제거함으로써 메모리 효율성을 높인다. 또한, 초기 단계에서는 MHSA 대신 3×3 depthwise convolution을 사용해 토큰 믹싱을 수행하고, 후기 단계에서는 SHSA를 통해 헤드 수를 1개로 줄여 연산 중복을 제거한다. SHSA는 입력 채널의 일부에만 self-attention을 적용하고 나머지는 그대로 유지함으로써 메모리 접근 비용을 줄이며, 전역과 로컬 정보를 병렬 처리하여 정확도를 유지한다.

기술적 접근법

**패치 임베딩**: 16×16 크기의 패치 임베딩을 사용하여 초기 단계에서 토큰 수를 196개로 줄임.
**단계 구조**: 3단계 구조를 채택하여 메모리 접근 비용을 16배 감소.
**토큰 믹싱**: 초기 단계에서는 3×3 depthwise convolution을 사용해 MHSA 대체.
**SHSA 모듈**: 단일 헤드로 MHSA 대체. 입력 채널의 일부에만 self-attention 적용.
**하이퍼파라미터**: SHSA의 partial ratio는 1/4.67로 설정하여 최적의 speed-accuracy tradeoff 달성.

주요 결과

**ImageNet-1k**: SHViT-S4는 MobileViTv2 ×1.0 대비 GPU에서 3.3×, CPU에서 8.1×, iPhone12에서 2.4× 빠르며 1.3% 더 정확.
**MS COCO**: Mask-RCNN 헤드를 사용할 때 FastViT-SA12 대비 GPU에서 3.8×, 모바일 기기에서 2.0× 더 빠른 백본 지연 시간.
**EfficientNet-B0 대비**: SHViT-S4는 GPU에서 69.4% 더 빠르고, CPU에서 90.6% 더 빠르며, 정확도는 2.3% 더 높음.

의의 및 한계

SHViT는 메모리 효율적인 설계를 통해 ViT의 연산 중복 문제를 해결하고, GPU, CPU, 모바일 기기 등 다양한 환경에서 뛰어난 speed-accuracy tradeoff를 달성함으로써 실시간 추론에 적합한 모델로 제시된다. 특히, SHSA 모듈은 헤드 수를 줄이는 동시에 정확도 저하를 최소화하여, 기존 MHSA 대체 기법보다 효율적이다. 그러나 SHViT는 패치 임베딩 크기를 16×16으로 늘림으로써 초기 단계의 토큰 수를 줄이는 방식을 채택했기 때문에, 고해상도 이미지 처리 시 공간 정보 손실이 발생할 수 있는 한계가 있다. 또한, SHSA는 헤드 수를 1개로 줄였기 때문에 복잡한 패턴 인식 능력이 제한될 수 있다.

실용적 활용

SHViT는 모바일 기기, IoT, 실시간 비디오 분석 등 메모리와 연산 자원이 제한된 환경에서 활용 가능하다. 특히, SHViT-S4는 iPhone12에서도 2.4× 빠른 추론 속도를 보이므로 모바일 앱, 자율주행, 스마트 카메라 등에 적용할 수 있다. 또한, SHViT는 이미지 분류, 객체 탐지, 인스턴스 세그멘테이션 등 다양한 CV 작업에 사용 가능하며, 높은 정확도와 빠른 속도를 동시에 요구하는 산업 분야에서 유용하게 사용될 수 있다.