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 была представлена модель Пространства Дуальности Состояний (SSD), улучшенная версия SSM, которая способствует повышению производительности и эффективности модели. Однако внутренняя причинная природа SSD/SSM ограничивает их применение в непричинных задачах зрения. Для преодоления этого ограничения мы представляем модель Визуальной Дуальности Пространства Состояний (VSSD), которая имеет непричинный формат SSD. Конкретно, мы предлагаем отбросить величину взаимодействий между скрытым состоянием и токенами, сохраняя при этом их относительные веса, что уменьшает зависимость вклада токена от предыдущих токенов. Совместно с использованием стратегий многократного сканирования мы показываем, что результаты сканирования могут быть интегрированы для достижения непричинности, что не только улучшает производительность 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.