V2M: Visual 2-Dimensional Mamba for Image Representation Learning
作者: Chengkun Wang, Wenzhao Zheng, Yuanhui Huang, Jie Zhou, Jiwen Lu
分类: cs.CV
发布日期: 2024-10-14
💡 一句话要点
提出V2M：一种用于图像表示学习的视觉二维Mamba模型
🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)
关键词: 视觉Mamba 二维状态空间模型 图像表示学习 目标检测 语义分割
📋 核心要点
- 现有视觉Mamba方法将2D图像展平为1D序列，损失了重要的2D结构信息，如局部相似性。
- V2M通过将状态空间模型推广到2D空间，直接在2D空间中处理图像token，保留了2D局部性先验。
- 实验表明，V2M在ImageNet分类和COCO、ADE20K等下游任务上优于其他视觉骨干网络。
📝 摘要（中文）
Mamba因其灵活的设计和高效的硬件性能在处理基于状态空间模型（SSM）的一维序列方面受到了广泛关注。最近的研究试图通过将二维图像展平为图像块，然后将其视为一维序列，从而将Mamba应用于视觉领域。为了弥补原始图像的二维结构信息损失（例如，局部相似性），大多数现有方法侧重于设计不同的顺序来依次处理token，但这只能在一定程度上缓解这个问题。在本文中，我们提出了一个视觉二维Mamba（V2M）模型作为一个完整的解决方案，它直接在二维空间中处理图像token。我们首先将SSM推广到二维空间，该空间生成下一个状态时考虑两个维度（例如，列和行）上的相邻状态。然后，我们基于二维SSM公式构建我们的V2M，并结合Mamba来实现硬件高效的并行处理。所提出的V2M有效地结合了二维局部性先验，同时继承了Mamba的效率和输入相关的可扩展性。在ImageNet分类以及下游视觉任务（包括COCO上的目标检测和实例分割以及ADE20K上的语义分割）上的大量实验结果表明，与其他视觉骨干网络相比，我们的V2M的有效性。
🔬 方法详解
问题定义：现有方法将二维图像展平为一维序列，导致原始图像的二维结构信息丢失，例如局部相似性。虽然一些方法尝试设计不同的token处理顺序来缓解这个问题，但效果有限，无法充分利用图像的二维空间信息。
核心思路：V2M的核心思路是将状态空间模型（SSM）从一维推广到二维，从而可以直接在二维空间中处理图像token。通过考虑相邻行列的状态信息，V2M能够更好地捕捉图像的局部结构，保留二维空间信息。
技术框架：V2M模型基于二维状态空间模型构建，其整体架构包含以下几个主要模块：1）二维状态空间模型：将传统的一维SSM扩展到二维，允许模型同时考虑行和列方向上的相邻状态。2）Mamba集成：将Mamba的硬件高效并行处理能力融入到二维SSM中，加速计算过程。3）图像token嵌入：将原始图像分割成token，并进行嵌入操作，作为V2M的输入。4）输出层：根据具体任务（如分类、检测、分割）设计相应的输出层。
关键创新：V2M最重要的技术创新点在于将状态空间模型推广到二维空间，使其能够直接处理二维图像数据，避免了展平操作带来的信息损失。与现有方法相比，V2M能够更好地捕捉图像的局部结构和空间关系。
关键设计：二维状态空间模型的关键设计在于如何定义状态转移方程，使其能够有效地融合相邻行列的状态信息。具体而言，模型需要学习两个方向上的状态转移矩阵，并设计合适的融合机制。此外，为了保证计算效率，V2M采用了Mamba的并行处理策略，并针对二维数据进行了优化。损失函数根据具体任务选择，例如交叉熵损失用于分类任务，IoU损失用于分割任务。
🖼️ 关键图片
📊 实验亮点
V2M在ImageNet分类任务上取得了显著的性能提升，并且在COCO目标检测和实例分割以及ADE20K语义分割等下游任务上均优于其他视觉骨干网络。实验结果表明，V2M能够有效地捕捉图像的二维结构信息，并具有良好的泛化能力。具体的性能数据需要在论文中查找。
🎯 应用场景
V2M作为一种通用的视觉骨干网络，可以广泛应用于各种计算机视觉任务，例如图像分类、目标检测、实例分割和语义分割。其高效的硬件性能和输入相关的可扩展性使其在资源受限的设备上也能表现良好。未来，V2M有望在自动驾驶、智能监控、医学图像分析等领域发挥重要作用。
📄 摘要（原文）
Mamba has garnered widespread attention due to its flexible design and efficient hardware performance to process 1D sequences based on the state space model (SSM). Recent studies have attempted to apply Mamba to the visual domain by flattening 2D images into patches and then regarding them as a 1D sequence. To compensate for the 2D structure information loss (e.g., local similarity) of the original image, most existing methods focus on designing different orders to sequentially process the tokens, which could only alleviate this issue to some extent. In this paper, we propose a Visual 2-Dimensional Mamba (V2M) model as a complete solution, which directly processes image tokens in the 2D space. We first generalize SSM to the 2-dimensional space which generates the next state considering two adjacent states on both dimensions (e.g., columns and rows). We then construct our V2M based on the 2-dimensional SSM formulation and incorporate Mamba to achieve hardware-efficient parallel processing. The proposed V2M effectively incorporates the 2D locality prior yet inherits the efficiency and input-dependent scalability of Mamba. Extensive experimental results on ImageNet classification and downstream visual tasks including object detection and instance segmentation on COCO and semantic segmentation on ADE20K demonstrate the effectiveness of our V2M compared with other visual backbones.