VSSD: 비캐주얼 상태 공간 이중성을 갖는 비전 맘바

VSSD: Vision Mamba with Non-Casual State Space Duality

July 26, 2024

저자: Yuheng Shi, Minjing Dong, Mingjia Li, Chang Xu

cs.AI

초록

비전 트랜스포머(Vision Transformers)는 컴퓨터 비전 분야에서 글로벌 수용 필드와 강력한 모델링 능력을 제공하며 큰 발전을 이루었습니다. 그러나 긴 시퀀스 처리에서의 높은 계산 요구량으로 인해 적용 가능성이 제한됩니다. 이를 해결하기 위해 상태 공간 모델(State Space Models, SSMs)이 선형 계산 복잡도를 제공하며 비전 작업에서 주목받고 있습니다. 최근 Mamba2에서 SSM의 개선된 변형인 상태 공간 이중성(State Space Duality, 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.