VMamba: Modelo de Espaço de Estados Visual

VMamba: Visual State Space Model

January 18, 2024

Autores: Yue Liu, Yunjie Tian, Yuzhong Zhao, Hongtian Yu, Lingxi Xie, Yaowei Wang, Qixiang Ye, Yunfan Liu

cs.AI

Resumo

Redes Neurais Convolucionais (CNNs) e Vision Transformers (ViTs) destacam-se como os dois modelos fundamentais mais populares para aprendizado de representação visual. Enquanto as CNNs exibem uma escalabilidade notável com complexidade linear em relação à resolução da imagem, os ViTs superam-nas em capacidade de ajuste, apesar de lidarem com complexidade quadrática. Uma análise mais detalhada revela que os ViTs alcançam um desempenho superior em modelagem visual por meio da incorporação de campos receptivos globais e pesos dinâmicos. Essa observação nos motiva a propor uma nova arquitetura que herda esses componentes enquanto melhora a eficiência computacional. Para isso, inspiramo-nos no recentemente introduzido modelo de espaço de estados e propomos o Visual State Space Model (VMamba), que alcança complexidade linear sem sacrificar os campos receptivos globais. Para abordar o problema de sensibilidade à direção encontrado, introduzimos o Módulo de Varredura Cruzada (Cross-Scan Module - CSM) para percorrer o domínio espacial e converter qualquer imagem visual não causal em sequências ordenadas de patches. Resultados experimentais extensivos corroboram que o VMamba não apenas demonstra capacidades promissoras em diversas tarefas de percepção visual, mas também exibe vantagens mais pronunciadas em relação a benchmarks estabelecidos à medida que a resolução da imagem aumenta. O código-fonte está disponível em 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.