LocalMamba: Visueel State Space Model met Venstergebaseerde Selectieve Scan

LocalMamba: Visual State Space Model with Windowed Selective Scan

March 14, 2024

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

cs.AI

Samenvatting

Recente vooruitgang in state space-modellen, met name Mamba, heeft aanzienlijke vooruitgang getoond in het modelleren van lange sequenties voor taken zoals taalbegrip. Toch heeft hun toepassing in visuele taken de prestaties van traditionele Convolutional Neural Networks (CNN's) en Vision Transformers (ViT's) niet duidelijk overtroffen. Dit artikel stelt dat de sleutel tot het verbeteren van Vision Mamba (ViM) ligt in het optimaliseren van scanrichtingen voor sequentiemodellering. Traditionele ViM-benaderingen, die ruimtelijke tokens afvlakken, negeren het behoud van lokale 2D-afhankelijkheden, waardoor de afstand tussen aangrenzende tokens wordt verlengd. We introduceren een nieuwe lokale scanstrategie die afbeeldingen verdeelt in afzonderlijke vensters, waardoor lokale afhankelijkheden effectief worden vastgelegd terwijl een globaal perspectief behouden blijft. Daarnaast, erkennend dat de voorkeuren voor scanpatronen variëren tussen verschillende netwerklagen, stellen we een dynamische methode voor om onafhankelijk naar de optimale scankeuzes voor elke laag te zoeken, wat de prestaties aanzienlijk verbetert. Uitgebreide experimenten in zowel eenvoudige als hiërarchische modellen onderstrepen de superioriteit van onze aanpak in het effectief vastleggen van beeldrepresentaties. Ons model overtreft bijvoorbeeld Vim-Ti significant met 3,1% op ImageNet bij dezelfde 1,5G FLOPs. Code is beschikbaar op: 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.