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

TL;DR

EfficientViM introduces a novel architecture leveraging hidden state mixer-based state space duality to efficiently capture global dependencies, achieving superior speed-accuracy trade-offs on ImageNet-1k with scalable improvements.

Contribution

The paper proposes EfficientViM, a new vision model that reduces computational cost and enhances performance by redesigning the SSD layer and introducing multi-stage hidden state fusion.

Findings

Achieves up to 0.7% better accuracy than SHViT with faster speed.

Significant throughput and accuracy improvements when scaling images or using distillation.

Sets new state-of-the-art speed-accuracy trade-off on ImageNet-1k.

Abstract

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 (SSM) has emerged as an effective operation for global interaction with its favorable linear computational cost in the number of tokens. To harness the efficacy of SSM, 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. With the observation that the runtime of the SSD layer is driven by the linear projections on the input sequences, we redesign the original SSD layer to perform the channel mixing operation within compressed hidden states in the HSM-SSD layer. Additionally, we propose multi-stage hidden state fusion to reinforce the representation power of hidden states and provide the design to alleviate 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.