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

📄 arXiv: 2405.18679v3 📥 PDF

作者: Juntao Zhang, Shaogeng Liu, Jun Zhou, Kun Bian, You Zhou, Jianning Liu, Pei Zhang, Bingyan Liu

分类: cs.CV

发布日期: 2024-05-29 (更新: 2025-09-25)

🔗 代码/项目: GITHUB


💡 一句话要点

提出Vim-F,利用频域学习增强视觉状态空间模型,提升图像建模能力。

🎯 匹配领域: 支柱二:RL算法与架构 (RL & Architecture) 支柱七:动作重定向 (Motion Retargeting)

关键词: 视觉状态空间模型 频域学习 Mamba 图像分类 全局感受野 快速傅里叶变换 卷积神经网络 视觉Transformer

📋 核心要点

  1. Vision Mamba (ViM) 在图像建模中性能不足,主要原因是将2D图像展平为1D序列时忽略了重要的2D局部依赖关系。
  2. Vim-F 通过快速傅里叶变换(FFT)将频域信息融入特征图,使模型能够同时在空间域和频域进行建模,从而获得全局感受野。
  3. Vim-F 移除了位置嵌入,并重新设计了patch embedding,利用卷积 stem 捕获局部相关性,进一步提升了模型性能。

📝 摘要(中文)

近年来,具有高效硬件感知设计的状态空间模型(SSM),即Mamba深度学习模型,在诸如语言理解等长序列建模方面取得了显著进展。因此,构建基于SSM的高效通用视觉骨干网络是一个很有前景的方向。与传统的卷积神经网络(CNN)和视觉Transformer(ViT)相比,Vision Mamba(ViM)方法的性能尚未完全具有竞争力。为了使SSM能够处理图像数据,ViM通常将2D图像展平为1D序列,不可避免地忽略了一些2D局部依赖关系,从而削弱了模型从全局角度解释空间关系的能力。我们使用快速傅里叶变换(FFT)来获得特征图的频谱,并将其添加到原始特征图中,使ViM能够在频域和空间域中建模统一的视觉表示。频域信息的引入使ViM在扫描期间具有全局感受野。我们提出了一种名为Vim-F的新模型,该模型采用纯Mamba编码器,并在频域和空间域中进行扫描。此外,我们质疑ViM中位置嵌入的必要性,并在Vim-F中相应地将其移除,这有助于充分利用ViM高效的长序列建模能力。最后,我们为Vim-F重新设计了一个patch embedding,利用卷积stem来捕获更多的局部相关性,进一步提高了Vim-F的性能。

🔬 方法详解

问题定义:现有的Vision Mamba (ViM) 模型在处理图像数据时,通常将二维图像展平为一维序列,这导致模型忽略了图像中重要的二维局部依赖关系,从而限制了模型对空间关系的全局理解能力。因此,如何有效地利用状态空间模型处理图像数据,同时保留图像的二维结构信息,是一个亟待解决的问题。

核心思路:Vim-F 的核心思路是将图像的频域信息引入到模型中,通过快速傅里叶变换(FFT)将图像从空间域转换到频域,并将频域信息与原始的空间域特征图融合。这样,模型就可以同时利用空间域和频域的信息,从而更好地理解图像的全局结构和局部细节。这种设计借鉴了信号处理中的思想,即频域分析可以揭示信号的全局特性。

技术框架:Vim-F 的整体架构基于纯 Mamba 编码器,主要包括以下几个模块:1) Patch Embedding:使用卷积 stem 提取局部特征。2) Mamba Encoder:在空间域和频域上进行扫描,建模长序列依赖关系。3) FFT Module:将特征图转换到频域。4) Feature Fusion:将空间域和频域的特征图融合。整个流程是,首先通过 Patch Embedding 提取图像的局部特征,然后将特征图输入到 Mamba Encoder 中进行处理。在 Mamba Encoder 中,特征图首先通过 FFT Module 转换到频域,然后将频域特征与原始的空间域特征融合,最后进行 Mamba 层的计算。

关键创新:Vim-F 的关键创新在于将频域信息引入到 Vision Mamba 模型中。通过 FFT 将图像转换到频域,并与空间域特征融合,使得模型能够同时利用空间域和频域的信息。此外,Vim-F 还移除了位置嵌入,并重新设计了 Patch Embedding,进一步提升了模型的性能。与现有方法的本质区别在于,Vim-F 不仅仅依赖于空间域的信息,而是同时利用了空间域和频域的信息,从而更好地理解图像的全局结构和局部细节。

关键设计:Vim-F 的关键设计包括:1) 使用卷积 stem 作为 Patch Embedding,以更好地捕获局部相关性。2) 在 Mamba Encoder 中,使用 FFT 将特征图转换到频域,并将频域特征与空间域特征融合。3) 移除位置嵌入,以充分利用 Mamba 模型的长序列建模能力。4) 损失函数采用标准的交叉熵损失函数。具体的参数设置在论文中有详细描述。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

论文提出的 Vim-F 模型在图像分类任务上取得了显著的性能提升。实验结果表明,Vim-F 在 ImageNet 数据集上优于现有的 Vision Mamba 模型,并且在参数量和计算复杂度方面具有优势。具体的性能数据和对比基线可以在论文中找到。

🎯 应用场景

Vim-F 具有广泛的应用前景,可以应用于图像分类、目标检测、图像分割等计算机视觉任务。其高效的长序列建模能力和全局感受野使其在处理高分辨率图像和视频数据时具有优势。未来,Vim-F 可以进一步扩展到其他领域,如医学图像分析、遥感图像处理等。

📄 摘要(原文)

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: https://github.com/yws-wxs/Vim-F.