VSSD : Vision Mamba avec Dualité d'Espace d'États Non Causal

VSSD: Vision Mamba with Non-Casual State Space Duality

July 26, 2024

Auteurs: Yuheng Shi, Minjing Dong, Mingjia Li, Chang Xu

cs.AI

Résumé

Les vision transformers ont considérablement fait progresser le domaine de la vision par ordinateur, offrant des capacités de modélisation robustes et un champ réceptif global. Cependant, leurs exigences computationnelles élevées limitent leur applicabilité dans le traitement de longues séquences. Pour résoudre ce problème, les modèles d'état (State Space Models, SSMs) ont gagné en importance dans les tâches de vision, car ils offrent une complexité computationnelle linéaire. Récemment, la Dualité d'État (State Space Duality, SSD), une variante améliorée des SSMs, a été introduite dans Mamba2 pour améliorer les performances et l'efficacité du modèle. Cependant, la nature causale inhérente de SSD/SSMs limite leurs applications dans les tâches de vision non causales. Pour pallier cette limitation, nous introduisons le modèle de Dualité d'État Visuelle (Visual State Space Duality, VSSD), qui adopte un format non causal de SSD. Plus précisément, nous proposons de supprimer l'amplitude des interactions entre l'état caché et les tokens tout en préservant leurs poids relatifs, ce qui réduit les dépendances de la contribution des tokens sur les tokens précédents. Avec l'intégration de stratégies de balayage multiple, nous montrons que les résultats du balayage peuvent être combinés pour atteindre la non-causalité, ce qui améliore non seulement les performances de SSD dans les tâches de vision, mais aussi son efficacité. Nous menons des expériences approfondies sur divers benchmarks, incluant la classification d'images, la détection et la segmentation, où VSSD surpasse les modèles basés sur SSM existants de pointe. Le code et les poids sont disponibles à l'adresse https://github.com/YuHengsss/VSSD.

English

Vision transformers have significantly advanced the field of computer vision, offering robust modeling capabilities and global receptive field. However, their high computational demands limit their applicability in processing long sequences. To tackle this issue, State Space Models (SSMs) have gained prominence in vision tasks as they offer linear computational complexity. Recently, State Space Duality (SSD), an improved variant of SSMs, was introduced in Mamba2 to enhance model performance and efficiency. However, the inherent causal nature of SSD/SSMs restricts their applications in non-causal vision tasks. To address this limitation, we introduce Visual State Space Duality (VSSD) model, which has a non-causal format of SSD. Specifically, we propose to discard the magnitude of interactions between the hidden state and tokens while preserving their relative weights, which relieves the dependencies of token contribution on previous tokens. Together with the involvement of multi-scan strategies, we show that the scanning results can be integrated to achieve non-causality, which not only improves the performance of SSD in vision tasks but also enhances its efficiency. We conduct extensive experiments on various benchmarks including image classification, detection, and segmentation, where VSSD surpasses existing state-of-the-art SSM-based models. Code and weights are available at https://github.com/YuHengsss/VSSD.