高效视觉模型——EfficientViM 的研究进展
为了在资源受限的环境中部署神经网络,早期的研究通过结合卷积(用于捕捉局部依赖)和注意力机制(用于捕捉全局依赖),构建了轻量级架构。然而,这种注意力机制的计算复杂度呈现二次方增长,成为了效率提升的主要障碍。
最近,状态空间模型 (State Space Model, SSM) 因其在线性计算复杂度下处理长序列的能力而受到广泛关注。尽管如此,在实际应用中,现有的视觉 Mamba 模型的速度仍然不及顶级的轻量级模型,主要瓶颈在于对输入序列进行线性投射的操作。
EfficientViM 的提出与核心贡献
针对这一挑战,本研究提出了一个全新的高效视觉 Mamba 架构——EfficientViM。该架构的核心创新点在于引入了一种名为“基于隐藏状态混合器的状态空间对偶 (Hidden State Mixer-based State Space Duality, HSM-SSD)”的新层设计。
HSM-SSD 的主要改进是将高计算量的通道混合操作从原始图像特征空间转移到了压缩后的隐藏状态空间中进行,从而有效地降低了计算成本。此外,研究还引入了“多阶段隐藏状态融合 (Multi-stage Hidden-State Fusion, MSF)”机制来增强隐藏状态的表征能力,并通过优化宏观设计减少内存密集型操作,在 ImageNet-1k 等任务上实现了当前最先进的速度-精度权衡。
论文基本信息
标题: EfficientViM: Efficient Vision Mamba with Hidden State Mixer based State Space Duality
- 核心模块:
- 基于隐藏状态混合器的状态空间对偶 (HSM-SSD)
- 多阶段隐藏状态融合 (MSF)
- 深度可分离卷积 (DWConv)
算法框架与核心模块
EfficientViM 是一个分层的视觉骨干网络。其算法框架如下:输入图像首先通过由卷积层组成的 Stem 模块,进行初步特征提取和 4 倍下采样。随后,特征图依次经过三个阶段(Stage)的 EfficientViM 模块堆叠处理。每个阶段末尾通过一个下采样层来缩小特征图尺寸并增加通道数,构建层级化的特征表示。最后,通过分类头输出结果,并在训练时利用 MSF 机制融合各阶段的隐藏状态信息以辅助最终预测。
核心模块介绍
1. 基于隐藏状态混合器的状态空间对偶 (HSM-SSD)
核心功能: 旨在解决标准非因果状态空间对偶 (Non-Causal State Space Duality, NC-SSD) 层中线性投射带来的性能瓶颈。
实现逻辑: HSM-SSD 的设计是将计算密集型的通道混合操作转移到维度更低的隐藏状态空间中。具体而言,输入序列首先被压缩为中间隐藏状态,然后在“隐藏状态混合器 (Hidden State Mixer, HSM)”模块内进行门控和线性投射,得到更新后的隐藏状态。最后,使用该隐藏状态和选择性矩阵重建最终输出。
其核心近似逻辑如下:
\[ \mathbf{x}_{\text{out}} = \text{Linear}(\mathbf{y} \odot \sigma(\mathbf{z})) \approx \mathbf{C}((\mathbf{h} \odot \sigma(\mathbf{h}_{\text{in}}\mathbf{W}_z))\mathbf{W}_{\text{out}}) = \mathbf{C}f(\mathbf{h}) \]其中,\(\mathbf{x}_{\text{out}}\) 是传统 NC-SSD 的输出,而 \(f(\mathbf{h})\) 表示在隐藏状态空间中进行的门控和线性投射操作。
优势: 该设计将主要计算复杂度从 \(\mathcal{O}(LD)\)(L 为序列长度,D 为通道数)降低到 \(\mathcal{O}(ND)\)(N 为状态数)。由于 N 远小于 L,这一改变使得层的主要开销变得可控,显著提高了模型的计算效率和可扩展性,特别是在处理高分辨率图像时效果更加明显。
2. 多阶段隐藏状态融合 (MSF)
核心功能: 在训练过程中,通过显式利用网络不同阶段的隐藏状态来增强模型的最终预测能力和泛化性。
实现逻辑: MSF 机制从每个阶段的最后一个模块中提取隐藏状态,并在训练时将这些信息融合到最终预测中。这一过程不仅提高了模型的表征能力,还增强了其对复杂任务的适应性。
O(LD?)
x
y
x_in
h_in
h
C
x_out
f
O(ND?)
N
L
h^(s)
点击阅读原文,获取更多模块信息
通过计算隐藏状态的平均值,得到一个全局表征。
?^(s)
之后,将这个全局表征通过线性层映射成分类所需的logits。
z^(s)
最终模型输出的logits是由原始主干网络产生的logits与所有阶段辅助logits加权求和得到的。
z
这里,原始主干网络提供的logits作为输入之一参与到最终输出中。
z^(0)
具体计算公式如下:
\(\mathbf{z} = \sum_{s=0}^{S} \hat{\beta}^{(s)}\mathbf{z}^{(s)}, \quad \hat{\beta}^{(s)} = \frac{\exp(\beta^{(s)})}{\sum_{i=0}^{S}\exp(\beta^{(i)})}\)
其中,\(\hat{\beta}^{(s)}\)是一个可训练的标量权重。
β^(s)
优势
这种类似于深度监督的技术促使网络在较浅层次学习有区分性的特征。它融合了从底层细节到高层语义信息,实现了类似模型集成的效果,从而显著提高了模型的整体性能。
核心应用场景
本文提出的方法主要用于图像分类任务,并展示了其在目标检测、实例分割及语义分割等密集预测任务上的有效性和可扩展性。特别适用于需要高吞吐量和低延迟的资源受限场景。
方法论核心
该方法的核心思想是计算重分配:通过识别模型中的计算瓶颈(即对高维数据的线性变换),将其重新分布到一个维度更低的隐藏空间中执行,从而在保持全局视角的同时提高计算效率。
启发性拓展
- 递归隐藏状态混合:HSM-SSD模块可作为其内部结构的一个实例,形成一种深度展开的递归架构,可能进一步降低计算复杂度。
- 高分辨率图像应用:由于HSM-SSD具备线性复杂度的优势,在处理超高分辨率图像时表现出巨大潜力,例如在医学影像和卫星图像分析领域。实验已初步验证了这一点。
点击阅读原文,获取更多模块信息
核心实验与结论
本文最能体现其贡献的核心实验是在ImageNet-1K数据集上对现有高效视觉骨干模型进行的速度-精度对比。目的是验证EfficientViM系列模型是否能在实际硬件(如NVIDIA RTX 3090 GPU)上,相比其他顶级轻量级CNN和Transformer模型,实现更优的速度-精度平衡。
关键结果显示,在相似的精确度水平下,EfficientViM模型家族展现出了明显更高的处理速度。例如,EfficientViM-M2在达到75.4% Top-1准确率的同时,吞吐量达到了17,005 img/s,远超相同精度级别的其他模型如SHViT-S2(15,899 img/s)和MobileNetV3-L 1.0(9,493 img/s)。
与之前的顶级模型SHViT相比,EfficientViM在所有规模的模型上都实现了更高的吞吐量和相当甚至更高的精度。
如图所示,在速度-精度散点图中,EfficientViM系列(红色和蓝色五角星)形成了一个新的前沿线,明确展示了其作为当前轻量级模型中的领导地位。
作者结论
基于这些实验结果,研究者总结道:通过HSM-SSD模块及一系列宏观设计优化,EfficientViM成功构建了一个新的、高效的视觉架构。它在速度和精度之间取得了最前沿的平衡,并且特别适合实际部署。
即插即用模块
隐藏状态混合器状态空间对偶(HSM-SSD)
主要功能:将高维特征空间中的通道混合及门控运算转换到低维度的隐藏状态空间中执行,以降低计算和内存成本。
优势:复杂度从O(L·D)降低至O(N·D),其中N(隐藏状态数)远小于L(序列长度),对实际硬件吞吐量和显存占用更加友好。
# 低维隐藏状态空间的投影与混合(简化示意)
def forward(self, x):
B, _, L = x.shape
H = int(math.sqrt(L))
# 计算 B, C, dt 三个状态通道参数
BCdt = self.dw(self.BCdt_proj(x).view(B, -1, H, H)).flatten(2)
Bm, Cm, dt = torch.split(BCdt, [self.state_dim]*3, dim=1)
# 非因果SSD的状态权重 A(随位置软最大)
A = (dt + self.A.view(1, -1, 1)).softmax(-1)
# 隐藏状态混合:在低维空间做通道混合,替代原始高维 y 上的重运算
h = x @ (A * Bm).transpose(-2, -1)
# 门控与输出投影(隐藏状态维度内完成)
h, z = torch.split(self.hz_proj(h), [self.d_inner]*2, dim=1)
h = self.out_proj(h * self.act(z) + h * self.D)
# 回投到原始空间得到输出 y
y = h @ Cm
return y.view(B, -1, H, H), h
多阶段隐藏状态融合(MSF)
主要功能:从不同阶段提取隐藏状态,通过全局聚合与独立分类头产生各阶段预测,并使用可学习的软权重进行整合,最终得到模型logits。
优势:以极低的成本实现深度监督和多尺度语义集成,大幅提升了性能和稳定性。
# 多阶段隐藏状态融合(简化示意)
def forward(self, x):
x = self.patch_embed(x)
weights = self.weights.softmax(-1) # 学习的融合权重
z = torch.zeros((x.shape[0], self.num_classes), device=x.device)
for i, stage in enumerate(self.stages):
x, x_out, h = stage(x) # 取出阶段末尾的隐藏状态 h
h = self.norm[i](h)
h = torch.nn.functional.adaptive_avg_pool1d(h, 1).flatten(1)
z = z + weights[i] * self.heads[i](h) # 阶段预测加权累加
# 最后阶段的空间特征也参与融合(主干末尾监督)
x = self.norm[3](x)
x = torch.nn.functional.adaptive_avg_pool2d(x, 1).flatten(1)
z = z + weights[3] * self.heads[3](x)
return z
单头 HSM-SSD(Single-head HSM-SSD)
主要功能:在HSM-SSD基础上,采用单头配置(ssd_expand=1),保持了最小的参数量和计算规模,同时保留了隐藏状态空间中的高效混合特性。
优势:特别适用于移动设备和边缘场景,在速度、显存占用以及实现复杂度之间取得了最佳平衡,易于集成与迁移。
隐藏状态混合器# 在 Block 中以单头(ssd_expand=1)方式实例化 HSM-SSD(简化示意)
class EfficientViMBlock(nn.Module):
def __init__(self, dim, mlp_ratio=4., ssd_expand=1, state_dim=64):
# ...
self.mixer = HSMSSD(d_model=dim, ssd_expand=ssd_expand, state_dim=state_dim)
# LayerScale 等宏观设计提升稳定性与性能
self.alpha = nn.Parameter(1e-4 * torch.ones(4, dim), requires_grad=True)
# 以单头配置注册不同规模的模型族(均 ssd_expand=1)
def EfficientViM_M1(pretrained=False, **kwargs):
return EfficientViM(embed_dim=[128, 192, 320], depths=[2, 2, 2],
mlp_ratio=4., ssd_expand=1., state_dim=[49, 25, 9], **kwargs)
def EfficientViM_M2(pretrained=False, **kwargs):
return EfficientViM(embed_dim=[128, 256, 512], depths=[2, 2, 2],
mlp_ratio=4., ssd_expand=1., state_dim=[49, 25, 9], **kwargs)
def EfficientViM_M3(pretrained=False, **kwargs):
return EfficientViM(embed_dim=[224, 320, 512], depths=[2, 2, 2],
mlp_ratio=4., ssd_expand=1., state_dim=[49, 25, 9], **kwargs)
def EfficientViM_M4(pretrained=False, **kwargs):
return EfficientViM(embed_dim=[224, 320, 512], depths=[3, 4, 2],
mlp_ratio=4., ssd_expand=1., state_dim=[64, 32, 16], **kwargs)
说明与使用建议:
在您的项目中,如果需要快速集成高效的通道混合能力,可以将 HSM-SSD 模块直接插入到骨干网络的卷积或注意力层之间。确保其输入输出形状和接口保持一致即可。
为了增强表征力同时保持对硬件的友好性,建议首先使用单头配置(ssd_expand=1)进行验证。根据实际性能需求,您可以适当增加 state_dim 的值或者引入多头设计(提高 ssd_expand)以进一步优化性能。
更多模块信息,请参阅原文。


雷达卡


京公网安备 11010802022788号







