VSSD: Vision Mamba con Dualità di Spazio degli Stati Non-Casuale

VSSD: Vision Mamba with Non-Casual State Space Duality

July 26, 2024

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

cs.AI

Abstract

I transformer per la visione hanno significativamente avanzato il campo della visione artificiale, offrendo robuste capacità di modellazione e un campo recettivo globale. Tuttavia, le loro elevate esigenze computazionali ne limitano l'applicabilità nell'elaborazione di sequenze lunghe. Per affrontare questo problema, i modelli a spazio di stato (SSM) hanno guadagnato popolarità nei compiti di visione, poiché offrono una complessità computazionale lineare. Recentemente, la Dualità dello Spazio di Stato (SSD), una variante migliorata degli SSM, è stata introdotta in Mamba2 per migliorare le prestazioni e l'efficienza del modello. Tuttavia, la natura intrinsecamente causale di SSD/SSM ne limita le applicazioni nei compiti di visione non causali. Per superare questa limitazione, introduciamo il modello Visual State Space Duality (VSSD), che presenta un formato non causale di SSD. Nello specifico, proponiamo di scartare l'ampiezza delle interazioni tra lo stato nascosto e i token, preservandone i pesi relativi, il che riduce le dipendenze del contributo dei token dai token precedenti. Insieme all'impiego di strategie di scansione multipla, dimostriamo che i risultati della scansione possono essere integrati per ottenere la non-causalità, migliorando non solo le prestazioni di SSD nei compiti di visione, ma anche la sua efficienza. Abbiamo condotto esperimenti approfonditi su vari benchmark, tra cui classificazione di immagini, rilevamento e segmentazione, dove VSSD supera i modelli basati su SSM esistenti all'avanguardia. Codice e pesi sono disponibili all'indirizzo 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.