EfficientViM : Mamba de Vision Efficace avec un Mélangeur d'États Cachés basé sur la Dualité de l'Espace d'États

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

November 22, 2024

Auteurs: Sanghyeok Lee, Joonmyung Choi, Hyunwoo J. Kim

cs.AI

Résumé

Pour le déploiement de réseaux neuronaux dans des environnements aux ressources limitées, des travaux antérieurs ont construit des architectures légères avec des convolutions et de l'attention pour capturer respectivement les dépendances locales et globales. Récemment, le modèle d'espace d'états a émergé comme une interaction de jeton global efficace avec son coût computationnel linéaire favorable en fonction du nombre de jetons. Cependant, les architectures de vision efficaces construites avec SSM ont été moins explorées. Dans cet article, nous introduisons Efficient Vision Mamba (EfficientViM), une architecture novatrice construite sur la dualité de l'espace d'états basée sur le mélangeur d'états cachés (HSM-SSD) qui capture efficacement les dépendances globales avec un coût computationnel réduit. Dans la couche HSM-SSD, nous redessinons la couche SSD précédente pour permettre l'opération de mélange de canaux au sein des états cachés. De plus, nous proposons une fusion d'états cachés multi-étapes pour renforcer davantage la puissance de représentation des états cachés, et fournissons la conception atténuant le goulot d'étranglement causé par les opérations liées à la mémoire. En conséquence, la famille EfficientViM atteint un nouvel équilibre vitesse-précision de pointe sur ImageNet-1k, offrant jusqu'à une amélioration de performance de 0,7% par rapport au deuxième meilleur modèle SHViT avec une vitesse plus rapide. De plus, nous observons des améliorations significatives en termes de débit et de précision par rapport aux travaux antérieurs, lors de la mise à l'échelle des images ou de l'utilisation de l'entraînement par distillation. Le code est disponible sur 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.