VSSD: Visión Mamba con Dualidad de Espacio de Estados No Casual

VSSD: Vision Mamba with Non-Casual State Space Duality

July 26, 2024

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

cs.AI

Resumen

Los transformadores de visión han avanzado significativamente en el campo de la visión por computadora, ofreciendo capacidades de modelado robustas y un campo receptivo global. Sin embargo, sus altas demandas computacionales limitan su aplicabilidad en el procesamiento de secuencias largas. Para abordar este problema, los Modelos de Espacio de Estados (SSMs, por sus siglas en inglés) han ganado prominencia en tareas de visión al ofrecer una complejidad computacional lineal. Recientemente, el modelo de Dualidad de Espacio de Estados (SSD, por sus siglas en inglés), una variante mejorada de los SSMs, fue introducido en Mamba2 para mejorar el rendimiento y la eficiencia del modelo. Sin embargo, la naturaleza causal inherente de SSD/SSMs restringe sus aplicaciones en tareas de visión no causales. Para abordar esta limitación, presentamos el modelo de Dualidad de Espacio de Estados Visual (VSSD), que tiene un formato no causal de SSD. Específicamente, proponemos desechar la magnitud de las interacciones entre el estado oculto y los tokens mientras se preservan sus pesos relativos, lo que alivia las dependencias de la contribución del token en los tokens anteriores. Junto con la participación de estrategias de múltiples escaneos, mostramos que los resultados del escaneo pueden integrarse para lograr la no causalidad, lo que no solo mejora el rendimiento de SSD en tareas de visión, sino que también mejora su eficiencia. Realizamos experimentos extensos en varios benchmarks, incluida la clasificación, detección y segmentación de imágenes, donde VSSD supera a los modelos existentes basados en SSM líderes en el estado del arte. El código y los pesos están disponibles en 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.