VSSD: Visão Mamba com Dualidade de Espaço de Estado Não-Casual

VSSD: Vision Mamba with Non-Casual State Space Duality

July 26, 2024

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

cs.AI

Resumo

Os transformadores de visão avançaram significativamente o campo da visão computacional, oferecendo capacidades de modelagem robustas e um campo receptivo global. No entanto, suas altas demandas computacionais limitam sua aplicabilidade no processamento de sequências longas. Para lidar com esse problema, os Modelos de Espaço de Estado (SSMs) ganharam destaque em tarefas de visão, pois oferecem complexidade computacional linear. Recentemente, o modelo de Dualidade de Espaço de Estado (SSD), uma variante aprimorada dos SSMs, foi introduzido no Mamba2 para aprimorar o desempenho e eficiência do modelo. No entanto, a natureza causal inerente do SSD/SSMs restringe suas aplicações em tarefas de visão não causais. Para abordar essa limitação, apresentamos o modelo de Dualidade de Espaço de Estado Visual (VSSD), que possui um formato não causal do SSD. Especificamente, propomos descartar a magnitude das interações entre o estado oculto e os tokens, preservando seus pesos relativos, o que alivia as dependências da contribuição do token em tokens anteriores. Juntamente com o envolvimento de estratégias de multi-varredura, mostramos que os resultados da varredura podem ser integrados para alcançar a não causalidade, o que não apenas melhora o desempenho do SSD em tarefas de visão, mas também aprimora sua eficiência. Realizamos experimentos extensivos em vários benchmarks, incluindo classificação, detecção e segmentação de imagens, onde o VSSD supera os modelos baseados em SSMs existentes de última geração. O código e os pesos estão disponíveis em 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.