VSSD：非因果的状態空間双対性を備えたVision Mamba

VSSD: Vision Mamba with Non-Casual State Space Duality

July 26, 2024

著者: Yuheng Shi, Minjing Dong, Mingjia Li, Chang Xu

cs.AI

要旨

ビジョントランスフォーマーは、コンピュータビジョンの分野を大きく進歩させ、堅牢なモデリング能力とグローバルな受容野を提供しています。しかし、その高い計算コストは、長いシーケンスを処理する際の適用性を制限しています。この問題に対処するため、状態空間モデル（SSM）がビジョンタスクで注目を集めており、線形計算複雑性を提供します。最近、Mamba2において、SSMの改良版である状態空間双対性（SSD）が導入され、モデルの性能と効率が向上しました。しかし、SSD/SSMの本質的な因果性は、非因果的なビジョンタスクでの応用を制限します。この制限を解決するため、我々はSSDの非因果形式を持つVisual State Space Duality（VSSD）モデルを提案します。具体的には、隠れ状態とトークン間の相互作用の大きさを捨てながら、それらの相対的な重みを保持することを提案し、これによりトークンの寄与が前のトークンに依存することを緩和します。マルチスキャン戦略の導入とともに、スキャン結果を統合して非因果性を達成できることを示し、これによりビジョンタスクにおけるSSDの性能が向上するだけでなく、効率も向上します。画像分類、検出、セグメンテーションを含むさまざまなベンチマークで広範な実験を行い、VSSDが既存の最先端のSSMベースのモデルを凌駕することを示します。コードと重みは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.