CVPR2025 vGamba：融合选择性状态空间与自注意力的高效视觉骨干

最新推荐文章于 2025-09-03 11:57:40 发布

原创

最新推荐文章于 2025-09-03 11:57:40 发布 · 1.1k 阅读

19 ·

CC 4.0 BY-SA版权

文章标签：

#人工智能

CVPR2025 vGamba：融合选择性状态空间与自注意力的高效视觉骨干

本文发表在CVPR2025上，开源地址：https://github.com/yunusa2k2/vGamba

论文概述

vGamba 是一种混合视觉骨干模型，旨在高效捕捉视觉识别任务中的长距离依赖关系。该模型结合了状态空间模型（SSMs）和注意力机制，以增强效率和表达能力。

核心创新

1.Gamba Cell: 基于Mamba扩展而来，适应二维空间结构，并结合位置嵌入和空间上下文感知，以增强特征表示。通过结合 2D 位置编码、空间上下文感知掩码和 Mamba 块，有效地捕捉视觉任务中的长距离依赖关系，同时保持计算效率。其核心思想是：

去除因果约束: 原始 Mamba 模型采用因果结构，适用于自回归任务，但在视觉任务中效率低下。Gamba Cell 去除了这种约束，允许像素在整个图像范围内灵活交互，从而更好地捕捉长距离依赖关系。(这点个人感觉有点牵强，代码中还是使用原始mamba方式)
2D 位置编码: 为了有效地编码空间先验知识，Gamba Cell 引入可学习的相对位置嵌入，分别沿高度和宽度维度进行编码，确保模型能够区分不同空间位置的像素。
空间上下文感知掩码: Gamba Cell 使用 1D 卷积和 sigmoid 激活函数生成空间上下文感知掩码，有效地调制特征交互，同时保持计算效率。（其实就是简单注意力分数计算，估计为了和下一个创新点用词不重复换了一种说法）
Mamba 块: Gamba Cell 将输入特征序列通过 Mamba 块进行处理，该块能够高效地捕捉长距离依赖关系。
门控机制: Gamba Cell 使用门控机制将空间上下文感知掩码与 Mamba 块的输出进行元素级交互，从而生成最终的特征表示。（其实就是注意力）

2.2D-MHSA (2D Multi-Head Self-Attention)

通过全局 2D 注意力和 2D 相对位置编码，有效地捕捉视觉任务中的全局关系和空间依赖关系，从而提升模型的整体性能。这里就是把QKV方式的注意力引入，没有什么特别的，图片已经很清楚。

3.门控融合模块

无缝集成Gamba Cell和MHSA模块的输出，促进两者之间的有效交互，确保最佳的特征融合。步骤如下：

特征图分割: 将输入特征图沿通道维度分割成两部分。
分别处理: 将分割后的特征图分别输入 Gamba Cell 和 2D-MHSA 进行处理。
门控融合: 使用可学习的卷积投影将 Gamba Cell 和 2D-MHSA 的输出进行变换，并应用 sigmoid 激活函数生成门控权重。然后，根据门控权重将两个输出进行加权融合，得到最终的特征表示。

这步很像很多YOLO系列中常用的沿着通道劈开后分别提取特征，只是最后用Sigmoid方式动态产生加权权重融合而不是直接拼接特征。

$Gatem=σ(Wm∗xmamba)\mathrm{Gate}_{m}=\sigma(W_{m}*x_{\mathrm{mamba}})$

$Gateh=σ(Wh∗xmhsa)\mathrm{Gate}_{h}=\sigma(W_{h}*x_{\mathrm{mhsa}})$

$xfused=Gatem⋅Wm∗xmamba+Gateh⋅Wh∗xmhsax_{\mathrm{fused}}=\mathrm{Gate}_{m}\cdot W_{m}*x_{\mathrm{mamba}}+\mathrm{Gate}_{h}\cdot W_{h}*x_{\mathrm{mhsa}}$

实验结果

在三大CV任务上进行了实验：

分类: 在ImageNet-1K上，vGamba-L在保持较低计算成本（6.32G FLOPs，31.89M参数）的同时，实现了82.8%的Top-1准确率，优于ViTs和其他SSMs。
语义分割: 在ADE20K数据集上，vGamba-L达到了51.4（SS）和52.2（MS）的mIoU，表现出色。****
目标检测: 在COCO 2017数据集上，vGamba-L在Mask R-CNN和Cascade Mask R-CNN框架下展示了高精度和效率。

局限性与未来工作

尽管vGamba表现出色，但仍存在一些局限性，如对超参数选择的敏感性以及在实时应用中的计算开销（论文中使用的是8 NVIDIA Titan XP 12GB GPUs）。未来的研究将致力于优化计算效率和扩展模型以处理更复杂的场景。

源码解读

源码比较简单，主要就是一个model.py文件，我分块加上了详细注释。

1.基础导入与环境设置部分

import torch
import torch.nn as nn
# 用于模型FLOPs计算的库
from thop import profile
# 引入 Mamba 状态空间模型模块
from mamba_ssm import Mamba
# 模型结构信息摘要工具
from torchinfo import summary
import torch.nn.functional as F
# 用于变形张量结构的模块，如 flatten、rearrange 等
from einops import pack, repeat, unpack

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

2.GambaCell 模块

# GambaCell是一个适用于图像的状态空间模块，整合位置编码与门控机制
class GambaCell(nn.Module):
    def __init__(self, n_dims, height, width):
        super().__init__()
        # 可学习的相对位置编码，用于高和宽方向
        self.rel_h = nn.Parameter(torch.randn([1, n_dims, 1, height]), requires_grad=True)
        self.rel_w = nn.Parameter(torch.randn([1, n_dims, width