효율적인 비전 맘바: 숨겨진 상태 믹서를 기반으로 한 상태 공간 이중성

EfficientViM: Efficient Vision Mamba with Hidden State Mixer based State Space Duality

November 22, 2024

저자: Sanghyeok Lee, Joonmyung Choi, Hyunwoo J. Kim

cs.AI

초록

신경망을 자원 제한된 환경에 배포하기 위해, 이전 연구들은 각각 지역 및 전역 의존성을 캡처하기 위해 컨볼루션과 어텐션을 사용한 가벼운 아키텍처를 구축해왔다. 최근에는 상태 공간 모델이 토큰 수에 대한 선형 계산 비용이 유리한 글로벌 토큰 상호 작용으로 효과적으로 등장했다. 그러나 SSM으로 구축된 효율적인 비전 백본은 덜 탐구되어 왔다. 본 논문에서는 효율적인 비전 Mamba(EfficientViM)를 소개한다. 이는 HSM-SSD(숨겨진 상태 믹서 기반 상태 공간 이중성)에 기반한 혁신적인 아키텍처로, 계산 비용을 더욱 줄이면서 글로벌 의존성을 효율적으로 캡처한다. HSM-SSD 레이어에서 이전 SSD 레이어를 재설계하여 숨겨진 상태 내에서 채널 믹싱 작업을 가능하게 한다. 또한, 숨겨진 상태 퓨전을 위한 다단계 숨겨진 상태 융합을 제안하고, 메모리 제한 작업으로 인한 병목 현상을 완화하는 설계를 제공한다. 결과적으로, EfficientViM 패밀리는 ImageNet-1k에서 새로운 최고 성능 속도-정확도 균형을 달성하여, 더 빠른 속도로 SHViT 두 번째 최고 모델보다 0.7%의 성능 향상을 제공한다. 더불어, 이미지 확대 또는 증류 훈련을 적용할 때, 이전 연구들과 비교하여 처리량과 정확도에서 상당한 개선을 관찰할 수 있다. 코드는 https://github.com/mlvlab/EfficientViM에서 확인할 수 있다.

English

For the deployment of neural networks in resource-constrained environments, prior works have built lightweight architectures with convolution and attention for capturing local and global dependencies, respectively. Recently, the state space model has emerged as an effective global token interaction with its favorable linear computational cost in the number of tokens. Yet, efficient vision backbones built with SSM have been explored less. In this paper, we introduce Efficient Vision Mamba (EfficientViM), a novel architecture built on hidden state mixer-based state space duality (HSM-SSD) that efficiently captures global dependencies with further reduced computational cost. In the HSM-SSD layer, we redesign the previous SSD layer to enable the channel mixing operation within hidden states. Additionally, we propose multi-stage hidden state fusion to further reinforce the representation power of hidden states, and provide the design alleviating the bottleneck caused by the memory-bound operations. As a result, the EfficientViM family achieves a new state-of-the-art speed-accuracy trade-off on ImageNet-1k, offering up to a 0.7% performance improvement over the second-best model SHViT with faster speed. Further, we observe significant improvements in throughput and accuracy compared to prior works, when scaling images or employing distillation training. Code is available at https://github.com/mlvlab/EfficientViM.