Vim-F: Visual State Space Model Benefiting from Learning in the Frequency Domain

2024-05-29

arXiv.org

AI Summary

To address the loss of 2D structural information due to image flattening and the difficulty in jointly modeling local and global contexts in Vision Mamba (ViM), this paper proposes a frequency-spatial joint modeling architecture. Methodologically, it introduces three key innovations: (1) a novel frequency-domain enhancement mechanism leveraging Fast Fourier Transform (FFT) to inject frequency-domain priors, enabling synergistic modeling with spatial-domain Mamba; (2) elimination of positional encoding to fully unleash the long-range dependency modeling capacity of state-space models; and (3) replacement of patch embedding with a convolutional stem to preserve local structure and improve feature initialization quality. Evaluated on benchmarks including ImageNet, the proposed method achieves accuracy comparable to Vision Transformers (ViTs) while significantly outperforming the original ViM—demonstrating superior performance at low computational cost. These results validate the effectiveness and generalizability of dual-domain (frequency + spatial) modeling for visual state-space models.

Technology Category

• Computer Vision: Diffusion Models for Vision
• Machine Learning: Large Multimodal Models (LMMs)
• Cognitive Modeling & Cognitive Systems: Neural Spike Coding

Application Category

• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• Search and Retrieval-Augmented AI: Vertical and domain-specific search
• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs

Abstract

In recent years, State Space Models (SSMs) with efficient hardware-aware designs, known as the Mamba deep learning models, have made significant progress in modeling long sequences such as language understanding. Therefore, building efficient and general-purpose visual backbones based on SSMs is a promising direction. Compared to traditional convolutional neural networks (CNNs) and Vision Transformers (ViTs), the performance of Vision Mamba (ViM) methods is not yet fully competitive. To enable SSMs to process image data, ViMs typically flatten 2D images into 1D sequences, inevitably ignoring some 2D local dependencies, thereby weakening the model's ability to interpret spatial relationships from a global perspective. We use Fast Fourier Transform (FFT) to obtain the spectrum of the feature map and add it to the original feature map, enabling ViM to model a unified visual representation in both frequency and spatial domains. The introduction of frequency domain information enables ViM to have a global receptive field during scanning. We propose a novel model called Vim-F, which employs pure Mamba encoders and scans in both the frequency and spatial domains. Moreover, we question the necessity of position embedding in ViM and remove it accordingly in Vim-F, which helps to fully utilize the efficient long-sequence modeling capability of ViM. Finally, we redesign a patch embedding for Vim-F, leveraging a convolutional stem to capture more local correlations, further improving the performance of Vim-F. Code is available at: url{https://github.com/yws-wxs/Vim-F}.

Problem

• Mamba model
• image processing
• performance enhancement

Innovation

• Vim-F Model
• Frequency Domain Learning
• FFT-based Feature Extraction