【论文阅读|V2M: VISUAL 2-DIMENSIONAL MAMBA FOR IMAGE REPRESENTATION LEARNING】
V2M: Visual two-dimensional MAMBA for image representation learning
年份：2024
期刊会议： arXiv
代码链接：https://github.com/wangck20/V2M
目录
当前面临的主要挑战
-
关键信息丢失
-
该架构在某些方面存在不足
-
优化视觉性能
- 相关研究
- 方法
-
-
-
- 二维状态空间模型设计：
-
-
该方法的整体流程设计包含以下七个关键环节：
- 首先进行输入数据预处理阶段；
- 空间四方向旋转操作随后展开；
- 二维状态空间模型（2D SSM）构建过程随后进行；
- 水平隐状态分解析及并行计算策略随后实施；
- 输入特征变换及融合机制随后完成；
- 二维状态空间模型输出整合及旋转复原步骤随后执行；
- 最终系统输出阶段随后实现。
在实验验证环节中： - 分类性能评估框架设计作为核心内容；
- 消融性分析模块构建作为辅助验证手段。
现阶段存在的问题
Visual 2-Dimensional Mamba（V2M ）的核心目标在于为视觉任务构建一种既能保持二维结构特征又具备强大状态空间建模能力的新一代视觉骨干网络。该方法主要针对的问题包括：降低计算复杂度、提升模型性能以及优化资源利用率等关键挑战。
1. 二维结构信息丢失
- 问题 ：传统的 Mamba 方法属于基于一维序列模型的设计，在视觉任务中，图像通常被划分为若干块（patch），然后将其展平为一维序列输入的过程会导致原始图像中重要的二维局部结构信息（例如邻域之间的相关性）被破坏。
- 解决方案 ：V2M 方法通过扩展 State Space Model（SSM）的空间维度，将其从一维扩展到二维，在这种情况下直接在二维网格结构上进行状态转移计算。该方法分别从行方向和列方向考虑相邻的状态转移关系，并最终实现了对空间信息的有效捕捉与利用。
2. 一维 Mamba 架构的局限性
- 问题 ：即便采用多种扫描策略, 1D Mamba 仍难以重建原始二维空间关系。
- 解决方案：基于二维状态方程, V2M能够以更加自然的方式在二维栅格上传播状态, 自四周起始点出发, 并防止长路径扫描对空间连贯性带来的破坏。
3. 提升视觉任务表现
- 问题 ：然而，在现有研究中发现1D Mamba模型在处理图像时未能充分挖掘其二维空间特征。
- 解决方案 ：本研究提出了一种基于二维SSM建模策略的新方法，在ImageNet、COCO等基准数据集上均展现了超越传统方法的性能。
相关研究
- CNN系列：包括ResNet、RegNet等
- Vision Transformers：涵盖ViT、Swin Transformer
- Visual-Mamba系列：Vim（Zhu等人, 2024年）；LocalMamba（Huang等人, 2024年）
- 状态空间模型（SSM）：S4 (Gu等人, 2021a)、Mamba (Gu & Dao, 2023)
- 二维SSM基础：Roesser模型（Kung等人,1977）
- 视觉表示学习：
- 监督学习：包括ResNet、MoCo等模型
- 自监督学习：涵盖MAE、SimCLR、BYOL等方法
方法
二维状态空间模型设计：
二维状态更新方程如下：
\begin{cases} h_1^{i,j+1} = A_1 h_1^{i,j} + A_3 h_2^{i,j} + B_1 x_{i,j} \\ h_2^{i+1,j} = A_2 h_1^{i,j} + A_4 h_2^{i,j} + B_2 x_{i,j} \\ y_{i,j} = C_1 h_1^{i,j} + C_2 h_2^{i,j} \end{cases}
简化后的状态更新方式为两条一维链：
h_1^{i,j+1} = A_1 h_1^{i,j} + A_2 h_2^{i,j} + B_1 x_{i,j}
h_2^{i+1,j} = A_2 h_2^{i,j} + I x'_{i,j}
V2M流程
1. 输入与预处理
- 输入是形状为 B \times H \times W \times C 的图像张量，比如 ( B=1, H=4, W=4, C=2 )：
x[0,:,:,:] = \begin{bmatrix} [1,1] & [2,2] & [3,3] & [4,4] \\ [5,5] & [6,6] & [7,7] & [8,8] \\ [9,9] & [10,10] & [11,11] & [12,12] \\ [13,13] & [14,14] & [15,15] & [16,16] \\ \end{bmatrix}
经由Patch Embedding技术将输入转换为一个具有维度结构B \times H \times W \times D的特征张量。
2. 四向旋转
- 对输入分别做四向旋转：
x_0 = 原始输入0^\circ旋转
x_1 = 逆时针旋转90^\circ
x_2 = 逆时针旋转180^\circ
x_3 = 逆时针旋转 270^\circ
3. 2D 状态空间模型 (2D SSM) 计算
- 对每个旋转后的输入x_i，使用二维状态空间模型计算隐藏状态和输出：
\begin{cases} h_1^{i,j+1} = A_1 h_1^{i,j} + A_2 h_2^{i,j} + B_1 x^{i,j} \\ h_2^{i+1,j} = A_3 h_1^{i,j} + A_4 h_2^{i,j} + B_2 x^{i,j} \end{cases}
h_1^{i,j} 和 h_2^{i,j} 分别表示水平方向和垂直方向的隐状态。
4. 水平方向隐状态拆解与并行计算
- 拆分第一条状态转移方程为两部分：
\begin{aligned} h_1^{i,j+1} &= A_1 h_1^{i,j} + B_1 x^{i,j} \\ h_1^{i,j+1} &= A_2 h_2^{i,j} \end{aligned}
- 第一部分和 1D SSM 格式一致，允许使用 Mamba 的并行扫描方法进行计算。
5. 输入变换与拼接
将输入 x \in \mathbb{R}^{B \times H \times W \times D} 变形为
水平方向处理时：x' \in \mathbb{R}^{(B \times H) \times W \times D}
垂直方向处理时：x' \in \mathbb{R}^{(B \times W) \times H \times D}
使用各自对应的参数集 (A_1, B_1) 和(A_2, I)计算。
6. 2D SSM输出拼接与旋转复原
四个旋转方向的输出分别计算后，反旋转回原始空间方向。
将四个方向的输出特征在 batch 维度拼接，形成最终特征：
z_{k+1} = \sum_{i=0}^3 \text{rot}_{-90^\circ \times i}(z_{k}^i)
- 使用 MLP 层进一步映射特征。
7. 最终输出
- 最后一个 V2M Block 的输出用于分类或下游任务。
实验
分类任务
消融实验
理论分析
【探讨二位SSM模型的稳定性特征及其泛化能力边界