VSSD: Vision Mamba mit nicht-kausaler Zustandsraum-Dualität

VSSD: Vision Mamba with Non-Casual State Space Duality

July 26, 2024

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

cs.AI

Zusammenfassung

Vision-Transformer haben das Gebiet der Computer Vision erheblich vorangetrieben, indem sie robuste Modellierungsfähigkeiten und einen globalen Rezeptionsbereich bieten. Allerdings begrenzen ihre hohen Rechenanforderungen ihre Anwendbarkeit bei der Verarbeitung langer Sequenzen. Um dieses Problem anzugehen, haben Zustandsraummodelle (SSMs) in Vision-Aufgaben an Bedeutung gewonnen, da sie eine lineare Rechenkomplexität bieten. Kürzlich wurde in Mamba2 der Zustandsraum-Dualität (SSD), eine verbesserte Variante von SSMs, eingeführt, um die Leistung und Effizienz des Modells zu verbessern. Die inhärente kausale Natur von SSD/SSMs beschränkt jedoch ihre Anwendungen auf nicht-kausale Vision-Aufgaben. Um diese Einschränkung zu überwinden, führen wir das Modell der visuellen Zustandsraum-Dualität (VSSD) ein, das ein nicht-kausales Format von SSD aufweist. Konkret schlagen wir vor, die Größenordnung der Interaktionen zwischen dem versteckten Zustand und Tokens zu verwerfen, während ihre relativen Gewichte beibehalten werden, was die Abhängigkeiten des Token-Beitrags von vorherigen Tokens entlastet. Zusammen mit der Einbeziehung von Multi-Scan-Strategien zeigen wir, dass die Scan-Ergebnisse integriert werden können, um Nicht-Kausalität zu erreichen, was nicht nur die Leistung von SSD in Vision-Aufgaben verbessert, sondern auch seine Effizienz steigert. Wir führen umfangreiche Experimente an verschiedenen Benchmarks durch, darunter Bildklassifizierung, Detektion und Segmentierung, bei denen VSSD bestehende SSM-basierte Modelle der Spitzenklasse übertrifft. Der Code und die Gewichte sind unter https://github.com/YuHengsss/VSSD verfügbar.

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.