A Separable Self-attention Inspired by the State Space Model for Computer Vision

TL;DR

This paper introduces VMINet, a novel architecture that incorporates a separable self-attention mechanism inspired by state space models, achieving competitive performance in image classification and dense prediction tasks.

Contribution

It presents a new separable self-attention method inspired by Mamba state space models and a simple prototype architecture, VMINet, for computer vision tasks.

Findings

Achieved competitive results on image classification.

Demonstrated effectiveness in high-resolution dense prediction.

Introduced a novel attention mechanism inspired by state space models.

Abstract

Mamba is an efficient State Space Model (SSM) with linear computational complexity. Although SSMs are not suitable for handling non-causal data, Vision Mamba (ViM) methods still demonstrate good performance in tasks such as image classification and object detection. Recent studies have shown that there is a rich theoretical connection between state space models and attention variants. We propose a novel separable self attention method, for the first time introducing some excellent design concepts of Mamba into separable self-attention. To ensure a fair comparison with ViMs, we introduce VMINet, a simple yet powerful prototype architecture, constructed solely by stacking our novel attention modules with the most basic down-sampling layers. Notably, VMINet differs significantly from the conventional Transformer architecture. Our experiments demonstrate that VMINet has achieved competitive results on image classification and high-resolution dense prediction tasks.Code is available at: https://github.com/yws-wxs/VMINet.