LocalMamba: Modello a Spazio di Stati Visivo con Scansione Selettiva a Finestra

LocalMamba: Visual State Space Model with Windowed Selective Scan

March 14, 2024

Autori: Tao Huang, Xiaohuan Pei, Shan You, Fei Wang, Chen Qian, Chang Xu

cs.AI

Abstract

I recenti progressi nei modelli a spazio di stato, in particolare Mamba, hanno dimostrato significativi avanzamenti nella modellazione di sequenze lunghe per compiti come la comprensione del linguaggio. Tuttavia, la loro applicazione nei compiti visivi non ha superato in modo marcato le prestazioni delle tradizionali Reti Neurali Convoluzionali (CNN) e dei Vision Transformer (ViT). Questo articolo sostiene che la chiave per migliorare il Vision Mamba (ViM) risiede nell'ottimizzazione delle direzioni di scansione per la modellazione delle sequenze. Gli approcci tradizionali di ViM, che appiattiscono i token spaziali, trascurano la preservazione delle dipendenze locali 2D, allungando così la distanza tra token adiacenti. Introduciamo una nuova strategia di scansione locale che divide le immagini in finestre distinte, catturando efficacemente le dipendenze locali mantenendo una prospettiva globale. Inoltre, riconoscendo le diverse preferenze per i pattern di scansione tra i vari strati della rete, proponiamo un metodo dinamico per cercare in modo indipendente le scelte di scansione ottimali per ciascuno strato, migliorando sostanzialmente le prestazioni. Esperimenti estesi su modelli sia semplici che gerarchici evidenziano la superiorità del nostro approccio nel catturare efficacemente le rappresentazioni delle immagini. Ad esempio, il nostro modello supera significativamente Vim-Ti del 3.1% su ImageNet con gli stessi 1.5G FLOP. Il codice è disponibile all'indirizzo: https://github.com/hunto/LocalMamba.

English

Recent advancements in state space models, notably Mamba, have demonstrated significant progress in modeling long sequences for tasks like language understanding. Yet, their application in vision tasks has not markedly surpassed the performance of traditional Convolutional Neural Networks (CNNs) and Vision Transformers (ViTs). This paper posits that the key to enhancing Vision Mamba (ViM) lies in optimizing scan directions for sequence modeling. Traditional ViM approaches, which flatten spatial tokens, overlook the preservation of local 2D dependencies, thereby elongating the distance between adjacent tokens. We introduce a novel local scanning strategy that divides images into distinct windows, effectively capturing local dependencies while maintaining a global perspective. Additionally, acknowledging the varying preferences for scan patterns across different network layers, we propose a dynamic method to independently search for the optimal scan choices for each layer, substantially improving performance. Extensive experiments across both plain and hierarchical models underscore our approach's superiority in effectively capturing image representations. For example, our model significantly outperforms Vim-Ti by 3.1% on ImageNet with the same 1.5G FLOPs. Code is available at: https://github.com/hunto/LocalMamba.