Happy · 1月19日

华科王兴刚团队首次将Mamba引入ViT,更高精度、更快速度、更低显存!

在Transformer如日中天时,一个称之为“Mamba”的架构横冲出世,在语言建模上与Transformers不相上下,具有线性复杂度,同时具有5倍的推理吞吐量!一时之间,被给予厚望“下一代架构”~

就在今日,华中科技大学王兴刚团队首次将“Mamda”里面引入到CV领域而得到Vim公开了,比DeiT精度更高、速度更快,还节省GPU显存,神了,估计后续会出现各种变种,DeiM、PvM,哈哈~

image.png

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处理高分辨率图像时的计算和内存限制的潜力,并有可能成为下一代视觉基础模型。

image.png

本文方案

本文所提Vim(Vision Mamba)旨在将陷阱的SSM模型(如Mamba)引入到CV领域,故先对SSM进行简要介绍,然后再介绍如何结合CV任务特性进行Vim模块设计,最后呈现所提Vim架构细节。

Preliminaries

image.png

Vision Mamba

image.png

Vim Block

原始Mamba模块针对1D序列而设计,不适合需要空域位置感知的视觉任务。针对此,我们引入了Vim模块,它为视觉任务引入了双向序列建模机制。Vim处理流程还是看代码吧~

image.png

架构细节

总体而言,所提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分类

image.png

上表给出了所提Vim与其他方案的性能对比,可以看到:

  • 相比ResNet,Vim具有优异性能。如Vim-Small取得了80.3%,高出4.1%;
  • 相比ViT,Vim在参数量与分类精度方面以显著优势胜出;
  • 相比DeiT,Vim以相当参数量取得更高精度,如Vim-Tiny比DeiT-Tiny高出0.9%;
  • 相比SSM类方案,Vim取得了相当精度,但参数量少3倍。

ADE20K语义分割

image.png

上表为ADE20K语义分割任务上不同方案性能对比,可以看到:

  • Vim在不同尺度均优于DeiT,Vim-Ti比DeiT-Ti高出1.0mIoU、Vim-S比DeiT-S高出0.9mIoU。
  • 相比ResNet101,Vim-S以更少的参数量取得了相当的分割精度。
  • 下图从FPS与GPU占用维度进一步验证了所提方案的高效性。

image.png

COCO检测分割

image.png

上表给出了COCO实例分割任务上的性能对比,可以看到:

  • Vim-Ti比DeiT-Ti高出1.3boxAp、1.1maskAP。
  • 对于中等大小目标与Vim-Ti指标高出1.6boxAP、1.3maskAP;对与大尺度目标,Vim-Ti指标高出1.4boxAP、1.8maskAP。这进一步验证了Vim具有比DeiT更优的长程上下文建模能力,可参考下图。

image.png

THE END

作者:AIWalker
来源:AIWalker

推荐阅读

本文章著作权归作者所有,任何形式的转载都请注明出处。更多动态滤波,图像质量,超分辨相关请关注我的专栏AIWalker。欢迎添加极术小姐姐微信(id:aijishu20)加入技术交流群,请备注研究方向。
推荐阅读
关注数
6197
内容数
191
夯实深度学习知识基础, 涵盖动态滤波,超分辨,轻量级框架等
目录
极术微信服务号
关注极术微信号
实时接收点赞提醒和评论通知
安谋科技学堂公众号
关注安谋科技学堂
实时获取安谋科技及 Arm 教学资源
安谋科技招聘公众号
关注安谋科技招聘
实时获取安谋科技中国职位信息