在Transformer如日中天时,一个称之为“Mamba”的架构横冲出世,在语言建模上与Transformers不相上下,具有线性复杂度,同时具有5倍的推理吞吐量!一时之间,被给予厚望“下一代架构”~
就在今日,华中科技大学王兴刚团队首次将“Mamda”里面引入到CV领域而得到Vim公开了,比DeiT精度更高、速度更快,还节省GPU显存,神了,估计后续会出现各种变种,DeiM、PvM,哈哈~
https://arxiv.org/abs/2401.09417
https://github.com/hustvl/Vim
最近,具有高效硬件感知设计的状态空间模型(State Space Models, SSM)曼巴,已显示出在长序列建模方面的巨大潜力。虽然基于SSM构建高效和通用的视觉骨干是一个有吸引力的方向。但是,由于视觉数据的位置敏感性、视觉理解所需要的全局上下文依赖性,表示视觉数据对于SSM而言是一项具有挑战性的任务。因此,我们提出了一个新的基于双向曼巴块(Vim)的通用视觉骨干,该模型通过位置嵌入标记图像序列并通过双向状态空间模型压缩视觉表示。 在ImageNet分类、COCO检测和ADE20K语义分割任务上,Vim实现了与成熟ViT(如DeiT)相比更高的性能,同时显著提高了计算和内存效率。例如,与DeiT相比,在对分辨率为1248×1248的图像进行批量推理时,Vim快了2.8倍同时可节省86.8%的GPU内存。这些结果表明,Vim具有克服ViT处理高分辨率图像时的计算和内存限制的潜力,并有可能成为下一代视觉基础模型。
本文方案
本文所提Vim(Vision Mamba)旨在将陷阱的SSM模型(如Mamba)引入到CV领域,故先对SSM进行简要介绍,然后再介绍如何结合CV任务特性进行Vim模块设计,最后呈现所提Vim架构细节。
Preliminaries
Vision Mamba
Vim Block
原始Mamba模块针对1D序列而设计,不适合需要空域位置感知的视觉任务。针对此,我们引入了Vim模块,它为视觉任务引入了双向序列建模机制。Vim处理流程还是看代码吧~
架构细节
总体而言,所提Vim架构有四个超参:模块数L,隐状态维度D、扩展状态维度E以及SSM维度N。
延续ViT与DeiTy,我们首先采用核投影层得到1D非重叠块嵌入序列,然后直接堆叠L个Vim模块。默认信息,设置L=24,N=16。为对齐DeiT系列,Tiny版本的D=192,E=384;Small版本的D=384,E=768.
本文实验
ImageNet分类
上表给出了所提Vim与其他方案的性能对比,可以看到:
- 相比ResNet,Vim具有优异性能。如Vim-Small取得了80.3%,高出4.1%;
- 相比ViT,Vim在参数量与分类精度方面以显著优势胜出;
- 相比DeiT,Vim以相当参数量取得更高精度,如Vim-Tiny比DeiT-Tiny高出0.9%;
- 相比SSM类方案,Vim取得了相当精度,但参数量少3倍。
ADE20K语义分割
上表为ADE20K语义分割任务上不同方案性能对比,可以看到:
- Vim在不同尺度均优于DeiT,Vim-Ti比DeiT-Ti高出1.0mIoU、Vim-S比DeiT-S高出0.9mIoU。
- 相比ResNet101,Vim-S以更少的参数量取得了相当的分割精度。
- 下图从FPS与GPU占用维度进一步验证了所提方案的高效性。
COCO检测分割
上表给出了COCO实例分割任务上的性能对比,可以看到:
- Vim-Ti比DeiT-Ti高出1.3boxAp、1.1maskAP。
- 对于中等大小目标与Vim-Ti指标高出1.6boxAP、1.3maskAP;对与大尺度目标,Vim-Ti指标高出1.4boxAP、1.8maskAP。这进一步验证了Vim具有比DeiT更优的长程上下文建模能力,可参考下图。
THE END
作者:AIWalker
来源:AIWalker
推荐阅读
- 语义分割新SOTA:华科与美团联合提出单分支推理分割架构SCTNet,即将开源!
- BFRffusion | 面向真实世界的盲人脸复原,中山大学提出全新盲人脸复原基准数据集
- 通用检测大模型 | 华科白翔团队提出以对象为中心的基础模型GLEE
- 端侧显著性检测新高度,OPPO提出面向真实场景的PSUNet
- 38.7fps!EdgeSAM = RepViT + SAM,移动端超强变种,已开源!
本文章著作权归作者所有,任何形式的转载都请注明出处。更多动态滤波,图像质量,超分辨相关请关注我的专栏AIWalker。欢迎添加极术小姐姐微信(id:aijishu20)加入技术交流群,请备注研究方向。