VMamba : Modèle d'espace d'état visuel

VMamba: Visual State Space Model

January 18, 2024

Auteurs: Yue Liu, Yunjie Tian, Yuzhong Zhao, Hongtian Yu, Lingxi Xie, Yaowei Wang, Qixiang Ye, Yunfan Liu

cs.AI

Résumé

Les réseaux de neurones convolutifs (CNN) et les Vision Transformers (ViT) constituent les deux modèles de base les plus populaires pour l'apprentissage de représentations visuelles. Alors que les CNN présentent une scalabilité remarquable avec une complexité linéaire par rapport à la résolution de l'image, les ViT les surpassent en termes de capacité d'ajustement malgré une complexité quadratique. Une analyse plus approfondie révèle que les ViT obtiennent des performances supérieures en modélisation visuelle grâce à l'intégration de champs réceptifs globaux et de poids dynamiques. Cette observation nous motive à proposer une nouvelle architecture qui hérite de ces composants tout en améliorant l'efficacité computationnelle. Pour ce faire, nous nous inspirons du modèle d'état récemment introduit et proposons le Visual State Space Model (VMamba), qui atteint une complexité linéaire sans sacrifier les champs réceptifs globaux. Pour résoudre le problème de sensibilité directionnelle rencontré, nous introduisons le Cross-Scan Module (CSM) pour parcourir le domaine spatial et convertir toute image visuelle non causale en séquences de patches ordonnées. Les résultats expérimentaux approfondis confirment que VMamba démontre non seulement des capacités prometteuses dans diverses tâches de perception visuelle, mais présente également des avantages plus marqués par rapport aux benchmarks établis à mesure que la résolution de l'image augmente. Le code source est disponible à l'adresse suivante : https://github.com/MzeroMiko/VMamba.

English

Convolutional Neural Networks (CNNs) and Vision Transformers (ViTs) stand as the two most popular foundation models for visual representation learning. While CNNs exhibit remarkable scalability with linear complexity w.r.t. image resolution, ViTs surpass them in fitting capabilities despite contending with quadratic complexity. A closer inspection reveals that ViTs achieve superior visual modeling performance through the incorporation of global receptive fields and dynamic weights. This observation motivates us to propose a novel architecture that inherits these components while enhancing computational efficiency. To this end, we draw inspiration from the recently introduced state space model and propose the Visual State Space Model (VMamba), which achieves linear complexity without sacrificing global receptive fields. To address the encountered direction-sensitive issue, we introduce the Cross-Scan Module (CSM) to traverse the spatial domain and convert any non-causal visual image into order patch sequences. Extensive experimental results substantiate that VMamba not only demonstrates promising capabilities across various visual perception tasks, but also exhibits more pronounced advantages over established benchmarks as the image resolution increases. Source code has been available at https://github.com/MzeroMiko/VMamba.