VSSD: Vision Mamba met Niet-Causale State Space Dualiteit

VSSD: Vision Mamba with Non-Casual State Space Duality

July 26, 2024

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

cs.AI

Samenvatting

Vision transformers hebben het vakgebied van computer vision aanzienlijk vooruitgebracht, door robuuste modelleermogelijkheden en een globaal receptief veld te bieden. Hun hoge rekenkundige eisen beperken echter hun toepasbaarheid bij het verwerken van lange sequenties. Om dit probleem aan te pakken, hebben State Space Models (SSMs) aan populariteit gewonnen in visuele taken, omdat ze lineaire rekenkundige complexiteit bieden. Onlangs werd State Space Duality (SSD), een verbeterde variant van SSMs, geïntroduceerd in Mamba2 om de modelprestaties en efficiëntie te verbeteren. De inherent causale aard van SSD/SSMs beperkt echter hun toepassingen in niet-causale visuele taken. Om deze beperking aan te pakken, introduceren we het Visual State Space Duality (VSSD) model, dat een niet-causale vorm van SSD heeft. Specifiek stellen we voor om de grootte van de interacties tussen de verborgen toestand en tokens te verwijderen, terwijl hun relatieve gewichten behouden blijven, wat de afhankelijkheid van de bijdrage van tokens op vorige tokens vermindert. Samen met de betrokkenheid van multi-scan strategieën laten we zien dat de scanresultaten geïntegreerd kunnen worden om niet-causaliteit te bereiken, wat niet alleen de prestaties van SSD in visuele taken verbetert, maar ook de efficiëntie ervan verhoogt. We voeren uitgebreide experimenten uit op verschillende benchmarks, waaronder beeldclassificatie, detectie en segmentatie, waarbij VSSD de bestaande state-of-the-art SSM-gebaseerde modellen overtreft. Code en gewichten zijn beschikbaar op 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.