VMamba: Визуальная модель пространства состояний

VMamba: Visual State Space Model

January 18, 2024

Авторы: Yue Liu, Yunjie Tian, Yuzhong Zhao, Hongtian Yu, Lingxi Xie, Yaowei Wang, Qixiang Ye, Yunfan Liu

cs.AI

Аннотация

Сверточные нейронные сети (CNN) и Vision Transformers (ViT) являются двумя наиболее популярными базовыми моделями для обучения визуальных представлений. В то время как CNN демонстрируют замечательную масштабируемость с линейной сложностью относительно разрешения изображения, ViT превосходят их в способности к аппроксимации, несмотря на квадратичную сложность. Более детальное изучение показывает, что ViT достигают превосходной производительности в визуальном моделировании благодаря включению глобальных рецептивных полей и динамических весов. Это наблюдение мотивирует нас предложить новую архитектуру, которая наследует эти компоненты, одновременно повышая вычислительную эффективность. Вдохновляясь недавно представленной моделью пространства состояний, мы предлагаем Visual State Space Model (VMamba), которая достигает линейной сложности без ущерба для глобальных рецептивных полей. Для решения проблемы чувствительности к направлению мы вводим Cross-Scan Module (CSM), который обходит пространственную область и преобразует любое некаузальное визуальное изображение в упорядоченные последовательности патчей. Обширные экспериментальные результаты подтверждают, что VMamba не только демонстрирует перспективные возможности в различных задачах визуального восприятия, но и проявляет более выраженные преимущества по сравнению с установленными эталонами по мере увеличения разрешения изображения. Исходный код доступен по адресу https://github.com/MzeroMiko/VMamba.

English

Convolutional Neural Networks (CNNs) and Vision Transformers (ViTs) stand as the two most popular foundation models for visual representation learning. While CNNs exhibit remarkable scalability with linear complexity w.r.t. image resolution, ViTs surpass them in fitting capabilities despite contending with quadratic complexity. A closer inspection reveals that ViTs achieve superior visual modeling performance through the incorporation of global receptive fields and dynamic weights. This observation motivates us to propose a novel architecture that inherits these components while enhancing computational efficiency. To this end, we draw inspiration from the recently introduced state space model and propose the Visual State Space Model (VMamba), which achieves linear complexity without sacrificing global receptive fields. To address the encountered direction-sensitive issue, we introduce the Cross-Scan Module (CSM) to traverse the spatial domain and convert any non-causal visual image into order patch sequences. Extensive experimental results substantiate that VMamba not only demonstrates promising capabilities across various visual perception tasks, but also exhibits more pronounced advantages over established benchmarks as the image resolution increases. Source code has been available at https://github.com/MzeroMiko/VMamba.