点击下方卡片,关注「AI视界引擎」公众号
( 添加时备注:方向+学校/公司+昵称/姓名 )
状态空间模型(SSMs)通过利用线性递归降低了 Transformer 的二次复杂度。最近,VMamba已成为一种强大的基于SSM的视觉 Backbone 网络,但其四方向扫描仍受限于空间冗余。作者提出了QuarterMap,这是一种在扫描前移除冗余空间激活并在后续通过最近邻上采样恢复维度的后训练激活剪枝方法。yinqing-QuarterMap_2507无需重新训练即可提高吞吐量。在ImageNet1K上,QuarterMap在VMamba上实现了高达11%的加速,准确率下降不到0.9%,并在ADE20K分割任务上取得了相似的增益。除了VMamba,作者在MedMamba上验证了QuarterMap,这是一个特定领域的模型,具有相同的四方向扫描结构,它在多个医学图像任务中始终如一地提高了吞吐量,同时保持了准确率。
与ToMe等token合并方法相比,QuarterMap专为SSM设计,避免了昂贵的合并-解合并操作。yinqing-QuarterMap_2507为部署时效率提供了一种即插即用的工具,同时不牺牲迁移性。
然而,卷积神经网络在捕捉长距离依赖关系方面存在局限性。视觉Transformer(ViT),凭借其自注意力机制,有效地克服了这些局限性,但由于其二次复杂度,导致了高昂的计算成本。为了缓解这些计算需求,近期研究集中于降低ViT的复杂度、应用模型压缩技术,并研究了替代架构,如RWKV 和状态空间模型 。
unsetunset1. 引言unsetunset
计算机视觉领域的进展在很大程度上得益于深度学习和大规模数据集的可用性。卷积神经网络(CNN)已成为图像分类和目标检测 等任务的基础。状态空间模型(SSM)最初在自然语言处理(NLP)领域被引入,以减少解码阶段维护隐藏状态的计算成本。相比之下,计算机视觉任务通常将隐藏状态解释为整个图像信息的表示。
最近,SSM已成为计算机视觉中ViT的有效替代方案,在多个任务中展现出具有竞争力的性能。例如,VMamba 在ImageNet-1K基准(Deng等人,2009年)上实现了82.6%的Top-1准确率,其FLOPs与Swin Transformer(Liu等人,2021a)相当,但性能高出1.3%。
然而,在VMamba中,负责选择性扫描的核心,该操作类似于Transformer模型中的注意力机制,两者都作为捕获全局上下文的机制,仍占总核执行时间的18.3%,凸显了一个显著的效率 Bottleneck 。为解决这一问题,作者探索了Transformer优化技术,但发现缺乏专门为SSM设计的优化方法。作者证明,传统的优化方法,如广泛用于Transformer的token合并(Bolya等人,2023年),由于频繁的合并和拆分操作带来的计算权衡,对VMamba来说是不理想的,如图1所示。其他近期方法,如Top-ViM(Zhan等人,2024年)和R-MeeTo(Shi等人,2025年),引入了针对基于ViM(Zhu等人,2024年)模型的token剪枝和合并策略。然而,两者都依赖于重新训练来维持性能,因此在训练后部署场景中有限的计算资源下适用性较低。
受这些挑战的启发,作者探索了是否可以将类似于token合并的技术应用于SSM的激活剪枝,以及更广泛地,如何在已经高效的线性SSM中进一步提高效率而不需要重新训练。作者首先分析了VMamba的跨扫描和跨合并机制,以及VMamba中的有效感受野(ERF)(Liu等人,2024)。作者的分析揭示,VMamba中的四向遍历引入了大量的空间冗余,其中一些token积累了过多且可能不必要的冗余信息。这使作者假设,与扫描结构相一致的专门激活剪枝可以减少延迟同时保持准确性。
为此,作者提出了QuarterMap,一种无需训练的激活剪枝方法,专门设计用于通过在扫描前将特征图的空间尺寸缩小至原始尺寸的四分之一来提高VMamba的效率。如图2所示,QuarterMap引入了一个轻量级的两阶段流程:在交叉扫描模块之前应用空间剪枝,并在交叉合并后使用最近邻上采样来恢复分辨率。在剪枝过程中,QuarterMap在两个空间维度上保留每隔一个元素,利用VMamba四向扫描中固有的空间冗余性。作者基于相邻空间位置传递相似信息的假设采用最近邻插值,将其作为一种高效且有效的重建策略。该设计显著降低了交叉扫描、选择性扫描和交叉合并过程中的计算量,所有这些都不需要修改模型权重或进行重新训练。
作者在图像分类和语义分割任务上评估了QuarterMap,展示了在精度损失最小的情况下通过提升吞吐量。在ImageNet-1K上,它实现了高达1.11倍的速度提升,同时top-1精度仅下降不到0.9%。
在语义分割和医学图像基准测试中也观察到类似的趋势。QuarterMap对VMamba及其变体特别有效,这一点通过与CNN、ViT以及其他SSM(如PlainMamba(Yang等人,2024)和ViM(Zhu等人,2024))的比较得到证实。注意力图和有效感受野(ERF)的可视化显示,QuarterMap去除了冗余激活,同时保留了关键的空间信号。
全面的消融研究进一步探索了剪枝间隔、层选择策略和上采样方法,为在现实世界中进行无训练部署提供了实用指导。
unsetunset2. 方法unsetunset
作者介绍了QuarterMap,这是一种专门设计的后训练剪枝函数,旨在通过减少激活特征图中的空间冗余来提升VMamba的效率,如图2所示。形式上,作者将QuarterMap定义为一个作用于输入激活图
的函数
,通过选择性地保留空间信息进行操作。给定一个激活图
,其中
和
表示空间维度,
表示通道维度,QuarterMap通过以下阶段进行操作:
QuarterMap在VMamba架构中对特定模块进行选择性剪枝,该模块由一个块选择间隔
确定。这种High-Level策略控制剪枝的频率,在计算效率和准确性之间取得平衡。将QuarterMap应用于每三个模块(即
),并排除第一层,可获得最佳的准确率-延迟权衡(yinqing-QuarterMap_2507E)。早期层对于编码基本特征至关重要,而深层对剪枝更具鲁棒性,因此适合作为优化候选。
剪枝阶段在每个选定的块内,QuarterMap在交叉扫描前对
的空间维度执行下采样操作。对于指定的间隔
,函数
在
和
维度上保留每
个元素,如图3所示,从而生成剪枝后的激活图
。该过程利用了VMamba的交叉扫描机制,该机制从四个方向聚合信息,并结合SSM递归函数。这些机制共同确保特征图中的每个元素
包含了来自邻近空间位置的信息,从而在提高剪枝有效性的同时最小化精度损失。作者发现设置
和
在计算效率和精度之间实现了最佳权衡。由于剪枝阶段在交叉扫描之前应用,因此计算节省主要来源于输入
在交叉扫描和Mamba操作中的尺寸减小。值得注意的是,在Mamba中,这种减少还导致SSM(公式(3))的输入长度节省,以及通过选择性机制(公式(5))的线性计算节省。
上采样阶段
在经过交叉扫描、选择性扫描和交叉合并处理后,QuarterMap使用上采样函数
对
(交叉合并输出)进行处理,以恢复激活图的空间维度。采用最近邻插值方法重建原始空间维度,生成输出
RHWD。这种方法基于相邻空间元素包含相似信息的假设,使QuarterMap能够在保证计算效率的同时,实现最小的精度损失。
unsetunset3. 实验unsetunset
3.1. 分类与分割
表1中的结果表明,将QuarterMap应用于基础VMamba模型时,当
时,准确率略有下降
,同时吞吐量提升了
。这种提升主要源于选择性扫描机制中序列长度的减少,而作者的剪枝和上采样阶段引入的额外开销极小。性能分析显示,QuarterMap将扫描核时间从1.4毫秒减少到0.6毫秒,并仅引入了0.2毫秒的额外开销。相比之下,ToMe虽然扫描核运行时间相似(0.7毫秒),但由于合并和拆分操作的成本,其引入的额外开销显著更高(9.7毫秒)。这些比较突显了QuarterMap在最小化不必要的计算同时保持准确率方面的效率。关于小型和微型VMamba配置的额外结果以及详细的分割指标,请参见附录D。
3.2. QuarterMap在医学影像中的MedMamba应用
为评估QuarterMap在VMamba之外的适用性,作者将其应用于MedMamba-T(Yue & Li,2024),这是一个基于VMamba架构构建的特定领域模型。作者在四个MedMNIST分类数据集上进行性能基准测试:BoodMNIST、OrganMNIST、RetinaMNIST和PathMNIST(Yang等人,2021;2023)。如表2所示,QuarterMap始终提供1.21倍的吞吐量提升(从854到1034张图像/秒),且分类精度未下降。这些结果表明QuarterMap对基于VMamba的模型的泛化能力,并强化了其在特定领域部署中的实用性,尤其是在训练后效率和精度保持至关重要的情况下。关于BloodMNIST上的类别性能分解,请参见附录D。
3.3. QuarterMap在其他架构上的应用
作者在CNN、ViT和SSM上评估QuarterMap,以评估其设计是否特别适用于VMamba。作者关注ConvNeXtv2(Woo等人,2023)、DeiT(Touvron等人,2021)、Swin Transformer(Liu等人,2021a)和ViM(Zhu等人,2024)的基础变体,所有模型均在ImageNet-1K上预训练,权重来自Hugging Face(Wolf等人,2020)。QuarterMap应用于每隔一个(即前两个之后的3个I块)。在CNN中,剪枝会破坏空间连续性,显著影响精度。ViT和ViM更具韧性,但仍表现出非平凡的精度下降。如表3所示,yinqing-QuarterMap_2507通过其四向扫描机制降低了VMamba的延迟,同时保持了精度。QuarterMap与CNN和类似ViM的一维扫描SSM的兼容性较差,因为这些模型缺乏VMamba中呈现的冗余激活模式。其他变体的结果见附录D。
3.4. 注意力图和有效感受野 (ERF)
作者可视化了应用QuarterMap前后注意力图和ERF,以分析其空间影响。遵循VMamba的公式,作者在多个前置块被剪枝后,提取了第
块(即最深层)的注意力图。如图4所示,注意力模式基本保持不变,表明QuarterMap保留了关键的上下文行为。对于ERF的可视化,灰度区域表示QuarterMap移除的激活区域。这些区域通常与空间冗余的感受野热点重叠。这支持了作者的假设,即QuarterMap能够有效消除冗余信息,同时保留模型的功能结构。
unsetunset4. 结论unsetunset
作者提出了QuarterMap,一种针对VMamba的模型后训练剪枝方法,该方法在最小化精度损失的同时提升了运行时效率,且无需重新训练。作者的实验表明,QuarterMap与VMamba的四向扫描特性高度契合,同时也适用于基于其 Backbone 网络衍生出的应用。
尽管QuarterMap是为VMamba设计的,但它与量化等技术正交,从而能够实现进一步的效率提升。虽然本研究聚焦于VMamba,但作者的发现为理解和扩展Sparse状态空间模型(SSMs)中的剪枝策略开辟了新的方向。
点击上方卡片,关注「AI视界引擎」公众号